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

from functools import wraps
from datetime import datetime, timedelta
import sys
from mapies.util.func_tools import *
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.pyplot as plt
import cartopy
import cartopy.crs as ccrs
import os
import yaml
import json
import matplotlib.cm as cm
from matplotlib.colors import ListedColormap
import matplotlib.colors as mcolors

from mapies.util.variables_names import VARS_ATTRIBUTES
from numpy.typing import DTypeLike, NDArray

#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.start_date = start_date
        self.end_date = end_date
        self.apply_qa = True
        self.geospatial_crop = None

        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"]
        
        self.qualityFlags = qa_dict["qa_variables"]


    @timeit
    def process_avg_data(self, monthly_avg=False, batch_size=100, save=False, apply_qa=True, geospatial_crop: NDArray | None = None):

        """
        Process the data.

        Parameters:geospatial_crop: NDArray | None = None
        - monthly_avg: bool, whether to compute single monthly averages.
        - batch_size: int, the number of files to process in each batch.

        Returns:
        - if monthly_avg --> updated monthly dictorionary for the processed month.
        - if not monthly_avg --> updated self.lon_values, self.lat_values, self.obs, self.count_obs.
        """

        #Check if the user wants qa flags
        if apply_qa:
            self.apply_qa = True
        else:
            self.apply_qa = False
        if geospatial_crop is not None:
            self.geospatial_crop = np.array(geospatial_crop)
        if monthly_avg:
            for month in self.monthly_dict.keys():

                # Catch if they haven't selected a grid
                try:
                    print(f"Running average on {self.grid}")
                except AttributeError:
                    raise Exception("Please specify a grid in the initialisation of MAPIES, using grid_repr=''")
                try:
                    # Process the data for the current month
                    final_lon, final_lat, final_obs, cumulative_count = self.process_data_mp(
                     batch_size=batch_size, month=month
                    )
                    
                    self.monthly_dict[month] = {
                        "lon": final_lon,
                        "lat": final_lat,
                        "obs": final_obs,
                        "count": cumulative_count,
                    }
                    print(f'Updated monthly dictionary for month {month}')
                    print(f'Lon shape: {final_lon.shape}, Lat shape: {final_lat.shape}, Obs shape: {final_obs.shape}')
                    if np.all(final_lon == 0) or np.all(np.isnan(final_lon)):
                        print(f"All longitude values for month {month} are 0 or NaN.")
                    if np.all(final_lat == 0) or np.all(np.isnan(final_lat)):
                        print(f"All latitude values for month {month} are 0 or NaN.")
                    if np.all(final_obs == 0) or np.all(np.isnan(final_obs)):
                        print(f"All observation values for month {month} are 0 or NaN.")
                
                except Exception as e:
                    print(f"Error processing data for month {month}: {e}")
                    continue
        else:
            try:
                final_lon, final_lat, final_obs, cumulative_count = self.process_data_mp(
                    batch_size=batch_size, month=None
                )

                self.obs, self.lat, self.lon, self.count_obs = filter_observations(final_obs, final_lat, final_lon, cumulative_count)

                
                if save: 
                    try: 
                        mid_time_array = generate_mid_time_array(self.start_date, self.end_date, self.obs)
                        
                        self.to_netCDF(self.obs, self.lat, self.lon, mid_time_array, self.start_date, self.end_date)
                        print(f'Saved processed data to netCDF')
                    except Exception as e:
                        print(f"Error saving data into netCDF: {e}")
                        return
                    
            except Exception as e:
                print(f"Error processing data for the time period from {self.start_date} to {self.end_date}: {e}")
                return


    def process_data_mp(self, batch_size, month=None):

        """
        Process the data with multiprocessing.
        """
        flag = False
        if month is None:
            month_files = self.files
            print(f'Processing all files from {self.start_date} to {self.end_date}')
        else:
            month_files = self.monthly_dict[month]["files"]
            print(f"Processing {len(month_files)} files for month {month}")


        # Split files into batches
        batches = create_batches(month_files, batch_size)
     
        # Initialize cumulative arrays
        cumulative_obs = None
        cumulative_lat = None
        cumulative_lon = None
        cumulative_count = None
        failed_files = []

        for batch in batches:
            print(f"Processing a batch of {len(batch)} files...")

            # Use multiprocessing to process files within the batch
            
            viirs_mp = self.mp_type(
                self.grid,
                self.obs_var,
                self.lat_var,
                self.lon_var,
                self.time_var,
                self.start_date,
                self.end_date,
                self.apply_qa,
                self.qualityFlags,
                batch,
                self.geospatial_crop,
                flag,
            )
            valid_results, failed = viirs_mp.process_batch()
            failed_files.extend(failed)

            if not valid_results:
                print("All files in this batch failed. Skipping...")
                continue

            # Aggregate results incrementally
            cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count = combine_arrays(valid_results, cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count)
        
        if failed_files:
            print(f"The following files failed to process: {failed_files}")

        # Check if any valid data was processed
        if cumulative_obs is None or cumulative_count is None:
            print(f"No valid data processed for month {month}. Returning empty arrays.")
            return None, None, None, None

        # Compute final averages
        final_obs, final_lat, final_lon = compute_final_averages(cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count)

        print(f"Completed processing for month {month}.")
        
        return final_lon, final_lat, final_obs, cumulative_count


    @timeit
    def yearly_average(self):
        """
        Compute the yearly average for the processed data.
        """
        cumulative_obs = None
        cumulative_lat = None
        cumulative_lon = None
        cumulative_count = None

        for key, month_data in self.monthly_dict.items():
            if not all(var in month_data for var in ["lon", "lat", "obs", "count"]):
                print(f"Skipping month {key}: Missing data.")
                continue
            if any(month_data[var] is None for var in ["lon", "lat", "obs", "count"]):
                print(f"Skipping month {key}: No valid data.")
                continue

            lon_values = month_data["lon"]
            lat_values = month_data["lat"]
            obs_values = month_data["obs"]
            count_values = month_data["count"]

            results = [(obs_values, lat_values, lon_values, count_values)]
            cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count = combine_arrays(results, cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count)

        if cumulative_obs is None or cumulative_count is None or np.all(cumulative_count == 0):
            print("No valid data for the entire year. Returning empty arrays.")
            self.yearly_lon = None
            self.yearly_lat = None
            self.yearly_obs = None
            self.yearly_count = None
            return

        yearly_obs, yearly_lat, yearly_lon = compute_final_averages(cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count)

        self.yearly_lon = yearly_lon
        self.yearly_lat = yearly_lat
        self.yearly_obs = yearly_obs
        self.yearly_count = cumulative_count

        print("Yearly average computation completed.")


    @timeit
    def process_lazy_data(self, apply_qa=True, geospatial_crop: NDArray | None = None, save=True):
        """
        Process the data for the specified time range. It only retrieves a daily dictionary using dask lazy loading.
        """
        flag = True
        #Check if the user wants qa flags
        if apply_qa:
            self.apply_qa = True
        else:
            self.apply_qa = False
        if geospatial_crop is not None:
            self.geospatial_crop = np.array(geospatial_crop)
        for day in self.daily_dict.keys():
            try:
                files = self.daily_dict[day]["files"]

                # Use multiprocessing to process files within the batch
                
                viirs_mp = self.mp_type(
                    self.grid,
                    self.obs_var,
                    self.lat_var,
                    self.lon_var,
                    self.time_var,
                    self.start_date,
                    self.end_date,
                    self.apply_qa,
                    self.qualityFlags,
                    files,
                    self.geospatial_crop,
                    flag,
                )
                valid_results, failed_files = viirs_mp.process_batch()

                final_obs, final_lat, final_lon, final_time = append_and_concatenate(valid_results, day)

                # Skip the day if no valid data was processed
                if final_obs is None:
                    continue 

                # Check if any valid data exists
                if np.all(final_obs == 0) or np.all(np.isnan(final_obs)):
                    print(f"All observation values for day {day} are 0 or NaN. Skipping...")
                    continue
                
                # Update day dict 
                self.daily_dict[day]["obs"] = final_obs 
                self.daily_dict[day]["lon"] = final_lon 
                self.daily_dict[day]["lat"] = final_lat
                self.daily_dict[day]["time"] = final_time 
                
                # Get day start and day end but only if they are within object start and end date
                day_start, day_end = day_boundaries(day, self.start_date, self.end_date)
                if save:
                    self.to_netCDF(final_obs, final_lat, final_lon, final_time, day_start, day_end)

            except Exception as e:
                print(f"Error processing data for day {day}: {e}")
                continue


    def to_netCDF(self, obs, lat, lon, time, start_time, end_time):
        try:
            # Convert time to seconds since TIME_ORIGIN
            final_time_seconds = (time - self.time_orig) / np.timedelta64(1, "s")
            final_time_seconds = np.round(final_time_seconds, 0)

            grid_representation_str = json.dumps(self.grid_config, ensure_ascii=False)
            qa_flags_str = json.dumps(self.qualityFlags, ensure_ascii=False)
            
            # Create xarray Dataset
            ds = xr.Dataset(
                coords={"time": ("time", final_time_seconds, {
                    "units": "seconds since 1900-01-01 00:00:00",
                    "calendar": "gregorian"
                })},
                
                data_vars={
                    "lon": (["time"], lon, VARS_ATTRIBUTES['lon'].copy()),
                    "lat": (["time"], lat, VARS_ATTRIBUTES['lat'].copy()),
                    VARS_ATTRIBUTES[self.obs_var]["mapies_variable"]: (["time"], obs, VARS_ATTRIBUTES[self.obs_var].copy()),
                },
                attrs={
                    "title": f"{VARS_ATTRIBUTES['title']['description']} from {start_time} to {end_time}",
                    "institution": "Barcelona Supercomputing Center",
                    "grid": f"{self.grid_repr}: {grid_representation_str}",
                    "Developers": "MAPIES team",
                    "QA": f"Quality Assurance: {self.apply_qa}, quality assurance applied: {qa_flags_str}",
                    "history": f"Created {datetime.now()}",
                },
            )
            
            # Define encoding for compact storage
            encoding = {
                "time": {"dtype": "float32"},  
                "lon": {"dtype": "float32"},  
                "lat": {"dtype": "float32"},  
                VARS_ATTRIBUTES[self.obs_var]["mapies_variable"]: {"dtype": "float32"},  
            }

            # Save to NetCDF file
            filename = f"{self.dest}/Data_{self.datatype}_{start_time}_{end_time}.nc"
            ds.to_netcdf(filename, mode='w', encoding=encoding)

            print(f"Saved data to {filename}")

        except Exception as e:
            print(f"Error saving data into netCDF: {e}")
            return


