<kaggle_start><code># # Pattern Recognition in Financial Market
# ## Initialization
# In this section, we initialize the environments with the following steps:
# 1. check all dependencies has been installed
# 2. if not installed, install the dependency with pip
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
# --------------------------
# DO NOT MODIFY!!!
# --------------------------

# ### Common dependencies
# The following cell imports the dependencies required for the basic pipeline.
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
import os
import logging
import functools
import random
from datetime import datetime
from pprint import PrettyPrinter
from collections import OrderedDict
from typing import Tuple, Union, List, Iterable, Dict, Optional
import torch
import darts
import pandas as pd
import torch.nn as nn
import pytorch_lightning as pl
import matplotlib.pyplot as plt
from tqdm import tqdm
from darts.utils.missing_values import fill_missing_values, extract_subseries
from darts.utils.data import (
    PastCovariatesTrainingDataset,
    PastCovariatesInferenceDataset,
)
from darts.dataprocessing.transformers.scaler import Scaler

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] line %(lineno)s, %(funcName)s: %(message)s",
)
os.makedirs("./output", exist_ok=True)

# ### Feature-model-specific dependencies
# The following cell imports the dependencies required for the feature engineering and for the algorithms. For instance, Pytorch for the RNN-CNN-based model.
# Generated by assemble.py (imports)
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import talib
import pandas as pd

# #### General Parameters
# --------------------------
# MODIFY CAREFULLY BASED ON RESPECITIVE USAGE!
# --------------------------
# Preprocessing Parameters
# -------------------------------------------------------------------------------------------------------------
# - GAP_SIZE: Remove weekend, holiday gaps with GAP_SIZE.
#             Usually on weekends, data does not change, hence weekend time slot has no meaning
#             Since most weekend gaps is ~48, GAP_SIZE is default as 40.
GAP_SIZE = 40
# - KEEP_MARKETS: A set containing the expected markets for testing. It could be left empty for using all markets.
#                 Each item is a string which could be a part of the market name,
#                 e.g. `Australia` for Australia 200 Cash -2022-04-26.csv
KEEP_MARKETS = ["Australia"]
# -------------------------------------------------------------------------------------------------------------
# Feature Engineering Parameters
# -------------------------------------------------------------------------------------------------------------
# - USE_FEATURE: whether to calculate features from the original dataset, default `True`.
USE_FEATURE = True
# -------------------------------------------------------------------------------------------------------------
# Model Parameters
# -------------------------------------------------------------------------------------------------------------
# - WINDOW_SIZE, HORIZON_SIZE: Determine the window and horizon for prediction, use window to predict horizon,
#                              e.g. window(10:00-16:00) --predict-> horizon(16:00-18:00),
#                              with WINDOW_SIZE as 6 and HORIZON_SIZE as 2
WINDOW_SIZE, HORIZON_SIZE = 64, 16
# - PRED_COLUMNS: the column(s) to be predicted by the model, usually `bid_close` or `ask_close`.
PRED_COLUMNS = ["bid_close"]
# - EPOCHS: the number of times to iterate on the dataset, default as 10
EPOCH = 300
# - STRIDE: the spacing of sampling steps. For instance, stride=1 means that slices will be like
#           [(9:00-11:00), (10:00-12:00), ..., (22:00-24:00)]
#           It could be set to `1` for full sampling, `HORIZON_SIZE` for prediction-covering sampling, and
#           `WINDOW_SIZE+HORIZON_SIZE` for non-overlapping sampling
STRIDE = 1
# - USE_DIFF: whether to predict the difference instead of predicting the price itself, default `False`
USE_DIFF = False
# - SLICE_NORM: whether to normalize a slice to (0,1) based on window,
#               the same scale is applied to the horizon as well, default `False`
SLICE_NORM = False
# - RANDOM_SPLIT: whether to split data randomly,
#                 if True, a slice(a pair of window & horizon) is randomly assigned to train/valid/test set;
#                 else, a slice is assigned to train/valid/test by its starting time;
RANDOM_SPLIT = False
# - FIX_RANDOM: whether to fix the random seed for fixed dataset split indices.
FIX_RANDOM = False
if FIX_RANDOM:
    random.seed(0)
# - GENERAL_MODEL_PARAMS: general parameters for initialising the model, only modify the value when using
GENERAL_MODEL_PARAMS = {
    "loss_fn": nn.MSELoss(),
    "optimizer_cls": torch.optim.SGD,
    "optimizer_kwargs": {"lr": 2e-3, "momentum": 0.5},
    "lr_scheduler_cls": torch.optim.lr_scheduler.ReduceLROnPlateau,
    "lr_scheduler_kwargs": {"factor": 0.5},
}
# -------------------------------------------------------------------------------------------------------------
# Pipeline Parameters
# -------------------------------------------------------------------------------------------------------------
# - MODE_LOAD_MODEL: whether to load model for prediction instead of training from scratch.
#                    if true, will not train model, the path of the model should be specified as well.
#                    if false, will train model
MODE_LOAD_MODEL = False
NAME_LOAD_MODEL = "./output/blockrnnmodel-0501-0601/epoch=74-step=9075.ckpt"

# ## Preprocessing
# Preprocessing in vital for the algorithms and features to work. The preprocessing step includes:
# 1. copy the data into designated folder and read it.
# 2. clean up missing values and outliers in the time-series data.
# ### Stage 1. Read data
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
DATA_SOURCE = "./data"
KAGGLE_DATASET_SOURCE = "../input/stock-indices-202103-202202"
if os.path.exists(DATA_SOURCE) and all(
    fname.endswith(".csv") for fname in os.listdir(DATA_SOURCE)
):
    logging.info(
        f'Data already loaded into designated folder: [{", ".join(os.listdir(DATA_SOURCE)[:5])} ...]'
    )
else:
    DATA_SOURCE = KAGGLE_DATASET_SOURCE
    logging.info(
        f'Data copied into designated folder: [{", ".join(os.listdir(DATA_SOURCE)[:5])} ...]'
    )

# #### Read Data
# Use pandas to read data
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
raw_data = OrderedDict()


def remove_gaps(ts: darts.TimeSeries) -> darts.TimeSeries:
    tss = extract_subseries(ts, min_gap_size=GAP_SIZE)
    return functools.reduce(lambda x, y: x.concatenate(y, ignore_time_axes=True), tss)


for fname in os.listdir(DATA_SOURCE):
    if fname.endswith(".csv"):
        market = fname[:-20].strip()
        if KEEP_MARKETS and not any(i in market for i in KEEP_MARKETS):
            continue
        try:
            ts = darts.TimeSeries.from_csv(
                os.path.join(DATA_SOURCE, fname), time_col="time", freq="H"
            )
            raw_data[market] = remove_gaps(ts).pd_dataframe()
        except:
            logging.warning(f"{market} with incompatible data, skipping")
    else:
        logging.warning(f"unsupported file exists in data folder: {fname}")
logging.info(f'Testing markets: [{", ".join(raw_data.keys())}]')


# ### Stage 2. Missing value / outlier processing / normalisation
# If missing values are used in the prediction, it would cause the calculation of the error to become NaN. The missing values would be filled by two methods:
# 1. For prices, it would be filled by forefill method which use most-recent data on the left side. The consideration is that in this part, the
# 2. For volumes, average of a window is used to fill-in.
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
preprocessing_funcs = []


# Generated by assemble.py (preprocessing)
def normalize(df: pd.DataFrame) -> pd.DataFrame:
    scaler = MinMaxScaler(feature_range=(0, 1))
    scale_cols = list(filter(lambda x: "volume" not in x, df.columns))
    df[scale_cols] = scaler.fit_transform(df[scale_cols])
    return df


def volume_fillna(df: pd.DataFrame) -> pd.DataFrame:
    fill_cols = list(filter(lambda x: "volume" not in x, df.columns))
    df[fill_cols] = df[fill_cols].fillna(0.0)
    return df


