#!/usr/bin/env python

import numpy as np
import math
from cfunits import Units
from nes.nc_projections.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, **kwargs):
        """
        Initialize the RotatedNes class

        Parameters
        ----------
        comm: MPI.COMM
            Path to the CSV file that contains all the information.
        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 over Y axis
            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,
                                         **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 over Y axis
            accepted values: ['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 get_projection_data(self, create_nes, **kwargs):
        """
        Read the projection data

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

        if create_nes:
            projection = {'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:
            projection = self.variables['rotated_pole']
            self.free_vars('rotated_pole')

        return projection

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

        Parameters
        ----------
        netcdf : Dataset
            NetCDF object.
        """
        super(RotatedNes, self)._create_dimensions(netcdf)

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

        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"
        rlat.units = Units("degrees", formatted=True).units
        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"
        rlon.units = Units("degrees", formatted=True).units
        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_lats = 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_lons = np.linspace(kwargs['west_boundary'], kwargs['west_boundary'] +
                                       (kwargs['inc_rlon'] * (self.n_lon - 1)), self.n_lon)

        return {'data': self.rotated_lats}, {'data': self.rotated_lons}

    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)
        """

        degrees_to_radians = math.pi / 180.

        tph0 = kwargs['centre_lat'] * degrees_to_radians
        tlm = lon_deg * degrees_to_radians
        tph = lat_deg * degrees_to_radians
        tlm0d = -180 + kwargs['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_centroids(self, **kwargs):
        """
        Calculate center latitudes and longitudes from grid details.

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

        # Complete dimensions
        self._rlat, self._rlon = self._create_rotated_coordinates(**kwargs)

        # Calculate center latitudes and longitudes (1D to 2D)
        self.center_lons, self.center_lats = self.rotated2latlon(np.array([self.rotated_lons] * len(self.rotated_lats)),
                                                                 np.array([self.rotated_lats] * len(self.rotated_lons)).T,
                                                                 **kwargs)

        return {'data': self.center_lats}, {'data': self.center_lons}

    @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'

        return None

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

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

        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
