from datetime import datetime, timedelta
import numpy as np
from nltk.util import ngrams

class SentimentAlphaModel(AlphaModel):
    
    def __init__(self, algo):
        self.algo = algo
        self.CandleSize = timedelta(minutes = 5)
        self.MaxHoldTime = 20*self.CandleSize
        self.Symbols = set()
        
        self.Prices = {}
        self.ExitPrices = {}
        self.BuyTimes = {}
        
        self.Holdings = set()
        self.MaxHoldings = 5
        
        self.CandlesInTraining = {}
        self.Candles = {}
        self.LastCandleTime = self.algo.Time
        
        self.TradingTime = False
        return
    
    def UpdateCandlesInTraining(self, data):
        for symbol in self.Symbols:
            if symbol in data and data[symbol] is not None:
                o, h, l, c = data[symbol].Open, data[symbol].High, data[symbol].Low, data[symbol].Close
                candle = self.CandlesInTraining[symbol]
                if candle["Open"] is None:
                    candle["Open"] = o
                if candle["High"] is None or h > candle["High"]:
                    candle["High"] = h
                if candle["Low"] is None or l < candle["Low"]:
                    candle["Low"] = l
                candle["Close"] = c
                
    def AddCandles(self):
        #Candles in training are finished; HA time!
        for symbol in self.Symbols:
            o, h, c, l = self.CandlesInTraining[symbol]["Open"], self.CandlesInTraining[symbol]["High"], self.CandlesInTraining[symbol]["Low"], self.CandlesInTraining[symbol]["Close"]
            if o is None or h is None or c is None or l is None:
                continue
            self.Prices[symbol] = c
            newCandle = {"Open": o, "High": h, "Low": l, "Close": c}
            newCandle["High"] = max([newCandle["Open"], h, newCandle["Close"]])
            newCandle["Low"] = min([newCandle["Open"], l, newCandle["Close"]])
            if o > c:
                newCandle["Color"] = "Red"
            else:
                newCandle["Color"] = "Green"
                
            newCandle["Top"] = max([newCandle["Open"], newCandle["Close"]])
            newCandle["Bottom"] = min([newCandle["Open"], newCandle["Close"]])
            self.Candles[symbol].append(newCandle)
                
    def UpdateCandles(self, data):
        self.UpdateCandlesInTraining(data)
        if self.algo.Time >= self.LastCandleTime + self.CandleSize:
            self.AddCandles()
            for symbol in self.Symbols:
                self.CandlesInTraining[symbol] = self.FreshCandle()
            self.LastCandleTime = self.algo.Time
            self.TradingTime = True
    
    def FreshCandle(self):
        return {"Open": None, "High": None, "Low": None, "Close": None}
    
    def BuySignal(self, symbol):
        if symbol not in self.Candles or len(self.Candles[symbol]) < 1:
            return 0
        return self.Candles[symbol][-1]["High"] - self.Candles[symbol][-1]["Low"]
        
    def SellSignal(self, symbol):
        return (symbol not in self.ExitPrices) or (self.Prices[symbol] < self.ExitPrices[symbol][0] or self.Prices[symbol] > self.ExitPrices[symbol][1])
    
    def OnSecuritiesChanged(self, algo, changes):
        for symbol in changes.AddedSecurities:
            self.Symbols.add(symbol.Symbol)
            self.CandlesInTraining[symbol.Symbol] = self.FreshCandle()
            self.Candles[symbol.Symbol] = []
        for symbol in changes.RemovedSecurities:
            self.Symbols.discard(symbol.Symbol)
            self.CandlesInTraining.pop(symbol.Symbol, None)
            self.Candles.pop(symbol.Symbol, None)
            self.ExitPrices.pop(symbol.Symbol, None)
            self.Holdings.discard(symbol.Symbol)
            self.algo.Liquidate(symbol.Symbol)
    
    def Update(self, algo, data):
        self.UpdateCandles(data)
        res = []
        if not self.TradingTime:
            for symbol in self.Symbols:
                if symbol in data and data[symbol] is not None and (symbol not in self.ExitPrices or data[symbol].Close < self.ExitPrices[symbol][0] or data[symbol].Close > self.ExitPrices[symbol][1]):
                    # self.algo.Log(f"Selling {symbol}")
                    res.append(Insight.Price(symbol, self.MaxHoldTime, InsightDirection.Flat))
                    self.Holdings.discard(symbol)
            return res
        
        buySignals = {}
        for symbol in self.Symbols:
            buySignals[symbol] = self.BuySignal(symbol)
                
        thingsWeShallBuy = set(sorted([symbol for symbol in buySignals.keys()], key = lambda x: buySignals[x], reverse = True)[:self.MaxHoldings - len(self.Holdings)])
        for symbol in self.Symbols:
            if self.SellSignal(symbol) and (symbol not in thingsWeShallBuy):
                # self.algo.Log(f"Selling {symbol}")
                res.append(Insight.Price(symbol, self.MaxHoldTime, InsightDirection.Flat))
                self.Holdings.discard(symbol)
        
        for symbol in thingsWeShallBuy:
            # self.algo.Log(f"Buying {symbol}")
            res.append(Insight.Price(symbol, self.MaxHoldTime, InsightDirection.Up))
            self.Holdings.add(symbol)
            p = self.Prices[symbol]
            bottom = self.Candles[symbol][-1]["Bottom"]
            self.ExitPrices[symbol] = (p - 0.5*(p - bottom), p + 2*(p - bottom))
        assert(len(self.Holdings) <= self.MaxHoldings)
        self.TradingTime = False
        return res

class CreativeVioletMule(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)
        
        self.MinEarningTime = timedelta(days = 1)
        self.MaxEarningTime = timedelta(days = 30)
        
        self.Symbols = set()
        self.SecuritiesSet = set()
        self.Holdings = set()
        self.Coarse_Count = 500
        self.Current_Day = None
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.SetAlpha(SentimentAlphaModel(self))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        
    def CoarseSelectionFunction(self, coarse):
        if self.Current_Day == self.Time.day:
            return Universe.Unchanged
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData],
            key=lambda x: x.DollarVolume, reverse=True)[:self.Coarse_Count]
        return [i.Symbol for i in sortedByDollarVolume]
        
    def FineSelectionFunction(self, fine):
        if self.Current_Day == self.Time.day:
            return Universe.Unchanged
        self.Current_Day = self.Time.day
        res = [x for x in fine if (x.EarningReports.FileDate > self.Time - self.MaxEarningTime and x.EarningReports.FileDate < self.Time - self.MinEarningTime)]
        symbolRes = [x.Symbol for x in res]
        return symbolRes

    def OnData(self, data):
        return