preprocessing_funcs = [normalize]
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
for market, df in raw_data.items():
    for func in preprocessing_funcs:
        raw_data[market] = func(df)


# ## Feature engineering
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
feature_funcs = []


# Generated by assemble.py (features)
def get_talib_overlap(df: pd.DataFrame) -> pd.DataFrame:
    high = (df["bid_high"] + df["ask_high"]) / 2
    low = (df["bid_low"] + df["ask_low"]) / 2
    close = (df["bid_close"] + df["ask_close"]) / 2
    open_ = (df["bid_open"] + df["ask_open"]) / 2
    volume = df["last_volume"]
    (
        df["BBANDS.up.5.2.2"],
        df["BBANDS.mid.5.2.2"],
        df["BBANDS.low.5.2.2"],
    ) = talib.BBANDS(close, timeperiod=5, nbdevup=2, nbdevdn=2, matype=0)
    df["DMA.30"] = talib.DEMA(close, timeperiod=30)
    df["EMA.30"] = talib.EMA(close, timeperiod=30)
    df["HT_TRENDLINE"] = talib.HT_TRENDLINE(close)
    df["KAMA.30"] = talib.KAMA(close, timeperiod=30)
    df["MA.30"] = talib.MA(close, timeperiod=30, matype=0)
    df["MAMA.ma"], df["MAMA.fa"] = talib.MAMA(close)
    df["MIDPOINT.14"] = talib.MIDPOINT(close, timeperiod=14)
    df["MIDPRICE.14"] = talib.MIDPRICE(high, low, timeperiod=14)
    df["SAR"] = talib.SAR(high, low, acceleration=0, maximum=0)
    df["SAREXT"] = talib.SAREXT(
        high,
        low,
        startvalue=0,
        offsetonreverse=0,
        accelerationinitlong=0,
        accelerationlong=0,
        accelerationmaxlong=0,
        accelerationinitshort=0,
        accelerationshort=0,
        accelerationmaxshort=0,
    )
    df["SMA.30"] = talib.SMA(close, timeperiod=30)
    df["T3.5"] = talib.T3(close, timeperiod=5, vfactor=0)
    df["TEMA.30"] = talib.TEMA(close, timeperiod=30)
    df["T3.5"] = talib.T3(close, timeperiod=5, vfactor=0)
    df["TRIMA.30"] = talib.TRIMA(close, timeperiod=30)
    df["WMA.30"] = talib.WMA(close, timeperiod=30)
    df["ADX.14"] = talib.ADX(high, low, close, timeperiod=14)
    df["ADXR.14"] = talib.ADXR(high, low, close, timeperiod=14)
    df["APO.12.26"] = talib.APO(close, fastperiod=12, slowperiod=26, matype=0)
    df["AROON.down.14"], df["AROON.up.14"] = talib.AROON(high, low, timeperiod=14)
    df["AROONOSC.14"] = talib.AROONOSC(high, low, timeperiod=14)
    df["BOP"] = talib.BOP(open_, high, low, close)
    df["CCI.14"] = talib.CCI(high, low, close, timeperiod=14)
    df["CMO.14"] = talib.CMO(close, timeperiod=14)
    df["DX.14"] = talib.DX(high, low, close, timeperiod=14)
    df["MACD.12.26"], df["MACD.12.26.sig"], df["MACD.12.26.hist"] = talib.MACD(
        close, fastperiod=12, slowperiod=26, signalperiod=9
    )
    (
        df["MACDEXT.12.26"],
        df["MACDEXT.12.26.sig"],
        df["MACDEXT.12.26.hist"],
    ) = talib.MACDEXT(
        close,
        fastperiod=12,
        fastmatype=0,
        slowperiod=26,
        slowmatype=0,
        signalperiod=9,
        signalmatype=0,
    )
    (
        df["MACDFIX.12.26"],
        df["MACDFIX.12.26.sig"],
        df["MACDFIX.12.26.hist"],
    ) = talib.MACDFIX(close, signalperiod=9)
    df["MFI.14"] = talib.MFI(high, low, close, volume, timeperiod=14)
    df["MINUS_DI.14"] = talib.MINUS_DI(high, low, close, timeperiod=14)
    df["MINUS_DM.14"] = talib.MINUS_DM(high, low, timeperiod=14)
    df["MOM.10"] = talib.MOM(close, timeperiod=10)
    df["PLUS_DI.14"] = talib.PLUS_DI(high, low, close, timeperiod=14)
    df["PLUS_DM.14"] = talib.PLUS_DM(high, low, timeperiod=14)
    df["PPO.12.26"] = talib.PPO(close, fastperiod=12, slowperiod=26, matype=0)
    df["ROC.10"] = talib.ROC(close, timeperiod=10)
    df["ROCP.10"] = talib.ROCP(close, timeperiod=10)
    df["ROCR.10"] = talib.ROCR(close, timeperiod=10)
    df["ROCR100.10"] = talib.ROCR100(close, timeperiod=10)
    df["RSI.14"] = talib.RSI(close, timeperiod=14)
    df["STOCH.5.3.k"], df["STOCH.5.3.d"] = talib.STOCH(
        high,
        low,
        close,
        fastk_period=5,
        slowk_period=3,
        slowk_matype=0,
        slowd_period=3,
        slowd_matype=0,
    )
    df["STOCHF.5.3.k"], df["STOCHF.5.3.d"] = talib.STOCHF(
        high, low, close, fastk_period=5, fastd_period=3, fastd_matype=0
    )
    df["STOCHRSI.5.3.k"], df["STOCHRSI.5.3.d"] = talib.STOCHRSI(
        close, timeperiod=14, fastk_period=5, fastd_period=3, fastd_matype=0
    )
    df["TRIX.30"] = talib.TRIX(close, timeperiod=30)
    df["ULTOSC.7.14.28"] = talib.ULTOSC(
        high, low, close, timeperiod1=7, timeperiod2=14, timeperiod3=28
    )
    df["WILLR.14"] = talib.WILLR(high, low, close, timeperiod=14)
    return df


feature_funcs = [get_talib_overlap]
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
if USE_FEATURE:
    for market, df in raw_data.items():
        for func in feature_funcs:
            raw_data[market] = func(df)


# ## Prediction Models
# #### Convert data to time series
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
time_series = OrderedDict()
for market, df in raw_data.items():
    ts = darts.timeseries.TimeSeries.from_dataframe(df, freq="1h")
    if USE_DIFF:
        ts = fill_missing_values(ts).diff(dropna=False)
    time_series[market] = fill_missing_values(ts).astype(np.float32)
sample = time_series.values().__iter__().__next__()


