import time
import logging
import numpy as np
import pandas as pd
import xarray as xr

from functools import wraps
from numpy import cos, sin, arctan, pi, nan
from glob import glob
from typing import List
from datetime import datetime
from abc import ABC, abstractmethod  
import multiprocessing
from dataclasses import dataclass
from numpy.typing import DTypeLike, NDArray

from mapies.util.func_tools import *
from mapies.grids.monarch import Grid, RotatedGrid, IrregularRotatedGrid, RegularGrid, regridding_function


def process_single_file_tropomi(
        r, 
        file, 
        obs_var, 
        lat_var, 
        lon_var, 
        time_var, 
        start_date, 
        end_date, 
        apply_qa, 
        qualityFlags, 
        flag=False, 
        geospatial_crop=None
    ):
    """
    Process a single file and return aggregated results or flattened arrays.
    """
    try:
        ds = xr.open_dataset(file, engine="h5netcdf", group="PRODUCT")
    except Exception as e:
        print(f"Reading file {file}: {e}")
        return None

    # Flatten arrays
    try:
        obs, lat, lon, time_values = new_preprocess_vars(ds, obs_var, lat_var, lon_var, time_var, start_date, end_date)
    except Exception as e:
        print(f"Error processing file {file}: {e}")
        return None
    
    # Apply quality assurance if applicable
    if apply_qa:
        try:
            obs, lat, lon, time_values = quality_assurance(ds, obs, lat, lon, time_values, qualityFlags)
        except Exception as e:
            print(f"Error applying QA Flags for file {file}: {e}")
            return None
    # Perform time domain selection beteen start and end date
    try:
        obs, lon, lat, time_values = time_domain_selection(obs, lat, lon, time_values, start_date, end_date)
    except Exception as e:
        print(f"Error performing time domain selection for file {file}: {e}")
        return None
    
    # Perform station domain selection if specified    
    if geospatial_crop is not None:
        try:
            obs, lat, lon, time_values = spatial_domain_selection(obs, lat, lon, time_values, geospatial_crop)
        except Exception as e:
            print(f"Error performing spatial domain selection for file {file}: {e}")
            return None
        
    if not flag:
        try:
            # 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
        except Exception as e:
            print(f"Error performing regridding for  file {file}: {e}")
            return None
    else:
        ds.close()
        return obs, lat, lon, time_values 



def process_single_file_viirs(
        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")  
    except Exception as e:
        print(f"Reading file {file}: {e}")
        return None

    # Flatten arrays
    try:
        obs, lat, lon, time_values = new_preprocess_vars(ds, obs_var, lat_var, lon_var, time_var, start_date, end_date)
    except Exception as e:
        print(f"Error processing file {file}: {e}")
        return None
    
    # Apply quality assurance if applicable
    if apply_qa:
        try:
            obs, lat, lon, time_values = quality_assurance(ds, obs, lat, lon, time_values, qualityFlags)
        except Exception as e:
            print(f"Error applying QA Flags for file {file}: {e}")
            return None
    # Perform time domain selection beteen start and end date
    try:
        obs, lon, lat, time_values = time_domain_selection(obs, lat, lon, time_values, start_date, end_date)
    except Exception as e:
        print(f"Error performing time domain selection for file {file}: {e}")
        return None
    
    # Perform station domain selection if specified
    if geospatial_crop is not None:
        try:
            obs, lat, lon, time_values = spatial_domain_selection(obs, lat, lon, time_values, geospatial_crop)
        except Exception as e:
            print(f"Error performing spatial domain selection for file {file}: {e}")
            return None
        
    if not flag:
        try:
            # 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
        except Exception as e:
            print(f"Error performing regridding for  file {file}: {e}")
            return None
    else:
        ds.close()
        return obs, lat, lon, time_values 
    


class GenericMultiProcessing(ABC):
    """
    Base method for multiprocessing
    """

    @abstractmethod
    def process_batch(self):
        pass


@dataclass
class VIIRSMultiprocessing(GenericMultiProcessing):
    """
    Multiprocessing method for VIIRS
    """
    r: Grid
    obs_var: NDArray 
    lat_var: NDArray 
    lon_var: NDArray
    time_var: NDArray
    start_date: pd.Timestamp
    end_date: pd.Timestamp 
    apply_qa: bool
    qualityFlags: dict
    batch: List[str]
    geospatial_crop: NDArray = None
    flag: bool = False
    

    
    def process_batch(self):
        """
        Batch processing for VIIRS
        """

        args = [
            (
            self.r, 
            file, 
            self.obs_var, 
            self.lat_var, 
            self.lon_var, 
            self.time_var, 
            self.start_date, 
            self.end_date, 
            self.apply_qa, 
            self.qualityFlags, 
            self.flag, 
            self.geospatial_crop
            ) 
            for file in self.batch
        ]



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

        # Filter valid results and log failed files
        valid_results = [res for res in results if res is not None]
        print(f"Valid results: {len(valid_results)}")
        failed_files= [file for file, res in zip(self.batch, results) if res is None]

        return valid_results, failed_files


@dataclass
class TROPOMIMultiprocessing(GenericMultiProcessing):
    """
    Multiprocessing method for TROPOMI
    """
    r: Grid
    obs_var: NDArray 
    lat_var: NDArray 
    lon_var: NDArray
    time_var: NDArray
    start_date: pd.Timestamp
    end_date: pd.Timestamp 
    apply_qa: bool
    qualityFlags: dict
    batch: List[str]
    geospatial_crop: NDArray = None
    flag: bool = False
    

    
    def process_batch(self):
        """
        Batch processing for TROPOMI
        """

        args = [
            (
            self.r, 
            file, 
            self.obs_var, 
            self.lat_var, 
            self.lon_var, 
            self.time_var, 
            self.start_date, 
            self.end_date, 
            self.apply_qa, 
            self.qualityFlags, 
            self.flag,
            self.geospatial_crop
            ) 
            for file in self.batch
        ]


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

        # Filter valid results and log failed files
        valid_results = [res for res in results if res is not None]
        print(f"Valid results: {len(valid_results)}")
        failed_files= [file for file, res in zip(self.batch, results) if res is None]

        return valid_results, failed_files
