#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
여러 개의 데이터셋(예: 4개)을 하나의 데이터셋으로 합치는 스크립트입니다.
- 각 데이터셋의 이미지 파일들을 하나의 폴더에 평면화하여 복사하고,
- 각 데이터셋의 JSON 어노테이션 정보를 합쳐 하나의 JSON 파일로 생성합니다.

주요 기능:
    1. 사용자가 설정한 여러 데이터셋의 정보를 기반으로 JSON 파일을 로드하여,
       이미지 정보와 어노테이션 정보를 추출합니다.
       - JSON 파일이 COCO 형식("annotations" 키 존재)일 경우와 그렇지 않은 경우 모두 지원합니다.
    2. 각 데이터셋의 이미지에 대해 새 파일명을 부여합니다.
       - 오직 데이터셋별 접두어(prefix)만 앞에 붙여 충돌을 방지합니다.
    3. 이미지와 어노테이션의 ID를 전체 데이터셋 기준으로 새롭게 순차 부여하여 중복 문제를 해결합니다.
    4. 어노테이션의 클래스(카테고리) 정보는 각 데이터셋마다 다를 수 있으므로,
       사용자가 최종 데이터셋에 포함할 클래스들을 지정하고, 이름 변경이나 병합, 삭제가 가능하도록 지원합니다.
    5. 모든 이미지 파일들을 지정한 대상 폴더로 복사하며, 복사 진행 상황(퍼센트)과 전체 소요시간을 출력합니다.
    6. 최종 합쳐진 JSON 파일을 저장하며, 이미지의 file_name은 절대경로로 업데이트되고,
       **추가로 'img_path' 키를 포함**하여 mmyolo 학습 파이프라인에서 요구하는 필드를 제공합니다.

사용 방법:
    1. dataset_configs 리스트에 각 데이터셋 정보를 입력합니다.
    2. combined_dataset_dir 변수에 최종 합친 데이터셋을 저장할 경로를 지정합니다.
    3. final_category_mapping 변수에서 최종 데이터셋에 포함할 클래스(카테고리)를 지정합니다.
    4. 스크립트를 실행하면, 이미지 파일은 combined_dataset_dir/images 폴더에 복사되고,
       JSON 파일은 combined_dataset_dir/combined.json에 저장됩니다.

예시 실행:
    python merge_multiple_datasets.py

