from unittest import TestCase
from unittest.mock import ANY, Mock, call, patch

import pandas as pd
from mlblocks import MLPipeline

from orion import benchmark
from orion.evaluation import CONTEXTUAL_METRICS as METRICS
from orion.evaluation import contextual_confusion_matrix


def test__sort_leaderboard_rank():
    rank = 'f1'
    metrics = METRICS
    score = pd.DataFrame({
        'pipeline': range(5),
        'f1': range(5),
    })

    expected_return = pd.DataFrame({
        'pipeline': range(5)[::-1],
        'rank': range(1, 6),
        'f1': range(5)[::-1],
    })

    returned = benchmark._sort_leaderboard(score, rank, metrics)
    pd.testing.assert_frame_equal(returned, expected_return)


def test__sort_leaderboard_rank_does_not_exist():
    rank = 'does not exist'
    metrics = {'f1': METRICS['f1']}
    score = pd.DataFrame({
        'pipeline': range(5),
        'f1': range(5),
    })

    expected_return = pd.DataFrame({
        'pipeline': range(5)[::-1],
        'rank': range(1, 6),
        'f1': range(5)[::-1],
    })

    returned = benchmark._sort_leaderboard(score, rank, metrics)
    pd.testing.assert_frame_equal(returned, expected_return)


def test__sort_leaderboard_no_rank():
    rank = None
    metrics = METRICS
    score = {k: range(5) for k in metrics.keys()}

    score['pipeline'] = range(5)
    score = pd.DataFrame(score)

    expected_return = score.iloc[::-1].reset_index(drop=True)
    expected_return['rank'] = range(1, 6)

    returned = benchmark._sort_leaderboard(score, rank, metrics)

    assert len(returned.columns) == len(expected_return.columns)
    assert sorted(returned.columns) == sorted(expected_return.columns)
    pd.testing.assert_frame_equal(returned, expected_return[returned.columns], check_dtype=False)


def test__detrend_signal_trend():
    df = pd.DataFrame({
        'timestamp': range(5),
        'value': range(5)
    })

    expected_return = pd.DataFrame({
        'timestamp': range(5),
        'value': [0.0] * 5,
    })

    returned = benchmark._detrend_signal(df, 'value')
    pd.testing.assert_frame_equal(returned, expected_return)


def test__detrend_signal_no_trend():
    df = pd.DataFrame({
        'timestamp': range(5),
        'value': [0.0] * 5
    })

    expected_return = df.copy()

    returned = benchmark._detrend_signal(df, 'value')
    pd.testing.assert_frame_equal(returned, expected_return)


def test__get_pipeline_hyperparameter():
    hyperparameters = {
        "pipeline1": "pipeline1.json",
        "pipeline2": "pipeline2.json",
    }
    pipeline = "pipeline1"

    expected_return = "pipeline1.json"
    returned = benchmark._get_pipeline_hyperparameter(hyperparameters, None, pipeline)
    assert returned == expected_return


def test__get_pipeline_hyperparameter_dataset():
    hyperparameters = {
        "dataset1": {
            "pipeline1": "pipeline1.json",
            "pipeline2": "pipeline2.json",
        }
    }
    dataset = "dataset1"

    expected_return = {
        "pipeline1": "pipeline1.json",
        "pipeline2": "pipeline2.json",
    }
    returned = benchmark._get_pipeline_hyperparameter(hyperparameters, dataset, None)
    assert returned == expected_return


def test__get_pipeline_hyperparameter_does_not_exist():
    hyperparameters = None
    pipeline = "pipeline1"

    expected_return = None
    returned = benchmark._get_pipeline_hyperparameter(hyperparameters, None, pipeline)
    assert returned == expected_return


@patch('orion.benchmark.load_signal')
def test__load_signal_test_split_true(load_signal_mock):
    train = Mock(autospec=pd.DataFrame)
    test = Mock(autospec=pd.DataFrame)

    load_signal_mock.return_value = (train, test)

    test_split = True
    returned = benchmark._load_signal('signal-name', test_split)

    assert isinstance(returned, tuple)
    assert len(returned) == 2

    expected_calls = [
        call('signal-name-train'),
        call('signal-name-test')
    ]
    assert load_signal_mock.call_args_list == expected_calls


