"""M2 — Extended Thinking Budget for ARCHON.

Pattern Anthropic Claude 3.7+: separate `<thinking>...</thinking>` budget,
not shown to user but computed via extra forward passes.

ARCHON ChatML v3 already has `task_type` special token (32005). We add
`<thinking>` and `</thinking>` as user-prompt-side delimiters, OR train
SFT with thinking traces from DeepSeek R1 distillation.

Bénéfice (Anthropic): GPQA 84.8% @ 64K thinking budget on Sonnet 4.6.
Pour ARCHON 282M: cible amélioration log-scale +5-10% sur math/code @ 8K thinking.
"""
from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass
class ExtendedThinkingConfig:
    """Thinking budget hyper-params."""

    thinking_budget_tokens: int = 8192  # ARCHON-scale (vs 64K Claude)
    final_answer_max_tokens: int = 1024
    thinking_start_token: str = "<thinking>"
    thinking_end_token: str = "</thinking>"
    answer_start_token: str = "<final>"
    answer_end_token: str = "</final>"
    # Sampling for thinking (more exploration)
    thinking_temperature: float = 0.95
    thinking_top_p: float = 0.95
    # Sampling for final (more deterministic)
    final_temperature: float = 0.3
    final_top_p: float = 0.8


@torch.no_grad()
def generate_with_thinking(
    model,
    tokenizer,
    prompt: str,
    cfg: ExtendedThinkingConfig = ExtendedThinkingConfig(),
) -> dict:
    """Two-phase generation: thinking + final answer.

    Phase 1: generate `<thinking>...</thinking>` (up to budget tokens)
    Phase 2: generate `<final>...</final>` from prompt + thinking

    Returns:
        {"thinking": str, "final": str, "thinking_tokens": int, "final_tokens": int}
    """
    # Build prompt
    full_prompt = f"{prompt}\n{cfg.thinking_start_token}\n"
    input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.to(next(model.parameters()).device)

    # Phase 1: thinking
    think_out = model.generate(
        input_ids,
        max_new_tokens=cfg.thinking_budget_tokens,
        temperature=cfg.thinking_temperature,
        top_p=cfg.thinking_top_p,
    )
    think_text = tokenizer.decode(think_out[0][input_ids.shape[-1]:], skip_special_tokens=False)
    # Truncate at </thinking>
    end_marker = cfg.thinking_end_token
    if end_marker in think_text:
        think_text = think_text.split(end_marker)[0]
    n_think_tokens = think_out.shape[-1] - input_ids.shape[-1]

    # Phase 2: final answer conditioned on thinking
    final_prompt = (full_prompt + think_text + cfg.thinking_end_token + "\n" +
                    cfg.answer_start_token + "\n")
    final_input_ids = tokenizer(final_prompt, return_tensors="pt").input_ids.to(input_ids.device)
    final_out = model.generate(
        final_input_ids,
        max_new_tokens=cfg.final_answer_max_tokens,
        temperature=cfg.final_temperature,
        top_p=cfg.final_top_p,
    )
    final_text = tokenizer.decode(
        final_out[0][final_input_ids.shape[-1]:], skip_special_tokens=False
    )
    if cfg.answer_end_token in final_text:
        final_text = final_text.split(cfg.answer_end_token)[0]
    n_final_tokens = final_out.shape[-1] - final_input_ids.shape[-1]

    return {
        "thinking": think_text.strip(),
        "final": final_text.strip(),
        "thinking_tokens": n_think_tokens,
        "final_tokens": n_final_tokens,
    }


def thinking_sft_data_recipe() -> str:
    """Doc: how to SFT ARCHON to use thinking properly."""
    return """
    SFT recipe for extended thinking:
    1. Distill 10K DeepSeek R1 traces on math/code/reasoning prompts
       (DeepSeek-R1 Distill-Qwen-1.5B traces are CC-BY freely available)
    2. Reformat as ChatML:
       <|im_start|>user
       {prompt}
       <|im_end|>
       <|im_start|>assistant
       <thinking>
       {R1_trace_truncated_8K}
       </thinking>
       <final>
       {final_answer}
       </final>
       <|im_end|>
    3. Mix 15-20% with regular SFT v2 data (avoid catastrophic forgetting)
    4. Train 1-2K extra steps on ARCHON SFT v2 -> M2-enabled v2.1
    """


if __name__ == "__main__":
    cfg = ExtendedThinkingConfig()
    print(f"[M2 ExtThink] budget={cfg.thinking_budget_tokens} tokens")
    print(thinking_sft_data_recipe())