# 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
import sys 
import os
import collections 


# 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


## YAMLUtil

# Read the DataCoupler YAML

def read_DataCoupler_yaml_input(yaml_file):
    '''
    Reads the experiment yaml file 
    '''

    with open(yaml_file) as file:
        read_in = yaml.load(file, Loader=SafeLoader)
        file.close()
    
    # Model Information
    ModelNameSend       = read_in.get('ModelNameSend')
    ModelNameReceive      = read_in.get('ModelNameReceive')
    LogFileName     = read_in.get('LogFileName')

    ## Run Information 

    RunLengthSec    = read_in.get('RunLengthSec')
    TimeStepSec     = read_in.get('TimeStepSec')
    StartYear       = read_in.get('StartYear')
    StartMonth      = read_in.get('StartMonth')
    StartDay        = read_in.get('StartDay')
    FixYear         = read_in.get('FixYear')
    GridInfo        = read_in.get('GridInfo')
    LDebug          = read_in.get('LDebug')

    # Create the log_file name and sanity-check 

    # Check the type of input in yaml file 
    if "FileInputVars" in read_in : 
        FileInputVars_readin = {}
        FileInputVars_readin   = read_in.get('FileInputVars')
        # Name of the Variable for this experiment 
        FileInputVarNames  = list(FileInputVars_readin.keys())
        # Variable configuration for this experiment 
        FileInputVarFields_conf = list(FileInputVars_readin.values())

    if "OasisOutputVars" in read_in :
        OasisOutputVars_readin = {} 
        OasisOutputVars_readin = read_in.get('OasisOutputVars')
        print("NOTE: OasisOutputVars_readin : {}".format(OasisOutputVars_readin), file=sys.stderr)
        # Name of the Variable for this experiment 
        OasisOutputVarsNames  = list(OasisOutputVars_readin.keys())

        # Variable configuration for this experiment 
        OasisOutputVars_conf = list(OasisOutputVars_readin.values())

    return ModelNameSend, ModelNameReceive, LogFileName, RunLengthSec, TimeStepSec, StartYear, StartMonth,\
             StartDay, FixYear, GridInfo, LDebug, FileInputVarFields_conf, OasisOutputVars_conf, OasisOutputVars_readin
             
