1️⃣ CoT Infrastructure & Sentence Taxonomy

# TODO - actually use these! or delete them

# Configuration
MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
DATASET_NAME = "uzaymacar/math-rollouts"  # Pre-computed rollouts from paper
SIMILARITY_THRESHOLD = 0.8  # Median threshold from paper
EMBEDDING_MODEL = "all-MiniLM-L6-v2"  # Sentence embedding model used in paper
N_ROLLOUTS = 100  # Number of rollouts per sentence

# # Paths (adjust these to match your setup)
# ROLLOUTS_DIR = Path("math_rollouts")  # Directory with generated rollouts
# ANALYSIS_DIR = Path("analysis")  # Directory to save analysis results
# ANALYSIS_DIR.mkdir(exist_ok=True, parents=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

embedding_model = SentenceTransformer(EMBEDDING_MODEL)
print(embedding_model)

prompts = [
    "Wait, I think I made an error in my reasoning and need to backtrack",
    "Hold on, I believe I made a mistake in my logic and should reconsider",
    "After careful analysis, I've determined the correct answer is 42",
    "Time is an illusion. Lunchtime doubly so.",
]
labels = [x[:35] + "..." for x in prompts]

embedding = embedding_model.encode(prompts)

cosine_sims = embedding @ embedding.T

px.imshow(
    cosine_sims,
    color_continuous_scale="RdBu",
    color_continuous_midpoint=0,
    labels=dict(x="Prompt", y="Prompt", color="Cosine Similarity"),
    x=labels,
    y=labels,
)

# dataset = load_dataset(DATASET_NAME, split="default", streaming=True)

PROBLEM_ID = 4682

path = f"deepseek-r1-distill-llama-8b/temperature_0.6_top_p_0.95/correct_base_solution/problem_{PROBLEM_ID}/base_solution.json"

local_path = hf_hub_download(repo_id=DATASET_NAME, filename=path, repo_type="dataset")

with open(local_path, "r") as f:
    problem_data = json.load(f)

problem_data

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype=torch.bfloat16,  # Use bfloat16 for efficiency
    device_map="auto",  # Automatically distribute across available GPUs
    trust_remote_code=True,
)
model = model.to(device)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

print(f"Model loaded on: {model.device}")
print(f"Model has {model.config.num_hidden_layers} layers")
print(f"Model has {model.config.num_attention_heads} attention heads per layer")

problem_data["chunk_solutions"][10][0].keys()

easy_problem_text = "A rectangle has a length of 8 cm and a width of 5 cm. What is its perimeter?"

easy_prompt = f"""Solve this math problem step by step.

Problem: {easy_problem_text}

Solution:
<think>"""

inputs = tokenizer(easy_prompt, return_tensors="pt").to(device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.6,
        top_p=0.95,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )

generated_text = tokenizer.decode(
    outputs[0][len(inputs["input_ids"][0]) :], skip_special_tokens=False
)
easy_problem_full_text = easy_prompt + generated_text

print(easy_problem_full_text)

# TODO(mcdougallc) maybe we don't want this, because we want the actual answer after the CoT? idk


class StopOnThink(StoppingCriteria):
    def __call__(self, input_ids, scores, **kwargs):
        # YOUR CODE HERE: return True iff the model has generated "</think>"
        think_token_id = tokenizer.encode("</think>", add_special_tokens=False)[0]
        return input_ids[0, -1] == think_token_id


with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.6,
        top_p=0.95,
        do_sample=True,
        stopping_criteria=[StopOnThink()],
        pad_token_id=tokenizer.eos_token_id,
    )

generated_text = tokenizer.decode(
    outputs[0][len(inputs["input_ids"][0]) :], skip_special_tokens=False
)
easy_problem_full_text = easy_prompt + generated_text

print(easy_problem_full_text)

