# ==============================================================================
# Multiplication SLM Benchmark Suite (Interactive TUI Wizard)
# Dynamic Dataset Generator corresponding to 'cot_partial', 'cot_direct', 'cot_cot'
# ==============================================================================

import json
import math
import os
import random
import subprocess
import sys
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

# ------------------------------------------------------------------------------
# 1. Install & Import Rich
# ------------------------------------------------------------------------------
try:
    from rich.align import Align
    from rich import box
    from rich.console import Console
    from rich.panel import Panel
    from rich.progress import (
        BarColumn,
        Progress,
        SpinnerColumn,
        TextColumn,
        TimeRemainingColumn,
    )
    from rich.prompt import Prompt
    from rich.table import Table
except ImportError:
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "rich", "-q"]
    )
    from rich.align import Align
    from rich import box
    from rich.console import Console
    from rich.panel import Panel
    from rich.progress import (
        BarColumn,
        Progress,
        SpinnerColumn,
        TextColumn,
        TimeRemainingColumn,
    )
    from rich.prompt import Prompt
    from rich.table import Table

from transformers import GPT2LMHeadModel

console = Console()

# ------------------------------------------------------------------------------
# 2. Tokenizer Definition
# ------------------------------------------------------------------------------
DEFAULT_CHARS = sorted(
    list(
        set(
            "0123456789 *+=:\nQStepsAnsLet'sthinkstepby"
            + "<|endoftext|>"
            + "<|pad|>"
        )
    )
)

class CharTokenizer:
    def __init__(self, chars=None):
        if chars is None:
            chars = DEFAULT_CHARS
        self.chars = sorted(list(set(chars)))
        if "<|pad|>" not in self.chars:
            self.chars.append("<|pad|>")
        self.c2i = {ch: i for i, ch in enumerate(self.chars)}
        self.i2c = {i: ch for i, ch in enumerate(self.chars)}
        self.pad_id = self.c2i["<|pad|>"]
        self.vocab_size = len(self.chars)

    @classmethod
    def load(cls, filepath):
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        return cls(data["chars"])

    def encode(self, text):
        return [self.c2i[c] for c in text if c in self.c2i]

    def decode(self, ids):
        return "".join(
            [self.i2c[i] for i in ids if i in self.i2c and i != self.pad_id]
        )

# ------------------------------------------------------------------------------
# 3. Interactive TUI Wizard
# ------------------------------------------------------------------------------
console.clear()
console.print()
console.print(
    Panel(
        Align.center(
            "[bold yellow]🤗 Hugging Face Micro-SLM Benchmark Harness[/bold yellow]\n"
            "[dim]10,000 Qs Log-Likelihood & Exact Match  |  500 Qs Perplexity (PPL)\n"
            "Rule: Strictly for Small Models (≤ 250M Parameters)[/dim]"
        ),
        box=box.ROUNDED,
        border_style="yellow",
        padding=(1, 2),
    )
)

# Step 1: Model Directory
candidate_paths = [
    d
    for d in ["./slm_mult_model", "./slm_mult_checkpoints", "."]
    if os.path.exists(os.path.join(d, "config.json"))
    or os.path.exists(os.path.join(d, "tokenizer.json"))
]
detected_hint = (
    f"[green]Found model at:[/green] [bold]{candidate_paths[0]}[/bold]"
    if candidate_paths
    else "[dim]No model detected in current folder.[/dim]"
)

input_dialog = Table.grid(padding=1)
input_dialog.add_column(style="bold cyan")
input_dialog.add_column()
input_dialog.add_row("Auto Detected:", detected_hint)
input_dialog.add_row("Default Path:", "[bold yellow]./slm_mult_model[/bold yellow]")

console.print(
    Panel(
        input_dialog,
        title="[bold white] 📁 Step 1/4: Model Directory [/bold white]",
        title_align="left",
        border_style="cyan",
        box=box.ROUNDED,
    )
)

while True:
    model_dir = Prompt.ask(
        " [bold cyan]➤ Model Directory Path[/bold cyan]",
        default=candidate_paths[0] if candidate_paths else "./slm_mult_model",
    )
    if os.path.exists(model_dir):
        break
    console.print("[bold red]Directory not found! Please try again.[/bold red]")

