linalg_zero/generator/utils.py

import ast
import json
import random
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

import numpy as np
from datasets import Dataset, DatasetDict
from sympy.core.random import seed

from linalg_zero.generator.models import Question
from linalg_zero.grpo.verify import parse_string, verify_answers
from linalg_zero.shared.lib import get_lib
from linalg_zero.shared.utils import get_logger

logger = get_logger(__name__)


def load_entropy_settings(path: str) -> dict[str, Any]:
    """
    Load entropy analysis results from JSON file.
    """
    results_file = Path(path)

    if not results_file.exists():
        raise FileNotFoundError(
            f"Entropy analysis results not found at {results_file}. "
            "Run the entropy analysis first using linalg_zero.generator.analysis.analyse"
        )

    with results_file.open("r", encoding="utf-8") as f:
        settings = json.load(f)

    return settings


def print_entropy_settings(settings: dict[str, Any]) -> None:
    """Print the loaded entropy settings in a readable format."""

    logger.info("=" * 60)
    logger.info("RECOMMENDED ENTROPY SETTINGS")
    logger.info("=" * 60)

    for problem_type, config in settings.items():
        combination = config["combination"]
        score = config["score"]
        logger.info(f"{problem_type}: {combination} (score: {score:.2f})")

    logger.info("=" * 60)


def _verify_step_result(step: dict[str, Any], lib: dict[str, Any]) -> None:
    """Verify a single step's result against library function output."""
    step_id = step["step_id"]

    if "result" not in step:
        raise ValueError(f"Step {step_id} has no result - implementation is bugged")

    result_value = parse_string(step["result"])
    if result_value is None:
        raise ValueError(f"Step {step_id}: invalid result - implementation is bugged")

    fn_type = step["tool"]
    lib_fn = lib[fn_type]
    input_data = json.loads(step["verification"]["input"])
    fn_result = lib_fn(**input_data)

    if not verify_answers(result_value, fn_result):
        raise ValueError(f"Step mismatch - step - {json.dumps(step)} - lib_fn - {fn_type}")


def _verify_step_dependencies(step: dict[str, Any], question_stepwise: list[dict[str, Any]]) -> None:
    """Verify step dependencies against referenced steps."""
    step_id = step.get("step_id", "unknown")
    dependent_on = step["verification"].get("dependent_on", None)

    if dependent_on is None:
        return

    if not isinstance(dependent_on, dict):
        raise TypeError(f"Step {step_id}: dependent_on must be a dict, got {type(dependent_on)}")

    # Verify each input_* field against its corresponding referenced step's result
    for input_name, input_value in step["verification"].items():
        if input_name.startswith("input_"):
            expected_step_index = dependent_on[input_name]

            # Validate the reference step exists
            if not isinstance(expected_step_index, int):
                raise TypeError(
                    f"Step {step_id}: dependency index for '{input_name}' must be an integer, got {type(expected_step_index)}"
                )

            if expected_step_index < 0 or expected_step_index >= len(question_stepwise):
                raise ValueError(
                    f"Step {step_id}: dependent_on index {expected_step_index} for '{input_name}' out of bounds "
                    f"(stepwise has {len(question_stepwise)} steps)"
                )

            referenced_step = question_stepwise[expected_step_index]
            referenced_result = parse_string(referenced_step["result"])

            if referenced_result is None:
                raise ValueError(f"Step {step_id}: referenced step {expected_step_index} has invalid result")

            field_value = json.loads(input_value)
            if not verify_answers(field_value, referenced_result) or field_value != referenced_result:
                raise ValueError(
                    f"Step {step_id}: dependency verification failed - "
                    f"{input_name} ({field_value}) does not match referenced step {expected_step_index} result ({referenced_result})"
                )


def _verify_golden_answer(question: Question, question_index: int) -> None:
    """Verify the golden answer matches the final stepwise result."""
    if not question.golden or "final_answer" not in question.golden:
        raise ValueError(f"Question {question_index} has no golden final answer - implementation is bugged")

    golden_value = parse_string(question.golden["final_answer"])
    answer_value = parse_string(question.stepwise[-1]["result"])

    if golden_value is None:
        raise ValueError(f"Question {question_index}: invalid golden answer - implementation is bugged")
    if answer_value is None:
        raise ValueError(f"Question {question_index}: invalid formatted answer - implementation is bugged")

    if not verify_answers(golden_value, answer_value):
        raise ValueError(
            f"Question {question_index}: Golden answer mismatch - implementation is bugged. "
            f"Golden={golden_value}, Answer={answer_value}"
        )


