from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from datetime import datetime, timedelta
from QuantConnect.Data.Custom.CBOE import CBOE
import numpy as np


class WhoPredictsWho(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2011, 11, 1)
        self.SetEndDate(datetime.today())
        self.SetCash(1000000)
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        # Leveraged ETF:
        self.uvxy = self.AddEquity("UVXY", Resolution.Daily).Symbol
        self.SetBenchmark("SPY")
        # The algorithm does not perform intraday operations:
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(
                                    InsightWeightingPortfolioConstructionModel(
                                                            lambda time: None))

        # VIX data:
        self.vix = self.AddData(CBOE, 'VIX', Resolution.Daily).Symbol

        # Initialize Alpha and Universe Selection
        self.AddAlpha(PredictWho(self, self.vix, self.spy, self.uvxy))
        
class PredictWho(AlphaModel):
    """
    A very simple volatility-predicted VIX model.
    """

    def __init__(self, algorithm, vix_symbol, spy_symbol,
                 uvxy_symbol):

        # Symbols:
        self.vix = vix_symbol
        self.spy = spy_symbol
        self.uvxy = uvxy_symbol
        #Holding Period:
        self.holding_period = 22
        self.gain_target = 0.10 #Placeholder

        
    def Update(self, algorithm, data):
        insights = []
        time = timedelta(days=self.holding_period)
        i_type = InsightType.Price
        if not data.HasData:
            return []
        
        # History calls do load the backtest a lot. It is better to use
        # "deque" objects to hold this short (22 day) closing data. In this 
        # case the test is very short and is not taking up much time:
        history = algorithm.History(self.spy, 22, Resolution.Daily)
        vol_22D = np.sqrt(252)*history['close'].rolling(22).std()[-1]
        
        if vol_22D > 150: 
            for security in algorithm.Securities.Values:
                if security.Invested:
                    direction = InsightDirection.Flat
                    insight = Insight(security.Symbol,
                              time, i_type, direction, self.gain_target,
                              1, "PredictWho", 1)
                    insights.append(insight)
            return insights
        
        if vol_22D > 80:
            if algorithm.Securities[self.spy].Invested:
                direction = InsightDirection.Flat
                insight = Insight(self.spy,
                                  time, i_type, direction, self.gain_target,
                                  1, "PredictWho", 1)
                insights.append(insight)
            if algorithm.Securities[self.uvxy].Invested:
                    return insights
            direction = InsightDirection.Up
            insight = Insight(self.uvxy,
                              time, i_type, direction, self.gain_target,
                              1, "PredictWho", 1)
            insights.append(insight)
            return insights

        if vol_22D < 80:
            if algorithm.Securities[self.spy].Invested:
                    return []
            direction = InsightDirection.Up
            insight = Insight(self.spy,
                              time, i_type, direction, self.gain_target,
                              1, "PredictWho", 1)
            insights.append(insight)
                
        return insights