# ### Creating train set, validation set and test set
# Since the dataset is rolling window based and behaves based on the model parameters, it should be split here.
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
class MarketDataset:
    def __init__(
        self,
        data: Dict[str, darts.TimeSeries],
        target_cols: Union[str, List[str]],
        random_: bool = False,
        window_size: int = WINDOW_SIZE,
        horizon_size: int = HORIZON_SIZE,
        stride: int = 1,
        test_ratio: float = 0.25,
        valid_ratio: float = 0.1,
    ) -> None:
        super().__init__()
        self.scaler = Scaler(MinMaxScaler())
        self.rand = random_
        self.window, self.horizon = window_size, horizon_size
        self.data = data
        self._init_ts_cov(target_cols)
        self._build_indices(data, stride)
        self._train_test_split(test_ratio, valid_ratio)
        self._gen_idx_map()

    def _init_ts_cov(self, target: Union[str, List[str]]) -> None:
        sample = self.data.values().__iter__().__next__()
        cols = set(sample.columns)
        if isinstance(target, str) and target in set(ts.columns):
            cols.remove(target)
            target = [target]
        elif isinstance(target, Iterable):
            cols = cols.difference(target)
        else:
            raise ValueError(f"target time series has no column: {target}")
        self.pred, self.covs = target, list(cols)

    def _get_series_and_covariates(
        self, ts: darts.TimeSeries
    ) -> Tuple[darts.TimeSeries, darts.TimeSeries]:
        return ts[self.pred], ts[self.covs]

    def _build_indices(self, data: Dict[str, darts.TimeSeries], stride: int) -> None:
        rolling = self.window + self.horizon
        indices = {}
        for market, ts in data.items():
            length = len(ts)
            if length < rolling:
                logging.error(
                    f"{market} does not have the length for the dataset, \
                                length is {length}, but should be larger than {rolling}"
                )
            elif stride > length - rolling:
                logging.warning(
                    f"For {market}, stride({stride}) is bigger than a rolling window, will instead use {rolling}"
                )
                stride = rolling
            indices[market] = set(range(0, length - rolling, stride))
        self.indices = indices

    def _train_test_split(
        self, test_ratio: float = 0.25, valid_ratio: float = 0.1
    ) -> None:
        train, test, valid = {}, {}, {}
        # split dataset by market
        for market, index in self.indices.items():
            if self.rand:
                _test = set(random.sample(index, int(len(index) * test_ratio)))
                train_val = index.difference(test)
                _valid = set(random.sample(train_val, int(len(index) * valid_ratio)))
                _train = train_val.difference(_valid)
            else:  # traditional split method by entire segment slicing
                idx_len, idx_sorted = len(index), sorted(index)
                train_valid_boundary, valid_test_boundary = (
                    int(idx_len * (1 - valid_ratio - test_ratio)),
                    int(idx_len * (1 - valid_ratio)),
                )
                _train, _valid, _test = (
                    idx_sorted[:train_valid_boundary],
                    idx_sorted[train_valid_boundary:valid_test_boundary],
                    idx_sorted[valid_test_boundary:],
                )
            train[market], test[market], valid[market] = _train, _test, _valid
        self.train, self.test, self.valid = train, test, valid

    def _gen_idx_map(self):
        def get_map(idx_dict: Dict[str, set]) -> Dict[int, Tuple[str, int]]:
            count, res = 0, {}
            for market, indices in idx_dict.items():
                for idx in indices:
                    res[count] = (market, idx)
                    count += 1
            return res

        self._idx_map = {}
        self._idx_map["train"] = get_map(self.train)
        self._idx_map["test"] = get_map(self.test)
        self._idx_map["valid"] = get_map(self.valid)

    def inspect_dataset(self):
        return (
            self.indices,
            (self.pred, self.covs),
            self._idx_map,
            {"train": self.train, "test": self.test, "valid": self.valid},
        )

    @functools.lru_cache()
    def get_len(self, type_: str) -> int:
        return len(self._idx_map[type_])

    @functools.lru_cache(maxsize=None)
    def get_sample(
        self, idx: int, type_: str
    ) -> Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]:
        market, start = self._idx_map[type_][idx]
        ts = self.data[market]
        window, horizon = start + self.window, start + self.window + self.horizon
        slice = ts[start:window]
        past, covs, truth = (
            slice[self.pred],
            slice[self.covs],
            ts[window:horizon][self.pred],
        )
        if SLICE_NORM:
            past = self.scaler.fit_transform(past)
            truth = self.scaler.transform(truth)
        return past, covs, truth


class PartialDataset(PastCovariatesTrainingDataset):
    TRAIN, TEST, VALID = TYPES = ("train", "test", "valid")

    def __init__(self, real_dataset: MarketDataset, type_: str) -> None:
        super().__init__()
        assert (
            type_ in self.TYPES
        ), f"Type must be either of `train, test, valid` but {type_} was given"
        self.type_ = type_
        self.real_dataset = real_dataset

    def __len__(self):
        return self.real_dataset.get_len(self.type_)

    def __getitem__(self, idx: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        return tuple(i.values() for i in self.real_dataset.get_sample(idx, self.type_))


class PartialInferenceDataset(PastCovariatesInferenceDataset):
    def __init__(self, real_dataset: MarketDataset) -> None:
        self.type_ = "test"
        self.real_dataset = real_dataset

    def __len__(self):
        return self.real_dataset.get_len(self.type_)

    def __getitem__(
        self, idx: int
    ) -> Tuple[
        np.ndarray, Optional[np.ndarray], Optional[np.ndarray], darts.TimeSeries
    ]:
        past, covs, target = self.real_dataset.get_sample(idx, self.type_)
        return past.values(), covs.values(), None, target

    def get_ts_by_idx(
        self, idx: int
    ) -> Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]:
        return self.real_dataset.get_sample(idx, self.type_)


# Initialize the dataset
dataset = MarketDataset(
    time_series, target_cols=PRED_COLUMNS, stride=STRIDE, random_=RANDOM_SPLIT
)
train = PartialDataset(dataset, PartialDataset.TRAIN)
valid = PartialDataset(dataset, PartialDataset.VALID)
test = PartialInferenceDataset(dataset)


# ### Model instantiating
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
from darts.models import NaiveDrift
from darts.models.forecasting.forecasting_model import ForecastingModel

model = NaiveDrift()
model.__not_initialized__ = True
# Generated by assemble.py (models)
from typing import List, Optional
from darts.models import BlockRNNModel
from darts.models.forecasting.pl_forecasting_module import PLPastCovariatesModule
import torch
import torch.nn as nn


class _ResidualBlock(nn.Module):
    def __init__(self, input_dim: int, output_dim: int = None) -> None:
        super(_ResidualBlock, self).__init__()
        output_dim = output_dim or input_dim
        self.cnn = nn.Conv1d(input_dim, output_dim, kernel_size=3, padding="same")
        self.relu = nn.ReLU()
        self.bn = nn.BatchNorm1d(output_dim)

    def forward(self, x: torch.Tensor):
        """
        args:
            - x: Tensor(batch_size, signal_length, channels)
        """
        residual = x
        x = x.permute(0, 2, 1)
        conv = self.cnn(x)
        bn = self.bn(conv)
        act = self.relu(bn)
        act = act.permute(0, 2, 1)
        return self.relu(act + residual)

    @classmethod
    def build_residual(
        cls, input_dim: int, output_dim: int = None, block_nums: int = 4
    ) -> nn.Sequential:
        blocks = [nn.Linear(input_dim, output_dim)]
        for _ in range(block_nums - 1):
            blocks.append(_ResidualBlock(output_dim, output_dim))
        return nn.Sequential(*blocks)


class RNNCNNModel(PLPastCovariatesModule):
    def __init__(
        self,
        name: str,
        input_size: int,
        hidden_dim: int,
        num_layers: int,
        target_size: int,
        nr_params: int = None,
        num_layers_out_fc: Optional[List] = None,
        dropout: float = 0.0,
        **kwargs,
    ) -> None:
        super(RNNCNNModel, self).__init__(**kwargs)
        self.name = name or "RNN-CNN"
        self.target_size = target_size
        self.lc_hidden_dim = 1024
        self.residual = _ResidualBlock.build_residual(input_size, hidden_dim)
        self.enc = nn.LSTM(input_size, hidden_dim, num_layers, dropout=dropout)
        self.relu = nn.ReLU()
        self.dec = nn.LSTM(hidden_dim * 2, hidden_dim, num_layers, dropout=dropout)
        self.lc1 = nn.Linear(hidden_dim, self.lc_hidden_dim)
        self.lc2 = nn.Linear(self.lc_hidden_dim, target_size, bias=False)
        self.bias_lc = nn.Linear(input_size, target_size)
        self.bias_acti = nn.Hardsigmoid()
        self.bias_pool = nn.AdaptiveMaxPool1d(1)

    def forward(self, x: torch.Tensor):
        """
        args:
            - x: Tensor(batch_size, signal_length, channels)
        """
        # encoding
        conv = self.residual(x)
        recr, _ = self.enc(x)
        hidden = self.relu(torch.cat((conv, recr), dim=-1))
        # decoding
        out, _ = self.dec(hidden)
        relu = self.relu(out[:, : self.output_chunk_length, :])
        lc_hidden = self.lc1(relu)
        wb = self.relu(self.lc2(lc_hidden))
        bias = self.bias_pool(self.bias_acti(self.bias_lc(x)).permute(0, 2, 1))
        predictions = wb + bias
        batch_size, out_len = x.size(0), self.output_chunk_length
        return predictions.view(batch_size, out_len, self.target_size, 1)


