#!/usr/bin/env python3
"""
Infer small-integer coupling constraints Ay ≈ Psi (mod P) in T^k from sampled phase vectors.

Input
- CSV rows of k comma-separated values in [0,P) for a fixed period P.

Output
- If a statistically rare, holdout-stable small-integer constraint family is detected in the declared search,
  returns a CST target with (A, Psi), residual statistics, a stability margin, and a null-model upper bound.
- Otherwise returns a full-torus target selection inside the declared search, with the same diagnostics.
- If excitation or sample sufficiency fails, returns a structured WOB with evidence.

The null-model bound is conservative and is used as a gate against spurious relations:
under independent uniform phases, the probability that n points lie inside some arc of length 2τ
is at most n*(2τ/P)^(n-1). A union bound over the declared search space yields an explicit upper bound.
"""

import argparse
import itertools
import json
import math
from fractions import Fraction
from pathlib import Path

import numpy as np


def _percentile(sorted_vals, q):
    if not sorted_vals:
        return 0.0
    if q <= 0.0:
        return float(sorted_vals[0])
    if q >= 1.0:
        return float(sorted_vals[-1])
    idx = q * (len(sorted_vals) - 1)
    lo = int(math.floor(idx))
    hi = int(math.ceil(idx))
    if lo == hi:
        return float(sorted_vals[lo])
    w = idx - lo
    return float((1.0 - w) * sorted_vals[lo] + w * sorted_vals[hi])


def rank_integer_matrix(A):
    """Rank over Q (exact) for a small integer matrix A (list of rows)."""
    if not A:
        return 0
    m = len(A)
    n = len(A[0])
    M = [[Fraction(int(A[i][j]), 1) for j in range(n)] for i in range(m)]
    r = 0
    c = 0
    while r < m and c < n:
        pivot = None
        for i in range(r, m):
            if M[i][c] != 0:
                pivot = i
                break
        if pivot is None:
            c += 1
            continue
        if pivot != r:
            M[r], M[pivot] = M[pivot], M[r]
        pv = M[r][c]
        M[r] = [x / pv for x in M[r]]
        for i in range(m):
            if i == r:
                continue
            if M[i][c] != 0:
                factor = M[i][c]
                M[i] = [M[i][j] - factor * M[r][j] for j in range(n)]
        r += 1
        c += 1
    return r


def circ_dist(x, P):
    """Circular distance on R/PZ."""
    y = x % P
    return min(y, P - y)


def minimal_covering_arc(points, P):
    """
    Given points in [0,P), find the shortest arc length L that contains all points,
    and return (L, center) where center is a phase in [0,P) that centers that arc.
    """
    pts = np.asarray(points, dtype=float)
    if pts.size == 0:
        return 0.0, 0.0
    pts = np.mod(pts, P)
    pts.sort()
    if pts.size == 1:
        return 0.0, float(pts[0])
    gaps = np.diff(pts)
    wrap_gap = (pts[0] + P) - pts[-1]
    gaps = np.concatenate([gaps, np.array([wrap_gap])])
    j = int(np.argmax(gaps))
    max_gap = float(gaps[j])
    L = float(P - max_gap)
    start = float(pts[(j + 1) % pts.size])
    center = (start + L / 2.0) % P
    return L, center


