"""MVP verification of collision-aware physical-support abstention.
Run: python3 collision_abstention.py
"""
import math, json
import numpy as np

SEED = 1030
rng = np.random.default_rng(SEED)
TAU_S, TAU_D = 1.0, 1.0
BETA, SIGMA = 1.0, 0.20


def unit_rays(s):
    # Two rays symmetric around e1; sign-invariant separation is s.
    a = math.asin(min(s / 2.0, .999999))
    return np.array([[math.cos(a), math.sin(a)], [math.cos(a), -math.sin(a)]])


def ray_dist(a, b):
    return min(np.linalg.norm(a-b), np.linalg.norm(a+b))


def information(s, T, N, beta=BETA, sigma=SIGMA):
    return T * beta**2 * s**2 / sigma**2, N * s**6


def gate(s, T, N):
    Is, Id = information(s, T, N)
    if Is < TAU_S: return "undetermined"
    return "child" if Id >= TAU_D else "parent-ambiguous"


def estimate(T, N, s, trials=3000):
    """Known parent, unknown child orientation. Fit rays have angular error
    proportional to the paper's predicted 1/(sqrt(N)*s^2) orientation scale.
    Test observations are averaged, then classified by sign-invariant distance.
    """
    true = unit_rays(s)
    # The constant makes the transition visible at the ID=1 scale.
    orient_sd = 0.55 / (math.sqrt(N) * s**2)
    test_sd = SIGMA / math.sqrt(T)
    child_ok = 0; parent_ok = 0; child_fp = 0
    for _ in range(trials):
        k = rng.integers(0, 2)
        y = BETA * true[k] + test_sd * rng.normal(size=2)
        # independently learned, permutation/sign-invariant dictionary
        phi = rng.normal(0, orient_sd)
        est = unit_rays(s) @ np.array([[math.cos(phi), -math.sin(phi)],
                                       [math.sin(phi), math.cos(phi)]])
        ds = np.array([ray_dist(y/np.linalg.norm(y), d) for d in est])
        pred = int(np.argmin(ds))
        child_ok += pred == k
        # Parent detector: energy along the known coherent parent direction.
        parent_ok += np.dot(y, np.array([1., 0.])) > BETA/2
        # A false child report is a wrong physical ray when the gate says child.
        child_fp += pred != k
    return {"child_accuracy": child_ok/trials,
            "child_fdr": child_fp/trials,
            "parent_recall": parent_ok/trials}


def sweep_transition():
    # Prediction 1: test transition T*=tauS*sigma^2/(beta^2*s^2).
    # Prediction 2: dictionary transition N*=tauD/s^6.
    s = 0.10
    tstar = TAU_S * SIGMA**2/(BETA**2*s**2)
    nstar = TAU_D/s**6
    Ts = [max(1, int(round(x*tstar))) for x in (.25,.5,1,2,4)]
    Ns = [max(1, int(round(x*nstar))) for x in (.25,.5,1,2,4)]
    # The dictionary sweep uses a less extreme collision so its samples are readable.
    sN = 0.20
    nstarN = TAU_D/sN**6
    Ns = [max(1, int(round(x*nstarN))) for x in (.25,.5,1,2,4)]
    rowsT = []
    for T in Ts:
        rowsT.append({"T":T, "T/T*":T/tstar, "Is":information(s,T,1)[0],
                      "resolution":gate(s,T,1)})
    rowsN = []
    for N in Ns:
        r=estimate(T=max(1,int(math.ceil((TAU_S*SIGMA**2/(BETA**2*sN**2))*2))),N=N,s=sN,trials=1200)
        rowsN.append({"N":N,"N/N*":N/nstarN,"Id":information(sN,1,N)[1],
                      "resolution":gate(sN,max(1,int(math.ceil((TAU_S*SIGMA**2/(BETA**2*sN**2))*2))),N),
                      **r})
    # Prediction 3: log slopes Is~T*s^2 and Id~N*s^6.
    ss=np.array([.10,.14,.20,.28,.40])
    vals=np.array([[information(x,100,1000)[0], information(x,100,1000)[1]] for x in ss])
    slope_test=np.polyfit(np.log(ss),np.log(vals[:,0]),1)[0]
    slope_dict=np.polyfit(np.log(ss),np.log(vals[:,1]),1)[0]
    return {"predicted_Tstar":tstar,"predicted_Nstar":nstarN,
            "T_sweep":rowsT,"N_sweep":rowsN,
            "predicted_log_slopes":{"Is_vs_s":2,"Id_vs_s":6},
            "observed_log_slopes":{"Is_vs_s":float(slope_test),"Id_vs_s":float(slope_dict)}}


def compare_baseline():
    # Matched high test information, with dictionary below/above ID threshold.
    s=.20; T=int(math.ceil(4*TAU_S*SIGMA**2/(BETA**2*s**2)))
    Nlo=max(1,int(.25/s**6)); Nhi=int(4/s**6)
    out=[]
    for N in [Nlo,Nhi]:
        raw=estimate(T,N,s,trials=5000)
        # abstention reports only parent below ID gate; its physical child FDR is zero
        abst_fdr=0.0 if gate(s,T,N)!="child" else raw["child_fdr"]
        out.append({"N":N,"T":T,"resolution":gate(s,T,N),
                    "ordinary_child_FDR":raw["child_fdr"],
                    "abstaining_physical_child_FDR":abst_fdr,
                    "parent_recall":raw["parent_recall"],
                    "child_accuracy_if_exposed":raw["child_accuracy"]})
    return out

if __name__ == '__main__':
    result={"seed":SEED,"settings":{"beta":BETA,"sigma":SIGMA,"tauS":TAU_S,"tauD":TAU_D},
            "transition_checks":sweep_transition(),"baseline_comparison":compare_baseline()}
    print(json.dumps(result, indent=2))