"""
ML-based algorithm with regime features
"""

from AlgorithmImports import *

import numpy as np
import pandas as pd
pd.set_option('mode.use_inf_as_na', True)
from sklearn.ensemble import GradientBoostingClassifier


class MLRegimes(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)
        #self.SetEndDate(2022, 1, 1)
        self.SetCash(100000)
        self.lookbacks = [1, 4, 13, 26, 52]  # Different lookbacks used for indicators
        self.training_len = 52 * 10  # Total weeks of data used for training
        self.period = 5  # Trading frequency in days
        self.resolution = Resolution.Daily
        self.tickers = ["NVDA","XOM","AMZN","GOOGL","JPM","NEE","PFE","LUV","HD","DIS"]

        [self.AddEquity(t, self.resolution) for t in self.tickers]
        self.vix = self.AddData(CBOE, "VIX", Resolution.Daily).Symbol
        self.tlt = self.AddData(CBOE, "TLT", Resolution.Daily).Symbol
        #self.gld = self.AddData(CBOE, "GLD", Resolution.Daily).Symbol
        #self.dbc = self.AddData(CBOE, "DBC", Resolution.Daily).Symbol
        #self.qqq = self.AddData(CBOE, "QQQ", Resolution.Daily).Symbol
        self.SetBenchmark(self.tickers[0])
        self.model = GradientBoostingClassifier(n_iter_no_change=3)

        every_week = self.DateRules.WeekStart(self.tickers[0])
        self.Train(every_week, self.TimeRules.At(9, 0), self.train)
        at_market_open = self.TimeRules.AfterMarketOpen(self.tickers[0], 0)
        self.Schedule.On(every_week, at_market_open, self.trade)

    def train(self):
        features, log_returns = self.get_data(self.training_len)
        target = log_returns > 0  # Up/Down target
        self.model.fit(features, target, sample_weight=abs(log_returns))
        self.Debug(f"{self.Time} Training - Pts.: {target.value_counts()}\n")

    def trade(self):
        self.Transactions.CancelOpenOrders()
        n_points = max(self.lookbacks) + 1
        x_pred = self.get_data(n_points, include_y=False).groupby("symbol").last()
        proba_up = self.model.predict_proba(x_pred)[:, 1]  # Selecting probability of 1
        x_pred_symbols = x_pred.index.get_level_values("symbol")
        positions = pd.Series(proba_up, index=x_pred_symbols)
        if positions.sum() > 1: positions /= positions.sum()  # Capping positions at 100% (no leverage)

        for sym, pos in positions.items():
            self.SetHoldings(sym, pos)
            self.Debug(f"{self.Time} Trading\n{sym}\nPos: {pos:.1%}")
            self.Plot("Positions", sym, pos)

    def get_data(self, n_points, include_y=True):
        data = self.History(self.tickers, n_points * self.period, self.resolution)
        x = pd.DataFrame()
        for l in self.lookbacks:
            roll_high = data["high"].rolling(l*self.period).max()
            roll_low = data["low"].rolling(l*self.period).min()
            x[f"range_{l}"] = (data["close"] - roll_low) / (roll_high - roll_low)

        vix_data = self.History(self.vix, n_points * self.period, self.resolution)
        tlt_data = self.History(self.tlt, n_points * self.period, self.resolution)  # Getting historical data for TLT
        #comment out both of the TLT lines if you want this to run successfully
        x = x.join(vix_data.droplevel("symbol")["close"].rename("vix"))
        x = x.join(tlt_data.droplevel("symbol")["close"].rename("tlt"))  #Possbily add this as 50-day moving average of TLT
        
        x.dropna(inplace=True)
        if include_y:
            y = data["close"].groupby("symbol").pct_change(self.period)
            y = y.apply(np.log1p).shift(-self.period).reindex(index=x.index).dropna()
            return x.loc[y.index], y
        else:
            return x