# 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.

import logging
import shutil
from functools import partial
from pathlib import Path
from typing import Any, Dict, Optional

from gluonts_customized.dataset import DatasetWriter
from gluonts_customized.dataset.artificial import ConstantDataset
from gluonts_customized.dataset.common import TrainDatasets, load_datasets
from gluonts_customized.dataset.jsonl import JsonLinesWriter

from ._artificial import generate_artificial_dataset
from ._airpassengers import generate_airpassengers_dataset
from ._ercot import generate_ercot_dataset
from ._ett_small import generate_ett_small_dataset
from ._gp_copula_2019 import generate_gp_copula_dataset
from ._lstnet import generate_lstnet_dataset
from ._m3 import generate_m3_dataset
from ._m4 import generate_m4_dataset
from ._m5 import generate_m5_dataset
from ._tsf_datasets import generate_forecasting_dataset
from ._uber_tlc import generate_uber_dataset


def get_download_path() -> Path:
    """

    Returns
    -------
    Path
        default path to download datasets or models of gluon-ts.
        The path is $HOME/.gluonts/
    """
    return Path.home() / ".gluonts"


dataset_recipes = {
    # each recipe generates a dataset given a path
    "constant": partial(
        generate_artificial_dataset, dataset=ConstantDataset()
    ),
    "exchange_rate": partial(
        generate_lstnet_dataset, dataset_name="exchange_rate"
    ),
    "solar-energy": partial(
        generate_lstnet_dataset, dataset_name="solar-energy"
    ),
    "electricity": partial(
        generate_lstnet_dataset, dataset_name="electricity"
    ),
    "traffic": partial(generate_lstnet_dataset, dataset_name="traffic"),
    "exchange_rate_nips": partial(
        generate_gp_copula_dataset, dataset_name="exchange_rate_nips"
    ),
    "electricity_nips": partial(
        generate_gp_copula_dataset, dataset_name="electricity_nips"
    ),
    "traffic_nips": partial(
        generate_gp_copula_dataset, dataset_name="traffic_nips"
    ),
    "solar_nips": partial(
        generate_gp_copula_dataset, dataset_name="solar_nips"
    ),
    "wiki2000_nips": partial(
        generate_gp_copula_dataset, dataset_name="wiki2000_nips"
    ),
    "wiki-rolling_nips": partial(
        generate_gp_copula_dataset, dataset_name="wiki-rolling_nips"
    ),
    "taxi_30min": partial(
        generate_gp_copula_dataset, dataset_name="taxi_30min"
    ),
    "kaggle_web_traffic_with_missing": partial(
        generate_forecasting_dataset,
        dataset_name="kaggle_web_traffic_with_missing",
    ),
    "kaggle_web_traffic_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="kaggle_web_traffic_without_missing",
    ),
    "kaggle_web_traffic_weekly": partial(
        generate_forecasting_dataset,
        dataset_name="kaggle_web_traffic_weekly",
    ),
    "m1_yearly": partial(
        generate_forecasting_dataset, dataset_name="m1_yearly"
    ),
    "m1_quarterly": partial(
        generate_forecasting_dataset, dataset_name="m1_quarterly"
    ),
    "m1_monthly": partial(
        generate_forecasting_dataset, dataset_name="m1_monthly"
    ),
    "nn5_daily_with_missing": partial(
        generate_forecasting_dataset, dataset_name="nn5_daily_with_missing"
    ),
    "nn5_daily_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="nn5_daily_without_missing",
    ),
    "nn5_weekly": partial(
        generate_forecasting_dataset, dataset_name="nn5_weekly"
    ),
    "tourism_monthly": partial(
        generate_forecasting_dataset, dataset_name="tourism_monthly"
    ),
    "tourism_quarterly": partial(
        generate_forecasting_dataset, dataset_name="tourism_quarterly"
    ),
    "tourism_yearly": partial(
        generate_forecasting_dataset, dataset_name="tourism_yearly"
    ),
    "cif_2016": partial(generate_forecasting_dataset, dataset_name="cif_2016"),
    "london_smart_meters_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="london_smart_meters_without_missing",
    ),
    "wind_farms_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="wind_farms_without_missing",
    ),
    "car_parts_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="car_parts_without_missing",
    ),
    "dominick": partial(generate_forecasting_dataset, dataset_name="dominick"),
    "fred_md": partial(generate_forecasting_dataset, dataset_name="fred_md"),
    "pedestrian_counts": partial(
        generate_forecasting_dataset, dataset_name="pedestrian_counts"
    ),
    "hospital": partial(generate_forecasting_dataset, dataset_name="hospital"),
    "covid_deaths": partial(
        generate_forecasting_dataset, dataset_name="covid_deaths"
    ),
    "kdd_cup_2018_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="kdd_cup_2018_without_missing",
    ),
    "weather": partial(generate_forecasting_dataset, dataset_name="weather"),
    "m3_monthly": partial(generate_m3_dataset, m3_freq="monthly"),
    "m3_quarterly": partial(generate_m3_dataset, m3_freq="quarterly"),
    "m3_yearly": partial(generate_m3_dataset, m3_freq="yearly"),
    "m3_other": partial(generate_m3_dataset, m3_freq="other"),
    "m4_hourly": partial(
        generate_m4_dataset,
        m4_freq="Hourly",
        pandas_freq="H",
        prediction_length=48,
    ),
    "m4_daily": partial(
        generate_m4_dataset,
        m4_freq="Daily",
        pandas_freq="D",
        prediction_length=14,
    ),
    "m4_weekly": partial(
        generate_m4_dataset,
        m4_freq="Weekly",
        pandas_freq="W",
        prediction_length=13,
    ),
    "m4_monthly": partial(
        generate_m4_dataset,
        m4_freq="Monthly",
        pandas_freq="M",
        prediction_length=18,
    ),
    "m4_quarterly": partial(
        generate_m4_dataset,
        m4_freq="Quarterly",
        pandas_freq="Q",
        prediction_length=8,
    ),
    "m4_yearly": partial(
        generate_m4_dataset,
        m4_freq="Yearly",
        pandas_freq="Y",
        prediction_length=6,
    ),
    "m5": partial(
        generate_m5_dataset,
        pandas_freq="D",
        prediction_length=28,
        m5_file_path=get_download_path() / "m5",
    ),
    "uber_tlc_daily": partial(
        generate_uber_dataset, uber_freq="Daily", prediction_length=7
    ),
    "uber_tlc_hourly": partial(
        generate_uber_dataset, uber_freq="Hourly", prediction_length=24
    ),
    "airpassengers": partial(generate_airpassengers_dataset),
    "australian_electricity_demand": partial(
        generate_forecasting_dataset,
        dataset_name="australian_electricity_demand",
    ),
    "electricity_hourly": partial(
        generate_forecasting_dataset,
        dataset_name="electricity_hourly",
    ),
    "electricity_weekly": partial(
        generate_forecasting_dataset,
        dataset_name="electricity_weekly",
    ),
    "rideshare_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="rideshare_without_missing",
    ),
    "saugeenday": partial(
        generate_forecasting_dataset,
        dataset_name="saugeenday",
    ),
    "solar_10_minutes": partial(
        generate_forecasting_dataset,
        dataset_name="solar_10_minutes",
    ),
    "solar_weekly": partial(
        generate_forecasting_dataset,
        dataset_name="solar_weekly",
    ),
    "sunspot_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="sunspot_without_missing",
    ),
    "temperature_rain_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="temperature_rain_without_missing",
    ),
    "vehicle_trips_without_missing": partial(
        generate_forecasting_dataset,
        dataset_name="vehicle_trips_without_missing",
    ),
    "ercot": partial(generate_ercot_dataset),
    "ett_small_15min": partial(
        generate_ett_small_dataset,
        base_file_name="ETTm",
        freq="15min",
        prediction_length=24,
    ),
    "ett_small_1h": partial(
        generate_ett_small_dataset,
        base_file_name="ETTh",
        freq="1h",
        prediction_length=24,
    ),
}


