# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.


from pathlib import Path
from typing import List, NamedTuple
import gluonts_customized
import logging
import numpy as np
import os
import pandas as pd
import pytest
import random
import sys
import tempfile
import warnings

from gluonts_customized.model import Predictor
from gluonts_customized.dataset.common import ListDataset

try:
    import mxnet as mx
except (ImportError, OSError):
    mx = None

try:
    import statsmodels
except ImportError:
    statsmodels = None


class HierarchicalMetaData(NamedTuple):
    S: np.ndarray
    freq: str
    nodes: List


class HierarchicalTrainDatasets(NamedTuple):
    train: ListDataset
    test: ListDataset
    metadata: HierarchicalMetaData


@pytest.fixture
def sine7():
    def _sine7(
            seq_length: int = 100,
            prediction_length: int = 10,
            nonnegative: bool = False,
            bias: float = 0,
    ):
        x = np.arange(0, seq_length)

        # Bottom layer (4 series)
        amps = [0.8, 0.9, 1, 1.1]
        freqs = [1 / 20, 1 / 30, 1 / 50, 1 / 100]

        b = np.zeros((4, seq_length))
        for i, f in enumerate(freqs):
            omega = 0
            if i == 3:
                np.random.seed(0)
                omega = np.random.uniform(0, np.pi)  # random phase shift
            b[i, :] = amps[i] * np.sin(2 * np.pi * x * f + omega) + bias

        if nonnegative:
            b = abs(b)

        # Aggregation matrix S
        S = np.array(
            [
                [1, 1, 1, 1],
                [1, 1, 0, 0],
                [0, 0, 1, 1],
                [1, 0, 0, 0],
                [0, 1, 0, 0],
                [0, 0, 1, 0],
                [0, 0, 0, 1],
            ],
        )

        Y = S @ b

        # Indices and timestamps
        index = pd.period_range(
            start=pd.Period("2020-01-01", freq="D"),
            periods=Y.shape[1],
            freq="D",
        )

        metadata = HierarchicalMetaData(
            S=S, freq=index.freqstr, nodes=[2, [2] * 2]
        )

        train_dataset = ListDataset(
            [
                {
                    "start": index[0],
                    "item_id": "all_items",
                    "target": Y[:, :-prediction_length],
                }
            ],
            freq=index.freqstr,
            one_dim_target=False,
        )

        test_dataset = ListDataset(
            [{"start": index[0], "item_id": "all_items", "target": Y}],
            freq=index.freqstr,
            one_dim_target=False,
        )

        assert Y.shape[0] == S.shape[0]
        return HierarchicalTrainDatasets(
            train=train_dataset, test=test_dataset, metadata=metadata
        )

    return _sine7


@pytest.fixture(scope="function", autouse=True)
def function_scope_seed(request):
    """A function scope fixture that manages rng seeds.
    This fixture automatically initializes the python, numpy and mxnet random
    number generators randomly on every test run.
    def test_ok_with_random_data():
        ...
    To fix the seed used for a test case mark the test function with the
    desired seed:
    @pytest.mark.seed(1)
    def test_not_ok_with_random_data():
        '''This testcase actually works.'''
        assert 17 == random.randint(0, 100)
    When a test fails, the fixture outputs the seed used. The user can then set
    the environment variable MXNET_TEST_SEED to the value reported, then rerun
    the test with:
        pytest --verbose -s <test_module_name.py> -k <failing_test>
    To run a test repeatedly, install pytest-repeat and add the --count argument:
        pip install pytest-repeat
        pytest --verbose -s <test_module_name.py> -k <failing_test> --count 1000
    """

    seed = request.node.get_closest_marker("seed")
    env_seed_str = os.getenv("MXNET_TEST_SEED")

    if seed is not None:
        seed = seed.args[0]
        assert isinstance(seed, int)
    elif env_seed_str is not None:
        seed = int(env_seed_str)
    else:
        seed = int.from_bytes(os.urandom(4), "big")

    post_test_state = np.random.get_state()
    np.random.seed(seed)
    random.seed(seed)

    if mx is not None:
        mx.random.seed(seed)

    seed_message = (
        "np/mx/python random seeds are set to "
        f"{seed}, use MXNET_TEST_SEED={seed} to reproduce."
    )

    # Always log seed on DEBUG log level. This makes sure we can find out the
    # value of the seed even if the test case causes a segfault and subsequent
    # teardown code is not run.
    logging.debug(seed_message)

    yield  # run the test

    np.random.set_state(post_test_state)


@pytest.fixture(autouse=True)
def doctest(doctest_namespace):
    doctest_namespace["np"] = np
    doctest_namespace["gluonts"] = gluonts_customized

    if mx is not None:
        doctest_namespace["mx"] = mx
        doctest_namespace["gluon"] = mx.gluon

    import doctest

    doctest.ELLIPSIS_MARKER = "-etc-"


