# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import importlib
from operator import attrgetter
from typing import Any
from typing import Dict
from typing import List
from typing import Union
from typing import cast
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pandas as pd
from parameterized.parameterized import parameterized

import kats.tsfeatures.tsfeatures
from kats.compat import statsmodels
from kats.consts import TimeSeriesData
from kats.data.utils import load_air_passengers
from kats.tsfeatures.tsfeatures import _FEATURE_GROUP_MAPPING
from kats.tsfeatures.tsfeatures import TsFeatures


def _univariate_features(feats: Union[Dict[str, float], List[Dict[str, float]]]) -> Dict[str, float]:
    return cast(Dict[str, float], feats)


def _multivariate_features(feats: Union[Dict[str, float], List[Dict[str, float]]]) -> List[Dict[str, float]]:
    return cast(List[Dict[str, float]], feats)


class TSfeaturesTest(TestCase):
    def setUp(self) -> None:
        DATA = load_air_passengers(return_ts=False)
        self.TSData = TimeSeriesData(DATA)
        self.TSData_multivariate = TimeSeriesData(
            pd.DataFrame({"time": DATA.time, "value1": DATA.y, "value2": DATA.y})
        )

        self.TSData_short = TimeSeriesData(DATA.iloc[:8, :])
        self.TSData_mini = TimeSeriesData(DATA.iloc[:2, :])

    def assertDictAlmostEqual(self, expected: Dict[str, Any], features: Dict[str, Any], places: int = 4) -> None:
        """Compares that two dictionaries are floating-point almost equal.

        Note: the dictionaries may or may not contain floating-point values.

        Args:
          expected: the expected dictionary of values.
          actual: the actual dictionary of values.
          places: the number of decimal places for floating point comparisons.
        """

        self.assertEqual(expected.keys(), features.keys())
        for k, v in expected.items():
            if isinstance(v, float):
                if np.isnan(v):
                    self.assertTrue(np.isnan(features[k]), msg=f"{k} differ")
                else:
                    self.assertAlmostEqual(v, features[k], places=4, msg=f"{k} differ")
            else:
                self.assertEqual(v, features[k], msg=f"{k} differ")

    def test_all_feature_names_unique(self) -> None:
        features = set()
        for _, feats in _FEATURE_GROUP_MAPPING.items():
            for feat in feats:
                self.assertFalse(feat in features, f"duplicate feature name {feat}")
                features.add(feat)

    def test_tsfeatures(self) -> None:
        feature_vector = _univariate_features(TsFeatures().transform(self.TSData))

        feature_vector_round = {key: round(feature_vector[key], 6) for key in feature_vector}

        # test there is no nan in feature vector
        self.assertEqual(
            np.isnan(np.asarray(list(feature_vector.values()))).any(),
            False,
        )

        # test there are 40 features in the feature vector now
        self.assertEqual(
            len(np.asarray(list(feature_vector.values()))) == 40,
            True,
        )

        # test feature vector value
        rounded_truth = {
            "length": 144,
            "mean": 280.298611,
            "var": 14291.973331,
            "entropy": 0.428737,
            "lumpiness": 3041164.562906,
            "stability": 12303.627267,
            "flat_spots": 2,
            "hurst": -0.080233,
            "std1st_der": 27.206288,
            "crossing_points": 7,
            "binarize_mean": 0.444444,
            "unitroot_kpss": 0.128475,
            "heterogeneity": 126.064506,
            "histogram_mode": 155.8,
            "linearity": 0.853638,
            "trend_strength": 0.9681,
            "seasonality_strength": 0.440863,
            "spikiness": 33.502886,
            "peak": 6,
            "trough": 3,
            "level_shift_idx": 118,
            "level_shift_size": 15.6,
            "y_acf1": 0.948047,
            "y_acf5": 3.392072,
            "diff1y_acf1": 0.302855,
            "diff1y_acf5": 0.259459,
            "diff2y_acf1": -0.191006,
            "diff2y_acf5": 0.134207,
            "y_pacf5": 1.003288,
            "diff1y_pacf5": 0.219412,
            "diff2y_pacf5": 0.26101,
            "seas_acf1": 0.662904,
            "seas_pacf1": 0.15617,
            "firstmin_ac": 8,
            "firstzero_ac": 52,
            "holt_alpha": 1.0,
            "holt_beta": 0.0,
            "hw_alpha": 0.842106,
            "hw_beta": 0.052631,
            "hw_gamma": 0.157901,
        }
        if statsmodels.version >= "0.12":
            rounded_truth["trend_strength"] = 0.93833
            rounded_truth["seasonality_strength"] = 0.329934
            rounded_truth["spikiness"] = 111.697325
            feature_vector_round["holt_alpha"] = np.round(feature_vector_round["holt_alpha"], 1)
            feature_vector_round["holt_beta"] = np.round(feature_vector_round["holt_beta"], 1)
            rounded_truth["holt_alpha"] = 1.0
            rounded_truth["holt_beta"] = 0.0
            rounded_truth["hw_alpha"] = 1.0
            rounded_truth["hw_beta"] = 0.0
            rounded_truth["hw_gamma"] = 0.0
        self.assertEqual(feature_vector_round, rounded_truth)

    # pyre-fixme[56]: Pyre was not able to infer the type of the decorator `parameter...
    @parameterized.expand([("univariate", "TSData"), ("multivariate", "TSData_multivariate")])
    def test_feature_selections(self, test_name: str, ts_name: str) -> None:
        # test disabling functions

        ts = attrgetter(ts_name)(self)
        features = TsFeatures(
            unitroot_kpss=False,
            histogram_mode=False,
            diff2y_pacf5=False,
            firstmin_ac=False,
        ).transform(ts)

        if ts.is_univariate():
            feature_list = [_univariate_features(features)]
        else:
            feature_list = _multivariate_features(features)

        rounded_truth = {
            "length": 144,
            "mean": 280.298611,
            "var": 14291.973331,
            "entropy": 0.428737,
            "lumpiness": 3041164.562906,
            "stability": 12303.627267,
            "flat_spots": 2,
            "hurst": -0.080233,
            "std1st_der": 27.206288,
            "crossing_points": 7,
            "binarize_mean": 0.444444,
            "heterogeneity": 126.064506,
            "linearity": 0.853638,
            "trend_strength": 0.9681,
            "seasonality_strength": 0.440863,
            "spikiness": 33.502886,
            "peak": 6,
            "trough": 3,
            "level_shift_idx": 118,
            "level_shift_size": 15.6,
            "y_acf1": 0.948047,
            "y_acf5": 3.392072,
            "diff1y_acf1": 0.302855,
            "diff1y_acf5": 0.259459,
            "diff2y_acf1": -0.191006,
            "diff2y_acf5": 0.134207,
            "y_pacf5": 1.003288,
            "diff1y_pacf5": 0.219412,
            "seas_acf1": 0.662904,
            "seas_pacf1": 0.15617,
            "firstzero_ac": 52,
            "holt_alpha": 1.0,
            "holt_beta": 0.0,
            "hw_alpha": 0.842106,
            "hw_beta": 0.052631,
            "hw_gamma": 0.157901,
        }

        for feature_vector in feature_list:
            if statsmodels.version >= "0.12":
                rounded_truth["trend_strength"] = 0.93833
                rounded_truth["seasonality_strength"] = 0.329934
                rounded_truth["spikiness"] = 111.697325
                feature_vector["holt_alpha"] = np.round(feature_vector["holt_alpha"], 1)
                feature_vector["holt_beta"] = np.round(feature_vector["holt_beta"], 1)
                rounded_truth["holt_alpha"] = 1.0
                rounded_truth["holt_beta"] = 0.0
                rounded_truth["hw_alpha"] = 1.0
                rounded_truth["hw_beta"] = 0.0
                rounded_truth["hw_gamma"] = 0.0
            self.assertDictAlmostEqual(rounded_truth, feature_vector, places=6)

        # test selecting features
        features = [
            "var",
            "linearity",
            "spikiness",
            "trough",
            "holt_alpha",
            "hw_gamma",
            "level_shift_idx",
        ]

        features = TsFeatures(selected_features=features).transform(ts)

        if ts.is_univariate():
            feature_list = [_univariate_features(features)]
        else:
            feature_list = _multivariate_features(features)

        # test feature vector value
        rounded_truth = {
            "var": 14291.973331,
            "linearity": 0.853638,
            "spikiness": 33.502886,
            "trough": 3,
            "level_shift_idx": 118,
            "holt_alpha": 1.0,
            "hw_gamma": 0.157901,
        }

        for feature_vector in feature_list:
            if statsmodels.version >= "0.12":
                rounded_truth["spikiness"] = 111.697325
                feature_vector["holt_alpha"] = np.round(feature_vector["holt_alpha"], 1)
                rounded_truth["holt_alpha"] = 1.0
                rounded_truth["hw_gamma"] = 0.0

            self.assertDictAlmostEqual(rounded_truth, feature_vector, places=6)

        # test selecting extension features
        extension_features = [
            "cusum_num",
            "cusum_conf",
            "cusum_cp_index",
            "cusum_delta",
            "cusum_llr",
            "cusum_regression_detected",
            "cusum_stable_changepoint",
            "cusum_p_value",
            "robust_num",
            "robust_metric_mean",
            "outlier_num",
            "trend_num",
            "trend_num_increasing",
            "trend_avg_abs_tau",
            "nowcast_roc",
            "nowcast_ma",
            "nowcast_mom",
            "nowcast_lag",
            "nowcast_macd",
            "nowcast_macdsign",
            "nowcast_macddiff",
            "seasonal_period",
            "seasonality_mag",
            "trend_mag",
            "residual_std",
        ]

        features = TsFeatures(selected_features=extension_features).transform(ts)

        if ts.is_univariate():
            feature_list = [_univariate_features(features)]
        else:
            feature_list = _multivariate_features(features)

        # test feature vector value
        rounded_truth = {
            "cusum_num": 1,
            "cusum_conf": 1.0,
            "cusum_cp_index": 0.527778,
            "cusum_delta": 199.098856,
            "cusum_llr": 168.663483,
            "cusum_regression_detected": 1,
            "cusum_stable_changepoint": 1,
            "cusum_p_value": 0.0,
            "robust_num": 3,
            "robust_metric_mean": -31.866667,
            "outlier_num": 0,
            "trend_num": 2,
            "trend_num_increasing": 2,
            "trend_avg_abs_tau": 0.821053,
            "nowcast_roc": 0.062858,
            "nowcast_ma": 280.417143,
            "nowcast_mom": 12.841727,
            "nowcast_lag": 273.136691,
            "nowcast_macd": 11.032608,
            "nowcast_macdsign": 10.985509,
            "nowcast_macddiff": 0.527714,
            "seasonal_period": 7,
            "trend_mag": 2.404464,
            "seasonality_mag": 35.0,
            "residual_std": 21.258429,
        }
        for feature_vector in feature_list:
            if statsmodels.version >= "0.12":
                rounded_truth["trend_mag"] = 2.318814
                rounded_truth["seasonality_mag"] = 36.0
                rounded_truth["residual_std"] = 29.630087
            self.assertDictAlmostEqual(rounded_truth, feature_vector, places=6)

    def test_others(self) -> None:
        # test there is nan in feature vector because the length of TS is too short
        feature_vector = _univariate_features(TsFeatures().transform(self.TSData_short))

        self.assertEqual(
            np.isnan(np.asarray(list(feature_vector.values()))).any(),
            True,
        )

    def test_errors(self) -> None:
        # test input error (time series is too short)
        self.assertRaises(
            ValueError,
            TsFeatures().transform,
            self.TSData_mini,
        )
        with self.assertRaises(ValueError):
            TsFeatures(selected_features=["mango"])

    def test_acfpacf_warning(self) -> None:
        value = np.arange(30)
        time_ = pd.date_range(end="2022-08-01", freq="1D", periods=30)
        tsd = TimeSeriesData(pd.DataFrame({"value": value, "time": time_}))
        features = _univariate_features(TsFeatures(selected_features=["acfpacf_features"]).transform(tsd))
        acfpacf_features = [features[f] for f in _FEATURE_GROUP_MAPPING["acfpacf_features"]]
        self.assertEqual(acfpacf_features, [np.nan] * len(acfpacf_features))

    def test_IntegerArrays(self) -> None:
        if statsmodels.version < "0.12":
            df = pd.DataFrame(
                {
                    "time": range(15),
                    "value": [1, 4, 9, 4, 5, 5, 7, 2, 5, 1, 6, 3, 6, 5, 5],
                }
            )
        elif statsmodels.version >= "0.12":
            df = pd.DataFrame(
                {
                    "time": range(20),
                    "value": [
                        1,
                        4,
                        9,
                        4,
                        5,
                        5,
                        7,
                        2,
                        5,
                        1,
                        6,
                        3,
                        6,
                        5,
                        5,
                        6,
                        9,
                        10,
                        5,
                        6,
                    ],
                }
            )
        df["value"] = df["value"].astype(dtype=pd.Int64Dtype())
        # pyre-fixme[61]: `df` may not be initialized here.
        ts = TimeSeriesData(df)

        ts_features = TsFeatures(
            selected_features=[
                "length",
                "mean",
                "entropy",
                "hurst",
                "y_acf1",
                "seas_acf1",
                "hw_gamma",
            ]
        )
        feats = _univariate_features(ts_features.transform(ts))
        feats = {key: round(feats[key], 3) for key in feats}
        if statsmodels.version < "0.12":
            self.assertEqual(
                {
                    "length": 15,
                    "mean": 4.533,
                    "entropy": 0.765,
                    "hurst": -0.143,
                    "y_acf1": -0.298,
                    "seas_acf1": -0.121,
                    "hw_gamma": 0.947,
                },
                feats,
            )
        elif statsmodels.version >= "0.12":
            self.assertEqual(
                {
                    "length": 20,
                    "mean": 5.2,
                    "entropy": 0.894,
                    "hurst": -0.12,
                    "y_acf1": 0.041,
                    "seas_acf1": -0.125,
                    "hw_gamma": 0.0,
                },
                feats,
            )

    def test_tsfeatures_time(self) -> None:
        dates = pd.date_range("2021-01-22", "2021-01-31", tz="US/Pacific").tolist()
        dates += pd.date_range("2022-01-30", "2022-02-04", tz="US/Pacific").tolist()
        ts = TimeSeriesData(
            df=pd.DataFrame(
                {
                    "time": dates,
                    "value": [10, 2, 11, 5, 13, 18, 4, 14, 17, 9] + [np.nan] * 6,
                }
            )
        )
        features = _univariate_features(TsFeatures(selected_features=["time"]).transform(ts))
        expected = {
            "time_years": 2,
            "time_months": 3,
            "time_monthsofyear": 2,
            "time_weeks": 4,
            "time_weeksofyear": 3,
            "time_days": 16,
            "time_daysofyear": 14,
            "time_avg_timezone_offset": -28800.0,
            "time_length_days": 378,
            "time_freq_Monday": 0.125,
            "time_freq_Tuesday": 0.125,
            "time_freq_Wednesday": 0.125,
            "time_freq_Thursday": 0.125,
            "time_freq_Friday": 0.1875,
            "time_freq_Saturday": 0.125,
            "time_freq_Sunday": 0.1875,
        }
        self.assertDictAlmostEqual(expected, features)
