from typing import Callable, Optional

import torch
from torch.func import jacrev

from vicentin.utils.logging import info, debug, warn


def prox_step(
    x: torch.Tensor,
    grad_x: torch.Tensor,
    prox_g: Callable,
    gamma: float,
):
    z = x - gamma * grad_x
    x_new = prox_g(z, gamma)
    return x_new, gamma


def prox_line_search_step(
    x: torch.Tensor,
    f: Callable,
    grad_x: torch.Tensor,
    prox_g: Callable,
    gamma: float,
):
    f_x = f(x)
    iters = 0

    while True:
        z = x - gamma * grad_x
        next_x = prox_g(z, gamma)

        diff = next_x - x
        sq_dist = torch.sum(diff**2)

        rhs = f_x + torch.dot(grad_x, diff.view(-1)) + sq_dist / (2 * gamma)

        if f(next_x) <= rhs or gamma < 1e-12:
            if gamma < 1e-12:
                warn("Proximal line search reached minimum step size limit.")
            else:
                debug(
                    f"Line search converged: gamma={gamma:.4e} after {iters} shrinks"
                )
            break

        gamma /= 2
        iters += 1

    gamma *= 1.2
    return next_x, gamma


def proximal_gradient(
    f,
    g,
    prox_g,
    x0: torch.Tensor,
    step_size: Optional[float] = None,
    max_iter: int = 100,
    tol: float = 1e-6,
    epsilon: float = 1e-8,
    return_loss: bool = False,
):
    x = x0.clone().detach()
    loss = []

    if step_size is None:
        gamma = 1.0
        info(
            f"Starting Proximal Gradient Descent [Torch] with Line Search on {x.device}"
        )
    else:
        gamma = step_size
        info(
            f"Starting Proximal Gradient Descent [Torch] with Fixed Step: {gamma} on {x.device}"
        )

    grad_f_func = jacrev(f)

    for i in range(max_iter):
        grad_x = grad_f_func(x)

        if isinstance(grad_x, (list, tuple)):
            grad_x = grad_x[0]

        with torch.no_grad():
            if step_size is None:
                x_new, gamma = prox_line_search_step(
                    x, f, grad_x, prox_g, gamma
                )
            else:
                x_new, gamma = prox_step(x, grad_x, prox_g, gamma)

            f_val, g_val = f(x_new), g(x_new)
            F_new = f_val + g_val
            loss.append(F_new.item())

            F_old = f(x) + g(x)
            norm_F = torch.abs(F_old - F_new)

            norm_grad = torch.linalg.norm((x - x_new) / gamma)

            info(
                f"Iter {i}: F(x)={F_new.item():.6e} [f={f_val:.4e}, g={g_val:.4e}] | G-Norm={norm_grad:.4e}"
            )

            x = x_new

            if norm_F < tol:
                info(
                    f"Success: Convergence reached via objective tolerance ({tol})."
                )
                break
            if norm_grad < epsilon:
                info(
                    f"Success: Convergence reached via proximal gradient norm ({epsilon})."
                )
                break
    else:
        warn(f"Reached max_iter ({max_iter}) without full convergence.")

    if return_loss:
        return x, loss

    return x
