"""Exact stage-one certificate for the FMS family-II, n=5 lattice.

Run with Sage's Python. No numerical linear algebra or paid inference.
The specialized Vinberg enumeration handles norm-two roots, not all possible
integral reflection norms. Its full-root scope is proved in README.md.
"""
from sage.all import (QQ, ZZ, matrix, vector, diagonal_matrix, block_diagonal_matrix,
                      Polyhedron, Graph, PolynomialRing, identity_matrix)
from sage.env import SAGE_VERSION
from itertools import combinations, product
from pathlib import Path
import json
import hashlib

ROOT = Path(__file__).resolve().parent
Q = diagonal_matrix(QQ, [-16, 1, 1, 1, 1])
I = identity_matrix(QQ, 5)


def rows(m):
    return [[str(v) for v in row] for row in m.rows()]


def vec(v):
    return [str(z) for z in v]


def is_lattice_vector(v):
    return all(z in ZZ for z in v) and sum(v[1:]) % 2 == 0


def reflection(r):
    assert r*Q*r == 2 and is_lattice_vector(r)
    return I-r.column()*(Q*r).row()


def companion(p):
    n = p.degree()
    m = matrix(QQ, n)
    for j in range(n-1):
        m[j+1,j] = 1
    for j in range(n):
        m[j,n-1] = -p[j]
    return m


def discover():
    # Negatives of a standard D4 simple system; inequalities use <x,r> <= 0.
    finite = [vector(QQ, r) for r in [(0,-1,1,0,0),(0,0,-1,1,0),
                                    (0,0,0,-1,1),(0,0,0,-1,-1)]]
    zero_roots = [vector(QQ,[0]+list(v)) for v in product(range(-1,2),repeat=4)
                  if sum(z*z for z in v)==2]
    assert len(zero_roots)==24 and all(is_lattice_vector(r) for r in zero_roots)
    # Enumerate ALL roots at t=1, then select the unique D4-dominant orbit members.
    level_one = [vector(QQ,[1]+list(v)) for v in product(range(-4,5),repeat=4)
                 if sum(z*z for z in v)==18 and sum(v)%2==0]
    dominant = sorted([r for r in level_one if all(r*Q*s<=0 for s in finite)])
    assert len(dominant)==4
    assert all(r*Q*s<=0 for r,s in combinations(dominant,2))
    roots = finite+dominant
    P = Polyhedron(ieqs=[[16*r[0]]+[-r[i] for i in range(1,5)] for r in roots],base_ring=QQ)
    assert P.dim()==4 and P.is_compact()
    assert len(P.inequalities_list())==8
    vertices = sorted([vector(QQ,v) for v in P.vertices_list()])
    assert len(vertices)==16 and all(v*v<=16 for v in vertices)
    assert sum(v*v==16 for v in vertices)==6
    # Independently enumerate all rational intersections of four candidate walls.
    independent = set()
    for inds in combinations(range(8),4):
        a=matrix(QQ,[list(roots[i][1:]) for i in inds])
        if a.det()==0:
            continue
        v=a.solve_right(vector(QQ,[16*roots[i][0] for i in inds]))
        if all(sum(v[j]*r[j+1] for j in range(4))<=16*r[0] for r in roots):
            independent.add(tuple(v))
    assert independent=={tuple(v) for v in vertices}
    gram=matrix(QQ,[[r*Q*s for s in roots] for r in roots])
    assert all(gram[i,j] in [0,-1,-3] for i,j in combinations(range(8),2))
    reflections=[reflection(r) for r in roots]
    for r,s in zip(roots,reflections):
        assert s*s==I and s.transpose()*Q*s==Q and s*r==-r
    for i,j in combinations(range(8),2):
        if gram[i,j]>=-1:
            order=2 if gram[i,j]==0 else 3
            assert (reflections[i]*reflections[j])**order==I
    return zero_roots, level_one, roots, P, vertices, gram, reflections


