#!/usr/bin/env python3
"""
Capitoli 10 e 11: selezione di portafoglio con vincolo di cardinalità come
problema binario quadratico (QUBO), e i termini di confronto classici che un
campionatore quantistico deve battere.

Istanza (costruita, seme dichiarato): n titoli, modello a un fattore
    Sigma = beta beta^T s_m^2 + diag(s_i^2),
beta_i ~ U(0,5; 1,5), s_m = 0,18, s_i ~ U(0,10; 0,35), mu_i ~ U(0,02; 0,12).
Variabili x_i in {0,1} (titolo incluso a peso uguale 1/K). Funzione obiettivo
(da minimizzare):
    E(x) = q (x/K)^T Sigma (x/K) - mu^T (x/K) + P (sum x - K)^2,
con K = n/4 arrotondato, q = 2 (avversione al rischio) e penalità
P = 2 (max mu / K + q max_i sum_j |Sigma_ij| / K^2), commisurata al massimo
beneficio marginale di un titolo (con P = 1 la simulated annealing a un bit resta congelato nel
primo insieme di cardinalità K che incontra: risultato delle prove di
progettazione, non riprodotto da questo script).

Metodi:
    esatto      enumerazione di tutti i 2^n vettori (n <= 22);
    uniforme    campionamento uniforme (riferimento: un dispositivo con rumore
                depolarizzante completo; il rumore T1 reale sbilancia invece verso 0);
    uniforme+LS campionamento uniforme seguito da local search a un bit
                (best improvement) fino a un minimo locale: lo schema di
                post-elaborazione usato negli esperimenti su hardware;
    SA          simulated annealing con mosse di swap (preserva la cardinalità),
                sweep di n proposte, beta geometrico da 0,1 a 1000 in unità di
                1/(deviazione standard dell'obiettivo su 2000 portafogli
                ammissibili casuali); TTS ottimizzato sul numero di sweep
                (griglia 1, 2, 5, 10, 20, 50) e R99 >= 1.

Tempo per la soluzione al 99% (TTS): numero di tentativi R99 = ln(0,01)/ln(1-p)
moltiplicato per il costo di un tentativo, misurato in valutazioni di
variazione di energia a un bit (unità indipendente dalla macchina).

Produce:
    c10_tts.dat              mediana su 8 istanze del log10 del TTS per metodo, n = 8..22
    c10_tabella_tts.txt      tabella dei valori e del fattore di crescita per titolo
    c10_distribuzione.dat    istogramma del divario relativo dall'ottimo, n = 20: casuale, casuale+LS, SA
"""

import math

import numpy as np

from comune import scrivi_dat, scrivi_txt, it, it_sci

RNG_SEME = 7
ISTANZE = 8
CAMPIONI = 4000


def istanza(n, rng):
    beta = rng.uniform(0.5, 1.5, n)
    s_i = rng.uniform(0.10, 0.35, n)
    mu = rng.uniform(0.02, 0.12, n)
    sigma = np.outer(beta, beta) * 0.18 ** 2 + np.diag(s_i ** 2)
    K = max(2, round(n / 4))
    q = 2.0
    # penalità commisurata al massimo beneficio marginale di un titolo:
    # abbastanza grande da rendere ammissibile l'ottimo, non tanto da congelare il simulated annealing
    P = 2.0 * (mu.max() / K + q * np.abs(sigma).sum(axis=1).max() / K ** 2)
    # E(x) = x^T Qm x + c^T x + cost, con x binario
    Qm = q * sigma / K ** 2 + P * np.ones((n, n))
    c = -mu / K - 2 * P * K * np.ones(n)
    cost = P * K * K
    return Qm, c, cost, K


def energia(X, Qm, c, cost):
    return np.einsum("ij,jk,ik->i", X, Qm, X) + X @ c + cost


def esatto(Qm, c, cost):
    n = len(c)
    best = math.inf
    tutte = []
    passo = 1 << 16
    for start in range(0, 1 << n, passo):
        idx = np.arange(start, min(start + passo, 1 << n))
        X = ((idx[:, None] >> np.arange(n)) & 1).astype(float)
        e = energia(X, Qm, c, cost)
        tutte.append(e)
        best = min(best, e.min())
    tutte = np.concatenate(tutte)
    mult = int(np.sum(tutte <= best + 1e-9))
    return best, mult, tutte


