import pathlib
import numpy as np
import pandas as pd
from datetime import datetime
from scipy import integrate
from scipy.interpolate import interp1d
from loguru import logger

from gfatpy.utils.io import read_yaml

# AERONET SYSTEM INFO
AERONET_INFO_FILE = pathlib.Path(__file__).parent.absolute() / "info.yml"
AERONET_INFO = read_yaml(AERONET_INFO_FILE)


def find_available_aod_wavelengths(header: dict, df: pd.DataFrame) -> list:
    """find available aod wavelengths in AERONET dataframe.

    Args:
        header (dict): AERONET header
        df (pd.DataFrame): AERONET dataframe
    Returns:
        list[int]: available wavelenghts
    """

    available_wavelengths = []
    aeronet_data = header["aeronet_data"].replace(" ", "_")
    for wave_ in AERONET_INFO["wavelengths"][aeronet_data]:
        property2check = AERONET_INFO["aod_string_type"][aeronet_data].replace(
            "&", f"{wave_}"
        )
        if property2check in df.keys():
            available_wavelengths.append(wave_)
    available_wavelengths = np.sort(np.asarray(available_wavelengths)).tolist()
    return available_wavelengths


def find_nearest_wavelength(
    header: dict, df: pd.DataFrame, target_wavelength: float
) -> list[float]:
    """Find the available wavelength above [`larger_wavelength`] and below [`lower_wavelength`] of the `target_wavelength`.

    Args:
        header (dict): AERONET header.
        df (pd.DataFrame): AERONET dataframe.
        target_wavelength (float): wavelength to be interpolated.

    Returns:
        list[float]: lower_wavelength, larger_wavelength
    """

    available_wavelengths = find_available_aod_wavelengths(header, df)
    idx_to_be_placed = np.searchsorted(
        available_wavelengths, [target_wavelength], side="left"
    )[0]
    if idx_to_be_placed == 0:  # Extrapolation is required
        bottom_wave = available_wavelengths[0]
        up_wave = available_wavelengths[1]
    elif idx_to_be_placed == len(available_wavelengths):  # Extrapolation is required
        bottom_wave = available_wavelengths[-2]
        up_wave = available_wavelengths[-1]
    else:
        bottom_wave = available_wavelengths[idx_to_be_placed - 1]
        up_wave = available_wavelengths[idx_to_be_placed + 1]
    return bottom_wave, up_wave


def add_interpol_aod(
    header: dict,
    df: pd.DataFrame,
    target_wavelength: int,
    bottom_wavelength: int | None = None,
    up_wavelength: int | None = None,
    allow_search: bool = True,
) -> pd.DataFrame:
    """Add the interpolated aod of `target_wavelength` to the AERONET dataframe.

    Args:
        header (dict): AERONET header.
        df (pd.DataFrame): AERONET dataframe.
        target_wavelength (int): Target wavelength
        bottom_wavelength (int | None, optional): lower wavelength to make the interpolation. Defaults to None. With None, the nearest lower wavelength available is used.
        up_wavelength (int | None, optional): lower wavelength to make the interpolation. Defaults to None. With None, the nearest larger wavelength available is used.
        allow_search (bool | optional): allow to search the best lower and larger wavelength if the provided are not available.
    Raises:
        ValueError: Interpolation not possible since bottom/up wavelengths are not available.

    Returns:
        pd.DataFrame: AERONET dataframe with new interpolated AOD at `target_wavelength`.
    """
    available_wavelengths = find_available_aod_wavelengths(header, df)

    if (
        bottom_wavelength not in available_wavelengths
        or up_wavelength not in available_wavelengths
    ) and allow_search is False:
        logger.error(
            "Interpolation not possible since bottom/up wavelengths are not available."
        )
        raise ValueError

    if bottom_wavelength is None or up_wavelength is None:
        bottom_wavelength, up_wavelength = find_nearest_wavelength(
            header, df, target_wavelength
        )
    # Property strings
    bottom_property = AERONET_INFO["aod_string_type"][
        header["aeronet_data"].replace(" ", "_")
    ].replace("&", f"{bottom_wavelength}")

    up_property = AERONET_INFO["aod_string_type"][
        header["aeronet_data"].replace(" ", "_")
    ].replace("&", f"{up_wavelength}")

    target_property = bottom_property.replace(
        f"{bottom_wavelength}", f"{target_wavelength}"
    )

    # Retrieve angstrom exponent
    angstrom_exponent = -np.log(
        df[up_property].values / df[bottom_property].values
    ) / np.log(up_wavelength / bottom_wavelength)

    # Interpolation
    if target_wavelength > bottom_wavelength and target_wavelength < up_wavelength:
        df[target_property] = df[bottom_property].values * (
            target_wavelength / bottom_wavelength
        ) ** (-angstrom_exponent)

    # Left-side extrapolation
    if target_wavelength < bottom_wavelength:
        logger.warning("Extrapolation is used.")
        df[target_property] = df[bottom_property].values * (
            target_wavelength / bottom_wavelength
        ) ** (-angstrom_exponent)

    # Right-side extrapolation
    if target_wavelength > up_wavelength:
        logger.warning("Extrapolation is used.")
        df[target_property] = df[up_property].values * (
            target_wavelength / up_wavelength
        ) ** (-angstrom_exponent)

    return df