# Step 2: Model Display Name
default_name = f"YourOrg/{os.path.basename(os.path.abspath(model_dir))}"
console.print()
console.print(
    Panel(
        "[dim]Enter the model name for the leaderboard (e.g. OrgName/ModelName).[/dim]",
        title="[bold white] 🏷️ Step 2/4: Model ID / Display Name [/bold white]",
        title_align="left",
        border_style="cyan",
        box=box.ROUNDED,
    )
)
model_name = Prompt.ask(
    " [bold cyan]➤ Model Name[/bold cyan]",
    default=default_name
)

# Step 3: CoT Mode Selection
cot_table = Table(show_header=False, box=None, padding=(0, 1))
cot_table.add_row("[bold yellow]1. cot_partial[/bold yellow]", ": Step-by-step arithmetic decomposition (筆算分解)")
cot_table.add_row("[bold yellow]2. cot_direct [/bold yellow]", ": Direct answer output without intermediate steps (直接出力)")
cot_table.add_row("[bold yellow]3. cot_cot    [/bold yellow]", ": Chain-of-Thought prompt -> Final Answer only (プロンプトCoT/最終回答のみ)")

console.print()
console.print(
    Panel(
        cot_table,
        title="[bold white] 🧠 Step 3/4: CoT (Reasoning) Mode [/bold white]",
        title_align="left",
        border_style="cyan",
        box=box.ROUNDED,
    )
)

cot_choice = Prompt.ask(
    " [bold cyan]➤ Select CoT Mode[/bold cyan]",
    choices=["cot_partial", "cot_direct", "cot_cot"],
    default="cot_partial"
)

# Step 4: Highlight Tag
console.print()
console.print(
    Panel(
        "[dim]Highlight this entry as a BASE model on the leaderboard?[/dim]",
        title="[bold white] ⭐ Step 4/4: Base Model Highlight [/bold white]",
        title_align="left",
        border_style="cyan",
        box=box.ROUNDED,
    )
)
is_highlight_str = Prompt.ask(
    " [bold cyan]➤ Highlight as Base Model?[/bold cyan]",
    choices=["y", "n"],
    default="n"
)
is_highlight = (is_highlight_str.lower() == "y")

# ------------------------------------------------------------------------------
# 4. Model Loading & Verification
# ------------------------------------------------------------------------------
tokenizer_path = os.path.join(model_dir, "tokenizer.json")
if os.path.exists(tokenizer_path):
    tokenizer = CharTokenizer.load(tokenizer_path)
    tokenizer_status = f"[green]Loaded from ({tokenizer_path})[/green]"
else:
    tokenizer = CharTokenizer(DEFAULT_CHARS)
    tokenizer_status = "[yellow]Standard Vocab (Fallback)[/yellow]"

device = "cuda" if torch.cuda.is_available() else "cpu"

with console.status(f"[bold cyan]Loading weights onto [bold green]{device.upper()}[/bold green]...[/bold cyan]"):
    model = GPT2LMHeadModel.from_pretrained(model_dir).to(device)
    model.eval()

total_params = sum(p.numel() for p in model.parameters())
params_str = (
    f"{total_params / 1e6:.1f}M"
    if total_params >= 1e6
    else f"{total_params / 1e3:.1f}K"
)
is_qualified = total_params <= 250_000_000

info_table = Table(show_header=False, box=None, padding=(0, 2))
info_table.add_row("Model Name:", f"[bold white]{model_name}[/bold white]")
info_table.add_row("Directory:", f"[dim]{os.path.abspath(model_dir)}[/dim]")
info_table.add_row("CoT Mode:", f"[bold yellow]{cot_choice}[/bold yellow]")
info_table.add_row("Highlighted:", f"[bold magenta]{is_highlight}[/bold magenta]")
info_table.add_row("Total Parameters:", f"[bold cyan]{params_str}[/bold cyan] ({total_params:,})")
info_table.add_row(
    "Eligibility (≤ 250M):",
    "[bold green]✔ QUALIFIED[/bold green]" if is_qualified else "[bold red]✖ DISQUALIFIED (>250M)[/bold red]"
)

console.print()
console.print(
    Panel(
        info_table,
        title="[bold green] 🤖 Model Specification Confirmed [/bold green]",
        title_align="left",
        border_style="green",
        box=box.ROUNDED,
    )
)

