import json
import os
import random
from collections import defaultdict

def split_coco_json(input_json, train_json, val_json, val_ratio=0.1, seed=42):
    random.seed(seed)
    
    with open(input_json, 'r') as f:
        coco = json.load(f)
    
    # 이미지와 어노테이션 매핑
    annotations = coco['annotations']
    images = coco['images']
    categories = coco['categories']
    
    # 이미지 ID와 이미지 정보를 매핑
    image_id_to_image = {image['id']: image for image in images}
    
    # 카테고리별로 이미지 ID를 수집
    category_id_to_image_ids = defaultdict(set)
    image_id_to_category_ids = defaultdict(set)
    
    for ann in annotations:
        category_id = ann['category_id']
        image_id = ann['image_id']
        category_id_to_image_ids[category_id].add(image_id)
        image_id_to_category_ids[image_id].add(category_id)
    
    # 모든 이미지 ID 가져오기
    all_image_ids = set(image_id_to_image.keys())
    
    # 카테고리별로 이미지를 그룹화하여 분할
    train_image_ids = set()
    val_image_ids = set()
    
    for category in categories:
        category_id = category['id']
        image_ids = list(category_id_to_image_ids[category_id])
        random.shuffle(image_ids)
        
        num_val = max(1, int(len(image_ids) * val_ratio))  # 각 카테고리마다 최소 하나의 이미지를 val로
        val_ids = set(image_ids[:num_val])
        train_ids = set(image_ids[num_val:])
        
        val_image_ids.update(val_ids)
        train_image_ids.update(train_ids)
    
    # 중복된 이미지가 있을 수 있으므로 조정
    overlapping_image_ids = train_image_ids & val_image_ids
    if overlapping_image_ids:
        for img_id in overlapping_image_ids:
            val_image_ids.remove(img_id)
    
    # 나머지 이미지를 train으로
    remaining_image_ids = all_image_ids - (train_image_ids | val_image_ids)
    train_image_ids.update(remaining_image_ids)
    
    # 이미지 정보 분할
    train_images = [image_id_to_image[img_id] for img_id in train_image_ids]
    val_images = [image_id_to_image[img_id] for img_id in val_image_ids]
    
    # 어노테이션 분할
    train_annotations = [ann for ann in annotations if ann['image_id'] in train_image_ids]
    val_annotations = [ann for ann in annotations if ann['image_id'] in val_image_ids]
    
    # 새로운 COCO 딕셔너리 생성
    train_coco = {
        "info": coco.get("info", {}),
        "licenses": coco.get("licenses", []),
        "images": train_images,
        "annotations": train_annotations,
        "categories": categories
    }
    
    val_coco = {
        "info": coco.get("info", {}),
        "licenses": coco.get("licenses", []),
        "images": val_images,
        "annotations": val_annotations,
        "categories": categories
    }
    
    # JSON 파일로 저장
    with open(train_json, 'w') as f:
        json.dump(train_coco, f, indent=4)
    
    with open(val_json, 'w') as f:
        json.dump(val_coco, f, indent=4)
    
    print(f"Total images: {len(all_image_ids)}")
    print(f"Train images: {len(train_images)}")
    print(f"Val images: {len(val_images)}")
    
    # 각 세트에서 카테고리별 어노테이션 수 출력
    def print_category_counts(annotations, dataset_name):
        category_counts = defaultdict(int)
        for ann in annotations:
            category_counts[ann['category_id']] += 1
        print(f"{dataset_name} annotations per category:")
        for category in categories:
            cid = category['id']
            count = category_counts.get(cid, 0)
            print(f"  Category {cid} ({category['name']}): {count}")
    
    print_category_counts(train_annotations, "Train")
    print_category_counts(val_annotations, "Val")

if __name__ == "__main__":
    input_json = '/DATA2/ltb/mmyolo/it1_worker_helmet_resize/merged_resized.json'
    train_json = '/DATA2/ltb/mmyolo/it1_worker_helmet_resize/train.json'
    val_json = '/DATA2/ltb/mmyolo/it1_worker_helmet_resize/val.json'
    split_coco_json(input_json, train_json, val_json, val_ratio=0.3, seed=54)