# ===== GOOGLE COLAB - CELULA COMPLETA (FARA ERORI) =====
# CHSH Bell Inequality Test pe IBM Quantum Hardware

# ===== INSTALARE PACHETE =====
print("📦 Instalare pachete necesare...")
!pip install -q qiskit qiskit-ibm-runtime matplotlib pylatexenc

# ===== IMPORT-URI =====
import os, math, json, sys
import getpass
from pathlib import Path
from datetime import datetime, timezone

from qiskit import QuantumCircuit
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit_ibm_runtime import QiskitRuntimeService
from qiskit_ibm_runtime import SamplerV2 as Sampler

from IPython.display import display, Image, JSON

print("✓ Pachete instalate cu succes!\n")

# ===== SETUP DIRECTOARE =====
OUTDIR = Path("/content/chsh_outputs")
OUTDIR.mkdir(exist_ok=True)

# ===== FUNCTII AUXILIARE =====
def dump_qasm(circ: QuantumCircuit) -> str:
    """Export QASM3 dacă e disponibil, altfel fallback la QASM2."""
    try:
        from qiskit.qasm3 import dumps as qasm3_dumps
        return qasm3_dumps(circ)
    except Exception:
        return circ.qasm()

def save_artifacts(circ: QuantumCircuit, stem: str):
    """Salvează circuitul ca PNG și QASM."""
    try:
        circ.draw("mpl").savefig(OUTDIR / f"{stem}.png", dpi=180, bbox_inches="tight")
    except Exception:
        pass
    try:
        (OUTDIR / f"{stem}.qasm").write_text(dump_qasm(circ), encoding="utf-8")
    except Exception:
        pass

def chsh_rot_to_z(qc: QuantumCircuit, q: int, angle: float):
    """Rotație pentru măsurare în bază rotită (planul X-Z)."""
    qc.ry(-2.0 * angle, q)

def build_entangled_chsh_circuit(a_angle: float, b_angle: float, name="E") -> QuantumCircuit:
    """
    Circuit ENTANGLED: |Φ+⟩ = (|00⟩ + |11⟩)/√2
    H → CNOT → Rotații → Măsurare
    """
    qc = QuantumCircuit(2, 2, name=name)
    qc.h(0)                      # Superpoziție pe q0
    qc.cx(0, 1)                  # Entanglement (CNOT)
    chsh_rot_to_z(qc, 0, a_angle)  # Rotație Alice
    chsh_rot_to_z(qc, 1, b_angle)  # Rotație Bob
    qc.measure([0, 1], [0, 1])
    return qc

def build_separable_chsh_circuit(a_angle: float, b_angle: float, name="S") -> QuantumCircuit:
    """
    Circuit SEPARABLE (control): Același setup ca entangled, dar FĂRĂ CNOT!
    H → H → [NU CNOT] → Rotații → Măsurare
    """
    qc = QuantumCircuit(2, 2, name=name)
    qc.h(0)                      # Superpoziție pe q0 (la fel ca entangled)
    qc.h(1)                      # Superpoziție pe q1 (la fel ca entangled)
    # FĂRĂ CNOT! - singura diferență
    chsh_rot_to_z(qc, 0, a_angle)  # Rotație Alice
    chsh_rot_to_z(qc, 1, b_angle)  # Rotație Bob
    qc.measure([0, 1], [0, 1])
    return qc

def expectation_from_counts(counts: dict) -> float:
    """
    Calculează E = P(00) + P(11) - P(01) - P(10)
    """
    shots = max(1, sum(counts.values()))
    p00 = counts.get("00", 0) / shots
    p01 = counts.get("01", 0) / shots
    p10 = counts.get("10", 0) / shots
    p11 = counts.get("11", 0) / shots
    return (p00 + p11) - (p01 + p10)

def compute_chsh_from_four(counts_ab, counts_abp, counts_apb, counts_apbp) -> float:
    """
    Calculează S = E(a,b) + E(a,b') + E(a',b) - E(a',b')
    """
    Eab   = expectation_from_counts(counts_ab)
    Eabp  = expectation_from_counts(counts_abp)
    Eapb  = expectation_from_counts(counts_apb)
    Eapbp = expectation_from_counts(counts_apbp)
    return Eab + Eabp + Eapb - Eapbp

def counts_of(idx, prim_result):
    """Extrage counts din rezultatul SamplerV2."""
    try:
        return prim_result[idx].join_data().get_counts()
    except Exception:
        d = getattr(prim_result[idx], "data", None)
        bs = getattr(d, "bitstrings", None) if d else None
        if bs is not None:
            acc = {}
            for bits in bs:
                s = "".join(str(int(b)) for b in bits)
                acc[s] = acc.get(s, 0) + 1
            return acc
        raise RuntimeError("Nu pot obține counts din rezultat!")

