from __future__ import annotations

import functools
import operator
import pickle
from copy import deepcopy

import cloudpickle
import numpy as np
import pytest

from pydvl.utils.status import Status
from pydvl.value import ValuationResult


@pytest.fixture
def dummy_values(values, names):
    return ValuationResult(
        algorithm="dummy_valuator",
        status=Status.Converged,
        values=np.array(values),
        variances=np.zeros_like(values),
        data_names=names,
        sort=True,
    )


@pytest.mark.parametrize(
    "values, names, ranks_asc", [([], [], []), ([2, 3, 1], ["a", "b", "c"], [2, 0, 1])]
)
def test_sorting(values, names, ranks_asc, dummy_values):
    dummy_values.sort(key="value")
    assert np.alltrue([it.value for it in dummy_values] == sorted(values))
    assert np.alltrue(dummy_values.indices == ranks_asc)
    assert np.alltrue(
        [it.value for it in reversed(dummy_values)] == sorted(values, reverse=True)
    )

    dummy_values.sort(reverse=True)
    assert np.alltrue([it.value for it in dummy_values] == sorted(values, reverse=True))
    assert np.alltrue(dummy_values.indices == list(reversed(ranks_asc)))

    dummy_values.sort(key="index")
    assert np.alltrue(dummy_values.indices == list(range(len(values))))
    assert np.alltrue([it.value for it in dummy_values] == values)


@pytest.mark.parametrize(
    "values, names, ranks_asc", [([], [], []), ([2, 3, 1], ["a", "b", "c"], [2, 0, 1])]
)
def test_dataframe_sorting(values, names, ranks_asc, dummy_values):
    sorted_names = [names[r] for r in ranks_asc]
    try:
        import pandas

        df = dummy_values.to_dataframe(use_names=False)
        assert np.alltrue(df.index.values == ranks_asc)

        df = dummy_values.to_dataframe(use_names=True)
        assert np.alltrue(df.index.values == sorted_names)
        assert np.alltrue(df["dummy_valuator"].values == sorted(values))

        dummy_values.sort(reverse=True)
        df = dummy_values.to_dataframe(use_names=True)
        assert np.alltrue(df.index.values == list(reversed(sorted_names)))
        assert np.alltrue(df["dummy_valuator"].values == sorted(values, reverse=True))
    except ImportError:
        pass


@pytest.mark.parametrize(
    "values, names, ranks_asc", [([], [], []), ([2, 3, 1], ["a", "b", "c"], [2, 0, 1])]
)
def test_iter(names, ranks_asc, dummy_values):
    for rank, it in enumerate(dummy_values):
        assert it.index == ranks_asc[rank]

    for rank, it in enumerate(dummy_values):
        assert it.name == names[ranks_asc[rank]]


@pytest.mark.parametrize(
    "values, names, ranks_asc", [([], [], []), ([2, 3, 1], ["a", "b", "c"], [2, 0, 1])]
)
def test_todataframe(ranks_asc, dummy_values):
    df = dummy_values.to_dataframe()
    assert "dummy_valuator" in df.columns
    assert "dummy_valuator_stderr" in df.columns
    assert np.alltrue(df.index.values == ranks_asc)

    df = dummy_values.to_dataframe(column="val")
    assert "val" in df.columns
    assert "val_stderr" in df.columns
    assert np.alltrue(df.index.values == ranks_asc)

    df = dummy_values.to_dataframe(use_names=True)
    assert np.alltrue(df.index.values == [it.name for it in dummy_values])


@pytest.mark.parametrize(
    "values, names, ranks_asc",
    [([], [], []), ([2.0, 3.0, 1.0, 6.0], ["a", "b", "c", "d"], [2, 0, 1, 3])],
)
def test_indexing(ranks_asc, dummy_values):
    if len(ranks_asc) == 0:
        with pytest.raises(IndexError):
            dummy_values[1]  # noqa
        dummy_values[:2]  # noqa
    else:
        assert ranks_asc[:] == [it.index for it in dummy_values[:]]
        assert ranks_asc[0] == dummy_values[0].index
        assert [ranks_asc[0]] == [it.index for it in dummy_values[[0]]]
        assert ranks_asc[:2] == [it.index for it in dummy_values[:2]]
        assert ranks_asc[:2] == [it.index for it in dummy_values[[0, 1]]]
        assert ranks_asc[:-2] == [it.index for it in dummy_values[:-2]]
        assert ranks_asc[-2:] == [it.index for it in dummy_values[-2:]]
        assert ranks_asc[-2:] == [it.index for it in dummy_values[[-2, -1]]]


