import pandas as pd
import numpy as np
from io import StringIO
import json
import pickle
from datetime import datetime
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

'''
Would it be beneficial to run a universe selection every day based on which stocks 
have a trading halt? Yes, it increases spead 10 fold
'''

class TradingHalt(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2021, 1, 30)
        self.SetCash(100000)
        
        serializedModel = bytes( self.ObjectStore.ReadBytes("HaltModel") )
        self.haltModel = pickle.loads(serializedModel)
        
        self.activeSymbols = {}
        self.symbolDataBySymbol = {}
        self.tickersT = []
        self.tickers = []
        
        # self.Resolution = Resolution.Hour
        # self.UniverseSettings.Resolution = self.Resolution
        # self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction))
        
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose("SPY", 10),
                 self.SellAll)
                 
        self.todaysHalts = pd.DataFrame()
        
        self.prepareTrading = []
        self.maxPositions = 10
        self.tickers = ["GME"]
        for ticker in self.tickers:
            self.AddEquity(ticker, Resolution.Minute)
            
        self.lastHour = -1
        self.options = {}
        self.gme = None
        self.count = 0
        
    def TakeProfit(self):
        activePositions = [x.Key for x in self.Portfolio if x.Value.Invested]
        for symbol in activePositions:
            if self.Portfolio.UnrealizedProfitPercent >= 0.05:
                self.Liquidate()
        
    def OnData(self, slice):
        if self.Portfolio.TotalUnrealizedProfit >= (self.Portfolio.TotalPortfolioValue*0.01):
            self.Debug("Profit")
            self.Liquidate()
            
        if self.Portfolio["GME"].Quantity != 0:
            self.Debug("Liquidating")
            self.Liquidate()
        
        if self.Time.day == self.lastHour:
            return
        self.lastHour = self.Time.day
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if not slice.ContainsKey(symbol): continue
            df = symbolData.Structured()
            prediction = self.haltModel.predict(df)
            # if prediction == -1:
            #     if self.count < 0:
            #         self.count = 0
            #     self.count += 1
            # elif prediction == 1:
            #     if self.count > 0:
            #         self.count = 0
            #     self.count -= 1
                
            # if self.count >= 3 and not self.Portfolio.Invested:
            if prediction == -1 and not self.Portfolio.Invested:
                self.Debug(df)
                self.Debug(self.Time)
                self.Debug(prediction)
                #self.SetHoldings(symbol, 0.05)
                self.TradeOptions(slice)
            # if self.count <= -3:
            #     self.Liquidate()

            
                
    def TradeOptions(self,slice):
        # If there is undelying assets in portfolio at expiration, liquidate the stocks in order to roll into new contracts
        if self.Portfolio["GME"].Quantity != 0:
            self.Debug("Liquidating 2")
            self.Liquidate()
        
        if not self.Portfolio.Invested and self.Time.hour != 0 and self.Time.minute != 0: 
            for i in slice.OptionChains:
                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)[-1].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 : continue
    
                # otm_put_lower = put_contracts[3]
                # otm_put = put_contracts[1]
                # otm_call = call_contracts[-1]
                # otm_call_higher = call_contracts[-7]
                # self.trade_contracts = [otm_call.Symbol,otm_call_higher.Symbol,otm_put.Symbol,otm_put_lower.Symbol]
                
                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.Buy(atm_call.Symbol ,1)
                self.Buy(atm_put.Symbol ,1)
                self.Debug(f"{atm_call.Expiry} {atm_put.Expiry}")
    
                # if there is no securities in portfolio, trade the options 
                # self.Sell(otm_put_lower.Symbol ,1)
                # self.Buy(otm_put.Symbol ,1)
                # self.Buy(otm_call.Symbol ,1)
                # self.Sell(otm_call_higher.Symbol ,1)
                # self.Debug(f"{otm_put_lower.Expiry} {otm_put.Expiry} {otm_call.Expiry} {otm_call_higher.Expiry}")
                
    def TradeButter(self, slice):
        if self.Portfolio["GME"].Quantity != 0:
            self.Debug("Liquidating 2")
            self.Liquidate()
        
        if not self.Portfolio.Invested and self.Time.hour != 0 and self.Time.minute != 0: 
            for i in slice.OptionChains:
                #if i.Key != self.symbol: continue
                chain = i.Value
                # sorted the optionchain by expiration date and choose the furthest date
                expiry = sorted(chain,key = lambda x: x.Expiry, reverse=True)[0].Expiry
                # filter the call options from the contracts expires on that date
                call = [i for i in chain if i.Expiry == expiry and i.Right == 0]
                # sorted the contracts according to their strike prices 
                call_contracts = sorted(call,key = lambda x: x.Strike)    
                if len(call_contracts) == 0: continue
                # choose OTM call 
                self.otm_call = call_contracts[-1]
                # choose ITM call 
                self.itm_call = call_contracts[0]
                # choose ATM call
                self.atm_call = sorted(call_contracts,key = lambda x: abs(chain.Underlying.Price - x.Strike))[0]
    
                self.Sell(self.atm_call.Symbol ,2)
                self.Buy(self.itm_call.Symbol ,1)
                self.Buy(self.otm_call.Symbol ,1)

    def SellAll(self):
        self.Liquidate()
        
    def HourBarHandler(self, consolidated):
        self.symbolDataBySymbol[self.gme].UpdateVol(consolidated)
        
    def OnSecuritiesChanged(self, changes):
        for x in changes.AddedSecurities:
            if x.Type==SecurityType.Option: continue
            symbol = x.Symbol
            if symbol == self.spy: continue
            self.gme = symbol
            self.Securities[symbol].SetSlippageModel(CustomSlippageModel(self))
            self.Securities[symbol].SetDataNormalizationMode(DataNormalizationMode.Raw)
            self.activeSymbols[symbol.Value] = symbol
            symbolData = SymbolData(self, symbol)
            self.symbolDataBySymbol[symbol] = symbolData
            
            option = self.AddOption(symbol)
            option.SetFilter(self.UniverseFunc)
            self.Consolidate(symbol, Resolution.Daily, self.HourBarHandler)


            
            
        for x in changes.RemovedSecurities:
            symbol = x.Symbol
            if x.Type==SecurityType.Option: continue
            self.Liquidate(symbol)
            self.activeSymbols.pop(symbol.Value, None)
            if symbol in self.symbolDataBySymbol:
                self.symbolDataBySymbol.pop(symbol, None)
                
    def UniverseFunc(self, universe):
        #-15 15
        #-9 9
        return universe.IncludeWeeklys().Strikes(-5, 5).Expiration(timedelta(0), timedelta(60))
            
