#!/usr/bin/env python3
"""Compare pure spatial spinor pencils, their stabilizers and charge characters."""
from __future__ import annotations

import argparse
from collections import defaultdict
from datetime import datetime, timezone
from hashlib import sha256
from itertools import combinations
import importlib.util
import json
import math
from pathlib import Path

import numpy as np
import sympy as sp


HERE = Path(__file__).resolve().parent
spec = importlib.util.spec_from_file_location('spinor_source', HERE/'spinor-hodge-source-checks.py')
source = importlib.util.module_from_spec(spec)
spec.loader.exec_module(source)
CHECKS = 0
CASES = []


def check(value, message):
    global CHECKS
    if not value:
        raise AssertionError(message)
    CHECKS += 1


def close(a, b, message):
    check(np.allclose(a, b, atol=3e-10, rtol=3e-10), message)


def basis(m, occupied=()):
    value = np.zeros(2**m, dtype=complex)
    value[sum(1 << (m-1-j) for j in occupied)] = 1
    return value


def annihilator_dimension(gamma, psi):
    return int(len(gamma)-np.linalg.matrix_rank(np.column_stack([g@psi for g in gamma]), tol=1e-9))


def complex_structure(gamma, psi):
    psi = psi/np.linalg.norm(psi)
    j = np.zeros((len(gamma), len(gamma)))
    for a, b in combinations(range(len(gamma)), 2):
        entry = -1j*np.vdot(psi, gamma[a]@gamma[b]@psi)
        close(entry.imag, 0, 'Hermitian spinor bilinear defines a real skew matrix')
        j[a, b] = entry.real
        j[b, a] = -entry.real
    return j


def realify(matrix):
    return np.block([[matrix.real, -matrix.imag], [matrix.imag, matrix.real]])


def real_vector(vector):
    return np.concatenate((vector.real, vector.imag))


def purity_defect(gamma, psi):
    norm = np.vdot(psi, psi).real
    entries = [-1j*np.vdot(psi, gamma[a]@gamma[b]@psi)
               for a, b in combinations(range(len(gamma)), 2)]
    return float(len(gamma)*norm**2-2*sum(abs(value)**2 for value in entries))


def defect_hessian(gamma, keep, psi):
    rank = len(keep)
    d = len(gamma)
    norm = np.vdot(psi, psi).real
    grad_norm = 2*real_vector(psi)
    matrices = [(-1j*gamma[a]@gamma[b])[np.ix_(keep, keep)]
                for a, b in combinations(range(d), 2)]
    moments = np.array([np.vdot(psi, b@psi).real for b in matrices])
    gradients = np.array([2*real_vector(b@psi) for b in matrices])
    return (2*d*np.outer(grad_norm, grad_norm)+4*d*norm*np.eye(2*rank)
            -4*gradients.T@gradients
            -8*sum(j*realify(b) for j, b in zip(moments, matrices)))


