from AlgorithmImports import *
from QuantConnect.DataSource import *

class BrainMLRankingDataAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2021, 7, 8)
        self.set_cash(100000) 
        
        # We cherry picked 5 largest stocks, high trading volume provides better information and credibility for ML ranking
        tickers = ["AAPL", "TSLA", "MSFT", "F", "KO"]
        self.symbol_by_dataset_symbol = {}
        for ticker in tickers:
            # Requesting data to get 2 days estimated relative ranking
            symbol = self.add_equity(ticker, Resolution.DAILY).symbol
            dataset_symbol = self.add_data(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 on_data(self, slice: Slice) -> None:
        # Collect rankings for all selected symbols for ranking them
        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 for positional sizing
        if len(ranks) == 0:
            return
        ranks = [sorted(ranks).index(rank) + 1 for rank in ranks]
        
        # Place orders according to the ML ranking, the better the rank, the higher the estimated return and hence weight
        for i, rank in enumerate(ranks):
            weight = rank / sum(ranks)
            self.set_holdings(symbols[i], weight)

using QuantConnect.DataSource;

namespace QuantConnect
{
    public class BrainMLRankingDataAlgorithm : QCAlgorithm
    {
        private Dictionary<Symbol, Symbol> _symbolByDatasetSymbol = new Dictionary<Symbol, Symbol>();
        
        public override void Initialize()
        {
            SetStartDate(2021, 1, 1);
            SetEndDate(2021, 7, 8);
            SetCash(100000);
            
            // We cherry picked 5 largest stocks, high trading volume provides better information and credibility for ML ranking
            var tickers = new List<string>() {"AAPL", "TSLA", "MSFT", "F", "KO"};
            foreach (var ticker in tickers)
            {
                // Requesting data to get 2 days estimated relative ranking
                var symbol = AddEquity(ticker, Resolution.Daily).Symbol;
                var datasetSymbol = AddData<BrainStockRanking2Day>(symbol).Symbol;
                _symbolByDatasetSymbol.Add(datasetSymbol, symbol);
                
                // Historical data
                var history = History<BrainStockRanking2Day>(datasetSymbol, 365, Resolution.Daily);
                Debug($"We got {history.Count()} items from our history request for {symbol}");
            }
        }

        public override void OnData(Slice slice)
        {
            // Collect rankings for all symbols for ranking them
            var points = slice.Get<BrainStockRanking2Day>();
            if (points == null)
            {
                return;
            }
            var symbols = new List<Symbol>();
            var ranks = new List<decimal>();
            foreach (var point in points.Values)
            {
                symbols.Add(_symbolByDatasetSymbol[point.Symbol]);
                ranks.Add(point.Rank);
            }
            
            // Rank each symbol's Brain ML ranking relative to the other symbols for positional sizing
            if (ranks.Count() == 0) return;
            var sortedRanksTemp = new List<decimal>(ranks); 
            sortedRanksTemp.Sort();
            var sortedRanks = new List<decimal>();
            for (var i = 0; i < symbols.Count(); i++)
            {
                sortedRanks.Add(sortedRanksTemp.IndexOf(ranks[i]) + 1);
            }
            
            // Place orders, the better the rank, the higher the estimated return and hence weight
            for (var i = 0; i < symbols.Count(); i++)
            {
                var rank = sortedRanks[i];
                var weight = rank / sortedRanks.Sum();
                SetHoldings(symbols[i], weight);
            }
        }
    }
}

from AlgorithmImports import *
from QuantConnect.DataSource import *

class BrainMLRankingDataAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2021, 7, 8)
        self.set_cash(100000) 
        
        self.universe_settings.resolution = Resolution.DAILY
        # Select based on ML ranking data
        self.add_universe(BrainStockRankingUniverse, self.universe_selection)
        
        self.add_alpha(BrainMLRankingAlphaModel())
        
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel())
         
        self.add_risk_management(NullRiskManagementModel())
        
        self.set_execution(ImmediateExecutionModel())

    def universe_selection(self, alt_coarse: List[BrainStockRankingUniverse]) -> List[Symbol]:
        # Select the ones that expected to out-perform the median of the whole market in 2-3 days
        return [d.symbol for d in alt_coarse \
                if d.Rank2Days > 0.2  \
                and d.Rank3Days > 0.2]

