# class TradeManagement:
#     '''Manager for trades for each security'''
#     def __init__(self, algorithm, symbol):
        
#         '''
#         stoploss = None
#         takeProfit = None'''
        
        
        
#     def EnterTrade(self, quantity,  takeprofit, stoploss):
        
        
#         '''''
        
        
#         '''''
    
    
#     # do this minutely/hourly
#     def UpdatePosition(self):
        
#         #current_price > takeprofit
#         #liquidate
        
    
    
#     # close out 
#     def timelimits(self):
        
        

#https://www.quantconnect.com/forum/discussion/8399/optionstrategies-limit-orders
import numpy as np

class QuantumParticleAtmosphericScrubbers(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 2, 8)  # Set Start Date #2019, 12,8
        self.SetEndDate(2018,5,8)
        self.SetCash(100000)  # Set Strategy Cash
        
        self.spy_symbol = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.Securities[self.spy_symbol].SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.SetSecurityInitializer(self.CustomSecurityInitializer) #TO WARM SHIT UP WHEN ADDED!!
        
        self.straddleTickets = {}
        
        self.max_util = .5
        self.lmts = False
        
        self.last_option_exp = None
        
        # credit/max_risk

    def OnData(self, data):
        
        if not data.Bars.ContainsKey(self.spy_symbol):
            return
        
        ## To handle any early assignment (Could loop through univ tickers here.)
        #if self.Portfolio[self.spy_symbol].Invested:
        #    self.Liquidate(self.spy_symbol)
        
        ''' Original (Straddle)
        if not self.Portfolio.Invested:
            
            price = data.Bars[self.spy_symbol].Close;
            expiry = self.Time + timedelta(days = 14)
            
            call_symbol, put_symbol = self.GetStraddleContracts(self.spy_symbol, price, expiry)
            
            #Why made into a call / put here (vs call / put symbol?) ***************
            call = self.AddOptionContract(call_symbol, Resolution.Minute)
            put = self.AddOptionContract(put_symbol, Resolution.Minute)
            
            call_tickets, put_tickets = self.TradeStraddle(call, put, 2) #(Tgt, Stop) format
        
            self.straddleTickets[self.spy_symbol] = [call_tickets, put_tickets] #[(Tgt, Stp), (Put Tgt, Put Stp)]
            #Addit
            self.last_option_exp = call_symbol.ID.Date
        '''
        
        if not self.Portfolio.Invested:
            #HERE is where you would loop (thorugh securities)
            
            price = data.Bars[self.spy_symbol].Close
            expiry = self.Time + timedelta(days = 30)
            
            #Just using 10% strikes for IC here -- NEXT UP: greeks
            icp = price * 1.10
            ocp = price * 1.20
            
            ipp = price * .9
            opp = price * .8
            
            #Lookup Option Contracts
            o_call, call, put, o_put = self.IronCondor(self.spy_symbol, expiry, icp, ocp, ipp, opp)
            
            # #Add Respective Symbols (To be Traded) -- DO THIS inside Execute
            # oc_sym = self.AddOptionContract(o_call, Resolution.Minute)
            # c_sym = self.AddOptionContract(call, Resolution.Minute)
            if o_call and call and put and o_put:
                to_buy =  [o_call, o_put]
                to_sell = [call,     put]
                
                self.ExecuteSpread(to_buy, to_sell, use_limits = self.lmts)
                
            
        if self.Portfolio.Invested:
            self.CloseExpiring(data, False)
                
            
            
            
        ''' 
        SAMPLE (Getting Options data from Securities or Portfolio -- once added/held):
            ## https://www.quantconnect.com/forum/discussion/3549/getting-option-contract-info-from-portfolio-holding-objects-in-python/p1
            ## So the OPTION uses the KEY aspect, really KEY.ID !! (CHECK if Option first.)
            for kvp in self.Securities:
                symbol = kvp.Key
                id = symbol.ID
                security_type = id.SecurityType
                if security_type = SecurityType.Option
                    expiry = id.Date
                    strike_price = id.StrikePrice
        
        #ZO Initial Attempt (Cleaned up Below in CloseExpiring)
        if self.last_option_exp and self.Portfolio.Invested:
            
            days_to_exp = (self.last_option_exp - self.Time).days
            self.Debug(f'DTE -- {days_to_exp}')
            
            if days_to_exp < 1:
                
                self.Debug(f'Closing Prior to Exit.') #Should we check for ITM? 
                option_posits = [symbol for symbol in self.Portfolio.Keys if self.Portfolio[symbol].Invested and \
                                            symbol.SecurityType == SecurityType.Option]
                                            
                for symbol in option_posits:
                    self.Liquidate(symbol)
        '''
    
    def GetICMargin(self, buys, sells):
        #calls = [i for i in buys + sells if i.Right == OptionRight.Call]
        #puts = [i for i in buys + sells if i.Right == OptionRight.Put]
        #self.GetVerticalMargin() #Weird to now sort, etc. Faster to just use abs diff
        
        #MUST BE ADDED (in Securities) PRIOR TO CALLING THIS!
    
        credits = np.sum([self.Securities[symbol].BidPrice for symbol in sells])
        debits = np.sum([self.Securities[symbol].AskPrice for symbol in buys]) #[symbol.ID.AskPrice
        
        credit = credits - debits
        
        call_stks = [i.ID.StrikePrice for i in buys + sells if i.ID.OptionRight == OptionRight.Call]
        put_stks = [i.ID.StrikePrice for i in buys + sells if i.ID.OptionRight == OptionRight.Put]
        

        self.Debug(f'Calls: {[i for i in call_stks]}')
        self.Debug(f'Puts: {[i for i in put_stks]}')
        call_dist = abs(max(call_stks) - min(call_stks))
        put_dist = abs(max(put_stks) - min(put_stks))
        
        #call_dist = abs(np.diff([i.ID.StrikePrice for i in buys + sells if i.ID.OptionRight == OptionRight.Call]))
        #put_dist = abs(np.diff([i.ID.StrikePrice for i in buys + sells if i.ID.OptionRight == OptionRight.Put]))
        
        distance = max(call_dist, put_dist)
        self.Debug(f'credit {credit} -- dist {distance}')
        return (distance - credit) * 100
        
                
            
                
    #Not going to be used here... Easier to track buys / sells  ^^          
    def GetVerticalMargin(self, buy, sell):
        credit = sell.BidPrice - buy.AskPrice #Ex: $1.50
        distance = abs(buy.Strike - sell.Strike) #Ex: 5 
        
        return (distance - credit) * 100 #100 shares per contract
        
        
        
    def ExecuteSpread(self, buys, sells, use_limits = False):

        added = [self.AddOptionContract(c) for c in buys + sells]
        
        if buys[0] not in self.Securities: return 
        #if not set(buys+sells).issubset(self.Securities): return  #More accurate, confusing.
    
        #usd_per_ct = self.GetICMargin(buys, sells)
        #self.Debug(f'USD per ct -- {usd_per_ct}')
        #num_cts = int(self.Portfolio.MarginRemaining / usd_per_ct * self.max_util)
        pf = (self.Portfolio.MarginRemaining * self.max_util)
        
        num_cts = 2 #max(num_cts,5)
        
        ## ---------------- ADDIT: APPROX of Margin
        
        call_dist = abs(buys[0].ID.StrikePrice - sells[0].ID.StrikePrice)
        put_dist = abs(buys[1].ID.StrikePrice - sells[1].ID.StrikePrice)

        distance = max(call_dist, put_dist)
        credit = abs(self.Securities[sells[0]].BidPrice - self.Securities[buys[0]].AskPrice) + \
                abs(self.Securities[sells[1]].BidPrice - self.Securities[buys[1]].AskPrice)
        
        self.Debug(f'Distance - credit * 100 --- {distance} - {credit} * 100 -- > {abs(distance - credit) * 100} ')
        if max(call_dist, put_dist) == 0:
            #Approximate
            ul = self.Securities[sells[0]].Underlying.Price * .2 
            k = sells[0].ID.StrikePrice
            m1 = .2 * ul - abs(k - ul)
            usd_per_cts = (m1 - credit) * 100
            self.Debug(f'Strangle Calc ----  per_ct: {m1}')
        else:
            usd_per_ct = (distance - credit) * 100
            
        num_cts = int(pf / usd_per_ct)
        
        ##  ---------------------- End Addit 
        self.Debug(f'Num Contracts -- {num_cts}')
        for buy in buys:
            if use_limits:
                ask = self.Securities[buy].AskPrice #Does this need to be buy.Symbol ? 
                self.LimitOrder(buy, num_cts, ask)
            else:
                self.MarketOrder(buy, num_cts)
        
        for sell in sells:
            if use_limits:
                bid = self.Securities[sell].BidPrice
                self.LimitOrder(sell, -num_cts, bid)
            else:
                self.MarketOrder(sell, -num_cts)
                
                
            
    def IronCondor(self, symbol, expiry, call, otm_call, put, otm_put):
        #Returns oc, ic, ip, op
        inner_call, outer_call = self.GetVertical(symbol, OptionRight.Call, call, otm_call, expiry)
        inner_put, outer_put = self.GetVertical(symbol, OptionRight.Put, put, otm_put, expiry)
        
        return outer_call, inner_call, inner_put, outer_put
        

        ''' Original -- Cleaned this up
        if inner_call and outer_call and inner_put and outer_put:
            
            ## For Limits ***
            #Selling inners (need bid)
            # call_price = self.Securities[inner_call].BidPrice
            # put_price = self.Securities[inner_put].BidPrice
            
            # #Buying Outers (need ask)
            # otm_call_p = self.Securities[outer_call].AskPrice
            # otm_put_p = self.Securities[outer_put].AskPrice
            
            self.MarketOrder(outer_call, 1)
            self.MarketOrder(inner_call, -1)
            
            self.MarketOrder(inner_put, -1)
            self.MarketOrder(outer_put, 1)
        '''
            
        
        
    def GetVertical(self, symbol, right, inner, outer, expiry):
        #Get Contracts (OCP, bc no greeks, iv)
        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time)
        
        #Get Expiry
        expiry_sorted = sorted(contracts, key=lambda k: abs(k.ID.Date - expiry), reverse=False) #Asc
        closest = expiry_sorted[0].ID.Date
        correct_exp = [symbol for symbol in contracts if symbol.ID.Date == closest]
        
        #Select Side (Call / Put)
        rights = [symbol for symbol in correct_exp if symbol.ID.OptionRight == right]
        
        if len(rights) == 0: return None, None
        
        #self.Debug(f'Inner: {inner} ----- Outer: {outer}') #THIS is correct... 
        #Select Strike
        inner_ct = sorted(rights, key=lambda k: abs(k.ID.StrikePrice - inner), reverse=False)[0]
        outer_ct = sorted(rights,key=lambda k: abs(k.ID.StrikePrice - outer), reverse=False)[0]
        
        #Not sure if I need to split this? (Maybe Just SORT by strike, and take FIRST, then NEXT?)
        
        #self.Debug(f' (EOVertical) -- Inner: {inner_ct.ID.StrikePrice} ---- Outer: {outer_ct.ID.StrikePrice}')
        return inner_ct, outer_ct
            
        
        
        
    def CloseExpiring(self, data, only_itm=False):
        ## QUESTION -- how to get to OCP type contract from PORTFOLIO, or SECURITIES
        #open_pos_exps = [symbol.ID.Date for symbol, val in self.Portfolio.items() \
        #                if val.Invested and symbol.SecurityType == SecurityType.Option ][0]

        #dte = (open_pos_exps - self.Time).days
        
        #ul_price = data.Bars[self.spy_symbol].Close
        
        to_close = []
        for kvp in self.Portfolio:
            symbol = kvp.Key #OPTION symbol I assume? -- HOW to get Underlying?
            value = kvp.Value
            if symbol.ID.SecurityType == SecurityType.Option: 
                
                #Check if Expiring (If not, Ignore)
                expiry = symbol.ID.Date
                dte = (expiry - self.Time).days
                if dte > 1: continue
                
                #Close ALL expiring
                if not only_itm:
                    self.Liquidate(symbol)
                    continue
                
                #Close expiring IF ITM
                elif only_itm:
                    
                    strike = symbol.ID.StrikePrice #Can also check ITM here!!
                    
                    ##How to get Underlying Price from OPTION?
                    ul = symbol.Underlying
                    ul_price = data.Bars[ul].Close
                    
                    if symbol.ID.OptionRight == OptionRight.Call:
                        if strike > ul_price:
                            self.Liquidate(symbol)

                    else:
                        if strike < ul_price:
                            self.Liquidate(symbol)
                


    def GetStraddleContracts(self, symbol, strike, expiry):
        
        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time)
        
        expiry_sorted = sorted(contracts, key=lambda k: abs(k.ID.Date - expiry), reverse=False)
        closest_expiry = expiry_sorted[0].ID.Date
        
        contracts_with_desired_expiry = [symbol for symbol in contracts if symbol.ID.Date == closest_expiry]
        
        calls = [symbol for symbol in contracts_with_desired_expiry if symbol.ID.OptionRight == OptionRight.Call]
        
        puts = [symbol for symbol in contracts_with_desired_expiry if symbol not in calls]
        
        sorted_calls = sorted(calls, key=lambda k: abs(k.ID.StrikePrice - strike), reverse=False)
        
        sorted_puts = sorted(puts, key=lambda k: abs(k.ID.StrikePrice - strike), reverse=False)
        
        return sorted_calls[0], sorted_puts[0]
        
    ## Interesting use of Stops / Targets (With OnOrder to make them OCO!!) VERY clever...
    def TradeStraddle(self, call, put, quantity):
        
        call_entry= call.AskPrice
        call_entry_ticket = self.MarketOrder(call.Symbol, quantity)
        
        call_take_profit = call_entry * 1.50
        call_stop_loss = call_entry * 0.90
        
        call_profit_ticket = self.LimitOrder(call.Symbol, -quantity, call_take_profit)
        call_loss_ticket = self.StopMarketOrder(call.Symbol, -quantity, call_stop_loss)
        
        put_entry = put.AskPrice
        put_entry_ticket = self.MarketOrder(put.Symbol, quantity)
        
        put_take_profit = put_entry * 1.50
        put_stop_loss = put_entry * 0.90
        
        put_profit_ticket = self.LimitOrder(put.Symbol, -quantity, put_take_profit)
        put_loss_ticket = self.StopMarketOrder(put.Symbol, -quantity, put_stop_loss)
        
        return (call_profit_ticket, call_loss_ticket), (put_profit_ticket, put_loss_ticket)
        
    
    def CustomSecurityInitializer(self, security):
        bar = self.GetLastKnownPrice(security)
        security.SetMarketPrice(bar)
        
    # OCO Logic (For Exits) --- CAN we get FILL price in here? I suppose it's in PF object anyway, but still.
    #Make the Stop / Tgt OCO's basically.
    def OnOrderEvent(self, orderevent):
        
        if orderevent.Status != OrderStatus.Filled:
            return
        
        #FILLED orderID
        orderId = orderevent.OrderId
        
        for underlying, order_pairs in self.straddleTickets.items():
            for order_pair in order_pairs:
                order_one = order_pair[0]
                order_two = order_pair[1]
                #IF tgt filled, cancel Stop (Vice Versa)
                if orderId == order_one.OrderId:
                    order_two.Cancel()
                elif orderId == order_two.OrderId:
                    order_one.Cancel()