#!/usr/bin/env python3
"""
Split aes_all.jsonl into train/val sets (90/10).
Shuffles with fixed seed for reproducibility.

Output:
  /workspace/elinnos/aes_training/data/aes_train.jsonl
  /workspace/elinnos/aes_training/data/aes_val.jsonl
"""
import json
import random
from pathlib import Path

WORKSPACE = Path("/workspace/elinnos")
INPUT = WORKSPACE / "aes_all.jsonl"
OUTPUT_DIR = WORKSPACE / "aes_training" / "data"
VAL_RATIO = 0.10
SEED = 42


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print(f"Loading {INPUT}...")
    samples = []
    with INPUT.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                samples.append(json.loads(line))

    n = len(samples)
    print(f"Total samples: {n}")

    random.seed(SEED)
    random.shuffle(samples)

    n_val = max(1, int(n * VAL_RATIO))
    n_train = n - n_val

    train_samples = samples[:n_train]
    val_samples = samples[n_train:]

    train_path = OUTPUT_DIR / "aes_train.jsonl"
    val_path = OUTPUT_DIR / "aes_val.jsonl"

    with train_path.open("w", encoding="utf-8") as f:
        for s in train_samples:
            f.write(json.dumps(s, ensure_ascii=False) + "\n")

    with val_path.open("w", encoding="utf-8") as f:
        for s in val_samples:
            f.write(json.dumps(s, ensure_ascii=False) + "\n")

    print(f"Train: {len(train_samples)} → {train_path}")
    print(f"Val:   {len(val_samples)} → {val_path}")
    print("Done.")


if __name__ == "__main__":
    main()