#!/usr/bin/env python3
"""
Computation 36 -- Optimal-alpha closure diagnostic for the holomorphic bridge
==============================================================================
Computation 35 established that the holomorphic Toeplitz subalgebra is the
right target for the bridge image (abelian preservation), but two gaps
remain: fidelity decay with weight, and residual spurious commutators in
the "a not in S" sector under the simplest Clifford analog T_{bar z_a}.

This computation asks the diagnostic question:

  Does there exist a SINGLE alpha = alpha(D) such that the bridge
  Toeplitz commutator approximately matches the substrate commutator
  across ALL test cases (a, S) simultaneously?

If yes, the framework is close to closure and the residual is the
fidelity rescaling.  If no -- if different (a, S) pairs require
different alpha values -- the holomorphic bridge with T_{bar z_a} as the
Clifford analog is structurally insufficient and a more elaborate
construction (Berezin derivative with alpha-correction; spinor extension
on a tensor C^2 factor) is needed.

METHOD.
At D = 4, sweep alpha over a fine grid.  For each (a, S) test case,
compute the substrate commutator norm (known: 2 or 0) and the
Toeplitz bridge commutator norm.  Define the per-case match score:

    score(a, S, alpha) = || ||Toep_comm(alpha)|| - ||sub_comm|| ||  (closer to 0 is better)

Sum these over all test cases to get an aggregate alpha-score.  The
optimal alpha minimises the aggregate.  If the optimal alpha gives a
small per-case score for ALL cases, closure is reachable; if some
cases are at the optimal alpha but others have large gaps, the bridge
is insufficient.

OUTPUT.
For each (a, S) test case, the alpha that BEST matches that case,
and the aggregate optimal alpha.  Identifies whether closure is
reachable in this framework.
"""
import math
from itertools import combinations

import numpy as np
import numpy.linalg as la


sx = np.array([[0, 1], [1, 0]], dtype=complex)
sy = np.array([[0, -1j], [1j, 0]], dtype=complex)
sz = np.array([[1, 0], [0, -1]], dtype=complex)
I2 = np.eye(2, dtype=complex)


def kron_chain(ops):
    out = ops[0]
    for op in ops[1:]:
        out = np.kron(out, op)
    return out


def chi_walsh(D, S):
    return kron_chain([sz if a in S else I2 for a in range(D)])


def chi_clifford(D, a):
    return kron_chain([sz] * a + [sx] + [I2] * (D - 1 - a))


def op_norm(M):
    return float(la.norm(M, ord=2))


def basis_indices(N):
    return [(m1, m2) for m1 in range(N + 1) for m2 in range(N + 1 - m1)]


def log_norm_sq(m1, m2, alpha):
    return (math.lgamma(m1 + 1) + math.lgamma(m2 + 1)
            + math.lgamma(alpha + 3) - math.lgamma(m1 + m2 + alpha + 3))


def Tz_matrix(a, basis, alpha):
    n = len(basis)
    idx = {J: i for i, J in enumerate(basis)}
    M = np.zeros((n, n), dtype=complex)
    for j, J in enumerate(basis):
        Jp = list(J)
        Jp[a] += 1
        Jp = tuple(Jp)
        if Jp in idx:
            log_ratio = 0.5 * (log_norm_sq(Jp[0], Jp[1], alpha)
                               - log_norm_sq(J[0], J[1], alpha))
            M[idx[Jp], j] = math.exp(log_ratio)
    return M


def Tzbar_matrix(a, basis, alpha):
    return Tz_matrix(a, basis, alpha).conj().T


def site_to_holo(D, basis, alpha, Tz):
    return {
        0: Tz[0],
        1: Tz[1],
        2: Tz[0] @ Tz[0],
        3: Tz[1] @ Tz[1],
    }


def bridge_holo(D, S, basis, alpha, Tz):
    site_map = site_to_holo(D, basis, alpha, Tz)
    n = len(basis)
    out = np.eye(n, dtype=complex)
    for a in sorted(S):
        out = out @ site_map[a]
    return out


