#!/usr/bin/env python
from datetime import datetime, timedelta
from functools import wraps

import logging
import os
import time
from tqdm import tqdm
import yaml
import multiprocessing
import gc
import dask.array as da
from glob import glob as glob_module

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt
import cartopy
import cartopy.crs as ccrs
from netCDF4 import Dataset

from .mapies import MAPIES
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 pathlib import Path
from glob import glob as glob_module
from memory_profiler import profile

os.environ["OMP_NUM_THREADS"] = "8"

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

TIME_ORIGIN = np.datetime64("1900-01-01T00:00:00")
       



def process_single_file(r, file, obs_var, lat_var, lon_var, time_var, start_date, end_date, apply_qa, qualityFlags, flag=False):
    """
    Process a single file and return aggregated results.
    """
    try:
        ds = xr.open_dataset(file, engine="h5netcdf")  # Open the file

        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 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  


def filter_file(args):
    root, filename, start_dt, end_dt = args
    try:
        parts = filename.split('.')
        date = parts[1][1:]
        hhmm = parts[2]
        file_dt = datetime.strptime(date + hhmm, "%Y%j%H%M")

        if start_dt <= file_dt <= end_dt:
            return os.path.join(root, filename), file_dt.month
    except Exception as e:
        print(f"Error processing file {filename}: {e}")
        return None

