#!/usr/bin/env python
from dataclasses import dataclass, field
from typing import List, Tuple
from datetime import datetime, timedelta

from functools import wraps

from mapies.mapies import MAPIES
from mapies.util.func_tools import timeit, get_file_list, time_converter, inflate_array
import time
import multiprocessing
import numpy as np
import pandas as pd
import xarray as xr
import os

from mapies.util.func_tools import *
#timeit, get_file_list, time_converter, frequency_int, error_estimation, time_domain_selection, geo_to_rot
from mapies.grids.monarch import RotatedGrid, IrregularRotatedGrid, RegularGrid, regridding_function
from numpy.typing import DTypeLike, NDArray


def process_single_file(
        r, 
        file, 
        obs_var, 
        lat_var, 
        lon_var, 
        time_var, 
        start_date, 
        end_date, 
        apply_qa, 
        qualityFlags, 
        flag=False, 
        geospatial_crop=None
    ):
    """
    Process a single file and return aggregated results or flattened arrays.
    """
    try:
        ds = xr.open_dataset(file, engine="h5netcdf", group="PRODUCT")

        obs, lat, lon, time_values, time_values_index = new_preprocess_vars(ds, obs_var, lat_var, lon_var, time_var, start_date, end_date)
        # obs = ds[obs_var].values.flatten()
        # lat = ds[lat_var].values.flatten()

        if apply_qa:
            obs, lat, lon, time_values = quality_assurance(ds, obs, lat, lon, time_values, qualityFlags)
        if geospatial_crop is not None:
            obs, lat, lon, time_values = spatial_domain_selection(obs, lat, lon, time_values, geospatial_crop)
            
        if not flag:
            # Regrid  and aggregate
            if isinstance(r, RotatedGrid):
                lon_agg, lat_agg, rlon_agg, rlat_agg, obs_agg, count_obs = regridding_function(r, lon, lat, obs)
            elif isinstance(r, RegularGrid):
                lon_agg, lat_agg, obs_agg, count_obs = regridding_function(r, lon, lat, obs)
            ds.close()
            return obs_agg, lat_agg, lon_agg, count_obs
        else:
            ds.close()
            return obs, lat, lon, time_values 
    
    except Exception as e:
        print(f"Error processing file {file}: {e}")
        return None