def lattice_and_generators():
    # Same isometry as the preceding audit; no import or mutation of that script.
    K=matrix(QQ,5,lambda i,j:[-2,-3,-4,-5,-6][abs(i-j)])
    M=matrix(QQ,[[-5,-1,-1,1,1],[8,1,2,-2,-1],[0,0,0,1,-1],
                 [0,0,0,0,1],[3,1,0,-1,-1]])
    D=matrix(QQ,[[0,1,0,0],[1,-1,0,0],[-1,0,1,1],[0,0,-1,1]])
    X=block_diagonal_matrix(matrix(QQ,[[1]]),D)
    G=X.transpose()*Q*X
    assert abs(M.det())==1 and abs(D.det())==2 and M.transpose()*K*M==G
    assert G==block_diagonal_matrix(matrix(QQ,[[-16]]),
        matrix(QQ,[[2,-1,-1,-1],[-1,2,0,0],[-1,0,2,0],[-1,0,0,2]]))
    x=PolynomialRing(QQ,'x').gen()
    A=companion((x+1)*(x**4+1)); B=companion((x-1)**2*(x**3-1)); C=A.inverse()*B
    u=(C-I).column(4)
    O=matrix(QQ,[( (-A)**j*u) for j in range(5)]).transpose()
    J=X*M.inverse()*O.inverse()
    generators={name:J*g*J.inverse() for name,g in [('A',A),('B',B),('C',C)]}
    for g in generators.values():
        assert g.transpose()*Q*g==Q
        assert all(z in ZZ for z in (X.inverse()*g*X).list())
        assert abs(g.det())==1
    return X,G,M,O,J,generators


def chamber_symmetries(roots,gram,X):
    graph=Graph(8)
    graph.add_edges([(i,j,int(gram[i,j])) for i,j in combinations(range(8),2)])
    aut=graph.automorphism_group(edge_labels=True)
    assert aut.order()==48
    inds=next(inds for inds in combinations(range(8),5)
              if matrix(QQ,[roots[i] for i in inds]).det()!=0)
    basis=matrix(QQ,[roots[i] for i in inds]).transpose()
    symmetries=[]
    for perm in aut:
        images=matrix(QQ,[roots[perm(i)] for i in inds]).transpose()
        g=images*basis.inverse()
        assert all(g*roots[i]==roots[perm(i)] for i in range(8))
        assert g.transpose()*Q*g==Q and g[0,0]>0
        assert all(z in ZZ for z in (X.inverse()*g*X).list()) and abs(g.det())==1
        symmetries.append({'permutation':[int(perm(i)) for i in range(8)],'matrix':rows(g)})
    symmetries.sort(key=lambda a:a['permutation'])
    return symmetries


def decompose(g, roots, reflections):
    # g=s_word[0] ... s_word[-1] * chamber_symmetry, up to central -I.
    sign=1 if g[0,0]>0 else -1
    h=sign*g
    point=vector(QQ,[10,4,3,2,1])
    assert all(point*Q*r<0 for r in roots)
    word=[]
    for count in range(10000):
        image=h*point
        viol=next((i for i,r in enumerate(roots) if image*Q*r>0),None)
        if viol is None:
            break
        h=reflections[viol]*h
        word.append(viol)
    else:
        raise RuntimeError('bounded chamber reduction exhausted; no conclusion')
    permutation=[next((j for j,s in enumerate(roots) if h*r==s),None) for r in roots]
    assert None not in permutation and sorted(permutation)==list(range(8))
    rebuilt=I
    for i in word:
        rebuilt=rebuilt*reflections[i]
    assert sign*rebuilt*h==g
    return {'central_sign':sign,'reflection_word':word,
            'chamber_symmetry_permutation':permutation,'chamber_symmetry_matrix':rows(h),
            'matrix':rows(g)}


