from AlgorithmImports import *

class CryingFluorescentYellowHamster(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 6, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
    
        self.SetSecurityInitializer(
        lambda x: x.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted) \
                and x.SetLeverage(1) and x.SetFillModel(CustomFillModel(self)))
                
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.SplitAdjusted
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.Leverage = 1
        self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(self.SelectCoarse, self.SelectFine))
        
        ''' Universe Metadata '''
        self.coarseMinPrice = 4
        self.coarseMaxPrice = 10
        self.coarseMaxSymbols = 30
        
        '''Init'''
        self.uniData = {}
        self.activeStocks = {}
        self.maxPeriod = 250  # max history lookback

    def SelectCoarse(self, coarse):

        for c in coarse:
            if c.Symbol not in self.uniData:
                self.uniData[c.Symbol] = UniverseData(c.Symbol, c.HasFundamentalData)
                
            # c.Volume/c.SplitFactor - Adjust volume for splits
            self.uniData[c.Symbol].update(c.Price, c.Volume / c.SplitFactor)

        # Filter the values of the dict
        values = [x for x in self.uniData.values() if x.Price >= self.coarseMinPrice
                  and x.hasFundamental
                  and x.Price <= self.coarseMaxPrice]

        # sort by the largest in volume.
        values = sorted(values, key=lambda x: x.Price, reverse=True)[:self.coarseMaxSymbols]

            # we need to return only the symbol objects
        return [x.Symbol for x in values]

    def SelectFine(self, fine):
        fineSymbols = [x for x in fine if x.CompanyReference.PrimaryExchangeID in ["NYS", "NAS", "ASE"]]

        # we need to return only the symbol objects
        return [x.Symbol for x in fineSymbols]
        
        
    def OnSecuritiesChanged(self, changes):
        self.changes = changes

        # load history for new symbols in uni
        for security in self.changes.AddedSecurities:
            if not security.Symbol.Value in self.activeStocks:
                self.LoadHistory(security)
                
        for security in self.changes.RemovedSecurities:
            if security.Symbol.Value in self.activeStocks and not security.Invested:
                del self.activeStocks[security.Symbol.Value]
                
    def LoadHistory(self, security):

        sd = SymbolData(security)

        sd.avgVol = self.SMA(security.Symbol, 30, Resolution.Daily, Field.Volume)

        sd.close = self.Identity(security.Symbol, Resolution.Daily, Field.Close)
        sd.close.Updated += sd.CloseUpdated
        sd.closeWindow = RollingWindow[IndicatorDataPoint](self.maxPeriod)

        sd.emaTwentyOne = self.EMA(security.Symbol, 21, Resolution.Daily, Field.Close)
        sd.emaTwentyOne.Updated += sd.EmaTwentyOneUpdated
        sd.emaTwentyOneWindow = RollingWindow[IndicatorDataPoint](self.maxPeriod)

        # warmup our indicators by pushing history through the indicators
        # rolling window must use warmup !IMPORTANT
        history = self.History(security.Symbol, self.maxPeriod, Resolution.Daily)

        if (history.empty or history.size < self.maxPeriod):
            return

        for index, row in history.loc[security.Symbol].iterrows():
            if index == self.Time:
                continue

            sd.avgVol.Update(index, row['volume'])
            sd.close.Update(index, row['close'])
            sd.emaTwentyOne.Update(index, row['close'])

        # add sd to active dict with all history
        self.activeStocks[security.Symbol.Value] = sd
        
    def IsReady(self, sd):

        if (sd.closeWindow.Count < self.maxPeriod or
                sd.emaTwentyOneWindow.Count < self.maxPeriod):
            return False

        if not (sd.closeWindow.IsReady and
                sd.emaTwentyOneWindow.IsReady):
            return False

        return True


    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # sort by avg volume
        self.activeStocks = dict(
            sorted(self.activeStocks.items(), key=lambda item: item[1].avgVol.Current.Value, reverse=True))

        for key, sd in self.activeStocks.items():

            if sd.Security.Fundamentals is None:
                continue

            if not (data.ContainsKey(sd.Symbol) and data.Bars.ContainsKey(sd.Symbol)):
                continue

            if not self.IsReady(sd):
                continue

            if not sd.Security.Invested:

                if sd.closeWindow[0].Value < sd.emaTwentyOneWindow[0].Value:
                    continue

                self.SetHoldings(sd.Symbol, 0.1)

            elif sd.Security.Invested and sd.closeWindow[0].Value > sd.emaTwentyOneWindow[0].Value:
                    self.Liquidate(sd.Symbol)


class CustomFillModel(FillModel):
    '''
    Implements a custom fill model that inherit from FillModel.
    Override the MarketFill method to simulate partially fill orders
    '''

    def __init__(self, algorithm, sd=None):
        self.algorithm = algorithm
        self.absoluteRemainingByOrderId = {}
        self.sd = sd

    def MarketFill(self, asset, order):
        absoluteRemaining = order.AbsoluteQuantity

        if order.Id in self.absoluteRemainingByOrderId.keys():
            absoluteRemaining = self.absoluteRemainingByOrderId[order.Id]

        fill = super().MarketFill(asset, order)
        absoluteFillQuantity = int(min(absoluteRemaining, order.AbsoluteQuantity))
        fill.FillQuantity = np.sign(order.Quantity) * absoluteFillQuantity
        
        if absoluteRemaining == absoluteFillQuantity:
            fill.Status = OrderStatus.Filled
            if self.absoluteRemainingByOrderId.get(order.Id):
                self.absoluteRemainingByOrderId.pop(order.Id)
        else:
            absoluteRemaining = absoluteRemaining - absoluteFillQuantity
            self.absoluteRemainingByOrderId[order.Id] = absoluteRemaining
            fill.Status = OrderStatus.PartiallyFilled
            
        self.algorithm.Debug(f"CustomFillModel: {fill}")
        return fill

class UniverseData:

    def __init__(self, symbol, fundamental):
        self.Symbol = symbol
        self.Volume = 0
        self.Price = 0
        self.hasFundamental = fundamental

    def update(self, price, volume):
        self.Volume = volume
        self.Price = price
        
        
class SymbolData:

    def __init__(self, security):
        self.Security = security
        self.Symbol = security.Symbol
        self.avgVol = 0

        ''' Candle '''
        self.close = None
        self.closeWindow = None

        ''' Moving Averages '''
        self.emaTwentyOne = None
        self.emaTwentyOneWindow = None

    def CloseUpdated(self, sender, updated):
        self.closeWindow.Add(updated)

    def EmaTwentyOneUpdated(self, sender, updated):
        self.emaTwentyOneWindow.Add(updated)