import unittest
import warnings
import pandas as pd
import numpy as np
import pathlib
from testfixtures import TempDirectory
import gc

import tensorflow.keras.backend as K

from ml4ir.applications.ranking.pipeline import RankingPipeline
from ml4ir.applications.ranking.config.parse_args import get_args

warnings.filterwarnings("ignore")


def ml4ir_sanity_pipeline(df, working_dir, log_dir, n_features):
    """
    Train ml4ir on the passed data and read the MRR for ml4ir
    """
    df.to_csv(working_dir / 'train' / 'data.csv')
    df.to_csv(working_dir / 'validation' / 'data.csv')
    df.to_csv(working_dir / 'test' / 'data.csv')

    fconfig_name = "feature_config_sanity_tests_" + str(n_features) + "_features.yaml"
    feature_config_file = pathlib.Path(__file__).parent / "data" / "configs" / fconfig_name
    model_config_file = pathlib.Path(__file__).parent / "data" / "configs" / "model_config_sanity_tests.yaml"
    train_ml4ir(working_dir.as_posix(), feature_config_file.as_posix(), model_config_file.as_posix(), log_dir.as_posix())
    ml4ir_results = pd.read_csv(log_dir / 'test_command_line'/ '_SUCCESS', header=None)
    ml4ir_mrr = float(ml4ir_results.loc[ml4ir_results[0] == 'test_new_MRR'][1])
    return ml4ir_mrr


def train_ml4ir(data_dir, feature_config, model_config, logs_dir):
    """
    Train a pointwise ranker, listwise loss model using ml4ir
    """
    argv = ["--data_dir", data_dir,
            "--feature_config", feature_config,
            "--loss_type", "listwise",
            "--scoring_type", "listwise",
            "--run_id", "test_command_line",
            "--data_format", "csv",
            "--execution_mode", "train_evaluate",
            "--loss_key", "rank_one_listnet",
            "--num_epochs", "150",
            "--model_config", model_config,
            "--batch_size", "1",
            "--logs_dir", logs_dir,
            "--max_sequence_size", "25",
            "--train_pcent_split", "1.0",
            "--val_pcent_split", "-1",
            "--test_pcent_split", "-1",
            "--early_stopping_patience", "25",
            "--metrics_keys", "MRR", "categorical_accuracy",
            "--monitor_metric", "categorical_accuracy"]
    args = get_args(argv)
    rp = RankingPipeline(args=args)
    rp.run()


def run_sanity_test(n_features, fname, perceptron_mrr, log_regression_mrr, working_dir, log_dir):
    """
    Runs sanity test for linear models.
    Parameters
    ----------
    n_features : int
        Number of features in the dataset.
    fname : string
        Dataset file name
    perceptron_mrr : float
        MRR of perceptron
    log_regression_mrr : float
        MRR of logistic regression
    working_dir : string
        Path of the working directory
    log_dir : string
        Path of the log directory
    """
    df = pd.read_csv(pathlib.Path(__file__).parent / "data" / "linear_sanity_tests" / fname)
    ml4ir_mrr = ml4ir_sanity_pipeline(df, working_dir, log_dir, n_features)

    assert np.isclose(ml4ir_mrr, perceptron_mrr, atol=0.001) or ml4ir_mrr >= perceptron_mrr
    assert np.isclose(ml4ir_mrr, log_regression_mrr, atol=0.001) or ml4ir_mrr >= log_regression_mrr