class BrainMLRankingAlphaModel(AlphaModel):
    
    symbol_data_by_symbol = {}
    symbol_by_dataset_symbol = {}
    
    def update(self, algorithm: QCAlgorithm, slice: Slice) -> List[Insight]:
        insights = []
        
        # Collect rankings for all selected symbols for ranking them
        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 for positional sizing
        if len(ranks) == 0:
            return []
        ranks = [sorted(ranks).index(rank) + 1 for rank in ranks]
        
        # Place orders according to the ML ranking, the better the rank, the higher the estimated return and hence weight
        for i, rank in enumerate(ranks):
            weight = rank / sum(ranks)
            insights.append(Insight.price(symbols[i], timedelta(days=7), InsightDirection.UP, None, None, None, weight))
        
        return insights
        
    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            symbol = security.symbol
            symbol_data = SymbolData(algorithm, symbol)
            self.symbol_data_by_symbol[symbol] = symbol_data
            self.symbol_by_dataset_symbol[symbol_data.dataset_symbol] = symbol
        
        for security in changes.removed_securities:
            symbol_data = self.symbol_data_by_symbol.pop(security.symbol, None)
            if symbol_data:
                symbol_data.dispose()
                
            for dataset_symbol, symbol in self.symbol_by_dataset_symbol.items():
                if symbol == security.symbol:
                    self.symbol_by_dataset_symbol.pop(dataset_symbol)
                    break


class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        
        # Requesting data to get 2 days estimated relative ranking
        self.dataset_symbol = algorithm.add_data(BrainStockRanking2Day, symbol).symbol
        
        # Historical data
        history = algorithm.history(self.dataset_symbol, 365, Resolution.DAILY)
        algorithm.debug(f"We got {len(history)} items from our history request for {symbol}")
        
    def dispose(self):
        # Unsubscribe from the Brain ML Ranking feed for this security to release computation resources
        self.algorithm.remove_security(self.dataset_symbol)

using QuantConnect.DataSource;

namespace QuantConnect
{
    public class BrainMLRankingDataAlgorithm : QCAlgorithm
    {
        public override void Initialize()
        {
            SetStartDate(2021, 1, 1);
            SetEndDate(2021, 7, 8);
            SetCash(100000);
            
            // Select based on ML ranking data
            AddUniverse<BrainStockRankingUniverse>(altCoarse =>
            {
                // Select the ones that expected to out-perform the median of the whole market in 2-3 days
                return from d in altCoarse.OfType<BrainStockRankingUniverse>() 
                    where d.Rank2Days > 0.2m && d.Rank3Days > 0.2m
                    select d.Symbol;
            });
            
            AddAlpha(new BrainMLRankingAlphaModel());
            
            SetPortfolioConstruction(new InsightWeightingPortfolioConstructionModel());
            
            AddRiskManagement(new NullRiskManagementModel());
            
            SetExecution(new ImmediateExecutionModel());
        }
    }
        
        
    public class BrainMLRankingAlphaModel : AlphaModel
    {
        private Dictionary<Symbol, SymbolData> _symbolDataBySymbol = new Dictionary<Symbol, SymbolData>();
        private Dictionary<Symbol, Symbol> _symbolByDatasetSymbol = new Dictionary<Symbol, Symbol>();

