from AlgorithmImports import *
from QuantConnect.DataSource import *

class BrainSentimentDataAlgorithm(QCAlgorithm):
    def Initialize(self) -> None:
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2021, 7, 8)
        self.SetCash(100000) 
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(BrainSentimentIndicatorUniverse, "BrainSentimentIndicatorUniverse", Resolution.Daily, self.UniverseSelection)
        
        self.AddAlpha(BrainSentimentAlphaModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
         
        self.AddRiskManagement(NullRiskManagementModel())
        
        self.SetExecution(ImmediateExecutionModel())

    def UniverseSelection(self, alt_coarse: List[BrainSentimentIndicatorUniverse]) -> List[Symbol]:
        # define our selection criteria
        return [d.Symbol for d in alt_coarse \
                    if d.TotalArticleMentions7Days > 0]

        
class BrainSentimentAlphaModel(AlphaModel):
    
    symbol_data_by_symbol = {}

    def Update(self, algorithm: QCAlgorithm, slice: Slice) -> List[Insight]:
        insights = []
        
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if slice.ContainsKey(symbol_data.dataset_symbol):
                sentiment = slice[symbol_data.dataset_symbol].Sentiment
                symbol_data.update(sentiment)
               
            # Ensure we have security data in the current slice
            if not (slice.ContainsKey(symbol) and slice[symbol] is not None):
                continue
                
            if symbol_data.target_direction == InsightDirection.Up != algorithm.Portfolio[symbol].Invested:
                insights.append(Insight.Price(symbol, timedelta(days=100), symbol_data.target_direction))
        return insights
        
    def OnSecuritiesChanged(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.symbol_data_by_symbol[symbol] = SymbolData(algorithm, symbol)
        
        for security in changes.RemovedSecurities:
            symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
            if symbol_data:
                symbol_data.dispose()


class SymbolData:
    
    target_direction = InsightDirection.Flat
    _latest_sentiment_value = None
    
    def __init__(self, algorithm: QCAlgorithm, symbol: Symbol) -> None:
        self.algorithm = algorithm
        
        # Requesting data
        self.dataset_symbol = algorithm.AddData(BrainSentimentIndicator30Day, symbol).Symbol
        
        # Historical data
        history = algorithm.History(self.dataset_symbol, 100, Resolution.Daily)
        algorithm.Debug(f"We got {len(history)} items from our history request for {self.dataset_symbol}")
        if history.empty:
            return
        
        # Warm up historical sentiment values
        previous_sentiment_values = history.loc[self.dataset_symbol].sentiment.values
        for sentiment in previous_sentiment_values:
            self.update(sentiment)
        
    def dispose(self) -> None:
        # Unsubscribe from the Brain Sentiment feed for this security
        self.algorithm.RemoveSecurity(self.dataset_symbol)
        
    def update(self, sentiment: float) -> None:
        if self._latest_sentiment_value is not None:
            if sentiment > self._latest_sentiment_value:
                self.target_direction = InsightDirection.Up
            else:
                self.target_direction = InsightDirection.Flat
        self._latest_sentiment_value = sentiment