#!/user/bin/env python
# -*- coding: utf-8 -*-
"""
Copyright (C) 2024 Galaxy Technologies
Licensed under the Apache License, Version 2.0
"""
import gc
import datetime
from abc import ABCMeta
from contextlib import suppress
import importlib
from typing import List
import numpy as np
from scipy.stats import norm, qmc
from ..processes import StochProcessBase
from ..utilities.enums import RandsMethod, LdMethod, ProcessType, EngineType
from ..date_handler import global_evaluation_date
from ..utilities.utility import logging
from .engine_base import PricingEngineBase


class McEngine(PricingEngineBase, metaclass=ABCMeta):
    """蒙特卡洛模拟定价引擎基类
    观察者，观察随机过程process对象，当process对象的属性变化时，自动更新状态，会重新生成价格路径"""
    engine_type = EngineType.McEngine

    def __init__(self, stoch_process: StochProcessBase = None, n_path=100000, rands_method=RandsMethod.LowDiscrepancy,
                 antithetic_variate=True, ld_method=LdMethod.Sobol, seed=0, *, rands=None, s_paths=None, var_paths=None,
                 s=None, r=None, q=None, vol=None):
        """构造函数
        Args:
            stoch_process: 随机过程StochProcessBase对象
            n_path: int，MC模拟路径数
            rands_method: 生成随机数方法，RandsMethod枚举类，Pseudorandom伪随机数/LowDiscrepancy低差异序列
            antithetic_variate: bool，是否使用对立变量法
            ld_method: 若使用了低差异序列，指定低差异序列方法，LdMethod枚举类，Sobol序列/Halton序列
            seed: int，随机数种子
        在未设置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)
        # MC模拟路径数
        if n_path % 2 != 0 and antithetic_variate:
            n_path += 1
        self.n_path = n_path
        self.rands_method = rands_method  # 生成随机数方法
        self._ld_method = ld_method  # 低差异序列方法
        self._antithetic_variate = antithetic_variate  # 是否使用对立变量
        self.seed = seed  # 随机数种子
        # 如果传入了路径矩阵，则不再生成新的路径矩阵
        if rands is None:
            self.rands = None  # 随机数矩阵
            self.__reset_rands = True  # 重置随机数标志位
        else:
            self.rands = rands  # 随机数矩阵
            self.__reset_rands = False  # 重置随机数标志位
        if s_paths is None and var_paths is None:
            self.s_paths = None  # 价格路径矩阵
            self.var_paths = None  # 方差路径矩阵
            self.__reset_paths = True  # 重置路径标志位
        else:  # 用户只要传入了一个路径矩阵，就视为用户已经设置了路径矩阵，其他矩阵也不再重新生成
            self.s_paths = s_paths  # 价格路径矩阵
            self.var_paths = var_paths  # 方差路径矩阵
            self.__reset_paths = False  # 重置路径标志位
            self.__reset_rands = False  # 重置随机数标志位

    def set_stoch_process(self, stoch_process):
        """设置随机过程，先将自己从原来的随机过程的观察者列表中移除，再将自己加入新的随机过程的观察者列表"""
        with suppress(AttributeError, ValueError):
            self.process.remove_observer(self)
        self.process = stoch_process
        stoch_process.add_observer(self)  # 将本实例加入stoch_process观察者列表
        self.rands = None  # 更换随机过程之后需要重新生成随机数
        self.__reset_rands = True  # 更换随机过程之后需要重新生成随机数
        self.__reset_paths = True  # 更换随机过程之后需要重新生成路径

    def set_s_paths(self, s_paths):
        """设置价格路径矩阵"""
        self.s_paths = s_paths
        self.__reset_paths = False  # 用户设置了路径矩阵，不需要重新生成路径
        self.__reset_rands = False  # 重置随机数标志位

    def remove_self(self):
        """删除对象自己，del自己之前，先将自己从被观察者的列表中移除"""
        self.process.remove_observer(self)
        del self

    def update(self, observable, *args, **kwargs):
        """被观察者的值发生变化时，自动调用update方法，通知观察者"""
        self.__reset_paths = True  # process自动向观察者发送通知, 参数已改变, 重置路径标志位，重新生成路径

    def reset_paths_flag(self):
        """重置路径标志位，重新生成路径"""
        self.__reset_paths = True

    @property
    def antithetic_variate(self):
        return self._antithetic_variate

    @antithetic_variate.setter
    def antithetic_variate(self, new_value):
        self._antithetic_variate = new_value
        self.__reset_paths = True  # 重置路径标志位，重新生成路径
        self.__reset_rands = True  # 重置随机数标志位，重新生成随机数

    @property
    def ld_method(self):
        return self._ld_method

    @ld_method.setter
    def ld_method(self, new_value):
        self._ld_method = new_value
        self.__reset_paths = True  # 重置路径标志位，重新生成路径
        self.__reset_rands = True  # 重置随机数标志位，重新生成随机数

    def randoms_generator(self, shape):  # 参数rands_method控制随机数获取方法，LD_method低差异序列生成方法；
        """生成随机数，返回shape形状的随机数矩阵
        如果需要的随机数矩阵的形状大于已有随机数矩阵的形状，则重新生成随机数矩阵，否则复用已有随机数矩阵
            rands_method: 控制随机数生成方法，包括伪随机数和低差异序列两种方法
            _ld_method: 控制低差异序列生成方法，包括Sobol和Halton两种低差异序列类型
        Args:
            shape: np.ndarray，(n_row, n_col)，随机数矩阵的形状
        Returns: self.rands，shape形状的随机数矩阵
        """
        if (self.rands is not None and (not self.__reset_rands)
                and self.rands.shape[0] >= shape[0] and self.rands.shape[1] >= shape[1]):
            logging.info("复用已有随机数")
            return self.rands[:shape[0], :shape[1]]
        else:
            np.random.seed(self.seed)
            if self.rands_method == RandsMethod.Pseudorandom:
                self.rands = np.random.standard_normal(shape)
            elif self.rands_method == RandsMethod.LowDiscrepancy:
                if self._ld_method == LdMethod.Sobol:
                    sampler = qmc.Sobol(d=shape[0], scramble=True, seed=self.seed)
                elif self._ld_method == LdMethod.Halton:
                    sampler = qmc.Halton(d=shape[0], scramble=True, seed=self.seed)
                else:
                    raise ValueError(
                        f'随机数生成方法输入错误，应为（Halton, Sobol）二者之一， 当前输入为{RandsMethod}')
                sampler.fast_forward(30)  # 跳过前三十项
                uniform_rands = sampler.random(shape[1]).transpose()
                np.random.shuffle(uniform_rands)
                self.rands = norm.ppf(uniform_rands)
            else:
                raise ValueError(
                    f'随机数生成方法输入错误，应为（RandsMethod.LowDiscrepancy, RandsMethod.Pseudorandom）二者之一， 当前输入为{RandsMethod}')
            self.__reset_rands = False  # 重置随机数标志位
            return self.rands

    def path_generator(self, n_step, spot=None, t_step_per_year=243):
        """根据标的种类和标的参数，生成价格路径，返回价格路径矩阵
        如果标的参数不变，且需要的价格路径矩阵的形状 ≤ 已有价格路径矩阵的形状，则复用已有价格路径矩阵，否则重新生成价格路径矩阵
        Args:
            n_step: int，价格路径的时间步数
            spot: float，标的期初价格
            t_step_per_year: int，每年的时间步数
        Returns: self.s_paths，价格路径矩阵
        """
        if (self.s_paths is not None and (not self.__reset_paths)
                and self.s_paths.shape[0] >= n_step + 1 and self.s_paths.shape[1] >= self.n_path
                and spot == self.process.spot()):
            logging.info("复用已有价格路径")
            return self.s_paths[:n_step + 1, :self.n_path]
        # else:
        return self._regenerate_paths(n_step, self.n_path, spot, t_step_per_year)

    def _regenerate_paths(self, n_step, n_path, spot, t_step_per_year):
        """生成价格路径的执行函数
        Args:
            n_step: int，价格路径的时间步数
            n_path: int，模拟路径数量
            spot: float，标的期初价格
            t_step_per_year: int，每年的时间步数
        Returns:
        """
        # self.s_paths = np.empty(shape=(n_step + 1, n_path))
        # self.s_paths[0] = spot
        dt = 1 / t_step_per_year

        if self.process() == ProcessType.BSProcess1D:
            if self.antithetic_variate:
                if n_path % 2 != 0:
                    n_path += 1
                half_n_path = n_path // 2
                self.rands = self.randoms_generator(shape=(n_step, half_n_path))
                rand_s = np.concatenate((self.rands, -self.rands), axis=1)
            else:
                self.rands = self.randoms_generator(shape=(n_step, n_path))
                rand_s = self.rands.copy()

            self.s_paths = np.empty(shape=(n_step + 1, n_path))
            self.s_paths[0] = spot
            for step in range(1, n_step + 1):
                self.s_paths[step] = self.process.evolve(dt * step, self.s_paths[step - 1], dt, rand_s[step - 1])

        elif self.process() == ProcessType.Heston:
            self.var_paths = np.empty(shape=(n_step + 1, n_path))
            if self.antithetic_variate:
                self.rands = self.randoms_generator(shape=(n_step, n_path))
                rand_s = self.rands[:, :self.rands.shape[1] // 2].copy()
                rand_v = self.rands[:, self.rands.shape[1] // 2:].copy()
                rand_v = self.process.var_rho * rand_s + np.sqrt(1 - self.process.var_rho ** 2) * rand_v
                rand_s = np.concatenate((rand_s, -rand_s), axis=1)
                rand_v = np.concatenate((rand_v, -rand_v), axis=1)
            else:
                self.rands = self.randoms_generator(shape=(n_step, n_path * 2))
                rand_s = self.rands[:, :self.rands.shape[1] // 2].copy()
                rand_v = self.rands[:, self.rands.shape[1] // 2:].copy()
                rand_v = self.process.var_rho * rand_s + np.sqrt(1 - self.process.var_rho ** 2) * rand_v

            self.s_paths = np.empty(shape=(n_step + 1, n_path))
            self.s_paths[0] = spot
            self.var_paths[0] = self.process.v0

            for step in range(1, n_step + 1):
                [self.s_paths[step], self.var_paths[step]] = self.process.evolve(dt * step, [self.s_paths[step - 1],
                                                                                             self.var_paths[step - 1]],
                                                                                 dt,
                                                                                 [rand_s[step - 1], rand_v[step - 1]])
            self.var_paths = np.maximum(self.var_paths, 0)  # heston方差非负修正：部分截断形式
        else:
            raise ValueError(f'随机过程类型输入错误，应为（ProcessType.BSProcess1D, ProcessType.Heston）二者之一，'
                             f'当前输入为{self.process()}')

        self.__reset_paths = False  # 重置路径标志位
        return self.s_paths

    def generate_reusable_price_paths(self, n_step=800, s_step=0.01, vol_step=0.01, r_step=0.01):
        """生成可复用的价格路径列表，包含计算希腊值需要的参数经过微小扰动的模拟价格路径. 仅支持BSM模型常数波动率、常数分红率、常数无风险利率
        对于一个指定的估值日，每个标的资产需要生成一个复用价格路径列表，因为每个标的资产有自己的波动率、分红率，
        两个标的资产相同的期权，其标的价格、波动率、分红率应该是相同的。另外，同一个估值日，所有期权的无风险利率应该是统一的。
        Args:
            n_step: int，价格路径的时间步数，默认生成800个交易日，因为大部分衍生品的到期时间不超过3年.
            s_step: float，价格步长。由于蒙特卡洛模拟使用有限差分计算希腊值时，步长较大时delta抖动较大，步长较小时gamma抖动较大，默认取0.01
            vol_step: float，默认vega计算波动率上升1%的期权价值变化
            r_step: float，默认rho计算无风险利率上升1%的期权价值变化
        Returns:
            s_paths_list: List，复用价格路径列表，用于有限差分法计算希腊值，其顺序为：
                          [原始价格路径，s-Δs价格路径，s+Δs价格路径，vol+0.01价格路径，r+0.01价格路径， # 价格路径的数据类型都是np.ndarray
                          s的原始数值，r的原始数值，s的步长（百分比），r的步长（百分比）]  # 这一行变量的数据类型都是float
        """
        s = self.process.spot.data
        r = self.process.interest.data
        vol = self.process.vol.data
        # 生成复用的随机数矩阵
        reusable_rands = self.randoms_generator(shape=(n_step, self.n_path))
        # 1. 生成原始价格路径
        s_paths_list = [self.path_generator(n_step=n_step, spot=s)]
        # 2. 生成s-Δs价格路径
        s_paths_list.append(self.path_generator(n_step=n_step, spot=(1 - s_step) * s))
        # 3. 生成s+Δs价格路径
        s_paths_list.append(self.path_generator(n_step=n_step, spot=(1 + s_step) * s))
        # 4. 生成vol+0.01价格路径
        self.process.vol.data = vol + vol_step
        s_paths_list.append(self.path_generator(n_step=n_step, spot=s))
        self.process.vol.data = vol
        # 4. 生成r+0.01价格路径
        self.process.interest.data = r + r_step
        s_paths_list.append(self.path_generator(n_step=n_step, spot=s))
        self.process.interest.data = r
        s_paths_list += [s, r, s_step, r_step]
        return s_paths_list

    def pv_and_greeks(self, prod, spot, s_paths_list: List=None, parallel=False):
        """根据复用的模拟价格路径计算期权希腊值
        Args:
            prod: Product类，期权产品对象
            spot: float, 标的价格，无实际用处，仅为了保持api参数统一
            s_paths_list: List，复用价格路径列表，其顺序为：
                          [原始价格路径，s-Δs价格路径，s+Δs价格路径，vol+0.01价格路径，r+0.01价格路径，
                          s的原始数值，r的原始数值，s的步长（百分比），r的步长（百分比）]  # 这一行变量的数据类型都是float
                          当串行计算时，路径数组是np.ndarray；当ray并行时，路径数组是ray.put存储到共享内存后的ObjectID
            parallel: bool，是否使用ray并行计算，默认False
        Returns:
            Dict[str:float]: {'pv': pv, 'delta': delta, 'gamma': gamma, 'theta': theta, 'vega': vega, 'rho': rho}
        """
        if s_paths_list is None:
            raise ValueError("McEngine.pv_and_greeks()必须输入复用价格路径列表s_paths_list.")
        assert len(s_paths_list) == 9, "复用价格路径列表s_paths_list应为长度9的列表，请使用McEngine的generate_reusable_price_paths方法生成"
        spot = s_paths_list[-4]
        r = s_paths_list[-3]
        s_step = s_paths_list[-2]
        r_step = s_paths_list[-1]
        if parallel:  # ray并行计算
            ray = importlib.import_module('ray')
        # 1.估值列表，存储不同参数微扰的估值结果，最后用于有限差分法计算希腊值。
        v_list = []  # 用于存储每种价格路径的估值结果
        # 2. 依次计算 s-Δs价格路径，s+Δs价格路径，vol+0.01价格路径，r+0.01价格路径 的估值
        for i in range(5):
            if i == 4:  # 计算rho的微扰后估值时，需要将无风险利率+1%，因为估值过程中，除了价格路径，折现也要用到r
                self.process.interest.data = r + r_step
            if parallel:  # ray并行计算，需要及时释放内存
                tmp_path = ray.get(s_paths_list[i])
                self.set_s_paths(tmp_path)
                del tmp_path
                gc.collect()
            else:  # 串行计算
                self.set_s_paths(s_paths_list[i])
            v_list.append(prod.price())
            if i == 0:  # 原始路径，既可用于计算pv，也可计算theta
                # 计算theta，假设标的价格不变，交易日推后一天的期权价值变化
                next_date = prod.trade_calendar.advance(global_evaluation_date(), datetime.timedelta(days=1))
                theta = prod.price(t=next_date, spot=spot) - v_list[0]
            if i == 4:  # 还原原始r
                self.process.interest.data = r
        # 4. 有限差分法计算希腊值
        ds = spot * s_step
        result = {"pv": v_list[0],
                  "delta": (v_list[2] - v_list[1]) / (2 * ds),
                  "gamma": (v_list[2] - 2 * v_list[0] + v_list[1]) / ds ** 2,
                  "vega": (v_list[3] - v_list[0]),
                  "theta": theta,
                  "rho": (v_list[4] - v_list[0]),
                  }
        return result


class McReusablePaths(McEngine):
    """专门用于生成可复用的Monte Carlo模拟价格路径的类"""

    def calc_present_value(self, prod, t: datetime.date = None, spot: float = None):
        """计算期权的现值，本来是抽象基类McEngine的抽象方法，
        这里重写calc_present_value，使子类McReusablePathGenerator可以被实例化"""
        raise NotImplementedError(
            "McReusablePathGenerator未实现计算期权价值的功能，该类仅用于生成可复用的Monte Carlo模拟价格路径！")
