import SimpleITK as sitk
import numpy as np
import os

def create_new_nifti_mask(output_nii_path, size=(256, 256, 128), mask_type="sphere"):
    """
    创建一个指定大小的新NIFTI掩码文件
    
    参数:
        output_nii_path: 输出的伪掩码路径
        size: 掩码的尺寸，默认为(256, 256, 128)
        mask_type: 掩码类型，可选"sphere"、"cube"或"random"
    """
    print(f"创建尺寸为{size}的伪掩码")
    
    # 创建一个指定尺寸的零数组
    mask_array = np.zeros(size, dtype=np.uint8)
    
    if mask_type == "sphere":
        # 创建一个球形掩码
        center = [s // 2 for s in size]
        radius = min(size) // 4
        
        # 创建坐标网格
        z, y, x = np.ogrid[:size[0], :size[1], :size[2]]
        
        # 计算到中心的距离
        dist_from_center = np.sqrt(
            (z - center[0])**2 + 
            (y - center[1])**2 + 
            (x - center[2])**2
        )
        
        # 创建掩码
        mask_array = (dist_from_center <= radius).astype(np.uint8)
        
    elif mask_type == "cube":
        # 创建一个立方体掩码
        center = [s // 2 for s in size]
        half_width = min(size) // 6
        
        mask_array[
            center[0]-half_width:center[0]+half_width,
            center[1]-half_width:center[1]+half_width,
            center[2]-half_width:center[2]+half_width
        ] = 1
        
    elif mask_type == "random":
        # 创建一个具有随机形状的掩码
        np.random.seed(42)  # 设置随机数种子以便结果可复现
        
        # 创建几个随机椭球
        for _ in range(3):
            center = [np.random.randint(s//4, 3*s//4) for s in size]
            radii = [s//10 + np.random.randint(s//8) for s in size]
            
            z, y, x = np.ogrid[:size[0], :size[1], :size[2]]
            
            dist_from_center = np.sqrt(
                ((z - center[0])/radii[0])**2 + 
                ((y - center[1])/radii[1])**2 + 
                ((x - center[2])/radii[2])**2
            )
            
            mask_array = np.logical_or(mask_array, dist_from_center <= 1)
    
    # 转换为SimpleITK图像
    mask_image = sitk.GetImageFromArray(mask_array.astype(np.uint8))
    
    # 设置默认的元数据
    mask_image.SetSpacing((1.0, 1.0, 1.0))  # 1mm 各向同性间距
    mask_image.SetOrigin((0.0, 0.0, 0.0))    # 默认原点
    direction = tuple([1.0 if i == j else 0.0 for i in range(3) for j in range(3)])  # 默认方向
    mask_image.SetDirection(direction)
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_nii_path), exist_ok=True)
    
    # 保存掩码
    sitk.WriteImage(mask_image, output_nii_path)
    print(f"伪掩码已保存至: {output_nii_path}")
    
    return mask_array

# 使用相对路径
input_path = "test_data/kidney_right/AMOS/imagesVal/test.nii.gz"
output_path = "test_data/kidney_right/AMOS/labelsVal/test.nii.gz"

# 确保父目录存在
os.makedirs(os.path.dirname(output_path), exist_ok=True)

# 创建新的掩码（尺寸可根据需要调整）
create_new_nifti_mask(output_path, size=(128, 128, 128), mask_type="sphere")