#!/usr/bin/env bash
set -euo pipefail

# Contest-style Vidal/MPI MPS cases.
#
# Usage:
#   tools/run_vidal_mpi_contest_cases.sh main1
#   tools/run_vidal_mpi_contest_cases.sh main2
#   tools/run_vidal_mpi_contest_cases.sh strong
#   tools/run_vidal_mpi_contest_cases.sh all
#
# Common overrides:
#   PYTHON_BIN=.venv/bin/python
#   MPIEXEC=mpiexec
#   MPIEXEC_FULL="mpirun -np 4 -hostfile /home/yx/qibotn/hostfile -perhost 2"
#   HOSTFILE=hostfile                 # optional; used only if the file exists
#   RANKS=8
#   TORCH_THREADS=8
#   CUT_RATIO=1e-12
#   OBS_FILTER="boundary_ZZ_q2 ring_xz dense3_spread complex_iZ0"
#
# Per-case overrides:
#   MAIN1_NQ=128 MAIN1_LAYERS=50 MAIN1_BOND=1024 MAIN1_SEED=31001
#   MAIN2_NQ=128 MAIN2_LAYERS=64 MAIN2_BOND=2048 MAIN2_SEED=31002
#   STRONG_NQ=256 STRONG_LAYERS=64 STRONG_BOND=2048 STRONG_SEED=41001

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"

PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
MPIEXEC="${MPIEXEC:-mpiexec}"
HOSTFILE="${HOSTFILE:-}"
RANKS="${RANKS:-4}"
TORCH_THREADS="${TORCH_THREADS:-1}"
CUT_RATIO="${CUT_RATIO:-1e-12}"
OBS_FILTER="${OBS_FILTER:-}"

RUNNER_DIR="$ROOT_DIR/.tmp"
mkdir -p "$RUNNER_DIR"
RUNNER="$(mktemp "$RUNNER_DIR/qibotn_vidal_contest.XXXXXX.py")"
cleanup() {
  rm -f "$RUNNER"
}
trap cleanup EXIT

cat > "$RUNNER" <<'PY'
from __future__ import annotations

import argparse
import math
import time

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

from qibotn.backends.vidal import VidalBackend


def set_torch_threads(nthreads):
    try:
        import torch

        torch.set_num_threads(nthreads)
    except Exception:
        pass


def build_circuit(kind, nqubits, nlayers, seed):
    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 in ("rxx_rzz", "scramble"):
                circuit.add(gates.RX(q, theta=rng.uniform(-math.pi, math.pi)))

        if kind == "reversed_cnot":
            for q in range(0, nqubits - 1, 2):
                circuit.add(gates.CNOT(q + 1, q) if layer % 2 else gates.CNOT(q, q + 1))
            for q in range(1, nqubits - 1, 2):
                circuit.add(gates.CNOT(q + 1, q) if layer % 2 == 0 else gates.CNOT(q, q + 1))
        elif kind == "rxx_rzz":
            for q in range(layer % 2, nqubits - 1, 2):
                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
        elif kind == "scramble":
            for q in range(layer % 2, nqubits - 1, 2):
                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
                if layer % 5 == 4:
                    circuit.add(gates.SWAP(q, q + 1))
        else:
            raise ValueError(f"Unknown circuit kind {kind!r}.")

    return circuit


def ring_xz(nqubits):
    form = 0
    for q in range(nqubits):
        form += 0.5 * X(q) * Z((q + 1) % nqubits)
    return hamiltonians.SymbolicHamiltonian(form=form)


def open_zz(nqubits):
    form = 0
    for q in range(nqubits - 1):
        form += (1.0 / (nqubits - 1)) * Z(q) * Z(q + 1)
    return hamiltonians.SymbolicHamiltonian(form=form)


def range2_xx(nqubits):
    form = 0
    for q in range(nqubits - 2):
        form += (1.0 / (nqubits - 2)) * X(q) * X(q + 2)
    return hamiltonians.SymbolicHamiltonian(form=form)


