"""
Shor-ov algoritam (demonstracija) — Qiskit implementacija sa
potpunom (matrix-based) modularnom eksponencijalizacijom.

Napomena:
- Ovaj pristup konstruiše jedinični (permutacioni) matrix za operator
  U : |y> -> |(a * y) mod N> koji deluje na m = ceil(log2(N)) kubita.
"""

import math
import numpy as np

from math import gcd
from fractions import Fraction
from qiskit import QuantumCircuit
from qiskit.circuit.library import QFT, UnitaryGate

import quantum_utils_v1 as q

# ===================================================================
# PARAMETRI — LAKO PODEŠAVO
# ===================================================================
n_count = 12      # merni kubiti
N = 55
a = 3

# ===================================================================
# POMOĆNE MATEMATIČKE FUNKCIJE
# ===================================================================


def continued_fraction(x, max_denominator):
    frac = Fraction(x).limit_denominator(max_denominator)
    return int(frac.numerator), int(frac.denominator)


def get_order_from_measurement(meas, n_count, N):
    fraction = meas / (2**n_count)
    p, q = continued_fraction(fraction, N)
    return q

# ===================================================================
# KOREKCIJA #1 — PRAVA PERMUTACIONA MATRICA
# ===================================================================


def build_mult_a_mod_N(a, N):
    """
    |y> → |(a*y) mod N>, y < N
    Za y >= N ostaje identitet.
    """
    m = math.ceil(math.log2(N))
    dim = 2**m

    # identitet
    M = np.zeros((dim, dim), dtype=complex)

    # prava permutacija — SVE kolone imaju TAČNO JEDNU jedinicu
    for y in range(dim):
        if y < N:
            new = (a * y) % N
        else:
            new = y  # identitet izvan domena

        M[new, y] = 1.0

    return M, m

# ===================================================================
# KONTROLISANA modularna eksponencijacija
# ===================================================================


def apply_controlled_modexp(qc, control_qubit, target_qubits, a, N, power):
    M, m = build_mult_a_mod_N(a, N)
    Mpow = np.linalg.matrix_power(M, power)

    gate = UnitaryGate(Mpow, label=f"U^{power}")
    cgate = gate.control()

    qc.append(cgate, [control_qubit] + target_qubits)

# ===================================================================
# GLAVNI KOD — KONSTRUKCIJA KOLA
# ===================================================================


def shor_order_finding_circuit(N, a, n_count):
    g = gcd(a, N)
    if g != 1:
        raise ValueError(f"Trivijalni faktor: gcd({a},{N}) = {g}")

    _, m = build_mult_a_mod_N(a, N)
    qc = QuantumCircuit(n_count + m, n_count)

    target_qubits = list(range(n_count, n_count + m))

    qc.barrier()

    # -------------------------------------------------------------------
    # Izaberi početno stanje drugog registra:
    # -------------------------------------------------------------------

    # y = 1
    qc.x(target_qubits[0])

    # y = 5  = 00101  (LSB = qubit[0])
    # qc.x(target_qubits[0])
    # qc.x(target_qubits[2])

    # y = 10 = 01010
    # qc.x(target_qubits[1])
    # qc.x(target_qubits[3])

    qc.barrier()

    # H-gejtovi na mernim kubitima
    for i in range(n_count):
        qc.h(i)
    qc.barrier()

    # primena kontrolisanih U^{2^k}
    for k in range(n_count):
        repetitions = 2**k
        apply_controlled_modexp(
            qc,
            control_qubit=k,
            target_qubits=target_qubits,
            a=a,
            N=N,
            power=repetitions
        )

    qc.barrier()

    # inverzna QFT
    qc = qc.compose(QFT(num_qubits=n_count, inverse=True),
                    list(range(n_count)))
    qc.barrier()

    # merenje
    for i in range(n_count):
        qc.measure(i, i)

    return qc, m


# ===================================================================
# POKRETANJE
# ===================================================================
if __name__ == "__main__":
    # ... (prethodni deo koda: inicijalizacija, shor_circuit poziv) ...

    qc, m = shor_order_finding_circuit(N, a, n_count)
    
    # 1. Pokrećemo simulaciju da dobijemo 'counts' (rečnik {stanje: broj_pojavljivanja})
    # Pretpostavljam da tvoja biblioteka 'q' ima funkciju koja vraća counts objekat
    # Ako ne, koristi: counts = q.run_sim(qc) ili ekvivalent
    # Ovde ću koristiti standardni qiskit poziv radi sigurnosti, prilagodi ako koristiš svoju biblioteku:
    q.show_qc(qc)
    
    counts=q.show_measurement(qc)

    # --- KLJUČNI DEO: FILTRIRANJE NULE ---
    # Pravimo novi rečnik koji NE SADRŽI stanje '0' (fazu 0)
    counts_without_zero = {}
    for bitstring, count in counts.items():
        # Uklanjamo razmake ako ih Qiskit ubaci i konvertujemo u int
        value = int(bitstring.replace(" ", ""), 2)
        if value != 0:
            counts_without_zero[bitstring] = count

    print(f"\n--- ANALIZA MERENJA) ---")

    if not counts_without_zero:
        print("Greška: Sva merenja su bila 0. Pokreni ponovo sa više shot-ova.")
    else:
        # Sada tražimo MAX u podacima bez nule
        best_bitstring = max(counts_without_zero, key=counts_without_zero.get)
        measured_int = int(best_bitstring.replace(" ", ""), 2)
        
        print(f"Odabrano stanje (izbačena nula): {measured_int}")
        
        # Izračunavanje faze
        phase_val = measured_int / (2**n_count)
        print(f"Izmerena faza: {phase_val:.4f}")
        
        # Verižni razlomci
        r_quantum = get_order_from_measurement(measured_int, n_count, N)
        print(f"Period izračunat iz kvantnog merenja: r = {r_quantum}")

        # Provera
        result_mod = pow(a, r_quantum, N)
        if result_mod == 1:
            print(f"USPEH: {a}^{r_quantum} ≡ 1 (mod {N})")
        else:
            print(f"REZULTAT ORBITE: {a}^{r_quantum} ≡ {result_mod} (mod {N})")