def verify_dataset(dataset: list[Question]) -> dict[str, Any]:
    """
    Verify a dataset of questions by checking constituent ground truths and target values.
    """
    # NOTE: this function is temporary

    verification_results = {
        "total_questions": len(dataset),
        "verified_questions": 0,
        "stepwise_verifications": 0,
        "golden_verifications": 0,
    }
    lib = get_lib()

    for i, question in enumerate(dataset):
        if len(question.stepwise) == 0:
            raise ValueError(f"Question {i} has no stepwise results - implementation is bugged")

        # Verify stepwise results
        for step in question.stepwise:
            _verify_step_result(step, lib)
            _verify_step_dependencies(step, question.stepwise)
            verification_results["stepwise_verifications"] += 1

        # Verify golden answer
        _verify_golden_answer(question, i)
        verification_results["golden_verifications"] += 1
        verification_results["verified_questions"] += 1

    logger.info(
        "Dataset verification complete: All %d questions verified successfully (%d stepwise checks, %d golden checks)",
        verification_results["total_questions"],
        verification_results["stepwise_verifications"],
        verification_results["golden_verifications"],
    )

    return verification_results


def check_constraints(dataset: list[Question], config: dict[str, Any], statistics: dict[str, Any]) -> None:
    """Check the constraints for the given config."""

    for problem_type, stats in statistics.get("per_problem_type", {}).items():
        if problem_type in config:
            actual_min = stats.get("min")
            actual_max = stats.get("max")
            expected_min = config[problem_type]["metadata"].get("target_min_value")
            expected_max = config[problem_type]["metadata"].get("target_max_value")

            if actual_min is not None and expected_min is not None and actual_min < expected_min:
                raise ValueError(f"{problem_type}: min {actual_min} < expected {expected_min}")
            if actual_max is not None and expected_max is not None and actual_max > expected_max:
                raise ValueError(f"{problem_type}: max {actual_max} > expected {expected_max}")


def set_seed(seed_val: int = 42) -> None:
    """Set the seed for the deterministic generation."""
    random.seed(seed_val)
    np.random.seed(seed_val)
    seed(seed_val)


def print_dataset(questions: list[Question], include_invalid: bool = False) -> None:  # pragma: no cover
    """Display a formatted dataset of questions."""

    questions_to_print = questions if include_invalid else [q for q in questions if q.is_valid]

    if not questions_to_print:
        logger.info("No questions to display.")
        return

    logger.info("=" * 30)
    logger.info("GENERATED DATASET")
    logger.info("=" * 30)

    # Questions
    for i, question in enumerate(questions_to_print, 1):
        status = " [INVALID]" if not question.is_valid else ""
        logger.info("Question %d:%s", i, status)
        logger.info("Q: %s", question.question)
        logger.info("A: %s", ast.literal_eval(question.answer))
        logger.info("")

    # Metadata
    topics = {q.topic for q in questions_to_print}
    problem_types = {q.problem_type for q in questions_to_print}
    difficulties = {q.difficulty for q in questions_to_print}
    entropy_values = [q.entropy_used for q in questions_to_print]
    tool_calls = [q.tool_calls_required for q in questions_to_print]

    # Summary
    logger.info("Dataset Summary:")
    logger.info("  Total Questions: %d", len(questions_to_print))
    logger.info("  Topics: %s", ", ".join(sorted(topic.value for topic in topics)))
    logger.info("  Problem Types: %s", ", ".join(sorted(pt.value for pt in problem_types)))
    logger.info("  Difficulties: %s", ", ".join(sorted(str(difficulty) for difficulty in difficulties)))
    logger.info(
        "  Entropy Used: %.2f - %.2f (avg: %.2f)",
        min(entropy_values),
        max(entropy_values),
        sum(entropy_values) / len(entropy_values),
    )
    logger.info(
        "  Tool Calls Required: %d - %d (avg: %.1f)",
        min(tool_calls),
        max(tool_calls),
        sum(tool_calls) / len(tool_calls),
    )
    # Distributions
    by_difficulty = Counter(q.difficulty for q in questions_to_print)
    logger.info("  By Difficulty:")
    for diff, count in sorted(
        by_difficulty.items(), key=lambda x: x[0].value if hasattr(x[0], "value") else str(x[0])
    ):
        logger.info("    %s: %d", str(diff), count)

    # Per-difficulty averages (entropy and tool calls)
    buckets: dict = defaultdict(list)
    for q in questions_to_print:
        buckets[q.difficulty].append(q)
    logger.info("  Per-Difficulty Averages:")
    for diff, qs in sorted(buckets.items(), key=lambda x: x[0].value if hasattr(x[0], "value") else str(x[0])):
        avg_entropy = sum(q.entropy_used for q in qs) / len(qs)
        avg_tool_calls = sum(q.tool_calls_required for q in qs) / len(qs)
        logger.info("    %s -> entropy avg: %.2f, tool calls avg: %.2f", str(diff), avg_entropy, avg_tool_calls)
    logger.info("=" * 30)