def split_solution_into_chunks(text: str) -> list[str]:
    """Split solution into sentence-level chunks."""

    # YOUR CODE HERE - fill in the rest of the function

    # Remove thinking tags
    if "<think>" in text:
        text = text.split("<think>")[1]
    if "</think>" in text:
        text = text.split("</think>")[0]
    text = text.strip()

    # Replace the "." characters which I don't want to split on
    text = re.sub(r"(\d)\.(\d)", r"\1<DECIMAL>\2", text)  # e.g. "4.5" -> "4<DECIMAL>5"
    text = re.sub(r"\n(\d)\.(\s)", r"\n\1<DECIMAL>\2", text)  # e.g. "1.\n" -> "\n1<DECIMAL>"

    # Split on sentence endings, and combine the endings with the previous chunk
    sentences = re.split(r"([!?:\n]|(?<!\n\d)\.)", text)
    chunks = []
    for i in range(0, len(sentences) - 1, 2):
        chunks.append((sentences[i] + sentences[i + 1]).replace("\n", " "))

    # Replace <DECIMAL> back with "."
    chunks = [re.sub(r"<DECIMAL>", ".", c) for c in chunks]

    # Merge chunks that are too short
    while len(chunks) > 1 and min([len(x) for x in chunks[:-1]]) < 10:
        for i, chunk in enumerate(chunks[:-1]):
            if len(chunk) < 10:
                chunks = chunks[:i] + [chunk + chunks[i + 1]] + chunks[i + 2 :]
                break

    return [c.strip() for c in chunks if c.strip()]


test_cases = [
    # (input_text, expected_chunks)
    (
        "<think>First, I understand the problem. Next, I'll solve for x. Finally, I verify!</think>",
        ["First, I understand the problem.", "Next, I'll solve for x.", "Finally, I verify!"],
    ),
    (
        "<think>Let me break this down: 1. Convert to decimal. 2. Calculate log. 3. Apply formula.</think>",
        [
            "Let me break this down:",
            "1. Convert to decimal.",
            "2. Calculate log.",
            "3. Apply formula.",
        ],
    ),
    (
        "<think>The formula is A = πr². Wait. No. Actually, it's different.</think>",
        ["The formula is A = πr².", "Wait. No.", "Actually, it's different."],
    ),
    (
        "<think>Convert 66666₁₆ to decimal. This equals 419,430. How many bits? We need log₂(419,430) ≈ 18.7. So 19 bits!</think>",
        [
            "Convert 66666₁₆ to decimal.",
            "This equals 419,430.",
            "How many bits?",
            "We need log₂(419,430) ≈ 18.7.",
            "So 19 bits!",
        ],
    ),
    ("<think>The answer is 42. Done.</think>", ["The answer is 42.", "Done."]),
]

for input_text, expected_chunks in test_cases:
    chunks = split_solution_into_chunks(input_text)
    assert chunks == expected_chunks, f"Expected {expected_chunks}, got {chunks}"

CATEGORIES = {
    "problem_setup": "Problem Setup",
    "plan_generation": "Plan Generation",
    "fact_retrieval": "Fact Retrieval",
    "active_computation": "Active Computation",
    "uncertainty_management": "Uncertainty Management",
    "result_consolidation": "Result Consolidation",
    "self_checking": "Self Checking",
    "final_answer_emission": "Final Answer Emission",
    "unknown": "Unknown",
}

CATEGORY_WORDS = {
    # Note: we put the most definitive categories first, so they override the later ones
    "final_answer_emission": ["\\boxed", "final answer"],
    "problem_setup": ["need to", "problem is", "given"],
    "fact_retrieval": ["remember", "formula", "know that", "recall"],
    "active_computation": ["calculate", "compute", "solve", "=", "equals", "result", "giving"],
    "uncertainty_management": ["wait", "let me", "double check", "hmm", "actually", "reconsider"],
    "result_consolidation": ["summarize", "so", "therefore", "in summary"],
    "self_checking": ["verify", "check", "confirm", "correct"],
    "plan_generation": ["plan", "approach", "strategy", "will", "i'll", "try"],
}


