from Alphas.RsiAlphaModel import RsiAlphaModel

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Risk.TrailingStopRiskManagementModel import TrailingStopRiskManagementModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from IchimokuAlphaModel import IchimokuAlphaModel

class CryptoGenie(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 12, 31)  # Set End Date
        self.SetCash(1000000)  # Set Strategy Cash
        self.SetWarmUp(78)
        # Symbol.Create("ETHUSD", SecurityType.Crypto, Market.Bitfinex)
        self.symbols = [ Symbol.Create("BTCUSD", SecurityType.Crypto, Market.Bitfinex) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(self.symbols) )
        self.UniverseSettings.Resolution = Resolution.Minute
        self.SetBrokerageModel(BrokerageName.Bitfinex)
        
        self.tenkan_period=9
        self.kijun_period=26 
        self.senkouA_period=26
        self.senkouB_period=52
        self.chikou_period=26
        self.chikou_threshold=0
        self.kijun_threshold=0
        self.resolution = Resolution.Minute
        
        self.Debug("Started Alpha Model")
        
        self.AddAlpha(IchimokuAlphaModel(self.tenkan_period, self.kijun_period, 
                                        self.senkouA_period, self.senkouB_period,
                                        self.chikou_period, self.chikou_threshold,
                                        self.kijun_threshold, self.resolution))
                                        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        #self.SetRiskManagement(TrailingStopRiskManagementModel(0.1)) # 
        self.SetExecution(ImmediateExecutionModel())



    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # if not self.Portfolio.Invested:
        #     self.SetHoldings("BTCUSD", 1)

from clr import AddReference
AddReference("QuantConnect.Indicators")
from QuantConnect.Indicators  import IchimokuKinkoHyo

