import json
import math
import numpy as np
from scipy.linalg import expm

SEED = 1468
rng = np.random.default_rng(SEED)
LAMBDA = 1e-8


def symvec(S):
    return np.array([S[0, 0], S[1, 1], S[2, 2], math.sqrt(2)*S[0, 1],
                     math.sqrt(2)*S[0, 2], math.sqrt(2)*S[1, 2]], dtype=float)


def symmat(c):
    S = np.zeros((3, 3), dtype=float)
    S[0, 0], S[1, 1], S[2, 2] = c[:3]
    S[0, 1] = S[1, 0] = c[3]/math.sqrt(2)
    S[0, 2] = S[2, 0] = c[4]/math.sqrt(2)
    S[1, 2] = S[2, 1] = c[5]/math.sqrt(2)
    return S


def axes_icosa():
    p = (1 + math.sqrt(5)) / 2
    raw = [(0,1,p), (1,p,0), (p,0,1), (0,1,-p),
           (1,-p,0), (p,0,-1)]
    return np.asarray(raw, dtype=float) / math.sqrt(1 + p*p)


def axes_coordinate():
    return np.eye(3)


def axes_random(n=6):
    x = rng.normal(size=(n, 3))
    return x / np.linalg.norm(x, axis=1, keepdims=True)


def frame(v, F):
    w = v @ F.T
    # Each row is the Frobenius-orthonormal vectorization of ww^T.
    A = np.stack([symvec(np.outer(q, q)) for q in w], axis=0)
    return A


def sl_deformation(anisotropy, rotation=True):
    H = np.diag([anisotropy, -anisotropy/2, -anisotropy/2])
    if rotation:
        Q, _ = np.linalg.qr(rng.normal(size=(3, 3)))
        if np.linalg.det(Q) < 0: Q[:, 0] *= -1
        H = Q @ H @ Q.T
    F = expm(H)
    # exp(trace-free H) has determinant one up to floating point precision.
    return F


def random_tracefree():
    S = rng.normal(size=(3, 3)); S = (S + S.T)/2
    return S - np.trace(S)*np.eye(3)/3


def reconstruct(v, F, S, noise_std=0.0, ridge=LAMBDA):
    A = frame(v, F)
    y = A @ symvec(S) + rng.normal(0, noise_std, size=len(v))
    c = np.linalg.solve(A.T @ A + ridge*np.eye(6), A.T @ y)
    # This is the stated optional trace-free postprocessing.
    Sh = symmat(c)
    Sh -= np.trace(Sh)*np.eye(3)/3
    return Sh, A


def rank_and_svals(v, F):
    A = frame(v, F)
    return np.linalg.matrix_rank(A, tol=1e-10), np.linalg.svd(A, compute_uv=False, full_matrices=False)


