# Copyright 2022 - Barcelona Supercomputing Center
# Authors: Etienne Tourigny, Rodrigo Martín Posada
# MIT License

import pyoasis
from pyoasis import OASIS  # pylint: disable=no-name-in-module
from netCDF4 import Dataset  # pylint: disable=no-name-in-module
import numpy as np
import os
import logging
import math
import sys
import argparse 

from DataCouplerUtils import read_namelist, days_since_refyear, is_leap, earth_area
from DataCouplerUtils import read_DataCoupler_yaml_input

from DataCouplerFileVar import FileInputVar
from DataCouplerOasisVar import OasisOutputVar

class DataCoupler:

    def __init__(self,yaml_conf_file):        
        
        ModelNameSend, ModelNameReceive, LogFileName, RunLengthSec, TimeStepSec, StartYear, StartMonth,   \
            StartDay, FixYear, GridInfo, LDebug, fileIn_vars_config, oasisOut_vars_config, OasisOutputVars_readin ,   \
            = read_DataCoupler_yaml_input(yaml_conf_file)

        if LDebug:
            logging.basicConfig(filename=LogFileName, format='%(message)s', level=logging.DEBUG)

        date_yref_min = max([conf['yref_min'] for conf in fileIn_vars_config])
        date_yref_max = min([conf['yref_max'] for conf in fileIn_vars_config])

        # valid FixYear to get a seasonal cycle, or set one if none and running outside time window
        if (FixYear > 0):
            FixYear = min(date_yref_max, max(date_yref_min, FixYear))
        else:
            if ((StartYear < date_yref_min) or (StartYear > date_yref_max)):
                # set pseudo fixyear,
                FixYear = min(date_yref_max, max(date_yref_min, StartYear))
                # but do not overshoot
                if ((StartYear < date_yref_min) and (RunLengthSec > (date_yref_min-StartYear+1)*3600*24*365)):
                    raise Exception("RunLengthSec is too long for StartYear less than date_yref_min")

        # Oasis Compiler description 
        self.OasisCPLNG = ModelNameSend

        #  time axis unit = days
        self.delta_t = TimeStepSec/86400.
        self.fix_year = FixYear
        self.start_year = StartYear
        self.start_month = StartMonth
        self.start_day = StartDay
        self.run_length_sec = RunLengthSec
        self.time_step_sec = TimeStepSec
        # Flags for discordant leap years
        self.lleapfx1 = False
        self.lleapfx2 = False
        self._fileIn_vars_config = fileIn_vars_config
        self._oasisOut_vars_config = oasisOut_vars_config
        self.oasisOut_vars = {}

        # Setup input vars
        self.fileIn_vars = {}
        for conf in fileIn_vars_config:
            self.fileIn_vars[conf['id']] = FileInputVar(conf, self.fix_year, self.start_year)
        logging.debug('fileIn_vars: {}'.format(self.fileIn_vars.keys()))
        



    def setup_oasis(self, write_grid=True, write_grid_mask=True, write_grid_corners= True, write_grid_areas = True):
        # OASIS init
        self.oasis_comp = pyoasis.Component(self.OasisCPLNG)

        self.oasisOut_vars = {}
        # Setup grids and write if requested
        grid_names = np.unique([v.grid_name for v in self.fileIn_vars.values()])
        logging.debug('grid_names: {}'.format(str(grid_names)))
        for grid_name in grid_names:
            # Get FileInputVars names within this grid_name
            grid_vars = list(filter(lambda x: x.grid_name == grid_name, self.fileIn_vars.values()))
            v = grid_vars[0]
            grid_nx = v.grid_nx
            grid_ny = v.grid_ny
            # Create partition
            partition = pyoasis.SerialPartition(grid_nx*grid_ny)
            # Define the grid and write to files
            if write_grid:
                grid = pyoasis.Grid(grid_name, grid_nx, grid_ny, v.grid_lon, v.grid_lat, partition)
                # add grid mask
                if write_grid_mask:
                    grid.set_mask(np.zeros(v.grid_lon.shape))
                # add grid corners - only if grid has many points
                if write_grid_corners and grid_nx>1 and grid_ny>1:
                    ncrn = 4
                    clo = pyoasis.asarray(np.zeros((grid_nx, grid_ny, ncrn), dtype=np.float64))
                    clo[:, :, 0] = v.grid_lon[:, :] - v.dx/2.0
                    clo[:, :, 1] = v.grid_lon[:, :] + v.dx/2.0
                    clo[:, :, 2] = clo[:, :, 1]
                    clo[:, :, 3] = clo[:, :, 0]
                    cla = pyoasis.asarray(np.zeros((grid_nx, grid_ny, ncrn), dtype=np.float64))
                    cla[:, :, 0] = v.grid_lat[:, :] - v.dy/2.0
                    cla[:, :, 1] = cla[:, :, 0]
                    cla[:, :, 2] = v.grid_lat[:, :] + v.dy/2.0
                    cla[:, :, 3] = cla[:, :, 2]
                    grid.set_corners(clo, cla)
                # add grid areas - only if grid has many points
                if write_grid_areas and grid_nx>1 and grid_ny>1:
                    dp_conv = math.pi/180.
                    area = np.zeros((grid_nx, grid_ny), dtype=np.float64)
                    # multiply by total earth area in m2 to have areas.nc in m2 and not in rad
                    area[:, :] = earth_area * dp_conv * \
                                 np.abs(np.sin(cla[:, :, 2] * dp_conv) -
                                        np.sin(cla[:, :, 0] * dp_conv)) * \
                                 np.abs(clo[:, :, 1] - clo[:, :, 0])
                    grid.set_area(area)
                # write grid to files
                grid.write()

            # Setup oasis vars using this grid
            for v in grid_vars:
                # look for this oasis var in OasisVars config
                configs = self._oasisOut_vars_config
                config = list(filter(lambda configs: configs['send_id'] == v.oasis_name, configs))
                if len(config)!=1:
                    raise Exception("oasis var {} not defined in OasisOutputVars".format(v.oasis_name))
                if not v.oasis_name in self.oasisOut_vars:
                    logging.debug('creating oasis var {} with grid {} and fileIn_var {}'.format(v.oasis_name, grid_name,v.id))
                    self.oasisOut_vars[v.oasis_name] = OasisOutputVar(config[0], partition, v.grid_lon.shape, {v.id:v})

                else:
                    logging.debug('adding to oasis var {} with grid {} fileIn_var {}'.format(v.oasis_name, grid_name,v.id))
                    self.oasisOut_vars[v.oasis_name].add_file_input_var(v.id,v)

            logging.debug('oasisOut_vars: {}'.format(str(self.oasisOut_vars.keys())))
            for var_name, oasisOut_var in self.oasisOut_vars.items():
                logging.debug('oasisOut_var {} is linked to fileIn vars {}'.format(var_name, oasisOut_var.file_input_vars() ))
                print('oasisOut_var {} is linked to fileIn vars {}'.format(var_name, oasisOut_var.file_input_vars() ),file=sys.stderr)

        # check that all OasisVars are actually defined
        for config in self._oasisOut_vars_config:
            if not config['send_id'] in self.oasisOut_vars:
                raise Exception("oasis var {} in OasisOutputVars not defined in FileInputVars".format(config['send_id']))

        self.oasis_comp.enddef()


    def setup_data_coupler(self, yy=None, mm=None, dd=None):
        '''
        Init method
        '''
        # default arguments
        if yy is None: yy=self.start_year
        if mm is None: mm=self.start_month
        if dd is None: dd=self.start_day
        # Simulated year
        self.yr_sim = yy
        # data year
        yyref = self.fix_year if self.fix_year > 0 else self.yr_sim

        # Timesteps [days since refdate]
        # Running step
        self.t_global = float(days_since_refyear(yyref, mm, dd, yyref))
        self.start_t = self.t_global

        # Step when next simulated year starts
        self.tnewyear = 365.
        if (is_leap(self.yr_sim)):
            self.tnewyear += 1

        logging.debug('DataCoupler INIT - simulated/ref/fixed years: {} {} {}'.format(self.yr_sim, yyref, self.fix_year))
        logging.debug('DataCoupler INIT - t_global  : {}'.format(self.t_global))

        # Initialise with dummy values
        # Step at last day of February
        self.tlastfeb = -9999.

        # Setup for FixYear
        if self.fix_year > 0:
            self.__discordant_leap_years()

        # Setup variables
        [v.setup(yyref, self.t_global, self.tnewyear) for v in self.fileIn_vars.values()]


    def run_time_loop(self):
        # run time loop
        # Note that according to OASIS convention, the first time step is at time=0
        # and the last one is at runlength-timestep (see OASIS documentation)
        for time in range(0, self.run_length_sec, self.time_step_sec):
            # update each fileIn
            self.update_file_input_vars()

            # OASIS put for each oasisOut_var
            for oasisOut_varname, oasisOut_var in self.oasisOut_vars.items():
                logging.debug('DataCoupler processing oasis var {} at time {}'.format(oasisOut_varname,time))
                oasisOut_var.update(time)
                oasisOut_var.put(time)

    def update_file_input_vars(self):
        '''
        Scan the data until next time step is covered.
        '''
        if self.fix_year > 0:
            if self.t_global == self.tlastfeb:
                # Fix for discordant leap years
                if self.lleapfx1:
                    logging.debug('Skipping 29th of February')
                    self.t_global += self.delta_t
                    self.lleapfx1 = False
                elif self.lleapfx2:
                    logging.debug('Repeating 28th of February')
                    self.t_global -= self.delta_t
                    self.lleapfx2 = False
            elif self.t_global == self.tnewyear:
                # If FixYear and a new year, reset everything
                logging.debug('NEWYEAR: t, tnewyear {:<018} {:<018}'.format(self.t_global, self.tnewyear))
                # Reset everything for fixed year
                self.yr_sim = self.yr_sim + 1
                self.finalise_data_coupler()
                self.setup_data_coupler(self.yr_sim, 1, 1)
                self.update_file_input_vars()
                return

        #[v.update(self.t_global) for v in self.fileIn_vars.values()]
        for v in self.fileIn_vars.values():
            if self.t_global == self.start_t or v.update_timestep:
                v.update(self.t_global)

        self.t_global += self.delta_t
        logging.debug('updating t_global= {}'.format(self.t_global))

    def finalise_data_coupler(self):
        '''
        Done method
        '''
        [v.finalise() for v in self.fileIn_vars.values()]

    def __del__(self):
        del self.oasis_comp

    def __discordant_leap_years(self):
        '''
        Check for discordant leap years. This function activates the lleapfx flags to avoid time mismatch of SST/SIC.

        In fix year condition the Gregorian calendar of EC-Earth may conflict
        with the calendar of the selected fixyear.

        See issue #820.
        '''
        # Define the last day of February
        # Redundant here but clearer inside the subroutine
        # Step at last day of February
        if is_leap(self.fix_year):
            self.tlastfeb = float(days_since_refyear(self.fix_year, 2, 29, self.fix_year))
        else:
            self.tlastfeb = float(days_since_refyear(self.fix_year, 2, 28, self.fix_year))

        # If the FixYear is a leap year and the model year is not
        if ((not is_leap(self.yr_sim)) and is_leap(self.fix_year)):
            self.lleapfx1 = True
            self.tnewyear = self.tnewyear + 1

        # If the FixYear is not a leap year and the model year is a leap year
        if (is_leap(self.yr_sim) and (not is_leap(self.fix_year))):
            self.lleapfx2 = True
            self.tnewyear = self.tnewyear - 1

        logging.debug('  Discordant leap years, {} {}'.format(self.yr_sim, self.fix_year))
        logging.debug('  lleapfx1, lleapfx2 {} {}'.format('T' if self.lleapfx1 else 'F',
                                                          'T' if self.lleapfx2 else 'F'))
        logging.debug('UPDATED DataCoupler INIT - tnewyear: {}'.format(self.tnewyear))
        logging.debug('DataCoupler INIT - tlastfeb: {}'.format(self.tlastfeb))



if __name__ == '__main__':
    # receive the yaml experiment input file is it exists , otherwise pass 
    try:
        parser=argparse.ArgumentParser()
        parser.add_argument("--conf",type=str)
        args = parser.parse_args()
        yaml_conf_file = args.conf 
    except:
        print("NOTE: Could note read in the input yaml file : {}".format(yaml_conf_file), file=sys.stderr)
    else:
        print("NOTE: input yaml file is : {}".format(yaml_conf_file), file=sys.stderr)

    # Setup
    datacoupler = DataCoupler(yaml_conf_file)
    datacoupler.setup_data_coupler()
    datacoupler.setup_oasis(write_grid=True, write_grid_mask=True,
                     write_grid_corners= True, write_grid_areas = True)
    # Time loop
    datacoupler.run_time_loop()
    # Finalize
    datacoupler.finalise_data_coupler()