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

from mapies.grids.monarch import RotatedGrid, IrregularRotatedGrid, RegularGrid, regridding_function

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

#Timing decorator
def timeit(func):
    @wraps(func)
    def timeit_wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        total_time = end_time - start_time
        print(f'Function {func.__name__}{args} {kwargs} Took {total_time:.4f} seconds')
        return result
    return timeit_wrapper


def get_filepaths(pattern):
    """
    Get file paths from a pattern
    """
    return glob(pattern, recursive=True)

# Function from Guillaume's script could be useful
# Only supported in python3.10
def get_file_list(patterns) -> List[str]:
    """
    param patterns: one or several glob patterns (or exact file names) pointing to the location of files on disk.
    return: the matching list of files
    """

    if isinstance(patterns, str):
        patterns = [patterns]

    files_final = []
    args = [pattern for pattern in patterns]
    with multiprocessing.Pool(processes=8) as pool: 
        for file_list in pool.imap_unordered(get_filepaths, args):
            files_final.extend(file_list)  # Flatten the lists into files_final


    return sorted(files_final)

def exception_factory(exception, message):
    return exception(message)


def time_converter(date:str) -> pd.Timestamp:
    """
    :param date: date that you wish convert to standardised format
    """
    # Remove any special characters
    for ch in ['\\','`','*','_','{','}','[',']','(',')','%','>','#','+','-','.','!','$','\'', ':', ' ']:
        if ch in date:
            date = date.replace(ch,"")
            logging.info(f"Special character {ch} was replaced in the date")
    
    if len(date) == 14: # Probable str format %Y%m%d%H%M%S
        date = pd.to_datetime(date, format="%Y%m%d%H%M%S")
    elif len(date) == 12: # Probable str format %Y%m%d%H%M
        date = pd.to_datetime(date, format="%Y%m%d%H%M")
    elif len(date) == 10: # Probable str format %Y%m%d%H
        date = pd.to_datetime(date, format="%Y%m%d%H")
    elif len(date) == 8: # Probable str format %Y%m%d
        date = pd.to_datetime(date, format="%Y%m%d")
    elif len(date) == 7: # Probable str format %Y%J
        date =pd.to_datetime(date, format="%Y%j")
    else:
        raise exception_factory(ValueError, "Invalid date format")

    return date



def quality_assurance(ds, obs, lat, lon, time_values, qualityFlags):
    """
    Applying quality assurance flags to mapies data
    """
    mask = np.zeros_like(obs, dtype=bool)
    # Iterate through each qa flag
    for i, (flag, info) in enumerate(qualityFlags.items()):
        # TODO add logger of flag being applied with condition and value
        
        flag_mask = np.zeros_like(obs, dtype=bool)
        arr = ds[info["qa_name"]].values.flatten()
        # Get qa conditions
        try:
            condition = info["qa_condition"]
        except KeyError:
            # print("No QA condition passed, putting equal")
            condition = "equal"
        # Get qa values
        try:
            values = info["qa_values"]
        except KeyError:
            raise Exception("No QA values passed, putting equal")

        if isinstance(values, (float, int)):
            values = [values]

        for value in values:
            if condition == "greater_than":
                flag_mask |= (arr > value)
            elif condition == "less_than":
                flag_mask |= (arr < value)
            elif condition == "greater_than_or_equal":
                flag_mask |= (arr >= value)
            elif condition == "less_than_or_equal":
                flag_mask |= (arr <= value)
            elif condition == "equal": 
                flag_mask |= (arr == value)
        
        if i == 0:
            mask = flag_mask
        else:
            # If we want to perform an and or or between conditions
            try:
                combination = info["qa_combined"]
            except KeyError:
                # print("No combination of and or or passed for the condition, presuming or(False)")
                combination = False
            
            if combination:
                mask &= flag_mask
            else:
                mask |= flag_mask

    obs = obs[mask]
    lat = lat[mask]
    lon = lon[mask]
    time_values = time_values[mask]


    return obs, lat, lon, time_values



