"""
420 Friendly

Take positions based on legalization terms in news articles.

Backtest Report: https://www.quantconnect.com/reports/8589b1fe98375a450d74889521470e38
Backtest: https://www.quantconnect.com/terminal/processCache/?request=embedded_backtest_8589b1fe98375a450d74889521470e38.html
"""

from QuantConnect.Data.Custom.Tiingo import *
from datetime import timedelta


class FourTwentyFriendlyAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 12, 9)
        self.SetCash(100000)

        weed_symbols = [
            Symbol.Create(sym, SecurityType.Equity, Market.USA) for sym in [
                'ABBV', 'BUD', 'MO', 'TAP', 'WEED.TO', 'SMG', 'CGC', 'TLRY',
                'GWPH', 'CRON', 'ACB', 'NBEV', 'CRBP', 'TGOD.TO', 'TGODF',
                'CANN', 'MJ', 'SNDL'
            ]
        ]

        self.SetUniverseSelection(ManualUniverseSelectionModel(weed_symbols))
        self.SetAlpha(FourTwentyFriendlyAlphaModel())
        self.SetExecution(StandardDeviationExecutionModel())
        self.SetPortfolioConstruction(AccumulativeInsightPortfolioConstructionModel())
        self.SetRiskManagement(MaximumUnrealizedProfitPercentPerSecurity(maximumUnrealizedProfitPercent=0.1168))

    def OnData(self, data):
        """ OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        """
        pass


class FourTwentyFriendlyAlphaModel(AlphaModel):
    def __init__(self):
        self.Name = '420FriendlyAlphaModel'
        self.newsData = {}
        self.wordScores = {
            "legalization": 5, "legal": 5
        }

    def Update(self, algorithm, data=None):
        """
            Updates this alpha model with the latest data from the algorithm.
            This is called each time the algorithm receives data for subscribed securities
            Generate insights on the securities in the universe.
        """
        insights = []

        if data is None:
            return insights

        news = data.Get(TiingoNews)
        for article in news.Values:
            words = article.Description.lower().split(" ")
            score = sum(
                [
                    self.wordScores[word] for word in words
                    if word in self.wordScores
                ]
            )

            symbol = article.Symbol.Underlying
            # Add scores to the rolling window associated with its newsData symbol
            self.newsData[symbol].Window.Add(score)

            # Sum the rolling window scores, save to sentiment
            # If sentiment aggregate score for the time period is greater than 5, emit an up insight
            sentiment = sum(self.newsData[symbol].Window)
            if sentiment > 5:
                insights.append(Insight.Price(symbol, timedelta(1), InsightDirection.Up, None, None))
            else:
                insights.append(Insight.Price(symbol, timedelta(1), InsightDirection.Flat, None, None))

        return insights

    def OnSecuritiesChanged(self, algorithm, changes=None):
        if changes is None:
            return

        for security in changes.AddedSecurities:
            symbol = security.Symbol
            news_asset = algorithm.AddData(TiingoNews, symbol)
            self.newsData[symbol] = NewsData(news_asset.Symbol)

        for security in changes.RemovedSecurities:
            newsData = self.newsData.pop(security.Symbol, None)
            if newsData is not None:
                algorithm.RemoveSecurity(newsData.Symbol)


class NewsData:
    def __init__(self, symbol):
        self.Symbol = symbol
        self.Window = RollingWindow[float](100)


class SymbolData:
    """Contains data specific to a symbol required by this model"""

    def __init__(self, security):
        self.Security = security
        self.Symbol = security.Symbol