# ===== PROGRAM PRINCIPAL =====
print("=" * 70)
print("🔬 TESTUL CHSH - Verificare Inegalitate Bell pe IBM Quantum Hardware")
print("=" * 70)
print("\n📝 Acest program demonstrează că natura este cuantică, nu clasică!")
print("   Vom compara circuite ENTANGLED vs SEPARABLE.\n")

# ===== INPUT TOKEN =====
token = getpass.getpass("🔑 Introduceți IBM Quantum API token: ").strip()
if not token:
    raise ValueError("❌ Token gol! Reporniți programul cu un token valid.")

# ===== CONECTARE LA IBM QUANTUM =====
print("\n⚡ Conectare la IBM Quantum Platform...")
instance = os.getenv("IBM_QUANTUM_INSTANCE")
service = QiskitRuntimeService(channel="ibm_quantum_platform", token=token, instance=instance)

# ===== SELECTARE BACKEND =====
print("🔍 Căutare backend disponibil...")
try:
    backend = service.least_busy(operational=True, simulator=False, min_num_qubits=2)
except Exception:
    backends = service.backends(operational=True, simulator=False, min_num_qubits=2)
    if not backends:
        raise RuntimeError("❌ Nu s-a găsit niciun backend hardware disponibil!")
    backend = backends[0]

print(f"✓ Backend selectat: {backend.name}\n")

# ===== TRANSPILER =====
pm = generate_preset_pass_manager(target=backend.target, optimization_level=3)

# ===== UNGHIURI CHSH (OPTIME PENTRU VIOLARE MAXIMA) =====
A0 = 0.0                  # Alice: unghi 0°
A1 = math.pi / 4.0        # Alice: unghi 45°
B0 = math.pi / 8.0        # Bob: unghi 22.5°
B1 = -math.pi / 8.0       # Bob: unghi -22.5°

print("📐 Unghiuri măsurare:")
print(f"   Alice: A0={A0:.3f} rad (0°), A1={A1:.3f} rad (45°)")
print(f"   Bob:   B0={B0:.3f} rad (22.5°), B1={B1:.3f} rad (-22.5°)\n")

# ===== CONSTRUIRE CIRCUITE ENTANGLED =====
print("🔗 Construire circuite ENTANGLED (cu CNOT)...")
circ_E_ab    = build_entangled_chsh_circuit(A0, B0, "E_ab")
circ_E_abp   = build_entangled_chsh_circuit(A0, B1, "E_abp")
circ_E_apb   = build_entangled_chsh_circuit(A1, B0, "E_apb")
circ_E_apbp  = build_entangled_chsh_circuit(A1, B1, "E_apbp")

# ===== CONSTRUIRE CIRCUITE SEPARABLE =====
print("📊 Construire circuite SEPARABLE (fără CNOT - control)...")
circ_S_ab    = build_separable_chsh_circuit(A0, B0, "S_ab")
circ_S_abp   = build_separable_chsh_circuit(A0, B1, "S_abp")
circ_S_apb   = build_separable_chsh_circuit(A1, B0, "S_apb")
circ_S_apbp  = build_separable_chsh_circuit(A1, B1, "S_apbp")

# ===== SALVARE ARTEFACTE =====
print("💾 Salvare scheme și fișiere QASM...")
for stem, c in [
    ("entangled_ab", circ_E_ab), ("entangled_abp", circ_E_abp),
    ("entangled_apb", circ_E_apb), ("entangled_apbp", circ_E_apbp),
    ("separable_ab", circ_S_ab), ("separable_abp", circ_S_abp),
    ("separable_apb", circ_S_apb), ("separable_apbp", circ_S_apbp),
]:
    save_artifacts(c, stem)

# ===== TRANSPILARE =====
print("⚙️  Transpilare circuite pentru backend...")
entangled_list = [pm.run(c) for c in [circ_E_ab, circ_E_abp, circ_E_apb, circ_E_apbp]]
separable_list = [pm.run(c) for c in [circ_S_ab, circ_S_abp, circ_S_apb, circ_S_apbp]]

# ===== EXECUTIE PE HARDWARE =====
SHOTS = 2000
sampler = Sampler(mode=backend)

