"""Exact fixed-magnetisation sector propagation for U(1)-symmetric spin models.

simulate(instance, times) builds the Hamiltonian from the instance's Pauli
terms restricted to the total-Z sector of the initial product state and
propagates it with a Chebyshev expansion.  No offline training, no reference
data, no lookup.  When the sector is too large to store (e.g. the 6 x 6
scaling instance), it falls back to overlapping exact cluster patches.
"""
import itertools
import math

import numpy as np
from scipy.sparse import csr_matrix, identity
from scipy.special import jv

PRIORS = {
    "family": "exact-sector-chebyshev",
    "parameters": {
        "max_sector_dim": 1000000,
        "propagator": "Chebyshev expansion of exp(-iH dt) between requested times, Gershgorin spectral bounds, Bessel tail < 1e-16",
        "fallback": "overlapping exact clusters (up to 4 x 4) with total-Z drift correction",
        "fallback_cluster": [4, 4],
        "offline_training": "none",
    },
    "notes": (
        "Adaptive/hybrid: if the fixed total-Z sector of the initial product "
        "state has dimension <= max_sector_dim, the dynamics are computed "
        "exactly in that sector (no approximation beyond floating point and "
        "the Chebyshev truncation, Bessel coefficients below 1e-16). Otherwise (6 x 6 "
        "scaling instance) each site's trace is taken from an exactly "
        "simulated open-boundary cluster of at most 4 x 4 sites in which it is "
        "most interior; this ignores correlations beyond the cluster and is "
        "only a sanity-level approximation. The fallback subtracts the mean "
        "total-Z drift per time and clips to [-1, 1]. No offline training, "
        "preprocessing or use of reference traces."
    ),
}


def _parse_terms(instance):
    terms = []
    for coef, ops in instance["terms"]:
        ops = {int(k): str(v).upper() for k, v in ops.items()}
        c = complex(coef)
        if c != 0:
            terms.append((c, ops))
    return terms


