#!/usr/bin/env python
from __future__ import annotations
import argparse
from itertools import repeat
import json
import math
import random
from dataclasses import dataclass
from enum import IntEnum
from multiprocessing import Pool
from typing import override
import numpy as np
from qiskit import QuantumCircuit as QiskitCircuit, transpile
from qiskit.circuit import ParameterVector, ParameterVectorElement
from qiskit_aer import AerSimulator
from tqdm import tqdm

from qas_flow import Stream

# ----------------------------
# Paper hyperparameters/defaults
# ----------------------------

TWO_QUBIT_GATE_PROBABILITY = 0.5
CHEMICAL_ACCURACY = 1.6e-3  # Ha, success if E - E0 <= this

# ----------------------------
# Circuit representation
# ----------------------------


class GateType(IntEnum):
    RX = 1
    RY = 2
    RZ = 3
    XX = 4
    YY = 5
    ZZ = 6


@dataclass(frozen=True)
class Gate:
    type: GateType
    qubits: int | tuple[int, int]
    param_idx: int

    def to_json(self):
        return {
            "type": int(self.type),
            "qubits": self.qubits,
            "param_idx": self.param_idx,
        }


def haar_fidelity_pdf(F: np.ndarray, d: int) -> np.ndarray:
    # p(F) = (d-1) * (1-F)^(d-2), for F in [0,1]
    return (d - 1) * np.power(1.0 - F, d - 2)


def kl_divergence(p: np.ndarray, q: np.ndarray) -> float:
    # assumes p, q are normalized and strictly positive
    return float(np.sum(p * np.log(p / q)))


