"""Build your own shell model -- step 5: 18O, complete.

Every piece is now on the table: seats (step 1), bitstring determinants
(step 2), the interaction file (step 3), and the sign rules (step 4). This
step assembles them into the full machinery -- the same code, line for line,
as the validated gen_shell_model.py -- and solves 18O.

The assembly loop is the whole idea of configuration interaction: for every
determinant, remove every occupied pair (with its popcount sign), and for
every pair the interaction can create from there (with its sign), connect the
two determinants by the pair-space matrix element. The result is a 14 x 14
symmetric matrix; numpy diagonalizes it exactly. No iteration, no
self-consistency, no convergence question -- within this valence space the
answer is EXACT. That is what "CI" buys compared to the mean-field
calculation of the companion tutorial.

The printed check: cosmo (Volya's production shell-model code), run on this
machine with the same interaction file, prints these same ten energies to
every digit shown.
"""
import math
import time
import numpy as np
from math import factorial, sqrt
from itertools import combinations
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import eigsh

# ── sd-shell single-particle space ──────────────────────────────────────────
# .int labels: 1 = d3/2, 2 = d5/2, 3 = s1/2   (l, 2j per label)
ORBITS = {1: (2, 3), 2: (2, 5), 3: (0, 1)}   # label -> (l, 2j)

def build_sp_states():
    """All (label, 2m, 2tz) single-particle m-states. tz=+1/2 proton, -1/2 neutron
    (only relative signs matter here). Returns list and index lookup."""
    states = []
    for tz2 in (+1, -1):
        for label, (l, j2) in ORBITS.items():
            for m2 in range(-j2, j2 + 1, 2):
                states.append((label, m2, tz2))
    return states, {s: i for i, s in enumerate(states)}

SP, SP_IDX = build_sp_states()
NSP = len(SP)   # 24