collect_ignore = []


def pytest_configure(config):
    test_folder = Path(__file__).parent.resolve()

    targets = config.getoption("file_or_dir")

    if not targets:
        targets.append(".")

    requirements = set()

    for target in targets:
        target = Path(target).resolve()

        if target.is_file():
            target = target.parent
        else:
            requirements.update(target.glob("**/require-packages.txt"))

        try:
            # .is_relative_to was only added in Py 3.9
            target.relative_to(test_folder)

            while True:
                require = target / "require-packages.txt"
                if require.exists():
                    requirements.add(require)

                if target == test_folder:
                    break
                target = target.parent
        except ValueError:
            pass

    old_path = sys.path
    sys.path = [
        path for path in sys.path if Path(path).resolve() != test_folder
    ]

    excludes = []

    for path in requirements:
        with path.open() as requirement:
            for requirement in map(str.strip, requirement):
                try:
                    __import__(requirement)
                except ImportError:
                    collect_ignore.append(
                        str(path.parent.relative_to(test_folder))
                    )
                    break

    if collect_ignore:
        warnings.warn(
            f"Skipping tests because some packages are not installed: {excludes}"
        )

    sys.path = old_path


class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self


def pytest_runtest_setup(item):
    skip_datasets = [
        mark.args[0] for mark in item.iter_markers(name="skip_dataset")
    ]

    if skip_datasets:
        ds_name = item._request.getfixturevalue("dsinfo")["name"]
        if ds_name in skip_datasets:
            pytest.skip(f"Skip test on dataset {ds_name}")


@pytest.fixture(scope="session", params=["synthetic", "constant"])
def dsinfo(request):
    from gluonts_customized import time_feature
    from gluonts_customized.dataset.artificial import constant_dataset, default_synthetic

    if request.param == "constant":
        ds_info, train_ds, test_ds = constant_dataset()

        return AttrDict(
            name="constant",
            cardinality=int(ds_info.metadata.feat_static_cat[0].cardinality),
            freq=ds_info.metadata.freq,
            num_parallel_samples=2,
            prediction_length=ds_info.prediction_length,
            # FIXME: Should time features should not be needed for GP
            time_features=[time_feature.day_of_week, time_feature.hour_of_day],
            train_ds=train_ds,
            test_ds=test_ds,
        )
    elif request.param == "synthetic":
        ds_info, train_ds, test_ds = default_synthetic()

        return AttrDict(
            name="synthetic",
            batch_size=32,
            cardinality=int(ds_info.metadata.feat_static_cat[0].cardinality),
            context_length=2,
            freq=ds_info.metadata.freq,
            prediction_length=ds_info.prediction_length,
            num_parallel_samples=2,
            train_ds=train_ds,
            test_ds=test_ds,
            time_features=None,
        )


def from_hyperparameters(Forecaster, hyperparameters, dsinfo):
    return Forecaster.from_hyperparameters(
        freq=dsinfo.freq,
        **{
            "prediction_length": dsinfo.prediction_length,
            "num_parallel_samples": dsinfo.num_parallel_samples,
        },
        **hyperparameters,
    )


@pytest.fixture()
def accuracy_test(dsinfo):
    from gluonts_customized.evaluation import backtest_metrics, Evaluator

    def test_accuracy(Forecaster, hyperparameters, accuracy):
        forecaster = from_hyperparameters(Forecaster, hyperparameters, dsinfo)

        if isinstance(forecaster, Predictor):
            predictor = forecaster
        else:
            predictor = forecaster.train(training_data=dsinfo.train_ds)

        agg_metrics, item_metrics = backtest_metrics(
            test_dataset=dsinfo.test_ds,
            predictor=predictor,
            evaluator=Evaluator(
                calculate_owa=statsmodels is not None, num_workers=0
            ),
        )

        if dsinfo.name == "synthetic":
            accuracy = 10.0

        assert agg_metrics["ND"] <= accuracy

    return test_accuracy


@pytest.fixture()
def serialize_test(dsinfo):
    from gluonts_customized.model.predictor import Predictor

    def test_serialize(Forecaster, hyperparameters):
        forecaster = from_hyperparameters(Forecaster, hyperparameters, dsinfo)

        if isinstance(forecaster, Predictor):
            predictor_act = forecaster
        else:
            predictor_act = forecaster.train(dsinfo.train_ds)

        with tempfile.TemporaryDirectory() as temp_dir:
            predictor_act.serialize(Path(temp_dir))
            predictor_exp = Predictor.deserialize(Path(temp_dir))
            # TODO: DeepFactorEstimator does not pass this assert
            assert predictor_act == predictor_exp

    return test_serialize