def test_get_idx():
    """Test getting by data index"""
    values = np.array([5.0, 2.0, 3.0])
    indices = np.array([3, 4, 2])
    result = ValuationResult(values=values, indices=indices)
    with pytest.raises(IndexError):
        result.get(5)
    for v, idx in zip(values, indices):
        assert v == result.get(idx).value
    result.sort()
    for v, idx in zip(values, indices):
        assert v == result.get(idx).value


def test_updating():
    # Test simple updating
    v = ValuationResult(values=np.array([1.0, 2.0]))
    v.update(0, 1.0)
    assert v.values[0] == 1.0
    assert v.counts[0] == 2

    v.update(1, 4.0)
    assert v.values[1] == 3.0
    assert v.variances[1] == 1.0

    v.update(1, 3.0)
    assert v.values[1] == 3.0
    assert np.isclose(v.variances[1], 2 / 3)

    # Test after sorting
    v = ValuationResult(values=np.array([3.0, 1.0]))
    v.sort()
    v.update(0, 1.0)
    assert v.values[0] == 2.0

    # Test data indexing
    v = ValuationResult(values=np.array([3.0, 1.0]), indices=np.array([3, 4]))
    v.update(4, 1.0)
    assert v.values[1] == 1.0
    assert v.counts[1] == 2


@pytest.mark.parametrize(
    "serialize, deserialize",
    [(pickle.dumps, pickle.loads), (cloudpickle.dumps, cloudpickle.loads)],
)
@pytest.mark.parametrize("values, names", [([], None), ([2.0, 3.0, 1.0], None)])
def test_serialization(serialize, deserialize, dummy_values):
    serded = deserialize(serialize(dummy_values))
    assert dummy_values == serded  # Serialization OK (if __eq__ ok...)
    dummy_values.sort(reverse=True)
    assert dummy_values != serded  # Order checks


@pytest.mark.parametrize("values, names", [([], []), ([2, 3, 1], ["a", "b", "c"])])
def test_equality(values, names, dummy_values):
    assert dummy_values == dummy_values

    c = deepcopy(dummy_values)
    dummy_values.sort(reverse=True)
    assert c != dummy_values

    c2 = ValuationResult(
        algorithm="dummy",
        status=c.status,
        values=c.values,
        variances=c._variances,
        data_names=c._names,
    )
    assert c != c2

    c2 = ValuationResult(
        algorithm=c._algorithm,
        status=Status.Failed,
        values=c.values,
        variances=c.variances,
        data_names=c.names,
    )
    assert c != c2

    c2 = ValuationResult(
        algorithm=c._algorithm,
        status=c._status,
        values=c._values,
        variances=c._variances,
        data_names=c._names,
    )
    c2.sort(c._sort_order)

    assert c == c2

    if len(c) > 0:
        c2 = ValuationResult(
            algorithm=c._algorithm,
            status=c._status,
            values=c._values + 1.0,
            variances=c._variances,
            data_names=c._names,
        )
        assert c != c2


@pytest.mark.parametrize(
    "extra_values", [{"test_value": 1.2}, {"test_value1": 1.2, "test_value2": "test"}]
)
def test_extra_values(extra_values):
    kwargs = dict(
        algorithm="test",
        status=Status.Converged,
        values=np.random.rand(10),
        sort=True,
        test_value=1.2,
    )
    kwargs.update(extra_values)
    result = ValuationResult(**kwargs)
    for k, v in extra_values.items():
        assert getattr(result, k) == v
    # Making sure that the repr dunder method works when using extra values
    repr_string = repr(result)
    for k, v in extra_values.items():
        assert k in repr_string


@pytest.mark.parametrize("size", [1, 10])
@pytest.mark.parametrize("total", [None, 1.0, -1.0])
def test_from_random_creation(size: int, total: float | None):
    result = ValuationResult.from_random(size=size, total=total)
    assert len(result) == size
    assert result.status == Status.Converged
    assert result.algorithm == "random"
    if total is not None:
        assert np.isclose(np.sum(result.values), total)


def test_from_random_creation_errors():
    with pytest.raises(ValueError):
        ValuationResult.from_random(size=0)


