import numpy as np
import mlfinlab as ml
import pandas as pd


class CalibratedResistanceAtmosphericScrubbers(QCAlgorithm):


    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2019, 3, 1) 
        self.SetCash(100000)
        self.spy = self.AddEquity("SPY", Resolution.Hour)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        
        # init close prices
        self.open = np.array([])
        self.high = np.array([])
        self.low = np.array([])
        self.close = pd.Series()
        self.volume = np.array([])
        self.lookback = 50
        self.SetWarmUp(self.lookback * 2)


    def OnData(self, data):
        if "SPY" not in data.Bars:
            return
        
        open_ = data["SPY"].Open
        high_ = data["SPY"].High
        low_ = data["SPY"].Low
        close_ = data["SPY"].Close
        volume_ = data["SPY"].Volume
        self.open = np.append(self.open, close_)[-self.lookback*2:]
        self.high = np.append(self.high, close_)[-self.lookback*2:]
        self.low = np.append(self.low, close_)[-self.lookback*2:]
        self.close = self.close.append(pd.Series([close_], index=[self.Time]))[-self.lookback*2:]
        self.volume = np.append(self.volume, close_)[-self.lookback*2:]
        self.time = self.Time
        
        if self.IsWarmingUp:
            return
        
        df = pd.DataFrame({'open': self.open, 'high': self.high, 'low': self.low, 'close': self.close, 'volume': self.volume})
        daily_vol = ml.util.get_daily_vol(self.close, lookback=100)