def finite_presentation(symmetries, gram, decomposed):
    # Cube coordinates exhibit Aut(chamber)=C2^3 semidirect S3 = W(B3).
    bits=[(0,0,0),(1,0,0),(1,0,1),(1,1,0),
          (0,1,0),(0,0,1),(0,1,1),(1,1,1)]
    transforms=[lambda v:(1-v[0],v[1],v[2]),
                lambda v:(v[1],v[0],v[2]),lambda v:(v[0],v[2],v[1])]
    by_perm={tuple(s['permutation']):s for s in symmetries}
    gens=[tuple(bits.index(f(v)) for v in bits) for f in transforms]
    assert all(g in by_perm for g in gens)
    identity=tuple(range(8)); words={identity:[]}; queue=[identity]
    for p in queue:
        for k,g in enumerate(gens):
            new=tuple(p[g[i]] for i in range(8))
            if new not in words:
                words[new]=words[p]+[k]; queue.append(new)
    assert len(words)==48 and set(words)==set(by_perm)
    gm=[matrix(QQ,by_perm[g]['matrix']) for g in gens]
    assert all(h*h==I for h in gm)
    for i,j,m in [(0,1,4),(1,2,3),(0,2,2)]:
        assert (gm[i]*gm[j])**m==I
    for p,word in words.items():
        rebuilt=I
        for k in word:
            rebuilt=rebuilt*gm[k]
        assert rebuilt==matrix(QQ,by_perm[p]['matrix'])
    for d in decomposed.values():
        d['symmetry_word_abc']=words[tuple(d['chamber_symmetry_permutation'])]
    # This file only declares the verified presentation and target subgroup.
    # It deliberately performs no coset enumeration or index computation.
    lines=['# Generated by chamber.py. Stage one only: no index computation.',
        '# Projective O(L) is identified with the time-cone-preserving group.',
        'FMSFree := FreeGroup("s0","s1","s2","s3","s4","s5","s6","s7","a","b","c");',
        'fms_g := GeneratorsOfGroup(FMSFree);',
        'fms_rel := List(fms_g, g -> g^2);']
    for i,j in combinations(range(8),2):
        if gram[i,j] in [0,-1]:
            m=2 if gram[i,j]==0 else 3
            lines.append(f'Add(fms_rel, (fms_g[{i+1}]*fms_g[{j+1}])^{m});')
    for i,j,m in [(8,9,4),(9,10,3),(8,10,2)]:
        lines.append(f'Add(fms_rel, (fms_g[{i+1}]*fms_g[{j+1}])^{m});')
    for k,p in enumerate(gens):
        for i in range(8):
            lines.append(f'Add(fms_rel, fms_g[{9+k}]*fms_g[{i+1}]*fms_g[{9+k}]^-1*fms_g[{p[i]+1}]^-1);')
    lines+=['FMSAmbient := FMSFree/fms_rel;', 'fms_g := GeneratorsOfGroup(FMSAmbient);']
    for name,d in decomposed.items():
        indexes=[i+1 for i in d['reflection_word']]+[9+i for i in d['symmetry_word_abc']]
        expr='*'.join(f'fms_g[{i}]' for i in indexes) or 'One(FMSAmbient)'
        lines.append(f'FMSProjective{name} := {expr};')
    lines+=['FMSProjectiveH := Subgroup(FMSAmbient,[FMSProjectiveA,FMSProjectiveB]);',
            '# Do not infer an index from this declaration; stage two remains open.']
    (ROOT/'ambient-presentation.g').write_text('\n'.join(lines)+'\n')
    return {'finite_factor':'W(B3) = C2^3 semidirect S3, order 48',
            'symmetry_generator_permutations':[list(g) for g in gens],
            'symmetry_coxeter_orders':[[1,4,2],[4,1,3],[2,3,1]],
            'semidirect_action':'a_i s_j a_i^-1 = s_(permutation_i(j))',
            'gap_file':'ambient-presentation.g', 'generators':11,
            'central_extension':'O(L) = {+I,-I} times O^+(L); GAP uses O^+(L)'}