class VIIRS(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.datatype="viirs"
        frequency = kwargs.get("frequency")
        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.start_day = int(start_date[6:8])  
        self.end_day = int(end_date[6:8]) 
        self.monthly_dict = {}
        self.daily_dict = {} # Dictionary to produce daily netcdf files 

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

        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.read_config(**kwargs)


    @timeit
    def read_config(self, **kwargs):
        """
        Read yaml config file
        """
        super().read_config(**kwargs)
        logger.info("Reading VIIRS specific config variables")

        da_dict = self.config[self.datatype]["da"]
        self.obsid_da = da_dict["obsid"] #int

        unc_const = kwargs.get("unc_const")
        if unc_const:
            self.unc_const = unc_const
        else:
            self.unc_const = da_dict["uncertainty_constants"] #float, int, list, dict
        

    @timeit
    def preprocess_vars(self):
        """
        Preprocessing of the dataset
        """
        print(f'Inside preprocess_vars')
        super().preprocess_vars()


        # Lon and lat values
        lon_dims = self.ds[self.lon_var].dims
        lon_shape = self.ds[self.lon_var].shape
        lon_attrs = self.ds[self.lon_var].attrs
        self.lon_values = self.ds[self.lon_var].values.flatten()

        # Reindex lon column values
        self.lon_values = self.reindex(
            self.time_values_index, 
            self.lon_values,
            )

        lat_dims = self.ds[self.lat_var].dims
        lat_shape = self.ds[self.lat_var].shape
        lat_attrs = self.ds[self.lat_var].attrs
        self.lat_values = self.ds[self.lat_var].values.flatten()

        # Reindex aod column values
        self.lat_values = self.reindex(
            self.time_values_index, 
            self.lat_values,
            )


        # AOD column values, default "Aerosol_Optical_Thickness_550_Land_Ocean_Best_Estimate"
        aod_dims = self.ds[self.obs_var].dims

        aod_shape = self.ds[self.obs_var].shape
        aod_attrs = self.ds[self.obs_var].attrs
        self.obs = self.ds[self.obs_var].values.flatten()

        # Reindex aod column values
        self.obs = self.reindex(
            self.time_values_index, 
            self.obs,
            )

        # Apply QA
        if self.apply_qa:
            logger.info("Applying quality assurance filters")
            self.quality_assurance()
        


    @timeit
    def regridding_function(self):
        """
        Perform Rotation of Grid representation
        """
        # Calculate the grid representation
        if isinstance(self.grid, RotatedGrid):
             # Aggregate the the observations to the grid representation
            lon_agg, lat_agg, rlon_agg, rlat_agg, obs_agg, count_obs = self.grid.aggregate(self.lon_values, self.lat_values, self.obs)

        elif isinstance(self.grid, RegularGrid):
             # Aggregate the the observations to the grid representation
            lon_agg, lat_agg, obs_agg, count_obs = self.grid.aggregate(self.lon_values, self.lat_values, self.obs)
        
        else:
            raise ValueError("Invalid grid representation")
        
        self.lon_values = lon_agg
        self.lat_values = lat_agg
        self.obs = obs_agg
        self.count_obs = count_obs

    @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%j')

            filepaths = f'{self.indir}/**/AERDB_L2_VIIRS_NOAA20.A{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)}")

        files_filtered = []
        for file in files:
            parts = file.split('.')
            date = parts[1][1:]
            hhmm = parts[2]

            file_dt = datetime.strptime(date + hhmm, "%Y%j%H%M")
            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[1][1:]
            # print(f'Julian day: {julian_day}')
            date = datetime.strptime(julian_day, "%Y%j")
            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, # Scan Start 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")
            # for filepath in data['files']:
            #     print(f"  {filepath}")


    
    @timeit
    def process_data(self, monthly_avg=False, batch_size=100):
        """
        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.

        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 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) 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

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

        for key, month_data in self.monthly_dict.items():
            # Skip months with missing or invalid data
            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

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


            # Initialize final arrays on the first valid month
            if cumulative_obs is None:
                cumulative_obs = obs_values * count_values
                cumulative_lat = lat_values * count_values
                cumulative_lon = lon_values * count_values
                cumulative_count = count_values
            else:
                cumulative_obs += obs_values * count_values
                cumulative_lat += lat_values * count_values
                cumulative_lon += lon_values * count_values
                cumulative_count += count_values


        # Check if any valid data exists
        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

        # Compute final yearly averages
        valid_mask = cumulative_count > 0
        yearly_obs = np.zeros_like(cumulative_obs)
        yearly_lat = np.zeros_like(cumulative_lat)
        yearly_lon = np.zeros_like(cumulative_lon)

        yearly_obs[valid_mask] = cumulative_obs[valid_mask] / cumulative_count[valid_mask]
        yearly_lat[valid_mask] = cumulative_lat[valid_mask] / cumulative_count[valid_mask]
        yearly_lon[valid_mask] = cumulative_lon[valid_mask] / cumulative_count[valid_mask]

        # Store results
        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.")

    def process_lazy_data(self):
        """
        Process the data for the specified time range. It only retrieves a daily dictionary using dask lazy loading.
        """
        flag = True
        for day in self.daily_dict.keys():
            try:
                files = self.daily_dict[day]["files"]
                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
                    ) 
                    for file in files
                ]

                daily_obs = []
                daily_lat = []
                daily_lon = []
                daily_time = []

                with multiprocessing.Pool(processes=8) as pool:  # Adjust number of processes as needed
                    results = pool.starmap(process_single_file, args)
            
                for obs, lat, lon, time_values in results:
                    if len(obs) > 0:  # Skip empty arrays
                        daily_lon.append(lon)
                        daily_lat.append(lat)
                        daily_obs.append(obs)
                        daily_time.append(time_values)

                # Skip saving if all arrays are empty
                if not daily_obs:
                    print(f"No valid data for day {day}")
                    continue

                # Concatenate all data in one step
                final_obs = np.concatenate(daily_obs, axis=0)
                final_lon = np.concatenate(daily_lon, axis=0)
                final_lat = np.concatenate(daily_lat, axis=0)
                final_time = np.concatenate(daily_time, axis=0)

                # 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
                # print first 100 values of final_obs
                print(f"First 100 values of final_obs: {final_obs[:100]}")
                print(f"First 100 values of final_lon: {final_lon[:100]}")
                print(f"First 100 values of final_lat: {final_lat[:100]}")

                # Save directly to NetCDF
                ds = xr.Dataset(
                    coords={"time": final_time},
                    data_vars={
                        "lon": (["time"], final_lon),
                        "lat": (["time"], final_lat),
                        "obs": (["time"], final_obs),
                    },
                )
                filename = f"{self.dest}/daily_{day}.nc"
                ds.to_netcdf(filename, encoding={})
            
            except Exception as e:
                print(f"Error processing data for day {day}: {e}")
                continue


# =============================================================================
# Supporting functions
# =============================================================================
    @timeit
    def to_da(self, frequency="D", save_nc= True, save_figs=False):
        """
        Function that returns all the needed variables for the DA
        """

        # Calculate error
        self.obserr = error_estimation(self.datatype, self.obs, self.unc_const)

        outfiles = []
        
        self.frequency = frequency_int(frequency)
        da_dates_slice = pd.date_range(
            self.start_date, 
            self.end_date,
            freq=frequency
            ).strftime('%Y%m%d%H%M')

        for date in da_dates_slice:
            l_border = (datetime.strptime(date, "%Y%m%d%H%M") - timedelta(hours=self.frequency/2)).strftime('%Y%m%d%H%M')
            r_border = (datetime.strptime(date, "%Y%m%d%H%M") + timedelta(hours=self.frequency/2)).strftime('%Y%m%d%H%M')
            # Included by Guillaume in his DA function so I believe he needs this to be true
            filename = Path(self.dest).joinpath(f'obs{date}.nc')
            # Cut down the time interval again then reindex obs, obserr, lat and lon
            # Returning also arrays for type, lev, id and n with the same index as the time var
            time_values, frequency_index = time_domain_selection(
                self.time_values, 
                l_border, 
                r_border,
                )

            # Reindex
            obs = self.reindex(frequency_index, self.obs)
            lon = self.reindex(frequency_index, self.lon_values)
            lat = self.reindex(frequency_index, self.lat_values)
            obserr = self.reindex(frequency_index, self.obserr)
            # Run aggregation with grid representation
            try:
                if self.grid_repr == "rotated":
                    lon, lat, rlon, rlat, obs, obserr, count_obs = regridding_function(self.grid, lon, lat, obs, obserr)
                elif self.grid_repr == "regular":
                    lon, lat, obs, obserr, count_obs = regridding_function(self.grid, lon, lat, obs, obserr)
            except ValueError:
                raise Exception("Have you specified a grid_repr")

            

            # Create new arrays with same length as obs
            obsid = np.full(shape=obs.shape, fill_value=self.obsid_da, dtype=self.int_type)
            obstype = np.full(shape=obs.shape, fill_value=30, dtype=self.int_type)
            obslev = np.full(shape=obs.shape, fill_value=-99999, dtype=self.int_type)
            obsn = np.full(shape=obs.shape, fill_value=count_obs, dtype=self.int_type)
            time_values = np.full(shape=obs.shape, fill_value=date, dtype=str)

            # Coords equals index

            assert obs.shape == obserr.shape
            assert lon.shape == time_values.shape
            coords = dict(index=("index", np.indices(obs.shape)[0, :]))
            data_vars = dict(
                time=(["index",], time_values),
                lon=(["index",], lon), 
                lat=(["index",], lat),
                obs=(["index",], obs),
                obserr=(["index", ], obserr),
                obsid=(["index", ], obsid),
                obstype=(["index", ], obstype),
                obslev=(["index", ], obslev),
                obsn=(["index", ], obsn),
            )
            if self.grid_repr == "rotated":
                rotated_dict = dict(
                    rlon=(["index",], rlon), 
                    rlat=(["index",], rlat), #Also return rlon/rlat values if available
                )
                data_vars.update(rotated_dict)

            
            #save the plots if specified for each time frequency
            if save_figs:
                self.plot_2D_obs_custom(lon, lat, obs, outdir=self.dest, filename=f"da_obs_{date}.png")

            ds = self.to_xarray(coords=coords, data_vars=data_vars)
            if save_nc:
                logger.info(f"Outputting da data with {filename}")
                ds.to_netcdf(filename, encoding={})
            outfiles.append(filename)

        return outfiles
    
    
    
    
    
       

        

# =============================================================================
# Plotting
# =============================================================================

    @timeit
    def plot_2D_observations(self, months=[0], 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 = f"{outdir}/{self.datatype}_2D_obs_from_{self.start_date}_to_{self.end_date}_dada.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]  # Access the dictionary for this month
                    # files = month_data["files"]  # List of files for the month (not used for plotting)
                    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 {self.datatype.upper()} data from {self.start_date} to {self.end_date}"
                title = f"Observation 2D plot of month {m} from {self.start_date} to {self.end_date}"
                filename = f"{outdir}/{self.datatype}_2D_obs_month{m}_{self.year}_dada.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} from {self.start_date} to {self.end_date}"
            filename = f"{outdir}/{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], 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 = f"{outdir}/{self.datatype}_2D_obs_count_from_{self.start_date}_to_{self.end_date}_dada.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]  # Access the dictionary for this month
                    # files = month_data["files"]  # List of files for the month (not used for plotting)
                    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"Observation 2D plot of {self.datatype.upper()} data from {self.start_date} to {self.end_date}"
                title = f"Plot of the number of valid observations for month {m} from {self.start_date} to {self.end_date}"
                filename = f"{outdir}/{self.datatype}_2D_obs_count_month{m}_{self.year}_dada.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} from {self.start_date} to {self.end_date}"
            filename = 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,
            )
