#!/user/bin/env python
# -*- coding: utf-8 -*-
"""
Copyright (C) 2024 Galaxy Technologies
Licensed under the Apache License, Version 2.0
"""
import numpy as np
from scipy.stats import norm
from pricelib.common.utilities.enums import UpDown, BarrierType, PaymentType, InOut, ExerciseType
from pricelib.common.processes import StochProcessBase
from pricelib.common.date_handler import global_evaluation_date
from pricelib.common.pricing_engine_base.analytic_engine_base import (AnalyticEngine, bs_formula,
                                                                      euro_asset_or_nothing, euro_cash_or_nothing)


class AnalyticBarrierEngine(AnalyticEngine):
    """障碍期权闭式解定价引擎
    欧式障碍期权：
        欧式香草与欧式资产或无、欧式现金或无的解析解的组合
    美式障碍期权：
        不支持同时输入非零的现金返还和不等于1的参与率，二者只有其一时支持。
        敲入期权现在支付期权费，但是当到期前资产价格触及障碍水平时，期权才生效。若一直没有发生敲入，则到期时支付现金返还rebate
        敲出期权现在支付期权费，但是当到期前资产价格触及障碍水平时，期权就失效了。若到期前发生敲出事件，则立刻支付现金返还rebate
        Merton(1973), Reiner & Rubinstein(1991a)提出障碍期权解析解，
        Broadie, Glasserman和Kou(1995)提出均匀离散观察障碍期权近似解"""

    def __init__(self, stoch_process: StochProcessBase = None, for_haug=False, *,
                 s=None, r=None, q=None, vol=None):
        """
        Args:
            stoch_process: 随机过程StochProcessBase对象
            for_haug: bool，Haug(1998)的双边障碍解析解，需要使用单边障碍期权解析解的公式直接计算出来的值，这种情况下不能应用
                            已敲入时payoff变为香草收益、或者已敲出时payoff变为现金返还的修正.
        在未设置stoch_process时，(stoch_process=None)，会默认创建BSMprocess，需要输入以下变量进行初始化
            s: float，标的价格
            r: float，无风险利率
            q: float，分红/融券率
            vol: float，波动率
        """
        super().__init__(stoch_process=stoch_process, s=s, r=r, q=q, vol=vol)
        self.for_haug = for_haug

    # pylint: disable=invalid-name, too-many-locals, missing-docstring, too-many-branches, too-many-statements
    def calc_present_value(self, prod, t=None, spot=None):
        """计算现值
        Args:
            prod: Product产品对象
            t: datetime.date，估值日; 如果是None，则使用全局估值日globalEvaluationDate
            spot: float，估值日标的价格，如果是None，则使用随机过程的当前价格
        Returns: float，现值
        """
        # 参数检查
        self._check_product_params(prod)

        calculate_date = global_evaluation_date() if t is None else t
        tau = prod.trade_calendar.business_days_between(calculate_date, prod.end_date) / prod.t_step_per_year  # 从估值日到到期日的交易日时间（年）
        # maturity = (prod.end_date - calculate_date).days / prod.annual_days.value  # 从估值日到到期日的自然日时间（年）
        tenor = (prod.stop_date - prod.start_date).days / prod.annual_days.value  # 从起始日到终止日的自然日时间（年）
        # 保底收益率
        _floor_yield = 0
        if prod.floor_yield != 0:  # 保底收益率需要输入期初价格s0作为名义本金
            _floor_yield = prod.floor_yield * tenor * prod.s0
        # 敲出/未敲入现金返还
        if prod.rebate_annual:
            _rebate = prod.rebate * tenor * prod.s0  # 年化百分比（仅适用于到期支付rebate的情形，敲出立即支付需自行按敲出时间计算）
        else:
            _rebate = prod.rebate  # 绝对数值
        if spot is None:
            spot = self.process.spot()
        r = self.process.interest(tau)
        q = self.process.div(tau)
        vol = self.process.vol(tau, self.process.spot())

        if calculate_date == prod.stop_date:  # 如果估值日就是终止日，则直接计算payoff
            value = self.calc_maturity_payoff(prod, np.array([spot]), _rebate, _floor_yield)[0]
            return value

        if prod.obs_type == ExerciseType.European:  # 如果是欧式障碍期权（到期观察）
            return self._calc_european_barrier(prod, spot, tau, r, q, vol, _rebate, _floor_yield)
        elif prod.obs_type == ExerciseType.American:  # 如果是美式障碍期权
            return self._calc_american_barrier(prod, spot, tau, r, q, vol, _rebate)
        else:
            raise ValueError("obs_type must be European or American")

    def _check_product_params(self, prod):
        """检查产品参数是否合法"""
        if prod.obs_type == ExerciseType.European:  # 如果是欧式障碍期权（到期观察）
            assert prod.payment_type == PaymentType.Expire, "ValueError: 欧式障碍期权（到期观察），payment_type只支持Expire到期支付"
        elif prod.obs_type == ExerciseType.American:  # 如果是美式障碍期权
            assert not prod.rebate_annual and prod.floor_yield == 0, "ValueError: 美式障碍解析解只支持绝对数值的现金返还，非百分比，非年化; 不支持保底收益率"
            assert prod.parti == 1 or prod.rebate == 0, "ValueError: 障碍期权解析解，由于公式限制，参与率parti不等于1、现金返还rebate不等于0不能共存，否则结果错误"
            if prod.inout == InOut.In:
                assert prod.payment_type == PaymentType.Expire, "ValueError: 敲入期权解析解，一直未敲入，到期时支付现金返还。payment_type只支持Expire"
            if prod.inout == InOut.Out:
                assert prod.payment_type == PaymentType.Hit, "ValueError: 敲出期权解析解，发生敲出，立刻支付现金返还。payment_type只支持Hit"

    @staticmethod
    def calc_maturity_payoff(prod, s_vec, _rebate, _floor_yield):
        """到期时的期权价值，在障碍价格内侧，默认设定未敲入，未敲出
        Args:
            prod: Product产品对象
            s_vec: np.ndarray, 网格的价格格点
            _rebate: float, 现金补偿绝对数值 (若为年化，计入时间为从起始日到终止日)
            _floor_yield: float, 保底收益绝对数值，计入时间为从起始日到终止日
        Returns:
            payoff: np.ndarray, 期末价值向量
        """
        if prod.inout == InOut.Out:
            v_vec = np.maximum(prod.callput.value * (s_vec - prod.strike), 0) * prod.parti + _floor_yield
            if prod.updown == UpDown.Up:
                v_vec = np.where(s_vec > prod.barrier, _rebate, v_vec)
            elif prod.updown == UpDown.Down:
                v_vec = np.where(s_vec < prod.barrier, _rebate, v_vec)
            else:
                raise ValueError("updown must be Up or Down")
        elif prod.inout == InOut.In:
            v_vec = _rebate * np.ones(len(s_vec))
            if prod.updown == UpDown.Up:
                v_vec = np.where(s_vec > prod.barrier, np.maximum(prod.callput.value * (s_vec - prod.strike),
                                                                  0) * prod.parti + _floor_yield, v_vec)
            elif prod.updown == UpDown.Down:
                v_vec = np.where(s_vec < prod.barrier, np.maximum(prod.callput.value * (s_vec - prod.strike),
                                                                  0) * prod.parti + _floor_yield, v_vec)
            else:
                raise ValueError("updown must be Up or Down")
        else:
            raise ValueError("inout must be In or Out")
        return v_vec

    def _calc_european_barrier(self, prod, spot, tau, r, q, vol, _rebate=0., _floor_yield=0.):
        """欧式障碍期权的解析解
            欧式香草与欧式资产或无、欧式现金或无的解析解的组合
            欧式现金或无的解析解可用于复制现金返还。
        Args:
            prod: Product产品对象
            spot: float，估值日标的价格
            tau: float，交易日到期时间（年）
            r: float，无风险利率（常数）
            q: float，连续分红率（常数）
            vol: float，波动率（常数）
            _rebate: float, 现金补偿绝对数值 (若为年化，计入时间为从起始日到终止日)
            _floor_yield: float, 保底收益绝对数值，计入时间为从起始日到终止日
        Returns:
            欧式障碍期权的解析解
        """
        if prod.barrier_type in [BarrierType.UIC, BarrierType.DOC]:  # 向上敲入看涨 或 向下敲出看涨
            if prod.barrier <= prod.strike:  # 此时相当于call的价值
                value = bs_formula(s=spot, k=prod.strike, t=tau, r=r, q=q, sigma=vol, sign=1) * prod.parti
            else:  # 此时 v = 方向向上、执行价为障碍价的 [ 资产或无 - 执行价 * 现金或无 ]
                value = (euro_asset_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=1) -
                         euro_cash_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=1) * prod.strike
                         ) * prod.parti
            # 现金返还rebate相当于看跌的现金或无; 如果存在保底收益率，则保底收益相当于看涨的现金或无.
            value += self._extra_returns(prod, spot, tau, r, q, vol, _rebate, _floor_yield, flag=1)
        elif prod.barrier_type in [BarrierType.DIC, BarrierType.UOC]:  # 向下敲入看涨 或 向上敲出看涨
            if prod.barrier <= prod.strike:  # 此时价值是0
                value = 0
            else:  # 此时 v = 多一份 call；空1份 方向向上、执行价为障碍价的 [ 资产或无 - 执行价 * 现金或无 ]
                value = (bs_formula(s=spot, k=prod.strike, t=tau, r=r, q=q, sigma=vol, sign=1) - (
                        euro_asset_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=1) -
                        euro_cash_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=1) * prod.strike
                )) * prod.parti
            # 现金返还rebate相当于看涨的现金或无; 如果存在保底收益率，则保底收益相当于看跌的现金或无.
            value += self._extra_returns(prod, spot, tau, r, q, vol, _rebate, _floor_yield, flag=-1)
        elif prod.barrier_type in [BarrierType.UIP, BarrierType.DOP]:  # 向上敲入看跌 或 向下敲出看跌
            if prod.barrier >= prod.strike:  # 此时价值是0
                value = 0
            else:  # 此时 v = 多一份 put；空1份 方向向下、执行价为障碍价的 [ 执行价 * 现金或无 - 资产或无 ]
                value = (bs_formula(s=spot, k=prod.strike, t=tau, r=r, q=q, sigma=vol, sign=-1) - (
                        euro_cash_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=-1) * prod.strike
                        - euro_asset_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=-1)
                )) * prod.parti
            # 现金返还rebate相当于看跌的现金或无; 如果存在保底收益率，则保底收益相当于看涨的现金或无.
            value += self._extra_returns(prod, spot, tau, r, q, vol, _rebate, _floor_yield, flag=1)
        elif prod.barrier_type in [BarrierType.DIP, BarrierType.UOP]:  # 向下敲入看跌 或 向上敲出看跌
            if prod.barrier >= prod.strike:  # 此时相当于put的价值
                value = bs_formula(s=spot, k=prod.strike, t=tau, r=r, q=q, sigma=vol, sign=-1) * prod.parti
            else:  # 此时 v = 方向向下、执行价为障碍价的 [ 执行价 * 现金或无 - 资产或无 ]
                value = (euro_cash_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=-1) * prod.strike
                         - euro_asset_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=-1)
                         ) * prod.parti
            # 现金返还rebate相当于看涨的现金或无; 如果存在保底收益率，则保底收益相当于看跌的现金或无.
            value += self._extra_returns(prod, spot, tau, r, q, vol, _rebate, _floor_yield, flag=-1)
        else:
            raise ValueError("Invalid barrier type")
        return value

    def _extra_returns(self, prod, spot, tau, r, q, vol, _rebate, _floor_yield, flag=1):
        """处理欧式障碍期权可能存在的额外payoff：现金返还rebate、保底收益率floor_yield
        Args:
            prod: Product产品对象
            spot: float，估值日标的价格
            tau: float，交易日到期时间（年）
            r: float，无风险利率（常数）
            q: float，连续分红率（常数）
            vol: float，波动率（常数）
            flag: int, 若 flag = 1， 现金返还rebate相当于看跌的现金或无; 如果存在保底收益率，则保底收益相当于看涨的现金或无.
                       若 flag = -1， 现金返还rebate相当于看涨的现金或无; 如果存在保底收益率，则保底收益相当于看跌的现金或无.
        Returns:

        """
        extra = 0
        if prod.floor_yield != 0:  # 保底收益率
            extra += euro_cash_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=flag) * _floor_yield
        if prod.rebate != 0:  # 现金返还
            extra += euro_cash_or_nothing(s=spot, k=prod.barrier, t=tau, r=r, q=q, sigma=vol, sign=-flag) * _rebate
        return extra

    def _calc_american_barrier(self, prod, spot, tau, r, q, vol, _rebate):
        """美式障碍期权的解析解
            Merton(1973), Reiner & Rubinstein(1991a)提出障碍期权解析解，公式中的现金返还是绝对数值，不支持保底收益率floor_yield
            支持Broadie, Glasserman和Kou(1995)提出的均匀离散观察障碍期权近似解.
        Args:
            prod: Product产品对象
            spot: float，估值日标的价格
            tau: float，交易日到期时间（年）
            r: float，无风险利率（常数）
            q: float，连续分红率（常数）
            vol: float，波动率（常数）
            _rebate: float, 现金补偿绝对数值 (若为年化，计入时间为从起始日到终止日)
        Returns:
            美式障碍期权的解析解
        """
        drift = r - q
        if not self.for_haug:
            # 如果估值时的标的价格已经触碰障碍价，直接返回现金返还或者香草期权的价值
            if (prod.updown == UpDown.Up and spot >= prod.barrier) or (
                    prod.updown == UpDown.Down and spot <= prod.barrier):
                if prod.inout == InOut.Out:  # 发生敲出，立刻支付现金返还rebate
                    return prod.rebate
                elif prod.inout == InOut.In:  # 发生敲入，返回香草期权的价值
                    return bs_formula(s=spot, k=prod.strike, t=tau, r=r, q=q, sigma=vol, sign=prod.callput.value)
                else:
                    raise ValueError("inout只能是InOut.In或InOut.Out")

        if prod.discrete_obs_interval is not None:
            # 均匀离散观察障碍期权，M. Broadie, P. Glasserman, S.G. Kou(1997) 在连续障碍期权解析解上加调整项，调整障碍价格水平
            # 指数上的beta = -zeta(1/2) / sqrt(2*pi) = 0.5826, 其中zeta是黎曼zeta函数
            if prod.updown == UpDown.Up:  # 向上
                barrier = prod.barrier * np.exp(0.5826 * vol * np.sqrt(prod.discrete_obs_interval))
            elif prod.updown == UpDown.Down:  # 向下
                barrier = prod.barrier * np.exp(-0.5826 * vol * np.sqrt(prod.discrete_obs_interval))
            else:
                raise ValueError("self.updown 只能是Up或Down")
        else:  # 连续观察障碍期权
            barrier = prod.barrier

        mu = drift / vol ** 2 - 0.5
        la = np.sqrt(mu ** 2 + 2 * r / vol ** 2)
        a = (barrier / spot) ** (2 * mu)
        b = (barrier / spot) ** (2 * mu + 2)
        c = (barrier / spot) ** (mu + la)
        d = (barrier / spot) ** (mu - la)
        a1 = np.log(spot / prod.strike)
        a2 = np.log(spot / barrier)
        a3 = np.log(spot * prod.strike / barrier ** 2)
        a4 = drift + 0.5 * vol ** 2
        a5 = drift - 0.5 * vol ** 2
        a6 = tau
        a7 = vol * np.sqrt(a6)
        d1 = (a1 + a4 * a6) / a7
        d2 = (a1 + a5 * a6) / a7
        d3 = (a2 + a4 * a6) / a7
        d4 = (a2 + a5 * a6) / a7
        d5 = (a2 - a5 * a6) / a7
        d6 = (a2 - a4 * a6) / a7
        d7 = (a3 - a5 * a6) / a7
        d8 = (a3 - a4 * a6) / a7
        d9 = -a2 / a7 + la * a7
        d10 = -a2 / a7 - la * a7

        def A(phi):
            return phi * spot * np.exp(-q * a6) * norm.cdf(phi * d1) - \
                phi * prod.strike * np.exp(-r * a6) * norm.cdf(phi * d2)

        def B(phi):
            return phi * spot * np.exp(-q * a6) * norm.cdf(phi * d3) - \
                phi * prod.strike * np.exp(-r * a6) * norm.cdf(phi * d4)

        def C(phi, eta):
            return phi * spot * np.exp(-q * a6) * b * norm.cdf(-eta * d8) - \
                phi * prod.strike * np.exp(-r * a6) * a * norm.cdf(-eta * d7)

        def D(phi, eta):
            return phi * spot * np.exp(-q * a6) * b * norm.cdf(-eta * d6) - \
                phi * prod.strike * np.exp(-r * a6) * a * norm.cdf(-eta * d5)

        def E(eta):
            return prod.rebate * np.exp(-r * a6) * (norm.cdf(eta * d4) - a * norm.cdf(-eta * d5))

        def F(eta):
            return prod.rebate * (c * norm.cdf(eta * d9) + d * norm.cdf(eta * d10))

        if prod.barrier_type == BarrierType.UIC:  # 向上敲入看涨
            phi = 1
            eta = -1
            if prod.strike >= barrier:
                price = A(phi) + E(eta)
            else:
                price = (B(phi) - C(phi, eta) + D(phi, eta)) + E(eta)
        elif prod.barrier_type == BarrierType.UIP:  # 向上敲入看跌
            phi = -1
            eta = -1
            if prod.strike >= barrier:
                price = (A(phi) - B(phi) + D(phi, eta)) + E(eta)
            else:
                price = C(phi, eta) + E(eta)
        elif prod.barrier_type == BarrierType.UOC:  # 向上敲出看涨
            phi = 1
            eta = -1
            if prod.strike >= barrier:
                price = F(eta)
            else:
                price = (A(phi) - B(phi) + C(phi, eta) - D(phi, eta)) + F(eta)
        elif prod.barrier_type == BarrierType.UOP:  # 向上敲出看跌
            phi = -1
            eta = -1
            if prod.strike >= barrier:
                price = (B(phi) - D(phi, eta)) + F(eta)
            else:
                price = (A(phi) - C(phi, eta)) + F(eta)
        elif prod.barrier_type == BarrierType.DIC:  # 向下敲入看涨
            phi = 1
            eta = 1
            if prod.strike >= barrier:
                price = C(phi, eta) + E(eta)
            else:
                price = (A(phi) - B(phi) + D(phi, eta)) + E(eta)
        elif prod.barrier_type == BarrierType.DIP:  # 向下敲入看跌
            phi = -1
            eta = 1
            if prod.strike >= barrier:
                price = (B(phi) - C(phi, eta) + D(phi, eta)) + E(eta)
            else:
                price = A(phi) + E(eta)
        elif prod.barrier_type == BarrierType.DOC:  # 向下敲出看涨
            phi = 1
            eta = 1
            if prod.strike >= barrier:
                price = (A(phi) - C(phi, eta)) + F(eta)
            else:
                price = (B(phi) - D(phi, eta)) + F(eta)
        elif prod.barrier_type == BarrierType.DOP:  # 向下敲出看跌
            phi = -1
            eta = 1
            if prod.strike >= barrier:
                price = (A(phi) - B(phi) + C(phi, eta) - D(phi, eta)) + F(eta)
            else:
                price = F(eta)
        else:
            raise ValueError("self.barrier_type 障碍类型错误")
        return price * prod.parti
