"""
ML Technical Algorithm for SPY with random signal generator and Kelly sizing

@email: info@beawai.com
@creation date: 25/11/2022
"""

from AlgorithmImports import *

import sklearn
import numpy as np
import pandas as pd
pd.set_option('mode.use_inf_as_na', True)
from sklearn.dummy import DummyClassifier
from sklearn.model_selection import train_test_split


class E2E(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)
        self.SetEndDate(2022, 11, 1)
        self.lookback = self.GetParameter("lookback", 21)

        self.resolution = Resolution.Daily
        self.ticker = "SPY"
        self.AddEquity(self.ticker, self.resolution)
        self.model = None
        self.kelly_size = 0

        every_day = self.DateRules.EveryDay(self.ticker)
        self.Train(every_day, self.TimeRules.At(0, 0), self.train)
        self.Schedule.On(every_day,
                         self.TimeRules.AfterMarketOpen(self.ticker, 0),
                         self.trade)

    def train(self):
        """ Train model and calculate kelly position daily """
        if self.model is None: self.model = DummyClassifier(strategy="uniform")  # Random binary generator
        features, returns = self.get_data(252)  # Use last year of data for training
        target = returns >= 0  # Up/Down binary target
        model_temp = sklearn.base.clone(self.model)
        x_train, x_test, y_train, y_test, r_train, r_test = \
            train_test_split(features, target, returns, train_size=0.5, shuffle=False)
        model_temp.fit(x_train, y_train)
        y_pred = model_temp.predict(x_test)
        self.kelly_size = kelly_size(y_test, y_pred, r_test)  # Calculate kelly position on test data
        self.kelly_size = np.clip(self.kelly_size, 0, 1)  # Applies fractional kelly and clips between 0 and 1
        self.model.fit(features, target)
        self.Debug(f"{self.Time} Training - Kelly: {self.kelly_size:.1%}\n")
        self.Plot("ML", "Score", self.kelly_size)

    def trade(self):
        """ Trades based on prediction at market open """
        if self.model is None: return  # Don't trade until the model is trained
        self.Transactions.CancelOpenOrders()
        x_pred = self.get_data(self.lookback, include_y=False)
        if len(x_pred) == 0: return

        y_pred = self.model.predict(x_pred)[0]
        position = y_pred * self.kelly_size  # Sizing based on Kelly and individual probabilty
        self.Plot("ML", "Prediction", y_pred.mean())
        self.Debug(f"{self.Time} Trading\tPos: {position:.1%}")
        self.SetHoldings(self.ticker, position)

    def get_data(self, datapoints, include_y=True):
        """ Calculate features and target data """
        data = self.History([self.ticker], datapoints, self.resolution)
        features = data.eval("close/open - 1").to_frame("returns")
        x = pd.concat([features.shift(s) for s in range(self.lookback)],
                      axis=1).dropna()  # Sequence of last "lookback" returns
        if include_y:
            y = features["returns"].shift(-1).reindex_like(x).dropna()
            return x.loc[y.index], y
        else:
            return x


def kelly_size(y_true, y_pred, returns):
    """ Calculate Kelly position based on the prediction accuracy """
    trades = y_pred!=0
    wins = y_true[trades]==y_pred[trades]
    win_rate = wins.mean()
    loss_rate = 1-win_rate
    avg_win = abs(returns[trades][wins].mean())
    avg_loss = abs(returns[trades][~wins].mean())
    return win_rate/avg_loss - loss_rate/avg_win