from QuantConnect.Data.Custom.Tiingo import *
from SentimentByPhrase import SentimentByPhrase
from nltk.util import ngrams

class DrugNewsSentimentAlphaModel(AlphaModel):
    symbol_data_by_symbol = {}
    sentiment_by_phrase = SentimentByPhrase.dictionary
    max_phrase_words = max([len(phrase.split()) for phrase in sentiment_by_phrase.keys()])
    sign = lambda _, x: int(x and (1, -1)[x < 0])
    
    
    def __init__(self, bars_before_insight=30):
        self.bars_before_insight = bars_before_insight
    
    
    def Update(self, algorithm, data):
        """
        Called each time our alpha model receives a new data slice.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - data
            A data structure for all of an algorithm's data at a single time step
        
        Returns a list of Insights to the portfolio construction model
        """
        insights = []
        
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
        
            # If it's after-hours or within 30-minutes of the open, update
            # cumulative sentiment for each symbol    
            if symbol_data.bars_seen_today < self.bars_before_insight:
                tiingo_symbol = symbol_data.tiingo_symbol
                if data.ContainsKey(tiingo_symbol) and data[tiingo_symbol] is not None:
                    article = data[tiingo_symbol]
                    symbol_data.cumulative_sentiment += self.CalculateSentiment(article)
        
            if data.ContainsKey(symbol) and data[symbol] is not None:
                symbol_data.bars_seen_today += 1

                # 30-mintes after the open, emit insights in the direction of the cumulative sentiment
                if symbol_data.bars_seen_today == self.bars_before_insight:
                    
                    next_close_time = symbol_data.exchange.Hours.GetNextMarketClose(data.Time, False)
                    mins_to_close = int((next_close_time - data.Time).total_seconds() / 60)
                    
                    direction = self.sign(symbol_data.cumulative_sentiment)
                    if direction == 0:
                        continue
                    insight = Insight.Price(symbol, 
                                            timedelta(minutes=mins_to_close-2), 
                                            direction)
                    insights.append(insight)
        
                # At the close, reset the sentiment
                if not symbol_data.exchange.DateTimeIsOpen(data.Time):
                    symbol_data.cumulative_sentiment = 0
                    symbol_data.bars_seen_today = 0
        
        return insights
        
        
    def OnSecuritiesChanged(self, algorithm, changes):
        """
        Called each time our universe has changed.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - changes
            The additions and subtractions to the algorithm's security subscriptions
        """
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(security, algorithm)

        for security in changes.RemovedSecurities:
            self.symbol_data_by_symbol.pop(security.Symbol, None)
            
    
    def CalculateSentiment(self, article):
        sentiment = 0
        for content in (article.Title, article.Description):
            words = content.lower().split()
            for num_words in range(1, self.max_phrase_words + 1):
                for gram in ngrams(words, num_words):
                    phrase = ' '.join(gram)
                    if phrase in self.sentiment_by_phrase.keys():
                        sentiment += self.sentiment_by_phrase[phrase]
        return sentiment


class SymbolData:
    cumulative_sentiment = 0
    bars_seen_today = 0
    
    def __init__(self, security, algorithm):
        self.exchange = security.Exchange
        self.tiingo_symbol = algorithm.AddData(TiingoNews, security.Symbol).Symbol

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class DrugManufacturerUniverseSelection(FundamentalUniverseSelectionModel):
    """
    This universe selection model contain securities in the drug manufacturing
    industry group.
    """
    def __init__(self, coarse_size=500, fine_size=5):
        self.coarse_size = coarse_size
        self.fine_size = fine_size
        super().__init__(True)

    def SelectCoarse(self, algorithm, coarse):
        """
        Coarse universe selection is called each day at midnight.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - coarse
            List of CoarseFundamental objects
            
        Returns the symbols that have fundamental data.
        """
        has_fundamentals = [c for c in coarse if c.HasFundamentalData]
        sorted_by_dollar_volume = sorted(has_fundamentals, key=lambda c: c.DollarVolume, reverse=True)
        return [ x.Symbol for x in sorted_by_dollar_volume[:self.coarse_size] ]
    
        
    def SelectFine(self, algorithm, fine):
        """
        Fine universe selection is performed each day at midnight after `SelectCoarse`.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - fine
            List of FineFundamental objects that result from `SelectCoarse` processing
        
        Returns a list of symbols that are in the drug manufacturing industry.
        """
        drug_manufacturers = [f for f in fine if f.AssetClassification.MorningstarIndustryGroupCode == MorningstarIndustryGroupCode.DrugManufacturers]
        sorted_by_pe = sorted(drug_manufacturers, key=lambda f: f.ValuationRatios.PERatio, reverse=True)
        return [ x.Symbol for x in sorted_by_pe[:self.fine_size] ]

