"""
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(2021, 1, 1)
        self.SetEndDate(2022, 11, 1)
        self.lookback = self.GetParameter("lookback", 21)

        self.resolution = Resolution.Daily
        self.ticker = ["QQQ", "BAC"]
        self.symbolDataBySymbol = {}

        self.model = None
        self.kelly_size = 0
        self.spy = self.AddEquity("SPY", Resolution.Hour).Symbol

        for symbol in self.ticker:
            self.AddEquity(symbol, Resolution.Hour).Symbol
            ema7 = self.EMA(symbol, 7, Resolution.Daily, Field.Close)
            sma20 = self.SMA(symbol, 20, Resolution.Daily, Field.Close)

        symbolData = SymbolData(symbol, ema7, sma20) 
        self.symbolDataBySymbol[symbol] = symbolData

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

    def trainkelly(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%}")
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if self.Portfolio[symbol].Invested and (symbolData.sma20.Current.Value < self.Securities[symbol].Close):
                self.SetHoldings(symbol, .1, False, "Buy Signal")

    def get_data(self, datapoints, include_y=True):

        for symbol, symbolData in self.symbolDataBySymbol.items():     
            """ 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


class SymbolData:
    def __init__(self, symbol, ema7, sma20):
        self.Symbol = symbol
        self.ema7 = ema7
        self.sma20 = sma20