import gradio as gr
import os
import yaml
import json
import random
import re
from datasets import load_dataset, get_dataset_config_names, get_dataset_split_names
from openai import OpenAI
from openevolve import run_evolution
from typing import Dict, List, Tuple, Optional
import tempfile
import shutil
import requests
import glob

# Free models from OpenRouter - Curated selection (verified as of 2025)
# IMPORTANT: The :free suffix is REQUIRED to use the free tier. Without it, requests are charged!
FREE_MODELS = [
    "qwen/qwen-2.5-72b-instruct:free",  # 72B - Strong in coding/math/multilingual (default - better rate limits)
    "meta-llama/llama-3.3-70b-instruct:free",  # 70B - Advanced reasoning
    "google/gemma-3-27b-it:free",  # 27B - Strong instruction-tuned
    "mistralai/mistral-small-3.1-24b-instruct:free",  # 24B - Efficient and capable
    "deepseek/deepseek-r1:free",  # 671B (37B active) - Top-tier but heavily rate-limited
    "meta-llama/llama-3.2-3b-instruct",  # 3B - PAID but very cheap fallback when free models hit rate limits
]


def validate_dataset(dataset_name: str, split: str, input_field: str, target_field: str) -> Tuple[bool, str]:
    """
    Validate that the dataset exists and has the required fields.

    Returns:
        Tuple of (is_valid, error_message)
    """
    try:
        # Check if dataset name has correct format (should be org/name or just name)
        if not dataset_name or dataset_name.strip() == "":
            return False, "❌ Dataset name cannot be empty"

        dataset_name = dataset_name.strip()

        # Try to get dataset info from HuggingFace API
        hf_token = os.environ.get("HF_TOKEN", None)
        headers = {}
        if hf_token:
            headers["Authorization"] = f"Bearer {hf_token}"

        # Check if dataset exists on HuggingFace Hub
        api_url = f"https://huggingface.co/api/datasets/{dataset_name}"
        response = requests.get(api_url, headers=headers, timeout=10)

        if response.status_code == 404:
            return False, f"❌ Dataset '{dataset_name}' not found on HuggingFace Hub. Please use the full dataset name (e.g., 'stanfordnlp/imdb' or 'gsm8k')"
        elif response.status_code != 200:
            # Try to load anyway - might be a private dataset or API issue
            print(f"Warning: Could not verify dataset via API (status {response.status_code}), attempting to load...")

        # Try to load a small sample to verify it works and check fields
        print(f"Loading dataset {dataset_name} with split {split}...")

        # First, check if the split exists
        try:
            available_splits = get_dataset_split_names(dataset_name)
            if split not in available_splits:
                return False, f"❌ Split '{split}' not found. Available splits: {', '.join(available_splits)}"
        except Exception as e:
            print(f"Could not get split names: {e}. Will try to load anyway...")

        # Load a small sample to check fields
        # Try loading with just dataset name first
        try:
            dataset = load_dataset(dataset_name, split=split, streaming=True)
        except ValueError as e:
            # If it fails with config error, try common configs
            if "config" in str(e).lower() or "Config name is missing" in str(e):
                # Try common configs based on dataset name
                default_config = "main"
                if dataset_name.lower() == "glue":
                    default_config = "sst2"

                print(f"Dataset requires config, trying with '{default_config}' config...")
                try:
                    dataset = load_dataset(dataset_name, default_config, split=split, streaming=True)
                except:
                    # If default config doesn't work, raise the original error
                    raise e
            else:
                raise

        # Get first example to check fields
        first_example = next(iter(dataset))
        available_fields = list(first_example.keys())

        # Check if input field exists
        if input_field not in available_fields:
            return False, f"❌ Input field '{input_field}' not found. Available fields: {', '.join(available_fields)}"

        # Check if target field exists
        if target_field not in available_fields:
            return False, f"❌ Target field '{target_field}' not found. Available fields: {', '.join(available_fields)}"

        # All validations passed
        return True, f"✅ Dataset validated successfully! Fields '{input_field}' and '{target_field}' found."

    except Exception as e:
        error_msg = str(e)
        if "404" in error_msg or "not found" in error_msg.lower():
            return False, f"❌ Dataset '{dataset_name}' not found. Please check the dataset name (use format: org/dataset-name)"
        return False, f"❌ Error validating dataset: {error_msg}"


def validate_inputs(dataset_name: str, split: str, input_field: str, target_field: str,
                   initial_prompt: str) -> Tuple[bool, str]:
    """
    Validate all inputs before starting optimization.

    Returns:
        Tuple of (is_valid, message)
    """
    # Check API key
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        return False, "❌ OPENAI_API_KEY environment variable not set. Please set it in the Space secrets."

    # Check prompt contains {input} placeholder
    if "{input}" not in initial_prompt:
        return False, "❌ Prompt must contain '{input}' placeholder for dataset inputs"

    # Check dataset name format
    dataset_name = dataset_name.strip()
    if not dataset_name:
        return False, "❌ Dataset name cannot be empty"

    # Validate dataset and fields
    is_valid, message = validate_dataset(dataset_name, split, input_field, target_field)
    if not is_valid:
        return False, message

    return True, message