dataset_names = list(dataset_recipes.keys())

default_dataset_path = get_download_path() / "datasets"

default_dataset_writer = JsonLinesWriter()


def materialize_dataset(
    dataset_name: str,
    path: Path = default_dataset_path,
    regenerate: bool = False,
    dataset_writer: DatasetWriter = default_dataset_writer,
    prediction_length: Optional[int] = None,
) -> Path:
    """
    Ensures that the dataset is materialized under the `path / dataset_name`
    path.

    Parameters
    ----------
    dataset_name
        Name of the dataset, for instance "m4_hourly".
    regenerate
        Whether to regenerate the dataset even if a local file is present.
        If this flag is False and the file is present, the dataset will not
        be downloaded again.
    path
        Where the dataset should be saved.
    prediction_length
        The prediction length to be used for the dataset. If None, the default
        prediction length will be used. The prediction length might not be
        available for all datasets.

    Returns
    -------
        The path where the dataset is materialized
    """
    assert dataset_name in dataset_recipes.keys(), (
        f"{dataset_name} is not present, please choose one from "
        f"{sorted(dataset_recipes)}."
    )

    path.mkdir(parents=True, exist_ok=True)
    dataset_path = path / dataset_name

    dataset_recipe = dataset_recipes[dataset_name]

    if not dataset_path.exists() or regenerate:
        logging.info(f"downloading and processing {dataset_name}")
        if dataset_path.exists():
            # If regenerating, we need to remove the directory contents
            shutil.rmtree(dataset_path)
            dataset_path.mkdir()

        # Optionally pass prediction length to not override any non-None
        # defaults (e.g. for M4)
        kwargs: Dict[str, Any] = {"dataset_path": dataset_path}
        if prediction_length is not None:
            kwargs["prediction_length"] = prediction_length
        kwargs["dataset_writer"] = dataset_writer
        dataset_recipe(**kwargs)
    else:
        logging.info(
            f"using dataset already processed in path {dataset_path}."
        )

    return dataset_path


