import datetime
import math

from QuantConnect.Securities.Option import OptionPriceModels

from datetime import timedelta
from QuantConnect.Data.UniverseSelection import * 

SYMBOL = "SPY"
TARGET_CASH_UTILIZATION = 1
REBALANCE_THRESHOLD = 0.01
STARTING_CASH = 100000
TARGET_CONTRACT_DAYS_IN_FUTURE = 35
SHOULD_LOG = True

class CoveredCallAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)
        # self.SetEndDate(2017, 4, 1)
        self.SetCash(STARTING_CASH)
        self.TargetDelta = float(self.GetParameter("target_delta"))
        self.LimitOrderRatio = float(self.GetParameter("limit_order_ratio"))
        self.InitialSpyValue = None
        
        equity = self.AddEquity(SYMBOL, Resolution.Minute)
        option = self.AddOption(SYMBOL, Resolution.Minute)
        self.symbol = option.Symbol
        self.UniverseSettings.Resolution = Resolution.Daily

        # Look for options that are within 30 days of the target contract date.
        # Options should be at or 60 strikes above present price.
        option.SetFilter(0, 60, timedelta(max(0, TARGET_CONTRACT_DAYS_IN_FUTURE - 30)), timedelta(TARGET_CONTRACT_DAYS_IN_FUTURE + 30))
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        
        # Give a generous warm-up to allow the options price model to initialize.
        self.SetWarmUp(TimeSpan.FromDays(10))
        
        self.SetBenchmark(equity.Symbol)
        self.LastDay = None
 
    def OnData(self, slice):
        """Processes the slice up to 1x per day."""
        if self.IsWarmingUp:
            return
        
        if self.LastDay == self.Time.day:
            return
        self.LastDay = self.Time.day
        
        self.BalancePortfolio(slice)
        self.TradeOptions(slice)
        
    def BalancePortfolio(self, slice):
        """Transacts the underlying symbol as needed to arrive at the target cash utilization.
        
        Buys or sells as necessary to keep the underlying symbol's value at TARGET_CASH_UTILIZATION
        of the total portoflio value. Does not transact unless the imbalance is greater than REBALANCE_THRESHOLD.
        """
        if SYMBOL in slice and slice[SYMBOL] is not None:
            # Get the current value and plot to main figure window.
            current_spy = slice[SYMBOL].Close
            if self.InitialSpyValue is None:
                self.InitialSpyValue = current_spy
            self.Plot("Strategy Equity", "SPY", STARTING_CASH * current_spy / self.InitialSpyValue)
            
            # Calculate how many shares should be owned to hit target cash utilization.
            value = self.Portfolio.TotalPortfolioValue
            desired_n_shares = math.floor(value / current_spy)
            actual_n_shares = self.Portfolio[SYMBOL].Quantity
            delta_shares = desired_n_shares - actual_n_shares
            
            # Transact if needed.
            if SHOULD_LOG:
                self.Log(f"Portfolio value: {value:.2f} Current spy: {current_spy:.2f} desired_n_shares: {desired_n_shares} "
                         f"actual_n_shares: {actual_n_shares} delta_shares: {delta_shares}")
            if abs(delta_shares * current_spy) > value * REBALANCE_THRESHOLD:
                self.MarketOrder(SYMBOL, delta_shares)
 
    def TradeOptions(self, slice):
        """If enough shares exist to use as collateral, sells a call and places a limit order for that call."""
            
        number_of_contracts_available = math.floor(self.Portfolio[SYMBOL].Quantity / 100)
        number_of_contracts = -sum([x.Value.Quantity for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option])
        options_to_buy = number_of_contracts_available - number_of_contracts
        
        if options_to_buy < 1:
            return
        
        if SHOULD_LOG:
            self.Log(f"I presently own {self.Portfolio[SYMBOL].Quantity} shares of spy.")
            self.Log(f"I presently have {number_of_contracts_available} that I can hold and "
                     f"only own {number_of_contracts}. I plan to buy {options_to_buy} more")

        chain = slice.OptionChains.GetValue(self.symbol)
        contract = self.GetClosestContract(chain, self.Time + datetime.timedelta(TARGET_CONTRACT_DAYS_IN_FUTURE))
        if contract is None:
            return
        
        if SHOULD_LOG:
            self.Log(f"Selling contract: Price: {contract.AskPrice} Underlying Price: {contract.UnderlyingLastPrice:1.2f} "
                     f"Strike: {contract.Strike:1.2f} Expiry: {contract.Expiry} Delta: {contract.Greeks.Delta:1.2f}")
        self.Sell(contract.Symbol, options_to_buy)
        self.LimitOrder(contract.Symbol, options_to_buy, round(contract.AskPrice * self.LimitOrderRatio, 2))
        
    def GetClosestContract(self, chain, desired_expiration):
        """Gets the contract nearest the target delta and expiry."""
        if not chain:
            return None
            
        calls = [contract for contract in chain if contract.Right == OptionRight.Call]
        if not calls:
            return None
        
        # Calculate the option expiry date nearest the target.
        available_expirations = list({contract.Expiry for contract in calls})
        nearest_expiration = sorted(available_expirations, key=lambda expiration: abs(expiration-desired_expiration))[0]
        
        # For all contracts that match the target expiration, find the one with delta nearest target.
        calls_at_target_expiration = [contract for contract in calls if contract.Expiry == nearest_expiration]
        calls_at_target_expiration = sorted(calls_at_target_expiration, key = lambda contract: abs(contract.Greeks.Delta - self.TargetDelta))
        if not calls_at_target_expiration:
            return None
        return calls_at_target_expiration[0]

from Selection.OptionUniverseSelectionModel import OptionUniverseSelectionModel
from datetime import date, timedelta

class OptionsUniverseSelectionModel(OptionUniverseSelectionModel):
    def __init__(self, select_option_chain_symbols):
        super().__init__(timedelta(1), select_option_chain_symbols)

    def Filter(self, filter):
        ## Define options filter -- strikes +/- 3 and expiry between 0 and 180 days away
        return (filter.Strikes(-2, +2)
                      .Expiration(timedelta(0), timedelta(180)))