#!/user/bin/env python
# -*- coding: utf-8 -*-
"""
Copyright (C) 2024 Galaxy Technologies
Licensed under the Apache License, Version 2.0
"""
import math
import numpy as np
from pricelib.common.utilities.enums import ExerciseType, PaymentType
from pricelib.common.date_handler import global_evaluation_date
from pricelib.common.pricing_engine_base import McEngine


class MCDoubleSharkEngine(McEngine):
    """双鲨期权 Monte Carlo 模拟定价引擎"""

    def calc_present_value(self, prod, t=None, spot=None):
        """计算现值
        Args:
            prod: Product产品对象
            t: datetime.date，估值日; 如果是None，则使用全局估值日globalEvaluationDate
            spot: float，估值日标的价格，如果是None，则使用随机过程的当前价格
        Returns: float，现值
        """
        if spot is None:
            spot = self.process.spot()
        calculate_date = global_evaluation_date() if t is None else t

        _maturity = (prod.end_date - calculate_date).days / prod.annual_days.value
        _maturity_business_days = prod.trade_calendar.business_days_between(calculate_date, prod.end_date)

        _elapsed_days = prod.trade_calendar.business_days_between(prod.start_date, calculate_date)  # 从起始日到估值日的交易日时间（天）
        _tenor = (prod.end_date - prod.start_date).days / prod.annual_days.value  # 从起始日到到期日的自然日时间（年）
        _floor_yield = self._preprocess_floor_yield(prod, _tenor)

        r = self.process.interest(_maturity)

        # 当估值日就是终止日时，直接用spot计算payoff，无需随机数矩阵
        if (calculate_date == prod.end_date) and isinstance(spot, np.ndarray):
            # 如果spot是一个np.ndarray一维价格向量，返回盈亏测算向量，即每个价格对应的现值，而不是平均值
            _tenor = (prod.end_date - prod.start_date).days / prod.annual_days.value  # 从起始日到终止日的自然日时间（年）
            _floor_yield = self._preprocess_floor_yield(prod, _tenor)
            value = self.calc_maturity_payoff(prod, np.array([spot, ]), _tenor, _floor_yield)
            return value
        elif calculate_date == prod.stop_date:  # spot是一个数值，计算到期日payoff
            _tenor = _tenor = (prod.stop_date - prod.start_date).days / prod.annual_days.value  # 从起始日到终止日的自然日时间（年）
            _floor_yield = self._preprocess_floor_yield(prod, _tenor)
            value = self.calc_maturity_payoff(prod, np.array([spot, ]), _tenor, _floor_yield)[0]
            return value
        else:
            if spot != self.process.spot():
                self.reset_paths_flag()  # 重置路径标志位，重新生成路径
            paths = self.path_generator(n_step=_maturity_business_days, spot=spot,
                                        t_step_per_year=prod.t_step_per_year).copy()

        call_payoff = (paths[-1] - prod.strike[1]) * prod.parti[1]
        put_payoff = (prod.strike[0] - paths[-1]) * prod.parti[0]
        # 美式：
        if prod.obs_type == ExerciseType.American:
            if prod.discrete_obs_interval is None:  # 每日观察
                obs_points = np.arange(0, _maturity_business_days + 1, 1)
            else:  # 均匀离散观察
                dt_step = prod.discrete_obs_interval * prod.t_step_per_year
                obs_points = np.flip(np.round(np.arange(_maturity_business_days, 0, -dt_step)).astype(int))
                obs_points = np.concatenate((np.array([0]), obs_points))
                obs_points[obs_points > _maturity_business_days] = _maturity_business_days  # 防止闰年导致的下标越界

            # 双鲨结构是双边敲出
            step_index = np.tile(obs_points, (self.n_path, 1)).T
            hit_lower = np.min(
                np.where(paths[obs_points] <= prod.bound[0], step_index.astype(int), _maturity_business_days + 1),
                axis=0)
            hit_upper = np.min(
                np.where(paths[obs_points] >= prod.bound[1], step_index.astype(int), _maturity_business_days + 1),
                axis=0)
            if prod.payment_type == PaymentType.Hit:  # 发生敲出时，立即支付现金返还
                _rebate_lower = self._preprocess_extra_returns(prod,
                                                               (_elapsed_days + hit_lower[hit_lower < hit_upper]
                                                                ) / prod.t_step_per_year, [prod.rebate[0], ])
                _rebate_upper = self._preprocess_extra_returns(prod,
                                                               (_elapsed_days + hit_upper[hit_lower > hit_upper]
                                                                ) / prod.t_step_per_year, [prod.rebate[1], ])
                value = np.mean(
                    np.where(hit_lower < hit_upper, _rebate_lower * np.exp(-r * hit_lower / prod.t_step_per_year),
                             np.where(hit_lower > hit_upper,
                                      _rebate_upper * np.exp(-r * hit_upper / prod.t_step_per_year),
                                      (np.where(call_payoff > 0, call_payoff, np.where(put_payoff > 0, put_payoff, 0))
                                       + _floor_yield) * math.exp(-r * _maturity))))
            elif prod.payment_type == PaymentType.Expire:
                _rebate = self._preprocess_extra_returns(prod, _tenor, prod.rebate)
                value = np.mean(np.where(hit_lower < hit_upper, _rebate[0],
                                         np.where(hit_lower > hit_upper, _rebate[1],
                                                  (np.where(call_payoff > 0, call_payoff,
                                                            np.where(put_payoff > 0, put_payoff, 0)) + _floor_yield
                                                   )))) * math.exp(-r * _maturity)
            else:
                raise ValueError("PaymentType must be Hit or Expire")

        # 欧式：
        elif prod.obs_type == ExerciseType.European:
            # 双鲨结构是双边敲出
            value = self.calc_maturity_payoff(prod, paths[-1], _tenor, _floor_yield)[0]
            value = np.mean(value) * math.exp(-r * _maturity)
        else:
            raise ValueError("ExerciseType must be American or European")
        return value

    def calc_maturity_payoff(self, prod, s_vec, _tenor, _floor_yield=0.):
        """初始化终止时的期权价值，在障碍价格内侧，默认设定未敲入，未敲出
        Args:
            prod: Product产品对象
            s_vec: np.ndarray, 网格的价格格点
            _tenor: float, 从起始日到终止日的自然日时间（年）
            _floor_yield: float，处理后的保底收益绝对数值
        Returns:
            value: np.ndarray, 期末价值向量
        """
        call_payoff = (s_vec - prod.strike[1]) * prod.parti[1]
        put_payoff = (prod.strike[0] - s_vec) * prod.parti[0]
        _rebate = self._preprocess_extra_returns(prod, _tenor, prod.rebate)
        value = np.where(s_vec <= prod.bound[0], _rebate[0],
                         np.where(s_vec >= prod.bound[1], _rebate[1],
                                  np.where(call_payoff > 0, call_payoff,
                                           np.where(put_payoff > 0, put_payoff, 0)) + _floor_yield))
        return value

    @staticmethod
    def _preprocess_extra_returns(prod, period, rebate_list):
        """预处理可能存在的年化现金返还rebate
        Args:
            prod: Product产品对象
            period: float 或 np.ndarray, 从起始到期权终止的时间（年）
                    对欧式障碍、美式敲入而言是整个存续期；对美式敲出而言，期权可能会因为敲出提前终止，此时period是起始日到敲出日的时间
            rebate_list: List[float], 现金返还率或绝对数值
        Returns:
            _rebate: float，处理后的现金返还绝对数值
        """
        if prod.rebate_annual:  # 年化百分比
            _rebate = []
            for r in rebate_list:
                _rebate.append(r * period * prod.s0)
        else:
            _rebate = rebate_list  # 绝对数值
        return _rebate

    @staticmethod
    def _preprocess_floor_yield(prod, period):
        """预处理可能存在的年化保底收益floor_yield
        Args:
            prod: Product产品对象
            period: float 或 np.ndarray, 从起始到期权终止的时间（年）
                    对欧式障碍、美式敲入而言是整个存续期；对美式敲出而言，期权可能会因为敲出提前终止，此时period是起始日到敲出日的时间
        Returns:
            _floor_yield: float，处理后的保底收益绝对数值
        """
        if prod.floor_yield != 0:  # 保底收益率
            _floor_yield = prod.floor_yield * period * prod.s0
        else:
            _floor_yield = 0.0
        return _floor_yield