class CustomSlippageModel:
    def __init__(self, algorithm):
        self.algorithm = algorithm

    def GetSlippageApproximation(self, asset, order):
        slippage = asset.Price * float(0.0001 * np.log10(2*float(order.AbsoluteQuantity)))
        return slippage
        
class SymbolData:
    def __init__(self, algo, symbol):
        self.Symbol = symbol
        self.algo = algo
        self.rollingData = RollingWindow[TradeBar](300)
        self.gkhv = 0
        self.pctGkhv = 0
        history = algo.History(symbol, 300, Resolution.Daily)
        history = history.loc[symbol]
        for tuple in history.itertuples():
            bar = TradeBar(tuple.Index, symbol, tuple.open, tuple.high, tuple.low, tuple.close, tuple.volume)
            self.rollingData.Add(bar)
            
        self.UpdateVol(None)
        
    def UpdateVol(self, data):
        if data != None:
            self.rollingData.Add(data)
        _open = []
        _high = []
        _low = []
        _close = []
        for index in range(0, self.rollingData.Count-1):
            bar = self.rollingData[index]
            _open.append(bar.Open)
            _high.append(bar.High)
            _low.append(bar.Low)
            _close.append(bar.Close)
        history = pd.DataFrame()
        history["open"] = _open
        history["high"] = _high
        history["low"] = _low
        history["close"] = _close
        
        gkhv = np.sqrt(1/252 * 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=252).sum())

        self.gkhv = gkhv.iloc[-1]                    
        self.pct10 = gkhv.pct_change(3).iloc[-1]
        
    def Structured(self):
        df = pd.DataFrame()
        df["PCT"] = [self.pct10]
        df["GKHV"] = [self.gkhv]
        return df

        
        
#-------------------------------------------------------------------------------------------------------------------------------
        
# import pandas as pd
# import numpy as np
# from io import StringIO
# import json
# import pickle
# from datetime import datetime
# from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

# '''
# Would it be beneficial to run a universe selection every day based on which stocks 
# have a trading halt? Yes, it increases spead 10 fold
# '''

# class TradingHalt(QCAlgorithm):

#     def Initialize(self):
#         self.SetStartDate(2018, 1, 1)
#         #self.SetEndDate(2021, 8, 21)
#         self.SetCash(100000)
#         csv = self.Download("https://www.dropbox.com/s/nsirge1gachgjg2/HaltData.csv?dl=1")
#         self.HaltData = pd.read_csv(StringIO(csv))
        
#         if self.ObjectStore.ContainsKey("TradingHaltObject"):
#             self.ObjectStore.Delete("TradingHaltObject")
        
#         result = self.HaltData.to_json(orient="index")
#         parsed = json.loads(result)
#         dump = json.dumps(parsed)#, indent=4)  
        
#         # dump = json.dumps(self.dict)
        
#         self.ObjectStore.Save("TradingHaltObject", dump)
        
#         serializedModel = bytes( self.ObjectStore.ReadBytes("HaltModel") )
#         self.haltModel = pickle.loads(serializedModel)

#         self.priorityTypes = ["T1", "T2", "T5", "T6", "LUDP", "LUDS"] #, "MWC1", "MWC2", "MWC3", "MWCO"]
#         self.ConvertTimes()
        
#         self.activeSymbols = {}
#         self.symbolDataBySymbol = {}
#         self.tickersT = []
#         self.tickers = []
        
#         self.Resolution = Resolution.Hour
#         self.UniverseSettings.Resolution = self.Resolution
#         self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction))
        
#         self.spy = self.AddEquity("SPY", self.Resolution).Symbol
                 
