# -*- coding: iso-8859-1 -*-
# Created : S.Sénési - nov 2015

""" 
Base pour le traitement des données SE de LMDZ : 
  - variables, saisons et grilles habituellement gérées, 
  - organisation des données SE
  - renommage et re-scaling des variables
  - dictionnaire de fichiers d'exemples pour les grilles traitées
et une fonction pour rendre un objet CliMAF pour tuple (simu, variable, saison, grille)
"""

from climaf.api import *

# Liste des variables LMDZ objet d'une cmorisation ou d'une projection+moyenne saisonniere
#########################################################################################
"""
variables_list=['sfcWind','hfls','hfss','huss','hurs','pr',\
    'rldscs','rlds','rlus','rsdscs','rsds','rsuscs','rsus','rsutcs','rsut','rlut','rlutcs','rsdt',\
    'sfcWind','tas','ts','tauu','tauv','psl','zg500','hfns','ta','ua','va','wap','hus','hur',\
    'cllcalipso','clmcalipso','clhcalipso','cltcalipso','clt','prw','tasmax','tasmin']
"""
variables_list=['sfcWind','hfls','hfss','bil','huss','hurs','pr',\
    'rldscs','rlds','rlus','rsdscs','rsds','rsuscs','rsus','rsutcs','rsut','rlut','rlutcs','rsdt',\
    'sfcWind','tas','sst','tauu','tauv','psl','zg500','hfns','ta','ua','va','wap','hus','hur',\
    'cllcalipso','clmcalipso','clhcalipso','cltcalipso','clt','prw','tasmax','tasmin']

#Dictionnaire des grilles connues et de fichiers d'exemple
#########################################################################################
grid_remapfiles={ 'VLR' : '/data/ssenesi/stable/VLR.nc'}

# Saisons gérées et leur traduction CDO
#########################################################################################
seasons={ 'YEAR':"", 'DJF' : "timavg -selmon,1,2,12", 'JJA' : "timavg -selmon,6,7,8" , 'JJAS' : "timavg -selmon,6,7,8,9", 'MAM' : "timavg -selmon,3,4,5", 'SON' : "timavg -selmon,9,10,11" }


#########################################################################################
# Definitions pour l'acces aux obs gérées par Ionela
#########################################################################################
cproject('OBS', ('period','fx'),   ('root','ROOT_OBS_DIR'))  #, '/data/musat/LMDZ6/SE/CMOR/OBS'))
dataloc(project='OBS',url='${root}/${variable}.nc')
# Définiton de variables dérivées pour ce projet

#derive('OBS','rstt','minus','rsdt','rsut')
#derive('OBS','rstscs','minus','rsdscs','rsuscs')
#derive('OBS','rsts','minus','rsds','rsus')

#derive('OBS','cress','minus','rsds','rsdscs')
#derive('OBS','crels','minus','rlds','rldscs')
#derive('OBS','crets','plus','crels','cress')
#derive('OBS','crest','minus','rsutcs','rsut')
#derive('OBS','crelt','minus','rlutcs','rlut')
#derive('OBS','crett','plus','crest','crelt')

# Variables dérivées
derive('OBS','rstt','minus','rsdt','rsut')
derive('OBS','rsts','minus','rsds','rsus')
#BUH here: IMderive('OBS','rlts','minus','rlds','rlus')
derive('OBS','rlts','minus','rlus','rlds')
derive('OBS','rtt','minus','rstt','rlut')
derive('OBS','rts','plus','rsts','rlts')

derive('OBS','rstscs','minus','rsdscs','rsuscs')
derive('OBS','rsttcs','minus','rsdt'  ,'rsutcs')

derive('OBS','cress','minus','rsds','rsdscs')
derive('OBS','crels','minus','rlds','rldscs')
derive('OBS','crets','plus','cress','crels')

derive('OBS','crest','minus','rsutcs','rsut')
derive('OBS','crelt','minus','rlutcs','rlut')
derive('OBS','crett','plus','crest','crelt')

derive('OBS','hfns','plus','hfls','hfss')
derive('OBS','bil' ,'minus','rts','hfns')
derive('OBS','tsmtas','minus','ts','tas')

derive('OBS','rlah','minus','rlts','rlut')
derive('OBS','rtmp','plus','rldscs','rlutcs')
derive('OBS','rlahcs','minus','rlus','rtmp')

