from AlgorithmImports import *


class TechnologyPremarketNewsAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.set_universe_selection(TechnologyETFUniverse())
        self.set_alpha(NewsSentimentAlphaModel())
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel())
        self.set_risk_management(NullRiskManagementModel())
        
        
class NewsSentimentAlphaModel(AlphaModel):
    
    def __init__(self):
        self._word_sentiment = {
            "bad": -0.5, "good": 0.5, "negative": -0.5, 
            "great": 0.5, "growth": 0.5, "fail": -0.5, 
            "failed": -0.5, "success": 0.5, "nailed": 0.5,
            "beat": 0.5, "missed": -0.5, "profitable": 0.5,
            "beneficial": 0.5, "right": 0.5, "positive": 0.5, 
            "large": 0.5, "attractive": 0.5, "sound": 0.5, 
            "excellent": 0.5, "wrong": -0.5, "unproductive": -0.5, 
            "lose": -0.5, "missing": -0.5, "mishandled": -0.5, 
            "un_lucrative": -0.5, "up": 0.5, "down": -0.5,
            "poor": -0.5, "worthwhile": 0.5, "lucrative": 0.5, "solid": 0.5
        }
        self._news_start_time = time(8)
        self._news_end_time = time(9, 30)
        self._securities = []
    
    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Accumulate news articles between 8:00 AM and 9:30 AM.
        if self._news_start_time <= algorithm.time.time() <= self._news_end_time:
            for dataset_symbol, article in data.get(TiingoNews).items():
                algorithm.securities[dataset_symbol.underlying].news.append(article.description)
        # Wait until the market opens.
        if not data.bars:
            return []
        # Calculate sentiment score for each symbol based on pre-market news articles.
        score_by_security = {}
        for security in self._securities:
            score = sum(
                sum(
                    self._word_sentiment[word] 
                    for word in set(article.lower().split()) 
                    if word in self._word_sentiment
                )
                for article in security.news
            )
            if score > 0:
                score_by_security[security] = score
            security.news = []
        if not score_by_security:
            return []
        # Select top 10 symbols by sentiment score.
        return [
            Insight.price(security, Expiry.END_OF_DAY, InsightDirection.UP, weight=abs(weight))
            for security, weight in sorted(score_by_security.items(), key=lambda x: x[1])[-10:]
        ]
    
    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            # Add the Tiingo News feed.
            security.tiingo_symbol = algorithm.add_data(TiingoNews, security).symbol
            security.news = []
            self._securities.append(security)
        for security in changes.removed_securities:
            # Remove the Tiingo News feed.
            algorithm.remove_security(security.tiingo_symbol)
            self._securities.remove(security)