import numpy as np
import pandas as pd
from datetime import datetime, timedelta

class ZIVTradingStrategy(QCAlgorithm):

    ## order ticket for stop order
    stopMarketTicket = None
    stopMarketFillTime = datetime.min
    highestPrice = 0
    
    def Initialize(self):
       self.SetStartDate(2015, 1, 2)  # Set Start Date
       self.SetCash(10000)  # Set Strategy Cash
               
       self.ziv = self.AddEquity("ZIV", Resolution.Daily)
       self.ziv.SetDataNormalizationMode(DataNormalizationMode.Raw)
       
       # Creates an indicator and adds to a rolling window when it is updated
       self.rocfiveday = self.ROC("ZIV", 5, Resolution.Daily)
       self.rocfiveday.Updated += self.ROCFiveDayUpdated
       self.rocfiveDayWindow = RollingWindow[IndicatorDataPoint](5)
       
       self.roctenday = self.ROC("ZIV", 10, Resolution.Daily)
       self.roctenday.Updated += self.ROCTenDayUpdated
       self.roctenDayWindow = RollingWindow[IndicatorDataPoint](5)
       
       self.momfiveday = self.MOM("ZIV", 5, Resolution.Daily)
       self.momfiveday.Updated += self.MOMFiveDayUpdated
       self.momfiveDayWindow = RollingWindow[IndicatorDataPoint](5)
       
       self.momtenday = self.MOM("ZIV", 10, Resolution.Daily)
       self.momtenday.Updated += self.MOMTenDayUpdated
       self.momtenDayWindow = RollingWindow[IndicatorDataPoint](5)
       
       # get ATR to use for stop loss.
       self.atr = self.ATR("ZIV", 20, Resolution.Daily)
       # get volatility lookback
       self.lookback = 30
       self.SetWarmUp(self.lookback)
       # set up volatility calculation. this is more efficient instead of calling historical data in onData
       self.roc1 = self.ROC("ZIV", 1, Resolution.Daily)
       self.std = IndicatorExtensions.Of(StandardDeviation(self.lookback), self.roc1)
       # define portfolio volatility
       self.target_portfolio_sigma = 0.15
       
       # this is for listening to order status
       self.lastOrderEvent = None 
       self.ticket = None # Flag for position status

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
                    Arguments:
                        data: Slice object keyed by symbol containing the stock data
                '''
        if self.IsWarmingUp:
            return
        
        # check that 1 day have passed since we hit the stop loss.
        if(self.Time - self.stopMarketFillTime).days < 2:
            return ## this means wait
        
        long_signal_1 = self.momfiveDayWindow[0].Value > float(0)
    
        if not self.Portfolio.Invested:
            if long_signal_1 == True:
                volatility = volatility = self.std.Current.Value * np.sqrt(252) # get vol.
                weight = self.target_portfolio_sigma / volatility # get weight.
                quantity = (self.Portfolio.TotalPortfolioValue * weight)/self.Securities["ZIV"].Close # get quantity.
                self.ticket = self.MarketOrder("ZIV", quantity)
                self.stopMarketTicket = self.StopMarketOrder("ZIV", -quantity, self.Securities["ZIV"].Close - 2 * self.atr.Current.Value)
                
        elif self.ticket is not None and (self.UtcTime - self.ticket.Time).days < 3: ## to update stop loss.
            if self.Securities["ZIV"].Close > self.highestPrice:
                self.highestPrice = self.Securities["ZIV"].Close # save the new highestPrice
                updateFields = UpdateOrderFields()
                updateFields.StopPrice = self.highestPrice - 2 * self.atr.Current.Value ## updating stop price
                self.stopMarketTicket.Update(updateFields)
                self.Debug("ZIV: " + str(self.highestPrice) + " Stop: " + str(updateFields.StopPrice))
        
        elif self.ticket is not None and self.UtcTime >= self.ticket.Time + timedelta(days = 3):
            self.Liquidate("ZIV")
            self.highestPrice = 0
            self.ticket = None
                
                
    # define methods to add indicator values to a rolling window.
    def MOMFiveDayUpdated(self, sender, updated):
        self.momfiveDayWindow.Add(updated)
        
    def MOMTenDayUpdated(self, sender, updated):
        self.momtenDayWindow.Add(updated)
        
    # Adds updated values to rolling window
    def ROCFiveDayUpdated(self, sender, updated):
        self.rocfiveDayWindow.Add(updated)
         
    def ROCTenDayUpdated(self, sender, updated):
        self.roctenDayWindow.Add(updated)
         
    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status != OrderStatus.Filled:
            return ## if no orders were filled then wait and dont worry.
        
        # this checks that the stop order ticket was submitted
        if self.stopMarketTicket is not None and self.stopMarketTicket.OrderId == orderEvent.OrderId:
            self.stopMarketFillTime = self.Time
            self.Debug(self.stopMarketFillTime)