# =============================================================================
# Plotting
# =============================================================================
    def plot_2D_observations(self, months=[0], filename=None, outdir=None):

        if months == [0]:
            try:
                lon_values = self.lon
                lat_values = self.lat
                obs_values = self.obs
            except Exception as e:
                print(f"Error plotting all data: {e}")
                return
            title = f"Observation 2D plot from {self.start_date} to {self.end_date}"
            filename = filename or f"{self.datatype}_2D_obs_from_{self.start_date}_to_{self.end_date}.png"
            self.plot_2D_obs_custom(
                lon_values=lon_values,
                lat_values=lat_values,
                obs_values=obs_values,
                outdir=outdir,
                title=title,
                filename=filename,
            )
            return
        elif months is not None:
            for m in months:
                try: 
                    month_data = self.monthly_dict[m]  
                    lon_values = month_data["lon"]
                    lat_values = month_data["lat"]
                    obs_values = month_data["obs"]
                except Exception as e:
                    print(f"Error for month {m}: {e}")
                    continue
                title = f"Observation 2D plot of month {m} from {self.start_date} to {self.end_date}"
                filename = filename or f"{self.datatype}_2D_obs_month_{m}_year_{self.year}.png"   
                super().plot_2D_obs_custom(
                    lon_values=lon_values,
                    lat_values=lat_values,
                    obs_values=obs_values,
                    outdir=outdir,
                    title=title,
                    filename=filename,
                )
        else: 
            try: 
                lon_values = self.yearly_lon
                lat_values = self.yearly_lat 
                obs_values = self.yearly_obs
            except Exception as e:
                print(f"Error for yearly data: {e}")
                return
            title = f"Observation 2D plot of yearly average of year {self.year}"
            filename = filename or f"{self.datatype}_2D_yearly_obs_{self.year}.png"
            self.plot_2D_obs_custom(
                lon_values=lon_values,
                lat_values=lat_values,
                obs_values=obs_values,
                outdir=outdir,
                title=title,
                filename=filename,
            )
    
    def plot_2D_num_obs(self, months=[0], filename=None, outdir=None):

        if months == [0]:
            try:
                lon_values = self.lon
                lat_values = self.lat
                cum_count = self.count_obs
            except Exception as e:
                print(f"Error plotting all data: {e}")
                return
            title = f"Plot of the number of valid observations from {self.start_date} to {self.end_date}"
            filename = filename or f"{self.datatype}_2D_obs_count_from_{self.start_date}_to_{self.end_date}.png"
            self.plot_2D_num_obs_custom(
                lon_values=lon_values,
                lat_values=lat_values,
                num_obs=cum_count,
                outdir=outdir,
                title=title,
                filename=filename,
            )
            return
        elif months is not None:
            for m in months:
                try: 
                    month_data = self.monthly_dict[m]  
                    lon_values = month_data["lon"]
                    lat_values = month_data["lat"]
                    cum_count = month_data["count"]
                except Exception as e:
                    print(f"Error for month {m}: {e}")
                    continue
                title = f"Plot of the number of valid observations for month {m} from {self.start_date} to {self.end_date}"
                filename = filename or f"{self.datatype}_2D_obs_count_month_{m}_year_{self.year}.png"   
                self.plot_2D_num_obs_custom(
                    lon_values=lon_values,
                    lat_values=lat_values,
                    num_obs=cum_count,
                    outdir=outdir,
                    title=title,
                    filename=filename,
                )
        else: 
            try: 
                lon_values = self.yearly_lon
                lat_values = self.yearly_lat 
                cum_count = self.yearly_count
            except Exception as e:
                print(f"Error for yearly data: {e}")
                return
            title = f"Plot of the number of valid observations for yearly average of {self.year}"
            filename = filename or f"{outdir}/{self.datatype}_2D_obs_count_yearly_{self.year}.png"
            self.plot_2D_num_obs_custom(
                lon_values=lon_values,
                lat_values=lat_values,
                num_obs=cum_count,
                outdir=outdir,
                title=title,
                filename=filename,
            )
    

    # TODO this can be moved to a plotting function python file
    @timeit
    def plot_2D_obs_custom(self, lon_values=None, lat_values=None, obs_values=None, outdir=None, title=None, filename=None, display_fig=False):
        """
        General method for plotting observations with CB-RdYlBu colormap and arrow-like ends for out-of-bound values.
        """

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

        # Define discrete levels for the color bar
        if self.datatype == "viirs":
            label = "AOD"
            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]
            level_format = [f"{level:.2f}" for level in levels]
        elif self.datatype == "tropomi":
            #TODO remove hard coding
            label = "NO2"
            levels = [0.00001, 0.00002, 0.00003, 0.00004, 0.00005, 0.00006, 0.00007, 0.00008, 0.00009, 0.00010, 0.00015, 0.00020, 0.00030, 0.00040]
            level_format = [f"{level:.5f}" for level in levels]
        # 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 color bar
        cbar = fig.colorbar(im, ax=ax, orientation="vertical", pad=0.05, ticks=levels, extend="both")
        cbar.set_label(label)
        cbar.ax.set_yticklabels(level_format)  # Format tick labels

        # Add title and save the plot
        ax.set_title(title)

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

        filepath = os.path.join(outdir, filename)
        plt.savefig(filepath, format="png")
        if display_fig:
            plt.show()
        print(f"Saved plot: {filepath}")
        

    # TODO this can be moved to a plotting function python file
    @timeit 
    def plot_2D_num_obs_custom(
        self, lon_values=None, 
        lat_values=None, 
        num_obs=None, 
        outdir="./",
        title=None, 
        filename=None,
        display_fig=False
        ):
        """
        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")
        if display_fig:
            plt.show()
        print(f"Saved plot: {filepath}")
        
        #return fig


    @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