"""From-scratch m-scheme shell-model CI for the sd shell (USDB interaction).

Validated against cosmo (Volya, github.com/alvolya/cosmo) run locally on the
same interaction file: 18O and 24Mg energies agree to the solver tolerance.

The three study codes each contributed one idea, rebuilt here from scratch:
  - CENS FCI  (Engeland & Hjorth-Jensen): the m-scheme basis as bitstrings —
    one integer per Slater determinant, one bit per single-particle m-state.
  - cosmo     (Volya): prune during enumeration — never generate a determinant
    whose remaining capacity can't reach the target Jz (MakeMBS.cxx:313).
  - OpenFCI   (Kvaal): the two-body machinery as 'which orbitals differ',
    with fermionic phases from occupation counts (MatrixMachine.hpp).

Conventions are the oxbash/.int file's: normalized isospin-coupled TBME
<ab;JT|V|cd;JT>, single-particle energies in the header, and a mass scaling
V -> V*(18/A)^0.3 (the header's 16.0 18.0 0.3 fields). Rather than trusting
any remembered phase convention, the coupled pair states |ab;J M T Tz> are
constructed NUMERICALLY as vectors in the two-particle m-scheme space and
normalized numerically — the only convention left is the file's own.

Usage: python3 gen_shell_model.py path/to/usdb.int [nucleus]
       nucleus in {18O, 20Ne, 24Mg}; default 24Mg.
"""
import sys
import math
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()

# ── driver ───────────────────────────────────────────────────────────────────
NUCLEI = {  # name -> (valence protons, valence neutrons)
    '18O': (0, 2), '20Ne': (2, 2), '24Mg': (4, 4),
}

def main():
    int_path = sys.argv[1] if len(sys.argv) > 1 else 'usdb.int'
    nuc = sys.argv[2] if len(sys.argv) > 2 else '24Mg'
    n_p, n_n = NUCLEI[nuc]
    n_val = n_p + n_n
    spe, tbme, tail = read_int(int_path)
    a_core, a_ref, power = tail[-3], tail[-2], tail[-1]
    scale = ((a_ref / (a_core + n_val)) ** power)
    print(f"{nuc}: {n_p}p + {n_n}n in sd shell; TBME scale = ({a_ref:.0f}/{a_core + n_val:.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(n_p, n_n, 0)
    print(f"m-scheme basis dimension (Jz=0): {len(basis)}")
    H = assemble(basis, spe_vec, V2)
    k = min(10, len(basis) - 2)
    E = np.sort(eigsh(H, k=k, which='SA', return_eigenvectors=False))
    print("lowest eigenvalues (MeV):")
    for e in E:
        print(f"  {e:12.4f}")

if __name__ == '__main__':
    main()