# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from AlgorithmImports import *
from scipy.optimize import brentq

class OptionIndicatorsRegressionAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2014, 6, 5)
        self.set_end_date(2014, 6, 7)
        self.set_cash(100000)

        equity = self.add_equity("AAPL", Resolution.DAILY).symbol
        option = Symbol.create_option("AAPL", Market.USA, OptionStyle.AMERICAN, OptionRight.PUT, 650, datetime(2014, 6, 21))
        self.add_option_contract(option, Resolution.DAILY)
        # add the call counter side of the mirrored pair
        mirror_option = Symbol.create_option("AAPL", Market.USA, OptionStyle.AMERICAN, OptionRight.CALL, 650, datetime(2014, 6, 21))
        self.add_option_contract(mirror_option, Resolution.DAILY)

        self.delta = self.d(option, mirror_option, option_model = OptionPricingModelType.BINOMIAL_COX_ROSS_RUBINSTEIN, iv_model = OptionPricingModelType.BLACK_SCHOLES)
        self.gamma = self.g(option, mirror_option, option_model = OptionPricingModelType.FORWARD_TREE, iv_model = OptionPricingModelType.BLACK_SCHOLES)
        self.vega = self.v(option, mirror_option, option_model = OptionPricingModelType.FORWARD_TREE, iv_model = OptionPricingModelType.BLACK_SCHOLES)
        self.theta = self.t(option, mirror_option, option_model = OptionPricingModelType.FORWARD_TREE, iv_model = OptionPricingModelType.BLACK_SCHOLES)
        self.rho = self.r(option, mirror_option, option_model = OptionPricingModelType.FORWARD_TREE, iv_model = OptionPricingModelType.BLACK_SCHOLES)

        # A custom IV indicator with custom calculation of IV
        risk_free_rate_model = InterestRateProvider()
        dividend_yield_model = DividendYieldProvider(equity)
        self.implied_volatility = CustomImpliedVolatility(option, mirror_option, risk_free_rate_model, dividend_yield_model)
        self.register_indicator(option, self.implied_volatility, QuoteBarConsolidator(timedelta(1)))
        self.register_indicator(mirror_option, self.implied_volatility, QuoteBarConsolidator(timedelta(1)))
        self.register_indicator(equity, self.implied_volatility, TradeBarConsolidator(timedelta(1)))

        # custom IV smoothing function: assume the lower IV is more "fair"
        smoothing_func = lambda iv, mirror_iv: min(iv, mirror_iv)
        # set the smoothing function
        self.delta.implied_volatility.set_smoothing_function(smoothing_func)
        self.gamma.implied_volatility.set_smoothing_function(smoothing_func)
        self.vega.implied_volatility.set_smoothing_function(smoothing_func)
        self.theta.implied_volatility.set_smoothing_function(smoothing_func)
        self.rho.implied_volatility.set_smoothing_function(smoothing_func)

    def on_end_of_algorithm(self):
        if self.implied_volatility.current.value == 0 or self.delta.current.value == 0 or self.gamma.current.value == 0 \
        or self.vega.current.value == 0 or self.theta.current.value == 0 or self.rho.current.value == 0:
            raise Exception("Expected IV/greeks calculated")

        self.debug(f"""Implied Volatility: {self.implied_volatility.current.value},
Delta: {self.delta.current.value},
Gamma: {self.gamma.current.value},
Vega: {self.vega.current.value},
Theta: {self.theta.current.value},
Rho: {self.rho.current.value}""")

class CustomImpliedVolatility(ImpliedVolatility):
    def __init__(self, option, mirror_option, risk_free_rate_model, dividend_yield_model):
        super().__init__(option, risk_free_rate_model, dividend_yield_model, mirror_option, period=2)
        self.set_smoothing_function(lambda iv, mirror_iv: iv)

    def calculate_iv(self, time_till_expiry: float) -> float:
        try:
            return brentq(self.f, 1e-7, 2.0, args=(time_till_expiry), xtol=1e-4, maxiter=100)
        except:
            print("ImpliedVolatility.calculate_i_v(): Fail to converge, returning 0.")
            return 0.0

    # we demonstate put-call parity calculation here, but note that it is not suitable for American options
    def f(self, vol: float, time_till_expiry: float) -> float:
        call_black_price = OptionGreekIndicatorsHelper.black_theoretical_price(
            vol, UnderlyingPrice.current.value, self.strike, time_till_expiry, RiskFreeRate.current.value, DividendYield.current.value, OptionRight.CALL)
        put_black_price = OptionGreekIndicatorsHelper.black_theoretical_price(
            vol, UnderlyingPrice.current.value, self.strike, time_till_expiry, RiskFreeRate.current.value, DividendYield.current.value, OptionRight.PUT)
        return Price.current.value + OppositePrice.current.value - call_black_price - put_black_price