#region imports
from AlgorithmImports import *
#endregion
from sklearn.model_selection import train_test_split
from optimize import *
import numpy as np
import pandas as pd

class EnsembleModel():

    def __init__(self, algo):
        self.algo = algo
        self.models = None
        self.optimizer = None

    # Build a set of models
    def Train(self, X, y, returns):
        self.optimizer = ModelOptimizer(self.algo)
        try:
            self.models = self.optimizer.BuildModels(X, y, returns)
        except:
            pass
        

    # Return prediction as a +/- Z-Score of the prediction from the the threshold
    def Predict(self, X):
        z = 0

        if self.models:
            total_weight = sum(model.test_score for model in self.models)
            

            for model in self.models:
                try:
                    result = model.Predict(X)
                    # Weight the result based on the models test score
                    z = z + result * (model.test_score/total_weight)
                except:
                    pass
            
        return z

#region imports
from AlgorithmImports import *
#endregion
import talib
from talib import func
import numpy as np
import pandas as pd

#
# This helper class allows us to easily add features to our dataset
# add_features(df, ['EMA_7'])
#
# Other examples
# EMA_7_28 gives you the ratio of these two EMA
# EMA_7_28_diff gives you the difference of the current ratio and the last ratio for those two EMA
#

def ichimoku_a(high, low, n1 = 9, n2 = 26, n3 = 52):
    conv = 0.5 * (high.rolling(n1, min_periods=0).max()
                  + low.rolling(n1, min_periods=0).min())
    base = 0.5 * (high.rolling(n2, min_periods=0).max()
                  + low.rolling(n2, min_periods=0).min())
    spana = 0.5 * (conv + base)
    return pd.Series(spana)

def ichimoku_b(high, low, n1 = 9, n2 = 26, n3 = 52):
    spanb = 0.5 * (high.rolling(n3, min_periods=0).max()
       + low.rolling(n3, min_periods=0).min())
    return pd.Series(spanb)

def obv_rolling_percent(close, volume, n, fillna=False):
    df = pd.DataFrame([close, volume]).transpose()
    df['OBV'] = 0
    c1 = close < close.shift(1)
    c2 = close > close.shift(1)
    if c1.any():
        df.loc[c1, 'OBV'] = - volume
    if c2.any():
        df.loc[c2, 'OBV'] = volume
    obv = df['OBV'].rolling(n).sum() / df[volume.name].rolling(n).sum()
    if fillna:
        obv = obv.replace([np.inf, -np.inf], np.nan).fillna(0)
    return pd.Series(obv, name='obv')

def kst(close, r1 = 3, r2 = 7, r3 = 14, r4 = 28, n1 = 10, n2 = 3, n3 = 7, n4 = 14):
    rocma1 = ((close - close.shift(r1))
            / close.shift(r1)).rolling(n1, min_periods=0).mean()
    rocma2 = ((close - close.shift(r2))
              / close.shift(r2)).rolling(n2, min_periods=0).mean()
    rocma3 = ((close - close.shift(r3))
              / close.shift(r3)).rolling(n3, min_periods=0).mean()
    rocma4 = ((close - close.shift(r4))
              / close.shift(r4)).rolling(n4, min_periods=0).mean()
    return 100 * (rocma1 + 2 * rocma2 + 3 * rocma3 + 4 * rocma4)
    
def kst_sig(close, r1 = 3, r2 = 7, r3 = 14, r4 = 28, n1 = 10, n2 = 3, n3 = 7, n4 = 14, nsig = 6):
    kst_val = kst(clse, r1, r2, r3, r4, n1, n2, n3, n4)
    kstsig = kst_val.rolling(nsig, minperiods=0).mean()
    return kstsig

def check_feature_cache(cache, name, f, args):
    if not name in cache:
        cache[name] = f(*args)
    return cache[name]