_rnn_cnn = RNNCNNModel(
    "RNN-CNN",
    input_size=len(sample.columns),
    input_chunk_length=WINDOW_SIZE,
    output_chunk_length=HORIZON_SIZE,
    hidden_dim=128,
    num_layers=4,
    target_size=len(PRED_COLUMNS),
)
model = BlockRNNModel(
    input_chunk_length=WINDOW_SIZE,
    output_chunk_length=HORIZON_SIZE,
    model=_rnn_cnn,
    **GENERAL_MODEL_PARAMS,
)
from darts.models.forecasting.nbeats import _Block, _TrendGenerator, _Stack


class NBeatsRNNModule(PLPastCovariatesModule):
    def __init__(
        self,
        name: str,
        input_size: int,
        hidden_dim: int,
        num_layers: int,
        target_size: int,
        nr_params: int = None,
        num_layers_out_fc: Optional[List] = None,
        dropout: float = 0.0,
        **kwargs,
    ) -> None:
        pass

    def forward(self, x: torch.Tensor):
        """
        args:
            - x: Tensor(batch_size, signal_length, channels)
        """
        return


# ### Model training
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
model_name = (
    f'{model.__class__.__name__.lower()}-{datetime.utcnow().strftime("%m%d-%H%M")}'
)


def init_trainer() -> pl.Trainer:
    train_args = {}
    # Allow GPU is cuda available
    if torch.cuda.is_available():
        train_args["accelerator"] = "gpu"
        train_args["gpus"] = torch.cuda.device_count()
    # Add early stopping strategy and controlled checkpoints to the trainer
    es = pl.callbacks.early_stopping.EarlyStopping(
        monitor="val_loss", min_delta=0.001, patience=100
    )
    check_point = pl.callbacks.ModelCheckpoint(
        dirpath=f"./output/{model_name}", save_top_k=1, monitor="val_loss"
    )
    train_args["callbacks"] = [es, check_point]
    # Set max_epochs and precision
    train_args["max_epochs"] = EPOCH
    train_args["precision"] = 32
    # Create the trainer with kwargs
    trainer = pl.Trainer(**train_args)
    return trainer


trainer = init_trainer()


# #### Start training
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
model_path = f"./output/{model_name}"
if MODE_LOAD_MODEL:
    model.train_sample, model.output_dim = train[0], train[0][-1].shape[1]
    model._init_model(trainer)
    model.model = model.model.load_from_checkpoint(NAME_LOAD_MODEL)
    logging.info(f"Successfully load model from: {NAME_LOAD_MODEL}")
else:
    if isinstance(model, ForecastingModel) and not getattr(
        model, "__not_initialized__", False
    ):
        GENERAL_MODEL_PARAMS.update(
            {
                "work_dir": model_path,
            }
        )
        model = model.fit_from_dataset(
            train_dataset=train, val_dataset=valid, trainer=trainer, epochs=EPOCH
        )
    else:
        logging.error("Model not defined")

# #### Save the model as ONNX format for future use
if (
    isinstance(model, ForecastingModel)
    and not getattr(model, "__not_initialized__", False)
    and not MODE_LOAD_MODEL
):
    try:
        model.model.to_onnx(
            os.path.join(model_path, "model.onnx"),
            (torch.Tensor(np.concatenate(test[0][:2], axis=-1)).unsqueeze(0)),
            export_params=True,
            input_names=["input"],
            output_names=["output"],
            opset_version=11,
            dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        )
    except RuntimeError:
        logging.warning("Model not able to be saved as onnx format")

# ## Evaluation
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
evaluation_funcs = []
# Generated by assemble.py (evaluation)
evaluation_funcs = [darts.metrics.rmse, darts.metrics.r2_score]


# --------------------------
# Aggregation Functions
# Could be defined here
# --------------------------
def avg(x):
    return sum(x) / len(x)


agg_methods = [avg, np.std]


# ### Get Predictions
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
def get_prediction_pairs(
    model: ForecastingModel, test_set: PartialInferenceDataset
) -> List[Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]]:
    tester = pl.Trainer(
        False, enable_progress_bar=False, accelerator="cpu", max_epochs=1
    )
    res = []
    for i in tqdm(range(len(test_set))):
        ts_, cov, truth = test.get_ts_by_idx(i)
        pred = model.predict(HORIZON_SIZE, ts_, cov, trainer=tester)
        res.append((ts_, truth, pred))
    return res


predictions = get_prediction_pairs(model, test)
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
eval_res = {}
evaluation_funcs.append(darts.metrics.mae)
for metric in evaluation_funcs:
    func_name = metric.__name__
    losses = []
    eval_res[func_name] = {}
    for _, truth, pred in predictions:
        losses.append(metric(truth, pred))
    for agg in agg_methods:
        eval_res[func_name][agg.__name__] = agg(losses)


# --------------------------
# DO NOT MODIFY!!!
# --------------------------
def plot_predictions(
    predictions: List[Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]],
    sample_size: int = 10,
):
    if sample_size > len(predictions):
        sample_size = len(predictions)
    sample_size = sample_size or len(predictions)
    import time

    random.seed(time.time())
    predictions = random.sample(predictions, sample_size)
    for i, (past, truth, pred) in enumerate(predictions):
        past.plot(new_plot=True, label="truth")
        truth.plot(label="truth")
        pred.plot(label="pred")
        if not MODE_LOAD_MODEL:
            plt.savefig(os.path.join(model_path, f"pred-{i}.png"))
    if FIX_RANDOM:
        random.seed(0)


# ## Conclusion
# ### 1. Show the results of the experiment
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
pp = PrettyPrinter(indent=2, width=80)
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
pp.pprint(eval_res)

# ### 2. Plot Typical Prediction Results
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
plot_predictions(predictions, sample_size=20)

# ### 3. Save Hyperparameters to yaml
if not MODE_LOAD_MODEL:
    import yaml

    hyper_params = {
        "KEEP_MARKETS": KEEP_MARKETS,
        "USE_FEATURE": USE_FEATURE,
        "USE_DIFF": USE_DIFF,
        "WINDOW_SIZE": WINDOW_SIZE,
        "HORIZON_SIZE": HORIZON_SIZE,
        "PRED_COLUMNS": PRED_COLUMNS,
        "EPOCH": EPOCH,
        "STRIDE": STRIDE,
        "USE_DIFF": USE_DIFF,
        "SLICE_NORM": SLICE_NORM,
        "RANDOM_SPLIT": RANDOM_SPLIT,
        "FIX_RANDOM": FIX_RANDOM,
        "GENERAL_MODEL_PARAMS": GENERAL_MODEL_PARAMS,
    }
    hyper_params.update(model.model.hparams)
    with open(os.path.join(model_path, "hyper-params.yml"), "w") as f:
        yaml.dump(hyper_params, f)
    with open(os.path.join(model_path, "result.txt"), "w") as f:
        PrettyPrinter(indent=2, stream=f).pprint(eval_res)
