#!/usr/bin/env python
from dataclasses import dataclass
from pandas import DataFrame
import xarray as xr
import pandas as pd
import numpy as np
from typing import Tuple, Dict, Optional
from numpy import cos, sin, arctan, pi, nan
from functools import partial, wraps
import geopandas as gpd
from geopandas import GeoDataFrame
from shapely.geometry import Polygon, Point
from mapies.util.func_tools import *
import logging

logger = logging.getLogger(__name__)
logging.basicConfig(filename='../logs/mapies.log', level=logging.INFO)


def regridding_function(r, lon, lat, obs, obserr=None):

    # If obserr is selected return that
    if obserr is not None:
        if isinstance(r, RotatedGrid):
            lon_agg, lat_agg, rlon_agg, rlat_agg, obs_agg, obserr_agg, count_obs = r.aggregate(lon, lat, obs, obserr)
            return lon_agg, lat_agg, rlon_agg, rlat_agg, obs_agg, obserr_agg, count_obs
        elif isinstance(r, RegularGrid):
            lon_agg, lat_agg, obs_agg, obserr_agg, count_obs = r.aggregate(lon, lat, obs, obserr)
            return lon_agg, lat_agg, obs_agg, obserr_agg, count_obs
        else:
            raise ValueError("Invalid grid representation")

    else:
        if isinstance(r, RotatedGrid):
            lon_agg, lat_agg, rlon_agg, rlat_agg, obs_agg, count_obs = r.aggregate(lon, lat, obs)
            return lon_agg, lat_agg, rlon_agg, rlat_agg, obs_agg, count_obs
        elif isinstance(r, RegularGrid):
            lon_agg, lat_agg, obs_agg, count_obs = r.aggregate(lon, lat, obs)
            return lon_agg, lat_agg, obs_agg, count_obs
        else:
            raise ValueError("Invalid grid representation")


@dataclass
class Grid:
    dlon: float = None
    dlat: float = None
    nlon: float = None
    nlat: float = None
    west: float = None
    south: float = None
    centre_lat: Optional[float] = None
    centre_lon: Optional[float] = None


    def calculate_grid_coords(self):
        """
        This method does two things:
        1. It calculates the centre values of the grid
        2. Before Calculating the boundary points of the regular grid 

        This can then be used in many grid representations, for example in a rotated grid used in the ihouse Monarch model
        """
        logger.info("Creating Grid")
        if (self.nlon == None) and (self.nlat == None):
            self.nlat = int((abs(self.south) / self.dlat) * 2)
            self.nlon = int((abs(self.west) / self.dlon) * 2)
        elif (self.dlon == None) and (self.dlat == None):
            self.dlat = abs(self.south) * 2 / (self.nlat)
            self.dlon = abs(self.west) * 2 / (self.nlon)
        else:
            print("Enter resolution or number of grid cells")

        self.cenlat = self.south + self.dlat/2
        self.cenlon = self.west + self.dlon/2
        self.center_latitudes = np.linspace(self.cenlat, self.cenlat + self.dlat * (self.nlat) - self.dlat, self.nlat, dtype=float)
        self.center_longitudes = np.linspace(self.cenlon, self.cenlon + self.dlon * (self.nlon) - self.dlat, self.nlon, dtype=float)

        
        counter = 0
        self.coords = np.empty((self.center_latitudes.shape[0] * self.center_longitudes.shape[0], 2))
        for i in range(self.center_longitudes.shape[0]):
            for j in range(self.center_latitudes.shape[0]):
                self.coords[counter, 0] = self.center_longitudes[i]
                self.coords[counter, 1] = self.center_latitudes[j]
                counter += 1


        self.b_lats = self.create_bounds(self.coords[:, 1], self.dlat, number_vertices=4, inverse=True)
        self.b_lons = self.create_bounds(self.coords[:, 0], self.dlon, number_vertices=4)

        return True


    @staticmethod
    def create_bounds(coordinates, inc, number_vertices=2, inverse=False):
        """
        Calculate the boundaries of each grid points
        """
        # Create new arrays moving the centers half increment less and more.
        coords_left = coordinates - inc / 2
        coords_right = coordinates + inc / 2

        # Defining the number of corners needed. 2 to regular grids and 4 for irregular ones.
        if number_vertices == 2:
            # Create an array of N arrays of 2 elements to store the floor and the ceil values for each cell
            bound_coords = np.dstack((coords_left, coords_right))
            bound_coords = bound_coords.reshape((len(coordinates), number_vertices))
        elif number_vertices == 4:
            # Create an array of N arrays of 4 elements to store the corner values for each cell
            # It can be stored in clockwise starting form the left-top element, or in inverse mode.
            if inverse:
                bound_coords = np.dstack((coords_left, coords_left, coords_right, coords_right))
            else:
                bound_coords = np.dstack((coords_left, coords_right, coords_right, coords_left))
        else:
            print('The number of vertices of the boundaries must be 2 or 4.')
        return bound_coords


    def building_grid(self):
        """
        Here we build the grid using Shapely polygons. And attribute it to a geopandas dataframe

        This is also where we decide whether to use a regular grid for global or a rotated grid for regional

        """


        # May need to add centre lats and centre lons to this dataframe as well

        # Re-shape
        aux_b_lats = self.b_lats.squeeze()
        aux_b_lons = self.b_lons.squeeze()
        

        geometry = []
        # Create one dataframe with 8 columns, 4 points with two coordinates each one
        for i in range(aux_b_lons.shape[0]):
            geometry.append(Polygon([(aux_b_lons[i, 0], aux_b_lats[i, 0]),
                                     (aux_b_lons[i, 1], aux_b_lats[i, 1]),
                                     (aux_b_lons[i, 2], aux_b_lats[i, 2]),
                                     (aux_b_lons[i, 3], aux_b_lats[i, 3]),
                                     (aux_b_lons[i, 0], aux_b_lats[i, 0])]))

        self.gdf = GeoDataFrame(index=range(aux_b_lons.shape[0]), geometry=geometry, crs='epsg:4326')
        self.gdf["grid_cell"] = ["grid_cell_{}".format(i+1) for i in range(len(self.gdf))]
        
        # Append the center lat and lon values to this dataframe
        self.gdf["lon_cen"] = self.coords[:, 0]
        self.gdf["lat_cen"] = self.coords[:, 1]

        
        return True

