
#DISPLAY BRAIN PHANTOM
'''
在医学成像（如 MRI、CT）和信号处理中，幻影（Phantom）是指用于测试、校准或研究的人工数据或物理模型。
它通常是一个合成（模拟）图像或实验装置，用来模拟真实世界的结构，比如大脑、心脏或其他人体组织。
'''
# 1. 导入必要的库

%matplotlib inline

import os.path as op
import os
import math
import cmath
import sys

import numpy as np
import pywt as pw
import matplotlib.pyplot as plt

from skimage import data, io, filters       # 用于图像处理。
import brainweb_dl as bwdl                  # 用于下载和处理脑部 MRI 幻影数据。

# 2. 设置 Matplotlib 显示参数
plt.rcParams["image.origin"]="lower"        # 设置图像的原点在左下角（默认是左上角）
plt.rcParams["image.cmap"]='Greys_r'        # 设置图像颜色映射为灰度反转（黑白）。

# get current working dir
#cwd = os.getcwd()
#dirimg_2d = op.join(cwd,"../data")
#img_size = 512   #256

#load data file corresponding to the target resolution
#filename = "BrainPhantom%s.png" % img_size
#mri_filename = op.join(dirimg_2d, filename)
#mri_img = io.imread(mri_filename)

# 3. 下载并处理 MRI 图像
mri_img = bwdl.get_mri(4, "T1")[70, ...].astype(np.float32)
'''
bwdl.get_mri(4, "T1") : 从 BrainWeb 数据库下载一个 T1 加权 MRI 扫描
70 : 获取 MRI 体数据（3D）中的第 70 层切片，即 2D 切片图像。
.astype(np.float32) : 转换数据类型为 float32，以便进行计算。
''' 
#mri_img = bwdl.get_mri(4, "T2")[120, ...].astype(np.float32)

# 4. 显示 MRI 图像
print(mri_img.shape)                        # (256, 256)
img_size = mri_img.shape[0]                 # 获取图像的宽度（假设是正方形图像）。

# 5. 绘制 MRI 图像
plt.figure()
plt.imshow(abs(mri_img))                    # 显示 MRI 图像，abs() 主要用于防止负像素值影响显示。
#plt.imshow(abs(mri_2D),origin="lower")         # lower : 图像的原点（0,0）位于左下角（默认）。
#plt.imshow(abs(mri_2D),origin="upper")         # upper : 图像的原点（0,0）位于左上角（常见于 Matplotlib 默认设置）
plt.title("Original brain image")
plt.show()

(256, 256)


# 1. 导入必要的库
# Load target sampling distribution (precalculated in Matlab)
import numpy as np
from scipy.io import loadmat            # 加载 .mat 文件，读取预先计算好的目标采样分布（Matlab 文件格式）。
import mrinufft.trajectories as mt      # 用于生成核磁共振成像（MRI）中的采样轨迹与优化采样密度。

# 2. 左图 : 生成优化采样密度 (opt_density)
'''
(img_size, img_size)：图像大小，二维矩阵表示。
"db4"：Daubechies 小波的 db4 作为基函数。
3：小波分解的分辨率层级。
'''
# Generate optimal sampoling density for a given wavelet transform, number of resolution levels and img size 
# see details Chauffert et al, IEEE ISBI 2013 for the computation of optimal sampling densities
#opt_density = mt.create_fast_chauffert_density((img_size,img_size),"sym10",3)
opt_density = mt.create_fast_chauffert_density((img_size,img_size),"db4",3)

# 3. 生成基于优化采样密度的 k-space 采样掩码 (kspace_mask)
# Generate Cartesian variable density mask
# change the value below if you want to change the final subsampling mask
threshold = 10. * opt_density.min()  # sys.float_info.epsilon \simeq 2e-16      # 阈值 (threshold)
'''
np.zeos 函数：创建 kspace_mask，
np.where 函数：将 opt_density 值大于阈值的位置设为 1，其余位置设为 0。 右图 : 最终生成一个二值采样掩码。
'''
kspace_mask = np.zeros((img_size, img_size), dtype="float64")
kspace_mask = np.where(opt_density > threshold, 1, kspace_mask)                 

# 4. 可视化 优化的采样密度 (opt_density) 与 二值采样掩码 (kspace_mask)
fig, axs = plt.subplots(1, 2, figsize=(7, 4) )
axs[0].imshow(opt_density)
axs[0].set_title("Optimal variable density\nfor Shannon wavelets")
axs[1].imshow(kspace_mask)
axs[1].set_title("Variable density sampling mask")

