import os
from collections import OrderedDict
from datetime import datetime, timedelta
from typing import Iterable, List, Union

try:
    from netCDF4 import Dataset, Variable, chartostring
except ImportError:
    raise ImportError("NetCDF4 is required when using NetCDF/NetCDFMixin")

try:
    from cftime import num2pydate as num2date
except ImportError:
    from cftime import num2date

import numpy as np


class Stations:

    def __init__(self, dataset: Dataset, station_variable: Variable):
        self.__station_variable = station_variable

        station_dimension = station_variable.dimensions[0]

        # todo make this a bit smarter, right now variables like station_name would be forgotten
        self.__attribute_variables = {}
        for variable_name in dataset.variables:
            variable = dataset.variables[variable_name]
            if variable != station_variable and variable.dimensions == (station_dimension,):
                self.__attribute_variables[variable_name] = variable

        self.__attributes = OrderedDict()
        for i in range(station_variable.shape[0]):
            id = str(chartostring(station_variable[i]))

            values = {}
            for variable_name in self.__attribute_variables.keys():
                values[variable_name] = dataset.variables[variable_name][i]

            self.__attributes[id] = values

    @property
    def station_ids(self) -> Iterable:
        """
        :return: An ordered iterable of the station ids (location ids) for which station data is available.
        """
        return self.__attributes.keys()

    @property
    def attributes(self) -> OrderedDict:
        """
        :return: An OrderedDict containing dicts containing the values for all station attributes of the input dataset.
        """
        return self.__attributes

    @property
    def attribute_variables(self) -> dict:
        """
        :return: A dict containing the station attribute variables of the input dataset.
        """
        return self.__attribute_variables


class ImportDataset:
    """
    A class used to open and import the data from a NetCDF file.
    Uses the NetCDF4 library. Contains various methods for reading the data in the file.
    """

    def __init__(self, folder: str, basename: str):
        """
        :param folder:    Folder the file is located in.
        :param basename:  Basename of the file, extension ".nc" will be appended to this
        """

        self.__ensemble_size = 1

        # Load the content of a NetCDF file into a Dataset.
        self.__filename = os.path.join(
            folder,
            basename + ".nc"
        )
        self.__dataset = Dataset(self.__filename)

        # Find the number of ensemble members and the time and station id variables
        self.__time_variable = self.__find_time_variable()
        if self.__time_variable is None:
            raise Exception('No time variable found in file ' + self.__filename + '. '
                            'Please ensure the file contains a time variable with standard_name "time" and axis "T".')

        self.__ensemble_member_variable = self.__find_ensemble_member_variable()
        if self.__ensemble_member_variable:
            self.__ensemble_size = self.__dataset.dimensions['realization'].size

        self.__station_variable = self.__find_station_variable()
        if self.__station_variable is None:
            raise Exception('No station variable found in file ' + self.__filename + '. '
                            'Please ensure the file contains a variable with cf_role "timeseries_id".')

    def __str__(self):
        return self.__filename

    def __find_time_variable(self) -> Union[Variable, None]:
        """
        Find the variable containing the times in the given Dataset.

        :param dataset: The Dataset to be searched.
        :return: a netCDF4.Variable object of the time variable (or None if none found)
        """
        for variable in self.__dataset.variables.values():
            if ('standard_name' in variable.ncattrs() and 'axis' in variable.ncattrs()
                    and variable.standard_name == 'time' and variable.axis == 'T'):
                return variable

        return None

    def __find_ensemble_member_variable(self) -> Union[Variable, None]:
        """
                Find the variable containing the ensemble member index in the given Dataset.

                :param dataset: The Dataset to be searched.
                :return: a netCDF4.Variable object of the ensemble member index variable (or None if none found)
        """
        for variable in self.__dataset.variables.values():
            if 'standard_name' in variable.ncattrs() and variable.standard_name == "realization":
                return variable

        return None

    def __find_station_variable(self) -> Union[Variable, None]:
        """
        Find the variable containing station id's  (location id's) in the given Dataset.

        :param dataset: The Dataset to be searched.
        :return: a netCDF4.Variable object of the station id variable (or None if none found)
        """
        for variable in self.__dataset.variables.values():
            if 'cf_role' in variable.ncattrs() and variable.cf_role == 'timeseries_id':
                return variable

        return None

    def read_import_times(self) -> np.ndarray:
        """
        Reads the import times in the time variable of the dataset.

        :param time_variable: The time variable containing input times
        :return: an array containing the input times as datetime objects
        """
        time_values = self.__time_variable[:]
        time_unit = self.__time_variable.units
        try:
            time_calendar = self.__time_variable.calendar
        except AttributeError:
            time_calendar = u'gregorian'

        return num2date(time_values, units=time_unit, calendar=time_calendar)

    def read_station_data(self) -> Stations:
        return Stations(self.__dataset, self.__station_variable)

    def find_timeseries_variables(self) -> List[str]:
        """
        Find the keys of all 2-D or 3-D variables with dimensions {station, time} or {station, time, realization}
        where station is the dimension of the station_variable, time the dimension of the time_variable and
        realization the dimension for ensemble_member_index.

        :param dataset:           The Dataset to be searched.
        :param station_variable:  The station id variable.
        :param time_variable:     The time variable.
        :return: a list of strings containing all keys found.
        """
        station_dim = self.__station_variable.dimensions[0]
        time_dim = self.__time_variable.dimensions[0]
        if self.__ensemble_member_variable is not None:
            ensemble_dim = self.__ensemble_member_variable.dimensions[0]
            expected_dims = [
                (time_dim, station_dim, ensemble_dim), (time_dim, ensemble_dim, station_dim),
                (station_dim, time_dim, ensemble_dim), (station_dim, ensemble_dim, time_dim),
                (ensemble_dim, time_dim, station_dim), (ensemble_dim, station_dim, time_dim)] \
                + [(station_dim, time_dim), (time_dim, station_dim)]
        else:
            expected_dims = [(station_dim, time_dim), (time_dim, station_dim)]

        timeseries_variables = []
        for var_key, variable in self.__dataset.variables.items():
            if variable.dimensions in expected_dims:
                timeseries_variables.append(var_key)

        return timeseries_variables

    def read_timeseries_values(self, station_index: int, variable_name: str, ensemble_member: int = 0) -> np.ndarray:
        """
        Reads the specified timeseries from the input file.

        :param station_index: The index of the station for which the values should be read
        :param variable_name: The name of the variable for which the values should be read
        :return: an array of values
        """

        station_dim = self.__station_variable.dimensions[0]
        timeseries_variable = self.__dataset.variables[variable_name]

        # possibly usefull for in a debugger mode
        # assert set(timeseries_variable.dimensions)==set(('time', 'station')) \
        #        or set(timeseries_variable.dimensions)==set(('time', 'station', 'realization'))

        if self.__ensemble_member_variable is not None and 'realization' in timeseries_variable.dimensions:
            ensemble_member_dim = self.__ensemble_member_variable.dimensions[0]
            for i in range(3):
                if timeseries_variable.dimensions[i] == station_dim:
                    station_arg_Index = i
                elif timeseries_variable.dimensions[i] == ensemble_member_dim:
                    ensemble_arg_Index = i
            time_arg_Index = set(range(3)) - {station_arg_Index, ensemble_arg_Index}
            time_arg_Index = time_arg_Index.pop()
            argument = [None]*3
            argument[station_arg_Index] = station_index
            argument[ensemble_arg_Index] = ensemble_member
            argument[time_arg_Index] = slice(None)
            values = timeseries_variable[tuple(argument)]
        else:
            if timeseries_variable.dimensions[0] == station_dim:
                values = timeseries_variable[station_index, :]
            else:
                values = timeseries_variable[:, station_index]

        # NetCDF4 reads the values as a numpy masked array,
        # convert to a normal array with nan where mask == True
        return np.ma.filled(values, np.nan)

    def variable_dimensions(self, variable):
        return self.__dataset.variables[variable].dimensions

    @property
    def time_variable(self):
        return self.__time_variable

    @property
    def station_variable(self):
        return self.__station_variable

    @property
    def ensemble_member_variable(self):
        return self.__ensemble_member_variable

    @property
    def ensemble_size(self):
        """
        Ensemble size.
        """
        return self.__ensemble_size


