from __future__ import annotations

import numpy as np


def payoff_stock(s_t: np.ndarray, spot0: float, units: float = 1.0) -> np.ndarray:
    return units * (s_t - spot0)


def payoff_protective_put(
    s_t: np.ndarray,
    spot0: float,
    strike_put: float,
    put_premium: float,
    units: float = 1.0,
) -> np.ndarray:
    stock = units * (s_t - spot0)
    put = units * np.maximum(strike_put - s_t, 0.0) - units * put_premium
    return stock + put


def payoff_collar(
    s_t: np.ndarray,
    spot0: float,
    strike_put: float,
    put_premium: float,
    strike_call: float,
    call_premium: float,
    units: float = 1.0,
) -> np.ndarray:
    pp = payoff_protective_put(s_t, spot0, strike_put, put_premium, units)
    short_call = -units * np.maximum(s_t - strike_call, 0.0) + units * call_premium
    return pp + short_call


def risk_metrics(
    payoffs: np.ndarray, probs: np.ndarray | None = None, floor: float | None = None
) -> dict[str, float]:
    p = np.asarray(payoffs, dtype=float)
    if probs is None:
        w = np.full_like(p, 1.0 / len(p), dtype=float)
    else:
        w = np.asarray(probs, dtype=float)
        w = w / w.sum()
    expected = float(np.sum(p * w))
    q05 = float(np.quantile(p, 0.05))
    out = {"expected_payoff": expected, "q05": q05}
    if floor is not None:
        out["floor_gap_min"] = float(np.min(p - floor))
    return out