"""
Consolidate the actual training data (litdata binary chunks) into single folders.

Pretraining:  litdata_3b + litdata_english -> consolidated_pretrain_litdata/
Finetuning:   finetune/ + finetune_english/ -> consolidated_finetune/

Token counts are computed directly from the litdata index.json metadata.
"""

import json
import shutil
import os
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent.parent  # LUNA root
DATA = ROOT / "Base" / "data"
DATASETS = ROOT / "Base" / "Datasets"


def read_litdata_index(litdata_dir):
    """Read index.json and return chunks list + config."""
    with open(litdata_dir / "index.json", "r") as f:
        index = json.load(f)
    return index["chunks"], index["config"]


def count_tokens_from_index(chunks):
    """Sum up all tokens from chunk dim fields."""
    return sum(c["dim"] for c in chunks)


# ══════════════════════════════════════════════════════════════════
# 1. PRETRAINING: litdata_3b + litdata_english
# ══════════════════════════════════════════════════════════════════
print("=" * 65)
print("  PRETRAINING DATA CONSOLIDATION")
print("=" * 65)

# Read both indexes
chunks_3b, config = read_litdata_index(DATA / "litdata_3b")
chunks_en, _ = read_litdata_index(DATA / "litdata_english")

tokens_3b = count_tokens_from_index(chunks_3b)
tokens_en = count_tokens_from_index(chunks_en)
tokens_pretrain = tokens_3b + tokens_en

print(f"\n  litdata_3b:      {len(chunks_3b):>4} chunks | {tokens_3b:>15,} tokens")
print(f"  litdata_english: {len(chunks_en):>4} chunks | {tokens_en:>15,} tokens")
print(f"  {'─' * 55}")
print(f"  TOTAL PRETRAIN:  {len(chunks_3b)+len(chunks_en):>4} chunks | {tokens_pretrain:>15,} tokens")
print(f"  ({tokens_pretrain / 1e9:.3f} B tokens)")

# Create consolidated litdata folder
out_pretrain = DATA / "consolidated_pretrain_litdata"
out_pretrain.mkdir(parents=True, exist_ok=True)

# Copy chunks from litdata_3b (keep original names as chunk-0-{0..N})
print(f"\n  Copying litdata_3b chunks ({len(chunks_3b)} files)...")
new_chunks = []
for i, chunk in enumerate(chunks_3b):
    src = DATA / "litdata_3b" / chunk["filename"]
    new_name = f"chunk-0-{i}.bin"
    dst = out_pretrain / new_name
    if not dst.exists():
        shutil.copy2(str(src), str(dst))
    new_chunks.append({**chunk, "filename": new_name})
    if (i + 1) % 50 == 0 or i == len(chunks_3b) - 1:
        print(f"    {i+1}/{len(chunks_3b)} copied")

# Copy chunks from litdata_english (renumber continuing from 3b)
offset = len(chunks_3b)
print(f"\n  Copying litdata_english chunks ({len(chunks_en)} files)...")
for i, chunk in enumerate(chunks_en):
    src = DATA / "litdata_english" / chunk["filename"]
    new_name = f"chunk-0-{offset + i}.bin"
    dst = out_pretrain / new_name
    if not dst.exists():
        shutil.copy2(str(src), str(dst))
    new_chunks.append({**chunk, "filename": new_name})
    print(f"    {i+1}/{len(chunks_en)} copied")

# Write combined index.json
combined_index = {"chunks": new_chunks, "config": config}
with open(out_pretrain / "index.json", "w") as f:
    json.dump(combined_index, f, indent=2)
print(f"\n  Saved: {out_pretrain}")


# ══════════════════════════════════════════════════════════════════
# 2. FINETUNING: finetune/ + finetune_english/
# ══════════════════════════════════════════════════════════════════
print(f"\n{'=' * 65}")
print("  FINETUNING DATA CONSOLIDATION")
print("=" * 65)