def purity_action_checks(m, gamma, first, second, rng):
    keep = np.array([j for j in range(2**m) if j.bit_count()%2 == 0])
    rank = len(keep)
    u, v = first[keep], second[keep]
    size = 2*rank
    full_hessian = np.zeros((2*size, 2*size))
    # Positive quartic defects enforce purity without an independently fitted gap.
    full_hessian[:size,:size] += defect_hessian(gamma, keep, u)
    full_hessian[size:,size:] += defect_hessian(gamma, keep, v)
    for phase in (1, 1j, -1, -1j):
        mapping = np.concatenate((np.eye(size), realify(phase*np.eye(rank))), axis=1)
        full_hessian += mapping.T@defect_hessian(gamma, keep, u+phase*v)@mapping/4
        close(purity_defect(gamma, first+phase*second), 0, 'phase-averaged purity action vanishes on the pencil')
    gu = np.concatenate((2*real_vector(u), np.zeros(size)))
    gv = np.concatenate((np.zeros(size), 2*real_vector(v)))
    full_hessian += 2*(np.outer(gu,gu)+np.outer(gv,gv))
    # At u^dagger v=0, the orthogonality penalty has only its linear-square Hessian.
    cz = []
    for index in range(2*size):
        direction = np.eye(2*size)[index]
        du = direction[:rank]+1j*direction[rank:size]
        dv = direction[size:size+rank]+1j*direction[size+rank:]
        cz.append(np.vdot(du,v)+np.vdot(u,dv))
    cz = np.array(cz)
    full_hessian += 2*(np.outer(cz.real,cz.real)+np.outer(cz.imag,cz.imag))
    close(full_hessian, full_hessian.T, 'complete positive pencil potential has a symmetric Hessian')
    exact = sp.Matrix(np.rint(full_hessian).astype(int))
    close(full_hessian, np.array(exact,dtype=float), 'quarter-phase average gives an exact integer Hessian')
    eigenvalues = np.linalg.eigvalsh(full_hessian)
    check(eigenvalues[0] > -1e-8, 'pencil potential has no negative local directions')
    nullity = int(sum(abs(eigenvalues)<1e-8))
    expected_nullity = 4 if m == 2 else m*m+3*m-6
    check(nullity == expected_nullity, 'pencil vacuum has precisely the homogeneous-orbit flat directions')
    check(2*size-exact.rank() == nullity, 'exact rational Hessian rank matches its numerical spectrum')
    positive_spectrum = defaultdict(int)
    for value in eigenvalues[eigenvalues>1e-8]:
        integer = round(float(value))
        close(value, integer, 'positive purity-action Hessian values are exact integers at this preparation')
        positive_spectrum[str(integer)] += 1
    # Independently differentiate the entire quartic along random real directions.
    def potential(uu,vv):
        uf, vf = np.zeros(2**m, complex), np.zeros(2**m, complex)
        uf[keep], vf[keep] = uu,vv
        return (purity_defect(gamma,uf)+purity_defect(gamma,vf)
                +sum(purity_defect(gamma,uf+p*vf) for p in (1,1j,-1,-1j))/4
                +(np.vdot(uu,uu).real-1)**2+(np.vdot(vv,vv).real-1)**2
                +abs(np.vdot(uu,vv))**2)
    for _ in range(4):
        direction = rng.normal(size=2*size)
        direction /= np.linalg.norm(direction)
        du = direction[:rank]+1j*direction[rank:size]
        dv = direction[size:size+rank]+1j*direction[size+rank:]
        step = 1e-3
        second_difference = (potential(u+step*du,v+step*dv)+potential(u-step*du,v-step*dv)
                             -2*potential(u,v))/step**2
        fourth_difference = (potential(u+2*step*du,v+2*step*dv)+potential(u-2*step*du,v-2*step*dv)
                             -2*potential(u,v))/(2*step)**2
        extrapolated = (4*second_difference-fourth_difference)/3
        check(abs(extrapolated-direction@full_hessian@direction)<3e-6,
              'independent quartic finite difference reproduces the exact action Hessian')
    for _ in range(6):
        test = rng.normal(size=2**m)+1j*rng.normal(size=2**m)
        test[np.setdiff1d(np.arange(2**m),keep)] = 0
        test /= np.linalg.norm(test)
        check(purity_defect(gamma,test)>-1e-10, 'unrestricted chiral comparison has nonnegative purity defect')
    CASES.append({'real_depth': 2*m, 'geometry': 'positive local mode-polarization potential',
                  'parameters': {'unit_norm_control_only': True, 'same_common_coefficient': 1},
                  'measurements': {'real_variables': 4*rank, 'hessian_nullity': nullity,
                                   'positive_hessian_spectrum': dict(positive_spectrum),
                                   'particle_mass_identified_with_hessian': False,
                                   'local_color_gauge_dynamics_derived': False}})


