from typing import Optional

import torch
from torch import Tensor

from pfhedge._utils.doc import _set_attr_and_docstring
from pfhedge._utils.doc import _set_docstring
from pfhedge._utils.str import _format_float
from pfhedge.nn.functional import european_payoff

from ..primary.base import BasePrimary
from .base import BaseDerivative
from .base import OptionMixin


class EuropeanOption(BaseDerivative, OptionMixin):
    r"""European option.

    The payoff of a European call option is given by:

    .. math::
        \mathrm{payoff} = \max(S - K, 0) ,

    where
    :math:`S` is the underlier's spot price at maturity and
    :math:`K` is the strike.

    The payoff of a European put option is given by:

    .. math::
        \mathrm{payoff} = \max(K - S, 0) .

    .. seealso::
        - :func:`pfhedge.nn.functional.european_payoff`

    Args:
        underlier (:class:`BasePrimary`): The underlying instrument of the option.
        call (bool, default=True): Specifies whether the option is call or put.
        strike (float, default=1.0): The strike price of the option.
        maturity (float, default=20/250): The maturity of the option.

    Attributes:
        dtype (torch.dtype): The dtype with which the simulated time-series are
            represented.
        device (torch.device): The device where the simulated time-series are.

    Examples:
        >>> import torch
        >>> from pfhedge.instruments import BrownianStock
        >>> from pfhedge.instruments import EuropeanOption
        >>>
        >>> _ = torch.manual_seed(42)
        >>> derivative = EuropeanOption(BrownianStock(), maturity=5/250)
        >>> derivative.simulate(n_paths=2)
        >>> derivative.underlier.spot
        tensor([[1.0000, 1.0016, 1.0044, 1.0073, 0.9930, 0.9906],
                [1.0000, 0.9919, 0.9976, 1.0009, 1.0076, 1.0179]])
        >>> derivative.payoff()
        tensor([0.0000, 0.0179])

        Using custom ``dtype`` and ``device``.

        >>> derivative = EuropeanOption(BrownianStock())
        >>> derivative.to(dtype=torch.float64, device="cuda:0")
        EuropeanOption(
          ...
          (underlier): BrownianStock(..., dtype=torch.float64, device='cuda:0')
        )

        Make ``self`` a listed derivative.

        >>> from pfhedge.nn import BlackScholes
        >>>
        >>> pricer = lambda derivative: BlackScholes(derivative).price(
        ...     log_moneyness=derivative.log_moneyness(),
        ...     time_to_maturity=derivative.time_to_maturity(),
        ...     volatility=derivative.ul().volatility)
        >>> derivative = EuropeanOption(BrownianStock(), maturity=5/250)
        >>> derivative.list(pricer, cost=1e-4)
        >>> derivative.simulate(n_paths=2)
        >>> derivative.ul().spot
        tensor([[1.0000, 0.9788, 0.9665, 0.9782, 0.9947, 1.0049],
                [1.0000, 0.9905, 1.0075, 1.0162, 1.0119, 1.0220]])
        >>> derivative.spot
        tensor([[0.0113, 0.0028, 0.0006, 0.0009, 0.0028, 0.0049],
                [0.0113, 0.0060, 0.0130, 0.0180, 0.0131, 0.0220]])

        Add a knock-out clause with a barrier at 1.03:

        >>> _ = torch.manual_seed(42)
        >>> derivative = EuropeanOption(BrownianStock(), maturity=5/250)
        >>> derivative.simulate(n_paths=8)
        >>> derivative.payoff()
        tensor([0.0000, 0.0000, 0.0113, 0.0414, 0.0389, 0.0008, 0.0000, 0.0000])
        >>>
        >>> def knockout(derivative, payoff):
        ...     max = derivative.underlier.spot.max(-1).values
        ...     return payoff.where(max < 1.03, torch.zeros_like(max))
        >>>
        >>> derivative.add_clause("knockout", knockout)
        >>> derivative
        EuropeanOption(
          strike=1., maturity=0.0200
          clauses=['knockout']
          (underlier): BrownianStock(sigma=0.2000, dt=0.0040)
        )
        >>> derivative.payoff()
        tensor([0.0000, 0.0000, 0.0113, 0.0000, 0.0000, 0.0008, 0.0000, 0.0000])
    """

    def __init__(
        self,
        underlier: BasePrimary,
        call: bool = True,
        strike: float = 1.0,
        maturity: float = 20 / 250,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
    ) -> None:
        super().__init__()
        self.register_underlier("underlier", underlier)
        self.call = call
        self.strike = strike
        self.maturity = maturity

        # TODO(simaki): Remove later. Deprecated for > v0.12.3
        if dtype is not None or device is not None:
            self.to(dtype=dtype, device=device)
            raise DeprecationWarning(
                "Specifying device and dtype when constructing a Derivative is deprecated."
                "Specify them in the constructor of the underlier instead."
            )

    def extra_repr(self) -> str:
        params = []
        if not self.call:
            params.append("call=" + str(self.call))
        params.append("strike=" + _format_float(self.strike))
        params.append("maturity=" + _format_float(self.maturity))
        return ", ".join(params)

    def payoff_fn(self) -> Tensor:
        return european_payoff(self.ul().spot, call=self.call, strike=self.strike)


# Assign docstrings so they appear in Sphinx documentation
_set_attr_and_docstring(EuropeanOption, "simulate", BaseDerivative.simulate)
_set_attr_and_docstring(EuropeanOption, "to", BaseDerivative.to)
_set_attr_and_docstring(EuropeanOption, "ul", BaseDerivative.ul)
_set_attr_and_docstring(EuropeanOption, "list", BaseDerivative.list)
_set_docstring(EuropeanOption, "payoff", BaseDerivative.payoff)
_set_attr_and_docstring(EuropeanOption, "moneyness", OptionMixin.moneyness)
_set_attr_and_docstring(EuropeanOption, "log_moneyness", OptionMixin.log_moneyness)
_set_attr_and_docstring(
    EuropeanOption, "time_to_maturity", OptionMixin.time_to_maturity
)