def dense_observable(nqubits, qubits, seed, dim):
    rng = np.random.default_rng(seed)
    raw = rng.normal(size=(dim, dim)) + 1j * rng.normal(size=(dim, dim))
    matrix = (raw + raw.conj().T) / 2.0
    matrix = matrix / np.linalg.norm(matrix)
    return {"matrix": matrix, "qubits": list(qubits)}


def observables_for_case(nqubits, seed):
    q1 = nqubits // 4
    q2 = nqubits // 2
    q3 = (3 * nqubits) // 4
    last = nqubits - 1

    return [
        ("boundary_ZZ_q1", hamiltonians.SymbolicHamiltonian(form=Z(q1 - 1) * Z(q1))),
        ("boundary_ZZ_q2", hamiltonians.SymbolicHamiltonian(form=Z(q2 - 1) * Z(q2))),
        ("boundary_ZZ_q3", hamiltonians.SymbolicHamiltonian(form=Z(q3 - 1) * Z(q3))),
        (
            "long_Z_5_sites",
            hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(q1) * Z(q2) * Z(q3) * Z(last)),
        ),
        (
            "mixed_XZYZX",
            hamiltonians.SymbolicHamiltonian(form=X(0) * Z(q1) * Y(q2) * Z(q3) * X(last)),
        ),
        ("ring_xz", ring_xz(nqubits)),
        ("open_zz", open_zz(nqubits)),
        ("range2_xx", range2_xx(nqubits)),
        ("complex_iZ0", hamiltonians.SymbolicHamiltonian(form=1.0j * Z(0))),
        ("dense2_mid", dense_observable(nqubits, (q2 - 1, q2), seed + 101, 4)),
        ("dense3_spread", dense_observable(nqubits, (q1, q2, q3), seed + 202, 8)),
    ]


def run_case(args):
    set_torch_threads(args.torch_threads)
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()

    circuit = build_circuit(args.kind, args.nqubits, args.nlayers, args.seed)
    observables = observables_for_case(args.nqubits, args.seed)
    if args.obs_filter:
        wanted = set(args.obs_filter.split(","))
        observables = [(name, obs) for name, obs in observables if name in wanted]
        if not observables:
            raise ValueError(f"OBS_FILTER matched no observables: {args.obs_filter!r}")

    if rank == 0:
        print("=" * 88, flush=True)
        print(
            "case "
            f"label={args.label} kind={args.kind} ranks={size} "
            f"nqubits={args.nqubits} nlayers={args.nlayers} gates={len(circuit.queue)} "
            f"bond={args.bond} cut_ratio={args.cut_ratio:g} "
            f"torch_threads={args.torch_threads} seed={args.seed} "
            f"obs_filter={args.obs_filter or 'all'}",
            flush=True,
        )
        print(
            "observable value seconds trunc_sum trunc_max status",
            flush=True,
        )

    for obs_name, observable in observables:
        backend = VidalBackend()
        backend.configure_tn_simulation(
            max_bond_dimension=args.bond,
            cut_ratio=args.cut_ratio,
            tensor_module="torch",
            mpi_approach="CT",
            mpi_num_procs=size,
            fallback=False,
        )

        comm.Barrier()
        start = time.perf_counter()
        try:
            value = backend.expectation(
                circuit,
                observable,
                preprocess=True,
                compile_circuit=False,
            )
            status = "ok"
        except Exception as exc:  # pragma: no cover - printed for manual runs
            value = np.nan
            status = type(exc).__name__ + ":" + str(exc).split("\n", 1)[0]
        seconds = time.perf_counter() - start

        if rank == 0:
            print(
                f"{obs_name} {value!r} {seconds:.3f} "
                f"{backend.last_truncation_error:.6e} "
                f"{backend.last_max_truncation_error:.6e} {status}",
                flush=True,
            )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--label", required=True)
    parser.add_argument("--kind", choices=("reversed_cnot", "rxx_rzz", "scramble"), required=True)
    parser.add_argument("--nqubits", type=int, required=True)
    parser.add_argument("--nlayers", type=int, required=True)
    parser.add_argument("--bond", type=int, required=True)
    parser.add_argument("--cut-ratio", type=float, required=True)
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--torch-threads", type=int, required=True)
    parser.add_argument("--obs-filter", default="")
    run_case(parser.parse_args())


