from AlgorithmImports import *
from QuantConnect.DataSource import *
import numpy as np

class BrainMLRankingDataAlgorithm(QCAlgorithm):
    
    percentile = 2.5
    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))
        self.knownSymbols = set()
        
    def OnData(self, 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)
                
        ranks = [symbol_data.brain_rank for _, symbol_data in self.symbol_data_by_symbol.items() if symbol_data.IsReady]
        if len(ranks) == 0:
            return
        short_threshold = np.percentile(ranks, self.percentile)
        long_threshold = np.percentile(ranks, 100 - self.percentile)
        weight = self.percentile / len(self.symbol_data_by_symbol)
        
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            # Ensure we have data to place orders
            if not (data.ContainsKey(symbol) and data[symbol] is not None):
                continue

            if not symbol_data.IsReady:
                continue
            
            if symbol_data.brain_rank <= short_threshold:
                self.SetHoldings(symbol, -weight)
            elif symbol_data.brain_rank >= long_threshold:
                self.SetHoldings(symbol, weight)
            else:
                self.Liquidate(symbol)
            
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.symbol_data_by_symbol[symbol] = SymbolData(self, symbol)
            self.knownSymbols.add(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()
    
    def OnEndOfAlgorithm(self):
        self.Debug(f"Seen {len(self.knownSymbols)} symbols")
                
                
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
        for time, row in history.loc[self.ranking_symbol].iterrows():
            self.brain_rank = row['rank']
        
    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)