def main():
    zero,one,roots,P,vertices,gram,reflections=discover()
    X,G,M,O,J,generators=lattice_and_generators()
    assert all(z in ZZ for z in (X.inverse()*matrix(QQ,roots).transpose()).list())
    assert matrix(ZZ,X.inverse()*matrix(QQ,roots).transpose()).elementary_divisors()==[1]*5
    symmetries=chamber_symmetries(roots,gram,X)
    decomposed={name:decompose(g,roots,reflections) for name,g in generators.items()}
    presentation=finite_presentation(symmetries,gram,decomposed)
    # Boundary link certificates: finite vertices have positive-definite active
    # Gram; ideal vertices have rank-three positive-semidefinite active Gram.
    vertex_records=[]
    for v in vertices:
        active=[i for i,r in enumerate(roots) if vector(QQ,[1]+list(v))*Q*r==0]
        assert len(active)==4
        sub=gram.matrix_from_rows_and_columns(active,active)
        ideal=v*v==16
        cp=sub.charpoly()
        assert all(sub.matrix_from_rows_and_columns(ids,ids).det()>=0
                   for size in range(1,5) for ids in combinations(range(4),size))
        assert sub.rank()==(3 if ideal else 4)
        if ideal:
            null=sub.right_kernel().basis()[0]
            assert all(v>0 for v in null) or all(v<0 for v in null)
        vertex_records.append({'point':vec(v),'euclidean_squared_norm':str(v*v),
            'type':'ideal' if ideal else 'finite','active_walls':active,
            'link_gram_characteristic_polynomial':str(cp),'link_rank':int(sub.rank())})
    result={'sage_version':SAGE_VERSION,'arithmetic':'QQ and ZZ; no floating point',
        'scope':'complete norm-two reflection group R2(L), not yet the hypergeometric subgroup index',
        'form_Q':rows(Q),'lattice':'Z times D4; D4={v in Z^4: sum(v) even}',
        'lattice_basis_X':rows(X),'gram_in_integral_basis':rows(G),
        'preceding_orbit_isometry_M':rows(M),'companion_to_orbit_basis_O':rows(O),
        'companion_to_euclidean_coordinates_J':rows(J),
        'all_height_zero_roots':[vec(r) for r in sorted(zero)],
        'all_height_one_roots':[vec(r) for r in sorted(one)],
        'walls':[{'id':i,'root':vec(r),'reflection_matrix':rows(s)}
                 for i,(r,s) in enumerate(zip(roots,reflections))],
        'wall_gram':rows(gram),
        'coxeter_matrix':[[1 if i==j else (2 if gram[i,j]==0 else 3 if gram[i,j]==-1 else 'infinity')
                           for j in range(8)] for i in range(8)],
        'vertices':vertex_records,'polytope_bounded':bool(P.is_compact()),
        'rays':P.rays_list(),'lines':P.lines_list(),
        'chamber_symmetries':symmetries,
        'ambient_presentation':presentation,
        'original_generators':decomposed,
        'not_computed':['hypergeometric subgroup index','arithmeticity versus thinness',
                        'absolute hyperbolic volume']}
    (ROOT/'certificate.json').write_text(json.dumps(result,indent=2)+'\n')
    print(f'Certified candidate: {len(roots)} walls, {len(vertices)} vertices '
          f'({sum(v*v==16 for v in vertices)} ideal), {len(symmetries)} chamber symmetries.')
    print(f'Enumerated {len(zero)} height-zero and {len(one)} height-one roots.')
    for name,d in decomposed.items():
        print(name,'sign',d['central_sign'],'word',d['reflection_word'],
              'symmetry',d['chamber_symmetry_permutation'])
    print('Exact completeness proof and classical presentation input: see README.md.')


if __name__=='__main__':
    main()