#!/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

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 or flattened arrays.
    """
    try:
        ds = xr.open_dataset(file, engine="h5netcdf")  

        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}')
        if self.indir:
            for date in self.dates_slice:
                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) 

            files = sorted(get_file_list(file_patterns))
            print(f"Total number of files: {len(files)}")
        else:
            raise Exception("Please specify an input directory")
       
        
        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:]
            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
                }

            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())}")
        
        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
                    )
                    
                    # 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
        """
        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...")

            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:
                    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.
        """
        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"]

            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

        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

        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]

        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):
        """
        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:  
                        daily_lon.append(lon)
                        daily_lat.append(lat)
                        daily_obs.append(obs)
                        daily_time.append(time_values)

                if not daily_obs:
                    print(f"No valid data for day {day}")
                    continue

                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

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

            try:
                # 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={})

                print(f"Saved daily data for day {day} to {filename}")
            
            except Exception as e:
                print(f"Error saving daily 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
# =============================================================================
    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,
            )
