#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Défaut 2 de la méthode du sous-gradient : le pas doit tendre vers 0.

Même problème LASSO que pour le défaut 3 :
    min_x  F(x) = 1/2 ||A x - z||^2 + lambda ||x||_1
avec A in R^{30 x 10} aléatoire et une solution à 3 composantes non nulles.

On compare la méthode du sous-gradient avec
  - un pas constant gamma_k = gamma ;
  - un pas décroissant gamma_k = gamma_0 / sqrt(k+1).

Figure : à gauche, une composante nulle à la solution au fil des itérations ;
à droite, l'écart F(x^[k]) - F* en échelle log-log.
"""
import numpy as np
import matplotlib.pyplot as plt


# ------------------------------------------------------------------
# Données (identiques au script du défaut 3)
# ------------------------------------------------------------------
rng = np.random.default_rng(0)
M, N = 30, 10
A = rng.standard_normal((M, N)) / np.sqrt(M)

x_vrai = np.zeros(N)
x_vrai[[1, 4, 7]] = [1.5, -2.0, 1.0]
z = A @ x_vrai + 0.05 * rng.standard_normal(M)

lam = 0.3
L = np.linalg.norm(A, 2) ** 2


def F(x):
    return 0.5 * np.sum((A @ x - z) ** 2) + lam * np.sum(np.abs(x))


def grad_f(x):
    return A.T @ (A @ x - z)


def proxl1(v, t):
    return np.sign(v) * np.maximum(np.abs(v) - t, 0.0)


# Solution de référence et valeur optimale F*
x_star = np.zeros(N)
for _ in range(20000):
    x_star = proxl1(x_star - grad_f(x_star) / L, lam / L)
F_star = F(x_star)
n = np.flatnonzero(x_star == 0)[0]  # une composante nulle à la solution


# ------------------------------------------------------------------
# Méthode du sous-gradient
# ------------------------------------------------------------------
def sous_gradient(pas, K, x0):
    """Renvoie l'historique des itérées et des valeurs F(x^[k])."""
    x = x0.copy()
    hist_x = np.zeros((K, N))
    hist_F = np.zeros(K)
    for k in range(K):
        hist_x[k] = x
        hist_F[k] = F(x)
        u = grad_f(x) + lam * np.sign(x)  # un sous-gradient de F en x
        x = x - pas(k) * u
    return hist_x, hist_F


def gradient_proximal(K, x0):
    """Gradient proximal (ISTA), pas 1/L : seuillage doux après un pas de gradient."""
    x = x0.copy()
    hist_x = np.zeros((K, N))
    hist_F = np.zeros(K)
    for k in range(K):
        hist_x[k] = x
        hist_F[k] = F(x)
        x = proxl1(x - grad_f(x) / L, lam / L)
    return hist_x, hist_F


K = 5000
x0 = 0.5 * np.ones(N)
gamma = 0.1   # pas constant
gamma0 = 0.5  # pas décroissant : gamma0 / sqrt(k+1)

x_cst, F_cst = sous_gradient(lambda k: gamma, K, x0)
x_dec, F_dec = sous_gradient(lambda k: gamma0 / np.sqrt(k + 1), K, x0)
x_prox, F_prox = gradient_proximal(K, x0)

ecart_cst = F_cst - F_star
ecart_dec = F_dec - F_star
# l'écart du gradient proximal atteint la précision machine : on le borne pour l'échelle log
ecart_prox = np.maximum(F_prox - F_star, 1e-16)
print(f"F* = {F_star:.6f}")
print(f"pas constant    : F(x^[K]) - F* = {ecart_cst[-1]:.2e}, "
      f"min_k = {ecart_cst.min():.2e}")
print(f"pas décroissant : F(x^[K]) - F* = {ecart_dec[-1]:.2e}, "
      f"min_k = {ecart_dec.min():.2e}")
k_prox = np.argmax(F_prox - F_star < 1e-10)
print(f"gradient proximal : F(x^[k]) - F* < 1e-10 dès k = {k_prox}")


# ------------------------------------------------------------------
# Figure
# ------------------------------------------------------------------
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 4))

ks = np.arange(1, K + 1)
ax1.axhline(0, color="gray", lw=0.8)
ax1.plot(ks, x_cst[:, n], color="tab:red", lw=0.8,
         label=rf"pas constant $\gamma={gamma}$")
ax1.plot(ks, x_dec[:, n], color="tab:blue", lw=0.8,
         label=rf"pas décroissant $\gamma_k={gamma0}/\sqrt{{k+1}}$")
