"""
QFT: KODIRANJE BROJA U FAZE
===========================
Ovaj program pokazuje šta Kvantna Furijeova Transformacija (QFT) zapravo radi sa
običnim brojem upisanim u kubite.

Ako kubite postavimo u bazno stanje |j⟩ (npr. |101⟩ = broj 5), QFT ga pretvara u
ravnomernu superpoziciju SVIH baznih stanja, ali sa različitim FAZAMA:

        QFT|j⟩ = (1/√N) · Σ_k  e^{2πi·j·k/N} · |k⟩

Verovatnoće svih stanja postaju jednake (svako 1/N), pa se broj j više ne vidi u
amplitudama — on je sada "sakriven" u brzini kojom faze rotiraju po stanjima |k⟩.
Veći broj j ⇒ brža rotacija faze. Program prolazi kroz svih 8 stanja (000…111) i
za svako prikazuje rezultat na Blohovoj sferi, da se uoči taj fazni "potpis".

Funkcija apply_qft ispod implementira QFT (i opciono inverznu iQFT) iz osnovnih
gejtova: H i kontrolisanih faznih rotacija cp().
"""

import math
from qiskit import QuantumCircuit
from typing import Iterable

import quantum_utils_v1 as q

# =============================================================================
# FUNKCIJA: Kvantna Furijeova transformacija
# =============================================================================


def apply_qft(qc, qubit_range: Iterable[int], inverse: bool = False, approximation: float = 0.0):
    """
    Primeni (ili inverznu) diskretnu kvantnu Furijeovu transformaciju na dati opseg kubita.
    
    Args:
        qc: QuantumCircuit - kvantno kolo na koje se primenjuju gejtovi.
        qubit_range: Iterable[int] - iterabilni niz indeksa kubita (npr. range(n) ili [0,1,2]).
                    Redosled je Vaš logički redosled kubita (najmanje značajan -> najmanje ili kako vi želite).
        inverse: bool - ako je True, primeni inverznu QFT (iQFT). Default: False.
        approximation: float - prag ispod kojeg se preskaču kontrolisane fazne rotacije (ugao u radijanima).
                    Koristno za aproksimativnu QFT. Default: 0.0 (ne preskače se ništa).
    
    Povratna vrednost:
        qc: QuantumCircuit - vraća isti QuantumCircuit nakon primene QFT-a.
    
    Napomena:
        - Koristi qc.h, qc.cp i qc.swap (Qiskit API).
        - Ovaj kod pretpostavlja standardne Qiskit metode (qc.cp dostupno).
    """
    # Normalizuj listu indeksa kubita u listu (ako je prosleđen range ili generator)
    qubits = list(qubit_range)
    n = len(qubits)
    if n == 0:
        return qc

    # Ako želimo inverznu QFT, možemo primeniti QFT sa obrnutim redosledom i zameniti znak uglova
    if inverse:
        # Inverzna QFT: izvrši korake QFT obrnutim redosledom i sa negativnim uglovima
        # 1) Swap-ovi (redosled zadržavamo isti kao i kod QFT - swap na kraju)
        # 2) Za svaki i od n-1 do 0: primeni kontrolisane faze sa negativnim uglom, pa H
        qc.barrier()
        for i in range(n - 1, -1, -1):
            target = qubits[i]
            # kontrolisane rotacije (od narednih kubita prema kraju)
            for j in range(i + 1, n):
                control = qubits[j]
                k = j - i + 0  # odn. eksponent
                angle = - math.pi / (2 ** (j - i))
                if abs(angle) < approximation:
                    continue
                # cp(angle, control, target) - kontrolisani fazni gejt
                qc.cp(angle, control, target)
            qc.h(target)
            qc.barrier()
        # swap da se vrati redosled bitova (QFT obrće bit-order)
        for i in range(n // 2):
            qc.swap(qubits[i], qubits[n - 1 - i])
        qc.barrier()
        return qc

    # --- Normalna QFT ---
    qc.barrier()
    # Za svaki cilj i od 0 do n-1:
    #  - primeni H na i
    #  - za svaki j = i+1..n-1 primeni kontrolisanu rotaciju cp(angle) sa kontrolom j i ciljem i
    # (ovo enkoduje faze na "manje značajne" kubite)
    for i in range(n):
        target = qubits[i]
        qc.h(target)
        # kontrolisane rotacije sa kubitima koji su "viši" u indeksu
        for j in range(i + 1, n):
            control = qubits[j]
            angle = math.pi / (2 ** (j - i))
            if abs(angle) < approximation:
                continue
            qc.cp(angle, control, target)
        qc.barrier()

    # Swap-ovi da obrnemo redosled kubita (QFT tipično zahteva reverse)
    # for i in range(n // 2):
    #     qc.swap(qubits[i], qubits[n - 1 - i])
    # qc.barrier()

    return qc

# =============================================================================
# FUNKCIJA: Enkodiranje binarnog broja u početno stanje 3 kubita
# =============================================================================


def prepare_basis_state(qc, binary_string):
    """
    Postavlja 3-kubitni sistem u određeno binarno stanje npr. '101' -> |101⟩
    """
    for i, bit in enumerate(binary_string):
        if bit == '1':
            qc.x(i)  
            # Napomena:
            # Ako koristiš drugačiji redosled kubita u Svojim primerima,
            # samo obrni indeks. Ovde je |q2 q1 q0⟩ konvencija.
    return qc


# =============================================================================
# GLAVNI DEO: Prolaz kroz svih 8 binarnih stanja
# =============================================================================
states = [
    "000", "001", "010", "011",
    "100", "101", "110", "111"
]

for state in states:

    print(f"\n=======================================")
    print(f"Binarno stanje: |{state}⟩")
    print(f"=======================================\n")

    # Kreiramo kvantno kolo sa 3 kubita
    qc = QuantumCircuit(3)
    
    # Postavljanje sistema u početno registrovano stanje
    qc = prepare_basis_state(qc, state)
    qc.barrier()

    # Prikaz početnog stanja na Blohovim sferama
    # q.show_bloch_sphere(qc, title=f"Početno stanje |{state}⟩")

    # Primena kvantne Furijeove transformacije
    apply_qft(qc, range(3))
    qc.barrier()

    # Prikaz stanja nakon QFT-a
    q.show_bloch_sphere(qc, title=f"Stanje nakon QFT za |{state}⟩")
    q.print_state(qc, show_amplitudes=True)