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, 7)
        self.SetEndDate(2020, 1, 12)
        self.SetCash(10000)


        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))

        # Add the option
        self.equity = self.AddEquity("SPY")
        option = self.AddOption("SPY")
        self.contract_multiplier = option.ContractMultiplier
        self.optionSymbol = option.Symbol

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


    def OnData(self,slice):
        
        for i in slice.OptionChains:
            if i.Key != self.optionSymbol: continue
            chain = i.Value
            calls = [x for x in chain if x.Right == OptionRight.Call]
            if len(calls) < 2:
                return
            
            latest_expiry = sorted(calls, key=lambda x: x.Expiry)[-1].Expiry
            calls = [x for x in calls if x.Expiry == latest_expiry]
            if len(calls) < 2:
                return
            
            sorted_by_strike = sorted(calls, key=lambda x: x.Strike)
            
            contract_1 = sorted_by_strike[0]
            contract_2 = sorted_by_strike[1]
            
            self.Order(contract_2.Symbol, 1)
            self.Order(contract_1.Symbol, -1)
                
            #strategy = OptionStrategies.BearCallSpread(self.optionSymbol, contract_1.Strike, contract_2.Strike, latest_expiry)
            #self.Order(strategy, 1)
            
            self.Quit(f"Contract_1 strike: {contract_1.Strike}; Contract_2 strike: {contract_2.Strike} Margin used: {self.Portfolio.TotalMarginUsed} Margin Remaining: {self.Portfolio.MarginRemaining}; Cash: {self.Portfolio.Cash}")