"""`plot` module contains many useful _plotting_ functions."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/P4_HPD_cg/90_plot.ipynb.

# %% auto 0
__all__ = ['epsilon', 'show_image_grid', 'show_compared_images', 'show_trajectory', 'gen_radial_traj']

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

# %% ../nbs/P4_HPD_cg/90_plot.ipynb 4
import numpy as np
import matplotlib.pyplot as plt

# %% ../nbs/P4_HPD_cg/90_plot.ipynb 5
plt.rcParams.update({'font.size': 13.25})
epsilon = 1e-10

# %% ../nbs/P4_HPD_cg/90_plot.ipynb 7
def show_image_grid(
    images:np.ndarray,      # 3D-array (Nimgs, Nx, Ny) images   
    figsize:float=(10,5),   # size of the figure (hor, ver)
    rows:int=2,             # number of rows
    cols:int=5,             # number of columns
    cmap:str='gray',        # colormap type
    vmin:float=None,        # minimum value for the colorbar()
    vmax:float=None,        # maximum value for the colorbar()
    labels:str=None,        # (Nimgs)-size list of labels 
    title:str=None,         # title of the figure
    axis_on:bool=False,     # flag to turn the axis on/off
    colorbar_on:bool=True,  # flag to turn the colorbar on/off
    calc_flag:str='abs',    # flag to run abs, real, imag, or angle
    log_flag:str=None,      # flag to display as log or log10
    last_image:bool=False): # flag to only display the last image
    """
    Display multiple images in a (rows,cols) grid.
    """
    
    if labels == None:
        labels = ['#'+str(idx+1) for idx in (range(rows*cols))]
    
    images = np.squeeze(images)
    if images.ndim==2: images=images[np.newaxis]
    if images.shape[0]==1: last_image=True
    if last_image: 
        images = images[-1]
        images = images[np.newaxis]
        figsize = (10,10)
        labels = None
        rows = 1
        cols=1
    
    if images.shape[0]>rows*cols:
        print(f"Warning: number of images ({images.shape[0]}) is larger than number of panels ({rows}x{cols})!")    
    fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=figsize)
    if title == None: pass
    else: fig.suptitle(title)
    
    if calc_flag=='abs': images = np.abs(images)
    elif calc_flag=='real': images = np.real(images)
    elif calc_flag=='imag': images = np.imag(images)
    elif calc_flag=='angle': images = np.angle(images)
    
    if log_flag=='log10': images = np.log10(images)
    elif log_flag=='log': images = np.log(images)

    # Calculate the min and max values of all images
    if vmin is None: vmin = np.min(images)
    if vmax is None: vmax = np.max(images)
    
    if rows==1 and cols==1: idx_ax_iter = enumerate([axes])
    else: idx_ax_iter =  enumerate(axes.flatten())
    
    for i, ax in idx_ax_iter:
        if i < len(images):
            img = ax.imshow(images[i], cmap=cmap, vmin=vmin, vmax=vmax)
            if labels is not None: ax.set_title(labels[i])
            if axis_on==False: ax.axis('off')
        else:
            if axis_on==False: ax.axis('off')
    if colorbar_on:
        cax = fig.add_axes([0.95, 0.12, 0.025, 0.75])  # [left, bottom, width, height]
        plt.colorbar(img, cax=cax)
    else:
        plt.tight_layout()

# %% ../nbs/P4_HPD_cg/90_plot.ipynb 16
def show_compared_images(
    X:np.ndarray,           # float-type 2D reference image
    Y:np.ndarray,           # float-type 2D compared image 
    figsize:float=(10,4),   # size of the figure (hor, ver)
    cmap:str='gray',        # colormap type
    vmin:float=None,        # minimum value for the colorbar()
    vmax:float=None,        # maximum value for the colorbar()
    labels:str=None,        # labels for X, Y, and X-Y
    title:str=None,         # title of the figure
    axis_on:bool=False,     # flag to turn the axis on/off
    colorbar_on:bool=True,  # flag to turn the colorbar on/off
    abs_calc:bool=False,     # flag to calculate abs or not
    diff_fac:float= 10):    # multiplication factor for (Y-X)
    """ Compare X and Y images by showing X, Y, and Y-X"""
    # concatinate individual image to images
    if abs_calc:
        images = np.stack((np.abs(X), np.abs(Y), diff_fac*np.abs(Y-X)))
    else:
        images = np.stack((X, Y, diff_fac*(Y-X)))
        
    # Calculate the min and max values of all images
    if vmin is None:
        vmin = np.min(images)
    if vmax is None:
        vmax = np.max(images)
    
    if labels == None:
        labels = ["X", "Y", "diff"+" (x"+str(diff_fac)+")"]
    else:
        labels[2] = labels[2] +" (x"+str(diff_fac)+")"
    
    rows, cols = 1,3
    fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=figsize)
    if title == None:
        pass
    else: 
        fig.suptitle(title)

    for i, ax in enumerate(axes.flatten()):
        if i < len(images):
            img = ax.imshow(images[i], cmap=cmap, vmin=vmin, vmax=vmax)
            ax.set_title(labels[i])
            if axis_on==False: ax.axis('off')
        else:
            if axis_on==False: ax.axis('off')
    if colorbar_on:
        cax = fig.add_axes([0.95, 0.15, 0.02, 0.7])  # [left, bottom, width, height]
        plt.colorbar(img, cax=cax)
    else:
        plt.tight_layout()

# %% ../nbs/P4_HPD_cg/90_plot.ipynb 21
def show_trajectory(
    traj:np.ndarray,        # kspace trajectory shaped Nspoke x Nreadout x 3 (x,y,z)
    figsize:float=(10,10),   # size of the figure (hor, ver)
    golden:bool=False):     # flag to plot golden angle radial trajectory   
    """show radial kspace trajectory using plt.plot"""
    fig, ax = plt.subplots(figsize=figsize)
    num_spokes = traj.shape[0]
    colormap = plt.cm.plasma
    for ns in range(num_spokes):
        kx = np.squeeze(traj[ns,:,0])
        ky = np.squeeze(traj[ns,:,1])
        if golden: 
            ax.plot(kx, ky, color=colormap(ns/num_spokes), lw=1.5)
            ax.plot(kx, ky, color=colormap(ns/num_spokes), lw=1.5)
        else:
            angle = np.arctan2(ky[-1] - ky[0], kx[-1] - kx[0])
            ax.plot(kx, ky, color=colormap(angle/np.pi), lw=1.5)
    if golden:
        cbar = plt.colorbar(plt.cm.ScalarMappable(norm=plt.Normalize(0, num_spokes), cmap=colormap), ax=ax)
        cbar.set_label('Trajectory Number')
    else:
        cbar = plt.colorbar(plt.cm.ScalarMappable(norm=plt.Normalize(0, np.pi), cmap=colormap), ax=ax)
        cbar.set_label('Angle (radians)')

# %% ../nbs/P4_HPD_cg/90_plot.ipynb 23
def gen_radial_traj(
    coords:int=(64, 256),       # coordinates (#spokes, #readout)
    golden:bool=False,          # flag to generate golden angle trajectory
    full_spoke:bool=True):      # flag to generate full/half spoke 
    "generating radial trajectories"
    nspokes, nreads = coords
    traj = np.zeros((nspokes, nreads, 2))
    if full_spoke: 
        if golden: dtheta = np.pi*(111.25/180) # Li Feng et al
        else: dtheta = np.pi/nspokes
        idxs, radius = np.mgrid[0:nspokes, -0.5:0.5:1.0/nreads]
    else:
        if golden: dtheta = 2*np.pi*(137.51/360) # Li Feng et al
        else:dtheta = 2*np.pi/nspokes
        idxs, radius = np.mgrid[0:nspokes, 0.0:0.5:0.5/nreads]
    thetas = idxs * dtheta    
    traj[:, :, 0] = radius * np.cos(thetas) # x = r*cos(th)
    traj[:, :, 1] = radius * np.sin(thetas) # y = r*sin(th)
    return traj
