import numpy as np
from qiskit import QuantumCircuit

import quantum_utils_v1 as q

# ============================================================================
# QUANTUM FOURIER TRANSFORM (QFT)
# ============================================================================
# QFT je kvantni analog klasične diskretne Fourier transformacije.
# Za N-dimenzionalni prostor (N = 2^n gde je n broj kubita), QFT transformiše
# bazno stanje |j⟩ u superpoziciju:
#
#   |j⟩ → (1/√N) Σₖ e^(2πijk/N) |k⟩
#
# QFT je fundamentalna komponenta mnogih kvantnih algoritama, uključujući:
# - Šorov algoritam za faktorizaciju
# - Kvantna procena faze (Quantum Phase Estimation)
# - Algoritmi za rešavanje sistema linearnih jednačina


def qft_rotations(qc, n):
    """
    Primenjuje rotacione gejte za QFT na n kubita.
    
    Ova funkcija implementira jezgro QFT algoritma kroz niz Adamar gejta
    i kontrolisanih faznih rotacija. Svaki kubit prolazi kroz:
    1. H-gejt - kreira ravnomernu superpoziciju
    2. Seriju kontrolisanih R gejta - dodaje fazne relacije
    
    Args:
        qc: QuantumCircuit na koji primenjujemo QFT
        n: Broj kubita
        
    Matematički, za kubit j primenjujemo:
    - H gejt
    - CR_k gejte gde je R_k = [[1, 0], [0, e^(2πi/2^k)]]
    """

    if n == 0:  # Bazni slučaj rekurzije
        return qc
    
    n -= 1  # Indeksiramo od 0
    
    # Primeni H-gejt na poslednji kubit
    # H transformiše |0⟩ → (|0⟩ + |1⟩)/√2 i |1⟩ → (|0⟩ - |1⟩)/√2
    qc.h(n)
    qc.barrier()
    q.show_bloch_sphere(qc)

    # Primeni kontrolisane fazne rotacije
    # Svaka rotacija dodaje fazni odnos između baznih stanja
    # Ugao za k-tu rotaciju: theta = π / 2^{k}, gde k zavisi od razlike indeksa.

    for qubit in range(n):
        # Kontrolisani-P gejt: dodaje fazu e^(iθ) kada je kontrolni kubit |1⟩
        # Ugao rotacije: θ = 2π/2^(n-qubit+1) = π/2^(n-qubit)
        qc.cp(np.pi/2**(n-qubit), qubit, n)
        qc.barrier()
        q.show_bloch_sphere(qc)
    
    # Rekurzivno primeni na preostale kubite
    qft_rotations(qc, n)


def swap_registers(qc, n):
    """
    Obrće redosled kubita u QFT kolu.
    
    QFT prirodno proizvodi izlaz u obrnutom redosledu (bit-reversal).
    Ova funkcija vraća kubite u originalni redosled primenom SWAP gejta.
    
    Args:
        qc: QuantumCircuit
        n: Broj kubita
        
    Napomena:
        Za n kubita potrebno je n//2 SWAP operacija.
        SWAP(i, n-1-i) zamenjuje kubit i sa kubitom n-1-i.
    """

    for qubit in range(n//2):
        qc.swap(qubit, n-qubit-1)

    return qc


def qft(qc, n):
    """
    Kompletna implementacija Quantum Fourier Transform-a.
    
    Args:
        qc: QuantumCircuit
        n: Broj kubita na kojima se primenjuje QFT
        
    Returns:
        QuantumCircuit sa primenjenim QFT
    """
    qft_rotations(qc, n)
    swap_registers(qc, n)
    return qc


# ----------------------------------------------------------------------------
# KORAK 1: INICIJALIZACIJA POČETNOG STANJA |ψ₀⟩ = |110⟩
# ----------------------------------------------------------------------------
# Napomena o konvenciji indeksa:
# - Ovde koristimo konvenciju da je q[2] najviše značajan bit (MSB),
#   a q[0] najmanje značajan bit (LSB).
#   Dakle |011⟩ znači q[0]=1, q[1]=1, q[2]=0.
#
# Kreiranje kvantnog kola sa 3 kubita i 3 klasična bita
# Početno stanje je |000⟩
qc = QuantumCircuit(3, 3)

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

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 transformiše bazno stanje |j⟩ u frekventni domen.
#
# Za |011⟩ (decimalno j = 3) očekujemo:
# QFT|3⟩ = (1/√8) Σ_k e^(2π i · 3 · k / 8) |k⟩
#
# Rezultat je ravnomerna superpozicija svih baznih stanja sa
# specifičnim faznim odnosima koji kodiraju frekvenciju

qft(qc, 3)
qc.barrier(label="Posle QFT")

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

# ----------------------------------------------------------------------------
# KORAK 3: MERENJE
# ----------------------------------------------------------------------------
# Merenjem dobijamo rezultat koji odgovara jednom od baznih stanja iz superpoziciju.
# Verovatnoća svakog stanja je |amplituda|^2 
# Amplituda za svako |k⟩ jednaka je 1/√N, pa je verovatnoća 1/N = 1/8 = 12.5%.

qc.measure(range(3), range(3))
q.show_qc(qc)
q.show_measurement(qc, shots=1024)