def _initial_bits(instance):
    n = int(instance["n_sites"])
    st = instance.get("initial_state", {})
    bits = st.get("bits")
    if bits is None:
        typ = str(st.get("type", "neel")).lower()
        geo = instance.get("geometry", {})
        if typ == "neel" and "Ly" in geo:
            ly = int(geo["Ly"])
            bits = [((i // ly) + (i % ly)) % 2 for i in range(n)]
        elif typ == "neel":
            bits = [i % 2 for i in range(n)]
        else:
            bits = [0] * n
    return [int(b) for b in bits]


def _sector_states(n, n_up_bits):
    """All n-bit integers with popcount n_up_bits (bit i = site i, 1 = down)."""
    m = math.comb(n, n_up_bits)
    states = np.empty(m, dtype=np.int64)
    for k, combo in enumerate(itertools.combinations(range(n), n_up_bits)):
        v = 0
        for b in combo:
            v |= 1 << b
        states[k] = v
    return np.sort(states)


def _sector_states_fast(n, k):
    if math.comb(n, k) < 200000:
        return _sector_states(n, k)
    # Gosper-free vectorised construction via popcount filtering in blocks.
    out = []
    block = 1 << 22
    total = 1 << n
    for start in range(0, total, block):
        x = np.arange(start, min(start + block, total), dtype=np.int64)
        pc = np.zeros_like(x)
        y = x.copy()
        while True:
            nz = y != 0
            if not nz.any():
                break
            pc += y & 1
            y >>= 1
        out.append(x[pc == k])
    return np.concatenate(out)


def _build_sector_hamiltonian(n, terms, states):
    dim = states.size
    index = {}
    use_dict = n > 26
    if not use_dict:
        lookup = np.full(1 << n, -1, dtype=np.int64)
        lookup[states] = np.arange(dim, dtype=np.int64)
    else:
        index = {int(s): i for i, s in enumerate(states)}
    rows, cols, vals = [], [], []
    diag = np.zeros(dim, dtype=np.complex128)
    cols_all = np.arange(dim, dtype=np.int64)
    for c, ops in terms:
        flip = 0
        amp = np.full(dim, c, dtype=np.complex128)
        for site, p in ops.items():
            b = (states >> site) & 1
            sign = 1 - 2 * b
            if p == "Z":
                amp *= sign
            elif p == "X":
                flip |= 1 << site
            elif p == "Y":
                flip |= 1 << site
                amp *= 1j * sign
            elif p != "I":
                raise ValueError("unknown Pauli " + p)
        if flip == 0:
            diag += amp
            continue
        tgt = states ^ flip
        if not use_dict:
            ti = lookup[tgt]
        else:
            ti = np.array([index.get(int(t), -1) for t in tgt], dtype=np.int64)
        ok = ti >= 0
        rows.append(ti[ok])
        cols.append(cols_all[ok])
        vals.append(amp[ok])
    rows.append(cols_all)
    cols.append(cols_all)
    vals.append(diag)
    r = np.concatenate(rows)
    cc = np.concatenate(cols)
    v = np.concatenate(vals)
    H = csr_matrix((v, (r, cc)), shape=(dim, dim))
    H.sum_duplicates()
    H.eliminate_zeros()
    if np.abs(H.data.imag).max(initial=0.0) < 1e-14:
        H = csr_matrix((H.data.real, H.indices, H.indptr), shape=H.shape)
    return H


def _is_u1(terms):
    """True if every term combination conserves total Z (checked pairwise XX/YY)."""
    pend = {}
    for c, ops in terms:
        letters = set(ops.values()) - {"I", "Z"}
        if not letters:
            continue
        flips = tuple(sorted(s for s, p in ops.items() if p in "XY"))
        if len(flips) != 2:
            return False
        key = (flips, tuple(sorted((s, p) for s, p in ops.items() if p == "Z")))
        a, b = flips
        pa, pb = ops[a], ops[b]
        pend.setdefault(key, {})
        pend[key][(pa, pb)] = pend[key].get((pa, pb), 0) + c
    for d in pend.values():
        if abs(d.get(("X", "X"), 0) - d.get(("Y", "Y"), 0)) > 1e-12:
            return False
        if abs(d.get(("X", "Y"), 0) + d.get(("Y", "X"), 0)) > 1e-12:
            return False
    return True


def _evolve_sector(n, terms, bits, times):
    k = int(sum(bits))
    states = _sector_states_fast(n, k)
    H = _build_sector_hamiltonian(n, terms, states)
    init = 0
    for i, b in enumerate(bits):
        if b:
            init |= 1 << i
    psi0 = np.zeros(states.size, dtype=np.complex128)
    psi0[int(np.searchsorted(states, init))] = 1.0
    zsign = np.empty((n, states.size), dtype=np.float64)
    for i in range(n):
        zsign[i] = 1.0 - 2.0 * ((states >> i) & 1)
    times = np.asarray(times, dtype=float)
    out = np.empty((times.size, n))
    prop = _ChebyshevPropagator(H)
    order = np.argsort(times, kind="stable")
    psi = np.column_stack([psi0.real, psi0.imag])
    tprev = 0.0
    for idx in order:
        t = float(times[idx])
        if t != tprev:
            psi = prop.step(psi, t - tprev)
            tprev = t
        out[idx] = zsign @ (psi[:, 0] ** 2 + psi[:, 1] ** 2)
    return out


class _ChebyshevPropagator:
    """exp(-i H dt) on psi stored as (dim, 2) real [Re, Im] for real symmetric H.

    Spectral bounds from Gershgorin discs (rigorous), Bessel-series truncation
    once the coefficients drop below 1e-16.
    """

    def __init__(self, H):
        if np.iscomplexobj(H.data):
            raise ValueError("complex H not supported by Chebyshev path")
        absH = abs(H)
        d = H.diagonal()
        radius = np.asarray(absH.sum(axis=1)).ravel() - np.abs(d)
        lo = float((d - radius).min())
        hi = float((d + radius).max())
        self.c = 0.5 * (hi + lo)
        self.r = 0.5 * (hi - lo) * 1.01 + 1e-12
        n = H.shape[0]
        self.Hs = (H - self.c * identity(n, format="csr")) * (1.0 / self.r)
        self.Hs = csr_matrix(self.Hs)

    def _apply(self, v):
        return self.Hs @ v

    def step(self, psi, dt):
        x = self.r * dt
        kmax = int(x + 40 + 10 * np.log10(x + 10))
        coef = jv(np.arange(kmax + 1), x)
        # truncate where the tail is negligible
        big = np.nonzero(np.abs(coef) > 1e-16)[0]
        kmax = int(big.max()) if big.size else 1
        # (-i)^k factors applied to a real-pair state: multiply by -i => (Re, Im) -> (Im, -Re)
        def mul_minus_i_pow(v, k):
            k %= 4
            if k == 0:
                return v
            if k == 1:
                return np.column_stack([v[:, 1], -v[:, 0]])
            if k == 2:
                return -v
            return np.column_stack([-v[:, 1], v[:, 0]])
        t0 = psi
        t1 = self._apply(psi)
        acc = coef[0] * t0 + 2.0 * coef[1] * mul_minus_i_pow(t1, 1)
        for k in range(2, kmax + 1):
            t2 = 2.0 * self._apply(t1) - t0
            acc += 2.0 * coef[k] * mul_minus_i_pow(t2, k)
            t0, t1 = t1, t2
        # global phase exp(-i c dt) leaves |psi|^2 unchanged but apply it for correctness
        ph = -self.c * dt
        cs, sn = np.cos(ph), np.sin(ph)
        re, im = acc[:, 0], acc[:, 1]
        return np.column_stack([cs * re - sn * im, sn * re + cs * im])


def _cluster_fallback(instance, terms, bits, times):
    n = int(instance["n_sites"])
    geo = instance.get("geometry", {})
    lx, ly = int(geo.get("Lx", n)), int(geo.get("Ly", 1))
    cx, cy = min(4, lx), min(4, ly)
    times = np.asarray(times, dtype=float)
    out = np.zeros((times.size, n))
    best = np.full(n, -1.0)
    for x0 in range(0, lx - cx + 1):
        for y0 in range(0, ly - cy + 1):
            sites = [x * ly + y for x in range(x0, x0 + cx) for y in range(y0, y0 + cy)]
            loc = {s: j for j, s in enumerate(sites)}
            sub = [(c, {loc[s]: p for s, p in ops.items()})
                   for c, ops in terms if all(s in loc for s in ops)]
            # interior score: distance to cut edges (lattice edges are not cuts)
            score = {}
            need = False
            for s in sites:
                x, y = divmod(s, ly)
                d = []
                if x0 > 0:
                    d.append(x - x0)
                if x0 + cx < lx:
                    d.append(x0 + cx - 1 - x)
                if y0 > 0:
                    d.append(y - y0)
                if y0 + cy < ly:
                    d.append(y0 + cy - 1 - y)
                score[s] = min(d) if d else 99.0
                if score[s] > best[s]:
                    need = True
            if not need:
                continue
            res = _evolve_sector(len(sites), sub, [bits[s] for s in sites], times)
            for s in sites:
                if score[s] > best[s]:
                    best[s] = score[s]
                    out[:, s] = res[:, loc[s]]
    z0 = np.array([1.0 - 2.0 * b for b in bits])
    drift = out.sum(axis=1) - z0.sum()
    out -= drift[:, None] / n
    return np.clip(out, -1.0, 1.0)


def simulate(instance, times):
    n = int(instance["n_sites"])
    terms = _parse_terms(instance)
    bits = _initial_bits(instance)
    times = np.asarray(times, dtype=float)
    if not _is_u1(terms):
        raise ValueError("Hamiltonian does not conserve total Z")
    k = int(sum(bits))
    dim = math.comb(n, k)
    if dim <= PRIORS["parameters"]["max_sector_dim"]:
        out = _evolve_sector(n, terms, bits, times)
    else:
        out = _cluster_fallback(instance, terms, bits, times)
    obs = instance.get("observables")
    if obs:
        cols = [int(o["site"]) for o in obs]
        if cols != list(range(n)):
            out = out[:, cols]
    return np.clip(out, -1.0, 1.0)