def evaluate_prompt(prompt: str, dataset_name: str, split: str, num_samples: int,
                    model: str, input_field: str, target_field: str,
                    fixed_indices: List[int] = None) -> Dict:
    """
    Evaluate a prompt on a dataset using the selected model.

    Args:
        fixed_indices: Optional list of dataset indices to use. If provided,
                      ensures we evaluate on the SAME samples every time.
    """
    try:
        # Get API key from environment
        api_key = os.environ.get("OPENAI_API_KEY")
        if not api_key:
            return {
                "error": "OPENAI_API_KEY not set in environment",
                "accuracy": 0,
                "correct": 0,
                "total": 0,
                "results": []
            }

        # Load dataset
        # Try loading with just dataset name first
        try:
            dataset = load_dataset(dataset_name, split=split, streaming=False)
        except ValueError as e:
            # If it fails with config error, try common configs
            if "config" in str(e).lower() or "Config name is missing" in str(e):
                # Try common configs based on dataset name
                default_config = "main"
                if dataset_name.lower() == "glue":
                    default_config = "sst2"
                dataset = load_dataset(dataset_name, default_config, split=split, streaming=False)
            else:
                raise

        # Sample examples - use fixed indices if provided to ensure consistency
        if fixed_indices is not None:
            # Use the provided indices (ensures same samples for initial/final eval)
            indices = fixed_indices
            samples = [dataset[i] for i in indices]
        elif len(dataset) > num_samples:
            # First time: use fixed seed for reproducible sampling
            random.seed(42)  # Fixed seed ensures same samples across runs
            indices = random.sample(range(len(dataset)), num_samples)
            samples = [dataset[i] for i in indices]
        else:
            indices = list(range(min(num_samples, len(dataset))))
            samples = list(dataset)[:num_samples]

        # Initialize OpenAI client with OpenRouter
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )

        correct = 0
        total = 0
        results = []
        errors = []

        for idx, sample in enumerate(samples):
            try:
                # Get input and target
                input_text = sample.get(input_field, "")
                if isinstance(input_text, dict):
                    input_text = str(input_text)

                target = sample.get(target_field, "")
                if isinstance(target, dict):
                    target = str(target)

                # Format the prompt with the input
                formatted_prompt = prompt.replace("{input}", str(input_text))

                # Call the model
                response = client.chat.completions.create(
                    model=model,
                    messages=[
                        {"role": "system", "content": "You are a helpful assistant."},
                        {"role": "user", "content": formatted_prompt}
                    ],
                    temperature=0.0,
                    max_tokens=500,
                )

                prediction = response.choices[0].message.content.strip()

                # IMDB labels: 0 = negative, 1 = positive
                true_label = int(target)  # 0 or 1

                # FORMAT REQUIREMENT: Need "sentiment" keyword + positive/negative in first 150 chars
                # This is strict enough to fail conversational responses, but learnable through evolution
                pred_lower = prediction.lower()
                pred_start = pred_lower[:150]  # First 150 chars

                # Must mention "sentiment" to get credit (helps evolution learn to add this keyword)
                has_sentiment_keyword = "sentiment" in pred_start

                # Check for positive/negative indicators
                has_positive = "positive" in pred_start
                has_negative = "negative" in pred_start

                # Only count as correct if sentiment keyword present AND unambiguous positive/negative
                if has_sentiment_keyword and has_positive and not has_negative:
                    predicted_label = 1
                elif has_sentiment_keyword and has_negative and not has_positive:
                    predicted_label = 0
                else:
                    predicted_label = -1

                is_correct = (predicted_label == true_label)

                if is_correct:
                    correct += 1
                total += 1

                results.append({
                    "input": str(input_text)[:100] + "..." if len(str(input_text)) > 100 else str(input_text),
                    "target": str(target),
                    "prediction": prediction[:100] + "..." if len(prediction) > 100 else prediction,
                    "correct": is_correct
                })

            except Exception as e:
                error_msg = f"Sample {idx+1}: {str(e)}"
                print(f"Error evaluating sample {idx+1}: {e}")
                errors.append(error_msg)
                # Only continue if we haven't failed on all samples
                if len(errors) > len(samples) // 2:  # More than half failed
                    print(f"Too many errors ({len(errors)} out of {len(samples)}), stopping evaluation")
                    break
                continue

        accuracy = (correct / total * 100) if total > 0 else 0

        result_dict = {
            "accuracy": accuracy,
            "correct": correct,
            "total": total,
            "results": results,
            "indices": indices  # Return indices so we can reuse them for final eval
        }

        # Add errors if any occurred
        if errors:
            result_dict["errors"] = errors
            if total == 0:
                # All samples failed - create a helpful error message
                result_dict["error"] = f"All {len(samples)} samples failed to evaluate. First few errors:\n" + "\n".join(errors[:3])

        return result_dict

    except Exception as e:
        return {
            "error": str(e),
            "accuracy": 0,
            "correct": 0,
            "total": 0,
            "results": []
        }


