"""Disordered XXZ chain: site <Z_i>(t) by exact Krylov evolution in the
conserved total-Z sector, with an MPS/TEBD fallback for chains whose
sector is too large to store.

Conventions: bit 0 = spin up = Z eigenvalue +1; site i <-> bit i of the
basis-state integer; Pauli operators; hbar = 1.

No training, no preprocessing, no use of reference traces: every number
comes from the Hamiltonian terms and initial state given in the instance.
"""
import itertools
import math
import time

import numpy as np

PRIORS = {
    "family": "xxz-chain-disordered",
    "parameters": {
        "exact_method": "Lanczos/Krylov propagation of the state restricted to "
                        "its fixed total-Z sector (exact up to krylov_tol)",
        "krylov_dim": 40,
        "krylov_tol": 1e-12,
        "max_sector_dim": 2_000_000,
        "fallback_method": "second-order TEBD (Trotter-Suzuki) on an MPS, used "
                           "only when the sector exceeds max_sector_dim "
                           "(e.g. the 32-site runtime/sanity instance)",
        "tebd_dt": 0.1,
        "tebd_chi_max": 128,
        "tebd_svd_cutoff": 1e-10,
        "tebd_wall_seconds": 120.0,
    },
    "notes": (
        "Non-adaptive in physics, hybrid in size: (1) Hamiltonian is built "
        "from the instance's Pauli terms; because XX+YY, ZZ and Z terms "
        "conserve total Z and the initial state is a Z product state, the "
        "state is evolved exactly inside one magnetisation sector using a "
        "short-iteration Lanczos propagator with a-posteriori error control. "
        "(2) For sectors larger than max_sector_dim (the 32-site sanity "
        "instance) a truncated-MPS TEBD is used instead; it is approximate "
        "(bond-dimension and Trotter error grow with entanglement) and has a "
        "wall-clock guard that holds the last computed profile if time runs "
        "out. Non-sector-conserving terms fall back to full-Hilbert-space "
        "Krylov for small systems. No offline training, fitting, or lookup "
        "of reference traces; numpy only (scipy.sparse used if importable)."
    ),
}

_PAULI = {
    "I": np.eye(2, dtype=complex),
    "X": np.array([[0, 1], [1, 0]], dtype=complex),
    "Y": np.array([[0, -1j], [1j, 0]], dtype=complex),
    "Z": np.array([[1, 0], [0, -1]], dtype=complex),
}


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


def _initial_bits(instance):
    st = instance["initial_state"]
    bits = [int(b) for b in st["bits"]]
    return bits


def _classify(n, terms):
    """Split into Z fields, ZZ couplings, and XX/YY couplings per pair.
    Returns None if any term is outside this set or XX != YY on a pair."""
    hz = np.zeros(n)
    zz = {}
    xx = {}
    yy = {}
    const = 0.0
    for c, ops in terms:
        if not ops:
            const += c
            continue
        items = sorted(ops.items())
        paulis = "".join(p for _, p in items)
        if len(items) == 1 and paulis == "Z":
            hz[items[0][0]] += c
        elif len(items) == 2 and paulis == "ZZ":
            key = (items[0][0], items[1][0])
            zz[key] = zz.get(key, 0.0) + c
        elif len(items) == 2 and paulis == "XX":
            key = (items[0][0], items[1][0])
            xx[key] = xx.get(key, 0.0) + c
        elif len(items) == 2 and paulis == "YY":
            key = (items[0][0], items[1][0])
            yy[key] = yy.get(key, 0.0) + c
        else:
            return None
    for key in set(xx) | set(yy):
        if abs(xx.get(key, 0.0) - yy.get(key, 0.0)) > 1e-14:
            return None
    return hz, zz, xx, const


# ---------------------------------------------------------------------------
# sparse operator helpers
# ---------------------------------------------------------------------------

class _Op:
    def __init__(self, diag, rows, cols, vals):
        self.dim = diag.shape[0]
        self.diag = diag
        self._sp = None
        try:
            import scipy.sparse as sp
            if rows.size:
                self._sp = sp.csr_matrix((vals, (rows, cols)), shape=(self.dim, self.dim))
            else:
                self._sp = sp.csr_matrix((self.dim, self.dim), dtype=vals.dtype)
        except Exception:
            self.rows, self.cols, self.vals = rows, cols, vals

    def __call__(self, v):
        out = self.diag * v
        if self._sp is not None:
            out = out + self._sp @ v
        elif self.rows.size:
            w = self.vals * v[self.cols]
            out = out + (np.bincount(self.rows, w.real, self.dim)
                         + 1j * np.bincount(self.rows, w.imag, self.dim))
        return out