print(f"\n⏳ Trimitere job ENTANGLED pe {backend.name} (2000 shots × 4 circuite)...")
print("   Aceasta poate dura câteva minute...")
job_ent = sampler.run(entangled_list, shots=SHOTS)
res_ent = job_ent.result()
print("✓ Job ENTANGLED completat!")

print(f"\n⏳ Trimitere job SEPARABLE pe {backend.name} (2000 shots × 4 circuite)...")
job_sep = sampler.run(separable_list, shots=SHOTS)
res_sep = job_sep.result()
print("✓ Job SEPARABLE completat!")

# ===== EXTRAGERE COUNTS =====
print("\n📊 Procesare rezultate...")
E_counts = [counts_of(i, res_ent) for i in range(4)]
S_counts = [counts_of(i, res_sep) for i in range(4)]

# ===== CALCUL S (CHSH) =====
S_entangled = compute_chsh_from_four(*E_counts)
S_separable = compute_chsh_from_four(*S_counts)

# ===== SALVARE JSON =====
out = {
    "timestamp": datetime.now(timezone.utc).isoformat(),
    "backend": backend.name,
    "shots": SHOTS,
    "angles_rad": {"A0": A0, "A1": A1, "B0": B0, "B1": B1},
    "S_entangled": S_entangled,
    "S_separable": S_separable,
    "bounds": {"classical": 2.0, "tsirelson": 2.0 * math.sqrt(2.0)},
    "efficiency_percent": (abs(S_entangled) / (2.0 * math.sqrt(2.0))) * 100,
    "notes": "Așteptat: |S_entangled| > 2 (violare Bell), |S_separable| ≤ 2 (control clasic)"
}
(OUTDIR / "results.json").write_text(json.dumps(out, indent=2), encoding="utf-8")

# ===== AFISARE REZULTATE =====
print("\n" + "=" * 70)
print("🎯 REZULTATE FINALE - TESTUL CHSH")
print("=" * 70)
print(f"\n🖥️  Backend:         {backend.name}")
print(f"🔬 Shots/circuit:   {SHOTS}")
print(f"\n📈 S (entangled):   {S_entangled:.3f}  ", end="")
if abs(S_entangled) > 2.0:
    print("✅ ÎNCĂLCARE BELL! (> 2.0)")
else:
    print("❌ Nu încalcă (≤ 2.0)")

print(f"📉 S (separable):   {S_separable:.3f}  ", end="")
if abs(S_separable) <= 2.0:
    print("✅ Control OK (≤ 2.0)")
else:
    print("⚠️  Neașteptat (> 2.0)")

print(f"\n📏 Limite teoretice:")
print(f"   • Clasică (Einstein):  |S| ≤ 2.0")
print(f"   • Cuantică (Tsirelson): |S| ≤ 2√2 ≈ 2.828")
print(f"\n⚡ Eficiență cuantică: {out['efficiency_percent']:.1f}%")
print("=" * 70)

# ===== INTERPRETARE =====
print("\n💡 INTERPRETARE:")
if abs(S_entangled) > 2.0 and abs(S_separable) <= 2.0:
    print("✅ SUCCES COMPLET! Experimentul demonstrează:")
    print("   1. Încurcarea cuantică încalcă limita clasică (S > 2)")
    print("   2. Controlul separabil rămâne clasic (S ≤ 2)")
    print("   3. Diferența provine DOAR din entanglement (CNOT)")
    print("\n🏆 Ați reprodus un experiment Premiu Nobel (2022)!")
    print("   Natura este fundamental CUANTICĂ, nu clasică! 🎉")
elif abs(S_entangled) > 2.0:
    print("⚠️  Încălcare Bell confirmată, dar controlul este neașteptat.")
    print("   Posibile cauze: zgomot hardware sever sau erori de calibrare.")
else:
    print("❌ Nu s-a observat încălcare Bell.")
    print("   Posibile cauze: zgomot hardware excesiv, erori de transpilare.")

print("\n" + "=" * 70)

# ===== AFISARE CIRCUITE =====
print("\n📷 CIRCUIT ENTANGLED (a,b) - cu CNOT:")
display(Image(OUTDIR / "entangled_ab.png"))

print("\n📷 CIRCUIT SEPARABLE (a,b) - fără CNOT:")
display(Image(OUTDIR / "separable_ab.png"))

# ===== AFISARE JSON =====
print("\n📋 Rezultate complete (JSON):")
display(JSON(out))

print(f"\n✅ Toate fișierele salvate în: {OUTDIR}")
print("   • entangled_*.png / *.qasm")
print("   • separable_*.png / *.qasm")
print("   • results.json")
print("\n🎓 Program completat cu succes! 🚀")