#!/usr/bin/env python3
"""
repetition_probe.py — A builder tool for tiny model developers.

Takes a HuggingFace model and probes how different repetition_penalty values
affect generation quality. Outputs a comparison table with:
  - Repetition rate (fraction of tokens already seen in the output)
  - Unique token ratio (diversity)
  - Loop detection (repeated n-grams)
  - Sample text

Usage:
    python3 repetition_probe.py --model kevin-bretz/Small-Language-Model --prompt "Once upon a time"
    python3 repetition_probe.py --model kevin-bretz/Small-Language-Model --prompts "Once upon a time|The cat sat|I want to"

Why this matters:
  Tiny models (<50M params) are extremely prone to degenerate repetition loops.
  The default repetition_penalty=1.0 means NO penalty — the model can output the
  same token forever. Most training scripts don't test generation quality, so
  you discover this at inference time with no idea what to do.

  This tool gives you a quick diagnostic: sweep the penalty, see where quality
  peaks, and know what to set in your inference config.
"""

import argparse
import json
import sys
import time
from collections import Counter
from dataclasses import dataclass, field
from typing import List, Optional

import torch
import torch.nn.functional as F


@dataclass
class ProbeResult:
    prompt: str
    penalty: float
    text: str
    n_tokens: int
    repetition_rate: float  # fraction of output tokens that appeared before
    unique_token_ratio: float  # unique tokens / total tokens
    max_loop_len: int  # longest repeated trigram
    n_repeated_trigrams: int  # count of trigrams appearing 2+ times
    elapsed_s: float


def compute_repetition_rate(token_ids: List[int]) -> float:
    """Fraction of tokens (excluding first) that have appeared earlier in the sequence."""
    if len(token_ids) < 2:
        return 0.0
    seen = set()
    repeated = 0
    for t in token_ids[1:]:
        if t in seen:
            repeated += 1
        seen.add(t)
    return repeated / (len(token_ids) - 1)


def compute_unique_ratio(token_ids: List[int]) -> float:
    """Unique tokens / total tokens."""
    if not token_ids:
        return 0.0
    return len(set(token_ids)) / len(token_ids)


def find_repeated_trigrams(token_ids: List[int]) -> tuple:
    """Find the longest repeated trigram and count of repeated trigrams."""
    if len(token_ids) < 6:
        return 0, 0
    trigrams = []
    for i in range(len(token_ids) - 2):
        trigrams.append(tuple(token_ids[i:i+3]))
    counts = Counter(trigrams)
    repeated = sum(1 for c in counts.values() if c >= 2)
    # Longest repeated trigram: find the longest n-gram that repeats
    max_len = 0
    for n in range(3, min(len(token_ids), 20)):
        ngrams = [tuple(token_ids[i:i+n]) for i in range(len(token_ids) - n + 1)]
        ngram_counts = Counter(ngrams)
        if any(c >= 2 for c in ngram_counts.values()):
            max_len = n
        else:
            break
    return max_len, repeated


def generate_with_penalty(
    model,
    tokenizer,
    prompt: str,
    max_new_tokens: int = 64,
    repetition_penalty: float = 1.0,
    temperature: float = 0.8,
    top_k: int = 50,
    device: str = "cpu",
) -> tuple:
    """Generate text with a given repetition penalty. Returns (text, token_ids, elapsed)."""
    model.eval()
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    gen_start = input_ids.shape[1]

    with torch.no_grad():
        generated = input_ids
        for _ in range(max_new_tokens):
            outputs = model(generated)
            logits = outputs.logits[0, -1, :]

            # Apply repetition penalty
            if repetition_penalty != 1.0:
                for i in range(generated.shape[1]):
                    tok = generated[0, i].item()
                    if logits[tok] > 0:
                        logits[tok] /= repetition_penalty
                    else:
                        logits[tok] *= repetition_penalty

            # Temperature
            if temperature != 1.0:
                logits = logits / temperature

            # Top-k
            if top_k > 0:
                top_k_val = min(top_k, logits.shape[-1])
                indices_to_remove = logits < torch.topk(logits, top_k_val)[0][..., -1:]
                logits[indices_to_remove] = float("-inf")

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1).unsqueeze(0)
            generated = torch.cat([generated, next_token], dim=1)

            # Stop on EOS
            if next_token.item() == tokenizer.eos_token_id:
                break

    elapsed = None  # measured outside
    output_ids = generated[0, gen_start:].tolist()
    text = tokenizer.decode(output_ids, skip_special_tokens=True)
    return text, output_ids