# ── Clebsch-Gordan (Condon–Shortley), all arguments doubled ─────────────────
def cg(j1, m1, j2, m2, J, M):
    if m1 + m2 != M or J > j1 + j2 or J < abs(j1 - j2):
        return 0.0
    if abs(m1) > j1 or abs(m2) > j2 or abs(M) > J:
        return 0.0
    def f(x2):  # factorial of a doubled-integer/2, must be integral
        assert x2 % 2 == 0
        return factorial(x2 // 2)
    pre = (J + 1) * f(J + j1 - j2) * f(J - j1 + j2) * f(j1 + j2 - J) / f(j1 + j2 + J + 2)
    pre *= f(J + M) * f(J - M) * f(j1 - m1) * f(j1 + m1) * f(j2 - m2) * f(j2 + m2)
    s = 0.0
    for k2 in range(0, j1 + j2 + J + 2, 2):
        d = [j1 + j2 - J - k2, j1 - m1 - k2, j2 + m2 - k2,
             J - j2 + m1 + k2, J - j1 - m2 + k2]
        if any(x < 0 for x in d):
            continue
        s += (-1) ** (k2 // 2) / (f(k2) * f(d[0]) * f(d[1]) * f(d[2]) * f(d[3]) * f(d[4]))
    return sqrt(pre) * s

# ── read the .int file ───────────────────────────────────────────────────────
def read_int(path):
    spe = {}
    tbme = {}   # (a,b,c,d,J,T) -> V, canonical a<=b, c<=d, (ab)<=(cd)
    header_tail = None
    n_lines = None
    with open(path) as fh:
        for raw in fh:
            line = raw.strip()
            if not line or line.startswith(('!', '#')):
                continue
            w = line.split()
            if n_lines is None:
                n_lines = int(w[0])
                spe = {1: float(w[1]), 2: float(w[2]), 3: float(w[3])}
                header_tail = [float(x) for x in w[4:]]
                continue
            a, b, c, d, J, T = (int(x) for x in w[:6])
            V = float(w[6])
            # canonicalize pair order inside bra/ket with the standard phase
            def canon(x, y):
                if x <= y:
                    return x, y, 1.0
                jx, jy = ORBITS[x][1], ORBITS[y][1]
                # |ba;JT> = (-1)^{(jx+jy)/2 - J - T} |ab;JT>  (j's doubled):
                # (-1)^{ja+jb-J} from the angular CG swap, (-1)^{1-T} from the
                # isospin CG swap, and one more (-1) from anticommuting a†a†.
                ph = (-1.0) ** ((jx + jy) // 2 - J - T)
                return y, x, ph
            a, b, pab = canon(a, b)
            c, d, pcd = canon(c, d)
            key = (a, b, c, d, J, T) if (a, b) <= (c, d) else (c, d, a, b, J, T)
            tbme[key] = pab * pcd * V
    return spe, tbme, header_tail

# ── numerically built coupled pair states ───────────────────────────────────
def pair_index_maps():
    pairs = list(combinations(range(NSP), 2))
    return pairs, {p: i for i, p in enumerate(pairs)}

PAIRS, PAIR_IDX = pair_index_maps()

def coupled_pair_vector(a, b, J2, M2, T2, Tz2):
    """|ab; J M T Tz> as a vector over ordered pairs (alpha<beta) of m-states.
    Built by brute expansion; normalized numerically (None if it vanishes)."""
    ja, jb = ORBITS[a][1], ORBITS[b][1]
    v = np.zeros(len(PAIRS))
    for ma in range(-ja, ja + 1, 2):
        mb = M2 - ma
        if abs(mb) > jb:
            continue
        cj = cg(ja, ma, jb, mb, J2, M2)
        if cj == 0.0:
            continue
        for ta in (+1, -1):
            tb = Tz2 - ta
            if abs(tb) > 1:
                continue
            ct = cg(1, ta, 1, tb, T2, Tz2)
            if ct == 0.0:
                continue
            ia, ib = SP_IDX[(a, ma, ta)], SP_IDX[(b, mb, tb)]
            if ia == ib:
                continue
            # a†_ia a†_ib |0> = |ia ib> ordered: sign if ia > ib
            if ia < ib:
                v[PAIR_IDX[(ia, ib)]] += cj * ct
            else:
                v[PAIR_IDX[(ib, ia)]] -= cj * ct
    n = np.linalg.norm(v)
    return (v / n, n) if n > 1e-12 else (None, 0.0)

def build_pair_hamiltonian(tbme, scale):
    """V2[p, q]: two-body matrix in the ordered-pair basis, from the TBME file.
    V2 = sum_JT V_JT sum_{M Tz} |ab;JMTTz><cd;JMTTz| with numerically
    normalized projectors — conventions cannot drift."""
    V2 = np.zeros((len(PAIRS), len(PAIRS)))
    for (a, b, c, d, J, T), V in tbme.items():
        J2, T2 = 2 * J, 2 * T
        for M2 in range(-J2, J2 + 1, 2):
            for Tz2 in range(-T2, T2 + 1, 2):
                va, na = coupled_pair_vector(a, b, J2, M2, T2, Tz2)
                if va is None:
                    continue
                vc, nc = (va, na) if (a, b, J, T) == (c, d, J, T) else \
                    coupled_pair_vector(c, d, J2, M2, T2, Tz2)
                if (a, b) == (c, d):
                    vc = va
                elif vc is None:
                    continue
                outer = np.outer(va, vc) * (V * scale)
                V2 += outer
                if (a, b) != (c, d):
                    V2 += outer.T
    return V2

# ── many-body basis: bitstrings with pruned enumeration ─────────────────────
def build_basis(n_protons, n_neutrons, M2_target):
    """All determinants (ints, bit i = m-state i occupied) with the right
    particle numbers per species and total 2*Jz = M2_target. Enumerated per
    species with cosmo-style pruning on the reachable Jz range."""
    def species_states(tz2):
        return [i for i, (lab, m2, t) in enumerate(SP) if t == tz2]

    def enum(states, n, m2_needed_min, m2_needed_max):
        # recursive with bound pruning: sort by m2 so prefix sums bound reach
        out = []
        ms = [SP[i][1] for i in states]
        # suffix min/max attainable sums for k picks from tail
        def rec(start, left, acc_bits, acc_m):
            if left == 0:
                out.append((acc_bits, acc_m))
                return
            for k in range(start, len(states) - left + 1):
                rem = left - 1
                tail = ms[k + 1:]
                lo = acc_m + ms[k] + sum(sorted(tail)[:rem])
                hi = acc_m + ms[k] + sum(sorted(tail)[-rem:] if rem else [])
                if lo > m2_needed_max or hi < m2_needed_min:
                    continue
                rec(k + 1, rem, acc_bits | (1 << states[k]), acc_m + ms[k])
        rec(0, n, 0, 0)
        return out

    prot = species_states(+1)
    neut = species_states(-1)
    # protons can carry any m2p; neutrons must supply M2_target - m2p
    m2_all = [SP[i][1] for i in prot]
    span = sum(sorted(m2_all)[-n_protons:]) if n_protons else 0
    plist = enum(prot, n_protons, -span, span) if n_protons else [(0, 0)]
    from collections import defaultdict
    by_need = defaultdict(list)
    nlist = enum(neut, n_neutrons, -span, span) if n_neutrons else [(0, 0)]
    for bits, m in nlist:
        by_need[m].append(bits)
    basis = []
    for pbits, pm in plist:
        for nbits in by_need.get(M2_target - pm, ()):
            basis.append(pbits | nbits)
    basis.sort()
    return basis

# ── Hamiltonian assembly ─────────────────────────────────────────────────────
def popcount_between(bits, i, j):
    """number of set bits strictly between positions i<j"""
    mask = ((1 << j) - 1) & ~((1 << (i + 1)) - 1)
    return bin(bits & mask).count('1')

def assemble(basis, spe_vec, V2):
    """Sparse H in the determinant basis.

    Pair convention: for p = (i<j), P_p = a_j a_i so P†_p|0> = a†_i a†_j|0> = |p>,
    and H2 = sum_{q,p} V2[q,p] P†_q P_p. Signs, measured with bit counts:
      P_p|D>   = (-1)^{n_<i(D) + n_<j(D) - 1} |D \ {i,j}>      (i<j, both occupied)
      P†_q|B>  = (-1)^{n_<r(B) + n_<s(B)}     |B + {r,s}>       (r<s, both empty)
    (OpenFCI's MatrixMachine does exactly this bookkeeping, in 512-bit registers.)
    """
    idx = {d: k for k, d in enumerate(basis)}
    dim = len(basis)
    conn = [np.nonzero(np.abs(V2[p]) > 1e-12)[0] for p in range(len(PAIRS))]
    rows, cols, vals = [], [], []
    for k, det in enumerate(basis):
        occ = [i for i in range(NSP) if det >> i & 1]
        acc = {k: sum(spe_vec[i] for i in occ)}
        for (i, j) in combinations(occ, 2):
            p = PAIR_IDX[(i, j)]
            s_ket = (-1) ** (bin(det & ((1 << i) - 1)).count('1')
                             + bin(det & ((1 << j) - 1)).count('1') - 1)
            base = det & ~(1 << i) & ~(1 << j)
            for q in conn[p]:
                r, s = PAIRS[q]
                if (base >> r & 1) or (base >> s & 1):
                    continue
                newdet = base | (1 << r) | (1 << s)
                kk = idx.get(newdet)
                if kk is None or kk < k:
                    continue
                s_bra = (-1) ** (bin(base & ((1 << r) - 1)).count('1')
                                 + bin(base & ((1 << s) - 1)).count('1'))
                v = V2[q, p] * s_ket * s_bra
                if abs(v) > 1e-14:
                    acc[kk] = acc.get(kk, 0.0) + v
        for kk, v in acc.items():
            rows.append(k); cols.append(kk); vals.append(v)
            if kk != k:
                rows.append(kk); cols.append(k); vals.append(v)
    return coo_matrix((vals, (rows, cols)), shape=(dim, dim)).tocsr()


COSMO_18O = [-11.9318, -9.9333, -8.4046, -7.5717, -7.3393,
             -6.5054, -2.9124, -2.0505, -1.1620, -0.9908]

spe, tbme, tail = read_int('usdb.int')
a_core, a_ref, power = tail[-3], tail[-2], tail[-1]
scale = (a_ref / (a_core + 2)) ** power
print(f"18O: 2 valence neutrons; TBME scale = ({a_ref:.0f}/{a_core + 2:.0f})^{power} = {scale:.6f}")

spe_vec = np.array([spe[SP[i][0]] for i in range(NSP)])
V2 = build_pair_hamiltonian(tbme, scale)
basis = build_basis(0, 2, 0)
print(f"m-scheme basis dimension (Jz = 0): {len(basis)}")

H = assemble(basis, spe_vec, V2).toarray()
E = np.sort(np.linalg.eigvalsh(H))[:10]
print(f"\n{'this code':>12}  {'cosmo':>10}")
for e, ref in zip(E, COSMO_18O):
    mark = "ok" if abs(e - ref) < 5e-5 else "MISMATCH"
    print(f"{e:12.4f}  {ref:10.4f}   {mark}")