# ===== GOOGLE COLAB - CELULA COMPLETA =====
# Instalare pachete
!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

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

# Setup directoare
OUTDIR = Path("/content/chsh_outputs")
OUTDIR.mkdir(exist_ok=True)

# Functii auxiliare
def dump_qasm(circ: QuantumCircuit) -> str:
    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):
    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):
    qc.ry(-2.0 * angle, q)

def build_entangled_chsh_circuit(a_angle: float, b_angle: float, name="E") -> QuantumCircuit:
    qc = QuantumCircuit(2, 2, name=name)
    qc.h(0)
    qc.cx(0, 1)
    chsh_rot_to_z(qc, 0, a_angle)
    chsh_rot_to_z(qc, 1, b_angle)
    qc.measure([0, 1], [0, 1])
    return qc

def build_separable_chsh_circuit(a_angle: float, b_angle: float, name="S") -> QuantumCircuit:
    qc = QuantumCircuit(2, 2, name=name)
    chsh_rot_to_z(qc, 0, a_angle)
    chsh_rot_to_z(qc, 1, b_angle)
    qc.measure([0, 1], [0, 1])
    return qc

def expectation_from_counts(counts: dict) -> float:
    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:
    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

# ===== PROGRAM PRINCIPAL =====
print("=== CHSH pe IBM Quantum Hardware (Google Colab) ===\n")

# Input token (securizat în Colab)
token = getpass.getpass("Introduceți IBM Quantum API token: ").strip()
if not token:
    raise ValueError("Token gol!")

# Conectare
service = QiskitRuntimeService(channel="ibm_quantum_platform", token=token)

# Backend
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!")
    backend = backends[0]

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

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

# Unghiuri CHSH
A0 = 0.0
A1 = math.pi / 4.0
B0 = math.pi / 8.0
B1 = -math.pi / 8.0

# Circuite ENTANGLED
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")

# Circuite SEPARABLE
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
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
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("⏳ Trimitere job ENTANGLED...")
job_ent = sampler.run(entangled_list, shots=SHOTS)
res_ent = job_ent.result()

print("⏳ Trimitere job SEPARABLE...")
job_sep = sampler.run(separable_list, shots=SHOTS)
res_sep = job_sep.result()

# Extragere counts
def counts_of(idx, prim_result):
    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!")

E_counts = [counts_of(i, res_ent) for i in range(4)]
S_counts = [counts_of(i, res_sep) for i in range(4)]

S_entangled = compute_chsh_from_four(*E_counts)
S_separable = compute_chsh_from_four(*S_counts)

# Rezultate
out = {
    "timestamp": datetime.utcnow().isoformat() + "Z",
    "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)},
}
(OUTDIR / "results.json").write_text(json.dumps(out, indent=2), encoding="utf-8")

# Afișare rezultate
print("\n" + "="*60)
print("📊 REZULTATE CHSH (2000 shots/circuit)")
print("="*60)
print(f"Backend:         {backend.name}")
print(f"S (entangled):   {S_entangled:.3f}  {'✓ ÎNCALCARE!' if abs(S_entangled) > 2.0 else ''}")
print(f"S (separable):   {S_separable:.3f}  {'✓ Control OK' if abs(S_separable) <= 2.0 else ''}")
print(f"Limita clasică:  2.0")
print(f"Limita Tsirelson: {2*math.sqrt(2):.3f}")
print("="*60)

# Afișare exemple de circuite
print("\n📷 Circuit ENTANGLED (a,b):")
display(Image(OUTDIR / "entangled_ab.png"))

print("\n📷 Circuit SEPARABLE (a,b):")
display(Image(OUTDIR / "separable_ab.png"))

print("\n📋 JSON Results:")
display(JSON(out))

print(f"\n✓ Toate fișierele salvate în: {OUTDIR}")