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

import torch
import statsmodels.api as sm
import numpy as np
import pandas as pd


from midas2 import model as md
from midas2.utils import combine


## Data generation as per King et al (2001)
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


if __name__ == "__main__":
    ## simulation
    B = 100
    m = 20

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

    results = pd.DataFrame(columns=["type", "coef1", "coef2", "coef3"])

    se_results = pd.DataFrame(columns=["type", "se1", "se2", "se3"])

    for b in range(B):

        print("Simulation: " + str(b))

        train_data = MAR1(10000)
        missing_data = make_missing(train_data.iloc[:2000,].copy())

        fdata_endog = train_data.iloc[:, 0]
        fdata_exog = sm.add_constant(train_data.iloc[:, 1:3], prepend=False)

        full_lm = sm.OLS(fdata_endog, fdata_exog)
        full_res = full_lm.fit()
        full_coefs = full_res.params.tolist()
        results.loc[len(results)] = ["full"] + full_coefs

        # train imputer
        MIDAS_model = md.MIDAS(hidden_layers=[256, 256, 256])
        MIDAS_model.fit(missing_data, epochs=5, verbose=False)
        imputed_data = MIDAS_model.transform(m=m)

        midas_res = combine(
            imputed_data,
            y="x1",
            ind_vars=["x2", "x3"],
        )

        results.loc[len(results)] = ["midas"] + midas_res.loc[
            [1, 2, 0], "estimate"
        ].tolist()

        se_results.loc[len(se_results)] = ["midas"] + midas_res.loc[
            [1, 2, 0], "std.error"
        ].tolist()

    results_full = results[results["type"] == "full"]
    results_full.mean(axis=0, numeric_only=True)

    results_midas = results[results["type"] == "midas"]
    results_midas.loc[:, "iteration"] = np.arange(1, len(results_midas) + 1)
    se_results.loc[:, "iteration"] = np.arange(1, len(se_results) + 1)

    results_midas_w = results_midas.melt(
        id_vars=["iteration"], value_vars=["coef1", "coef2", "coef3"]
    )
    results_midas_w.columns = ["iteration", "variable", "beta"]

    se_results_w = se_results.melt(
        id_vars=["iteration"], value_vars=["se1", "se2", "se3"]
    )
    se_results_w.variable = [
        "coef" + d for d in se_results_w["variable"].str.extract("(\\d)")[0]
    ]
    se_results_w.columns = ["iteration", "variable", "se"]

    midas_results = pd.merge(
        results_midas_w, se_results_w, how="left", on=["iteration", "variable"]
    )

    # calculate whether the coefficient +/- 1.96*se contains the full model coefficient
    midas_results["in_ci"] = (
        midas_results["beta"] - 1.96 * midas_results["se"]
        <= results_full.mean(axis=0, numeric_only=True)[
            midas_results["variable"]
        ].reset_index(drop=True)
    ) & (
        midas_results["beta"] + 1.96 * midas_results["se"]
        >= results_full.mean(axis=0, numeric_only=True)[
            midas_results["variable"]
        ].reset_index(drop=True)
    )

    midas_results.groupby("variable")["in_ci"].mean()
