#!/user/bin/env python
# -*- coding: utf-8 -*-
"""
Copyright (C) 2024 Galaxy Technologies
Licensed under the Apache License, Version 2.0
"""
from functools import partial
import numpy as np
from pricelib.common.utilities.enums import StatusType
from pricelib.common.date_handler import CN_CALENDAR, AnnualDays
from .autocallable_base import AutocallableBase


class BothDownSnowball(AutocallableBase):
    """双降雪球结构产品类"""

    def __init__(self, s0, barrier_out_start, barrier_out_step, barrier_in, coupon_ko_start=None, coupon_ko_step=None,
                 coupon_out_stairs=None, coupon_stair_ends=None,
                 coupon_div=None, lock_term=3, maturity=None, start_date=None, end_date=None,
                 obs_dates=None, pay_dates=None, parti_in=1, margin_lvl=1, status=StatusType.NoTouch, engine=None,
                 trade_calendar=CN_CALENDAR, annual_days=AnnualDays.N365, t_step_per_year=243,
                 stoch_process=None, s=None, r=None, q=None, vol=None, engine_method=None):
        """继承自动赎回基类AutocallableBase的参数，详见AutocallableBase的__init__方法
        Args: 以下为区别于经典雪球的参数
            barrier_out_start: float, 敲出障碍价起始值，绝对值/百分比
            barrier_out_step: float, 降敲步长，绝对值/百分比
            lock_term: int，锁定期，单位为月，锁定期内不触发敲出

            降票息有两种方式，选填其中一组参数：
                1. 每月等差下降
                coupon_ko_start: float, 敲出票息起始值，百分比，年化
                coupon_ko_step: float, 降票息步长，百分比，年化
                2. 分阶段下降
                coupon_out_stairs: tuple[float], 分段票息，百分比，年化。例如(0.18, 0.05)
                coupon_stair_ends: tuple[int], 分段票息节点, 通常为(12, 24), 代表每个阶段到第几个月末为止
        """
        super().__init__(s0=s0, maturity=maturity, start_date=start_date, end_date=end_date, lock_term=lock_term,
                         trade_calendar=trade_calendar, obs_dates=obs_dates, pay_dates=pay_dates, status=status,
                         annual_days=annual_days, parti_in=parti_in, margin_lvl=margin_lvl,
                         t_step_per_year=t_step_per_year, engine=engine, stoch_process=stoch_process,
                         s=s, r=r, q=q, vol=vol, engine_method=engine_method)
        len_obs_dates = len(self.obs_dates.date_schedule)

        barrier_out = np.arange(barrier_out_start, barrier_out_start - barrier_out_step * len_obs_dates,
                                -barrier_out_step)[:len_obs_dates]
        self.barrier_out = barrier_out
        self.barrier_in = np.ones(len_obs_dates) * barrier_in

        self.coupon_ko_start = coupon_ko_start
        self.coupon_ko_step = coupon_ko_step
        self.coupon_out_stairs = coupon_out_stairs
        self.coupon_stair_ends = coupon_stair_ends
        step_true = (coupon_ko_start is not None) and (coupon_ko_step is not None)
        step_false = (coupon_ko_start is None) and (coupon_ko_step is None)
        stair_true = (coupon_out_stairs is not None) and (coupon_stair_ends is not None)
        stair_false = (coupon_out_stairs is None) and (coupon_stair_ends is None)
        if step_true and stair_false:
            self.coupon_type = "step"
            coupon_out = np.arange(coupon_ko_start, coupon_ko_start - coupon_ko_step * len_obs_dates, -coupon_ko_step)[
                         :len_obs_dates]
            self.coupon_out = coupon_out
            # 反算敲出票息起点
            self.calc_coupon = partial(self.calc_param, reset_param_method=self.reset_coupon_ko_start)
        elif stair_true and step_false:
            assert len(coupon_out_stairs) == len(coupon_stair_ends), ("输入的分段票息coupon_out_stairs"
                                                                      "与分段票息节点coupon_stair_ends的长度必须相等！")
            self.coupon_type = "stair"
            coupon_out = np.zeros(len_obs_dates + lock_term - 1)
            coupon_stair_ends = [0] + list(coupon_stair_ends)
            for i in range(len(coupon_out_stairs)):
                coupon_out[coupon_stair_ends[i]:coupon_stair_ends[i + 1]] = coupon_out_stairs[i]
            self.coupon_out = coupon_out[lock_term - 1:]
            # 反算分段敲出票息，额外参数 nth：int, 反算第几个阶段的票息
            self.calc_coupon = partial(self.calc_param, reset_param_method=self.reset_coupon_nth)
        else:
            raise ValueError("只支持两种票息输入方式:\n"
                             "1.等差下降，仅输入敲出票息起始值coupon_ko_start和降票息步长coupon_ko_step;\n"
                             "2.分阶段下降，仅输入分段票息列表coupon_out_stairs和分段票息节点列表coupon_stair_ends.")

        self.coupon_div = coupon_div if coupon_div is not None else self.coupon_out[-1]
        self.parti_out = 0
        self.strike_upper = s0
        self.strike_lower = 0

    def __repr__(self):
        """返回期权的描述"""
        return "双降雪球"

    def reset_coupon_ko_start(self, coupon_ko_start: float, coupon_div=None, *args, **kwargs):
        """重新设置敲出票息起点"""
        len_obs_dates = len(self.obs_dates.date_schedule)
        self.coupon_out = np.arange(coupon_ko_start, coupon_ko_start - self.coupon_ko_step * len_obs_dates,
                                    -self.coupon_ko_step)[:len_obs_dates]
        if coupon_div is None:  # 未指定红利票息时，将红利票息的值更新为最后一个敲出票息
            self.coupon_div = self.coupon_out[-1]

    def reset_coupon_nth(self, coupon: float, nth: int = 1, coupon_div=None, *args, **kwargs):
        """重新设置指定分段的敲出票息
        Args:
            coupon: float, 重设票息值
            coupon_div: float, 是否指定了红利票息的值
            nth: int, 修改的是第几个阶段的票息
        """
        if nth == 1:
            self.coupon_out[:self.coupon_stair_ends[0] - self.lock_term + 1] = coupon
        else:
            self.coupon_out[self.coupon_stair_ends[nth - 2] - self.lock_term + 1:
                            self.coupon_stair_ends[nth - 1] - self.lock_term + 1] = coupon
        # 如果反算的是最后一段票息，并且没有指定红利票息，同时更新红利票息的值
        if (len(self.coupon_stair_ends) == nth) and (coupon_div is None):
            self.coupon_div = coupon
