import pandas as pd
from netCDF4 import Dataset

BGC_EUR0275_8daily = {'name': 'CLM5_DETECT',
                      'version': (0, 1, 0),
                      'path': '/p/data1/jibg31/CLM5_EUR-0275/BGC_outputs/carbon_control_DETECT/',
                      'type_file': 'netcdf',
                      'year_start': 1960,
                      'month_start': 1,
                      'year_end': 1989,
                      'month_end': 12,
                      'leapday': False,
                      'resolution_time': '8D',
                      'grid': 'EU3',
                      'variables': ['NEE', 
                                    'GPP', 
                                    'TLAI',
                                    'TOTECOSYSC',
                                    'TOTSOMC',
                                    'TOTVEGC',
                                    'LEAFC',
                                    'QFLX_EVAP_TOT',
                                    'QFLX_EVAP_VEG',
                                    'ZWT'],
                      'variable_names': {'NEE': 'NEE', 
                                         'GPP': 'GPP', 
                                         'TLAI': 'TLAI',
                                         'TOTECOSYSC': 'TOTECOSYSC',
                                         'TOTSOMC': 'TOTSOMC',
                                         'TOTVEGC': 'TOTVEGC',
                                         'LEAFC': 'LEAFC',
                                         'QFLX_EVAP_TOT': 'QFLX_EVAP_TOT',
                                         'QFLX_EVAP_VEG': 'QFLX_EVAP_VEG',
                                         'ZWT': 'ZWT'},
                      'variable_dimensions': {'NEE': ['time', 'lat', 'lon'], 
                                              'GPP': ['time', 'lat', 'lon'], 
                                              'TLAI': ['time', 'lat', 'lon'],
                                              'TOTECOSYSC': ['time', 'lat', 'lon'],
                                              'TOTSOMC': ['time', 'lat', 'lon'],
                                              'TOTVEGC': ['time', 'lat', 'lon'],
                                              'LEAFC': ['time', 'lat', 'lon'],
                                              'QFLX_EVAP_TOT': ['time', 'lat', 'lon'],
                                              'QFLX_EVAP_VEG': ['time', 'lat', 'lon'],
                                              'ZWT': ['time', 'lat', 'lon']},
                      'variable_units': {'NEE': 'g C / m^2 / s', 
                                         'GPP': 'g C / m^2 / s', 
                                         'TLAI': 'm^2 / m^2',
                                         'TOTECOSYSC': 'g C / m^2',
                                         'TOTSOMC': 'g C / m^2',
                                         'TOTVEGC': 'g C / m^2',
                                         'LEAFC': 'g C / m^2',
                                         'QFLX_EVAP_TOT': 'kg / m^2 / s',
                                         'QFLX_EVAP_VEG': 'mm H2O / s',
                                         'ZWT': 'm'},
                      'mask_value': None}

