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


def pairs(n):
    return [(i, j) for i in range(n) for j in range(i + 1, n)]


def wedge2(M):
    """Second multiplicative compound in the ordered Plucker basis."""
    p = pairs(M.shape[0])
    W = np.empty((len(p), len(p)))
    for a, (i, j) in enumerate(p):
        for b, (k, l) in enumerate(p):
            W[a, b] = M[i, k] * M[j, l] - M[i, l] * M[j, k]
    return W


def additive_compound(A):
    """A^[2] = d/dh wedge2(exp(h A)) at h=0."""
    n = A.shape[0]
    p = len(pairs(n))
    B = np.zeros((p, p))
    ij = pairs(n)
    for a, (i, j) in enumerate(ij):
        B[a, a] = A[i, i] + A[j, j]
        for b, (k, l) in enumerate(ij):
            if a != b:
                # A reliable direct derivative avoids sign-convention mistakes.
                h = 1e-6
                B[a, b] = (wedge2(expm(h * A))[a, b]) / h
    return B


def hilbert(u, v):
    u = np.maximum(np.asarray(u), 1e-300)
    v = np.maximum(np.asarray(v), 1e-300)
    r = u / v
    return float(np.log(r.max()) - np.log(r.min()))


def birkhoff_diameter(P):
    # Exact finite-dimensional formula for a strictly positive matrix.
    best = 0.0
    for i in range(P.shape[0]):
        for j in range(P.shape[0]):
            for k in range(P.shape[0]):
                for l in range(P.shape[0]):
                    best = max(best, abs(math.log(P[i, k] * P[j, l] /
                                                (P[i, l] * P[j, k]))))
    return best


def cycle_generator(N, alpha, diagonal=-1.0):
    B = diagonal * np.eye(N)
    # Fixed strongly connected directed cycle, with an additional reverse edge.
    for i in range(N):
        B[(i + 1) % N, i] = alpha
        B[(i - 1) % N, i] = alpha
    return B


def random_positive(rng, N):
    return np.exp(rng.normal(0.0, 1.0, N))


def main():
    rng = np.random.default_rng(1478)
    # Core algebra check: wedge(exp(hA)) = I + h A^[2] + O(h^2).
    A = rng.normal(size=(5, 5))
    B = additive_compound(A)
    hs = np.array([2e-2, 1e-2, 5e-3, 2.5e-3])
    errs = []
    for h in hs:
        errs.append(np.linalg.norm(wedge2(expm(h * A)) - np.eye(10) - h * B))
    slope = np.polyfit(np.log(hs), np.log(errs), 1)[0]

    # Toy mechanism: uniformly Metzler, fixed irreducible support, varying alpha.
    N, T = 6, 0.8
    alphas = [0.05, 0.1, 0.2, 0.4, 0.8]
    rows = []
    for alpha in alphas:
        Bp = cycle_generator(N, alpha)
        P = expm(T * Bp)
        delta = birkhoff_diameter(P)
        q_bound = math.tanh(delta / 4.0)
        # Empirical contraction over many positive vectors.
        ratios = []
        for _ in range(300):
            u, v = random_positive(rng, N), random_positive(rng, N)
            d0 = hilbert(u, v)
            if d0 > 1e-10:
                ratios.append(hilbert(P @ u, P @ v) / d0)
        # The shortest three-edge path controls the worst small-window entry on this 6-cycle.
        tsmall = 0.08
        minentry = expm(tsmall * Bp).min()
        predicted_path = (alpha * tsmall) ** 3 / 6.0
        # Repeated identical windows test geometric decay.
        u, v = random_positive(rng, N), random_positive(rng, N)
        ds = [hilbert(u, v)]
        for _ in range(5):
            u, v = P @ u, P @ v
            ds.append(hilbert(u, v))
        observed_window_ratios = [ds[k + 1] / ds[k] for k in range(5)]
        rows.append({
            "alpha": alpha,
            "min_P": float(P.min()),
            "q_birkhoff_bound": q_bound,
            "max_sampled_ratio": float(max(ratios)),
            "path_minentry_at_t0.08": float(minentry),
            "path_prediction_alpha3_t3_over_6": predicted_path,
            "repeated_window_ratios": observed_window_ratios,
        })

    # Unconstrained comparison: signed cycle has no positive-cone guarantee.
    Bu = cycle_generator(N, 0.4)
    Bu[1, 0] = -0.4  # break Metzler condition while retaining comparable scale
    Pu = expm(T * Bu)
    signed_min = float(Pu.min())
    unconstrained = []
    for _ in range(300):
        u, v = random_positive(rng, N), random_positive(rng, N)
        d0 = hilbert(u, v)
        if np.all(Pu @ u > 0) and np.all(Pu @ v > 0):
            unconstrained.append(hilbert(Pu @ u, Pu @ v) / d0)

    # Quantitative sweep summaries.
    loga = np.log(alphas)
    logpath = np.log([r["path_minentry_at_t0.08"] for r in rows])
    path_slope = float(np.polyfit(loga, logpath, 1)[0])
    q_by_alpha = [r["q_birkhoff_bound"] for r in rows]
    result = {
        "compound_identity": {
            "h": hs.tolist(), "errors": errs, "loglog_slope_expected_2": float(slope)
        },
        "predictions": {
            "P1_compound_identity_error_is_O_h2": {"observed_slope": float(slope), "expected": 2.0},
            "P2_shortest_path_min_entry_scales_alpha3": {
                "observed_log_slope": path_slope, "expected": 3.0,
                "note": "For the bidirectional 6-cycle, the worst off-diagonal entries require three edges; the leading path term is (alpha*t)^3/3!."
            },
            "P3_stronger_uniform_coupling_improves_projective_contraction": {
                "alphas": alphas, "q_bounds": q_by_alpha,
                "q_at_alpha_0.05": q_by_alpha[0], "q_at_alpha_0.8": q_by_alpha[-1]
            }
        },
        "rows": rows,
        "baseline_signed_generator": {
            "min_transition_entry": signed_min,
            "valid_positive_output_sample_ratio_max": (max(unconstrained) if unconstrained else None),
            "metzler_violation": 0.8
        },
        "interpretation": "Positive irreducible compound dynamics produced strictly positive finite-window transitions and geometric Hilbert contraction. The signed baseline has a negative transition entry, so the positive-cone theorem is unavailable."
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()