"""
Original from: https://github.com/jonhanke/nam_causal-head-gating
Adapted and extended for Causal Summer project by:
- Bowen
- 2026-01-27

Few-shot prompt formatting utilities.

This module replicates the original `format_data.py` functionality for creating
few-shot prompts with proper train/validation/example splits.
"""

import itertools
import random
from collections.abc import Callable
from typing import Any

import pandas as pd
from torch import Tensor

from .tokenisation import PromptTokenizer


def generate_prompt_permutations(
    prompt: str,
    options: dict[str, int],
    max_permutations: int | None = None,
) -> list[dict[str, str]]:
    """
    Generate all permutations of multiple choice options.

    For MCQ tasks, this creates different orderings of the answer choices
    to prevent the model from learning positional biases.

    Args:
        prompt: The question text.
        options: Dict mapping option text to correctness (1 for correct, 0 for wrong).
        max_permutations: Maximum number of permutations to generate.

    Returns:
        List of dicts with 'question', 'options', and 'target' keys.
    """
    items = list(options.items())
    all_perms = list(itertools.permutations(items))
    perms = (
        random.sample(all_perms, max_permutations)
        if max_permutations is not None and max_permutations < len(all_perms)
        else all_perms
    )
    results = []
    for perm in perms:
        letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        labeled = [f"{letters[i]}. {k}" for i, (k, _) in enumerate(perm)]
        correct_index = next(i for i, (_, v) in enumerate(perm) if v == 1)
        results.append({
            "question": prompt.strip(),
            "options": "\n".join(labeled),
            "target": letters[correct_index],
        })
    return results


