'''Tests for wrangler.py'''

# Copyright (c) 2024 Carnegie Mellon University
# This code is subject to the license terms contained in the LICENSE file.

# pylint: disable=missing-function-docstring,duplicate-code,protected-access,too-many-ancestors
# pylint: disable=missing-class-docstring,redefined-outer-name,unused-variable

import glob
import json
import os
from pathlib import Path
from typing import Dict, Optional, Tuple

from numpy.random import RandomState
import pandas as pd
import pytest

from ..aggregators.impl.aggregator import AggregatorStub
from ..algorithms.impl.algorithm import Algorithm
from ..algorithms.impl.algorithm_auto import FakeAlgorithmCatalogAuto
from ..algorithms.impl.algorithm_instance import AlgorithmInstance
from ..algorithms.impl.fake_algorithm import FakeInstance
from ..executor.cucumber import Cucumber
from ..executor.executor import ExecutorStub
from ..executor.simple.simple_executor import SimpleExecutor
from ..generator.designator import Designator
from ..instantiator.executable_pipeline import (ExecutablePipeline,
                                                PipelineResult)
from ..problem_def.problem_def import ProblemDefinition
from ..templates.impl.pipeline_step import PipelineStep
from ..templates.impl.pipeline_template import PipelineTemplate
from ..templates.impl.template import TemplateCatalog

from .constants import Defaults
from .dataset import Dataset, RoleName
from .wrangler import Wrangler


class FakeExecutor(ExecutorStub):

    def predict(self,
                dataset: Dataset,
                pipelines: Dict[Designator, ExecutablePipeline]
                ) -> Dict[Designator, PipelineResult]:
        '''Use a list of pipelines to predict from a dataset.'''
        predictions = dataset.output()
        # create a dataframe with the shape we expect predictions to be:
        # 1 column (the target) and a number of rows equal to the size of the input
        predictions.predictions = pd.DataFrame({
            dataset.metadata.roles[RoleName.TARGET][0].name:
                range(0, dataset.dataframe.shape[0])})
        return {d: PipelineResult(executable_pipeline=p, prediction=predictions)
                for d, p in pipelines.items()}


MEMORY_PROBLEM_DEF = '''
{
    "dataset": {
        "config": "memory",
        "column_roles": {
            "target": {
                "name": "c"
            }
        },
        "params": {
            "train_data": "train_data",
            "test_data": "test_data"
        }
    },
    "problem_type": {
        "task": "regression"
    },
    "cross_validation": {
        "k": 10
    },
    "metrics" : {
        "root_mean_squared_error" : {}
    },
    "hyperparams": ["disable_grid_search"]
}
'''

AGGREGATION_PROBLEM_DEF = '''
{
    "dataset": {
        "config": "memory",
        "column_roles": {
            "target": {
                "name": "c"
            }
        },
        "params": {
            "train_data": "train_data",
            "test_data": "test_data"
        }
    },
    "problem_type": {
        "task": "regression"
    },
    "cross_validation": {
        "k": 10
    },
    "metrics": {
        "root_mean_squared_error" : {},
        "mean_squared_error": {},
        "mean_absolute_error": {}
    },
    "aggregation": {
        "method": ["method1", "method2"]
    },
    "hyperparams": ["disable_grid_search"]
}
'''


def get_data() -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
    data = pd.DataFrame(
        {
            'a': range(1, 21),
            'b': range(1, 21),
            'c': range(1, 21)
        }
    )
    return (data, data)


def memory_problem_def(has_aggregation: bool = False) -> ProblemDefinition:
    problem_def = MEMORY_PROBLEM_DEF
    if has_aggregation:
        problem_def = AGGREGATION_PROBLEM_DEF

    return ProblemDefinition(clause=problem_def)