        public override IEnumerable<Insight> Update(QCAlgorithm algorithm, Slice slice)
        {
            var insights = new List<Insight>();
            
            // Collect rankings for all selected symbols for ranking them
            var points = slice.Get<BrainStockRanking2Day>();
            if (points == null)
            {
                return insights;
            }
            var symbols = new List<Symbol>();
            var ranks = new List<decimal>();
            foreach (var point in points.Values)
            {
                symbols.Add(_symbolByDatasetSymbol[point.Symbol]);
                ranks.Add(point.Rank);
            }
            
            // Rank each symbol's Brain ML ranking relative to the other symbols for positional sizing
            if (ranks.Count() == 0) 
            {
                return insights; 
            }
            var sortedRanksTemp = new List<decimal>(ranks); 
            sortedRanksTemp.Sort();
            var sortedRanks = new List<decimal>();
            for (var i = 0; i < symbols.Count(); i++)
            {
                sortedRanks.Add(sortedRanksTemp.IndexOf(ranks[i]) + 1);
            }
            
            // Place orders according to the ML ranking, the better the rank, the higher the estimated return and hence weight
            for (var i = 0; i < symbols.Count(); i++)
            {
                var rank = sortedRanks[i];
                var weight = rank / sortedRanks.Sum();
                insights.Add(Insight.Price(symbols[i], TimeSpan.FromDays(7), InsightDirection.Up, null, null, null, (double)weight));
            }
            
            return insights;
        }

        public override void OnSecuritiesChanged(QCAlgorithm algorithm, SecurityChanges changes)
        {
            foreach (var security in changes.AddedSecurities)
            {
                var symbol = security.Symbol;
                var symbolData = new SymbolData(algorithm, symbol);
                _symbolDataBySymbol.Add(symbol, symbolData);
                _symbolByDatasetSymbol.Add(symbolData.datasetSymbol, symbol);
            }

            foreach (var security in changes.RemovedSecurities)
            {
                var symbol = security.Symbol;
                if (_symbolDataBySymbol.ContainsKey(symbol))
                {
                    _symbolDataBySymbol[symbol].dispose();
                    _symbolDataBySymbol.Remove(symbol);
                }
                
                foreach (var entry in _symbolByDatasetSymbol)
                {
                    if (entry.Value == symbol)
                    {
                        _symbolByDatasetSymbol.Remove(entry.Key);
                    }
                }
            }
        }
    }
    
    
    public class SymbolData
    {
        private QCAlgorithm _algorithm;
        public Symbol datasetSymbol;

        public SymbolData(QCAlgorithm algorithm, Symbol symbol)
        {
            _algorithm = algorithm;
            
            // Requesting data to get 2 days estimated relative ranking
            datasetSymbol = algorithm.AddData<BrainStockRanking2Day>(symbol).Symbol;
            
            // Historical data
            var history = algorithm.History<BrainStockRanking2Day>(datasetSymbol, 365, Resolution.Daily);
            algorithm.Debug($"We got {history.Count()} items from our history request for {symbol}");
        }
        
        public void dispose()
        {
            // Unsubscribe from the Brain ML Ranking feed for this security to release computation resources
            _algorithm.RemoveSecurity(datasetSymbol);
        }
    }
}

#r "../QuantConnect.DataSource.BrainSentiment.dll"
using QuantConnect.DataSource;

var qb = new QuantBook();

// Requesting data
var aapl = qb.AddEquity("AAPL", Resolution.Daily).Symbol;
var symbol = qb.AddData<BrainStockRanking2Day>(aapl).Symbol;

// Historical data
var history = qb.History<BrainStockRanking2Day>(symbol, 30, Resolution.Daily);
foreach (BrainStockRanking2Day rank in history)
{
    Console.WriteLine($"{rank} at {rank.EndTime}");
}

// Add Universe Selection
IEnumerable<Symbol> UniverseSelection(IEnumerable<BaseData> altCoarse)
{
    return (from d in altCoarse.OfType<BrainStockRankingUniverse>()
        orderby d.Rank2Days descending select d.Symbol).Take(10);
}
var universe = qb.AddUniverse<BrainStockRankingUniverse>(UniverseSelection);

// Historical Universe data
var universeHistory = qb.UniverseHistory(universe, qb.Time.AddDays(-5), qb.Time);
foreach (var ranks in universeHistory)
{
    foreach (BrainStockRankingUniverse rank in ranks)
    {
        Console.WriteLine($"{rank.Symbol} 2-day rank at {rank.EndTime}: {rank.Rank2Days}");
    }
}