import os
import pickle
import shutil

import pandas as pd
import pytest
from pandas.testing import assert_frame_equal

from ablation.experiment import Config, Experiment

TEST_FILE_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "test_files"
)


def clean_dir(path):
    shutil.rmtree(path)


@pytest.fixture
def config_file():
    return os.path.join(TEST_FILE_PATH, "test_config.yml")


@pytest.fixture
def config_file_reload():
    return os.path.join(TEST_FILE_PATH, "test_config_reload.yml")


@pytest.fixture
def config_file_new():
    return os.path.join(TEST_FILE_PATH, "test_config_new.yml")


def test_config_equal(config_file, config_file_new):
    a = Config.from_yaml_file(config_file)
    b = Config.from_yaml_file(config_file)
    c = Config.from_yaml_file(config_file_new)

    assert a == b
    assert a != c
    assert a.diff(c) == [
        "explanation_methods",
        "ablation_args",
        "dataset_sample_perc",
        "n_trials",
    ]


def test_config_save(config_file):
    config = Config.from_yaml_file(config_file)

    if not os.path.exists(config.path):
        os.makedirs(config.path)

    config.save()

    saved_config = Config.from_yaml_file(
        os.path.join(config.path, "config.yml")
    )
    assert config == saved_config

    shutil.rmtree(config.path, ignore_errors=True)


def test_experiment_reproducability(config_file):

    config = Config.from_yaml_file(config_file)

    exp1 = Experiment(config)
    exp2 = Experiment(config)

    results1 = exp1.run_exp()
    results2 = exp2.run_exp()

    # Results
    assert_frame_equal(results1, results2)

    # Sanity checks
    for check in exp1.model_sanity_checks:
        for score_name in exp1.model_sanity_checks[check]:
            assert (
                exp1.model_sanity_checks[check][score_name]
                == exp2.model_sanity_checks[check][score_name]
            )

    shutil.rmtree(config.path, ignore_errors=True)


def test_experiment_load_reproducability(config_file):

    config = Config.from_yaml_file(config_file)

    exp1 = Experiment(config)
    exp1.run_exp()

    exp2 = Experiment(config)
    results2 = exp2.run_exp()

    results1 = pd.read_pickle(os.path.join(config.path, "results-local.pkl"))
    sanity_checks1 = pickle.load(
        open(os.path.join(config.path, "model_sanity_checks.pkl"), "rb")
    )

    assert_frame_equal(results1, results2)
    assert sanity_checks1 == exp2.model_sanity_checks

    shutil.rmtree(config.path, ignore_errors=True)


def test_experiment_reload_reproducability(config_file, config_file_reload):

    config = Config.from_yaml_file(config_file)
    config_reload = Config.from_yaml_file(config_file_reload)

    exp1 = Experiment(config)
    results1 = exp1.run_exp()

    exp2 = Experiment(config_reload)
    results2 = exp2.run_exp()

    assert_frame_equal(results1, results2)

    shutil.rmtree(config.path, ignore_errors=True)


def test_experiment_result(config_file):

    config = Config.from_yaml_file(config_file)

    exp = Experiment(config)
    results = exp.run_exp()

    test_results = pd.read_csv(os.path.join(TEST_FILE_PATH, "test_result.csv"))

    assert_frame_equal(results, test_results, atol=1e-4)

    shutil.rmtree(config.path, ignore_errors=True)
