from AlgorithmImports import *
import QuantConnect.Securities.Option as option
import QuantConnect.Data.Custom as iqfeed
from QuantConnect.DataSource import *

class OptionTradingAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 12, 31)
        self.SetCash(100000)
        self.AddSecurity("SPY")

        # Set up the 13 and 48 EMA indicators
        self.ema13 = self.EMA("SPY", 13)
        self.ema48 = self.EMA("SPY", 48)

        # Set up the RSI indicator
        self.rsi = self.RSI("SPY", 14)

        # Set up a trailing stoploss of 10%
        #self.SetTrailingStop(0.1)
        self.stopLoss = 0.1

        # Enable option trading
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        #self.SetOptionChainProvider(QuantConnect.Interfaces.IOptionChainProvider)
        # Subscribe to the SPY option chain
        self.AddOption("SPY")
        self.option_chain = None
        #option.SetFilter(-2, +2, 0, 180)

class MySecurityInitializer(SecurityInitializer):
    def Initialize(self, security):
        security.SetLeverage(1000)    

        self.SetSecurityInitializer(SecurityInitializer())

    def OnData(self, data):
        if not self.IsWarmingUp:
                if not self.ema13.IsReady:
                    return  
                if not self.ema48.IsReady:
                    return
                if not self.RSI.IsReady: 
                    return  
    # Update the option chain data
        if self.option_chain is None or not self.option_chain.IsValid:
            self.option_chain = self.OptionChainProvider.GetOptionContractList(self.Time, "SPY")

    # Filter the option chain by strike price and expiration date
            filtered_chain = self.option_chain.Filter((lambda c: c.Right == option.OptionRight.Call)
                                              and (lambda c: c.Expiry > self.Time)
                                              and (lambda c: c.StrikePrice > 200))
    # Check if the 13 EMA has crossed above the 48 EMA
            if self.ema13.Current.Value > self.ema48.Current.Value:
            # Check if the RSI is above 54
                if self.rsi.Current.Value > 54:
            # Calculate the number of option contracts to buy
            # using 10% of holdings
                    num_contracts = int(self.Portfolio.TotalHoldingsValue * 0.1 / filtered_chain.UnderlyingLastPrice)
            # Buy call options with the nearest expiration date
                    self.Buy(option.OptionSymbol.Create(filtered_chain.UnderlyingSymbol, filtered_chain.Expiry, filtered_chain.StrikePrice, option.OptionRight.Call), num_contracts)

    # Filter the option chain by strike price and expiration date
            filtered_chain = self.option_chain.Filter((lambda c: c.Right == option.OptionRight.Put)
                                              and (lambda c: c.Expiry > self.Time)
                                              and (lambda c: c.StrikePrice < 200))
    # Check if the 13 EMA has crossed under the 48 EMA
            if self.ema13.Current.Value < self.ema48.Current.Value:
        # Check if the RSI is below 46
                if self.rsi.Current.Value < 46:
            # Calculate the number of option contracts to buy
            # using 10% of holdings
                    num_contracts = int(self.Portfolio.TotalHoldingsValue * 0.1 / filtered_chain.UnderlyingLastPrice)
            # Buy put options with the nearest expiration date
                    self.Buy(option.OptionSymbol.Create(filtered_chain.UnderlyingSymbol, filtered_chain.Expiry, filtered_chain.StrikePrice, option.OptionRight.Put), num_contracts)