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

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)

        for i, (flag, info) in enumerate(qualityFlags.items()):
            
            flag_mask = np.zeros_like(obs, dtype=bool)
            arr = ds[info["qa_name"]].values.flatten()
            for value in info["qa_values"]:
                flag_mask |= (arr == value)
            
            if i == 0:
                mask = flag_mask
            else:
                if info["qa_combined"]:
                    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

# Time Domain Selection function based off of xarray Will this work with dask too, a question for later
def time_domain_selection(
        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
    """

    # But the format of the time column may change aaaaaaaghh

    # 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_values_index = np.where(is_between_times(time_values))[0] # np.where returns tuple
    
    time_values = time_values[time_values_index]

    return time_values, time_values_index


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
    """
    time_values, time_values_index = preprocess_vars_mapies(ds, time_var, start_date, end_date)


    # 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()

    # Reindex lon column values
    lon_values = reindex(
        time_values_index, 
        lon_values,
        )

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

    # Reindex aod column values
    lat_values = reindex(
        time_values_index, 
        lat_values,
        )


    # AOD column values, default "Aerosol_Optical_Thickness_550_Land_Ocean_Best_Estimate"
    aod_dims = ds[obs_var].dims

    aod_shape = ds[obs_var].shape
    aod_attrs = ds[obs_var].attrs
    obs = ds[obs_var].values.flatten()

    # Reindex aod column values
    obs = reindex(
        time_values_index, 
        obs,
        )

    return obs, lat_values, lon_values, time_values, time_values_index

def preprocess_vars_mapies(ds, time_var, start_date, end_date):
    """
    Preprocessing of the dataset
    """
    # 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)
    
    
    # Time domain selection
    # print(f'Time values before selection: {time_values}')
    time_values, time_values_index = time_domain_selection(time_values, start_date, end_date)
    # print(f'Time values after selection: {time_values}')

    return time_values, time_values_index

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