"""Tạo split CÂN BẰNG 1:1 từ các CSV split đầy đủ (data/splits/*.csv).

Cân bằng bằng cách chỉ giữ:
  - real = wiki (toàn bộ ảnh thật)
  - fake = 1 loại duy nhất (mặc định: insight — swap face)

Identity vẫn tách biệt giữa train/val/test (kế thừa từ split gốc theo nhân vật).

Output: data/splits/{train,val,test}_{fake_type}.csv
"""
import argparse
import csv
import os
import re

EXTENSIONS = (".jpg", ".jpeg", ".png")
REAL_CAT = "wiki"


def keep_row(path: str, fake_type: str) -> bool:
    """Giữ ảnh nếu thuộc real (wiki) hoặc đúng loại fake được chọn."""
    return f"/hug/{REAL_CAT}/" in path or f"/hug/{fake_type}/" in path


def main():
    parser = argparse.ArgumentParser(description="Tạo split cân bằng 1:1")
    parser.add_argument("--splits-dir", default="data/splits")
    parser.add_argument("--fake-type", default="insight", choices=["insight", "inpainting", "text2img"])
    parser.add_argument("--out-dir", default="data/splits")
    args = parser.parse_args()

    for name in ("train", "val", "test"):
        src = os.path.join(args.splits_dir, f"{name}.csv")
        out = os.path.join(args.out_dir, f"{name}_{args.fake_type}.csv")

        kept, dropped = [], 0
        with open(src, newline="") as f:
            reader = csv.reader(f)
            header = next(reader, None)
            for row in reader:
                if len(row) < 2:
                    continue
                path, label = row[0], int(row[1])
                if not path.lower().endswith(EXTENSIONS):
                    continue
                if keep_row(path, args.fake_type):
                    kept.append((path, label))
                else:
                    dropped += 1

        with open(out, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(["path", "label"])
            w.writerows(kept)

        n_real = sum(1 for _, lbl in kept if lbl == 0)
        n_fake = sum(1 for _, lbl in kept if lbl == 1)
        print(f"  {name:>5}: {len(kept):>6} ảnh (giữ) | real={n_real} fake={n_fake} | bỏ {dropped} loại fake khác")

    print(f"\nĐã tạo split cân bằng (real=wiki, fake={args.fake_type}) vào {args.out_dir}/")


if __name__ == "__main__":
    main()