def collect_prompt_history(output_dir: str, initial_score: float = 0.0) -> List[Dict]:
    """
    Collect only the prompts that were "best" at some point during evolution.
    Returns only programs that improved upon the initial score (deduplicated).

    Args:
        output_dir: Directory containing checkpoint data
        initial_score: Score of the initial prompt (baseline to beat)

    Returns a list of dicts with: {prompt, score, iteration, id}
    """
    try:
        all_programs = []
        seen_prompts = set()  # Track unique prompts

        # OpenEvolve saves programs in checkpoint directories as JSON files
        # Structure: output_dir/checkpoints/checkpoint_{iteration}/programs/{program_id}.json
        checkpoints_dir = os.path.join(output_dir, "checkpoints")

        if not os.path.exists(checkpoints_dir):
            return []

        # Find all checkpoint directories
        checkpoint_dirs = sorted(glob.glob(os.path.join(checkpoints_dir, "checkpoint_*")))

        # Collect all programs from all checkpoints
        for checkpoint_dir in checkpoint_dirs:
            programs_dir = os.path.join(checkpoint_dir, "programs")
            if not os.path.exists(programs_dir):
                continue

            # Read all program JSON files
            program_files = glob.glob(os.path.join(programs_dir, "*.json"))

            for pfile in program_files:
                try:
                    with open(pfile, 'r') as f:
                        program_data = json.load(f)

                    # Extract the code (prompt) from the program data
                    prompt_content = program_data.get("code", "").strip()
                    prog_id = program_data.get("id", os.path.basename(pfile).replace(".json", ""))
                    iteration = program_data.get("iteration_found", 0)
                    metrics = program_data.get("metrics", {})

                    # Get combined score for comparison
                    combined_score = metrics.get("combined_score", 0.0)

                    all_programs.append({
                        "prompt": prompt_content,
                        "id": prog_id,
                        "file": pfile,
                        "iteration": iteration,
                        "metrics": metrics,
                        "score": combined_score
                    })
                except Exception as e:
                    print(f"Error reading program file {pfile}: {e}")
                    continue

        # Sort all programs by iteration (chronological order)
        all_programs.sort(key=lambda x: x.get("iteration", 0))

        # Filter to keep only programs that improved the best score
        # Start from the initial score as the baseline
        best_programs = []
        current_best_score = initial_score

        for program in all_programs:
            prompt_content = program["prompt"]
            score = program["score"]
            iteration = program["iteration"]

            # Skip iteration 0 (that's the initial prompt, already added separately)
            if iteration == 0:
                continue

            # Create a normalized version for duplicate detection (ignore whitespace differences)
            normalized_prompt = " ".join(prompt_content.split())

            # Skip duplicates
            if normalized_prompt in seen_prompts:
                continue

            # Only keep if this program improved the best score
            if score > current_best_score:
                seen_prompts.add(normalized_prompt)
                best_programs.append(program)
                improvement = score - current_best_score
                print(f"  ✓ Best program at iteration {iteration}: score={score:.2%} (improved by +{improvement:.2%})")
                current_best_score = score

        return best_programs

    except Exception as e:
        print(f"Error collecting prompt history: {e}")
        return []


