"""Baseline MPS expectation scan with the qmatchatea backend."""

import argparse
import logging
import math
import time

from qibo import Circuit, gates, hamiltonians
from qibo.symbols import X, Z

from qibotn.backends.qmatchatea import QMatchaTeaBackend


def parse_bonds(value):
    return [int(item) for item in value.split(",") if item.strip()]


def build_circuit(nqubits, nlayers, seed):
    import numpy as np

    rng = np.random.default_rng(seed)
    circuit = Circuit(nqubits)
    for _ in range(nlayers):
        for qubit in range(nqubits):
            circuit.add(gates.RY(qubit, theta=rng.uniform(-math.pi, math.pi)))
            circuit.add(gates.RZ(qubit, theta=rng.uniform(-math.pi, math.pi)))
        for qubit in range(0, nqubits - 1, 2):
            circuit.add(gates.CNOT(qubit, qubit + 1))
        for qubit in range(1, nqubits - 1, 2):
            circuit.add(gates.CNOT(qubit, qubit + 1))
    return circuit


def build_observable(nqubits):
    form = 0
    for qubit in range(nqubits - 1):
        form += 0.5 * Z(qubit) * Z(qubit + 1)
    form += 0.25 * X(0)
    return hamiltonians.SymbolicHamiltonian(form=form)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nqubits", type=int, default=20)
    parser.add_argument("--nlayers", type=int, default=8)
    parser.add_argument("--bonds", type=parse_bonds, default=parse_bonds("2,4,8,16,32"))
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--cut-ratio", type=float, default=1e-12)
    parser.add_argument("--svd-control", default="V")
    parser.add_argument("--no-exact", action="store_true")
    args = parser.parse_args()
    logging.getLogger("qibo.config").setLevel(logging.ERROR)
    logging.getLogger("qtealeaves").setLevel(logging.ERROR)

    circuit = build_circuit(args.nqubits, args.nlayers, args.seed)
    observable = build_observable(args.nqubits)
    exact = None
    if not args.no_exact:
        exact = float(observable.expectation_from_state(circuit().state()).real)

    print(f"nqubits={args.nqubits} nlayers={args.nlayers} seed={args.seed}")
    if exact is not None:
        print(f"exact={exact:.16e}")
    print("bond_dim expval abs_error rel_error seconds")

    backend = QMatchaTeaBackend()
    for bond in args.bonds:
        backend.configure_tn_simulation(
            ansatz="MPS",
            max_bond_dimension=bond,
            cut_ratio=args.cut_ratio,
            svd_control=args.svd_control,
        )
        start = time.perf_counter()
        value = float(backend.expectation(circuit, observable, preprocess=False).real)
        elapsed = time.perf_counter() - start
        abs_error = float("nan") if exact is None else abs(value - exact)
        rel_error = float("nan") if exact is None else abs_error / max(abs(exact), 1e-15)
        print(f"{bond:d} {value:.16e} {abs_error:.6e} {rel_error:.6e} {elapsed:.3f}")


if __name__ == "__main__":
    main()