#!/usr/bin/env python
# MAPIES base

from functools import wraps
from datetime import datetime, timedelta
import sys
from mapies.util.func_tools import timeit, time_domain_selection, frequency_int, error_estimation
from mapies.grids.monarch import RotatedGrid, IrregularRotatedGrid, RegularGrid
import time
import logging
import numpy as np
import pandas as pd
import xarray as xr
#import matplotlib
#matplotlib.use("TkAgg") # Use tinker to perform plt.show() as it is GUI supported 
import matplotlib.pyplot as plt
import cartopy
import cartopy.crs as ccrs
import os
import yaml
import matplotlib.cm as cm
from matplotlib.colors import ListedColormap
import matplotlib.colors as mcolors

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

class MAPIES:
    """
    Base class for only base functions used in every Mapies run
    
    For example Reading and writing functions
    """
    def __init__(self, start_date, end_date, **kwargs):
        """
        INITIATE the MAPIES class with the variables used in every run
        """
        self.np_typetime = "datetime64[s]"
        self.int_type = 'int64'
        self.time_orig = np.datetime64("1900-01-01T00:00:00")
        self.start_date = start_date
        self.end_date = end_date


        self.grid_dict = {
            "rotated":RotatedGrid, 
            "regular":RegularGrid, 
            "ireg_rotated":IrregularRotatedGrid
        }


    def read_config(self, **kwargs):
        """
        Read yaml config file
        """
        module_dir = os.path.dirname(__file__)


        # Let the user read in their own config
        config_file = kwargs.get("config_file")
        if config_file:
            try:
                self.config = yaml.safe_load(open(config_file))
            except FileNotFoundError:
                logger.error("Config file passed does not exist: Using default config instead")
                self.config = yaml.safe_load(open(os.path.join(module_dir, "config/satellite_config.yaml")))
        else:
            self.config = yaml.safe_load(open(os.path.join(module_dir, "config/satellite_config.yaml")))


        variable_dict = self.config[self.datatype]["variables"]
        self.time_var = variable_dict["time_variable"]
        self.lon_var = variable_dict["lon_variable"]
        self.lat_var = variable_dict["lat_variable"]

        # If obs_var passed then run data analysis on that variable if not pull default
        obs_var = kwargs.get("obs_var")
        if obs_var:
            self.obs_var = obs_var
        else:
            self.obs_var = variable_dict["obs_variable"]
        

        grid_dict = self.config[self.datatype]["grids"]
        grid_repr = kwargs.get("grid_repr")
        if grid_repr:
            self.grid_repr = grid_repr
        else:
            self.grid_repr = grid_dict["grid_repr"] #str

        if self.grid_repr:
            # Cread grid object if grid repr is selected in the input
            self.grid_config = grid_dict[self.grid_repr]
            self.grid = self.grid_dict[self.grid_repr](**self.grid_config)
            print(self.grid)
        

        qa_dict = self.config[self.datatype]["qa"]
        #Check if the user wants qa flags
        apply_qa = kwargs.get("apply_qa")

        if apply_qa is None:
            self.apply_qa = True
        else:
            self.apply_qa = apply_qa
        self.qualityFlags = qa_dict["qa_variables"]



    def preprocess_vars(self):
        """
        Preprocessing of the dataset
        """
        # Get all info about time columns
        self.time_dims = self.ds[self.time_var].dims
        self.time_shape = self.ds[self.time_var].shape
        self.time_attrs = self.ds[self.time_var].attrs # Useful if we need to convert to datetime
        

        # Get time values flatten and convert to datetime values
        self.time_values = self.ds[self.time_var].values
        self.time_values = self.time_values.flatten()
        
        if self.time_values.dtype == "timedelta64[ns]":
            logger.info("Adding time origin to time values as the time variable is in timedelta")
            self.time_values = np.add(self.time_orig, self.time_values)
            self.time_values = pd.to_datetime(self.time_values)
        else:
            self.time_values = pd.to_datetime(self.time_values)
        
        
        # Time domain selection
        self.time_values, self.time_values_index = time_domain_selection(self.time_values, self.start_date, self.end_date)



    @timeit
    def plot_2D_obs(self, **kwargs):
        """
        Plotting the observations
        """
        # Regrid if specified as argument
        if self.grid_repr:
            logger.info("Performing regridding of data")
            self.regridding_function()
        # TODO: Make these arguments
        figsize = (15,10) 
        markersize = 2.5

        # Set Basemap projection
        proj = ccrs.PlateCarree()

        # Create the plot and add features, TODO make this adjustable in arguments
        fig, ax = plt.subplots(subplot_kw={"projection": proj}, figsize=figsize)
        ax.gridlines()
        #ax.add_feature(cartopy.feature.BORDERS, linestyle=':', alpha=1)
        #ax.add_feature(cartopy.feature.OCEAN,facecolor=("lightblue"))
        #ax.add_feature(cartopy.feature.LAND)
        ax.coastlines(resolution='10m')

        

        print("Plotting observations")
        

        x, y = self.lon_values, self.lat_values

        im = ax.scatter(x,y,markersize,c=self.obs, transform=proj)
        
        fig.colorbar(im, ax=ax)
        
        ax.set_title(f'Observation 2D plot of {self.datatype.upper()} data from {self.start_date} to {self.end_date}')
        print("Saving Figure")
        plt.savefig(f"{self.dest}{self.datatype}_2D_obs.png", format="png")
        plt.close(fig)


    @timeit
    def plot_2D_obs_custom(self, lon_values=None, lat_values=None, obs_values=None, outdir="/home/cgile/Documents/mapies/mapies/figures", title=None, filename=None):
        """
        General method for plotting observations with CB-RdYlBu colormap and arrow-like ends for out-of-bound values.
        """

        if lon_values is None:
            lon_values = self.lon_values
        if lat_values is None:
            lat_values = self.lat_values
        if obs_values is None:
            obs_values = self.obs


        # Set Basemap projection
        figsize = (15, 10)
        markersize = 2.5
        proj = ccrs.PlateCarree()

        # Define discrete levels for the color bar
        levels = [0.02, 0.04, 0.06, 0.08, 0.10, 0.15, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 1.00]

        # Define the custom color list
        colors = [
            '#290AD8', '#2645FA', '#3A91FF', '#66CBFF', '#99EEFF', 
            '#CDFCFF', '#F2FFD9', '#FFEEAA', '#FFCA88', '#FC966A', 
            '#ED5750', '#CB1C2D', '#9A001E', '#68001D', '#690067'
        ]

        # Create the custom colormap
        cmap = ListedColormap(colors)
        norm = mcolors.BoundaryNorm(boundaries=levels, ncolors=len(colors), extend="both")

        # Create the plot and add features
        fig, ax = plt.subplots(subplot_kw={"projection": proj}, figsize=figsize)
        ax.gridlines()
        ax.coastlines(resolution="10m")

        print("Plotting observations (new method)")

        # Scatter plot with discrete color mapping
        im = ax.scatter(
            lon_values, lat_values, markersize, c=obs_values, cmap=cmap, norm=norm, transform=proj
        )

        # Add the discrete color bar with arrow-like ends
        cbar = fig.colorbar(im, ax=ax, orientation="vertical", pad=0.05, ticks=levels, extend="both")
        cbar.set_label("AOD")
        cbar.ax.set_yticklabels([f"{level:.2f}" for level in levels])  # Format tick labels

        # Add title and save the plot
        ax.set_title(title or "Observation 2D Plot")

        # Ensure the output directory exists
        os.makedirs(outdir, exist_ok=True)

        filepath = os.path.join(outdir, filename or "observation_2D_plot.png")
        plt.savefig(filepath, format="png")
        plt.close(fig)
        print(f"Saved plot: {filepath}")

    @timeit
    def plot_2D_cumulative_count(self, lon_values, lat_values, cumulative_count, outdir="/home/cgile/Documents/mapies/mapies/figures", title=None, filename=None):
        """
        Plot the cumulative count of observations contributing to the final averages.

        Parameters:
        - lon_values: 1D array of longitude values.
        - lat_values: 1D array of latitude values.
        - cumulative_count: 1D array of the count of valid observations per grid cell.
        - outdir: Directory to save the plot.
        - title: Title for the plot.
        - filename: Name of the output file.
        """

        # Set Basemap projection
        figsize = (15, 10)
        markersize = 2.5
        proj = ccrs.PlateCarree()

        # Use continuous colormap
        cmap = plt.cm.YlOrBr  # Colormap for cumulative count
        norm = mcolors.Normalize(vmin=0, vmax=np.max(cumulative_count))  # Continuous normalization

        # Create the plot and add features
        fig, ax = plt.subplots(subplot_kw={"projection": proj}, figsize=figsize)
        ax.gridlines()
        ax.coastlines(resolution="10m")

        print("Plotting cumulative count")

        # Scatter plot for cumulative count
        im = ax.scatter(
            lon_values, lat_values, markersize, c=cumulative_count, cmap=cmap, norm=norm, transform=proj
        )

        # Add the continuous color bar
        cbar = fig.colorbar(im, ax=ax, orientation="vertical", pad=0.05)
        cbar.set_label("Cumulative Count")

        # Add title and save the plot
        ax.set_title(title or "Cumulative Count Plot")

        # Ensure the output directory exists
        os.makedirs(outdir, exist_ok=True)

        filepath = os.path.join(outdir, filename or "cumulative_count_plot.png")
        plt.savefig(filepath, format="png")
        plt.close(fig)
        print(f"Saved plot: {filepath}")



    @staticmethod
    def reindex(
        dependent_var_index,
        independent_var_values,
        ):
        """"
        Recutting the data whenever a selction of the data has been made along one of the dimenisons
        
        Based off of how it has been done in Providentia
        """
        # Maybe add a checker try/except or assert
        independent_var_values = independent_var_values[dependent_var_index]

        # return reindexed values
        return independent_var_values

    def quality_assurance(self):
        """
        Applying quality assurance flags to mapies data
        """
        mask = np.zeros_like(self.obs, dtype=bool)
        for i, (flag, info) in enumerate(self.qualityFlags.items()):
            
            flag_mask = np.zeros_like(self.obs, dtype=bool)
            arr = self.ds[info["qa_name"]].values.flatten()
            arr = self.reindex(self.time_values_index, arr)
            

            for value in info["qa_values"]:
                flag_mask |= (arr == value)
            if i == 0:
                mask = flag_mask
            else:
                if info["qa_combined"]:
                    mask &= flag_mask
                else:
                    mask |= flag_mask

        self.obs = self.obs[mask]
        self.lon_values = self.lon_values[mask]
        self.lat_values = self.lat_values[mask]
        self.time_values = self.time_values[mask]


    @staticmethod
    def to_xarray(coords:dict, data_vars:dict, **kwargs):
        """
        Method to convert numpy arrays to xarray opject
        """
        attrs = kwargs.get("attrs")
        if attrs is None:
            attrs=dict()

        ds = xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs)
        return ds


    @timeit
    def plot_2D_obs_custom(self, lon_values=None, lat_values=None, obs_values=None, outdir="./", title=None, filename=None):
        """
        General method for plotting observations with CB-RdYlBu colormap and arrow-like ends for out-of-bound values.
        """
        if lon_values is None:
            lon_values = self.lon_values
        if lat_values is None:
            lat_values = self.lat_values
        if obs_values is None:
            obs_values = self.obs

        # Set Basemap projection
        figsize = (15, 10)
        markersize = 2.5
        proj = ccrs.PlateCarree()

        # Define discrete levels for the color bar
        levels = [0.02, 0.04, 0.06, 0.08, 0.10, 0.15, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 1.00]

        # Define the custom color list
        colors = [
            '#290AD8', '#2645FA', '#3A91FF', '#66CBFF', '#99EEFF', 
            '#CDFCFF', '#F2FFD9', '#FFEEAA', '#FFCA88', '#FC966A', 
            '#ED5750', '#CB1C2D', '#9A001E', '#68001D', '#690067'
        ]

        # Create the custom colormap
        cmap = ListedColormap(colors)
        norm = mcolors.BoundaryNorm(boundaries=levels, ncolors=len(colors), extend="both")

        # Create the plot and add features
        fig, ax = plt.subplots(subplot_kw={"projection": proj}, figsize=figsize)
        ax.gridlines()
        ax.coastlines(resolution="10m")

        # Scatter plot with discrete color mapping
        im = ax.scatter(
            lon_values, lat_values, markersize, c=obs_values, cmap=cmap, norm=norm, transform=proj
        )

        # Add the discrete color bar with arrow-like ends
        cbar = fig.colorbar(im, ax=ax, orientation="vertical", pad=0.05, ticks=levels, extend="both")
        cbar.set_label("AOD")
        cbar.ax.set_yticklabels([f"{level:.2f}" for level in levels])  # Format tick labels

        # Add title and save the plot
        ax.set_title(title or "Observation 2D Plot")

        # Ensure the output directory exists
        os.makedirs(outdir, exist_ok=True)

        filepath = os.path.join(self.dest, filename or "observation_2D_plot.png")
        plt.savefig(filepath, format="png")
        plt.close(fig)
        print(f"Saved plot: {filepath}")

    @timeit 
    def plot_2D_num_obs_custom(
        self, lon_values=None, 
        lat_values=None, 
        num_obs=None, 
        outdir="./", 
        title=None, 
        filename=None,
        ):
        """
        General method for plotting the number of observations contributing to the final averages.
        """

        # Set Basemap projection
        figsize = (15, 10)
        markersize = 2.5
        proj = ccrs.PlateCarree()

        # Define discrete levels for the color bar
        unique_values = np.unique(num_obs)  # Unique values in num_obs
        levels = np.linspace(unique_values.min(), unique_values.max(), unique_values.max()-unique_values.min()+1)


        # Use a discrete colormap

        cmap = plt.cm.rainbow
        norm = mcolors.Normalize(vmin=unique_values.min(), vmax=unique_values.max())
        # norm = mcolors.BoundaryNorm(np.arange(unique_values.min()-0.5, unique_values.max()+1, 1), cmap.N)

        # Create the plot and add features
        fig, ax = plt.subplots(subplot_kw={"projection": proj}, figsize=figsize)
        ax.gridlines(draw_labels=True)
        ax.coastlines(resolution="10m")

        # Scatter plot for the number of observations
        im = ax.scatter(
            lon_values, lat_values, markersize, c=num_obs, cmap=cmap, norm=norm, transform=proj
        )

        # Add the discrete colorbar without arrow-like ends
        cbar = fig.colorbar(im, ax=ax, orientation="vertical")
        cbar.set_label("Number of Observations")
        # cbar.ax.set_yticklabels([f"{level}" for level in range(unique_values)]) 

        # Add title and save the plot
        ax.set_title(title or "Number of Observations Plot")

        # Ensure the output directory exists
        os.makedirs(outdir, exist_ok=True)

        filepath = os.path.join(self.dest, filename or "num_observation_2D_plot.png")
        plt.savefig(filepath, format="png")
        plt.close(fig)
        print(f"Saved plot: {filepath}")


    