def parse_evolution_history(output_dir: str) -> str:
    """
    Parse evolution history from OpenEvolve output directory.

    Returns a markdown string with visualization of the evolution process.
    """
    try:
        evolution_viz = "## 🧬 Evolution Progress\n\n"

        # Look for generation files or logs
        generation_files = sorted(glob.glob(os.path.join(output_dir, "generation_*.txt")))
        log_file = os.path.join(output_dir, "evolution.log")

        # Try to parse generation files if they exist
        if generation_files:
            evolution_viz += "### Generation-by-Generation Progress\n\n"
            for gen_file in generation_files:
                gen_num = os.path.basename(gen_file).replace("generation_", "").replace(".txt", "")
                try:
                    with open(gen_file, 'r') as f:
                        content = f.read()
                    evolution_viz += f"**Generation {gen_num}:**\n```\n{content[:200]}{'...' if len(content) > 200 else ''}\n```\n\n"
                except:
                    pass

        # Try to parse log file
        elif os.path.exists(log_file):
            evolution_viz += "### Evolution Log\n\n"
            try:
                with open(log_file, 'r') as f:
                    log_content = f.read()
                evolution_viz += f"```\n{log_content[-1000:]}\n```\n\n"
            except:
                pass

        # Look for scores or history file
        scores_file = os.path.join(output_dir, "scores.json")
        if os.path.exists(scores_file):
            try:
                with open(scores_file, 'r') as f:
                    scores = json.load(f)

                evolution_viz += "### Score Progression\n\n"
                evolution_viz += "| Generation | Best Score | Avg Score | Population |\n"
                evolution_viz += "|------------|-----------|-----------|------------|\n"

                for gen in scores:
                    evolution_viz += f"| {gen['generation']} | {gen['best']:.3f} | {gen['avg']:.3f} | {gen['population']} |\n"

                evolution_viz += "\n"
            except:
                pass

        # Look for all program variants
        program_files = sorted(glob.glob(os.path.join(output_dir, "program_*.txt")))
        if program_files:
            evolution_viz += f"### Explored Variants\n\n"
            evolution_viz += f"OpenEvolve explored {len(program_files)} different prompt variants during evolution.\n\n"

            # Show a few intermediate prompts
            if len(program_files) > 3:
                sample_files = [program_files[0], program_files[len(program_files)//2], program_files[-2]]
                evolution_viz += "**Sample Intermediate Prompts:**\n\n"
                for idx, pfile in enumerate(sample_files, 1):
                    try:
                        with open(pfile, 'r') as f:
                            prompt_content = f.read()
                        evolution_viz += f"**Variant {idx}:**\n```\n{prompt_content[:150]}{'...' if len(prompt_content) > 150 else ''}\n```\n\n"
                    except:
                        pass

        # If no specific files found, show directory contents
        if not generation_files and not os.path.exists(log_file) and not os.path.exists(scores_file):
            evolution_viz += "### Evolution Complete\n\n"
            evolution_viz += "OpenEvolve ran 5 iterations of evolutionary optimization using:\n"
            evolution_viz += "- **Population Size**: 10 prompts per generation\n"
            evolution_viz += "- **Selection Strategy**: 10% elite, 30% explore, 60% exploit\n"
            evolution_viz += "- **Islands**: 1 population with mutation and crossover\n"
            evolution_viz += "- **Evaluation**: 50 samples per prompt variant\n\n"

            # Count files in output directory
            all_files = os.listdir(output_dir)
            evolution_viz += f"Generated {len(all_files)} files during evolution process.\n\n"

        return evolution_viz

    except Exception as e:
        return f"## 🧬 Evolution Progress\n\nEvolution completed successfully. Unable to parse detailed history: {str(e)}\n\n"


def create_evaluator_file(dataset_name: str, split: str, model: str,
                         input_field: str, target_field: str, work_dir: str):
    """Create an evaluator.py file for OpenEvolve that uses 150 samples for better signal."""
    evaluator_code = f'''
import os
import random
from datasets import load_dataset
from openai import OpenAI

def evaluate(prompt: str) -> dict:
    """
    Evaluate a prompt using 150 fixed samples for stronger evolution signal.

    Using more samples (150 vs 50) gives evolution better signal to distinguish
    good prompts from bad ones. Final comparison still uses the same 50 samples.
    """
    try:
        # IMPORTANT: Use fixed seed for consistent sampling across all evaluations
        random.seed(42)

        # Load dataset
        try:
            dataset = load_dataset("{dataset_name}", split="{split}", streaming=False)
        except ValueError as e:
            if "config" in str(e).lower() or "Config name is missing" in str(e):
                default_config = "main"
                if "{dataset_name}".lower() == "glue":
                    default_config = "sst2"
                dataset = load_dataset("{dataset_name}", default_config, split="{split}", streaming=False)
            else:
                raise

        # Sample 150 samples with seed 42 for stronger signal during evolution
        num_samples = 150
        if len(dataset) > num_samples:
            # Use SAME sampling logic as initial/final eval
            indices = random.sample(range(len(dataset)), num_samples)
            samples = [dataset[i] for i in indices]
        else:
            indices = list(range(min(num_samples, len(dataset))))
            samples = list(dataset)[:num_samples]

        # Initialize OpenAI client
        api_key = os.environ.get("OPENAI_API_KEY")
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )

        correct = 0
        total = 0

        print(f"Evaluating on {{len(samples)}} samples...")

        for idx, sample in enumerate(samples):
            try:
                # Get input and target
                input_text = sample.get("{input_field}", "")
                if isinstance(input_text, dict):
                    input_text = str(input_text)

                target = sample.get("{target_field}", "")
                if isinstance(target, dict):
                    target = str(target)

                # Format the prompt
                formatted_prompt = prompt.replace("{{input}}", str(input_text))

                # Call the model
                response = client.chat.completions.create(
                    model="{model}",
                    messages=[
                        {{"role": "system", "content": "You are a helpful assistant."}},
                        {{"role": "user", "content": formatted_prompt}}
                    ],
                    temperature=0.0,
                    max_tokens=500,
                )

                prediction = response.choices[0].message.content.strip()

                # IMDB labels: 0 = negative, 1 = positive
                true_label = int(target)  # 0 or 1

                # FORMAT REQUIREMENT: Need "sentiment" keyword + positive/negative in first 150 chars
                # This is strict enough to fail conversational responses, but learnable through evolution
                pred_lower = prediction.lower()
                pred_start = pred_lower[:150]  # First 150 chars

                # Must mention "sentiment" to get credit (helps evolution learn to add this keyword)
                has_sentiment_keyword = "sentiment" in pred_start

                # Check for positive/negative indicators
                has_positive = "positive" in pred_start
                has_negative = "negative" in pred_start

                # Only count as correct if sentiment keyword present AND unambiguous positive/negative
                if has_sentiment_keyword and has_positive and not has_negative:
                    predicted_label = 1
                elif has_sentiment_keyword and has_negative and not has_positive:
                    predicted_label = 0
                else:
                    predicted_label = -1

                is_correct = (predicted_label == true_label)

                if is_correct:
                    correct += 1
                total += 1

                if (idx + 1) % 50 == 0:
                    print(f"  Progress: {{idx + 1}}/{{len(samples)}} - Current accuracy: {{correct/total:.2%}}")

            except Exception as e:
                print(f"Error evaluating sample {{idx+1}}: {{e}}")
                continue

        accuracy = (correct / total) if total > 0 else 0.0

        print(f"Final: {{correct}}/{{total}} = {{accuracy:.2%}}")

        return {{
            "combined_score": accuracy,
            "accuracy": accuracy,
            "correct": correct,
            "total": total
        }}

    except Exception as e:
        print(f"Error in evaluation: {{e}}")
        return {{
            "combined_score": 0.0,
            "accuracy": 0.0,
            "correct": 0,
            "total": 0,
            "error": str(e)
        }}
'''

    evaluator_path = os.path.join(work_dir, "evaluator.py")
    with open(evaluator_path, "w") as f:
        f.write(evaluator_code)

    return evaluator_path


def create_config_file(model: str, work_dir: str):
    """Create a config.yaml file for OpenEvolve."""

    # Create custom templates directory for prompt optimization
    templates_dir = os.path.join(work_dir, "templates")
    os.makedirs(templates_dir, exist_ok=True)

    # Create custom system template for PROMPT optimization (not code)
    system_template = """You are an expert prompt engineer tasked with iteratively improving prompts for language models.
Your job is to analyze the current prompt and suggest improvements based on performance feedback.

CRITICAL RULES:
1. Keep prompts BRIEF and DIRECT - shorter is usually better
2. Preserve the EXACT output format that the evaluation expects
3. Do NOT make prompts conversational or verbose
4. Do NOT ask for explanations - just ask for the answer
5. Maintain all placeholder variables like {input}, {text}, etc.
6. Focus on clarity and directness, not linguistic elegance
7. Avoid prompts that might cause the model to discuss multiple possibilities

For classification tasks:
- Ask for direct classification (e.g., "The sentiment is positive")
- Avoid asking "what", "why", or "explain" - just ask for the label
- Ensure the response will include the label word (positive/negative/neutral)
- Keep prompts short enough that responses stay focused
- IMPORTANT: The prompt should naturally cause the model to echo the task type in its response
  (e.g., if classifying sentiment, the response should include the word "sentiment")

Good examples for sentiment:
- "Review sentiment {input}" → model responds "The sentiment is positive"
- "Classify sentiment: {input}" → model responds "Sentiment: positive"
- "Determine the sentiment of: {input}" → model responds "The sentiment is negative"

Bad examples for sentiment:
- "Is this positive or negative: {input}" → model might respond just "Positive" (missing "sentiment" keyword)
- "Classify: {input}" → too vague, unclear what to classify
- "What sentiment: {input}" → conversational, might get verbose response
- "Analyze the following text and provide a detailed explanation of its sentiment: {input}" → way too verbose
"""

    with open(os.path.join(templates_dir, "system_message.txt"), "w") as f:
        f.write(system_template)

    # Create custom user template for prompt rewriting
    user_template = """# Current Prompt Performance
- Current metrics: {metrics}
- Areas for improvement: {improvement_areas}

{artifacts}

# Prompt Evolution History
{evolution_history}

# Current Prompt
```text
{current_program}
```

# Task
Rewrite the prompt above to improve its performance on the specified metrics.

REMEMBER:
- SHORTER is usually BETTER - avoid adding unnecessary words
- Keep the EXACT same output format (especially placeholder variables like {{input}})
- Focus on DIRECTNESS - what's the clearest way to ask for what we need?
- Avoid conversational language that might confuse the model
- For classification: ask directly for the label, don't ask for explanations

Provide a complete new version of the prompt that:
1. Maintains the same input/output format (keep placeholders like {{input}}, {{text}}, etc.)
2. Is brief and direct
3. Clearly asks for the classification/answer without asking for reasoning
4. Will cause the model to output the label word in its response

Output ONLY the new prompt text between ```text markers:

```text
Your improved prompt here
```
"""

    with open(os.path.join(templates_dir, "full_rewrite_user.txt"), "w") as f:
        f.write(user_template)

    config = {
        "llm": {
            "primary_model": model,
            "api_base": "https://openrouter.ai/api/v1",  # Use OpenRouter endpoint
            "temperature": 1.0,  # Higher temperature for more diverse prompt variations
        },
        "max_iterations": 5,
        "checkpoint_interval": 1,  # Save checkpoints every iteration to preserve prompt history
        "diff_based_evolution": False,  # Use full rewrite mode for prompts (not diff/patch mode)
        "language": "text",  # CRITICAL: Optimize text/prompts, not Python code!
        "max_code_length": 40000,  # Allow long prompts (default 10000 is too short)
        "num_islands": 1,  # IMPORTANT: Use only 1 island (not 5) for simpler evolution
        "prompt": {
            "template_dir": templates_dir,  # Use our custom prompt engineering templates
        },
        "evolution": {
            "population_size": 15,  # Increased from 10 for more exploration
            "num_islands": 1,  # Single island for simpler evolution
            "elite_ratio": 0.1,  # Keep top 10% (1-2 best prompts)
            "explore_ratio": 0.4,  # Increased exploration (was 0.3)
            "exploit_ratio": 0.5,  # Reduced exploitation (was 0.6)
        },
        "database": {
            "log_prompts": True,  # Save prompts used to generate each program
            "num_islands": 1,  # CRITICAL: This is where island count is actually read from!
        },
        "evaluator": {
            "timeout": 3600,  # 1 hour timeout (effectively disabled, but prevents NoneType arithmetic errors)
            "cascade_evaluation": False,  # Disable cascade to prevent signal errors
            "parallel_evaluations": 1,  # Single worker to avoid multiprocessing complexity
            "distributed": False,  # No distributed processing
        }
    }

    config_path = os.path.join(work_dir, "config.yaml")
    with open(config_path, "w") as f:
        yaml.dump(config, f)

    return config_path


def optimize_prompt(initial_prompt: str, dataset_name: str, dataset_split: str,
                   model: str, input_field: str, target_field: str,
                   progress=gr.Progress()) -> Tuple[str, str, str]:
    """Run OpenEvolve to optimize the prompt."""

    progress(0, desc="Validating inputs...")

    # Validate all inputs
    is_valid, validation_message = validate_inputs(
        dataset_name, dataset_split, input_field, target_field, initial_prompt
    )

    if not is_valid:
        return f"## Validation Failed\n\n{validation_message}", "", ""

    progress(0.05, desc=f"Validation passed: {validation_message}")

    # Create temporary working directory
    work_dir = tempfile.mkdtemp(prefix="openevolve_")

    try:
        # Save initial prompt
        initial_prompt_path = os.path.join(work_dir, "initial_prompt.txt")
        with open(initial_prompt_path, "w") as f:
            f.write(initial_prompt)

        # Create evaluator
        progress(0.1, desc="Creating evaluator...")
        evaluator_path = create_evaluator_file(dataset_name, dataset_split, model,
                                               input_field, target_field, work_dir)

        # Create config
        progress(0.15, desc="Creating configuration...")
        config_path = create_config_file(model, work_dir)

        # Run initial evaluation with 50 samples
        # IMPORTANT: We save the indices to ensure final eval uses THE SAME samples
        progress(0.2, desc="Running initial evaluation on 50 samples...")
        initial_eval = evaluate_prompt(
            initial_prompt, dataset_name, dataset_split, 50,
            model, input_field, target_field
        )

        if "error" in initial_eval:
            return f"## Error\n\n❌ Initial evaluation failed: {initial_eval['error']}", "", ""

        if initial_eval["total"] == 0:
            return f"## Error\n\n❌ Initial evaluation failed: No samples could be evaluated. This usually means:\n- API key is invalid or has no credits\n- Model is unavailable or rate-limited\n- Dataset fields are incorrect\n- Network connectivity issues\n\nPlease check your configuration and try again.", "", ""

        # Save the indices for final evaluation (ensures fair comparison)
        eval_indices = initial_eval.get("indices", [])

        initial_results = f"""
### Initial Prompt Evaluation

**Prompt:**
```
{initial_prompt}
```

**Results:**
- Accuracy: {initial_eval['accuracy']:.2f}%
- Correct: {initial_eval['correct']}/{initial_eval['total']}

**Sample Results:**
"""
        for i, result in enumerate(initial_eval['results'][:5], 1):
            initial_results += f"\n{i}. Input: {result['input']}\n"
            initial_results += f"   Target: {result['target']}\n"
            initial_results += f"   Prediction: {result['prediction']}\n"
            initial_results += f"   ✓ Correct\n" if result['correct'] else f"   ✗ Incorrect\n"

        # Run OpenEvolve
        progress(0.3, desc="Starting OpenEvolve optimization (5 iterations)...")

        output_dir = os.path.join(work_dir, "output")
        os.makedirs(output_dir, exist_ok=True)

        try:
            # Comprehensive fix for "signal only works in main thread" in Gradio
            # We need to prevent OpenEvolve from using signal handlers entirely

            # Step 1: Set environment variable to disable process pool
            import os as os_env
            os_env.environ['OPENEVOLVE_NO_PARALLEL'] = '1'

            # Step 2: Monkey-patch signal module to ignore signal calls in threads
            import signal
            import threading

            original_signal = signal.signal

            def safe_signal(signum, handler):
                """Only set signal handlers in main thread"""
                if threading.current_thread() is threading.main_thread():
                    return original_signal(signum, handler)
                else:
                    # Return a dummy handler in non-main threads
                    return signal.SIG_DFL

            signal.signal = safe_signal

            # Run evolution with signal patch in place
            result = run_evolution(
                initial_program=initial_prompt_path,
                evaluator=evaluator_path,
                config=config_path,
                output_dir=output_dir
            )

            # Restore signal handler
            signal.signal = original_signal

            progress(0.80, desc="Parsing evolution history...")

            # Parse evolution history for visualization
            evolution_viz = parse_evolution_history(output_dir)

            progress(0.85, desc="Evaluating best evolved prompt...")

            # Get the best prompt (OpenEvolve saves to output_dir/best/best_program.txt)
            best_prompt_path = os.path.join(output_dir, "best", "best_program.txt")
            if os.path.exists(best_prompt_path):
                with open(best_prompt_path, "r") as f:
                    best_prompt = f.read()
            else:
                # Fallback: try without the "best" subdirectory
                best_prompt_path_alt = os.path.join(output_dir, "best_program.txt")
                if os.path.exists(best_prompt_path_alt):
                    with open(best_prompt_path_alt, "r") as f:
                        best_prompt = f.read()
                else:
                    best_prompt = initial_prompt

            # Final evaluation: Use same 50 samples as initial eval for fair comparison
            progress(0.85, desc="Evaluating best prompt on 50 samples (same as initial)...")
            final_eval = evaluate_prompt(
                best_prompt, dataset_name, dataset_split, 50,
                model, input_field, target_field,
                fixed_indices=eval_indices  # Use same 50 samples as initial eval!
            )

            progress(0.95, desc=f"Evaluation complete: {final_eval['correct']}/{final_eval['total']} = {final_eval['accuracy']:.1f}%")

            final_results = f"""
### Evolved Prompt Evaluation

**Prompt:**
```
{best_prompt}
```

**Validation:**
- Contains {{input}} placeholder: {'✓ Yes' if '{input}' in best_prompt else '❌ NO - This will break evaluation!'}
- Prompt length: {len(best_prompt)} characters

**Results:**
- Accuracy: {final_eval['accuracy']:.2f}%
- Correct: {final_eval['correct']}/{final_eval['total']}
- Improvement: {final_eval['accuracy'] - initial_eval['accuracy']:+.2f}%

**Sample Results:**
"""
            for i, result in enumerate(final_eval['results'][:5], 1):
                final_results += f"\n{i}. Input: {result['input']}\n"
                final_results += f"   Target: {result['target']}\n"
                final_results += f"   Prediction: {result['prediction']}\n"
                final_results += f"   ✓ Correct\n" if result['correct'] else f"   ✗ Incorrect\n"

            summary = f"""
## 🎉 Optimization Complete!

### Summary
- **Dataset**: {dataset_name} ({dataset_split} split)
- **Model**: {model}
- **Initial Eval**: 50 samples
- **Final Eval**: 50 samples (same samples for fair comparison)
- **Evolution**: 150 samples per variant (more data for stronger signal)
- **Iterations**: 5

### Results
- **Initial Accuracy**: {initial_eval['accuracy']:.2f}% ({initial_eval['correct']}/{initial_eval['total']})
- **Final Accuracy**: {final_eval['accuracy']:.2f}% ({final_eval['correct']}/{final_eval['total']})
- **Improvement**: {final_eval['accuracy'] - initial_eval['accuracy']:+.2f}%

{validation_message}
"""

            progress(1.0, desc="Complete!")

            return summary, initial_results, final_results

        except Exception as e:
            return f"## Error During Evolution\n\n❌ {str(e)}", initial_results, ""

    finally:
        # Don't clean up - keep prompts for browsing
        # User can manually clean /tmp if needed
        pass


# Create Gradio interface
with gr.Blocks(title="OpenEvolve Prompt Optimizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🧬 OpenEvolve Prompt Optimizer

    Automatically evolve and optimize your prompts using evolutionary algorithms!

    This space uses [OpenEvolve](https://github.com/algorithmicsuperintelligence/openevolve) to iteratively improve prompts
    by testing them on real datasets and evolving better versions.

    ## How it works:
    1. Enter an initial prompt (use `{input}` as a placeholder for dataset inputs)
    2. Default dataset is **IMDB** (movie review sentiment classification) - great for showing prompt improvement!
    3. Specify the dataset split and field names (or use other datasets like `gsm8k`, `dair-ai/emotion`)
    4. Choose a free model from OpenRouter
    5. Click "Optimize Prompt" - the system will validate everything first!
    6. Watch the evolution progress in real-time
    7. Compare initial vs. best prompt side-by-side (both evaluated on same 200 samples)!

    **Note**: API key is read from `OPENAI_API_KEY` environment variable (set in Space secrets)
    """)

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Configuration")

            model = gr.Dropdown(
                choices=FREE_MODELS,
                value=FREE_MODELS[0],
                label="Select Model",
                info="Choose from 5 curated free models on OpenRouter (24B to 671B parameters)"
            )

            dataset_name = gr.Textbox(
                label="HuggingFace Dataset (Full Name)",
                value="stanfordnlp/imdb",
                placeholder="e.g., stanfordnlp/imdb, gsm8k, MathArena/aime_2025",
                info="Dataset name from HuggingFace Hub. Default: IMDB (sentiment classification)"
            )

            dataset_split = gr.Textbox(
                label="Dataset Split",
                value="test",
                placeholder="e.g., train, test, validation"
            )

            input_field = gr.Textbox(
                label="Input Field Name",
                value="text",
                placeholder="e.g., text, question, sentence",
                info="The field containing inputs to process"
            )

            target_field = gr.Textbox(
                label="Target Field Name",
                value="label",
                placeholder="e.g., label, answer, target",
                info="The field containing expected outputs"
            )

            initial_prompt = gr.TextArea(
                label="Initial Prompt",
                value="Review sentiment {input}",
                lines=5,
                info="Use {input} as placeholder for dataset inputs. Start simple - evolution will improve it!"
            )

    # Button outside the column for better visibility
    with gr.Row():
        with gr.Column():
            optimize_btn = gr.Button("🚀 Validate & Optimize Prompt", variant="primary", size="lg")

    # Results section - clearly separated
    gr.Markdown("---")
    gr.Markdown("## 📊 Results")

    with gr.Row():
        with gr.Column():
            summary = gr.Markdown("Click 'Validate & Optimize Prompt' to start optimization...", visible=True)

    # Side-by-side comparison: Initial vs Best Prompt
    gr.Markdown("---")
    gr.Markdown("## 🔍 Prompt Comparison: Initial vs Best")

    with gr.Row():
        with gr.Column():
            initial_results = gr.Markdown("### Initial Prompt\nWill appear here after validation...", visible=True)
        with gr.Column():
            final_results = gr.Markdown("### Best Prompt\nWill appear here after optimization...", visible=True)

    # Documentation section - in collapsible accordion
    gr.Markdown("---")
    with gr.Accordion("📚 Documentation & Examples", open=False):
        gr.Markdown("""
        ### Example Datasets & Fields:

        | Dataset | Split | Input Field | Target Field | Task | Size |
        |---------|-------|-------------|--------------|------|------|
        | **stanfordnlp/imdb** | test | text | label | Sentiment Analysis | 25,000 |
        | dair-ai/emotion | test | text | label | Emotion Classification | 2,000 |
        | gsm8k | train | question | answer | Grade School Math | 7,473 |
        | fancyzhx/ag_news | test | text | label | News Classification | 7,600 |
        | MathArena/aime_2025 | train | problem | answer | Hard Math (AIME) | 30 |

        ### About This Demo Space:

        **This is a demonstration space** showcasing OpenEvolve's prompt optimization capabilities.
        The interface shows you how the system works, but **you'll need to set up your own instance to run optimizations**.

        ### How to Run This Yourself:

        1. **Clone this Space**: Click "⋮" (three dots) at top-right → "Duplicate this Space"
        2. **Set Environment Variables** in your cloned Space's settings:
           - `OPENAI_API_KEY`: Your OpenRouter API key (get free key at [openrouter.ai/keys](https://openrouter.ai/keys))
           - `HF_TOKEN`: (Optional) HuggingFace token for private datasets
        3. **Configure Your Optimization**:
           - Dataset: Use full name format (e.g., `stanfordnlp/imdb` or `openai/gsm8k`)
           - Fields: Specify exact field names from the dataset schema
           - Model: Choose from 5 curated free models (larger models = better results but slower/rate-limited)
        4. **Run & Monitor**:
           - All inputs are validated before starting
           - **Evaluation strategy**:
             - Initial evaluation: 200 samples
             - Final evaluation: Same 200 samples (fair comparison)
             - Evolution: Each variant tested on same 200 samples (ensures optimization aligns with test set)
           - Compare initial vs best prompt side-by-side with identical test sets

        ### About OpenEvolve:
        OpenEvolve is an open-source evolutionary optimization framework. Learn more at:
        - [GitHub Repository](https://github.com/algorithmicsuperintelligence/openevolve)
        - [Documentation](https://github.com/algorithmicsuperintelligence/openevolve#readme)
        """)

    # Wire up the optimize button
    optimize_btn.click(
        fn=optimize_prompt,
        inputs=[initial_prompt, dataset_name, dataset_split, model,
                input_field, target_field],
        outputs=[summary, initial_results, final_results]
    )

if __name__ == "__main__":
    demo.launch()