@dataclass()
class QuantumCircuit:
    qubits: int
    gates: list[list[Gate]]
    single_qubit_gates: int
    two_qubit_gates: int
    params: int
    paths: int = 0
    expressibility: float = float("-inf")

    # ground-truth fields (filled later)
    gt_energy: float | None = None
    gt_error: float | None = None
    gt_steps: int | None = None
    gt_success: bool | None = None

    def calculate_paths(self):
        path_counts = [1 for _ in range(self.qubits)]
        for layer in self.gates:
            for gate in layer:
                if gate.type <= 3 or isinstance(gate.qubits, int):
                    continue
                (q1, q2) = gate.qubits
                # same logic as your existing file
                path_counts[q1] = path_counts[q1] + path_counts[q2]
                path_counts[q2] = path_counts[q1]
        self.paths = sum(path_counts)

    def to_json(self):
        return {
            "qubits": self.qubits,
            "gates": [[gate.to_json() for gate in layer] for layer in self.gates],
            "params": self.params,
            "paths": self.paths,
            "expressibility": self.expressibility,
            "gt_energy": self.gt_energy,
            "gt_error": self.gt_error,
            "gt_steps": self.gt_steps,
            "gt_success": self.gt_success,
        }

    def to_qiskit_ansatz(self) -> tuple[QiskitCircuit, ParameterVector]:
        """Return *just* the parametrized ansatz defined by this circuit."""
        qc = QiskitCircuit(self.qubits)
        thetas = ParameterVector("theta", self.params)

        for layer in self.gates:
            for gate in layer:
                theta = thetas[gate.param_idx]
                if gate.type == GateType.RX:
                    qc.rx(theta, gate.qubits)
                elif gate.type == GateType.RY:
                    qc.ry(theta, gate.qubits)
                elif gate.type == GateType.RZ:
                    qc.rz(theta, gate.qubits)
                elif gate.type == GateType.XX:
                    qc.rxx(theta, *gate.qubits)
                elif gate.type == GateType.YY:
                    qc.ryy(theta, *gate.qubits)
                elif gate.type == GateType.ZZ:
                    qc.rzz(theta, *gate.qubits)
                else:
                    raise ValueError(f"Unknown gate type: {gate.type}")
        return qc, thetas

    def to_qiskit_for_expressibility(self) -> tuple[QiskitCircuit, ParameterVector]:
        """Expressibility uses |0...0> input (no TFIM-specific H layer)."""
        qc, thetas = self.to_qiskit_ansatz()
        qc.save_statevector()
        return qc, thetas

    def to_qiskit_for_tfim_vqe(self) -> tuple[QiskitCircuit, ParameterVector]:
        """TFIM VQE begins with a layer of Hadamards on all qubits."""
        ans, thetas = self.to_qiskit_ansatz()
        qc = QiskitCircuit(self.qubits)
        qc.h(range(self.qubits))
        qc.compose(ans, inplace=True)
        qc.save_statevector()
        return qc, thetas

    def expressibility_estimate(
        self, samples: int, seed: int, bins: int = 75, eps: float = 1e-12
    ) -> "QuantumCircuit":
        qc, thetas = self.to_qiskit_for_expressibility()

        if self.params <= 0:
            self.expressibility = float("-inf")
            return self

        rng = random.Random(seed)
        d = 1 << self.qubits

        backend = AerSimulator(method="statevector", seed_simulator=seed)
        tqc = transpile(qc, backend, optimization_level=0)

        nexp = 2 * samples
        # vectorized binds: one circuit, many parameter values
        binds: list[dict[ParameterVectorElement, list[float]]] = [
            {
                param: [rng.random() * 4 * math.pi - 2 * math.pi for _ in range(nexp)]
                for param in thetas.params
            }
        ]

        job = backend.run([tqc], parameter_binds=binds)
        result = job.result()

        sv = [
            np.asarray(result.get_statevector(i), dtype=np.complex128)
            for i in range(nexp)
        ]
        left = sv[:samples]
        right = sv[samples:]

        inners = np.array(
            [np.vdot(l, r) for l, r in zip(left, right)], dtype=np.complex128
        )
        F = (inners.conjugate() * inners).real

        hist, edges = np.histogram(F, bins=bins, range=(0.0, 1.0), density=False)
        p = hist.astype(np.float64) + eps
        p = p / p.sum()

        mids = 0.5 * (edges[:-1] + edges[1:])
        widths = edges[1:] - edges[:-1]
        q = haar_fidelity_pdf(mids, d) * widths
        q = q + eps
        q = q / q.sum()

        kl = kl_divergence(p, q)
        self.expressibility = -kl
        return self

    @override
    def __str__(self):
        strs = ["-" for _ in range(self.qubits)]
        for layer in self.gates:
            for i in range(self.qubits):
                strs[i] += "---"

            idx = 0
            for gate in layer:
                match gate.type:
                    case GateType.RX:
                        strs[gate.qubits] = strs[gate.qubits][:-2] + "RX"
                    case GateType.RY:
                        strs[gate.qubits] = strs[gate.qubits][:-2] + "RY"
                    case GateType.RZ:
                        strs[gate.qubits] = strs[gate.qubits][:-2] + "RZ"
                    case GateType.XX:
                        (q1, q2) = gate.qubits
                        strs[q1] = strs[q1][:-2] + f"X{idx}"
                        strs[q2] = strs[q2][:-2] + f"X{idx}"
                        idx += 1
                    case GateType.YY:
                        (q1, q2) = gate.qubits
                        strs[q1] = strs[q1][:-2] + f"Y{idx}"
                        strs[q2] = strs[q2][:-2] + f"Y{idx}"
                        idx += 1
                    case GateType.ZZ:
                        (q1, q2) = gate.qubits
                        strs[q1] = strs[q1][:-2] + f"Z{idx}"
                        strs[q2] = strs[q2][:-2] + f"Z{idx}"
                        idx += 1
        gt = ""
        if self.gt_energy is not None:
            gt = f"\nGT: E={self.gt_energy:.12f}, err={self.gt_error:.3e}, steps={self.gt_steps}, success={self.gt_success}"
        return (
            "\n".join(strs)
            + f"\npaths: {self.paths}, expressibility: {self.expressibility}"
            + gt
        )


# ----------------------------
# Search space sampling (layerwise)
# ----------------------------


def even_parity(qubits: int):
    return [(x, x + 1) for x in range(0, qubits, 2)]