def real_constraints(columns):
    value = np.column_stack(columns)
    return sp.Matrix(np.rint(np.vstack((value.real, value.imag))*2).astype(int))


def stabilizer(gamma, first, second, condition):
    pairs = list(combinations(range(len(gamma)), 2))
    spin_generators = [gamma[a]@gamma[b]/2 for a, b in pairs]
    columns = [np.concatenate((g@first, g@second)) for g in spin_generators]
    size = len(first)
    if condition in ('fixed-projective', 'rays'):
        columns.append(np.concatenate((np.zeros(size), -1j*second)))
    if condition == 'rays':
        columns.append(np.concatenate((-1j*first, np.zeros(size))))
    nullspace = real_constraints(columns).nullspace()
    ordinary = []
    for vector in nullspace:
        coefficients = np.array(vector[:len(pairs)], dtype=float).reshape(-1)
        check(np.linalg.norm(coefficients) > 0, 'phase variables add no spurious stabilizer dimension')
        x = sum(c*g for c, g in zip(coefficients, spin_generators))
        r = np.zeros((len(gamma), len(gamma)))
        for c, (a, b) in zip(coefficients, pairs):
            r[a, b] = c
            r[b, a] = -c
        for i in range(len(gamma)):
            close(x@gamma[i]-gamma[i]@x,
                  sum(r[j, i]*gamma[j] for j in range(len(gamma))),
                  'spin stabilizer uses the corresponding ordinary spatial rotation')
        ordinary.append(r)
    return len(nullspace), ordinary


def pencil_checks(m, rng):
    gamma = [source.dense(g) for g in source.gamma_monomials(m)]
    first, second = basis(m), basis(m, (0, 1))
    d, k = 2*m, m-2
    for psi in (first, second, first+second, first+2j*second,
                (1+.3j)*first+(.4-.7j)*second):
        check(annihilator_dimension(gamma, psi) == m, 'every sampled pure-pencil member has maximal annihilator')
        j = complex_structure(gamma, psi)
        close(j@j, -np.eye(d), 'pencil member gives an orthogonal spatial complex structure')
    j1, j2 = complex_structure(gamma, first), complex_structure(gamma, second)
    close(j1@j2, j2@j1, 'orthogonal pure-pencil endpoints give commuting complex structures')
    product = j1@j2
    base = (np.eye(d)+product)/2
    normal = (np.eye(d)-product)/2
    close(base@base, base, 'four-plane is an orthogonal projector')
    close(normal@normal, normal, 'complement is an orthogonal projector')
    check(np.linalg.matrix_rank(base, tol=1e-9) == 4, 'pencil selects four ordinary spatial directions')
    check(np.linalg.matrix_rank(normal, tol=1e-9) == d-4, 'pencil complement retains all other spatial directions')
    expected = {'fixed-projective': 3 if k == 0 else k*k+3,
                'fixed-fixed': 3 if k == 0 else k*k+2,
                'rays': 4 if k == 0 else k*k+4}
    dimensions = {}
    for condition in expected:
        dimension, rotations = stabilizer(gamma, first, second, condition)
        check(dimension == expected[condition], 'exact real stabilizer dimension agrees with block unitary group')
        for r in rotations:
            close(r@j1, j1@r, 'stabilizer preserves first spatial complex structure')
            close(r@j2, j2@r, 'stabilizer preserves second spatial complex structure')
            close(r@base, base@r, 'stabilizer preserves the derived four-plane')
        dimensions[condition] = dimension
    # Covariantly rotate the field pair; no preferred coordinate split is inserted.
    a, b, angle = 0, d-1, .371
    spin = math.cos(angle/2)*np.eye(2**m)+math.sin(angle/2)*gamma[a]@gamma[b]
    rotation = np.eye(d)
    rotation[a, a] = rotation[b, b] = math.cos(angle)
    rotation[a, b], rotation[b, a] = math.sin(angle), -math.sin(angle)
    rj1 = complex_structure(gamma, spin@first)
    rj2 = complex_structure(gamma, spin@second)
    close(rj1, rotation@j1@rotation.T, 'first complex structure transforms under physical rotation')
    close(rj2, rotation@j2@rotation.T, 'second complex structure transforms under physical rotation')
    close((np.eye(d)+rj1@rj2)/2, rotation@base@rotation.T, 'derived four-plane is spatially covariant')
    controls = []
    for degree in range(4, m+1, 2):
        alternate = basis(m, range(degree))
        check(annihilator_dimension(gamma, alternate) == m, 'higher even-grade endpoint is itself pure')
        value = annihilator_dimension(gamma, first+alternate)
        check(value == m-degree, 'higher-grade sum fails the pure-pencil premise')
        controls.append({'second_grade': degree, 'annihilator_dimension_of_sum': value})
    if m >= 4:
        psi = basis(m, (0, 1))+basis(m, (2, 3))
        check(annihilator_dimension(gamma, psi) < m, 'nondecomposable grade-two state is not pure')
        j = complex_structure(gamma, psi)
        close(j[:8, :8], 0, 'nonpure polarization loses a full spatial rotation structure')
    purity_action_checks(m, gamma, first, second, rng)
    CASES.append({'real_depth': d, 'geometry': 'orthogonal pure pencil in the even spatial carrier',
                  'parameters': {'m': m, 'chiral_rank': 2**(m-1)},
                  'measurements': {'base_rank': 4, 'normal_rank': d-4,
                                   'stabilizer_dimensions': dimensions, 'higher_grade_controls': controls,
                                   'normal_universal_frame_rule_passes': k == 3}})