def main():
    D = 4
    N = 5
    basis = basis_indices(N)

    print("=" * 90)
    print("  Computation 36  --  optimal-alpha closure diagnostic for holomorphic bridge")
    print("=" * 90)
    print()
    print(f"  D = {D}, N = {N}")
    print()

    # Test cases from Computation 35
    cases = [
        (0, {0}),
        (0, {1}),
        (0, {0, 1}),
        (0, {1, 2}),
        (1, {0, 1}),
        (1, {0, 2}),
        (2, {0, 2}),
        (3, {0, 3}),
    ]

    # Precompute substrate commutator norms (alpha-independent)
    sub_norms = {}
    for a_idx, S in cases:
        chi_a = chi_clifford(D, a_idx)
        chi_S = chi_walsh(D, S)
        sub_norms[(a_idx, frozenset(S))] = op_norm(chi_a @ chi_S - chi_S @ chi_a)

    # Alpha sweep
    alphas_fine = np.linspace(-0.5, 30.0, 100)

    # For each case, find alpha that minimises |Toep_norm - sub_norm|
    print("  Per-case optimal alpha (where Toeplitz commutator matches substrate)")
    print(f"  {'case':>20}  {'||sub||':>10}  {'best alpha':>14}  {'Toep at best alpha':>22}  {'residual':>10}")
    per_case_optimal = {}
    for a_idx, S in cases:
        sub_norm = sub_norms[(a_idx, frozenset(S))]
        best_alpha = None
        best_diff = float('inf')
        best_toep = None
        for alpha in alphas_fine:
            Tz = [Tz_matrix(0, basis, alpha), Tz_matrix(1, basis, alpha)]
            tzbar_a = Tzbar_matrix(a_idx % 2, basis, alpha)
            b_S = bridge_holo(D, S, basis, alpha, Tz)
            t_norm = op_norm(tzbar_a @ b_S - b_S @ tzbar_a)
            diff = abs(t_norm - sub_norm)
            if diff < best_diff:
                best_diff = diff
                best_alpha = alpha
                best_toep = t_norm
        per_case_optimal[(a_idx, frozenset(S))] = best_alpha
        case_str = f"a={a_idx}, S={set(S)}"
        print(f"  {case_str:>20}  {sub_norm:>10.4f}  {best_alpha:>14.4f}  {best_toep:>22.4f}  {best_diff:>10.4f}")
    print()

    # Aggregate optimal alpha: alpha that minimises sum of per-case residuals
    print("  Aggregate optimal alpha (single alpha across all cases)")
    print()
    best_aggregate_alpha = None
    best_aggregate_score = float('inf')
    for alpha in alphas_fine:
        Tz = [Tz_matrix(0, basis, alpha), Tz_matrix(1, basis, alpha)]
        score = 0.0
        for a_idx, S in cases:
            sub_norm = sub_norms[(a_idx, frozenset(S))]
            tzbar_a = Tzbar_matrix(a_idx % 2, basis, alpha)
            b_S = bridge_holo(D, S, basis, alpha, Tz)
            t_norm = op_norm(tzbar_a @ b_S - b_S @ tzbar_a)
            score += abs(t_norm - sub_norm)
        if score < best_aggregate_score:
            best_aggregate_score = score
            best_aggregate_alpha = alpha

    print(f"  Aggregate optimal alpha = {best_aggregate_alpha:.4f}")
    print(f"  Aggregate residual = {best_aggregate_score:.4f}  (sum of |Toep - sub|, lower better)")
    print()
    print("  Per-case Toeplitz commutator at the aggregate optimal alpha")
    print(f"  {'case':>20}  {'||sub||':>10}  {'Toep at agg alpha':>20}  {'gap':>10}")
    Tz = [Tz_matrix(0, basis, best_aggregate_alpha), Tz_matrix(1, basis, best_aggregate_alpha)]
    for a_idx, S in cases:
        sub_norm = sub_norms[(a_idx, frozenset(S))]
        tzbar_a = Tzbar_matrix(a_idx % 2, basis, best_aggregate_alpha)
        b_S = bridge_holo(D, S, basis, best_aggregate_alpha, Tz)
        t_norm = op_norm(tzbar_a @ b_S - b_S @ tzbar_a)
        case_str = f"a={a_idx}, S={set(S)}"
        gap = abs(t_norm - sub_norm)
        print(f"  {case_str:>20}  {sub_norm:>10.4f}  {t_norm:>20.4f}  {gap:>10.4f}")
    print()

    # Consistency check: per-case optimal alphas vs aggregate
    print("  Consistency of per-case optimal alphas")
    print()
    alphas_chosen = list(set(round(a, 2) for a in per_case_optimal.values()))
    print(f"  Distinct per-case optimal alpha values (rounded): {sorted(alphas_chosen)}")
    spread = max(per_case_optimal.values()) - min(per_case_optimal.values())
    print(f"  Spread of per-case alphas:  {spread:.4f}")
    print()

    print("=" * 90)
    print("  Verdict")
    print("=" * 90)
    print()
    if spread < 0.5:
        print("  CONSISTENT.  Per-case optimal alphas are within a narrow band; a single")
        print("  alpha = alpha(D) closely matches all test cases.  Closure is reachable.")
    elif spread < 5:
        print("  PARTIAL.  Per-case optimal alphas have moderate spread; a single alpha")
        print("  achieves approximate closure but not exact.  An alpha-correction term")
        print("  in the Clifford analog may absorb the residual.")
    else:
        print("  INCONSISTENT.  Per-case optimal alphas span a wide range; no single")
        print("  alpha simultaneously matches all cases.  The simple T_{bar z_a} Clifford")
        print("  analog is structurally insufficient.  An alpha-twisted Berezin")
        print("  derivative or spinor extension is required to close (A.ii).")
    print()
    print("  STRUCTURAL OBSERVATION.")
    print("  The fundamental issue is that the substrate commutator takes only TWO values")
    print("  (0 or 2) while the Toeplitz commutator under the simplest bridge is a")
    print("  CONTINUOUS function of alpha.  A single scalar alpha cannot enforce the")
    print("  bimodal target distribution.  A multi-parameter Clifford construction")
    print("  (alpha plus internal degrees of freedom in the bridge spinor structure)")
    print("  is needed for closure.")


if __name__ == "__main__":
    main()