from QuantConnect.Data.Custom.Benzinga import *

from datetime import datetime, timedelta

### <summary>
### Benzinga is a provider of news data. Their news is made in-house
### and covers stock related news such as corporate events.
### </summary>
class BenzingaNewsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.words = {
            "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
        }

        self.lastTrade = datetime(1, 1, 1)

        self.SetStartDate(2018, 6, 5)
        self.SetEndDate(2018, 8, 4)
        self.SetCash(100000)

        aapl = self.AddEquity("AAPL", Resolution.Hour).Symbol
        ibm = self.AddEquity("IBM", Resolution.Hour).Symbol

        self.AddData(BenzingaNews, aapl)
        self.AddData(BenzingaNews, ibm)

    def OnData(self, data):
        if (self.Time - self.lastTrade) < timedelta(days=5):
            return

        # Get rid of our holdings after 5 days, and start fresh
        self.Liquidate()

        # Get all Benzinga data and loop over it
        for article in data.Get(BenzingaNews).Values:
            selectedSymbol = None

            # Use loop instead of list comprehension for clarity purposes

            # Select the same Symbol we're getting a data point for
            # from the articles list so that we can get the sentiment of the article
            # We use the underlying Symbol because the Symbols included in the `Symbols` property
            # are equity Symbols.
            for symbol in article.Symbols:
                if symbol == article.Symbol.Underlying:
                    selectedSymbol = symbol
                    break

            if selectedSymbol is None:
                raise Exception(f"Could not find current Symbol {article.Symbol.Underlying} even though it should exist")

            # The intersection of the article contents and the pre-defined words are the words that are included in both collections
            intersection = set(article.Contents.lower().split(" ")).intersection(list(self.words.keys()))
            # Get the words, then get the aggregate sentiment
            sentimentSum = sum([self.words[i] for i in intersection])

            if sentimentSum >= 0.5:
                self.Log(f"Longing {article.Symbol.Underlying} with sentiment score of {sentimentSum}")
                self.SetHoldings(article.Symbol.Underlying, sentimentSum / 5)

                self.lastTrade = self.Time

            if sentimentSum <= -0.5:
                self.Log(f"Shorting {article.Symbol.Underlying} with sentiment score of {sentimentSum}")
                self.SetHoldings(article.Symbol.Underlying, sentimentSum / 5)

                self.lastTrade = self.Time

    def OnSecuritiesChanged(self, changes):
        for r in changes.RemovedSecurities:
            # If removed from the universe, liquidate and remove the custom data from the algorithm
            self.Liquidate(r.Symbol)
            self.RemoveSecurity(Symbol.CreateBase(BenzingaNews, r.Symbol, Market.USA))