# Copyright 2022 - Barcelona Supercomputing Center
# Author: Etienne Tourigny, Rodrigo Martín Posada
# MIT License
from datetime import date
import re
import f90nml
import numpy as np
import yaml
from yaml.loader import SafeLoader


# constants
earth_area = 6371000.**2

# CO2 model fraction values taken from TM5 code
xmair = 28.964
xmc  = 12.01115
xmo  = 15.9994
xmco2 = xmc + xmo*2.0
fscale = xmair/xmco2
grav = 9.80665

# total airmass of atmosphere is 5.1480e18 kg
# Trenberth, K.E., Smith, L., 2005. The mass of the atmosphere: a constraint on global analyses. J. Clim. https://doi.org/10.1175/JCLI-3299.1.
total_air_mass = 5.1480e18
# this is valid for a NH/SH box model only!
# TODO add a global parameter to define this in case we use a global box model!
#total_air_mass = 5.1480e18 / 2.0

def scale_factor_from_str(scale_str):
    if scale_str == 'co2_mass_to_ppm':
        return 1 / 1e-6 / total_air_mass * fscale
    elif scale_str == 'co2_ppm_to_mass':
        return 1 * 1e-6 * total_air_mass / fscale
    elif scale_str == 'second2day':
        return 24 * 60 * 60
    else:
        raise Exception('scale_factor_from_str() not defined for {}'.format(scale_str))

def days_since_refyear(yy, mm, dd, ref_year):
    '''
    Returns the number of days since a reference year.
    '''
    d0 = date(ref_year, 1, 1)
    d1 = date(yy, mm, dd)
    delta = d1 - d0
    return delta.days


def is_leap(year):
    '''
    Leap year test
    '''
    return year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)

def extract_timeframe_from_date(name, curr_date):
    '''
    Extract `name` part from `curr_date`.
    '''
    if name == 'year':
        return curr_date.year
    else:
        raise Exception('Timeframe not defined for ' + name)


def filename_from_pattern(file_pattern, curr_date):
    '''
    Returns a filename out of a filename pattern and the current date.

    I.e: from `HadISST2_prelim_0to360_alldays_sst_[year].nc` to `HadISST2_prelim_0to360_alldays_sst_2013.nc`.
    '''
    re_pattern = re.compile(r'\[[A-Za-z]+\]')
    results = re.finditer(re_pattern, file_pattern)

    start = 0
    file_name = ''
    for m in results:
        m_start, m_end = m.span()
        file_name += file_pattern[start:m_start]
        timeframe_name = m.group(0)[1:-1]
        file_name += str(extract_timeframe_from_date(timeframe_name, curr_date))
        start = m_end
    file_name += file_pattern[start:]
    return file_name


