# %% [cell]
# [초기 설정] - 필요한 라이브러리 임포트 및 경로 설정
import os
import json

# (!!!) 원본 JSON 파일 경로를 절대경로로 지정합니다.
original_json_path = "/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/데이터셋/1819_resize_3x_8x_10x_helmet_cutmix_shiftscalerotate_bigcanvas/resized_3x_8x_10x.json"

# 결과 JSON 파일이 저장될 폴더 경로 (기존 경로를 기준으로 exp, exp2 등 자동 증가 가능하도록 설정)
# 여기서는 편의상 같은 폴더 내에 저장하도록 하되 파일명에 스케일 정보를 포함합니다.
output_dir = os.path.dirname(original_json_path)

# %% [cell]
# [데이터 불러오기] - 원본 JSON 파일 읽기
try:
    with open(original_json_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    print("원본 JSON 파일 로드 완료!")
except Exception as e:
    print("JSON 파일 로드 중 오류 발생:", e)
    raise

# %% [cell]
# [필터 함수 정의] - 접미사에 따라 images와 annotations 필터링 함수

def filter_json_by_suffix(data, suffix):
    """
    data: 원본 JSON 데이터 (dict)
    suffix: 필터링에 사용할 접미사 예) "_1p0.jpg", "_2p0.jpg", "_0p5.jpg"
    
    반환: 접미사에 해당하는 이미지와 해당하는 annotation만 포함한 새로운 JSON 데이터 (dict)
    """
    # 원본 JSON 구조에서 변경하지 않는 info, licenses, categories를 그대로 복사
    new_data = {
        "info": data.get("info", {}),
        "licenses": data.get("licenses", []),
        "categories": data.get("categories", []),
    }
    
    # images 필터링: file_name이 지정된 접미사로 끝나는 것들만 선택
    images = data.get("images", [])
    filtered_images = [img for img in images if img.get("file_name", "").endswith(suffix)]
    print(f"'{suffix}' 조건에 해당하는 이미지 개수: {len(filtered_images)}")
    
    new_data["images"] = filtered_images
    
    # 선택된 이미지들의 id 리스트 추출
    image_ids = set(img.get("id") for img in filtered_images)
    
    # annotations 필터링: annotation의 image_id가 위 이미지 id 리스트에 포함되어 있는 경우만 선택
    annotations = data.get("annotations", [])
    filtered_annotations = [ann for ann in annotations if ann.get("image_id") in image_ids]
    print(f"'{suffix}' 조건에 해당하는 annotation 개수: {len(filtered_annotations)}")
    
    new_data["annotations"] = filtered_annotations
    
    return new_data

# %% [cell]
# [스케일별 데이터 분리 및 파일 저장]
# 각 접미사에 따라 새로운 JSON 데이터를 생성하고 저장합니다.
scale_suffixes = {
    "3x": "_3p0.jpg",
    "8x": "_8p0.jpg",
    "10x": "_10p0.jpg",
    # "0p3x": "_0p3.jpg"
}

for scale, suffix in scale_suffixes.items():
    # 필터 함수 호출로 새로운 데이터 생성
    filtered_data = filter_json_by_suffix(data, suffix)
    
    # 저장할 파일명 구성: 원본 파일명 앞부분에 scale 정보를 추가하여 구분
    output_filename = f"1819_resize_3x_8x_10x_helmet_cutmix_shiftscalerotate_bigcanvas_{scale}.json"
    output_path = os.path.join(output_dir, output_filename)
    
    try:
        # JSON 파일 저장 (indent=2로 가독성 있게 저장)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(filtered_data, f, ensure_ascii=False, indent=2)
        print(f"{scale} 데이터셋 JSON 파일 저장 완료: {output_path}")
    except Exception as e:
        print(f"{scale} 데이터셋 JSON 파일 저장 중 오류 발생:", e)