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

from functools import wraps

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

from mapies.util.func_tools import *
#timeit, get_file_list, time_converter, frequency_int, error_estimation, time_domain_selection, geo_to_rot
from mapies.grids.monarch import RotatedGrid, IrregularRotatedGrid, RegularGrid, regridding_function
from mapies.util.multiprocessing import TROPOMIMultiprocessing
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)




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


        # Add quality value filter number


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

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

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


        self.read_config(**kwargs)


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

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

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

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

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

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

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

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