import configs as cfg

class FatApricotElephant(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2010, 2, 4)
        self.SetEndDate(2015, 2, 4)
        self.SetCash(100000) 
        self.AddUniverseSelection(QC500UniverseSelectionModel())
        self.UniverseSettings.Resolution = cfg.resolution
        self.symbolDataDict = {}
        self.next_rebalance = self.Time
        
        self.divisor = {
            Resolution.Daily: 1,
            Resolution.Hour: 24,
            Resolution.Minute: 24 * 60,
            Resolution.Second: 24 * 60 * 60
        }
        
        self.mkt = self.AddEquity(cfg.mkt, cfg.resolution).Symbol
        self.mkt_sma = self.SMA(self.mkt, cfg.mkt_sma_period, cfg.resolution)
        self.SetWarmUp(cfg.warmup_period)
        
    def OnData(self, data):
        if self.next_rebalance >= self.Time or self.IsWarmingUp:
            return
        
        self.Log('Beginning rebalance logic')
        
        if not self.mkt_sma.IsReady:
            self.Log('Market is bearished, ejecting')
            return
        
        self.next_rebalance = self.Time + timedelta(cfg.rebalance_period / self.divisor[cfg.resolution])
        
        if self.GetPrice(self.mkt) < self.mkt_sma.Current.Value:
            self.Liquidate()
            return
        
        filtered_rdy = [sd for sd in self.symbolDataDict.values() if sd.IsReady]
        topROC = sorted(filtered_rdy, key=lambda sd: sd.get_roc(), reverse=True)
        filtered_sma = [sd for sd in topROC if sd.is_buy(self.GetPrice(sd.symbol))][:cfg.num_stocks]
        weights = {sd.symbol:self.GetPrice(sd.symbol)/sd.get_atr() for sd in filtered_sma}
        total_weight = sum(weights.values())
        adj_weights = {symbol: weight/total_weight for symbol, weight in weights.items()}
        
        invested = [kvp.Key for kvp in self.Portfolio if kvp.Value.Invested]
        to_buy = adj_weights.keys()
        for symbol in invested:
            if symbol not in to_buy:
                self.Liquidate(symbol)
        
        if len(adj_weights) > 0:
            self.Log('Going LONG')
        
        if cfg.enable_atr:
            for symbol, weight in adj_weights.items():
                self.SetHoldings(symbol, weight)
        else:
            weight = 1/len(adj_weights)
            for symbol in adj_weights:
                self.SetHoldings(symbol, weight)
        
    def OnSecuritiesChanged(self, changed):
        for security in changed.AddedSecurities:
            if security.Symbol == self.mkt:
                continue
            self.symbolDataDict[security.Symbol] = SymbolData(self, security.Symbol)
        for security in changed.RemovedSecurities:
            self.symbolDataDict.pop(security.Symbol, None)
    
    def GetPrice(self, symbol):
        return self.Securities[symbol].Price
    
class SymbolData:
    def __init__(self, algo, symbol):
        self.symbol = symbol
        self.roc = algo.ROCP(symbol, cfg.roc_period, cfg.resolution)
        self.sma = algo.SMA(symbol, cfg.sma_period, cfg.resolution)
        self.atr = algo.ATR(symbol, cfg.atr_period, cfg.resolution)
        
        hist = algo.History(symbol, cfg.warmup_period, cfg.resolution).loc[symbol]
        
        if hist.empty or len(hist)==0:
            return
        
        for time, row in hist.iterrows():
            tradeBar = TradeBar(time, symbol, row.open, row.high, row.low, row.close, row.volume, timedelta(1))
            self.atr.Update(tradeBar)
            self.roc.Update(time, row.close)
            self.sma.Update(time, row.close)
    
    @property
    def IsReady(self):
        return self.atr.IsReady and self.roc.IsReady and self.sma.IsReady
    
    def get_roc(self):
        return self.roc.Current.Value
        
    def get_sma(self):
        return self.sma.Current.Value
        
    def get_atr(self):
        return self.atr.Current.Value
        
    def is_buy(self, price):
        return price > self.get_sma()

# ---configurables that won't be changed often---

# resolution for data/indicators
resolution = Resolution.Daily 

# wait how many periods of above resolution to rebalance
rebalance_period = 30

# market proxy - used to turn on/off bull
mkt = 'SPY'


# ---configurables you might want to change---

# SMA period of the market proxy chosen
mkt_sma_period = 100

# ROC period for sorting
roc_period = 30

# SMA period of individual stocks to determine uptrend
sma_period = 30

# ATR for weighting, where lower ATR means higher weight
atr_period = 14

# if enabled, use inverse adj ATR weight
# if disabled, use equal weighting
enable_atr = True

# how many stocks you want
# used to choose how many of the top ROC stocks
num_stocks = 20

# probably don't touch this
warmup_period = max(mkt_sma_period, roc_period, sma_period, atr_period)