class TROPOMI(MAPIES):
    """
    Class for VIIRS specific data
    """
    
    def __init__(self, start_date, end_date,**kwargs):
        """
        Inherited init class with new variables
        """
        super().__init__(start_date, end_date)
        
        self.time_orig = np.datetime64("1993-01-01T00:00:00")
        self.dest = kwargs.get("dest")
        if not os.path.exists(self.dest):
            raise Exception("Output directory doesn't exist")
        self.indir = kwargs.get("indir")
        self.year = int(start_date[:4])  
        #self.quality_flag_limit
        self.monthly_dict = {}
        self.daily_dict = {} # Dictionary to produce daily netcdf files 
        self.datatype = "tropomi"

        self.geospatial_crop = None

        # Add quality value filter number


        if isinstance(self.start_date, str):
            self.start_date = time_converter(self.start_date)

        if isinstance(self.end_date, str):
            self.end_date = time_converter(self.end_date)

        self.dates_slice = pd.date_range(
            self.start_date, 
            self.end_date, 
            ).strftime('%Y%m%d%H%M')


        self.read_config(**kwargs)


    @timeit
    def read_nc(self):
        """
        Returns list of files needed to be processed by the application
        """
        file_patterns = []
        print(f' Date slice = {self.dates_slice}')
        for date in self.dates_slice:
            # Convert to Julian day
            date = datetime.strptime(date, '%Y%m%d%H%M').strftime('%Y%m%d')

            #TODO add something that takes in the 
            filepaths = f'{self.indir}/**/S5P_OFFL_L2__NO2____{date}*'
            file_patterns.append(filepaths)  # Add pattern to the set

        files = sorted(get_file_list(file_patterns))
        print(f"Total number of files: {len(files)}")
        

        # I'm not sure we need this part
        files_filtered = []
        for file in files:
            parts = file.split('_')
            date = parts[-6]

            file_dt = datetime.strptime(date, "%Y%m%dT%H%M%S")
            if self.start_date <= file_dt <= self.end_date:
                files_filtered.append(file)

        
        first_idx = max(0, files.index(files_filtered[0]) - 2)
        last_idx = min(len(files) - 1, files.index(files_filtered[-1]) + 2)
        self.files = files[first_idx : last_idx + 1]
        
        print(f"Total number of filtered files: {len(self.files)}")
        
        for file in self.files:
            parts = file.split('_')
            julian_day = parts[-6]

            date = datetime.strptime(julian_day, "%Y%m%dT%H%M%S")
            if date.month not in self.monthly_dict:
                self.monthly_dict[date.month] = {
                "files": [],  # List to store file paths
                "lon": None,  # Longitudes array
                "lat": None,  # Latitudes array
                "obs": None,  # Observations array
                "count": None,  # Count of observations used to compute the average
                } 
            if date.day not in self.daily_dict:
                self.daily_dict[date.day] = {
                "files": [],  # List to store file paths
                "lon": None,  # Longitudes array
                "lat": None,  # Latitudes array
                "obs": None,  # Observations array
                "time": None, # Time array
                }

            # Append the file to the list of files for the corresponding month
            self.monthly_dict[date.month]["files"].append(file)
            self.daily_dict[date.day]["files"].append(file)
            
        print(f"Discovered files for months: {list(self.monthly_dict.keys())}")
        for month, data in self.monthly_dict.items():
            print(f"  Month {month:02d}: {len(data['files'])} files")
        
        print(f"Discovered files for days: {list(self.daily_dict.keys())}")
        
        # Print the daily disctionary to see if files are stored in chronological order 
        for day, data in self.daily_dict.items():
            print(f"Day {day:02d}: {len(data['files'])} files")


    @timeit
    def process_data(self, monthly_avg=False, batch_size=100, geospatial_crop: NDArray | None = None):
        """
        Process the data for the specified year and months.

        Parameters:
        - monthly_avg: bool, whether to compute single monthly averages.
        - batch_size: int, the number of files to process in each batch.
        - chunk_size: int, the size of chunks for processing within each batch.
        - geospatial_crop, 2D numpy array

        Returns:
        - if monthly_avg --> lon_final, lat_final, obs_final: Aggregated longitude, latitude, and observation arrays.
        - if not monthly_avg --> preprocess_vars() is called.
        """
        if geospatial_crop is not None:
            self.geospatial_crop = np.array(geospatial_crop)
        if monthly_avg:
            for month in self.monthly_dict.keys():
                # print(f'Month: {self.monthly_dict[month]}')
                # 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_monthly_data(
                     batch_size=batch_size, month=month
                    )
                    # final_lon, final_lat, final_obs, cumulative_count = self.process_files_in_blocks(month, block_size=300, batch_size=batch_size)
                    
                    # Update the monthly dictionary
                    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_monthly_data(
                    batch_size=batch_size, month=None
                )
                self.lon_values = final_lon
                self.lat_values = final_lat
                self.obs = final_obs
                self.count_obs = cumulative_count

            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_monthly_data(self, batch_size, month=None):

        """
        Function that returns all the needed variables for the DA
        """
        # r = RotatedGrid(centre_lon=20, centre_lat=35, dlon=.1, dlat=.1, west=-51, south=-35)
        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 = [
            month_files[i:i + batch_size] for i in range(0, len(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...")
            print(f"Batch: {batch}")

            # Use multiprocessing to process files within the batch
            args = [(self.grid, file, self.obs_var, self.lat_var, self.lon_var, self.time_var, self.start_date, self.end_date, self.apply_qa, self.qualityFlags, flag, self.geospatial_crop) for file in batch]

            with multiprocessing.Pool(processes=8) as pool:  # Adjust number of processes as needed
                results = pool.starmap(process_single_file, args)

            # Filter valid results and log failed files
            valid_results = [res for res in results if res is not None]
            failed_files.extend([file for file, res in zip(batch, results) if res is None])

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

            # Aggregate results incrementally
            for obs_agg, lat_agg, lon_agg, count_obs in valid_results:
                if cumulative_obs is None:
                    # Initialize cumulative arrays
                    cumulative_obs = obs_agg * count_obs
                    cumulative_lat = lat_agg * count_obs
                    cumulative_lon = lon_agg * count_obs
                    cumulative_count = count_obs
                else:
                    cumulative_obs += obs_agg * count_obs
                    cumulative_lat += lat_agg * count_obs
                    cumulative_lon += lon_agg * count_obs
                    cumulative_count += count_obs

        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
        valid_mask = cumulative_count > 0
        final_obs = np.zeros_like(cumulative_obs)
        final_lat = np.zeros_like(cumulative_lat)
        final_lon = np.zeros_like(cumulative_lon)

        final_obs[valid_mask] = cumulative_obs[valid_mask] / cumulative_count[valid_mask]
        final_lat[valid_mask] = cumulative_lat[valid_mask] / cumulative_count[valid_mask]
        final_lon[valid_mask] = cumulative_lon[valid_mask] / cumulative_count[valid_mask]

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



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

        if months == [0]:
            try:
                lon_values = self.lon_values
                lat_values = self.lat_values
                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"
            super().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"
            super().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_values
                lat_values = self.lat_values
                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"
            super().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"   
                super().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"
            super().plot_2D_num_obs_custom(
                lon_values=lon_values,
                lat_values=lat_values,
                num_obs=cum_count,
                outdir=outdir,
                title=title,
                filename=filename,
            )
