#!/usr/bin/env python

import numpy as np
from cfunits import Units
from pyproj import Proj
from nes.nc_projections.default_nes import Nes


class MercatorNes(Nes):
    """

    Attributes
    ----------
    _y : dict
        Y coordinates dictionary with the complete 'data' key for all the values and the rest of the attributes.
    _x : dict
        X coordinates dictionary with the complete 'data' key for all the values and the rest of the attributes.
    y : dict
        Y coordinates dictionary with the portion of 'data' corresponding to the rank values.
    x : dict
        X coordinates dictionary with the portion of 'data' corresponding to the rank values.
    _var_dim : tuple
        Tuple with the name of the Y and X dimensions for the variables.
        ('y', 'x') for a Mercator projection.
    _lat_dim : tuple
        Tuple with the name of the dimensions of the Latitude values.
        ('y', 'x') for a Mercator projection.
    _lon_dim : tuple
        Tuple with the name of the dimensions of the Longitude values.
        ('y', 'x') for a Mercator 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 MercatorNes 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(MercatorNes, self).__init__(comm=comm, path=path, info=info, dataset=dataset,
                                          xarray=xarray, parallel_method=parallel_method, balanced=balanced,
                                          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._y = self._get_coordinate_dimension('y')
            self._x = self._get_coordinate_dimension('x')
        
        # Dimensions screening
        self.y = self._get_coordinate_values(self._y, 'Y')
        self.x = self._get_coordinate_values(self._x, '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 = ('y', 'x')
        self._lat_dim = ('y', 'x')
        self._lon_dim = ('y', 'x')

        self.free_vars('crs')


    @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: ['X', 'Y', 'T']
        """
        new = MercatorNes(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': 'mercator',
                          'standard_parallel':  [kwargs['lat_ts']], # TODO: Check if True
                          'longitude_of_projection_origin': kwargs['lon_0'],
                          }

        else:
            projection = self.variables['mercator']
            print(projection)
            self.free_vars('mercator')

        return projection

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

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

        netcdf.createDimension('y', len(self._y['data']))
        netcdf.createDimension('x', len(self._x['data']))

        return None

    def _create_dimension_variables(self, netcdf):
        """
        Create the 'y' and 'x' variables.

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

        super(MercatorNes, self)._create_dimension_variables(netcdf)

        # MERCATOR Y COORDINATES
        y = netcdf.createVariable('y', self._y['data'].dtype, ('y',))
        y.long_name = "y coordinate of projection"
        if 'units' in self._y.keys():
            y.units = Units(self._y['units'], formatted=True).units
        else:
            y.units = 'm'
        y.standard_name = "projection_y_coordinate"
        if self.size > 1:
            y.set_collective(True)
        y[:] = self._y['data']

        # MERCATOR X COORDINATES
        x = netcdf.createVariable('x', self._x['data'].dtype, ('x',))
        x.long_name = "x coordinate of projection"
        if 'units' in self._x.keys():
            x.units = Units(self._x['units'], formatted=True).units
        else:
            x.units = 'm'
        x.standard_name = "projection_x_coordinate"
        if self.size > 1:
            x.set_collective(True)
        x[:] = self._x['data']

        return None

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

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

        # Create a regular grid in metres (1D)
        self._x = {'data': np.linspace(kwargs['x_0'] + (kwargs['inc_x'] / 2), 
                                       kwargs['x_0'] + (kwargs['inc_x'] / 2) +
                                       (kwargs['inc_x'] * (kwargs['nx'] - 1)), kwargs['nx'],
                                       dtype=np.float)}


        self._y = {'data': np.linspace(kwargs['y_0'] + (kwargs['inc_y'] / 2), 
                                       kwargs['y_0'] + (kwargs['inc_y'] / 2) +
                                       (kwargs['inc_y'] * (kwargs['ny'] - 1)), kwargs['ny'],
                                       dtype=np.float)}

        # Create a regular grid in metres (1D to 2D)
        x = np.array([self._x['data']] * len(self._y['data']))
        y = np.array([self._y['data']] * len(self._x['data'])).T

        projection = Proj(
                          proj='merc',
                          a=6370000.00,
                          b=6370000.00,
                          lat_ts=kwargs['lat_ts'],
                          lon_0=kwargs['lon_0'],
                         )

        # Calculate center latitudes and longitudes (UTM to Mercator)
        self.center_lons, self.center_lats = projection(x, y, inverse=True)

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

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

        Parameters
        ----------
        var : Variable
            netCDF4-python variable object.
        """
        var.grid_mapping = 'mercator'

        return None

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

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

        mapping = netcdf.createVariable('mercator', 'c')
        mapping.grid_mapping_name = self.projection_data['grid_mapping_name']
        mapping.standard_parallel = self.projection_data['standard_parallel']
        mapping.longitude_of_projection_origin = self.projection_data['longitude_of_projection_origin']
       
        return None