if __name__ == "__main__":
    main()
PY

if [[ -n "${MPIEXEC_FULL:-}" ]]; then
  read -r -a mpi_prefix <<< "$MPIEXEC_FULL"
else
  mpi_prefix=("$MPIEXEC")
  if [[ -n "$HOSTFILE" && -f "$HOSTFILE" ]]; then
    mpi_prefix+=("-hostfile" "$HOSTFILE")
  fi
  mpi_prefix+=("-n" "$RANKS")
fi

run_case() {
  local label="$1"
  local kind="$2"
  local nq="$3"
  local layers="$4"
  local bond="$5"
  local seed="$6"

  echo
  echo "Running $label: kind=$kind nqubits=$nq layers=$layers bond=$bond seed=$seed"
  echo "MPI: ${mpi_prefix[*]}"
  "${mpi_prefix[@]}" "$PYTHON_BIN" -u "$ROOT_DIR/tools/vidal_mpi_contest_runner.py" \
    --label "$label" \
    --kind "$kind" \
    --nqubits "$nq" \
    --nlayers "$layers" \
    --bond "$bond" \
    --cut-ratio "$CUT_RATIO" \
    --seed "$seed" \
    --torch-threads "$TORCH_THREADS" \
    --obs-filter "$(tr ' ' ',' <<< "$OBS_FILTER")"
}

case "${1:-help}" in
  main1)
    run_case \
      "main1-reversed-cnot" \
      "reversed_cnot" \
      "${MAIN1_NQ:-128}" \
      "${MAIN1_LAYERS:-50}" \
      "${MAIN1_BOND:-1024}" \
      "${MAIN1_SEED:-31001}"
    ;;
  main2)
    run_case \
      "main2-rxx-rzz" \
      "rxx_rzz" \
      "${MAIN2_NQ:-128}" \
      "${MAIN2_LAYERS:-64}" \
      "${MAIN2_BOND:-2048}" \
      "${MAIN2_SEED:-31002}"
    ;;
  strong)
    run_case \
      "strong-scramble" \
      "scramble" \
      "${STRONG_NQ:-256}" \
      "${STRONG_LAYERS:-64}" \
      "${STRONG_BOND:-2048}" \
      "${STRONG_SEED:-41001}"
    ;;
  all)
    "$0" main1
    "$0" main2
    "$0" strong
    ;;
  smoke)
    MAIN1_NQ="${MAIN1_NQ:-32}" \
    MAIN1_LAYERS="${MAIN1_LAYERS:-6}" \
    MAIN1_BOND="${MAIN1_BOND:-128}" \
    "$0" main1
    ;;
  help|*)
    cat >&2 <<'EOF'
Usage: tools/run_vidal_mpi_contest_cases.sh [main1|main2|strong|all|smoke]

Cases:
  main1   128 qubits, 50 layers, reversed-CNOT brickwall, chi=1024
  main2   128 qubits, 64 layers, RXX/RZZ brickwall, chi=2048
  strong  256 qubits, 64 layers, RXX/RZZ + periodic SWAP scramble, chi=2048
  smoke   Small syntax/runtime check of main1

Common overrides:
  PYTHON_BIN=.venv/bin/python
  MPIEXEC=mpiexec
  MPIEXEC_FULL="mpirun -np 4 -hostfile /home/yx/qibotn/hostfile -perhost 2"
  HOSTFILE=hostfile
  RANKS=8
  TORCH_THREADS=8
  CUT_RATIO=1e-12
  OBS_FILTER="boundary_ZZ_q2 ring_xz dense3_spread complex_iZ0"

Per-case overrides:
  MAIN1_NQ=128 MAIN1_LAYERS=50 MAIN1_BOND=1024 MAIN1_SEED=31001
  MAIN2_NQ=128 MAIN2_LAYERS=64 MAIN2_BOND=2048 MAIN2_SEED=31002
  STRONG_NQ=256 STRONG_LAYERS=64 STRONG_BOND=2048 STRONG_SEED=41001
EOF
    exit 2
    ;;
esac