from __future__ import annotations

from unittest.mock import MagicMock

import numpy as np
import pandas as pd
import pytest
import torch
from sklearn.preprocessing import LabelEncoder

from tabpfn import TabPFNClassifier
from tabpfn.constants import NA_PLACEHOLDER
from tabpfn.inference_config import InferenceConfig
from tabpfn.preprocessors.preprocessing_helpers import get_ordinal_encoder
from tabpfn.utils import (
    balance_probas_by_class_counts,
    fix_dtypes,
    infer_categorical_features,
    infer_devices,
    process_text_na_dataframe,
    validate_Xy_fit,
)


def test_infer_categorical_with_str_and_nan_provided_included():
    X = np.array([[np.nan, "NA"]], dtype=object).reshape(-1, 1)
    out = infer_categorical_features(
        X,
        provided=[0],
        min_samples_for_inference=0,
        max_unique_for_category=2,
        min_unique_for_numerical=5,
    )
    assert out == [0]


def test_infer_categorical_with_str_and_nan_multiple_rows_provided_included():
    X = np.array([[np.nan], ["NA"], ["NA"]], dtype=object)
    out = infer_categorical_features(
        X,
        provided=[0],
        min_samples_for_inference=0,
        max_unique_for_category=2,
        min_unique_for_numerical=5,
    )
    assert out == [0]


def test_infer_categorical_auto_inference_blocked_when_not_enough_samples():
    X = np.array([[1.0], [1.0], [np.nan]])
    out = infer_categorical_features(
        X,
        provided=None,
        min_samples_for_inference=3,
        max_unique_for_category=2,
        min_unique_for_numerical=4,
    )
    assert out == []


def test_infer_categorical_auto_inference_enabled_with_enough_samples():
    X = np.array([[1.0, 0.0], [1.0, 1.0], [2.0, 2.0], [2.0, 3.0], [np.nan, 9.0]])
    out = infer_categorical_features(
        X,
        provided=None,
        min_samples_for_inference=3,
        max_unique_for_category=3,
        min_unique_for_numerical=4,
    )
    assert out == [0]


def test_infer_categorical_provided_column_excluded_if_exceeds_max_unique():
    X = np.array([[0], [1], [2], [3], [np.nan]], dtype=float)
    out = infer_categorical_features(
        X,
        provided=[0],
        min_samples_for_inference=0,
        max_unique_for_category=3,
        min_unique_for_numerical=2,
    )
    assert out == []


def test_infer_categorical_with_dict_raises_error():
    X = np.array([[{"a": 1}], [{"b": 2}]], dtype=object)
    with pytest.raises(TypeError):
        infer_categorical_features(
            X,
            provided=None,
            min_samples_for_inference=0,
            max_unique_for_category=2,
            min_unique_for_numerical=2,
        )


def test__infer_devices__auto__cuda_and_mps_not_available__selects_cpu(
    mocker: MagicMock,
) -> None:
    mocker.patch("torch.cuda").is_available.return_value = False
    mocker.patch("torch.backends.mps").is_available.return_value = False
    assert infer_devices(devices="auto") == (torch.device("cpu"),)


def test__infer_devices__auto__single_cuda_gpu_available__selects_it(
    mocker: MagicMock,
) -> None:
    mock_cuda = mocker.patch("torch.cuda")
    mock_cuda.is_available.return_value = True
    mock_cuda.device_count.return_value = 1
    mocker.patch("torch.backends.mps").is_available.return_value = True
    assert infer_devices(devices="auto") == (torch.device("cuda:0"),)


def test__infer_devices__auto__multiple_cuda_gpus_available__selects_first(
    mocker: MagicMock,
) -> None:
    mock_cuda = mocker.patch("torch.cuda")
    mock_cuda.is_available.return_value = True
    mock_cuda.device_count.return_value = 3
    mocker.patch("torch.backends.mps").is_available.return_value = True

    assert infer_devices(devices="auto") == (torch.device("cuda:0"),)