def add_features(df, features, close="close", high="high", low="low", volume="volume"):

    fmap = {'MFI':'money_flow_index','ICHA':'ichimoku_a','ICHB':'ichimoku_b','BOLLH':'bollinger_hband', \
            'BOLLL':'bollinger_lband','KCC':'keltner_channel_central','NVI':'negative_volume_index', \
            'OBVP' : 'obv_rolling_percent', 'KST' : 'kst', 'KST_SIG' : 'kst_sig'}

    cache = {}


    for feature in features:
        # parse string.  Style is func_period1_period2:COL=ColumnName1 (column name optional)

        col = None
        col_idx = feature.find(':')
        if col_idx > -1:
            col = feature[col_idx+1:]
            feature = feature[0:col_idx]

        p = feature.split('_')
        if not col is None:
            feature = feature + "_" + col

        fname = p[0].upper()

        # If DM, DI, or HT function will have underscore in name, need to add suffix back and shift params back
        if len(p) > 1 and (p[1] in ['DM','DI'] or p[0] in ['HT']):
            fname += '_' + p[1]
            for i in range(2,len(p)):
                p[i-1] = p[i]
            del p[-1]

        p1 = p[1].upper() if len(p) > 1 else None
        p2 = p[2].upper() if len(p) > 2 else None
        p3 = p[3].upper() if len(p) > 3 else None

        if fname in fmap:
            if fmap[fname].find('.') > -1:
                s = fmap[fname].split('.')
                cls = s[0]
                f = getattr(globals()[cls], s[1])
            else:
                f = globals()[fmap[fname].lower()]
        elif fname in talib.__TA_FUNCTION_NAMES__:
            f = getattr(func, fname)
        elif fname.lower() in globals():
            f = globals()[fname.lower()]
        else:
            raise Exception(f'Could not find function. fname: {fname} feature: {feature}')


        if fname.endswith('MA') or fname == 'T3':
            args = [df[close], int(p1)]
            if p2 is None:
                df[feature] = (check_feature_cache(cache, feature, f, args ))
            elif p2 == 'DIFF':
                ma1 = check_feature_cache(cache, fname + '_' + p1, f, args)
                df[feature] = ma1.pct_change()
                p2 = None # So it doesn't get diffed at end of method
            else:
                ma1 = check_feature_cache(cache, fname + '_' + p1, f, args)
                args = [df[close], int(p2)]
                ma2 = check_feature_cache(cache, fname + '_' + p2, f, args)
                df[feature] = (ma1 - ma2) / ma1
                df[feature].replace([np.inf, -np.inf], 0, inplace=True)


        elif fname in ['CMO','TRIX','ROCP','ROCR100','ROC','BOLLH','BOLLL', 'RSI','MOM']:
            df[feature] = f(df[close], int(p1))


        elif fname in ['MINUS_DM', 'PLUS_DM']:
            df[feature] = f(df[high], df[low], int(p1))

        elif fname in ['ADX','ATR','KCC','MINUS_DI','PLUS_DI','WILLR']:
            try:
                if df[close].isnull().all():
                    df[feature] = np.nan
                elif p2 is None or p2.upper() == 'DIFF':
                    df[feature] = f(df[high], df[low], df[close], int(p1))
                    # Change all to percent, except WILLR
                    if not fname in ['WILLR','ADX']:
                        df[feature] = df[feature]/df[close]
                else:
                    #params = {'close' : df[close], 'high' : df[high],'low' : df[low], 'timeperiod' : int(p1)}
                    args = [df[high], df[low], df[close], int(p1)]
                    f1 = check_feature_cache(cache, fname + '_' + p1, f, args)
                    args = [df[high], df[low], df[close], int(p2)]
                    f2 = check_feature_cache(cache, fname + '_' + p2, f, args)
                    df[feature] = (f1-f2)/f1
                    df[feature].replace([np.inf, -np.inf], 0, inplace=True)
            except Exception as ex:
                if str(ex) == 'inputs are all NaN':
                    df[feature] = np.nan
                else:
                    raise


        elif fname in ['MFI']:
            df[feature] = f(df[high], df[low], df[close], df[volume], int(p1))

        elif fname in ['MACD', 'PPO', 'TSI']:
             df[feature] = f(df[close], int(p1), int(p2))

        elif fname in ['ICHA', 'ICHB', 'AO']:
            # Normalize
            df[feature] = f(df[high], df[low], int(p1), int(p2))
