import json
import os
import random
import shutil
from collections import defaultdict

# 경로 설정
roboflow_coco_path = "/mmyolo/roboflow_fall/train/_annotations.coco.json"
roboflow_img_dir = "/mmyolo/roboflow_fall/train"
output_img_dir = "/mmyolo/it1-worker-helmet-1721-detail/images/roboflow_fall"
merged_coco_path = "/mmyolo/it1-worker-helmet-1721-detail/lying_dataset_worker_background.json"

# 제외할 클래스 id (원본 데이터셋 상의 category_id)
exclude_ids = [1, 4, 5, 7, 8, 10]  
# 여기서는 원래 이미지의 category_id 값 기준으로 제외합니다.
# 주의: 이 exclude_ids는 로보플로우 데이터셋의 original category_id에 해당하는 값이어야 합니다.
# lying 데이터셋의 category_id들이 실제 어떤 값인지 확인이 필요하지만,
# 여기서는 질문자가 이전에 언급한 exclude_ids를 그대로 사용한다고 가정합니다.

# 최종 선택할 이미지 수 (예: 300장 -> 여기서는 3739장이라 명시)
target_count = 3739

# 디렉토리 생성
os.makedirs(output_img_dir, exist_ok=True)

# 로보플로우 coco 로드
with open(roboflow_coco_path, 'r') as f:
    rf_data = json.load(f)

rf_images = rf_data['images']
rf_annotations = rf_data['annotations']

image_id_to_ann = defaultdict(list)
for ann in rf_annotations:
    image_id_to_ann[ann['image_id']].append(ann)

def print_progress(prefix, current, total):
    perc = (current/total)*100
    print(f"{prefix}: {current}/{total} ({perc:.2f}%)")

all_img_ids = set(img['id'] for img in rf_images)
all_img_ids = list(all_img_ids)
random.shuffle(all_img_ids)
selected_img_ids = all_img_ids[:target_count]

total_imgs = len(selected_img_ids)
selected_images_data = []
selected_annotations_data = []

for i, img_id in enumerate(selected_img_ids, start=1):
    # 이미지 정보 획득
    img_info = [im for im in rf_images if im['id']==img_id][0]
    src_path = os.path.join(roboflow_img_dir, img_info['file_name'])
    dst_path = os.path.join(output_img_dir, img_info['file_name'])
    
    # 이미지 복사
    shutil.copy2(src_path, dst_path)
    selected_images_data.append(img_info)

    # 해당 이미지의 ann 추가 (exclude_ids에 해당하는 category_id는 제외)
    for ann in image_id_to_ann[img_id]:
        if ann['category_id'] in exclude_ids:
            # 제외대상 카테고리는 skip
            continue
        # 제외하지 않는다면 Worker(1)로 통일
        ann['category_id'] = 1
        selected_annotations_data.append(ann)

    print_progress("Copying images", i, total_imgs)

print("Image copy done. Total selected images:", len(selected_img_ids))
print("Selected annotations:", len(selected_annotations_data))

# file_name에 roboflow_fall/ prefix 추가
for im in selected_images_data:
    im['file_name'] = f"roboflow_fall/{im['file_name']}"

# ID 재할당: 새로 생성하는 coco이므로 이미지 id 1부터 시작
old_id_to_new_id = {}
for idx, im in enumerate(selected_images_data, start=1):
    old_id = im['id']
    im['id'] = idx
    old_id_to_new_id[old_id] = idx

# 어노테이션 id 재할당
for idx, ann in enumerate(selected_annotations_data, start=1):
    ann['id'] = idx
    # image_id 갱신
    ann['image_id'] = old_id_to_new_id[ann['image_id']]

# categories 정의: Worker(1), background(2)
categories = [
    {"id":1, "name":"Worker", "supercategory":""},
    {"id":2, "name":"background", "supercategory":""}
]

# 새로운 COCO 딕셔너리 생성
new_coco = {
    "images": selected_images_data,
    "annotations": selected_annotations_data,
    "categories": categories
}

with open(merged_coco_path, 'w') as f:
    json.dump(new_coco, f, indent=4)

print("All done.")
print("Copied images:", len(selected_images_data))
print("Appended annotations:", len(selected_annotations_data))
print("Saved new JSON to:", merged_coco_path)

# 누락 여부 체크
missing = []
for im in selected_images_data:
    actual_file_name = im['file_name'].replace("roboflow_fall/","")
    if not os.path.exists(os.path.join(output_img_dir, actual_file_name)):
        missing.append(actual_file_name)

if missing:
    print("Missing images:", missing)
else:
    print("No missing images. All good!")