"""
S-GEJT (FAZNI GEJT ZA π/2)
==========================
S-gejt je fazni gejt koji baznom stanju |1⟩ dodaje fazu i = e^{iπ/2}, dok |0⟩
ostavlja nepromenjeno (matrica diag(1, i)). Na Blohovoj sferi to je rotacija
za +90° oko Z ose.

Program demonstrira njegovo dejstvo na dva karakteristična stanja:
  - |+⟩  (na +X osi)  →  S|+⟩ = (|0⟩ + i|1⟩)/√2 = |i⟩  (rotira na +Y osu)
  - |0⟩  (severni pol) →  S|0⟩ = |0⟩  (nema promene, jer je |0⟩ na osi rotacije)

Napomena: S = T², a S† (sdg) radi rotaciju za −90°.
"""

from qiskit import QuantumCircuit
from qiskit.quantum_info import Statevector

import quantum_utils_v1 as q

# =============================================================================
# POMOĆNA FUNKCIJA: Blohova sfera (koristi quantum_utils ako postoji, inače fallback)
# =============================================================================


def show_bloch(statevec, title=""):
    """
    statevec: qiskit.quantum_info.Statevector ili niz kompleksnih amplituda
    """
    # Pokušaj tvoj helper (najčešće postoji u tvojoj biblioteci)
    if hasattr(q, "show_bloch_sphere"):
        q.show_bloch_sphere(statevec, title=title)
        return

    # Fallback: Qiskit-ov prikaz ako tvoja funkcija nije dostupna
    from qiskit.visualization import plot_bloch_multivector
    fig = plot_bloch_multivector(statevec)
    if title:
        try:
            fig.suptitle(title)
        except Exception:
            pass
    import matplotlib.pyplot as plt
    plt.show()


# =============================================================================
# DEMO: Dejstvo S-gejta na kubit na Blohovoj sferi
# =============================================================================
# Dobar primer je da krenemo iz |+> (na +X osi), jer S je Z-rotacija za pi/2,
# pa |+> prelazi u |i> (na +Y osi).
qc_before = QuantumCircuit(1)
qc_before.h(0)  # |0> -> |+>

qc_after = QuantumCircuit(1)
qc_after.h(0)   # |0> -> |+>
qc_after.s(0)   # S: diag(1, i) => rotacija oko Z ose za +pi/2

# Stanja (statevector) pre i posle S
psi_before = Statevector.from_instruction(qc_before)
psi_after  = Statevector.from_instruction(qc_after)

print("Stanje pre S (|+>):", psi_before.data)
print("Stanje posle S (S|+> ~ |i>):", psi_after.data)

# Prikaz na Blohovoj sferi
show_bloch(psi_before, title="Pre S-gejta: |+> (na +X osi)")
show_bloch(psi_after,  title="Posle S-gejta: S|+> = (|0> + i|1>)/√2 = |i> (na +Y osi)")

# (Opcionalno) Ako želiš da vidiš i da S ne menja |0> na Blohovoj sferi:
qc0 = QuantumCircuit(1)
psi0_before = Statevector.from_instruction(qc0)  # |0>
qc0.s(0)
psi0_after = Statevector.from_instruction(qc0)   # S|0> = |0>
show_bloch(psi0_before, title="Pre S-gejta: |0> (Severni pol)")
show_bloch(psi0_after,  title="Posle S-gejta: S|0> = |0> (nema promene)")