def odd_parity(qubits: int):
    return [(x, x + 1) for x in range(1, qubits, 2)]


def sample_circuit(rng: random.Random, qubits: int, depth: int) -> QuantumCircuit:
    even = even_parity(qubits)
    odd = odd_parity(qubits)

    total_single = 0
    total_double = 0
    params = 0

    gates: list[list[Gate]] = []
    for _ in range(depth):
        if total_single + total_double >= 36:
            break
        gate_type_offset = 3 if rng.random() < TWO_QUBIT_GATE_PROBABILITY else 0
        gate_type = rng.randint(1, 3)
        gate_locations = even if rng.random() < 0.5 else odd

        if gate_type_offset == 0:
            gates.append(
                [Gate(GateType(gate_type), x, params) for (x, _) in gate_locations]
            )
            total_single += len(gate_locations)
        else:
            gates.append(
                [
                    Gate(GateType(gate_type + gate_type_offset), xy, params)
                    for xy in gate_locations
                    if xy[1] != qubits
                ]
            )
            total_double += len(gate_locations)
        params += 1

    return QuantumCircuit(qubits, gates, total_single, total_double, params)


def sample_generator(rng: random.Random, qubits: int, depth: int):
    while True:
        yield sample_circuit(rng, qubits, depth)


def more_single_than_double(qc: QuantumCircuit) -> bool:
    return qc.single_qubit_gates >= qc.two_qubit_gates


# ----------------------------
# TFIM ground-truth VQE
# ----------------------------


def _pauli_x():
    return np.array([[0.0, 1.0], [1.0, 0.0]], dtype=np.float64)


def _pauli_z():
    return np.array([[1.0, 0.0], [0.0, -1.0]], dtype=np.float64)


def _eye():
    return np.eye(2, dtype=np.float64)


def kron_n(ops: list[np.ndarray]) -> np.ndarray:
    out = ops[0]
    for op in ops[1:]:
        out = np.kron(out, op)
    return out


def tfim_hamiltonian_matrix(n: int, periodic: bool = True) -> np.ndarray:
    """
    HTFIM = sum_{i=1}^n Z_i Z_{i+1} + X_i  (paper definition) :contentReference[oaicite:5]{index=5}
    with Z_{n+1}=Z_1 if periodic.
    """
    X = _pauli_x()
    Z = _pauli_z()
    I = _eye()
    dim = 1 << n
    H = np.zeros((dim, dim), dtype=np.float64)

    # Xi terms
    for i in range(n):
        ops = [I] * n
        ops[i] = X
        H += kron_n(ops)

    # Zi Zi+1 terms
    for i in range(n):
        j = (i + 1) % n
        if (not periodic) and (j == 0):
            continue
        ops = [I] * n
        ops[i] = Z
        ops[j] = Z
        H += kron_n(ops)

    return H


def exact_ground_energy(H: np.ndarray) -> float:
    # H is real symmetric here
    w = np.linalg.eigvalsh(H)
    return float(w[0])


def statevector_from_bound_circuit(
    backend: AerSimulator, tqc: QiskitCircuit, bind: dict
) -> np.ndarray:
    res = backend.run([tqc], parameter_binds=[bind]).result()
    sv = np.asarray(res.get_statevector(0), dtype=np.complex128)
    return sv


def energy_expectation_from_sv(H: np.ndarray, sv: np.ndarray) -> float:
    # E = <psi|H|psi>
    # H real symmetric; sv complex
    return float(np.real(np.vdot(sv, H @ sv)))


