"""`optimizer` module contains the steepest descent and conjugate gradient methods."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/P4_HPD_cg/95_optim.ipynb.

# %% auto 0
__all__ = ['frobenius_norm', 'naive_steepest_descent', 'steepest_descent', 'conjugate_gradient']

# %% ../nbs/P4_HPD_cg/95_optim.ipynb 2
from nbdev.showdoc import *
from fastcore.test import *
from fastcore.utils import *

# %% ../nbs/P4_HPD_cg/95_optim.ipynb 3
import numpy as np
import h5py
import matplotlib.pyplot as plt

# %% ../nbs/P4_HPD_cg/95_optim.ipynb 4
from .plot import *
from .metrics import *
from .math import *
from .io import *
import cgsense2023 as cgs2003
from .mri import *

# %% ../nbs/P4_HPD_cg/95_optim.ipynb 6
def frobenius_norm(
    A:np.ndarray,   # input matrix A 
)->float:           # Frobenius norm of matrix A
    "Frobenius norm of a matrix A defined as the square root of the sum of the absolute squares of its elements"
    return np.sqrt(np.sum(np.square(np.abs(A))))

# %% ../nbs/P4_HPD_cg/95_optim.ipynb 16
def naive_steepest_descent(
    mriop,                  # MRI operator NuFFT, etc.
    guess:np.ndarray,       # Initial Guess of Zeros
    data:np.ndarray,        # Dens_compensated kspace
    iters:int=10,           # Number of Iterations
    lr:float=0.4,           # Learning Rate
    tol:float=0.0,          # Tolerance to terminate
    lambd:float=0.0,        # Tikhonov Regularization Param
    kfilter:bool=True,      # flag to turn kspace filter True/False,
    dstore: bool=False,     # flag to store recon at every iter
    ref:np.ndarray=None,    # Ground truth reference image
)-> tuple:                  # return (final_im, res_norms, ref_norms)
    "naive gradient descent method"
    
    # number of readouts point & image dimension
    nreads, image_dim = data.shape[-1], guess.shape[-1]
    # We start with an initial guess of zeros..
    xs = []
    x = guess
    # calculate b (lhs)
    b = mriop.operator_rhs(data)
    # residual == gradient w/o regularization
    res = mriop.operator_lhs(x) - b
    # residual norms
    res_norms = []
    # reference norms
    if ref is not None: ref_norms = []
    for i in range(iters):
        # regularization
        gamma = res + lambd * x
        # update
        x = x - lr * gamma
        xs.append(x)
        # residual
        res = mriop.operator_lhs(x) - b
        delta = frobenius_norm(res)
        res_norms.append(delta)
        # reference norms
        if ref is not None: ref_norms.append(frobenius_norm(x-ref))
        if delta < tol:
            print("Converged after %i iterations to %1.6e" % (i + 1, delta))
            break
        if not np.mod(i, 1):
            print("Residuum at iter %i : %1.6e" % (i + 1, delta), end='\r')
    
    xs = np.array(xs)      
    if kfilter: xs = kspace_filter(xs, nreads, image_dim)
    
    if ref is not None:
        return xs if dstore else xs[-1], res_norms, ref_norms
    else:
        return xs if dstore else xs[-1], res_norms

# %% ../nbs/P4_HPD_cg/95_optim.ipynb 19
def steepest_descent(
    mriop,                  # MRI operator NuFFT, etc.
    guess:np.ndarray,       # Initial Guess of Zeros
    data:np.ndarray,        # Dens_compensated kspace
    iters:int=10,           # Number of Iterations
    tol:float=0.0,          # Tolerance to terminate
    lambd:float=0.0,        # Tikhonov Regularization Param
    kfilter:bool=True,      # flag to turn kspace filter True/False
    dstore: bool=False,     # flag to store recon at every iter
    ref:np.ndarray=None,    # Ground truth reference image
)-> tuple:                  # return (final_im, res_norms, ref_norms)
    "naive gradient descent method"
    # number of readouts point & image dimension
    nreads, image_dim = data.shape[-1], guess.shape[-1]
    # We start with an initial guess of zeros..
    xs = []
    x = guess
    # calculate b (lhs)
    b = mriop.operator_rhs(data)
    # residual norms
    res_norms = []
    # res_norms.append(frobenius_norm(mriop.operator_lhs(x)-b))
    # reference norms
    if ref is not None: ref_norms = []
    for i in range(iters):
        # gamma - gradient
        gamma = mriop.operator_lhs(x) - b
        # H-gamma
        Hgamma = mriop.operator_lhs(gamma)
        # alpha - optimal learning rate
        alpha = np.vdot(gamma, gamma) / (np.vdot(gamma, Hgamma))
        # update
        x = x - alpha * (gamma + lambd*x)
        xs.append(x)
        # residual
        delta = frobenius_norm(mriop.operator_lhs(x) - b)
        res_norms.append(delta)
        # reference norms
        if ref is not None: ref_norms.append(frobenius_norm(x-ref))
        if delta < tol:
            print("Converged after %i iterations to %1.6e" % (i + 1, delta))
            break
        if not np.mod(i, 1):
            print("Residuum at iter %i : %1.6e" % (i + 1, delta), end='\r')
    xs = np.array(xs)
    if kfilter: xs = kspace_filter(xs, nreads, image_dim)
    if ref is not None:
        return xs if dstore else xs[-1], res_norms, ref_norms
    else:
        return xs if dstore else xs[-1], res_norms

# %% ../nbs/P4_HPD_cg/95_optim.ipynb 22
def conjugate_gradient(
    mriop,                  # MRI operator NuFFT, etc.
    guess:np.ndarray,       # Initial Guess of Zeros
    data:np.ndarray,        # Dens_compensated kspace
    iters:int=10,           # Number of Iterations
    tol:float=0.0,          # Tolerance to terminate
    lambd:float=0.0,        # Tikhonov Regularization Param
    kfilter:bool=True,      # flag to turn kspace filter True/False
    dstore: bool=False,     # flag to store recon at every iter
    ref:np.ndarray=None,    # Ground truth reference image
)-> tuple:                  # return (final_im, res_norms, ref_norms)
    "naive gradient descent method"
    # number of readouts point & image dimension
    nreads, image_dim = data.shape[-1], guess.shape[-1]
    # We start with an initial guess of zeros..
    xs = []
    x = guess
    # calculate b (lhs)
    b = mriop.operator_rhs(data)
    # residual norms
    gamma = mriop.operator_lhs(x) - b
    d = gamma
    res_norms = []
    # res_norms.append(frobenius_norm(gamma))
    # reference norms
    if ref is not None: ref_norms = []
    for i in range(iters):
        # Ad
        Ad = mriop.operator_lhs(d)
        # regularization
        Ad = Ad + lambd * d
        # alpha
        alpha = -np.vdot(gamma, d) / np.vdot(d, Ad)
        # update
        x = x + alpha * d
        xs.append(x)
        # new gamma
        gamma = mriop.operator_lhs(x) - b
        # residual
        delta = frobenius_norm(gamma)
        res_norms.append(delta)
        # reference norms
        if ref is not None: ref_norms.append(frobenius_norm(x-ref))
        if delta < tol:
            print("Converged after %i iterations to %1.6e" % (i + 1, delta))
            break
        if not np.mod(i, 1):
            print("Residuum at iter %i : %1.6e" % (i + 1, delta), end='\r')
        # beta
        beta = -np.vdot(gamma, Ad) / np.vdot(d, Ad)
        # new d
        d = gamma + beta * d
    xs = np.array(xs)
    if kfilter: xs = kspace_filter(xs, nreads, image_dim)
    if ref is not None:
        return xs if dstore else xs[-1], res_norms, ref_norms
    else:
        return xs if dstore else xs[-1], res_norms
