#!/usr/bin/env python3
"""
Infer two-cycle affine diagonal coupling:

b ≈ α a + Ψ (mod P)

Given samples (a_i, b_i), estimates a small rational α and checks whether a holdout-stable
affine coupling exists at the declared tolerance. Returns either a CST target or a full-torus
selection, together with a conservative null-model upper bound.
"""

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


def unwrap(seq, P):
    if not seq:
        return []
    out = [seq[0]]
    for x in seq[1:]:
        prev = out[-1]
        k = round((prev - x) / P)
        cand = x + k * P
        if abs((x + (k + 1) * P) - prev) < abs(cand - prev):
            cand = x + (k + 1) * P
        if abs((x + (k - 1) * P) - prev) < abs(cand - prev):
            cand = x + (k - 1) * P
        out.append(cand)
    return out


def circ_dist(x, P):
    y = x % P
    return min(y, P - y)


def minimal_covering_arc(points, P):
    pts = sorted([p % P for p in points])
    if len(pts) <= 1:
        center = pts[0] if pts else 0.0
        return 0.0, center
    gaps = [pts[i + 1] - pts[i] for i in range(len(pts) - 1)]
    gaps.append((pts[0] + P) - pts[-1])
    j = max(range(len(gaps)), key=lambda i: gaps[i])
    max_gap = gaps[j]
    L = P - max_gap
    start = pts[(j + 1) % len(pts)]
    center = (start + L / 2.0) % P
    return L, center


def null_upper_bound(n, tau, P):
    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(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 columns a,b (no header)")
    ap.add_argument("--tol", type=float, default=0.05, help="tolerance as fraction of P")
    ap.add_argument("--max_denom", type=int, default=12)
    ap.add_argument("--holdout", type=float, default=0.5)
    ap.add_argument("--p_threshold", type=float, default=1e-6)
    ap.add_argument("--seed", type=int, default=12345)
    args = ap.parse_args()

    P = float(args.period)
    tau = float(args.tol) * P
    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
        a, b = line.split(",")
        rows.append((float(a), float(b)))

    n = len(rows)
    idx = list(range(n))
    rnd = random.Random(int(args.seed))
    rnd.shuffle(idx)
    n_hold = int(round(n * holdout_frac))
    n_hold = max(0, min(n - 1, n_hold))
    hold_idx = set(idx[:n_hold])
    train = [rows[i] for i in range(n) if i not in hold_idx]
    hold = [rows[i] for i in range(n) if i in hold_idx]

    a_tr = [r[0] for r in train]
    b_tr = [r[1] for r in train]

    au = unwrap(a_tr, P)
    bu = unwrap(b_tr, P)

    mean_a = sum(au) / len(au)
    mean_b = sum(bu) / len(bu)
    num = sum((au[i] - mean_a) * (bu[i] - mean_b) for i in range(len(au)))
    den = sum((au[i] - mean_a) ** 2 for i in range(len(au))) or 1.0
    alpha = num / den

    frac = Fraction(alpha).limit_denominator(int(args.max_denom))
    alpha_r = float(frac)

    # Choose psi by minimizing the covering arc of residual phases on training data
    resid_tr = [(bu[i] - alpha_r * au[i]) % P for i in range(len(au))]
    L, psi = minimal_covering_arc(resid_tr, P)
    max_err_tr = 0.5 * L

    max_err_ho = None
    if hold:
        a_ho = unwrap([r[0] for r in hold], P)
        b_ho = unwrap([r[1] for r in hold], P)
        resid_ho = [(b_ho[i] - alpha_r * a_ho[i]) for i in range(len(a_ho))]
        max_err_ho = max(circ_dist(resid_ho[i] - psi, P) for i in range(len(a_ho))) if a_ho else 0.0

    p_train = null_upper_bound(len(train), tau, P)
    p_hold = null_upper_bound(len(hold), tau, P) if max_err_ho is not None else 1.0
    p_combined = float(min(1.0, p_train * p_hold))

    coupled = (max_err_tr <= tau) and (max_err_ho is None or max_err_ho <= tau) and (p_combined <= p_thr)

    out = {
        "period_P": P,
        "inferred_alpha": f"{frac.numerator}/{frac.denominator}",
        "inferred_psi": float(psi),
        "max_error_train": float(max_err_tr),
        "max_error_holdout": None if max_err_ho is None else float(max_err_ho),
        "tol_abs": float(tau),
        "target_type": "cst" if coupled else "torus",
        "significance": {
            "model": "independent_uniform_residual_arc_bound",
            "p_upper_train": float(p_train),
            "p_upper_holdout": float(p_hold),
            "p_upper_combined": float(p_combined),
            "p_threshold": float(p_thr),
        },
        "split": {"n_train": len(train), "n_holdout": len(hold)},
    }

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


if __name__ == "__main__":
    main()