def local_search(X, Qm, c, cost):
    """Best-improvement a un bit; restituisce energie finali e valutazioni consumate."""
    X = X.copy()
    n = X.shape[1]
    valut = np.zeros(X.shape[0])
    diag = np.diag(Qm)
    attivi = np.ones(X.shape[0], dtype=bool)
    while attivi.any():
        Xa = X[attivi]
        campo = 2 * Xa @ Qm - 2 * Xa * diag + c        # dE/dx_i senza termine diagonale
        delta = (1 - 2 * Xa) * (campo + diag)          # variazione di energia del flip di i
        valut[attivi] += n
        j = np.argmin(delta, axis=1)
        migliora = delta[np.arange(len(j)), j] < -1e-12
        righe = np.where(attivi)[0]
        X[righe[migliora], j[migliora]] = 1 - X[righe[migliora], j[migliora]]
        attivi[righe[~migliora]] = False
    return energia(X, Qm, c, cost), valut


def simulated_annealing_swap(Qm, c, cost, R, rng, scala, K, sweep):
    """Simulated annealing con mosse di swap: parte da portafogli ammissibili
    casuali e propone di scambiare un titolo incluso con uno escluso, così che il
    vincolo di cardinalità resti soddisfatto e la penalità non crei barriere."""
    n = len(c)
    X = np.zeros((R, n))
    for r in range(R):
        X[r, rng.choice(n, K, replace=False)] = 1.0
    diag = np.diag(Qm)
    betas = np.geomspace(0.1, 1000.0, sweep) / scala
    ar = np.arange(R)
    for b in betas:
        for _ in range(n):
            # un titolo incluso i e uno escluso j per ogni replica
            pi = rng.random((R, n)) * X
            i = np.argmax(pi, axis=1)
            pj = rng.random((R, n)) * (1 - X)
            j = np.argmax(pj, axis=1)
            campo = 2 * X @ Qm - 2 * X * diag + c
            dE = -(campo[ar, i] + diag[i]) + (campo[ar, j] + diag[j]) - 2 * Qm[i, j]
            acc = (dE <= 0) | (rng.random(R) < np.exp(-np.clip(b * dE, 0, 700)))
            X[ar[acc], i[acc]] = 0.0
            X[ar[acc], j[acc]] = 1.0
    return energia(X, Qm, c, cost), sweep * n


def r99(p):
    """Tentativi per il 99% di successo, mai sotto 1 (un tentativo è il minimo eseguibile)."""
    if p <= 0:
        return math.inf
    if p >= 1:
        return 1.0
    return max(1.0, math.log(0.01) / math.log(1 - p))


def scala_campionata(Qm, c, cost, n, K, rng, m=2000):
    """Deviazione standard dell'obiettivo su portafogli ammissibili casuali
    (nessuna informazione dall'enumerazione esatta)."""
    X = np.zeros((m, n))
    for r in range(m):
        X[r, rng.choice(n, K, replace=False)] = 1.0
    return float(np.std(energia(X, Qm, c, cost)))


SWEEP = (1, 2, 5, 10, 20, 50)