def charge_checks():
    k_symbol = sp.symbols('k', positive=True, integer=True)
    yq, yu, yd = sp.Rational(1,2)-1/k_symbol, -1+1/k_symbol, 1/k_symbol
    cubic_minimal = sp.factor(2*k_symbol*yq**3+k_symbol*yu**3+k_symbol*yd**3
                             +2*(-sp.Rational(1,2))**3+1)
    mixed_minimal = sp.factor((k_symbol*yq-sp.Rational(1,2))/2)
    check(sp.simplify(cubic_minimal+3*(k_symbol-3)/4) == 0, 'minimal geometric family cubic trace selects three')
    check(sp.simplify(mixed_minimal-(k_symbol-3)/4) == 0, 'minimal geometric family weak mixed trace selects three')
    for k in range(1, 11):
        m = k+2
        rows = []
        tr_y = tr_y3 = mixed = sp.S.Zero
        doublets = 0
        for p in range(k+1):
            for q in range(3):
                if (p+q) % 2:
                    continue
                y = -sp.Rational(p, k)+sp.Rational(q, 2)
                color_rank = math.comb(k, p)
                weak_rank = math.comb(2, q)
                multiplicity = color_rank*weak_rank
                charges = [y-sp.Rational(1, 2), y+sp.Rational(1, 2)] if q == 1 else [y]
                rows.append({'color_grade': p, 'weak_grade': q,
                             'color_rank': color_rank, 'weak_rank': weak_rank,
                             'hypercharge': str(y), 'electric_weights': list(map(str, charges))})
                tr_y += multiplicity*y
                tr_y3 += multiplicity*y**3
                if q == 1:
                    doublets += color_rank
                    mixed += color_rank*y/2
        check(sum(r['color_rank']*r['weak_rank'] for r in rows) == 2**(k+1),
              'branching exhausts every even exterior spinor component')
        check(tr_y == 0, 'full chiral exterior carrier has zero hypercharge trace')
        check(tr_y3 == (sp.Rational(3, 4) if k == 1 else 0), 'full carrier cubic trace includes its small-rank exception')
        check(mixed == (-sp.Rational(1, 4) if k == 1 else 0),
              'full weak mixed trace includes its small-rank exception')
        check(doublets == 2**(k-1), 'full weak doublet count follows the odd-binomial identity')
        geometry = [1-sp.Rational(1, k), -sp.Rational(1, k)]
        minimal = [sp.Rational(k+1, 2*k), sp.Rational(1-k, 2*k)]
        check((geometry == minimal) == (k == 3), 'geometric fundamental doublet matches minimal family only at three colors')
        if k == 3:
            check(len(rows) == 6, 'six ordinary one-family representation blocks at parent ten')
            expected = {(0,0):'0', (0,2):'1', (1,1):'1/6',
                        (2,0):'-2/3', (2,2):'1/3', (3,1):'-1/2'}
            check({(r['color_grade'],r['weak_grade']):r['hypercharge'] for r in rows} == expected,
                  'canonical determinant-trivial charge table at parent ten')
        minimal_family = []
        extra_rank = None
        if k >= 3 and k % 2:
            minimal_family = [r for r in rows if (r['color_grade'],r['weak_grade'])
                              in ((0,0),(0,2),(1,1),(k-1,0),(k-1,2),(k,1))]
            family_rank = sum(r['color_rank']*r['weak_rank'] for r in minimal_family)
            check(family_rank == 4*(k+1), 'minimal selected subfamily has its full fundamental and singlet count')
            extra_rank = 2**(k+1)-family_rank
            check((extra_rank == 0) == (k == 3), 'no exterior exotics in this subfamily only at parent ten')
            family_y3 = sum(r['color_rank']*r['weak_rank']*sp.Rational(r['hypercharge'])**3
                            for r in minimal_family)
            family_mixed = sum(r['color_rank']*sp.Rational(r['hypercharge'])/2
                              for r in minimal_family if r['weak_grade'] == 1)
            check(family_y3 == cubic_minimal.subs(k_symbol,k), 'selected subfamily cubic trace matches exact general formula')
            check(family_mixed == mixed_minimal.subs(k_symbol,k), 'selected subfamily weak mixed trace matches exact general formula')
        has_colorless_doublet = any(r['color_rank'] == 1 and r['weak_rank'] == 2 for r in rows)
        check(has_colorless_doublet == bool(k%2), 'positive chirality has a colorless weak doublet precisely at odd normal rank')
        CASES.append({'real_depth': 2*m, 'geometry': 'complete even exterior carrier for a derived two-plus-k complex split',
                      'parameters': {'complex_normal_rank': k, 'normal_real_depth': 2*k},
                      'measurements': {'branching': rows, 'trace_y': str(tr_y),
                                       'trace_y_cubed': str(tr_y3), 'su2_squared_y_trace': str(mixed),
                                       'number_of_weak_doublets': doublets,
                                       'minimal_family_match': k == 3,
                                       'geometric_fundamental_doublet': list(map(str, geometry)),
                                       'minimal_family_doublet': list(map(str, minimal)),
                                       'minimal_subfamily_cubic_trace': str(cubic_minimal.subs(k_symbol,k)) if minimal_family else None,
                                       'minimal_subfamily_weak_trace': str(mixed_minimal.subs(k_symbol,k)) if minimal_family else None,
                                       'additional_full_carrier_components': extra_rank,
                                       'colorless_weak_doublet_present': has_colorless_doublet,
                                       'anomaly_traces_exclude_higher_complete_carriers': False}})


