"""
PAULIJEVI GEJTOVI: I, X, Y, Z
=============================
Tri Paulijeva gejta (X, Y, Z) su osnovne rotacije za 180° oko odgovarajućih osa
Blohove sfere, a I je identitet (ništa ne menja):

  - X  (bit-flip)    : |0⟩ ↔ |1⟩            — rotacija oko X ose
  - Z  (phase-flip)  : |1⟩ → −|1⟩           — rotacija oko Z ose
  - Y  = iXZ         : kombinuje bit- i phase-flip — rotacija oko Y ose
  - I  (identitet)   : ostavlja stanje nepromenjeno

Da bi se efekat svakog gejta jasno video, program svaki od njih primenjuje na
šest karakterističnih stanja koja "pokrivaju" Blohovu sferu (|0⟩, |1⟩, |+⟩, |−⟩,
|i+⟩, |i−⟩) i prikazuje stanje pre i posle, na sferi i kao verovatnoće.
"""

from qiskit import QuantumCircuit

import quantum_utils_v1 as q

# ============================================================================
# POMOĆNE FUNKCIJE: priprema kanonskih stanja na Bloch sferi
# ============================================================================


def prep_state(qc: QuantumCircuit, label: str, qubit: int = 0) -> QuantumCircuit:
    """
    Priprema 1-kubitna stanja koja lepo “pokrivaju” Bloch sferu:
      |0>, |1>, |+>, |->, |i+>, |i->
    gde:
      |+>  = (|0>+|1>)/sqrt(2)
      |->  = (|0>-|1>)/sqrt(2)
      |i+> = (|0>+i|1>)/sqrt(2)
      |i-> = (|0>-i|1>)/sqrt(2)
    """
    if label == "|0>":
        pass
    elif label == "|1>":
        qc.x(qubit)
    elif label == "|+>":
        qc.h(qubit)
    elif label == "|->":
        qc.x(qubit)
        qc.h(qubit)
    elif label == "|i+>":
        qc.h(qubit)
        qc.s(qubit)
    elif label == "|i->":
        qc.h(qubit)
        qc.sdg(qubit)
    else:
        raise ValueError(f"Nepoznato stanje: {label}")
    return qc


def apply_pauli(qc: QuantumCircuit, gate: str, qubit: int = 0) -> QuantumCircuit:
    """Primena Pulijevih gejtova: I, X, Y, Z."""
    gate = gate.upper()
    if gate == "I":
        qc.id(qubit)
    elif gate == "X":
        qc.x(qubit)
    elif gate == "Y":
        qc.y(qubit)
    elif gate == "Z":
        qc.z(qubit)
    else:
        raise ValueError(f"Nepoznat Pauli gejt: {gate}")
    return qc

# ============================================================================
# DEMO: Pulijevi gejtovi na više ulaznih stanja + Bloch sfera posle svakog koraka
# ============================================================================


def demo_pauli_gate(gate: str, input_states):
    print("\n" + "=" * 80)
    print(f"DEMO: Pauli {gate} gejt")
    print("=" * 80)

    for st in input_states:
        qc = QuantumCircuit(1)  # jedan kubit je idealan za Bloch sferu
        qc.barrier()

        # 1) Priprema ulaznog stanja
        prep_state(qc, st, 0)
        qc.barrier()
        print(f"\nUlazno stanje: {st}")
        q.show_bloch_sphere(qc)
        q.print_state(qc, "Stanje (pre gejta):", True)
        q.print_probabilities(qc, "Verovatnoće (pre gejta):")

        # 2) Primena Pulijevog gejta
        apply_pauli(qc, gate, 0)
        qc.barrier()
        print(f"Primena gejta: {gate}")
        q.show_bloch_sphere(qc)
        q.print_state(qc, "Stanje (posle gejta):", True)
        q.print_probabilities(qc, "Verovatnoće (posle gejta):")

        # (Opcionalno) prikaži kolo
        q.show_qc(qc)


def main():
    # Kanonska stanja koja eksplicitno pokazuju:
    # - X flipuje z-polove, ostavlja x osu
    # - Z flipuje fazu (x osu), ostavlja z osu
    # - Y kombinuje bit-flip + phase-flip (rotacija oko y ose za pi, uz globalnu fazu)
    states = ["|0>", "|1>", "|+>", "|->", "|i+>", "|i->"]

    for gate in ["X", "Y", "Z"]:
        demo_pauli_gate(gate, states)


if __name__ == "__main__":
    main()