#!/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 .fdm_barrier_engine import FdmBarrierEngine


class FdmAirbagEngine(FdmBarrierEngine):
    """安全气囊PDE有限差分法定价引擎
    安全气囊是一个向下敲入的期权，到期支付，敲入前是call或牛市价差，敲入后的收益是标的涨跌幅
    如果存在敲入前或敲入后封顶收益率，或者封底收益率，则需要改写终止条件和边界条件
    支持连续观察/离散观察(默认为每日观察)"""

    def _init_boundary_condition(self, smax, maturity, elapsed_time=None, _floor_yield=0):
        """初始化边界条件
        安全气囊，未敲入是call，已敲入是标的资产
        Args:
            smax: float, 有限差分网格上边界的价格最大值
            maturity: float, 到期时间
        Returns:
            self.fn_callput: List[Callable], 0~smax网格的边界条件
            self.fn_bound: List[Callable], 障碍价格网格的边界条件
        """
        # 生成已敲入网格的边界条件： 0~smax，此时收益为标的资产涨跌幅（可能存在敲入后封顶收益率，或者封底收益率）
        if self.prod.in_strike_floor is not None:  # 若存在敲入后封底价格
            _in_min_s = self.prod.in_strike_floor
        else:
            _in_min_s = 0
        if self.prod.in_strike_cap is not None:  # 若存在敲入后封顶价格
            _in_max_s = self.prod.in_strike_cap
        else:
            _in_max_s = smax
        self.fn_callput = [lambda u: (self.prod.knockin_parti * (_in_min_s - self.prod.s0) + _floor_yield
                                      ) * self.process.interest.disc_factor(maturity, u),
                           lambda u: (self.prod.reset_call_parti * (_in_max_s - self.prod.s0) + _floor_yield
                                      ) * self.process.interest.disc_factor(maturity, u)]
        # 生成未敲入网格的边界条件： bound[0]~smax, 此时收益为向下敲出看涨期权或者价差
        if self.prod.init_strike_cap is not None:  # 若存在敲入后封顶价格
            _max_s = self.prod.init_strike_cap
        else:
            _max_s = smax
        # 返回一个包含两个lambda函数的列表，第u天触碰障碍产生的payoff，折现到第u天的价值。 列表索引分别为[0, 1]，每个lambda函数的形参r=self.rebate中的r补偿值
        self.fn_bound = []
        for i, _bound in enumerate(self.bound):
            if _bound in [-np.inf, np.inf]:  # 安全气囊没有上方的障碍，即障碍价无穷。未敲入call的上边界是smax（若封顶则为init_strike_cap）
                self.fn_bound.append(lambda u: (self.prod.call_parti * (_max_s - self.prod.strike) + _floor_yield
                                                ) * self.process.interest.disc_factor(maturity, u))
            else:  # 安全气囊有下方的敲入障碍价格，触碰下方的障碍后，到期时支付，安全气囊已敲入是标的资产跌幅
                if self.prod.discrete_obs_interval is None:  # 连续观察, 价格下边界是self.bound[0]
                    self.fn_bound.append(
                        lambda u: (self.prod.knockin_parti * (max(_bound, _in_min_s) - self.prod.s0)
                                   + _floor_yield) * self.process.interest.disc_factor(maturity, u))
                else:  # 离散观察，价格下边界是0
                    self.fn_bound.append(
                        lambda u: (self.prod.knockin_parti * (_in_min_s - self.prod.s0) + _floor_yield
                                   ) * self.process.interest.disc_factor(maturity, u))

    def _init_terminal_condition(self, s_vec, maturity, elapsed_time=None, _floor_yield=0.):
        """初始化终止时已敲入网格终止时的期权价值，安全气囊敲入后是标的资产涨跌幅
        _init_no_touch_terminal_condition返还未敲入网格的终止条件
        连续观察：网格上下界是障碍价格
        离散观察/到期观察：网格上下界是0和smax
        """
        prod = self.prod
        # 已敲入，且标的上涨
        if prod.in_strike_cap is None:
            in_payoff_upper = s_vec - prod.s0
        else:  # 若存在敲入后封顶价格
            in_payoff_upper = np.where(s_vec - prod.in_strike_cap > 0,
                                       prod.in_strike_cap - prod.s0, s_vec - prod.s0)
        # 已敲入，且标的下跌
        if prod.in_strike_floor is None:
            in_payoff_lower = s_vec - prod.s0
        else:  # 若存在敲入后封底价格
            in_payoff_lower = np.where(s_vec - prod.in_strike_floor < 0,
                                       prod.in_strike_floor - prod.s0, s_vec - prod.s0)
        yv = (np.where(s_vec - prod.s0 > 0,
                       # 敲入且标的上涨
                       in_payoff_upper * prod.reset_call_parti,  # 敲入重置后看涨参与率
                       # 敲入且标的下跌
                       in_payoff_lower * prod.knockin_parti)  # 敲入下跌参与率
              + _floor_yield)  # 保底收益
        return yv

    def _init_no_touch_terminal_condition(self, s_vec, p=0, maturity=0, elapsed_time=None, _floor_yield=0):
        """初始化安全气囊未敲入网格终点的期权价值，
        连续观察：网格上下界是障碍价格，因此都在敲入价格以上，只有未敲入的call（或牛市价差）
        离散观察/到期观察：网格上下界是0和smax，敲入价上方，未敲入是call（或牛市价差）；敲入价下方是标的资产跌幅（或熊市价差）"""
        prod = self.prod
        # 未敲入，期末价值是call或牛市价差
        if prod.init_strike_cap is None:
            init_payoff = s_vec - prod.strike
        else:  # 若存在敲入前封顶价格
            init_payoff = np.where(s_vec - prod.init_strike_cap > 0,
                                   prod.init_strike_cap - prod.strike, s_vec - prod.strike)
        # 已敲入，且标的下跌
        if prod.in_strike_floor is None:
            in_payoff_lower = s_vec - prod.s0
        else:  # 若存在敲入后封底价格
            in_payoff_lower = np.where(s_vec - prod.in_strike_floor < 0,
                                       prod.in_strike_floor - prod.s0, s_vec - prod.s0)
        if self.prod.discrete_obs_interval is None:  # 连续观察
            yv = (_floor_yield + np.maximum(init_payoff, 0) * self.prod.call_parti  # 未敲入的看涨参与率
                  ) * self.process.interest.disc_factor(maturity, p)
        else:  # 离散观察
            yv = np.where(s_vec - self.bound[0] > 0,
                          np.maximum(init_payoff, 0) * self.prod.call_parti,  # 未敲入的看涨参与率
                          in_payoff_lower * self.prod.knockin_parti  # 敲入的下跌参与率
                          ) + _floor_yield  # 保底收益
        return yv
