import numpy as np
import pytest
import xarray as xr
from scipy.stats import norm
from xarray.tests import assert_allclose

from xskillscore.core.probabilistic import (
    brier_score,
    crps_ensemble,
    crps_gaussian,
    crps_quadrature,
    discrimination,
    rank_histogram,
    reliability,
    rps,
    threshold_brier_score,
)


@pytest.fixture
def threshold():
    return 0.5


@pytest.mark.parametrize("outer_bool", [False, True])
def test_crps_gaussian_accessor(o, f_prob, outer_bool):
    mu = f_prob.mean("member")
    sig = f_prob.std("member")
    actual = crps_gaussian(o, mu, sig)
    ds = xr.Dataset()
    ds["o"] = o
    ds["mu"] = mu
    ds["sig"] = sig
    if outer_bool:
        ds = ds.drop_vars("mu")
        expected = ds.xs.crps_gaussian("o", mu, sig)
    else:
        expected = ds.xs.crps_gaussian("o", "mu", "sig")
    assert_allclose(actual, expected)


@pytest.mark.parametrize("outer_bool", [False, True])
def test_crps_ensemble_accessor(o, f_prob, outer_bool):
    actual = crps_ensemble(o, f_prob)
    ds = xr.Dataset()
    ds["o"] = o
    ds["f_prob"] = f_prob
    if outer_bool:
        ds = ds.drop_vars("f_prob")
        expected = ds.xs.crps_ensemble("o", f_prob)
    else:
        expected = ds.xs.crps_ensemble("o", "f_prob")
    assert_allclose(actual, expected)


@pytest.mark.slow
def test_crps_quadrature_accessor(o):
    # to speed things up
    o = o.isel(time=0, drop=True)
    cdf_or_dist = norm
    actual = crps_quadrature(o, cdf_or_dist)
    ds = xr.Dataset()
    ds["o"] = o
    expected = ds.xs.crps_quadrature("o", cdf_or_dist)
    assert_allclose(actual, expected)


@pytest.mark.parametrize("outer_bool", [False, True])
def test_threshold_brier_score_accessor(o, f_prob, threshold, outer_bool):
    actual = threshold_brier_score(o, f_prob, threshold)
    ds = xr.Dataset()
    ds["o"] = o
    ds["f_prob"] = f_prob
    if outer_bool:
        ds = ds.drop_vars("f_prob")
        expected = ds.xs.threshold_brier_score("o", f_prob, threshold)
    else:
        expected = ds.xs.threshold_brier_score("o", "f_prob", threshold)
    assert_allclose(actual, expected)


@pytest.mark.parametrize("outer_bool", [False, True])
def test_brier_score_accessor(o, f_prob, threshold, outer_bool):
    actual = brier_score(o > threshold, (f_prob > threshold).mean("member"))
    ds = xr.Dataset()
    ds["o"] = o > threshold
    ds["f_prob"] = (f_prob > threshold).mean("member")
    if outer_bool:
        ds = ds.drop_vars("f_prob")
        expected = ds.xs.brier_score("o", (f_prob > threshold).mean("member"))
    else:
        expected = ds.xs.brier_score("o", "f_prob")
    assert_allclose(actual, expected)


@pytest.mark.parametrize("outer_bool", [False, True])
def test_rank_histogram_accessor(o, f_prob, outer_bool):
    actual = rank_histogram(o, f_prob)
    ds = xr.Dataset()
    ds["o"] = o
    ds["f_prob"] = f_prob
    if outer_bool:
        ds = ds.drop_vars("f_prob")
        expected = ds.xs.rank_histogram("o", f_prob)
    else:
        expected = ds.xs.rank_histogram("o", "f_prob")
    assert_allclose(actual, expected)


@pytest.mark.parametrize("outer_bool", [False, True])
def test_discrimination_accessor(o, f_prob, threshold, outer_bool):
    actual = discrimination(o > threshold, (f_prob > threshold).mean("member"))
    ds = xr.Dataset()
    ds["o"] = o > threshold
    ds["f_prob"] = (f_prob > threshold).mean("member")
    if outer_bool:
        ds = ds.drop_vars("f_prob")
        expected = ds.xs.discrimination("o", (f_prob > threshold).mean("member"))
    else:

        expected = ds.xs.discrimination("o", "f_prob")
    assert_allclose(actual, expected)


@pytest.mark.parametrize("outer_bool", [False, True])
def test_reliability_accessor(o, f_prob, threshold, outer_bool):
    actual = reliability(o > threshold, (f_prob > threshold).mean("member"))

    ds = xr.Dataset()
    ds["o"] = o > threshold
    ds["f_prob"] = (f_prob > threshold).mean("member")
    if outer_bool:
        ds = ds.drop_vars("f_prob")
        expected = ds.xs.reliability("o", (f_prob > threshold).mean("member"))
    else:
        expected = ds.xs.reliability("o", "f_prob")
    assert_allclose(actual, expected)


@pytest.mark.parametrize("outer_bool", [False, True])
def test_rps_accessor(o, f_prob, outer_bool):
    category_edges = np.linspace(0, 1 + 1e-8, 6)
    actual = rps(o, f_prob, category_edges=category_edges)

    ds = xr.Dataset()
    ds["o"] = o
    ds["f_prob"] = f_prob
    if outer_bool:
        ds = ds.drop_vars("f_prob")
        expected = ds.xs.rps("o", f_prob, category_edges=category_edges)
    else:
        expected = ds.xs.rps("o", "f_prob", category_edges=category_edges)
    assert_allclose(actual, expected)
