"""Check that the current output of the model matches reference predictions.

This ensures that the output of a model does not change once it has been released, and
allows us to catch unintended changes to the model.

The predictions can vary slightly between platforms, thus we maintain separate
reference predictions for all the platforms in `ENABLED_PLATFORMS`. The tests are
skipped on other platforms.

If the predictions change legitimately, the reference predictions can be regenerated by
running this file: `python -m tests.test_consistency`.
"""

from __future__ import annotations

import json
import logging
import pathlib
import platform
from dataclasses import dataclass
from functools import partial
from typing import Callable, Union

import numpy as np
import pytest
import torch
from sklearn.datasets import load_iris
from sklearn.utils import check_random_state

from tabpfn import TabPFNClassifier, TabPFNRegressor
from tabpfn.constants import ModelVersion

from .utils import get_pytest_devices

logger = logging.getLogger(__name__)


def _get_tiny_classification_data() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    random_state = check_random_state(0)
    X = random_state.rand(10, 5)  # 10 samples, 5 features
    y = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])  # Binary classification

    # Split into train/test
    X_train, X_test = X[:7], X[7:]
    y_train = y[:7]

    return X_train, y_train, X_test


def _get_tiny_regression_data() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    random_state = check_random_state(0)
    X = random_state.rand(10, 5)  # 10 samples, 5 features
    y = random_state.rand(10) * 10  # Continuous target

    # Split into train/test
    X_train, X_test = X[:7], X[7:]
    y_train = y[:7]

    return X_train, y_train, X_test


