"""Site-Z dynamics for the open Ising chain with transverse and longitudinal fields.

n <= EXACT_MAX_SITES: exact state-vector propagation. H is read from the instance
`terms`; diagonal (Z-string) terms become an energy vector, single-site X terms are
applied by flipping one axis of the reshaped state. exp(-i H dt) between consecutive
output times is applied with a Chebyshev expansion whose spectral window comes from
a Lanczos estimate of the extremal eigenvalues plus a safety margin.

Larger n: second-order TEBD on an MPS (right-canonical form with singular values,
Hastings' update) with a capped bond dimension. This is an approximation used only
for chains too large for the state vector.

No training, fitting, preprocessing or lookup of reference data is used.
"""
import time

import numpy as np
from scipy.special import jv

PRIORS = {
    "family": "exact state-vector Chebyshev propagation (n<=20); bond-capped 2nd-order TEBD MPS fallback (n>20)",
    "parameters": {
        "exact_max_sites": 20,
        "chebyshev_bessel_cutoff": 1e-16,
        "norm_guard": "redo window with Gershgorin bounds if |norm-1| > 1e-8",
        "threads": 4,
        "lanczos_iterations": 30,
        "chebyshev_window": 2.0,
        "spectral_margin": "5% of half-width + 1.0",
        "tebd_max_bond": 64,
        "tebd_dt": 0.05,
        "tebd_svd_cutoff": 1e-10,
        "tebd_trotter_order": 2,
    },
    "notes": (
        "Physics-only method: the Hamiltonian is built from the instance terms (Z/ZZ diagonal, X fields) "
        "and the initial product state from instance bits (bit 0 -> Z=+1). For n<=20 the dynamics is exact up "
        "to Chebyshev truncation (Bessel coefficients < 1e-16) and floating point; no adaptive switching other than the fixed size "
        "threshold, and a wall-clock guard on the TEBD fallback that lowers the bond cap if a run is slow. "
        "For n>20 (only the 32-site runtime/sanity run) TEBD with bond dimension <=64 is an uncontrolled "
        "approximation at late times in this chaotic chain. No offline training, preprocessing, or use of "
        "reference traces."
    ),
}

EXACT_MAX_SITES = 20
_NTHREADS = 4


def _parse(instance):
    n = int(instance["n_sites"])
    diag_terms = []  # (coeff, [sites]) for Z strings
    x_terms = []  # (coeff, site)
    for coeff, ops in instance["terms"]:
        ops = {int(k): v for k, v in ops.items()}
        paulis = set(ops.values())
        if paulis <= {"Z"}:
            diag_terms.append((float(coeff), sorted(ops)))
        elif len(ops) == 1 and paulis == {"X"}:
            x_terms.append((float(coeff), next(iter(ops))))
        else:
            raise ValueError("unsupported term %r" % (ops,))
    init = instance["initial_state"]
    if "bits" not in init:
        raise ValueError("initial state needs explicit bits")
    bits = [int(b) for b in init["bits"]]
    return n, diag_terms, x_terms, bits


# ---------------------------------------------------------------- exact

def _zvals(n):
    # site i is axis i of the state reshaped to (2,)*n; index 0 -> Z=+1
    idx = np.arange(2 ** n, dtype=np.int64)
    return [(1 - 2 * ((idx >> (n - 1 - i)) & 1)).astype(np.float64) for i in range(n)]


class _ExactH:
    def __init__(self, n, diag_terms, x_terms):
        self.n = n
        zs = _zvals(n)
        d = np.zeros(2 ** n)
        for c, sites in diag_terms:
            v = np.full(2 ** n, c)
            for s in sites:
                v = v * zs[s]
            d += v
        self.diag = d
        self.zs = zs
        xc = {}
        for c, s in x_terms:
            xc[s] = xc.get(s, 0.0) + c
        self.x = [(s, c) for s, c in sorted(xc.items()) if c != 0.0]
        self.xnorm = sum(abs(c) for _, c in self.x)

    def apply(self, v, out):
        np.multiply(self.diag, v, out=out)
        n = self.n
        for s, c in self.x:
            shp = (2 ** s, 2, 2 ** (n - s - 1))
            o = out.reshape(shp)
            w = v.reshape(shp)
            if c == 1.0:
                o += w[:, ::-1, :]
            else:
                o += c * w[:, ::-1, :]
        return out