def main():
    ico = axes_icosa(); coord = axes_coordinate()
    # A fixed random frame is used as a secondary practical comparator.
    rand6 = axes_random()
    results = {"seed": SEED, "lambda": LAMBDA}

    # Prediction 1: at identity, the six icosa projectors span all Sym(3),
    # while coordinate projectors have rank exactly three.
    r_i, sv_i = rank_and_svals(ico, np.eye(3))
    r_c, sv_c = rank_and_svals(coord, np.eye(3))
    results["prediction_1_identity_rank"] = {
        "predicted": "ico rank=6; coordinate rank=3",
        "observed": {"ico_rank": r_i, "coordinate_rank": r_c,
                     "ico_sigma_min": float(sv_i[-1]),
                     "coordinate_smallest_singular_value": 0.0}}

    # Prediction 2: congruence by any invertible F preserves rank. Sweep SL(3)
    # deformation strengths and record ranks/minimum singular values.
    strengths = [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0]
    rank_sweep = []
    for a in strengths:
        F = sl_deformation(a)
        ri, si = rank_and_svals(ico, F)
        rc, sc = rank_and_svals(coord, F)
        rank_sweep.append({"anisotropy": a, "ico_rank": ri,
                           "coordinate_rank": rc,
                           "ico_sigma_min": float(si[-1]),
                           "coordinate_smallest_nonzero_singular_value": float(sc[-1]),
                           "coordinate_has_zero_singular_values": True})
    results["prediction_2_rank_preservation"] = {
        "predicted": "ico rank remains 6 for every invertible F; coordinate remains rank 3",
        "observed": rank_sweep}

    # Prediction 3: at fixed measurement noise, reconstruction error tracks
    # inverse conditioning and rises under strong anisotropy; compare frames.
    noise = 1e-3
    trials = 100
    error_sweep = []
    # reset independent RNG sequence is not needed; all results deterministic.
    for a in strengths:
        row = {"anisotropy": a}
        for name, v in [("icosahedral", ico), ("random6", rand6), ("coordinate", coord)]:
            errs, invmins, conds = [], [], []
            for _ in range(trials):
                F = sl_deformation(a)
                S = random_tracefree()
                Sh, A = reconstruct(v, F, S, noise_std=noise)
                errs.append(np.linalg.norm(Sh-S) / np.linalg.norm(S))
                ss = np.linalg.svd(A, compute_uv=False)
                invmins.append(1.0 / max(ss[-1], 1e-30))
                conds.append(float(ss[0] / max(ss[-1], 1e-30)))
            row[name] = {"relative_error_mean": float(np.mean(errs)),
                         "relative_error_std": float(np.std(errs)),
                         "inv_sigma_min_mean": float(np.mean(invmins)),
                         "condition_mean": float(np.mean(conds))}
        error_sweep.append(row)
    results["prediction_3_noise_conditioning"] = {
        "predicted": "noise error increases with inverse smallest singular value; coordinate is non-identifiable",
        "noise_std": noise, "trials": trials, "observed": error_sweep}

    # Exact noiseless reconstruction sanity check at random SL deformations.
    exact = {}
    for name, v in [("icosahedral", ico), ("random6", rand6)]:
        es = []
        for _ in range(100):
            F = sl_deformation(rng.uniform(0, 3)); S = random_tracefree()
            Sh, _ = reconstruct(v, F, S, noise_std=0.0, ridge=0.0)
            es.append(np.linalg.norm(Sh-S)/np.linalg.norm(S))
        exact[name] = float(np.max(es))
    results["exact_reconstruction_max_relative_error"] = exact

    # Compact verdict-relevant aggregates.
    ico_err = [x["icosahedral"]["relative_error_mean"] for x in error_sweep]
    ico_inv = [x["icosahedral"]["inv_sigma_min_mean"] for x in error_sweep]
    results["quantitative_predictions"] = {
        "P1_identity_rank": {"predicted_icosa_rank": 6, "observed_icosa_rank": int(r_i),
                             "predicted_coordinate_rank": 3, "observed_coordinate_rank": int(r_c)},
        "P2_SL3_rank_invariance": {"predicted_icosa_rank_at_all_strengths": 6,
                                    "observed_icosa_ranks": [int(x["ico_rank"]) for x in rank_sweep],
                                    "predicted_coordinate_rank_at_all_strengths": 3,
                                    "observed_coordinate_ranks": [int(x["coordinate_rank"]) for x in rank_sweep]},
        "P3_noise_amplification": {"prediction": "error rises as inverse smallest singular value rises",
                                    "strengths": strengths,
                                    "ico_error_at_strength_0_and_3": [ico_err[0], ico_err[-1]],
                                    "ico_inverse_sigma_at_strength_0_and_3": [ico_inv[0], ico_inv[-1]],
                                    "error_growth_factor": ico_err[-1] / ico_err[0],
                                    "inverse_sigma_growth_factor": ico_inv[-1] / ico_inv[0]}}
    results["summary"] = {
        "identity_ico_full_rank": bool(r_i == 6),
        "identity_coordinate_deficient": bool(r_c < 6),
        "all_ico_swept_full_rank": bool(all(x["ico_rank"] == 6 for x in rank_sweep)),
        "all_coordinate_swept_deficient": bool(all(x["coordinate_rank"] < 6 for x in rank_sweep)),
        "exact_ico_error": exact["icosahedral"],
        "exact_random6_error": exact["random6"]}
    with open("results.json", "w") as f:
        json.dump(results, f, indent=2, default=lambda x: x.item() if isinstance(x, np.generic) else x)
    print(json.dumps(results, indent=2, default=lambda x: x.item() if isinstance(x, np.generic) else x))

if __name__ == "__main__":
    main()