#!/usr/bin/env python

import numpy as np
from nes.nc_projections.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, **kwargs):
        """
        Initialize the LatLonNes 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(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,
                                        **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 over Y axis
            accepted values: ['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_centroids(self, **kwargs):
        """
        Calculate center latitudes and longitudes from grid details.

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

        # Calculate center latitudes
        lat_c_orig = kwargs['lat_orig'] + (kwargs['inc_lat'] / 2)
        self.center_lats = np.linspace(
            lat_c_orig, lat_c_orig + (kwargs['inc_lat'] * (kwargs['n_lat'] - 1)), kwargs['n_lat'])

        # Calculate center longitudes
        lon_c_orig = kwargs['lon_orig'] + (kwargs['inc_lon'] / 2)
        self.center_lons = np.linspace(
            lon_c_orig, lon_c_orig + (kwargs['inc_lon'] * (kwargs['n_lon'] - 1)), kwargs['n_lon'])

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

    def _create_bounds(self, **kwargs):

        # This function is not being used
        spatial_nv = 2
        boundary_lats = self.create_bounds(self.center_lats, kwargs['inc_lat'], spatial_nv)
        boundary_lons = self.create_bounds(self.center_lons, kwargs['inc_lon'], spatial_nv)

        return boundary_lats, boundary_lons

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

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

        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('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 _create_metadata(self, netcdf):
        """
        Create metadata variables

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

        # Set crs
        self._set_crs(netcdf)

        return None