"""Run MAR-1 simulation from Lall and Robinson (2022)"""

import numpy as np
import pandas as pd
import pytest
import torch
from sklearn.linear_model import LinearRegression

from midas2 import model as md


def MAR1(n: int):
    """
    Generate data from a multivariate normal distribution with MAR-1 covariances.
    See King et al (2001) for details.
    """
    dta = np.random.multivariate_normal(
        mean=[0, 0, 0, 0, 0],
        cov=[
            [
                1,
                -0.12,
                -0.1,
                0.5,
                0.1,
            ],
            [
                -0.12,
                1,
                0.1,
                -0.6,
                0.1,
            ],
            [
                -0.1,
                0.1,
                1,
                -0.5,
                0.1,
            ],
            [
                0.5,
                -0.6,
                -0.5,
                1,
                0.1,
            ],
            [0.1, 0.1, 0.1, 0.1, 1],
        ],
        size=n,
    )
    return pd.DataFrame(dta, columns=["x" + str(i) for i in range(1, 6)])


def make_missing(data: pd.DataFrame):
    """
    Add MAR-1 missingness to generated data.
    """
    n = data.shape[0]
    M = np.ndarray(data.shape, dtype=bool)
    U1 = np.random.uniform(0, 1, n)

    # Y and X4 are MCAR:
    M[:, 0] = U1 < 0.85
    M[:, 4] = U1 < 0.85

    # X3 is always observed:
    M[:, 3] = True

    # X1 and X2 are MAR:
    U2 = np.random.uniform(0, 1, n)
    M[:, 1] = ~np.all([data.iloc[:, 3] < -1, U2 < 0.9], axis=0)

    U3 = np.random.uniform(0, 1, n)
    M[:, 2] = ~np.all([data.iloc[:, 3] < -1, U3 < 0.9], axis=0)

    data[~M] = np.nan

    return data


@pytest.mark.slow
def test_mar_single_run():
    """Basic MAR pipeline sanity check (reduced size for test runtime)."""
    np.random.seed(89)
    torch.manual_seed(89)

    train_data = MAR1(500)
    missing_data = make_missing(train_data.copy())

    full_lm = LinearRegression(fit_intercept=True)
    full_lm.fit(train_data.iloc[:, [1, 2]], train_data.iloc[:, 0])
    full_rmse = np.sqrt(
        np.mean(
            (full_lm.predict(train_data.iloc[:, [1, 2]]) - train_data.iloc[:, 0]) ** 2
        )
    )

    midas_model = md.MIDAS()
    midas_model.fit(missing_data, epochs=5, batch_size=128, lr=0.01, verbose=False)
    imputed_data = midas_model.transform(m=3)

    rmses = []
    for imp in imputed_data:
        imp_lm = LinearRegression(fit_intercept=True)
        imp_lm.fit(imp.iloc[:, [1, 2]], imp.iloc[:, 0])
        rmses.append(
            np.sqrt(
                np.mean((imp_lm.predict(imp.iloc[:, [1, 2]]) - imp.iloc[:, 0]) ** 2)
            )
        )

    # Imputation should not degrade RMSE too severely relative to full data.
    assert np.mean(rmses) < full_rmse * 1.5
