"""Compute and cache a qibojit state-vector reference for the ring-XZ observable."""

import argparse
import json
import math
import time
from pathlib import Path

import numpy as np
import qibo
from qibo import Circuit, gates


def build_circuit(nqubits, nlayers, seed):
    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 ring_xz_expectation(state, nqubits, chunk_size):
    value = 0.0
    for qubit in range(nqubits):
        next_qubit = (qubit + 1) % nqubits
        x_flip = 1 << (nqubits - 1 - qubit)
        z_shift = nqubits - 1 - next_qubit
        term = 0.0
        for start in range(0, state.size, chunk_size):
            stop = min(start + chunk_size, state.size)
            indices = np.arange(start, stop, dtype=np.int64)
            z_bit = (indices >> z_shift) & 1
            z_phase = 1 - 2 * z_bit
            term += np.vdot(state[indices ^ x_flip], z_phase * state[start:stop]).real
        value += 0.5 * term
    return float(value)


def default_output_path(nqubits, nlayers, seed):
    return Path("references") / (
        f"qibojit_ring_xz_n{nqubits}_l{nlayers}_seed{seed}.json"
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nqubits", type=int, default=32)
    parser.add_argument("--nlayers", type=int, default=3)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output")
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--allow-large", action="store_true")
    parser.add_argument("--max-state-gb", type=float, default=32.0)
    parser.add_argument("--chunk-size", type=int, default=1 << 20)
    args = parser.parse_args()

    output = Path(args.output) if args.output else default_output_path(
        args.nqubits, args.nlayers, args.seed
    )
    if output.exists() and not args.force:
        with open(output, "r", encoding="utf-8") as f:
            data = json.load(f)
        print(f"loaded {output}")
        print(f"expectation={float(data['expectation']):.16e}")
        return

    state_gb = (2**args.nqubits) * np.dtype(np.complex128).itemsize / (1024**3)
    if state_gb > args.max_state_gb and not args.allow_large:
        raise MemoryError(
            f"Estimated state vector alone is {state_gb:.1f} GiB. "
            "Pass --allow-large after confirming the node has enough memory."
        )

    qibo.set_backend("qibojit")
    circuit = build_circuit(args.nqubits, args.nlayers, args.seed)

    start = time.perf_counter()
    state = circuit().state(numpy=True).reshape(-1)
    expectation = ring_xz_expectation(state, args.nqubits, args.chunk_size)
    elapsed = time.perf_counter() - start

    data = {
        "backend": "qibojit",
        "observable": "0.5 * sum_i X_i Z_((i+1) mod n)",
        "nqubits": args.nqubits,
        "nlayers": args.nlayers,
        "seed": args.seed,
        "expectation": expectation,
        "seconds": elapsed,
        "state_vector_gib_estimate": state_gb,
    }
    output.parent.mkdir(parents=True, exist_ok=True)
    with open(output, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, sort_keys=True)
        f.write("\n")

    print(f"saved {output}")
    print(f"expectation={expectation:.16e}")
    print(f"seconds={elapsed:.3f}")


if __name__ == "__main__":
    main()