"""
Precision is a resource — experiments for the article
"Precision Is a Resource: How Much Arithmetic Does a Covariance Matrix Deserve?" (numericalmethod.com, 2026).

Requirements: Python >= 3.9, numpy >= 1.24. No other dependencies.
Run:          python3 experiments.py            (writes results.json next to this file)

Low precision is *emulated*: every arithmetic result is rounded to t significant bits
(t = 8 bfloat16, 11 binary16, 24 binary32, 53 binary64). Only the significand is emulated;
the exponent range is not (no overflow/underflow). Rounding is round-to-nearest-even
unless stated otherwise.
"""
import json
import math
import os
import time

import warnings

import numpy as np

# numpy 2.x + Apple Accelerate can emit spurious 'divide by zero/overflow in matmul' warnings.
# The results were checked against exactly rounded sums (math.fsum); the warnings are silenced here.
warnings.filterwarnings('ignore', message='.*encountered in matmul')

PRECISIONS = {"bf16": 8, "fp16": 11, "fp32": 24, "fp64": 53}


# ---------------------------------------------------------------------------
# Rounding to t significant bits
# ---------------------------------------------------------------------------
def rnd(x, t):
    """Round to nearest (ties to even) with a t-bit significand. Works on scalars and arrays."""
    if t >= 53:
        return x
    m, e = np.frexp(x)                     # x = m * 2**e, 0.5 <= |m| < 1
    return np.ldexp(np.round(m * 2.0 ** t) / 2.0 ** t, e)  # np.round: ties to even


def rnd_stochastic(x, t, rng):
    """Stochastic rounding with a t-bit significand: round up with probability equal to the fractional part."""
    m, e = math.frexp(x)
    s = m * 2.0 ** t
    lo = math.floor(s)
    q = lo + (1.0 if rng.random() < (s - lo) else 0.0)
    return math.ldexp(q / 2.0 ** t, e)


def rnd_nearest_scalar(x, t):
    m, e = math.frexp(x)
    return math.ldexp(round(m * 2.0 ** t) / 2.0 ** t, e)  # Python round: ties to even


# ---------------------------------------------------------------------------
# Cholesky and triangular solves in emulated precision (right-looking, every op rounded)
# ---------------------------------------------------------------------------
def cholesky_t(A, t):
    """Lower-triangular L with A ~= L L^T, every operation rounded to t bits. Returns None on breakdown."""
    A = rnd(np.array(A, dtype=np.float64), t)
    n = A.shape[0]
    L = np.zeros_like(A)
    for k in range(n):
        d = A[k, k]
        if not d > 0:
            return None
        lkk = rnd(math.sqrt(d), t)
        L[k, k] = lkk
        col = rnd(A[k + 1:, k] / lkk, t)
        L[k + 1:, k] = col
        # Schur complement update: product rounded, then difference rounded
        A[k + 1:, k + 1:] = rnd(A[k + 1:, k + 1:] - rnd(np.outer(col, col), t), t)
    return L


def forward_t(L, b, t):
    n = len(b)
    y = np.zeros(n)
    for i in range(n):
        s = b[i]
        row = L[i, :i]
        for j in range(i):
            s = rnd(s - rnd(row[j] * y[j], t), t)
        y[i] = rnd(s / L[i, i], t)
    return y


def backward_t(L, y, t):
    n = len(y)
    x = np.zeros(n)
    for i in range(n - 1, -1, -1):
        s = y[i]
        for j in range(i + 1, n):
            s = rnd(s - rnd(L[j, i] * x[j], t), t)
        x[i] = rnd(s / L[i, i], t)
    return x


def solve_t(A, b, t, L=None):
    L = cholesky_t(A, t) if L is None else L
    if L is None:
        return None, None
    return backward_t(L, forward_t(L, rnd(b, t), t), t), L


def min_variance_weights(x):
    return x / x.sum()


# ---------------------------------------------------------------------------
# Synthetic market: k-factor model, Sigma = B B^T + D (annualised-scale daily covariances)
# ---------------------------------------------------------------------------
def factor_model(n, k, rng, idio=(0.10, 0.30), market_strength=1.0):
    B = rng.normal(0.0, 0.15, size=(n, k)) / math.sqrt(252)
    B[:, 0] = market_strength * rng.uniform(0.5, 1.5, size=n) * 0.18 / math.sqrt(252)  # market factor
    D = np.diag((rng.uniform(*idio, size=n) / math.sqrt(252)) ** 2)
    return B @ B.T + D


def sample_cov(Sigma, T, rng):
    n = Sigma.shape[0]
    C = np.linalg.cholesky(Sigma)
    R = rng.standard_normal((T, n)) @ C.T
    R -= R.mean(axis=0)
    return R.T @ R / (T - 1)


def rel(a, b):
    return float(np.linalg.norm(a - b) / np.linalg.norm(b))


