import numpy as np
from qiskit import QuantumCircuit

import quantum_utils_v1 as q

# ============================================================================
# QUANTUM FOURIER TRANSFORM (QFT)
# ============================================================================


def qft_rotations(qc, n):
    """
    Primenjuje rotacione gejte za QFT na n kubita.
    """
    if n == 0:
        return qc

    n -= 1
    qc.h(n)
    qc.barrier()
    q.show_bloch_sphere(qc)

    for qubit in range(n):
        qc.cp(np.pi/2**(n-qubit), qubit, n)
        qc.barrier()
        q.show_bloch_sphere(qc)

    qft_rotations(qc, n)


def qft(qc, n):
    """
    Kompletna implementacija Quantum Fourier Transform-a.
    """
    qft_rotations(qc, n)
    return qc

# ============================================================================
# INVERSE QUANTUM FOURIER TRANSFORM (IQFT)
# ============================================================================


def iqft_rotations(qc, n):
    """
    Primenjuje inverzne rotacione gejte za IQFT na n kubita.
    
    Ova funkcija radi obrnutim redosledom u odnosu na qft_rotations:
    - Za svaki kubit (od najmanje značajnog prema najznačajnijem)
      primenjujemo kontrolisane fazne rotacije sa negativnim uglovima
      (inverzno od QFT) pa zatim Hadamard gejt.
    """
    # Iterativna implementacija (prozaičnija za čitanje)
    for target in range(n):
        # Kontrolisane rotacije u rastućem nizu kontrolnih kubita
        for control in range(target):
            # Inverzni ugao: negativan od originalnog
            # U originalu za ciljni kubit i, ugao je π/2^{i-control}
            qc.cp(-np.pi/2**(target-control), control, target)
            qc.barrier()
            q.show_bloch_sphere(qc)
        # Nakon inverznih kontrolisanih rotacija, primeni H
        qc.h(target)
        qc.barrier()
        q.show_bloch_sphere(qc)

    return qc


def iqft(qc, n):
    """
    Kompletna implementacija Inverse Quantum Fourier Transform-a.
    
    Napomena:
      QFT = (rotations)  
      IQFT = (rotations)^{-1}
      Dakle primenjujemo inverzne rotacije (u suprotnom redosledu),
    """
    iqft_rotations(qc, n)
    return qc


# ----------------------------------------------------------------------------
# KORAK 1: INICIJALIZACIJA POČETNOG STANJA |ψ₀⟩ = |011⟩ (kako je u originalnom kodu)
# ----------------------------------------------------------------------------
qc = QuantumCircuit(3, 3)

# Postavljamo kubite tako da dobijemo |011⟩ = |0⟩ (MSB) ⊗ |1⟩ ⊗ |1⟩ (LSB)
qc.x(0)
qc.x(1)
qc.barrier()

print("\n--- Početno stanje ---")
q.print_state(qc, "Stanje", show_amplitudes=True)
q.show_bloch_sphere(qc)

# ----------------------------------------------------------------------------
# KORAK 2: PRIMENA QUANTUM FOURIER TRANSFORM
# ----------------------------------------------------------------------------
qft(qc, 3)

print("\n--- Stanje nakon QFT ---")
q.print_state(qc, "Stanje", show_amplitudes=True)
q.print_probabilities(qc, "Verovatnoće")

# ----------------------------------------------------------------------------
# KORAK 2.5: PRIMENA INVERSE QFT (IQFT) koja bi trebala da vrati stanje nazad
# ----------------------------------------------------------------------------
# Primena IQFT odmah nakon QFT kao demonstracija inverzije
iqft(qc, 3)

print("\n--- Stanje nakon IQFT (trebalo bi da se vrati na početno) ---")
q.print_state(qc, "Stanje", show_amplitudes=True)
q.print_probabilities(qc, "Verovatnoće")

# ----------------------------------------------------------------------------
# KORAK 3: MERENJE
# ----------------------------------------------------------------------------
qc.measure(range(3), range(3))
q.show_qc(qc)
q.show_measurement(qc, shots=1024)