# Copyright 2021 - Barcelona Supercomputing Center
# Author: Rodrigo Martín Posada
# MIT License
from netCDF4 import Dataset  # pylint: disable=no-name-in-module
import numpy as np
import os
import logging

from amip_utils import days_since_refyear, is_leap
from amip_var import AMIPVar

OASIS_DATA_FILE = 'amip-oasis-data.nc'


class AMIPForcing:
    def __init__(self, deltaT, FixYear, vars_config):
        self.fix_year = FixYear
        self.delta_t = deltaT
        # Flags for discordant leap years
        self.lleapfx1 = False
        self.lleapfx2 = False
        # Read grid from one of the input netCDF file
        self.vars = []
        ds = Dataset(OASIS_DATA_FILE, 'r')
        for conf in vars_config:
            grid_name = conf['grid_name']
            grid_lat = ds[grid_name+'.lat'][:]
            grid_lon = ds[grid_name+'.lon'][:]
            grid_mask = ds[grid_name+'.msk'][:]
            grid_area = ds[grid_name+'.srf'][:]
            self.vars.append(AMIPVar(conf, grid_lat, grid_lon, grid_mask, grid_area, self.fix_year))
        ds.close()

    def setup_amip_forcing(self, yy, mm, dd):
        '''
        Init method
        '''
        # Simulated year
        self.yr_sim = yy
        # AMIP reader 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))

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

        logging.debug('AMIP INIT - simulated/amip/fixed years: {} {} {}'.format(self.yr_sim, yyref, self.fix_year))
        logging.debug('AMIP INIT - tglobal  : {}'.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.vars]

    def update_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_amip_forcing()
                self.setup_amip_forcing(self.yr_sim, 1, 1)
                self.update_vars()
                return

        [v.update(self.t_global) for v in self.vars]

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

    def finalise_amip_forcing(self):
        '''
        Done method
        '''
        [v.finalise() for v in self.vars]

    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 AMIP INIT - tnewyear: {}'.format(self.tnewyear))
        logging.debug('AMIP INIT - tlastfeb: {}'.format(self.tlastfeb))