def test_wrangler_split() -> None:
    (train_data, test_data) = get_data()  # noqa: F841
    problem_def = memory_problem_def()
    dut = Wrangler(
        problem_definition=problem_def,
        algorithm_catalog=FakeAlgorithmCatalogAuto,
        executor=FakeExecutor)

    train_result = dut.fit_predict_rank()

    # we expect 10 folds because k=10 was set in the problem definition
    assert len(train_result.split_dataset.folds) == 10
    # we expect each fold's validate set to have 2 rows,
    #    because there are 10 folds and the dataset is 20 rows long,
    #    and 2 columns, because it contains the attributes but not the target.
    assert train_result.split_dataset.folds[0].validate.dataframe.shape == (2, 2)
    # we expect each fold's train set to have 18 rows,
    #    because it is composed of 9 folds of 2 rows each,
    #    and 3 columns, becase it contains the attributes and the target.
    assert train_result.split_dataset.folds[0].train.dataframe.shape == (18, 3)
    # we expect the ground truth to have 20 rows,
    #    because it covers the whole dataset,
    #    and 1 col, because it only has the target.
    assert train_result.split_dataset.ground_truth is not None
    assert train_result.split_dataset.ground_truth.ground_truth.shape == (20, 1)


def test_fit_predict_rank() -> None:
    (train_data, test_data) = get_data()  # noqa: F841
    problem_def = memory_problem_def()

    dut = Wrangler(
        problem_definition=problem_def,
        algorithm_catalog=FakeAlgorithmCatalogAuto)

    got = dut.fit_predict_rank()

    # we expect pipelines to be named according to thier template (tabular regression)
    #    and the model that was chosen from thier query step.
    #    in the case of random forest regressor, wrangler will do grid search,
    #    and thus hyperparam values will also be in the designator.
    random_forest_des = Designator('tabular_regression@sklearn.ensemble.randomforestregressor:'
                                   'max_depth=none,min_samples_split=2')
    linear_regression_des = Designator('tabular_regression@sklearn.linear_model.linearregression')
    assert {random_forest_des, linear_regression_des} == set(got.executable_pipelines.keys())

    assert len(got.train_results) == 2  # 2 pipelines
    got_result = got.train_results.predictions[random_forest_des]
    assert got_result is not None
    assert got_result.predictions is not None
    assert got_result.predictions.shape == (20, 2)

    assert len(got.rankings) == 1  # 1 metric
    assert 'root_mean_squared_error' in got.rankings

    assert got.test_results is not None
    assert len(got.test_results) == 2  # 2 pipelines

    # executable pipelines associated with the test results should be identical
    assert set(got.executable_pipelines.keys()) == set(got.test_results.executable_pipelines.keys())
    got_result = got.test_results.predictions[random_forest_des]
    assert got_result is not None
    assert got_result.predictions is not None
    assert got_result.predictions.shape == (20, 1)


def test_fit_predict_rank_with_aggregation() -> None:
    (train_data, test_data) = get_data()  # noqa: F841
    problem_def = memory_problem_def(has_aggregation=True)

    dut = Wrangler(
        problem_definition=problem_def,
        algorithm_catalog=FakeAlgorithmCatalogAuto)

    dut.aggregator_catalog.register(AggregatorStub(name="method1"))
    dut.aggregator_catalog.register(AggregatorStub(name="method2"))

    got = dut.fit_predict_rank()

    assert len(got.rankings) == 5  # 3 metrics, and 2 aggregate metrics
    want = {
        'root_mean_squared_error', 'mean_squared_error', 'mean_absolute_error',
        'method1', 'method2'
    }
    assert want == set(got.rankings.keys())


class FakeTemplateCatalog(TemplateCatalog):
    '''fake, no templates'''


def memory_test_df() -> pd.DataFrame:
    return pd.DataFrame(
        {
            'a': range(1, 11),
            'b': range(1, 11)
        }
    )


def memory_train_df() -> pd.DataFrame:
    return pd.DataFrame(
        {
            'a': range(1, 21),
            'b': range(1, 21),
            'c': range(1, 21)
        }
    )


