from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Data import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from QuantConnect.Python import PythonQuandl
# from QuantConnect.Data.Custom import Quandl

from datetime import timedelta

import pandas as pd


class GaboVolatilityZscoreCubedWithTimer(QCAlgorithm):
    def Initialize(self):
        # Backtest initialization
        self.SetStartDate(2012, 1, 1)
        self.SetEndDate(2018, 3, 11)
        self.SetCash(15000)

        # Constant definitions
        self.ASSET_UNDERLYING = 'SPY'
        self.ASSET_TRADE_LONG = 'SVXY'
        self.ASSET_TRADE_SHORT = 'TZA'
        self.BENCHMARK = 'SVXY'

        self.INITIAL_TIME = 20
        self.MIN_TIME = 0
        self.MAX_TIME = 40
        self.TIME_INCREASE = 1
        self.TIME_DECREASE = .5
        

        self.DAYS_DMA_LOW = 50
        self.DAYS_DMA_HIGH = 200
        self.DAYS_DMA_SHORT_TERM = 3

        self.PCT_INVEST_SHORT = .99
        self.PCT_INVEST_LONG = .99

        self.OPEN_TRADNG_HOURS = [30, 60, 180, 210, 360]
        self.CLOSE_TRADING_HOURS = [15]

        self.ZSCORE_ADJ_SELL_LONG = 0
        self.ZSCORE_ADJ_SELL_SHORT = -0.05
        self.ZSCORE_ADJ_BUY_LONG = -.01
        
        self.ZSCORE_EXP_LONG_EXP = 3
        self.ZSCORE_EXP_LONG_MULT = 1.2
        self.ZSCORE_EXP_LONG_TIMER_MULT = -1/60
        self.ZSCORE_EXP_LONG_STDDEV_MULT = 1/20
        
        self.ZSCORE_EXP_SHORT_EXP = 3
        self.ZSCORE_EXP_SHORT_MULT = .9
        self.ZSCORE_EXP_SHORT_TIMER_MULT = -1/45
        self.ZSCORE_EXP_SHORT_STDDEV_MULT = -1/20

        self._daily_history = {}

        # Initialization process
        # Add the equities to universe
        self._asset_underlying = self.AddEquity(self.ASSET_UNDERLYING, Resolution.Minute)
        self._asset_underlying.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)
        
        self._asset_trade_long = self.AddEquity(self.ASSET_TRADE_LONG, Resolution.Minute)
        self._asset_trade_short = self.AddEquity(self.ASSET_TRADE_SHORT, Resolution.Minute)

        # Work with split adjustment (not dividend adjustment)
        # self._asset_underlying.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)
        # self._asset_trade_long.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)
        # self._asset_trade_short.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)
        
        # variable initialization
        self._time = self.INITIAL_TIME

        self.SetBenchmark(self.BENCHMARK)
        
        oneDayConsolidator = self.ResolveConsolidator(self.ASSET_UNDERLYING, Resolution.Daily)
        oneDayConsolidator.DataConsolidated += self.daily_history_consolidator
        self.SubscriptionManager.AddConsolidator(self.ASSET_UNDERLYING, oneDayConsolidator)

        # Store the underlying history prices in memory, to avoid waiting the
        # warmup period
        underlying_id = str(self._asset_underlying.Symbol.ID)
        self._daily_history[underlying_id] = self.History(
            [self.ASSET_UNDERLYING],
            self.DAYS_DMA_HIGH,
            Resolution.Daily
        )

        # Set trading events
        for minutes in self.OPEN_TRADNG_HOURS:
            self.Schedule.On(
                self.DateRules.EveryDay(self.ASSET_UNDERLYING),
                self.TimeRules.AfterMarketOpen(self.ASSET_UNDERLYING, minutes),
                Action(self._trade)
            )

        for minutes in self.CLOSE_TRADING_HOURS:
            self.Schedule.On(
                self.DateRules.EveryDay(self.ASSET_UNDERLYING),
                self.TimeRules.BeforeMarketClose(self.ASSET_UNDERLYING, minutes),
                Action(self._trade)
            )

        # Timer adjustment
        self.Schedule.On(
            self.DateRules.EveryDay(self.ASSET_UNDERLYING),
            self.TimeRules.BeforeMarketClose(self.ASSET_UNDERLYING, 1),
            Action(self._timer_management)
        )

    def daily_history_consolidator(self, sender, bar):
        # ticker = bar.Symbol.Value
        ticker_id = str(bar.Symbol.ID)
        current_quotes_df = pd.DataFrame(
            {
                'open': [float(bar.Open)],
                'high': [float(bar.High)],
                'low': [float(bar.Low)],
                'close': [float(bar.Close)],
                'volume': [int(bar.Volume)]
            },
            index=[pd.to_datetime(self.Time)]
        )

        if ticker_id in self._daily_history:
            self._daily_history[ticker_id] = self._daily_history[ticker_id] \
                .append(current_quotes_df)
        else:
            self._daily_history[ticker_id] = current_quotes_df

    def _trade(self):
        underlying_id = str(self._asset_underlying.Symbol.ID)

        if underlying_id not in self._daily_history:
            # No history enough
            return

        history_size = self._daily_history[underlying_id].shape[0]
        if history_size < self.DAYS_DMA_HIGH:
            # No history enough
            return

        # Get all historic quotes
        underlying_price = float(self.Securities[self.ASSET_UNDERLYING].Close)

        history_low = self._daily_history[underlying_id].iloc[-self.DAYS_DMA_LOW:]
        history_high = self._daily_history[underlying_id].iloc[-self.DAYS_DMA_HIGH:]
        history_short_term = self._daily_history[underlying_id].iloc[-self.DAYS_DMA_SHORT_TERM:]

        # Compute all the z-scores
        sma_low = history_low['close'].mean()
        sma_high = history_high['close'].mean()
        sma_short_term = history_short_term['close'].mean()

        std_low = history_low['close'].std()
        std_high = history_high['close'].std()

        z_score = (sma_low - sma_high) / std_high

        z_score_current = (underlying_price - sma_high) / std_high
        z_score_short_term = (sma_short_term - sma_high) / std_high

        z_score_exp_long = (
            z_score ** self.ZSCORE_EXP_LONG_EXP * self.ZSCORE_EXP_LONG_MULT +
            self._time * self.ZSCORE_EXP_LONG_TIMER_MULT + 
            std_low * self.ZSCORE_EXP_LONG_STDDEV_MULT
        )
        
        z_score_exp_short = (
            z_score ** self.ZSCORE_EXP_SHORT_EXP * self.ZSCORE_EXP_SHORT_MULT + 
            self._time * self.ZSCORE_EXP_SHORT_TIMER_MULT +
            std_low * self.ZSCORE_EXP_SHORT_STDDEV_MULT
        )
        
        # self.Log('date: %s - zscore: %.02f - current zscore: %.2f' % (
        #     str(self.Time), 
        #     z_score, 
        #     z_score_current
        #     )
        # )
        # self.Log('underlying price: %.04f' % underlying_price)

        if not self.Transactions.GetOpenOrders():
            cash_pct = self.Portfolio.Cash / self.Portfolio.TotalPortfolioValue
            
            if z_score_current + self.ZSCORE_ADJ_SELL_LONG <= z_score_exp_long:
                self.SetHoldings(self.ASSET_TRADE_LONG, 0)

            if z_score_current + self.ZSCORE_ADJ_SELL_SHORT > z_score_exp_short:
                self.SetHoldings(self.ASSET_TRADE_SHORT, 0)

            if z_score_current < z_score_exp_short and z_score_current < z_score_exp_long:
                if cash_pct > 0.05 * self.PCT_INVEST_SHORT:
                    self.SetHoldings(self.ASSET_TRADE_SHORT, self.PCT_INVEST_SHORT)

            # if z_score_current < z_score_short_term - 0.01 and z_score_current > z_score_exp_long:
            if (z_score_current < z_score_short_term + self.ZSCORE_ADJ_BUY_LONG 
                    and z_score_current > z_score_exp_long):
                if cash_pct > 0.05 * self.PCT_INVEST_LONG:
                    self.SetHoldings(self.ASSET_TRADE_LONG, self.PCT_INVEST_LONG)
                    # self.Log('zs current: %.04f - zs short: %0.4f - zse long: %.04f' % (
                    #     z_score_current, z_score_current, z_score_exp_long))
                    

    def _timer_management(self):
        cash_pct = self.Portfolio.Cash / self.Portfolio.TotalPortfolioValue

        if cash_pct > 0.3 or (
                self.Portfolio[self.ASSET_TRADE_SHORT].Quantity > 0 and
                self._time < self.MAX_TIME):
            self._time += self.TIME_INCREASE

        if cash_pct < 0.1:
            self._time -= self.TIME_DECREASE

        if self._time < self.MIN_TIME:
            self._time = self.MIN_TIME