<|endoftext|>
[{"Id": 94659766, "ScriptId": 26185548, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 1452696, "CreationDate": "05/03/2022 12:56:18", "VersionNumber": 11.0, "Title": "Experiment of Stock Price Prediction", "EvaluationDate": "05/03/2022", "IsChange": true, "TotalLines": 804.0, "LinesInsertedFromPrevious": 195.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 609.0, "LinesInsertedFromFork": NaN, "LinesDeletedFromFork": NaN, "LinesChangedFromFork": NaN, "LinesUnchangedFromFork": NaN, "TotalVotes": 0}]
# # Pattern Recognition in Financial Market
# ## Initialization
# In this section, we initialize the environments with the following steps:
# 1. check all dependencies has been installed
# 2. if not installed, install the dependency with pip
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
# --------------------------
# DO NOT MODIFY!!!
# --------------------------

# ### Common dependencies
# The following cell imports the dependencies required for the basic pipeline.
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
import os
import logging
import functools
import random
from datetime import datetime
from pprint import PrettyPrinter
from collections import OrderedDict
from typing import Tuple, Union, List, Iterable, Dict, Optional
import torch
import darts
import pandas as pd
import torch.nn as nn
import pytorch_lightning as pl
import matplotlib.pyplot as plt
from tqdm import tqdm
from darts.utils.missing_values import fill_missing_values, extract_subseries
from darts.utils.data import (
    PastCovariatesTrainingDataset,
    PastCovariatesInferenceDataset,
)
from darts.dataprocessing.transformers.scaler import Scaler

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] line %(lineno)s, %(funcName)s: %(message)s",
)
os.makedirs("./output", exist_ok=True)

# ### Feature-model-specific dependencies
# The following cell imports the dependencies required for the feature engineering and for the algorithms. For instance, Pytorch for the RNN-CNN-based model.
# Generated by assemble.py (imports)
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import talib
import pandas as pd

# #### General Parameters
# --------------------------
# MODIFY CAREFULLY BASED ON RESPECITIVE USAGE!
# --------------------------
# Preprocessing Parameters
# -------------------------------------------------------------------------------------------------------------
# - GAP_SIZE: Remove weekend, holiday gaps with GAP_SIZE.
#             Usually on weekends, data does not change, hence weekend time slot has no meaning
#             Since most weekend gaps is ~48, GAP_SIZE is default as 40.
GAP_SIZE = 40
# - KEEP_MARKETS: A set containing the expected markets for testing. It could be left empty for using all markets.
#                 Each item is a string which could be a part of the market name,
#                 e.g. `Australia` for Australia 200 Cash -2022-04-26.csv
KEEP_MARKETS = ["Australia"]
# -------------------------------------------------------------------------------------------------------------
# Feature Engineering Parameters
# -------------------------------------------------------------------------------------------------------------
# - USE_FEATURE: whether to calculate features from the original dataset, default `True`.
USE_FEATURE = True
# -------------------------------------------------------------------------------------------------------------
# Model Parameters
# -------------------------------------------------------------------------------------------------------------
# - WINDOW_SIZE, HORIZON_SIZE: Determine the window and horizon for prediction, use window to predict horizon,
#                              e.g. window(10:00-16:00) --predict-> horizon(16:00-18:00),
#                              with WINDOW_SIZE as 6 and HORIZON_SIZE as 2
WINDOW_SIZE, HORIZON_SIZE = 64, 16
# - PRED_COLUMNS: the column(s) to be predicted by the model, usually `bid_close` or `ask_close`.
PRED_COLUMNS = ["bid_close"]
# - EPOCHS: the number of times to iterate on the dataset, default as 10
EPOCH = 300
# - STRIDE: the spacing of sampling steps. For instance, stride=1 means that slices will be like
#           [(9:00-11:00), (10:00-12:00), ..., (22:00-24:00)]
#           It could be set to `1` for full sampling, `HORIZON_SIZE` for prediction-covering sampling, and
#           `WINDOW_SIZE+HORIZON_SIZE` for non-overlapping sampling
STRIDE = 1
# - USE_DIFF: whether to predict the difference instead of predicting the price itself, default `False`
USE_DIFF = False
# - SLICE_NORM: whether to normalize a slice to (0,1) based on window,
#               the same scale is applied to the horizon as well, default `False`
SLICE_NORM = False
# - RANDOM_SPLIT: whether to split data randomly,
#                 if True, a slice(a pair of window & horizon) is randomly assigned to train/valid/test set;
#                 else, a slice is assigned to train/valid/test by its starting time;
RANDOM_SPLIT = False
# - FIX_RANDOM: whether to fix the random seed for fixed dataset split indices.
FIX_RANDOM = False
if FIX_RANDOM:
    random.seed(0)
# - GENERAL_MODEL_PARAMS: general parameters for initialising the model, only modify the value when using
GENERAL_MODEL_PARAMS = {
    "loss_fn": nn.MSELoss(),
    "optimizer_cls": torch.optim.SGD,
    "optimizer_kwargs": {"lr": 2e-3, "momentum": 0.5},
    "lr_scheduler_cls": torch.optim.lr_scheduler.ReduceLROnPlateau,
    "lr_scheduler_kwargs": {"factor": 0.5},
}
# -------------------------------------------------------------------------------------------------------------
# Pipeline Parameters
# -------------------------------------------------------------------------------------------------------------
# - MODE_LOAD_MODEL: whether to load model for prediction instead of training from scratch.
#                    if true, will not train model, the path of the model should be specified as well.
#                    if false, will train model
MODE_LOAD_MODEL = False
NAME_LOAD_MODEL = "./output/blockrnnmodel-0501-0601/epoch=74-step=9075.ckpt"

# ## Preprocessing
# Preprocessing in vital for the algorithms and features to work. The preprocessing step includes:
# 1. copy the data into designated folder and read it.
# 2. clean up missing values and outliers in the time-series data.
# ### Stage 1. Read data
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
DATA_SOURCE = "./data"
KAGGLE_DATASET_SOURCE = "../input/stock-indices-202103-202202"
if os.path.exists(DATA_SOURCE) and all(
    fname.endswith(".csv") for fname in os.listdir(DATA_SOURCE)
):
    logging.info(
        f'Data already loaded into designated folder: [{", ".join(os.listdir(DATA_SOURCE)[:5])} ...]'
    )
else:
    DATA_SOURCE = KAGGLE_DATASET_SOURCE
    logging.info(
        f'Data copied into designated folder: [{", ".join(os.listdir(DATA_SOURCE)[:5])} ...]'
    )

# #### Read Data
# Use pandas to read data
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
raw_data = OrderedDict()


def remove_gaps(ts: darts.TimeSeries) -> darts.TimeSeries:
    tss = extract_subseries(ts, min_gap_size=GAP_SIZE)
    return functools.reduce(lambda x, y: x.concatenate(y, ignore_time_axes=True), tss)


for fname in os.listdir(DATA_SOURCE):
    if fname.endswith(".csv"):
        market = fname[:-20].strip()
        if KEEP_MARKETS and not any(i in market for i in KEEP_MARKETS):
            continue
        try:
            ts = darts.TimeSeries.from_csv(
                os.path.join(DATA_SOURCE, fname), time_col="time", freq="H"
            )
            raw_data[market] = remove_gaps(ts).pd_dataframe()
        except:
            logging.warning(f"{market} with incompatible data, skipping")
    else:
        logging.warning(f"unsupported file exists in data folder: {fname}")
logging.info(f'Testing markets: [{", ".join(raw_data.keys())}]')


# ### Stage 2. Missing value / outlier processing / normalisation
# If missing values are used in the prediction, it would cause the calculation of the error to become NaN. The missing values would be filled by two methods:
# 1. For prices, it would be filled by forefill method which use most-recent data on the left side. The consideration is that in this part, the
# 2. For volumes, average of a window is used to fill-in.
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
preprocessing_funcs = []


# Generated by assemble.py (preprocessing)
def normalize(df: pd.DataFrame) -> pd.DataFrame:
    scaler = MinMaxScaler(feature_range=(0, 1))
    scale_cols = list(filter(lambda x: "volume" not in x, df.columns))
    df[scale_cols] = scaler.fit_transform(df[scale_cols])
    return df