def add_logAOD(header: dict, df: pd.DataFrame) -> pd.DataFrame:
    """Add colums in the DataFrame with natural logarithyms of AODs.

    Args:
        df (pd.DataFrame): Aeronet DataFrame from gfatpy.aeronet.reader()

    Returns:
        df (pd.DataFrame):  Aeronet DataFrame with new columns
    """

    # Ln applied to AOD
    waves = find_available_aod_wavelengths(header, df)
    for wave_ in waves:
        aod_var = AERONET_INFO["aod_string_type"][
            header["aeronet_data"].replace(" ", "_")
        ].replace("&", f"{wave_}")

        df["lnAOD%d" % wave_] = np.log(df[aod_var])
    return df


def resample_logradius_distribution(
    distribution: np.ndarray | pd.Series,
    factor: int = 10,
    interpolation_kind: str = "quadratic",
) -> tuple[np.ndarray]:
    """Resample a logaritmic distribution.

    Args:
        distribution (np.ndarray): Aeronet DataFrame from `gfatpy.aeronet.reader()`
        factor (int, optional): Increase the resolution by dividing the current resolution by this factor. Defaults to 10.
        interpolation_kind (str, optional): Method of interpolation. Defaults to "quadratic".

    Returns:
        list[np.ndarray]:
            - Resampled distribution.
            - Resampled radius.
            - Resampled logarithmic radius.
            - New logarithmic radius resolution.
    """
    radius = np.asarray(AERONET_INFO["radius"])
    lnr = np.log(radius)
    resol_lnr = np.diff(lnr).mean()
    resample_resol_lnr = resol_lnr / factor
    resample_lnr = np.arange(lnr[0], lnr[-1] + resol_lnr / factor, resol_lnr / factor)
    resample_radius = np.exp(resample_lnr)

    # Increase resolution to improve the fitting
    resample_function = interp1d(
        radius, distribution, kind=interpolation_kind, bounds_error=True
    )
    resample_distribution_scipy = resample_function(resample_radius)
    resample_distribution_numpy = np.interp(
        resample_radius, radius, distribution.values.astype(float)
    )
    resample_distribution = np.concatenate(
        (
            resample_distribution_scipy[resample_radius < 11.432287],
            resample_distribution_numpy[resample_radius >= 11.432287],
        )
    )
    return resample_distribution, resample_radius, resample_lnr, resample_resol_lnr


def filter_df(df: pd.DataFrame, conditions_dict: dict) -> pd.DataFrame:
    """Filter the aeronet dataframe according to the conditions specified in `conditions_dict`

    Args:
        df (pd.DataFrame): Aeronet DataFrame from `gfatpy.aeronet.reader()`
        conditions_dict (dict): condition dictionary.
            Keys: Values
                'aerosol_type': ['DUST', 'HA', 'MA', 'SA', 'MIXTURE'] from `gfatpy.aeronet.classification.aerosol_type_classification_Lee()`
                'mode_predominance': [HYGHLY_COARSE, COARSE, BIMODAL, FINE, HIGHLY_FINE, UNKNOWN] from `gfatpy.aeronet.classification.mode_predominance_classification()`
                'If_Retrieval_is_L2(without_L2_0.4_AOD_440_threshold)': [1, 0]

    Returns:
        pd.DataFrame: Filtered Aeronet DataFrame
    """

    conditions = True * np.ones(len(df.index))
    for key_ in conditions_dict.keys():
        condition_ = df[key_] == conditions_dict[key_]
        conditions = np.logical_and(conditions, condition_)
    df_ = df[conditions]
    return df_