@dataclass
class RegularGrid(Grid):
    """
    Regridding if the Grid we want is regular, but also want to aggregate the observations to a regular grid.
    
    Useful for a global grid representation
    """
    projection: str = "regular"

    def __post_init__(self):
        self.calculate_grid_coords()
        
    def aggregate(self, lon, lat, obs, obserr=None):
        """
        Aggregate the observations to the grid cells
        
        Parameters:
        lon, lat, obs, obserr (optional) -- Longitude, Latitude, Observations, and Errors
        
        Returns:
        Aggregated lon, lat, obs values
        """

        # Make sure lon and lat are within the borders of the grid
        # Filter out values that are outside the grid bounds
        mask = (lon >= self.west) & (lon <= self.west + (self.dlon * (self.nlon - 1))) & \
               (lat >= self.south) & (lat <= self.south + (self.dlat * (self.nlat - 1))) & \
               (~np.isnan(obs))


        # Apply the mask to filter out the out-of-bounds points
        lon = lon[mask]
        lat = lat[mask]
        obs = obs[mask]
        if obserr is not None:
            obserr = obserr[mask]
        
        # Vectorized calculation of grid indices
        lon_idx = np.searchsorted(self.center_longitudes, lon) - 1  # Index of longitude grid
        lat_idx = np.searchsorted(self.center_latitudes, lat) - 1  # Index of latitude grid
        
        # Ensure indices are within bounds
        lon_idx = np.clip(lon_idx, 0, self.nlon - 1)
        lat_idx = np.clip(lat_idx, 0, self.nlat - 1)
        
        # Initialize arrays for aggregation
        aggregated_obs = np.zeros((self.nlat, self.nlon))
        aggregated_lon = np.zeros((self.nlat, self.nlon))
        aggregated_lat = np.zeros((self.nlat, self.nlon))
        aggregated_obs_err = np.zeros((self.nlat, self.nlon)) if obserr is not None else None
        count_obs = np.zeros((self.nlat, self.nlon), dtype=int)
        
        # Efficient accumulation using np.add.at
        np.add.at(aggregated_obs, (lat_idx, lon_idx), obs)  # Accumulate observations
        np.add.at(aggregated_lon, (lat_idx, lon_idx), lon)  # Accumulate lon values
        np.add.at(aggregated_lat, (lat_idx, lon_idx), lat)  # Accumulate lat values
        np.add.at(count_obs, (lat_idx, lon_idx), 1)  # Count observations per grid cell
        
        if obserr is not None:
            np.add.at(aggregated_obs_err, (lat_idx, lon_idx), obserr)  # Accumulate errors
        
        # Compute averages by dividing by counts
        non_zero_counts = count_obs > 0
        aggregated_obs[non_zero_counts] /= count_obs[non_zero_counts]
        aggregated_lon[non_zero_counts] /= count_obs[non_zero_counts]
        aggregated_lat[non_zero_counts] /= count_obs[non_zero_counts]
        
        if aggregated_obs_err is not None:
            aggregated_obs_err[non_zero_counts] /= count_obs[non_zero_counts]
        
        # TODO: Add a function that takes the center or aggreagated value
        # Return the results
        lon_out = aggregated_lon.flatten() #self.center_longitudes
        lat_out = aggregated_lat.flatten() #self.center_latitudes
        obs_out = aggregated_obs.flatten()
        count_obs = count_obs.flatten()
        
        if aggregated_obs_err is not None:
            obs_err_out = aggregated_obs_err.flatten()
            return lon_out, lat_out, obs_out, obs_err_out, count_obs
        else:
            return lon_out, lat_out, obs_out, count_obs


