#!/usr/bin/env python

from nes.nc_projections.default_nes import Nes
from cfunits import Units


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,
                 **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
        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']
        """
        super(RotatedNes, self).__init__(comm=comm, path=path,
                                         info=info, dataset=dataset,
                                         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)

        self._rlat = self._get_coordinate_dimension('rlat')
        self._rlon = self._get_coordinate_dimension('rlon')

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

        self.projection_data = self.get_projection_data()

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

    def get_projection_data(self):
        """
        Read the projection data

        Returns
        -------
        projection : dict
            Dictionary with the projection data
        """
        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)

        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 Longitude
        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

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