import types

import pandas as pd
import pytest

pytest.importorskip("torch")
pytest.importorskip("optuna")
pytest.importorskip("xgboost")
pytest.importorskip("statsmodels")

from ins_pricing.modelling.bayesopt.trainers import FTTrainer


class DummyCtx:
    def __init__(self, train_df: pd.DataFrame, test_df: pd.DataFrame):
        self.task_type = "regression"
        self.config = types.SimpleNamespace(use_ft_ddp=False, geo_feature_nmes=["geo"])
        self.train_data = train_df
        self.test_data = test_df
        self.train_geo_tokens = None
        self.test_geo_tokens = None
        self.geo_token_cols = []
        self.geo_gnn_model = None
        self._build_calls = []
        self._last_params = None

    def _build_geo_tokens(self, _params=None):
        self._last_params = _params
        self._build_calls.append(
            (self.train_data.copy(deep=True), self.test_data.copy(deep=True))
        )
        return self.train_data.copy(deep=True), self.test_data.copy(deep=True), ["geo_token"], None


def test_geo_token_split_uses_fold_and_restores_context():
    train = pd.DataFrame({"geo": ["a", "b", "c", "d"], "x": [1, 2, 3, 4]})
    test = pd.DataFrame({"geo": ["e"], "x": [5]})
    ctx = DummyCtx(train, test)
    trainer = FTTrainer(ctx)

    X_train = train.iloc[[0, 1]]
    X_val = train.iloc[[2, 3]]

    orig_train = ctx.train_data
    orig_test = ctx.test_data

    result = trainer._build_geo_tokens_for_split(X_train, X_val, geo_params={"k": 1})

    assert ctx.train_data is orig_train
    assert ctx.test_data is orig_test
    assert result is not None

    train_snapshot, test_snapshot = ctx._build_calls[0]
    assert train_snapshot.equals(train.loc[X_train.index])
    assert test_snapshot.equals(train.loc[X_val.index])


def test_apply_model_params_separates_geo_token_keys():
    train = pd.DataFrame({"geo": ["a", "b"], "x": [1, 2]})
    test = pd.DataFrame({"geo": ["c"], "x": [3]})
    ctx = DummyCtx(train, test)
    trainer = FTTrainer(ctx)

    class DummyModel:
        def __init__(self):
            self.calls = []

        def set_params(self, params):
            self.calls.append(dict(params))
            if any(key.startswith("geo_token_") for key in params):
                raise AssertionError("geo params must not be passed directly")
            return self

    model = DummyModel()
    geo_params = trainer._apply_model_params(
        model,
        {
            "d_model": 64,
            "dropout": 0.1,
            "geo_token_hidden_dim": 32,
            "geo_token_layers": 2,
        },
    )

    assert geo_params == {"geo_token_hidden_dim": 32, "geo_token_layers": 2}
    assert model.calls[0] == {"d_model": 64, "dropout": 0.1}
    assert model.calls[1] == {"_geo_params": geo_params}


def test_resolve_final_geo_tokens_rebuilds_context_with_tuned_geo_params():
    train = pd.DataFrame({"geo": ["a", "b", "c"], "x": [1, 2, 3]})
    test = pd.DataFrame({"geo": ["d"], "x": [4]})
    ctx = DummyCtx(train, test)
    trainer = FTTrainer(ctx)

    tuned_geo = {"geo_token_hidden_dim": 16, "geo_token_layers": 1}
    geo_train, geo_test = trainer._resolve_final_geo_tokens(tuned_geo)

    assert ctx._last_params == tuned_geo
    assert geo_train is not None and geo_train.equals(train)
    assert geo_test is not None and geo_test.equals(test)
    assert ctx.train_geo_tokens is not None and ctx.train_geo_tokens.equals(train)
    assert ctx.test_geo_tokens is not None and ctx.test_geo_tokens.equals(test)
    assert ctx.geo_token_cols == ["geo_token"]