@dataclass
class RotatedGrid(Grid):
    """
    Regridding if the Grid we want is regular, but also want to aggregate the observations to a regular grid.
    
    Useful for a global grid representation
    """

    projection: str = "rotated"

    def __post_init__(self):
        self.calculate_grid_coords()
        
    def aggregate(self, lon, lat, obs, obserr=None):
        """
        Aggregate the observations to the grid cells
        
        Parameters:
        lon, lat, obs, obserr (optional) -- Longitude, Latitude, Observations, and Errors
        
        Returns:
        Aggregated lon, lat, rlon, rlat, obs values
        """
        # Perform the rotation
        rlon, rlat = geo_to_rot(lon, lat, centre_lon=self.centre_lon, centre_lat=self.centre_lat)

        # Make sure rlon and rlat are within the borders of the grid
        # Filter out values that are outside the grid bounds
        mask = (rlon >= self.west) & (rlon <= self.west + (self.dlon * (self.nlon - 1))) & \
               (rlat >= self.south) & (rlat <= self.south + (self.dlat * (self.nlat - 1))) & \
               (~np.isnan(obs))


        # Apply the mask to filter out the out-of-bounds points
        rlon = rlon[mask]
        rlat = rlat[mask]
        lon = lon[mask]
        lat = lat[mask]
        obs = obs[mask]
        if obserr is not None:
            obserr = obserr[mask]


        # Vectorized calculation of grid indices
        rlon_idx = np.searchsorted(self.center_longitudes, rlon) - 1  # Index of longitude grid
        rlat_idx = np.searchsorted(self.center_latitudes, rlat) - 1  # Index of latitude grid

        
        # Ensure indices are within bounds
        rlon_idx = np.clip(rlon_idx, 0, self.nlon - 1)
        rlat_idx = np.clip(rlat_idx, 0, self.nlat - 1)

        
        # Initialize arrays for aggregation
        aggregated_obs = np.zeros((self.nlat, self.nlon))
        aggregated_lon = np.zeros((self.nlat, self.nlon))
        aggregated_lat = np.zeros((self.nlat, self.nlon))
        aggregated_rlon = np.zeros((self.nlat, self.nlon))
        aggregated_rlat = np.zeros((self.nlat, self.nlon))
        aggregated_obs_err = np.ma.zeros((self.nlat, self.nlon)) if obserr is not None else None
        count_obs = np.zeros((self.nlat, self.nlon), dtype=int)



        # Efficient accumulation using np.add.at
        np.add.at(aggregated_obs, (rlat_idx, rlon_idx), obs)  # Accumulate observations
        np.add.at(aggregated_lon, (rlat_idx, rlon_idx), lon)  # Accumulate lon values
        np.add.at(aggregated_lat, (rlat_idx, rlon_idx), lat)  # Accumulate lat values
        np.add.at(aggregated_rlon, (rlat_idx, rlon_idx), rlon)  # Accumulate lon values
        np.add.at(aggregated_rlat, (rlat_idx, rlon_idx), rlat)  # Accumulate lat values
        np.add.at(count_obs, (rlat_idx, rlon_idx), 1)  # Count observations per grid cell
        

        if obserr is not None:
            np.add.at(aggregated_obs_err, (rlat_idx, rlon_idx), obserr)  # Accumulate errors

        
        # Compute averages by dividing by counts
        non_zero_counts = count_obs > 0
        aggregated_obs[non_zero_counts] /= count_obs[non_zero_counts]
        aggregated_lon[non_zero_counts] /= count_obs[non_zero_counts]
        aggregated_lat[non_zero_counts] /= count_obs[non_zero_counts]
        aggregated_rlon[non_zero_counts] /= count_obs[non_zero_counts]
        aggregated_rlat[non_zero_counts] /= count_obs[non_zero_counts]
        
        if aggregated_obs_err is not None:
            aggregated_obs_err[non_zero_counts] /= count_obs[non_zero_counts]

        # Return the results
        lon_out = aggregated_lon.flatten()
        lat_out = aggregated_lat.flatten()
        rlon_out = aggregated_lon.flatten() #self.center_longitudes
        rlat_out = aggregated_lat.flatten() #self.center_latitudess
        obs_out = aggregated_obs.flatten()
        count_obs = count_obs.flatten()


        if aggregated_obs_err is not None:
            obs_err_out = aggregated_obs_err.flatten()
            return lon_out, lat_out, rlon_out, rlat_out, obs_out, obs_err_out, count_obs
        else:
            return lon_out, lat_out, rlon_out, rlat_out, obs_out, count_obs


