import math
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister

import quantum_utils_v1 as q

# ============================================================================ 
# KVANTNA TELEPORTACIJA - PROGRAM
# ============================================================================ 
# Cilj: teleportovati stanje |ψ> sa kubita 0 (Alisa) na kubit 2 (Bob),
# koristeći pred-deljeni Bell par između kubita 1 i 2 i dve klasične poruke.

# ---------------------------------------------------------------------------
# KREIRANJE REGISTARA
# ---------------------------------------------------------------------------
# Koristimo:
# - 3 kvantna kubita: 0 = |ψ> (stanje koje teleportujemo), 1 = Alisin Bell kubit,
#   2 = Bob-ov kubit (ciljni)
# - 3 klasična registra po 1 bitu svaki (m0, m1 za merenja Alisa; m2 za
#   konačno merenje Bob-ovog kubita radi verifikacije)

qr = QuantumRegister(3, 'q')

cr0 = ClassicalRegister(1, 'm0')   # rezultat merenja kubita 0 (Alisa)
cr1 = ClassicalRegister(1, 'm1')   # rezultat merenja kubita 1 (Alisa)
cr2 = ClassicalRegister(1, 'm2')   # rezultat završnog merenja kubita 2 (Bob)
qc = QuantumCircuit(qr, cr0, cr1, cr2)

qc.barrier()
q.show_bloch_sphere(qc)  # Početno stanje |000>

# ---------------------------------------------------------------------------
# KORAK 1: KREIRANJE BELL PARA (DELJENI RESURS ALISE I BOBA)
# kubiti: 1 (Alisa), 2 (Bob)
# ---------------------------------------------------------------------------
qc.h(1)         # Adamar na kubitu 1
qc.cx(1, 2)     # CNOT(1->2) -> Bell par |Φ+> između 1 i 2
qc.barrier()
q.print_state(qc, "Nakon kreiranja Bell para (1-2):", True)
q.show_bloch_sphere(qc)

# ---------------------------------------------------------------------------
# KORAK 2: PRIPREMA STANJA |ψ> NA KUBITU 0 (stanje koje želimo teleportovati)
# ---------------------------------------------------------------------------
# Možeš da menjaš theta i phi da bi teleportovao drugačija stanja.
theta = math.pi / 3      # ugao polarne koordinate (primer)
phi = math.pi / 2        # fazni ugao (primer)

# State vector: [cos(theta/2), e^{i phi} sin(theta/2)]
alpha = math.cos(theta / 2)
beta = math.sin(theta / 2) * complex(math.cos(phi), math.sin(phi))
qc.initialize([alpha, beta], 0)   # inicijalizujemo kubit 0 u željeno stanje
qc.barrier()
q.print_state(qc, "Nakon pripreme |ψ> na kubitu 0:", True)
q.show_bloch_sphere(qc)
# ---------------------------------------------------------------------------
# KORAK 3: ALISA: BELL-MERENJE (entanglovanje stanja |ψ> sa njenim Bell kubitom)
# Operacije: CNOT(0,1) zatim H(0), pa merenja kubita 0 i 1 (Alisa)
# ---------------------------------------------------------------------------
qc.cx(0, 1)
qc.h(0)
qc.barrier()
q.print_state(qc, "Pre merenja (posle CNOT i H):", True)
q.show_bloch_sphere(qc)

# Merenja (Alisa meri svoj kubit q1 iz kvantnog para)
qc.measure(0, cr0[0])  # m0 <= mjerenje kubita 0
qc.measure(1, cr1[0])  # m1 <= mjerenje kubita 1
qc.barrier()
q.print_state(qc, "Nakon Alisinog merenja (posle CNOT i H):", True)
q.show_bloch_sphere(qc, from_instruction=False, qubit_index=2)
q.show_measurement(qc, selected_qubits=[0, 1])
# ---------------------------------------------------------------------------
# KORAK 4: BOB: KOREKCIJE NA OSNOVU KLASIČNIH BITOVA (klasična komunikacija)
# Ako je m1 == 1 -> primeni X na kubit 2
# Ako je m0 == 1 -> primeni Z na kubit 2
# Implementirano koristeći odvojene ClassicalRegister(1) registre i c_if.
# ---------------------------------------------------------------------------
# Napomena: upotreba .c_if sa jednogbitnim classical registrom: vrednost 1 znači
# da je taj bit 1, tako da se uslovna gejt primeni ispravno.

# Uslovna primena X na Bob-ov kubit ako je Alisin drugi merni bit 1
with qc.if_test((cr1[0], 1)):
    qc.x(2)

# Uslovna primena Z na Bob-ov kubit ako je Alisin prvi merni bit 1
with qc.if_test((cr0[0], 1)):
    qc.z(2)

qc.barrier()
q.print_state(qc, "Nakon klasičnih korekcija (pre verifikacionog merenja):", True)
q.show_bloch_sphere(qc, from_instruction=False, qubit_index=2)

# ---------------------------------------------------------------------------
# KORAK 5: VERIFIKACIJA - MERENJE BOB-OVOG KUBITA
# Očekivanje: Bob-ov kubit (2) treba da bude u istom stanju kao inicijalno |ψ>.
# ---------------------------------------------------------------------------
qc.measure(2, cr2[0])   # Merimo kubit 2 i upisujemo u m2 (radi verifikacije)
qc.barrier()
q.show_qc(qc)
q.show_measurement(qc, marginal_counts_flag=True, selected_qubits=[2])    # Prikaži histogram rezultata

# Ispis detalja
q.print_state(qc, "Finalno stanje (posle svih koraka):", True)
q.print_probabilities(qc, "Verovatnoće finalnog stanja:")

# ============================================================================ 
# ZAKLJUČAK:
# - Ako sve radi kako treba, stanje koje je bilo na kubitu 0 pre teleportacije
#   (parametri theta, phi) sada je na kubitu 2 (pre merenja).
# - Merenjem kubita 2 i poređenjem sa očekivanom distribucijom iz alpha,beta
#   možete verifikovati uspeh teleportacije.
# ============================================================================