def _lanczos_bounds(H, iters):
    # real start vector: H is real, so the Krylov space stays real
    dim = H.diag.size
    rng = np.random.default_rng(12345)
    q = rng.standard_normal(dim)
    q /= np.linalg.norm(q)
    q_prev = np.zeros_like(q)
    alphas, betas = [], []
    beta = 0.0
    for _ in range(iters):
        w = H.apply(q, np.empty_like(q))
        a = float(q @ w)
        w -= a * q + beta * q_prev
        alphas.append(a)
        beta = float(np.linalg.norm(w))
        if beta < 1e-12:
            break
        betas.append(beta)
        q_prev, q = q, w / beta
    k = len(alphas)
    T = np.diag(alphas) + np.diag(betas[: k - 1], 1) + np.diag(betas[: k - 1], -1)
    ev = np.linalg.eigvalsh(T)
    return ev[0], ev[-1]


def _gersh(H):
    return H.diag.min() - H.xnorm, H.diag.max() + H.xnorm


class _Cheb:
    """Chunk-parallel Chebyshev recurrence for Hs = (H - c)/a."""

    def __init__(self, H, lo, hi, pool):
        n = H.n
        self.n, self.dim = n, 2 ** n
        self.c = 0.5 * (hi + lo)
        self.a = 0.5 * (hi - lo)
        self.d2 = 2.0 * (H.diag - self.c) / self.a
        groups = {}
        for s, coef in H.x:
            groups.setdefault(2.0 * coef / self.a, []).append(s)
        self.groups = list(groups.items())
        self.nb = 2 if n >= 6 else 0
        self.nchunk = 2 ** self.nb
        self.csz = self.dim // self.nchunk
        self.bufs = [np.empty(self.csz, dtype=np.complex128) for _ in range(self.nchunk)]
        self.tmps = [np.empty(self.csz, dtype=np.complex128) for _ in range(self.nchunk)]
        self.pool = pool

    def _chunk(self, k, u1, u0, out, accs, first):
        n, nb, csz = self.n, self.nb, self.csz
        sl = slice(k * csz, (k + 1) * csz)
        o = out[sl]
        np.multiply(self.d2[sl], u1[sl], out=o)
        f = self.bufs[k]
        for xc, sites in self.groups:
            started = False
            for s in sites:
                if s < nb:
                    src = k ^ (1 << (nb - 1 - s))
                    w = u1[src * csz:(src + 1) * csz]
                    fv = f
                else:
                    shp = (2 ** (s - nb), 2, 2 ** (n - s - 1))
                    w = u1[sl].reshape(shp)[:, ::-1, :]
                    fv = f.reshape(shp)
                if started:
                    np.add(fv, w, out=fv)
                else:
                    np.copyto(fv, w)
                    started = True
            np.multiply(f, xc, out=f)
            np.add(o, f, out=o)
        if first:
            np.multiply(o, 0.5, out=o)
        else:
            np.subtract(o, u0[sl], out=o)
        t = self.tmps[k]
        for acc, coef in accs:
            if coef == 0.0:
                continue
            np.multiply(o, coef, out=t)
            a_ = acc[sl]
            np.add(a_, t, out=a_)

    def rec(self, u1, u0, out, accs, first=False):
        if self.pool is None:
            for k in range(self.nchunk):
                self._chunk(k, u1, u0, out, accs, first)
        else:
            list(self.pool.map(lambda k: self._chunk(k, u1, u0, out, accs, first), range(self.nchunk)))

    def propagate(self, psi, dts, tol=1e-16):
        """Return [exp(-i H dt) psi for dt in dts] from one shared expansion."""
        a, c = self.a, self.c
        xs = [a * dt for dt in dts]
        xmax = max(xs)
        kmax = int(xmax + 10 * xmax ** (1 / 3) + 30)
        J = np.array([jv(np.arange(kmax + 1), x) for x in xs])  # (W, K)
        last = kmax
        while last > xmax and np.all(np.abs(J[:, last]) < tol):
            last -= 1
        ph = (-1j) ** np.arange(last + 1)
        C = 2.0 * J[:, : last + 1] * ph[None, :]
        C[:, 0] *= 0.5
        for j, x in enumerate(xs):  # drop negligible tail per output
            C[j, (np.abs(J[j, : last + 1]) < tol) & (np.arange(last + 1) > x)] = 0.0
        accs = [np.multiply(psi, C[j, 0]) for j in range(len(dts))]
        u0 = psi.copy()
        u1 = np.empty_like(psi)
        u2 = np.empty_like(psi)
        self.rec(u0, None, u1, [(accs[j], C[j, 1]) for j in range(len(dts))], first=True)
        for k in range(2, last + 1):
            self.rec(u1, u0, u2, [(accs[j], C[j, k]) for j in range(len(dts))])
            u0, u1, u2 = u1, u2, u0
        outs = []
        for j, dt in enumerate(dts):
            outs.append(accs[j] * np.exp(-1j * c * dt))
        return outs