#             if fname.startswith('ICH'):
#                 df[feature] = df[feature]/ df[close]

        elif fname in ['NVI']:
            df[feature] = f(df[close], df[volume])

        elif fname in ['SAR', 'SAREXT']:
            # Normalize
            df[feature] = f(df[high], df[low]) #/ df[close]

        elif fname == 'KST':
            df[feature] = f(df[close], r1=10, r2=15, r3=20, r4=30, n1=10, n2=10, n3=10, n4=15)

        elif fname == 'KST_SIG':
            df[feature] = f(df[close], r1=10, r2=15, r3=20, r4=30, n1=10, n2=10, n3=10, n4=15, nsig=9)

        elif fname == 'AROON':
            df[feature] = f(df[high], df[low], int(p1))[int(p2)-1]

        elif fname == 'ULTOSC':
            df[feature] = f(df[high], df[low], df[close])

        elif fname in ['OBVP']:
            df[feature] = f(df[close],df[volume],int(p1))

        elif fname in ['OBTP']:
            df[feature] = f(df[close],df["Trades"],int(p1))

        elif fname in ['HT_DCPERIOD','HT_DCPHASE','HT_PHASOR','HT_SINE', 'HT_TRENDMODE','STOCHRSI' ]:
            df[feature] = f(df[close])

        else:
            raise Exception('Feature not found: ' + feature + ' ' + fname )

        # Difference those that are 1 period diffs
        if p3 == 'DIFF' or p2 == 'DIFF':
            df[feature] = df[feature].diff()


    return df
    
def get_exhaustive_ma_features_list():

    features = []
    periods = [7, 14, 28, 42, 70]

    mas = ["SMA", "DEMA", "KAMA", "EMA", "TEMA", "T3", "TRIMA", "WMA"]
    for p in periods:
        for ma in mas:
            features.append(ma + "_" + str(p))

    for p1 in periods:
        for p2 in periods:
            if p2 <= p1:
                continue

            for ma in mas:
                ma1 = ma + str(p1)
                ma2 = ma + str(p2)
                features.append(ma + "_" + str(p1) + "_" + str(p2))
                features.append(features[-1] + "_diff")

    return features

def get_exhaustive_features_list(diffs=True):

    features = get_exhaustive_ma_features_list()

    periods = [3,7,14,20,30]
    band_periods = [7,14,21,28]
    hi_lo_periods = [[4,9],[7,14],[10,21],[12,26]]

    pfs = ['CMO','MINUS_DI','MINUS_DM','PLUS_DI','PLUS_DM','ROCP','ATR','ADX',
                       'TRIX','WILLR']
    # Left out: MOM, VIN, VIP, ROC,ROCR,OBVM,EOM,DPI,RSI,STOCH,STOCH_SIG,FI
    for p in periods:

        features.append('AROON_' + str(p) + '_1')
        features.append('AROON_' + str(p) + '_2')
        for pf in pfs:
            if pf == 'ADX' and p == 3:
                continue
            features.append(pf + '_' + str(p))
            if diffs:
                features.append(features[-1] + "_diff")
    if diffs:
        for p1 in periods:
            for p2 in periods:
                if p2 <= p1:
                    continue

                for pf in pfs:
                    pf1 = pf + str(p1)
                    pf2 = pf + str(p2)
                    features.append(pf + "_" + str(p1) + "_" + str(p2))
                    features.append(features[-1] + "_diff")


    hlfs = ['ICHA','ICHB']
    # Exclude PPO, same as our MA setups
    # hi lo Indicators
    for hl in hi_lo_periods:
        for hlf in hlfs:
            features.append(hlf + '_' + str(hl[0]) + '_' + str(hl[1]))

    # other indicators
    otfs = ['SAR','SAREXT','HT_DCPERIOD','ULTOSC','OBVP_10','OBVP_20','OBVP_40','OBVP_80']
    for otf in otfs:
        features.append(otf)

    return features

