from AlgorithmImports import *
from QuantConnect.DataSource import *

class BrainMLRankingDataAlgorithm(QCAlgorithm):
    
    num_longs_shorts = 1
    symbol_data_by_symbol = {}
    
    def Initialize(self):
        self.SetStartDate(2020, 7, 1)
        self.SetEndDate(2021, 6, 30)
        self.SetCash(1000000)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.Universe.DollarVolume.Top(5))
        
    def OnData(self, data):
        # Update Symbol Data objects with the latest Brain Rank data
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if data.ContainsKey(symbol_data.ranking_symbol) and data[symbol_data.ranking_symbol] is not None:
                rank = data[symbol_data.ranking_symbol].Rank
                symbol_data.update(rank)
                
        # Select the Symbol Data objects that are ready to be traded
        ready_symbol_datas = []
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if symbol_data.IsReady and data.ContainsKey(symbol) and data[symbol] is not None:
                ready_symbol_datas.append((symbol, symbol_data))
        if len(ready_symbol_datas) < 2 * self.num_longs_shorts:
            return
        
        # Sort the SymbolData objects by their Brain Rank
        sorted_by_brain_rank = sorted(ready_symbol_datas, key=lambda x: x[1].brain_rank)
        
        # Place orders to form a long-short portfolio
        weight = 0.5 / self.num_longs_shorts
        for i, (symbol, symbol_data) in enumerate(sorted_by_brain_rank):
            if i < self.num_longs_shorts:
                self.SetHoldings(symbol, -weight)
            elif i >= len(sorted_by_brain_rank) - self.num_longs_shorts:
                self.SetHoldings(symbol, weight)
            elif self.Portfolio[symbol].Invested:
                self.Liquidate(symbol)
            
    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 SymbolData:
    
    brain_rank = None
    
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.ranking_symbol = algorithm.AddData(BrainStockRanking2Day, symbol).Symbol
        
        # Warm up Brain rank
        history = algorithm.History(self.ranking_symbol, 3, Resolution.Daily)
        if history.empty or 'rank' not in history.columns:
            return
        self.brain_rank = history.loc[self.ranking_symbol]['rank'].values[-1]
        
    def update(self, rank):
        self.brain_rank = rank
        
    @property
    def IsReady(self):
        return self.brain_rank is not None
        
    def dispose(self):
        self.algorithm.RemoveSecurity(self.ranking_symbol)