# region imports
from AlgorithmImports import *
# endregion

import re
class TestCallSell(QCAlgorithm):

    def Initialize(self):
        # set start/end date for backtest
        self.SetStartDate(2013, 12, 1)
        self.SetEndDate(2014, 5, 31)
        # set starting balance for backtest
        self.SetCash(100000)
        # add the underlying asset
        self.equity = self.AddEquity("QQQ", Resolution.Daily)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.symbol = self.equity.Symbol
        self.callContract = str()
        self.contractsAdded = set()

        # parameters ------------------------------------------------------------
        self.callDaysB4Exp = 2 # number of days before expiry to exit calls
        self.callDTE = 25 # target days till expiration
        self.callOTM = 1.00 # target percentage OTM of call
        self.percentage = 0.9 # percentage of portfolio for underlying asset
        self.options_alloc = 100 # 1 option for X num of shares (balanced would be 100)
        # ------------------------------------------------------------------------

    def OnData(self, data):
        if(self.IsWarmingUp):
            return
        
        # buy underlying asset
        if not self.Portfolio[self.symbol].Invested:
            self.Debug(f"Buying {self.symbol} {self.Time}")
            self.SetHoldings(self.symbol, self.percentage)
        
        # buy calls
        self.BuyCall(data)

        # close call before it expires
        if self.callContract:
            db4 = self.callContract.ID.Date - self.Time
            if self.Time > datetime(2014,3,1) and self.Time < datetime(2014,3,11):
                self.Debug(f"Call exp {self.callContract.ID.Date} db4e {db4}")
            if (self.callContract.ID.Date - self.Time) <= timedelta(self.callDaysB4Exp):
                self.Liquidate(self.callContract)
                self.Debug(f"Closed call: too close to expiration {db4} {self.Time.date()}")
                self.callContract = str()

    def BuyCall(self, data):
        # get option data
        if self.callContract == str():
            self.callContract = self.CallOptionsFilter(data)
            return
        
        # if not invested and option data added successfully, buy option
        elif not self.Portfolio[self.callContract].Invested and data.ContainsKey(self.callContract):
            nc = round(self.Portfolio[self.symbol].Quantity / self.options_alloc)
            q = self.Portfolio[self.symbol].Quantity
            mr = self.Portfolio.MarginRemaining
            mu = self.Portfolio.TotalMarginUsed
            dt = self.callContract.ID.Date.date()
            self.Log(f"BC Buying calls {self.Time.date()} nc {nc} q {q} mu {mu} X {dt}")
            self.Buy(self.callContract, nc)

    def CallOptionsFilter(self, data):
        ''' OptionChainProvider gets a list of option contracts for an underlying symbol at requested date.
            Then you can manually filter the contract list returned by GetOptionContractList.
            The manual filtering will be limited to the information included in the Symbol
            (strike, expiration, type, style) and/or prices from a History call '''

        contracts = self.OptionChainProvider.GetOptionContractList(self.symbol, data.Time)
        self.underlyingPrice = self.Securities[self.symbol].Price
        # filter the out-of-money call options from the contract list which expire close to self.callDTE num of days from now
        otm_calls = [i for i in contracts if i.ID.OptionRight == OptionRight.Call and
                                            i.ID.StrikePrice >= self.callOTM * self.underlyingPrice and
                                            self.callDTE - 8 < (i.ID.Date - data.Time).days < self.callDTE + 8]
        if len(otm_calls) > 0:
            # sort options by closest to self.callDTE days from now and desired strike, and pick first
            contract = sorted(sorted(otm_calls, key = lambda x: abs((x.ID.Date - self.Time).days - self.callDTE)),
             key = lambda x: self.underlyingPrice - x.ID.StrikePrice,reverse=True)[0]
#            self.Debug(f"call selected {contract.ID.Date.date()} str {contract.ID.StrikePrice}")
            if contract not in self.contractsAdded:
#                self.Debug(f"Buy call {contract.ID.StrikePrice} {self.symbol} {'%.2f'%self.underlyingPrice}")
                self.contractsAdded.add(contract)
                # use AddOptionContract() to subscribe the data for specified contract
                self.AddOptionContract(contract, Resolution.Daily)
            return contract
        else:
            return str()

    def OnOrderEvent(self, orderEvent):
        if self.Time > datetime(2014,1,1) and self.Time < datetime(2014,5,31):
            self.Debug(f" Order {str(orderEvent)}")

        if 0 and orderEvent.Status == OrderStatus.Filled:
            self.Debug(f"Cash {self.Portfolio.Cash} Holdings {self.Portfolio.TotalHoldingsValue} Val {self.Portfolio.TotalPortfolioValue}")
            self.Debug(str(orderEvent))

    def OnAssignmentOrderEvent(self, assignmentEvent: OrderEvent) -> None:
        self.Debug(f"Assignment event: {str(assignmentEvent)}")
# just set a flag to reset holdings in OnData
        self.NeedReset = True

    def OnEndOfAlgorithm(self):
        self.Debug(f"Final Report")
#        self.Liquidate()