derive('OBS','rlahcre','minus','rlah','rlahcs')

derive('OBS','rsah','minus','rstt','rsts')
derive('OBS','rsahcs','minus','rsttcs','rstscs')
derive('OBS','rsahcre','minus','rsah','rsahcs')

derive('OBS','rah','plus','rsah','rlah')
derive('OBS','rahcs','plus','rsahcs','rlahcs')
derive('OBS','rahcre','minus','rah','rahcs')

cscript('mask_val_inf_1','cdo setrtomiss,0,1 ${in} ${out}')
derive('OBS','rsdt_mask','mask_val_inf_1','rsdt')
derive('OBS','albt','divide','rsut','rsdt_mask')
derive('OBS','albtcs','divide','rsutcs','rsdt_mask')


derive('OBS','albs','divide','rsus','rsds')
derive('OBS','dtr','minus','tasmax','tasmin')

calias('OBS', 'hflsevap', 'hfls', scale=1./2.5e6)

#########################################################################################
# Définition logique des donnees de type 'SE' de l'IPSL, et de leur localisation
# La periode n'est pas traitee comme un objet 'period' de CliMAF, mais comme une chaine 
# de caracteres dont le nom de facette est 'years'

# Le nommage des données découvert chez F.Hourdin est traité. On utilise la facette 'root' 
# pour indiquer le répertoire de base
# Les re-scaling habituels sont decrits a l'aide de la fonction calias
# Les variables dérivées sont décrites aussi
#########################################################################################

cproject('SE', 
         ('frequency','seasonnal'), 
         ('period','fx'), 
         'years',
         ('root','ROOT_SIM_DIR'), 
         separator='|')
#exemple de nom de fichier : NPv3.1ada_SE_1982_1991_1M_histmthCOSP.nc
pattern='${root}/${simulation}_SE_${years}_1M_histmth*.nc'
dataloc(project='SE',url=pattern, organization='generic')
    
    
calias('SE','hfls','flat',scale=-1.)
calias('SE','hfss','sens',scale=-1.)
calias('SE','pr','precip')
calias('SE','sfcWind','wind10m')
calias('SE','rldscs','LWdnSFCclr')
calias('SE','rlds','LWdnSFC')
calias('SE','rlus','LWupSFC')
calias('SE','rsdscs','SWdnSFCclr')
calias('SE','rsds','SWdnSFC')
calias('SE','rsuscs','SWupSFCclr')
calias('SE','rsus','SWupSFC')
calias('SE','rsutcs','SWupTOAclr')
calias('SE','rsut','SWupTOA')
calias('SE','rsdt','SWdnTOA')
calias('SE','rlut','topl')
calias('SE','rlutcs','topl0')
calias('SE','sfcWind','wind10m')
calias('SE','tas','t2m')
calias('SE','tasmax','ave_t2m_daily_max')
calias('SE','tasmin','ave_t2m_daily_min')
calias('SE','ts','tsol')
calias('SE','sst','tsol_oce')
calias('SE','huss','q2m')
calias('SE','hurs','rh2m')
calias('SE','tauu','taux_oce')
calias('SE','tauv','tauy_oce')
#calias('SE','psl','slp')
calias('SE','psl','slp', scale=0.01)
calias('SE','zg500','z500')
#calias('SE','pslhPa','slp', scale=0.01)
#calias('OBS','pslhPa','psl', scale=0.01)

#calias('SE','hfns','bils')
calias('SE','ta','temp')
calias('SE','ua','vitu')
calias('SE','va','vitv')
calias('SE','wap','vitw')
calias('SE','hus','ovap')
calias('SE','hur','rhum', scale=100.)
calias('SE','clt','cldt', scale=100.)
calias('SE','cltcalipso','cltcalipso', scale=100.)
calias('SE','clhcalipso','clhcalipso', scale=100.)
calias('SE','clmcalipso','clmcalipso', scale=100.)
calias('SE','cllcalipso','cllcalipso', scale=100.)

# Variables dérivées
derive('SE','rstt','minus','rsdt','rsut')
derive('SE','rsts','minus','rsds','rsus')
#BUG HERE derive('SE','rlts','minus','rlds','rlus')
derive('SE','rlts','minus','rlus','rlds')
derive('SE','rltscs','minus','rldscs','rluscs')

