from datetime import timedelta, time
import math

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
Symbol Data

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
class SymbolData(object):
    
    def __init__(self, algorithm, symbol, criteriaDict):
        
        self.symbol = symbol
        self.symbolValue = symbol.Value
        
        self.algorithm = algorithm
        self.criteriaDict = criteriaDict
        
        #-- Trailing Stop
        self.lastPeriodMax = 0
        self.trailStopOrderTicket = None
        self.trailStopPrice = None


class ErrorSubmission_TrailStop(QCAlgorithm):
    
    
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    Iniitalization
    
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""    
    def Initialize(self):
        
        #-- Date Range
        self.SetStartDate(2020, 10, 1) # Set Start Date
        
        #-- Starting Equity
        self.SetCash(1000000)  # Set Strategy Cash # check if will be replaced

        #---------- Universe Selection -------------
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.ExtendedMarketHours = True  # test live if work
        self.maxSymbolsFromUniverseFilter = 30
        
        self.handpicked_SymbolList = [ 
                    Symbol.Create("SPY", SecurityType.Equity, Market.USA),
                    Symbol.Create("DNMR", SecurityType.Equity, Market.USA),
                    Symbol.Create("PLTR", SecurityType.Equity, Market.USA),
                    Symbol.Create("GME", SecurityType.Equity, Market.USA),
                    ]

        self.AddUniverse(self.CoarseSelectionFilter)


        #---------- Alpha Generation -------------

        #-- Criteria Settings
        self.criteriaDict = {
        }


        #---------- Portfolio Construction -------------
        
        self.maxSymbolsToTrade = 10  # maxSymbols to trade
        self.symbolDict = {} # list which stores symbols (refreshes with universe updates + handpick selection)
        self.symbolWithPositionDict = {}  # list which stores symbols with position
        self.potentialTradeDict = {}  # temp list to store potential symbols to trade
        
        for symbol in self.handpicked_SymbolList:
            if symbol not in self.symbolDict:
                self.AddEquity(symbol.Value, Resolution.Minute)
                self.symbolDict[symbol] = SymbolData(self,symbol,self.criteriaDict)

        
        #---------- Risk Management -------------
        self.useTrailStop = True
        self.trailStopPercentage = 0.25
        
        
        #---------- Schedule to run calculation -------------
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(TimeSpan.FromMinutes(2)), self.runAlgo)



    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    # Coarse Selection Filter
       
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    def CoarseSelectionFilter(self, coarse):
        
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        finalFilter = sortedByDollarVolume[:5]

        for coarseObject in finalFilter:
            symbol = coarseObject.Symbol
        
        return [x.Symbol for x in finalFilter]
        

    
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    runAlgo - Main execution function for the system
        - function is triggered every x minute to check for signal and execute trades
    
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""    
    def runAlgo(self):
        
        #------------------------ Exit Signals ------------------------
        #-- Trailing Stop
        if self.useTrailStop:

            for symbol in self.symbolWithPositionDict:            
                
                symbolData = self.symbolDict[symbol]
                symbolPrice = self.Securities[symbol].Close
        
                # Trailstop order not yet set
                if symbolData.trailStopOrderTicket is None:
                    # self.Debug(symbol.Value + "  set trailstop")
                    if self.Portfolio[symbol].Invested:
                        quantity = -self.Portfolio[symbol].Quantity
                        price = (1 - self.trailStopPercentage) * symbolPrice
                        symbolData.trailStopOrderTicket = self.StopMarketOrder(symbol.Value, quantity, price)  
                        self.Debug(f"Order for {symbol} quantity {quantity} at price {price}; status: {symbolData.trailStopOrderTicket.Status}")
                    
                # Trailstop order exist, adjust it if price moves higher
                else:
                    if symbolPrice > symbolData.lastPeriodMax:
                        symbolData.lastPeriodMax = symbolPrice
                        updateFields = UpdateOrderFields()
                        updateFields.StopPrice = (1 - self.trailStopPercentage) * symbolPrice
                        self.Debug(f"Stop: {updateFields.StopPrice}; OrderTicket Status: {symbolData.trailStopOrderTicket.Status == OrderStatus.Invalid}")
                        self.Quit()
                        return
                        #self.Debug(symbol.Value + " moving symbolLastPeriodMax up " + str(symbolPrice) + "   "  + str(updateFields.StopPrice))
                        symbolData.trailStopOrderTicket.Update(updateFields) 
                        return
        
            
        #------------------------ Entry Signals ------------------------
        self.potentialTradeDict.clear()  # reset potentialTradeDict
        
        # for symbol in self.Portfolio.Keys:
        for symbol in self.symbolDict:
            if symbol not in self.symbolWithPositionDict:  # run check only if symbol has no position
                currentClose = self.Securities[symbol].Close

                finalEntrySignal = 1
    
                #-- append symbol to potential tradelist
                if finalEntrySignal > 0:
                    self.potentialTradeDict[symbol] = finalEntrySignal


        #------------------------ Portfolio Construction + Execution ------------------------

        #-- remove symbols from potentialTradeDict that already has a position
        items_to_be_removed = []
        for symbol in self.potentialTradeDict:
            if symbol in self.symbolWithPositionDict.keys():
                items_to_be_removed.append(symbol)
        
        for symbol in items_to_be_removed:
                del self.potentialTradeDict[symbol] 

        #-- Total symbols (currentPosition + newly shortlisted) < maxSymbolsToTrade
        if len(self.symbolWithPositionDict) + len(self.potentialTradeDict) <= self.maxSymbolsToTrade:
            # trade all symbols in the potentialTradeDict and add them to the symbolInPositionList
            for symbol, signal in self.potentialTradeDict.items():
                # calculate posSize
                posSize = self.getPosSize(symbol) 
            
                if(posSize != 0):
                    tag = "Sig: {:.2f}".format(signal)
                    ticket = self.MarketOrder(symbol, posSize,False,tag)
                    # add to positionDict
                    
                    self.symbolWithPositionDict[symbol] = posSize
                    
            
        else:
            pass
        


    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    # getPosSize

    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    def getPosSize(self, symbol):

        dollarAllocation = self.Portfolio.TotalPortfolioValue//self.maxSymbolsToTrade  

        symbolPrice = self.Securities[symbol].Price
        posSize = math.floor(dollarAllocation/symbolPrice) if symbolPrice>0 else 0

        return posSize
            
        
        
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    OnSecuritiesChanged
       - this event fires whenever we have changes to our universe
       
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    def OnSecuritiesChanged(self, changes):
        # self.Log("OnSecuritiesChanged: " + str(self.Time))
        
        for security in changes.RemovedSecurities:
            if ((security.Symbol in self.symbolDict) and (security.Symbol not in self.handpicked_SymbolList) and (security.Symbol not in self.symbolWithPositionDict)):
                # self.Log("Removing " + security.Symbol.Value + " from symbolDict")
                del self.symbolDict[security.Symbol] 

        for security in changes.AddedSecurities:
            if security.Symbol not in self.symbolDict:
                # self.Log("Adding " + security.Symbol.Value + " to symbolDict")
                self.AddEquity(security.Symbol.Value, Resolution.Minute)
                self.symbolDict[security.Symbol] = SymbolData(self,security.Symbol,self.criteriaDict)



    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    OnData
    
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""    
    def OnData(self, data):
        pass