# Copyright 2021 - Barcelona Supercomputing Center
# Author: Rodrigo Martín Posada
# MIT License
from datetime import date
import re
import f90nml
import numpy as np

NAMELIST_FILE_NAME = 'namelist.amip'


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():
    '''
    Reads the AMIP namelist.
    '''
    nml = f90nml.read(NAMELIST_FILE_NAME)
    nml_dict = nml['NAMAMIP']
    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 'Vars' in nml_dict:
        vars_data = np.transpose(nml_dict['Vars'])
        vars_config = [None] * vars_data.shape[0]
        for i in range(len(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],
            }
            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
        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
        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('AMIP: Invalid namelist.amip configuration')

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

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