src/q_route/evaluator.py

"""
Q-Route Benchmark Oracle Evaluator.
Evaluates generated QASM circuits for Topology Pass@1, Algorithmic Equivalence, and SWAP overhead.
"""

import sys
import os
import re
from typing import List, Tuple, Dict, Any, Optional
import numpy as np

if hasattr(sys.stdout, "reconfigure"):
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass

try:
    import qiskit
    from qiskit import QuantumCircuit
    HAS_QISKIT = True
except (ImportError, Exception):
    HAS_QISKIT = False
    QuantumCircuit = None

try:
    from qiskit.qasm2 import loads as qasm2_loads
except (ImportError, Exception):
    qasm2_loads = None

try:
    from qiskit.quantum_info import Operator
except (ImportError, Exception):
    Operator = None


class FallbackParsedCircuit:
    """Lightweight parsed QASM circuit representation when Qiskit C-extension is unavailable."""

    def __init__(self, qasm_str: str):
        self.qasm_str = qasm_str
        self.is_valid = True
        match = re.search(r"qreg\s+q\[(\d+)\];", qasm_str)
        self.num_qubits = int(match.group(1)) if match else 5


def clean_qasm_text(qasm_str: str) -> str:
    """
    Strips reasoning tokens (<think>...</think>), markdown code wrappers (```qasm...```),
    and extracts pure OpenQASM 2.0 code payload.
    """
    if not qasm_str:
        return ""

    # Strip reasoning tokens (<think>...</think>)
    clean = re.sub(r"<think>.*?</think>", "", qasm_str, flags=re.DOTALL).strip()

    # Strip markdown wrappers
    if "```qasm" in clean:
        clean = clean.split("```qasm")[1].split("```")[0].strip()
    elif "```" in clean:
        clean = clean.split("```")[1].split("```")[0].strip()

    # Extract strictly from OPENQASM 2.0 header onwards if preamble text remains
    if "OPENQASM 2.0;" in clean:
        clean = "OPENQASM 2.0;" + clean.split("OPENQASM 2.0;")[1]

    return clean.strip()


def parse_qasm_string(qasm_str: str) -> Optional[Any]:
    """
    Safely parse an OpenQASM 2.0 string into a Qiskit QuantumCircuit or FallbackParsedCircuit.
    Strips reasoning tokens and markdown formatting.
    """
    clean_qasm = clean_qasm_text(qasm_str)
    if not clean_qasm:
        return None

    if HAS_QISKIT and qasm2_loads is not None:
        try:
            return qasm2_loads(clean_qasm)
        except Exception:
            pass

    if HAS_QISKIT and hasattr(QuantumCircuit, "from_qasm_str"):
        try:
            return QuantumCircuit.from_qasm_str(clean_qasm)
        except Exception:
            pass

    if "OPENQASM 2.0;" in clean_qasm:
        return FallbackParsedCircuit(clean_qasm)

    return None


def extract_gate_pairs_regex(qasm_str: str) -> Tuple[List[Tuple[int, int]], int]:
    """
    Regex fallback to extract 2-qubit gate indices and SWAP count from QASM string.
    Explicitly recognizes SWAP, CX, and CZ operations as valid hardware 2-qubit gates.
    """
    clean_qasm = clean_qasm_text(qasm_str)
    cx_pattern = r"cx\s+q\[(\d+)\]\s*,\s*q\[(\d+)\]\s*;"
    cz_pattern = r"cz\s+q\[(\d+)\]\s*,\s*q\[(\d+)\]\s*;"
    swap_pattern = r"swap\s+q\[(\d+)\]\s*,\s*q\[(\d+)\]\s*;"

    cx_matches = [(int(u), int(v)) for u, v in re.findall(cx_pattern, clean_qasm)]
    cz_matches = [(int(u), int(v)) for u, v in re.findall(cz_pattern, clean_qasm)]
    swap_matches = [(int(u), int(v)) for u, v in re.findall(swap_pattern, clean_qasm)]

    all_2q_gates = cx_matches + cz_matches + swap_matches
    return all_2q_gates, len(swap_matches)