Text(0.5, 1.0, 'Variable density sampling mask')


# 5. 定义 Fourier 变换和逆变换

#import numpy.fft as fft

norm = "ortho"      # 正交归一化，保持傅里叶变换的幅值与图像幅值一致
#norm = None

def fft(x):
    return np.fft.fft2(x, norm=norm)

def ifft(x):
    return np.fft.ifft2(x, norm=norm)


# 6. 计算 k-space 数据并加噪
# Generate the kspace data: first Fourier transform the image
kspace_data = np.fft.fftshift(fft(mri_img))
#add Gaussian complex-valued random noise
signoise = 10       # 添加高斯复值噪声（独立添加到实部和虚部）
kspace_data += np.random.randn(*mri_img.shape) * signoise * (1+1j)

# 7. 使用采样掩码进行次采样
# Mask data to perform subsampling
kspace_data *= kspace_mask

# 8. 进行 MRI 重建
# Zero order solution
image_rec0 = ifft(np.fft.ifftshift(kspace_data))

# 9. 可视化结果
fig, axs = plt.subplots(2, 2, figsize=(7, 7) )
# 显示原始 MRI 图像
axs[0, 0].imshow(mri_img)
axs[0, 0].set_title("True image")
# 显示采样掩码
axs[0, 1].imshow(kspace_mask)
axs[0, 1].set_title("Sampling mask")
# 显示被噪声影响的 k-space 数据
axs[1,0].imshow(np.abs(kspace_data), vmax=0.01*np.abs(kspace_data).max())
axs[1, 0].set_title("k-space noisy data")
# 显示重建的图像
axs[1, 1].imshow(np.abs(image_rec0))
axs[1, 1].set_title("Zero-order recon")
plt.show()


def Reconstruction_v1 (mask, threshold_level, signoise):
    norm = "ortho"
    #norm = None

    def fft(x):
        return np.fft.fft2(x, norm=norm)

    def ifft(x):
        return np.fft.ifft2(x, norm=norm)
    
    if mask == "db4":
        # 2. 左图 : 生成优化采样密度 (opt_density)
        '''
        (img_size, img_size)：图像大小，二维矩阵表示。
        "db4"：Daubechies 小波的 db4 作为基函数。
        3：小波分解的分辨率层级。
        '''
        # Generate optimal sampoling density for a given wavelet transform, number of resolution levels and img size 
        # see details Chauffert et al, IEEE ISBI 2013 for the computation of optimal sampling densities
        #opt_density = mt.create_fast_chauffert_density((img_size,img_size),"sym10",3)
        opt_density = mt.create_fast_chauffert_density((img_size,img_size),"db4",3)
        
        # 3. 生成基于优化采样密度的 k-space 采样掩码 (kspace_mask)
        # Generate Cartesian variable density mask
        # change the value below if you want to change the final subsampling mask
        threshold = threshold_level * opt_density.min()  # sys.float_info.epsilon \simeq 2e-16      # 阈值 (threshold)
        ## threshold = 10. * opt_density.min()  # sys.float_info.epsilon \simeq 2e-16      # 阈值 (threshold)
        '''
        np.zeos 函数：创建 kspace_mask，
        np.where 函数：将 opt_density 值大于阈值的位置设为 1，其余位置设为 0。 右图 : 最终生成一个二值采样掩码。
        '''
        kspace_mask = np.zeros((img_size, img_size), dtype="float64")
        kspace_mask = np.where(opt_density > threshold, 1, kspace_mask) 

        # Generate the kspace data: first Fourier transform the image
        kspace_data = np.fft.fftshift(fft(mri_img))

        #add Gaussian complex-valued random noise
        # signoise = 10       # 添加高斯复值噪声（独立添加到实部和虚部）
        kspace_data += np.random.randn(*mri_img.shape) * signoise * (1+1j)

        # 7. 使用采样掩码进行次采样
        # Mask data to perform subsampling
        kspace_data *= kspace_mask

        # 8. 进行 MRI 重建
        # Zero order solution
        image_rec0 = ifft(np.fft.ifftshift(kspace_data))

        return kspace_data, image_rec0


mask="db4"
vector_threshold_level = [0.1, 0.2, 14, 15, 20, 30]
vector_signoise = [0, 10, 20, 100]

fig, axs = plt.subplots(len(vector_threshold_level), 2, figsize=(12, 12))

