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):
    """
    COCO 포맷의 통합 JSON 파일을 train/val 데이터셋으로 분할하는 함수입니다.
    
    매개변수:
      - input_json: 원본 COCO JSON 파일 경로 (통합 annotation 파일)
      - train_json: 분할된 train JSON 파일 저장 경로
      - val_json: 분할된 val JSON 파일 저장 경로
      - val_ratio: 각 카테고리별로 val에 배정할 이미지 비율 (기본 0.1)
      - seed: 무작위 분할을 위한 시드 값 (재현성을 위해)
    """
    random.seed(seed)
    
    # 원본 COCO JSON 파일 로드 (encoding='utf-8' 필수)
    with open(input_json, 'r', encoding='utf-8') as f:
        coco = json.load(f)
    
    annotations = coco.get('annotations', [])
    images = coco.get('images', [])
    categories = coco.get('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:
        image_id = ann.get('image_id')
        category_id = ann.get('category_id')
        if image_id is not None and category_id is not None:
            category_id_to_image_ids[category_id].add(image_id)
            image_id_to_category_ids[image_id].add(category_id)
    
    all_image_ids = set(image_id_to_image.keys())
    
    # 카테고리별 이미지 분할: 각 카테고리마다 최소 1개 이상 val에 배정
    train_image_ids = set()
    val_image_ids = set()
    
    for category in categories:
        category_id = category.get('id')
        image_ids = list(category_id_to_image_ids.get(category_id, []))
        random.shuffle(image_ids)
        num_val = max(1, int(len(image_ids) * val_ratio))
        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)
    
    # train, val에 중복 포함된 이미지가 있다면 val에서 제거
    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)
    
    # 분할된 이미지 및 annotation 정보 생성
    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.get('image_id') in train_image_ids]
    val_annotations = [ann for ann in annotations if ann.get('image_id') in val_image_ids]
    
    # 새 COCO JSON 구조 생성 (info, licenses 정보는 원본에서 그대로 가져옴)
    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 파일 저장 (ensure_ascii=False 설정으로 한글이 그대로 출력되도록 함)
    with open(train_json, 'w', encoding='utf-8') as f:
        json.dump(train_coco, f, ensure_ascii=False, indent=4)
    
    with open(val_json, 'w', encoding='utf-8') as f:
        json.dump(val_coco, f, ensure_ascii=False, indent=4)
    
    print(f"Total images: {len(all_image_ids)}")
    print(f"Train images: {len(train_images)}")
    print(f"Val images: {len(val_images)}")
    
    # 카테고리별 annotation 수를 출력하는 함수 (디버깅용)
    def print_category_counts(annotations, dataset_name):
        category_counts = defaultdict(int)
        for ann in annotations:
            cid = ann.get('category_id')
            if cid is not None:
                category_counts[cid] += 1
        print(f"{dataset_name} annotations per category:")
        for category in categories:
            cid = category.get('id')
            count = category_counts.get(cid, 0)
            print(f"  Category {cid} ({category.get('name', '')}): {count}")
    
    print_category_counts(train_annotations, "Train")
    print_category_counts(val_annotations, "Val")

if __name__ == "__main__":
    # 경로는 절대경로로 지정합니다.
    input_json = '/data/ai-hub/fire-dataset/106.화재_발생_예측_데이터(연기_동영상)/01.데이터/fire_dataset_val_clean.json'
    train_json = '/data/ai-hub/fire-dataset/106.화재_발생_예측_데이터(연기_동영상)/01.데이터/train.json'
    val_json = '/data/ai-hub/fire-dataset/106.화재_발생_예측_데이터(연기_동영상)/01.데이터/val.json'
    # 예제에서는 val_ratio=0.3, seed=54 로 설정
    split_coco_json(input_json, train_json, val_json, val_ratio=0.3, seed=54)