def spatial_domain_selection(obs, lat, lon, time, spatial_bounds):
    """
    A function for cutting the data spatially

    :param spatial_bounds: - the boundaries that you would like to crop te data by [[lon], [lat]] e.g. [[0, 20], [10, 12]]
    """
    # Check that the spatial bounds are 2 dimensional
    if len(spatial_bounds) != 2:
        raise Exception("Spatial bounds passed are not 2 dimensional, make sure to pass [[left_lon, right_lon], [left_lat, right_lat]]")
    
    # Get the bounds
    lon_l_bound, lon_r_bound = spatial_bounds[0]
    lat_l_bound, lat_r_bound = spatial_bounds[1]

    # Crop the spatial variabes
    lon_mask = np.logical_and(lon>=lon_l_bound, lon<=lon_r_bound)
    lat_mask = np.logical_and(lat>=lat_l_bound, lat<=lat_r_bound)
    
    # Masking the data
    lon = lon[lon_mask & lat_mask]
    lat = lat[lon_mask & lat_mask]
    time = time[lon_mask & lat_mask]
    obs = obs[lon_mask & lat_mask]

    assert len(lon) == len(lat)
    assert len(obs) == len(time)
    assert len(lon) == len(obs)
    
    return obs, lat, lon, time
    

# Time Domain Selection function based off of xarray Will this work with dask too, a question for later
def time_domain_selection(
        obs, 
        lat, 
        lon, 
        time_values, 
        start_date, 
        end_date, 
        closed = None) -> np.array:
    """
    :param ds: one or several xarray datasets to cut down to correct size.
    :param start_date: the start date that we want to start with for our time domain selection
    :param end_date: the end date that we want to end with for our time domain seleciton
    :param time_column: the column for the time variables
    :param closed: like the pandas old param on the pandas date_range method (None, left, right), default:None
    :return: a cut down numpy array
    """


    # For these if it is a string we need to know what the format is (maybe pandas figures that out)
    if isinstance(start_date, str):
        start_date = pd.to_datetime(start_date)

    if isinstance(end_date, str):
        end_date = pd.to_datetime(end_date)

    
    # For the time column we can either centralise it in the seperate Instrument objects or pass it here,
    # I will start by passing it here

    def is_between_times(array):
        
        if closed == "left":
            return (array > start_date) & (array <= end_date)
        elif closed == "right":
            return (array >= start_date) & (array < end_date)
        else:
            return (array >= start_date) & (array <= end_date)
    
    time_mask = np.where(is_between_times(time_values))[0]


    time_values = time_values[time_mask]
    lon = lon[time_mask]
    lat = lat[time_mask]
    obs = obs[time_mask]

    return obs, lon, lat, time_values


def inflate_array(var_array, shape1, shape2):
    """
    A function to inflate an array by repeating it along one dimension
    This is useful if you have a flattened array and need to repeat it due to the original array having less 
    dimensions that the target array
    """
    dif1 = np.setdiff1d(shape1, shape2)
    dif2 = np.setdiff1d(shape2, shape1)

    repeats = np.concatenate((dif1, dif2))

    return np.repeat(var_array, repeats)


def frequency_int(time_interval):
    """
    Convert frequency to an integer value, used in the DA
    """
    if time_interval == "H":
        frequency = 1
    elif time_interval == "D":
        frequency = 24       
    else:
        frequency = int(time_interval[0])
    return frequency


def error_estimation(datatype:str, obs:np.typing.NDArray, unc_const:list) -> np.typing.NDArray:
    """
    Error estimation function
    """
    if datatype == "viirs":
        obserr = np.sqrt((unc_const[0]*obs+unc_const[1])**2 + unc_const[2]**2)
    elif datatype == "in-situ":
        # TODO: Need to work out how to do this for each different element
        #obserr = np.maximum(unc_const[e][1], obs*unc_const[e][0])
        pass
    return obserr




# def geo_to_rot(lons, lats, centre_lon: float, centre_lat: float):
#     """
#     Rotating coordinates from cartesian lat/lon to rotated rlon/rlat
#     """
#     distance_lons = np.radians(lons - centre_lon)
#     lons = np.radians(lons)
#     lats = np.radians(lats)
#     centre_lon = np.radians(centre_lon)
#     centre_lat = np.radians(centre_lat)
    
#     x = cos(centre_lat) * sin(lats) - sin(centre_lat) * cos(lats) * cos(distance_lons)
#     y = cos(lats) * sin(distance_lons)
#     z = cos(centre_lat) * cos(lats) * cos(distance_lons) + sin(centre_lat) * sin(lats)
#     # Arctan2 used 
#     # Explanation of the difference https://geo.libretexts.org/Courses/University_of_California_Davis/GEL_056%3A_Introduction_to_Geophysics/Geophysics_is_everywhere_in_geology.../zz%3A_Back_Matter/Arctan_vs_Arctan2
#     rlon = np.arctan2(y, z)
#     rlat = np.arcsin(x)
#     #rlon[x < 0] += pi
#     # Convert back to degrees
#     rlon = np.degrees(rlon)
#     rlat = np.degrees(rlat)
#     return rlon, rlat

