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


class AMIPVar:
    def __init__(self, conf, grid_lat, grid_lon, mask, area, fix_year):
        self.id = conf['id']
        self.oasis_name = conf['oasis_name']
        self.grid_name = conf['grid_name']
        self.file_pattern = conf['file_pattern']
        self.var_name = conf['netcdf_variable']
        self.timestep = conf['timestep']
        self.scale_factor = conf['scale_factor']
        self.offset = conf['offset']
        self.min = conf['min'] if conf['min'] is not None else np.NINF
        self.max = conf['max'] if conf['max'] is not None else np.Inf
        self.interpolate = conf['interpolate']
        self.fix_year = fix_year
        # Initialization
        self.grid_lat = grid_lat
        self.grid_lon = grid_lon
        self.mask = mask
        self.area = area
        grid_shape = self.grid_lat.shape
        self.field = np.empty((2,) + grid_shape)
        self.send_field = np.empty(grid_shape)
        self.current_file = ''
        self.dataset = None
        self.time_array = None
        self.t_index = 0
        # Dummy values
        self.t1 = -9999
        self.t2 = -9999
        self.t0101 = -9999
        self.t1231 = -9999
        self.tnewyear = -9999

    def setup(self, t, t0101, t1231, tnewyear):
        self.t1 = -9999
        self.t2 = -9999
        self.t_index = 0
        self.t0101 = t0101
        self.t1231 = t1231
        self.tnewyear = tnewyear
        self.__open_file(self.file_pattern)
        if t >= self.time_array[0]:
            self.t_index = np.searchsorted(self.time_array, t, side='right')-1
        self.t2 = self.time_array[self.t_index]
        self.t1 = self.t2
        self.__read_file()

    def update(self, t):
        while t > self.t2:
            self.__read_next_timestep_from_files(t)

        var_max = self.max
        var_min = self.min
        # Monthly interpolation
        if self.timestep == 'monthly' and self.interpolate:
            if self.t1 < self.t2:
                logging.debug('{}: time interpolation t1,t,t2 {} {} {}'.format(self.id, self.t1, t, self.t2))
                raw_field = ((t-self.t1)*self.field[1, :, :] + (self.t2-t)*self.field[0, :, :])/(self.t2-self.t1)
            else:
                logging.debug('{}: no time interpolation t,t2 {} {}'.format(self.id, t, self.t2))
                # Hardcoded to match Fortran's AMIP clipping avoidance
                if self.oasis_name == 'AMIP_sst':
                    var_min = np.NINF
                raw_field = self.field[1, :, :]
        else:
            # Nearest neighbour
            if self.t1 < self.t2:
                if np.floor(t)-np.floor(self.t1) < np.floor(self.t2)-np.floor(t):
                    nn = 0
                else:
                    nn = 1
            else:
                nn = 1
            logging.debug('{}: nearest neigbour nn,t1,t,t2 {} {} {} {}'.format(self.id, nn, self.t1, t, self.t2))
            # Hardcoded to match Fortran's AMIP clipping avoidance
            if self.oasis_name == 'AMIP_sst':
                var_min = np.NINF
            raw_field = self.field[nn, :, :]
        # Clipping required for CMIP6 BCs
        self.send_field = np.maximum(var_min, np.minimum(var_max, raw_field))

    def finalise(self):
        self.__close_current_file()

    def __read_next_timestep_from_files(self, t):
        self.t1 = self.t2
        self.field[0, :, :] = self.field[1, :, :]

        if t > np.max(self.time_array):
            self.__close_current_file()
            # TODO
            self.__open_file(self.file_pattern)
            self.old_t_index = -1

        self.t_index += 1

        self.t2 = self.time_array[self.t_index]
        self.__read_file()

    def __open_file(self, file_name):
        self.current_file = file_name
        logging.debug('{}: open file  {}'.format(self.id, self.current_file))
        self.dataset = Dataset(file_name, 'r')
        self.time_array = self.dataset['time'][:]

    def __read_file(self):
        if self.fix_year > 0 and self.time_array[self.t_index] >= self.tnewyear:
            logging.debug('{}:...use first BC data for timestep {}'.format(self.id, self.time_array[self.t_index]))
            # Assign BC only if values have been assigned previously
            first_t_index = np.searchsorted(self.time_array, self.t0101, side='right')-1
            raw_field = self.dataset[self.var_name][first_t_index+1, :, :]
        elif self.fix_year > 0 and self.time_array[self.t_index] < self.t0101:
            logging.debug('{}:...use last BC data for timestep {}'.format(self.id, self.time_array[self.t_index]))
            # Assign BC only if values have been assigned previously
            last_t_index = np.searchsorted(self.time_array, self.t1231, side='right')-1
            raw_field = self.dataset[self.var_name][last_t_index, :, :]
        else:
            logging.debug('{}: ...read timestep {} {:<018}'.format(
                self.id, self.t_index, self.time_array[self.t_index]))
            raw_field = self.dataset[self.var_name][self.t_index, :, :]

        self.field[1] = raw_field.astype('float64')*self.scale_factor+self.offset

    def __close_current_file(self):
        logging.debug('{}: close file {}'.format(self.id, self.current_file))
        self.dataset.close()
