"""
Step 2 — Split crops into train/val/test BY CROP (so overlapping patches from
one crop never straddle two splits), copying them into
<out>/dataset/{split}/{features,masks}/.
"""

import os
import glob
import shutil
import random


def split_dataset(cfg):
    out = cfg["out_path"]
    feat_dir = os.path.join(out, "features")
    mask_dir = os.path.join(out, "masks")
    dataset_dir = os.path.join(out, "dataset")
    tr_r, va_r, seed = cfg["train_ratio"], cfg["val_ratio"], cfg["seed"]

    ids = []
    for fp in sorted(glob.glob(os.path.join(feat_dir, "feature_*.tif"))):
        idx = os.path.basename(fp).split("_")[1].split(".")[0]
        if os.path.exists(os.path.join(mask_dir, f"mask_{idx}.tif")):
            ids.append(idx)
    print(f"[split] matched pairs: {len(ids)}")

    random.seed(seed)
    random.shuffle(ids)
    n = len(ids)
    n_tr = int(n * tr_r)
    n_va = int(n * (tr_r + va_r)) - n_tr
    splits = {"train": ids[:n_tr],
              "val":   ids[n_tr:n_tr + n_va],
              "test":  ids[n_tr + n_va:]}

    for name, sub in splits.items():
        fo = os.path.join(dataset_dir, name, "features")
        mo = os.path.join(dataset_dir, name, "masks")
        os.makedirs(fo, exist_ok=True)
        os.makedirs(mo, exist_ok=True)
        for idx in sub:
            shutil.copy2(os.path.join(feat_dir, f"feature_{idx}.tif"), fo)
            shutil.copy2(os.path.join(mask_dir, f"mask_{idx}.tif"), mo)
        print(f"[split]   {name}: {len(sub)}")

    tr, va, te = (set(splits[k]) for k in ("train", "val", "test"))
    assert not (tr & va) and not (tr & te) and not (va & te), "SPLIT LEAK!"
    print("[split] leak check passed — no crop in two splits.")