for i in range(len(vector_threshold_level)) :
    for j in range(2):
        threshold_level = vector_threshold_level[i]
        signoise = vector_signoise[1]

        kspace_data, image_rec0 = Reconstruction_v1(mask, threshold_level, signoise)

        # 显示被噪声影响的 k-space 数据
        axs[i, 0].imshow(np.abs(kspace_data), vmax=0.01*np.abs(kspace_data).max())
        
        # 显示重建的图像
        axs[i, 1].imshow(np.abs(image_rec0))

        axs[i, j].set_xlabel(f"signoise = {signoise}")
        axs[i, j].xaxis.set_label_position('top')  # 移动 xlabel 到顶部
        axs[i, j].set_ylabel(f"threshold_level = {threshold_level}")
    

axs[0, 0].set_title("k-space noisy data")
axs[0, 1].set_title("Zero-order recon")

# Add a main title (suptitle) to the figure
fig.suptitle(f"k-space noisy data and its Reconstruction Results for Mask: {mask}", fontsize=16)

# Display the figure with adjusted layout
plt.tight_layout(rect=[0, 0, 1, 0.96])  # Adjust layout to fit suptitle
plt.show()


mask="db4"
vector_threshold_level = np.linspace(0, 20, 100)
vector_signoise = [0, 10, 20, 100]

vector_Frobenius_norm = []
vector_image_rec0 = []

#  mri_img : True image

for i in range(len(vector_threshold_level)) :
    
        threshold_level = vector_threshold_level[i]
        signoise = vector_signoise[1]

        kspace_data, image_rec0 = Reconstruction_v1(mask, threshold_level, signoise)
        
        vector_image_rec0.append(image_rec0)
        vector_Frobenius_norm.append(np.linalg.norm(image_rec0 - mri_img, 'fro'))

plt.figure(figsize=(10, 6))
plt.plot(vector_threshold_level, vector_Frobenius_norm, linestyle='-', color='red', marker='.', markersize=8, markerfacecolor='red', label='Frobenius Norm')

plt.xlabel('Threshold Level')
plt.ylabel('Frobenius Norm')
plt.title('Frobenius Norm vs Threshold Level')
plt.grid(True)
plt.legend()
plt.show()


# Find the index and value of the minimum Frobenius norm.
min_index = np.argmin(vector_Frobenius_norm)
image_corresponding_min_index = vector_image_rec0[min_index]
threshold_corresponding_level_min_index = vector_threshold_level[min_index]
# print(threshold_corresponding_level_min_index)

# Find the index corresponding to the Threshold Level closest to 15.
closest15_index = (np.abs(vector_threshold_level - 15)).argmin()
image_corresponding_closest15_index = vector_image_rec0[closest15_index]
threshold_level_corresponding_closest15_index = vector_threshold_level[closest15_index]
# print(threshold_level_corresponding_closest15_index)

fig, axes = plt.subplots(1, 3, figsize=(12, 4))

axes[0].imshow(np.abs(image_corresponding_min_index))
axes[0].set_title(f'Threshold Level (for Min Frobenius Norm) : {threshold_corresponding_level_min_index:.2f}')
axes[0].axis('off')

axes[1].imshow(mri_img)
axes[1].set_title("True image")
axes[1].axis('off')

axes[2].imshow(np.abs(image_corresponding_closest15_index))
axes[2].set_title(f'Threshold Level (≈ 15) : {threshold_level_corresponding_closest15_index:.2f}')
axes[2].axis('off')

plt.tight_layout()
plt.show()


