from arch import arch_model
import numpy as np
import pandas as pd
from sklearn.model_selection import KFold
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from sklearn.metrics import mean_squared_error
import pickle

class KNNVolModel(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)  # Set Start Date
        #self.SetEndDate(2018, 3, 1)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash
        # serializedModel = bytes( self.ObjectStore.ReadBytes("VolModel") )
        # self.volModel = pickle.loads(serializedModel)
        self.symbolDataBySymbol = {}
        
        self.tickers = ["MSFT", "WFC", "HD"]
        
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        
        for ticker in self.tickers:
            self.AddEquity(ticker, Resolution.Minute)
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen(self.spy, 2),
                 self.EveryDayAfterMarketOpen)
                 
        self.AutomaticIndicatorWarmup = True
                 
        self.slice = None
        
        self.lastYear = self.Time.year
        
    def CheckYear(self):
        if self.lastYear == self.Time.year:
            return
        self.lastYear = self.Time.year
        for symbol, symbolData in self.symbolDataBySymbol.items():
            symbolData.TrainModel()

    def OnData(self, slice):
        self.slice = slice
        self.CheckYear()
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
            symbolData.CheckExercise()
        
    def EveryDayAfterMarketOpen(self):
        for symbol, symbolData in self.symbolDataBySymbol.items():
            symbolData.EveryDay()
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Symbol.SecurityType != SecurityType.Equity: continue
            symbol = security.Symbol
            if symbol == self.spy: continue
            option = self.AddOption(symbol)
            option.SetFilter(self.UniverseFunc)
            atr = self.ATR(symbol, 14, MovingAverageType.Simple, Resolution.Daily)
            history = self.History(symbol, 100, Resolution.Daily)
            history = history.loc[symbol]
            symbolData = SymbolData(symbol, option, atr, history, self)
            self.Debug(f"Created symbolData for {symbol.Value}")
            self.symbolDataBySymbol[symbol] = symbolData
            
    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys().Strikes(-15, 15).Expiration(timedelta(10), timedelta(30))
                