def _simulate_exact(n, diag_terms, x_terms, bits, times):
    from concurrent.futures import ThreadPoolExecutor

    H = _ExactH(n, diag_terms, x_terms)
    dim = 2 ** n
    psi = np.zeros(dim, dtype=np.complex128)
    psi[int("".join(str(b) for b in bits), 2)] = 1.0
    gl, gu = _gersh(H)
    lo, hi = _lanczos_bounds(H, PRIORS["parameters"]["lanczos_iterations"])
    margin = 0.05 * 0.5 * (hi - lo) + 1.0
    lo, hi = max(lo - margin, gl), min(hi + margin, gu)
    zmat = np.stack(H.zs)
    window = PRIORS["parameters"]["chebyshev_window"]
    pool = None
    if n >= 6:
        try:
            pool = ThreadPoolExecutor(max_workers=_NTHREADS)
        except Exception:  # threads unavailable: run the chunks serially
            pool = None
    try:
        cheb = _Cheb(H, lo, hi, pool)
        safe = None
        res = np.empty((len(times), n))
        tcur = 0.0
        j = 0
        while j < len(times):
            # group output times into windows spanning at most `window` time units
            grp = [j]
            while grp[-1] + 1 < len(times) and float(times[grp[-1] + 1]) - tcur <= window + 1e-12:
                grp.append(grp[-1] + 1)
            if float(times[grp[0]]) - tcur > window:
                # long gap: advance without output first
                nsub = int(np.ceil((float(times[grp[0]]) - tcur) / window)) - 1
                for _ in range(nsub):
                    psi = cheb.propagate(psi, [window])[0]
                    tcur += window
                continue
            dts = [float(times[i]) - tcur for i in grp]
            outs = cheb.propagate(psi, dts)
            if max(abs(np.linalg.norm(o) - 1.0) for o in outs) > 1e-8:
                # spectral window too narrow for this state: redo with Gershgorin bounds
                if safe is None:
                    safe = _Cheb(H, gl, gu, pool)
                outs = safe.propagate(psi, dts)
            for i, o in zip(grp, outs):
                o = o / np.linalg.norm(o)
                p = o.real ** 2 + o.imag ** 2
                res[i] = zmat @ p
            psi = outs[-1] / np.linalg.norm(outs[-1])
            tcur = float(times[grp[-1]])
            j = grp[-1] + 1
    finally:
        if pool is not None:
            pool.shutdown()
    return res


# ---------------------------------------------------------------- TEBD

_X = np.array([[0, 1], [1, 0]], dtype=np.complex128)
_Z = np.array([[1, 0], [0, -1]], dtype=np.complex128)
_I = np.eye(2, dtype=np.complex128)