class IchimokuAlphaModel(AlphaModel):
    ''' 
    Enhanced IchimokuKinkoHyo Indicator is used to generate trend signal.
    We use tenkansen, kijunsen, spanA,spanB and Chikou span to calculate the 
    deltas. We then combine then to get a bullish or bearish signal
    
    '''
    
    def __init__(self, tenkan_period=9, kijun_period=26, 
                senkouA_period=26,senkouB_period=52,
                chikou_period=26,chikou_threshold=0,kijun_threshold=0,resolution=Resolution.Daily):

                    
        self.TenkanPeriod = tenkan_period
        self.KijunPeriod = kijun_period
        self.SenkouAPeriod = senkouA_period
        self.SenkouBPeriod = senkouB_period
        self.ChikouPeriod = chikou_period
        self.ChikouThreshold = chikou_threshold
        self.KijunThreshold = kijun_threshold
        self.resolution = resolution
        self.symbolData = {};
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(self.resolution), self.SenkouBPeriod)
        
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{},{},{},{},{},{},{})'.format(self.__class__.__name__, 
                                                tenkan_period, kijun_period, senkouA_period, 
                                                senkouB_period, chikou_period, 
                                                chikou_threshold, kijun_threshold, 
                                                resolutionString)
                                                

    def Update(self, algorithm, data):
        ''' Determines an insight for each security based on it's current Ichimoku signal
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []
        
        for symbol, sd in self.symbolData.items():
            
            direction = InsightDirection.Flat
            if sd.strong_bullish() and sd.price_above_kijunsen() and sd.price_above_kumo():
                direction = InsightDirection.Up
            if sd.strong_bearish() and sd.price_below_kijunsen() and sd.price_below_kumo():
                direction = InsightDirection.Down
            
            if direction != InsightDirection.Flat and direction != sd.PreviousDirection:
                algorithm.Debug("Added  Insight for Symbol : " + str(symbol) + " Direction : " + str(direction))
                insights.append(Insight.Price(symbol, self.insightPeriod, direction))
                sd.PreviousDirection = direction
                    
        return insights
            
    
    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed.
        This initializes the Ichimoku for each added security and cleans up the indicator for each removed security.
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        
        # Handle Added Securities
        for security in changes.AddedSecurities:
            self.symbolData[security.Symbol] = SymbolData(security, algorithm, 
                                                            self.TenkanPeriod, 
                                                            self.KijunPeriod, 
                                                            self.SenkouAPeriod,
                                                            self.SenkouBPeriod,
                                                            self.ChikouPeriod,
                                                            self.ChikouThreshold,
                                                            self.KijunThreshold,
                                                            self.resolution)
            algorithm.Debug("Added Symbol : " + str(security.Symbol))
        for security in changes.RemovedSecurities:
            algorithm.Debug("Removed Symbol : " + str(security.Symbol))
            data = self.symbolData.pop(security.Symbol, None)
            if data is not None:
                pass # Do any clean up here if required

class SymbolData():
    def __init__(self,
                security,
                algorithm,
                tenkan_period, 
                kijun_period, 
                senkouA_period,
                senkouB_period,
                chikou_period,
                chikou_threshold,
                kijun_threshold,
                resolution):
        
    
        self.Security = security
        self.ichimoku = IchimokuKinkoHyo(security.Symbol, tenkan_period, 
                                        kijun_period, senkouA_period, senkouB_period,
                                        chikou_period, chikou_period)
                                        
        self.thirtyMinuteConsolidator = TradeBarConsolidator(timedelta(minutes=30))
        #self.thirtyMinuteConsolidator.DataConsolidated += self.OnThirtyMinuteBarHandler
        algorithm.SubscriptionManager.AddConsolidator(security.Symbol, self.thirtyMinuteConsolidator)
        algorithm.RegisterIndicator(security.Symbol, self.ichimoku, self.thirtyMinuteConsolidator)
        
        self.algorithm = algorithm
        self.tenkansen = self.ichimoku.Tenkan
        self.kijunsen = self.ichimoku.Kijun
        self.senkouA = self.ichimoku.SenkouA
        self.senkouB = self.ichimoku.SenkouB
        self.chikou = self.ichimoku.Chikou
        self.low = self.Security.Low
        self.high = self.Security.High
        self.PreviousDirection = InsightDirection.Flat
        
    def OnThirtyMinuteBarHandler(self, sender, bar):
        pass
        #self.algorithm.Debug( "Updated 30 minute bar : " + str(bar) )
    
    def weak_bullish(self):
        tenkan_above_span_a_b = (self.tenkansen.Current.Value > self.senkouA.Current.Value) and (self.tenkansen.Current.Value > self.senkouB.Current.Value)
        kijun_above_span_a_b = (self.kijunsen.Current.Value > self.senkouA.Current.Value) and (self.kijunsen.Current.Value > self.senkouB.Current.Value)
        return (tenkan_above_span_a_b) and (kijun_above_span_a_b)
        
    def weak_bearish(self):
        tenkan_below_span_a_b = (self.tenkansen.Current.Value < self.senkouA.Current.Value) and (self.tenkansen.Current.Value < self.senkouB.Current.Value)
        kijun_below_span_a_b = (self.kijunsen.Current.Value < self.senkouA.Current.Value) and (self.kijunsen.Current.Value < self.senkouB.Current.Value)
        return (tenkan_below_span_a_b) and (kijun_below_span_a_b)
        
    def tenkansen_cross(self):
        retVal = (self.tenkansen.Current.Value > self.kijunsen.Current.Value)
        return retVal
        
    def kijunsen_cross(self):
        retVal = (self.tenkansen.Current.Value < self.kijunsen.Current.Value)
        return retVal
        
    def strong_bullish(self):
        return self.tenkansen_cross() and self.weak_bullish()
        
    def strong_bearish(self):
        return self.kijunsen_cross() and self.weak_bearish()
        
    def tenkansen_delta(self):
        return ((self.tenkansen-self.kijunsen)/self.kijunsen)*100
        
    def kijunsen_delta(self):
        return ((self.kijunsen-self.tenkansen)/self.tenkansen)*100
        
    def price_above_kijunsen(self):
        return self.low > self.kijunsen.Current.Value
    
    def price_below_kijunsen(self):
        return self.high > self.kijunsen.Current.Value
        
    def price_above_kumo(self):
        return self.low > self.senkouA.Current.Value and self.low > self.senkouB.Current.Value
        
    def price_below_kumo(self):
        return self.high < self.senkouA.Current.Value and self.high < self.senkouB.Current.Value
        
    def price_above_kijun_delta(self):
        return ((self.low - self.kijunsen.Current.Value)/self.kijunsen.Current.Value)*100
        
    def price_below_kijun_delta(self):
        return ((self.kijunsen.Current.Value-self.high)/self.high)*100