# tools/prepare_samples.py
"""
예제:
    cd waca_unet_space
    python tools/prepare_samples.py --root /data/ICCAD_2023/hidden-real-circuit-data \
        --configs_path configs/began_iccad_fake \
        --dataset iccad_hidden \
        --indices 0 1 2 3 4

- build_dataset_iccad_hidden / build_dataset_iccad_real 을 이용해서
  (input, target, casename)을 가져온 뒤, input만 samples/ 폴더에 저장.
- input은 IRDropDataset에서 반환하는 (C,H,W) 텐서(정규화 완료)를 그대로 npy로 저장.
"""
import os
import argparse
import numpy as np
import torch

from config import get_config
from ir_dataset import (
    build_dataset_iccad_hidden,
    build_dataset_iccad_real,
)


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument(
        "--root",
        type=str,
        help="ICCAD hidden/real root path. 예: /data/ICCAD_2023/hidden-real-circuit-data or /data/ICCAD_2023/real-circuit-data",
        default='/data/ICCAD_2023/hidden-real-circuit-data',
    )
    p.add_argument(
        "--dataset",
        type=str,
        choices=["iccad_hidden", "iccad_real"],
        default="iccad_real",
    )
    p.add_argument(
        "--img_size",
        type=int,
        default=384,
    )
    p.add_argument(
        "--in_ch",
        type=int,
        default=25,
    )
    p.add_argument(
        "--configs_path",
        type=str,
        default="/workspace/IR_Drop_prior_study/XICCAD/configs/cfirst/began_iccad_fake/stats_1um.json",
        help="stats_1um.json 이 들어있는 폴더",
    )
    p.add_argument(
        "--unit",
        type=str,
        default="1um",
    )
    p.add_argument(
        "--indices",
        type=int,
        nargs="+",
        default=[0, 1, 2, 3, 4],
        help="예제 샘플로 뽑을 dataset 인덱스들",
    )
    p.add_argument(
        "--out_dir",
        type=str,
        default="samples",
        help="npy를 저장할 디렉토리",
    )
    return p.parse_args()


def main():
    args = parse_args()
    os.makedirs(args.out_dir, exist_ok=True)

    # 학습 때 쓰던 통계 config 로드 (당신의 get_config 시그니처에 맞게 조정)
    norm_config = get_config(
        args.unit,
        configs_path=args.configs_path,
        dataset_name="began_iccad_fake",
    )

    common_kwargs = dict(
        img_size=args.img_size,
        in_ch=args.in_ch,
        train=False,
        use_raw=False,              # 학습 때와 동일하게 z-score 등으로 정규화된 입력 사용
        input_norm_type="z_score",
        target_norm_type="raw",
        target_layers=[],           # 전체 종합 IR-drop (이미 npy에 통합된 경우)
        use_pdn_density=True,
        use_pad_distance=True,
        use_comprehensive_feature=True,
        norm_config=norm_config,
        return_case=True,           # (x, y, casename) 반환
        interpolation="lanczos",
    )

    if args.dataset == "iccad_hidden":
        dataset = build_dataset_iccad_hidden(
            root_path=args.root,
            **common_kwargs,
        )
    else:
        dataset = build_dataset_iccad_real(
            root_path=args.root,
            **common_kwargs,
        )

    # build_dataset_* 가 (train,val) 튜플을 반환하는 경우 대응
    if isinstance(dataset, (tuple, list)) and len(dataset) == 2:
        dataset = dataset[1]  # val set을 예제로 사용하거나, 필요에 따라 수정

    print(f"Dataset length: {len(dataset)}")

    for idx in args.indices:
        if idx >= len(dataset):
            print(f"[WARN] index {idx} is out of range, skip.")
            continue

        sample = dataset[idx]
        if len(sample) == 3:
            x, y, casename = sample
        else:
            x, y = sample
            casename = f"idx{idx}"

        if isinstance(x, torch.Tensor):
            x_np = x.detach().cpu().numpy()
        else:
            x_np = np.asarray(x)

        out_name = f"{casename}_input.npy"
        out_path = os.path.join(args.out_dir, out_name)
        np.save(out_path, x_np)
        print(f"Saved: {out_path} (shape={x_np.shape})")

    print("Done.")


if __name__ == "__main__":
    main()