import numpy as np
import xarray as xr


def _last_item_cond_true(cond, dim):
    """Return the final item from cond that evaluates to True.

    Args:
        cond (xr.DataArray, xr.Dataset): User-defined boolean array where True means
            the system is predictable at the given lead. E.g., this could be based on
            the dynamical forecast beating a reference forecast, p values, confidence
            intervals, etc. cond should contain the dimension dim at the minimum.
        dim (str): Dimension to check for condition == True over.

    Returns:
        xr.DataArray, xr.Dataset: ``dim`` value until condition is True.

    """
    # force DataArray because isel (when transforming to dim space) requires DataArray
    if isinstance(cond, xr.Dataset):
        was_dataset = True
        cond = cond.to_array()
    else:
        was_dataset = False
    # index last True
    reached = cond.argmin(dim)
    # fix below one
    reached = reached.where(reached >= 1, np.nan)
    # reset where always true to len(lead)
    reached = reached.where(~cond.all("lead"), other=cond[dim].size)
    # fix locations where always nan to nan
    mask = cond.notnull().all("lead")
    reached = reached.where(mask, other=np.nan)
    # shift back into coordinate space
    # problem: cannot convert nan to idx in isel
    # therefore set to dim:0 and mask again afterwards
    reached_notnull = reached.notnull()  # remember where not masked
    reached = reached.where(
        reached.notnull(), other=cond.isel({dim: 0})
    )  # set nan to dim:0
    # take one index before calculated by argmin
    reached_dim_space = cond[dim].isel(
        {dim: reached.astype(int) - 1}
    )  # to not break conversion to dim space
    reached_dim_space = reached_dim_space.where(
        reached_notnull, other=np.nan
    )  # cleanup replace dim:0 with nan again
    if was_dataset:
        reached_dim_space = reached_dim_space.to_dataset(dim="variable").squeeze(
            drop=True
        )
        if "lead" in reached_dim_space.coords:
            reached_dim_space = reached_dim_space.drop_vars("lead")
    return reached_dim_space


def horizon(cond):
    """Calculate the predictability horizon based on a condition ```cond``.

    Args:
        cond (xr.DataArray, xr.Dataset): User-defined boolean array where True means
            the system is predictable at the given lead. E.g., this could be based on
            the dynamical forecast beating a reference forecast, p values, confidence
            intervals, etc. cond should contain the dimension lead at the minimum.

    Returns:
        xr.DataArray, xr.Dataset: predictability horizon reduced by ``lead`` dimension.

    Example:
        >>> skill = PerfectModelEnsemble.verify(
        ...     metric="acc",
        ...     comparison="m2e",
        ...     dim=["init", "member"],
        ...     reference=["persistence"],
        ... )
        >>> horizon(skill.sel(skill="initialized") > skill.sel(skill="persistence"))
        <xarray.Dataset>
        Dimensions:  ()
        Data variables:
            tos      float64 15.0
        Attributes:
            units:          years
            standard_name:  forecast_period
            long_name:      Lead
            description:    Forecast period is the time interval between the forecast...


        >>> bskill = PerfectModelEnsemble.bootstrap(
        ...     metric="acc",
        ...     comparison="m2e",
        ...     dim=["init", "member"],
        ...     reference="uninitialized",
        ...     iterations=201,
        ...     resample_dim="init",
        ... )
        >>> horizon(bskill.sel(skill="uninitialized", results="p") <= 0.05)
        <xarray.Dataset>
        Dimensions:  ()
        Coordinates:
            skill    <U13 'uninitialized'
            results  <U12 'p'
        Data variables:
            tos      float64 2.0
        Attributes:
            units:          years
            standard_name:  forecast_period
            long_name:      Lead
            description:    Forecast period is the time interval between the forecast...

    """
    ph = _last_item_cond_true(cond, "lead")
    if isinstance(ph, xr.DataArray):
        ph.attrs["units"] = cond["lead"].attrs["units"]
    elif isinstance(ph, xr.Dataset):
        for v in ph.data_vars:
            ph[v].attrs["units"] = cond["lead"].attrs["units"]
    return ph
