#region imports
from AlgorithmImports import *
#endregion

from QuantConnect.DataSource import *

import pandas as pd
import numpy as np
from scipy import stats

class VX_TS_Algorithm(QCAlgorithm):


    # ------------------------------------------------------------------------------------------------
    def Init_system_part(self) -> None:

        # self.SetStartDate(2012,1,1) 
        self.SetStartDate(2020,1,1) 
        self.SetEndDate(datetime.now().date() - timedelta(1))
        self.SetCash(10000)

    # ------------------------------------------------------------------------------------------------
    def Init_equity_part(self) -> None:

        self.SetBenchmark(self.AddEquity("SPY", self.var_resolution).Symbol)

        # Define the Equity 
        self.sym_svxy = self.AddEquity("SVXY", self.var_resolution).Symbol

        self.equity_vxx = self.AddEquity("VXX", self.var_resolution)
        self.equity_vxx.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.sym_vxx = self.equity_vxx.Symbol

        self.option_vxx = self.AddOption("VXX", self.var_resolution)
        self.option_vxx.SetFilter(-3, 3, timedelta(40), timedelta(60))



    # ------------------------------------------------------------------------------------------------
    def UpdateRollingWindow(self, sender, updated) -> None:
        if self.vix3m_vix_RY_over_sma.IsReady:
            self.vix3m_vix_RY_over_sma_window.Add(self.vix3m_vix_RY_over_sma.Current.Value)
            self.vix_window.Add(self.vix_sma.Current.Value)

    # ------------------------------------------------------------------------------------------------
    def Init_vix_part(self) -> None:

        # Define the symbol and "type" of our generic data
        self.sym_vix = self.AddData(CBOE, 'VIX', self.var_resolution).Symbol
        self.sym_vix3m = self.AddData(CBOE, 'VIX3M', self.var_resolution).Symbol

        # Set up default Indicators, these are just 'identities' of the closing price
        self.vix_sma = self.SMA(self.sym_vix, 1, self.var_resolution)
        self.vix3m_sma = self.SMA(self.sym_vix3m, 1, self.var_resolution)
        
    # ------------------------------------------------------------------------------------------------
    def Init_vx_futures_part(self) -> None:

        self.var_vxData = None

        # VIX Future data symbol
        self.sym_vx = self.AddData(VIXCentralContango, "VX", self.var_resolution).Symbol

    # ------------------------------------------------------------------------------------------------
    def Initialize(self) -> None:

        self.Init_system_part()

        # Define the internal variable
        self.var_resolution = Resolution.Daily

        self.Init_equity_part()
        self.Init_vix_part()
        self.Init_vx_futures_part()

    # ------------------------------------------------------------------------------------------------
    def OnData(self, slice: Slice) -> None:

        # Check data readiness ---------------------------------------------
        data_missing = False

        if not slice.ContainsKey(self.sym_vxx): 
            self.Debug(f"vx_ts> {self.Time} ERROR VXX data missing..." )
            data_missing = True

        if not slice.ContainsKey(self.sym_svxy): 
            self.Debug(f"vx_ts> {self.Time} ERROR SVXY data missing..." )
            data_missing = True

        # Get the VX data
        if (not slice.ContainsKey(self.sym_vx)) or (slice[self.sym_vx] is None): 
            self.Debug(f"vx_ts> {self.Time} ERROR VX data of VIXCentralContango missing..." )
            data_missing = True

        # Wait for all indicators to fully initialize
        if not (self.vix_sma.IsReady and self.vix3m_sma.IsReady ): 
            self.Debug(f"vx_ts> {self.Time} ERROR VIX, VIX3M Indicator etc. data missing..." )
            data_missing = True

        if data_missing:
            return

        # Data is ready, start data process ----------------------------------