import xarray as xr
import numpy as np
from typing import Callable
from pathlib import Path
from strataer.utils import regridding_weights, planck, linear_weight


BASE_FILENAME = "input4MIPs_aerosolProperties_CMIP_UOEXETER-CMIP-1-3-0_gnz_175001-202312.nc"
clim_filename = "input4MIPs_aerosolProperties_CMIP_UOEXETER-CMIP-1-3-0_gnz_185001-202112-clim.nc"


def planck_weights(
    wavelength_1: float | np.ndarray, wavelength_2: float | np.ndarray, temperature: float
) -> float | np.ndarray:
    """
    compute the area under of the plank curve between wavelength_1 and wavelength_2.
    If inputs are numpy arrays they should be of the same size.
    """
    if isinstance(wavelength_2, np.ndarray) and isinstance(wavelength_1, np.ndarray):
        weights = []
        for w1, w2 in zip(wavelength_1, wavelength_2):
            wavel = np.sort(np.linspace(w1, w2, 1000).flatten())
            weights.append(np.trapezoid(planck(wavel, temperature), wavel))
        return np.array(weights)
    else:
        wavel = np.sort(np.linspace(wavelength_1, wavelength_2, 1000).flatten())
        return np.trapezoid(planck(wavel, temperature), wavel)


def solar_planck_weights(wavelength_1: float | np.ndarray, wavelength_2: float | np.ndarray):
    return planck_weights(wavelength_1=wavelength_1, wavelength_2=wavelength_2, temperature=5777.0)


def terrestrial_planck_weights(wavelength_1: float | np.ndarray, wavelength_2: float | np.ndarray):
    return planck_weights(wavelength_1=wavelength_1, wavelength_2=wavelength_2, temperature=254.356)


def wavel_interpolation_weights(
    input_wavel_bnds: xr.DataArray,
    output_wavel_bnds: xr.DataArray,
    input_wavel: xr.DataArray = None,
    output_wavel: xr.DataArray = None,
    weight_function: Callable = solar_planck_weights,
) -> xr.DataArray:
    """
    Determine the wavelength interpolation matrix given an input and output grid.

    Parameters
    ----------
    input_wavel_bnds : xr.DataArray
        wavelength bounds of input grid
    output_wavel_bnds : xr.DataArray
        wavelength bounds of output grid
    input_wavel : xr.DataArray, optional
        midpoints of input grid, by default None
    output_wavel : xr.DataArray, optional
        midpoints of output grid, by default None

    Returns
    -------
    xr.DataArray
        2d matrix of weights
    """

    input_bounds = np.sort(np.unique(input_wavel_bnds.values.flatten()))
    output_bounds = np.sort(np.unique(output_wavel_bnds.values.flatten()))
    weights = regridding_weights(input_bounds, output_bounds, weight_function=weight_function)

    if input_wavel is None:
        input_wavel = input_bounds[0:-1] + np.diff(input_bounds) / 2

    if output_wavel is None:
        output_wavel = output_bounds[0:-1] + np.diff(output_bounds) / 2

    return xr.DataArray(
        weights,
        dims=["new_wavelength", "wavelength"],
        coords=[output_wavel, input_wavel],
    )


def interpolate_variable(
    folder: Path,
    output_wavel_bnds: xr.DataArray,
    variable: str,
    base_filename: str = BASE_FILENAME,
    weight_function=linear_weight,
    weights: xr.DataArray = None,
):
    """
    no additional weighting is required
    """
    file = folder / f"{variable}_{base_filename}"
    data = xr.open_mfdataset(file)
    w = wavel_interpolation_weights(
        data.wavelength_bnds,
        output_wavel_bnds,
        data.wavelength.values,
        output_wavel_bnds.wavelength.values,
        weight_function=weight_function,
    )

    if weights is None:
        ds = (data[variable] @ w).rename({"new_wavelength": "wavelength"})
    else:
        ds = (((data[variable] * weights) @ w) / (weights @ w)).rename(
            {"new_wavelength": "wavelength"}
        )

    ds = xr.merge([ds.rename(variable), data[["time_bnds", "lat_bnds", "height_bnds"]]])
    ds.attrs = data.attrs
    return ds


def interpolate_ext(
    folder: Path,
    output_wavel_bnds: xr.DataArray,
    base_filename: str = BASE_FILENAME,
    weight_function=solar_planck_weights,
):
    """
    no additional weighting is required
    """
    return interpolate_variable(
        folder, output_wavel_bnds, "ext", base_filename, weight_function=weight_function
    )


def interpolate_asy(
    folder,
    output_wavel_bnds: xr.DataArray,
    base_filename: str = BASE_FILENAME,
    weight_function=solar_planck_weights,
):
    """
    assymetry factor is weighted by scattering extinction
    """
    file = folder / f"ext_{base_filename}"
    ext = xr.open_mfdataset(file)

    file = folder / f"ssa_{base_filename}"
    ssa = xr.open_mfdataset(file)

    weights = ext.ext * ssa.ssa
    return interpolate_variable(
        folder,
        output_wavel_bnds,
        "asy",
        base_filename,
        weight_function=weight_function,
        weights=weights,
    )


def interpolate_ssa(
    folder,
    output_wavel_bnds: xr.DataArray,
    base_filename: str = BASE_FILENAME,
    weight_function=solar_planck_weights,
):
    """single scatter albedo is weighted by total extinction"""
    file = folder / f"ext_{base_filename}"
    ext = xr.open_mfdataset(file)

    weights = ext.ext
    return interpolate_variable(
        folder,
        output_wavel_bnds,
        "ssa",
        base_filename,
        weight_function=weight_function,
        weights=weights,
    )