ax1.plot(ks, x_prox[:, n], color="tab:green", lw=1.5,
         label=r"gradient proximal, pas $1/L$")
ax1.set_xscale("log")
ax1.set_ylim(-0.08, 0.2)
ax1.set_xlabel("$k$")
ax1.set_ylabel(f"$x_{{{n + 1}}}^{{[k]}}$")
ax1.set_title("Evolution d'une composante nulle à la solution")
ax1.legend()

ax2.loglog(ks, ecart_cst, color="tab:red", lw=1, alpha=0.8,
           label=rf"pas constant $\gamma={gamma}$")
ax2.loglog(ks, ecart_dec, color="tab:blue", lw=1, alpha=0.8,
           label=rf"pas décroissant $\gamma_k={gamma0}/\sqrt{{k+1}}$")
ax2.loglog(ks, ecart_prox, color="tab:green", lw=1.5,
           label=r"gradient proximal, pas $1/L$")
ax2.set_xlabel("$k$")
ax2.set_ylim(1e-10, 10)
ax2.set_ylabel(r"$F(x^{[k]}) - F^\star$")
ax2.set_title("Écart à la valeur optimale")
ax2.legend(loc="lower right")

# Parcimonie : nombre de composantes non nulles des itérées
nnz_cst = np.count_nonzero(x_cst, axis=1)
nnz_dec = np.count_nonzero(x_dec, axis=1)
nnz_prox = np.count_nonzero(x_prox, axis=1)
nnz_star = np.count_nonzero(x_star)
ax3.step(ks, nnz_cst, where="post", color="tab:red", lw=1.3,
         label=rf"pas constant $\gamma={gamma}$")
ax3.step(ks, nnz_dec, where="post", color="tab:blue", lw=1.3, ls="--",
         label=rf"pas décroissant $\gamma_k={gamma0}/\sqrt{{k+1}}$")
ax3.step(ks, nnz_prox, where="post", color="tab:green", lw=1.5,
         label=r"gradient proximal, pas $1/L$")
ax3.axhline(nnz_star, color="gray", ls=":", lw=1, label=f"support de $x^\\star$ ({nnz_star})")
ax3.set_xscale("log")
ax3.set_ylim(0, N + 1)
ax3.set_yticks([0, nnz_star, N])
ax3.set_xlabel("$k$")
ax3.set_ylabel(r"$\#\{n : x_n^{[k]} \neq 0\}$")
ax3.set_title("Parcimonie des itérées")
ax3.legend(loc="center right", fontsize=8)

fig.tight_layout()
fig.savefig("defaut2_pas.png", dpi=150)
plt.show()

# ------------------------------------------------------------------
# Signal original et estimations après 5, 10 et 100 itérations
# ------------------------------------------------------------------
iters = [5, 10, 100]
idx = np.arange(1, N + 1)
methodes = [
    (None, -0.27, "k", "o", "signal original"),
    (x_prox, -0.09, "tab:green", "s", r"gradient proximal, pas $1/L$"),
    (x_dec, 0.09, "tab:blue", "v", r"sous-gradient, pas décroissant"),
    (x_cst, 0.27, "tab:red", "^", rf"sous-gradient, pas constant $\gamma={gamma}$"),
]

fig3, axes = plt.subplots(1, len(iters), figsize=(15, 4), sharey=True)
for ax, k in zip(axes, iters):
    ax.axhline(0, color="gray", lw=0.8)
    for hist, dx, col, mk, lab in methodes:
        x_aff = x_vrai if hist is None else hist[k]
        ml, sl, _ = ax.stem(idx + dx, x_aff, basefmt=" ", label=lab)
        plt.setp(ml, color=col, marker=mk, markersize=5)
        plt.setp(sl, color=col, linewidth=1.1)
    ax.set_xticks(idx)
    ax.set_xlabel("indice $n$")
    ax.set_title(f"$k = {k}$")
axes[0].set_ylabel("$x_n^{[k]}$")
handles, labels = axes[0].get_legend_handles_labels()
fig3.legend(handles, labels, loc="lower center", ncol=4, fontsize=9, frameon=False)
fig3.tight_layout(rect=(0, 0.08, 1, 1))
fig3.savefig("defaut2_signaux.png", dpi=150)
plt.show()

print("\nSignal original :", np.round(x_vrai, 3))
for k in iters:
    print(f"\nk = {k}")
    print("  gradient proximal         :", np.round(x_prox[k], 3))
    print("  sous-gradient décroissant :", np.round(x_dec[k], 3))
    print("  sous-gradient constant    :", np.round(x_cst[k], 3))