def _simulate_tebd(n, diag_terms, x_terms, bits, times, budget=150.0):
    t_start = time.time()
    hz = np.zeros(n)
    hx = np.zeros(n)
    Jb = np.zeros(n - 1)
    for c, sites in diag_terms:
        if len(sites) == 1:
            hz[sites[0]] += c
        elif len(sites) == 2 and sites[1] == sites[0] + 1:
            Jb[sites[0]] += c
        elif len(sites) == 0:
            pass
        else:
            raise ValueError("TEBD needs nearest-neighbour terms")
    for c, s in x_terms:
        hx[s] += c
    # bond Hamiltonians: single-site terms assigned to the bond on their right (last site to the left bond)
    hb = []
    for i in range(n - 1):
        h = Jb[i] * np.kron(_Z, _Z)
        wl = 1.0 if i == 0 else 0.5
        wr = 1.0 if i == n - 2 else 0.5
        h = h + wl * np.kron(hx[i] * _X + hz[i] * _Z, _I) + wr * np.kron(_I, hx[i + 1] * _X + hz[i + 1] * _Z)
        hb.append(h)

    def gates(dt):
        out = []
        for h in hb:
            e, V = np.linalg.eigh(h)
            out.append(((V * np.exp(-1j * dt * e)) @ V.conj().T).reshape(2, 2, 2, 2))
        return out

    chi_max = PRIORS["parameters"]["tebd_max_bond"]
    dt0 = PRIORS["parameters"]["tebd_dt"]
    cutoff = PRIORS["parameters"]["tebd_svd_cutoff"]
    B = []
    for b in bits:
        t = np.zeros((1, 2, 1), dtype=np.complex128)
        t[0, b, 0] = 1.0
        B.append(t)
    lam = [np.ones(1) for _ in range(n + 1)]

    def apply(i, g, chi):
        th = np.tensordot(B[i], B[i + 1], axes=(2, 0))  # (l, s1, s2, r)
        th = np.tensordot(th, g, axes=([1, 2], [2, 3]))  # (l, r, s1', s2')
        th = th.transpose(0, 2, 3, 1)
        l, _, _, r = th.shape
        full = (lam[i][:, None, None, None] * th).reshape(l * 2, 2 * r)
        try:
            U, S, Vh = np.linalg.svd(full, full_matrices=False)
        except np.linalg.LinAlgError:
            U, S, Vh = np.linalg.svd(full + 1e-14, full_matrices=False)
        keep = max(1, min(chi, int(np.sum(S > cutoff * S[0]))))
        S = S[:keep]
        Vh = Vh[:keep]
        S = S / np.linalg.norm(S)
        Bn = Vh.reshape(keep, 2, r)
        Bi = th.reshape(l * 2, 2 * r) @ Vh.conj().T
        B[i] = Bi.reshape(l, 2, keep)
        B[i + 1] = Bn
        lam[i + 1] = S

    def measure():
        z = np.empty(n)
        for i in range(n):
            w = lam[i][:, None, None] * B[i]
            p = np.sum(np.abs(w) ** 2, axis=(0, 2))
            z[i] = (p[0] - p[1]) / (p[0] + p[1])
        return z

    res = np.empty((len(times), n))
    tcur = 0.0
    chi = chi_max
    even = list(range(0, n - 1, 2))
    odd = list(range(1, n - 1, 2))
    cache = {}
    for j, t in enumerate(times):
        dt_tot = float(t) - tcur
        nst = int(np.ceil(abs(dt_tot) / dt0 - 1e-9)) if dt_tot != 0 else 0
        if nst:
            dt = dt_tot / nst
            key = round(dt, 14)
            if key not in cache:
                cache[key] = (gates(dt / 2), gates(dt))
            gh, gf = cache[key]
            for _ in range(nst):
                for i in even:
                    apply(i, gh[i], chi)
                for i in odd:
                    apply(i, gf[i], chi)
                for i in even:
                    apply(i, gh[i], chi)
            # wall-clock guard: lower the bond cap if the run would exceed the budget
            el = time.time() - t_start
            frac = (j + 1) / len(times)
            if el / frac > budget and chi > 8:
                chi = max(8, chi // 2)
        tcur = float(t)
        res[j] = measure()
    return res


def simulate(instance, times):
    n, diag_terms, x_terms, bits = _parse(instance)
    times = [float(t) for t in times]
    if n <= EXACT_MAX_SITES:
        out = _simulate_exact(n, diag_terms, x_terms, bits, times)
    else:
        out = _simulate_tebd(n, diag_terms, x_terms, bits, times)
    out = np.clip(np.nan_to_num(out, nan=0.0), -1.0, 1.0)
    return out