from __future__ import annotations

import pandas as pd
import pytest

torch = pytest.importorskip("torch")

from ins_pricing.modelling.bayesopt.models.model_ft_trainer import FTTransformerSklearn
from ins_pricing.modelling.bayesopt.models.model_resn import _LazyTabularDataset


def test_ft_unsupervised_lazy_path_skips_full_tensorize(monkeypatch):
    X = pd.DataFrame(
        {
            "x_num": [1.0, 2.0, 3.0],
            "x_cat": ["a", "b", "a"],
        }
    )
    model = FTTransformerSklearn(
        model_nme="ft",
        num_cols=["x_num"],
        cat_cols=["x_cat"],
        epochs=0,
        use_gpu=False,
        use_data_parallel=False,
        use_ddp=False,
    )
    model.use_lazy_dataset = True

    def _fail(*args, **kwargs):
        raise AssertionError("_tensorize_split should not be used in lazy unsupervised mode")

    monkeypatch.setattr(model, "_tensorize_split", _fail)

    loss = model.fit_unsupervised(X)

    assert loss == 0.0


def test_resn_lazy_dataset_caches_dense_feature_matrix():
    X = pd.DataFrame({"x1": [1.0, 2.0], "x2": [3.0, 4.0]})
    y = pd.Series([10.0, 20.0])
    dataset = _LazyTabularDataset(X, y)

    assert dataset._x_values is not None
    x_item, y_item, w_item = dataset[1]
    assert tuple(x_item.shape) == (2,)
    assert float(y_item.item()) == 20.0
    assert float(w_item.item()) == 1.0


def test_resn_lazy_dataset_keeps_sparse_frame_lazy():
    X = pd.DataFrame(
        {
            "x1": pd.arrays.SparseArray([0.0, 1.0]),
            "x2": pd.arrays.SparseArray([1.0, 0.0]),
        }
    )
    y = pd.Series([10.0, 20.0])
    dataset = _LazyTabularDataset(X, y)

    assert dataset._x_values is None
