from typing import Callable, Optional

import numpy as np

from vicentin.utils.logging import info, debug


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

    info(f"Starting Conjugate Gradient [NumPy] (max_iter={max_iter})")

    r = b - A(x)
    p = r.copy()
    rs_old = np.dot(r, r)

    loss = [rs_old]

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

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

        rs_new = np.dot(r, r)
        loss.append(rs_new)

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

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

        beta = rs_new / (rs_old + 1e-12)
        p = r + beta * p
        rs_old = rs_new
    else:
        info(f"CG reached max_iter ({max_iter}) without full convergence")

    return (x, loss) if return_loss else x
