"""MPS time-evolving block decimation (TeNPy) with no imposed U(1) symmetry, bond dimension capped at 16.
Prior: low entanglement, without assuming a conserved charge."""
import numpy as np
from tenpy.networks.site import SpinHalfSite
from tenpy.models.model import CouplingModel, NearestNeighborModel, MPOModel
from tenpy.models.lattice import Chain
from tenpy.networks.mps import MPS
from tenpy.algorithms import tebd

PRIORS = {"family": "MPS-TEBD", "parameters": {"max_bond": 16, "svd_min": 1e-10, "dt": 0.25, "trotter_order": 2, "conserve": None},
          "notes": "low-entanglement prior; no imposed symmetry; suitable for transverse fields; exact for product states at t=0"}

class PauliChain(CouplingModel, NearestNeighborModel, MPOModel):
    """H = sum_terms coef * prod Pauli, for nearest-neighbour chains, built from the instance's term list. Pauli = 2 S."""
    def __init__(self, instance):
        L = instance["n_sites"]
        site = SpinHalfSite(conserve=None)
        lat = Chain(L, site, bc="open", bc_MPS="finite")
        CouplingModel.__init__(self, lat)
        for coef, ops in instance["terms"]:
            sites = sorted(int(s) for s in ops)
            if len(sites) == 1:
                s = sites[0]; self.add_onsite_term(2*coef,s,"S"+ops[str(s)].lower())
            elif len(sites)==2 and sites[1]==sites[0]+1:
                a,b=sites
                self.add_coupling_term(4*coef,a,b,"S"+ops[str(a)].lower(),"S"+ops[str(b)].lower())
            else: raise ValueError("Nearest-neighbour chain required")
        MPOModel.__init__(self, lat, self.calc_H_MPO())
        NearestNeighborModel.__init__(self, lat, self.calc_H_bond())

def simulate(instance, times):
    model = PauliChain(instance)
    # bit 0 = up = |0> = <Z> = +1; TeNPy's SpinHalfSite states are ["down", "up"]
    psi = MPS.from_product_state(model.lat.mps_sites(), ["down" if b else "up" for b in instance["initial_state"]["bits"]], bc="finite")
    p = PRIORS["parameters"]
    eng = tebd.TEBDEngine(psi, model, {"order": p["trotter_order"], "dt": p["dt"], "N_steps": 1, "trunc_params": {"chi_max": p["max_bond"], "svd_min": p["svd_min"]}})
    out, t = [], 0.0
    for tt in times:
        while t < tt - 1e-12:
            eng.run(); t += p["dt"]
        psi.canonical_form()  # truncation leaves the canonical form slightly off; restore it so the U(1) sector reads exactly
        out.append([2 * float(v) for v in psi.expectation_value("Sz")])
    return out