"""MPS time-evolving block decimation (TeNPy) with total-Sz conservation built into the tensors, bond dimension capped at 16.
Prior: low entanglement, and the family's U(1) symmetry respected exactly."""
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-U1", "parameters": {"max_bond": 16, "svd_min": 1e-10, "dt": 0.25, "trotter_order": 2, "conserve": "Sz"},
          "notes": "low-entanglement prior; total Sz conserved by construction; 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="Sz")
        lat = Chain(L, site, bc="open", bc_MPS="finite")
        CouplingModel.__init__(self, lat)
        pair = {"X": ("Sp", "Sm"), "Y": None, "Z": ("Sz", "Sz")}
        for coef, ops in instance["terms"]:
            sites = sorted(int(s) for s in ops)
            if len(sites) == 1:
                s, p = sites[0], ops[str(sites[0])]
                if p != "Z": raise ValueError("U(1) method: single-site terms must be Z")
                self.add_onsite_term(2 * coef, s, "Sz")
            elif len(sites) == 2 and sites[1] == sites[0] + 1:
                a, b = sites; pa, pb = ops[str(a)], ops[str(b)]
                if pa == pb == "Z": self.add_coupling_term(4 * coef, a, b, "Sz", "Sz")
                elif pa == pb in ("X", "Y"):  # XX + YY = 2 (S+S- + S-S+); each of the pair contributes half
                    self.add_coupling_term(coef, a, b, "Sp", "Sm"); self.add_coupling_term(coef, a, b, "Sm", "Sp")
                else: raise ValueError("U(1) method: XX and YY must appear with equal weight")
            else:
                raise ValueError("this method handles nearest-neighbour chains only")
        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