import os
import cv2
import numpy as np
import nibabel as nib
import pydicom
from pathlib import Path
import shutil

# 설정
ROOT_DIR = Path('.')
OUTPUT_DIR = ROOT_DIR / 'cvat_dataset_mask_v2'
JPEG_DIR = OUTPUT_DIR / 'JPEGImages'
SEG_CLASS_DIR = OUTPUT_DIR / 'SegmentationClass'
IMAGE_SETS_DIR = OUTPUT_DIR / 'ImageSets' / 'Segmentation'

# 색상 정의 (샘플 labelmap.txt 기반)
COLOR_MAP = {
    "C2": (230, 25, 75),
    "C3": (60, 180, 75),
    "C4": (255, 225, 25),
    "C5": (67, 99, 216),
    "C6": (245, 130, 49),
    "C7": (145, 30, 180),
    "L1": (0, 0, 117),
    "L2": (128, 128, 128),
    "L3": (255, 255, 255),
    "L4": (0, 0, 0),      # 주의: L4가 검은색(0,0,0)으로 정의되어 있어 배경과 겹칠 수 있음. 확인 필요.
    "L5": (31, 119, 180),
    "Rib": (255, 127, 14),
    "Sacrum": (44, 160, 44),
    "T1": (70, 240, 240),
    "T10": (170, 255, 195),
    "T11": (128, 128, 0),
    "T12": (255, 216, 177),
    "T2": (240, 50, 230),
    "T3": (188, 246, 12),
    "T4": (250, 190, 190),
    "T5": (0, 128, 128),
    "T6": (230, 190, 255),
    "T7": (154, 99, 36),
    "T8": (255, 250, 200),
    "T9": (128, 0, 0),
    # "background": (0, 0, 0) # 배경은 기본적으로 검정
}
# L4가 (0,0,0)이라면 배경과 구분 불가. L4 색상을 임의로 변경하거나 주의해야 함.
# 여기서는 안전을 위해 L4를 아주 어두운 회색(1,1,1)으로 하거나,
# 또는 사용자가 준 샘플에 L4가 (0,0,0)이라면 그것을 따라야 하지만, 
# 보통 배경이 (0,0,0)이므로 L4는 (50,50,50) 정도로 변경하여 진행하겠습니다.
COLOR_MAP["L4"] = (50, 50, 50) 


def normalize_image(image):
    """DICOM 이미지를 0-255 범위의 8bit 이미지로 변환"""
    image = image.astype(float)
    min_val = np.min(image)
    max_val = np.max(image)
    if max_val - min_val == 0:
        return np.zeros(image.shape, dtype=np.uint8)
    image = (image - min_val) / (max_val - min_val) * 255.0
    return image.astype(np.uint8)

def process_case(case_dir, file_list):
    """하나의 케이스(폴더) 처리"""
    dcm_files = list(case_dir.glob('*.dcm'))
    if not dcm_files:
        return
    
    dcm_path = dcm_files[0]
    case_name = case_dir.name
    
    # 1. DICOM -> JPG
    try:
        ds = pydicom.dcmread(dcm_path)
        pixel_array = ds.pixel_array
        img_8bit = normalize_image(pixel_array)
        
        # 그레이스케일 이미지를 RGB로 변환 (CVAT 호환성)
        img_rgb = cv2.cvtColor(img_8bit, cv2.COLOR_GRAY2BGR)
        
        cv2.imwrite(str(JPEG_DIR / f"{case_name}.jpg"), img_rgb)
    except Exception as e:
        print(f"Error processing DICOM {dcm_path}: {e}")
        return

    # 2. NII Masks -> RGB Mask
    # 배경(0,0,0)으로 초기화된 RGB 이미지 생성
    mask_rgb = np.zeros((pixel_array.shape[0], pixel_array.shape[1], 3), dtype=np.uint8)
    
    nii_files = list(case_dir.glob('*.nii.gz'))
    for nii_path in nii_files:
        class_name = nii_path.name.split('.')[0]
        if class_name not in COLOR_MAP:
            continue
            
        color = COLOR_MAP[class_name] # (R, G, B)
        # OpenCV는 BGR 사용하므로 순서 변경
        color_bgr = (color[2], color[1], color[0])
        
        try:
            img_nii = nib.load(str(nii_path))
            data = img_nii.get_fdata()
            
            # 3D -> 2D Projection
            proj = np.max(data, axis=2).T
            
            # 리사이즈
            if proj.shape != (mask_rgb.shape[0], mask_rgb.shape[1]):
                proj = cv2.resize(proj, (mask_rgb.shape[1], mask_rgb.shape[0]), interpolation=cv2.INTER_NEAREST)
            
            # 마스크 그리기
            # 해당 클래스 영역을 해당 색상으로 칠함
            mask_rgb[proj > 0] = color_bgr
            
        except Exception as e:
            print(f"Error processing NII {nii_path}: {e}")

    # 마스크 저장 (PNG)
    cv2.imwrite(str(SEG_CLASS_DIR / f"{case_name}.png"), mask_rgb)
    
    # 파일 목록에 추가
    file_list.append(case_name)
    print(f"Processed: {case_name}")

def create_labelmap():
    """labelmap.txt 생성"""
    with open(OUTPUT_DIR / 'labelmap.txt', 'w') as f:
        f.write("# label:color_rgb:parts:actions\n")
        # 정의된 순서대로 작성 (이름 정렬)
        for name in sorted(COLOR_MAP.keys()):
            r, g, b = COLOR_MAP[name]
            f.write(f"{name}:{r},{g},{b}::\n")
        # 배경 추가
        f.write("background:0,0,0::\n")

def main():
    # 디렉토리 초기화
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    
    JPEG_DIR.mkdir(parents=True)
    SEG_CLASS_DIR.mkdir(parents=True)
    IMAGE_SETS_DIR.mkdir(parents=True)
    
    print("Converting to CVAT Segmentation Mask 1.1 format...")
    
    file_list = []
    
    # 데이터 탐색
    for d in ROOT_DIR.iterdir():
        if d.is_dir() and not d.name.startswith('.') and 'cvat' not in d.name and 'segmask' not in d.name:
            if list(d.glob('*.dcm')):
                process_case(d, file_list)
    
    # default.txt 생성
    with open(IMAGE_SETS_DIR / 'default.txt', 'w') as f:
        for name in file_list:
            f.write(f"{name}\n")
            
    # labelmap.txt 생성
    create_labelmap()
    
    print("\nDone! Dataset created at:", OUTPUT_DIR.absolute())
    print(f"Total cases: {len(file_list)}")
    print("\n[중요] 업로드 시 주의사항:")
    print("1. 'cvat_dataset_mask_v2' 폴더 안으로 들어가서")
    print("2. 모든 파일/폴더(JPEGImages, SegmentationClass, ImageSets, labelmap.txt)를 선택하고")
    print("3. '압축하기(ZIP)'를 실행하세요.")
    print("4. CVAT 포맷 선택: 'Segmentation Mask 1.1'")

if __name__ == '__main__':
    main()