def test_predict() -> None:
    (train_data, test_data) = get_data()  # noqa: F841
    problem_def = memory_problem_def()
    executor = SimpleExecutor
    alg_catalog = FakeAlgorithmCatalogAuto
    dut = Wrangler(
        problem_definition=problem_def,
        executor=executor,
        algorithm_catalog=alg_catalog)
    got = dut.fit_predict_rank()

    fake_des = Designator('fake_des')

    testset = dut.dataset(memory_test_df())
    got_result = list(got.train_results.values())[0]
    pipeline = got_result.executable_pipeline
    assert pipeline is not None
    test_got = dut.predict(
        new_data=testset,
        trained_pipelines={fake_des: pipeline})
    got_result = test_got[fake_des]
    assert got_result is not None
    assert got_result.prediction is not None
    prediction_df = got_result.prediction.predictions

    # we expect 10 rows (length of test dataframe) and 1 col (target)
    assert prediction_df.shape == (10, 1)
    # we expect to be able to predict perfectly since test data is identical to train data.
    assert prediction_df['c'][9] == 10.0


def test_predict_defaulted() -> None:
    (train_data, test_data) = get_data()  # noqa: F841
    problem_def = memory_problem_def()
    executor = SimpleExecutor
    alg_catalog = FakeAlgorithmCatalogAuto
    dut = Wrangler(
        problem_definition=problem_def,
        executor=executor,
        algorithm_catalog=alg_catalog)
    _ = dut.fit_predict_rank()

    testset = dut.dataset(memory_test_df())
    test_got = dut.predict(
        new_data=testset)
    assert {
        'tabular_regression@sklearn.linear_model.linearregression',
        'tabular_regression@sklearn.ensemble.randomforestregressor:'
        'max_depth=none,min_samples_split=2',
    } == set(test_got.keys())


def test_set_current_predict() -> None:
    (train_data, test_data) = get_data()  # noqa: F841
    problem_def = memory_problem_def()
    executor = SimpleExecutor
    alg_catalog = FakeAlgorithmCatalogAuto
    dut = Wrangler(
        problem_definition=problem_def,
        executor=executor,
        algorithm_catalog=alg_catalog)
    _ = dut.fit_predict_rank()

    testset = dut.dataset(memory_test_df())

    # Pick one pipeline.
    dut.set_current('tabular_regression@sklearn.ensemble.randomforestregressor:'
                    'max_depth=none,min_samples_split=2')
    test_got = dut.predict(new_data=testset)
    assert set(test_got.keys()) == {
        'tabular_regression@sklearn.ensemble.randomforestregressor:'
        'max_depth=none,min_samples_split=2',
    }

    # Reset to all known pipelines.
    dut.set_current()

    test_got = dut.predict(new_data=testset)
    assert {
        'tabular_regression@sklearn.linear_model.linearregression',
        'tabular_regression@sklearn.ensemble.randomforestregressor:'
        'max_depth=none,min_samples_split=2',
    } == set(test_got.keys())


def test_set_current_fit() -> None:
    (train_data, test_data) = get_data()
    problem_def = memory_problem_def()
    executor = SimpleExecutor
    alg_catalog = FakeAlgorithmCatalogAuto
    dut = Wrangler(
        problem_definition=problem_def,
        executor=executor,
        algorithm_catalog=alg_catalog)
    _ = dut.fit_predict_rank()

    testset = dut.dataset(memory_train_df())

    # Pick one pipeline.
    dut.set_current('tabular_regression@sklearn.ensemble.randomforestregressor:'
                    'max_depth=none,min_samples_split=2')

    test_got = dut.fit(dataset=testset)
    assert set(test_got.keys()) == {
        'tabular_regression@sklearn.ensemble.randomforestregressor:'
        'max_depth=none,min_samples_split=2',
    }

    # Reset to all known pipelines.
    dut.set_current()

    test_got = dut.fit(dataset=testset)
    assert {
        'tabular_regression@sklearn.linear_model.linearregression',
        'tabular_regression@sklearn.ensemble.randomforestregressor:'
        'max_depth=none,min_samples_split=2',
    } == set(test_got.keys())


