#!/usr/bin/env python

import numpy as np
from pyproj import Proj
from .default_nes import Nes
from .latlon_nes import LatLonNes
import rasterio as rio
import datetime
from rasterio.windows import from_bounds
from pyproj import Proj, transform
from rasterio.features import shapes


class RasterNes(Nes):
    """
    Attributes
    ----------
    dataset : rasterio.io.DatasetReader
    """

    def __init__(self, comm=None, path=None, info=False, parallel_method='Y', **kwargs):
        raise NotImplementedError("RasterNes is not implemented yet")
        self.window = None

        super().__init__(comm=comm, path=path, info=info, dataset=None, parallel_method=parallel_method,
                         balanced=False, avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None,
                         create_nes=False, times=None, **kwargs)

    def open(self):
        """
        Open the NetCDF.
        """

        self.dataset = self.__open_raster()

        return None

    def __open_raster(self):
        """

        Returns
        -------
        rasterio.io.DatasetReader
        """
        raster = rio.open(self.get_ini_path(), 'r')
        self.dataset = raster

        return raster

    def _get_strlen(self):
        """
        Get the strlen

        Returns
        -------
        int
            Max length of the string data
        """
        return None

    def _get_lazy_variables(self):
        """
        Get all the variables' information.

        Returns
        -------
        variables : dict
            Dictionary with the variable name as key and another dictionary as value.
            De value dictionary will have the 'data' key with None as value and all the variable attributes as the
            other keys.
            e.g.
            {'var_name_1': {'data': None, 'attr_1': value_1_1, 'attr_2': value_1_2, ...},
             'var_name_2': {'data': None, 'attr_1': value_2_1, 'attr_2': value_2_2, ...},
             ...}
        """

        if self.master:
            variables = {}
            # Initialise data
            for band, units in enumerate(self.dataset.units):
                var_name = "band_{0}".format(band + 1)
                variables[var_name] = {}
                variables[var_name]['data'] = None
                variables[var_name]['dimensions'] = ('lat', 'lon')
                variables[var_name]['dtype'] = self.dataset.dtypes[band]
                if units is None:
                    units = '-'
                variables[var_name]['units'] = units

        else:
            variables = None
        variables = self.comm.bcast(variables, root=0)

        return variables

    def _get_time(self):
        """
        It is assumed that the Raster has no time information.
        Setting time to 01/01/2000.

        Returns
        -------
        List[datetime.datetime]
            A list of datetime objects corresponding to the time values in the Raster data.
        """
        time = [datetime.datetime(year=2000, month=1, day=1, hour=0, minute=0)]

        return time

    def _get_time_bnds(self, create_nes=False):
        """
        Get the Raster time bounds values.

        Parameters
        ----------
        create_nes : bool, optional
            If True, indicates that the method will create the object from scratch.
            If False (default), the method will use an existing file to obtain time bounds.

        Returns
        -------
        None
        """
        return None

    def _get_coordinates_bnds(self, create_nes=False):
        return None, None

    def _get_cell_measures(self, create_nes=False):
        """
        Get the Raster cell measures values.

        Parameters
        ----------
        create_nes : bool
            Indicates if you want to create the object from scratch (True) or through an existing file.

        Returns
        -------
        dict
            Dictionary of cell measures of the Raster data.
        """

        return {}

    def _get_global_attributes(self, create_nes=False):
        """
        Read the netcdf global attributes.

        Parameters
        ----------
        create_nes : bool
            Indicates if you want to create the object from scratch (True) or through an existing file.

        Returns
        -------
        gl_attrs : dict
            Dictionary with the netCDF global attributes.
        """

        gl_attrs = self.dataset.meta

        # Removing None
        for key, value in gl_attrs.items():
            if value is None:
                gl_attrs[key] = '-'
                gl_attrs[key] = '-'

        return gl_attrs

    def sel(self, hours_start=None, time_min=None, hours_end=None, time_max=None, lev_min=None, lev_max=None,
            lat_min=None, lat_max=None, lon_min=None, lon_max=None):
        """
        Select a slice of time, lev, lat or lon given a minimum and maximum limits.
        """
        if (hours_start is not None or time_min is not None or hours_end is not None or time_max is not None or
                lev_min is not None or lev_max is not None):
            raise TypeError("Raster only can be selected by lat lon values.")

        # Rasterio reads from top to bottom, contrary to NetCDF. Thus, the latitudes and data need to be flipped.
        # Calculate the inverse rank to ensure the first rank loads the bottom part of the raster instead of the top
        inverse_rank = self.size - self.rank - 1
        # Define the geographic coordinates of the entire dataset
        left, bottom, right, top = (lon_min, lat_min, lon_max, lat_max)

        # Create a window using the bounds for the entire dataset
        full_dataset_window = from_bounds(left, bottom, right, top, self.dataset.transform)

        # Calculate the width and height of each sub-window
        if self.parallel_method == 'Y':
            sub_window_width = full_dataset_window.width
            sub_window_height = full_dataset_window.height // self.size

            start_col = full_dataset_window.col_off
            start_row = (inverse_rank * sub_window_height) + full_dataset_window.row_off

            self.read_axis_limits = {'x_min': 0, 'x_max': None,
                                     'y_min': int(self.rank * sub_window_height),
                                     'y_max': int((self.rank + 1) * sub_window_height),
                                     'z_min': 0, 'z_max': None, 't_min': 0, 't_max': None}
            self.write_axis_limits = self.read_axis_limits.copy()
        elif self.parallel_method == 'X':
            sub_window_width = full_dataset_window.width // self.size
            sub_window_height = full_dataset_window.height

            start_col = (inverse_rank * sub_window_width) + full_dataset_window.col_off
            start_row = full_dataset_window.row_off

            self.read_axis_limits = {'x_min': int(self.rank * sub_window_width),
                                     'x_max': int((self.rank + 1) * sub_window_width),
                                     'y_min': 0, 'y_max': None,
                                     'z_min': 0, 'z_max': None, 't_min': 0, 't_max': None}
            self.write_axis_limits = self.read_axis_limits.copy()
        else:
            raise RuntimeError("Parallel method only accepted X or Y")

        # Create a window for each MPI process
        subset_window = rio.windows.Window(start_col, start_row, sub_window_width, sub_window_height)
        full_transform = self.dataset.window_transform(full_dataset_window)

        # Generate arrays of latitude and longitude values using np.linspace
        subset_y = np.linspace(full_transform.f,
                               full_transform.f + full_transform.e * int(full_dataset_window.height),
                               int(full_dataset_window.height))
        subset_x = np.linspace(full_transform.c,
                               full_transform.c + full_transform.a * int(full_dataset_window.width),
                               int(full_dataset_window.width))

        subset_y = np.flip(subset_y, axis=0)

        print(full_transform)
        print(self.dataset.transform)
        # Set the window and coordinate arrays as attributes
        self.window = subset_window
        print(self.dataset.crs.to_string())
        print(self.dataset.crs.to_epsg())
        if self.dataset.crs.to_string() == 'EPSG:4326':
            self._lat = {'data': subset_y}
            self._lon = {'data': subset_x}

            self.lat = self._get_coordinate_values(self._lat, 'Y')
            self.lon = self._get_coordinate_values(self._lon, 'X')

            self._var_dim = ('lat', 'lon')
            self._lat_dim = ('lat',)
            self._lon_dim = ('lon',)
        else:
            raise RuntimeError("Only EPSG:4326 supported for this raster")

        return subset_window

    def _read_variable(self, var_name):
        """
        Read a portion of the variable data based on the current rank.

        Parameters
        ----------
        var_name : str
            Name of the variable to read.

        Returns
        -------
        np.array
            A portion of the variable data corresponding to the current rank.
        """
        var_data = self.dataset.read(int(var_name[-1]), window=self.window)

        # Flip data
        var_data = np.flipud(var_data)
        # 2D to 4D
        var_data = np.expand_dims(var_data, axis=(0, 1))

        return var_data

    def to_latlon_nes(self):
        return LatLonNes(
            comm=self.comm, path=self.get_ini_path(), info=self.info, dataset=None,
            parallel_method=self.parallel_method, create_nes=True, balanced=self.balanced, times=self.time,
            levels=self.lev, lat=self.lat, lat_bnds=self.lat_bnds, lon=self.lon, lon_bnds=self.lon_bnds,
            vars=self.variables)
