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

class  MomentumAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.VERBOSEMODE = True
        self.VERBOSEMODE2 = False          
        self.SetStartDate(2019,5,10)
        self.SetEndDate(2019,5,12)
        self.SetCash(100000)
        self.benchmark = "SPY" 
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.0172))
        self.symbols = ["SPY","TLT","GLD"]
        self.lastSymbol = None
        self.contract = None

        self.spy = self.AddEquity("SPY",Resolution.Minute)
        self.spy.SetDataNormalizationMode(DataNormalizationMode.Raw)  
        self.spy.SetDataNormalizationMode(DataNormalizationMode.Raw)  
        self.tlt = self.AddEquity("TLT",Resolution.Minute)
        self.protectedPositions = {'SPY': False  , 'TLT': False , 'GLD': False}
        self.myOptions={}
        self.pauseOneMinute = True
        
        for symbol in self.Securities.Keys:
            self.Securities[symbol].FeeModel = ConstantFeeModel(2.00)
            self.Securities[symbol].SetDataNormalizationMode(DataNormalizationMode.Raw)  
            self.myOptions[symbol] = self.AddOption(symbol.Value,Resolution.Minute)
            #limit options strikes to 10 strikes away  with 6 day minimum
            self.myOptions[symbol].SetFilter(-10, +10, timedelta(0), timedelta(30))

    def OnData(self, slice):
        #if self.VERBOSEMODE2:self.Log("Time: {}" . format(slice.Time) )
        if slice is None:
            if self.VERBOSEMODE2:self.Log("no slice " )
            return    
        
        if self.contract is not None:
            if self.Securities[self.contract].Price != 0:
                orderId =  self.MarketOrder(self.contract, 1,asynchronous = False)
                self.contract = None
        optionchain = slice.OptionChains            
        for element  in slice:
            symbol = element.Key
            symbolStr = symbol.Value
            #make sure only act on invested symbol
            if self.pauseOneMinute:
                if symbolStr in  self.protectedPositions and  self.protectedPositions[symbolStr] == False:  
                        self.MarketOrder(symbol,100)
                        self.PrepCall(symbol)
                        self.pauseOneMinute = False
            else:
               if len(symbol.ToString()) > 3 : #workaround for ensuging not already an option
                    #self.MarketOrder(self.contract, 1,asynchronous = False)
                    self.SetHoldings(self.contract, 0.1)
                    self.pauseOneMinute = True

    #   except:
    #         if self.VERBOSEMODE:self.Log('An unknown error occurred trying OnData' +  str(sys.exc_info()[0]) )
                                    
    def calcNumberOptionsContractsNeeded(self,symbol):

        symbolStr = symbol.Value
        current = float(self.Securities[symbolStr].Price)
        quantity = float(self.Portfolio[symbolStr].Quantity)  
        optionContractCount = round(current * quantity / 100)
        if self.VERBOSEMODE: self.Log("contract count for symbol {} current {} quantity {} is {} " . format(symbolStr,current, quantity, optionContractCount))
        return optionContractCount
                    
    def PrepCall(self,symbol):
        if len(symbol.ToString()) > 3 : #workaround for ensuging not already an option
            return
        symbolStr = symbol.Value
        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time.date()) # must be of type Symbol
        if len(contracts) == 0: return
        min_expiry = 0
        max_expiry = 14
        
        filtered_contracts = [i for i in contracts if min_expiry <= (i.ID.Date.date() - self.Time.date()).days <= max_expiry]
        call = [x for x in filtered_contracts if x.ID.OptionRight == 0] 
        
        if len(call) == 0: return
        # sorted the contracts according to their expiration dates and choose the ATM options
        self.contract = sorted(sorted(call, key = lambda x: abs(self.Securities[symbolStr].Price - x.ID.StrikePrice)), 
                                        key = lambda x: x.ID.Date, reverse=True)[0]
      
        self.AddOptionContract(self.contract, Resolution.Minute)
        optionContractCount = self.calcNumberOptionsContractsNeeded(symbol)
        if optionContractCount == 0:
            return
        optionContractCount = 1
        #option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and  x.Value.Type==SecurityType.Option]
        return self.contract
        # orderId =  self.MarketOrder(self.contract, optionContractCount,asynchronous = False)
        # self.protectedPositions[symbolStr] = True

            
        if self.VERBOSEMODE: self.Log(  "bought 1 call contract of {} {} " . format(symbolStr,self.contract ))