rng = np.random.default_rng(RNG_SEME)
righe = []
tab = []
dist = None
for n in range(8, 23, 2):
    tts = {"casuale": [], "ls": [], "sa": []}
    sweep_opt = []
    for k in range(ISTANZE):
        Qm, c, cost, K = istanza(n, rng)
        e0, mult, tutte = esatto(Qm, c, cost)
        idx_opt = int(np.argmin(tutte))
        assert bin(idx_opt).count("1") == K, "ottimo non ammissibile: penalità troppo bassa"
        scala = scala_campionata(Qm, c, cost, n, K, rng)
        # casuale: probabilità esatta mult/2^n, costo di un tentativo = n valutazioni
        p_r = mult / 2 ** n
        tts["casuale"].append(r99(p_r) * n)
        X0 = rng.integers(0, 2, size=(CAMPIONI, n)).astype(float)
        e_ls, val = local_search(X0, Qm, c, cost)
        p_ls = float(np.mean(e_ls <= e0 + 1e-9))
        tts["ls"].append(r99(p_ls) * (n + val.mean()) if p_ls > 0 else math.inf)
        # TTS ottimo sul numero di sweep: il minimo su una griglia dichiarata
        migliore, sw_opt = math.inf, None
        for sw in SWEEP:
            e_sa, costo_sa = simulated_annealing_swap(Qm, c, cost, 400, rng, scala, K, sw)
            p_sa = float(np.mean(e_sa <= e0 + 1e-9))
            t = r99(p_sa) * costo_sa if p_sa > 0 else math.inf
            if t < migliore:
                migliore, sw_opt = t, sw
        tts["sa"].append(migliore)
        sweep_opt.append(sw_opt)
        if n == 20 and k == 0:
            e_r = energia(X0, Qm, c, cost)
            rif = abs(e0)
            e_sa10, _ = simulated_annealing_swap(Qm, c, cost, CAMPIONI, rng, scala, K, 10)
            dist = (e_r - e0) / rif, (e_ls - e0) / rif, (e_sa10 - e0) / rif
    med = {m: float(np.median(v)) for m, v in tts.items()}
    righe.append([n] + [math.log10(med[m]) for m in ("casuale", "ls", "sa")])
    sw_med = int(np.median(sweep_opt))
    print(f"n={n}: TTS mediano (valutazioni) uniforme {med['casuale']:.3g}, "
          f"uniforme+local search {med['ls']:.3g}, SA swap {med['sa']:.3g} (sweep ottimi mediani {sw_med})")
    tab.append([str(n), it_sci(med["casuale"], 1), it_sci(med["ls"], 1), it_sci(med["sa"], 1), str(sw_med)])

scrivi_dat("c10_tts.dat",
           "n log10_tts_uniforme log10_tts_uniforme_ls log10_tts_sa_swap  (mediana su 8 istanze; unita: valutazioni; SA con sweep ottimi)",
           righe)
# fattore di crescita per titolo: pendenza della regressione di log10 TTS su n
cresc = []
for col in (1, 2, 3):
    pend = np.polyfit([r[0] for r in righe], [r[col] for r in righe], 1)[0]
    cresc.append(10 ** pend)
tab.append(["crescita per titolo", it(cresc[0], 2), it(cresc[1], 2), it(cresc[2], 2), "--"])
print("fattore di crescita per titolo (uniforme, uniforme+LS, SA):", [round(x, 3) for x in cresc])
scrivi_txt("c10_tabella_tts.txt",
           ["n", "campionamento uniforme", "uniforme + local search", "simulated annealing (swap)", "sweep ottimi"],
           tab, "TTS al 99% in valutazioni di energia, mediana su 8 istanze, R99 >= 1; SA: minimo su sweep 1, 2, 5, 10, 20, 50 (400 repliche ciascuno); ultima riga: 10^(pendenza) della regressione log10 TTS ~ n")

# istogramma del divario relativo, n = 20
bordi = np.linspace(0, 0.5, 26)
h = [np.histogram(np.clip(d, 0, 0.5 - 1e-12), bins=bordi)[0] / len(d) for d in dist]
scrivi_dat("c10_distribuzione.dat",
           "divario_relativo frazione_uniforme frazione_uniforme_ls frazione_sa10  (n=20, istanza 1; SA con 10 sweep; divario = (E-E0)/|E0|)",
           [[bordi[i], h[0][i], h[1][i], h[2][i]] for i in range(len(bordi) - 1)] +
           [[bordi[-1], 0, 0, 0]])
print(f"n=20: frazione all'ottimo uniforme {np.mean(dist[0] <= 1e-9):.4f}, uniforme+LS {np.mean(dist[1] <= 1e-9):.4f}, "
      f"SA(10 sweep) {np.mean(dist[2] <= 1e-9):.4f}")