# region imports
from AlgorithmImports import *
# endregion

class BrainMLRankingDataAlgorithm(QCAlgorithm):
    def Initialize(self) -> None:
        self.SetStartDate(2021, 1, 1)
        self.SetEndDate(2021, 7, 8)
        self.SetCash(100000) 
        
        tickers = ["AAPL", "TSLA", "MSFT", "F", "KO"]
        self.symbol_by_dataset_symbol = {}
        for ticker in tickers:
            
            # Requesting data
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            dataset_symbol = self.AddData(BrainStockRanking2Day, symbol).Symbol
            self.symbol_by_dataset_symbol[dataset_symbol] = symbol
            
            # Historical data
            history = self.History(dataset_symbol, 365, Resolution.Daily)
            self.Debug(f"We got {len(history)} items from our history request for {symbol}")
        
    def OnData(self, slice: Slice) -> None:
        
        # Collect rankings for all symbols
        points = slice.Get(BrainStockRanking2Day)
        if points is None:
            return
        symbols = []
        ranks = []
        for point in points.Values:
            symbols.append(self.symbol_by_dataset_symbol[point.Symbol])
            ranks.append(point.Rank)
        
        # Rank each symbol's Brain ML ranking relative to each other
        if len(ranks) == 0:
            return
        ranks = [sorted(ranks).index(rank) + 1 for rank in ranks]
        
        # Place orders
        for i, rank in enumerate(ranks):
            weight = rank / sum(ranks)
            self.SetHoldings(symbols[i], weight)