#         self.Schedule.On(self.DateRules.EveryDay("SPY"),
#                  self.TimeRules.BeforeMarketClose("SPY", 60),
#                  self.SellAll)
                 
#         self.todaysHalts = pd.DataFrame()
        
#         self.prepareTrading = []
#         self.maxPositions = 10
        
#     def TakeProfit(self):
#         activePositions = [x.Key for x in self.Portfolio if x.Value.Invested]
#         for symbol in activePositions:
#             if self.Portfolio[symbol].UnrealizedProfitPercent >= 0.05:
#                 self.Liquidate(symbol)
        
#     def OnData(self, data):
#         # self.TakeProfit()
#         if len(self.todaysHalts.index) == 0:
#             return
#         for tuple in self.todaysHalts.itertuples():
#             if tuple.Tickers not in list(self.activeSymbols.keys()):
#                 continue
            
#             if tuple.Tickers == "UNAM" or tuple.Tickers == "AGBAU":
#                 continue
            
#             symbol = self.activeSymbols[tuple.Tickers]
#             curHaltDate = str(tuple.StringHaltDate)
#             curHaltTime = str(tuple.HaltTime)
#             curResumeDate = str(tuple.StringResumedDate)
#             curResumeTime = str(tuple.ResumedTime)
            
#             try:
#                 curHalt = curHaltDate + " " + curHaltTime
#                 curHaltTime = datetime.strptime(curHalt, '%m/%d/%Y %H:%M:%S')
#                 curResume = curResumeDate + " " + curResumeTime
#                 curResumeTime = datetime.strptime(curResume, '%m/%d/%Y %H:%M:%S')
#             except:
#                 continue
#             haltDelta = self.Time - curHaltTime
#             resumeDelta = self.Time - curResumeTime
#             if haltDelta.total_seconds() >= 0 and resumeDelta.total_seconds() < 0 and tuple.HaltType in self.priorityTypes:# and data.ContainsKey(self.activeSymbols[tuple.Tickers]):
#                 self.HaltOccuring(symbol)
#                 self.symbolDataBySymbol[symbol].gatheredData = True
                
#             elif resumeDelta.total_seconds() >= (60*60)*2 and self.symbolDataBySymbol[symbol].gatheredData == True and tuple.HaltType in self.priorityTypes and not self.Portfolio[self.activeSymbols[tuple.Tickers]].Invested:# and data.ContainsKey(self.activeSymbols[tuple.Tickers]):
#                 self.AfterHaltTrades(symbol, resumeDelta)

#     def ConvertTimes(self):
#         self.HaltData["StringResumedDate"] = self.HaltData["ResumedDate"]
#         self.HaltData["ResumedDate"] = [datetime.strptime(x, '%m/%d/%Y') for x in self.HaltData["ResumedDate"]]
#         self.HaltData["StringHaltDate"] = self.HaltData["HaltDate"]
#         self.HaltData["HaltDate"] = [datetime.strptime(x, '%m/%d/%Y') for x in self.HaltData["HaltDate"]]
        
#     def CurrentHalts(self):
#         self.todaysHalts = self.HaltData.loc[self.HaltData['HaltDate'] == self.Time.date()]
        
#     def SellAll(self):
#         self.Liquidate()
        
#     def CoarseSelectionFunction(self, coarse):
#         self.CurrentHalts()
#         tickers = []
#         symbols = []
#         for tuple in self.todaysHalts.itertuples():
#             tickers.append(tuple.Tickers)
        
#         for x in coarse:
#             symbolTicker = x.Symbol.Value
#             if symbolTicker in tickers:
#                 symbols.append(x.Symbol)
        
#         return symbols
    
#     def FineSelectionFunction(self, fine):
#         return [x.Symbol for x in fine]
        
#     def OnSecuritiesChanged(self, changes):
#         for x in changes.AddedSecurities:
#             symbol = x.Symbol
#             self.Securities[symbol].SetSlippageModel(CustomSlippageModel(self))
#             self.Securities[symbol].SetDataNormalizationMode(DataNormalizationMode.Raw)
#             self.activeSymbols[symbol.Value] = symbol
#             symbolData = SymbolData(symbol)
#             self.symbolDataBySymbol[symbol] = symbolData
            
#         for x in changes.RemovedSecurities:
#             symbol = x.Symbol
#             self.Liquidate(symbol)
#             self.activeSymbols.pop(symbol.Value, None)
#             if symbol in self.symbolDataBySymbol:
#                 self.symbolDataBySymbol.pop(symbol, None)
            
# class CustomSlippageModel:
#     def __init__(self, algorithm):
#         self.algorithm = algorithm

#     def GetSlippageApproximation(self, asset, order):
#         slippage = asset.Price * float(0.0001 * np.log10(2*float(order.AbsoluteQuantity)))
#         return slippage
        
# class SymbolData:
#     def __init__(self, symbol):
#         self.Symbol = symbol
#         self.gatheredData = False
#         self.beforeHaltPriceMOMP = None
#         self.afterHaltPriceMOMP = None
#         self.beforeHaltSD = None
#         self.earlyIndicator = None