def flag_checks():
    projectors = {d: np.diag([1.]*d+[0.]*(10-d)) for d in (2,3,4,5,6,10)}
    for a, b in combinations(projectors, 2):
        close(projectors[a]@projectors[b], projectors[a], 'all physical sectors coexist in one nested flat flag')
    base = projectors[4]
    normal = np.eye(10)-base
    close(base@normal, 0, 'pencil normal six-plane is orthogonal to physical four-sector')
    check(not np.allclose(normal, projectors[6]), 'normal six-plane differs from the nested physical six-sector')
    close(projectors[6]-base, np.diag([0.]*4+[1.]*2+[0.]*4), 'physical six-sector needs an additional normal complex line')
    # SU3 acting on the normal plane must reduce to preserve that line.
    generators = []
    for a, b in combinations(range(3), 2):
        x = np.zeros((3,3), complex); x[a,b] = 1; x[b,a] = -1; generators.append(x)
        x = np.zeros((3,3), complex); x[a,b] = 1j; x[b,a] = 1j; generators.append(x)
    generators.extend((1j*np.diag([1.,-1.,0.]), 1j*np.diag([0.,1.,-1.])))
    line = np.diag([1.,0.,0.])
    constraints = real_constraints([(g@line-line@g).reshape(-1) for g in generators])
    check(len(constraints.nullspace()) == 4, 'preserving a normal complex line reduces ambient SU3 to U2')
    real_axis = np.zeros((6,6)); real_axis[0,0] = 1
    realified = [np.array([[g[a//2,b//2].real if a%2 == b%2 else
                           (-g[a//2,b//2].imag if a%2 == 0 else g[a//2,b//2].imag)
                           for b in range(6)] for a in range(6)]) for g in generators]
    constraint = real_constraints([(g@real_axis-real_axis@g).reshape(-1) for g in realified])
    check(len(constraint.nullspace()) == 3, 'preserving the actual normal real-line sector reduces ambient SU3 further to SU2')
    CASES.append({'geometry': 'one compatible physical nested flag',
                  'parameters': {'depths': [2,3,4,5,6,10]},
                  'measurements': {'pencil_normal_plane_is_physical_Q6': False,
                                   'extra_normal_line_needed': True,
                                   'ambient_su3_stabilizer_of_Q6_dimension': 4,
                                   'ambient_su3_stabilizer_of_Q5_dimension': 3}})


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    rng = np.random.default_rng(20261009)
    check(2**(1-1) == 1, 'depth two has no independent same-chirality spinor pair')
    CASES.append({'real_depth': 2, 'measurements': {'independent_chiral_pair_exists': False}})
    for m in range(2, 7):
        pencil_checks(m, rng)
    charge_checks()
    flag_checks()
    report = {'title': 'Pure spatial spinor pencils and their four-plane, stabilizer and charge character',
              'created_utc': datetime.now(timezone.utc).isoformat(),
              'checks_passed': CHECKS, 'cases': CASES,
              'scope': 'Bounded spatial Clifford algebra and exact stabilizers; classical field amplitudes, no particle simulations.',
              'assumptions': ['Two orthogonal same-chirality spinors span an everywhere-pure local pencil.',
                              'One spinor is fixed exactly while the other is fixed projectively.',
                              'The earlier universal nonzero alternating normal-vector channel is imposed on the complement.'],
              'limitations': ['Pure pencils exist at every even parent depth at least four; their presence alone does not select ten.',
                              'The normal channel criterion selects ten only in combination with the pencil premise.',
                              'The normal six-plane is not the nested physical depth-six sector.',
                              'The pair itself shares a normal vacuum and does not generate the additional normal frame.',
                              'Charge characters are representation weights, not derived occupied-particle electric charges.',
                              'Full higher-rank carriers retain extra representations and cancel the displayed anomaly traces.',
                              'Preserving fixed lower-sector flags reduces ambient spatial SU3; an independent color action requires its own dynamics.'],
              'sources': {'experiments/pure-pencil-geometry-checks.py': sha256(Path(__file__).read_bytes()).hexdigest(),
                          'experiments/spinor-hodge-source-checks.py': sha256((HERE/'spinor-hodge-source-checks.py').read_bytes()).hexdigest()}}
    args.output.write_text(json.dumps(report, indent=2)+'\n')
    print(json.dumps({'checks_passed': CHECKS, 'cases': len(CASES), 'output': str(args.output)}))


if __name__ == '__main__':
    main()