# ------------------------------------------------------------------------------
# 5. Dynamic Dataset Generator (Based on cot_choice)
# ------------------------------------------------------------------------------
def generate_benchmark_sample(cot_mode="cot_partial"):
    mode = random.choices(["2terms", "3terms"], weights=[0.7, 0.3])[0]

    if mode == "2terms":
        a = random.randint(2, 999)
        b = random.randint(2, 99)
        expr = f"{a} * {b}"
        ans = a * b

        if cot_mode == "cot_direct":
            prompt = f"Q: {expr}\nAns:"
            completion = f" {ans}<|endoftext|>"

        elif cot_mode == "cot_cot":
            prompt = f"Q: {expr}\nLet's think step by step.\nAns:"
            completion = f" {ans}<|endoftext|>"

        else:  # cot_partial (standard step-by-step)
            if b >= 10:
                steps = []
                sub_products = []
                for i, digit in enumerate(reversed(str(b))):
                    p_val = int(digit) * (10**i)
                    if p_val > 0:
                        part_ans = a * p_val
                        sub_products.append(part_ans)
                        steps.append(f"Step {len(steps)+1}: {a} * {p_val} = {part_ans}")
                if len(sub_products) > 1:
                    add_expr = " + ".join(map(str, sub_products))
                    steps.append(f"Step {len(steps)+1}: {add_expr} = {ans}")
                steps_text = "\n".join(steps)
            else:
                steps_text = f"Step 1: {expr} = {ans}"

            prompt = f"Q: {expr}\nSteps:\n"
            completion = f"{steps_text}\nAns: {ans}<|endoftext|>"

    else:  # 3terms
        a = random.randint(2, 50)
        b = random.randint(2, 20)
        c = random.randint(2, 20)
        expr = f"{a} * {b} * {c}"
        r1, r2 = a * b, a * b * c

        if cot_mode == "cot_direct":
            prompt = f"Q: {expr}\nAns:"
            completion = f" {r2}<|endoftext|>"

        elif cot_mode == "cot_cot":
            prompt = f"Q: {expr}\nLet's think step by step.\nAns:"
            completion = f" {r2}<|endoftext|>"

        else:  # cot_partial
            prompt = f"Q: {expr}\nSteps:\n"
            completion = f"Step 1: {a} * {b} = {r1}\nStep 2: {r1} * {c} = {r2}\nAns: {r2}<|endoftext|>"

    return prompt, completion

class EvalDataset(Dataset):
    def __init__(self, samples, max_len=160):
        self.items = []
        for prompt, comp in samples:
            full_text = prompt + comp
            p_tokens = tokenizer.encode(prompt)
            full_tokens = tokenizer.encode(full_text)
            prompt_len = len(p_tokens)
            seq_len = len(full_tokens)

            if seq_len > max_len:
                full_tokens = full_tokens[:max_len]

            labels = list(full_tokens)
            for i in range(min(prompt_len, len(labels))):
                labels[i] = -100

            pad_len = max_len - len(full_tokens)
            input_ids = full_tokens + [tokenizer.pad_id] * pad_len
            labels = labels + [-100] * pad_len

            self.items.append(
                {
                    "input_ids": torch.tensor(input_ids, dtype=torch.long),
                    "labels": torch.tensor(labels, dtype=torch.long),
                }
            )

    def __len__(self):
        return len(self.items)

    def __getitem__(self, idx):
        return self.items[idx]

with console.status(f"[bold blue]Generating 10,500 questions in '[yellow]{cot_choice}[/yellow]' mode...[/bold blue]"):
    random.seed(2025)
    eval_10k_data = [generate_benchmark_sample(cot_choice) for _ in range(10000)]
    eval_ppl_data = [generate_benchmark_sample(cot_choice) for _ in range(500)]
    dataset_10k = EvalDataset(eval_10k_data)
    dataset_ppl = EvalDataset(eval_ppl_data)

# Auto Batch Size Discovery
def find_optimal_batch_size(model, min_b=1, max_b=1024, seq_len=160, device="cuda"):
    if device == "cpu": return 64
    candidate_sizes = [2**i for i in range(int(math.log2(min_b)), int(math.log2(max_b)) + 1)]
    optimal_b = min_b
    with console.status("[bold magenta]Probing maximum safe batch size...[/bold magenta]"):
        for b in candidate_sizes:
            try:
                torch.cuda.empty_cache()
                dummy_input = torch.randint(0, tokenizer.vocab_size, (b, seq_len), device=device)
                with torch.no_grad():
                    _ = model(dummy_input)
                optimal_b = b
            except Exception:
                break
    return optimal_b

