"""
Build 5-fold cross-validation splits from the pooled train+val patches.

The TEST set (dataset/patches_v2/test) stays completely untouched and held out,
so the final ensemble test number is directly comparable to previous runs.

For each fold k:
  val   = fold k  (~114 patches)      → honest held-out validation
  train = other 4 folds (~456)        → training
  norm_stats.json computed from THIS fold's train patches only (no leakage)

Output (symlinks — no data duplicated):
  dataset/cv_folds/fold_{0..4}/train/features|masks/*.npy
  dataset/cv_folds/fold_{0..4}/val/features|masks/*.npy
  dataset/cv_folds/fold_{0..4}/metadata/norm_stats.json
"""

import os
import glob
import json
import numpy as np

SRC        = "dataset/patches_v2"
OUT        = "dataset/cv_folds"
N_FOLDS    = 5
SEED       = 42


def pooled_patches():
    """All train+val feature files (absolute paths). Test is excluded."""
    files = []
    for split in ("train", "val"):
        files += sorted(glob.glob(os.path.join(SRC, split, "features", "*.npy")))
    return [os.path.abspath(f) for f in files]


def link(src, dst):
    if os.path.lexists(dst):
        os.remove(dst)
    os.symlink(src, dst)


def compute_norm_stats(feat_paths, n_bands, band_names):
    s1 = np.zeros(n_bands, np.float64)
    s2 = np.zeros(n_bands, np.float64)
    cnt = np.zeros(n_bands, np.float64)
    for fp in feat_paths:
        arr = np.load(fp).astype(np.float64).reshape(-1, n_bands)
        ok  = np.isfinite(arr) & (arr != -9999)
        x   = np.where(ok, arr, 0.0)
        s1  += x.sum(axis=0)
        s2  += (x * x).sum(axis=0)
        cnt += ok.sum(axis=0)
    cnt  = np.maximum(cnt, 1)
    mean = s1 / cnt
    std  = np.sqrt(np.maximum(s2 / cnt - mean ** 2, 1e-12))
    return {"band_names": band_names, "mean": mean.tolist(), "std": std.tolist(),
            "num_train_patches": len(feat_paths)}


def main():
    rng   = np.random.default_rng(SEED)
    feats = pooled_patches()
    rng.shuffle(feats)
    n = len(feats)
    print(f"Pooled train+val patches: {n}")

    # band names + count from the source metadata
    src_meta = json.load(open(os.path.join(SRC, "metadata", "norm_stats.json")))
    band_names = src_meta["band_names"]
    n_bands = len(band_names)

    # Split indices into N_FOLDS contiguous chunks
    fold_idx = np.array_split(np.arange(n), N_FOLDS)

    for k in range(N_FOLDS):
        val_ids   = set(fold_idx[k].tolist())
        val_feats   = [feats[i] for i in range(n) if i in val_ids]
        train_feats = [feats[i] for i in range(n) if i not in val_ids]

        fold_dir = os.path.join(OUT, f"fold_{k}")
        for split in ("train", "val"):
            os.makedirs(os.path.join(fold_dir, split, "features"), exist_ok=True)
            os.makedirs(os.path.join(fold_dir, split, "masks"),    exist_ok=True)
        os.makedirs(os.path.join(fold_dir, "metadata"), exist_ok=True)

        def link_set(feat_list, split):
            for f in feat_list:
                base = os.path.basename(f)
                mask = os.path.abspath(f.replace("/features/", "/masks/"))
                link(f,    os.path.join(fold_dir, split, "features", base))
                link(mask, os.path.join(fold_dir, split, "masks",    base))

        link_set(train_feats, "train")
        link_set(val_feats,   "val")

        # per-fold norm stats from TRAIN patches only
        norm = compute_norm_stats(train_feats, n_bands, band_names)
        with open(os.path.join(fold_dir, "metadata", "norm_stats.json"), "w") as f:
            json.dump(norm, f, indent=2)

        print(f"  fold_{k}: train={len(train_feats)}  val={len(val_feats)}  → {fold_dir}")

    print("\nDone. 5 folds ready under", OUT)


if __name__ == "__main__":
    main()