QuantumCircuits.py

"""
QuantumCircuits.py
Native PyTorch Micro-VM for Heisenberg-picture quantum circuits.

UPDATE 24.4: Bypassed SymPy entirely. This module now compiles circuits 
into a direct, vectorized PyTorch execution graph. Evaluates the BBGKY 
Cumulant Expansion natively on the GPU/CPU for massive speedup and zero crashes.
"""

from __future__ import annotations
import math
import torch
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Any

GATE_TYPES = {"Rx", "Ry", "Rz", "H", "S", "T", "CNOT", "CZ"}
PAULI = ("x", "y", "z")
PAULI2 = [p1 + p2 for p1 in PAULI for p2 in PAULI]

@dataclass
class QGate:
    type:    str
    qubits:  List[int]
    params:  List[str] = field(default_factory=list)
    layer:   int = 0

    def validate(self):
        if self.type not in GATE_TYPES:
            raise ValueError(f"Unknown gate: {self.type}")

@dataclass
class QuantumCircuit:
    n_qubits:             int
    gates:                List[QGate]
    param_bounds:         Dict[str, Tuple[float, float]] = field(default_factory=dict)
    initial_state:        str = "zero"
    initial_obs:          Dict[str, float] = field(default_factory=dict)
    expected_observables: Dict[str, float] = field(default_factory=dict)
    include_2body:        bool = True
    name:                 str = "QuantumCircuit"
    description:          str = ""

    @property
    def n_layers(self) -> int: 
        return max([g.layer for g in self.gates] + [-1]) + 1
        
    @property
    def all_params(self) -> List[str]:
        seen = set()
        params = []
        for g in self.gates:
            for p in g.params:
                if p not in seen: 
                    seen.add(p); params.append(p)
        return params

def obs1(pauli: str, qubit: int, layer: int) -> str:
    return f"{pauli}{qubit}_L{layer}"

def obs2_safe(pauli1: str, q1: int, pauli2: str, q2: int, layer: int) -> str:
    if q1 == q2: raise ValueError(f"Invalid 2-body observable on same qubit: {q1}")
    if q1 > q2: return f"{pauli2}{pauli1}{q2}{q1}_L{layer}"
    return f"{pauli1}{pauli2}{q1}{q2}_L{layer}"

def _initial_obs_values(circuit: QuantumCircuit) -> Dict[str, float]:
    obs = {}
    if circuit.initial_state == "zero":
        for q in range(circuit.n_qubits):
            obs[obs1("z", q, 0)] = 1.0; obs[obs1("x", q, 0)] = 0.0; obs[obs1("y", q, 0)] = 0.0
        if circuit.include_2body:
            for q1 in range(circuit.n_qubits):
                for q2 in range(q1 + 1, circuit.n_qubits):
                    for p1, p2 in PAULI2:
                        obs[obs2_safe(p1, q1, p2, q2, 0)] = 1.0 if (p1 == "z" and p2 == "z") else 0.0
    return obs

