import numpy as np
import decimal

### <summary>
### RSI algorithm the buys on overbought and sells on oversold but only with trend.
### </summary>

class TrendRSIAlgorithm(QCAlgorithm):
    '''Trend following: if strong dailt RSI, up trend. If weak RSI: downtrend. Take entry on hourly overbought or oversold.'''

    def Initialize(self):
        '''Initializes cash and max trade size. Sets target currencies to run on.'''

        # What do we want to trade?
        self.currenciesToUse = {"EURUSD", "AUDUSD", "EURGBP", "USDNZD"}

        self.SetStartDate(2007,10, 7)  #Set Start Date
        #self.SetEndDate(2016,6,11)    #Set End Date
        
        # Strategy cash management rules
        self.SetCash(1000) # Set Strategy Cash
        
        self.numlots = 4 # number of mini lots to trade (all positions added up)
        
        # Total size of all trades added together:
        self.maxTradeSize = self.numlots*10000 # 10000 = 1 mini lot
        
        # Integer actually equals number of pips when priced this way....
        self.profitTarget = self.numlots*40 # profit per mini
        self.stopLoss = self.numlots*0 # stop per mini
        self.secondTrade = self.numlots*20 # when to try to get out of trouble if no stop
        self.criticalDeath = self.secondTrade*3 # when to abandon all hope and close it
        
        # Strategy RSI values
        self.overBought = 80
        self.overSold = 20
        self.uptrend = 65
        self.downtrend = 35
        
        self.SetBrokerageModel(BrokerageName.OandaBrokerage)
        
        self.rsiSlow = {}
        self.rsiFast = {}
        self.inTrouble = {}
        
        # Setup currencies
        for cur in self.currenciesToUse:
            # Add symbols to universe
            self.Debug("Adding symbol: " + cur)
            self.forex = self.AddForex(cur, Resolution.Hour, Market.Oanda)
            self.rsiFast[cur] = self.RSI(cur, 30, MovingAverageType.Simple, Resolution.Hour)
            self.rsiSlow[cur] = self.RSI(cur, 30, MovingAverageType.Simple, Resolution.Daily)
            self.inTrouble[cur] = 0 
            
        # If we don't warmup, we are trading randomly right on day one
        self.SetWarmUp(30, Resolution.Daily)
        
        # Adjust so portfolio total stays within max trade size (not counting getting out of trouble)
        self.tradesize = round(self.maxTradeSize / len(self.currenciesToUse))
        self.profitTarget = round(self.profitTarget / len(self.currenciesToUse))
        self.stopLoss = round(self.stopLoss / len(self.currenciesToUse))
        self.secondTrade = round(self.secondTrade / len(self.currenciesToUse))
        self.criticalDeath = round(self.criticalDeath / len(self.currenciesToUse))
        
    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
        '''
        
        # Probably not necessary, but good to be safe ;)
        if self.IsWarmingUp:
            return
        
        # Now loop through everyone we want to trade
        for symbol in self.ActiveSecurities.Keys:
            
            holdings = self.Portfolio[symbol.Value].Quantity
        
            if holdings == 0:
                
                # No current holdings, so initiate new position if appropriate
                self.InitiatePosition(symbol)
                
            else:
                
                # We own it, so figure out what to do
              
                profit = self.Portfolio[symbol.Value].UnrealizedProfit
                
                # Take profit if we have hit profit target
                if profit > self.profitTarget:
                    self.Liquidate(symbol.Value)
                    self.inTrouble[symbol.Value] = 0
                    continue # don't bother with anything else once we profit
                
                if self.stopLoss == 0:
                    
                    # No stop, we either die or take second trade
                    
                    # Critical death is a last resort stop, so must check for it
                    if profit < -self.criticalDeath:
                        # Things have gone horribly wrong, just throw in the towel
                        self.Debug("ZOMG hit critical death: " + symbol.Value)
                        self.Liquidate(symbol.Value)
                        continue # required since this likely means RSI is extreme so we'd liquide then immediate take a 2nd trade
                    
                    # Didn't stop out or take profit, so deal with 2nd trade situation
                    # Filter to let it move against us enough before adding to the position    
                    if profit < -self.secondTrade:
                        self.HandleSecondTrade(symbol, holdings)
                        continue # for completeness or adding code later, not necessary here
                
                else:
                    # We have a stop, so use it 
                    # (we check every bar instead of setting stops to give us some room to come around)
                    if profit < -self.stopLoss:
                        # get out!
                        self.Liquidate(symbol.Value)

    def InitiatePosition(self, symbol):
        
        # Filter on Trend
        if self.rsiSlow[symbol.Value].Current.Value > self.uptrend:
            # Up trend so only take long
                
            # Price below oversold RSI, so buy
            if self.rsiFast[symbol.Value].Current.Value < self.overSold:
                self.MarketOrder(symbol, self.tradesize)
                return
                
        if self.rsiSlow[symbol.Value].Current.Value < self.downtrend:
            # Downtrend so only take short 
            
            # Price above overbought, so sell
            if self.rsiFast[symbol.Value].Current.Value > self.overBought:
                self.MarketOrder(symbol, -self.tradesize)
                return

    def HandleSecondTrade(self, symbol, holdings):
        '''Essentially, we take a 2nd trade if the 1st has failed to profit 
            and we hit overbought or oversold again'''

        # in trouble, try to get out
        
        # inTrouble tracks whether we already added to the position
        if self.inTrouble[symbol.Value] == 0:
            
            # haven't initiated fix, so try to do better
            
            if holdings < 0:
                # We are short, so add short
                if self.rsiFast[symbol.Value].Current.Value > self.overBought:
                    self.MarketOrder(symbol, -self.tradesize)
                    self.inTrouble[symbol.Value] = 1
            else:
                # We are long, so try to add long
                if self.rsiFast[symbol.Value].Current.Value < self.overSold:
                    self.MarketOrder(symbol, self.tradesize)
                    self.inTrouble[symbol.Value] = 1