#region imports
from AlgorithmImports import *
#endregion
from scipy.stats import linregress
import numpy as np


class MeanReverse(QCAlgorithm):
    
    def Initialize(self):
        # Define backtest window and portfolio cash
        self.SetStartDate(2022, 1, 1)
        self.SetEndDate(2022, 8, 1)
        self.CashDispoInitial=10000
        self.SetCash(self.CashDispoInitial)
        self.sizelot=30000
        
        # Add the assets to be fed into the algorithm and save the symbol objects (to be referred later)
        self.NameAsset1 = "EURCAD"
        self.NameAsset2 = "GBPCAD"
        self.asset1 = self.AddForex(self.NameAsset1, Resolution.Hour, Market.Oanda).Symbol
        self.asset2 = self.AddForex(self.NameAsset2, Resolution.Hour, Market.Oanda).Symbol
        
        # set the starting flag as not invested and initiate paramaters
        self.entry1 = 0
        self.entry2 = 0
        self.entry_slope = 0
        self.entry1_second = 0
        self.entry2_second = 0
        self.entry_slope_second = 0
        self.is_invested = None

        # set paramaters that should be optimized
        self.lookback = 4
        bb_lookback = 48
        bb_std = 1.8
        self.out_pips=6

        # We then create a bollinger band with 120 steps for lookback period
        self.bb1 = BollingerBands(bb_lookback, bb_std, MovingAverageType.Simple)
        
        #Define charts showing when buy and sell
        stockPlot1 = Chart('Trade Plot Asset1')
        stockPlot1.AddSeries(Series('Buy1', SeriesType.Scatter, '$', Color.Green, ScatterMarkerSymbol.Triangle))
        stockPlot1.AddSeries(Series('Sell1', SeriesType.Scatter, '$', Color.Red, ScatterMarkerSymbol.TriangleDown))
        stockPlot1.AddSeries(Series('Liquidate1', SeriesType.Scatter, '$', Color.Blue, ScatterMarkerSymbol.Diamond))
        self.AddChart(stockPlot1)
        stockPlot2 = Chart('Trade Plot Asset2')
        stockPlot2.AddSeries(Series('Buy2', SeriesType.Scatter, '$', Color.Green, ScatterMarkerSymbol.Triangle))
        stockPlot2.AddSeries(Series('Sell2', SeriesType.Scatter, '$', Color.Red, ScatterMarkerSymbol.TriangleDown))
        stockPlot2.AddSeries(Series('Liquidate2', SeriesType.Scatter, '$', Color.Blue, ScatterMarkerSymbol.Diamond))
        self.AddChart(stockPlot2)

    def OnData(self, data):
        
        #Compute lot quantity
        self.CashDispo = self.Portfolio.TotalProfit
        self.QuantityAsset = self.sizelot #*(1+ 2*(self.CashDispo/self.CashDispoInitial-1))
    
        #Compute and refresh slope and intercept
        history = self.History([self.asset1, self.asset2], self.lookback, Resolution.Hour)
        history = history.unstack(level=0).dropna()

        asset1 = np.log(history['close', self.asset1.Value])
        asset2 = np.log(history['close', self.asset2.Value])
        
        reg = linregress(asset1, asset2)

        #Define portfolio and update BB
        portfolio = np.log(data[self.asset2].Close) - np.log(data[self.asset1].Close) * reg.slope - reg.intercept
        ass1=data[self.asset1].Close
        ass2=data[self.asset2].Close
        self.bb1.Update(self.Time, float(portfolio))
        
        # Plot the portfolio (to see if it is working, and the bb bands)
        self.Plot("Indicators", "Portfolio", float(portfolio))
        self.Plot("Indicators", "Upper", self.bb1.UpperBand.Current.Value)
        self.Plot("Indicators", "Lower", self.bb1.LowerBand.Current.Value)

        #self.Plot("Slope", "Slope", float(reg.slope))
        self.Plot("Cash", "Cash", float(self.CashDispo))
        
        # check if the bolllinger band indicator is ready 
        if not self.bb1.IsReady:
            self.Debug('BB not ready')
        
        upper_band = self.bb1.UpperBand.Current.Value
        lower_band = self.bb1.LowerBand.Current.Value
        middle_band = self.bb1.MiddleBand.Current.Value

        # Trading mechanism
        # if it is not invested, see if there is an entry point
        if not self.is_invested:
            # if our portfolio is bellow the lower band, enter long --> sell A1 and buy A2
            if portfolio < lower_band:
                self.MarketOrder(self.NameAsset1, -self.QuantityAsset*reg.slope)
                self.MarketOrder(self.NameAsset2, self.QuantityAsset)
                self.Plot('Trade Plot Asset1','Sell1',ass1)
                self.Plot('Trade Plot Asset2','Buy2',ass2)
                self.Debug('Entering Long')
                self.is_invested = 'long'
                self.entry1=ass1
                self.entry2=ass2
                self.entry_slope = reg.slope
                self.Debug(self.CashDispo/self.CashDispoInitial)

            # if our portfolio is above the upper band, go short --> buy A1 and sell A2
            if portfolio > upper_band:
                self.MarketOrder(self.NameAsset1, self.QuantityAsset*reg.slope)
                self.MarketOrder(self.NameAsset2, -self.QuantityAsset)
                self.Plot('Trade Plot Asset1','Buy1',ass1)
                self.Plot('Trade Plot Asset2','Sell2',ass2)
                self.Debug('Entering Short')
                self.is_invested = 'short'
                self.entry1=ass1
                self.entry2=ass2
                self.entry_slope = reg.slope
                self.Debug(self.CashDispo/self.CashDispoInitial)
        
        # if it is invested in something, check the exiting signal (when it crosses the mean)   
        else:
            #liquidate position
            if self.is_invested == 'long' and (ass2 - self.entry2) + self.entry_slope*(self.entry1-ass1)>self.out_pips/10000:
                self.Liquidate()
                self.Plot('Trade Plot Asset1','Liquidate1',ass1)
                self.Plot('Trade Plot Asset2','Liquidate2',ass2)
                self.Debug('Exiting Long')
                self.Debug(self.CashDispo/self.CashDispoInitial)
                self.is_invested = None
            elif self.is_invested == 'short' and (self.entry2-ass2) + self.entry_slope*(ass1 - self.entry1)>self.out_pips/10000:
                self.Liquidate()
                self.Plot('Trade Plot Asset1','Liquidate1',ass1)
                self.Plot('Trade Plot Asset2','Liquidate2',ass2)
                self.Debug('Exiting Short')
                self.Debug(self.CashDispo/self.CashDispoInitial)
                self.is_invested = None
            #liquidate position
            elif self.is_invested == 'longlong' and (ass2 - self.entry2_second) + (ass2 - self.entry2) + self.entry_slope_second*(self.entry1_second-ass1) +self.entry_slope*(self.entry1-ass1)>1.5*self.out_pips/10000:
                self.Liquidate()
                self.Plot('Trade Plot Asset1','Liquidate1',ass1)
                self.Plot('Trade Plot Asset2','Liquidate2',ass2)
                self.Debug('Exiting Second Long')
                self.Debug(self.CashDispo/self.CashDispoInitial)
                self.is_invested = None
            elif self.is_invested == 'shortshort' and (self.entry2_second-ass2) + (self.entry2-ass2) + self.entry_slope_second*(ass1 - self.entry1_second) + self.entry_slope*(ass1 - self.entry1)>1.5*self.out_pips/10000:
                self.Liquidate()
                self.Plot('Trade Plot Asset1','Liquidate1',ass1)
                self.Plot('Trade Plot Asset2','Liquidate2',ass2)
                self.Debug('Exiting Second Short')
                self.Debug(self.CashDispo/self.CashDispoInitial)
                self.is_invested = None
            #create second position
            elif self.is_invested == 'long' and portfolio < 1.3*lower_band:
                self.MarketOrder(self.NameAsset1, -2*self.QuantityAsset*reg.slope)
                self.MarketOrder(self.NameAsset2, 2*self.QuantityAsset)
                self.Debug('Entering Second Long')
                self.is_invested = 'longlong'
                self.entry1_second=ass1
                self.entry2_second=ass2
                self.entry_slope_second = reg.slope
            elif self.is_invested == 'short' and portfolio > 1.3*upper_band:
                self.MarketOrder(self.NameAsset1, 2*self.QuantityAsset*reg.slope)
                self.MarketOrder(self.NameAsset2, -2*self.QuantityAsset)
                self.Debug('Entering Second Short')
                self.is_invested = 'shortshort'
                self.entry1_second=ass1
                self.entry2_second=ass2
                self.entry_slope_second = reg.slope