"""Compare QMatchaTeaBackend with the VidalBackend fast path."""

from __future__ import annotations

import argparse
import json
import math
import time

import numpy as np
import torch
from qibo import Circuit, gates, hamiltonians
from qibo.symbols import X, Y, Z

from qibotn.backends.qmatchatea import QMatchaTeaBackend
from qibotn.backends.vidal import VidalBackend


def build_circuit(nqubits, nlayers, seed, kind):
    rng = np.random.default_rng(seed)
    circuit = Circuit(nqubits)
    for layer in range(nlayers):
        for q in range(nqubits):
            circuit.add(gates.RY(q, theta=rng.uniform(-math.pi, math.pi)))
            circuit.add(gates.RZ(q, theta=rng.uniform(-math.pi, math.pi)))
        if kind == "brickwall":
            for q in range(0, nqubits - 1, 2):
                circuit.add(gates.CNOT(q, q + 1))
            for q in range(1, nqubits - 1, 2):
                circuit.add(gates.CNOT(q, q + 1))
        elif kind == "shifted-cz":
            for q in range(layer % 2, nqubits - 1, 2):
                circuit.add(gates.CZ(q, q + 1))
        elif kind == "reversed-cnot":
            for q in range(0, nqubits - 1, 2):
                circuit.add(gates.CNOT(q + 1, q))
            for q in range(1, nqubits - 1, 2):
                circuit.add(gates.CNOT(q, q + 1))
        else:
            raise ValueError(f"Unknown circuit kind {kind!r}.")
    return circuit


def build_observable(nqubits, kind):
    form = 0
    if kind == "ring-xz":
        for q in range(nqubits):
            form += 0.5 * X(q) * Z((q + 1) % nqubits)
    elif kind == "open-zz":
        for q in range(nqubits - 1):
            form += Z(q) * Z(q + 1) / (nqubits - 1)
    elif kind == "mixed":
        form += 0.25 * X(0) - 0.5 * Z(nqubits - 1)
        for q in range(0, nqubits - 1, 3):
            form += 0.125 * Y(q) * Y(q + 1)
    else:
        raise ValueError(f"Unknown observable kind {kind!r}.")
    return hamiltonians.SymbolicHamiltonian(form=form)


def run_backend(backend, circuit, observable):
    start = time.perf_counter()
    value = backend.expectation(circuit, observable, preprocess=False, compile_circuit=True)
    return float(np.real(value)), time.perf_counter() - start


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nqubits", type=int, default=34)
    parser.add_argument("--nlayers", type=int, default=20)
    parser.add_argument("--bond", "--bonds", dest="bond", type=int, default=512)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--tensor-module", choices=("torch", "numpy"), default="torch")
    parser.add_argument("--torch-threads", type=int, default=32)
    parser.add_argument(
        "--circuit-kind",
        choices=("brickwall", "shifted-cz", "reversed-cnot"),
        default="brickwall",
    )
    parser.add_argument(
        "--observable-kind",
        choices=("ring-xz", "open-zz", "mixed"),
        default="ring-xz",
    )
    parser.add_argument("--reference-file")
    parser.add_argument("--skip-qmatchatea", action="store_true")
    args = parser.parse_args()

    torch.set_num_threads(args.torch_threads)
    circuit = build_circuit(args.nqubits, args.nlayers, args.seed, args.circuit_kind)
    observable = build_observable(args.nqubits, args.observable_kind)

    exact = None
    if args.reference_file:
        with open(args.reference_file, "r", encoding="utf-8") as f:
            exact = float(json.load(f)["expectation"])

    print(
        f"nqubits={args.nqubits} nlayers={args.nlayers} bond={args.bond} "
        f"circuit={args.circuit_kind} observable={args.observable_kind} "
        f"tensor_module={args.tensor_module} torch_threads={args.torch_threads}"
    )
    if exact is not None:
        print(f"exact={exact:.16e}")
    print("backend value abs_error seconds")

    if not args.skip_qmatchatea:
        qmt = QMatchaTeaBackend()
        qmt.configure_tn_simulation(
            ansatz="MPS",
            max_bond_dimension=args.bond,
            cut_ratio=1e-12,
            svd_control="E!",
            tensor_module=args.tensor_module,
            compile_circuit=True,
            track_memory=False,
        )
        value, seconds = run_backend(qmt, circuit, observable)
        error = float("nan") if exact is None else abs(value - exact)
        print(f"qmatchatea {value:.16e} {error:.6e} {seconds:.3f}")

    vidal = VidalBackend()
    vidal.configure_tn_simulation(
        ansatz="MPS",
        max_bond_dimension=args.bond,
        cut_ratio=1e-12,
        tensor_module=args.tensor_module,
        compile_circuit=True,
        fallback=True,
    )
    value, seconds = run_backend(vidal, circuit, observable)
    error = float("nan") if exact is None else abs(value - exact)
    print(f"vidal {value:.16e} {error:.6e} {seconds:.3f}")


if __name__ == "__main__":
    main()