import pandas as pd
import xarray as xr
import numpy as np
from scipy.integrate import cumulative_trapezoid as cumtrapz


from gfatpy.lidar.retrieval.klett import klett_rcs
from gfatpy.lidar.retrieval.raman import retrieve_backscatter, retrieve_extinction
from gfatpy.lidar.utils.utils import refill_overlap, signal_to_rcs
from gfatpy.atmo.rayleigh import molecular_properties


def overlap_iterative_raman(
    elastic_signal: xr.DataArray,
    raman_signal: xr.DataArray,
    meteo_profiles: pd.DataFrame,
    particle_lidar_ratio: float,
    wavelengths: tuple[float, float],
    reference: tuple[float, float],
    beta_aer_ref: float = 0,
    min_overlap_range: float = 200,  # minimum range of overlap function
    particle_angstrom_exponent: float = 0.0,
    iteration_limit: int = 10,
    debugging: bool = False,
    **kwargs: dict,
) -> xr.Dataset:
    """Retrieve overlap function by means of the Wandinger's method.

    Args:
        elastic_signal (xr.DataArray): Elastic lidar signal at wavelength `wavelengths[0]`.
        raman_signal (xr.DataArray): Raman lidar signal at wavelength `wavelengths[1]`.
        meteo_profiles (pd.DataFrame): from  `gfatpy.atmo.atmo.generate_meteo_profiles` with pressure and temperature data.
        particle_lidar_ratio (float): Particle lidar ratio to be used in the Klett inversion. It shall be chosen so Klett backscatter profile fits Raman backscatter profile once full overlap is reached.
        reference (tuple[float, float], optional): Reference range required by the Klett and Raman inversions. Defaults to (8000,8500).
        wavelengths (tuple[float, float], optional): Elastic and Raman wavelengths. Defaults to (532, 531).
        beta_aer_ref (float, optional): Particle backscatter coefficient at `reference range` in m^-1sr^-1. Defaults to 0 m^-1sr^-1.
        particle_angstrom_exponent (float, optional): Particle extinction-related Angstrom exponent. It is used in the Raman inversion. Defaults to 0.
        min_overlap_range (float, optional): Overlap function will be forced to zero below [m]. Defaults to 200 m.
        iteration_limit (int, optional): Maximum number of iterations. Defaults to 10.
        debugging (bool, optional): If True, it will return a dataset with all the intermediate variables. Defaults to False.
        **kwargs: Additional arguments such as `fulloverlap_height` and `output_dir`.

    Returns:
        xr.Dataset: overlap function. If debugging is True, it will return a dataset with all the intermediate variables.
    """
    # Check z in elastic and raman signals (the same for both signals)
    if len(elastic_signal.range) != len(raman_signal.range):
        raise ValueError(
            "elastic and Raman signals must have the same 'range' dimension."
        )

    if len(elastic_signal.range) != len(meteo_profiles["height"]):
        raise ValueError(
            "Signals range dimension and meteo_profiles['heights'] must have the same size."
        )

    z = elastic_signal.range.values.astype(np.float64)

    mol_properties = molecular_properties(
        wavelengths[0],
        meteo_profiles["pressure"],
        meteo_profiles["temperature"],
        heights=z,
    )
    molecular_backscatter = mol_properties["molecular_beta"].values.astype(np.float64)
    pressure = meteo_profiles["pressure"].values.astype(np.float64)
    temperature = meteo_profiles["temperature"].values.astype(np.float64)

    # Beta Klett
    elastic_backscater = klett_rcs(
        signal_to_rcs(elastic_signal, z).astype(np.float64),
        z,
        molecular_backscatter,
        lr_part=particle_lidar_ratio,
        reference=reference,
    )

    extinction = retrieve_extinction(
        raman_signal.values.astype(np.float64),
        z,
        wavelengths,
        pressure,
        temperature,
        reference=reference,
    )

    # Refill extinction
    extinction = refill_overlap(extinction, z, kwargs.get("fulloverlap_height", 600.0))

    raman_backscatter = retrieve_backscatter(
        raman_signal.values.astype(np.float64),
        elastic_signal.values.astype(np.float64),
        extinction,
        z,
        wavelengths,
        pressure,
        temperature,
        particle_angstrom_exponent=particle_angstrom_exponent,
    )  # Fine mode has been chosen in Amstrong exponent

    corrected_signal = np.zeros([iteration_limit, len(elastic_signal)])
    delta_overlap = np.zeros(corrected_signal.shape)
    elastic_backscatter_matrix = np.zeros(corrected_signal.shape)
    overlap_function = np.zeros(corrected_signal.shape)
    correction_factor = np.zeros(corrected_signal.shape)

    i = 0
    corrected_signal[i, :] = elastic_signal
    elastic_backscatter_matrix[i, :] = elastic_backscater
    correction_factor[i, :] = (
        2 * raman_backscatter + molecular_backscatter - elastic_backscatter_matrix[i, :]
    ) / (raman_backscatter + molecular_backscatter)

    # Iteraciones restantes
    for i in range(1, iteration_limit):
        corrected_signal[i, :] = (
            corrected_signal[i - 1, :] * correction_factor[i - 1, :]
        )
        rcs_ = signal_to_rcs(corrected_signal[i, :], z)
        elastic_backscatter_matrix[i, :] = klett_rcs(
            rcs_,
            z,
            molecular_backscatter,
            lr_part=particle_lidar_ratio,
            beta_aer_ref=beta_aer_ref,
            reference=reference,
        )

        correction_factor[i, :] = (
            2 * raman_backscatter
            + molecular_backscatter
            - elastic_backscatter_matrix[i, :]
        ) / (raman_backscatter + molecular_backscatter)

        overlap_function[i, :] = corrected_signal[0, :] / corrected_signal[i, :]

    final_overlap_function = overlap_function[-1, :]

    # Set to zero the values of overlap function below 200 m
    final_overlap_function[: np.where(z > min_overlap_range)[0][0]] = 0

    # Search first 1 in overlap function
    first_1 = np.where(final_overlap_function >= 1)[0][0]

    # Set constant profile to 1 form first 1 up
    final_overlap_function[first_1:] = 1

    final_overlap_function = xr.DataArray(
        final_overlap_function, dims=("range"), coords={"range": z}
    )

    if debugging:
        # Generate a xr.Dataset to save the results
        overlap_function = xr.DataArray(
            overlap_function,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        elastic_backscatter_matrix = xr.DataArray(
            elastic_backscatter_matrix,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        raman_backscatter = xr.DataArray(
            raman_backscatter, dims=("range"), coords={"range": z}
        )
        extinction = xr.DataArray(extinction, dims=("range"), coords={"range": z})
        correction_factor = xr.DataArray(
            correction_factor,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        delta_overlap = xr.DataArray(
            delta_overlap,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        corrected_signal = xr.DataArray(
            corrected_signal,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        dataset = xr.Dataset(
            {
                "overlap_function_matrix": overlap_function,
                "elastic_backscatter_matrix": elastic_backscatter_matrix,
                "raman_backscatter": raman_backscatter,
                "correction_factor": correction_factor,
                "delta_overlap": delta_overlap,
                "corrected_signal": corrected_signal,
                "extinction": extinction,
                "overlap_function": final_overlap_function,
            }
        )

    else:
        dataset = xr.Dataset(
            {
                "overlap_function": final_overlap_function,
            }
        )

    return dataset

def overlap_iterative_klett(
    ff_elastic_signal: xr.DataArray,
    nf_elastic_signal: xr.DataArray,
    meteo_profiles: pd.DataFrame,
    particle_lidar_ratio: float,
    wavelength: float,
    reference: tuple[float, float],
    beta_aer_ref: float = 0,
    min_overlap_range: float = 200,  # minimum range of overlap function
    iteration_limit: int = 10,
    debugging: bool = False,
    **kwargs: dict,
) -> xr.Dataset:
    """Retrieve overlap function by means of the iterative method with near-field Klett as reference.

    Args:
        ff_elastic_signal (xr.DataArray): Far-field elastic lidar signal.
        nf_elastic_signal (xr.DataArray): Near-field elastic lidar signal.
        meteo_profiles (pd.DataFrame): from  `gfatpy.atmo.atmo.generate_meteo_profiles` with pressure and temperature data.
        particle_lidar_ratio (float): Particle lidar ratio to be used in the Klett inversion. It shall be chosen so Klett backscatter profile fits Raman backscatter
        wavelength (float): Elastic wavelength.
        reference (tuple[float, float]): Reference range required by the Klett inversion.
        beta_aer_ref (float, optional): Particle backscatter coefficient at `reference range` in m^-1sr^-1. Defaults to 0 m^-1sr^-1.
        min_overlap_range (float, optional): Overlap function will be forced to zero below [m]. Defaults to 200 m.
        debugging (bool, optional): If True, it will return a dataset with all the intermediate variables. Defaults to False.

    Raises:
        ValueError: Both nf- and ff-elastic signals must have the same 'range' dimension.
        ValueError: Signals range dimension and meteo_profiles['heights'] must have the same size.

    Returns:
        xr.Dataset: overlap function. If debugging is True, it will return a dataset with all the intermediate variables.
    """    

    # Check z in elastic and raman signals (the same for both signals)
    if len(ff_elastic_signal.range) != len(nf_elastic_signal.range):
        raise ValueError(
            "Both nf- and ff-elastic signals must have the same 'range' dimension."
        )

    if len(ff_elastic_signal.range) != len(meteo_profiles["height"]):
        raise ValueError(
            "Signals range dimension and meteo_profiles['heights'] must have the same size."
        )

    z = ff_elastic_signal.range.values.astype(np.float64)

    mol_properties = molecular_properties(
        wavelength,
        meteo_profiles["pressure"],
        meteo_profiles["temperature"],
        heights=z,
    )
    molecular_backscatter = mol_properties["molecular_beta"].values.astype(np.float64)

    # Beta Klett
    ff_elastic_backscater = klett_rcs(
        signal_to_rcs(ff_elastic_signal, z).astype(np.float64),
        z,
        molecular_backscatter,
        lr_part=particle_lidar_ratio,
        reference=reference,
    )

    nf_elastic_backscater = klett_rcs(
        signal_to_rcs(nf_elastic_signal, z).astype(np.float64),
        z,
        molecular_backscatter,
        lr_part=particle_lidar_ratio,
        reference=reference,
    )

    corrected_signal = np.zeros([iteration_limit, len(ff_elastic_signal)])
    delta_overlap = np.zeros(corrected_signal.shape)
    ff_elastic_backscatter_matrix = np.zeros(corrected_signal.shape)
    overlap_function = np.zeros(corrected_signal.shape)
    correction_factor = np.zeros(corrected_signal.shape)

    i = 0
    corrected_signal[i, :] = ff_elastic_signal
    ff_elastic_backscatter_matrix[i, :] = ff_elastic_backscater
    correction_factor[i, :] = (
        2 * nf_elastic_backscater + molecular_backscatter - ff_elastic_backscatter_matrix[i, :]
    ) / (nf_elastic_backscater + molecular_backscatter)

    # Iteraciones restantes
    for i in range(1, iteration_limit):
        corrected_signal[i, :] = (
            corrected_signal[i - 1, :] * correction_factor[i - 1, :]
        )
        rcs_ = signal_to_rcs(corrected_signal[i, :], z)
        ff_elastic_backscatter_matrix[i, :] = klett_rcs(
            rcs_,
            z,
            molecular_backscatter,
            lr_part=particle_lidar_ratio,
            beta_aer_ref=beta_aer_ref,
            reference=reference,
        )

        correction_factor[i, :] = (
            2 * nf_elastic_backscater
            + molecular_backscatter
            - ff_elastic_backscatter_matrix[i, :]
        ) / (nf_elastic_backscater + molecular_backscatter)

        overlap_function[i, :] = corrected_signal[0, :] / corrected_signal[i, :]

    final_overlap_function = overlap_function[-1, :]

    # Set to zero the values of overlap function below 200 m
    final_overlap_function[: np.where(z > min_overlap_range)[0][0]] = 0

    # Search first 1 in overlap function
    first_1 = np.where(final_overlap_function >= 1)[0][0]

    # Set constant profile to 1 form first 1 up
    final_overlap_function[first_1:] = 1

    final_overlap_function = xr.DataArray(
        final_overlap_function, dims=("range"), coords={"range": z}
    )

    if debugging:
        # Generate a xr.Dataset to save the results
        overlap_function = xr.DataArray(
            overlap_function,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        ff_elastic_backscatter_matrix = xr.DataArray(
            ff_elastic_backscatter_matrix,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        nf_elastic_backscatter = xr.DataArray(
            nf_elastic_backscater, dims=("range"), coords={"range": z}
        )
        correction_factor = xr.DataArray(
            correction_factor,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        delta_overlap = xr.DataArray(
            delta_overlap,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        corrected_signal = xr.DataArray(
            corrected_signal,
            dims=("iteration", "range"),
            coords={"iteration": np.arange(iteration_limit), "range": z},
        )
        dataset = xr.Dataset(
            {
                "overlap_function_matrix": overlap_function,
                "elastic_backscatter_matrix": ff_elastic_backscatter_matrix,
                "nf_elastic_backscatter": nf_elastic_backscatter,
                "correction_factor": correction_factor,
                "delta_overlap": delta_overlap,
                "corrected_signal": corrected_signal,
                "overlap_function": final_overlap_function,
            }
        )

    else:
        dataset = xr.Dataset(
            {
                "overlap_function": final_overlap_function,
            }
        )

    return dataset


def overlap_function_explicit(
    elastic_rcs: xr.DataArray,
    raman_rcs: xr.DataArray,
    wavelengths: tuple[float, float],
    reference_heights: tuple[float, float],
    particle_lidar_ratio: float,
    meteo_profiles: pd.DataFrame,
    debugging=False,
    **kwargs: dict,
) -> np.ndarray:
    """Retrieve overlap function by means of the explicit method.

    Args:

        - elastic_rcs (xr.DataArray): Elastic range corrected signal at wavelength `wavelengths[0]`.
        - raman_rcs (xr.DataArray): Raman range corrected signal at wavelength `wavelengths[1]`.
        - wavelengths (tuple[float, float]): Elastic and Raman wavelengths.
        - reference_height (float): Aerosol-free region reference height in meters.
        - particle_lidar_ratio (float): Assumed particle lidar ratio.
        - meteo_profiles (pd.DataFrame): Meteo profiles with pressure and temperature data from `gfatpy.atmo.atmo.generate_meteo_profiles`.

    Returns:

        - xr.Dataset: overlap function. If debugging is True, it will return a dataset with all the intermediate variables.

    References:

        - Comerón et al., 2023: https://doi.org/10.5194/amt-16-3015-2023
    """

    r = elastic_rcs.range.values.astype(np.float64)
    dr = r[1] - r[0]

    # Create an array full of 1 of size r
    array_ones = np.ones(r.shape)
    elastic_mol_properties = molecular_properties(
        wavelengths[0],
        meteo_profiles["pressure"],
        meteo_profiles["temperature"],
        heights=r,
    )
    raman_mol_properties = molecular_properties(
        wavelengths[1],
        meteo_profiles["pressure"],
        meteo_profiles["temperature"],
        heights=r,
    )

    # Computation of beta molecular for elastic and Raman signal
    beta_mol_elastic = elastic_mol_properties["molecular_beta"]  # it has to be in m-1*sr-1
    beta_mol_raman = raman_mol_properties["molecular_beta"]

    # Cálculo inverso de beta molecular, a partir de la señal Rayleigh
    XE = elastic_rcs.values  # elastic signal
    X_Ray = XE  # molecular backscatter normalized to RCS

    XR = raman_rcs.values  # raman signal
    XR_Ray = XR  # molecular backscatter normalized to RCS

    ref_height = np.mean(np.array(reference_heights))  # Reference height (m)
    idx = np.round(ref_height / dr).astype(int)
    idx_min = np.round(reference_heights[0] / dr).astype(int)
    idx_max = np.round(reference_heights[1] / dr).astype(int)

    X_ref = np.nanmean(X_Ray[idx_min:idx_max])
    XR_ref = np.nanmean(XR_Ray[idx_min:idx_max])
    beta_mol_ref = np.nanmean(beta_mol_elastic[idx_min:idx_max])

    norm_XR = XR / XR_ref
    nbeta_molecular = beta_mol_elastic / beta_mol_ref

    Sa = particle_lidar_ratio  # Assumed Lidar Ratio
    Sm_el = elastic_mol_properties["molecular_lidar_ratio"]
    # Sm_ra = raman_mol_properties['molecular_lidar_ratio']

    beta_mol_el = beta_mol_elastic[:idx]
    beta_mol_ra = beta_mol_raman[:idx]
    nbeta_mol = nbeta_molecular[:idx]
    nXR = norm_XR[:idx]
    X = XE[:idx]

    fexp = np.exp(
        -2
        * (Sa - Sm_el.values)
        * np.flip(cumtrapz(np.flip(beta_mol_el), initial=0, dx=dr))
    )

    # Not used because its efect is negligible
    # fexp_dif_ext_mol = np.exp(- Sm_el * np.flip(cumtrapz(np.flip(beta_mol_el - beta_mol_ra), initial=0, dx=dr)))

    g = (nbeta_mol / nXR) * fexp

    phi = (2 * beta_mol_el / (nXR * X_ref)) * fexp

    psi = (X * Sa) / fexp

    phi_psi = phi * psi

    fexp_phi_psi = np.exp(np.flip(cumtrapz(np.flip(phi_psi), initial=0, dx=dr)))

    fint_g_psi = np.flip(cumtrapz(np.flip(g * psi / fexp_phi_psi), initial=0, dx=dr))

    f = g + phi * fexp_phi_psi * fint_g_psi

    Oe = 1 / f

    # Paste Oe into the array of ones so that it has the same length as r
    Oe = np.concatenate((Oe, array_ones[idx:]))

    # Create a dataset with all the intermediate parameters
    Oe = xr.DataArray(Oe, dims=("range"), coords={"range": r})

    if debugging:
        f_ = np.empty_like(r)
        f_[:idx] = f
        g_ = np.empty_like(r)
        g_[:idx] = g
        phi_ = np.empty_like(r)
        phi_[:idx] = phi
        psi_ = np.empty_like(r)
        psi_[:idx] = psi
        phi_psi_ = np.empty_like(r)
        fexp_ = np.empty_like(r)
        fexp_[:idx] = fexp
        phi_psi_[:idx] = phi_psi
        fexp_phi_psi_ = np.empty_like(r)
        fexp_phi_psi_[:idx] = fexp_phi_psi
        fint_g_psi_ = np.empty_like(r)
        fint_g_psi_[:idx] = fint_g_psi

        f = xr.DataArray(f_, dims=("range"), coords={"range": r})
        g = xr.DataArray(g_, dims=("range"), coords={"range": r})
        phi = xr.DataArray(phi_, dims=("range"), coords={"range": r})
        psi = xr.DataArray(psi_, dims=("range"), coords={"range": r})
        fexp = xr.DataArray(fexp_, dims=("range"), coords={"range": r})
        phi_psi = xr.DataArray(phi_psi_, dims=("range"), coords={"range": r})
        fexp_phi_psi = xr.DataArray(fexp_phi_psi_, dims=("range"), coords={"range": r})
        fint_g_psi = xr.DataArray(fint_g_psi_, dims=("range"), coords={"range": r})
        beta_mol_el = xr.DataArray(beta_mol_elastic, dims=("range"), coords={"range": r})
        beta_mol_ra = xr.DataArray(beta_mol_raman, dims=("range"), coords={"range": r})
        nbeta_mol = xr.DataArray(nbeta_molecular, dims=("range"), coords={"range": r})
        nXR = xr.DataArray(norm_XR, dims=("range"), coords={"range": r})
        X = xr.DataArray(XE, dims=("range"), coords={"range": r})
        XR = xr.DataArray(XR, dims=("range"), coords={"range": r})

        dataset = xr.Dataset(
            {
                "overlap_function": Oe,
                "f": f,
                "g": g,
                "phi": phi,
                "psi": psi,
                "phi_psi": phi_psi,
                "fexp_phi_psi": fexp_phi_psi,
                "fint_g_psi": fint_g_psi,
                "beta_mol_el": beta_mol_el,
                "beta_mol_ra": beta_mol_ra,
                "nbeta_mol": nbeta_mol,
                "nXR": nXR,
                "X": X,
                "fexp": fexp,
                "XR": XR,
                "XR_Ray": XR_Ray,
                "X_Ray": X_Ray,
            }
        )
    else:
        dataset = xr.Dataset({"overlap_function": Oe})

    return dataset