def Reconstruction (mask, threshold_level, signoise, decay):

    # 1. 定义快速傅里叶变换 (FFT) 和 逆快速傅里叶变换 (IFFT)
    norm = "ortho"
    #norm = None

    def fft(x):
        return np.fft.fft2(x, norm=norm)

    def ifft(x):
        return np.fft.ifft2(x, norm=norm)
    
    if mask == "db4":
        # 2. 左图 : 生成优化采样密度 (opt_density) (生成最优采样密度 (Optimal Sampling Density))
        '''
        (img_size, img_size)：图像大小，二维矩阵表示。
        "db4"：Daubechies 小波的 db4 作为基函数。
        3：小波分解的分辨率层级。
        '''
        # Generate optimal sampoling density for a given wavelet transform, number of resolution levels and img size 
        # see details Chauffert et al, IEEE ISBI 2013 for the computation of optimal sampling densities
        #opt_density = mt.create_fast_chauffert_density((img_size,img_size),"sym10",3)
        opt_density = mt.create_fast_chauffert_density((img_size,img_size),"db4",3)

        # 3. 生成变量采样分布 (Variable Sampling Distribution)
        # Now construct by hands a variable sampling distribution
        # You can change the decay to modify the decreasing behavior in the center of k-space
        # the larger the decay, the faster the decrease from low to high-frequencies
        # decay = 2
        x = np.linspace(-1. / (2. * np.pi), 1. / (2. * np.pi), img_size)
        X, Y = np.meshgrid(x, x)
        r = np.sqrt(X ** 2 + Y ** 2)    
        # print(r)
        p_decay = np.power(r,-decay)        # p_decay 表示随径向距离的衰减函数, decay 控制衰减速度
        p_decay = p_decay/np.sum(p_decay)   # 衰减越快 (decay 越大)，k 空间低频采样会显著减少。

        #print(p_decay.max())
        #print(p_decay.min())

        # 4. 生成采样掩码 (Sampling Mask)
        # change the value below if you want to change the final subsampling mask
        threshold = threshold_level * opt_density.min()  # sys.float_info.epsilon \simeq 2e-16      # 阈值 (threshold)
        # threshold = 2* opt_density.min()  # sys.float_info.epsilon \simeq 2e-16
        '''
        np.zeos 函数：创建 kspace_mask，
        np.where 函数：将 opt_density 值大于阈值的位置设为 1，其余位置设为 0。 右图 : 最终生成一个二值采样掩码。
        '''
        kspace_mask = np.zeros((img_size,img_size), dtype="float64")
        kspace_mask = np.where(p_decay > threshold, 1, kspace_mask)

        # 5. 添加复值高斯噪声
        # Generate the kspace data: first Fourier transform the image
        kspace_data = np.fft.fftshift(fft(mri_img))

        #add Gaussian complex-valued random noise
        # signoise = 10       # 添加高斯复值噪声（独立添加到实部和虚部）
        # Simulate independent noise realization on the real & imag parts
        kspace_data += (np.random.randn(*mri_img.shape) + 1j * np.random.randn(*mri_img.shape)) * signoise

        # 7. 使用采样掩码进行次采样
        # Mask data to perform subsampling
        kspace_data *= kspace_mask

        # 8. 进行 MRI 重建
        # Zero order solution
        image_rec0 = ifft(np.fft.ifftshift(kspace_data))

        return kspace_data, image_rec0


mask="db4"
vector_threshold_level = [0.1, 0.2, 2, 14, 15, 20, 30]
vector_signoise = [0, 10, 20, 100]
vector_decay = [1, 2, 4, 7, 10]

fig, axs = plt.subplots(len(vector_decay), 2, figsize=(12, 12))

for i in range(len(vector_decay)) :
    for j in range(2):

        threshold_level = vector_threshold_level[2]
        signoise = vector_signoise[1]
        decay = vector_decay[i]

        kspace_data, image_rec0 = Reconstruction (mask, threshold_level, signoise, decay)

        # 显示被噪声影响的 k-space 数据
        axs[i, 0].imshow(np.abs(kspace_data), cmap='gray', vmax=0.01*np.abs(kspace_data).max())
        
        # 显示重建的图像
        axs[i, 1].imshow(np.abs(image_rec0), cmap='Greys_r')

        axs[i, j].set_xlabel(f"signoise = {signoise}")
        axs[i, j].xaxis.set_label_position('top')  # 移动 xlabel 到顶部
        axs[i, j].set_ylabel(f"decay = {decay}")
    

axs[0, 0].set_title("k-space noisy data")
axs[0, 1].set_title("Zero-order recon")

# Add a main title (suptitle) to the figure
fig.suptitle(f"k-space noisy data and its Reconstruction Results for Mask: {mask}", fontsize=16)

# Display the figure with adjusted layout
plt.tight_layout(rect=[0, 0, 1, 0.96])  # Adjust layout to fit suptitle
plt.show()


mask="db4"
vector_threshold_level = [0.1, 0.2, 2, 14, 15, 20, 30]
vector_signoise = [0, 10, 20, 100]
vector_decay = np.linspace(0, 10, 100)

vector_Frobenius_norm = []
vector_image_rec0 = []

#  mri_img : True image

for i in range(len(vector_decay)) :
    
        threshold_level = vector_decay[2]
        signoise = vector_signoise[1]
        decay = vector_decay[i]

        kspace_data, image_rec0 = Reconstruction (mask, threshold_level, signoise, decay)
        
        vector_image_rec0.append(image_rec0)
        vector_Frobenius_norm.append(np.linalg.norm(image_rec0 - mri_img, 'fro'))

