from typing import Callable, Optional

import torch

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


def soft_thresholding(x: torch.Tensor, threshold: float) -> torch.Tensor:
    zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
    return torch.sign(x) * torch.maximum(torch.abs(x) - threshold, zero)


def ista(
    f: Callable,
    x0: torch.Tensor,
    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 [Torch] on {x0.device} with lambda={lambda_:.4e}")

    def g(x):
        return lambda_ * torch.norm(x, p=1)

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

    F = f
    G = (g, prox_g)

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