@dataclass
class IrregularRotatedGrid(Grid):

    projection: str = "rotated"
        
    def __post_init__(self):
        self.calculate_grid_coords()
        self.building_grid()
        
        
    def aggregate(self, lon, lat, obs, obserr=None):
        """
        Aggregate
        
        Parameters:
        lon, lat, obs values Optional obserr
        
        Returns:
        lon, lat, rlon, rlat, obs values
        """
        # Rotate the observations
        rlon, rlat = geo_to_rot(lon, lat, centre_lon=self.centre_lon, centre_lat=self.centre_lat)

        if obserr is not None:
            df = pd.DataFrame({"obs":obs, "lon":lon, "lat":lat, "rlon":rlon, "rlat":rlat, "obserr": obserr})
        else:
            df = pd.DataFrame({"obs":obs, "lon":lon, "lat":lat, "rlon":rlon, "rlat":rlat})

        df['coords'] = list(zip(df['rlon'],df['rlat']))
        df['coords'] = df['coords'].apply(Point)
        

        # Create a geodataframe of the rotated obseravtional lon/ lat points
        points = gpd.GeoDataFrame(df, geometry='coords', crs='epsg:4326')
        
        
        # Geopandas sjoin calculates if these shapely points are within the shapely grid pologons
        try:
            gdf = self.gdf.sjoin(points, how="left") # , predicate="within")
        except TypeError:
            print("Empty dataframe messing up some versions of geopandas")
            gdf = pd.DataFrame(columns=[
                'obstype', 'obslon', 'obslat', 'obsn', 'obsid', 'obslev', 'obs',
                'obserr', 'time', 'lon', 'lat', 'coords', 'index_right',
                'lon_cen', 'lat_cen', 'grid_cell'],
                )
        # print(f'gdf: {gdf}', flush=True)
        
        # Need to also return the original lat lon
        # Calculate the mean of all points within each grid cell
        counts = gdf[["grid_cell","obs"]].groupby("grid_cell").count().reset_index()
        gdf = gdf[["grid_cell", "lon_cen", "lat_cen", "obs", "lon", "lat", "rlon", "rlat"]].groupby("grid_cell").mean().reset_index()
        
        # Remove nan values
        gdf = gdf.dropna(subset="obs")
        counts = counts.dropna(subset="obs")
        
        
        # Return the data we are interested in 
        if obserr is not None:
            return gdf["lon"].values, gdf["lat"].values, gdf["rlon"].values, gdf["rlat"].values, gdf["obs"].values, gdf["obserr"].values, counts["obs"].values
        else:
            return gdf["lon"].values, gdf["lat"].values, gdf["rlon"].values, gdf["rlat"].values, gdf["obs"].values, counts["obs"].values