def _get_iris_multiclass_data() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Get a small subset of iris data for multiclass testing."""
    # Load iris dataset with 3 well-separated classes
    X, y = load_iris(return_X_y=True)

    # Use fixed test samples (first sample of each class)
    test_indices = [0, 50, 100]  # First sample of each class

    # Use fixed training samples (samples 1-6 from each class)
    train_indices = []
    for base in [0, 50, 100]:  # Start of each class
        train_indices.extend([base + i for i in range(1, 7)])  # 6 samples per class

    X_train, X_test = X[train_indices], X[test_indices]
    y_train, _y_test = y[train_indices], y[test_indices]

    return X_train, y_train, X_test


def _to_tensors(
    xs: tuple[np.ndarray, np.ndarray, np.ndarray],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    return (torch.as_tensor(xs[0]), torch.as_tensor(xs[1]), torch.as_tensor(xs[2]))


def _add_extra_devices(
    model: TabPFNClassifier | TabPFNRegressor,
) -> TabPFNClassifier | TabPFNRegressor:
    # The model does not allow specifying the same device twice, so override
    # after the fact instead.
    # Use lots of devices to maximise the chance of hitting a race condition.
    model.devices_ = (torch.device("cpu"),) * 10
    return model


TensorOrArray = Union[np.ndarray, torch.Tensor]


@dataclass(frozen=True)
class _ConsistencyCase:
    data: Callable[[], tuple[TensorOrArray, TensorOrArray, TensorOrArray]]
    model: Callable[[], TabPFNClassifier | TabPFNRegressor]


DEFAULT_CONFIG = {"n_estimators": 2, "random_state": 42, "device": "cpu"}

TEST_CASES = {
    **{
        f"classifier_tiny_dataset_{version.value}_{fit_mode}": _ConsistencyCase(
            data=_get_tiny_classification_data,
            model=partial(
                TabPFNClassifier.create_default_for_version,
                fit_mode=fit_mode,
                version=version,
                **DEFAULT_CONFIG,
            ),
        )
        for fit_mode in ["fit_preprocessors", "low_memory", "fit_with_cache"]
        for version in [ModelVersion.V2, ModelVersion.V2_5]
        # Save compute by only running all the tests for the latest model.
        if version == ModelVersion.V2_5 or fit_mode == "fit_preprocessors"
    },
    **{
        f"regressor_tiny_dataset_{version.value}_{fit_mode}": _ConsistencyCase(
            data=_get_tiny_regression_data,
            model=partial(
                TabPFNRegressor.create_default_for_version,
                fit_mode=fit_mode,
                version=version,
                **DEFAULT_CONFIG,
            ),
        )
        for fit_mode in ["fit_preprocessors", "low_memory", "fit_with_cache"]
        for version in [ModelVersion.V2, ModelVersion.V2_5]
        # Save compute by only running all the tests for the latest model.
        if version == ModelVersion.V2_5 or fit_mode == "fit_preprocessors"
    },
    "classifier_tiny_dataset_differentiable_input": _ConsistencyCase(
        data=lambda: _to_tensors(_get_tiny_classification_data()),
        model=lambda: TabPFNClassifier(**DEFAULT_CONFIG, differentiable_input=True),
    ),
    "classifier_iris_dataset": _ConsistencyCase(
        data=_get_iris_multiclass_data,
        model=lambda: TabPFNClassifier(**DEFAULT_CONFIG),
    ),
    "regressor_tiny_dataset_several_devices": _ConsistencyCase(
        data=_get_tiny_regression_data,
        model=lambda: _add_extra_devices(TabPFNRegressor(**DEFAULT_CONFIG)),
    ),
    "classifier_iris_dataset_several_devices": _ConsistencyCase(
        data=_get_iris_multiclass_data,
        model=lambda: _add_extra_devices(TabPFNClassifier(**DEFAULT_CONFIG)),
    ),
    "classifier_tiny_dataset_5_estimators": _ConsistencyCase(
        data=_get_tiny_classification_data,
        model=lambda: TabPFNClassifier(**DEFAULT_CONFIG | {"n_estimators": 5}),
    ),
}


ENABLED_PLATFORMS = ["darwin_arm64"]
"""The tests are enabled if _get_current_platform() returns a string in this set."""


def _get_current_platform_string() -> str:
    """Return a string identifying the current platform.

    For now we just test on MacOS, and simply identifying that seems to be enough to get
    consistent predictions. But, this could also be used to e.g. include the device
    string.
    """
    if platform.system() == "Darwin" and platform.machine() == "arm64":
        return "darwin_arm64"
    return "unknown"


HOW_TO_FIX_MESSAGE = (
    "If this is expected, generate the reference predictions by running: "
    "python -m tests.test_consistency"
)
REFERENCE_DIR = pathlib.Path(__file__).parent / "reference_predictions"
"""The directory that contains the reference predictions."""


@pytest.mark.skipif(
    _get_current_platform_string() not in ENABLED_PLATFORMS,
    reason="Current platform does not have consistency tests enabled.",
)
@pytest.mark.skipif("cpu" not in get_pytest_devices(), reason="Only runs on the CPU")
@pytest.mark.parametrize(("test_case_name", "test_case"), TEST_CASES.items())
def test__fit_predict__predictions_match_reference(
    test_case_name: str, test_case: _ConsistencyCase
) -> None:
    reference_predictions = _load_reference_predictions_or_fail(test_case_name)
    predictions = _predict(test_case)

    np.testing.assert_allclose(
        predictions,
        reference_predictions,
        rtol=1e-3,  # 0.1% relative tolerance
        atol=1e-3,  # 0.001 absolute tolerance
        err_msg=(f"Predictions have changed.\n{HOW_TO_FIX_MESSAGE}"),
    )


def _predict(test_case: _ConsistencyCase) -> np.ndarray:
    X_train, y_train, X_test = test_case.data()
    model = test_case.model()
    if model.differentiable_input and isinstance(model, TabPFNClassifier):
        model.fit_with_differentiable_input(X_train, y_train)
    else:
        model.fit(X_train, y_train)

    if isinstance(model, TabPFNClassifier):
        return model.predict_proba(X_test)
    return model.predict(X_test)


def _get_platform_dir() -> pathlib.Path:
    """Return a string that identifies the current platform.

    If two platforms return the same string, then their predictions should match.
    """
    return REFERENCE_DIR / _get_current_platform_string()


def _get_reference_path(test_name: str) -> pathlib.Path:
    return _get_platform_dir() / f"{test_name}.json"


def _load_reference_predictions_or_fail(test_name: str) -> np.ndarray:
    path = _get_reference_path(test_name)
    if not path.exists():
        raise AssertionError(
            f"Reference predictions were missing at {path}\n{HOW_TO_FIX_MESSAGE}"
        )
    with path.open("r") as f:
        return np.array(json.load(f))


def save_reference_predictions() -> None:
    """Generate and save the reference predictions for this platform."""
    _get_platform_dir().mkdir(parents=True, exist_ok=True)
    for test_case_name, test_case in TEST_CASES.items():
        predictions = _predict(test_case)
        path = _get_reference_path(test_case_name)
        with path.open("w") as f:
            json.dump(predictions.tolist(), f, indent=2)
        logger.info(f"Reference predictions saved for {test_case_name} at {path}.")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    save_reference_predictions()