auto_batch_size = find_optimal_batch_size(model, device=device)
console.print(f"\n[bold green]✔ Hardware Optimized:[/bold green] Batch Size set to [bold cyan]{auto_batch_size}[/bold cyan]\n")

# ------------------------------------------------------------------------------
# 6. Benchmark Execution Loop
# ------------------------------------------------------------------------------
loader_10k = DataLoader(dataset_10k, batch_size=auto_batch_size, shuffle=False)
loader_ppl = DataLoader(dataset_ppl, batch_size=auto_batch_size, shuffle=False)

total_log_likelihood = 0.0
total_tokens = 0
correct_tokens = 0
perfect_sequences = 0
total_sequences = 0

with Progress(
    SpinnerColumn(),
    TextColumn("[progress.description]{task.description}"),
    BarColumn(complete_style="cyan", finished_style="green"),
    TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
    TimeRemainingColumn(),
    console=console,
) as progress:

    task1 = progress.add_task(f"[cyan]Evaluating 10,000 Qs ({cot_choice})...", total=len(loader_10k))
    with torch.no_grad():
        for batch in loader_10k:
            input_ids = batch["input_ids"].to(device)
            labels = batch["labels"].to(device)
            outputs = model(input_ids)
            logits = outputs.logits

            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()

            log_probs = F.log_softmax(shift_logits, dim=-1)
            mask = shift_labels != -100

            gathered = torch.gather(log_probs, 2, shift_labels.unsqueeze(-1).clamp(min=0)).squeeze(-1)
            gathered = gathered * mask

            total_log_likelihood += gathered.sum().item()
            total_tokens += mask.sum().item()

            preds = torch.argmax(shift_logits, dim=-1)
            matches = (preds == shift_labels) & mask
            correct_tokens += matches.sum().item()

            seq_matches = (matches.sum(dim=1) == mask.sum(dim=1)).sum().item()
            perfect_sequences += seq_matches
            total_sequences += input_ids.size(0)

            progress.update(task1, advance=1)

    task2 = progress.add_task("[magenta]Calculating Perplexity on 500 Qs...", total=len(loader_ppl))
    total_loss = 0.0
    ppl_batches = 0
    loss_fn = torch.nn.CrossEntropyLoss(ignore_index=-100)

    with torch.no_grad():
        for batch in loader_ppl:
            input_ids = batch["input_ids"].to(device)
            labels = batch["labels"].to(device)
            outputs = model(input_ids)
            shift_logits = outputs.logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()

            loss = loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
            total_loss += loss.item()
            ppl_batches += 1
            progress.update(task2, advance=1)

# Metrics Calculation
token_accuracy = (correct_tokens / total_tokens) * 100
sequence_accuracy = (perfect_sequences / total_sequences) * 100
avg_loss = total_loss / ppl_batches
ppl_score = math.exp(avg_loss)

p = sequence_accuracy / 100.0
p_smooth = min(max(p, 1e-4), 1.0 - 1e-5)
logit_val = math.log(p_smooth / (1.0 - p_smooth))
rating_score = max(0.0, 1050.0 + (252.47 * logit_val))

# ------------------------------------------------------------------------------
# 7. HTML Leaderboard Ready JSON Output
# ------------------------------------------------------------------------------
hf_leaderboard_entry = {
    "name": model_name,
    "params_str": params_str,
    "params_num": total_params,
    "cot": cot_choice,
    "score": round(rating_score, 1),
    "exact": round(sequence_accuracy, 1),
    "token_acc": round(token_accuracy, 1),
    "ppl": round(ppl_score, 3),
    "highlight": is_highlight
}

json_str = json.dumps(hf_leaderboard_entry, indent=4)

console.print()
console.print(
    Panel(
        f"[bold green]{json_str}[/bold green]",
        title="[bold yellow] 📋 HTML Ready Leaderboard Entry (Copy & Paste to modelData) [/bold yellow]",
        title_align="left",
        border_style="yellow",
        box=box.ROUNDED,
    )
)
console.print("[dim]Copy the JSON above and paste it directly into the 'modelData' array inside index.html![/dim]\n")