from QuantConnect.Indicators import *
import decimal as d
from pykalman import KalmanFilter
import numpy as np


### <summary>
### In this example we are looking for price to breakout above the bollinger bands
### and look to buy when we see that. We hold our position until price touches the 
### middle band of the bollinger bands.
###
# ----------------------------------------
CRYPTO = 'BTCUSD'; PERIODSHORT = 10; PERIOD = 30; PERIODLONG = 30; LEV = 1.0;
# ----------------------------------------

class BollingerBreakoutAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2020, 1, 1)  #Set Start Date
        self.SetEndDate(2021, 11, 1)    #Set End Date
        self.SetCash(10000)             #Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.GDAX, AccountType.Cash)
        
        # define crypto we want to trade on
        # ETHUSD, LTCUSD or BTCUSD
        #self.target_crypto = "ETHUSD"
        self.target_crypto = "BTCUSD"
        
        self.crypto = self.AddCrypto(self.target_crypto, Resolution.Daily).Symbol
        self.rollingWindow = RollingWindow[TradeBar](PERIODLONG)
        self.Consolidate(self.crypto, Resolution.Daily, self.CustomBarHandler)
        self.SetWarmUp(1*PERIODLONG, Resolution.Daily)

 
        # create a momentum indicator over 3 days
        #self.mom = self.MOMP(self.target_crypto, 3)
        self.mom = self.MOM(self.target_crypto, 3)
        #self.fastavg = self.EMA(self.target_crypto, 3);
        self.slowmom = self.MOM(self.target_crypto, 30);
        #self.sslowmom = self.ROC(self.target_crypto, 60);
        self.sslowmom = self.MOM(self.target_crypto, 60);
        
        
        # Plot Momentum
        self.PlotIndicator(
            "Indicators",
            self.mom
        )

        # set warmup period
        #self.SetWarmUp(20)
        
    def CustomBarHandler(self, bar):
        self.rollingWindow.Add(bar)

    def OnData(self, data):
        if not self.rollingWindow.IsReady: return
    
        L = np.flipud(np.array([self.rollingWindow[i].Low for i in range(PERIOD)]))

        self.kf = KalmanFilter(transition_matrices = [1],
                 observation_matrices = [1],
                 initial_state_mean = 0,
                 initial_state_covariance = 1,
                 observation_covariance=1,
                 transition_covariance=.01)
        
        kf,_ = self.kf.filter(L)
        

        self.Plot(self.crypto, "Price", self.Securities[self.crypto].Price)        
        self.Plot(self.crypto, "Kalman", kf[-1])

    
        holdings = self.Portfolio[self.target_crypto].Quantity
        price = self.Securities[self.target_crypto].Close
        fastmom = self.mom.Current.Value
        slowmom = self.slowmom.Current.Value
        sslowmom = self.sslowmom.Current.Value
        mom = sslowmom + (slowmom) + (2 * fastmom)
        
        # buy if price closes above upper bollinger band
        if holdings <= 0:
             #if price > self.Bolband.UpperBand.Current.Value:
             if price > kf[-1] and mom > 0:
                    self.SetHoldings(self.target_crypto, 1.0)
        
        # sell if price closes below middle bollinger band
        #if holdings > 0 and price < self.Bolband.MiddleBand.Current.Value:
        #if holdings > 0 and price < fastavg:
        if holdings >= 0:
            if price < kf[-1] and mom < 0:
                self.Liquidate()