from datetime import timedelta
import numpy as np
import math
import pandas as pd
from sklearn.linear_model import LinearRegression
from scipy.odr import Model, Data, ODR
from scipy import stats
from statistics import stdev
from pykalman import KalmanFilter

class VentralUncoupledShield(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020,3,1)    # Set Start Date
        #self.SetEndDate(2020,3,15)      # Set End Date
        self.SetEndDate(2020,3,11)      # Set End Date
        self.SetCash(10000) 
        self.currency1='GBPUSD'
        self.currency2='USDJPY'
        
        self.AddForex(self.currency1, Resolution.Daily, Market.Oanda)
        self.AddForex(self.currency2, Resolution.Daily, Market.Oanda)
        
        #kalman filter initialization
        self.delta = 1e-7 #process noise
        self.observation_cov=1e-4 #observation covariance
        self.counter=1
        self.means=0
        self.covs=0
        self.spread=0
        self.kf=0
        self.std_upper=0
        self.std_lower=0
        self.no_of_std=3
        self.trade_direction=''
        self.allocation=0.5
        self.spread_inc_per = 0.30#%increase in spread from last trade 
        self.spread_last    = 0#last spread at trade
        self.pred_currency2 = 0
        
        #charting
        self.Charts = dict()
        pairplot = Chart("Trade Plot")
        pairplot.AddSeries(Series(self.currency1,       SeriesType.Line,0))
        #pairplot.AddSeries(Series("Pred_"+self.currency2,SeriesType.Line,0))
        pairplot.AddSeries(Series(self.currency2,       SeriesType.Line,1))
        pairplot.AddSeries(Series("Pred_"+self.currency2,SeriesType.Line,1))
        pairplot.AddSeries(Series("Spread",             SeriesType.Bar, 2))
        pairplot.AddSeries(Series("Hedge_ratio",        SeriesType.Line,3))
        pairplot.AddSeries(Series("Intercept",          SeriesType.Line,4))
        pairplot.AddSeries(Series("STD_upper_bound",    SeriesType.Line,2))
        pairplot.AddSeries(Series("STD_lower_bound",    SeriesType.Line,2))
        self.AddChart(pairplot)
        
    def Kalman_update(self,value1,value2):
        if self.counter == 1:
            trans_cov = self.delta / (1 - self.delta) * np.eye(2)
            obs_mat = [value1, 1]
            self.kf = KalmanFilter(n_dim_obs=1, n_dim_state = 2,
                                  initial_state_mean = np.zeros(2),
                                  initial_state_covariance = np.ones((2, 2)),
                                  transition_matrices = np.eye(2),
                                  observation_matrices = obs_mat,
                                  observation_covariance = self.observation_cov,
                                  transition_covariance = trans_cov)
                              
            means, covs = self.kf.filter(np.asarray(value2))
            self.means = means[0]#mean for slope and intercept
            self.covs = covs[0]#2X2 variance for slope and intercept
            self.counter = self.counter + 1
            
        else:
            obs_mat = np.asarray([[value1, 1]])
            means, covs = self.kf.filter_update(self.means, 
                                               self.covs, 
                                               observation = value2, 
                                               observation_matrix = obs_mat)
            self.means = means.data
            self.covs = covs
            
    def kf_spread(self,value1,value2):
        self.pred_currency2=self.means[0]*value1+self.means[1]
        self.spread =  value2 - self.pred_currency2 
        
    def pair_quantity(self,value1,value2,hedge_ratio,fund):
        if hedge_ratio < 1:
            self.Debug('Hedge ratio {} is negative, need to swap the pair'.format(hedge_ratio))
            value1_unit = 0
            value2_unit = 0
        else:
            lot=math.floor(fund/(value1*hedge_ratio+value2*1))#get lot
            value1_unit = math.floor(hedge_ratio*lot)
            value2_unit = math.floor(lot)
        return value1_unit,value2_unit #number of unit to buy in currency 1 , currency 2
            
    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
        '''
      
        self.Kalman_update(data[self.currency1].Close,data[self.currency2].Close)
        self.kf_spread(data[self.currency1].Close,data[self.currency2].Close)
        #self.std_upper=np.sqrt(self.covs[0,0]) * self.no_of_std
        #self.std_lower=-1*self.std_upper
        self.std_upper= self.no_of_std
        self.std_lower= -self.no_of_std
        
        self.Debug('time:{} spread:{} slope:{} intercept:{} STD:{}'.format(data[self.currency1].Time,self.spread,self.means[0],self.means[1],np.sqrt(self.covs[1,1])))
        
        self.Plot("Trade Plot",self.currency1,data[self.currency1].Close)
        #self.Plot("Trade Plot","Pred_"+self.currency2,float(self.pred_currency2))
        self.Plot("Trade Plot",self.currency2,data[self.currency2].Close)
        self.Plot("Trade Plot","Pred_"+self.currency2,float(self.pred_currency2))
        self.Plot("Trade Plot","Spread",float(self.spread))
        self.Plot("Trade Plot","Hedge_ratio",float(self.means[0]))
        self.Plot("Trade Plot","Intercept",float(self.means[1]))
        self.Plot("Trade Plot","STD_upper_bound",float(self.std_upper))
        self.Plot("Trade Plot","STD_lower_bound",float(self.std_lower))


        if self.trade_direction == 'short_spread' and self.spread < 0 and self.Portfolio.Invested:
            self.Liquidate()#liquidate all
            self.trade_direction = ''
            self.spread_last = 0
        elif self.trade_direction == 'long_spread' and self.spread > 0 and self.Portfolio.Invested:
            self.Liquidate()#liquidate all
            self.trade_direction = ''
            self.spread_last = 0
        elif self.spread > self.std_upper and self.spread > (self.spread_last * (1+self.spread_inc_per)):
            currency1_unit,currency2_unit = self.pair_quantity(data[self.currency1].Close,data[self.currency2].Close,self.means[0],self.Portfolio.Cash*self.allocation)
            self.MarketOrder(self.currency2, -currency2_unit, True)
            self.MarketOrder(self.currency1, currency1_unit, True)
            self.trade_direction = 'short_spread'
            self.spread_last = self.spread
        elif self.spread < self.std_lower and self.spread < (self.spread_last * (1+self.spread_inc_per)):
            currency1_unit,currency2_unit = self.pair_quantity(data[self.currency1].Close,data[self.currency2].Close,self.means[0],self.Portfolio.Cash*self.allocation)
            self.MarketOrder(self.currency1, -currency1_unit, True)
            self.MarketOrder(self.currency2, currency2_unit, True)
            self.trade_direction = 'long_spread'
            self.spread_last = self.spread