# ── PyTorch Native Micro-VM ──────────────────────────────────────────
class NativeQuantumVM:
    """Executes Heisenberg topological constraints directly on PyTorch Tensors."""
    def __init__(self, circuit: QuantumCircuit, var_idx: Dict[str, int]):
        self.circuit = circuit
        self.var_idx = var_idx
        self.n_layers = circuit.n_layers
        self.n = circuit.n_qubits
        self.include_2b = circuit.include_2body

    def idx(self, name: str) -> int:
        return self.var_idx[name]

    def execute(self, X: torch.Tensor, step_ratio: float, device: torch.device) -> torch.Tensor:
        B = X.shape[0] if X.dim() == 2 else 1
        loss = torch.zeros(B, device=device, dtype=torch.float32)
        
        # Helper to smoothly penalize differences (L2 MSE)
        def enforce(actual_idx, expected_tensor):
            actual = X[:, actual_idx] if X.dim() == 2 else X[actual_idx]
            return (actual - expected_tensor) ** 2

        # 1. Bloch Sphere Bounds (Purity)
        for layer in range(self.n_layers + 1):
            for q in range(self.n):
                x, y, z = X[..., self.idx(obs1('x', q, layer))], X[..., self.idx(obs1('y', q, layer))], X[..., self.idx(obs1('z', q, layer))]
                purity = 1.0 - x**2 - y**2 - z**2
                loss += (torch.relu(-purity) ** 2) * 2.0

        # 2. Layer Transitions (Native Tensor Math)
        for layer_in in range(self.n_layers):
            layer_out = layer_in + 1
            layer_gates = [g for g in self.circuit.gates if g.layer == layer_in]
            affected = set()

            for gate in layer_gates:
                if gate.type in ("Rz", "Rx"):
                    q = gate.qubits[0]; affected.add(q)
                    theta = X[..., self.idx(gate.params[0])]
                    c, s = torch.cos(theta), torch.sin(theta)
                    
                    xi, yi, zi = X[..., self.idx(obs1('x', q, layer_in))], X[..., self.idx(obs1('y', q, layer_in))], X[..., self.idx(obs1('z', q, layer_in))]
                    
                    if gate.type == "Rz":
                        loss += enforce(self.idx(obs1('x', q, layer_out)), c * xi - s * yi) * 5.0
                        loss += enforce(self.idx(obs1('y', q, layer_out)), s * xi + c * yi) * 5.0
                        loss += enforce(self.idx(obs1('z', q, layer_out)), zi) * 5.0
                    else: # Rx
                        loss += enforce(self.idx(obs1('x', q, layer_out)), xi) * 5.0
                        loss += enforce(self.idx(obs1('y', q, layer_out)), c * yi - s * zi) * 5.0
                        loss += enforce(self.idx(obs1('z', q, layer_out)), s * yi + c * zi) * 5.0

                elif gate.type == "H":
                    q = gate.qubits[0]; affected.add(q)
                    loss += enforce(self.idx(obs1('x', q, layer_out)), X[..., self.idx(obs1('z', q, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs1('y', q, layer_out)), -X[..., self.idx(obs1('y', q, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs1('z', q, layer_out)), X[..., self.idx(obs1('x', q, layer_in))]) * 5.0

                elif gate.type == "CNOT":
                    k, j = gate.qubits[0], gate.qubits[1]; affected.update([k, j])
                    # 1-body to 2-body Entanglement
                    loss += enforce(self.idx(obs1('x', k, layer_out)), X[..., self.idx(obs2_safe('x', k, 'x', j, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs1('y', k, layer_out)), X[..., self.idx(obs2_safe('y', k, 'x', j, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs1('z', k, layer_out)), X[..., self.idx(obs1('z', k, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs1('x', j, layer_out)), X[..., self.idx(obs1('x', j, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs1('y', j, layer_out)), X[..., self.idx(obs2_safe('z', k, 'y', j, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs1('z', j, layer_out)), X[..., self.idx(obs2_safe('z', k, 'z', j, layer_in))]) * 5.0
                    
                    # 2-body internal mixing
                    loss += enforce(self.idx(obs2_safe('x', k, 'x', j, layer_out)), X[..., self.idx(obs1('x', k, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs2_safe('z', k, 'z', j, layer_out)), X[..., self.idx(obs1('z', j, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs2_safe('x', k, 'z', j, layer_out)), -X[..., self.idx(obs2_safe('y', k, 'y', j, layer_in))]) * 5.0
                    loss += enforce(self.idx(obs2_safe('y', k, 'z', j, layer_out)), X[..., self.idx(obs2_safe('x', k, 'y', j, layer_in))]) * 5.0

                    # NATIVE BBGKY CUMULANT EXPANSION (Spectators)
                    for r in range(self.n):
                        if r in (k, j): continue
                        if self.include_2b:
                            for pr in PAULI:
                                # Example: X_k P_r -> X_k X_j P_r (3-body)
                                # Cumulant: <Xk Xj><Pr> + <Xk Pr><Xj> + <Xj Pr><Xk> - 2<Xk><Xj><Pr>
                                o_k_j = X[..., self.idx(obs2_safe('x', k, 'x', j, layer_in))]
                                o_k_r = X[..., self.idx(obs2_safe('x', k, pr, r, layer_in))]
                                o_j_r = X[..., self.idx(obs2_safe('x', j, pr, r, layer_in))]
                                o_k = X[..., self.idx(obs1('x', k, layer_in))]
                                o_j = X[..., self.idx(obs1('x', j, layer_in))]
                                o_r = X[..., self.idx(obs1(pr, r, layer_in))]
                                
                                c3 = (o_k_j * o_r) + (o_k_r * o_j) + (o_j_r * o_k) - (2.0 * o_k * o_j * o_r)
                                loss += enforce(self.idx(obs2_safe('x', k, pr, r, layer_out)), c3) * 4.0

            # Identity for unaffected qubits
            unaffected = [q for q in range(self.n) if q not in affected]
            for q in unaffected:
                for p in PAULI:
                    loss += enforce(self.idx(obs1(p, q, layer_out)), X[..., self.idx(obs1(p, q, layer_in))]) * 5.0
                if self.include_2b:
                    for q2 in range(self.n):
                        if q2 == q or q2 in affected: continue
                        for p1p2 in PAULI2:
                            loss += enforce(self.idx(obs2_safe(p1p2[0], q, p1p2[1], q2, layer_out)), X[..., self.idx(obs2_safe(p1p2[0], q, p1p2[1], q2, layer_in))]) * 5.0

        return loss.squeeze()

# ── Main builder ──────────────────────────────────────────────────────
def build_quantum_axl(circuit: QuantumCircuit, axl_problem_class: Any, axl_invariant_class: Any) -> Any:
    for g in circuit.gates: g.validate()
        
    n = circuit.n_qubits
    n_layers = circuit.n_layers
    include_2b = circuit.include_2body

    variables = []
    for pname in circuit.all_params:
        lo, hi = circuit.param_bounds.get(pname, (-math.pi, math.pi))
        variables.append({"name": pname, "lo": lo, "hi": hi})
        
    for layer in range(n_layers + 1):
        for q in range(n):
            for p in PAULI: 
                variables.append({"name": obs1(p, q, layer), "lo": -1.0, "hi": 1.0})
        if include_2b:
            for q1 in range(n):
                for q2 in range(q1 + 1, n):
                    for p1p2 in PAULI2: 
                        variables.append({"name": obs2_safe(p1p2[0], q1, p1p2[1], q2, layer), "lo": -1.0, "hi": 1.0})

    obs_fixed = _initial_obs_values(circuit)
    
    anchors = []
    for obs, val in circuit.expected_observables.items():
        if len(obs) == 2:   expr = obs1(obs[0], int(obs[1:]), n_layers)
        elif len(obs) == 4: expr = obs2_safe(obs[0], int(obs[2]), obs[1], int(obs[3]), n_layers)
        anchors.append(axl_invariant_class(name=f"verify_{obs}", expr=f"{expr} - ({val})", tolerance=0.05, mode="eq"))

    # Empty constraints/scopes because Native PyTorch VM will handle evaluation
    prob = axl_problem_class(
        name=circuit.name, description=circuit.description,
        axioms={"CONTINUOUS", "CONSERVED", "TRANSITIVE", "BILINEAR", "METRIC", "SUPERPOSITION", "SYMMETRIC"},
        variables=variables, constraints=[], scopes=[], anchors=anchors, observations=obs_fixed
    )
    
    prob.is_quantum = True
    prob.circuit = circuit
    return prob

# ══════════════════════════════════════════════════════════════════════
# EXAMPLE CIRCUITS
# ══════════════════════════════════════════════════════════════════════

def example_bell_state() -> QuantumCircuit:
    return QuantumCircuit(
        n_qubits=2, name="BellStatePrep",
        gates=[QGate("H", [0], layer=0), QGate("CNOT", [0, 1], layer=1)],
        initial_state="zero", include_2body=True,
        expected_observables={"zz01": 1.0, "xx01": 1.0} 
    )

def example_vqe_h2() -> QuantumCircuit:
    return QuantumCircuit(
        n_qubits=2, name="VQE_H2",
        gates=[
            QGate("Ry", [0], ["theta0"], layer=0),
            QGate("Ry", [1], ["theta1"], layer=0),
            QGate("CNOT", [0, 1], layer=1),
        ],
        param_bounds={"theta0": (-math.pi, math.pi), "theta1": (-math.pi, math.pi)},
        expected_observables={"z0": -0.5, "z1": -0.5, "zz01": 0.25, "xx01": -0.5} 
    )

def example_qaoa_maxcut_p1() -> QuantumCircuit:
    return QuantumCircuit(
        n_qubits=3, name="QAOA_MaxCut_Deep",
        gates=[
            QGate("H", [0], layer=0), QGate("H", [1], layer=0), QGate("H", [2], layer=0),
            QGate("CNOT", [0, 1], layer=1), QGate("Rz", [1], ["gamma"], layer=2), QGate("CNOT", [0, 1], layer=3),
            QGate("CNOT", [1, 2], layer=4), QGate("Rz", [2], ["gamma"], layer=5), QGate("CNOT", [1, 2], layer=6),
            QGate("CNOT", [0, 2], layer=7), QGate("Rz", [2], ["gamma"], layer=8), QGate("CNOT", [0, 2], layer=9),
            QGate("Rx", [0], ["beta"], layer=10), QGate("Rx", [1], ["beta"], layer=10), QGate("Rx", [2], ["beta"], layer=10),
        ],
        param_bounds={"gamma": (0, math.pi), "beta": (0, math.pi / 2)},
        expected_observables={"zz01": -0.5, "zz12": -0.5, "zz02": -0.5}
    )