def categorize_sentences_heuristic(chunks: list[str]) -> list[str]:
    """
    Categorize sentences using heuristics/keyword matching (simplified version).

    For full LLM-based labeling, see prompts.py DAG_PROMPT in the thought-anchors repo.

    Args:
        sentences: List of sentence strings

    Returns:
        List of tags.
    """

    categories = []

    for idx, chunk in enumerate(chunks):
        chunk_lower = chunk.lower()

        if idx == 0:
            tag = "problem_setup"
        else:
            for category, words in CATEGORY_WORDS.items():
                if any(word in chunk_lower for word in words):
                    tag = category
                    break
            else:
                tag = "unknown"

        categories.append(CATEGORIES.get(tag))

    return categories


example_problem_text = "What is the area of a circle with radius 5?"

example_sentences, example_categories = list(
    zip(
        *[
            ("I need to find the area of a circle with radius 5.", "Problem Setup"),
            ("The formula for circle area is A = πr².", "Fact Retrieval"),
            ("Substituting r = 5: A = π × 5² = 25π.", "Active Computation"),
            ("Wait, let me look again at that calculation.", "Uncertainty Management"),
            ("So the area is 25π square units.", "Result Consolidation"),
            ("Therefore, the answer is \\boxed{25π}.", "Final Answer Emission"),
        ]
    )
)

categories = categorize_sentences_heuristic(example_sentences)

for sentence, category, expected_category in zip(example_sentences, categories, example_categories):
    assert category == expected_category, (
        f"Expected {expected_category!r}, got {category!r} for sentence: {sentence!r}"
    )

easy_chunks = split_solution_into_chunks(easy_problem_full_text)

easy_categories = categorize_sentences_heuristic(easy_chunks)

for chunk, category in zip(easy_chunks, easy_categories):
    chunk_str = chunk if len(chunk) < 80 else chunk[:60] + " ... " + chunk[-20:]
    print(f"{category:>20} | {chunk_str!r}")

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

openai_client = OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None


def generate_api_response(
    model: str = "gpt-4.1-mini",
    messages: list[dict[str, str]] = [],
    max_tokens: int = 128,
    stop_sequences: list[str] | None = None,
    temperature: float = 0.0,
    max_retries: int = 3,
) -> str:
    """Helper function with retry logic and error handling."""

    for attempt in range(max_retries):
        try:
            # Generate response
            resp = openai_client.chat.completions.create(
                model=model,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                stop=stop_sequences if stop_sequences else None,
            )
            # Extract text from response
            return resp.choices[0].message.content or ""

        except Exception as e:
            if "rate_limit_error" in str(e) or "429" in str(e):
                if attempt < max_retries - 1:
                    # Exponential backoff: 2^attempt seconds (1, 2, 4, 8, 16...)
                    wait_time = 2**attempt
                    print(f"Rate limit hit, waiting {wait_time} seconds before retry...")
                    time.sleep(wait_time)
                    continue
                else:
                    print(f"Failed to get response after {max_retries} attempts, returning early.")
            raise e


resp = generate_api_response(messages=[{"role": "user", "content": "Who are you?"}])

print(textwrap.fill(resp))

DAG_SYSTEM_PROMPT = """
You are an expert in interpreting how language models solve math problems using multi-step reasoning. Your task is to analyze a Chain-of-Thought (CoT) reasoning trace, broken into discrete text chunks, and label each chunk with a tag that describes what this chunk is *doing* functionally in the reasoning process.

---

### Function Tags:

1. `problem_setup`: 
    Parsing or rephrasing the problem (initial reading or comprehension).

2. `plan_generation`: 
    Stating or deciding on a plan of action, or on the next step in the reasoning process.

3. `fact_retrieval`: 
    Recalling facts, formulas, problem details (without immediate computation).

4. `active_computation`: 
    Performing algebra, calculations, manipulations toward the answer.
    This only includes actual calculations being done & values computed, not stating formulas or plans.

5. `result_consolidation`: 
    Aggregating intermediate results, summarizing, or preparing final answer.

6. `uncertainty_management`: 
    Expressing confusion, re-evaluating, proposing alternative plans (includes backtracking).

7. `final_answer_emission`: 
    Explicit statement of the final boxed answer or earlier chunks that contain the final answer.

8. `self_checking`: 
    Verifying previous steps, Pythagorean checking, re-confirmations.
    This is at the object-level, whereas "uncertainty_management" is at the planning level.

9. `unknown`: 
    Use only if the chunk does not fit any of the above tags or is purely stylistic or semantic.

---

### Output Format:

Return a numbered list, one item for each chunk, consisting of the function tag that best describes the chunk.

For example:

1. problem_setup
...
5. final_answer_emission
"""

