import random

import litellm
from datasets import load_dataset

import mlflow
from mlflow.genai.optimize.optimizers import MetaPromptOptimizer
from mlflow.genai.scorers import Correctness


def load_aime_data(num_train=10, num_test=32, seed=42):
    """Load AIME dataset from HuggingFace and split into train/test."""
    dataset = load_dataset("gneubig/aime-1983-2024", split="train")

    # Shuffle with fixed seed for reproducibility
    indices = list(range(len(dataset)))
    random.seed(seed)
    random.shuffle(indices)

    # Split into train and test
    train_indices = indices[:num_train]
    test_indices = indices[num_train : num_train + num_test]

    train_dataset = dataset.select(train_indices)
    test_dataset = dataset.select(test_indices)

    # Convert to MLflow format
    def convert_to_mlflow_format(dataset):
        mlflow_data = []
        for example in dataset:
            mlflow_data.append(
                {
                    "inputs": {"question": example["Question"]},
                    "expectations": {"expected_response": str(example["Answer"])},
                }
            )
        return mlflow_data

    return convert_to_mlflow_format(train_dataset), convert_to_mlflow_format(
        test_dataset
    )


def run_benchmark(
    reflection_model="gpt-4o",
    eval_model="gpt-4o-mini",
    num_train_examples=10,
    num_test_examples=32,
):
    """
    Benchmark MetaPromptOptimizer on AIME dataset.

    Args:
        reflection_model: Model to use for prompt optimization
        eval_model: Model to use for predictions
        num_train_examples: Number of training examples
        num_test_examples: Number of test examples for final evaluation
    """

    print("=" * 80)
    print("MetaPromptOptimizer Benchmark on AIME")
    print("=" * 80)

    # Load data
    print(
        f"\nLoading {num_train_examples} train examples and {num_test_examples} test examples..."
    )
    train_data, test_data = load_aime_data(
        num_train=num_train_examples, num_test=num_test_examples
    )

    print(f"Train size: {len(train_data)}")
    print(f"Test size: {len(test_data)}")

    try:
        prompt = mlflow.genai.load_prompt("prompts:/aime_solver/6")
    except Exception:
        prompt = mlflow.genai.register_prompt(
            name="aime_solver",
            template="Solve the following math problem. Provide only the numerical answer:\n\n{{question}}",
        )

    print(f"\nInitial prompt: {prompt.template}")

    # Define predict function
    def predict_fn(question: str) -> str:
        """Prediction function that uses the registered prompt."""
        prompt = mlflow.genai.load_prompt("prompts:/aime_solver/6")
        formatted_prompt = prompt.format(question=question)
        response = litellm.completion(
            model=f"openai/{eval_model}",
            messages=[{"role": "user", "content": formatted_prompt}],
            temperature=1.0,
        )
        return response.choices[0].message.content

    # Optimize prompts
    print("\n" + "-" * 80)
    print("Optimizing Prompts")
    print("-" * 80)

    result = mlflow.genai.optimize_prompts(
        predict_fn=predict_fn,
        train_data=train_data,
        prompt_uris=[prompt.uri],
        optimizer=MetaPromptOptimizer(
            reflection_model="openai:/gpt-5-mini",
            lm_kwargs={"temperature": 1.0, "max_tokens": 4096},
        ),
        scorers=[Correctness(model="openai:/gpt-5-mini")],
    )

    # result = mlflow.genai.optimize_prompts(
    #     predict_fn=predict_fn,
    #     train_data=[],
    #     prompt_uris=[prompt.uri],
    #     optimizer=MetaPromptOptimizer(
    #         reflection_model="openai:/gpt-5-mini",
    #         lm_kwargs={"temperature": 1.0, "max_tokens": 4096},
    #     ),
    #     scorers=[],
    # )

    print(f"\n{'=' * 80}")
    print("Optimization Results")
    print("=" * 80)
    print(f"Initial prompt:\n{prompt.template}\n")
    print(f"Optimized prompt:\n{result.optimized_prompts[0].template}\n")
    print(
        f"Training score improvement: {result.initial_eval_score:.4f} → {result.final_eval_score:.4f}"
    )

    # Check if prompt changed
    if result.optimized_prompts[0].template == prompt.template:
        print("\n⚠️  Note: The optimized prompt is identical to the initial prompt.")
        print("No improvement was found during optimization.")


def run_evaluation(prompt_version: str, data, eval_model: str):
    """
    Evaluate the performance of a prompt on a dataset.
    """

    def predict_fn(question: str) -> str:
        prompt = mlflow.genai.load_prompt(f"prompts:/aime_solver/{prompt_version}")
        formatted_prompt = prompt.format(question=question)
        response = litellm.completion(
            model=f"openai/{eval_model}",
            messages=[{"role": "user", "content": formatted_prompt}],
            temperature=1.0,
        )
        return response.choices[0].message.content

    return mlflow.genai.evaluate(
        data=data,
        predict_fn=predict_fn,
        scorers=[Correctness(model=eval_model)],
    )


def main():
    # # Configure MLflow tracking
    mlflow.set_tracking_uri("http://localhost:5000")
    mlflow.set_experiment("metaprompt-optimizer-aime")

    print("Starting benchmark...")

    run_benchmark(
        reflection_model="gpt-5-mini",
        eval_model="gpt-5-nano",
        num_train_examples=12,
        num_test_examples=24,
    )


if __name__ == "__main__":
    main()

================================================================================
Optimization Results
================================================================================
Initial prompt:
Solve the following math problem. Provide only the numerical answer:

{{question}}

Optimized prompt:
You are an expert contest mathematician. Solve the following math problem and output only the final numerical answer. Do not show any solution steps, reasoning, or explanations — perform all reasoning internally. Before answering, verify your result by recomputing or checking algebra/arithmetic.

Rules for the output (must be followed exactly):
- If the answer is an integer, output only the digits of the integer (no commas, no words, no punctuation).
- If the answer is a simplified fraction, output it as numerator/denominator in lowest terms (e.g., 3/7).
- If the problem asks for a specific expression (such as m+n), output that final integer only.
- Do not output any extra text, labels, or whitespace — only the exact answer characters.