import nibabel as nib
import matplotlib.pyplot as plt
import numpy as np
import os

def load_nifti(file_path):
    """載入NIfTI檔案"""
    return nib.load(file_path).get_fdata()

def create_colored_mask(data):
    """創建彩色遮罩"""
    color_dict = {
        1: [144/255, 238/255, 144/255],  # Brainstem - 綠色
        2: [255/255, 218/255, 150/255],  # Right_Eye - 淺黃色
        3: [205/255, 170/255, 125/255],  # Left_Eye - 棕色
        4: [135/255, 206/255, 235/255],  # Optic_Chiasm - 藍色
        5: [255/255, 99/255, 71/255],    # Right_Optic_Nerve - 亮珊瑚紅
        6: [255/255, 160/255, 122/255]   # Left_Optic_Nerve - 淺珊瑚紅
    }
    
    mask = np.zeros((*data.shape, 4))
    
    for label, color in color_dict.items():
        organ_mask = (data == label)
        if np.any(organ_mask):
            mask[organ_mask] = [*color, 1.0]
    
    return mask

def get_brain_bounds(image):
    """獲取腦部區域的邊界，並調整顯示位置以避開底部基架"""
    # 基本閾值處理
    mask = image > np.percentile(image, 1)
    
    # 獲取非零區域的邊界
    rows = np.any(mask, axis=1)
    cols = np.any(mask, axis=0)
    
    # 確保有足夠的點
    if not np.any(rows) or not np.any(cols):
        return 0, image.shape[0], 0, image.shape[1]
        
    row_indices = np.where(rows)[0]
    col_indices = np.where(cols)[0]
    
    # 獲取初始邊界
    rmin, rmax = row_indices[[0, -1]]
    cmin, cmax = col_indices[[0, -1]]
    
    # 計算中心點，但略微上移以避開基架
    center_r = (rmin + rmax) // 2 - int(image.shape[0] * 0.15)  # 上移15%
    center_c = (cmin + cmax) // 2
    
    # 計算顯示範圍
    height = image.shape[0]
    width = image.shape[1]
    display_size = int(min(height, width) * 0.7)  # 使用70%的圖像大小
    
    # 設置新的邊界，確保不會超出圖像範圍
    rmin = max(center_r - display_size//2, 0)
    rmax = min(center_r + display_size//2, height - int(height * 0.2))  # 留出底部空間
    cmin = max(center_c - display_size//2, 0)
    cmax = min(center_c + display_size//2, width)
    
    # 如果上邊界太小，適當下移整個顯示區域
    if rmin < height * 0.1:
        shift = int(height * 0.1) - rmin
        rmin += shift
        rmax += shift
    
    return rmin, rmax, cmin, cmax

def apply_brain_window(image, window_width=80, window_level=40):
    """應用brain window設置"""
    window_min = window_level - window_width/2
    window_max = window_level + window_width/2
    windowed = np.clip(image, window_min, window_max)
    windowed = (windowed - window_min) / (window_max - window_min)
    return windowed

def check_organs_in_slice(slice_data):
    """檢查切片中存在的器官標籤"""
    unique_labels = set(np.unique(slice_data))
    if 0 in unique_labels:
        unique_labels.remove(0)
    organ_dict = {
        1: "Brainstem",
        2: "Right_Eye",
        3: "Left_Eye",
        4: "Optic_Chiasm",
        5: "Right_Optic_Nerve",
        6: "Left_Optic_Nerve"
    }
    return {organ_dict[label] for label in unique_labels if label in organ_dict}

def save_slice(image, ground_truth, prediction, slice_num, save_path):
    """保存指定切片的比較圖"""
    plt.clf()
    # Create figure with 1 row and 3 columns
    fig, axes = plt.subplots(1, 3, figsize=(24, 8))
    plt.subplots_adjust(wspace=0.01)
    
    img_slice = np.rot90(image[:, :, slice_num])
    rmin, rmax, cmin, cmax = get_brain_bounds(img_slice)
    
    # 使用brain window設置
    img_slice = apply_brain_window(img_slice)
    
    # Original Image
    axes[0].imshow(img_slice[rmin:rmax, cmin:cmax], cmap='gray')
    axes[0].set_title('Original Image', fontsize=32, pad=20, weight='bold')
    axes[0].axis('off')
    
    # Ground Truth
    gt_slice = np.rot90(ground_truth[:, :, slice_num])
    gt_mask = create_colored_mask(gt_slice)
    
    axes[1].imshow(img_slice[rmin:rmax, cmin:cmax], cmap='gray', zorder=1)
    axes[1].imshow(gt_mask[rmin:rmax, cmin:cmax], zorder=2)
    axes[1].set_title('Ground Truth', fontsize=32, pad=20, weight='bold')
    axes[1].axis('off')
    
    # Prediction
    pred_slice = np.rot90(prediction[:, :, slice_num])
    pred_mask = create_colored_mask(pred_slice)
    
    axes[2].imshow(img_slice[rmin:rmax, cmin:cmax], cmap='gray', zorder=1)
    axes[2].imshow(pred_mask[rmin:rmax, cmin:cmax], zorder=2)
    axes[2].set_title('Prediction', fontsize=32, pad=20, weight='bold')
    axes[2].axis('off')
    
    plt.savefig(save_path, bbox_inches='tight', dpi=300, pad_inches=0.05)
    plt.close()
    print(f"已保存切片 {slice_num} 至: {save_path}")

if __name__ == "__main__":
    base_dir = "/mnt/1248/onlylian"
    image_path = '/mnt/1248/onlylian/nnUNet/nnUNet_raw/Dataset666/imagesTs/CTC_001_0000.nii.gz'
    gt_path = os.path.join(base_dir, "nnUNet/nnUNet_raw/Dataset666/labelsTs/CTC_001.nii.gz")
    pred_path = '/mnt/1248/onlylian/CTC/output_predictions/2d/fold_0/CTC_001.nii.gz'
    save_dir = os.path.join(base_dir, "CTC")
    
    os.makedirs(save_dir, exist_ok=True)
    
    print("正在載入影像...")
    image = load_nifti(image_path)
    ground_truth = load_nifti(gt_path)
    prediction = load_nifti(pred_path)
    print("影像載入完成")
    
    print("\n分析切片中的器官...")
    total_slices = image.shape[2]
    for slice_num in range(total_slices):
        gt_slice = ground_truth[:, :, slice_num]
        pred_slice = prediction[:, :, slice_num]
        
        gt_organs = check_organs_in_slice(gt_slice)
        pred_organs = check_organs_in_slice(pred_slice)
        
        if gt_organs or pred_organs:
            print(f"\n切片 {slice_num}:")
            print(f"Ground Truth 包含: {', '.join(sorted(gt_organs))}")
            print(f"Prediction 包含: {', '.join(sorted(pred_organs))}")
    
    print("\n開始切片選擇和保存過程...")
    while True:
        try:
            slice_num = int(input("\n請輸入要保存的切片編號 (-1 退出): "))
            if slice_num == -1:
                print("程序結束")
                break
            if 0 <= slice_num < total_slices:
                save_path = os.path.join(save_dir, f'comparison_CTC_001_{slice_num:03d}.png')
                save_slice(image, ground_truth, prediction, slice_num, save_path)
            else:
                print(f"切片編號必須在 0 到 {total_slices-1} 之間")
        except ValueError:
            print("請輸入有效的數字")
        except Exception as e:
            print(f"發生錯誤: {str(e)}")