#!/usr/bin/env python3
"""
Analyze token-level sequence lengths in aes_all.jsonl using the actual tokenizer.
Applies the chat template to get real token counts.

Outputs statistics and a recommended max_seq_length.
"""
import json
import sys
from pathlib import Path

from transformers import AutoTokenizer

WORKSPACE = Path("/workspace/elinnos")
INPUT = WORKSPACE / "aes_all.jsonl"
TOKENIZER_SRC = WORKSPACE / "Qwen2.5-7B-Instruct"
CHAT_TEMPLATE = WORKSPACE / "elinnos-qwen2.5-7b-multi-ip-lora-v4" / "chat_template.jinja"


def main():
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(str(TOKENIZER_SRC), trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    if CHAT_TEMPLATE.is_file():
        tokenizer.chat_template = CHAT_TEMPLATE.read_text()
        print(f"Chat template loaded from {CHAT_TEMPLATE}")

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

    print(f"Total samples: {len(records)}")
    print("Tokenizing with chat template (this may take a minute)...")

    token_lengths = []
    for i, rec in enumerate(records):
        messages = rec["messages"]
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
        tokens = tokenizer(text, return_tensors=None, add_special_tokens=False)
        token_lengths.append(len(tokens["input_ids"]))

        if (i + 1) % 500 == 0:
            print(f"  Processed {i+1}/{len(records)}...")

    token_lengths.sort()
    n = len(token_lengths)

    print(f"\n{'='*60}")
    print("  TOKEN LENGTH STATISTICS")
    print(f"{'='*60}")
    print(f"  Samples:  {n}")
    print(f"  Min:      {token_lengths[0]}")
    print(f"  Max:      {token_lengths[-1]}")
    print(f"  Mean:     {sum(token_lengths)/n:.0f}")
    print(f"  Median:   {token_lengths[n//2]}")

    for p in [75, 90, 95, 97, 99, 100]:
        idx = min(int(n * p / 100), n - 1)
        print(f"  {p}th pct: {token_lengths[idx]}")

    # Recommend seq length
    p95 = token_lengths[min(int(n * 0.95), n - 1)]
    p99 = token_lengths[min(int(n * 0.99), n - 1)]
    p100 = token_lengths[-1]

    if p95 <= 4096:
        recommended = 4096
    elif p95 <= 8192:
        recommended = 8192
    elif p99 <= 12288:
        recommended = 12288
    else:
        recommended = min(p100, 16384)

    # Round up to nearest 512
    recommended = ((recommended + 511) // 512) * 512

    print(f"\n  Recommended max_seq_length: {recommended}")
    print(f"    (covers 95th pct={p95}, 99th pct={p99}, max={p100})")

    # Count how many samples would be truncated at various lengths
    for sl in [4096, 8192, 12288, 16384]:
        truncated = sum(1 for t in token_lengths if t > sl)
        print(f"  At seq_len={sl}: {truncated}/{n} ({truncated/n*100:.1f}%) samples truncated")

    out_path = WORKSPACE / "aes_training" / "data" / "recommended_seq_length.txt"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(str(recommended))
    print(f"\n  Written to: {out_path}")

    # Also save full stats
    stats_path = WORKSPACE / "aes_training" / "data" / "seq_length_stats.json"
    stats = {
        "total_samples": n,
        "min": token_lengths[0],
        "max": token_lengths[-1],
        "mean": sum(token_lengths) / n,
        "median": token_lengths[n // 2],
        "percentiles": {
            str(p): token_lengths[min(int(n * p / 100), n - 1)]
            for p in [50, 75, 90, 95, 97, 99, 100]
        },
        "recommended_max_seq_length": recommended,
    }
    with stats_path.open("w") as f:
        json.dump(stats, f, indent=2)
    print(f"  Stats saved to: {stats_path}")


if __name__ == "__main__":
    main()