"""样本加载：ml/data/samples/*.{X,Y,V,SG,TID}.npy + *.tasks.jsonl（ml/prep/build_dataset.cjs 产出）
  python -m ml.data --index ml/data/samples      # 汇总样本/任务数，写 index.json
全部载入内存（200 万样本约 1.5 GB），按任务 split 字段划分 train / val。"""
import argparse, glob, json, os
import numpy as np
import torch


def load_tasks(dirs):
    tasks = []
    for d in dirs:
        for f in sorted(glob.glob(os.path.join(d, '*.tasks.jsonl'))):
            for line in open(f, encoding='utf-8'):
                if line.strip():
                    tasks.append(json.loads(line))
    return tasks


class Samples:
    def __init__(self, dirs, device='cpu', split=None):
        if isinstance(dirs, str): dirs = [dirs]
        xs, ys, vs, ss, ts = [], [], [], [], []
        for d in dirs:
            for fx in sorted(glob.glob(os.path.join(d, '*.X.npy'))):
                b = fx[:-6]
                xs.append(np.load(fx)); ys.append(np.load(b + '.Y.npy')); vs.append(np.load(b + '.V.npy'))
                ss.append(np.load(b + '.SG.npy')); ts.append(np.load(b + '.TID.npy'))
        if not xs: raise SystemExit(f'没有样本：{dirs}（先跑 ml/prep/run_prep.sh）')
        self.tasks = load_tasks(dirs)
        val_tid = {t['tid'] for t in self.tasks if t['split'] == 'val'}
        X, Y, V, SG, TID = map(np.concatenate, (xs, ys, vs, ss, ts))
        is_val = np.isin(TID, np.fromiter(val_tid, dtype=np.int64)) if val_tid else np.zeros(len(TID), bool)
        sel = {'train': ~is_val, 'val': is_val, None: np.ones(len(TID), bool)}[split]
        t = lambda a: torch.from_numpy(np.ascontiguousarray(a[sel])).to(device)
        self.X, self.Y, self.V, self.SG, self.TID = t(X), t(Y), t(V), t(SG), t(TID)
        self.n = self.X.shape[0]

    def batch(self, bs, gen=None):
        i = torch.randint(0, self.n, (bs,), device=self.X.device, generator=gen)
        return self.X[i], self.Y[i], self.V[i], self.SG[i]


def main():
    ap = argparse.ArgumentParser(); ap.add_argument('--index', default='ml/data/samples'); a = ap.parse_args()
    tr, va = Samples(a.index, split='train'), Samples(a.index, split='val')
    tasks = load_tasks([a.index])
    cats = {}
    for t in tasks: cats[t['cat'] + '/' + t['mode']] = cats.get(t['cat'] + '/' + t['mode'], 0) + 1
    info = {'train_samples': tr.n, 'val_samples': va.n, 'tasks': len(tasks), 'val_tasks': sum(t['split'] == 'val' for t in tasks), 'by_cat_mode': cats}
    json.dump(info, open(os.path.join(a.index, 'index.json'), 'w'), indent=1)
    print(json.dumps(info, indent=1))


if __name__ == '__main__':
    main()