"""Test the inference engines."""

from __future__ import annotations

from typing import Literal, overload
from typing_extensions import override

import pytest
import torch
from numpy.random import default_rng
from torch import Tensor

from tabpfn.architectures.interface import Architecture
from tabpfn.inference import InferenceEngineCachePreprocessing, InferenceEngineOnDemand
from tabpfn.preprocessing import (
    ClassifierEnsembleConfig,
    EnsembleConfig,
    PreprocessorConfig,
)


class TestModel(Architecture):
    @overload
    def forward(
        self,
        x: Tensor | dict[str, Tensor],
        y: Tensor | dict[str, Tensor] | None,
        *,
        only_return_standard_out: Literal[True] = True,
        categorical_inds: list[list[int]] | None = None,
    ) -> Tensor: ...

    @overload
    def forward(
        self,
        x: Tensor | dict[str, Tensor],
        y: Tensor | dict[str, Tensor] | None,
        *,
        only_return_standard_out: Literal[False],
        categorical_inds: list[list[int]] | None = None,
    ) -> dict[str, Tensor]: ...

    @override
    def forward(
        self,
        x: Tensor | dict[str, Tensor],
        y: Tensor | dict[str, Tensor] | None,
        *,
        only_return_standard_out: bool = True,
        categorical_inds: list[list[int]] | None = None,
    ) -> Tensor | dict[str, Tensor]:
        """Perform a forward pass, see doc string of `Architecture`."""
        assert isinstance(x, Tensor)
        assert isinstance(y, Tensor)
        n_train_test, _, _ = x.shape
        n_train, _ = y.shape
        test_rows = n_train_test - n_train
        return x.sum(-2, keepdim=True).sum(-1, keepdim=True).reshape(-1, test_rows)

    @property
    def ninp(self) -> int:
        return 2

    @property
    def features_per_group(self) -> int:
        return 2

    def reset_save_peak_mem_factor(self, factor: int | None = None) -> None:
        pass


def test__cache_preprocessing__result_equal_in_serial_and_in_parallel() -> None:
    rng = default_rng(seed=0)
    n_train = 100
    n_features = 4
    n_classes = 3
    X_train = rng.standard_normal(size=(n_train, n_features))
    y_train = rng.integers(low=0, high=n_classes - 1, size=(n_train, 1))
    X_test = rng.standard_normal(size=(2, n_features))

    engine = InferenceEngineCachePreprocessing.prepare(
        X_train,
        y_train,
        cat_ix=[] * n_train,
        models=[TestModel()],
        ensemble_configs=_create_test_ensemble_configs(
            n_configs=5,
            n_classes=n_classes,
            num_models=1,
        ),
        # We want to test n_preprocessing_jobs>1 as this might mean the outputs are not
        # in the same order as the input configs, and we want to check that the parallel
        # evaluation code behaves correctly in this scenario.
        n_preprocessing_jobs=5,
        rng=rng,
        dtype_byte_size=4,
        force_inference_dtype=None,
        save_peak_mem=True,
        inference_mode=True,
    )

    outputs_sequential = list(
        engine.iter_outputs(X_test, devices=[torch.device("cpu")], autocast=False)
    )
    outputs_parallel = list(
        engine.iter_outputs(
            X_test, devices=[torch.device("cpu"), torch.device("cpu")], autocast=False
        )
    )

    assert len(outputs_sequential) == len(outputs_parallel)
    for par_output, par_config in outputs_parallel:
        seq_output = _find_seq_output(par_config, outputs_sequential)
        assert isinstance(seq_output, Tensor)
        assert isinstance(par_output, Tensor)
        assert torch.allclose(seq_output, par_output)


def test__on_demand__result_equal_in_serial_and_in_parallel() -> None:
    rng = default_rng(seed=0)
    n_train = 100
    n_features = 4
    n_classes = 3
    X_train = rng.standard_normal(size=(n_train, n_features))
    y_train = rng.integers(low=0, high=n_classes - 1, size=(n_train, 1))
    X_test = rng.standard_normal(size=(2, n_features))

    num_models = 3
    models = [TestModel() for _ in range(num_models)]
    engine = InferenceEngineOnDemand.prepare(
        X_train,
        y_train,
        cat_ix=[] * n_train,
        models=models,
        ensemble_configs=_create_test_ensemble_configs(
            n_configs=5,
            n_classes=3,
            num_models=num_models,
        ),
        # We want to test n_preprocessing_jobs>1 as this might mean the outputs are not
        # in the same order as the input configs, and we want to check that the parallel
        # evaluation code behaves correctly in this scenario.
        n_preprocessing_jobs=5,
        rng=rng,
        dtype_byte_size=4,
        force_inference_dtype=None,
        save_peak_mem=True,
    )

    outputs_sequential = list(
        engine.iter_outputs(X_test, devices=[torch.device("cpu")], autocast=False)
    )
    outputs_parallel = list(
        engine.iter_outputs(
            X_test, devices=[torch.device("cpu"), torch.device("cpu")], autocast=False
        )
    )

    assert len(outputs_sequential) == len(outputs_parallel)
    last_model_index = 0
    for par_output, par_config in outputs_parallel:
        # Test that models are executed in order.
        assert par_config._model_index >= last_model_index
        seq_output = _find_seq_output(par_config, outputs_sequential)
        assert isinstance(seq_output, Tensor)
        assert isinstance(par_output, Tensor)
        assert torch.allclose(seq_output, par_output)
        last_model_index = par_config._model_index


def _create_test_ensemble_configs(
    n_configs: int,
    n_classes: int,
    num_models: int,
) -> list[ClassifierEnsembleConfig]:
    preprocessor_configs = [
        PreprocessorConfig(
            "quantile_uni_coarse",
            append_original="auto",
            categorical_name="ordinal_very_common_categories_shuffled",
            global_transformer_name="svd",
            max_features_per_estimator=500,
        ),
        PreprocessorConfig(
            "none",
            categorical_name="numeric",
            max_features_per_estimator=500,
        ),
    ]
    return EnsembleConfig.generate_for_classification(
        num_estimators=n_configs,
        subsample_size=None,
        max_index=n_classes - 1,
        add_fingerprint_feature=True,
        polynomial_features="all",
        feature_shift_decoder="shuffle",
        preprocessor_configs=preprocessor_configs,
        class_shift_method=None,
        n_classes=n_classes,
        random_state=0,
        num_models=num_models,
    )


def _find_seq_output(
    config: EnsembleConfig,
    outputs_sequential: list[tuple[Tensor | dict, EnsembleConfig]],
) -> Tensor | dict:
    """Find the sequential output corresponding to the given config.

    The configs are not hashable, so we have to resort to this search method.
    """
    for output, trial_config in outputs_sequential:
        if trial_config == config:
            return output

    return pytest.fail(f"Parallel config was not found in sequential configs: {config}")
