#!/usr/bin/env python

import warnings
import numpy as np
import pandas as pd
import math
from cfunits import Units
from copy import deepcopy
import geopandas as gpd
from shapely.geometry import Polygon
from .default_nes import Nes


class RotatedNes(Nes):
    """

    Attributes
    ----------
    _rlat : dict
        Rotated latitudes dictionary with the complete 'data' key for all the values and the rest of the attributes.
    _rlon : dict
        Rotated longitudes dictionary with the complete 'data' key for all the values and the rest of the attributes.
    rlat : dict
        Rotated latitudes dictionary with the portion of 'data' corresponding to the rank values.
    rlon : dict
        Rotated longitudes dictionary with the portion of 'data' corresponding to the rank values.
    projection_data : dict
        Dictionary with the projection information.
        'grid_north_pole_latitude' and 'grid_north_pole_longitude' keys.
    _var_dim : tuple
        Tuple with the name of the Y and X dimensions for the variables.
        ('rlat', 'rlon') for a rotated projection.
    _lat_dim : tuple
        Tuple with the name of the dimensions of the Latitude values.
        ('rlat', 'rlon') for a rotated projection.
    _lon_dim : tuple
        Tuple with the name of the dimensions of the Longitude values.
        ('rlat', 'rlon') for a rotated projection.
    """
    def __init__(self, comm=None, path=None, info=False, dataset=None, xarray=False, parallel_method='Y',
                 avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, create_nes=False,
                 balanced=False, times=None, **kwargs):
        """
        Initialize the RotatedNes class.

        Parameters
        ----------
        comm: MPI.COMM
            MPI Communicator.
        path: str
            Path to the NetCDF to initialize the object.
        info: bool
            Indicates if you want to get reading/writing info.
        dataset: Dataset
            NetCDF4-python Dataset to initialize the class.
        xarray: bool:
            (Not working) Indicates if you want to use xarray as default.
        parallel_method : str
            Indicates the parallelization method that you want. Default: 'Y'.
            Accepted values: ['X', 'Y', 'T'].
        avoid_first_hours : int
            Number of hours to remove from first time steps.
        avoid_last_hours : int
            Number of hours to remove from last time steps.
        """

        super(RotatedNes, self).__init__(comm=comm, path=path,
                                         info=info, dataset=dataset, balanced=balanced,
                                         xarray=xarray, parallel_method=parallel_method,
                                         avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours,
                                         first_level=first_level, last_level=last_level, create_nes=create_nes,
                                         times=times, **kwargs)

        if create_nes:
            # Dimensions screening
            self.lat = self._get_coordinate_values(self._lat, 'Y')
            self.lon = self._get_coordinate_values(self._lon, 'X')
        else:
            # Complete dimensions
            self._rlat = self._get_coordinate_dimension('rlat')
            self._rlon = self._get_coordinate_dimension('rlon')

        # Dimensions screening
        self.rlat = self._get_coordinate_values(self._rlat, 'Y')
        self.rlon = self._get_coordinate_values(self._rlon, 'X')

        # Get projection details
        self.projection_data = self.get_projection_data(create_nes, **kwargs)

        # Set axis limits for parallel writing
        self.write_axis_limits = self.get_write_axis_limits()

        self._var_dim = ('rlat', 'rlon')
        self._lat_dim = ('rlat', 'rlon')
        self._lon_dim = ('rlat', 'rlon')

    @staticmethod
    def new(comm=None, path=None, info=False, dataset=None, xarray=False, create=False, balanced=False,
            parallel_method='Y', avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None):
        """
        Initialize the Nes class.

        Parameters
        ----------
        comm: MPI.COMM
            MPI Communicator.
        path: str
            Path to the NetCDF to initialize the object.
        info: bool
            Indicates if you want to get reading/writing info.
        dataset: Dataset
            NetCDF4-python Dataset to initialize the class.
        xarray: bool:
            (Not working) Indicates if you want to use xarray as default.
        avoid_first_hours : int
            Number of hours to remove from first time steps.
        avoid_last_hours : int
            Number of hours to remove from last time steps.
        parallel_method : str
            Indicates the parallelization method that you want. Default: 'Y'.
            Accepted values: ['X', 'Y', 'T'].
        """

        new = RotatedNes(comm=comm, path=path, info=info, dataset=dataset, xarray=xarray, balanced=balanced,
                         parallel_method=parallel_method, avoid_first_hours=avoid_first_hours,
                         avoid_last_hours=avoid_last_hours, first_level=first_level, last_level=last_level)
        return new

    def filter_coordinates_selection(self):

        idx = self.get_idx_intervals()

        self.rlat = self._get_coordinate_values(self._rlat, 'Y')
        self.rlon = self._get_coordinate_values(self._rlon, 'X')

        self._rlat['data'] = self._rlat['data'][idx['idx_y_min']:idx['idx_y_max']]
        self._rlon['data'] = self._rlon['data'][idx['idx_x_min']:idx['idx_x_max']]

        super(RotatedNes, self).filter_coordinates_selection()

        return None

    def get_projection_data(self, create_nes, **kwargs):
        """
        Read the projection data.

        Returns
        -------
        projection_data : dict
            Dictionary with the projection data.
        """

        if create_nes:
            projection_data = {'data': None,
                               'dimensions': (),
                               'grid_mapping_name': 'rotated_latitude_longitude',
                               'grid_north_pole_latitude': 90 - kwargs['centre_lat'],
                               'grid_north_pole_longitude': -180 + kwargs['centre_lon'],
                               }
        else:
            if 'rotated_pole' in self.variables.keys():
                projection_data = self.variables['rotated_pole']
                self.free_vars('rotated_pole')
            else:
                msg = 'There is no variable called rotated_pole, projection has not been defined.'
                warnings.warn(msg)
                projection_data = None
        
        return projection_data

    def _create_dimensions(self, netcdf):
        """
        Create 'rlat', 'rlon'  and 'spatial_nv' dimensions and the super dimensions 'lev', 'time', 'time_nv', 'lon' and 'lat'.

        Parameters
        ----------
        netcdf : Dataset
            NetCDF object.
        """

        super(RotatedNes, self)._create_dimensions(netcdf)

        # Create rlat and rlon dimensions
        netcdf.createDimension('rlon', len(self._rlon['data']))
        netcdf.createDimension('rlat', len(self._rlat['data']))

        # Create spatial_nv (number of vertices) dimension
        if (self._lat_bnds is not None) and (self._lon_bnds is not None):
            netcdf.createDimension('spatial_nv', 4)
            pass

        return None

    def _create_dimension_variables(self, netcdf):
        """
        Create the 'rlat' and 'rlon' variables.

        Parameters
        ----------
        netcdf : Dataset
            NetCDF object.
        """

        super(RotatedNes, self)._create_dimension_variables(netcdf)

        # ROTATED LATITUDES
        rlat = netcdf.createVariable('rlat', self._rlat['data'].dtype, ('rlat',))
        rlat.long_name = "latitude in rotated pole grid"
        if 'units' in self._rlat.keys():
            rlat.units = Units(self._rlat['units'], formatted=True).units
        else:
            rlat.units = 'degrees'
        rlat.standard_name = "grid_latitude"
        if self.size > 1:
            rlat.set_collective(True)
        rlat[:] = self._rlat['data']

        # ROTATED LONGITUDES
        rlon = netcdf.createVariable('rlon', self._rlon['data'].dtype, ('rlon',))
        rlon.long_name = "longitude in rotated pole grid"
        if 'units' in self._rlon.keys():
            rlon.units = Units(self._rlon['units'], formatted=True).units
        else:
            rlon.units = 'degrees'
        rlon.standard_name = "grid_longitude"
        if self.size > 1:
            rlon.set_collective(True)
        rlon[:] = self._rlon['data']

        return None

    def _create_rotated_coordinates(self, **kwargs):
        """
        Calculate rotated latitudes and longitudes from grid details.

        Parameters
        ----------
        netcdf : Dataset
            NetCDF object.
        """

        # Calculate rotated latitudes
        self.n_lat = int((abs(kwargs['south_boundary']) / kwargs['inc_rlat']) * 2 + 1)
        self.rotated_lat = np.linspace(kwargs['south_boundary'], kwargs['south_boundary'] +
                                       (kwargs['inc_rlat'] * (self.n_lat - 1)), self.n_lat)
        # Calculate rotated longitudes
        self.n_lon = int((abs(kwargs['west_boundary']) / kwargs['inc_rlon']) * 2 + 1)
        self.rotated_lon = np.linspace(kwargs['west_boundary'], kwargs['west_boundary'] +
                                       (kwargs['inc_rlon'] * (self.n_lon - 1)), self.n_lon)

        return {'data': self.rotated_lat}, {'data': self.rotated_lon}

    def rotated2latlon(self, lon_deg, lat_deg, lon_min=-180, **kwargs):
        """
        Calculate the unrotated coordinates using the rotated ones.

        :param lon_deg: Rotated longitude coordinate.
        :type lon_deg: numpy.array

        :param lat_deg: Rotated latitude coordinate.
        :type lat_deg: numpy.array

        :param lon_min: Minimum value for the longitudes: -180 (-180 to 180) or 0 (0 to 360).
        :type lon_min: float

        :return: Unrotated coordinates. Longitudes, Latitudes.
        :rtype: tuple(numpy.array, numpy.array)
        """

        if 'centre_lat' in kwargs:
            centre_lat = kwargs['centre_lat']
        else:
            centre_lat = 90 - float(self.projection_data['grid_north_pole_latitude'])

        if 'centre_lon' in kwargs:
            centre_lon = kwargs['centre_lon']
        else:
            centre_lon = float(self.projection_data['grid_north_pole_longitude']) + 180

        degrees_to_radians = math.pi / 180.

        tph0 = centre_lat * degrees_to_radians
        tlm = lon_deg * degrees_to_radians
        tph = lat_deg * degrees_to_radians
        tlm0d = -180 + centre_lon
        ctph0 = np.cos(tph0)
        stph0 = np.sin(tph0)

        stlm = np.sin(tlm)
        ctlm = np.cos(tlm)
        stph = np.sin(tph)
        ctph = np.cos(tph)

        # LATITUDES
        sph = (ctph0 * stph) + (stph0 * ctph * ctlm)
        sph[sph > 1.] = 1.
        sph[sph < -1.] = -1.
        aph = np.arcsin(sph)
        aphd = aph / degrees_to_radians

        # LONGITUDES
        anum = ctph * stlm
        denom = (ctlm * ctph - stph0 * sph) / ctph0
        relm = np.arctan2(anum, denom) - math.pi
        almd = relm / degrees_to_radians + tlm0d
        almd[almd > (lon_min + 360)] -= 360
        almd[almd < lon_min] += 360

        return almd, aphd

    def _create_centre_coordinates(self, **kwargs):
        """
        Calculate centre latitudes and longitudes from grid details.

        Returns
        ----------
        centre_lat : dict
            Dictionary with data of centre coordinates for latitude in 2D (latitude, longitude).
        centre_lon : dict
            Dictionary with data of centre coordinates for longitude in 2D (latitude, longitude).
        """
        
        # Complete dimensions
        self._rlat, self._rlon = self._create_rotated_coordinates(**kwargs)

        # Calculate centre latitudes and longitudes (1D to 2D)
        centre_lon_data, centre_lat_data = self.rotated2latlon(np.array([self.rotated_lon] * len(self.rotated_lat)),
                                                               np.array([self.rotated_lat] * len(self.rotated_lon)).T,
                                                               **kwargs)
        centre_lon = {'data': centre_lon_data}
        centre_lat = {'data': centre_lat_data}

        return centre_lat, centre_lon

    def create_providentia_exp_centre_coordinates(self):
        """
        Calculate centre latitudes and longitudes from original coordinates and store as 2D arrays.

        Returns
        ----------
        model_centre_lat : dict
            Dictionary with data of centre coordinates for latitude in 2D (latitude, longitude).
        model_centre_lon : dict
            Dictionary with data of centre coordinates for longitude in 2D (latitude, longitude).
        """

        # Get centre latitudes
        model_centre_lat = self.lat

        # Get centre longitudes
        model_centre_lon = self.lon

        return model_centre_lat, model_centre_lon

    def create_providentia_exp_grid_edge_coordinates(self):
        """
        Calculate grid edge latitudes and longitudes and get model grid outline.

        Returns
        ----------
        grid_edge_lat : dict
            Dictionary with data of grid edge latitudes.
        grid_edge_lon : dict
            Dictionary with data of grid edge longitudes.
        """

        # Get grid resolution
        inc_rlon = np.abs(np.mean(np.diff(self.rlon['data'])))
        inc_rlat = np.abs(np.mean(np.diff(self.rlat['data'])))

        # Get bounds for rotated coordinates
        rlat_bounds = self.create_single_spatial_bounds(self.rlat['data'], inc_rlat)
        rlon_bounds = self.create_single_spatial_bounds(self.rlon['data'], inc_rlon)

        # Get rotated latitudes for grid edge
        left_edge_rlat = np.append(rlat_bounds.flatten()[::2], rlat_bounds.flatten()[-1])
        right_edge_rlat = np.flip(left_edge_rlat, 0)
        top_edge_rlat = np.repeat(rlat_bounds[-1][-1], len(self.rlon['data']) - 1)
        bottom_edge_rlat = np.repeat(rlat_bounds[0][0], len(self.rlon['data']))
        rlat_grid_edge = np.concatenate((left_edge_rlat, top_edge_rlat, right_edge_rlat, bottom_edge_rlat))

        # Get rotated longitudes for grid edge
        left_edge_rlon = np.repeat(rlon_bounds[0][0], len(self.rlat['data']) + 1)
        top_edge_rlon = rlon_bounds.flatten()[1:-1:2]
        right_edge_rlon = np.repeat(rlon_bounds[-1][-1], len(self.rlat['data']) + 1)
        bottom_edge_rlon = np.flip(rlon_bounds.flatten()[:-1:2], 0)
        rlon_grid_edge = np.concatenate((left_edge_rlon, top_edge_rlon, right_edge_rlon, bottom_edge_rlon))

        # Get edges for regular coordinates
        grid_edge_lon_data, grid_edge_lat_data = self.rotated2latlon(rlon_grid_edge, rlat_grid_edge)
        
        # Create grid outline by stacking the edges in both coordinates
        model_grid_outline = np.vstack((grid_edge_lon_data, grid_edge_lat_data)).T

        grid_edge_lat = {'data': model_grid_outline[:,1]}
        grid_edge_lon = {'data': model_grid_outline[:,0]}
        
        return grid_edge_lat, grid_edge_lon

    def create_spatial_bounds(self):
        """
        Calculate longitude and latitude bounds and set them.
        """

        # Calculate rotated coordinates bounds
        inc_rlat = np.abs(np.mean(np.diff(self._rlat['data'])))
        rlat_bnds = self.create_single_spatial_bounds(np.array([self._rlat['data']] * len(self._rlon['data'])).T, 
                                                      inc_rlat, spatial_nv=4, inverse=True)

        inc_rlon = np.abs(np.mean(np.diff(self._rlon['data'])))
        rlon_bnds = self.create_single_spatial_bounds(np.array([self._rlon['data']] * len(self._rlat['data'])), 
                                                      inc_rlon, spatial_nv=4)

        # Transform rotated bounds to regular bounds
        lon_bnds, lat_bnds = self.rotated2latlon(rlon_bnds, rlat_bnds)

        # Obtain regular coordinates bounds
        self._lat_bnds = {}
        self._lat_bnds['data'] = deepcopy(lat_bnds)
        self.lat_bnds = {}
        self.lat_bnds['data'] = lat_bnds[self.write_axis_limits['y_min']:self.write_axis_limits['y_max'],
                                         self.write_axis_limits['x_min']:self.write_axis_limits['x_max'], 
                                         :]

        self._lon_bnds = {}
        self._lon_bnds['data'] = deepcopy(lon_bnds)
        self.lon_bnds = {}
        self.lon_bnds['data']= lon_bnds[self.write_axis_limits['y_min']:self.write_axis_limits['y_max'],
                                        self.write_axis_limits['x_min']:self.write_axis_limits['x_max'], 
                                        :]

        return None

    @staticmethod
    def _set_var_crs(var):
        """
        Set the grid_mapping to 'rotated_pole'.

        Parameters
        ----------
        var : Variable
            netCDF4-python variable object.
        """

        var.grid_mapping = 'rotated_pole'
        var.coordinates = "lat lon"

        return None

    def _create_metadata(self, netcdf):
        """
        Create the 'crs' variable for the rotated latitude longitude grid_mapping.

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python Dataset.
        """

        if self.projection_data is not None:
            mapping = netcdf.createVariable('rotated_pole', 'c')
            mapping.grid_mapping_name = self.projection_data['grid_mapping_name']
            mapping.grid_north_pole_latitude = self.projection_data['grid_north_pole_latitude']
            mapping.grid_north_pole_longitude = self.projection_data['grid_north_pole_longitude']

        return None

    def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=False, info=False):
        """
        Write output file with grib2 format.

        Parameters
        ----------
        path : str
            Path to the output file.
        grib_keys : dict
            Dictionary with the grib2 keys.
        grib_template_path : str
            Path to the grib2 file to use as template.
        info : bool
            Indicates if you want to print extra information during the process.
        """

        raise NotImplementedError("Grib2 format cannot be written in a Rotated pole projection.")

    def create_shapefile(self):
        """
        Create spatial geodataframe (shapefile).
        """

        if self._lat_bnds is None or self._lon_bnds is None:
            self.create_spatial_bounds()

        # Reshape arrays to create geometry
        aux_b_lats = self.lat_bnds['data'].reshape((self.lat_bnds['data'].shape[0] * self.lat_bnds['data'].shape[1], 
                                                    self.lat_bnds['data'].shape[2]))
        aux_b_lons = self.lon_bnds['data'].reshape((self.lon_bnds['data'].shape[0] * self.lon_bnds['data'].shape[1], 
                                                    self.lon_bnds['data'].shape[2]))

        # Create dataframe cointaining all polygons
        geometry = []
        for i in range(aux_b_lons.shape[0]):
            geometry.append(Polygon([(aux_b_lons[i, 0], aux_b_lats[i, 0]),
                                     (aux_b_lons[i, 1], aux_b_lats[i, 1]),
                                     (aux_b_lons[i, 2], aux_b_lats[i, 2]),
                                     (aux_b_lons[i, 3], aux_b_lats[i, 3]),
                                     (aux_b_lons[i, 0], aux_b_lats[i, 0])]))
        fids = np.arange(self._lat['data'].shape[0] * self._lat['data'].shape[1])
        fids = fids.reshape((self._lat['data'].shape[0], self._lat['data'].shape[1]))
        fids = fids[self.write_axis_limits['y_min']:self.write_axis_limits['y_max'],
                    self.write_axis_limits['x_min']:self.write_axis_limits['x_max']]
        gdf = gpd.GeoDataFrame(index=pd.Index(name='FID', data=fids.ravel()), 
                               geometry=geometry, 
                               crs="EPSG:4326")
        self.shapefile = gdf

        return gdf