def print_split_statistics(splits: DatasetDict) -> None:  # pragma: no cover
    """Display statistics about train/validation/test splits."""
    logger.info("=" * 30)
    logger.info("DATASET SPLITS")
    logger.info("=" * 30)

    total = sum(len(splits[split_name]) for split_name in splits)

    for split_name in ["train", "validation", "test"]:
        if split_name in splits:
            count = len(splits[split_name])
            percentage = 100 * count / total if total > 0 else 0
            logger.info("  %s: %d samples (%.1f%%)", split_name.capitalize(), count, percentage)

    logger.info("  Total: %d samples", total)
    logger.info("=" * 30)


def _question_to_example(q: Question) -> dict[str, Any]:
    """Map a Question to a flat example for Hugging Face datasets."""
    stepwise_truths: list[dict[str, Any]] = []
    for step in q.stepwise:
        tool_name = step.get("tool")
        result_value = parse_string(step.get("result"))
        if tool_name is None or result_value is None:
            continue
        stepwise_truths.append({tool_name: result_value})

    # Derive composition metadata from stepwise verification and problem_type
    composition_type = "sequential" if len(q.stepwise) > 1 else "single"
    dependency_edges: list[tuple[int, int]] = []
    for idx, step in enumerate(q.stepwise):
        verification = step.get("verification", {})
        if isinstance(verification, dict):
            dependent_on = verification.get("dependent_on")
            if isinstance(dependent_on, dict):
                for _, from_idx in dependent_on.items():
                    if isinstance(from_idx, int):
                        dependency_edges.append((from_idx, idx))

    dependency_type = "strict"

    return {
        "query": q.question,
        "ground_truth": q.golden.get("final_answer", q.answer),
        "stepwise_ground_truths": json.dumps(stepwise_truths),
        "difficulty": getattr(q.difficulty, "name", str(q.difficulty)),
        "problem_type": getattr(q.problem_type, "value", str(q.problem_type)),
        "composition_type": composition_type,
        "composition_dependencies": dependency_type,
        "dependency_edges": json.dumps(dependency_edges) if dependency_edges else None,
    }


def convert_to_dataset_dict(questions: list[Question]) -> DatasetDict:
    """Convert questions to a single-split DatasetDict (train)."""
    examples = [_question_to_example(q) for q in questions if q.is_valid]
    return DatasetDict({"train": Dataset.from_list(examples)})


def convert_to_dataset_splits(
    questions: list[Question],
    test_size: float = 0.1,
    val_size: float = 0.1,
    seed: int = 42,
    stratify_by: str | None = None,
) -> DatasetDict:
    """Create train/validation/test DatasetDict using HF's split utilities."""

    examples = [_question_to_example(q) for q in questions if q.is_valid]
    ds = Dataset.from_list(examples).shuffle(seed=seed)

    stratify_column = stratify_by if stratify_by in ds.column_names else None

    # Convert stratification column to ClassLabel if needed
    if stratify_column and ds.features[stratify_column]._type != "ClassLabel":
        from datasets import ClassLabel

        unique_values = ds.unique(stratify_column)
        ds = ds.cast_column(stratify_column, ClassLabel(names=sorted(unique_values)))

    split = ds.train_test_split(test_size=test_size, seed=seed, stratify_by_column=stratify_column)

    if val_size and val_size > 0:
        # Adjust val proportion relative to remaining train portion
        relative_val = val_size / (1 - test_size)
        train_val = split["train"].train_test_split(
            test_size=relative_val, seed=seed, stratify_by_column=stratify_column
        )
        return DatasetDict(train=train_val["train"], validation=train_val["test"], test=split["test"])  # type: ignore[reportCallIssue]

    return DatasetDict(train=split["train"], test=split["test"])  # type: ignore[reportCallIssue]