BGC_EUR0275_8daily_control = {'name': 'CLM5_DETECT_control',
                              'version': (0, 1, 0),
                              'path': '/p/data1/jibg31/CLM5_EUR-0275/BGC_outputs/carbon_control_DETECT/control/',
                              'type_file': 'netcdf',
                              'year_start': 1960,
                              'month_start': 1,
                              'year_end': 1989,
                              'month_end': 12,
                              'leapday': False,
                              'resolution_time': '8D',
                              'grid': 'EU3',
                              'variables': ['NEE', 
                                            'GPP', 
                                            'TLAI',
                                            'TOTECOSYSC',
                                            'TOTSOMC',
                                            'TOTVEGC',
                                            'LEAFC',
                                            'QFLX_EVAP_TOT',
                                            'QFLX_EVAP_VEG',
                                            'ZWT'],
                              'variable_names': {'NEE': 'NEE', 
                                                 'GPP': 'GPP', 
                                                 'TLAI': 'TLAI',
                                                 'TOTECOSYSC': 'TOTECOSYSC',
                                                 'TOTSOMC': 'TOTSOMC',
                                                 'TOTVEGC': 'TOTVEGC',
                                                 'LEAFC': 'LEAFC',
                                                 'QFLX_EVAP_TOT': 'QFLX_EVAP_TOT',
                                                 'QFLX_EVAP_VEG': 'QFLX_EVAP_VEG',
                                                 'ZWT': 'ZWT'},
                              'variable_dimensions': {'NEE': ['time', 'lat', 'lon'], 
                                                      'GPP': ['time', 'lat', 'lon'], 
                                                      'TLAI': ['time', 'lat', 'lon'],
                                                      'TOTECOSYSC': ['time', 'lat', 'lon'],
                                                      'TOTSOMC': ['time', 'lat', 'lon'],
                                                      'TOTVEGC': ['time', 'lat', 'lon'],
                                                      'LEAFC': ['time', 'lat', 'lon'],
                                                      'QFLX_EVAP_TOT': ['time', 'lat', 'lon'],
                                                      'QFLX_EVAP_VEG': ['time', 'lat', 'lon'],
                                                      'ZWT': ['time', 'lat', 'lon']},
                              'variable_units': {'NEE': 'g C / m^2 / s', 
                                                 'GPP': 'g C / m^2 / s', 
                                                 'TLAI': 'm^2 / m^2',
                                                 'TOTECOSYSC': 'g C / m^2',
                                                 'TOTSOMC': 'g C / m^2',
                                                 'TOTVEGC': 'g C / m^2',
                                                 'LEAFC': 'g C / m^2',
                                                 'QFLX_EVAP_TOT': 'kg / m^2 / s',
                                                 'QFLX_EVAP_VEG': 'mm H2O / s',
                                                 'ZWT': 'm'},
                              'mask_value': None}


BGC_EU3_8daily = {
    'name': 'BGC_EU3',
    'version': (1, 1, 7),
    'path': '/p/scratch/cjibg31/jibg3105/data/CLM5EU3/007/join_8d/',
    'type_file': 'netcdf',
    'year_start': 1995,
    'month_start': 1,
    'year_end': 2018,
    'month_end': 12,
    'resolution_time': '8D',
    'leapday': False,
    'grid': 'EU3',
    'variables': [
                    'GPP', 
                    'ET',
                    'ESOIL',
                    'Tr',
                    'BTRAN',
                    'Gs',
                    'Runoff',
                    'SM',
                    'WTD',
                    'H',
                    'SoiH',
                    'Rnet'
                 ],
    'variable_names': {
                        'GPP': 'GPP', 
                        'ET': 'QFLX_EVAP_TOT',
                        'ESOIL': 'QFLX_EVAP_GRND',
                        'Tr': 'QFLX_EVAP_VEG',
                        'BTRAN': 'BTRANMN',
                        'Gs': 'GSSHA',
                        'Runoff': 'QOVER',
                        'SM': 'H2OSOI',
                        'WTD': 'ZWT',
                        'H': 'Qh',
                        'SoiH': 'FGR',
                        'Rnet': 'Rnet',
                       },
    'variable_dimensions': {
                            'GPP': ['time', 'lat', 'lon'], 
                            'ET': ['time', 'lat', 'lon'],
                            'ESOIL': ['time', 'lat', 'lon'],
                            'Tr': ['time', 'lat', 'lon'],  
                            'BTRAN': ['time', 'lat', 'lon'],
                            'Gs': ['time', 'lat', 'lon'],
                            'Runoff': ['time', 'lat', 'lon'],
                            'SM': ['time', 'layer', 'lat', 'lon'],
                            'WTD': ['time', 'lat', 'lon'],
                            'H': ['time', 'lat', 'lon'],
                            'SoiH': ['time', 'lat', 'lon'],
                            'Rnet': ['time', 'lat', 'lon'],
                            }, 
    'variable_units': {'GPP': 'g/s',
                       'ET': 'mm/s',
                       'ESOIL': 'mm/s',
                       'Tr': 'mm/s',
                       'BTRAN': 'dimensionless',
                       'Gs': 'umol/s',
                       'Runoff': 'mm/s',
                       'SM': 'mm^3/mm^3',
                       'WTD': 'm',
                       'H': 'W/m^2',
                       'SoiH': 'W/m^2',
                       'Rnet': 'W/m^2'},
    'mask_value': None
}



