from functools import partial
from itertools import product
import numpy as np
import pytest
import torch
import torch.nn as nn
from typing import Callable, List
from alibi_detect.cd.pytorch.mmd_online import MMDDriftOnlineTorch
from alibi_detect.cd.pytorch.preprocess import HiddenOutput, preprocess_drift
from alibi_detect.utils._random import fixed_seed

n, n_hidden, n_classes = 400, 10, 5


class MyModel(nn.Module):
    def __init__(self, n_features: int):
        super().__init__()
        self.dense1 = nn.Linear(n_features, 20)
        self.dense2 = nn.Linear(20, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = nn.ReLU()(self.dense1(x))
        return self.dense2(x)


def preprocess_list(x: List[np.ndarray]) -> np.ndarray:
    if len(x) > 1:  # test List[Any] reference data inputs to the detector with Any=np.ndarray
        return np.concatenate(x, axis=0)
    else:  # test Any inputs to the prediction function of the detector with Any=List[np.ndarray]
        return np.array(x)[0]


n_features = [10]
ert = [25]
window_size = [5]
preprocess = [
    (None, None),
    (preprocess_drift, {'model': HiddenOutput, 'layer': -1}),
    (preprocess_list, None)
]
n_bootstraps = [200]
tests_mmddriftonline = list(product(n_features, ert, window_size, preprocess, n_bootstraps))
n_tests = len(tests_mmddriftonline)


@pytest.fixture
def mmd_online_params(request):
    return tests_mmddriftonline[request.param]


@pytest.mark.parametrize('mmd_online_params', list(range(n_tests)), indirect=True)
def test_mmd_online(mmd_online_params, seed):
    n_features, ert, window_size, preprocess, n_bootstraps = mmd_online_params

    with fixed_seed(seed):
        x_ref = np.random.randn(n * n_features).reshape(n, n_features).astype(np.float32)
    preprocess_fn, preprocess_kwargs = preprocess
    to_list = False
    if hasattr(preprocess_fn, '__name__') and preprocess_fn.__name__ == 'preprocess_list':
        to_list = True
        x_ref = [_[None, :] for _ in x_ref]
    elif isinstance(preprocess_fn, Callable) and 'layer' in list(preprocess_kwargs.keys()) \
            and preprocess_kwargs['model'].__name__ == 'HiddenOutput':
        model = MyModel(n_features)
        layer = preprocess_kwargs['layer']
        preprocess_fn = partial(preprocess_fn, model=HiddenOutput(model=model, layer=layer))
    else:
        preprocess_fn = None

    with fixed_seed(seed):
        cd = MMDDriftOnlineTorch(
            x_ref=x_ref,
            ert=ert,
            window_size=window_size,
            preprocess_fn=preprocess_fn,
            n_bootstraps=n_bootstraps
        )
        x_h0 = np.random.randn(n * n_features).reshape(n, n_features).astype(np.float32)
        x_h1 = np.random.randn(n * n_features).reshape(n, n_features).astype(np.float32) + 1

    detection_times_h0 = []
    test_stats_h0 = []
    for x_t in x_h0:
        if to_list:
            x_t = [x_t]
        pred_t = cd.predict(x_t, return_test_stat=True)
        test_stats_h0.append(pred_t['data']['test_stat'])
        if pred_t['data']['is_drift']:
            detection_times_h0.append(pred_t['data']['time'])
            cd.reset_state()
    average_delay_h0 = np.array(detection_times_h0).mean()
    test_stats_h0 = [ts for ts in test_stats_h0 if ts is not None]
    assert ert/3 < average_delay_h0 < 3*ert

    cd.reset_state()

    detection_times_h1 = []
    test_stats_h1 = []
    for x_t in x_h1:
        if to_list:
            x_t = [x_t]
        pred_t = cd.predict(x_t, return_test_stat=True)
        test_stats_h1.append(pred_t['data']['test_stat'])
        if pred_t['data']['is_drift']:
            detection_times_h1.append(pred_t['data']['time'])
            cd.reset_state()
    average_delay_h1 = np.array(detection_times_h1).mean()
    test_stats_h1 = [ts for ts in test_stats_h1 if ts is not None]
    assert np.abs(average_delay_h1) < ert/2

    assert np.mean(test_stats_h1) > np.mean(test_stats_h0)


def test_mmd_online_state_online(tmp_path, seed):
    """
    Test save/load/reset state methods for MMDDriftOnlineTorch. State is saved, reset, and loaded, with
    prediction results and stateful attributes compared to original.
    """
    n = 100
    with fixed_seed(seed):
        x_ref = np.random.normal(0, 1, (n, n_classes))
        x = np.random.normal(0.1, 1, (n, n_classes))
        dd = MMDDriftOnlineTorch(x_ref, window_size=10, ert=20)
        # Store state for comparison
        state_dict_t0 = {}
        for key in dd.online_state_keys:
            state_dict_t0[key] = getattr(dd, key)

    # Run for 10 time steps
    test_stats_1 = []
    for t, x_t in enumerate(x):
        if t == 5:
            dd.save_state(tmp_path)
            # Store state for comparison
            state_dict_t5 = {}
            for key in dd.online_state_keys:
                state_dict_t5[key] = getattr(dd, key)
        preds = dd.predict(x_t)
        test_stats_1.append(preds['data']['test_stat'])

    # Reset and check state cleared
    dd.reset_state()
    for key, orig_val in state_dict_t0.items():
        np.testing.assert_array_equal(orig_val, getattr(dd, key))  # use np.testing here as it handles torch.Tensor etc

    # Repeat, check that same test_stats both times
    test_stats_2 = []
    for t, x_t in enumerate(x):
        preds = dd.predict(x_t)
        test_stats_2.append(preds['data']['test_stat'])
    np.testing.assert_array_equal(test_stats_1, test_stats_2)

    # Load state from t=5 timestep
    dd.load_state(tmp_path)

    # Compare stateful attributes to original at t=5
    for key, orig_val in state_dict_t5.items():
        np.testing.assert_array_equal(orig_val, getattr(dd, key))  # use np.testing here as it handles torch.Tensor etc

    # Compare predictions to original at t=5
    new_pred = dd.predict(x[5])
    assert new_pred['data']['test_stat'] == test_stats_1[5]