derive('SE','rtt','minus','rstt','rlut')
derive('SE','rts','plus','rsts','rlts')

derive('SE','rstscs','minus','rsdscs','rsuscs')
derive('SE','rsttcs','minus','rsdt'  ,'rsutcs')

derive('SE','cress','minus','rsds','rsdscs')
derive('SE','crels','minus','rlds','rldscs')
derive('SE','crets','plus','cress','crels')

derive('SE','crest','minus','rsutcs','rsut')
derive('SE','crelt','minus','rlutcs','rlut')
derive('SE','crett','plus','crest','crelt')

derive('SE','hfns','plus','hfls','hfss')
derive('SE','bil' ,'minus','rts','hfns')
derive('SE','tsmtas','minus','ts','tas')

derive('SE','rlah','minus','rlts','rlut')
derive('SE','rtmp','plus','rldscs','rlutcs')
derive('SE','rlahcs','minus','rlus','rtmp')

#let rlahcs=rlus-rldscs-rlutcs
derive('SE','rlahcre','minus','rlah','rlahcs')
#
derive('SE','rsah','minus','rstt','rsts')
derive('SE','rsahcs','minus','rsttcs','rstscs')
derive('SE','rsahcre','minus','rsah','rsahcs')

derive('SE','rah','plus','rsah','rlah')
derive('SE','rahcs','plus','rsahcs','rlahcs')
derive('SE','rahcre','minus','rah','rahcs')

cscript('mask_valLMDZ_inf_1','cdo setrtomiss,0,1 ${in} ${out}')
derive('SE','rsdt_maskLMDZ','mask_valLMDZ_inf_1','rsdt')
derive('SE','albt','divide','rsut','rsdt_maskLMDZ')
derive('SE','albtcs','divide','rsutcs','rsdt_maskLMDZ')
#derive('SE','albt','divide','rsut','rsdt')
#derive('SE','albtcs','divide','rsutcs','rsdt')

derive('SE','albs','divide','rsus','rsds')

derive('SE','dtr','minus','tasmax','tasmin')

# -- P - E
calias('SE', 'hflsevap', 'flat', scale=-1./2.5e6, filenameVar='histmth')
derive('SE', 'pme', 'minus', 'pr' ,'hflsevap')


# -- Atmospheric Variables on vertical levels
for tmpvar in ['ua', 'va', 'ta', 'hus', 'hur']:
    for tmplev in ['850', '700', '500', '200']:
        derive('OBS', tmpvar+tmplev, 'ccdo', tmpvar, operator='intlevel,'+tmplev)
        derive('SE', tmpvar+tmplev, 'ccdo', tmpvar, operator='intlevel,'+tmplev+'00')


# let albt=100*(rsut/rsdt)
# let albs=100*(rsus/rsds)
# let tsk=ts+273.18
# let pslhPa=psl/100.
# let tasc=tas-273.16

def all_SE_simulations():
    """
    Listage de toutes les simulations du projet par listage de leurs 
    fichiers de données et décodage de leur nom
    
    Il faudrait quelque chose de moins adhoc ...
    """
    import re
    simus=[]
    # Listage des fichiers de donnees
    a=ds(project='SE',years="*", simulation='*', variable='*')
    for f in a.baseFiles().split(' ') : 
        basename=f.split('/')[-1].replace('_SE','')
        #basename=re.sub(r'_1M_histmth.*.nc','',basename)
        basename=re.sub(r'_1M_histmth*.*.nc','',basename)
        if basename not in simus  : simus.append(basename)
    return simus