@pytest.fixture(scope='session')
def saver_problem_def(tmp_path_factory: pytest.TempPathFactory) -> ProblemDefinition:
    tmp_path = tmp_path_factory.mktemp('data')
    tmp_dir = tmp_path / 'sub'
    tmp_dir.mkdir()
    config = {
        "dataset": {
            "config": "memory",
            "column_roles": {
                "target": {
                    "name": "c"
                }
            },
            "params": {
                "train_data": "train_data",
                "test_data": "test_data"
            }
        },
        "problem_type": {
            "task": "regression"
        },
        "cross_validation": {
            "k": 10
        },
        "metrics": {
            "root_mean_squared_error": {}
        },
        "output": {
            "path": tmp_dir,
            "instantiations": [
                "json"
            ],
            "file_type": "csv"
        },
        "hyperparams": ["disable_grid_search"]
    }

    return ProblemDefinition(clause=config)


def test_save_filesystem(saver_problem_def: ProblemDefinition) -> None:
    PipelineStep.reset_serial_number()
    (train_data, test_data) = get_data()
    dut = Wrangler(
        problem_definition=saver_problem_def,
        algorithm_catalog=FakeAlgorithmCatalogAuto)
    _ = dut.fit_predict_rank()

    path = saver_problem_def.output.path
    assert path is not None
    filenames = glob.glob(str(path / 'models' / '*.linearregression' / '@connect_*@.pkl'))
    assert len(filenames) == 1
    with Path(filenames[0]).open(mode='rb') as filepointer:
        cucumber = Cucumber.deserialize(filepointer.read())
    assert cucumber.hyperparams == {
        'target': ['target', 'dataframe'],
        'covariates': ['attributes', 'dataframe']
    }

    pipelines = os.listdir(str(path / 'pipelines'))
    assert {'tabular_regression@sklearn.ensemble.randomforestregressor:'
            'max_depth=none,min_samples_split=2.json',
            'tabular_regression@sklearn.linear_model.linearregression.json',
            } == set(pipelines)

    with (path / 'pipelines' / 'tabular_regression@sklearn.linear_model.linearregression.json'
          ).open() as file_pointer:
        got = json.load(file_pointer)

    assert got['pipeline_designator'] == 'tabular_regression@sklearn.linear_model.linearregression'


@pytest.fixture(scope='session')
def search_problem_def() -> ProblemDefinition:
    config = '''
    {
        "dataset": {
            "config": "memory",
            "column_roles": {
                "target": {
                    "name": "c"
                }
            },
            "params": {
                "train_data": "train_data",
                "test_data": "test_data"
            }
        },
        "problem_type": {
            "task": "regression"
        },
        "cross_validation": {
            "k": 2
        },
        "metrics" : {
            "root_mean_squared_error" : {}
        },
        "hyperparams": [
            {
                "select": {
                    "algorithm": "sklearn.ensemble.RandomForestRegressor"
                },
                "params": {
                    "n_estimators": {
                        "list": [100, 200]
                    }
                }
            }
        ]
    }
    '''

    return ProblemDefinition(
        clause=config)


def test_hyperparam_search(search_problem_def: ProblemDefinition) -> None:
    (train_data, test_data) = get_data()
    dut = Wrangler(
        problem_definition=search_problem_def,
        algorithm_catalog=FakeAlgorithmCatalogAuto)
    got = dut.fit_predict_rank()
    assert len(got.executable_pipelines) == 25
    assert {
        'tabular_regression@sklearn.linear_model.linearregression',
        'tabular_regression@sklearn.ensemble.randomforestregressor:'
        'max_depth=8,min_samples_split=10,n_estimators=100',
        'tabular_regression@sklearn.ensemble.randomforestregressor:'
        'max_depth=15,min_samples_split=10,n_estimators=200',
    }.issubset(got.executable_pipelines.keys())


@pytest.fixture(scope='session')
def order_problem_def() -> ProblemDefinition:
    config = '''
    {
        "dataset": {
            "config": "memory",
            "column_roles": {
                "target": {
                    "name": "c"
                }
            },
            "params": {
                "train_data": "train_data",
                "test_data": "test_data"
            }
        },
        "problem_type": {
            "task": "test_task"
        },
        "cross_validation": {
            "k": 3
        },
        "metrics" : {
            "root_mean_squared_error" : {}
        },
        "hyperparams": ["disable_grid_search"]
    }
    '''

    return ProblemDefinition(clause=config)


