import json
import os
from collections import defaultdict

# 1. JSON 파일 로드
with open('/mmyolo/it1_1639_worker_helmet/it1-worker-helmet-1639-coco.json', 'r') as f:
    data = json.load(f)

# 2. COCO 포맷으로 정리
clean_data = {
    "images": [],
    "annotations": [],
    "categories": []
}

# images 정리
for image in data.get('images', []):
    clean_image = {
        "id": image.get('id'),
        "file_name": image.get('file_name'),
        "width": image.get('width'),
        "height": image.get('height')
    }
    clean_data['images'].append(clean_image)

# annotations 정리
for annotation in data.get('annotations', []):
    clean_annotation = {
        "id": annotation.get('id'),
        "image_id": annotation.get('image_id'),
        "category_id": annotation.get('category_id'),
        "bbox": annotation.get('bbox'),
        "area": annotation.get('area'),
        "iscrowd": annotation.get('iscrowd', 0)
    }
    clean_data['annotations'].append(clean_annotation)

# categories 정리
# 원래의 category_id와 name을 유지하면서 새로운 id를 부여
original_categories = data.get('categories', [])
category_id_map = {}
new_categories = []
new_category_id = 1

for category in original_categories:
    old_id = category['id']
    category_id_map[old_id] = new_category_id
    new_category = {
        "id": new_category_id,
        "name": category['name'],
        "supercategory": category.get('supercategory', '')
    }
    new_categories.append(new_category)
    new_category_id += 1

clean_data['categories'] = new_categories

# 3. ID 순차 재할당
# images ID 재할당
image_id_map = {}
for idx, image in enumerate(clean_data['images'], start=1):
    old_id = image['id']
    image['id'] = idx
    image_id_map[old_id] = idx

# annotations ID 재할당 및 image_id, category_id 매핑
for idx, annotation in enumerate(clean_data['annotations'], start=1):
    annotation['id'] = idx
    # image_id 매핑
    annotation['image_id'] = image_id_map.get(annotation['image_id'], annotation['image_id'])
    # category_id 매핑
    annotation['category_id'] = category_id_map.get(annotation['category_id'], annotation['category_id'])

# 4. 파일 이름 중복 처리
seen_files = set()  # (folder_name, file_name) 저장
folder_duplication_started = {}  # 폴더별로 duplication 시작 여부
duplicates = []  # 중복된 폴더 목록

for image in clean_data['images']:
    file_name = image['file_name']  # 예: '20220701_95512_ch07/0000029.png'
    folder_name, base_name = os.path.split(file_name)  # 폴더명과 파일명 분리

    file_key = (folder_name, base_name)

    if file_key in seen_files:
        # 중복된 파일 발견
        if not folder_duplication_started.get(folder_name, False):
            folder_duplication_started[folder_name] = True
            duplicates.append(folder_name)
    else:
        seen_files.add(file_key)

    # 해당 폴더에서 중복이 시작되었으면 폴더명에 '_2'를 추가
    if folder_duplication_started.get(folder_name, False):
        new_folder_name = f"{folder_name}_2"
        new_file_name = os.path.join(new_folder_name, base_name)
        image['file_name'] = new_file_name
    else:
        # 변경 없이 그대로 유지
        pass

# 중복된 폴더 이름 목록 출력
if duplicates:
    print("중복된 폴더 이름 목록:")
    for folder_name in duplicates:
        print(f"- {folder_name}")

# 5. 수정된 JSON 저장
with open('/mmyolo/it1_1639_worker_helmet/it1-worker-helmet-1639-coco_modify.json', 'w') as f:
    json.dump(clean_data, f, indent=2)