#!/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 pricelib.common.date_handler import global_evaluation_date
from pricelib.common.utilities.enums import InOut, UpDown, PaymentType, ExerciseType
from pricelib.common.pricing_engine_base import McEngine


class MCBarrierEngine(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，现值
                 如果输入spot = np.ndarray一维价格向量，t=终止日，则返回盈亏测算向量，即每个价格对应的现值，而不是平均值
        """
        calculate_date = global_evaluation_date() if t is None else t
        _maturity_business_days = prod.trade_calendar.business_days_between(calculate_date,
                                                                            prod.end_date)  # 从估值日到到期日的交易日时间（天）
        _maturity = (prod.end_date - calculate_date).days / prod.annual_days.value  # 从估值日到到期日的自然日时间（年）
        _elapsed_days = prod.trade_calendar.business_days_between(prod.start_date, calculate_date)  # 从起始日到估值日的交易日时间（天）

        if spot is None:
            spot = self.process.spot()
        # 当估值日就是终止日时，直接用spot计算payoff，无需随机数矩阵
        if (calculate_date == prod.end_date) and isinstance(spot, np.ndarray):
            paths = np.array([spot, ])  # spot是一个np.ndarray一维价格向量，返回盈亏测算向量，即每个价格对应的现值，而不是平均值
            _tenor = (prod.end_date - prod.start_date).days / prod.annual_days.value  # 从起始日到到期日的自然日时间（年）
        elif calculate_date == prod.stop_date:
            paths = np.array([[spot, ]])  # spot是一个数值，计算期末payoff
            _tenor = (prod.stop_date - prod.start_date).days / prod.annual_days.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()
            _tenor = (prod.end_date - prod.start_date).days / prod.annual_days.value  # 从起始日到到期日的自然日时间（年）

        if prod.obs_type == ExerciseType.European:  # 如果是欧式障碍期权（到期观察）
            value = self._calc_european_barrier(prod, paths, _maturity, _tenor)
        elif prod.obs_type == ExerciseType.American:  # 如果是美式障碍期权
            value = self._calc_american_barrier(prod, paths, _maturity, _maturity_business_days, _tenor, _elapsed_days)
        else:
            raise ValueError("obs_type must be European or American")

        if (calculate_date == prod.end_date) and isinstance(spot, np.ndarray):
            return value  # 终止日, spot是价格向量, 则返回盈亏测算向量，即每个价格对应的现值，而不是平均值
        else:
            return np.mean(value)

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

    def _calc_european_barrier(self, prod, paths, _maturity, _tenor):
        """欧式障碍期权(到期观察)的蒙特卡洛模拟
        Args:
            prod: Product产品对象
            paths: np.ndarray，模拟价格路径
            _maturity: float，从估值日到终止日的自然日时间（年）
            _tenor: float, 从起始日到终止日的自然日时间（年）
        Returns: float，欧式障碍期权(到期观察)的估值结果
        """
        if prod.updown == UpDown.Up:
            knock_inout = (paths[-1] >= prod.barrier)
        elif prod.updown == UpDown.Down:
            knock_inout = (paths[-1] <= prod.barrier)
        else:
            raise ValueError("不支持的UpDown类型")

        _floor_yield, _rebate = self._preprocess_extra_returns(prod, _tenor)

        if prod.inout == InOut.In:
            payoff = np.ones(paths[-1].size) * _rebate
            payoff[knock_inout] = np.maximum(prod.callput.value * (paths[-1, knock_inout] - prod.strike),
                                             0) * prod.parti + _floor_yield
        elif prod.inout == InOut.Out:
            payoff = np.maximum(prod.callput.value * (paths[-1] - prod.strike), 0) * prod.parti + _floor_yield
            payoff[knock_inout] = _rebate
        else:
            raise ValueError("不支持的InOut类型")
        payoff = payoff * self.process.interest.disc_factor(_maturity)
        return payoff

    def _calc_american_barrier(self, prod, paths, _maturity, _maturity_business_days, _tenor, _elapsed_days):
        """美式障碍期权的蒙特卡洛模拟
        Args:
            prod: Product产品对象
            paths: np.ndarray，模拟价格路径
            _maturity: float，从估值日到终止日的自然日到期时间（年）
            _maturity_business_days: int，从估值日到终止日的交易日到期时间（天）
            _tenor: float, 从起始日到终止日的自然日存续时间（年）
            _elapsed_days: int, 从起始日到估值日的交易日时间（天）
        Returns: float，美式障碍期权的估值结果
        """
        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  # 防止闰年导致的下标越界

        if prod.updown == UpDown.Up:
            knock_inout = np.any(paths[obs_points] >= prod.barrier, axis=0)
        elif prod.updown == UpDown.Down:
            knock_inout = np.any(paths[obs_points] <= prod.barrier, axis=0)
        else:
            raise ValueError("不支持的UpDown类型")

        _floor_yield, _rebate = self._preprocess_extra_returns(prod, _tenor)
        if prod.inout == InOut.In:
            payoff = np.ones(paths[-1].size) * _rebate
            payoff[knock_inout] = np.maximum(prod.callput.value * (paths[-1, knock_inout] - prod.strike),
                                             0) * prod.parti + _floor_yield
            payoff = payoff * self.process.interest.disc_factor(_maturity)
            return payoff
        elif prod.inout == InOut.Out:
            payoff = (np.maximum(prod.callput.value * (paths[-1] - prod.strike), 0
                                 ) * prod.parti + _floor_yield) * self.process.interest.disc_factor(_maturity)
            if prod.payment_type == PaymentType.Expire:  # 敲出时，到期再支付现金返还
                payoff[knock_inout] = _rebate * self.process.interest.disc_factor(_maturity)
            elif prod.payment_type == PaymentType.Hit:  # 敲出时，立即支付现金返还
                step_index = np.tile(obs_points, (paths.shape[1], 1)).T.astype(int)
                if prod.updown == UpDown.Up:
                    hit_time = np.min(np.where(paths[obs_points] >= prod.barrier, step_index, np.inf), axis=0)
                elif prod.updown == UpDown.Down:
                    hit_time = np.min(np.where(paths[obs_points] <= prod.barrier, step_index, np.inf), axis=0)
                else:
                    raise ValueError("不支持的UpDown类型")
                if prod.rebate != 0:  # 现金返还
                    if prod.rebate_annual:  # 年化百分比
                        _rebate = prod.s0 * prod.rebate * (_elapsed_days + hit_time[hit_time != np.inf]
                                                           ) / prod.t_step_per_year
                    else:
                        _rebate = prod.rebate  # 绝对数值
                else:
                    _rebate = 0.0
                payoff[knock_inout] = _rebate * self.process.interest.disc_factor(
                    hit_time[hit_time != np.inf] / prod.t_step_per_year)
            else:
                raise ValueError("PaymentType must be Hit or Expire")
            return payoff
        else:
            raise ValueError("不支持的InOut类型")
