#!/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
            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(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: 'Y'.
            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']
            self.free_vars('mercator')

        return projection

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

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

        super(MercatorNes, self)._create_dimensions(netcdf)

        # Create y and x dimensions
        netcdf.createDimension('y', len(self._y['data']))
        netcdf.createDimension('x', len(self._x['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)

        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_centre_coordinates(self, **kwargs):
        """
        Calculate centre 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=6378137.0, # WGS84_SEMIMAJOR_AXIS (as defined in Cartopy source code)
                          b=6356752.3142, # WGS84_SEMIMINOR_AXIS (as defined in Cartopy source code)
                          lat_ts=kwargs['lat_ts'],
                          lon_0=kwargs['lon_0'],
                         )

        # Calculate centre latitudes and longitudes (UTM to Mercator)
        centre_lon_data, centre_lat_data = projection(x, y, inverse=True)
        centre_lat = {'data': centre_lat_data}
        centre_lon = {'data': centre_lon_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_x = np.abs(np.mean(np.diff(self.x['data'])))
        inc_y = np.abs(np.mean(np.diff(self.y['data'])))

        # Get bounds for rotated coordinates
        y_bounds = self.create_bounds(self.y['data'], inc_y)
        x_bounds = self.create_bounds(self.x['data'], inc_x)

        # Get rotated latitudes for grid edge
        left_edge_y = np.append(y_bounds.flatten()[::2], y_bounds.flatten()[-1])
        right_edge_y = np.flip(left_edge_y, 0)
        top_edge_y = np.repeat(y_bounds[-1][-1], len(self.x['data']) - 1)
        bottom_edge_y = np.repeat(y_bounds[0][0], len(self.x['data']))
        y_grid_edge = np.concatenate((left_edge_y, top_edge_y, right_edge_y, bottom_edge_y))

        # Get rotated longitudes for grid edge
        left_edge_x = np.repeat(x_bounds[0][0], len(self.y['data']) + 1)
        top_edge_x = x_bounds.flatten()[1:-1:2]
        right_edge_x = np.repeat(x_bounds[-1][-1], len(self.y['data']) + 1)
        bottom_edge_x = np.flip(x_bounds.flatten()[:-1:2], 0)
        x_grid_edge = np.concatenate((left_edge_x, top_edge_x, right_edge_x, bottom_edge_x))

        # Get edges for regular coordinates
        projection = Proj(
                          proj='merc',
                          a=6378137.0, # WGS84_SEMIMAJOR_AXIS (as defined in Cartopy source code)
                          b=6356752.3142, # WGS84_SEMIMINOR_AXIS (as defined in Cartopy source code)
                          lat_ts=float(self.projection_data['standard_parallel']),
                          lon_0=float(self.projection_data['longitude_of_central_meridian']),
                         )
        grid_edge_lon_data, grid_edge_lat_data = projection(x_grid_edge, y_grid_edge, inverse=True)
        
        # 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

    @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

    def to_grib2(self, path, grib_keys, grib_template_path, 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 write in a Mercator projection.")