from AlgorithmImports import *
from risk import BracketRiskModel, TrailingStopRiskManagementModelLiquidation

from feature import add_features
from ensemble import EnsembleModel
from datetime import timedelta
import pandas as pd
import numpy as np
import statistics
import sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.utils import resample, shuffle
import math

class TransdimensionalDynamicThrustAssembly(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2010, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 1, 1)  # Set End Date
        self.InitCash = 100000
        self.SetCash(self.InitCash)  
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        
        # time frames and frequencies
        self.resolution = Resolution.Daily
        self.rebalancefunc = Expiry.EndOfQuarter
        self.lookback = 250
        self.metamodel_rolling_window = 50

        # ML parameters
        self.years = 4
        self.n_estimators = 10
        self.min_samples_split = 140
        self.num_iter = 10
        self.num_test_records = 60

        # UniverseSize determines the number of ETF constituents to select from each ETF. The top x by weight are selected.
        self.UniverseSize = 9
        
        self.kelly_multiplier = 1

        self.SetWarmup(self.lookback)

        # if this flag is set to false, none of the meta labeling code gets called, thus saving a lot of expensive compute
        self.meta_labelling = False

        if self.meta_labelling:
            import ht_auth
            ht_auth.SetToken(self.GetParameter("mlfinlab-api-key"))
            import mlfinlab as ml
            self.ml = ml
            from mlfinlab.bet_sizing import bet_sizing
            self.bet_sizing = bet_sizing

        # add SPY as the benchmark but don't trade it (that's what self.exclusions is for, it's being checked at various points)!
        self.SetBenchmark('SPY')
        self.MKT = self.AddEquity('SPY', self.resolution).Symbol
        self.mkt = []
        self.exclusions = ['SPY']

        self.UniverseSettings.Resolution = self.resolution
        
        features = ['HT_DCPERIOD', 'HT_DCPHASE', 'MOM_5', 'ADX_10_diff','MOM_5_diff']

        self.etfs = ["SPY"]

        for etf in self.etfs:
            self.AddUniverse(self.Universe.ETF(etf, Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))

        # Select the alpha model
        self.SetAlpha(WeightedClassifierAlpha(self, features, self.lookback, self.metamodel_rolling_window, self.meta_labelling, self.years))

        # Equally weigh securities in portfolio, based on insights
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())

        # Set VWAP Execution Model (Immidiate Execution worked well, but VWAP execution seems to produce slightly better results)
        self.SetExecution(VolumeWeightedAveragePriceExecutionModel())

        # Set Risk Management Model - none of the risk models imroved the performance over the alpha model's own flat insights, so for now NullRiskModel seems best
        # Bracket Risk Model (using SL and TP) and Trailing Stop (using SL only) are both modified to prevent the Alpha Model from re-opening positions that the risk model has closed
        # self.stopLoss = 0.02
        # self.takeProfit = 0.95
        self.SetRiskManagement(NullRiskManagementModel())
        # self.SetRiskManagement(BracketRiskModel(self.stopLoss, self.takeProfit))
        # self.SetRiskManagement(TrailingStopRiskManagementModelLiquidation(self.stopLoss))


    def ETFConstituentsFilter(self, constituents: List[ETFConstituentData]) -> List[Symbol]:
        # Get the UniverseSize securities with the largest weight in the index
        selected = sorted([c for c in constituents if c.Weight],
            key=lambda c: c.Weight, reverse=True)[:self.UniverseSize]

        if not selected:
            self.Log("ETF symbols without weight. Returning Universe Unchanged: "+str([c.Value for c in Universe.Unchanged]))
            return Universe.Unchanged
        else:
            self.Log("ETF symbols selected: "+str([c.Symbol.Value for c in selected]))
            return [c.Symbol for c in selected]

    def OnEndOfDay(self, symbol):
        # plot the benchmark on the main equities chart for comparison
        if not self.IsWarmingUp: 
            mkt_price = self.History(self.MKT, 2, Resolution.Daily)['close'].unstack(level= 0).iloc[-1]
            self.mkt.append(mkt_price)
            mkt_perf = self.InitCash * self.mkt[-1] / self.mkt[0] 
            self.Plot('Strategy Equity', self.MKT, mkt_perf)