def get_dataset(
    dataset_name: str,
    path: Path = default_dataset_path,
    regenerate: bool = False,
    dataset_writer: DatasetWriter = default_dataset_writer,
    prediction_length: Optional[int] = None,
) -> TrainDatasets:
    """
    Get a repository dataset.

    The datasets that can be obtained through this function have been used
    with different processing over time by several papers (e.g., [SFG17]_,
    [LCY+18]_, and [YRD15]_) or are obtained through the `Monash Time Series
    Forecasting Repository <https://forecastingdata.org/>`_.

    Parameters
    ----------
    dataset_name
        Name of the dataset, for instance "m4_hourly".
    regenerate
        Whether to regenerate the dataset even if a local file is present.
        If this flag is False and the file is present, the dataset will not
        be downloaded again.
    path
        Where the dataset should be saved.
    prediction_length
        The prediction length to be used for the dataset. If None, the default
        prediction length will be used. If the dataset is already materialized,
        setting this option to a different value does not have an effect.
        Make sure to set `regenerate=True` in this case. Note that some
        datasets from the Monash Time Series Forecasting Repository do not
        actually have a default prediction length -- the default then depends
        on the frequency of the data:
        - Minutely data --> prediction length of 60 (one hour)
        - Hourly data --> prediction length of 48 (two days)
        - Daily data --> prediction length of 30 (one month)
        - Weekly data --> prediction length of 8 (two months)
        - Monthly data --> prediction length of 12 (one year)
        - Yearly data --> prediction length of 4 (four years)

    Returns
    -------
        Dataset obtained by either downloading or reloading from local file.
    """
    dataset_path = materialize_dataset(
        dataset_name,
        path,
        regenerate,
        dataset_writer,
        prediction_length,
    )

    return load_datasets(
        metadata=dataset_path,
        train=dataset_path / "train",
        test=dataset_path / "test",
    )


if __name__ == "__main__":
    for dataset in dataset_names:
        print(f"generate {dataset}")
        ds = get_dataset(dataset, regenerate=True)
        print(ds.metadata)
        print(len(ds.train))