# Sentiment dicationary retrieved from:
# https://github.com/queensbamlab/NewsSentiments/blob/master/dict.csv

class SentimentByPhrase:
    dictionary = {
        'okay from fda' : 1,
        'fda approval' : 1,
        'usfda approval' : 1,
        'weaker rupee' : 1,
        'positive step' : 1,
        'resolution' : 1,
        'successful' : 1,
        'stellar' : 1,
        'better' : 1,
        'much better' : 1,
        'better margins' : 1,
        'favourable' : 1,
        'approval' : 1,
        'tough' : -1,
        'reported lower than expected sales' : -1,
        'lower than expected sales' : -1,
        'affecting sales growth' : -1,
        'difficult one' : -1,
        'pricing pressure' : -1,
        'sales declined' : -1,
        'dull' : -1,
        'significant violations' : -1,
        'warning letter' : -1,
        'issued warning letter' : -1,
        'adulterate' : -1,
        'potentially contaminate' : -1,
        'contaminate' : -1,
        'fail' : -1,
        'warn' : -1,
        'violation' : -1,
        'legal action' : -1,
        'drag' : -1,
        'sales decline' : -1,
        'margins decline' : -1,
        'weak' : -1,
        'offset price erosion' : 1,
        'price erosion' : -1,
        'slowdown' : -1,
        'sanction' : -1,
        'concern' : -1,
        'drag on sale' : -1,
        'drop' : -1,
        'challenge' : -1,
        'toll' : -1,
        'uncertain' : -1,
        'recall' : -1,
        'health' : 1,
        'stability' : 1,
        'mixed set' : -1,
        'shares declined' : 0,
        'major breakthrough' : 1,
        'good quarter' : 1,
        'appreciating rupee' : -1,
        'depreciating rupee' : 1,
        'heightened competition' : -1,
        'incorrect instructions' : -1,
        'shares decline' : 0,
        'zero observations' : 1,
        'strong us pipeline' : 1,
        'upgrade' : 1,
        'downgrade' : -1,
        'mixed bag' : -1,
        'disappointing year' : -1,
        'domestic challenges' : -1,
        'benefit' : 1,
        'percent growth' : 1,
        'flat revenue' : -1,
        'flat' : -1,
        'beat' : 1,
        'achieve' : 1,
        'steady margins' : 1,
        'rise' : 1,
        'expand' : 1,
        'ramp up' : 1,
        'launch' : 1,
        'not issued' : 1,
        'clear' : 1,
        'address' : 0,
        'observation' : 0,
        'procedural' : 0,
        'eir' : 1,
        'monetise' : 1,
        'outperform' : 1,
        'enhance' : 1,
        'form 483' : -1,
        'clarify' : 1,
        'facility' : 0,
        'starts' : 1,
        'stable' : 1,
        'initiative' : 1,
        'sold rights' : 1,
        'terminate' : -1,
        'strengthen' : 1,
        'sahpra approval' : 1,
        'nod' : 1,
        'acquire' : 1,
        'raise target' : 1,
        'scaling up' : 1,
        'raise' : 1,
        'subject to clearance' : 0
    }

from DrugManufacturerUniverseSelection import DrugManufacturerUniverseSelection
from DrugNewsSentimentAlphaModel import DrugNewsSentimentAlphaModel

class UncoupledVerticalInterceptor2(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetCash(100000)
        
        self.SetUniverseSelection(DrugManufacturerUniverseSelection())
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.SetAlpha(DrugNewsSentimentAlphaModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())