# Load all finetune JSONs
ft_v1_train = json.loads((DATASETS / "finetune" / "train.json").read_text(encoding="utf-8"))
ft_v1_val = json.loads((DATASETS / "finetune" / "val.json").read_text(encoding="utf-8"))
ft_en_train = json.loads((DATASETS / "finetune_english" / "train.json").read_text(encoding="utf-8"))
ft_en_val = json.loads((DATASETS / "finetune_english" / "val.json").read_text(encoding="utf-8"))

ft_v1_all = ft_v1_train + ft_v1_val
ft_en_all = ft_en_train + ft_en_val

# Tag sources
for item in ft_v1_all:
    item["source"] = "finetune_v1"
for item in ft_en_all:
    item["source"] = "finetune_english"

# Count tokens using the same tokenizer
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_file(
    str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
)

def finetune_text(item):
    parts = []
    if item.get("instruction"): parts.append(item["instruction"])
    if item.get("input"): parts.append(item["input"])
    if item.get("output"): parts.append(item["output"])
    return " ".join(parts)

def count_tokens_batch(texts, label=""):
    total = 0
    BATCH = 10000
    for i in range(0, len(texts), BATCH):
        batch = texts[i:i+BATCH]
        encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
        total += sum(len(e.ids) for e in encoded)
    return total

print(f"\n  finetune_v1:      train={len(ft_v1_train):>7,}  val={len(ft_v1_val):>6,}  total={len(ft_v1_all):>7,}")
print(f"  finetune_english: train={len(ft_en_train):>7,}  val={len(ft_en_val):>6,}  total={len(ft_en_all):>7,}")

print("\n  Counting finetune tokens...")
ft_v1_tokens = count_tokens_batch([finetune_text(x) for x in ft_v1_all], "v1")
ft_en_tokens = count_tokens_batch([finetune_text(x) for x in ft_en_all], "english")
ft_total = ft_v1_tokens + ft_en_tokens

print(f"\n  finetune_v1 tokens:      {ft_v1_tokens:>12,}")
print(f"  finetune_english tokens: {ft_en_tokens:>12,}")
print(f"  {'─' * 55}")
print(f"  TOTAL FINETUNE:          {ft_total:>12,} tokens")

# Save consolidated finetune
out_finetune = DATASETS / "consolidated_finetune"
out_finetune.mkdir(parents=True, exist_ok=True)

all_finetune = ft_v1_all + ft_en_all
with open(out_finetune / "all_finetune_data.json", "w", encoding="utf-8") as f:
    json.dump(all_finetune, f, ensure_ascii=False, indent=2)
print(f"\n  Saved: {out_finetune / 'all_finetune_data.json'}")
print(f"  Total samples: {len(all_finetune):,}")


# ══════════════════════════════════════════════════════════════════
# FINAL SUMMARY
# ══════════════════════════════════════════════════════════════════
print(f"\n{'=' * 65}")
print("  FINAL SUMMARY")
print("=" * 65)

summary = f"""
PRETRAINING (Base/data/consolidated_pretrain_litdata/)
  litdata_3b       : {len(chunks_3b):>4} chunks | {tokens_3b:>15,} tokens
  litdata_english  : {len(chunks_en):>4} chunks | {tokens_en:>15,} tokens
  ───────────────────────────────────────────────────────
  TOTAL            : {len(chunks_3b)+len(chunks_en):>4} chunks | {tokens_pretrain:>15,} tokens  ({tokens_pretrain/1e9:.3f}B)

FINETUNING (Base/Datasets/consolidated_finetune/)
  finetune_v1      : {len(ft_v1_all):>7,} samples | {ft_v1_tokens:>12,} tokens
  finetune_english : {len(ft_en_all):>7,} samples | {ft_en_tokens:>12,} tokens
  ───────────────────────────────────────────────────────
  TOTAL            : {len(all_finetune):>7,} samples | {ft_total:>12,} tokens

GRAND TOTAL TOKENS: {tokens_pretrain + ft_total:,}
"""
print(summary)

# Save summary
with open(out_pretrain / "SUMMARY.txt", "w", encoding="utf-8") as f:
    f.write(summary)
with open(out_finetune / "SUMMARY.txt", "w", encoding="utf-8") as f:
    f.write(summary)

print("Done! Summary saved to both consolidated folders.")