#!/usr/bin/env python3
"""Actual two-marked-point periodic Temperley--Lieb link-state matrices.

The basis is a noncrossing matching of N boundary vertices, together with
the two complementary faces containing the ordered punctures a and b.
Faces are encoded by their sets of boundary gaps.  No assumed Jordan form
is used in constructing any matrix.
"""
from __future__ import annotations
import argparse
import functools
import itertools
import json
import math
from pathlib import Path
import time
import numpy as np


@functools.lru_cache(None)
def matchings(n: int):
    """All disk noncrossing perfect matchings, mate arrays, lexicographic."""
    if n == 0:
        return ((),)
    out = []
    for j in range(1, n, 2):
        for left in matchings(j - 1):
            for right in matchings(n - j - 1):
                m = [None] * n
                m[0], m[j] = j, 0
                for i, a in enumerate(left, 1):
                    m[i] = a + 1
                for i, a in enumerate(right, j + 1):
                    m[i] = a + j + 1
                out.append(tuple(m))
    return tuple(sorted(out))


@functools.lru_cache(None)
def faces(m):
    """Face cycles g -> mate(g+1) on boundary gaps g=(g,g+1)."""
    n = len(m)
    if n == 0:
        return ((),)
    unseen = set(range(n))
    out = []
    while unseen:
        g = min(unseen)
        f = []
        while g not in f:
            f.append(g)
            unseen.remove(g)
            g = m[(g + 1) % n]
        out.append(tuple(sorted(f)))
    if not (len(out) == n // 2 + 1):
        raise AssertionError('Check failed: len(out) == n // 2 + 1')
    return tuple(sorted(out))


@functools.lru_cache(None)
def gap_faces(m):
    return {g: f for f in faces(m) for g in f}


@functools.lru_cache(None)
def basis(n):
    return tuple((m, a, b) for m in matchings(n)
                 for a in faces(m) for b in faces(m))


def ei_action(state, i):
    """Return output state and erased-loop puncture mask, or None if no loop."""
    m, a, b = state
    n = len(m)
    j = (i + 1) % n
    old = gap_faces(m)
    if m[i] == j:
        enclosed = old[i]
        if not (enclosed == (i,)):
            raise AssertionError('Check failed: enclosed == (i,)')
        other = old[(i - 1) % n]
        mask = int(a == enclosed) + 2 * int(b == enclosed)
        return (m, other if a == enclosed else a,
                other if b == enclosed else b), mask
    mm = list(m)
    u, v = m[i], m[j]
    mm[i], mm[j], mm[u], mm[v] = j, i, v, u
    mm = tuple(mm)
    new = gap_faces(mm)
    def moved(f):
        candidates = {new[g] for g in f if g != i}
        if not (len(candidates) == 1):
            raise AssertionError((state, i, f, candidates))
        return candidates.pop()
    return (mm, moved(a), moved(b)), None


def rotate_action(state, shift=-1):
    m, a, b = state
    n = len(m)
    mm = [None] * n
    for i, j in enumerate(m):
        mm[(i + shift) % n] = (j + shift) % n
    return (tuple(mm), tuple(sorted((g + shift) % n for g in a)),
            tuple(sorted((g + shift) % n for g in b)))


def add_aux_cup(state):
    """Insert an empty cup at boundary positions N,N+1."""
    m, a, b = state
    n = len(m)
    mm = m + (n + 1, n)
    new = gap_faces(mm)
    if n == 0:
        return (mm, new[1], new[1])
    return (mm, new[a[0]], new[b[0]])


def remove_seam_cap(state):
    """Cap the final and first boundary points, then retain points 1,...,M-2."""
    m, a, b = state
    state, mask = ei_action(state, len(m) - 1)
    m, a, b = state
    size = len(m)
    mm = tuple(m[i] - 1 for i in range(1, size - 1))
    if not mm:
        return (mm, (), ()), mask
    new = gap_faces(mm)
    def moved(f):
        candidates = {new[g - 1] for g in f if 1 <= g <= size - 2}
        if not (len(candidates) == 1):
            raise AssertionError((state, f, candidates))
        return candidates.pop()
    return (mm, moved(a), moved(b)), mask


def generator_transitions(n):
    states = basis(n)
    ix = {s: i for i, s in enumerate(states)}
    e = [[None] * len(states) for _ in range(n)]
    omega = []
    for col, s in enumerate(states):
        for i in range(n):
            out, mask = ei_action(s, i)
            e[i][col] = (ix[out], mask)
        omega.append(ix[rotate_action(s)])
    return e, omega


def depth_one_columns(n, eps=1):
    """An explicit basis for the selected depth-at-most-one submodule.

    First all same-face states; next v_(f,h)+eps*v_(h,f) for every edge of
    the dual tree of every noncrossing matching, oriented by f<h.
    Returns sparse column dictionaries and descriptors of the edge columns.
    """
    states = basis(n)
    ix = {s:i for i,s in enumerate(states)}
    columns, descriptors = [], []
    for matching in matchings(n):
        for f in faces(matching):
            columns.append({ix[(matching,f,f)]:1})
    for matching in matchings(n):
        gf = gap_faces(matching)
        for i,j in enumerate(matching):
            if i<j:
                f,h = sorted((gf[(i-1)%n],gf[i]))
                columns.append({ix[(matching,f,h)]:1,ix[(matching,h,f)]:eps})
                descriptors.append({'matching':matching,'face_f':f,'face_h':h,
                                    'diagonal_rows':[ix[(matching,f,f)],ix[(matching,h,h)]]})
    return columns, descriptors


def matrices(n, q, x, y, w, dtype=complex):
    states = basis(n)
    d = len(states)
    ix = {s: i for i, s in enumerate(states)}
    beta = -q - q**(-1)
    weights = {None: 1, 0: beta, 1: x + x**(-1),
               2: y + y**(-1), 3: w + w**(-1)}
    transitions, rotation = generator_transitions(n)
    es = []
    for transition in transitions:
        mat = np.zeros((d, d), dtype=dtype)
        for col, (row, mask) in enumerate(transition):
            mat[row, col] = weights[mask]
        es.append(mat)
    omega = np.zeros((d, d), dtype=dtype)
    for col, row in enumerate(rotation):
        omega[row, col] = 1
    hoops = []
    for inverse in [False, True]:
        # The N half-powers combine to integral powers since N is even.
        # Each branch below stores its number of e_i smoothings separately.
        mat = np.zeros((d, d), dtype=dtype)
        for col, s in enumerate(states):
            partial = {(add_aux_cup(s), 0): 1}
            for i in range(n - 1, -1, -1):
                nex = {}
                for (t, k), coeff in partial.items():
                    nex[(t, k)] = nex.get((t, k), 0) + coeff
                    u, mask = ei_action(t, i)
                    key = (u, k + 1)
                    nex[key] = nex.get(key, 0) + coeff * weights[mask]
                partial = nex
            for (t, k), coeff in partial.items():
                out, mask = remove_seam_cap(t)
                power = (n // 2 - k) * (-1 if inverse else 1)
                mat[ix[out], col] += coeff * weights[mask] * q**power
        hoops.append(mat)
    return states, es, omega, hoops[0], hoops[1]


def numerical_summary(n, p, x=1.0, y=2.0, m=1, eps=1):
    q = np.exp(1j * np.pi / p)
    w = eps * q**m
    states, es, om, f, fb = matrices(n, q, x, y, w)
    d = len(states)
    ident = np.eye(d)
    beta = -q - 1/q
    f0 = eps * (q**m + q**(-m))
    def norm(a): return float(np.max(np.abs(a)))
    checks = {
        'quadratic': max(norm(e@e-beta*e) for e in es),
        'adjacent': max(norm(es[i]@es[(i+1)%n]@es[i]-es[i])
                        for i in range(n)) if n > 2 else None,
        'distant': max([norm(es[i]@es[j]-es[j]@es[i])
                        for i in range(n) for j in range(n)
                        if (i-j)%n not in (0,1,n-1)] or [0.0]),
        'rotation': max(norm(om@es[i]-es[(i-1)%n]@om) for i in range(n)),
        'full_rotation': norm(np.linalg.matrix_power(om,n)-ident),
        'F_central_e': max(norm(f@e-e@f) for e in es),
        'F_central_Omega': norm(f@om-om@f),
        'Fbar_central_e': max(norm(fb@e-e@fb) for e in es),
        'F_Fbar': norm(f@fb-fb@f),
    }
    depth0 = [i for i, (_,a,b) in enumerate(states) if a == b]
    checks['F_depth0'] = norm((f-(w+1/w)*ident)[:,depth0])
    shift = f-f0*ident
    ranks = [int(np.linalg.matrix_rank(np.linalg.matrix_power(shift,k),tol=1e-8))
             for k in range(1,4)]
    vals = np.linalg.eigvals(f)
    out = {
        'parameters': {'N':n,'p':p,'m':m,'epsilon':eps,'x':x,'y':y},
        'dimension':d,'depth_zero_dimension':len(depth0),
        'residuals':checks,'ranks_F_minus_f0_powers':ranks,
        'eigenvalues_numerical':[[float(v.real),float(v.imag)] for v in vals],
    }
    if n == 2:
        out['basis'] = [{'matching':s[0],'face_a':s[1],'face_b':s[2]} for s in states]
        out['F'] = [[[float(z.real),float(z.imag)] for z in row] for row in f]
    return out


def exact_case(n, p, x=1, y=2, m=1, eps=1, output=None):
    """Compute and verify over Q(q), q=exp(i*pi/p), using SymPy domains.

    Optional dependency: SymPy >=1.13.  python-flint is recommended but not
    required.  Matrix entries are exported as rational coefficient vectors
    in the basis 1,q,...,q^(degree-1), with a shared coefficient dictionary.
    """
    from sympy import QQ, exp, I, pi, Rational, __version__ as sympy_version
    from sympy.polys.matrices import DomainMatrix
    started = time.time()
    field = QQ.algebraic_field(exp(I*pi/p))
    q = field.unit
    f0 = eps*(q**m+q**(-m))
    states, es0, omega0, f0mat, fb0mat = matrices(
        n, q, field.from_sympy(Rational(x)), field.from_sympy(Rational(y)),
        eps*q**m, dtype=object)
    d = len(states)
    def dm(mat):
        return DomainMatrix([[field.convert(z) for z in row] for row in mat],
                            mat.shape, field).to_sparse()
    es, om, f, fb = list(map(dm,es0)), dm(omega0), dm(f0mat), dm(fb0mat)
    identity = DomainMatrix.eye(d,field).to_sparse()
    beta = -q-q**(-1)
    checks = {}
    checks['quadratic_relations'] = all((e*e-beta*e).is_zero_matrix for e in es)
    checks['distant_commutation'] = all(
        (es[i]*es[j]-es[j]*es[i]).is_zero_matrix
        for i in range(n) for j in range(n)
        if (i-j)%n not in (0,1,n-1))
    checks['rotation_relations'] = all(
        (om*es[i]-es[(i-1)%n]*om).is_zero_matrix for i in range(n))
    checks['full_rotation'] = (om**n-identity).is_zero_matrix
    if n > 2:
        checks['adjacent_relations'] = all(
            (es[i]*es[(i+1)%n]*es[i]-es[i]).is_zero_matrix for i in range(n))
        checks['affine_closure_relation'] = (
            (om*es[0])**(n-1)-om**n*(om*es[0])).is_zero_matrix
    else:
        checks['size_two_relation'] = (om**2*es[0]-es[0]).is_zero_matrix
    checks['F_central'] = all((f*g-g*f).is_zero_matrix for g in es+[om])
    checks['Fbar_central'] = all((fb*g-g*fb).is_zero_matrix for g in es+[om])
    checks['hoops_commute'] = (f*fb-fb*f).is_zero_matrix
    depth0 = [i for i,(_,a,b) in enumerate(states) if a == b]
    checks['F_depth_zero_scalar'] = (
        f-(eps*q**m+(eps*q**m)**(-1))*identity).extract(range(d),depth0).is_zero_matrix
    a, b = f-f0*identity, fb-f0*identity
    apowers, bpowers = [identity], [identity]
    ranks, branks, joint_ranks = [], [], []
    stable_power = None
    for k in range(1,5):
        apowers.append(apowers[-1]*a)
        bpowers.append(bpowers[-1]*b)
        ranks.append(apowers[-1].rank())
        branks.append(bpowers[-1].rank())
        joint_ranks.append(apowers[-1].vstack(bpowers[-1]).rank())
        if k > 1 and ranks[-1] == ranks[-2] and stable_power is None:
            stable_power = k-1
    if not (stable_power is not None):
        raise AssertionError('Increase the maximum power.')
    kernel = apowers[stable_power].nullspace(divide_last=True).transpose().to_sparse()
    tdim = kernel.shape[1]
    pivot_rows = kernel.transpose().rref()[1]
    inverse = kernel.extract(pivot_rows,range(tdim)).inv()
    def restrict(mat):
        acted = mat*kernel
        result = inverse*acted.extract(pivot_rows,range(tdim))
        if not ((kernel*result-acted).is_zero_matrix):
            raise AssertionError('Check failed: (kernel*result-acted).is_zero_matrix')
        return result.to_sparse()
    nes, nom, nf, nfb = [restrict(e) for e in es], restrict(om), restrict(f), restrict(fb)
    ident = DomainMatrix.eye(tdim,field).to_sparse()
    nn, mm = nf-f0*ident, nfb-f0*ident
    dd = nn-mm
    nn2, mm2, nm = nn*nn, mm*mm, nn*mm
    internal_ranks = [tdim]+[(nn**j).rank() for j in range(1,5)]
    jordan = {str(j): internal_ranks[j-1]-2*internal_ranks[j]+internal_ranks[j+1]
              for j in range(1,4)
              if internal_ranks[j-1]-2*internal_ranks[j]+internal_ranks[j+1]}
    relations = {
        'N3_zero':(nn2*nn).is_zero_matrix,
        'M3_zero':(mm2*mm).is_zero_matrix,
        'NM_equals_N2':(nm-nn2).is_zero_matrix,
        'M2_equals_N2':(mm2-nn2).is_zero_matrix,
        'ND_zero':(nn*dd).is_zero_matrix,
        'D2_zero':(dd*dd).is_zero_matrix,
        'F_equals_Fbar_on_target':dd.is_zero_matrix,
    }
    monomials = [ident,nn,mm,nn2,nm,mm2,nn2*nn,nn2*mm,nn*mm2,mm2*mm]
    flattened = DomainMatrix([sum(z.to_list(),[]) for z in monomials],
                             (len(monomials),tdim*tdim),field)
    algebra_dimension = flattened.rank()
    dcols, edges = depth_one_columns(n,eps)
    zdim = len(dcols)
    zbase = DomainMatrix.from_dod(
        {r:{c:field.convert(col[r]) for c,col in enumerate(dcols) if r in col}
         for r in range(d)},(d,zdim),field).to_sparse()
    zpiv = zbase.transpose().rref()[1]
    zinv = zbase.extract(zpiv,range(zdim)).inv()
    zstable = True
    for gen in es+[om,f,fb]:
        acted = gen*zbase
        zstable &= (zbase*zinv*acted.extract(zpiv,range(zdim))-acted).is_zero_matrix
    gx = field.from_sympy(Rational(x))
    gy = field.from_sympy(Rational(y))
    g = gx+gx**(-1)+eps*(gy+gy**(-1))
    expected = DomainMatrix.zeros((d,zdim),field).to_sparse().to_dod()
    for c,ed in enumerate(edges,len(depth0)):
        r1,r2=ed['diagonal_rows']
        expected.setdefault(r1,{})[c]=g
        expected.setdefault(r2,{})[c]=eps*g
    expected=DomainMatrix.from_dod(expected,(d,zdim),field)
    depth_one_checks={
        'generator_stable':bool(zstable),
        'tree_incidence_formula_F':(a*zbase-expected).is_zero_matrix,
        'tree_incidence_formula_Fbar':(b*zbase-expected).is_zero_matrix,
        'nilpotent_square_zero':(a*a*zbase).is_zero_matrix,
        'F_equals_Fbar':((f-fb)*zbase).is_zero_matrix,
    }
    field_degree = len(field.mod.to_list())-1
    coeffs, coeff_lookup = [], {}
    def coeff_id(z):
        z = field.convert(z)
        v = list(reversed(z.to_list()))
        v += [field.dom.zero]*(field_degree-len(v))
        v = tuple(str(t) for t in v)
        if v not in coeff_lookup:
            coeff_lookup[v] = len(coeffs)
            coeffs.append(v)
        return coeff_lookup[v]
    def sparse(mat):
        entries = []
        for i,row in mat.to_dod().items():
            for j,z in row.items():
                if z:
                    entries.append([i,j,coeff_id(z)])
        return {'shape':list(mat.shape),'entries':sorted(entries)}
    all_matrices = {f'e_{i}':sparse(e) for i,e in enumerate(es)}
    all_matrices.update({'Omega':sparse(om),'F':sparse(f),'Fbar':sparse(fb)})
    target_matrices = {f'e_{i}':sparse(e) for i,e in enumerate(nes)}
    target_matrices.update({'Omega':sparse(nom),'F':sparse(nf),'Fbar':sparse(nfb),
                            'N':sparse(nn),'M':sparse(mm),'D':sparse(dd),
                            'kernel_basis':sparse(kernel)})
    rank_certificates=[]
    for power in range(1,4):
        operator=apowers[power]
        _,columns=operator.rref()
        _,rows=operator.extract(range(d),columns).transpose().rref()
        determinant=operator.extract(rows,columns).det()
        ker=operator.nullspace(divide_last=True).transpose().to_sparse()
        free=sorted(set(range(d))-set(columns))
        if not (determinant):
            raise AssertionError('Check failed: determinant')
        if not (ker.extract(free,range(len(free)))==DomainMatrix.eye(len(free),field).to_sparse()):
            raise AssertionError('Check failed: ker.extract(free,range(len(free)))==DomainMatrix.eye(len(free),field).to_sparse()')
        rank_certificates.append({'power':power,'rank':len(columns),
                                  'minor_rows':list(rows),'minor_columns':list(columns),
                                  'minor_determinant_coeff_id':coeff_id(determinant),
                                  'right_kernel':sparse(ker),'kernel_identity_rows':free})
    result = {
        'construction':'Disk noncrossing matchings with ordered puncture-face labels; hoops by auxiliary cup, N skein crossings, and seam cap.',
        'arithmetic':'Exact rational arithmetic in a cyclotomic number field; no floating point rank decisions.',
        'software':{'sympy':sympy_version},
        'parameters':{'N':n,'p':p,'x':str(x),'y':str(y),'m':m,'epsilon':eps},
        'field':{'generator':'q','embedding':f'exp(i*pi/{p})',
                 'modulus_descending':[str(z) for z in field.mod.to_list()],
                 'coefficient_order':'ascending powers 1,q,...', 'degree':field_degree},
        'dimension':d,'depth_zero_dimension':len(depth0),'exact_checks':checks,
        'ranks_F_minus_f0_powers':ranks,'ranks_Fbar_minus_f0_powers':branks,
        'joint_stacked_ranks':joint_ranks,
        'target':{'dimension':tdim,'kernel_power':stable_power,
                  'f0_coeff_id':coeff_id(f0),'nilpotent_ranks':internal_ranks[1:],
                  'Jordan_blocks_by_size':jordan,'rank_N_minus_M':dd.rank(),
                  'rank_N_plus_M':(nn+mm).rank(),'rank_NM':nm.rank(),
                  'hoop_algebra_dimension':algebra_dimension,'relations':relations,
                  'basis_pivot_rows':list(pivot_rows)},
        'depth_one':{'dimension':zdim,'same_face_dimension':len(depth0),
                     'edge_dimension':len(edges),'nilpotent_rank':(a*zbase).rank(),
                     'g_coeff_id':coeff_id(g),'checks':depth_one_checks,
                     'basis':sparse(zbase),'edge_descriptors':edges},
        'basis_states':[{'matching':s[0],'face_a':s[1],'face_b':s[2]} for s in states],
        'matrices':all_matrices,'target_matrices':target_matrices,
        'rank_certificates':rank_certificates,
        'coefficient_dictionary':coeffs,
        'elapsed_seconds':round(time.time()-started,3),
    }
    if not (all(checks.values())):
        raise AssertionError(checks)
    if not (all(depth_one_checks.values())):
        raise AssertionError(depth_one_checks)
    if output:
        Path(output).parent.mkdir(parents=True,exist_ok=True)
        Path(output).write_text(json.dumps(result,separators=(',',':'))+'\n')
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--n',type=int,default=6)
    parser.add_argument('--p',type=int,default=5)
    parser.add_argument('--x',default='1')
    parser.add_argument('--y',default='2')
    parser.add_argument('--m',type=int,default=1)
    parser.add_argument('--epsilon',type=int,choices=[-1,1],default=1)
    parser.add_argument('--exact',action='store_true')
    parser.add_argument('--out',default='research/marked_results.json')
    args = parser.parse_args()
    if args.exact:
        if args.m!=1:
            raise ValueError('The depth-one verification currently requires m=1.')
        result = exact_case(args.n,args.p,x=args.x,y=args.y,m=args.m,
                            eps=args.epsilon,output=args.out)
        print(json.dumps({k:result[k] for k in ['parameters','dimension','exact_checks',
                         'ranks_F_minus_f0_powers','joint_stacked_ranks','target',
                         'elapsed_seconds']},indent=2))
        return
    result = numerical_summary(args.n,args.p,x=float(args.x),y=float(args.y),
                               m=args.m,eps=args.epsilon)
    Path(args.out).parent.mkdir(parents=True,exist_ok=True)
    Path(args.out).write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps({k:v for k,v in result.items() if k!='eigenvalues_numerical'},indent=2))


if __name__ == '__main__':
    main()