class FakeInst(FakeInstance):
    '''The instance of a fake algorithm for order tests.'''
    def predict(self, dataset: Optional[Dataset]) -> Optional[Dataset]:
        if dataset is None:
            return None
        retval = dataset.output()
        df = pd.DataFrame({'c': dataset.dataframe['a'] + 1})
        retval.predictions = df
        return retval


class FakeAlg(Algorithm):
    '''A fake algorithm used for order tests.'''
    _name = 'fake_algorithm'

    def instantiate(self, **hyperparams) -> 'FakeInstance':
        return FakeInst(parent=self, **hyperparams)


def test_order(order_problem_def) -> None:
    train_data = pd.DataFrame(
        {
            'a': range(1, 22),
            'b': range(1, 22),
            'c': range(2, 23)
        }
    )
    test_data = train_data  # noqa: F841 pylint: disable=unused-variable
    dut = Wrangler(
        problem_definition=order_problem_def)

    fake_template = PipelineTemplate(
        name='test_template',
        tags={
            'task': ['test_task'],
            'data_type': ['tabular']},
        generator=dut.generator
    )
    fake_template.step(model=FakeAlg())
    dut.template_catalog.register(fake_template)

    got = dut.fit_predict_rank()
    results = got.train_results.predictions[Designator('test_template')]
    assert results is not None
    assert results.predictions is not None
    predict_df = results.predictions
    assert got.split_dataset.ground_truth is not None
    ground_truth_df = got.split_dataset.ground_truth.ground_truth
    pd.testing.assert_frame_equal(predict_df[['c']], ground_truth_df)
    assert all(predict_df['c'] == predict_df['ground truth'])


def non_seeded_problem_def() -> ProblemDefinition:
    config = '''
    {
        "dataset": {
            "config": "memory",
            "column_roles": {
                "target": {
                    "name": "c"
                }
            },
            "params": {
                "train_data": "train_data",
                "test_data": "test_data"
            }
        },
        "problem_type": {
            "task": "test_task"
        },
        "cross_validation": {
            "k": 2
        },
        "metrics" : {
            "root_mean_squared_error" : {}
        },
        "hyperparams": ["disable_grid_search"]
    }
    '''
    train_data = pd.DataFrame(  # noqa: F841 pylint: disable=unused-variable
        {
            'a': range(1, 21),
            'b': range(1, 21),
            'c': range(2, 22)
        }
    )

    return ProblemDefinition(clause=config)


def seeded_problem_def(seed: int) -> ProblemDefinition:
    config = {
        "dataset": {
            "config": "memory",
            "column_roles": {
                "target": {
                    "name": "c"
                }
            },
            "params": {
                "train_data": "train_data",
                "test_data": "test_data"
            }
        },
        "problem_type": {
            "task": "test_task"
        },
        "cross_validation": {
            "k": 2
        },
        "metrics": {
            "root_mean_squared_error": {}
        },
        "hyperparams": [
            "disable_grid_search",
            {
                "select": {
                    "tags": {
                        "supports_random_seed": "true"
                    }
                },
                "params": {
                    "random_seed": {
                        "fixed": seed
                    }
                }
            }
        ]
    }
    train_data = pd.DataFrame(  # noqa: F841 pylint: disable=unused-variable
        {
            'a': range(1, 21),
            'b': range(1, 21),
            'c': range(2, 22)
        }
    )

    return ProblemDefinition(clause=config)


class RandomAlgorithmInstance(AlgorithmInstance):
    _rng: RandomState

    def __init__(self, parent: Algorithm, **hyperparams):
        super().__init__(parent)
        if 'random_seed' in hyperparams:
            self._rng = RandomState(hyperparams['random_seed'])

    def predict(self, dataset: Optional[Dataset]) -> Optional[Dataset]:
        if dataset is None:
            return None
        retval = dataset.output()
        retval['predictions'] = pd.DataFrame({
            'c': list(self._rng.randint(0, 10, len(dataset.dataframe)))
        }).reset_index(drop=True)
        return retval