class TestML4IRSanity(unittest.TestCase):
    """
    Datasets used in these tests are generated by sklearn.datasets.make_classification by changing its parameters.
    Then queries and documents are grouped randomly (without replacement) so that each query has 2-25 documents with
    one clicked document.
    Baseline models (Perceptron and logistic regression are trained with sklearn)
    """

    def setUp(self):
        self.dir = pathlib.Path(__file__).parent
        self.working_dir = TempDirectory()
        self.log_dir = self.working_dir.makedir('logs')
        self.working_dir.makedir('train')
        self.working_dir.makedir('test')
        self.working_dir.makedir('validation')

    def tearDown(self):
        TempDirectory.cleanup_all()

        # Explicitly clear keras memory
        gc.collect()
        K.clear_session()

    def test_linear_ml4ir_sanity_1(self):
        """
        Weights of pre-trained models:
            perceptron=np.array([1.87212065, -0.00305068])
            log_regression=np.array([28.62071696, 1.18915853])
        """
        run_sanity_test(n_features=2, fname="dataset1.csv",
                        perceptron_mrr=1.0,
                        log_regression_mrr=1.0,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))

    def test_linear_ml4ir_sanity_2(self):
        """
        Weights of pre-trained models:
            perceptron=np.array([4.50209484, -0.80280452])
            log_reg=np.array([22.73585585, -3.94821153])
        """
        run_sanity_test(n_features=2, fname="dataset2.csv",
                        perceptron_mrr=1.0,
                        log_regression_mrr=1.0,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))

    def test_linear_ml4ir_sanity_3(self):
        """
        Weights of pre-trained models:
            perceptron=np.array([-1.27651475, -4.07647092, 8.23950305, 0.29241316, 3.24763417])
            log_reg=np.array([-1.67270377, -5.76088727, 8.36278576, -0.90878154, 3.47653204])
        """
        run_sanity_test(n_features=5, fname="dataset3.csv",
                        perceptron_mrr=1.0,
                        log_regression_mrr=1.0,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))

    def test_linear_ml4ir_sanity_4(self):
        """
        Weights of pre-trained models:
            perceptron=np.array([5.10535665, 1.44131417])
            log_reg=np.array([20.0954756, 4.69360163])
        """
        run_sanity_test(n_features=2, fname="dataset4.csv",
                        perceptron_mrr=0.975609756097561,
                        log_regression_mrr=0.9878048780487805,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))

    def test_linear_ml4ir_sanity_5(self):
        """
        Weights of pre-trained models:
            perceptron=np.array([0.57435291, -0.99437351])
            log_reg=np.array([1.15593505, -0.93317691])
        """
        run_sanity_test(n_features=2, fname="dataset5.csv",
                        perceptron_mrr=0.37021768881524975,
                        log_regression_mrr=0.4290204375570229,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))

    def test_linear_ml4ir_sanity_6(self):
        """
        Weights of pre-trained models:
        perceptron=np.array(
                        [4.59994733, -3.56373965, -6.15935686, 0.87523846, -0.64231058, 2.15971991, 5.79875003,
                         -7.70152594, -0.07521741, 2.8817456])
        log_reg=np.array(
                        [-0.38064406, -0.27970534, 0.02775136, 0.25641926, 0.15413321, 0.29194965, 0.72707686,
                         0.24791729, -0.39367192, 0.4882174])
        """
        run_sanity_test(n_features=10, fname="dataset6.csv",
                        perceptron_mrr=0.2630707011921534,
                        log_regression_mrr=0.29970690874564615,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))

    def test_linear_ml4ir_sanity_7(self):
        """
        Weights of pre-trained models:
            perceptron=np.array([0.40127356, -0.43773627])
            log_reg=np.array([4.15630544, -1.09111369])
        """
        run_sanity_test(n_features=2, fname="dataset7.csv",
                        perceptron_mrr=0.304679599925838,
                        log_regression_mrr=0.4938108725899423,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))

    def test_linear_ml4ir_sanity_8(self):
        """
        Weights of pre-trained models:
            perceptron=np.array(
                            [2.91798129, 4.24880336, 7.42919018, 2.49609694, -0.84988373, 0.43435823, -0.18953416,
                             2.23129287, -0.67951411, -0.63925108])
            log_reg=np.array(
                            [-0.14472192, -0.22594271, 0.62703883, 0.16002515, 0.17084088, -0.22872226, 0.89200279,
                             0.06297475, 0.70470567, -0.19396659])
        """
        run_sanity_test(n_features=10, fname="dataset8.csv",
                        perceptron_mrr=0.31185341006782435,
                        log_regression_mrr=0.3092000998257748,
                        working_dir=pathlib.Path(self.working_dir.path), log_dir=pathlib.Path(self.log_dir))


if __name__ == "__main__":
    unittest.main()
