import numpy as np
import matplotlib.pyplot as plt
import pywt
from scipy.fft import fft2, ifft2, fftshift, ifftshift
import os


def fft2c(img):
    """Centered 2D FFT."""
    return np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(img)))


def ifft2c(kspace):
    """Centered 2D IFFT."""
    return np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(kspace)))


def soft_threshold(x, lam):
    """软阈值操作"""
    return np.sign(x) * np.maximum(np.abs(x) - lam, 0)


def total_variation(x, alpha):
    """计算图像的总变分"""
    dx = np.diff(x, axis=0)
    dy = np.diff(x, axis=1)
    return alpha * (np.sum(np.abs(dx)) + np.sum(np.abs(dy)))


def reconstruct_image(kspace_data):
    im = fftshift(ifft2(kspace_data, (256, 256)), axes=0)
    # 取模得到实值图像
    magnitude_image = np.abs(im)
    # 归一化图像
    normalized_image = (magnitude_image - magnitude_image.min()) / (magnitude_image.max() - magnitude_image.min())
    return normalized_image


def split_bregman_mri_reconstruction(kspace_sampled, mask, lam, num_iters=2000, step_size=0.05, tv_alpha=0.05):
    """
    使用分裂布雷格曼算法（Split Bregman）进行MRI重建，并引入总变分约束
    :param kspace_sampled: 采样的k空间数据
    :param mask: 采样掩码
    :param lam: 正则化参数
    :param num_iters: 迭代次数
    :param step_size: 步长
    :param tv_alpha: 总变分约束强度
    :return: 重建的图像
    """
    # 初始重建（零填充）
    img_reconstructed = reconstruct_image(kspace_sampled)

    # 显示初始化图像
    plt.imshow(np.abs(img_reconstructed), cmap='gray')
    plt.title('Initial Reconstructed MRI Image')
    plt.show()
    plt.imsave('initial_reconstructed_image.png', np.abs(img_reconstructed), cmap='gray')

    # 初始化布雷格曼变量
    b1 = np.zeros_like(img_reconstructed)
    b2 = np.zeros_like(img_reconstructed)

    for i in range(num_iters):
        # 数据一致性步骤
        kspace_reconstructed = fft2c(img_reconstructed)
        kspace_reconstructed = kspace_sampled * mask + kspace_reconstructed * (1 - mask)
        img_temp = reconstruct_image(kspace_reconstructed)

        # 计算总变分
        grad_tv = np.gradient(img_temp)
        grad_tv = np.sqrt(grad_tv[0] ** 2 + grad_tv[1] ** 2)

        # 更新图像
        img_reconstructed = img_temp - step_size * (tv_alpha * grad_tv + b1 + b2)

        # 更新布雷格曼变量
        b1 = b1 + step_size * (grad_tv - soft_threshold(img_reconstructed + b1, lam))
        b2 = b2 + step_size * (grad_tv - soft_threshold(img_reconstructed + b2, lam))

        if (i + 1) % 10 == 0 or i == 0:
            print(f"Iteration {i + 1}/{num_iters}")

    return img_reconstructed


def main():
    # 设置当前工作目录
    current_dir = os.getcwd()

    # 加载完整的k空间数据
    kspace_path = os.path.join(current_dir, 'fid_20250418_163142.npy')
    if not os.path.exists(kspace_path):
        print(f"错误: 找不到文件 {kspace_path}")
        return

    try:
        kspace_full = np.load(kspace_path)
        print(f"成功加载k空间数据: {kspace_path}")
    except Exception as e:
        print(f"错误: 加载k空间数据时出错: {e}")
        return

    # 设置分裂布雷格曼算法参数
    lam = 0.05  # 正则化参数，可以根据需要调整
    num_iters = 10000  # 迭代次数
    step_size = 0.05  # 步长，可以根据需要调整
    tv_alpha = 0.1  # 总变分约束强度，可以根据需要调整

    # 遍历当前目录下的所有.npy文件，排除'mrd_data.npy'
    for filename in os.listdir(current_dir):
        if filename.lower().endswith('.npy') and filename!= 'fid_20250418_163142.npy':
            mask_path = os.path.join(current_dir, filename)
            try:
                mask = np.load(mask_path)
                print(f"处理文件: {mask_path}")
            except Exception as e:
                print(f"错误: 加载文件 {mask_path} 时出错: {e}")
                continue

            # 检查kspace_full和mask的形状是否匹配
            if kspace_full.shape!= mask.shape:
                print(f"警告: 文件 {filename} 的形状与k空间数据不匹配，跳过。")
                continue

            # 采样k空间
            kspace_sampled = kspace_full * mask

            # 重建图像
            try:
                img_reconstructed = split_bregman_mri_reconstruction(kspace_sampled, mask, lam, num_iters, step_size, tv_alpha)
                print(f"成功重建图像: {filename}")
            except Exception as e:
                print(f"错误: 重建文件 {filename} 时出错: {e}")
                continue

            # 构造保存的PNG文件名
            base_name = os.path.splitext(filename)[0]
            png_filename = f"{base_name}_split_bregman_reconstructed_fid_20250417_103329.png"
            png_path = os.path.join(current_dir, png_filename)
            plt.imshow(img_reconstructed, cmap='gray')  # 设置为灰度图显示
            plt.title('Reconstructed MRI Image')
            plt.show()
            plt.imsave(png_path, np.abs(img_reconstructed), cmap='gray')


if __name__ == "__main__":
    main()