def svsg(simulation,variable,season='YEAR',grid='',root='ROOT_SIM_DIR'):
    """
    Rend l'objet CliMAF pour une variable d'une simulation, pour les  données 
    - SE d'un intervalle d'annéees (forme YYY1_YYY2)
    - ou OBS 
    et une saison et une grille

    (Le nom de la fonction est la concaténation des intiales des arguments)

    Exemples : 
    >>> svsg('NPv3.1ada','1982_1991','hurs')
    >>> svsg(simulation='NPv3.1ada',years='1982_1991',variable='hurs',season='DJF, grid='')
    >>> svsg('OBS','','hurs','JJA')

    Il faut au préalable avoir déclaré les projets 'SE' et 'OBS'
    """
    if simulation != 'OBS' :
        # Il faut identifier les annees dans le nom de la simu
        yeardeb=simulation.split('_')[1]
        yearfin=simulation.split('_')[2]
        years=yeardeb+"_"+yearfin
        simulation=simulation.split('_')[0]
        data=ds(project='SE',variable=variable,years=years, simulation=simulation)
    else:
        data=ds(project='OBS',variable=variable)
    if season != 'YEAR' : 
        seas=ccdo(data,operator=seasons[season])
    else : 
        seas=time_average(data)
    rds=seas
    if grid != '' : 
        if grid not in grid_remapfiles : 
            print("La grille %s n'est pas connue"%options.grid) ; 
            return None
        rds=regrid(seas,fds(grid_remapfiles[grid],period='fx'))
    return rds

def bias_SE(simu,variable,season='YEAR',grid='',root='ROOT_SIM_DIR') :
    """
    Calcule le biais pour une variable d'une simu, vs 'OBS'
    Rend -999. en cas de probleme
    """
    try :
        sim=svsg(simu ,variable,season,grid,root)
        ref=svsg('OBS',variable,season,grid)
        if (grid == '' ) : ref=regrid(ref,sim)
        dif=minus(sim,ref)
        return cvalue(space_average(dif))
    except :
        return -999.

def is3d(variable) :
    if variable in ['ta','ua','va','hus','hur','wap','cl','clw','cli','mc','tro3'] :
        return True
    return False


