import os
import json
import cv2
from tqdm import tqdm
import numpy as np

# =============================================================================
# [사용자 설정]
# =============================================================================
# 통합된 JSON 파일 경로 (예: '/DATA2/ltb/mmyolo/combined_dataset_no_copy/combined.json')
json_file_path = '/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl_all_combined/val.json'

# 이미지 파일들은 여러 폴더에 흩어져 있으므로, JSON의 "file_name" 필드에
# 이미 절대경로가 기록되어 있다면 별도의 이미지 폴더 경로는 필요하지 않습니다.
# 만약 "file_name"이 상대경로라면, 아래 base_folder_path를 적절히 수정하세요.
base_folder_path = ''  # 비워두면 절대경로라고 가정함

# 라벨링된 이미지를 저장할 출력 폴더 (없으면 자동 생성)
output_folder = '/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl_all_combined/labeled_images_auto'
os.makedirs(output_folder, exist_ok=True)

# =============================================================================
# JSON 파일 로드 및 기본 정보 추출
# =============================================================================
with open(json_file_path, 'r', encoding='utf-8') as f:
    data = json.load(f)

# JSON 파일 내 "categories" 항목을 통해, 카테고리 ID와 이름을 매핑하는 딕셔너리 생성
category_id_to_name = {category['id']: category['name'] for category in data['categories']}

# =============================================================================
# 클래스별 색상 자동 할당
# =============================================================================
# JSON의 카테고리 수에 맞춰, 각 클래스에 대해 자동으로 색상을 부여합니다.
# 아래 함수는 카테고리 ID를 기반으로 한 간단한 수학 연산을 통해 deterministic(항상 동일한)
# 색상을 생성합니다.
def generate_color(category_id):
    # 각 채널의 색상값은 임의의 상수와 곱한 후 modulo 256으로 계산합니다.
    r = (category_id * 37) % 256
    g = (category_id * 73) % 256
    b = (category_id * 17) % 256
    return (r, g, b)

# category_id_to_name에 있는 모든 카테고리에 대해 색상을 할당합니다.
category_colors = {cat_id: generate_color(cat_id) for cat_id in category_id_to_name.keys()}

# =============================================================================
# JSON의 이미지 정보를 이용하여, 이미지 ID와 파일명을 매핑하는 딕셔너리 생성
# =============================================================================
# 주의: 통합된 JSON의 "images" 항목의 "file_name" 필드는
# 각 이미지의 원본 경로를 반영하고 있으므로, 그대로 사용합니다.
image_id_to_file = {image['id']: image['file_name'] for image in data['images']}

# 전체 이미지 개수
total_images = len(data['images'])

# =============================================================================
# 이미지별 어노테이션을 읽어와 바운딩 박스를 그리고, 라벨링된 이미지를 저장하는 부분
# =============================================================================
# tqdm 라이브러리를 이용하여 진행률을 표시합니다.
with tqdm(total=total_images, desc="Processing Images") as pbar:
    for image in data['images']:
        image_id = image['id']
        image_file_name = image['file_name']
        
        # 이미지 파일 경로 처리
        # 만약 image_file_name이 절대경로가 아니라면 base_folder_path와 결합
        if not os.path.isabs(image_file_name):
            image_path = os.path.join(base_folder_path, image_file_name)
        else:
            image_path = image_file_name
        
        # 이미지 로드 (cv2.imread는 이미지가 없거나 읽지 못할 경우 None 반환)
        img = cv2.imread(image_path)
        if img is None:
            print(f"이미지 로드 실패: {image_path}")
            pbar.update(1)
            continue
        
        # 해당 이미지의 모든 어노테이션(바운딩 박스) 정보를 가져옴
        annotations = [ann for ann in data['annotations'] if ann['image_id'] == image_id]
        
        # 각 어노테이션에 대해 바운딩 박스와 클래스 레이블 그리기
        for annotation in annotations:
            bbox = annotation['bbox']  # [x, y, width, height] 형식
            category_id = annotation['category_id']
            category_name = category_id_to_name.get(category_id, 'Unknown')
            # category_colors에서 해당 클래스의 색상을 가져오며, 기본값은 빨간색
            color = category_colors.get(category_id, (0, 0, 255))
            
            # bbox 좌표를 정수형으로 변환
            x, y, w, h = map(int, bbox)
            cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
            
            # 클래스 이름(레이블)을 바운딩 박스 위에 표시
            label = category_name
            label_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            # 텍스트 위치: 바운딩 박스 위쪽 혹은 아래쪽 (y 좌표가 너무 작으면 아래쪽에 표시)
            label_x = x
            label_y = y - 10 if y - 10 > 10 else y + 10
            # 텍스트 배경을 채운 사각형을 먼저 그림 (가독성을 위해)
            cv2.rectangle(img, (label_x, label_y - label_size[1] - baseline),
                          (label_x + label_size[0], label_y + baseline), color, -1)
            cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX,
                        0.5, (255, 255, 255), 1)
        
        # =============================================================================
        # 출력 이미지 저장
        # =============================================================================
        # 이미지 파일명에 "/" 문자가 포함되어 있으면 "_"로 대체하여 저장합니다.
        # (여러 폴더에 분산된 파일들을 한 폴더에 저장할 때, 폴더 구조를 유지하지 않도록 함)
        output_image_path = os.path.join(output_folder, image_file_name.replace("/", "_"))
        # 출력 폴더가 없으면 생성
        os.makedirs(os.path.dirname(output_image_path), exist_ok=True)
        cv2.imwrite(output_image_path, img)
        
        pbar.update(1)

print("라벨링된 이미지를 저장했습니다.")