def normalize_relation(a):
    """Normalize an integer relation vector up to sign and gcd so multiples collapse."""
    a = [int(x) for x in a]
    g = 0
    for x in a:
        g = math.gcd(g, abs(x))
    if g == 0:
        return tuple(a)
    a = [x // g for x in a]
    for x in a:
        if x != 0:
            if x < 0:
                a = [-y for y in a]
            break
    return tuple(a)


def null_upper_bound(num_candidates, n, tau, P):
    """
    Conservative union bound under independent uniform phases:
    P(exists candidate with all n points in some arc of length 2tau) <= num_candidates * n * (2tau/P)^(n-1)
    """
    if n <= 1:
        return 1.0
    frac = max(0.0, min(1.0, (2.0 * tau) / P))
    if frac == 0.0:
        return 0.0
    logp = math.log(max(1, num_candidates)) + math.log(float(n)) + (n - 1) * math.log(frac)
    if logp < -800:
        return 0.0
    return float(min(1.0, math.exp(logp)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--period", type=float, required=True)
    ap.add_argument("--data", required=True, help="CSV with k columns in [0,P) (no header)")
    ap.add_argument("--k", type=int, required=True, help="dimension")
    ap.add_argument("--max_coeff", type=int, default=4, help="max absolute coefficient per entry")
    ap.add_argument("--tol", type=float, default=0.08, help="absolute tolerance τ (same units as P)")
    ap.add_argument("--max_relations", type=int, default=2, help="max number of non-multiple relation rows to report")
    ap.add_argument("--holdout", type=float, default=0.5, help="fraction of samples reserved for holdout validation")
    ap.add_argument("--p_threshold", type=float, default=1e-6, help="null-model upper bound gate for accepting coupling")
    ap.add_argument("--seed", type=int, default=12345)
    args = ap.parse_args()

    P = float(args.period)
    dim = int(args.k)
    tau = float(args.tol)
    max_coeff = int(args.max_coeff)
    Lmax = int(args.max_relations)
    holdout_frac = float(args.holdout)
    p_thr = float(args.p_threshold)

    rows = []
    for line in Path(args.data).read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        parts = [p.strip() for p in line.split(",")]
        if len(parts) != dim:
            raise SystemExit(f"Expected {dim} columns, got {len(parts)}")
        rows.append([float(x) for x in parts])

    Y = np.asarray(rows, dtype=float)
    n = int(Y.shape[0])

    min_samples = max(25, 8 * dim)
    excitation_threshold = 4.0 * tau

    coord_ranges = []
    for j in range(dim):
        coord_ranges.append(float(np.max(Y[:, j]) - np.min(Y[:, j])))

    exc = {
        "coord_ranges": coord_ranges,
        "excitation_threshold": excitation_threshold,
        "weak_coords": [j for j, rng in enumerate(coord_ranges) if rng < excitation_threshold],
        "min_samples": min_samples,
        "num_samples": n,
    }

    if n < min_samples:
        print(json.dumps({
            "period_P": P,
            "target_type": "unknown",
            "status": "wob",
            "obstruction_code": "WOB-14-insufficient-excitation",
            "reason": "insufficient_samples_for_constraint_inference",
            "evidence": {"excitation": exc, "tol_abs": tau},
        }, indent=2))
        return

    weak = exc["weak_coords"]
    if len(weak) == dim:
        print(json.dumps({
            "period_P": P,
            "target_type": "unknown",
            "status": "wob",
            "obstruction_code": "WOB-14-insufficient-excitation",
            "reason": "insufficient_excitation_for_constraint_inference",
            "evidence": {"excitation": exc, "tol_abs": tau},
        }, indent=2))
        return

    if 0 < len(weak) < dim:
        print(json.dumps({
            "period_P": P,
            "target_type": "unknown",
            "status": "wob",
            "obstruction_code": "WOB-13-ambiguity",
            "reason": "partial_excitation_induces_rank_ambiguity",
            "evidence": {"excitation": exc, "tol_abs": tau},
            "recommended_next_moves": [
                "Increase excitation on the weak coordinates (longer run, larger drive variation).",
                "Restrict inference to the excited coordinates only, and treat the remaining coordinates as unconstrained.",
                "Reduce tolerance or increase sample count to separate competing small-integer hypotheses.",
            ],
        }, indent=2))
        return

    rng = np.random.default_rng(int(args.seed))
    idx = np.arange(n)
    rng.shuffle(idx)
    n_hold = int(round(n * holdout_frac))
    n_hold = max(0, min(n - 1, n_hold))
    hold_idx = idx[:n_hold]
    train_idx = idx[n_hold:]
    Ytr = Y[train_idx, :]
    Yho = Y[hold_idx, :] if n_hold > 0 else None

    # Enumerate candidate relation vectors
    cand = []
    for a in itertools.product(range(-max_coeff, max_coeff + 1), repeat=dim):
        if all(x == 0 for x in a):
            continue
        na = normalize_relation(a)
        cand.append(na)
    cand = sorted(set(cand))
    num_candidates = len(cand)

    # Evaluate each candidate on training data
    scored = []
    for a in cand:
        a_vec = np.asarray(a, dtype=float)
        z = np.mod(Ytr @ a_vec, P)
        L, center = minimal_covering_arc(z, P)
        max_res = 0.5 * L
        scored.append((max_res, a, center))

    scored.sort(key=lambda t: t[0])
    best_res, best_a, best_psi = scored[0]
    runner_res = scored[1][0] if len(scored) > 1 else float("inf")
    gap = max(0.0, float(runner_res - best_res))

    # Holdout validation for the best relation
    best_res_hold = None
    if Yho is not None and Yho.shape[0] > 0:
        a_vec = np.asarray(best_a, dtype=float)
        deltas = np.mod((Yho @ a_vec) - best_psi, P)
        d = np.minimum(deltas, P - deltas)
        best_res_hold = float(np.max(d)) if d.size > 0 else 0.0

    p_train = null_upper_bound(num_candidates, int(Ytr.shape[0]), tau, P)
    p_hold = null_upper_bound(num_candidates, int(Yho.shape[0]), tau, P) if best_res_hold is not None else 1.0
    p_combined = float(min(1.0, p_train * p_hold))

    accept = (best_res <= tau) and (best_res_hold is None or best_res_hold <= tau) and (p_combined <= p_thr)

    residuals = [s[0] for s in scored[: min(200, len(scored))]]
    residuals_sorted = sorted(residuals)

    diagnostics = {
        "num_candidates": num_candidates,
        "n_train": int(Ytr.shape[0]),
        "n_holdout": int(Yho.shape[0]) if Yho is not None else 0,
        "best_residual_train": float(best_res),
        "best_residual_holdout": None if best_res_hold is None else float(best_res_hold),
        "runner_up_residual_train": float(runner_res) if math.isfinite(runner_res) else None,
        "gap_best_to_runner_up": float(gap),
        "residual_percentiles_train": {
            "p10": _percentile(residuals_sorted, 0.10),
            "p50": _percentile(residuals_sorted, 0.50),
            "p90": _percentile(residuals_sorted, 0.90),
        },
    }

    significance = {
        "model": "independent_uniform_phases_union_bound",
        "p_upper_train": float(p_train),
        "p_upper_holdout": float(p_hold),
        "p_upper_combined": float(p_combined),
        "p_threshold": float(p_thr),
        "rule": "Accept coupling only when residual <= τ on training and holdout and p_upper_combined <= p_threshold.",
    }

    if not accept:
        out = {
            "period_P": P,
            "target_type": "torus",
            "status": "torus",
            "tol_abs": tau,
            "search_space": {"max_coeff": max_coeff, "max_relations": Lmax},
            "diagnostics": diagnostics,
            "excitation": exc,
            "significance": significance,
        }
        print(json.dumps(out, indent=2))
        return

    # Build a small row set (non-multiple) by greedily taking the best residuals
    rows = []
    psis = []
    used = set()
    for res, a, psi in scored:
        if res > tau:
            break
        if a in used:
            continue
        used.add(a)
        rows.append(list(a))
        psis.append(float(psi))
        if len(rows) >= Lmax:
            break

    r = rank_integer_matrix(rows)
    d = dim - r

    # Evaluate the selected matrix on training data
    A = np.asarray(rows, dtype=float)
    Psi = np.asarray(psis, dtype=float)
    Z = np.mod(Ytr @ A.T, P)
    # max residual per row
    row_max = []
    for j in range(A.shape[0]):
        deltas = np.mod(Z[:, j] - Psi[j], P)
        dists = np.minimum(deltas, P - deltas)
        row_max.append(float(np.max(dists)))
    rho = float(max(row_max)) if row_max else 0.0
    m_abs = float(max(0.0, tau - rho))

    out = {
        "period_P": P,
        "target_type": "cst",
        "status": "cst",
        "A": rows,
        "Psi": psis,
        "rank_r": int(r),
        "induced_target_rank": int(d),
        "residual_stats": {
            "row_max": row_max,
            "rho": rho,
            "best_relation_residual_train": float(best_res),
            "best_relation_residual_holdout": None if best_res_hold is None else float(best_res_hold),
        },
        "tol_abs": tau,
        "stability_margin_abs": m_abs,
        "stability_margin_fraction": float(m_abs / P) if P > 0 else 0.0,
        "search_space": {"max_coeff": max_coeff, "max_relations": Lmax},
        "diagnostics": diagnostics,
        "excitation": exc,
        "identifiability": {
            "eps_max_feasibility": float(m_abs / 2.0),
            "eps_max_stable_best_relation": float(max(0.0, gap) / 2.0),
            "rule": "If observation error <= eps_max_feasibility then feasibility at tolerance τ is stable. If observation error <= eps_max_stable_best_relation then the best-relation ranking is stable within the declared search.",
        },
        "significance": significance,
    }

    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()