class RandomAlgorithm(Algorithm):
    _name = 'random_algorithm'
    _tags = {
        'task': ['test_task'],
        'data_type': ['tabular'],
        'supports_random_seed': ['true']
    }
    _default_hyperparams = {
        'random_seed': Defaults.SEED,
    }

    def instantiate(self, **hyperparams) -> RandomAlgorithmInstance:
        return RandomAlgorithmInstance(parent=self, **self.hyperparams(**hyperparams))


def test_set_seed() -> None:
    (train_data, test_data) = get_data()
    dut_def = Wrangler(problem_definition=non_seeded_problem_def())
    dut5678 = Wrangler(problem_definition=seeded_problem_def(5678))

    pipe_def = PipelineTemplate(
        name='test_template',
        tags={
            'task': ['test_task'],
            'data_type': ['tabular']},
        generator=dut_def.generator
    )
    pipe_def.step(model=RandomAlgorithm())
    dut_def.template_catalog.register(pipe_def)

    pipe5678 = PipelineTemplate(
        name='test_template',
        tags={
            'task': ['test_task'],
            'data_type': ['tabular']},
        generator=dut5678.generator
    )
    pipe5678.step(model=RandomAlgorithm())
    dut5678.template_catalog.register(pipe5678)

    got_def = dut_def.fit_predict_rank()
    got5678 = dut5678.fit_predict_rank()

    des_def = Designator('test_template')
    des5678 = Designator('test_template@random_algorithm:random_seed=5678')

    got_def_train_results_des_def = got_def.train_results[des_def]
    got5678_train_results_des5678 = got5678.train_results[des5678]

    assert got_def_train_results_des_def.prediction is not None
    assert got5678_train_results_des5678.prediction is not None

    predictions_def = got_def_train_results_des_def.prediction.predictions[['c']]
    predictions5678 = got5678_train_results_des5678.prediction.predictions[['c']]

    # The k-fold-cross-validator instantiates the algorithm separately
    # for each fold, so in this special case, we expect to see the same
    # sequence twice.

    want_def = pd.DataFrame({'c': [3, 3, 6, 2, 1, 8, 5, 1, 7, 3, 3, 3, 6, 2, 1, 8, 5, 1, 7, 3]})
    want5678 = pd.DataFrame({'c': [5, 3, 6, 9, 3, 3, 9, 8, 3, 3, 5, 3, 6, 9, 3, 3, 9, 8, 3, 3]})
    pd.testing.assert_frame_equal(predictions_def, want_def)
    pd.testing.assert_frame_equal(predictions5678, want5678)


@pytest.fixture(scope='session')
def disable_grid_search_problem_def() -> ProblemDefinition:
    config = '''
    {
        "dataset": {
            "config": "memory",
            "column_roles": {
                "target": {
                    "name": "c"
                }
            },
            "params": {
                "train_data": "train_data",
                "test_data": "test_data"
            }
        },
        "problem_type": {
            "task": "regression"
        },
        "cross_validation": {
            "k": 2
        },
        "metrics" : {
            "root_mean_squared_error" : {}
        },
        "hyperparams": [
            "disable_grid_search",
            {
                "select": {
                    "algorithm": "sklearn.ensemble.RandomForestRegressor"
                },
                "params": {
                    "n_estimators": {
                        "list": [100, 200],
                        "default": 100
                    }
                }
            }
        ],
        "output": {}
    }
    '''

    return ProblemDefinition(clause=config)


def test_disable_grid_search(disable_grid_search_problem_def: ProblemDefinition) -> None:
    (train_data, test_data) = get_data()
    dut = Wrangler(
        problem_definition=disable_grid_search_problem_def,
        algorithm_catalog=FakeAlgorithmCatalogAuto)
    got = dut.fit_predict_rank()
    assert {
        'tabular_regression@sklearn.linear_model.linearregression',
        'tabular_regression@sklearn.ensemble.randomforestregressor'
        ':max_depth=none,min_samples_split=2,n_estimators=100'
    } == set(got.executable_pipelines.keys())
