import os
import cv2
from skimage.metrics import structural_similarity as ssim

# ===== 설정 =====
INPUT_DIR = "test1"            # 원본 프레임 폴더
OUTPUT_DIR = "test2"           # 중복/유사 제거 후 저장 폴더
IMAGE_FORMAT = "png"           # 저장 확장자 ('png', 'jpg' 등)
SSIM_THRESHOLD = 0.91          # 유사도 임계값 (값↑ = 더 엄격)
RESIZE_WIDTH = 320             # 비교용 리사이즈 너비 (작을수록 빠름)
TARGET_COUNT = 80              # 최종 선택할 이미지 개수
# =================

os.makedirs(OUTPUT_DIR, exist_ok=True)

# 사용할 이미지 확장자
VALID_EXT = (".png", ".jpg", ".jpeg", ".bmp", ".webp")

# test1 폴더의 파일 목록 정렬
file_list = sorted(
    f for f in os.listdir(INPUT_DIR)
    if f.lower().endswith(VALID_EXT)
)

if not file_list:
    print("입력 폴더에 유효한 이미지가 없습니다.")
    exit(1)

print(f"[INFO] 입력 이미지 수: {len(file_list)}")

prev_gray_small = None
candidate_paths = []   # 1차 SSIM 필터를 통과한 이미지 경로들

# ---------- 1단계: SSIM 기준으로 중복/유사 프레임 제거 ----------
for idx, filename in enumerate(file_list):
    input_path = os.path.join(INPUT_DIR, filename)
    img = cv2.imread(input_path)

    if img is None:
        print(f"[WARN] 이미지를 읽을 수 없음: {input_path}")
        continue

    # ---- 비교용으로 작은 그레이스케일 이미지 만들기 ----
    h, w = img.shape[:2]
    new_w = RESIZE_WIDTH
    new_h = int(h * (new_w / w))
    img_small = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
    gray_small = cv2.cvtColor(img_small, cv2.COLOR_BGR2GRAY)
    # ----------------------------------------------------

    if prev_gray_small is None:
        # 첫 이미지는 무조건 후보로 추가
        candidate_paths.append(input_path)
        prev_gray_small = gray_small
        print(f"[CAND] 첫 이미지 → {filename}")
    else:
        # 이전에 저장한 이미지와 SSIM 비교
        score = ssim(prev_gray_small, gray_small)

        if score < SSIM_THRESHOLD:
            # 충분히 다르면 새로운 이미지로 보고 후보에 추가
            candidate_paths.append(input_path)
            prev_gray_small = gray_small
            print(f"[CAND] {filename} (SSIM={score:.4f}) → 후보 추가")
        else:
            # 너무 비슷하면 스킵
            print(f"[SKIP] {filename} (SSIM={score:.4f}) - 이전과 거의 동일")

print(f"\n[INFO] SSIM 필터 통과 후보 수: {len(candidate_paths)}")

if not candidate_paths:
    print("SSIM 필터를 통과한 이미지가 없습니다.")
    exit(1)

# ---------- 2단계: 후보들 중에서 최종 80장 선택 ----------
if len(candidate_paths) <= TARGET_COUNT:
    # 후보가 80장 이하면 그냥 전부 사용
    selected_paths = candidate_paths
    print(f"[INFO] 후보가 {len(candidate_paths)}장이라 80장 미만, 모두 사용합니다.")
else:
    # 후보가 80장보다 많으면 전체 구간에서 균등하게 80장 샘플링
    step = len(candidate_paths) / TARGET_COUNT
    selected_paths = [
        candidate_paths[int(i * step)]
        for i in range(TARGET_COUNT)
    ]
    print(f"[INFO] 후보 {len(candidate_paths)}장 중에서 균등 간격으로 {TARGET_COUNT}장 선택.")

# ---------- 3단계: 선택된 이미지들을 test2에 저장 ----------
saved_idx = 0
for path in selected_paths:
    img = cv2.imread(path)
    if img is None:
        continue

    save_name = f"frame_{saved_idx:06d}.{IMAGE_FORMAT}"
    save_path = os.path.join(OUTPUT_DIR, save_name)
    cv2.imwrite(save_path, img)
    print(f"[SAVE] {os.path.basename(path)} → {save_name}")
    saved_idx += 1

print(f"\n[DONE] 최종 저장된 이미지 수: {saved_idx}")
print("결과 폴더:", os.path.abspath(OUTPUT_DIR))