import xarray as xr
from libs.data_utils import to_180, subset_latlon
from libs.pipeline_utils import normalize_probabilities


def load_raw_nmme_tercile_for_plot(
    filepath,
    *,
    lead,
    lat_min=None,
    lat_max=None,
    lon_min=None,
    lon_max=None,
):
    ds = xr.open_dataset(filepath, decode_times=False)

    target_index = int(lead)
    if target_index < 0 or target_index >= ds.sizes["target"]:
        raise ValueError(
            f"Requested lead={lead} but target dimension has size {ds.sizes['target']}."
        )

    below = ds["prob_below"].isel(target=target_index)
    near = ds["prob_norm"].isel(target=target_index)
    above = ds["prob_above"].isel(target=target_index)

    prob = xr.concat(
        [below, near, above],
        dim=xr.IndexVariable("category", ["below", "near", "above"]),
    ).transpose("category", "lat", "lon")

    prob = prob.sortby("lat").sortby("lon")
    prob = to_180(prob)

    if None not in (lat_min, lat_max, lon_min, lon_max):
        prob = subset_latlon(
            prob,
            lat_min=lat_min,
            lat_max=lat_max,
            lon_min=lon_min,
            lon_max=lon_max,
            lon_mode="180",
            name="raw_nmme_tercile",
        )

    prob = prob / 100.0
    prob = normalize_probabilities(prob, category_dim="category")
    prob.name = "raw_nmme_probability"
    return prob
