#!/usr/bin/env python

import numpy as np
from .default_nes import Nes


class LatLonNes(Nes):
    """

    Attributes
    ----------
    _var_dim : tuple
        Tuple with the name of the Y and X dimensions for the variables.
        ('lat', 'lon') for a regular latitude-longitude projection.
    _lat_dim : tuple
        Tuple with the name of the dimensions of the Latitude values.
        ('lat',) for a regular latitude-longitude projection.
    _lon_dim : tuple
        Tuple with the name of the dimensions of the Longitude values.
        ('lon',) for a regular latitude-longitude 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 LatLonNes 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(LatLonNes, 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')

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

        self._var_dim = ('lat', 'lon')
        self._lat_dim = ('lat',)
        self._lon_dim = ('lon',)

        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 = LatLonNes(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 _create_dimensions(self, netcdf):
        """
        Create 'spatial_nv' dimensions and the super dimensions 'lev', 'time', 'time_nv', 'lon' and 'lat'.

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

        super(LatLonNes, self)._create_dimensions(netcdf)

        # 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', 2)

        return None

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

        Returns
        ----------
        centre_lat : dict
            Dictionary with data of centre latitudes in 1D
        centre_lon : dict
            Dictionary with data of centre longitudes in 1D
        """

        inc_lat = kwargs['inc_lat']
        inc_lon = kwargs['inc_lon']

        # Global domain
        if len(kwargs) == 2:
            lat_orig = -90
            lon_orig = -180
            n_lat = int(180 // inc_lat)
            n_lon = int(360 // inc_lon)
        
        # Other domains
        else:
            lat_orig = kwargs['lat_orig']
            lon_orig = kwargs['lon_orig']
            n_lat = kwargs['n_lon']
            n_lon = kwargs['n_lon']

        # Calculate centre latitudes
        lat_c_orig = lat_orig + (inc_lat / 2)
        centre_lat = np.linspace(lat_c_orig, lat_c_orig + (inc_lat * (n_lat - 1)), n_lat)

        # Calculate centre longitudes
        lon_c_orig = lon_orig + (inc_lon / 2)
        centre_lon = np.linspace(lon_c_orig, lon_c_orig + (inc_lon * (n_lon - 1)), n_lon)

        return {'data': centre_lat}, {'data': 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).
        """

        model_centre_lon_data, model_centre_lat_data = np.meshgrid(self.lon['data'], self.lat['data'])

        # Calculate centre latitudes
        model_centre_lat = {'data': model_centre_lat_data}

        # Calculate centre longitudes
        model_centre_lon = {'data': model_centre_lon_data}

        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_lon = np.abs(np.mean(np.diff(self.lon['data'])))
        inc_lat = np.abs(np.mean(np.diff(self.lat['data'])))

        # Get bounds
        lat_bounds = self.create_single_spatial_bounds(self.lat['data'], inc_lat)
        lon_bounds = self.create_single_spatial_bounds(self.lon['data'], inc_lon)

        # Get latitudes for grid edge
        left_edge_lat = np.append(lat_bounds.flatten()[::2], lat_bounds.flatten()[-1])
        right_edge_lat = np.flip(left_edge_lat, 0)
        top_edge_lat = np.repeat(lat_bounds[-1][-1], len(self.lon['data']) - 1)
        bottom_edge_lat = np.repeat(lat_bounds[0][0], len(self.lon['data']))
        lat_grid_edge = np.concatenate((left_edge_lat, top_edge_lat, right_edge_lat, bottom_edge_lat))

        # Get longitudes for grid edge
        left_edge_lon = np.repeat(lon_bounds[0][0], len(self.lat['data']) + 1)
        top_edge_lon = lon_bounds.flatten()[1:-1:2]
        right_edge_lon = np.repeat(lon_bounds[-1][-1], len(self.lat['data']) + 1)
        bottom_edge_lon = np.flip(lon_bounds.flatten()[:-1:2], 0)
        lon_grid_edge = np.concatenate((left_edge_lon, top_edge_lon, right_edge_lon, bottom_edge_lon))

        # Create grid outline by stacking the edges in both coordinates
        model_grid_outline = np.vstack((lon_grid_edge, lat_grid_edge)).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 'crs'.

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

        var.grid_mapping = 'crs'
        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.
        """

        mapping = netcdf.createVariable('crs', 'c')
        # mapping = netcdf.createVariable('crs', 'i')
        mapping.grid_mapping_name = "latitude_longitude"
        mapping.semi_major_axis = 6371000.0
        mapping.inverse_flattening = 0

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

        return super(LatLonNes, self).to_grib2(path, grib_keys, grib_template_path, lat_flip=lat_flip, info=info)
