import json

# =============================
# 설정 섹션
# =============================
COMMON_DIR = "/ltb"
# JSON 파일 경로 설정
input_json_path = COMMON_DIR + "/media/ltb/90887173887158A4/Users/ltb/si/dl_all_combined/combined.json"  # 입력 JSON 파일 경로
output_json_path = COMMON_DIR + "/media/ltb/90887173887158A4/Users/ltb/si/dl_all_combined/class_modified_combined.json"  # 출력 JSON 파일 경로

# 클래스 재매핑 설정
# 
# 구조:
# new_categories = {
#     "새로운_클래스_이름": {
#         "id": 새로운_클래스_ID,
#         "source_ids": [병합할_기존_카테고리_ID들]
#     },
#     ...
# }
#
# 예시:
# {
#     "Worker": {
#         "id": 1,
#         "source_ids": [1, 3]  # worker_dl (1)과 signalman_dl (3)을 병합
#     },
#     "helmet": {
#         "id": 2,
#         "source_ids": [2]  # hardhat_dl (2)
#     }
# }
new_categories = {
    "worker": {
        "id": 1,
        "source_ids": [2]  # worker_dl (3)과 signalman_dl (6)을 Worker로 병합
    },
    "signalman": {
        "id": 2,
        "source_ids": [6]  # hardhat_dl (4)을 helmet으로 병합
    },
    "helmet": {
        "id": 3,
        "source_ids": [3]  # hardhat_dl (4)을 helmet으로 병합
    },
    "harness": {
        "id": 4,
        "source_ids": [5]  # hardhat_dl (4)을 helmet으로 병합
    },
    "mixer_truck": {
        "id": 5,
        "source_ids": [1]  # hardhat_dl (4)을 helmet으로 병합
    },
    "excavator": {
        "id": 6,
        "source_ids": [4]  # hardhat_dl (4)을 helmet으로 병합
    }
}

# =============================
# 처리 섹션
# =============================

def create_mapping(new_categories):
    """
    새로운 카테고리 설정을 기반으로 기존 카테고리 ID를 새로운 카테고리 ID로 매핑합니다.

    Args:
        new_categories (dict): 새로운 카테고리 설정.

    Returns:
        dict: 기존 카테고리 ID를 새로운 카테고리 ID로 매핑한 사전.
    """
    mapping = {}
    for new_class, details in new_categories.items():
        new_id = details["id"]
        for old_id in details["source_ids"]:
            mapping[old_id] = new_id
    return mapping

def get_new_categories(new_categories):
    """
    새로운 카테고리를 COCO 형식에 맞게 변환합니다.

    Args:
        new_categories (dict): 새로운 카테고리 설정.

    Returns:
        list: COCO 형식의 카테고리 사전 리스트.
    """
    categories = []
    for new_class, details in new_categories.items():
        categories.append({
            "id": details["id"],
            "name": new_class,
            "supercategory": "root"  # 필요에 따라 supercategory를 수정할 수 있습니다.
        })
    return categories

def remap_and_filter_annotations(json_data, mapping):
    """
    주석(annotation) 내의 category_id를 매핑에 따라 변경하고, 매핑되지 않은 어노테이션을 제거합니다.

    Args:
        json_data (dict): 원본 COCO JSON 데이터.
        mapping (dict): 기존 카테고리 ID를 새로운 카테고리 ID로 매핑한 사전.

    Returns:
        None: json_data가 직접 수정됩니다.
    """
    original_count = len(json_data['annotations'])
    # 새로운 어노테이션 리스트 생성
    filtered_annotations = []
    for annotation in json_data['annotations']:
        old_id = annotation['category_id']
        if old_id in mapping:
            annotation['category_id'] = mapping[old_id]
            filtered_annotations.append(annotation)
        else:
            # 매핑되지 않은 카테고리 ID의 어노테이션 제거
            print(f"경고: category_id {old_id}은(는) 매핑에 없어 제거됩니다.")
    json_data['annotations'] = filtered_annotations
    removed_count = original_count - len(filtered_annotations)
    print(f"어노테이션 필터링 완료: {removed_count}개 어노테이션이 제거되었습니다.")

def update_categories(json_data, new_categories):
    """
    JSON 데이터 내의 categories 리스트를 새로운 카테고리로 교체합니다.

    Args:
        json_data (dict): 원본 COCO JSON 데이터.
        new_categories (dict): 새로운 카테고리 설정.

    Returns:
        None: json_data가 직접 수정됩니다.
    """
    json_data['categories'] = get_new_categories(new_categories)
    print(f"카테고리 업데이트 완료: {len(json_data['categories'])}개의 카테고리가 설정되었습니다.")

def main():
    # JSON 파일 읽기
    with open(input_json_path, 'r', encoding='utf-8') as file:
        json_data = json.load(file)

    # 매핑 생성
    mapping = create_mapping(new_categories)

    # 매핑 확인 (선택사항)
    print("카테고리 ID 매핑:")
    for old_id, new_id in mapping.items():
        print(f"  기존 ID {old_id} -> 새로운 ID {new_id}")

    # 주석의 category_id 재매핑 및 필터링
    remap_and_filter_annotations(json_data, mapping)

    # 카테고리 리스트 업데이트
    update_categories(json_data, new_categories)

    # 결과 JSON 저장
    with open(output_json_path, 'w', encoding='utf-8') as file:
        json.dump(json_data, file, indent=4, ensure_ascii=False)
        print(f"통합된 JSON이 저장되었습니다: {output_json_path}")

if __name__ == "__main__":
    main()