@patch('orion.benchmark.load_signal')
def test__load_signal_test_split_false(load_signal_mock):
    df = pd.DataFrame({
        'timestamp': list(range(10)),
        'value': list(range(10, 20))
    })
    load_signal_mock.return_value = df

    test_split = False
    returned = benchmark._load_signal('signal-name', test_split)

    assert isinstance(returned, tuple)
    assert len(returned) == 2

    train, test = returned
    pd.testing.assert_frame_equal(train, test)

    expected_calls = [
        call('signal-name'),
    ]
    assert load_signal_mock.call_args_list == expected_calls


@patch('orion.benchmark.load_signal')
def test__load_signal_test_split_float(load_signal_mock):
    train = Mock(autospec=pd.DataFrame)
    test = Mock(autospec=pd.DataFrame)

    load_signal_mock.return_value = (train, test)

    test_split = 0.2
    returned = benchmark._load_signal('signal-name', test_split)

    assert isinstance(returned, tuple)
    assert len(returned) == 2

    expected_calls = [
        call('signal-name', test_size=test_split),
    ]
    assert load_signal_mock.call_args_list == expected_calls


class TestBenchmark(TestCase):

    @classmethod
    def setup_class(cls):
        cls.pipeline = Mock(autospec=MLPipeline)
        cls.name = 'pipeline-name'
        cls.dataset = 'dataset-name'
        cls.signal = 'signal-name'
        cls.hyper = None
        cls.distributed = False
        cls.rank = 'metric-name'
        cls.test_split = False
        cls.detrend = False
        cls.pipeline_path = None
        cls.cache_dir = None
        cls.iteration = 0
        cls.run_id = ANY
        cls.metrics = {
            'metric-name': Mock(autospec=METRICS['f1'], return_value=1)
        }

    def args(self):
        return (
            self.pipeline,
            self.name,
            self.dataset,
            self.signal,
            self.hyper,
            self.metrics,
            self.test_split,
            self.detrend,
            self.iteration,
            self.cache_dir,
            self.pipeline_path,
            self.run_id
        )

    def set_score(self, metric, elapsed, test_split):
        return {
            'metric-name': metric,
            'elapsed': elapsed,
            'split': test_split,
            'status': 'OK'
        }

    def set_output(self, metric, elapsed, test_split):
        return {
            'dataset': self.dataset,
            'pipeline': self.name,
            'signal': self.signal,
            'iteration': self.iteration,
            'metric-name': metric,
            'elapsed': elapsed,
            'split': test_split,
            'status': 'OK',
            'run_id': self.run_id
        }

    @patch('orion.benchmark.load_anomalies')
    @patch('orion.benchmark.analyze')
    @patch('orion.benchmark._load_pipeline')
    @patch('orion.benchmark.load_signal')
    def test__evaluate_signal(
            self, load_signal_mock, load_pipeline_mock, analyze_mock, load_anomalies_mock):
        train = Mock(autospec=pd.DataFrame)
        test = Mock(autospec=pd.DataFrame)
        load_signal_mock.side_effect = [train, test]
        load_pipeline_mock.return_value = self.pipeline

        anomalies = Mock(autospec=pd.DataFrame)
        analyze_mock.return_value = anomalies

        returned = benchmark._evaluate_signal(
            self.pipeline, self.signal, self.hyper, self.metrics, True)

        expected_return = self.set_score(1, ANY, ANY)
        assert returned == expected_return

        expected_calls = [
            call('signal-name-train'),
            call('signal-name-test')
        ]
        assert load_signal_mock.call_args_list == expected_calls

        load_pipeline_mock.assert_called_once_with(self.pipeline, self.hyper)
        analyze_mock.assert_called_once_with(self.pipeline, train, test)
        load_anomalies_mock.assert_called_once_with(self.signal)

    @patch('orion.benchmark.load_anomalies')
    @patch('orion.benchmark.analyze')
    @patch('orion.benchmark._load_pipeline')
    @patch('orion.benchmark.load_signal')
    def test__evaluate_signal_exception(
            self, load_signal_mock, load_pipeline_mock, analyze_mock, load_anomalies_mock):
        train = Mock(autospec=pd.DataFrame)
        test = Mock(autospec=pd.DataFrame)
        load_signal_mock.side_effect = [train, test]
        load_pipeline_mock.return_value = self.pipeline

        analyze_mock.side_effect = Exception("failed analyze.")

        returned = benchmark._evaluate_signal(
            self.pipeline, self.signal, self.hyper, self.metrics, True)

        expected_return = self.set_score(0, ANY, ANY)
        expected_return['status'] = 'ERROR'
        assert returned == expected_return

        expected_calls = [
            call('signal-name-train'),
            call('signal-name-test')
        ]
        assert load_signal_mock.call_args_list == expected_calls

        load_pipeline_mock.assert_called_once_with(self.pipeline, self.hyper)
        analyze_mock.assert_called_once_with(self.pipeline, train, test)

        assert load_anomalies_mock.called

    @patch('orion.benchmark.load_anomalies')
    @patch('orion.benchmark.analyze')
    @patch('orion.benchmark._load_pipeline')
    @patch('orion.benchmark.load_signal')
    def test__evaluate_signal_exception_confusion_matrix(
            self, load_signal_mock, load_pipeline_mock, analyze_mock, load_anomalies_mock):
        anomalies = pd.DataFrame({
            'start': [10, 35],
            'end': [20, 40]
        })

        train = Mock(autospec=pd.DataFrame)
        test = Mock(autospec=pd.DataFrame)
        load_signal_mock.side_effect = [train, test]
        load_pipeline_mock.return_value = self.pipeline
        load_anomalies_mock.return_value = anomalies
        analyze_mock.side_effect = Exception("failed analyze.")

        metrics = {'confusion_matrix': Mock(autospec=contextual_confusion_matrix)}
        metrics = {**metrics, **self.metrics}
        returned = benchmark._evaluate_signal(
            self.pipeline, self.signal, self.hyper, metrics, True)

        expected_return = self.set_score(0, ANY, ANY)
        expected_return['status'] = 'ERROR'
        expected_return['tn'] = None
        expected_return['fp'] = 0
        expected_return['fn'] = 2
        expected_return['tp'] = 0

        assert returned == expected_return

    @patch('orion.benchmark.load_anomalies')
    @patch('orion.benchmark.analyze')
    @patch('orion.benchmark._load_pipeline')
    @patch('orion.benchmark.load_signal')
    def test__evaluate_signal_test_split(
            self, load_signal_mock, load_pipeline_mock, analyze_mock, load_anomalies_mock):
        train = Mock(autospec=pd.DataFrame)
        test = Mock(autospec=pd.DataFrame)
        load_signal_mock.side_effect = [train, test]
        load_pipeline_mock.return_value = self.pipeline

        test_split = True

        returned = benchmark._evaluate_signal(
            self.pipeline, self.signal, self.hyper, self.metrics, test_split=test_split)

        expected_return = self.set_score(1, ANY, test_split)
        assert returned == expected_return

        expected_calls = [
            call('signal-name-train'),
            call('signal-name-test')
        ]
        assert load_signal_mock.call_args_list == expected_calls

        load_pipeline_mock.assert_called_once_with(self.pipeline, self.hyper)
        analyze_mock.assert_called_once_with(self.pipeline, train, test)
        load_anomalies_mock.assert_called_once_with(self.signal)

    @patch('orion.benchmark.load_anomalies')
    @patch('orion.benchmark.analyze')
    @patch('orion.benchmark._load_pipeline')
    @patch('orion.benchmark.load_signal')
    def test__evaluate_signal_no_test_split(
            self, load_signal_mock, load_pipeline_mock, analyze_mock, load_anomalies_mock):
        train = test = Mock(autospec=pd.DataFrame)
        load_signal_mock.side_effect = [train, test]
        load_pipeline_mock.return_value = self.pipeline

        test_split = False

        returned = benchmark._evaluate_signal(
            self.pipeline, self.signal, self.hyper, self.metrics, test_split=test_split)

        expected_return = self.set_score(1, ANY, test_split)
        assert returned == expected_return

        expected_calls = [
            call('signal-name')
        ]
        assert load_signal_mock.call_args_list == expected_calls

        load_pipeline_mock.assert_called_once_with(self.pipeline, self.hyper)
        analyze_mock.assert_called_once_with(self.pipeline, train, test)
        load_anomalies_mock.assert_called_once_with(self.signal)

    @patch('orion.benchmark.load_anomalies')
    @patch('orion.benchmark.analyze')
    @patch('orion.benchmark._load_pipeline')
    @patch('orion.benchmark.load_signal')
    def test__evaluate_signal_no_detrend(
            self, load_signal_mock, load_pipeline_mock, analyze_mock, load_anomalies_mock):
        train = Mock(autospec=pd.DataFrame)
        test = Mock(autospec=pd.DataFrame)
        load_signal_mock.side_effect = [train, test]
        load_pipeline_mock.return_value = self.pipeline

        detrend = False

        returned = benchmark._evaluate_signal(
            self.pipeline, self.signal, self.hyper, self.metrics, test_split=True, detrend=detrend)

        expected_return = self.set_score(1, ANY, ANY)
        assert returned == expected_return

        expected_calls = [
            call('signal-name-train'),
            call('signal-name-test')
        ]
        assert load_signal_mock.call_args_list == expected_calls

        load_pipeline_mock.assert_called_once_with(self.pipeline, self.hyper)
        analyze_mock.assert_called_once_with(self.pipeline, train, test)
        load_anomalies_mock.assert_called_once_with(self.signal)

    @patch('orion.benchmark.load_anomalies')
    @patch('orion.benchmark.analyze')
    @patch('orion.benchmark._load_pipeline')
    @patch('orion.benchmark.load_signal')
    @patch('orion.benchmark._detrend_signal')
    def test__evaluate_signal_detrend(self, detrend_signal_mock, load_signal_mock,
                                      load_pipeline_mock, analyze_mock, load_anomalies_mock):

        train = Mock(autospec=pd.DataFrame)
        test = Mock(autospec=pd.DataFrame)
        detrend_signal_mock.side_effect = [train, test]
        load_signal_mock.side_effect = [train, test]
        load_pipeline_mock.return_value = self.pipeline

        detrend = True

        returned = benchmark._evaluate_signal(
            self.pipeline, self.signal, self.hyper, self.metrics, test_split=True, detrend=detrend)

        expected_return = self.set_score(1, ANY, ANY)
        assert returned == expected_return

        expected_calls = [
            call('signal-name-train'),
            call('signal-name-test')
        ]
        assert load_signal_mock.call_args_list == expected_calls

        expected_calls = [
            call(train, 'value'),
            call(test, 'value')
        ]
        assert detrend_signal_mock.call_args_list == expected_calls

        load_pipeline_mock.assert_called_once_with(self.pipeline, self.hyper)
        analyze_mock.assert_called_once_with(self.pipeline, train, test)
        load_anomalies_mock.assert_called_once_with(self.signal)

    @patch('orion.benchmark._evaluate_signal')
    def test__run_job(self, evaluate_signal_mock):
        pass

        score = self.set_score(1, ANY, ANY)
        evaluate_signal_mock.return_value = score

        benchmark._run_job(self.args())

        expected_calls = [
            self.pipeline, self.signal, self.hyper, self.metrics,
            self.test_split, self.detrend, self.pipeline_path
        ]

        evaluate_signal_mock.assert_called_once_with(*expected_calls)

    @patch('orion.benchmark._run_job')
    def test_benchmark(self, run_job_mock):
        signals = [self.signal]
        datasets = {self.dataset: signals}
        pipelines = {self.name: self.pipeline}

        output = self.set_output(1, ANY, ANY)
        run_job_mock.return_value = pd.DataFrame.from_records([output])

        order = [
            'pipeline',
            'rank',
            'dataset',
            'signal',
            'iteration',
            'metric-name',
            'status',
            'elapsed',
            'split',
            'run_id']

        expected_return = pd.DataFrame.from_records([{
            'metric-name': 1,
            'rank': 1,
            'elapsed': ANY,
            'split': ANY,
            'status': 'OK',
            'iteration': self.iteration,
            'pipeline': self.name,
            'dataset': self.dataset,
            'signal': self.signal,
            'run_id': self.run_id
        }])[order]

        returned = benchmark.benchmark(
            pipelines, datasets, self.hyper, self.metrics, self.rank)[order]

        pd.testing.assert_frame_equal(returned, expected_return)

        run_job_mock.assert_called_once_with(self.args())

    @patch('orion.benchmark._run_job')
    def test_benchmark_dataset_list(self, run_job_mock):
        signals = [self.signal]
        pipelines = {self.name: self.pipeline}

        output = self.set_output(1, ANY, ANY)
        output['dataset'] = 'dataset'
        run_job_mock.return_value = pd.DataFrame.from_records([output])

        order = [
            'pipeline',
            'rank',
            'dataset',
            'signal',
            'iteration',
            'metric-name',
            'status',
            'elapsed',
            'split',
            'run_id']

        expected_return = pd.DataFrame.from_records([{
            'metric-name': 1,
            'rank': 1,
            'elapsed': ANY,
            'split': ANY,
            'status': 'OK',
            'iteration': self.iteration,
            'pipeline': self.name,
            'dataset': 'dataset',
            'signal': self.signal,
            'run_id': self.run_id
        }])[order]

        returned = benchmark.benchmark(
            pipelines, signals, self.hyper, self.metrics, self.rank)[order]

        pd.testing.assert_frame_equal(returned, expected_return)

        args = list(self.args())
        args[2] = 'dataset'
        run_job_mock.assert_called_once_with(tuple(args))

    @patch('orion.benchmark._run_job')
    def test_benchmark_pipelines_list(self, run_job_mock):
        signals = [self.signal]
        datasets = {self.dataset: signals}
        pipelines = [self.pipeline]

        hyper = {}

        output = self.set_output(1, ANY, ANY)
        output['pipeline'] = self.pipeline
        run_job_mock.return_value = pd.DataFrame.from_records([output])

        order = [
            'pipeline',
            'rank',
            'dataset',
            'signal',
            'iteration',
            'metric-name',
            'status',
            'elapsed',
            'split',
            'run_id']

        expected_return = pd.DataFrame.from_records([{
            'metric-name': 1,
            'rank': 1,
            'elapsed': ANY,
            'split': ANY,
            'status': 'OK',
            'iteration': self.iteration,
            'pipeline': self.pipeline,
            'dataset': self.dataset,
            'signal': self.signal,
            'run_id': self.run_id
        }])[order]

        returned = benchmark.benchmark(
            pipelines, datasets, hyper, self.metrics, self.rank)[order]

        pd.testing.assert_frame_equal(returned, expected_return)

        args = list(self.args())
        args[1] = self.pipeline
        args[4] = hyper

        run_job_mock.assert_called_once_with(tuple(args))

    @patch('orion.benchmark._run_job')
    def test_benchmark_metrics_list(self, run_job_mock):
        pass

        signals = [self.signal]
        datasets = {self.dataset: signals}
        pipelines = {self.name: self.pipeline}

        metric = Mock(autospec=METRICS['f1'], return_value=1)
        metric.__name__ = 'metric-name'
        metrics = [metric]
        metrics_ = {metric.__name__: metric}

        output = self.set_output(1, ANY, ANY)
        output[metric.__name__] = metric
        run_job_mock.return_value = pd.DataFrame.from_records([output])

        order = [
            'pipeline',
            'rank',
            'dataset',
            'signal',
            'iteration',
            'metric-name',
            'status',
            'elapsed',
            'split',
            'run_id']

        expected_return = pd.DataFrame.from_records([{
            'metric-name': metric,
            'rank': 1,
            'elapsed': ANY,
            'split': ANY,
            'status': 'OK',
            'iteration': self.iteration,
            'pipeline': self.name,
            'dataset': self.dataset,
            'signal': self.signal,
            'run_id': self.run_id
        }])[order]

        returned = benchmark.benchmark(
            pipelines, datasets, self.hyper, metrics, self.rank)[order]

        pd.testing.assert_frame_equal(returned, expected_return)

        args = list(self.args())
        args[5] = metrics_

        run_job_mock.assert_called_once_with(tuple(args))

    def test_benchmark_metrics_exception(self):
        pass

        signals = [self.signal]
        datasets = {self.dataset: signals}
        pipelines = {self.name: self.pipeline}

        metric = 'does-not-exist'
        metrics = [metric]

        with self.assertRaises(ValueError) as ex:
            benchmark.benchmark(pipelines, datasets, self.hyper, metrics, self.rank)

            self.assertTrue(metric in ex.exception)