def varlongname(variable,short=False):
   """
   Returns the long name of variable
   """
   shortvarname = longvarname = variable
   if variable=='tas':
        longvarname  = '2M Temperature'
        shortvarname = '2M Temp.'
   if variable=='tasmax':
        longvarname  = '2M daymax Temperature'
        shortvarname = '2M daymx Temp.'
   if variable=='tasmin':
        longvarname  = '2M daymin Temperature'
        shortvarname = '2M daymn Temp.'
   if variable=='DRT':
        longvarname  = '2M diurnal amplitude temperature'
        shortvarname = '2M dtr '
   if variable=='pr':
        longvarname  = 'Precipitation'
        shortvarname = 'Precip.'
   if variable=='psl':
        longvarname  = 'Sea Level Pressure'
        shortvarname = 'Sea Level Pres.'
   if variable=='ua':
        longvarname  = 'Zonal Wind'
        shortvarname = 'U-Wind'
   if variable=='va':
        longvarname  = 'Meridional Wind'
        shortvarname = 'V-Wind'
   if variable=='ta':
        longvarname  = 'Air Temperature'
        shortvarname = 'Air Temp.'
   if variable=='hus':
        longvarname  = 'Specific Humidity'
        shortvarname = 'Sp. Humidity'
   if variable=='hur':
        longvarname  = 'Relative Humidity'
        shortvarname = 'Rel. Humidity'
   if variable=='huss':
        longvarname  = 'Specific Humidity at Surface'
        shortvarname = 'Sp. Humidity (surf)'
   if variable=='rstt':
        longvarname  = 'Rad SW Total TOA'
        shortvarname = 'Rad SW Total TOA'
   if variable=='rsts':
        longvarname  = 'Total SW rad surface'
        shortvarname = 'Total SW rad surf.'
   if variable=='rtt':
        longvarname  = 'Total Radiation TOA'
        shortvarname = 'Total Rad. TOA'
   if variable=='crelt':
        longvarname  = 'Longwave Cloud Radiative Effect TOA'
        shortvarname = 'LW CRE TOA'
   if variable=='crest':
        longvarname  = 'Shortwave Cloud Radiative Effect TOA'
        shortvarname = 'SW CRE TOA'
   if variable=='crett':
        longvarname  = 'Total CRE TOA'
        shortvarname = 'Total CRE TOA'
   if variable=='hfls':
        longvarname  = 'Latent Heat Flux'
        shortvarname = 'Latent HF'
   if variable=='hfss':
        longvarname  = 'Sensible Heat Flux'
        shortvarname = 'Sensible HF'
   if variable=='hfns':
        longvarname  = 'Surface Total Heat Flux'
        shortvarname = 'Surf. Tot. HF'
   if variable=='zg500':
        longvarname  = '500mb geopotential height'
        shortvarname = ''
   if variable=='rsut':
        longvarname  = 'Upward SW rad TOA'
        shortvarname = ''
   if variable=='rlut':
        longvarname  = 'Outgoing Long Wave Radiation'
        shortvarname = 'OLR'
   if variable=='rlutcs':
        longvarname  = 'Clear Sky OLR'
        shortvarname = 'Clear sky OLR'
   if variable=='albs':
        longvarname  = 'Surface albedo'
        shortvarname = ''
   if variable=='albt':
        longvarname  = 'Planetary albedo'
        shortvarname = ''
   if variable=='albtcs':
        longvarname  = 'Clear Sky Planetary albedo'
        shortvarname = ''
   if variable=='cress':
        longvarname  = 'SW CRE surface'
        shortvarname = ''
   if variable=='crels':
        longvarname  = 'LW CRE surface'
        shortvarname = ''
   if variable=='crets':
        longvarname  = 'Total CRE surface'
        shortvarname = ''
   if variable=='rts':
        longvarname  = 'Total radiation surface'
        shortvarname = ''
   if variable=='rah':
        longvarname  = 'Atm. Rad. Heat.'
        shortvarname = ''
   if variable=='rahcs':
        longvarname  = 'Atm. Rad. Heat. - clear sky'
        shortvarname = ''
   if variable=='rahcre':
        longvarname  = 'Atm. rad. Heat. - CRE'
        shortvarname = ''
   if variable=='rsah':
        longvarname  = 'Atm. SW Heat.'
        shortvarname = ''
   if variable=='rsahcs':
        longvarname  = 'Atm. SW Heat. - Clear sky'
        shortvarname = ''
   if variable=='rsahcre':
        longvarname  = 'Atm. SW Heat. - CRE'
        shortvarname = ''
   if variable=='rlah':
        longvarname  = 'Atm. LW Heat.'
        shortvarname = ''
   if variable=='rlahcs':
        longvarname  = 'Atm. LW Heat. - Clear sky'
        shortvarname = ''
   if variable=='rlahcre':
        longvarname  = 'Atm. LW Heat. - CRE'
        shortvarname = ''
   if variable=='cltcalipso':
        longvarname  = 'Total Cloud Cover'
        shortvarname = ''
   if variable=='cllcalipso':
        longvarname  = 'Low Cloud Cover'
        shortvarname = ''
   if variable=='clmcalipso':
        longvarname  = 'Medium Cloud Cover'
        shortvarname = ''
   if variable=='clhcalipso':
        longvarname  = 'High Cloud Cover'
        shortvarname = ''
   if variable=='rlds':
        longvarname  = 'Downward LW rad at Surface'
        shortvarname = ''
   if variable=='rldscs':
        longvarname  = 'Downward LW rad at Surface - Clear Sky'
        shortvarname = ''
   if variable=='hurs':
        longvarname  = 'Relative Humidity at Surface'
        shortvarname = ''
   if variable=='rlus':
        longvarname  = 'Upward LW rad at Surface'
        shortvarname = ''
   if variable=='rsus':
        longvarname  = 'Upward SW rad at Surface'
        shortvarname = ''
   if variable=='rsuscs':
        longvarname  = 'Upward SW rad at Surface - Clear Sky'
        shortvarname = ''
   if variable=='rsdscs':
        longvarname  = 'Downward SW rad at Surface - Clear Sky'
        shortvarname = ''
   if variable=='rsds':
        longvarname  = 'Downward SW rad at Surface'
        shortvarname = ''
   if variable=='rsucs':
        longvarname  = 'Upward SW rad at Surface - Clear Sky'
        shortvarname = ''
   if variable=='rsutcs':
        longvarname  = 'Upward SW rad at TOA - Clear Sky'
        shortvarname = ''
   if variable=='tauu':
        longvarname  = 'Zonal Wind Stress'
        shortvarname = ''
   if variable=='tauv':
        longvarname  = 'Meridional Wind Stress'
        shortvarname = ''
   if variable=='pme':
        longvarname  = 'P-E Precip-Evap(hfls/2.5e6) mm/day'
        shortvarname = ''
   #
   if short:
        return shortvarname
   else:
        return longvarname