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

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 mapies.util.variables_names import VARS_ATTRIBUTES
from pathlib import Path
import json
from numpy.typing import DTypeLike, NDArray

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, 
#         geospatial_crop=None
#     ):
#     """
#     Process a single file and return aggregated results or flattened arrays.
#     """
#     try:
#         ds = xr.open_dataset(file, engine="h5netcdf")  
        
#         obs, lat, lon, time_values = new_preprocess_vars(ds, obs_var, lat_var, lon_var, time_var, start_date, end_date)

#         #Run a time domain selection 
#         if apply_qa:
#             obs, lat, lon, time_values = quality_assurance(ds, obs, lat, lon, time_values, qualityFlags)
        
#         obs, lon, lat, time_values = time_domain_selection(obs, lat, lon, time_values, start_date, end_date)
        
#         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 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"
        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 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 = filter_files(files, self.start_date, self.end_date)
        
        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")


    
    @timeit
    def process_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():
                # 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
                    # valid_obs = final_obs > 0
                    # self.monthly_dict[month] = {
                    # "lon": final_lon[valid_obs],
                    # "lat": final_lat[valid_obs],
                    # "obs": final_obs[valid_obs],
                    # "count": cumulative_count[valid_obs],
                    # }
                    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.obs, self.lat, self.lon, self.count_obs = filter_observations(final_obs, final_lat, final_lon, cumulative_count)
                # valid_obs = final_obs > 0
                # self.lon = final_lon[valid_obs]
                # self.lat = final_lat[valid_obs]
                # self.obs = final_obs[valid_obs]
                # self.count_obs = cumulative_count[valid_obs]
                
                if save: 
                    try: 
                        # mid_time = self.start_date + (self.end_date - self.start_date) / 2
                        # mid_time = np.datetime64(mid_time)
                        # N = len(self.obs)
                        # mid_time_array = np.repeat(mid_time, N)
                        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_monthly_data(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...")
            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
                    ]
            valid_results, failed = process_batch(args, 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"]
                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 files
                ]

                valid_results, failed_files = process_batch(args, files)

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


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

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

            filename = Path(self.dest).joinpath(f'obs{date}.nc')
            
            # concat all days that we need
            obs = np.array([])
            lon = np.array([])
            lat = np.array([])
            time_values = np.array([], dtype='datetime64')
            days = pd.date_range(l_border, r_border, freq="D")
            for pandas_day in days:
                day = pandas_day.day
                try:
                    day_obs = self.daily_dict[day]["obs"]
                    day_lon = self.daily_dict[day]["lon"]
                    day_lat = self.daily_dict[day]["lat"]
                    day_time = self.daily_dict[day]["time"]
                except KeyError:
                    print("No data for day selected")
                    continue
                obs = np.append(obs, day_obs)
                lon = np.append(lon, day_lon)
                lat = np.append(lat, day_lat)
                time_values = np.append(time_values, day_time)

            # Perform a further time cut
            try:
                obs, lon, lat, time_values = time_domain_selection(obs, lat, lon, time_values, l_border, r_border)
            except TypeError:
                print("No data, skipping")
                continue
            # Calculate obserr
            obserr = error_estimation(self.datatype, obs, self.unc_const)


            # 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}")
                print(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, display_fig=False):

        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"
            super().plot_2D_obs_custom(
                lon_values=lon_values,
                lat_values=lat_values,
                obs_values=obs_values,
                outdir=outdir,
                title=title,
                filename=filename,
                display_fig=display_fig
            )
            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,
                    display_fig=display_fig
                )
        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,
                display_fig=display_fig,
            )
    
    def plot_2D_num_obs(self, months=[0], filename=None, outdir=None, display_fig=False):

        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"
            super().plot_2D_num_obs_custom(
                lon_values=lon_values,
                lat_values=lat_values,
                num_obs=cum_count,
                outdir=outdir,
                title=title,
                filename=filename,
                display_fig=display_fig
            )
            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,
                    display_fig=display_fig
                )
        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,
                display_fig=display_fig
            )