import pandas as pd
import torch

from causalml.inference.nn import CEVAE
from causalml.dataset import simulate_hidden_confounder
from causalml.metrics import get_cumgain


def test_CEVAE():
    y, X, treatment, tau, b, e = simulate_hidden_confounder(n=10000, p=5, sigma=1.0, adj=0.)

    outcome_dist = "normal"
    latent_dim = 20
    hidden_dim = 200
    num_epochs = 50
    batch_size = 100
    learning_rate = 1e-3
    learning_rate_decay = 0.1

    cevae = CEVAE(outcome_dist=outcome_dist,
                  latent_dim=latent_dim,
                  hidden_dim=hidden_dim,
                  num_epochs=num_epochs,
                  batch_size=batch_size,
                  learning_rate=learning_rate,
                  learning_rate_decay=learning_rate_decay)

    cevae.fit(X=torch.tensor(X, dtype=torch.float),
              treatment=torch.tensor(treatment, dtype=torch.float),
              y=torch.tensor(y, dtype=torch.float))

    # check the accuracy of the ite accuracy
    ite = cevae.predict(X).flatten()

    auuc_metrics = pd.DataFrame({'ite': ite,
                                 'W': treatment,
                                 'y': y,
                                 'treatment_effect_col': tau})

    cumgain = get_cumgain(auuc_metrics,
                          outcome_col='y',
                          treatment_col='W',
                          treatment_effect_col='tau')

    # Check if the cumulative gain when using the model's prediction is
    # higher than it would be under random targeting
    assert cumgain['ite'].sum() > cumgain['Random'].sum()