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

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Securities.Option import OptionPriceModels
from QuantConnect.Securities.Option import OptionStrategies
import pandas as pd
import numpy as np
import math 

class QuantumVerticalInterceptor(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 1, 7)
        self.SetCash(10000)
        self.DELTA_TARGET=0.5
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        self.df_calls = None
        #set a risk limit of 1% of portfolio value
        self.investment_limit = self.Portfolio.TotalPortfolioValue * 0.01 
        # Add the option
        option = self.AddOption("SPY")
        self.optionSymbol = option.Symbol

        # Add the initial contract filter
        option.SetFilter(-5, +5, 5, 10)

        # Define the Option Price Model
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        #option.PriceModel = OptionPriceModels.BlackScholes()
        #option.PriceModel = OptionPriceModels.AdditiveEquiprobabilities()
        #option.PriceModel = OptionPriceModels.BaroneAdesiWhaley()
        #option.PriceModel = OptionPriceModels.BinomialCoxRossRubinstein()
        #option.PriceModel = OptionPriceModels.BinomialJarrowRudd()
        #option.PriceModel = OptionPriceModels.BinomialJoshi()
        #option.PriceModel = OptionPriceModels.BinomialLeisenReimer()
        #option.PriceModel = OptionPriceModels.BinomialTian()
        #option.PriceModel = OptionPriceModels.BinomialTrigeorgis()
        #option.PriceModel = OptionPriceModels.BjerksundStensland()
        #option.PriceModel = OptionPriceModels.Integral()

        # Set warm up with 30 trading days to warm up the underlying volatility model
        self.SetWarmUp(30, Resolution.Daily)


    def OnData(self,slice):
        self.Plot("Portfolio", "Margin Remaining", self.Portfolio.MarginRemaining)         # Remaining margin on the account
        self.Plot("Portfolio", "Margin Used", self.Portfolio.TotalMarginUsed)         # Sum of margin used across all securities

        if self.IsWarmingUp or not slice.OptionChains.ContainsKey(self.optionSymbol):
            return

        chain = slice.OptionChains[self.optionSymbol]
        #set float format so delta displays correctly
        pd.set_option('display.float_format', lambda x: '%.5f' % x)
        #put the relevant data into the dataframe
        df = pd.DataFrame([[x.Right,float(x.Strike),x.Expiry,float(x.BidPrice),float(x.AskPrice),x.Greeks.Delta,x.UnderlyingLastPrice] for x in chain],
        index=[x.Symbol.Value for x in chain],
        columns=['type', 'strike', 'expiry', 'ask price', 'bid price', 'delta','underlyinglast'])
        #ensure expiry column is in datetime format
        df['expiry'] = pd.to_datetime(df['expiry'])
        # sort by expiry, descending
        df.sort_values(by=['expiry'],ascending=False)
        # get the most future date
        furthest_date = df['expiry'].iloc[0]
        # keep only those rows which have that furthest date
        df = df[df.expiry == furthest_date]
        #split the dataframe into calls and puts (calls are 0, puts are 1)
        self.df_calls = df[df.type==0]
        #sort by delta
        self.df_calls.sort_values(by=['delta'],ascending=False)

        #select the closest two records to the DELTA TARGET
        #try:
        uppercall_ind = self.df_calls[self.df_calls.delta<self.DELTA_TARGET].delta.idxmax()
        lowercall_ind = self.df_calls[self.df_calls.delta>self.DELTA_TARGET].delta.idxmin()
        self.df_calls  = self.df_calls[self.df_calls.index.isin([lowercall_ind,uppercall_ind])]
        spread_value = self.df_calls.at[lowercall_ind,'bid price'] - self.df_calls.at[uppercall_ind,'ask price']
        max_risk = self.df_calls.at[uppercall_ind,'strike'] - self.df_calls.at[lowercall_ind,'strike']
        max_risk_contract = max_risk * 100
        max_investment = math.trunc(self.investment_limit / max_risk_contract)
        self.Sell(OptionStrategies.BearCallSpread(self.optionSymbol, self.df_calls.at[lowercall_ind,'strike'], self.df_calls.at[uppercall_ind,'strike'] , self.df_calls.at[uppercall_ind,'expiry']), max_investment)
        #except:
        #    return