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

class SentimentAlpha(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2020, 7, 14)
        self.SetCash(1000000)
        
        self.tickers = ["TSLA", "MSFT", "GS", "PEP", "NIO", "F", "MTCH","DDOG", "LUV", "TMUS", "AMD", "INO"]
         
        symbols = []
        
        for i in self.tickers:
            symbols.append(Symbol.Create(i, SecurityType.Equity, Market.USA))
        
        self.UniverseSettings.Resolution = Resolution.Hour
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.SetAlpha(SentimentWindow())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) 
        self.SetExecution(ImmediateExecutionModel()) 
        self.SetRiskManagement(takeprofit.TrailingStopRiskManagementModel())
        

class NewsData():
    def __init__(self, symbol):
        self.Symbol = symbol
        #two month rolling window if insights are emitted daily
        self.Window = RollingWindow[float](60)  
        
class SentimentWindow(AlphaModel):
    
    def __init__(self): 
        self.newsData = {}

        self.wordScores = {
            "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,
            "unproductive": -0.5, "poor": -0.5, "wrong": -0.5,
            "worthwhile": 0.5, "lucrative": 0.5, "solid": 0.5,
            "beat": 0.5, "missed": -0.5
            
        }
        
        self.lastday = -1
        
    def Update(self, algorithm, data):

        insights = []
        
        day = algorithm.Time.day
        if day == self.lastday:
            return insights
        self.lastday = day
        
        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])
                
            if score == 0.0:
                continue
            
            symbol = article.Symbol.Underlying 
    
            self.newsData[symbol].Window.Add(score)
            
            
            if self.newsData[symbol].Window.Count < 3:
                return insights
                
            currentWindow = self.newsData[symbol].Window[0]
                
            sentimentMean = sum(self.newsData[symbol].Window)/self.newsData[symbol].Window.Count
            algorithm.Debug("Mean "+ str(sentimentMean))
            algorithm.Debug("Window "+ str(currentWindow))
            
            if currentWindow > (sentimentMean*1.1):
                insights.append(Insight.Price(symbol, timedelta(days=2), InsightDirection.Up, None, None))
            elif currentWindow < -1:
                insights.append(Insight.Price(symbol, timedelta(days=1), InsightDirection.Down, None, None))
            else:
                insights.append(Insight.Price(symbol, timedelta(days=1), InsightDirection.Flat, None, None))
            
        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)

class TrailingStopRiskManagementModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the maximum possible loss
    measured from the highest unrealized profit'''
    
    def __init__(self):
        '''Initializes a new instance of the TrailingStopRiskManagementModel class'''

        self.liquidated = set()
        self.lastmonth=-1
        
    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
        month= algorithm.Time.month
        if month!= self.lastmonth:
            self.liquidated.clear()
            self.lastmonth= month
        #if different month clear all liquidations
        riskAdjustedTargets = list()

        for kvp in algorithm.Securities:
            symbol = kvp.Key
            security = kvp.Value

            if security.Holdings.UnrealizedProfitPercent>0.05 or security.Holdings.UnrealizedProfitPercent<-0.03 or security.Symbol in self.liquidated:
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))
                if algorithm.Securities[security.Symbol].Invested:
                    self.liquidated.add(security.Symbol)
        return riskAdjustedTargets