def test_adding_random():
    """Test adding multiple valuation results together.

    First we generate a matrix of values, then we split it into multiple subsets
    and create a valuation result for each subset. Then we add all the valuation
    results together and check that the resulting means and variances match with
    those of the original matrix.
    """
    n_samples, n_values, n_subsets = 10, 1000, 12
    values = np.random.rand(n_samples, n_values)
    split_indices = np.sort(np.random.randint(1, n_values, size=n_subsets - 1))
    splits = np.split(values, split_indices, axis=1)
    vv = [
        ValuationResult(
            algorithm="dummy",
            status=Status.Pending,
            values=np.average(s, axis=1),
            variances=np.var(s, axis=1),
            counts=s.shape[1] * np.ones(n_samples),
        )
        for s in splits
    ]
    result: ValuationResult = functools.reduce(operator.add, vv)

    true_means = values.mean(axis=1)
    true_variances = values.var(axis=1)

    assert np.allclose(true_means[result.indices], result.values)
    assert np.allclose(true_variances[result.indices], result.variances)


@pytest.mark.parametrize(
    "indices_1, names_1, values_1, indices_2, names_2, values_2,"
    "expected_indices, expected_names, expected_values",
    [  # Disjoint indices
        (
            [0, 1, 2],
            ["a", "b", "c"],
            [0, 1, 2],
            [3, 4, 5],
            ["d", "e", "f"],
            [3, 4, 5],
            [0, 1, 2, 3, 4, 5],
            ["a", "b", "c", "d", "e", "f"],
            [0, 1, 2, 3, 4, 5],
        ),
        # Overlapping indices (recall that values are running averages)
        (
            [0, 1, 2],
            ["a", "b", "c"],
            [0, 1, 2],
            [1, 2, 3],
            ["b", "c", "d"],
            [3, 4, 5],
            [0, 1, 2, 3],
            ["a", "b", "c", "d"],
            [0, 2, 3, 5],
        ),
        # Overlapping indices with different lengths
        (
            [0, 1, 2],
            ["a", "b", "c"],
            [0, 1, 2],
            [1, 2],
            ["b", "c"],
            [3, 4],
            [0, 1, 2],
            ["a", "b", "c"],
            [0, 2, 3],
        ),
    ],
)
def test_adding_different_indices(
    indices_1,
    names_1,
    values_1,
    indices_2,
    names_2,
    values_2,
    expected_indices,
    expected_names,
    expected_values,
):
    """Test adding valuation results with disjoint indices"""
    v1 = ValuationResult(
        indices=np.array(indices_1), values=np.array(values_1), data_names=names_1
    )
    v2 = ValuationResult(
        indices=np.array(indices_2), values=np.array(values_2), data_names=names_2
    )
    v3 = v1 + v2

    assert np.allclose(v3.indices, np.array(expected_indices))
    assert np.allclose(v3.values, np.array(expected_values))
    assert np.all(v3.names == expected_names)


@pytest.mark.parametrize(
    "indices, index_t, data_names, name_t",
    [
        ([0, 1, 2], np.int32, ["a", "b", "c"], "<U1"),
        ([4, 1, 7], np.int64, [4, 1, 7], np.int64),
        ([4, 1, 7], np.int32, [4, 1, 7], np.float_),
    ],
)
def test_types(indices, index_t, data_names, name_t):
    """Test that types for indices and names are correctly preserved when adding
    valuation results"""

    v = ValuationResult(
        indices=np.array(indices, dtype=index_t),
        values=np.ones(len(indices), dtype=np.float_),
        data_names=np.array(data_names, dtype=name_t),
    )
    assert v.indices.dtype == index_t
    assert v.names.dtype == name_t

    v2 = ValuationResult(
        indices=np.array(indices),
        values=np.ones(len(indices)),
        variances=np.zeros(len(indices)),
        data_names=data_names,
    )
    v += v2
    assert v.indices.dtype == index_t
    assert v.names.dtype == name_t


@pytest.mark.parametrize("data_names", [["a", "b", "c"]])
def test_names(data_names):
    """Test that data names are preserved after addition of results"""

    n = len(data_names)
    v = ValuationResult.from_random(size=n, data_names=data_names)
    v2 = ValuationResult.from_random(size=n, data_names=data_names)

    v += v2
    assert np.all(v.names == np.array(data_names))


@pytest.mark.parametrize("n", [0, 5])
def test_empty(n):
    v = ValuationResult.empty()
    assert len(v) == 0
    v2 = ValuationResult(values=np.arange(n))
    v += v2
    assert len(v2) == n


def test_empty_deprecation():
    with pytest.warns(DeprecationWarning):
        v3 = ValuationResult.empty(indices=[1, 2, 3])
