from typing import Callable, Optional

import torch

from vicentin.utils.logging import info, debug


def conjugate_gradient(
    A: Callable[[torch.Tensor], torch.Tensor],
    b: torch.Tensor,
    x0: Optional[torch.Tensor] = None,
    max_iter: int = 100,
    tol: float = 1e-8,
    return_loss: bool = False,
):
    if x0 is None:
        x = torch.zeros_like(b)
    else:
        x = x0.clone().detach()

    info(f"Starting Conjugate Gradient [Torch] on {b.device}")

    with torch.no_grad():
        r = b - A(x)
        p = r.clone()
        rs_old = torch.dot(r, r)

        loss = [rs_old.item()]

        for i in range(max_iter):
            Ap = A(p)

            alpha = rs_old / (torch.dot(p, Ap) + 1e-12)
            x = x + alpha * p
            r = r - alpha * Ap

            rs_new = torch.dot(r, r)
            loss.append(rs_new.item())

            res_norm = torch.sqrt(rs_new).item()
            debug(f"Iter {i}: Residual Norm = {res_norm:.4e}")

            if res_norm < tol:
                info(f"CG converged at iteration {i}")
                break

            beta = rs_new / (rs_old + 1e-12)
            p = r + beta * p
            rs_old = rs_new

    return (x, loss) if return_loss else x
