"""Benchmark: qibojit (reference) vs qibotn/quimb MPS, with error comparison."""
import time
import argparse
import os
import numpy as np
import qibo
from qibo import Circuit, gates

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")


def make_circuit(circuit_type, nqubits, nlayers=1):
    c = Circuit(nqubits)
    if circuit_type == "qft":
        from qibo.models import QFT
        return QFT(nqubits)
    elif circuit_type == "variational":
        for layer in range(nlayers):
            for q in range(nqubits):
                c.add(gates.RY(q, theta=np.random.uniform(0, 2 * np.pi)))
            offset = layer % 2
            for q in range(offset, nqubits - 1, 2):
                c.add(gates.CZ(q, q + 1))
    elif circuit_type == "ghz":
        c.add(gates.H(0))
        for q in range(nqubits - 1):
            c.add(gates.CNOT(q, q + 1))
    else:
        raise ValueError(f"Unknown circuit: {circuit_type}")
    return c


def run_qibojit(circuit):
    qibo.set_backend("qibojit", platform="numba")
    t0 = time.time()
    result = circuit()
    elapsed = time.time() - t0
    sv = result.state()
    return sv, elapsed


def run_quimb_mps(circuit, max_bond, svd_cutoff, optimizer):
    qibo.set_backend("qibotn", platform="quimb")
    b = qibo.get_backend()
    b.configure_tn_simulation(ansatz="mps", max_bond_dimension=max_bond, svd_cutoff=svd_cutoff)
    b.contractions_optimizer = optimizer

    t0 = time.time()
    result = b.execute_circuit(circuit, return_array=True)
    elapsed = time.time() - t0
    sv = result.state()
    return sv, elapsed


def compare(sv_ref, sv_mps):
    sv_ref = np.array(sv_ref, dtype=complex).flatten()
    sv_mps = np.array(sv_mps, dtype=complex).flatten()
    fidelity = abs(np.dot(sv_ref.conj(), sv_mps)) ** 2
    l2_err = np.linalg.norm(sv_ref - sv_mps)
    return fidelity, l2_err


def jit_cache_path(circuit_type, nqubits, nlayers):
    os.makedirs(DATA_DIR, exist_ok=True)
    return os.path.join(DATA_DIR, f"jit_{circuit_type}_n{nqubits}_l{nlayers}.npy")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nqubits", type=int, default=10)
    parser.add_argument("--circuit", type=str, default="ghz",
                        choices=["qft", "variational", "ghz"])
    parser.add_argument("--nlayers", type=int, default=3)
    parser.add_argument("--max-bond", type=int, default=None,
                        help="Max bond dimension for MPS (None = unlimited)")
    parser.add_argument("--svd-cutoff", type=float, default=1e-6)
    parser.add_argument("--optimizer", type=str, default="auto-hq")
    parser.add_argument("--skip-jit", action="store_true",
                        help="Skip qibojit run, load cached statevector if available")
    args = parser.parse_args()

    print(f"Circuit: {args.circuit}, nqubits={args.nqubits}, nlayers={args.nlayers}")
    print(f"MPS config: max_bond={args.max_bond}, svd_cutoff={args.svd_cutoff}, optimizer={args.optimizer}")

    cache_path = jit_cache_path(args.circuit, args.nqubits, args.nlayers)
    t_ref = None

    if args.skip_jit and os.path.exists(cache_path):
        sv_ref = np.load(cache_path)
        print(f"\n[qibojit]  loaded from cache: {cache_path}")
    else:
        np.random.seed(42)
        circuit_ref = make_circuit(args.circuit, args.nqubits, args.nlayers)
        sv_ref, t_ref = run_qibojit(circuit_ref)
        np.save(cache_path, sv_ref)
        print(f"\n[qibojit]  time={t_ref:.4f}s  (saved to {cache_path})")

    np.random.seed(42)
    circuit_mps = make_circuit(args.circuit, args.nqubits, args.nlayers)
    try:
        sv_mps, t_mps = run_quimb_mps(circuit_mps, args.max_bond, args.svd_cutoff, args.optimizer)
        fidelity, l2_err = compare(sv_ref, sv_mps)
        print(f"[quimb MPS] time={t_mps:.4f}s")
        print(f"\nFidelity : {fidelity:.8f}  (1=perfect)")
        print(f"L2 error : {l2_err:.2e}")
        if t_ref is not None and t_mps > 0:
            print(f"Speedup  : {t_ref/t_mps:.2f}x")
    except Exception as e:
        print(f"[quimb MPS] FAILED: {e}")
        raise


if __name__ == "__main__":
    main()