class ExportDataset:
    """
    A class used to write data to a NetCDF file.
    Creates a new file or overwrites an old file. The file metadata will be written upon initialization. Data such
    as times, station data and timeseries data should be presented to the ExportDataset through the various methods.
    When all data has been written, the close method must be called to flush the changes from local memory to the
    actual file on disk.
    """

    def __init__(self, folder: str, basename: str):
        """
        :param folder:   Folder the file will be located in.
        :param basename: Basename of the file, extension ".nc" will be appended to this
        """
        # Create the file and open a Dataset to access it
        self.__filename = os.path.join(
            folder,
            basename + ".nc"
        )
        # use same write format as FEWS
        self.__dataset = Dataset(self.__filename, mode='w', format='NETCDF3_CLASSIC')

        # write metadata to the file
        self.__dataset.title = 'RTC-Tools Output Data'
        self.__dataset.institution = 'Deltares'
        self.__dataset.source = 'RTC-Tools'
        self.__dataset.history = 'Generated on {}'.format(datetime.now())
        self.__dataset.Conventions = 'CF-1.6'
        self.__dataset.featureType = 'timeseries'

        # dimensions are created when writing times and station data, must be created before writing variables
        self.__time_dim = None
        self.__station_dim = None
        self.__station_id_to_index_mapping = None
        self.__ensemble_member_index_dim = None

        self.__timeseries_variables = {}

    def __str__(self):
        return self.__filename

    def write_times(self, times: np.ndarray, forecast_time: float, forecast_date: datetime) -> None:
        """
        Writes a time variable to the given dataset.

        :param dataset:        The NetCDF4.Dataset object that the times will be written to (must have write permission)
        :param times:          The times that are to be written in seconds.
        :param forecast_time:  The forecast time in seconds corresponding to the forecast date
        :param forecast_date:  The datetime corresponding with time in seconds at the forecast index.
        """

        # in a NetCDF file times are written with respect to a reference date
        # the written values for the times may never be negative, so use the earliest time as the reference date
        reference_date = forecast_date
        minimum_time = np.min(times)
        if minimum_time < 0:
            times = times - minimum_time
            reference_date = reference_date - timedelta(seconds=forecast_time - minimum_time)

        self.__time_dim = self.__dataset.createDimension('time', None)

        time_var = self.__dataset.createVariable('time', 'f8', ('time',))
        time_var.standard_name = 'time'
        time_var.units = 'seconds since {}'.format(reference_date)
        time_var.axis = 'T'
        time_var[:] = times

    def write_ensemble_data(self, ensemble_size):

        if ensemble_size > 1:
            self.__ensemble_member_dim = self.__dataset.createDimension('realization', ensemble_size)
            ensemble_member_var = self.__dataset.createVariable('realization', 'i', ('realization',))
            ensemble_member_var.standard_name = 'realization'
            ensemble_member_var.long_name = 'Index of an ensemble member within an ensemble'
            ensemble_member_var.units = 1

    def write_station_data(self, stations: Stations, output_station_ids: List[str]) -> None:
        """
        Writes the station ids and additional station information to the given dataset.

        :param stations:           The stations data read from the input file.
        :param output_station_ids: The set of station ids for which output will be written. Must be unique.
        """
        assert len(set(output_station_ids)) == len(output_station_ids)

        self.__station_dim = self.__dataset.createDimension('station', len(output_station_ids))

        # first write the ids
        max_id_length = max(len(id) for id in output_station_ids)
        self.__dataset.createDimension('char_leng_id', max_id_length)
        station_id_var = self.__dataset.createVariable('station_id', 'c', ('station', 'char_leng_id'))
        station_id_var.long_name = 'station identification code'
        station_id_var.cf_role = 'timeseries_id'

        # we must store the index we use for each station id, to be able to write the data at the correct index later
        self.__station_id_to_index_mapping = {}
        for i, id in enumerate(output_station_ids):
            station_id_var[i, :] = list(id)
            self.__station_id_to_index_mapping[id] = i

        # now write the stored attributes
        for var_name, attr_var in stations.attribute_variables.items():
            variable = self.__dataset.createVariable(var_name, attr_var.datatype, ('station',))
            # copy all attributes from the original input variable
            variable.setncatts(attr_var.__dict__)

            for station_id in output_station_ids:
                if station_id in stations.attributes:
                    station_index = self.__station_id_to_index_mapping[station_id]
                    variable[station_index] = stations.attributes[station_id][var_name]

    def create_variables(self, variable_names: List[str], ensemble_size: int) -> None:
        """
        Creates variables in the dataset for each of the provided parameter ids.
        The write_times and write_station_data methods must be called first, to ensure the necessary dimensions have
        already been created in the output NetCDF file.

        :param variable_names: The parameter ids for which variables must be created. Must be unique.
        :param ensemble_size: the number of members in the ensemble
        """
        assert len(set(variable_names)) == len(variable_names)

        assert self.__time_dim is not None, 'First call write_times to ensure the time dimension has been created.'
        assert self.__station_dim is not None, 'First call write_station_data to ensure ' \
                                               'the station dimension has been created'
        assert self.__station_id_to_index_mapping is not None  # should also be created in write_station_data

        if ensemble_size > 1:
            assert self.__ensemble_member_dim is not None, \
                'First call write_ensemble_data to ensure the realization dimension has been created.'

            for variable_name in variable_names:
                self.__dataset.createVariable(variable_name,
                                              'f8', ('time', 'station', 'realization'), fill_value=np.nan)
        else:
            for variable_name in variable_names:
                self.__dataset.createVariable(variable_name,
                                              'f8', ('time', 'station'), fill_value=np.nan)

    def write_output_values(self, station_id: str, variable_name: str,
                            ensemble_member_index: int, values: np.ndarray,
                            ensemble_size: int) -> None:
        """
        Writes the given data to the dataset. The variable must have already been created through the
        create_variables method. After all calls to write_output_values, the close method must be called to flush all
        changes.

        :param station_id: The id of the station the data is written for.
        :param variable_name: The name of the variable the data is written to (must have already been created).
        :param ensemble_member_index: The index associated to the ensemble member
        :param values:        The values that are to be written to the file
        :param ensemble_size: the number of members in the ensemble
        """
        assert self.__station_id_to_index_mapping is not None, 'First call write_station_data and create_variables.'

        station_index = self.__station_id_to_index_mapping[station_id]
        if ensemble_size > 1:
            self.__dataset.variables[variable_name][:, station_index, ensemble_member_index] = values
        else:
            self.__dataset.variables[variable_name][:, station_index] = values

    def close(self) -> None:
        """
        Closes the NetCDF4 Dataset to ensure all changes made are written to the file.
        This method must be called after writing all data through the various write method.
        """
        self.__dataset.close()
