"""
3 stage .bin dosyasının her birinden val kesimi yapar.

Her stage'in son %X'ini val olarak ayırır → ana train .bin'i kısaltır.
Tüm val parçalarını tek bir v5_val.bin'de birleştirir (training script bunu kullanır).
Ayrıca per-stage val dosyalarını da bırakır (debug/inceleme için).

Cikti:
    data/v5_stage1.bin    (kısaltıldı)
    data/v5_stage2.bin    (kısaltıldı)
    data/v5_stage3.bin    (kısaltıldı)
    data/v5_val_stage1.bin
    data/v5_val_stage2.bin
    data/v5_val_stage3.bin
    data/v5_val.bin       (3'ünün birleşimi — training kullanır)

Kullanim:
    python 04b_make_val.py                  # %1 val (varsayilan)
    python 04b_make_val.py --val-ratio 0.005  # %0.5
    python 04b_make_val.py --dry-run        # sadece raporla, dosya degistirme
"""

import argparse
from pathlib import Path

import numpy as np

DATA_DIR = Path(__file__).parent / "data"
STAGES = [1, 2, 3]


def fmt(n):
    if n >= 1e9:
        return f"{n/1e9:.2f}B"
    if n >= 1e6:
        return f"{n/1e6:.1f}M"
    if n >= 1e3:
        return f"{n/1e3:.1f}K"
    return str(n)


def split_stage(stage: int, val_ratio: float, dry_run: bool):
    src = DATA_DIR / f"v5_stage{stage}.bin"
    val_out = DATA_DIR / f"v5_val_stage{stage}.bin"

    if not src.exists():
        print(f"  ! {src.name} yok, atlandı")
        return None

    data = np.memmap(src, dtype=np.uint16, mode="r")
    n_total = len(data)
    n_val = int(n_total * val_ratio)
    n_train = n_total - n_val

    print(f"\nStage {stage}: {src.name}")
    print(f"  Toplam: {fmt(n_total)} token")
    print(f"  Val:    {fmt(n_val)} token ({val_ratio*100:.1f}%)")
    print(f"  Train:  {fmt(n_train)} token (kalan)")

    if dry_run:
        print(f"  [dry-run] dosya degistirilmedi")
        # Memmap'i kopyala (ana .bin'i etkilemeden)
        val_chunk = np.array(data[-n_val:], dtype=np.uint16)
        return val_chunk

    # Önce val parçasını ayrı dosyaya yaz
    val_chunk = np.array(data[-n_val:], dtype=np.uint16)
    train_chunk = np.array(data[:n_train], dtype=np.uint16)

    # Memmap'i kapat
    del data

    # Val ve train'i yaz (train üzerine yazıyoruz, atomik değil ama tek seferlik)
    val_chunk.tofile(val_out)
    print(f"  → {val_out.name} yazıldı ({fmt(len(val_chunk))} token)")

    # Train'i geri yaz (yeniden boyutlandırılmış)
    train_chunk.tofile(src)
    print(f"  → {src.name} kısaltıldı ({fmt(len(train_chunk))} token)")

    return val_chunk


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--val-ratio", type=float, default=0.01,
                        help="Her stage'den val'a ayırılacak oran (varsayılan 0.01 = %%1)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Sadece raporla, dosyaları değiştirme")
    parser.add_argument("--no-merge", action="store_true",
                        help="v5_val.bin birleşik dosyasını yapma")
    args = parser.parse_args()

    print(f"{'='*60}")
    print(f"V5 VAL KESIMI  (ratio={args.val_ratio*100:.1f}%)")
    print(f"{'='*60}")

    val_chunks = []
    for stage in STAGES:
        chunk = split_stage(stage, args.val_ratio, args.dry_run)
        if chunk is not None:
            val_chunks.append(chunk)

    if not val_chunks:
        print("\n! Hiçbir stage işlenmedi")
        return

    if args.dry_run:
        total_val = sum(len(c) for c in val_chunks)
        print(f"\n[dry-run] toplam val: {fmt(total_val)} token (yazılmadı)")
        return

    # Birleşik val dosyası
    if not args.no_merge:
        merged_path = DATA_DIR / "v5_val.bin"
        merged = np.concatenate(val_chunks).astype(np.uint16)
        merged.tofile(merged_path)
        print(f"\n{'='*60}")
        print(f"Birleşik val: {merged_path.name}")
        print(f"  Toplam: {fmt(len(merged))} token "
              f"({merged.nbytes/1e6:.1f} MB)")
        print(f"{'='*60}")

    print(f"\n[OK] Val kesimi tamam.")
    print(f"  Training script bu dosyaları bekliyor:")
    print(f"    data/v5_stage1.bin   (kısaltıldı)")
    print(f"    data/v5_stage2.bin   (kısaltıldı)")
    print(f"    data/v5_stage3.bin   (kısaltıldı)")
    print(f"    data/v5_val.bin      (3 stage'in karışımı)")


if __name__ == "__main__":
    main()