
"""
PKTron Finance — Quantum Monte Carlo (Amplitude Estimation)
===========================================================
Provides:
  simulate_returns()     — simulate portfolio return distribution
  estimate_var()         — Value at Risk (VaR) via amplitude estimation
  amplitude_estimation() — canonical QAE (iterative, classical-sim)

Memory: float32 throughout; max 12 GB safe.
"""
import numpy as np

_R = np.float32
_C = np.complex64

# ── Return simulation ────────────────────────────────────────────────────

def simulate_returns(mu, sigma, n_scenarios=50_000, horizon_days=252,
                     seed=42, use_gbm=True):
    """
    Simulate portfolio return paths.

    Parameters
    ----------
    mu           : float  annualised expected return
    sigma        : float  annualised volatility
    n_scenarios  : int    number of MC paths (max 100k for 12 GB)
    horizon_days : int    trading days
    use_gbm      : bool   Geometric Brownian Motion if True, else Normal

    Returns
    -------
    dict:
        terminal_returns : float32 array (n_scenarios,)
        paths            : float32 array (n_scenarios, horizon_days) — subset
        daily_mu, daily_sigma
    """
    n_scenarios = min(n_scenarios, 100_000)
    rng = np.random.default_rng(seed)
    dt  = _R(1.0 / 252)
    mu  = _R(mu);  sigma = _R(sigma)

    daily_mu    = _R(mu * dt)
    daily_sigma = _R(sigma * np.sqrt(dt))

    # Generate all shocks at once (memory-efficient block)
    shocks = rng.standard_normal((n_scenarios, horizon_days)).astype(_R)

    if use_gbm:
        log_ret = (daily_mu - 0.5*daily_sigma**2) + daily_sigma * shocks
        cum_ret = np.cumsum(log_ret, axis=1)
        paths   = np.exp(cum_ret) - 1
    else:
        log_ret = daily_mu + daily_sigma * shocks
        paths   = np.cumsum(log_ret, axis=1)

    terminal = paths[:, -1].astype(_R)

    return {'terminal_returns': terminal,
            'paths'           : paths[:, ::max(1, horizon_days//50)],
            'daily_mu'        : float(daily_mu),
            'daily_sigma'     : float(daily_sigma)}


def estimate_var(terminal_returns, confidence=0.95):
    """
    Value at Risk (VaR) and Conditional VaR (CVaR / Expected Shortfall).

    Parameters
    ----------
    terminal_returns : array-like  simulated returns
    confidence       : float       confidence level (e.g. 0.95)

    Returns
    -------
    dict: var, cvar, mean_return, std_return
    """
    r   = np.asarray(terminal_returns, dtype=_R)
    var = float(np.percentile(r, (1 - confidence) * 100))
    cvar = float(r[r <= var].mean()) if (r <= var).any() else var
    return {'var'         : var,
            'cvar'        : cvar,
            'mean_return' : float(r.mean()),
            'std_return'  : float(r.std()),
            'confidence'  : confidence}


def amplitude_estimation(p_target, n_shots=1024, n_qpe_qubits=5, seed=0):
    """
    Iterative Quantum Amplitude Estimation (IQAE) — classically simulated.
    Estimates probability p that a quantum oracle marks the "good" states.

    Used for: option pricing (p = normalised expected payoff),
              VaR threshold probability, etc.

    Parameters
    ----------
    p_target     : float   true probability to estimate (oracle encodes this)
    n_shots      : int     measurement budget
    n_qpe_qubits : int     precision qubits (resolution = 1/2^n_qpe_qubits)

    Returns
    -------
    dict: estimate, error, true_value, n_shots_used
    """
    rng       = np.random.default_rng(seed)
    theta_ref = float(np.arcsin(np.sqrt(max(0.0, min(1.0, p_target)))))

    # Canonical QAE phase: theta such that p = sin^2(theta)
    # QPE would estimate theta / pi.  Here we simulate measurement outcomes.
    resolution = 1.0 / 2**n_qpe_qubits
    n_grid     = 2**n_qpe_qubits

    # Probability distribution over QPE outcomes (Dirichlet kernel)
    grid   = np.arange(n_grid) / n_grid * np.pi   # candidate thetas
    kernel = np.sin(n_grid * (grid - theta_ref)) / (n_grid * np.sin(grid - theta_ref + 1e-12))
    kernel[np.isnan(kernel)] = 1.0
    probs  = kernel**2
    probs  = np.abs(probs); probs /= probs.sum()

    # Simulate n_shots QPE measurements
    outcomes   = rng.choice(n_grid, size=n_shots, p=probs)
    best_k     = int(np.bincount(outcomes).argmax())
    theta_est  = best_k / n_grid * np.pi
    p_est      = float(np.sin(theta_est)**2)

    return {'estimate'     : p_est,
            'true_value'   : p_target,
            'error'        : abs(p_est - p_target),
            'n_shots_used' : n_shots,
            'resolution'   : resolution}


def price_european_option(S0, K, r, sigma, T, n_scenarios=20_000,
                          option_type='call', seed=99):
    """
    Classical MC option pricing — baseline for QMC comparison.

    Returns: price, stderr
    """
    rng   = np.random.default_rng(seed)
    Z     = rng.standard_normal(n_scenarios).astype(_R)
    ST    = _R(S0) * np.exp((_R(r) - 0.5*_R(sigma)**2)*_R(T)
                             + _R(sigma)*np.sqrt(_R(T))*Z)
    if option_type == 'call':
        payoff = np.maximum(ST - _R(K), _R(0))
    else:
        payoff = np.maximum(_R(K) - ST, _R(0))
    disc  = _R(np.exp(-r * T))
    price = float(disc * payoff.mean())
    stderr= float(disc * payoff.std() / np.sqrt(n_scenarios))
    return {'price': price, 'stderr': stderr,
            'option_type': option_type, 'S0': S0, 'K': K}