class SymbolData:
    def __init__(self, symbol, option, atr, history, algo):
        self.Symbol = symbol
        self.Option = option
        self.Atr = atr
        self.history_ = history
        self.algo = algo
        self.previousPred = -1
        self.investedOptions = False
        self.volModel = None
        self.TrainModel()
        
    def volpred(self, returns,h):
        am = arch_model(returns, p=1,o=1, q=1, vol='Garch', dist='skewt')
        res = am.fit(disp='off')
        forecasts=res.forecast(horizon=h, method='simulation')
        return forecasts.variance.iloc[-1]
    
    def EveryDay(self):
        
        df = self.RetrieveVolatilityMetrics()
        
        if self.CheckDataIntegrity(df) == True:
            self.algo.Debug(f"Integrity failed for {self.Symbol.Value}")
            return

        prediction = self.volModel.predict(df)
        self.EnterPositions(prediction, df)
        self.previousPred = prediction
        
    def CheckExercise(self):
        if self.algo.Portfolio[self.Symbol].Quantity != 0:
            #self.algo.Debug(f"Liquidating {self.Symbol.Value}")
            self.algo.Liquidate(self.Symbol)
            self.investedOptions = False
        
    def CheckDataIntegrity(self, df):
        if pd.isnull(df["GKHV"][0]) == True or pd.isnull(df["GARCH"][0]) == True or pd.isnull(df["ATR"][0]) == True:
            #self.algo.Debug(f"NaN value detected {self.Symbol.Value}")
            self.history_ = self.algo.History(self.Symbol, 100, Resolution.Daily)
            self.history_ = self.history_.loc[self.Symbol]
            return True
        return False
        
    def RetrieveVolatilityMetrics(self):
        yesterday = self.algo.History(self.Symbol, 1, Resolution.Daily)
        yesterday = yesterday.loc[self.Symbol]
        self.history_.append(yesterday)
        
        
        self.history_ = self.history_.iloc[1: , :]
        returns = 100 * self.history_["close"].pct_change().dropna()
        
        gkhv = np.sqrt(252/10 * pd.DataFrame.rolling(0.5 * np.log(self.history_.loc[:, 'high'] / self.history_.loc[:, 'low']) ** 2 -
                                                (2 * np.log(2) - 1) *
                                                 np.log(self.history_.loc[:, 'close'] / self.history_.loc[:, 'open']) ** 2, window=10).sum())
                                                 
        self.gkhv = gkhv.pct_change(3).iloc[-1]
        
        df = pd.DataFrame()
        df["GKHV"] = [abs(self.gkhv)]
        df["GARCH"] = [self.volpred(returns, 10).mean()]
        df["ATR"] = [self.Atr.Current.Value]
        
        return df
        
    def CheckOptionPositions(self):
        option_invested = [x.Key for x in self.algo.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        for option in option_invested:
            if self.Symbol == option.ID.Underlying.Symbol:
                self.investedOptions = True
                return
        self.investedOptions = False
        
    def EnterPositions(self, prediction, df):
        
        self.CheckOptionPositions()
        
        if prediction == 1 and self.previousPred == 0:
            self.algo.Liquidate(self.Symbol)
            self.investedOptions = False

        if prediction == 1 and self.investedOptions == False:
            #self.algo.Debug(df)
            self.LongStraddle(self.algo.slice)
            #self.algo.Debug(f"Long {prediction} for {self.Symbol.Value}")
        elif prediction == 0 and self.investedOptions == False:
            #self.algo.Debug(df)
            self.ShortStraddle(self.algo.slice)
            #self.algo.Debug(f"Short {prediction} for {self.Symbol.Value}")
    
    # def CheckProfit(self):
    #     if self.Portfolio.TotalUnrealizedProfit >= (self.Portfolio.TotalUnrealizedProfit):
    
        
    def LongStraddle(self, slice):
        # If there is undelying assets in portfolio at expiration, liquidate the stocks in order to roll into new contracts
        self.CheckExercise()
        
        if self.algo.Time.hour != 0 and self.algo.Time.minute != 0: 
            for i in slice.OptionChains:
                if self.Symbol != i.Value.Underlying.Symbol: continue
                chain = i.Value
                contract_list = [x for x in chain]
                # if there is no optionchain or no contracts in this optionchain, pass the instance
                if (slice.OptionChains.Count == 0) or (len(contract_list) == 0): 
                    return   
    
                # sorted the optionchain by expiration date and choose the furthest date
                expiry = sorted(chain,key = lambda x: x.Expiry)[0].Expiry
                # filter the call and put options from the contracts
                call = [i for i in chain if i.Expiry == expiry and i.Right == 0]
                put = [i for i in chain if i.Expiry == expiry and i.Right == 1]
    
                # sorted the contracts according to their strike prices 
                call_contracts = sorted(call,key = lambda x: x.Strike)    
                put_contracts = sorted(put,key = lambda x: x.Strike)    
                if len(call_contracts) == 0 or len(put_contracts) == 0 : return
                
                atm_call = sorted(call_contracts,key = lambda x: abs(chain.Underlying.Price - x.Strike))[0]
                atm_put = sorted(put_contracts,key = lambda x: abs(chain.Underlying.Price - x.Strike))[0]
                self.algo.Buy(atm_call.Symbol, 2)
                self.algo.Buy(atm_put.Symbol, 2)
                self.investedOptions = True
                return
        self.algo.Debug(f"Failed Long for {self.Symbol.Value}")
                
    def ShortStraddle(self, slice):
        # If there is undelying assets in portfolio at expiration, liquidate the stocks in order to roll into new contracts
        self.CheckExercise()
        
        if self.algo.Time.hour != 0 and self.algo.Time.minute != 0: 
            for i in slice.OptionChains:
                if self.Symbol != i.Value.Underlying.Symbol: continue
                chain = i.Value
                contract_list = [x for x in chain]
                # if there is no optionchain or no contracts in this optionchain, pass the instance
                if (slice.OptionChains.Count == 0) or (len(contract_list) == 0): 
                    return   
    
                # sorted the optionchain by expiration date and choose the furthest date
                expiry = sorted(chain,key = lambda x: x.Expiry)[0].Expiry
                # filter the call and put options from the contracts
                call = [i for i in chain if i.Expiry == expiry and i.Right == 0]
                put = [i for i in chain if i.Expiry == expiry and i.Right == 1]
    
                # sorted the contracts according to their strike prices 
                call_contracts = sorted(call,key = lambda x: x.Strike)    
                put_contracts = sorted(put,key = lambda x: x.Strike)    
                if len(call_contracts) == 0 or len(put_contracts) == 0 : return
                
                if len(call_contracts) < 14:
                    return
                
                atm_call = sorted(call_contracts,key = lambda x: x.Strike)[len(call_contracts)-1]
                atm_put = sorted(put_contracts,key = lambda x: x.Strike)[0]
                self.algo.Sell(atm_call.Symbol, 4)
                self.algo.Sell(atm_put.Symbol, 4)
                #self.algo.Debug(f"{atm_call.Strike} {atm_put.Strike} {self.algo.Securities[self.Symbol].Price}")
                self.investedOptions = True
                return
        self.algo.Debug(f"Failed Short for {self.Symbol.Value}")
            
    def TrainModel(self):
        history = self.algo.History(self.Symbol, 1000, Resolution.Daily)
        history = history.loc[self.Symbol]
        
        atr = AverageTrueRange(14, MovingAverageType.Simple)
        atr_arr = []
        for tuple in history.itertuples():
            bar = TradeBar(tuple.Index, self.Symbol, tuple.open, tuple.high, tuple.low, tuple.close, tuple.volume)
            atr.Update(bar)
            if not atr.IsReady:
                atr_arr.append(np.nan)
                continue
            atr_arr.append(atr.Current.Value)
            
        history["ATR"] = atr_arr
        
        history["pct_change"] = abs(100 * history["close"].pct_change(10).dropna())
        
        Garch = [np.nan for x in range(0, 100)]
        for index in range(100, len(history.index)):
            returns = 100 * history["close"].iloc[index-100:index].pct_change().dropna()
            garchPred = self.volpred(returns, 10)
            Garch.append(garchPred.mean())
        history["Garch_Pred"] = Garch
        
        gkhv = np.sqrt(252/10 * pd.DataFrame.rolling(0.5 * np.log(history.loc[:, 'high'] / history.loc[:, 'low']) ** 2 -
                                                (2 * np.log(2) - 1) *
                                                 np.log(history.loc[:, 'close'] / history.loc[:, 'open']) ** 2, window=10).sum())
        
        history["GKHV"] = abs(gkhv.pct_change(3))# gkhv.iloc[-1]#
        history = history.dropna()
        
        history["pct_delayed"] = history["pct_change"].shift(-10).dropna()
        
        history.loc[abs(history['pct_delayed']) > (history["pct_change"].mean()*1.5), 'Strong_Change'] = 1
        history.loc[abs(history['pct_delayed']) <= (history["pct_change"].mean()*1.5), 'Strong_Change'] = 0
        
        history = history.dropna()
        
        y = np.array(history["Strong_Change"])
        X = history[["GKHV", "Garch_Pred", "ATR"]]
        
        # split_percentage = 0.8
        # split = int(split_percentage*len(history.index))
        
        # X_train = X[:split]
        # y_train = y[:split]
        # X_test = X[split:]
        # y_test = y[split:]
        
        cls = KNeighborsClassifier(n_neighbors=10)
        cls = cls.fit(X, y)

        self.volModel = cls