def volume_fillna(df: pd.DataFrame) -> pd.DataFrame:
    fill_cols = list(filter(lambda x: "volume" not in x, df.columns))
    df[fill_cols] = df[fill_cols].fillna(0.0)
    return df


preprocessing_funcs = [normalize]
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
for market, df in raw_data.items():
    for func in preprocessing_funcs:
        raw_data[market] = func(df)


# ## Feature engineering
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
feature_funcs = []


# Generated by assemble.py (features)
def get_talib_overlap(df: pd.DataFrame) -> pd.DataFrame:
    high = (df["bid_high"] + df["ask_high"]) / 2
    low = (df["bid_low"] + df["ask_low"]) / 2
    close = (df["bid_close"] + df["ask_close"]) / 2
    open_ = (df["bid_open"] + df["ask_open"]) / 2
    volume = df["last_volume"]
    (
        df["BBANDS.up.5.2.2"],
        df["BBANDS.mid.5.2.2"],
        df["BBANDS.low.5.2.2"],
    ) = talib.BBANDS(close, timeperiod=5, nbdevup=2, nbdevdn=2, matype=0)
    df["DMA.30"] = talib.DEMA(close, timeperiod=30)
    df["EMA.30"] = talib.EMA(close, timeperiod=30)
    df["HT_TRENDLINE"] = talib.HT_TRENDLINE(close)
    df["KAMA.30"] = talib.KAMA(close, timeperiod=30)
    df["MA.30"] = talib.MA(close, timeperiod=30, matype=0)
    df["MAMA.ma"], df["MAMA.fa"] = talib.MAMA(close)
    df["MIDPOINT.14"] = talib.MIDPOINT(close, timeperiod=14)
    df["MIDPRICE.14"] = talib.MIDPRICE(high, low, timeperiod=14)
    df["SAR"] = talib.SAR(high, low, acceleration=0, maximum=0)
    df["SAREXT"] = talib.SAREXT(
        high,
        low,
        startvalue=0,
        offsetonreverse=0,
        accelerationinitlong=0,
        accelerationlong=0,
        accelerationmaxlong=0,
        accelerationinitshort=0,
        accelerationshort=0,
        accelerationmaxshort=0,
    )
    df["SMA.30"] = talib.SMA(close, timeperiod=30)
    df["T3.5"] = talib.T3(close, timeperiod=5, vfactor=0)
    df["TEMA.30"] = talib.TEMA(close, timeperiod=30)
    df["T3.5"] = talib.T3(close, timeperiod=5, vfactor=0)
    df["TRIMA.30"] = talib.TRIMA(close, timeperiod=30)
    df["WMA.30"] = talib.WMA(close, timeperiod=30)
    df["ADX.14"] = talib.ADX(high, low, close, timeperiod=14)
    df["ADXR.14"] = talib.ADXR(high, low, close, timeperiod=14)
    df["APO.12.26"] = talib.APO(close, fastperiod=12, slowperiod=26, matype=0)
    df["AROON.down.14"], df["AROON.up.14"] = talib.AROON(high, low, timeperiod=14)
    df["AROONOSC.14"] = talib.AROONOSC(high, low, timeperiod=14)
    df["BOP"] = talib.BOP(open_, high, low, close)
    df["CCI.14"] = talib.CCI(high, low, close, timeperiod=14)
    df["CMO.14"] = talib.CMO(close, timeperiod=14)
    df["DX.14"] = talib.DX(high, low, close, timeperiod=14)
    df["MACD.12.26"], df["MACD.12.26.sig"], df["MACD.12.26.hist"] = talib.MACD(
        close, fastperiod=12, slowperiod=26, signalperiod=9
    )
    (
        df["MACDEXT.12.26"],
        df["MACDEXT.12.26.sig"],
        df["MACDEXT.12.26.hist"],
    ) = talib.MACDEXT(
        close,
        fastperiod=12,
        fastmatype=0,
        slowperiod=26,
        slowmatype=0,
        signalperiod=9,
        signalmatype=0,
    )
    (
        df["MACDFIX.12.26"],
        df["MACDFIX.12.26.sig"],
        df["MACDFIX.12.26.hist"],
    ) = talib.MACDFIX(close, signalperiod=9)
    df["MFI.14"] = talib.MFI(high, low, close, volume, timeperiod=14)
    df["MINUS_DI.14"] = talib.MINUS_DI(high, low, close, timeperiod=14)
    df["MINUS_DM.14"] = talib.MINUS_DM(high, low, timeperiod=14)
    df["MOM.10"] = talib.MOM(close, timeperiod=10)
    df["PLUS_DI.14"] = talib.PLUS_DI(high, low, close, timeperiod=14)
    df["PLUS_DM.14"] = talib.PLUS_DM(high, low, timeperiod=14)
    df["PPO.12.26"] = talib.PPO(close, fastperiod=12, slowperiod=26, matype=0)
    df["ROC.10"] = talib.ROC(close, timeperiod=10)
    df["ROCP.10"] = talib.ROCP(close, timeperiod=10)
    df["ROCR.10"] = talib.ROCR(close, timeperiod=10)
    df["ROCR100.10"] = talib.ROCR100(close, timeperiod=10)
    df["RSI.14"] = talib.RSI(close, timeperiod=14)
    df["STOCH.5.3.k"], df["STOCH.5.3.d"] = talib.STOCH(
        high,
        low,
        close,
        fastk_period=5,
        slowk_period=3,
        slowk_matype=0,
        slowd_period=3,
        slowd_matype=0,
    )
    df["STOCHF.5.3.k"], df["STOCHF.5.3.d"] = talib.STOCHF(
        high, low, close, fastk_period=5, fastd_period=3, fastd_matype=0
    )
    df["STOCHRSI.5.3.k"], df["STOCHRSI.5.3.d"] = talib.STOCHRSI(
        close, timeperiod=14, fastk_period=5, fastd_period=3, fastd_matype=0
    )
    df["TRIX.30"] = talib.TRIX(close, timeperiod=30)
    df["ULTOSC.7.14.28"] = talib.ULTOSC(
        high, low, close, timeperiod1=7, timeperiod2=14, timeperiod3=28
    )
    df["WILLR.14"] = talib.WILLR(high, low, close, timeperiod=14)
    return df


feature_funcs = [get_talib_overlap]
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
if USE_FEATURE:
    for market, df in raw_data.items():
        for func in feature_funcs:
            raw_data[market] = func(df)


# ## Prediction Models
# #### Convert data to time series
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
time_series = OrderedDict()
for market, df in raw_data.items():
    ts = darts.timeseries.TimeSeries.from_dataframe(df, freq="1h")
    if USE_DIFF:
        ts = fill_missing_values(ts).diff(dropna=False)
    time_series[market] = fill_missing_values(ts).astype(np.float32)
sample = time_series.values().__iter__().__next__()