class WeightedClassifierAlpha(AlphaModel):

    def __init__(self, algo, features, lookback, metamodel_rolling_window, meta_labelling, years):
        self.algo = algo
        self.models = {}
        self.features = features
        self.rebalanceTime = datetime.min
        self.symbols = []
        self.lookback = lookback
        self.metamodel_rolling_window = metamodel_rolling_window
        self.years = years
        self.kelly_size = {}
        self.avg_win = {}
        self.avg_loss = {}

        self.meta_labelling = meta_labelling
        self.primary_history = {}
        self.metamodel = {}
        self.triple_barrier_events = {}

    def Update(self, algo, data):

        insights = []

        if not self.symbols:
            return insights

        try:

            # This can be adjusted to only allow trading for symbols with a certain Kelly size (as determined during the last training of the model)
            min_kelly = 0

            # Load History and Features          
            end = self.algo.Time
            start = end - timedelta(days=self.lookback)

            X = self.GetXforPrimaryModel(self.symbols, start, end)
            time = X.index[-1][0]
            X = X.loc[pd.IndexSlice[time, :], :]

            insights = []

            if not self.algo.IsWarmingUp:

                # Predict each symbols direction
                for symbol in self.symbols:
                    symbol = SymbolCache.GetSymbol(symbol)

                    if symbol in self.kelly_size:
                        kelly_size = self.kelly_size[symbol]
                        if math.isnan(kelly_size):
                            kelly_size = 0
                            self.algo.Log("NaN Kelly Size for "+str(symbol.Value))
                    else: 
                        kelly_size = 0
                        if symbol.Value not in self.algo.exclusions:
                            self.algo.Log("No Kelly Size for "+str(symbol.Value))

                    if kelly_size >= min_kelly:

                        X_symbol = X.loc[pd.IndexSlice[:, [symbol.ID.ToString()]], :]

                        if self.models and data.ContainsKey(symbol) and data[symbol]:
                            if symbol in self.models:

                                model = self.models[symbol]
                                result = model.Predict(X_symbol)

                                direction = InsightDirection.Flat

                                metapred = 1
                                metapred_prob = 1
                                size = self.kelly_size[symbol] * self.algo.kelly_multiplier

                                if self.meta_labelling:
                                    if symbol in self.primary_history.keys() and symbol in self.metamodel:
                                        self.UpdateRollingHistoryWindow(symbol, data[symbol].EndTime, data[symbol].Close, result)
                                        X_meta, y_meta, li, tbi = self.GetXyforMetaModel(self.primary_history[symbol], symbol,0, False)

                                        metapred = self.metamodel[symbol].predict(X_meta.iloc[-1:])[0]
                                        metapred_prob = prob[-1][1]

                                        if self.avg_win[symbol] == 0 or self.avg_loss[symbol] == 0:
                                            size = 0
                                        else:
                                            size = metapred_prob - ((1-metapred_prob)/(self.avg_win[symbol]/self.avg_loss[symbol]))

                                        # Alternatibely: using MLFinLab's bet_size_probability instead of kelly to determine  size:
                                        # prob = self.metamodel[symbol].predict_proba(X_meta.fillna(0))
                                        # zlst = list(zip(*prob))
                                        # size = self.algo.bet_sizing.bet_size_probability(self.triple_barrier_events[symbol], pd.Series(zlst[1]), 2).iloc[-1]

                                if math.isnan(size):
                                    size = 0

                                # using result as measure of confidence and (Kelly) size as weight. 
                                # (For InsightWeightingPortfolioConstructionModel confidence has no effect, this is just in case it might be useful for aby other PCM)
                                if result > 0 and metapred and size > 0:
                                    insights.append(
                                        Insight.Price(symbol, timedelta(days=2), InsightDirection.Up, result, None, None, size))

                                elif result < 0:
                                    size = 0
                                    insights.append(
                                        Insight.Price(symbol, timedelta(days=1), InsightDirection.Flat, -result, None, None, size))

            return insights

        except Exception as e:
            self.algo.Log('Unexpected error during insights generation:' + str(e))
            raise

    def OnSecuritiesChanged(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:

        pred_hist = pd.DataFrame

        for security in changes.RemovedSecurities:
            if security.Symbol in self.symbols:
                self.symbols.remove(security.Symbol)

            self.metamodel.pop(security.Symbol, None)
            self.triple_barrier_events.pop(security.Symbol, None)
            if self.meta_labelling:
                self.primary_history.pop(security.Symbol, None)

            if security.Invested:
                self.algo.Liquidate(security.Symbol, "Removed from Universe")
        
        for security in changes.AddedSecurities:
            if security.Symbol.Value not in self.algo.exclusions:
                self.symbols.extend([security.Symbol])

        if self.algo.Time > self.rebalanceTime:
            # it's time to re-train all models (and not just the newly added ones)
            symbols = self.symbols
            self.rebalanceTime = self.algo.rebalancefunc(self.algo.Time)
            self.algo.Debug(str(self.algo.Time)+" Rebalancing. Training all "+str([symbol.Value for symbol in symbols]))

        else:
            symbols = [security.Symbol for security in changes.AddedSecurities]
            self.algo.Debug(str(self.algo.Time)+" Not yet time for Rebalancing. Training new symbols "+str([symbol.Value for symbol in symbols]))

        if symbols:
            for symbol in symbols:
                self.models[symbol] = self.TrainModel(symbol, self.years)

            # initial history of newly added securities for meta labelling
            if self.meta_labelling:
                for symbol in symbols:
                    if symbol.Value not in self.algo.exclusions:

                        try:
                            end = self.algo.Time - timedelta(days=1)
                            start = end - timedelta(days=self.lookback * 2)

                            X = self.GetXforPrimaryModel([symbol], start, end)

                            historystart = self.algo.Time - timedelta(days=self.lookback)
                            nearest = X.reset_index().time.searchsorted(historystart)
                            historystart = X.index[nearest-1][0]

                            pred_hist= self.algo.History(symbol, historystart, end, self.algo.resolution).droplevel(0, axis=0)
                            pred_hist.index.names = ['date_time']
                            pred_hist = pred_hist[['close']]
                            pred_hist['result']=0.0

                            iter_date = self.algo.Time - timedelta(days=self.lookback)

                            self.primary_history[symbol] = pd.DataFrame(columns=['time','close','result']).set_index('time')

                            while iter_date <= end:
                                nearest = X.reset_index().time.searchsorted(iter_date)
                                if 0 <= nearest < len(X.index):
                                    time = X.index[nearest][0]

                                    X_slice = X.loc[pd.IndexSlice[time, :], :]

                                    X_symbol = X_slice.loc[pd.IndexSlice[:, [symbol.ID.ToString()]], :]

                                    model = self.models[symbol]
                                    result = model.Predict(X_symbol)
                                    pred_hist.loc[time]['result'] = result
                                    if self.meta_labelling:
                                        self.UpdateRollingHistoryWindow(symbol, time, pred_hist.loc[time]['close'], result)
                                iter_date += timedelta(days=1)

                            self.metamodel[symbol], self.triple_barrier