def _sector_op(n, cls, bits):
    hz, zz, xx, const = cls
    k = sum(bits)  # number of "1" bits (down spins)
    combos = np.array(list(itertools.combinations(range(n), k)), dtype=np.int64)
    if k == 0:
        states = np.zeros(1, dtype=np.int64)
    else:
        states = np.sum(np.left_shift(np.int64(1), combos), axis=1)
    states = np.sort(states)
    dim = states.size
    zs = 1 - 2 * ((states[:, None] >> np.arange(n)[None, :]) & 1)  # dim x n
    diag = np.full(dim, const) + zs @ hz
    for (i, j), c in zz.items():
        diag = diag + c * zs[:, i] * zs[:, j]
    rows, cols, vals = [], [], []
    for (i, j), c in xx.items():
        differ = np.nonzero(zs[:, i] != zs[:, j])[0]
        if differ.size == 0 or c == 0.0:
            continue
        flipped = states[differ] ^ ((1 << i) | (1 << j))
        dst = np.searchsorted(states, flipped)
        rows.append(dst)
        cols.append(differ)
        vals.append(np.full(differ.size, 2.0 * c))  # (XX+YY)|ud> = 2|du>
    if rows:
        rows = np.concatenate(rows); cols = np.concatenate(cols); vals = np.concatenate(vals)
    else:
        rows = cols = np.zeros(0, dtype=np.int64); vals = np.zeros(0)
    op = _Op(diag.astype(float), rows, cols, vals.astype(complex))
    init = np.searchsorted(states, int(sum(b << i for i, b in enumerate(bits))))
    return op, zs.astype(float), init


def _full_op(n, terms, bits):
    dim = 1 << n
    states = np.arange(dim, dtype=np.int64)
    zs = 1 - 2 * ((states[:, None] >> np.arange(n)[None, :]) & 1)
    diag = np.zeros(dim)
    acc = {}
    for c, ops in terms:
        xmask = 0
        phase = np.ones(dim, dtype=complex)
        for s, p in ops.items():
            if p in ("X", "Y"):
                xmask |= 1 << s
            if p == "Z":
                phase = phase * zs[:, s]
            elif p == "Y":
                # Y|0> = i|1>, Y|1> = -i|0>  (acting on source state)
                phase = phase * (1j * zs[:, s])
        if xmask == 0:
            diag = diag + c * phase.real
        else:
            acc.setdefault(xmask, np.zeros(dim, dtype=complex))
            acc[xmask] = acc[xmask] + c * phase
    rows, cols, vals = [], [], []
    for xmask, amp in acc.items():
        nz = np.nonzero(np.abs(amp) > 0)[0]
        rows.append(states[nz] ^ xmask); cols.append(nz); vals.append(amp[nz])
    if rows:
        rows = np.concatenate(rows); cols = np.concatenate(cols); vals = np.concatenate(vals)
    else:
        rows = cols = np.zeros(0, dtype=np.int64); vals = np.zeros(0, dtype=complex)
    op = _Op(diag.astype(complex), rows, cols, vals)
    init = int(sum(b << i for i, b in enumerate(bits)))
    return op, zs.astype(float), init