def categorize_sentences_autorater(problem_text: str, chunks: list[str]) -> list[str]:
    """
    Categorize sentences using heuristics/keyword matching (simplified version).

    Args:
        sentences: List of sentence strings

    Returns:
        List of categories.
    """

    # YOUR CODE HERE

    chunk_str = ""
    for i, chunk in enumerate(chunks):
        chunk_str += f"{i + 1}. {chunk}\n"

    user_prompt = f"""
Here is the math problem:

[PROBLEM]
{problem_text}

Here is the full Chain of Thought, broken into chunks:

[CHUNKS]
{chunk_str.strip()}

Now label each chunk with function tags and dependencies."""

    raw_response = generate_api_response(
        messages=[
            {"role": "system", "content": DAG_SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ],
    )
    response = re.split(r"\n\d{1,2}\.", "\n" + raw_response.strip())
    response = [r.strip() for r in response if r.strip()]

    assert len(response) == len(chunks), f"Length mismatch: {len(response)} != {len(chunks)}"

    return [CATEGORIES.get(r, "Unknown") for r in response]

example_categories = categorize_sentences_autorater(example_problem_text, example_sentences)

for chunk, category in zip(example_sentences, example_categories):
    print(f"{category:>25} | {chunk!r}")

easy_categories_autorater = categorize_sentences_autorater(easy_problem_text, easy_chunks)

df = pd.DataFrame(
    {
        "Category (heuristics)": easy_categories,
        "Category (autorater)": easy_categories_autorater,
        "Chunk": easy_chunks,
    }
)
with pd.option_context("display.max_colwidth", None):
    display(df)

CATEGORY_COLORS = {
    "Problem Setup": "#4285F4",
    "Plan Generation": "#EA4335",
    "Fact Retrieval": "#FBBC05",
    "Active Computation": "#34A853",
    "Uncertainty Management": "#9C27B0",
    "Result Consolidation": "#00BCD4",
    "Self Checking": "#FF9800",
    "Final Answer Emission": "#795548",
    "Unknown": "#9E9E9E",
}


def visualize_trace_structure(chunks: list[str], categories: list[str], problem_text: str = None):
    """Visualize a reasoning trace with color-coded sentence categories."""

    n_chunks = len(chunks)
    fig, ax = plt.subplots(figsize=(12, 1 + int(0.5 * n_chunks)))

    for idx, (chunk, category) in enumerate(zip(chunks, categories)):
        color = CATEGORY_COLORS.get(category, "#9E9E9E")

        y = n_chunks - idx  # Start from top

        # Category label with colored background
        ax.barh(y, 0.15, left=0, height=0.8, color=color, alpha=0.6)
        ax.text(0.075, y, f"{category}", ha="center", va="center", fontsize=9, weight="bold")

        # Sentence text
        text = chunk[:100] + ("..." if len(chunk) > 100 else "")
        ax.text(0.17, y, f"[{idx}] {text}", va="center", fontsize=9)

    ax.set_xlim(0, 1)
    ax.set_ylim(0.5, n_chunks + 0.5)
    ax.axis("off")

    if problem_text:
        fig.suptitle(f"Problem: {problem_text[:120]}...", fontsize=11, y=0.98, weight="bold")

    plt.title("Reasoning Trace Structure", fontsize=13, pad=30)
    plt.tight_layout()
    plt.show()


visualize_trace_structure(easy_chunks, easy_categories, easy_problem_text)