def calculate_fine_mode_fraction(
    aod_fine: np.ndarray | pd.Series,
    aod_total: np.ndarray | pd.Series,
) -> np.ndarray | pd.Series:

    """calculate_fine_mode_fraction retrieves the fine mode fraction as the fine-to-total AOD ratio.

    Args:
        aod_fine (np.ndarray): Fine AOD
        aod_total (np.ndarray): Total AOD

    Raise:
        ValueError: 'Length of arrays must be the same.'

    Returns:
        np.ndarray: fine mode fraction
    """

    if len(aod_fine) != len(aod_total):
        logger.error("Length of arrays must be the same.")
        raise ValueError

    return aod_fine / aod_total


def distribution_to_total_concentration(
    df: pd.DataFrame,
    concentration_type: str = "number",
    reference_radius_nm: int | None = None,
    integral_towards: str = "positive",
) -> pd.DataFrame:
    """Retrieve the total concentration (and natural logarithm) in two radius ranges: [r_min, r_reference] and [r_reference, r_max].

    Args:
        df (pd.DataFrame): AERONET dataframe from *.all file.
        concentration_type (str): ['surface', 'number']
        reference_radius_nm (int | None): minimum radius value to retrieve the concentration. If None, inflection radius is used.
        integral_towards (str): Direction of the integration ['positive': towards the maximum radius, 'negative': towards the minimum radius , 'both': both are retrieved and saved]. Defaults 'positive'.

    Raises:
        ValueError: Dataframe is empty.
        ValueError: integral_towardds is not one of the following options: ['positive', 'negative', 'both'].

    Returns:
        df (pd.DataFrame): dataframe with number/surface concentration and its natural logarithm 'lnN'.
    """

    if len(df) == 0:
        logger.error("Dataframe is empty.")
        raise ValueError

    # Initializing variables
    # Radius
    radius = np.asarray(AERONET_INFO["radius"])  # um¡

    # Positive
    if integral_towards == "both":
        do_positive, do_negative = True, True
    elif integral_towards == "positive":
        do_positive, do_negative = True, False
    elif integral_towards == "negative":
        do_positive, do_negative = False, True
    else:
        logger.error("Value of integral_towards not allowed.")
        raise AttributeError

        # Variables to storage concentrations
    conc_up = np.zeros(len(df))
    conc_down = np.zeros(len(df))
    if reference_radius_nm is not None:
        # Convert from nanomter to micrometer
        reference_radius_used_microm = reference_radius_nm / 1000.0
        conc_up_str = (
            f"{concentration_type}_concentration_above_{reference_radius_nm:d}"
        )
        conc_down_str = (
            f"{concentration_type}_concentration_below_{reference_radius_nm:d}"
        )
    else:
        conc_up_str = "COARSE"
        conc_down_str = "FINE"

        # Indexes of the radius in the Dataframe.
    idx_min = df.columns.get_indexer(["0.050000"])
    idx_max = df.columns.get_indexer(["15.000000"])

    # Stablising the concentration_factor (number Vs surface concentration)
    if concentration_type == "number":
        concentration_factor = 1.0
    else:
        concentration_factor = (
            4 * np.pi * radius[radius >= reference_radius_used_microm] ** 2
        )

    # Selection of the VSD
    dV_dlnr = df.iloc[:, np.arange(idx_min, idx_max + 1, 1)]  # VSD um^3/um^2
    dV_dlnr = dV_dlnr.to_numpy()

    # Retrieve NSD
    Vradius = (4.0 / 3.0) * np.pi * np.power(radius, 3)  # i-bin volume  um^3
    delta_lnr = np.diff(np.log(radius)).mean()  # resolution um
    dN_dlnr_serie = dV_dlnr / Vradius  # NSD #/um^2

    if do_positive:
        for idx in np.arange(len(df)):
            if reference_radius_nm is None:
                inflection_radius = df[
                    "Inflection_Radius_of_Size_Distribution(um)"
                ].iloc[idx]
                reference_radius_used_microm = inflection_radius / 1000.0

            dN_dlnr_ = dN_dlnr_serie[idx, :]

            conc_up[idx] = integrate.simps(
                dN_dlnr_[radius >= reference_radius_used_microm] * concentration_factor,
                dx=delta_lnr,
            )

        # Save arrays in DATAFRAME
        df[conc_up_str] = conc_up
        df[f"ln_{conc_up_str}"] = np.log(conc_up)

    if do_negative:
        for idx in np.arange(len(df)):
            if reference_radius_nm is None:
                inflection_radius = df[
                    "Inflection_Radius_of_Size_Distribution(um)"
                ].iloc[idx]
                reference_radius_used_microm = inflection_radius / 1000.0

            dN_dlnr_ = dN_dlnr_serie[idx, :]

            conc_down[idx] = integrate.simps(
                dN_dlnr_[radius < reference_radius_used_microm] * concentration_factor,
                dx=delta_lnr,
            )

        # Save arrays in DATAFRAME
        df[conc_down_str] = conc_down
        df[f"ln_{conc_down_str}"] = np.log(conc_down)
    return df