def assign_split(
    df_questions: pd.DataFrame,
    num_examples: int,
    f_validation: float = 0.1,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Assign train/validation/example splits to questions.

    The shortest questions (by total length) are assigned to the 'example' split
    to be used as few-shot examples. Remaining questions are split between
    'train' and 'validation'.

    Args:
        df_questions: DataFrame with 'question' and 'target' columns.
        num_examples: Number of questions to reserve as few-shot examples.
        f_validation: Fraction of remaining questions for validation.
        seed: Random seed for reproducibility.

    Returns:
        DataFrame with added 'split' column ('example', 'train', or 'validation').
    """
    df = df_questions.copy()
    rng = random.Random(seed)

    # Calculate total length for each question
    if "options" in df.columns:
        grouped = df.groupby("question").agg(
            {"options": lambda x: sum(len(opt) for opt in x)}
        )
        grouped["total_len"] = grouped.index.str.len() + grouped["options"]
    elif "target" in df.columns:
        grouped = df.groupby("question").agg({"target": lambda x: len(x.iloc[0])})
        grouped["total_len"] = grouped.index.str.len() + grouped["target"]
    else:
        raise ValueError("DataFrame must contain either 'options' or 'target' column.")

    # Shortest questions become examples
    shortest = grouped.sort_values("total_len").head(num_examples).index
    df["split"] = "train"
    df.loc[df["question"].isin(shortest), "split"] = "example"

    # Split remaining into train/validation
    remaining = df.loc[~df["question"].isin(shortest), "question"].unique()
    num_val = int(len(remaining) * f_validation)
    val_qs = set(rng.sample(list(remaining), num_val))
    df.loc[df["question"].isin(val_qs), "split"] = "validation"

    return df


def create_prompts(
    df_questions: pd.DataFrame,
    num_examples: int,
    format_example_fn: Callable[[dict[str, Any]], str],
    instructions: str | None = None,
    seed: int | None = None,
    question_prefix: str = "Question: ",
    answer_prefix: str = "Answer: ",
    sep: str = "\n",
) -> pd.DataFrame:
    """
    Create few-shot prompts from a DataFrame of questions.

    Each prompt will contain:
    1. Optional instructions
    2. N randomly sampled examples (from the 'example' split)
    3. The actual question to answer

    Args:
        df_questions: DataFrame with 'question', 'target', and 'split' columns.
        num_examples: Number of few-shot examples to include in each prompt.
        format_example_fn: Function to format each example as a string.
        instructions: Optional instruction text to prepend.
        seed: Random seed for reproducibility.
        question_prefix: Prefix for questions (e.g., "Question: ").
        answer_prefix: Prefix for answers (e.g., "Answer: ").
        sep: Separator between examples.

    Returns:
        DataFrame with 'split', 'prompt', and 'target' columns.
    """
    rng = random.Random(seed)
    df = df_questions.copy()

    # Group examples by question
    example_groups = df[df["split"] == "example"].groupby("question")
    example_dict = {q: g.to_dict("records") for q, g in example_groups}

    def build_prompt(row: pd.Series) -> tuple[str, str]:
        # Sample examples
        sampled = rng.sample(list(example_dict), min(num_examples, len(example_dict)))
        examples = [rng.choice(example_dict[q]) for q in sampled]
        example_strs = [format_example_fn(ex) for ex in examples]

        # Build prompt
        text = (instructions + "\n") if instructions else ""
        text += sep.join(example_strs)

        # Add the actual question
        query_lines = [f"{question_prefix}{row['question']}"]
        if "options" in row and pd.notna(row.get("options")):
            query_lines.append(row["options"])
        query_lines.append(f"{answer_prefix}")
        query_block = "\n".join(query_lines)

        return sep.join([text, query_block]), row["target"]

    records = []
    for split in ["train", "validation"]:
        for _, row in df[df["split"] == split].iterrows():
            prompt, target = build_prompt(row)
            records.append({"split": split, "prompt": prompt, "target": target})

    return pd.DataFrame.from_records(records)


def create_fewshot_dataset(
    tokenizer: Any,
    questions: list[str],
    targets: list[str] | list[dict[str, int]],
    instructions: str | None = None,
    max_permutations: int | None = None,
    example_set_size: int = 50,
    num_examples: int = 1,
    f_validation: float = 0.1,
    seed: int = 0,
    question_prefix: str = "Question: ",
    answer_prefix: str = "Answer: ",
    sep: str = "\n",
    verbose: bool = True,
    device: str = "cpu",
) -> tuple[pd.DataFrame, Tensor, Tensor]:
    """
    Create a few-shot learning dataset with train/validation splits.

    This replicates the original `create_dataset()` function from `format_data.py`.
    It:
    1. Assigns the shortest questions to be few-shot examples
    2. Creates prompts with randomly sampled examples for each train/val question
    3. Tokenizes everything with proper loss masks

    Supports both free-response questions (FRQ) and multiple choice questions (MCQ).

    Args:
        tokenizer: HuggingFace tokenizer.
        questions: List of question strings.
        targets: Either:
            - List of target/answer strings (FRQ mode)
            - List of dicts mapping options to correctness {option: 0 or 1} (MCQ mode)
        instructions: Optional instruction text to prepend to all prompts.
        max_permutations: For MCQ, maximum answer orderings per question.
        example_set_size: Number of questions to reserve as few-shot examples.
        num_examples: Number of examples to include in each prompt.
        f_validation: Fraction of non-example questions for validation.
        seed: Random seed for reproducibility.
        question_prefix: Prefix for questions.
        answer_prefix: Prefix for answers.
        sep: Separator between examples.
        verbose: Whether to print progress.
        device: Device to place tensors on.

    Returns:
        Tuple of (df_prompts, input_ids, loss_masks):
        - df_prompts: DataFrame with 'split', 'prompt', 'target' columns
        - input_ids: Tensor of shape [num_samples, seq_len]
        - loss_masks: Boolean tensor of shape [num_samples, seq_len]

    Example (FRQ - math):
        >>> df_prompts, input_ids, loss_masks = create_fewshot_dataset(
        ...     tokenizer,
        ...     questions=problems,
        ...     targets=solutions,
        ...     instructions="Solve step by step.",
        ...     example_set_size=50,
        ...     num_examples=1,
        ... )

    Example (MCQ):
        >>> targets = [{"Paris": 1, "London": 0, "Berlin": 0}]
        >>> df_prompts, input_ids, loss_masks = create_fewshot_dataset(
        ...     tokenizer,
        ...     questions=["What is the capital of France?"],
        ...     targets=targets,
        ...     max_permutations=6,
        ... )
    """
    prompt_tokenizer = PromptTokenizer(tokenizer, device=device)

    # Handle MCQ (dict targets) vs FRQ (string targets)
    is_mcq = len(targets) > 0 and isinstance(targets[0], dict)

    if is_mcq:
        # Generate permutations for multiple choice questions
        all_questions = []
        for q, opts in zip(questions, targets, strict=False):
            all_questions.extend(generate_prompt_permutations(q, opts, max_permutations))
        if verbose:
            print(f"Generated {len(all_questions)} MCQ permutations.")
        df_questions = pd.DataFrame(all_questions)
    else:
        # Free-response questions
        df_questions = pd.DataFrame({"question": questions, "target": targets})

    df_questions = df_questions.drop_duplicates().reset_index(drop=True)

    # Assign splits
    df_questions = assign_split(
        df_questions,
        num_examples=example_set_size,
        f_validation=f_validation,
        seed=seed,
    )

    if verbose:
        print(
            f"Assigned splits: {df_questions['split'].value_counts().to_dict()}. "
            "Generating few-shot prompts."
        )

    # Create few-shot prompts with appropriate format function
    if is_mcq:
        # MCQ format includes options
        def format_fn(ex: dict[str, Any]) -> str:
            return (
                f"{question_prefix}{ex['question']}\n"
                f"{ex['options']}\n"
                f"{answer_prefix}{ex['target']}"
            )
    else:
        # FRQ format
        def format_fn(ex: dict[str, Any]) -> str:
            return (
                f"{question_prefix}{ex['question']}\n"
                f"{answer_prefix}{ex['target']}"
            )

    df_prompts = create_prompts(
        df_questions,
        num_examples=num_examples,
        format_example_fn=format_fn,
        instructions=instructions,
        seed=seed,
        question_prefix=question_prefix,
        answer_prefix=answer_prefix,
        sep=sep,
    )

    if verbose:
        print(f"Generated {len(df_prompts)} unique prompts. Tokenizing.")

    # Tokenize
    input_ids, loss_masks = prompt_tokenizer.tokenize_batch(
        df_prompts["prompt"].tolist(),
        df_prompts["target"].tolist(),
    )

    if verbose:
        # Find actual lengths (before padding)
        lengths = (input_ids != tokenizer.pad_token_id).sum(dim=1)
        print(
            f"Tokenized {len(input_ids)} prompts with "
            f"min length {lengths.min().item()} and max length {lengths.max().item()}."
        )

    return df_prompts, input_ids, loss_masks


def load_math_dataset(
    tokenizer: Any,
    dataset_name: str = "nvidia/OpenMathInstruct-2",
    split: str = "train_1M",
    num_examples: int = 50,
    num_train: int = 50000,
    num_validation: int = 5000,
    instructions: str | None = None,
    num_fewshot_examples: int = 1,
    seed: int = 0,
    lengths_cache_path: str | None = None,
    verbose: bool = True,
    device: str = "cpu",
) -> tuple[pd.DataFrame, Tensor, Tensor]:
    """
    Load and prepare the OpenMathInstruct-2 dataset for CHG training.

    This function automates the full original math notebook workflow:
    1. Load dataset from HuggingFace
    2. Compute sequence lengths (or load from cache)
    3. Filter to problems with non-empty answers
    4. Sort by length and select shortest problems
    5. Create few-shot prompts with `create_fewshot_dataset()`

    Args:
        tokenizer: HuggingFace tokenizer.
        dataset_name: HuggingFace dataset name.
        split: Dataset split to use.
        num_examples: Number of shortest problems to use as few-shot examples.
        num_train: Number of training problems (after examples).
        num_validation: Number of validation problems.
        instructions: Instruction text for prompts. If None, uses default math instructions.
        num_fewshot_examples: Number of few-shot examples per prompt.
        seed: Random seed for reproducibility.
        lengths_cache_path: Path to save/load precomputed lengths (parquet file).
            If file exists, loads from cache. If None, computes but doesn't cache.
        verbose: Whether to print progress.
        device: Device to place tensors on.

    Returns:
        Tuple of (df_prompts, input_ids, loss_masks) ready for CHG training.

    Length Selection:
        The function selects the **shortest** problems from the dataset. The total
        number of problems selected is `num_examples + num_train + num_validation`.

        With the default parameters (50 + 50,000 + 5,000 = 55,050), this matches the
        original paper's approach of selecting the 55,050 shortest problems, which
        typically results in sequences of ~170-340 tokens.

    Example:
        >>> # Full workflow in one call
        >>> df_prompts, input_ids, loss_masks = load_math_dataset(
        ...     tokenizer,
        ...     num_train=50000,
        ...     lengths_cache_path="math_lengths.parquet",
        ... )

        >>> # Create dataset for training
        >>> train_mask = df_prompts['split'] == 'train'
        >>> train_dataset = MaskedSequenceDataset(
        ...     pad_token_id=tokenizer.pad_token_id,
        ...     input_ids=input_ids[train_mask.values],
        ...     loss_masks=loss_masks[train_mask.values],
        ... )

    Note:
        Computing lengths for 1M samples takes significant time (~30-60 minutes).
        Use `lengths_cache_path` to save and reuse computed lengths.
    """
    from pathlib import Path

    from datasets import load_dataset
    from tqdm.auto import tqdm

    # Default instructions match original notebook
    if instructions is None:
        instructions = (
            "For each problem, explain your reasoning step by step and use LaTeX "
            "for all mathematical expressions. Indicate your final answer using \\boxed{...}."
        )

    # Define the number of shortest problems (by token length) to select from the dataset.
    # Defaults to 50 + 50,000 + 5,000 = 55,050, matching the number used in the paper.
    total_needed = num_examples + num_train + num_validation

    # Check for cached lengths
    cache_path = Path(lengths_cache_path) if lengths_cache_path else None
    if cache_path and cache_path.exists():
        if verbose:
            print(f"Loading precomputed lengths from {cache_path}")
        df_problems = pd.read_parquet(cache_path)
    else:
        # Load dataset from HuggingFace
        if verbose:
            print(f"Loading dataset {dataset_name} ({split})...")
        hf_dataset = load_dataset(dataset_name, split=split)

        # Compute lengths for all samples
        if verbose:
            print(f"Computing sequence lengths for {len(hf_dataset)} samples...")
            print("(This may take a while. Use lengths_cache_path to save results.)")

        prompt_tokenizer = PromptTokenizer(tokenizer, device=device)

        rows = []
        iterator = tqdm(hf_dataset) if verbose else hf_dataset
        for item in iterator:
            problem = item["problem"]
            solution = item["generated_solution"]
            answer = item.get("expected_answer", "")

            # Compute tokenized length
            text_tokens, _ = prompt_tokenizer.tokenize(problem, solution)
            rows.append({
                "problem": problem,
                "solution": solution,
                "answer": answer,
                "length": len(text_tokens),
            })

        df_problems = pd.DataFrame(rows)

        # Save to cache if path provided
        if cache_path:
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            df_problems.to_parquet(cache_path)
            if verbose:
                print(f"Saved lengths to {cache_path}")

    # Filter and select shortest problems (matching original notebook)
    if verbose:
        print("Filtering and selecting shortest problems...")

    # Filter to problems with non-empty answers
    df_problems = df_problems[df_problems["answer"].str.len() > 0]

    # Drop duplicate problems
    df_problems = df_problems.drop_duplicates("problem")

    # Sort by length and take shortest
    df_problems = df_problems.sort_values("length").head(total_needed).reset_index(drop=True)

    if verbose:
        print(f"Selected {len(df_problems)} shortest problems")
        print(f"  Length range: {df_problems['length'].min()} - {df_problems['length'].max()}")

    # Create few-shot dataset
    df_prompts, input_ids, loss_masks = create_fewshot_dataset(
        tokenizer=tokenizer,
        questions=df_problems["problem"].tolist(),
        targets=df_problems["solution"].tolist(),
        instructions=instructions,
        example_set_size=num_examples,
        num_examples=num_fewshot_examples,
        f_validation=num_validation / (num_train + num_validation),
        seed=seed,
        verbose=verbose,
        device=device,
    )

    return df_prompts, input_ids, loss_masks