작성자: [작성자 이름 또는 팀]
작성일: [작성일]
"""

import os
import json
import shutil
import time

def load_and_separate_json(json_path):
    """
    JSON 파일을 로드한 후, 이미지 정보와 어노테이션 정보를 분리하는 함수.
    
    COCO 형식처럼 "annotations" 키가 존재하면 해당 키의 내용을 사용하며,
    그렇지 않을 경우 "images" 리스트 내에서 "file_name" 유무로 이미지와 어노테이션 정보를 구분합니다.
    
    매개변수:
        json_path (str): JSON 파일의 절대/상대 경로.
    
    반환값:
        images (list): 이미지 정보 리스트.
        annotations (list): 어노테이션 정보 리스트.
        categories (list): 카테고리 정보 리스트 (존재 시).
        info (dict): info 필드 (존재 시).
        licenses (list): licenses 필드 (존재 시).
    """
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    if "annotations" in data:
        images = data.get("images", [])
        annotations = data.get("annotations", [])
    else:
        images = []
        annotations = []
        for item in data["images"]:
            if "file_name" in item:
                images.append(item)
            else:
                annotations.append(item)
    
    categories = data.get("categories", [])
    info = data.get("info", {})
    licenses = data.get("licenses", [])
    
    return images, annotations, categories, info, licenses

def main():
    # ================= 사용자 설정 변수 (필요 시 수정) =================
    dataset_configs = [
        {
            "name": "dl1",
            "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project/images",
            "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project/mm_coco_annotations.json",
            "prefix": "dl1_"
        },
        {
            "name": "dl2",
            "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_2/images",
            "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_2/mm_coco_annotations.json",
            "prefix": "dl2_"
        },
        {
            "name": "dl3",
            "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_3/images",
            "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_3/mm_coco_annotations.json",
            "prefix": "dl3_"
        },
        {
            "name": "dl4",
            "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_4/images",
            "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_4/mm_coco_annotations.json",
            "prefix": "dl4_"
        }
    ]
    
    # 2. 최종 합친 데이터셋 저장 경로
    combined_dataset_dir = "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl_all_combined"
    combined_images_dir = os.path.join(combined_dataset_dir, "images")
    combined_json_path = os.path.join(combined_dataset_dir, "combined.json")
    
    # 3. 최종 데이터셋에 포함할 클래스(카테고리) 설정
    final_category_mapping = {
        "worker_dl": "worker",
        "signalman_dl": "signalman",
        "hardhat_dl": "helmet",
        "worker_helmat_off_dl": "worker",
        "harness_dl": "harness",
        "mixer_truck_dl": "mixer_truck",
        "excavator_dl": "excavator"
    }
    # =============================================================================
    
    # 폴더 생성 (이미지 저장용)
    os.makedirs(combined_images_dir, exist_ok=True)
    
    # 데이터셋별 정보를 빠르게 조회하기 위한 딕셔너리 (키: 데이터셋 이름)
    dataset_info = {cfg["name"]: cfg for cfg in dataset_configs}
    
    # ---------------- 글로벌 변수 초기화 ----------------
    combined_images = []      # 최종 JSON의 "images" 항목
    combined_annotations = [] # 최종 JSON의 "annotations" 항목
    new_image_id = 1          # 전체 데이터셋 기준 이미지 ID 순차 부여
    new_annotation_id = 1     # 전체 데이터셋 기준 어노테이션 ID 순차 부여
    
    global_final_categories = {}
    final_cat_counter = 1
    
    dataset_image_id_mapping = {}      # 각 데이터셋의 이미지 id 매핑
    dataset_category_id_mapping = {}     # 각 데이터셋의 카테고리 id 매핑
    
    # ---------------- 각 데이터셋 처리 ----------------
    for cfg in dataset_configs:
        dataset_name = cfg["name"]
        json_path = cfg["json_path"]
        prefix = cfg["prefix"]
        print(f">> '{dataset_name}' 데이터셋 JSON 파일 로드 중... ({json_path})")
        
        images, annotations, categories, info, licenses = load_and_separate_json(json_path)
        
        dataset_image_id_mapping[dataset_name] = {}
        dataset_category_id_mapping[dataset_name] = {}
        
        # 1. 카테고리 정보 처리
        for cat in categories:
            orig_cat_name = cat.get("name", "").lower()
            if orig_cat_name not in final_category_mapping:
                continue
            final_name = final_category_mapping[orig_cat_name]
            if final_name is None:
                continue
            if final_name not in global_final_categories:
                global_final_categories[final_name] = final_cat_counter
                final_cat_counter += 1
            dataset_category_id_mapping[dataset_name][cat["id"]] = global_final_categories[final_name]
        
        # 2. 이미지 정보 처리
        print(f">> '{dataset_name}' 데이터셋 이미지 정보 처리 중...")
        for img in images:
            orig_file_name = img["file_name"]
            # [수정] 기존 전체 경로 대신 파일명(basename)만 추출하여 prefix 추가
            base_file_name = os.path.basename(orig_file_name)
            new_file_name = prefix + base_file_name
            new_img = img.copy()
            new_img["id"] = new_image_id
            # 초기에 단순 파일명으로 설정 (후에 절대경로로 업데이트)
            new_img["file_name"] = new_file_name
            new_img["orig_file_name"] = orig_file_name
            new_img["dataset"] = dataset_name
            combined_images.append(new_img)
            dataset_image_id_mapping[dataset_name][img["id"]] = new_image_id
            new_image_id += 1
        
        # 3. 어노테이션 정보 처리
        print(f">> '{dataset_name}' 데이터셋 어노테이션 정보 처리 중...")
        for ann in annotations:
            orig_cat_id = ann.get("category_id")
            if orig_cat_id not in dataset_category_id_mapping[dataset_name]:
                continue
            new_ann = ann.copy()
            old_img_id = ann["image_id"]
            if old_img_id in dataset_image_id_mapping[dataset_name]:
                new_ann["image_id"] = dataset_image_id_mapping[dataset_name][old_img_id]
            else:
                print(f"경고: '{dataset_name}' 데이터셋 어노테이션에서 이미지 id {old_img_id}를 찾을 수 없습니다.")
                continue
            new_ann["category_id"] = dataset_category_id_mapping[dataset_name][orig_cat_id]
            new_ann["id"] = new_annotation_id
            new_annotation_id += 1
            combined_annotations.append(new_ann)
    # -----------------------------------------------------
    
    # ---------------- 최종 JSON 파일 구성 ----------------
    combined_categories = [{"id": cid, "name": name} 
                           for name, cid in sorted(global_final_categories.items(), key=lambda x: x[1])]
    
    combined_info = {
        "description": "여러 데이터셋을 합친 데이터셋 (최종 클래스 매핑 적용)",
        "date_created": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    }
    combined_licenses = []  # 라이선스 정보 필요 시 추가
    
    combined_json = {
        "info": combined_info,
        "licenses": combined_licenses,
        "images": combined_images,
        "annotations": combined_annotations,
        "categories": combined_categories
    }
    # -----------------------------------------------------
    
    # ---------------- 이미지 파일 복사 ----------------
    print(">> 이미지 파일 복사 시작...")
    start_time = time.time()
    total_images = len(combined_images)
    copied_count = 0
    
    for img in combined_images:
        dataset_name = img["dataset"]
        orig_file_name = img["orig_file_name"]
        new_file_name = img["file_name"]
        source_images_dir = dataset_info[dataset_name]["images_dir"]
        source_path = os.path.join(source_images_dir, orig_file_name)
        dest_path = os.path.join(combined_images_dir, new_file_name)
        
        if os.path.exists(source_path):
            shutil.copy2(source_path, dest_path)
        else:
            print(f"\n경고: 소스 이미지 파일이 존재하지 않습니다: {source_path}")
        copied_count += 1
        
        if copied_count % max(1, total_images // 100) == 0 or copied_count == total_images:
            progress = (copied_count / total_images) * 100
            print(f"이미지 복사 진행 상황: {copied_count}/{total_images} ({progress:.2f}%)", end='\r')
    
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"\n>> 이미지 복사 완료: 총 {copied_count}개 파일 복사, 소요 시간: {elapsed_time:.2f}초")
    # -----------------------------------------------------
    
    # ---------------- JSON 파일 저장 전 임시 키 제거 및 절대경로/키 업데이트 ----------------
    for img in combined_images:
        # [수정] 이미지의 file_name을 절대경로로 업데이트
        abs_path = os.path.join(combined_images_dir, img["file_name"])
        img["file_name"] = abs_path
        # [수정] mmyolo 파이프라인에서 요구하는 'img_path' 키를 추가 (절대경로 사용)
        img["img_path"] = abs_path
        if "orig_file_name" in img:
            del img["orig_file_name"]
        if "dataset" in img:
            del img["dataset"]
    
    with open(combined_json_path, 'w', encoding='utf-8') as f:
        json.dump(combined_json, f, ensure_ascii=False, indent=4)
    print(f">> 합쳐진 JSON 파일 저장 완료: {combined_json_path}")

if __name__ == '__main__':
    main()