"""Core unit tests for MIDAS2 utilities and model behavior."""

import numpy as np
import pandas as pd
import torch
import pytest

from midas2 import model as md
from midas2.dataset import Dataset
from midas2.custom_loss import MixedLoss, _masked_loss


def test_omit_first_skips_column_in_forward():
    """When omit_first is True, the first column should be excluded from the encoder input."""
    df = pd.DataFrame(
        {
            "y": [0.0, 1.0, np.nan, 0.0],
            "x1": [1.0, 2.0, 3.0, 4.0],
            "x2": [0.0, 1.0, np.nan, 0.0],
        }
    )
    mod = md.MIDAS(hidden_layers=[4], device="cpu")
    mod.fit(df, epochs=0, batch_size=2, omit_first=True, seed=42)

    # Input dim should be full width, encoder input width should be minus one.
    assert mod.input_dim == 3
    assert mod.encoder[0].in_features == 2

    out = next(mod.transform(m=1))
    # Column order and names preserved.
    assert list(out.columns) == list(df.columns)


def test_transform_preserves_observed_values():
    """Observed entries should pass through unchanged after imputation."""
    df = pd.DataFrame(
        {
            "num": [1.0, np.nan, 3.0, 4.0],
            "bin": [0.0, 1.0, np.nan, 0.0],
            "cat": pd.Series(["r", "g", "b", np.nan], dtype="category"),
        }
    )

    mod = md.MIDAS(hidden_layers=[8])
    mod.fit(df, epochs=0, batch_size=2, seed=7)

    imp = next(mod.transform(m=1))

    for col in df.columns:
        observed = df[col].dropna()
        assert imp.loc[observed.index, col].tolist() == observed.tolist()


def test_unseen_category_flagged_missing():
    """Unseen categories in test data should be marked missing (mask false)."""
    train = pd.DataFrame({"cat": pd.Series(["a", "b", "c"], dtype="category")})
    ds_train = Dataset(train)

    test = pd.DataFrame({"cat": pd.Series(["d"], dtype="category")})
    ds_test = Dataset(
        test,
        col_types=ds_train.col_types,
        type_dict=ds_train.type_dict,
        col_names=ds_train.col_names,
        original_names=ds_train.original_col_names,
        test_format=True,
    )

    # First three columns correspond to the one-hot cat block; unseen "d" -> NaN -> mask False.
    assert not ds_test.mask_expand[0, :3].any()


def test_mixed_loss_shape_and_broadcast():
    """MixedLoss should return per-element losses aligned with predictions."""
    col_types = ["num", "bin", 3]
    pred = torch.tensor([[2.0, 0.0, 0.2, -0.1, 0.0]])
    tgt = torch.tensor([[1.0, 1.0, 0.0, 1.0, 0.0]])

    loss = MixedLoss(col_types)(pred, tgt)
    assert loss.shape == pred.shape

    # Categorical loss should be identical across its block.
    cat_slice = loss[0, 2:]
    assert torch.allclose(cat_slice, cat_slice[0])


def test_masked_loss_matches_matrix_path():
    """MixedLoss.masked should produce the same scalar as the matrix-based path."""
    col_types = ["num", "bin", 3]
    pred = torch.tensor([[2.0, 0.0, 0.2, -0.1, 0.0]])
    tgt = torch.tensor([[1.0, 1.0, 0.0, 1.0, 0.0]])
    mask = torch.tensor([[1, 1, 1, 0, 1]], dtype=torch.bool)

    loss_fn = MixedLoss(col_types)
    matrix = loss_fn(pred, tgt)
    old = _masked_loss(matrix, mask)
    new = loss_fn.masked(pred, tgt, mask)

    assert torch.allclose(old, new)


def test_masked_loss_raises_on_shape_mismatch():
    """_masked_loss should guard against shape mismatch."""
    a = torch.zeros(2, 3)
    b = torch.ones(2, 2)
    try:
        _masked_loss(a, b)
    except ValueError:
        return
    raise AssertionError("Expected ValueError on mismatched shapes.")


def test_masked_loss_raises_on_empty_mask():
    """_masked_loss should guard against zero observed entries."""
    losses = torch.ones(2, 2)
    mask = torch.zeros(2, 2, dtype=torch.bool)
    with pytest.raises(ValueError):
        _masked_loss(losses, mask)
