#!/usr/bin/env python3
"""
Nativ3 Topology-Aware Quantum Compiler v2
==========================================
Uses transfer matrix theory (25 theorems) with three optimization axes:
  r: radial (spectral) — minimize amplitude decay
  θ: angular (geometric) — minimize channel asymmetry
  E: exchange coupling — track symmetry breaking between consecutive gates

Replaces SWAP-count heuristics with exact topology cost.

Usage:
    from nativ3_compiler import compile_circuit, exchange_coupling
    mapping, Z, metrics = compile_circuit(circuit, hardware, calibration)
"""

import numpy as np
pi = np.pi

# ═══ EXCHANGE OPERATOR ═══
J = np.array([[0,0,0,1],[0,0,1,0],[0,1,0,0],[1,0,0,0]], dtype=complex)

def make_M(alpha, U):
    a,b,c,d = U[0,0],U[0,1],U[1,0],U[1,1]
    ca, sa = np.cos(alpha), np.sin(alpha)
    r0 = a*ca + b*sa; r1 = c*ca + d*sa
    return np.array([[ca**2,sa**2,0,0],[ca*r0,sa*r1,0,0],
                     [0,0,ca*np.conj(r0),sa*np.conj(r1)],
                     [0,0,abs(r0)**2,abs(r1)**2]])

def make_base(beta, U):
    a,b,c,d = U[0,0],U[0,1],U[1,0],U[1,1]
    cb, sb = np.cos(beta), np.sin(beta)
    xi = a*cb**2 + (b+c)*cb*sb + d*sb**2
    return np.array([1, xi, np.conj(xi), 1], dtype=complex)

def readout(alpha, state, U):
    a,b,c,d = U[0,0],U[0,1],U[1,0],U[1,1]
    ca, sa = np.cos(alpha), np.sin(alpha)
    r0 = a*ca + b*sa; r1 = c*ca + d*sa
    O0 = ca**2*state[0] + sa**2*state[1]
    O1 = abs(r0)**2*state[2] + abs(r1)**2*state[3]
    return O0, O1

def fidelity_to_alpha(F):
    """T2 inverted: F = cos^4(α) → α = arccos(F^(1/4))"""
    return np.arccos(np.clip(F, 1e-10, 1.0) ** 0.25)

def exchange_coupling(M):
    """T25c: normalized distance from exchange-invariance."""
    JMJ = J @ M @ J
    norm_M = np.sqrt(np.sum(np.abs(M)**2))
    if norm_M < 1e-15: return 0.0
    return float(np.sqrt(np.sum(np.abs(JMJ - M)**2)) / norm_M)

def is_exchange_invariant(U):
    """T25a: True iff U is anti-diagonal (a=d=0)."""
    return abs(U[0,0]) < 1e-10 and abs(U[1,1]) < 1e-10

def compute_Z_rte(circuit, mapping, hw, cal, gate_map):
    """Compute Z with full r⊗θ⊗E metrics."""
    Z_total = 1.0; r_total = 1.0
    theta_sum = 0.0; n_edges = 0
    E_sum = 0.0; n_pairs = 0
    prev_M = None

    for ctrl, targ, gate_type in circuit:
        U = gate_map.get(gate_type, np.array([[0,1],[1,0]], dtype=complex))
        hw_c, hw_t = mapping[ctrl], mapping[targ]
        dist, path = _bfs(hw, hw_c, hw_t)

        alpha_avg = np.mean([fidelity_to_alpha(cal.get(q,{}).get("fidelity",0.99))
                             for q in path])
        M_curr = make_M(alpha_avg, U)

        if dist <= 1:
            f = cal.get(hw_c,{}).get("fidelity",0.99) * cal.get(hw_t,{}).get("fidelity",0.99)
            Z_total *= f
        else:
            beta = fidelity_to_alpha(cal.get(hw_t,{}).get("fidelity",0.99))
            base = make_base(beta, U); state = base.copy()
            for q in path[1:-1]:
                aq = fidelity_to_alpha(cal.get(q,{}).get("fidelity",0.99))
                state = make_M(aq, U) @ state

            O0, O1 = readout(alpha_avg, state, U)
            Z_val = 0.5*(O0+O1); Delta = 0.5*(O0-O1)
            z = abs(Z_val)
            r = np.sqrt(abs(Z_val)**2 + abs(Delta)**2)
            theta = np.arctan2(abs(Delta), abs(Z_val))

            Z_total *= z; r_total *= r
            theta_sum += theta; n_edges += 1

        if prev_M is not None:
            E_sum += exchange_coupling(M_curr @ prev_M)
            n_pairs += 1
        prev_M = M_curr

    return Z_total, {
        'r': r_total,
        'theta': theta_sum / max(1, n_edges),
        'E': E_sum / max(1, n_pairs) if n_pairs > 0 else 0.0,
    }

def _bfs(hw, q1, q2):
    if q1 == q2: return 0, [q1]
    visited = {q1}; queue = [(q1, [q1])]
    while queue:
        node, path = queue.pop(0)
        for nb in hw[node]:
            if nb == q2: return len(path), path + [nb]
            if nb not in visited:
                visited.add(nb); queue.append((nb, path + [nb]))
    return float('inf'), []

def compile_circuit(circuit, hw, cal, gate_map=None, n_samples=10000,
                    E_weight=0.1):
    """
    Nativ3 r⊗θ⊗E compiler.
    
    Args:
        circuit: list of (control, target, gate_type) tuples
        hw: {qubit: [neighbors]} connectivity dict
        cal: {qubit: {"fidelity": float}} calibration dict
        gate_map: {gate_type: 2x2 unitary} mapping
        n_samples: number of random mappings to try
        E_weight: weight of exchange coupling in score (0=ignore, 1=equal to Z)
    
    Returns:
        (mapping, Z, metrics) where metrics = {r, theta, E, swaps}
    """
    if gate_map is None:
        gate_map = {
            "cx": np.array([[0,1],[1,0]], dtype=complex),
            "cz": np.array([[1,0],[0,-1]], dtype=complex),
            "cy": np.array([[0,-1j],[1j,0]], dtype=complex),
            "ch": np.array([[1,1],[1,-1]], dtype=complex)/np.sqrt(2),
            "cs": np.array([[1,0],[0,1j]], dtype=complex),
            "ct": np.array([[1,0],[0,np.exp(1j*pi/4)]], dtype=complex),
        }

    n_log = max(max(c,t) for c,t,_ in circuit) + 1
    phys = list(hw.keys())

    best_score = -1; best_map = None; best_Z = 0; best_metrics = {}

    for _ in range(n_samples):
        sel = list(np.random.choice(phys, n_log, replace=False))
        mapping = dict(enumerate(sel))

        Z, metrics = compute_Z_rte(circuit, mapping, hw, cal, gate_map)
        score = Z * (1 + E_weight * (1 - metrics['E']))

        if score > best_score:
            swaps = sum(max(0, _bfs(hw, mapping[c], mapping[t])[0]-1)
                       for c,t,_ in circuit)
            best_score = score; best_map = mapping; best_Z = Z
            best_metrics = {**metrics, 'swaps': swaps}

    return best_map, best_Z, best_metrics


if __name__ == "__main__":
    print("Nativ3 Compiler v2 — 25 theorems, r⊗θ⊗E optimization")
    print("Usage: from nativ3_compiler import compile_circuit, exchange_coupling")