plt.figure(figsize=(10, 6))
plt.plot(vector_decay, vector_Frobenius_norm, linestyle='-', color='red', marker='.', markersize=8, markerfacecolor='red', label='Frobenius Norm')

plt.xlabel('Decay')
plt.ylabel('Frobenius Norm')
plt.title('Frobenius Norm vs Decay')
plt.grid(True)
plt.legend()
plt.show()


# Find the index and value of the minimum Frobenius norm.
min_index = np.argmin(vector_Frobenius_norm)
image_corresponding_min_index = vector_image_rec0[min_index]
decay_corresponding_min_index = vector_decay[min_index]
# print(decay_corresponding_min_index)

fig, axes = plt.subplots(1, 2, figsize=(8, 4))

axes[0].imshow(np.abs(image_corresponding_min_index), cmap='Greys_r')
axes[0].set_title(f'Decay (for Min Frobenius Norm) : {decay_corresponding_min_index:.2f}')
axes[0].axis('off')

axes[1].imshow(mri_img, cmap='Greys_r')
axes[1].set_title("True image")
axes[1].axis('off')

plt.tight_layout()
plt.show()


# Now construct by hands a variable sampling distribution
# You can change the decay to modify the decreasing behavior in the center of k-space
# the larger the decay, the faster the decrease from low to high-frequencies
decay = 2
x = np.linspace(-1. / (2. * np.pi), 1. / (2. * np.pi), img_size)
X, Y = np.meshgrid(x, x)
r = np.sqrt(X ** 2 + Y ** 2)
print(r)
p_decay = np.power(r,-decay)
p_decay = p_decay/np.sum(p_decay)

#print(p_decay.max())
#print(p_decay.min())

# change the value below if you want to change the final subsampling mask
threshold = 2* opt_density.min()  # sys.float_info.epsilon \simeq 2e-16
kspace_mask = np.zeros((img_size,img_size), dtype="float64")
kspace_mask = np.where(p_decay > threshold, 1, kspace_mask)

fig, axs = plt.subplots(1, 2, figsize=(7, 4))
axs[0].imshow(p_decay, cmap='Greys_r', vmax=0.001 * np.abs(p_decay).max())
axs[0].set_title("Handcrafted VD")
axs[1].imshow(kspace_mask, cmap='Greys_r')
axs[1].set_title("VD sampling mask")

[[0.22507908 0.22419815 0.22332073 ... 0.22332073 0.22419815 0.22507908]
 [0.22419815 0.22331375 0.22243284 ... 0.22243284 0.22331375 0.22419815]
 [0.22332073 0.22243284 0.22154843 ... 0.22154843 0.22243284 0.22332073]
 ...
 [0.22332073 0.22243284 0.22154843 ... 0.22154843 0.22243284 0.22332073]
 [0.22419815 0.22331375 0.22243284 ... 0.22243284 0.22331375 0.22419815]
 [0.22507908 0.22419815 0.22332073 ... 0.22332073 0.22419815 0.22507908]]

Text(0.5, 1.0, 'VD sampling mask')


# Generate the kspace data: first Fourier transform the image
kspace_data = np.fft.fftshift(fft(mri_img))
#add Gaussian complex-valued random noise
signoise = 10
# Simulate independent noise realization on the real & imag parts
kspace_data += (np.random.randn(*mri_img.shape) + 1j * np.random.randn(*mri_img.shape)) * signoise
# Mask data to perform subsampling
kspace_data *= kspace_mask

# Zero order solution
image_rec0 = ifft(np.fft.ifftshift(kspace_data))

fig, axs = plt.subplots(2, 2, figsize=(7, 7))
axs[0, 0].imshow(mri_img, cmap='Greys_r')
axs[0, 0].set_title("True image")
axs[0, 1].imshow(kspace_mask, cmap='Greys_r')
axs[0, 1].set_title("Sampling mask")
axs[1, 0].imshow(np.abs(kspace_data),  cmap='gray', vmax=0.01*np.abs(kspace_data).max())
#axs[1].imshow(np.abs(np.fft.ifftshift(kspace_data)), cmap='Greys_r')
axs[1, 0].set_title("k-space noisy data")
axs[1, 1].imshow(np.abs(image_rec0), cmap='Greys_r')
axs[1, 1].set_title("Zero-order recon")
plt.show()

2. iid Variable Density Sampling¶