def adam_optimize_tfim_energy(
    circuit: QuantumCircuit,
    H: np.ndarray,
    seed: int,
    lr: float = 0.01,  # paper :contentReference[oaicite:6]{index=6}
    max_steps: int = 2000,
    tol: float = 1e-7,
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-8,
    param_shift: float = math.pi / 2,
) -> tuple[float, int]:
    """
    Optimize parameters with Adam using parameter-shift gradients.
    Returns: (best_energy, steps_taken)
    """
    qc, thetas = circuit.to_qiskit_for_tfim_vqe()

    p = circuit.params
    rng = np.random.default_rng(seed)
    theta = rng.uniform(
        -2 * math.pi, 2 * math.pi, size=p
    )  # paper init range :contentReference[oaicite:7]{index=7}

    backend = AerSimulator(method="statevector", seed_simulator=seed)
    tqc = transpile(qc, backend, optimization_level=0)

    def E(th: np.ndarray) -> float:
        bind = {param: [val] for param, val in zip(thetas.params, th)}
        sv = statevector_from_bound_circuit(backend, tqc, bind)
        return energy_expectation_from_sv(H, sv)

    def grad(th: np.ndarray) -> np.ndarray:
        # parameter-shift: d/dθ f(θ) = 0.5*(f(θ+π/2)-f(θ-π/2))
        g = np.zeros_like(th)
        for i in range(p):
            plus = th.copy()
            minus = th.copy()
            plus[i] += param_shift
            minus[i] -= param_shift
            g[i] = 0.5 * (E(plus) - E(minus))
        return g

    m = np.zeros_like(theta)
    v = np.zeros_like(theta)

    best_E = float("inf")
    prev_E = None

    for t in range(1, max_steps + 1):
        g = grad(theta)
        m = beta1 * m + (1 - beta1) * g
        v = beta2 * v + (1 - beta2) * (g * g)
        mhat = m / (1 - beta1**t)
        vhat = v / (1 - beta2**t)
        theta = theta - lr * mhat / (np.sqrt(vhat) + eps)

        cur_E = E(theta)
        if cur_E < best_E:
            best_E = cur_E

        if prev_E is not None and abs(prev_E - cur_E) < tol:
            return best_E, t
        prev_E = cur_E

    return best_E, max_steps


def ground_truth_tfim(
    circ: QuantumCircuit,
    H: np.ndarray,
    E0: float,
    seed: int,
    lr: float,
    max_steps: int,
    tol: float,
) -> QuantumCircuit:
    best_E, steps = adam_optimize_tfim_energy(
        circ, H, seed=seed, lr=lr, max_steps=max_steps, tol=tol
    )
    err = best_E - E0
    circ.gt_energy = best_E
    circ.gt_error = err
    circ.gt_steps = steps
    circ.gt_success = err <= CHEMICAL_ACCURACY
    return circ


# ----------------------------
# Main
# ----------------------------


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--qubits",
        type=int,
        default=6,
        help="TFIM qubit count (paper uses 6). :contentReference[oaicite:8]{index=8}",
    )
    ap.add_argument(
        "--depth",
        type=int,
        default=10,
        help="Max depth for TFIM layerwise (paper max depth 10). :contentReference[oaicite:9]{index=9}",
    )
    ap.add_argument(
        "--sample_amount",
        type=int,
        default=50000,
        help="Number of sampled circuits (paper uses 50,000). :contentReference[oaicite:10]{index=10}",
    )
    ap.add_argument(
        "--expressibility_samples",
        type=int,
        default=2000,
        help="Fidelity samples for expressibility (paper uses 2000). :contentReference[oaicite:11]{index=11}",
    )
    ap.add_argument(
        "--proxy_pass_amount",
        type=int,
        default=5000,
        help="R: top circuits kept by paths (paper uses 5000). :contentReference[oaicite:12]{index=12}",
    )
    ap.add_argument(
        "--topk",
        type=int,
        default=100,
        help="K: how many top circuits to output/store.",
    )
    ap.add_argument(
        "--seed", type=int, default=0, help="RNG seed for sampling/reproducibility."
    )
    ap.add_argument(
        "--periodic",
        action="store_true",
        help="Use periodic TFIM boundary (default true if set).",
    )
    ap.add_argument(
        "--no-periodic",
        dest="periodic",
        action="store_false",
        help="Use open boundary TFIM.",
    )
    ap.set_defaults(periodic=True)

    # ground-truth options
    ap.add_argument(
        "--do_ground_truth",
        action="store_true",
        help="Also evaluate ground-truth TFIM VQE performance.",
    )
    ap.add_argument(
        "--gt_budget",
        type=int,
        default=100,
        help="Max number of circuits to query for ground-truth (paper often reports within 100 queries). :contentReference[oaicite:13]{index=13}",
    )
    ap.add_argument(
        "--gt_lr",
        type=float,
        default=0.01,
        help="Adam learning rate (paper uses 0.01). :contentReference[oaicite:14]{index=14}",
    )
    ap.add_argument(
        "--gt_max_steps",
        type=int,
        default=2000,
        help="Max Adam steps per circuit (practical cap).",
    )
    ap.add_argument(
        "--gt_tol",
        type=float,
        default=1e-7,
        help="Convergence tolerance for |E_t - E_{t-1}|.",
    )
    ap.add_argument("--dump", type=str, default="dump.json", help="Output JSON file.")

    return ap.parse_args()


