from AlgorithmImports import *
from QuantConnect.DataSource import *

class QuiverWallStreetBetsDataAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetEndDate(2021, 7, 1)
        self.SetCash(100000)
        
        self.universe_data_by_symbol = {}
        self.symbol_data_by_symbol = {}
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)

    def CoarseSelectionFunction(self, coarse):
        
        sorted_by_dollar_volume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)
        selected = sorted_by_dollar_volume[:500]
        
        for c in selected:
            if c.Symbol not in self.universe_data_by_symbol:
                self.universe_data_by_symbol[c.Symbol] = UniverseData(self, c.Symbol)
            else:
                self.universe_data_by_symbol[c.Symbol].Update(self.Time, c.AdjustedPrice)
                
        selected = [c for c in selected if self.universe_data_by_symbol[c.Symbol].IsReady]
        sorted_by_std = sorted(selected, key=lambda c: self.universe_data_by_symbol[c.Symbol].std.Current.Value, reverse=True)
        return [ x.Symbol for x in sorted_by_std[:3] ]
            

    def OnData(self, data):
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            # Gather WSB mentions data
            if data.ContainsKey(symbol_data.wsb_symbol):
                wsb_mentions = data[symbol_data.wsb_symbol].Mentions
                symbol_data.Update(data.Time, wsb_mentions)
        
            # Ensure we have data to place orders
            if not (data.ContainsKey(symbol) and data[symbol] is not None):
                continue
            # Place orders
            if not data[symbol].IsFillForward and symbol_data.IsReady:
                target_holding = symbol_data.get_target_holding() / len(self.symbol_data_by_symbol)
                self.SetHoldings(symbol, target_holding)
                
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.symbol_data_by_symbol[symbol] = SymbolData(self, symbol)
            
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            self.Liquidate(symbol)
            symbol_data = self.symbol_data_by_symbol.pop(symbol, None)
            if symbol_data:
                symbol_data.dispose()
            
class UniverseData:
    def __init__(self, algorithm, symbol, std_period=30):
        self.std = StandardDeviation(std_period)
        
        # Warm up std indicator
        history = algorithm.History(symbol, std_period, Resolution.Daily)
        if history.empty or 'close' not in history.columns:
            return
        for time, row in history.loc[symbol].iterrows():
            self.std.Update(time, row.close)
    
    def Update(self, date, price):
        self.std.Update(date, price)
        
    @property
    def IsReady(self):
        return self.std.IsReady
        
            
class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        
        self.wsb_symbol = algorithm.AddData(QuiverWallStreetBets, symbol).Symbol
        
        self.max = Maximum(7)
        self.current = Maximum(1)
        
        # Warm up max indicator
        history = algorithm.History(QuiverWallStreetBets, self.wsb_symbol, 14, Resolution.Daily)
        if history.empty or 'mentions' not in history.columns:
            return
        
        for time, row in history.loc[self.wsb_symbol].iterrows():
            self.max.Update(time, row.mentions)
            self.current.Update(time, row.mentions)
        
    def Update(self, time, mentions):
        self.max.Update(time, mentions)
        self.current.Update(time, mentions)
        
    @property
    def IsReady(self):
        return self.max.IsReady
        
    def get_target_holding(self):
        return int(self.current >= self.max)
        
    def dispose(self):
        self.algorithm.RemoveSecurity(self.wsb_symbol)