#!/usr/bin/env python3
"""
Capitolo 3: rumore e mitigazione degli errori.

Modello. Circuito di Ising a calcio su catena aperta di N = 10 qubit,
theta_J = -pi/2, theta_h = 0,6, osservabile Z_5. Dopo ogni passo agisce
un canale depolarizzante locale su ogni qubit; nell'immagine di Heisenberg
ogni stringa di Pauli di peso w è moltiplicata per f^w a ogni passo, con
f = 1 - 4p/3. Il calcolo è la propagazione di Pauli di comune.py con soglia
1e-7 (errore di troncamento sotto 1e-5, verificato contro il vettore di
stato nel caso senza rumore).

Estrapolazione a rumore nullo (ZNE): il rumore è amplificato del fattore
lambda (f -> f^lambda) e il valore a lambda = 0 è stimato per
estrapolazione lineare (lambda = 1, 2), di Richardson (lambda = 1, 2, 3) ed
esponenziale (lambda = 1, 2, 3).

Cancellazione probabilistica (PEC): per un modello di rumore di Pauli–Lindblad
sparso con somma dei tassi Lambda su tutto il circuito, il fattore di
campionamento è gamma = exp(2 Lambda) e il numero di shot cresce di
gamma^2 = exp(4 Lambda). Ipotesi di lavoro: Lambda = G * eps, con G gate a
due qubit di errore eps (si ignorano i gate a un qubit e il readout).

Produce:
    c03_rumore.dat         <Z_5> ideale e rumoroso per passo (f = 0,99 e 0,97)
    c03_zne.dat            <Z_5> al passo 6 in funzione di lambda, con le tre estrapolazioni
    c03_tabella_zne.txt    stime ZNE e loro distorsione
    c03_pec.dat            log10(gamma^2) in funzione del numero di gate, tre errori per gate
    c03_tabella_pec.txt    gamma^2 e shot per errore standard 0,01
"""

import math

import numpy as np

from comune import (catena, kicked_ising_esatto, propaga_pauli, scrivi_dat,
                    scrivi_txt, it, it_sci)

N = 10
Q = 5
TJ, TH = -math.pi / 2, 0.6
PASSI = 6
SOGLIA = 1e-7


def valori(f):
    return [r[1] for r in propaga_pauli(N, catena(N), TJ, TH, PASSI, Q,
                                        soglia=SOGLIA, smorzamento=f)]


ideale = valori(1.0)
esatto = kicked_ising_esatto(N, catena(N), TJ, TH, PASSI, Q)
err_tronc = max(abs(a - b) for a, b in zip(ideale, esatto))
print(f"errore di troncamento (senza rumore) contro il vettore di stato: {err_tronc:.1e}")
assert err_tronc < 1e-5

r99 = valori(0.99)
r97 = valori(0.97)
scrivi_dat("c03_rumore.dat",
           "passo ideale f099 f097  (N=10, theta_J=-pi/2, theta_h=0,6, <Z_5>)",
           [[t, ideale[t], r99[t], r97[t]] for t in range(PASSI + 1)])
print("ideale:", [round(v, 4) for v in ideale])
print("f=0,99:", [round(v, 4) for v in r99])
print("f=0,97:", [round(v, 4) for v in r97])

# ------------------------------------------------ ZNE al passo finale
F0 = 0.97
lams = [1.0, 1.5, 2.0, 2.5, 3.0]
mis = {lam: valori(F0 ** lam)[PASSI] for lam in lams}
y1, y2, y3 = mis[1.0], mis[2.0], mis[3.0]
lin = 2 * y1 - y2
rich = 3 * y1 - 3 * y2 + y3
# esponenziale y = A exp(-k lambda) + C non è identificabile con tre punti senza C;
# si usa y = A exp(-k lambda), stimato su (1, y1) e (2, y2), controllato su lambda = 3
k = math.log(y1 / y2)
A = y1 * math.exp(k)
expo = A
vero = ideale[PASSI]
righe = []
for j in range(0, 31):
    lam = j * 0.1
    righe.append([lam,
                  valori(F0 ** lam)[PASSI] if j % 5 == 0 else float("nan"),
                  y1 + (y1 - y2) * (1 - lam),
                  y1 * (lam - 2) * (lam - 3) / 2 - y2 * (lam - 1) * (lam - 3) + y3 * (lam - 1) * (lam - 2) / 2,
                  A * math.exp(-k * lam)])
scrivi_dat("c03_zne.dat",
           "lambda misurato lineare richardson esponenziale  (passo 6, f0=0,97; nan = non calcolato)",
           [[r[0]] + [("nan" if (isinstance(v, float) and math.isnan(v)) else v) for v in r[1:]] for r in righe])
print(f"ZNE: vero {vero:.4f}; misurato a lambda=1 {y1:.4f}; lineare {lin:.4f}; "
      f"Richardson {rich:.4f}; esponenziale {expo:.4f}; esponenziale a lambda=3 {A*math.exp(-3*k):.4f} contro {y3:.4f}")


def amplif(c):
    """Fattore di amplificazione della varianza dell'estrapolatore (somma dei quadrati dei pesi)."""
    return sum(x * x for x in c)


scrivi_txt("c03_tabella_zne.txt",
           ["metodo", "punti di rumore", "stima", "distorsione", "amplificazione della varianza"],
           [["nessuna mitigazione", "1", it(y1, 4), it(y1 - vero, 4), it(1, 0)],
            ["lineare", "1, 2", it(lin, 4), it(lin - vero, 4), it(amplif([2, -1]), 0)],
            ["Richardson", "1, 2, 3", it(rich, 4), it(rich - vero, 4), it(amplif([3, -3, 1]), 0)],
            ["esponenziale", "1, 2", it(expo, 4), it(expo - vero, 4),
             it(amplif([2 * y1 / y2, -(y1 / y2) ** 2]), 1)]],
           f"N = 10, passo 6, f0 = 0,97; valore esatto {vero:.4f}. Amplificazione: somma dei quadrati dei "
           "coefficienti con cui le misure entrano nella stima (per l'esponenziale A = y1^2/y2, linearizzata: pesi 2r e -r^2 con r = y1/y2)")
print(f"valore esatto al passo 6: {vero:.4f}")

# ------------------------------------------------ PEC
EPS = [1e-3, 1.5e-3, 2e-3]
righe = []
for G in range(0, 5001, 100):
    righe.append([G] + [4 * G * e / math.log(10) for e in EPS])
scrivi_dat("c03_pec.dat",
           "porte_2q log10_gamma2_eps1e-3 log10_gamma2_eps1.5e-3 log10_gamma2_eps2e-3  (gamma^2 = exp(4 G eps))",
           righe)
tab = []
for G in (500, 1000, 1488, 2500, 5000):
    g2 = math.exp(4 * G * 1.5e-3)
    tab.append([it(G, 0), it(G * 1.5e-3, 2), it_sci(g2, 1) if g2 >= 1e4 else it(g2, 0),
                it_sci(g2 / 0.01 ** 2, 1)])
    print(f"G={G}: Lambda={G*1.5e-3:.3f}, gamma^2={g2:.3g}, shot per sigma=0,01: {g2/1e-4:.3g}")
scrivi_txt("c03_tabella_pec.txt",
           ["gate a due qubit G", "Lambda = G eps", "gamma^2", "shot per errore standard 0,01"],
           tab, "eps = 1,5e-3 per porta; varianza del singolo esito <= 1; gamma^2 = exp(4 Lambda)")