"""Statistical estimators used throughout the Q3 hardware analysis.

Every confidence interval, suppression ratio, correlation, and calibration
fit reported in the paper is produced by one of the functions below.  The
formulas are exactly those stated in the "Hardware platform and methods"
section:

  * binomial fractions p_hat = x / n with 95% Wilson score intervals,
  * suppression ratios R = p_a / p_b with log-ratio (delta-method) intervals,
  * Pearson correlation with a Fisher z-transform interval,
  * the floor-plus-gain calibration model p_hw(c) = p_floor + g c^2/(1+c^2).

The default critical value z = 1.96 corresponds to a two-sided 95% interval.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np

Z95 = 1.96


@dataclass(frozen=True)
class Interval:
    """A point estimate together with a lower/upper confidence bound."""

    point: float
    low: float
    high: float

    def __repr__(self) -> str:  # pragma: no cover - cosmetic
        return f"{self.point:.4g} [{self.low:.4g}, {self.high:.4g}]"


def wilson_ci(x: int, n: int, z: float = Z95) -> Interval:
    """Wilson score interval for a binomial proportion.

    Implements the interval quoted in the methods section,

        ( p_hat + z^2/2n  +/-  z sqrt( p_hat(1-p_hat)/n + z^2/4n^2 ) ) / (1 + z^2/n),

    and returns the *raw* point estimate p_hat = x/n alongside the Wilson
    lower/upper bounds (the paper tabulates p_hat and the interval separately).
    """
    if n <= 0:
        raise ValueError("n must be positive")
    if not (0 <= x <= n):
        raise ValueError("require 0 <= x <= n")
    p = x / n
    denom = 1.0 + z * z / n
    center = (p + z * z / (2 * n)) / denom
    half = (z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))) / denom
    return Interval(point=p, low=center - half, high=center + half)


def suppression_ratio(
    x_a: int, n_a: int, x_b: int, n_b: int, z: float = Z95
) -> Interval:
    """Suppression ratio R = p_a / p_b with a log-ratio (delta) interval.

    The standard error of log R is

        SE(log R) = sqrt( (1 - p_a)/x_a + (1 - p_b)/x_b ),

    exactly as stated in the methods section, and the interval is
    R * exp(+/- z SE).  Use a = control (large leakage) and b = protected
    (small leakage) so that R > 1 is a suppression factor.
    """
    p_a = x_a / n_a
    p_b = x_b / n_b
    if p_a == 0 or p_b == 0:
        raise ValueError("log-ratio undefined for a zero rate; use a count >= 1")
    R = p_a / p_b
    se = math.sqrt((1 - p_a) / x_a + (1 - p_b) / x_b)
    return Interval(point=R, low=R * math.exp(-z * se), high=R * math.exp(z * se))


def pearson_fisher_ci(x: np.ndarray, y: np.ndarray, z: float = Z95) -> Interval:
    """Pearson correlation r with a Fisher z-transform confidence interval."""
    x = np.asarray(x, dtype=float)
    y = np.asarray(y, dtype=float)
    n = x.size
    if n < 4:
        raise ValueError("Fisher interval needs n >= 4 points")
    r = float(np.corrcoef(x, y)[0, 1])
    zf = math.atanh(r)
    se = 1.0 / math.sqrt(n - 3)
    lo = math.tanh(zf - z * se)
    hi = math.tanh(zf + z * se)
    return Interval(point=r, low=lo, high=hi)


def dc_dump_theory(c: np.ndarray | float) -> np.ndarray | float:
    """Ideal DC-leakage response p(c) = c^2 / (1 + c^2) (Eq. for Theory dump)."""
    c2 = np.asarray(c, dtype=float) ** 2
    return c2 / (1.0 + c2)


def fit_floor_gain(c: np.ndarray, p_hw: np.ndarray) -> tuple[float, float]:
    """Least-squares fit of p_hw(c) = p_floor + g * c^2/(1+c^2).

    The model is linear in (p_floor, g), so the fit is a plain linear
    least-squares solve with design columns [1, c^2/(1+c^2)].
    Returns (p_floor, g).
    """
    c = np.asarray(c, dtype=float)
    p_hw = np.asarray(p_hw, dtype=float)
    shape = dc_dump_theory(c)
    design = np.column_stack([np.ones_like(shape), shape])
    (p_floor, g), *_ = np.linalg.lstsq(design, p_hw, rcond=None)
    return float(p_floor), float(g)


def fmt_pct(value: float, digits: int = 2) -> str:
    """Format a fraction as a percentage string (for table printing)."""
    return f"{100 * value:.{digits}f}%"


if __name__ == "__main__":  # pragma: no cover - smoke test
    # Reproduce the |0> - |1> neutral row of Table II: 70 / 10000.
    ci = wilson_ci(70, 10_000)
    print("p(|0>-|1>) =", ci, "(paper: 0.0073 [0.0058, 0.0092])")
