"""Test MIDAS2 pipeline"""

import pathlib
import random
import sys

import torch
import numpy as np
import pandas as pd

from midas2 import model as md


def test_basic_fit_transform():
    # Ensure local package importable when running from repo root.
    sys.path.append(str(pathlib.Path(__file__).resolve().parents[1]))

    np.random.seed(89)
    torch.manual_seed(89)
    random.seed(89)

    n = 20
    data = pd.DataFrame(
        {
            "xn": np.random.randn(n).astype("float32"),
            "xc": np.random.choice(["a", "b", "c"], n),
            "xb": np.random.choice(["x", "y"], n),
            "xb2": np.random.choice([0, 1], n).astype("float32"),
            "xn2": np.random.randint(0, 100, n).astype("float32"),
        }
    )
    data["xc2"] = np.where(data["xn"] > 0, "a", "b")

    for col in data.columns:
        data[col] = np.where(np.random.uniform(0, 1, n) < 0.1, np.nan, data[col])
    data.iloc[0, 1] = np.nan

    mod = md.MIDAS()
    mod.fit(data, epochs=5, batch_size=2, seed=89)

    imputations = list(mod.transform(m=3))

    # Basic structural checks on the first draw.
    df0 = imputations[0]
    assert list(df0.columns) == list(
        data.columns
    ), "Column names do not match the original data."
    assert (
        df0.shape == data.shape
    ), "Shape of the imputed data does not match the original data."
    assert (
        df0.isna().sum().sum() == 0
    ), "There are still missing values in the imputed data."

    # Categorical levels should be drawn from the original categories only.
    for col in ["xc", "xb", "xc2"]:
        assert set(df0[col].unique()) <= set(data[col].dropna().unique())

    # Binary and categorical draws should differ across imputations.
    assert any(
        not imp.equals(df0) for imp in imputations[1:]
    ), "All imputations are identical; expected variation across draws."


def test_categorical_round_trip():
    """Ensure categorical columns are one-hot encoded and decoded correctly."""
    cats = pd.Series(["a", "a", "b", "c"], dtype="category")
    df = pd.DataFrame({"cat": cats, "x": [0, 1, 2, 3]})

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

    ds = mod.dataset
    assert ds.col_types == [3, "pos"]
    assert list(ds.type_dict["cat"]) == ["a", "b", "c"]
    assert ds.original_col_names == ["cat", "x"]
    assert ds.col_names == ["cat_a", "cat_b", "cat_c", "x"]
    expected_matrix = np.array(
        [
            [1, 0, 0, 0],
            [1, 0, 0, 1],
            [0, 1, 0, 2],
            [0, 0, 1, 3],
        ],
        dtype=np.float32,
    )
    np.testing.assert_array_equal(ds.data.numpy(), expected_matrix)

    out_df = next(mod.transform(m=1))
    assert out_df["cat"].tolist() == ["a", "a", "b", "c"]
    assert out_df["x"].tolist() == [0, 1, 2, 3]


def test_transform_with_new_data_formatting():
    """Ensure transform(format_X=True) applies training formatting to new data."""
    train = pd.DataFrame(
        {
            "cat": pd.Series(["a", "b", "c", "a"], dtype="category"),
            "num": [1.0, 2.0, np.nan, 4.0],
        }
    )
    test = pd.DataFrame(
        {
            "cat": pd.Series(["b", "c", "a", "b"], dtype="category"),
            "num": [5.0, np.nan, 7.0, 8.0],
        }
    )

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

    # Should not raise when formatting new data to match training schema.
    imp = next(mod.transform(X=test, m=1, format_X=True))
    assert list(imp.columns) == list(train.columns)


def test_transform_logits_shape_without_decode():
    """revert_cols=False should return logits/one-hot blocks matching training schema."""
    train = pd.DataFrame(
        {
            "cat": pd.Series(["a", "b", "c", "a"], dtype="category"),
            "bin": [0.0, 1.0, np.nan, 0.0],
            "num": [1.0, 2.0, 3.0, 4.0],
        }
    )

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

    out = next(mod.transform(m=1, revert_cols=False))
    # Expect three cat logits + one bin logit + one num = 5 columns
    assert out.shape[1] == 5
    # No decoding means column names remain expanded
    assert list(out.columns) == mod.dataset.col_names


def test_transform_seed_behaviour():
    """With a fixed seed, repeated calls are identical; different seeds yield different draws."""
    data = pd.DataFrame(
        {
            "x": [0.0, 1.0, np.nan, 0.0],
            "cat": pd.Series(["a", "b", "a", "b"], dtype="category"),
        }
    )

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

    a1 = list(mod.transform(m=2))
    a2 = list(mod.transform(m=2))
    assert all(df1.equals(df2) for df1, df2 in zip(a1, a2))

    # Deterministic across repeated calls when seed is fixed
    mod.seed = 7
    b1 = list(mod.transform(m=2))
    b2 = list(mod.transform(m=2))
    assert all(df1.equals(df2) for df1, df2 in zip(b1, b2))


def test_omit_first_with_categorical_outcome():
    """Ensure omit_first works when the first column is categorical."""
    data = pd.DataFrame(
        {
            "y": pd.Series(["yes", "no", np.nan, "yes"], dtype="category"),
            "x": [1.0, np.nan, 3.0, 4.0],
        }
    )
    mod = md.MIDAS(hidden_layers=[4])
    mod.fit(data, epochs=0, batch_size=2, seed=7, omit_first=True)
    out = next(mod.transform(m=1))
    assert list(out.columns) == list(data.columns)
