from typing import Callable, Optional

import numpy as np

from vicentin.optimization.minimization import proximal_gradient_descent
from vicentin.utils.logging import info, debug


def soft_thresholding(x: np.ndarray, threshold: float) -> np.ndarray:
    return np.sign(x) * np.maximum(np.abs(x) - threshold, 0.0)


def ista(
    f: Callable,
    grad_f: Callable,
    x0: np.ndarray,
    lambda_: float,
    step_size: Optional[float] = None,
    max_iter: int = 100,
    tol: float = 1e-6,
    epsilon: float = 1e-8,
    return_loss: bool = False,
):
    info(f"Initializing ISTA [NumPy] with lambda={lambda_:.4e}")

    def g(x):
        return lambda_ * np.linalg.norm(x, ord=1)

    def prox_g(z, gamma):
        threshold = gamma * lambda_
        debug(f"Applying soft-thresholding (threshold={threshold:.4e})")
        return soft_thresholding(z, threshold)

    F = (f, grad_f)
    G = (g, prox_g)

    return proximal_gradient_descent(
        F, G, x0, step_size, max_iter, tol, epsilon, return_loss
    )