def new_preprocess_vars(ds, obs_var, lat_var, lon_var, time_var, start_date, end_date):
    """
    * Preprocessing of the dataset
    * Flattening arrays 
    * Applying time domain selection
    """
   # Get all info about time columns
    time_dims = ds[time_var].dims
    time_shape = ds[time_var].shape
    time_attrs = ds[time_var].attrs # Useful if we need to convert to datetime
    

    # Get time values flatten and convert to datetime values
    time_values = ds[time_var].values
    time_values = time_values.flatten()
    
    if time_values.dtype == "timedelta64[ns]":
        logging.info("Adding time origin to time values as the time variable is in timedelta")
        time_values = np.add(TIME_ORIGIN, time_values)
        time_values = pd.to_datetime(time_values)
    else:
        time_values = pd.to_datetime(time_values)
    
    
    obs_shape = ds[obs_var].shape
    obs_attrs = ds[obs_var].attrs
    
    # Duplicate values in the time array to get the same shape as the flattened variable array
    if time_shape != obs_shape:
        time_values = inflate_array(time_values, time_shape, obs_shape)
        
    # Lon and lat values
    lon_dims = ds[lon_var].dims
    lon_shape = ds[lon_var].shape
    lon_attrs = ds[lon_var].attrs
    lon_values = ds[lon_var].values.flatten()


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


    # AOD column values, default "Aerosol_Optical_Thickness_550_Land_Ocean_Best_Estimate"
    
    obs = ds[obs_var].values.flatten()

    return obs, lat_values, lon_values, time_values


def reindex(
    dependent_var_index,
    independent_var_values,
    ):
    """"
    Recutting the data whenever a selction of the data has been made along one of the dimenisons
    
    Based off of how it has been done in Providentia
    """
    # Maybe add a checker try/except or assert
    independent_var_values = independent_var_values[dependent_var_index]

    # return reindexed values
    return independent_var_values


def day_boundaries(day, start_date, end_date):
    """
    Get daily boundaries only if they are within start and end date
    If not return the start and end date as the boundaries
    """
    from datetime import datetime, timedelta, time
    date = start_date + timedelta(days=day) - timedelta(days=1)
    day_start = date.floor(freq="d")
    day_end = (date + timedelta(days=1)).floor(freq="d")

    # If day_start or day_end are outside the regions of start_date and end_date 
    # Return those instead
    if day_start < start_date:
        day_start = start_date

    if day_end > end_date:
        day_end = end_date

    return day_start, day_end

def filter_files(files, start_date, end_date):
    """
    Filter files based on the start_date and end_date
    """
    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 start_date <= file_dt <= end_date:
            files_filtered.append(file)

    return files_filtered

def generate_mid_time_array(start_date, end_date, obs):
    """
    Generate the mid time array for the observations
    (Used when saving into nectCDF averaged data)
    """

    mid_time = start_date + (end_date - start_date) / 2
    mid_time = np.datetime64(mid_time)
    N = len(obs)
    mid_time_array = np.repeat(mid_time, N)

    return mid_time_array

def create_batches(files, batch_size):
    """
    Create batches of files to pass to MP
    """

    batches = [
            files[i:i + batch_size] for i in range(0, len(files), batch_size)
        ]
    
    return batches


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

def process_batch(args, 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]
    print(f"Valid results: {len(valid_results)}")
    failed_files= [file for file, res in zip(batch, results) if res is None]

    return valid_results, failed_files


def combine_arrays(valid_results, cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count):
    """
    Combine the arrays from the valid results. In other words sum the arrays and count the number of observations.
    """

    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

    return cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count

def compute_final_averages(cumulative_obs, cumulative_lat, cumulative_lon, cumulative_count):
    """
    Compute the final averages from the cumulative arrays. In other words divide the cumulative_obs by the cumulative_count.
    """

    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]

    return final_obs, final_lat, final_lon

def append_and_concatenate(valid_results, day):
    """
    Append and concatenate arrays
    """
    daily_obs = []
    daily_lat = []
    daily_lon = []
    daily_time = []

    for obs, lat, lon, time_values in valid_results:
        #print(f'Shape of obs: {obs.shape}, Shape of lat: {lat.shape}, Shape of lon: {lon.shape}, Shape of time: {time_values.shape}')
        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}")
        return None, None, None, None  

    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)

    return final_obs, final_lat, final_lon, final_time

def filter_observations(obs, lat, lon, count):
    """
    Filter out observations with negative values
    """
    
    valid_obs = obs > 0
    final_lon = lon[valid_obs]
    final_lat = lat[valid_obs]
    final_obs = obs[valid_obs]
    final_count = count[valid_obs]

    return final_obs, final_lat, final_lon, final_count