BGC_EU3_daily = {
    'name': 'BGC_EU3',
    'version': (1, 1, 7),
    'path': '/p/data1/jibg31/CLM5.0_EU3km/outputs/007/join_d/',
    'type_file': 'netcdf',
    'year_start': 1995,
    'month_start': 1,
    'year_end': 2018,
    'month_end': 12,
    'resolution_time': 'D',
    'leapday': False,
    'grid': 'EU3',
    'variables': [
                    'GPP', 
                    'ET',
                    'ESOIL',
                    'Tr',
                    'BTRAN',
                    'Gs',
                    'Runoff',
                    'SM',
                    'WTD',
                    'H',
                    'SoiH',
                    'Rnet'
                 ],
    'variable_names': {
                        'GPP': 'GPP', 
                        'ET': 'QFLX_EVAP_TOT',
                        'ESOIL': 'QFLX_EVAP_GRND',
                        'Tr': 'QFLX_EVAP_VEG',
                        'BTRAN': 'BTRANMN',
                        'Gs': 'GSSHA',
                        'Runoff': 'QOVER',
                        'SM': 'H2OSOI',
                        'WTD': 'ZWT',
                        'H': 'Qh',
                        'SoiH': 'FGR',
                        'Rnet': 'Rnet',
                       },
    'variable_dimensions': {
                            'GPP': ['time', 'lat', 'lon'], 
                            'ET': ['time', 'lat', 'lon'],
                            'ESOIL': ['time', 'lat', 'lon'],
                            'Tr': ['time', 'lat', 'lon'],  
                            'BTRAN': ['time', 'lat', 'lon'],
                            'Gs': ['time', 'lat', 'lon'],
                            'Runoff': ['time', 'lat', 'lon'],
                            'SM': ['time', 'layer', 'lat', 'lon'],
                            'WTD': ['time', 'lat', 'lon'],
                            'H': ['time', 'lat', 'lon'],
                            'SoiH': ['time', 'lat', 'lon'],
                            'Rnet': ['time', 'lat', 'lon'],
                            }, 
    'variable_units': {'GPP': 'g/s',
                       'ET': 'mm/s',
                       'ESOIL': 'mm/s',
                       'Tr': 'mm/s',
                       'BTRAN': 'dimensionless',
                       'Gs': 'umol/s',
                       'Runoff': 'mm/s',
                       'SM': 'mm^3/mm^3',
                       'WTD': 'm',
                       'H': 'W/m^2',
                       'SoiH': 'W/m^2',
                       'Rnet': 'W/m^2'},
    'mask_value': None
}