def expr_worker(payload):
    circ, seed, samples = payload
    return circ.expressibility_estimate(samples, seed)


def main():
    args = parse_args()
    rng = random.Random(args.seed)

    # --- Stage 0: sample circuits
    circuits = (
        Stream(sample_generator(rng, args.qubits, args.depth))
        .filter(more_single_than_double)
        .take(args.sample_amount)
        .collect()
    )

    # --- Stage 1: paths
    for c in tqdm(circuits, desc="paths"):
        c.calculate_paths()
    circuits.sort(key=lambda qc: qc.paths, reverse=True)
    candidates = circuits[: args.proxy_pass_amount]

    # --- Stage 2: expressibility
    seeds = [rng.randint(0, 1_000_000_000) for _ in range(len(candidates))]

    final_circuits: list[QuantumCircuit] = []
    with Pool() as p:
        for circ in tqdm(
            p.imap_unordered(
                expr_worker, zip(candidates, seeds, repeat(args.expressibility_samples))
            ),
            total=len(candidates),
            desc="expressibility",
        ):
            final_circuits.append(circ)

    final_circuits.sort(key=lambda qc: qc.expressibility, reverse=True)

    # --- Ground-truth TFIM (optional): query from best to worst until success or budget
    success_idx = None
    if args.do_ground_truth:
        H = tfim_hamiltonian_matrix(args.qubits, periodic=args.periodic)
        E0 = exact_ground_energy(H)

        print(f"\nTFIM exact ground energy E0 = {E0:.12f}  (periodic={args.periodic})")

        gt_seed_stream = random.Random(args.seed + 1234567)
        queried = 0
        min_error = float("inf")
        for i, circ in enumerate(tqdm(final_circuits, desc="ground-truth queries")):
            if queried >= args.gt_budget:
                break
            seed = gt_seed_stream.randint(0, 1_000_000_000)
            ground_truth_tfim(
                circ,
                H,
                E0,
                seed=seed,
                lr=args.gt_lr,
                max_steps=args.gt_max_steps,
                tol=args.gt_tol,
            )
            if circ.gt_error is not None and circ.gt_error < min_error:
                print(f"new best error for {queried}: {circ.gt_error}")
            queried += 1
            if circ.gt_success:
                success_idx = i
                break

        if success_idx is None:
            print(f"\nNo circuit hit chemical accuracy within budget={args.gt_budget}.")
        else:
            print(f"\nFirst success at rank {success_idx} (0-based) within budget.")
            print(final_circuits[success_idx])

    # Print bottom few (worst) and top few (best) for sanity
    print("\nTop circuits (best expressibility):")
    for i, c in enumerate(final_circuits[: min(args.topk, 10)]):
        print(f"\nrank {i}:\n{c}")

    # Dump top-K (or all if you want)
    out = final_circuits[: args.topk]
    with open(args.dump, "w", encoding="utf-8") as fp:
        json.dump([c.to_json() for c in out], fp, indent=2)

    print(f"\nWrote {len(out)} circuits to {args.dump}.")


if __name__ == "__main__":
    main()