
import xarray as xr
import numpy as np
import xskillscore as xs
from sklearn.metrics import average_precision_score

import matplotlib.pyplot as plt
import cartopy.crs as ccrs
from matplotlib.colors import LinearSegmentedColormap
from matplotlib import ticker

import tinysunset as ts

## TODO: integrate with comp_metrics
def skill_scores(
    metrics: dict, 
    ref_metrics: dict):
    """
    Compute skill scores for metrics. Skill scores are metrics referenced 
    to some benchmark metrics e.g. climatology/persistence or another model.
    
    metrics: dict
    
    """
    
    _SKILL_METRICS = ["crps", "rps", "bs", "rmse", "mae", "mse"]
    
    for metric in metrics:
        
        if metric in _SKILL_METRICS:
            
            metrics["{0}s".format(metric)] = (1 - (metrics[metric] / ref_metrics[metric]))
            
        else:
            
            metrics["{0}s".format(metric)] = metrics[metric] - ref_metrics[metric]
    
    return metrics

def comp_metrics(
    fcst: xr.DataArray, 
    obs: xr.DataArray, 
    metrics: list,
    time_dim: str = "time",
    ensemble_dim: str = "ensemble",
    rps_thresholds: np.ndarray = np.array([-0.431, 0.431]),
    bs_threshold: float = -1,
    bootstrap: bool = False,
    ):
    
    """
    Compute forecast metrics.
    
    """

    METRICS = {
        "crps": lambda: xs.crps_ensemble(obs, fcst, dim=time_dim, member_dim=ensemble_dim),
        "rps": lambda: xs.rps(obs, fcst, dim=time_dim, member_dim=ensemble_dim, category_edges=rps_thresholds),
        "bs": lambda: xs.threshold_brier_score(obs, fcst, threshold=bs_threshold, dim=time_dim, member_dim=ensemble_dim),
        "rmse": lambda: xs.rmse(fcst.median(ensemble_dim), obs, dim=time_dim),
        "mae": lambda: xs.mae(fcst.median(ensemble_dim), obs, dim=time_dim),
        "mse": lambda: xs.mse(fcst.median(ensemble_dim), obs, dim=time_dim),
        "meanbias": lambda: (fcst.mean([time_dim, ensemble_dim]) - obs).mean(time_dim),
        "medianbias": lambda: (fcst.median(ensemble_dim) - obs).mean(time_dim),
        "ensmeanpearsoncorr": lambda: xr.corr(obs, fcst.mean(ensemble_dim), dim=time_dim),
        "ensmedianpearsoncorr": lambda: xr.corr(obs, fcst.median(ensemble_dim), dim=time_dim),
        "ensmeanspearmancorr": lambda: xs.spearman_r(obs, fcst.mean(ensemble_dim), dim=time_dim),
        "ensmedianspearmancorr": lambda: xs.spearman_r(obs, fcst.median(ensemble_dim), dim=time_dim),
        "climbias": lambda: (fcst.mean([time_dim, ensemble_dim]) - obs.mean(time_dim))
    }

    # Verify all requested metrics are available
    missing_metrics = set(metrics) - set(METRICS.keys())
    if missing_metrics:
        raise ValueError(f"The following metrics were not found: {missing_metrics}")

    
    METRICS_DS = []
    for metric in metrics:

        METRICS_DS.append(METRICS[metric]().rename(metric))
        
        if bootstrap:
        
            test = ts.verification._xbootstrap(fcst, obs, METRICS[metric], boot_dim=time_dim, n_sim=1000)
            
            test_random = ts.verification._xbootstrap(
                fcst.copy()*np.random.rand(*fcst.shape)/fcst, # random forecast
                obs, 
                METRICS[metric], boot_dim=time_dim, n_sim=1000
                )

            METRICS_DS.append(
                xr.apply_ufunc(
                    ts.verification._find_crossover_quantile,
                    test,
                    test_random,
                    input_core_dims=[["n_sim"], ["n_sim"]],
                    output_core_dims=[[]],
                    vectorize=True,
                    dask="parallelized",
                    output_dtypes=[float]
                ).rename(f"{metric}_pval")
            )
    
    return xr.merge(METRICS_DS)
            
CBAR_KWARGS = {
    "shrink": 0.8, 
    "pad": 0.05,
    "orientation": "horizontal", #"label": None
    }
    
