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(2021, 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 >= 20:
                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 HaltOccuring(self, symbol):
        history = self.History(symbol, 20, Resolution.Hour)
        mompClose = MomentumPercent(20)
        standardDeviation = history.loc[symbol]["close"].std()
        for historyTuple in history.loc[symbol].itertuples():
            mompClose.Update(historyTuple.Index, historyTuple.close)
        self.symbolDataBySymbol[symbol].beforeHaltPriceMOMP = mompClose.Current.Value
        self.symbolDataBySymbol[symbol].beforeHaltSD = standardDeviation

    def AfterHaltTrades(self, symbol, delta):
        history = self.History(symbol, 3, Resolution.Hour)
        closes = history.loc[symbol]["close"]
        self.symbolDataBySymbol[symbol].afterHaltPriceMOMP = (closes[-1] - closes[-3]) / closes[-3]
        if self.symbolDataBySymbol[symbol].beforeHaltPriceMOMP == None:
            return
        if (self.symbolDataBySymbol[symbol].beforeHaltPriceMOMP >= 0 and self.symbolDataBySymbol[symbol].afterHaltPriceMOMP) >= 0 or (self.symbolDataBySymbol[symbol].beforeHaltPriceMOMP < 0 and self.symbolDataBySymbol[symbol].afterHaltPriceMOMP < 0):
            self.symbolDataBySymbol[symbol].earlyIndicator = True
        else:
            self.symbolDataBySymbol[symbol].earlyIndicator = False
            
        direction = -1
        if self.symbolDataBySymbol[symbol].beforeHaltPriceMOMP >= 0:
            direction = 1
        
        tempDict = {}
        tempDict[symbol] = [self.symbolDataBySymbol[symbol].beforeHaltSD, self.symbolDataBySymbol[symbol].earlyIndicator]
        testDf = pd.DataFrame.from_dict(tempDict)

        prediction = self.haltModel.predict(testDf.T)
        positionSize = (1)/len(self.activeSymbols.keys())
        if prediction == 1:
            self.SetHoldings(symbol, positionSize)
        else:
            self.SetHoldings(symbol, -1/positionSize)

    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