# ### Creating train set, validation set and test set
# Since the dataset is rolling window based and behaves based on the model parameters, it should be split here.
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
class MarketDataset:
    def __init__(
        self,
        data: Dict[str, darts.TimeSeries],
        target_cols: Union[str, List[str]],
        random_: bool = False,
        window_size: int = WINDOW_SIZE,
        horizon_size: int = HORIZON_SIZE,
        stride: int = 1,
        test_ratio: float = 0.25,
        valid_ratio: float = 0.1,
    ) -> None:
        super().__init__()
        self.scaler = Scaler(MinMaxScaler())
        self.rand = random_
        self.window, self.horizon = window_size, horizon_size
        self.data = data
        self._init_ts_cov(target_cols)
        self._build_indices(data, stride)
        self._train_test_split(test_ratio, valid_ratio)
        self._gen_idx_map()

    def _init_ts_cov(self, target: Union[str, List[str]]) -> None:
        sample = self.data.values().__iter__().__next__()
        cols = set(sample.columns)
        if isinstance(target, str) and target in set(ts.columns):
            cols.remove(target)
            target = [target]
        elif isinstance(target, Iterable):
            cols = cols.difference(target)
        else:
            raise ValueError(f"target time series has no column: {target}")
        self.pred, self.covs = target, list(cols)

    def _get_series_and_covariates(
        self, ts: darts.TimeSeries
    ) -> Tuple[darts.TimeSeries, darts.TimeSeries]:
        return ts[self.pred], ts[self.covs]

    def _build_indices(self, data: Dict[str, darts.TimeSeries], stride: int) -> None:
        rolling = self.window + self.horizon
        indices = {}
        for market, ts in data.items():
            length = len(ts)
            if length < rolling:
                logging.error(
                    f"{market} does not have the length for the dataset, \
                                length is {length}, but should be larger than {rolling}"
                )
            elif stride > length - rolling:
                logging.warning(
                    f"For {market}, stride({stride}) is bigger than a rolling window, will instead use {rolling}"
                )
                stride = rolling
            indices[market] = set(range(0, length - rolling, stride))
        self.indices = indices

    def _train_test_split(
        self, test_ratio: float = 0.25, valid_ratio: float = 0.1
    ) -> None:
        train, test, valid = {}, {}, {}
        # split dataset by market
        for market, index in self.indices.items():
            if self.rand:
                _test = set(random.sample(index, int(len(index) * test_ratio)))
                train_val = index.difference(test)
                _valid = set(random.sample(train_val, int(len(index) * valid_ratio)))
                _train = train_val.difference(_valid)
            else:  # traditional split method by entire segment slicing
                idx_len, idx_sorted = len(index), sorted(index)
                train_valid_boundary, valid_test_boundary = (
                    int(idx_len * (1 - valid_ratio - test_ratio)),
                    int(idx_len * (1 - valid_ratio)),
                )
                _train, _valid, _test = (
                    idx_sorted[:train_valid_boundary],
                    idx_sorted[train_valid_boundary:valid_test_boundary],
                    idx_sorted[valid_test_boundary:],
                )
            train[market], test[market], valid[market] = _train, _test, _valid
        self.train, self.test, self.valid = train, test, valid

    def _gen_idx_map(self):
        def get_map(idx_dict: Dict[str, set]) -> Dict[int, Tuple[str, int]]:
            count, res = 0, {}
            for market, indices in idx_dict.items():
                for idx in indices:
                    res[count] = (market, idx)
                    count += 1
            return res

        self._idx_map = {}
        self._idx_map["train"] = get_map(self.train)
        self._idx_map["test"] = get_map(self.test)
        self._idx_map["valid"] = get_map(self.valid)

    def inspect_dataset(self):
        return (
            self.indices,
            (self.pred, self.covs),
            self._idx_map,
            {"train": self.train, "test": self.test, "valid": self.valid},
        )

    @functools.lru_cache()
    def get_len(self, type_: str) -> int:
        return len(self._idx_map[type_])

    @functools.lru_cache(maxsize=None)
    def get_sample(
        self, idx: int, type_: str
    ) -> Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]:
        market, start = self._idx_map[type_][idx]
        ts = self.data[market]
        window, horizon = start + self.window, start + self.window + self.horizon
        slice = ts[start:window]
        past, covs, truth = (
            slice[self.pred],
            slice[self.covs],
            ts[window:horizon][self.pred],
        )
        if SLICE_NORM:
            past = self.scaler.fit_transform(past)
            truth = self.scaler.transform(truth)
        return past, covs, truth


class PartialDataset(PastCovariatesTrainingDataset):
    TRAIN, TEST, VALID = TYPES = ("train", "test", "valid")

    def __init__(self, real_dataset: MarketDataset, type_: str) -> None:
        super().__init__()
        assert (
            type_ in self.TYPES
        ), f"Type must be either of `train, test, valid` but {type_} was given"
        self.type_ = type_
        self.real_dataset = real_dataset

    def __len__(self):
        return self.real_dataset.get_len(self.type_)

    def __getitem__(self, idx: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        return tuple(i.values() for i in self.real_dataset.get_sample(idx, self.type_))


class PartialInferenceDataset(PastCovariatesInferenceDataset):
    def __init__(self, real_dataset: MarketDataset) -> None:
        self.type_ = "test"
        self.real_dataset = real_dataset

    def __len__(self):
        return self.real_dataset.get_len(self.type_)

    def __getitem__(
        self, idx: int
    ) -> Tuple[
        np.ndarray, Optional[np.ndarray], Optional[np.ndarray], darts.TimeSeries
    ]:
        past, covs, target = self.real_dataset.get_sample(idx, self.type_)
        return past.values(), covs.values(), None, target

    def get_ts_by_idx(
        self, idx: int
    ) -> Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]:
        return self.real_dataset.get_sample(idx, self.type_)


# Initialize the dataset
dataset = MarketDataset(
    time_series, target_cols=PRED_COLUMNS, stride=STRIDE, random_=RANDOM_SPLIT
)
train = PartialDataset(dataset, PartialDataset.TRAIN)
valid = PartialDataset(dataset, PartialDataset.VALID)
test = PartialInferenceDataset(dataset)


# ### Model instantiating
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
from darts.models import NaiveDrift
from darts.models.forecasting.forecasting_model import ForecastingModel

model = NaiveDrift()
model.__not_initialized__ = True
# Generated by assemble.py (models)
from typing import List, Optional
from darts.models import BlockRNNModel
from darts.models.forecasting.pl_forecasting_module import PLPastCovariatesModule
import torch
import torch.nn as nn


class _ResidualBlock(nn.Module):
    def __init__(self, input_dim: int, output_dim: int = None) -> None:
        super(_ResidualBlock, self).__init__()
        output_dim = output_dim or input_dim
        self.cnn = nn.Conv1d(input_dim, output_dim, kernel_size=3, padding="same")
        self.relu = nn.ReLU()
        self.bn = nn.BatchNorm1d(output_dim)

    def forward(self, x: torch.Tensor):
        """
        args:
            - x: Tensor(batch_size, signal_length, channels)
        """
        residual = x
        x = x.permute(0, 2, 1)
        conv = self.cnn(x)
        bn = self.bn(conv)
        act = self.relu(bn)
        act = act.permute(0, 2, 1)
        return self.relu(act + residual)

    @classmethod
    def build_residual(
        cls, input_dim: int, output_dim: int = None, block_nums: int = 4
    ) -> nn.Sequential:
        blocks = [nn.Linear(input_dim, output_dim)]
        for _ in range(block_nums - 1):
            blocks.append(_ResidualBlock(output_dim, output_dim))
        return nn.Sequential(*blocks)


class RNNCNNModel(PLPastCovariatesModule):
    def __init__(
        self,
        name: str,
        input_size: int,
        hidden_dim: int,
        num_layers: int,
        target_size: int,
        nr_params: int = None,
        num_layers_out_fc: Optional[List] = None,
        dropout: float = 0.0,
        **kwargs,
    ) -> None:
        super(RNNCNNModel, self).__init__(**kwargs)
        self.name = name or "RNN-CNN"
        self.target_size = target_size
        self.lc_hidden_dim = 1024
        self.residual = _ResidualBlock.build_residual(input_size, hidden_dim)
        self.enc = nn.LSTM(input_size, hidden_dim, num_layers, dropout=dropout)
        self.relu = nn.ReLU()
        self.dec = nn.LSTM(hidden_dim * 2, hidden_dim, num_layers, dropout=dropout)
        self.lc1 = nn.Linear(hidden_dim, self.lc_hidden_dim)
        self.lc2 = nn.Linear(self.lc_hidden_dim, target_size, bias=False)
        self.bias_lc = nn.Linear(input_size, target_size)
        self.bias_acti = nn.Hardsigmoid()
        self.bias_pool = nn.AdaptiveMaxPool1d(1)

    def forward(self, x: torch.Tensor):
        """
        args:
            - x: Tensor(batch_size, signal_length, channels)
        """
        # encoding
        conv = self.residual(x)
        recr, _ = self.enc(x)
        hidden = self.relu(torch.cat((conv, recr), dim=-1))
        # decoding
        out, _ = self.dec(hidden)
        relu = self.relu(out[:, : self.output_chunk_length, :])
        lc_hidden = self.lc1(relu)
        wb = self.relu(self.lc2(lc_hidden))
        bias = self.bias_pool(self.bias_acti(self.bias_lc(x)).permute(0, 2, 1))
        predictions = wb + bias
        batch_size, out_len = x.size(0), self.output_chunk_length
        return predictions.view(batch_size, out_len, self.target_size, 1)