def run_probe(
    model_name: str,
    prompts: List[str],
    penalties: List[float],
    max_new_tokens: int = 64,
    temperature: float = 0.8,
    top_k: int = 50,
    device: str = "cpu",
    seed: int = 42,
) -> List[ProbeResult]:
    """Run the full probe sweep."""
    print(f"Loading model: {model_name}")
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32).to(device)

    torch.manual_seed(seed)
    results = []

    for prompt in prompts:
        for penalty in penalties:
            torch.manual_seed(seed)  # Same seed for fair comparison
            start = time.time()
            text, token_ids = generate_with_penalty(
                model, tokenizer, prompt,
                max_new_tokens=max_new_tokens,
                repetition_penalty=penalty,
                temperature=temperature,
                top_k=top_k,
                device=device,
            )
            elapsed = time.time() - start

            rep_rate = compute_repetition_rate(token_ids)
            unique_ratio = compute_unique_ratio(token_ids)
            max_loop, n_rep_tri = find_repeated_trigrams(token_ids)

            results.append(ProbeResult(
                prompt=prompt,
                penalty=penalty,
                text=text,
                n_tokens=len(token_ids),
                repetition_rate=rep_rate,
                unique_token_ratio=unique_ratio,
                max_loop_len=max_loop,
                n_repeated_trigrams=n_rep_tri,
                elapsed_s=elapsed,
            ))
            print(f"  [{prompt[:20]:20s}] rp={penalty:.1f}  rep={rep_rate:.1%}  "
                  f"uniq={unique_ratio:.1%}  max_loop={max_loop}  "
                  f"n_rep_tri={n_rep_tri}  ({len(token_ids)} tokens)")

    return results


def format_table(results: List[ProbeResult]) -> str:
    """Format results as a readable table."""
    lines = []
    lines.append(f"{'Prompt':<22} {'RP':>4} {'Tokens':>6} {'Repet':>7} {'Unique':>7} {'MaxLoop':>8} {'RepTri':>7}")
    lines.append("-" * 72)
    for r in results:
        lines.append(
            f"{r.prompt[:20]:<22} {r.penalty:>4.1f} {r.n_tokens:>6} "
            f"{r.repetition_rate:>7.1%} {r.unique_token_ratio:>7.1%} "
            f"{r.max_loop_len:>8} {r.n_repeated_trigrams:>7}"
        )
    lines.append("-" * 72)
    lines.append("")
    lines.append("Sample outputs (first 100 chars):")
    lines.append("")
    for r in results:
        lines.append(f"  [{r.prompt[:15]:15s}] rp={r.penalty:.1f}: {r.text[:100]}")
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Repetition penalty probe for tiny models")
    parser.add_argument("--model", required=True, help="HuggingFace model name or path")
    parser.add_argument("--prompt", default="Once upon a time", help="Single prompt")
    parser.add_argument("--prompts", default=None, help="Pipe-separated prompts")
    parser.add_argument("--penalties", default="1.0,1.1,1.2,1.3,1.5",
                        help="Comma-separated repetition penalty values")
    parser.add_argument("--max-tokens", type=int, default=64, help="Max new tokens to generate")
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-k", type=int, default=50)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", default=None, help="Save JSON results to file")
    args = parser.parse_args()

    if args.prompts:
        prompts = [p.strip() for p in args.prompts.split("|")]
    else:
        prompts = [args.prompt]

    penalties = [float(p) for p in args.penalties.split(",")]

    print(f"Model: {args.model}")
    print(f"Prompts: {prompts}")
    print(f"Penalties: {penalties}")
    print(f"Max tokens: {args.max_tokens}, Temp: {args.temperature}, Top-k: {args.top_k}")
    print()

    results = run_probe(
        args.model, prompts, penalties,
        max_new_tokens=args.max_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        device=args.device,
        seed=args.seed,
    )

    print()
    print(format_table(results))

    if args.output:
        out = {
            "model": args.model,
            "prompts": prompts,
            "penalties": penalties,
            "max_tokens": args.max_tokens,
            "temperature": args.temperature,
            "top_k": args.top_k,
            "seed": args.seed,
            "results": [
                {
                    "prompt": r.prompt,
                    "penalty": r.penalty,
                    "text": r.text,
                    "n_tokens": r.n_tokens,
                    "repetition_rate": round(r.repetition_rate, 4),
                    "unique_token_ratio": round(r.unique_token_ratio, 4),
                    "max_loop_len": r.max_loop_len,
                    "n_repeated_trigrams": r.n_repeated_trigrams,
                }
                for r in results
            ],
        }
        with open(args.output, "w") as f:
            json.dump(out, f, indent=2)
        print(f"\nResults saved to {args.output}")


if __name__ == "__main__":
    main()