def find_aod(
    header: dict,
    df: pd.DataFrame,
    target_wavelength: int,
    target_datetime: datetime,
    allowed_time_gap_hour: int = 1,
    allowed_interpolation: bool = True,
) -> tuple[float | None, float | None]:
    """Find the aod value at `wavelength` close to the `datetime` as much as determined by `allowed_time_gap_hour`.

    Args:
        header (dict): AERONET header from `gfatpy.aeronet.reader()`
        df (pd.DataFrame): AERONET dataframe from `gfatpy.aeronet.reader()`
        target_wavelength (float): wavelength of the aod requiered. If not found, it will be interpolated with the available ones.
        target_datetime (datetime): datetime
        allowed_time_gap_hour (int, optional): allowed time gap between the mesured and required datetime. Defaults to 1. Provides None if requirement not fulfilled.
        allowed_interpolation (bool, optional):it allows wavelength interpolation.

    Raises:
        RuntimeWarning: Time gap found to be larger than allowed by the user with `allowed_time_gap_hour`
        RuntimeWarning: "Extra- inter-polation required but not allowed by the user (`allowed_interpolation` is `False`).

    Returns:
        tuple[float | None, float | None]: measured_aod, measured_datetime
    """
    """

    Args:
    Returns:
        tuple[float, datetime]:
    """
    # Set default avaiable data flag
    go = False

    # Set result to None by default
    measured_datetime, measured_aod = None, None

    # Finde the nereast datetime
    nearest_datetime = (
        df.index[df.index.get_indexer([target_datetime], method="nearest")]
    ).to_pydatetime()[0]

    # Check distance between both
    if target_datetime > nearest_datetime:
        time_gap = target_datetime - nearest_datetime
    else:
        time_gap = nearest_datetime - target_datetime
    time_gap_hour = time_gap.seconds / 60.0 / 60.0

    if time_gap_hour > allowed_time_gap_hour:
        logger.warning(
            "Time gap found to be larger than allowed by the user with `allowed_time_gap_hour`."
        )
        raise RuntimeWarning
    property = AERONET_INFO["aod_string_type"][
        header["aeronet_data"].replace(" ", "_")
    ].replace("&", f"{target_wavelength}")

    if (
        property not in df.keys() and not allowed_interpolation
    ):  # extrapolation required
        logger.warning(
            "Extra- inter-polation required but not allowed by the user (`allowed_interpolation` is `False`)."
        )
        raise RuntimeWarning

    # Extrapolation allowed
    try:
        bottom_wave, up_wave = find_nearest_wavelength(header, df, target_wavelength)
        df = add_interpol_aod(header, df, target_wavelength, bottom_wave, up_wave)
        go = True
    except Exception as e:
        logger.error("Extra- inter-polatization not finished.")
        raise RuntimeError

    if go:
        row = df.iloc[df.index.get_indexer([target_datetime], method="nearest")]
        measured_aod = row[property].values[0]
        measured_datetime = nearest_datetime

    return measured_aod, measured_datetime
