from QuantConnect.Data.Custom.Tiingo import *
from datetime import datetime, timedelta
import numpy as np

class TiingoNS(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        #self.SetEndDate(2015, 11, 1)  
        symbols = [
            Symbol.Create("SPY",SecurityType.Equity, Market.USA),#on
            ]
        
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.SetAlpha(NewsSentimentAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) 
        self.SetExecution(ImmediateExecutionModel()) 
        #self.SetRiskManagement(NullRiskManagementModel())
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.SetCash(1000000) 

class NewsData():
    def __init__(self, symbol):
        self.Symbol = symbol
        self.Window = RollingWindow[float](100)  
        
class NewsSentimentAlphaModel(AlphaModel):
    
    def __init__(self): 
        self.newsData = {}
        self.wordScores = {
            "over":1, #1
        } 
                
    def Update(self, algorithm, data):

        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
            
            self.newsData[symbol].Window.Add(score)
            
            sentiment = sum(self.newsData[symbol].Window)
           
            if sentiment > 3:
                insights.append(Insight.Price(symbol, timedelta(1), InsightDirection.Up))
           
        return insights
    
    def OnSecuritiesChanged(self, algorithm, changes):

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

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