from netCDF4 import Dataset  # pylint: disable=no-name-in-module
import numpy as np
import os

from namelist_reader import read_namelist
from date_utils import days_since_refdate, is_leap

# Constants
LOG_FILE_NAME = 'amip.log'
# Reference years of the dataset
DATA_YREF_MIN = 1870
DATA_YREF_MAX = 2016


class AMIPForcing:
    def __init__(self):
        # Open log_file
        self.log_file = open(LOG_FILE_NAME, 'w', buffering=1)

        # Read RunLengthSec,TimeStepSec from namelist
        self.RunLengthSec, self.TimeStepSec, self.StartYear, self.StartMonth, \
            self.StartDay, self.FixYear, self.FileListSST, self.FileListSIC, self.LDebug, \
            self.LInterpolate = read_namelist()

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

        self.lstored = False
        # flags for discordant leap years
        self.lleapfx1 = False
        self.lleapfx2 = False
        # Get shapes and constants from netCDF files
        self.setup_from_netcdf()

    def setup_amip_forcing(self, yy, mm, dd):
        #
        # Init method
        #
        #  time axis unit = days
        self.deltaT = self.TimeStepSec/86400.
        # simulated year
        self.yr_sim = yy
        # amip_reader year
        yyref = self.FixYear if self.FixYear > 0 else yy

        # timesteps [days since refdate]
        # running step, step before and after with data
        self.t = float(days_since_refdate(yyref, mm, dd, DATA_YREF_MIN))
        self.t0101 = float(days_since_refdate(yyref, 1, 1, DATA_YREF_MIN))
        self.t1231 = float(days_since_refdate(yyref, 12, 31, DATA_YREF_MIN))

        # find out when we switch year in amip_year according to simulated year
        self.tnewyear = self.t0101 + 365  # step when next simulated year starts
        if (is_leap(yy)):
            self.tnewyear = self.t0101 + 366

        if (self.LDebug):
            print('AMIP INIT - simulated/amip/fixed years:', yy, yyref, self.FixYear, file=self.log_file)
            print('AMIP INIT - tstart  :', self.t, file=self.log_file)
            print('AMIP INIT - t0101   :', self.t0101, file=self.log_file)
            print('AMIP INIT - t1231   :', self.t1231, file=self.log_file)
            print('AMIP INIT - tnewyear:', self.tnewyear, file=self.log_file)

        # initialise with dummy values
        self.t1 = -9999.
        self.t2 = -9999.
        self.tlastfeb = 0

        # point to first file
        self.ifile = 0
        self.lfirst = True

        # setup BC for FixYear
        if self.FixYear > 0:
            if not self.lstored:
                self.setup_fixyear_bc()
                self.lstored = True
            self.discordant_leap_years(self.yr_sim, self.FixYear)

    def discordant_leap_years(self, simyr, fixyr):
        # Check for discordant leap years.
        # In fix year condition the Gregorian calendar of EC-Earth may conflict
        # with the calendar of the selected fixyear.
        # This function activates the lleapfx flags to avoid time mismatch of SST/SIC
        # 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(fixyr):
            self.tlastfeb = float(days_since_refdate(fixyr, 2, 29, DATA_YREF_MIN))
        else:
            self.tlastfeb = float(days_since_refdate(fixyr, 2, 28, DATA_YREF_MIN))

        # if the FixYear is a leap year and the model year is not
        if ((not is_leap(simyr)) and is_leap(fixyr)):
            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(simyr) and (not is_leap(fixyr))):
            self.lleapfx2 = True
            self.tnewyear = self.tnewyear - 1

        if (self.LDebug):
            print('  Discordant leap years,', self.yr_sim, self.FixYear, file=self.log_file)
            print('  lleapfx1, lleapfx2', 'T' if self.lleapfx1 else 'F',
                  'T' if self.lleapfx2 else 'F', file=self.log_file)
            print('UPDATED AMIP INIT - tnewyear:', self.tnewyear, file=self.log_file)
            print('AMIP INIT - tlastfeb:', self.tlastfeb, file=self.log_file)

    def setup_fixyear_bc(self):
        #
        # Scan the datasets in the file(s) and store the first and last
        # available data for current AMIP year. Used for FixYear case.
        #
        tinit = self.t
        FixYearInit = self.FixYear
        self.FixYear = -999

        istart = self.t0101
        iend = self.t1231

        # Scan must start on January 1st
        self.t = self.t0101

        for _ in range(int(istart), int(iend)+1):
            self.update_sst_and_sic()

            if self.t == self.t0101:
                if self.LDebug:
                    print(' stored first BC using data from t2=', self.t2, file=self.log_file)
                self.sst_first[:] = self.sst[1, :, :]
                self.sic_first[:] = self.sic[1, :, :]
            if self.t == self.t1231:
                if self.LDebug:
                    print(' stored last BC using data from t1=', self.t1, file=self.log_file)
                self.sst_last[:] = self.sst[0, :, :]
                self.sic_last[:] = self.sic[0, :, :]

        self.close_file(self.ifile)
        self.FixYear = FixYearInit
        self.t = tinit
        self.t1 = -9999.
        self.t2 = -9999.

        self.ifile = 0
        self.lfirst = True

        if self.LDebug:
            print("  Done storing BC", file=self.log_file)

    def update_sst_and_sic(self):
        #
        # Scan the data until next time step is covered.
        #
        if (self.lfirst):
            self.open_file(self.ifile)
            if self.t < self.t_sst[0]:
                self.t_index = 0
            else:
                for self.t_index in range(len(self.t_sst)-1):
                    if self.t_sst[self.t_index] <= self.t and self.t < self.t_sst[self.t_index+1]:
                        break

            if self.LDebug:
                print('updating t=', self.t, file=self.log_file)

            self.t2 = self.t_sst[self.t_index]
            self.t1 = self.t2
            self.old_t_index = self.t_index
            self.read_file()
            self.lfirst = False
        else:
            self.t = self.t+self.deltaT
            if self.LDebug:
                print('updating t=', self.t, file=self.log_file)

            # fix for discordant leap years
            if self.t == self.tlastfeb and self.lleapfx1:
                print('Skipping 29th of February', file=self.log_file)
                self.t = self.t+self.deltaT
                self.lleapfx1 = False
            elif self.t == self.tlastfeb and self.lleapfx2:
                print('Repeating 28th of February', file=self.log_file)
                self.t = self.t-self.deltaT
                self.lleapfx2 = False

            # if FixYear and a new year, reset everything
            if self.FixYear > 0 and self.t == self.tnewyear:
                if self.LDebug:
                    print('NEWYEAR: t, tnewyear', '{:<018}'.format(self.t),
                          '{:<018}'.format(self.tnewyear), file=self.log_file)

                self.yr_sim = self.yr_sim + 1
                self.close_file(self.ifile)
                self.setup_amip_forcing(self.yr_sim, 1, 1)
                self.update_sst_and_sic()
                return

        while self.t > self.t2 and self.file_found:
            self.read_next_timestep_from_file()

        if self.LInterpolate:
            # linear time interpolation
            if self.t1 < self.t2:
                if self.LDebug:
                    print('time interpolation t1,t,t2', self.t1, self.t, self.t2, file=self.log_file)

                # Clipping required for CMIP6 BCs
                self.AMIP_sst = np.maximum(
                    ((self.t-self.t1)*self.sst[1, :, :]+(self.t2-self.t)*self.sst[0, :, :])/(self.t2-self.t1), 271.38)
                self.AMIP_sic = np.maximum(0, np.minimum(
                    1, ((self.t-self.t1)*self.sic[1, :, :]+(self.t2-self.t)*self.sic[0, :, :])/(self.t2-self.t1)))
            else:
                if (self.LDebug):
                    print('no time interpolation t,t2', self.t, self.t2, file=self.log_file)

                self.AMIP_sst = self.sst[1, :, :]
                self.AMIP_sic = np.maximum(0, np.minimum(1, self.sic[1, :, :]))
        else:
            # nearest neighbour
            if self.t1 < self.t2:
                if np.floor(self.t)-np.floor(self.t1) < np.floor(self.t2)-np.floor(self.t):
                    self.nn = 0
                else:
                    self.nn = 1
            else:
                self.nn = 1

            if self.LDebug:
                print('nearest neigbour nn,t1,t,t2', self.nn+1, self.t1, self.t, self.t2, file=self.log_file)

            self.AMIP_sst = self.sst[self.nn, :, :]
            self.AMIP_sic = np.maximum(0, np.minimum(1, self.sic[self.nn, :, :]))

    def setup_from_netcdf(self):
        fname = self.FileListSST[0]
        ncid_sst = Dataset(fname, 'r')
        # x and y (i and j) dimensions of the AMIP grid
        _, grid_ny, grid_nx = ncid_sst['tosbcs'].shape

        self.sst = np.zeros((2, grid_ny, grid_nx), dtype='float64')
        self.sic = np.zeros((2, grid_ny, grid_nx), dtype='float64')
        self.sst_first = np.zeros((grid_ny, grid_nx), dtype='float64')
        self.sic_first = np.zeros((grid_ny, grid_nx), dtype='float64')
        self.sst_last = np.zeros((grid_ny, grid_nx), dtype='float64')
        self.sic_last = np.zeros((grid_ny, grid_nx), dtype='float64')
        self.AMIP_sst = np.zeros((grid_ny, grid_nx), dtype='float64')
        self.AMIP_sic = np.zeros((grid_ny, grid_nx), dtype='float64')
        ncid_sst.close()

    def open_file(self, i):
        #
        # Open files, read time steps
        #
        self.fname_sst = self.FileListSST[i]
        self.file_found = os.access(self.fname_sst, os.R_OK)
        if not self.file_found:
            return
        print('open SST file ', self.fname_sst, file=self.log_file)
        self.ncid_sst = Dataset(self.fname_sst, 'r')

        self.fname_sic = self.FileListSIC[i]
        print('open SIC file ', self.fname_sic, file=self.log_file)
        self.ncid_sic = Dataset(self.fname_sic, 'r')

        self.t_sst = self.ncid_sst['time'][:]

    def close_file(self, i):
        if self.file_found:
            print('close SST file', self.fname_sst, file=self.log_file)
            self.ncid_sst.close()
            print('close SIC file', self.fname_sic, file=self.log_file)
            self.ncid_sic.close()

    def read_file(self):
        if self.FixYear > 0 and self.t_sst[self.t_index] >= self.tnewyear:
            if self.LDebug:
                print('...use first BC data for timestep', self.t_sst[self.t_index], file=self.log_file)

            # Assign BC only if values have been assigned previously
            if np.sum(self.sst_first) > 1e-20:
                self.sst[1, :, :] = self.sst_first
                self.sic[1, :, :] = self.sic_first
            else:
                raise Exception('First BC are empty, what are you doing? Aborting!')

        elif self.FixYear > 0 and self.t_sst[self.t_index] < self.t0101:
            if self.LDebug:
                print('...use last BC data for timestep', self.t_sst[self.t_index], file=self.log_file)

            # Assign BC only if values have been assigned previously
            if np.sum(self.sst_last) > 1e-20:
                self.sst[1, :, :] = self.sst_last
                self.sic[1, :, :] = self.sic_last
            else:
                raise Exception('Last BC are empty, what are you doing? Aborting!')
        else:
            print('...read timestep', self.t_index+1, '{:<018}'.format(self.t_sst[self.t_index]), file=self.log_file)
            self.sst[1, :, :] = self.ncid_sst['tosbcs'][self.t_index, :, :]
            self.sic[1, :, :] = self.ncid_sic['siconcbcs'][self.t_index, :, :]

            # units change with AMIP forcing dataset v1.1.2
            self.sst[1, :, :] = self.sst[1, :, :]+273.15
            self.sic[1, :, :] = self.sic[1, :, :]*0.01

    def read_next_timestep_from_file(self):
        self.t1 = self.t2
        self.sst[0, :, :] = self.sst[1, :, :]
        self.sic[0, :, :] = self.sic[1, :, :]

        if self.t > np.max(self.t_sst):
            self.close_file(self.ifile)
            self.ifile = self.ifile+1
            self.open_file(self.ifile)

            if not self.file_found:
                self.ifile = self.ifile-1
                return

            self.old_t_index = -1

        self.t_index = self.old_t_index+1

        self.t2 = self.t_sst[self.t_index]
        self.read_file()

        self.old_t_index = self.t_index

    def finalise_amip_forcing(self):
        #
        # Done method
        #
        self.close_file(self.ifile)
        self.log_file.close()