BGC_EU3_weekly = {
    'name': 'BGC_EU3',
    'version': (1, 1, 7),
    'path': '/p/scratch/cjibg31/jibg3105/data/CLM5EU3/007/join_w/',
    'type_file': 'netcdf',
    'year_start': 1995,
    'month_start': 1,
    'year_end': 2018,
    'month_end': 12,
    'resolution_time': '7D',
    'leapday': False,
    'grid': 'EU3',
    'variables': [
                    'GPP', 
                    'ET',
                    'ESOIL',
                    'Tr',
                    'BTRAN',
                    'Gs',
                    'Runoff',
                    'SM',
                    'WTD',
                    'H',
                    'SoiH',
                    'Rnet'
                 ],
    'variable_names': {
                        'GPP': 'GPP', 
                        'ET': 'QFLX_EVAP_TOT',
                        'ESOIL': 'QFLX_EVAP_GRND',
                        'Tr': 'QFLX_EVAP_VEG',
                        'BTRAN': 'BTRANMN',
                        'Gs': 'GSSHA',
                        'Runoff': 'QOVER',
                        'SM': 'H2OSOI',
                        'WTD': 'ZWT',
                        'H': 'Qh',
                        'SoiH': 'FGR',
                        'Rnet': 'Rnet',
                       },
    'variable_dimensions': {
                            'GPP': ['time', 'lat', 'lon'], 
                            'ET': ['time', 'lat', 'lon'],
                            'ESOIL': ['time', 'lat', 'lon'],
                            'Tr': ['time', 'lat', 'lon'],  
                            'BTRAN': ['time', 'lat', 'lon'],
                            'Gs': ['time', 'lat', 'lon'],
                            'Runoff': ['time', 'lat', 'lon'],
                            'SM': ['time', 'layer', 'lat', 'lon'],
                            'WTD': ['time', 'lat', 'lon'],
                            'H': ['time', 'lat', 'lon'],
                            'SoiH': ['time', 'lat', 'lon'],
                            'Rnet': ['time', 'lat', 'lon'],
                            }, 
    'variable_units': {'GPP': 'g/s',
                       'ET': 'mm/s',
                       'ESOIL': 'mm/s',
                       'Tr': 'mm/s',
                       'BTRAN': 'dimensionless',
                       'Gs': 'umol/s',
                       'Runoff': 'mm/s',
                       'SM': 'mm^3/mm^3',
                       'WTD': 'm',
                       'H': 'W/m^2',
                       'SoiH': 'W/m^2',
                       'Rnet': 'W/m^2'},
    'mask_value': None
}


class CLM5_PFT:

    dir_surfdata: str = '/p/scratch/cjibg31/jibg3105/CESMDataRoot/InputData/lnd/clm2/surfdata_map/'
    file_surfdata: str = 'surfdata_CLM5EU3_v4_pos.nc'

    names = {1:'NET Temperate',
             2:'NET Boreal',
             3:'NDT Boreal',
             4:'BET Tropical',
             5:'BET Temperate',
             6:'BDT tropical',
             7:'BDT temperate',
             8:'BDT boreal',
             9:'BES temperate',
             10:'BDS temperate',
             11:'BDS boreal',
             12:'C3 grass arctic',
             13:'C3 grass',
             14:'C4 grass',
             15:'Crop R',
             16:'Crop I'}

    soil_layers = pd.DataFrame(data = [[0.0, 0.0, 0.0],
                                       [0.010,0.020,0.020],
                                       [0.040,0.040,0.060],
                                       [0.090,0.060,0.120],
                                       [0.160,0.080,0.200],
                                       [0.260,0.120,0.320],
                                       [0.400,0.160,0.480],
                                       [0.580,0.200,0.680],
                                       [0.800,0.240,0.920],
                                       [1.060,0.280,1.200],
                                       [1.360,0.320,1.520],
                                       [1.700,0.360,1.880],
                                       [2.080,0.400,2.280],
                                       [2.500,0.440,2.720],
                                       [2.990,0.540,3.260],
                                       [3.580,0.640,3.900],
                                       [4.270,0.740,4.640],
                                       [5.060,0.840,5.480],
                                       [5.950,0.940,6.420],
                                       [6.940,1.040,7.460],
                                       [8.030,1.140,8.600],
                                       [9.795,2.390,10.990],
                                       [13.328,4.676,15.666],
                                       [19.483,7.635,23.301],
                                       [28.871,11.140,34.441],
                                       [41.998,15.115,49.556]],
                               columns = ['node_depth',
                                          'thickness',
                                          'depth_interface'])
    
    parameter = {'betaroot': {1: 0.976,
                              2: 0.943,
                              3: 0.943,
                              4: 0.993,
                              5: 0.966,
                              6: 0.993,
                              7: 0.966,
                              8: 0.943,
                              9: 0.964,
                              10: 0.964,
                              11: 0.914,
                              12: 0.914,
                              13: 0.943,
                              14: 0.943,
                              15: 0.943,
                              16: 0.943}}

    def load_surface_var(self,
                         variables: str | list[str]):
        
        data = Dataset(f'{self.dir_surfdata}/{self.file_surfdata}')

        values = {v: data.variables[v][:] for v in variables}

        data.close()

        return values