def _krylov_step(op, v, dt, m, tol):
    """Advance v by exp(-i H dt) with adaptive substeps; returns new v."""
    t_done = 0.0
    tau = dt
    while t_done < dt - 1e-15:
        tau = min(tau, dt - t_done)
        beta0 = np.linalg.norm(v)
        V = np.zeros((m + 1, v.size), dtype=complex)
        alpha = np.zeros(m)
        beta = np.zeros(m)
        V[0] = v / beta0
        kdim = m
        for j in range(m):
            w = op(V[j])
            alpha[j] = np.vdot(V[j], w).real
            w = w - alpha[j] * V[j] - (beta[j - 1] * V[j - 1] if j > 0 else 0)
            # full reorthogonalisation for robustness
            w = w - V[: j + 1].T @ (V[: j + 1].conj() @ w)
            beta[j] = np.linalg.norm(w)
            if beta[j] < 1e-13:
                kdim = j + 1
                break
            V[j + 1] = w / beta[j]
        T = np.diag(alpha[:kdim]) + np.diag(beta[: kdim - 1], 1) + np.diag(beta[: kdim - 1], -1)
        evals, evecs = np.linalg.eigh(T)
        while True:
            c = evecs @ (np.exp(-1j * evals * tau) * evecs[0].conj())
            err = beta[kdim - 1] * abs(c[-1]) if kdim == m else 0.0
            if err <= tol * max(tau, 1e-3) or tau < 1e-6:
                break
            tau *= 0.5
        v = beta0 * (c @ V[:kdim])
        t_done += tau
        if err < 0.1 * tol * tau:
            tau *= 1.5
    return v


def _exact(op, zs, init, times):
    m = PRIORS["parameters"]["krylov_dim"]
    tol = PRIORS["parameters"]["krylov_tol"]
    v = np.zeros(op.dim, dtype=complex)
    v[init] = 1.0
    order = np.argsort(times)
    out = np.zeros((len(times), zs.shape[1]))
    t_cur = 0.0
    for idx in order:
        t = float(times[idx])
        if t > t_cur:
            v = _krylov_step(op, v, t - t_cur, m, tol)
            t_cur = t
        p = np.abs(v) ** 2
        p /= p.sum()
        out[idx] = p @ zs
    return out


# ---------------------------------------------------------------------------
# TEBD fallback (nearest-neighbour chain)
# ---------------------------------------------------------------------------

