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(2021, 5, 5)
        self.SetCash(100000)
        self.AddEquity("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)
        chart = Chart("Buys")
        self.AddChart(chart)
        chart.AddSeries(Series("Buy signal", SeriesType.Scatter, "$", Color.Green, ScatterMarkerSymbol.Triangle))
        chart.AddSeries(Series("Sell signal", SeriesType.Scatter, "$", Color.Red, ScatterMarkerSymbol.TriangleDown))
        self.SetWarmup(100)
        self.SetSecurityInitializer(self.customSecurityInitializer)

    def customSecurityInitializer(self, security):            
        bar = self.GetLastKnownPrice(security)
        security.SetMarketPrice(bar)
        security.SetLeverage(1000)
        '''
        class MySecurityInitializer(SecurityInitializer):
            def Initialize(self, security):
                security.SetLeverage(1000)    

                self.SetSecurityInitializer(SecurityInitializer())
        '''
    def OnData(self, data):
        if self.IsWarmingUp:
            return
        if not self.ema13.IsReady:
            return  
        if not self.ema48.IsReady:
            return
        if not self.rsi.IsReady: 
            return 
        #self.Plot("Buys", "ema13", self.ema13.Current.Value)
        #self.Plot("Buys", "rsi", self.rsi.Current.Value)
        #self.Plot("Buys", "ema48", self.ema48.Current.Value) 
    # Update the option chain data
        if self.option_chain is None:
            self.option_chain = self.OptionChainProvider.GetOptionContractList("SPY",self.Time )

    # Filter the option chain by strike price and expiration date
            
    # 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
                    filtered_chain = sorted(self.option_chain, key = lambda c: c.ID.OptionRight == OptionRight.Call and c.ID.Date > self.Time and c.ID.StrikePrice > 200)
                    self.Plot("Buys", "Buy Signal", self.ema48.Current.Value) 
                    self.Debug("DDD")
                    num_contracts = int(self.Portfolio.TotalHoldingsValue * 0.1 / self.Securities["SPY"].Price)
            # Buy call options with the nearest expiration date
                    self.AddOptionContract(filtered_chain[0], Resolution.Minute)
                    self.Buy(filtered_chain[0], 1)
                    

    # Filter the option chain by strike price and expiration date
           
    # 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
                    self.Debug("DDD")
                    filtered_chain = sorted(self.option_chain, lambda c: c.ID.OptionRight == OptionRight.Put and c.ID.Date > self.Time and c.ID.StrikePrice < 200)
                    self.Plot("Buys", "Sell Signal", self.ema48.Current.Value) 
                    num_contracts = int(self.Portfolio.TotalHoldingsValue * 0.1 / self.Securities["SPY"].Price)
            # Buy put options with the nearest expiration date
                    self.AddOptionContract(filtered_chain[0], Resolution.Minute)
                    self.Buy(filtered_chain[0], 1)

    def OnEndOfDay(self):
        self.Plot("Buys", "ema13", self.ema13.Current.Value)
        self.Plot("Buys", "rsi", self.rsi.Current.Value)
        self.Plot("Buys", "ema48", self.ema48.Current.Value)