#region imports
from AlgorithmImports import *
#endregion
from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from datetime import timedelta

class BasicTemplateOptionsAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        # self.SetEndDate(2020, 1, 1)
        self.SetCash(10000)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        equity = self.AddEquity("SPY", Resolution.Minute)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.underlyingsymbol = "SPY"
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(9, 32), Action(self.CheckExpiry))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(TimeSpan.FromMinutes(10)), Action(self.CheckRSI))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(TimeSpan.FromMinutes(5)), Action(self.TrailStop))
        
        self.rsi = self.RSI("SPY", 10)
        self.rsiWindow = RollingWindow[float](3)

        self.sellBucket = []
        self.buyBucket = []

        self.SetBenchmark("SPY")
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        self.SetWarmUp(timedelta(7))

    def CheckRSI(self):
        if not self.rsi.IsReady:
            return
        if not self.IsMarketOpen:
            return

        start_time = self.Time.replace(hour=9, minute=30, second=0)
        stop_time = self.Time.replace(hour=12, minute=0, second=0)

        RSI = self.rsi.Current.Value
        self.rsiWindow.Add(RSI)

        if not self.rsiWindow.IsReady:
            return
        
        if self.Time < start_time:
            return
        #if self.Time > stop_time:
            #return

        r1 = self.rsiWindow[0]
        r3 = self.rsiWindow[2]

        if r1 < 45 and r3 > 45:
            contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
            self.TradeOptions(contracts)

    def TradeOptions(self,contracts):
        if len(contracts) == 0:
            self.Debug("No contracts found")
            return
        filtered_contracts = self.InitialFilter(self.underlyingsymbol, contracts, 20, 70)
        if filtered_contracts == None:
            self.Debug("No filtered contracts found")
            return
        call = [i for i in filtered_contracts if i.ID.OptionRight == 0]
        put = [i for i in filtered_contracts if i.ID.OptionRight == 1]

        self.trade_contracts = sorted(call, key = lambda x: x.ID.StrikePrice, reverse = True)
        asset = self.AddOptionContract(self.trade_contracts[0], Resolution.Minute)
        symbol = asset.Symbol

        Margin = self.Portfolio.MarginRemaining
        if Margin < 1:
            return
        if symbol in self.buyBucket:
            return
        price = self.Securities[symbol].Price * 100
        if price < 1:
            return
        quantity = 5
        totalvalue = quantity * price
        if totalvalue > Margin:
            return
        self.ticketOne = self.MarketOrder(symbol, quantity)
        self.buyBucket.append(symbol)
        n = self.ticketOne
        a = str(n)
        b = n
        n = Ticket()
        n.Symbol = b.Symbol
        n.buyPrice = self.Securities[symbol].Price
        n.Entry = self.Time
        n.Quantity = b.Quantity
        n.Stop = n.buyPrice * .5
        n.Limit = n.buyPrice * 1.1
        myTicks.append(n)

    def InitialFilter(self, underlyingsymbol, symbol_list, min_expiry, max_expiry):
            
        if len(symbol_list) == 0:
            self.Debug("No Symbols Found")
            return
        
        contract_list = []
        
        for i in symbol_list:
            if i.ID.Date.date() > self.Time.date() + timedelta(20):
                if i.ID.Date.date() < self.Time.date() +timedelta(70):
                    contract_list.append(i)
        
        near_money = []
        price = self.Securities[underlyingsymbol].Price
        target = price * .98
        wiggle = price * 1.02

        for c in contract_list:
            if target <= c.ID.StrikePrice <= wiggle:
                near_money.append(c)

        filtered_contracts = near_money
        number = len(near_money)
        if number < 1:
            self.Debug("no contracts near money")
            return

        return filtered_contracts

    def TrailStop(self):
        if not self.IsMarketOpen:
            return
        
        if self.IsWarmingUp:
            return
        
        for t in myTicks:
            symbol = t.Symbol
                
            if self.Portfolio[symbol].Invested:
                if symbol in self.sellBucket:
                    return
                quantity, lastPrice = self.Portfolio[symbol].Quantity, self.Portfolio[symbol].Price
                if quantity < 1:
                    return
                
                if lastPrice >= t.Limit:
                    t.Limit = lastPrice * 1.05
                    t.Stop = lastPrice * .95
                
                if lastPrice >= t.buyPrice * 2:
                    self.MarketOrder(symbol, -quantity, False, "Doubled")
                    myTicks.remove(t)
                    self.sellBucket.append(symbol)

                if lastPrice <= t.Stop:
                    if lastPrice < t.buyPrice:
                        self.MarketOrder(symbol, -quantity, False, "Sold at Loss")
                        myTicks.remove(t)
                        self.sellBucket.append(symbol)
                    else:
                        self.MarketOrder(symbol, -quantity, False, "Stop Loss Trigger")
                        self.sellBucket.append(symbol)
                        myTicks.remove(t)

    def CheckExpiry(self):
        for t in myTicks:
            symbol = t.Symbol
            entry = t.Entry
            if symbol in self.sellBucket:
                return
            
            if self.Portfolio[symbol].Invested:
                quantity, expiry = self.Portfolio[symbol].Quantity, self.Securities[symbol].Expiry

                if expiry < self.Time + timedelta(2):
                    self.sellTicket = self.MarketOrder(symbol, -quantity, False, "Sold Before Expiration")
                    self.sellBucket.append(symbol)
                    myTicks.remove(t)

                if self.Time > t.Entry + timedelta(3):
                    self.sellTicket = self.MarketOrder(symbol, -quantity, False, "Sold After 3 Days")
                    self.sellBucket.append(symbol)
                    myTicks.remove(t)
        
        if len(self.sellBucket) >= 60:
            del self.sellBucket[0]

        if len(self.buyBucket) >= 60:
            del self.buyBucket[0]

    def DumpOptions(self):
        for t in myTicks:
            symbol = t.Symbol
            if symbol in self.sellBucket:
                return
            if self.Portfolio[symbol].Invested:
                self.Liquidate(symbol)
                self.sellBucket.append(symbol)
                myTicks.remove(t)

    def OnEndOfDay(self):
        self.Plot("Indicators","RSI", self.rsi.Current.Value)

class Ticket:
    def _init_(self, symbol, buyPrice, price, quantity, stop, limit, entry, oPrice):
        self.Symbol = symbol
        self.Price = price
        self.Quantity = quantity
        self.buyPrice = buyPrice
        self.Stop = stop
        self.Limit = limit
        self.Entry = entry
        self.oPrice = oPrice
myTicks = []