def test__infer_devices__auto__cuda_and_mps_available_but_excluded__selects_cpu(
    mocker: MagicMock, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setenv("TABPFN_EXCLUDE_DEVICES", "mps,cuda")
    mock_cuda = mocker.patch("torch.cuda")
    mock_cuda.is_available.return_value = True
    mock_cuda.device_count.return_value = 1
    mocker.patch("torch.backends.mps").is_available.return_value = True
    assert infer_devices(devices="auto") == (torch.device("cpu"),)


def test__infer_devices__device_specified__selects_it(
    mocker: MagicMock,
) -> None:
    mock_cuda = mocker.patch("torch.cuda")
    mock_cuda.is_available.return_value = True
    mock_cuda.device_count.return_value = 2
    mocker.patch("torch.backends.mps").is_available.return_value = True

    assert infer_devices(devices="cuda:0") == (torch.device("cuda:0"),)


def test__infer_devices__multiple_devices_specified___selects_them(
    mocker: MagicMock,
) -> None:
    mock_cuda = mocker.patch("torch.cuda")
    mock_cuda.is_available.return_value = True
    mock_cuda.device_count.return_value = 3
    mocker.patch("torch.backends.mps").is_available.return_value = False

    inferred = set(infer_devices(devices=["cuda:0", "cuda:1", "cuda:4"]))
    expected = {torch.device("cuda:0"), torch.device("cuda:1"), torch.device("cuda:4")}
    assert inferred == expected


def test__infer_devices__device_selected_twice__raises() -> None:
    with pytest.raises(
        ValueError,
        match="The list of devices for inference cannot contain the same device more ",
    ):
        infer_devices(devices=["cpu", "cpu"])


# --- Test Data for the "test_process_text_na_dataframe" test ---
test_cases = [
    {
        # Mixed dtypes & None / pd.Na
        "df": pd.DataFrame(
            {
                "ratio": [0.4, 0.5, 0.6],
                "risk": ["High", None, "Low"],
                "height": ["Low", "Low", "Low"],
                "amount": [10.2, 20.4, 20.5],
                "type": ["guest", "member", pd.NA],
            }
        ),
        "categorical_indices": [1, 2, 4],
        "ground_truth": np.array(
            [
                [0.4, 0, 0, 10.2, 0],
                [0.5, np.nan, 0, 20.4, 1],
                [0.6, 1, 0, 20.5, np.nan],
            ]
        ),
    },
    {
        # Mixed dtypes & no missing values
        "df": pd.DataFrame(
            {
                "ratio": [0.4, 0.5, 0.6],
                "risk": ["High", "Medium", "Low"],
                "height": ["Low", "Low", "High"],
                "amount": [10.2, 20.4, np.nan],
                "type": ["guest", "member", "vip"],
            }
        ),
        "categorical_indices": [1, 2, 4],
        "ground_truth": np.array(
            [
                [0.4, 0, 1, 10.2, 0],
                [0.5, 2, 1, 20.4, 1],
                [0.6, 1, 0, np.nan, 2],
            ]
        ),
    },
    {
        # All numerical no nan
        "df": pd.DataFrame(
            {
                "ratio": [0.1, 0.2, 0.3],
                "amount": [5.0, 15.5, 25.0],
                "score": [1.0, 2.5, 3.5],
            }
        ),
        "categorical_indices": [],
        "ground_truth": np.array(
            [
                [0.1, 5.0, 1.0],
                [0.2, 15.5, 2.5],
                [0.3, 25.0, 3.5],
            ]
        ),
    },
    {
        # all categorical no nan
        "df": pd.DataFrame(
            {
                "risk": ["High", "High", "High"],
                "height": ["Low", "Low", "Low"],
                "type": ["guest", "guest", "guest"],
            }
        ),
        "categorical_indices": [0, 1, 2],
        "ground_truth": np.array(
            [
                [0, 0, 0],
                [0, 0, 0],
                [0, 0, 0],
            ]
        ),
    },
]


# --- Fixture for the "test_process_text_na_dataframe" test ---
# prepare the DataFrame
@pytest.fixture(params=test_cases)
def prepared_tabpfn_data(request):
    temp_df = request.param["df"].copy()
    categorical_idx = request.param["categorical_indices"]
    # Dummy target, as tests do not need a target
    y = np.array([0, 1, 0])

    cls = TabPFNClassifier()

    X, y, feature_names_in, n_features_in = validate_Xy_fit(
        temp_df,
        y,
        estimator=cls,
        ensure_y_numeric=False,
        max_num_samples=InferenceConfig.MAX_NUMBER_OF_SAMPLES,
        max_num_features=InferenceConfig.MAX_NUMBER_OF_FEATURES,
        ignore_pretraining_limits=False,
    )

    if feature_names_in is not None:
        cls.feature_names_in_ = feature_names_in
    cls.n_features_in_ = n_features_in

    if not cls.differentiable_input:
        _, counts = np.unique(y, return_counts=True)
        cls.class_counts_ = counts
        cls.label_encoder_ = LabelEncoder()
        y = cls.label_encoder_.fit_transform(y)
        cls.classes_ = cls.label_encoder_.classes_
        cls.n_classes_ = len(cls.classes_)
    else:
        cls.label_encoder_ = None
        if not hasattr(cls, "n_classes_"):
            cls.n_classes_ = int(torch.max(torch.tensor(y)).item()) + 1
        cls.classes_ = torch.arange(cls.n_classes_)

    cls.inferred_categorical_indices_ = infer_categorical_features(
        X=X,
        provided=categorical_idx,
        min_samples_for_inference=InferenceConfig.MIN_NUMBER_SAMPLES_FOR_CATEGORICAL_INFERENCE,
        max_unique_for_category=InferenceConfig.MAX_UNIQUE_FOR_CATEGORICAL_FEATURES,
        min_unique_for_numerical=InferenceConfig.MIN_UNIQUE_FOR_NUMERICAL_FEATURES,
    )
    return (
        fix_dtypes(X, cat_indices=cls.inferred_categorical_indices_),
        cls.inferred_categorical_indices_,
        request.param["ground_truth"],
    )


# --- Actual test ---
# This is a test for the OrderPreservingColumnTransformer, which is not used currently
# But might be used in the future, therefore I'll leave it in.
@pytest.mark.skip
def test_process_text_na_dataframe(prepared_tabpfn_data):
    X, categorical_idx, ground_truth = prepared_tabpfn_data  # use the fixture

    ord_encoder = get_ordinal_encoder()
    X_out = process_text_na_dataframe(
        X,
        placeholder=NA_PLACEHOLDER,
        ord_encoder=ord_encoder,
        fit_encoder=True,
    )

    assert X_out.shape[0] == ground_truth.shape[0]
    assert X_out.shape[1] == ground_truth.shape[1]

    for col_name in X.columns:
        # col_name should already be a numeric index but using get_loc for safety
        col_idx = X.columns.get_loc(col_name)
        original_col = X[col_name].to_numpy()
        output_col = X_out[:, col_idx]
        gt_col = ground_truth[:, col_idx]
        if col_idx not in categorical_idx:
            # For numeric columns, values should be preserved (within float tolerance).
            # NaNs should also be in the same positions.
            np.testing.assert_allclose(
                output_col,
                original_col,
                equal_nan=True,
                rtol=1e-5,
            )
        else:
            # OrdinalEncoder does not guarante that element order is preserved:

            # First, check if np.nan are correctly positioned
            # ! use np.isnan on outputcol -> must be numerical
            # ! use pd.isna on original col -> can be any type
            np.testing.assert_array_equal(np.isnan(output_col), pd.isna(original_col))
            # Second, check if there are as many unique non-nan values, as expected
            # e.g.: ["high", "mid", "low"] -> [0,2,1] or [2,1,0],...
            assert len(np.unique(output_col[~pd.isna(output_col)])) == len(
                np.unique(gt_col[~pd.isna(gt_col)])
            )


def test_balance_probas_by_class_counts():
    """Test balancing probabilities by class counts."""
    probas = torch.tensor([[0.2, 0.8], [0.6, 0.4], [0.5, 0.5]])
    class_counts = np.array([1, 2])

    balanced = balance_probas_by_class_counts(probas, class_counts)

    # Check that each row sums to one
    sums = balanced.sum(dim=-1)
    assert torch.allclose(sums, torch.ones(3), rtol=1e-5, atol=1e-5)

    expected_balanced = torch.tensor([[1 / 3, 2 / 3], [0.75, 0.25], [2 / 3, 1 / 3]])
    assert torch.allclose(balanced, expected_balanced, rtol=1e-4, atol=1e-4)