def _tebd(n, terms, bits, times, deadline):
    P = PRIORS["parameters"]
    dt0, chi_max, cutoff = P["tebd_dt"], P["tebd_chi_max"], P["tebd_svd_cutoff"]
    hbond = [np.zeros((4, 4), dtype=complex) for _ in range(n - 1)]
    nb = [0] * n
    for b in range(n - 1):
        nb[b] += 1; nb[b + 1] += 1
    for c, ops in terms:
        sites = sorted(ops)
        if len(sites) == 1:
            s = sites[0]
            o = _PAULI[ops[s]]
            # split single-site terms evenly over the bonds touching the site
            for b in (s - 1, s):
                if 0 <= b < n - 1:
                    if b == s:
                        hbond[b] += c / nb[s] * np.kron(o, _PAULI["I"])
                    else:
                        hbond[b] += c / nb[s] * np.kron(_PAULI["I"], o)
        elif len(sites) == 2 and sites[1] == sites[0] + 1:
            hbond[sites[0]] += c * np.kron(_PAULI[ops[sites[0]]], _PAULI[ops[sites[1]]])
        elif len(sites) == 0:
            pass
        else:
            raise ValueError("TEBD fallback needs nearest-neighbour terms")

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

    # right-canonical MPS (Hastings form): B[s] (chi_l, 2, chi_r), S[s] on
    # the bond to the left of site s (S[0] = S[n] = [1]).
    B = []
    for b in bits:
        t = np.zeros((1, 2, 1), dtype=complex)
        t[0, b, 0] = 1.0
        B.append(t)
    S = [np.ones(1) for _ in range(n + 1)]
    # U(1) charge (number of down spins to the left) of every bond index;
    # SVDs are done block by block in charge so total Z is conserved exactly.
    Q = [np.array([sum(bits[:s])]) for s in range(n + 1)]

    def apply(b, G):
        # C = G . (B_b B_b+1), theta = S_b C
        C = np.tensordot(B[b], B[b + 1], axes=(2, 0))          # a i j c
        C = np.tensordot(C, G, axes=([1, 2], [2, 3]))          # a c i j
        C = C.transpose(0, 2, 3, 1)                            # a i j c
        th = S[b][:, None, None, None] * C
        a, _, _, c = th.shape
        M = th.reshape(a * 2, 2 * c)
        rq = (Q[b][:, None] + np.arange(2)[None, :]).reshape(-1)
        cq = (Q[b + 2][None, :] - np.arange(2)[:, None]).reshape(-1)
        blocks = []
        for q in np.intersect1d(rq, cq):
            ri = np.nonzero(rq == q)[0]
            ci = np.nonzero(cq == q)[0]
            u, y, vh = np.linalg.svd(M[np.ix_(ri, ci)], full_matrices=False)
            blocks.append((q, ci, y, vh))
        ys = np.concatenate([bl[2] for bl in blocks])
        order = np.argsort(-ys, kind="stable")
        keep = max(1, min(chi_max, int(np.sum(ys > cutoff * ys[order[0]]))))
        sel = np.zeros(ys.size, dtype=bool)
        sel[order[:keep]] = True
        Z = np.zeros((keep, 2 * c), dtype=complex)
        Y = np.zeros(keep)
        qn = np.zeros(keep, dtype=np.int64)
        k0 = off = 0
        for q, ci, y, vh in blocks:
            s = sel[off:off + y.size]
            kk = int(s.sum())
            if kk:
                Z[k0:k0 + kk][:, ci] = vh[s]
                Y[k0:k0 + kk] = y[s]
                qn[k0:k0 + kk] = q
                k0 += kk
            off += y.size
        Z = Z.reshape(keep, 2, c)
        S[b + 1] = Y / np.linalg.norm(Y)
        Q[b + 1] = qn
        B[b + 1] = Z
        B[b] = np.tensordot(C, Z.conj(), axes=([2, 3], [1, 2]))  # a i keep

    zop = np.array([1.0, -1.0])

    def measure():
        # full contraction (does not rely on exact canonical form)
        R = [None] * (n + 1)
        R[n] = np.ones((1, 1), dtype=complex)
        for s in range(n - 1, -1, -1):
            t = B[s]
            R[s] = np.einsum("aic,bid,cd->ab", t.conj(), t, R[s + 1], optimize=True)
        norm = R[0][0, 0].real
        L = np.ones((1, 1), dtype=complex)
        z = np.zeros(n)
        for s in range(n):
            t = B[s]
            z[s] = np.einsum("ab,aic,i,bid,cd->", L, t.conj(), zop, t, R[s + 1],
                             optimize=True).real / norm
            L = np.einsum("ab,aic,bid->cd", L, t.conj(), t, optimize=True)
        return np.clip(z, -1.0, 1.0)

    order = np.argsort(times)
    out = np.zeros((len(times), n))
    t_cur = 0.0
    last = measure()
    timed_out = False
    gcache = {}
    for idx in order:
        t = float(times[idx])
        if t > t_cur and not timed_out:
            nsteps = max(1, int(math.ceil((t - t_cur) / dt0 - 1e-9)))
            dt = (t - t_cur) / nsteps
            key = round(dt, 12)
            if key not in gcache:
                gcache[key] = (gates(dt / 2), gates(dt))
            Gh, Gf = gcache[key]
            for _ in range(nsteps):
                for b in range(0, n - 1, 2):
                    apply(b, Gh[b])
                for b in range(1, n - 1, 2):
                    apply(b, Gf[b])
                for b in range(0, n - 1, 2):
                    apply(b, Gh[b])
                if time.time() > deadline:
                    timed_out = True
                    break
            t_cur = t
            last = measure()
        out[idx] = last
    return out


# ---------------------------------------------------------------------------

def simulate(instance, times):
    t0 = time.time()
    times = np.asarray(times, dtype=float)
    n, terms = _parse_terms(instance)
    bits = _initial_bits(instance)
    cls = _classify(n, terms)
    maxdim = PRIORS["parameters"]["max_sector_dim"]
    if cls is not None:
        k = sum(bits)
        if math.comb(n, k) <= maxdim:
            op, zs, init = _sector_op(n, cls, bits)
            res = _exact(op, zs, init, times)
        else:
            res = _tebd(n, terms, bits, times,
                        t0 + PRIORS["parameters"]["tebd_wall_seconds"])
    elif n <= 20:
        op, zs, init = _full_op(n, terms, bits)
        res = _exact(op, zs, init, times)
    else:
        res = _tebd(n, terms, bits, times, t0 + PRIORS["parameters"]["tebd_wall_seconds"])
    # observables requested (default: Z on every site)
    obs = instance.get("observables")
    if obs:
        cols = [int(o["site"]) for o in obs]
        res = res[:, cols]
    return np.clip(res, -1.0, 1.0)