def read_namelist(NAMELIST_FILE_NAME,NAMELIST_SECTION):
    '''
    Reads the namelist.
    '''
    nml = f90nml.read(NAMELIST_FILE_NAME)
    nml_dict = nml[NAMELIST_SECTION]
    RunLengthSec = nml_dict['RunLengthSec']
    TimeStepSec = nml_dict['TimeStepSec']
    StartYear = nml_dict['StartYear']
    StartMonth = nml_dict['StartMonth']
    StartDay = nml_dict['StartDay']
    FixYear = nml_dict['FixYear']
    LDebug = nml_dict['LDebug']

    if 'FileInputVars' in nml_dict or 'Vars' in nml_dict:
        varname='FileInputVars' if 'FileInputVars' in nml_dict else 'Vars'
        vars_data = np.transpose(np.asarray(nml_dict[varname], dtype=object))
        #print(str(vars_data))
        fileIn_vars_config = [None] * vars_data.shape[0]
        for i in range(len(fileIn_vars_config)):
            curr_data = vars_data[i]
            conf = {
                'id': curr_data[0],
                'grid_name': curr_data[1],
                'oasis_name': curr_data[2],
                'file_pattern': curr_data[3],
                'netcdf_variable': curr_data[4],
                'yref_min': curr_data[5],
                'yref_max': curr_data[6],
                'timestep': curr_data[7],
                'interpolate': curr_data[8],
                'scale_factor': curr_data[9],
                'offset': curr_data[10],
                'min': curr_data[11],
                'max': curr_data[12],
                'update': curr_data[13],
                'accum': curr_data[14],
            }
            fileIn_vars_config[i] = conf
    elif 'FileListSST' in nml_dict and 'LInterpolate' in nml_dict and 'FileListSIC' in nml_dict:
        # BACKWARDS COMPATIBILITY: Old namelist configuration for AMIP forcing
        fileIn_vars_config = [
            {
                'id': 'AMIP_sst_monthly',
                'grid_name': 'AMIP',
                'oasis_name': 'AMIP_sst',
                'file_pattern': nml_dict['FileListSST'],
                'netcdf_variable': 'tosbcs',
                'yref_min': 1870,
                'yref_max': 2016,
                'timestep': 'monthly',
                'interpolate': nml_dict['LInterpolate'],
                'scale_factor': 1,
                'offset': 273.15,
                'min': 271.38,
                'max': None,
            },
            {
                'id': 'AMIP_sic_monthly',
                'grid_name': 'AMIP',
                'oasis_name': 'AMIP_sic',
                'file_pattern': nml_dict['FileListSIC'],
                'netcdf_variable': 'siconcbcs',
                'yref_min': 1870,
                'yref_max': 2016,
                'timestep': 'monthly',
                'interpolate': nml_dict['LInterpolate'],
                'scale_factor': 0.01,
                'offset': 0,
                'min': 0,
                'max': 1,
            }
        ]
    elif 'AmipFileRoot' in nml_dict:
        # BACKWARDS COMPATIBILITY: Old namelist configuration for AMIP primavera
        fileIn_vars_config = [
            {
                'id': 'AMIP_sst_daily',
                'grid_name': 'PSST',
                'oasis_name': 'AMIP_sst',
                'file_pattern': nml_dict['AmipFileRoot'] + '_sst_[year].nc',
                'netcdf_variable': 'sst',
                'yref_min': 1850,
                'yref_max': 2016,
                'timestep': 'daily',
                'interpolate': False,
                'scale_factor': 1,
                'offset': 0,
                'min': None,
                'max': None,
            },
            {
                'id': 'AMIP_sic_daily',
                'grid_name': 'PSIC',
                'oasis_name': 'AMIP_sic',
                'file_pattern': nml_dict['AmipFileRoot'] + '_sic_[year].nc',
                'netcdf_variable': 'sic',
                'yref_min': 1850,
                'yref_max': 2016,
                'timestep': 'daily',
                'interpolate': False,
                'scale_factor': 1,
                'offset': 0,
                'min': None,
                'max': None,
            },
        ]
    else:
        raise Exception('DataCoupler: Invalid namelist configuration')

    # Check all ids are unique
    ids = [conf['id'] for conf in fileIn_vars_config]
    if np.unique(ids).size != len(ids):
        raise Exception('DataCoupler: Invalid configuration ids, all ids must be unique')

    if 'OasisOutputVars' in nml_dict:
        vars_data = np.transpose(np.asarray(nml_dict['OasisOutputVars'], dtype=object))
        oasisOut_vars_config = [None] * vars_data.shape[0]
        for i in range(len(oasisOut_vars_config)):
            curr_data = vars_data[i]
            conf = {
                'id': curr_data[0],
                'grid_name': curr_data[1],
                'timestep': curr_data[2],
                'scale_factor': curr_data[3],
                'offset': curr_data[4],
                'reset': curr_data[5],
            }
            oasisOut_vars_config[i] = conf
        # Check all ids are unique
        ids = [conf['id'] for conf in oasisOut_vars_config]
        if np.unique(ids).size != len(ids):
            raise Exception('DataCoupler: Invalid configuration ids, all ids must be unique')
    else:
        oasisOut_vars_config = None


    return RunLengthSec, TimeStepSec, StartYear, StartMonth, StartDay, FixYear, LDebug, fileIn_vars_config, oasisOut_vars_config



def read_ymllist(YAML_NAMELIST_FILE_NAME):
    '''
    Reads the namelist.
    '''
    yamel_file=YAML_NAMELIST_FILE_NAME

    with open(yamel_file) as f:
        nmly = yaml.load(f, Loader=SafeLoader)
        f.close()
    
    RunLengthSec = nmly['RunLengthSec']
    TimeStepSec = nmly['TimeStepSec']
    StartYear = nmly['StartYear']
    StartMonth = nmly['StartMonth']
    StartDay = nmly['StartDay']
    FixYear = nmly['FixYear']
    #LDebug = nmly['LDebug']

    return RunLengthSec, TimeStepSec, StartYear, StartMonth, StartDay, FixYear

# TODO will be deleted
# ### Read in YAML file 
#     current_path = os.getcwd()
#     parser=argparse.ArgumentParser()
#     parser.add_argument("--yaml_file_path",type=str)
#     args = parser.parse_args()
#     yamel_file_path = os.path.join(current_path,args.yaml_file_path)
#     if os.path.isfile(yamel_file_path) == False: 
#         print("YAML file does not exist. Aborted")
#         raise NameError("YAML file does not exist.")
        
#     # Open the file and load the file 


# ### Read in file 
#     #read_yaml(read_in_data)
#     model = read_in_data.get("model")
#     model_name = model.get("model_name")
    
#     # Read in Simulation information
#     simulation = read_in_data.get("simulation")
#     simulation_coupling_interval = simulation.get("coupling_interval")
#     simulation_run_length_sec = simulation.get("run_length_sec")
#     simulation_nx = simulation.get("nx")

#     # Read in Coupling infomration
#     coupling = read_in_data.get("coupling")
#     coupling_out_vars = coupling.get("out_vars")
#     coupling_in_vars = coupling.get("in_vars")
#     restart_file = coupling.get("restart_file")  # This feature is not used for now. TODO 