# ---------------------------------------------------------------------------
# Experiment 1: numerical vs statistical error of minimum-variance weights
# ---------------------------------------------------------------------------
def experiment_1(seeds=12):
    out = []
    ones = None
    for n in (50, 200):
        for T in (250, 1000, 4000, 16000):
            if T <= n:
                continue
            for s in range(seeds):
                rng = np.random.default_rng(1000 * n + 10 * T + s)
                Sigma = factor_model(n, 5, rng)
                ones = np.ones(n)
                w_true = min_variance_weights(np.linalg.solve(Sigma, ones))
                S = sample_cov(Sigma, T, rng)
                w64 = min_variance_weights(np.linalg.solve(S, ones))
                stat = rel(w64, w_true)
                kappa = float(np.linalg.cond(S))
                row = {"n": n, "T": T, "seed": s, "kappa": kappa, "stat": stat}
                for name, t in PRECISIONS.items():
                    if name == "fp64":
                        continue
                    x, _ = solve_t(S, ones, t)
                    row[name] = None if x is None else rel(min_variance_weights(x), w64)
                out.append(row)
                print(f"E1 n={n:4d} T={T:6d} s={s:2d} kappa={kappa:9.2e} stat={stat:.3e} "
                      + " ".join(f"{k}={row[k]:.2e}" if row[k] is not None else f"{k}=FAIL" for k in ("bf16", "fp16", "fp32")))
    return out


# ---------------------------------------------------------------------------
# Experiment 2: conditioning decides existence; refinement decides accuracy
#   Low-precision Cholesky as a preconditioner, residuals in fp64 (classical iterative refinement).
# ---------------------------------------------------------------------------
def iterative_refinement(A, b, t_factor, iters=12):
    L = cholesky_t(A, t_factor)
    if L is None:
        return None
    x = backward_t(L, forward_t(L, rnd(b, t_factor), t_factor), t_factor)
    hist = [x.copy()]
    for _ in range(iters):
        r = b - A @ x                                          # residual in fp64
        scale = np.max(np.abs(r))
        if scale == 0:
            hist.append(x.copy())
            continue
        d = backward_t(L, forward_t(L, rnd(r / scale, t_factor), t_factor), t_factor) * scale
        x = x + d                                              # update in fp64
        hist.append(x.copy())
    return hist


def experiment_2():
    out = []
    n, T = 100, 1000
    for strength in (1.0, 4.0, 12.0):        # stronger market factor -> larger condition number
        for idio_lo in (0.10, 0.02):
            rng = np.random.default_rng(int(strength * 100 + idio_lo * 1000))
            Sigma = factor_model(n, 5, rng, idio=(idio_lo, 0.30), market_strength=strength)
            S = sample_cov(Sigma, T, rng)
            b = np.ones(n)
            x64 = np.linalg.solve(S, b)
            kappa = float(np.linalg.cond(S))
            row = {"kappa": kappa, "market_strength": strength, "idio_lo": idio_lo}
            for name in ("bf16", "fp16", "fp32"):
                t = PRECISIONS[name]
                u = 2.0 ** (-t)
                hist = iterative_refinement(S, b, t)
                row[name] = {
                    "kappa_u": kappa * u,
                    "errors": None if hist is None else [rel(h, x64) for h in hist],
                }
            out.append(row)
            print(f"E2 kappa={kappa:9.2e} " + " ".join(
                f"{k}: k*u={row[k]['kappa_u']:.1e} " + ("FAIL" if row[k]['errors'] is None else f"e0={row[k]['errors'][0]:.1e} e_end={row[k]['errors'][-1]:.1e}")
                for k in ("bf16", "fp16", "fp32")))
    return out


# ---------------------------------------------------------------------------
# Experiment 3: long accumulations — round-to-nearest stagnates, stochastic rounding does not
#   Accumulate N daily P&L-sized increments in binary16-like precision (t = 11).
# ---------------------------------------------------------------------------
def experiment_3(runs=8):
    out = []
    t = 11
    for N in (2 ** 10, 2 ** 12, 2 ** 14, 2 ** 16, 2 ** 18):
        rng = np.random.default_rng(N)
        x = rng.uniform(0.5e-3, 1.5e-3, size=N)              # positive increments, mean 1e-3
        exact = math.fsum(x)
        s = 0.0
        for v in x:
            s = rnd_nearest_scalar(s + rnd_nearest_scalar(v, t), t)
        err_rn = abs(s - exact) / exact
        errs_sr = []
        for r in range(runs):
            g = np.random.default_rng(10_000 + r)
            s = 0.0
            for v in x:
                s = rnd_stochastic(s + rnd_nearest_scalar(v, t), t, g)
            errs_sr.append(abs(s - exact) / exact)
        row = {"N": N, "rn": err_rn, "sr_median": float(np.median(errs_sr)), "sr_max": float(np.max(errs_sr)),
               "rn_final_over_exact": None}
        out.append(row)
        print(f"E3 N={N:7d} RN rel.err={err_rn:.3e}  SR median={row['sr_median']:.3e} max={row['sr_max']:.3e}")
    return out


if __name__ == "__main__":
    t0 = time.time()
    results = {"experiment_1": experiment_1(), "experiment_2": experiment_2(), "experiment_3": experiment_3(),
               "numpy": np.__version__}
    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results.json")
    with open(path, "w") as f:
        json.dump(results, f, indent=1)
    print(f"done in {time.time() - t0:.0f}s -> {path}")