import stocha
import numpy as np

# ── Random number generation ──────────────────────────────────────────────
rng = stocha.RNG(seed=42)
samples = rng.normal(size=10_000, loc=0.0, scale=1.0)

# ── GBM stock price simulation ────────────────────────────────────────────
paths = stocha.gbm(
    s0=100.0, mu=0.05, sigma=0.20,
    t=1.0, steps=252, n_paths=100_000, seed=42,
)
# paths.shape == (100_000, 253)

# ── Multi-asset correlated GBM ────────────────────────────────────────────
corr = np.array([[1.0, 0.6, 0.3],
                 [0.6, 1.0, 0.5],
                 [0.3, 0.5, 1.0]])
multi_paths = stocha.multi_gbm(
    s0=[100.0, 50.0, 200.0], mu=[0.05, 0.08, 0.03],
    sigma=[0.2, 0.3, 0.15], corr=corr,
    t=1.0, steps=252, n_paths=10_000, seed=42,
)
# multi_paths.shape == (10_000, 253, 3)
# Portfolio terminal value: (multi_paths[:, -1, :] * weights).sum(axis=1)

# ── Quasi-random sequences (low-discrepancy) ─────────────────────────────
pts = stocha.sobol(dim=2, n_samples=1024)    # (1024, 2), values in [0, 1)
pts = stocha.halton(dim=2, n_samples=1024)   # (1024, 2), values in (0, 1)

# ── Heston stochastic volatility ─────────────────────────────────────────
paths_h = stocha.heston(
    s0=100.0, v0=0.04, mu=0.05,
    kappa=2.0, theta=0.04, xi=0.3, rho=-0.7,
    t=1.0, steps=252, n_paths=10_000, seed=42,
)
# paths_h.shape == (10_000, 253)

# ── Merton jump-diffusion ─────────────────────────────────────────────────
paths_m = stocha.merton_jump_diffusion(
    s0=100.0, mu=0.05, sigma=0.20,
    lambda_=1.0, mu_j=-0.05, sigma_j=0.10,
    t=1.0, steps=252, n_paths=10_000, seed=42,
)
# paths_m.shape == (10_000, 253)

# ── VaR / CVaR ────────────────────────────────────────────────────────────
returns = paths[:, -1] / paths[:, 0] - 1
var, cvar = stocha.var_cvar(returns, confidence=0.95)
print(f"95% VaR={var:.4f}  CVaR={cvar:.4f}")

# ── Gaussian copula ───────────────────────────────────────────────────────
corr = np.array([[1.0, 0.8], [0.8, 1.0]])
u = stocha.gaussian_copula(corr, n_samples=10_000)
# u.shape == (10_000, 2),  values in (0, 1)

# ── Student-t copula (heavier joint tails) ────────────────────────────────
u_t = stocha.student_t_copula(corr, nu=5.0, n_samples=10_000)

# ── Hull-White short-rate model ───────────────────────────────────────────
rates = stocha.hull_white(
    r0=0.05, a=0.1, theta=0.005, sigma=0.01,
    t=1.0, steps=252, n_paths=10_000,
)
# rates.shape == (10_000, 253)

# ── SABR implied volatility ───────────────────────────────────────────────
iv = stocha.sabr_implied_vol(
    f=0.05, k=0.05, t=1.0,
    alpha=0.20, beta=0.5, rho=-0.3, nu=0.4,
)
print(f"SABR ATM implied vol: {iv:.4f}")

# ── SABR calibration: recover (alpha, rho, nu) from a market smile ───────
strikes = np.array([0.04, 0.045, 0.05, 0.055, 0.06])
market_vols = np.array([0.244, 0.218, 0.201, 0.190, 0.184])
fit = stocha.sabr_calibrate(strikes, market_vols, f=0.05, t=1.0, beta=0.5)
print(f"calibrated: alpha={fit['alpha']:.4f}, rho={fit['rho']:.4f}, nu={fit['nu']:.4f}")

# ── American option via LSMC ──────────────────────────────────────────────
price, std_err = stocha.lsmc_american_option(
    s0=100.0, k=100.0, r=0.05, sigma=0.20,
    t=1.0, steps=50, n_paths=50_000,
)
print(f"American put: {price:.4f} ± {std_err:.4f}")

# ── Heston analytical pricing (COS method) ───────────────────────────────
prices = stocha.heston_price(
    strikes=np.array([90.0, 100.0, 110.0]),
    is_call=[True, True, True],
    s0=100.0, v0=0.04, r=0.05,
    kappa=2.0, theta=0.04, xi=0.3, rho=-0.7, t=1.0,
)
print(f"Heston call prices: {prices}")

# ── Heston calibration ───────────────────────────────────────────────────
fit = stocha.heston_calibrate(
    strikes=np.array([90.0, 95.0, 100.0, 105.0, 110.0]),
    maturities=np.array([1.0, 1.0, 1.0, 1.0, 1.0]),
    market_prices=prices_from_market,  # your observed prices
    is_call=[True]*5,
    s0=100.0, r=0.05,
)
print(f"v0={fit['v0']:.4f}, kappa={fit['kappa']:.2f}, rho={fit['rho']:.2f}")

# ── Monte Carlo Greeks ───────────────────────────────────────────────────
greeks = stocha.greeks_fd(
    model="gbm",
    params={"s0": 100.0, "r": 0.05, "sigma": 0.2, "t": 1.0},
    payoff="call", strike=100.0,
    n_paths=100_000, n_steps=252,
    greeks=["delta", "gamma", "vega", "theta", "rho"],
)
print(f"Delta={greeks['delta']:.4f}  Gamma={greeks['gamma']:.4f}  Vega={greeks['vega']:.2f}")

# Pathwise method (GBM only, higher accuracy)
pw = stocha.greeks_pathwise(
    s0=100.0, r=0.05, sigma=0.2, t=1.0,
    strike=100.0, is_call=True,
    n_paths=100_000, n_steps=252,
    greeks=["delta", "vega"],
)
print(f"Pathwise Delta={pw['delta']:.4f}  Vega={pw['vega']:.2f}")

# ── Barrier option ──────────────────────────────────────────────────────
bp = stocha.barrier_price(
    s=100.0, k=100.0, r=0.05, sigma=0.2, t=1.0,
    barrier=120.0, barrier_type="up-and-out",
)
print(f"Up-and-out call: {bp:.4f}")

# ── Asian option (arithmetic average) ────────────────────────────────────
ap = stocha.asian_price(s=100.0, k=100.0, r=0.05, sigma=0.2, t=1.0)
print(f"Asian arithmetic call: {ap:.4f}")

# ── Lookback option (floating strike) ────────────────────────────────────
lp = stocha.lookback_price(s=100.0, r=0.05, sigma=0.2, t=1.0)
print(f"Lookback floating call: {lp:.4f}")