def plot_metric(
    metric: xr.DataArray, 
    ax: plt.Axes = None,
    mask: xr.DataArray = None,
    projection: ccrs.Projection = ccrs.PlateCarree(),
    add_colorbar: bool = True,
    cbar_kwargs: dict = CBAR_KWARGS, 
    subtitles_kwargs: dict = {"fontsize": 40},
    cbar_fontsize: int = 30,
    **kwargs
    ):
    
    """
    Plot forecast metrics.
    """

    # ---------------------------
    #  CUSTOM COLORMAPS

    _auprc_cmap = plt.cm.YlOrRd
    _auprc_cmap.set_under(color='white')

    _crpss_cmap = plt.cm.YlOrRd
    _crpss_cmap.set_under(color='white')

    # Define the number of color segments for the lower and upper parts of the colormap
    n_segments = 28

    # Define the midpoint of the colormap
    midpoint = 0.5

    # Create the lower part of the colormap (blue to white)
    lower_colors = plt.cm.Blues_r(np.linspace(0.2, 1, n_segments))

    # Create the middle (white) part of the colormap
    n_white = 80 - 2 * n_segments  # Calculate the number of white segments
    white_colors = plt.cm.bwr(np.ones(n_white) * midpoint)

    # Create the upper part of the colormap (white to red)
    upper_colors = plt.cm.YlOrRd(np.linspace(0.2, 1, n_segments))

    # Combine all parts of the colormap
    combined_colors = np.vstack((lower_colors, white_colors, upper_colors))

    # Create the custom colormap
    _div_cmap = LinearSegmentedColormap.from_list('custom_diverging', combined_colors)

    def create_custom_cmap():
        # Define the color segments
        colors = [
            (0.0, (0, 0, 0.5)),       # Dark blue
            (0.05, (0, 0, 0.75)),     # Medium blue
            (0.1, (0, 0.2, 1)),       # Blue
            (0.15, (0, 0.5, 1)),      # Light blue
            (0.2, (0.2, 0.7, 1)),     # Very light blue
            (0.3, (0.4, 0.8, 1)),     # Pale blue
            (0.4, (0.8, 0.8, 0.8)),   # Light gray
            (0.6, (0.8, 0.8, 0.8)),   # Light gray
            (0.7, (1, 1, 0)),         # Yellow
            (0.8, (1, 0.8, 0)),       # Light orange
            (0.85, (1, 0.6, 0)),      # Orange
            (0.9, (1, 0.4, 0)),       # Dark orange
            (0.95, (1, 0, 0)),        # Red
            (1.0, (0.5, 0, 0))        # Dark red
        ]
        
        # Create the colormap
        cmap = LinearSegmentedColormap.from_list("custom", colors, N=256)
        return cmap
        
    
    # Create the custom colormap
    _div_cmap = create_custom_cmap()
    METRIC_KWARGS = {

        "crpss": dict(cmap=_crpss_cmap, vmin=0, vmax=1, levels=21),
        "rpss": dict(cmap=_crpss_cmap, vmin=0, vmax=1, levels=21),
        "bss": dict(cmap=_crpss_cmap, vmin=0, vmax=1, levels=21),
        "mses": dict(cmap=_crpss_cmap, vmin=0, vmax=1, levels=21),
        "rmses": dict(cmap=_crpss_cmap, vmin=0, vmax=1, levels=21),
        "maes": dict(cmap="RdBu_r", vmin=-1, vmax=1, levels=21),
        "auprc": dict(cmap=_auprc_cmap, vmin=0.2, vmax=1, levels=11),
        "mse": dict(cmap="RdBu_r", vmin=0, vmax=3, levels=21),
        "rmse": dict(cmap="RdBu_r", vmin=0, vmax=3, levels=21),
        "mae": dict(cmap="viridis", vmin=0, vmax=3, levels=21),
        "ensmeanspearmancorr": dict(cmap=_div_cmap, vmin=-1, vmax=1, levels=21),
        "ensmedianspearmancorr": dict(cmap=_div_cmap, vmin=-1, vmax=1, levels=21),
        "ensmeanpearsoncorr": dict(cmap=_div_cmap, vmin=-1, vmax=1, levels=21),
        "ensmedianpearsoncorr": dict(cmap=_div_cmap, vmin=-1, vmax=1, levels=21),
        "climbias": dict(cmap="bwr", vmin=-3, vmax=3, levels=21),
        "meanbias": dict(cmap="bwr", vmin=-3, vmax=3, levels=13),
        "medianbias": dict(cmap="bwr", vmin=-3, vmax=3, levels=21),
        "fcst": dict(cmap="bwr", vmin=-3, vmax=3, levels=21),
        "obs": dict(cmap="bwr", vmin=-3, vmax=3, levels=21),
        "fcstensmean": dict(cmap="bwr", vmin=-3, vmax=3, levels=21),

    }

    METRIC_TITLE = {

        "crpss": "CRPSS",
        "rpss": "RPSS",
        "bss": "BSS",
        "auprc_neg1sig": "AUPRC [-1$\sigma$]",
        "auprc_1sig": "AUPRC [1$\sigma$]",
        "auprc_neg2sig": "AUPRC [-2$\sigma$]",
        "auprc_2sig": "AUPRC [2$\sigma$]",
        #"ensmeancorr": "ens. mean corr.",
        #"ensmediancorr": "ens. median corr.",
        "mse": "MSE",
        "rmse": "RMSE",
        "rmses": "RMSS",
        "mses": "MSS",
        "maes": "MAS",
        "mae": "MAE",
        "ensmeanspearmancorr": "ens. mean Spearman corr.",
        "ensmedianspearmancorr": "ens. median Spearman corr.",
        "ensmeanpearsoncorr": "ens. mean Pearson corr.",
        "ensmedianpearsoncorr": "ens. median Pearson corr.",
        "climbias": "Clim. bias",
        "meanbias": "Mean bias",
        "medianbias": "Median bias",
        "fcst": "fcst",
        "obs": "obs",
        "fcstensmean": "fcst ens. mean"
        
    }

    ax, im = ts.plots.plot_array_to_map(
        metric, 
        mask=mask,
        ax=ax,
        projection=projection, 
        cbar_kwargs=cbar_kwargs,
        add_colorbar=add_colorbar,
        **METRIC_KWARGS[metric.name.split("_")[0]],
        **kwargs
    )

    cbar = ax.collections[0].colorbar

    for l in cbar.ax.xaxis.get_ticklabels():
        l.set_weight("bold")
        l.set_fontsize(cbar_fontsize)

    # constrain the number of ticks labels plotted in the cbar to 5
    cbar.ax.locator_params(nbins=10)

    plt.title(METRIC_TITLE[metric.name], **subtitles_kwargs)

    return ax, im