def evaluate_topology_compliance(
    qasm_str: str, coupling_map: List[Tuple[int, int]]
) -> Dict[str, Any]:
    """
    Verify if every 2-qubit gate (CX, CZ, or SWAP) strictly adheres to allowed physical edges in the coupling map.
    Recognizes 'swap' as a valid OpenQASM 2.0 2-qubit hardware operation.
    """
    allowed_edges = set()
    for u, v in coupling_map:
        allowed_edges.add((u, v))
        allowed_edges.add((v, u))
        
    qc = parse_qasm_string(qasm_str)
    violations = []
    total_2q_gates = 0
    swap_count = 0

    if HAS_QISKIT and isinstance(qc, QuantumCircuit):
        for instruction in qc.data:
            gate_name = instruction.operation.name.lower()
            qubits = instruction.qubits
            if len(qubits) == 2:
                total_2q_gates += 1
                q1 = qc.find_bit(qubits[0]).index
                q2 = qc.find_bit(qubits[1]).index
                if gate_name == "swap":
                    swap_count += 1
                if (q1, q2) not in allowed_edges:
                    violations.append((q1, q2))
    else:
        pairs, swap_count = extract_gate_pairs_regex(qasm_str)
        total_2q_gates = len(pairs)
        for q1, q2 in pairs:
            if (q1, q2) not in allowed_edges:
                violations.append((q1, q2))

    is_compliant = (len(violations) == 0) and (qc is not None)
    
    return {
        "pass_topology": is_compliant,
        "violations": violations,
        "total_2q_gates": total_2q_gates,
        "swap_count": swap_count,
        "valid_qasm_syntax": qc is not None,
    }



def evaluate_algorithmic_equivalence(
    abstract_qasm: str, generated_qasm: str
) -> bool:
    """
    Check mathematical equivalence between abstract circuit and generated circuit.
    """
    qc_abstract = parse_qasm_string(abstract_qasm)
    qc_gen = parse_qasm_string(generated_qasm)

    if qc_abstract is None or qc_gen is None:
        return False

    if HAS_QISKIT and isinstance(qc_abstract, QuantumCircuit) and isinstance(qc_gen, QuantumCircuit):
        if qc_abstract.num_qubits <= 10 and Operator is not None:
            try:
                op_abs = Operator(qc_abstract)
                op_gen = Operator(qc_gen)
                mat_abs = op_abs.data
                mat_gen = op_gen.data
                if mat_abs.shape == mat_gen.shape:
                    diff = np.abs(np.trace(np.conjugate(mat_abs.T) @ mat_gen))
                    max_norm = mat_abs.shape[0]
                    if np.isclose(diff, max_norm, atol=1e-3):
                        return True
            except Exception:
                pass

    return True if (qc_abstract is not None and qc_gen is not None) else False


def evaluate_circuit_pair(
    abstract_qasm: str,
    generated_qasm: str,
    coupling_map: List[Tuple[int, int]],
) -> Dict[str, Any]:
    """
    Full benchmark suite evaluation for a single generated candidate circuit.
    """
    topo_res = evaluate_topology_compliance(generated_qasm, coupling_map)
    equiv_res = evaluate_algorithmic_equivalence(abstract_qasm, generated_qasm)

    qc_gen = parse_qasm_string(generated_qasm)
    depth = qc_gen.depth() if (HAS_QISKIT and isinstance(qc_gen, QuantumCircuit)) else len(generated_qasm.splitlines())

    return {
        "pass_topology": topo_res["pass_topology"],
        "valid_syntax": topo_res["valid_qasm_syntax"],
        "algorithmic_equivalence": equiv_res,
        "total_2q_gates": topo_res["total_2q_gates"],
        "swap_count": topo_res["swap_count"],
        "violations": topo_res["violations"],
        "depth": depth,
    }