import os
import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim
from concurrent.futures import ThreadPoolExecutor, as_completed

def load_images_from_folder(folder):
    """
    從指定文件夾加載灰度圖像。
    
    參數(shù):
    folder (str): 包含圖像的文件夾路徑。
    
    返回:
    dict: 以文件名為鍵，加載的圖像為值的字典。
    """
    images = {}
    for filename in os.listdir(folder):
        img_path = os.path.join(folder, filename)
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        if img is not None:
            images[filename] = img
    return images

def crop_center(img, crop_fraction=0.6):
    """
    裁剪圖像的中心區(qū)域。
    
    參數(shù):
    img (ndarray): 輸入圖像。
    crop_fraction (float): 要保留的圖像部分的比例。
    
    返回:
    ndarray: 裁剪后的圖像。
    """
    h, w = img.shape
    crop_h, crop_w = int(h * crop_fraction), int(w * crop_fraction)
    start_h, start_w = (h - crop_h) // 2, (w - crop_w) // 2
    return img[start_h:start_h + crop_h, start_w:start_w + crop_w]

def find_best_match(image, image_dict):
    """
    在圖像字典中找到與輸入圖像最相似的圖像。
    
    參數(shù):
    image (ndarray): 輸入圖像。
    image_dict (dict): 包含圖像的字典。
    
    返回:
    tuple: 最相似圖像的文件名和相似度值。
    """
    max_ssim = -1
    best_match = None
    cropped_image = crop_center(image)
    for filename, img in image_dict.items():
        cropped_img = crop_center(img)
        # 將裁剪后的輸入圖像調(diào)整到與裁剪后的比較圖像相同的大小
        resized_image = cv2.resize(cropped_image, (cropped_img.shape[1], cropped_img.shape[0]))
        current_ssim = ssim(resized_image, cropped_img)
        if current_ssim > max_ssim:
            max_ssim = current_ssim
            best_match = filename
    return best_match, max_ssim

def skip_black_background_image(image, threshold=0.06):
    """
    跳過黑色背景的圖像。
    
    參數(shù):
    image (ndarray): 輸入圖像。
    threshold (float): 非黑色像素的閾值比例。
    
    返回:
    bool: 如果圖像主要是黑色背景，則返回True，否則返回False。
    """
    # 計(jì)算非黑色像素的百分比
    num_non_black_pixels = np.sum(image > 0)
    total_pixels = image.size
    if num_non_black_pixels / total_pixels < threshold:
        return True
    return False

def process_patient_folder(patient_folder, train_folder, output_folder, similarity_threshold, file_index):
    """
    處理病人的文件夾，找到與每個(gè)T1w圖像最匹配的T1wCE圖像，并保存匹配對(duì)。
    
    參數(shù):
    patient_folder (str): 病人文件夾的名稱。
    train_folder (str): 訓(xùn)練數(shù)據(jù)的根文件夾。
    output_folder (str): 輸出文件夾路徑。
    similarity_threshold (float): 圖像相似度的閾值。
    file_index (int): 輸出文件的起始索引。
    
    返回:
    int: 更新后的文件索引。
    """
    patient_path = os.path.join(train_folder, patient_folder)
    t1w_folder = os.path.join(patient_path, 'T1w')
    t1wce_folder = os.path.join(patient_path, 'T1wCE')

    if not os.path.exists(t1w_folder) or not os.path.exists(t1wce_folder):
        return file_index

    t1w_images = load_images_from_folder(t1w_folder)
    t1wce_images = load_images_from_folder(t1wce_folder)

    total_t1w_images = len(t1w_images)
    unmatched_count = 0
    # 可選地，保存或處理匹配的圖像對(duì)
    patient_output_folder = os.path.join(output_folder, patient_folder)

    for t1w_filename, t1w_image in t1w_images.items():
        if skip_black_background_image(t1w_image):
            unmatched_count += 1
            continue

        best_match, max_ssim = find_best_match(t1w_image, t1wce_images)
        if max_ssim >= similarity_threshold:
            if not os.path.exists(patient_output_folder):
                os.makedirs(patient_output_folder)
            print(f'病人: {patient_folder}, T1w: {t1w_filename}, T1wCE: {best_match}, SSIM: {max_ssim}, 輸出文件: {file_index}_T1w.png')
            t1w_output_path = os.path.join(patient_output_folder, f"{file_index}_T1w.png")
            t1wce_output_path = os.path.join(patient_output_folder, f"{file_index}_T1wCE.png")
            cv2.imwrite(t1w_output_path, t1w_image)
            cv2.imwrite(t1wce_output_path, t1wce_images[best_match])
            file_index += 1

    return file_index

# 定義路徑
train_folder = './train'
output_folder = 'output_folder'
similarity_threshold = 0.75

if not os.path.exists(output_folder):
    os.makedirs(output_folder)

file_index = 1
patient_folders = [f for f in os.listdir(train_folder) if os.path.isdir(os.path.join(train_folder, f))]

with ThreadPoolExecutor(max_workers=6) as executor:
    future_to_patient = {
        executor.submit(process_patient_folder, patient_folder, train_folder, output_folder, similarity_threshold, file_index): patient_folder for patient_folder in patient_folders}
    for future in as_completed(future_to_patient):
        patient_folder = future_to_patient[future]
        try:
            file_index = future.result()
        except Exception as exc:
            print(f'{patient_folder} 生成異常: {exc}')
        else:
            print(f'{patient_folder} 處理完成。')