_rnn_cnn = RNNCNNModel(
    "RNN-CNN",
    input_size=len(sample.columns),
    input_chunk_length=WINDOW_SIZE,
    output_chunk_length=HORIZON_SIZE,
    hidden_dim=128,
    num_layers=4,
    target_size=len(PRED_COLUMNS),
)
model = BlockRNNModel(
    input_chunk_length=WINDOW_SIZE,
    output_chunk_length=HORIZON_SIZE,
    model=_rnn_cnn,
    **GENERAL_MODEL_PARAMS,
)
from darts.models.forecasting.nbeats import _Block, _TrendGenerator, _Stack


class NBeatsRNNModule(PLPastCovariatesModule):
    def __init__(
        self,
        name: str,
        input_size: int,
        hidden_dim: int,
        num_layers: int,
        target_size: int,
        nr_params: int = None,
        num_layers_out_fc: Optional[List] = None,
        dropout: float = 0.0,
        **kwargs,
    ) -> None:
        pass

    def forward(self, x: torch.Tensor):
        """
        args:
            - x: Tensor(batch_size, signal_length, channels)
        """
        return


# ### Model training
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
model_name = (
    f'{model.__class__.__name__.lower()}-{datetime.utcnow().strftime("%m%d-%H%M")}'
)


def init_trainer() -> pl.Trainer:
    train_args = {}
    # Allow GPU is cuda available
    if torch.cuda.is_available():
        train_args["accelerator"] = "gpu"
        train_args["gpus"] = torch.cuda.device_count()
    # Add early stopping strategy and controlled checkpoints to the trainer
    es = pl.callbacks.early_stopping.EarlyStopping(
        monitor="val_loss", min_delta=0.001, patience=100
    )
    check_point = pl.callbacks.ModelCheckpoint(
        dirpath=f"./output/{model_name}", save_top_k=1, monitor="val_loss"
    )
    train_args["callbacks"] = [es, check_point]
    # Set max_epochs and precision
    train_args["max_epochs"] = EPOCH
    train_args["precision"] = 32
    # Create the trainer with kwargs
    trainer = pl.Trainer(**train_args)
    return trainer


trainer = init_trainer()


# #### Start training
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
model_path = f"./output/{model_name}"
if MODE_LOAD_MODEL:
    model.train_sample, model.output_dim = train[0], train[0][-1].shape[1]
    model._init_model(trainer)
    model.model = model.model.load_from_checkpoint(NAME_LOAD_MODEL)
    logging.info(f"Successfully load model from: {NAME_LOAD_MODEL}")
else:
    if isinstance(model, ForecastingModel) and not getattr(
        model, "__not_initialized__", False
    ):
        GENERAL_MODEL_PARAMS.update(
            {
                "work_dir": model_path,
            }
        )
        model = model.fit_from_dataset(
            train_dataset=train, val_dataset=valid, trainer=trainer, epochs=EPOCH
        )
    else:
        logging.error("Model not defined")

# #### Save the model as ONNX format for future use
if (
    isinstance(model, ForecastingModel)
    and not getattr(model, "__not_initialized__", False)
    and not MODE_LOAD_MODEL
):
    try:
        model.model.to_onnx(
            os.path.join(model_path, "model.onnx"),
            (torch.Tensor(np.concatenate(test[0][:2], axis=-1)).unsqueeze(0)),
            export_params=True,
            input_names=["input"],
            output_names=["output"],
            opset_version=11,
            dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        )
    except RuntimeError:
        logging.warning("Model not able to be saved as onnx format")

# ## Evaluation
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
evaluation_funcs = []
# Generated by assemble.py (evaluation)
evaluation_funcs = [darts.metrics.rmse, darts.metrics.r2_score]


# --------------------------
# Aggregation Functions
# Could be defined here
# --------------------------
def avg(x):
    return sum(x) / len(x)


agg_methods = [avg, np.std]


# ### Get Predictions
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
def get_prediction_pairs(
    model: ForecastingModel, test_set: PartialInferenceDataset
) -> List[Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]]:
    tester = pl.Trainer(
        False, enable_progress_bar=False, accelerator="cpu", max_epochs=1
    )
    res = []
    for i in tqdm(range(len(test_set))):
        ts_, cov, truth = test.get_ts_by_idx(i)
        pred = model.predict(HORIZON_SIZE, ts_, cov, trainer=tester)
        res.append((ts_, truth, pred))
    return res


predictions = get_prediction_pairs(model, test)
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
eval_res = {}
evaluation_funcs.append(darts.metrics.mae)
for metric in evaluation_funcs:
    func_name = metric.__name__
    losses = []
    eval_res[func_name] = {}
    for _, truth, pred in predictions:
        losses.append(metric(truth, pred))
    for agg in agg_methods:
        eval_res[func_name][agg.__name__] = agg(losses)


# --------------------------
# DO NOT MODIFY!!!
# --------------------------
def plot_predictions(
    predictions: List[Tuple[darts.TimeSeries, darts.TimeSeries, darts.TimeSeries]],
    sample_size: int = 10,
):
    if sample_size > len(predictions):
        sample_size = len(predictions)
    sample_size = sample_size or len(predictions)
    import time

    random.seed(time.time())
    predictions = random.sample(predictions, sample_size)
    for i, (past, truth, pred) in enumerate(predictions):
        past.plot(new_plot=True, label="truth")
        truth.plot(label="truth")
        pred.plot(label="pred")
        if not MODE_LOAD_MODEL:
            plt.savefig(os.path.join(model_path, f"pred-{i}.png"))
    if FIX_RANDOM:
        random.seed(0)


# ## Conclusion
# ### 1. Show the results of the experiment
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
pp = PrettyPrinter(indent=2, width=80)
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
pp.pprint(eval_res)

# ### 2. Plot Typical Prediction Results
# --------------------------
# DO NOT MODIFY!!!
# --------------------------
plot_predictions(predictions, sample_size=20)

# ### 3. Save Hyperparameters to yaml
if not MODE_LOAD_MODEL:
    import yaml

    hyper_params = {
        "KEEP_MARKETS": KEEP_MARKETS,
        "USE_FEATURE": USE_FEATURE,
        "USE_DIFF": USE_DIFF,
        "WINDOW_SIZE": WINDOW_SIZE,
        "HORIZON_SIZE": HORIZON_SIZE,
        "PRED_COLUMNS": PRED_COLUMNS,
        "EPOCH": EPOCH,
        "STRIDE": STRIDE,
        "USE_DIFF": USE_DIFF,
        "SLICE_NORM": SLICE_NORM,
        "RANDOM_SPLIT": RANDOM_SPLIT,
        "FIX_RANDOM": FIX_RANDOM,
        "GENERAL_MODEL_PARAMS": GENERAL_MODEL_PARAMS,
    }
    hyper_params.update(model.model.hparams)
    with open(os.path.join(model_path, "hyper-params.yml"), "w") as f:
        yaml.dump(hyper_params, f)
    with open(os.path.join(model_path, "result.txt"), "w") as f:
        PrettyPrinter(indent=2, stream=f).pprint(eval_res)
