reasonkit-core 0.1.8

//! # Prompt Optimization Module
//!
//! This module provides DSPy-inspired prompt optimization capabilities for ReasonKit.
//!
//! ## Overview
//!
//! The optimization module implements automatic prompt optimization through:
//! - **Signature parsing**: Declarative input/output specifications
//! - **MIPRO optimization**: Multi-stage Instruction Proposal and Optimization
//! - **Bootstrapping**: Automatic demonstration generation
//!
//! ## Research Background
//!
//! This implementation draws heavily from the DSPy framework and associated research:
//!
//! - Khattab, O., Singhvi, A., Maheshwari, P., Zhang, Z., Santhanam, K., Vardhamanan, S.,
//!   Haq, S., Sharma, A., Joshi, T. T., Mober, H., Khandelwal, U., & Zaharia, M. (2023).
//!   "DSPy: Compiling Declarative Language Model Calls into State-of-the-Art Pipelines."
//!   arXiv:2310.03714. <https://arxiv.org/abs/2310.03714>
//!
//! - Opsahl-Ong, K., Ryan, M. J., Hardy, J., & Khattab, O. (2024).
//!   "Optimizing Instructions and Demonstrations for Multi-Stage Language Model Programs."
//!   arXiv:2406.11695 (MIPRO). <https://arxiv.org/abs/2406.11695>
//!
//! ## Example Usage
//!
//! ```rust,ignore
//! use reasonkit_core::optimization::{Signature, PromptOptimizer, OptimizerConfig};
//!
//! // Parse a simple signature
//! let sig = Signature::parse("question, context -> answer")?;
//!
//! // Parse with type annotations
//! let typed_sig = Signature::parse("query:str, docs:list -> summary:str, confidence:float")?;
//!
//! // Create optimizer with custom config
//! let config = OptimizerConfig::default()
//!     .with_num_trials(50)
//!     .with_minibatch_size(32);
//!
//! let optimizer = PromptOptimizer::new(config);
//! let result = optimizer.optimize(&sig, &training_data).await?;
//! ```
//!
//! ## Signatures
//!
//! Signatures define the input/output contract for language model calls using a declarative
//! syntax inspired by DSPy:
//!
//! - Simple: `"question -> answer"`
//! - Multiple inputs: `"question, context -> answer"`
//! - With types: `"query:str, count:int -> results:list"`
//! - With instructions: Set via `Signature::with_instructions()`
//!
//! ## Supported Field Types
//!
//! | Type | Syntax | Description |
//! |------|--------|-------------|
//! | String | `str` or default | Text content |
//! | Integer | `int` | Whole numbers |
//! | Float | `float` | Decimal numbers |
//! | Boolean | `bool` | True/false values |
//! | JSON | `json` | Structured JSON data |
//! | List | `list` or `list[type]` | Arrays of values |

use std::collections::HashMap;
use std::fmt;
use std::str::FromStr;

use serde::{Deserialize, Serialize};

/// Errors that can occur during signature parsing or optimization.
#[derive(Debug, Clone, PartialEq)]
pub enum OptimizationError {
    /// The signature string could not be parsed.
    InvalidSignature(String),
    /// A field type specification is invalid.
    InvalidFieldType(String),
    /// The signature is missing required components.
    MissingComponent(String),
    /// Optimization failed to converge.
    OptimizationFailed(String),
    /// Training data is insufficient.
    InsufficientData { required: usize, provided: usize },
}

impl fmt::Display for OptimizationError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidSignature(s) => write!(f, "Invalid signature: {s}"),
            Self::InvalidFieldType(s) => write!(f, "Invalid field type: {s}"),
            Self::MissingComponent(s) => write!(f, "Missing component: {s}"),
            Self::OptimizationFailed(s) => write!(f, "Optimization failed: {s}"),
            Self::InsufficientData { required, provided } => {
                write!(
                    f,
                    "Insufficient training data: required {required}, provided {provided}"
                )
            }
        }
    }
}

impl std::error::Error for OptimizationError {}

/// Result type for optimization operations.
pub type Result<T> = std::result::Result<T, OptimizationError>;

// ============================================================================
// Field Type System
// ============================================================================

/// Represents the data type of a signature field.
///
/// DSPy uses Python's type system; this provides an equivalent Rust representation
/// that can be used for validation and prompt formatting.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum FieldType {
    /// A string/text field (default type).
    #[default]
    String,
    /// An integer field.
    Integer,
    /// A floating-point number field.
    Float,
    /// A boolean field.
    Boolean,
    /// A JSON object field.
    Json,
    /// A list/array field with an optional inner type.
    List(Box<FieldType>),
}

impl fmt::Display for FieldType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::String => write!(f, "str"),
            Self::Integer => write!(f, "int"),
            Self::Float => write!(f, "float"),
            Self::Boolean => write!(f, "bool"),
            Self::Json => write!(f, "json"),
            Self::List(inner) => {
                if **inner == FieldType::String {
                    write!(f, "list")
                } else {
                    write!(f, "list[{inner}]")
                }
            }
        }
    }
}

impl FromStr for FieldType {
    type Err = OptimizationError;

    fn from_str(s: &str) -> Result<Self> {
        let s = s.trim().to_lowercase();

        // Handle list types with inner type
        if s.starts_with("list[") && s.ends_with(']') {
            let inner = &s[5..s.len() - 1];
            let inner_type = inner.parse::<FieldType>()?;
            return Ok(Self::List(Box::new(inner_type)));
        }

        match s.as_str() {
            "str" | "string" | "text" => Ok(Self::String),
            "int" | "integer" => Ok(Self::Integer),
            "float" | "number" | "decimal" => Ok(Self::Float),
            "bool" | "boolean" => Ok(Self::Boolean),
            "json" | "object" | "dict" => Ok(Self::Json),
            "list" | "array" => Ok(Self::List(Box::new(FieldType::String))),
            _ => Err(OptimizationError::InvalidFieldType(s)),
        }
    }
}

// ============================================================================
// Signature Field
// ============================================================================

/// A single field in a signature (either input or output).
///
/// Each field has a name, type, and optional description. The `prefix` is used
/// in prompt formatting to label the field (defaults to capitalized name with colon).
///
/// # Example
///
/// ```rust
/// use reasonkit_core::optimization::{SignatureField, FieldType};
///
/// let field = SignatureField::new("question", FieldType::String)
///     .with_description("The user's question to answer");
///
/// assert_eq!(field.name, "question");
/// assert_eq!(field.prefix, "Question:");
/// ```
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SignatureField {
    /// The name of the field (e.g., "question", "answer").
    pub name: String,
    /// The data type of the field.
    pub field_type: FieldType,
    /// The prefix used in prompts (e.g., "Question:").
    pub prefix: String,
    /// Optional description of the field's purpose.
    pub description: Option<String>,
}

impl SignatureField {
    /// Creates a new signature field with the given name and type.
    ///
    /// The prefix is automatically generated as the capitalized name followed by a colon.
    pub fn new(name: impl Into<String>, field_type: FieldType) -> Self {
        let name = name.into();
        let prefix = Self::generate_prefix(&name);
        Self {
            name,
            field_type,
            prefix,
            description: None,
        }
    }

    /// Creates a new string-typed field with the given name.
    pub fn string(name: impl Into<String>) -> Self {
        Self::new(name, FieldType::String)
    }

    /// Adds a description to the field.
    #[must_use]
    pub fn with_description(mut self, description: impl Into<String>) -> Self {
        self.description = Some(description.into());
        self
    }

    /// Sets a custom prefix for the field.
    #[must_use]
    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.prefix = prefix.into();
        self
    }

    /// Generates a default prefix from a field name.
    ///
    /// Converts snake_case to Title Case and appends a colon.
    fn generate_prefix(name: &str) -> String {
        let words: Vec<&str> = name.split('_').collect();
        let titled: Vec<String> = words
            .iter()
            .map(|w| {
                let mut chars = w.chars();
                match chars.next() {
                    Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
                    None => String::new(),
                }
            })
            .collect();
        format!("{}:", titled.join(" "))
    }

    /// Parses a field specification string like "name:type" or just "name".
    pub fn parse(spec: &str) -> Result<Self> {
        let spec = spec.trim();

        if let Some((name, type_str)) = spec.split_once(':') {
            let name = name.trim();
            let field_type = type_str.trim().parse::<FieldType>()?;
            Ok(Self::new(name, field_type))
        } else {
            Ok(Self::string(spec))
        }
    }
}

// ============================================================================
// Signature
// ============================================================================

/// A declarative specification of inputs and outputs for an LLM call.
///
/// Signatures are the core abstraction in DSPy-style programming. They define
/// what information goes into a language model call and what comes out, without
/// specifying how the model should process it.
///
/// # Parsing Syntax
///
/// Signatures can be parsed from strings using a simple syntax:
///
/// - `"input -> output"` - Single input, single output
/// - `"a, b, c -> x, y"` - Multiple inputs and outputs
/// - `"query:str -> answer:str"` - With type annotations
/// - `"items:list[str] -> count:int"` - Complex types
///
/// # Example
///
/// ```rust
/// use reasonkit_core::optimization::Signature;
///
/// let sig = Signature::parse("question, context -> answer, confidence:float")?;
///
/// assert_eq!(sig.input_fields.len(), 2);
/// assert_eq!(sig.output_fields.len(), 2);
/// assert_eq!(sig.input_fields[0].name, "question");
/// # Ok::<(), reasonkit_core::optimization::OptimizationError>(())
/// ```
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct Signature {
    /// Instructions or task description for the LLM.
    pub instructions: String,
    /// Input fields that will be provided to the LLM.
    pub input_fields: Vec<SignatureField>,
    /// Output fields expected from the LLM.
    pub output_fields: Vec<SignatureField>,
}

impl Signature {
    /// Creates a new empty signature.
    pub fn new() -> Self {
        Self::default()
    }

    /// Parses a signature from a string specification.
    ///
    /// The format is: `"input_fields -> output_fields"` where each side is a
    /// comma-separated list of field specifications.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use reasonkit_core::optimization::Signature;
    ///
    /// // Simple signature
    /// let sig = Signature::parse("question -> answer")?;
    ///
    /// // With multiple fields and types
    /// let sig = Signature::parse("query:str, docs:list -> summary:str, score:float")?;
    /// # Ok::<(), reasonkit_core::optimization::OptimizationError>(())
    /// ```
    pub fn parse(spec: &str) -> Result<Self> {
        let spec = spec.trim();

        // Split on "->" to separate inputs from outputs
        let parts: Vec<&str> = spec.split("->").collect();

        if parts.len() != 2 {
            return Err(OptimizationError::InvalidSignature(format!(
                "Expected 'inputs -> outputs' format, got: {spec}"
            )));
        }

        let input_spec = parts[0].trim();
        let output_spec = parts[1].trim();

        // Parse input fields
        let input_fields = Self::parse_field_list(input_spec)?;
        if input_fields.is_empty() {
            return Err(OptimizationError::MissingComponent(
                "At least one input field is required".to_string(),
            ));
        }

        // Parse output fields
        let output_fields = Self::parse_field_list(output_spec)?;
        if output_fields.is_empty() {
            return Err(OptimizationError::MissingComponent(
                "At least one output field is required".to_string(),
            ));
        }

        Ok(Self {
            instructions: String::new(),
            input_fields,
            output_fields,
        })
    }

    /// Parses a comma-separated list of field specifications.
    fn parse_field_list(spec: &str) -> Result<Vec<SignatureField>> {
        if spec.is_empty() {
            return Ok(Vec::new());
        }

        spec.split(',')
            .map(|s| SignatureField::parse(s.trim()))
            .collect()
    }

    /// Sets the instructions for this signature.
    #[must_use]
    pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
        self.instructions = instructions.into();
        self
    }

    /// Adds an input field to the signature.
    #[must_use]
    pub fn with_input(mut self, field: SignatureField) -> Self {
        self.input_fields.push(field);
        self
    }

    /// Adds an output field to the signature.
    #[must_use]
    pub fn with_output(mut self, field: SignatureField) -> Self {
        self.output_fields.push(field);
        self
    }

    /// Returns the names of all input fields.
    pub fn input_names(&self) -> Vec<&str> {
        self.input_fields.iter().map(|f| f.name.as_str()).collect()
    }

    /// Returns the names of all output fields.
    pub fn output_names(&self) -> Vec<&str> {
        self.output_fields.iter().map(|f| f.name.as_str()).collect()
    }

    /// Generates a prompt template from the signature.
    ///
    /// This creates a structured prompt with field prefixes and placeholders.
    pub fn to_prompt_template(&self) -> String {
        let mut parts = Vec::new();

        // Add instructions if present
        if !self.instructions.is_empty() {
            parts.push(self.instructions.clone());
            parts.push(String::new()); // Empty line
        }

        // Add input field placeholders
        for field in &self.input_fields {
            parts.push(format!("{} {{{}}}", field.prefix, field.name));
        }

        // Add output field markers
        for field in &self.output_fields {
            parts.push(field.prefix.to_string());
        }

        parts.join("\n")
    }

    /// Formats the signature for display (reverse of parsing).
    pub fn to_spec_string(&self) -> String {
        let inputs: Vec<String> = self
            .input_fields
            .iter()
            .map(|f| {
                if f.field_type == FieldType::String {
                    f.name.clone()
                } else {
                    format!("{}:{}", f.name, f.field_type)
                }
            })
            .collect();

        let outputs: Vec<String> = self
            .output_fields
            .iter()
            .map(|f| {
                if f.field_type == FieldType::String {
                    f.name.clone()
                } else {
                    format!("{}:{}", f.name, f.field_type)
                }
            })
            .collect();

        format!("{} -> {}", inputs.join(", "), outputs.join(", "))
    }
}

impl fmt::Display for Signature {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.to_spec_string())
    }
}

impl FromStr for Signature {
    type Err = OptimizationError;

    fn from_str(s: &str) -> Result<Self> {
        Self::parse(s)
    }
}

// ============================================================================
// Optimizer Configuration
// ============================================================================

/// Configuration for the MIPRO prompt optimizer.
///
/// MIPRO (Multi-stage Instruction Proposal and Optimization) is an algorithm
/// from DSPy that optimizes prompts through multiple stages:
///
/// 1. **Instruction proposal**: Generate candidate instructions
/// 2. **Demonstration selection**: Choose effective few-shot examples
/// 3. **Optimization**: Search for best prompt configuration
///
/// # Default Values
///
/// The defaults are based on the MIPRO paper (Opsahl-Ong et al., 2024):
///
/// | Parameter | Default | Description |
/// |-----------|---------|-------------|
/// | `num_trials` | 20 | Total optimization trials |
/// | `minibatch_size` | 25 | Examples per evaluation |
/// | `max_bootstrapped_demos` | 4 | Auto-generated demos |
/// | `max_labeled_demos` | 16 | Human-labeled demos |
/// | `num_instruction_candidates` | 10 | Instructions to consider |
///
/// # Example
///
/// ```rust
/// use reasonkit_core::optimization::OptimizerConfig;
///
/// let config = OptimizerConfig::default()
///     .with_num_trials(50)
///     .with_minibatch_size(32);
///
/// assert_eq!(config.num_trials, 50);
/// assert_eq!(config.minibatch_size, 32);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizerConfig {
    /// Number of optimization trials to run.
    pub num_trials: usize,
    /// Number of examples to use in each evaluation minibatch.
    pub minibatch_size: usize,
    /// Maximum number of bootstrapped (auto-generated) demonstrations.
    pub max_bootstrapped_demos: usize,
    /// Maximum number of labeled (human-provided) demonstrations.
    pub max_labeled_demos: usize,
    /// Number of instruction candidates to generate and evaluate.
    pub num_instruction_candidates: usize,
    /// Seed for random number generation (for reproducibility).
    pub seed: Option<u64>,
    /// Enable verbose logging during optimization.
    pub verbose: bool,
    /// Early stopping threshold (stop if no improvement for N trials).
    pub early_stopping_patience: Option<usize>,
    /// Minimum improvement threshold to continue optimization.
    pub min_improvement_threshold: f64,
}

impl Default for OptimizerConfig {
    fn default() -> Self {
        Self {
            num_trials: 20,
            minibatch_size: 25,
            max_bootstrapped_demos: 4,
            max_labeled_demos: 16,
            num_instruction_candidates: 10,
            seed: None,
            verbose: false,
            early_stopping_patience: Some(5),
            min_improvement_threshold: 0.01,
        }
    }
}

impl OptimizerConfig {
    /// Creates a new configuration with default values.
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the number of optimization trials.
    #[must_use]
    pub fn with_num_trials(mut self, num_trials: usize) -> Self {
        self.num_trials = num_trials;
        self
    }

    /// Sets the minibatch size for evaluation.
    #[must_use]
    pub fn with_minibatch_size(mut self, minibatch_size: usize) -> Self {
        self.minibatch_size = minibatch_size;
        self
    }

    /// Sets the maximum number of bootstrapped demonstrations.
    #[must_use]
    pub fn with_max_bootstrapped_demos(mut self, max_demos: usize) -> Self {
        self.max_bootstrapped_demos = max_demos;
        self
    }

    /// Sets the maximum number of labeled demonstrations.
    #[must_use]
    pub fn with_max_labeled_demos(mut self, max_demos: usize) -> Self {
        self.max_labeled_demos = max_demos;
        self
    }

    /// Sets the number of instruction candidates to generate.
    #[must_use]
    pub fn with_num_instruction_candidates(mut self, num_candidates: usize) -> Self {
        self.num_instruction_candidates = num_candidates;
        self
    }

    /// Sets the random seed for reproducibility.
    #[must_use]
    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }

    /// Enables verbose logging.
    #[must_use]
    pub fn with_verbose(mut self, verbose: bool) -> Self {
        self.verbose = verbose;
        self
    }

    /// Sets the early stopping patience.
    #[must_use]
    pub fn with_early_stopping(mut self, patience: usize) -> Self {
        self.early_stopping_patience = Some(patience);
        self
    }

    /// Disables early stopping.
    #[must_use]
    pub fn without_early_stopping(mut self) -> Self {
        self.early_stopping_patience = None;
        self
    }
}

// ============================================================================
// Optimization Result
// ============================================================================

/// A single trial result in the optimization process.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrialResult {
    /// Trial number (0-indexed).
    pub trial_number: usize,
    /// The prompt configuration tested.
    pub prompt: String,
    /// The instruction used.
    pub instruction: String,
    /// Number of demonstrations used.
    pub num_demos: usize,
    /// Score achieved on the evaluation set.
    pub score: f64,
    /// Duration of the trial in milliseconds.
    pub duration_ms: u64,
    /// Additional metadata about the trial.
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Metrics collected during optimization.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct OptimizationMetrics {
    /// Total optimization time in milliseconds.
    pub total_duration_ms: u64,
    /// Number of trials completed.
    pub trials_completed: usize,
    /// Number of LLM calls made.
    pub llm_calls: usize,
    /// Total tokens consumed.
    pub tokens_used: usize,
    /// Improvement from initial to best score.
    pub improvement: f64,
    /// Whether early stopping was triggered.
    pub early_stopped: bool,
    /// Score at each trial for plotting convergence.
    pub score_history: Vec<f64>,
}

/// The result of prompt optimization.
///
/// Contains the optimized prompt, the best score achieved, and detailed
/// information about the optimization process.
///
/// # Example
///
/// ```rust,ignore
/// let result = optimizer.optimize(&signature, &training_data).await?;
///
/// println!("Best score: {:.2}%", result.best_score * 100.0);
/// println!("Optimized prompt:\n{}", result.best_prompt);
/// println!("Trials run: {}", result.metrics.trials_completed);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizationResult {
    /// The best prompt found during optimization.
    pub best_prompt: String,
    /// The best instruction found.
    pub best_instruction: String,
    /// The demonstrations selected for the best prompt.
    pub best_demos: Vec<Demonstration>,
    /// The best score achieved during optimization.
    pub best_score: f64,
    /// History of all trials.
    pub trials_history: Vec<TrialResult>,
    /// Aggregated metrics about the optimization.
    pub metrics: OptimizationMetrics,
    /// The original signature that was optimized.
    pub signature: Signature,
    /// The configuration used for optimization.
    pub config: OptimizerConfig,
}

impl OptimizationResult {
    /// Returns the best trial from the optimization.
    pub fn best_trial(&self) -> Option<&TrialResult> {
        self.trials_history
            .iter()
            .max_by(|a, b| a.score.partial_cmp(&b.score).unwrap())
    }

    /// Returns the improvement from the first trial to the best trial.
    pub fn improvement(&self) -> f64 {
        self.metrics.improvement
    }

    /// Returns whether the optimization converged (found a stable solution).
    pub fn converged(&self) -> bool {
        self.metrics.early_stopped || self.metrics.trials_completed == self.config.num_trials
    }
}

// ============================================================================
// Demonstration
// ============================================================================

/// A demonstration (few-shot example) for the prompt.
///
/// Demonstrations provide input-output examples that guide the LLM's behavior.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Demonstration {
    /// Input values keyed by field name.
    pub inputs: HashMap<String, String>,
    /// Output values keyed by field name.
    pub outputs: HashMap<String, String>,
    /// Optional source identifier (e.g., "bootstrapped", "labeled").
    pub source: Option<String>,
    /// Quality score if evaluated.
    pub quality_score: Option<f64>,
}

impl Demonstration {
    /// Creates a new empty demonstration.
    pub fn new() -> Self {
        Self {
            inputs: HashMap::new(),
            outputs: HashMap::new(),
            source: None,
            quality_score: None,
        }
    }

    /// Adds an input field value.
    #[must_use]
    pub fn with_input(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.inputs.insert(name.into(), value.into());
        self
    }

    /// Adds an output field value.
    #[must_use]
    pub fn with_output(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.outputs.insert(name.into(), value.into());
        self
    }

    /// Sets the source of this demonstration.
    #[must_use]
    pub fn with_source(mut self, source: impl Into<String>) -> Self {
        self.source = Some(source.into());
        self
    }

    /// Formats the demonstration according to a signature.
    pub fn format(&self, signature: &Signature) -> String {
        let mut parts = Vec::new();

        // Format inputs
        for field in &signature.input_fields {
            if let Some(value) = self.inputs.get(&field.name) {
                parts.push(format!("{} {}", field.prefix, value));
            }
        }

        // Format outputs
        for field in &signature.output_fields {
            if let Some(value) = self.outputs.get(&field.name) {
                parts.push(format!("{} {}", field.prefix, value));
            }
        }

        parts.join("\n")
    }
}

impl Default for Demonstration {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Training Example
// ============================================================================

/// A training example for optimization.
///
/// Training examples are used to evaluate prompt candidates during optimization.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingExample {
    /// Input values keyed by field name.
    pub inputs: HashMap<String, String>,
    /// Expected output values keyed by field name (ground truth).
    pub expected_outputs: HashMap<String, String>,
    /// Optional weight for this example in evaluation.
    pub weight: Option<f64>,
}

impl TrainingExample {
    /// Creates a new empty training example.
    pub fn new() -> Self {
        Self {
            inputs: HashMap::new(),
            expected_outputs: HashMap::new(),
            weight: None,
        }
    }

    /// Adds an input field value.
    #[must_use]
    pub fn with_input(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.inputs.insert(name.into(), value.into());
        self
    }

    /// Adds an expected output field value.
    #[must_use]
    pub fn with_expected(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.expected_outputs.insert(name.into(), value.into());
        self
    }

    /// Sets the weight for this example.
    #[must_use]
    pub fn with_weight(mut self, weight: f64) -> Self {
        self.weight = Some(weight);
        self
    }
}

impl Default for TrainingExample {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Prompt Optimizer
// ============================================================================

/// MIPRO-based prompt optimizer.
///
/// The `PromptOptimizer` implements the Multi-stage Instruction Proposal and
/// Optimization (MIPRO) algorithm from DSPy. It automatically finds optimal
/// prompt configurations including instructions and demonstrations.
///
/// # Algorithm Overview
///
/// 1. **Instruction Proposal**: Generate diverse instruction candidates
/// 2. **Demo Bootstrap**: Create demonstrations from successful outputs
/// 3. **Joint Optimization**: Search for best instruction + demo combination
/// 4. **Evaluation**: Score candidates on held-out validation set
///
/// # Example
///
/// ```rust,ignore
/// use reasonkit_core::optimization::{PromptOptimizer, OptimizerConfig, Signature};
///
/// let config = OptimizerConfig::default()
///     .with_num_trials(30)
///     .with_verbose(true);
///
/// let optimizer = PromptOptimizer::new(config);
///
/// let signature = Signature::parse("question, context -> answer")?;
/// let training_data = vec![/* training examples */];
///
/// let result = optimizer.optimize(&signature, &training_data).await?;
///
/// println!("Best prompt:\n{}", result.best_prompt);
/// println!("Score: {:.2}%", result.best_score * 100.0);
/// ```
#[derive(Debug, Clone)]
pub struct PromptOptimizer {
    /// Configuration for the optimizer.
    config: OptimizerConfig,
}

impl PromptOptimizer {
    /// Creates a new optimizer with the given configuration.
    pub fn new(config: OptimizerConfig) -> Self {
        Self { config }
    }

    /// Creates a new optimizer with default configuration.
    pub fn with_defaults() -> Self {
        Self::new(OptimizerConfig::default())
    }

    /// Returns a reference to the optimizer configuration.
    pub fn config(&self) -> &OptimizerConfig {
        &self.config
    }

    /// Validates that the training data is sufficient for optimization.
    fn validate_training_data(
        &self,
        signature: &Signature,
        data: &[TrainingExample],
    ) -> Result<()> {
        let min_required = self.config.minibatch_size;

        if data.len() < min_required {
            return Err(OptimizationError::InsufficientData {
                required: min_required,
                provided: data.len(),
            });
        }

        // Validate that examples have required fields
        for (i, example) in data.iter().enumerate() {
            for field in &signature.input_fields {
                if !example.inputs.contains_key(&field.name) {
                    return Err(OptimizationError::InvalidSignature(format!(
                        "Training example {i} missing input field '{}'",
                        field.name
                    )));
                }
            }
            for field in &signature.output_fields {
                if !example.expected_outputs.contains_key(&field.name) {
                    return Err(OptimizationError::InvalidSignature(format!(
                        "Training example {i} missing expected output field '{}'",
                        field.name
                    )));
                }
            }
        }

        Ok(())
    }

    /// Generates instruction candidates for the signature.
    ///
    /// This is a placeholder for the actual instruction proposal logic.
    fn generate_instruction_candidates(&self, signature: &Signature) -> Vec<String> {
        let base_instruction = if signature.instructions.is_empty() {
            format!(
                "Given the input fields {}, produce the output fields {}.",
                signature.input_names().join(", "),
                signature.output_names().join(", ")
            )
        } else {
            signature.instructions.clone()
        };

        // Generate variations (placeholder - in full implementation this would use an LLM)
        vec![
            base_instruction.clone(),
            format!("You are an expert assistant. {base_instruction}"),
            format!("{base_instruction} Be concise and accurate."),
            format!("{base_instruction} Think step by step."),
            format!("Task: {base_instruction}\nProvide your response in the specified format."),
        ]
    }

    /// Selects demonstrations from available examples.
    ///
    /// This is a placeholder for the actual demo selection logic.
    fn select_demonstrations(
        &self,
        _signature: &Signature,
        data: &[TrainingExample],
        max_demos: usize,
    ) -> Vec<Demonstration> {
        // Simple selection: take first N examples as demos
        data.iter()
            .take(max_demos)
            .map(|ex| {
                let mut demo = Demonstration::new().with_source("labeled");
                for (k, v) in &ex.inputs {
                    demo.inputs.insert(k.clone(), v.clone());
                }
                for (k, v) in &ex.expected_outputs {
                    demo.outputs.insert(k.clone(), v.clone());
                }
                demo
            })
            .collect()
    }

    /// Builds a complete prompt from instruction and demonstrations.
    fn build_prompt(
        &self,
        signature: &Signature,
        instruction: &str,
        demos: &[Demonstration],
    ) -> String {
        let mut parts = Vec::new();

        // Add instruction
        parts.push(instruction.to_string());
        parts.push(String::new());

        // Add demonstrations
        if !demos.is_empty() {
            parts.push("---".to_string());
            for (i, demo) in demos.iter().enumerate() {
                if i > 0 {
                    parts.push(String::new());
                }
                parts.push(demo.format(signature));
            }
            parts.push("---".to_string());
            parts.push(String::new());
        }

        // Add input placeholders
        for field in &signature.input_fields {
            parts.push(format!("{} {{{}}}", field.prefix, field.name));
        }

        // Add output field markers
        for field in &signature.output_fields {
            parts.push(field.prefix.to_string());
        }

        parts.join("\n")
    }

    /// Evaluates a prompt candidate on the training data.
    ///
    /// This is a placeholder - actual evaluation would call an LLM.
    fn evaluate_prompt(
        &self,
        _prompt: &str,
        _signature: &Signature,
        _data: &[TrainingExample],
    ) -> f64 {
        // Placeholder: return random score for skeleton
        // In full implementation, this would:
        // 1. Sample a minibatch from data
        // 2. Run the prompt with each example's inputs
        // 3. Compare outputs to expected outputs
        // 4. Return accuracy/metric score
        0.5
    }

    /// Runs the optimization process.
    ///
    /// This is the main entry point for prompt optimization. It implements
    /// a simplified version of the MIPRO algorithm.
    ///
    /// # Arguments
    ///
    /// * `signature` - The signature defining the task
    /// * `training_data` - Examples for optimization and evaluation
    ///
    /// # Returns
    ///
    /// An `OptimizationResult` containing the best prompt and metrics.
    pub fn optimize(
        &self,
        signature: &Signature,
        training_data: &[TrainingExample],
    ) -> Result<OptimizationResult> {
        use std::time::Instant;

        // Validate training data
        self.validate_training_data(signature, training_data)?;

        let start_time = Instant::now();
        let mut trials_history = Vec::new();
        let mut best_score = f64::NEG_INFINITY;
        let mut best_prompt = String::new();
        let mut best_instruction = String::new();
        let mut best_demos = Vec::new();
        let mut score_history = Vec::new();
        let mut trials_without_improvement = 0;

        // Generate instruction candidates
        let instructions = self.generate_instruction_candidates(signature);

        // Select demonstrations
        let all_demos =
            self.select_demonstrations(signature, training_data, self.config.max_labeled_demos);

        // Main optimization loop
        for trial in 0..self.config.num_trials {
            let trial_start = Instant::now();

            // Select instruction (cycle through candidates)
            let instruction = &instructions[trial % instructions.len()];

            // Select demo subset
            let num_demos = (trial % (self.config.max_bootstrapped_demos + 1)).min(all_demos.len());
            let demos: Vec<Demonstration> = all_demos.iter().take(num_demos).cloned().collect();

            // Build and evaluate prompt
            let prompt = self.build_prompt(signature, instruction, &demos);
            let score = self.evaluate_prompt(&prompt, signature, training_data);

            let trial_duration = trial_start.elapsed().as_millis() as u64;

            // Record trial
            let trial_result = TrialResult {
                trial_number: trial,
                prompt: prompt.clone(),
                instruction: instruction.clone(),
                num_demos,
                score,
                duration_ms: trial_duration,
                metadata: HashMap::new(),
            };
            trials_history.push(trial_result);
            score_history.push(score);

            // Update best if improved
            if score > best_score + self.config.min_improvement_threshold {
                best_score = score;
                best_prompt = prompt;
                best_instruction = instruction.clone();
                best_demos = demos;
                trials_without_improvement = 0;

                if self.config.verbose {
                    eprintln!(
                        "[Trial {}/{}] New best score: {:.4}",
                        trial + 1,
                        self.config.num_trials,
                        best_score
                    );
                }
            } else {
                trials_without_improvement += 1;
            }

            // Early stopping check
            if let Some(patience) = self.config.early_stopping_patience {
                if trials_without_improvement >= patience {
                    if self.config.verbose {
                        eprintln!(
                            "[Trial {}/{}] Early stopping: no improvement for {} trials",
                            trial + 1,
                            self.config.num_trials,
                            patience
                        );
                    }
                    break;
                }
            }
        }

        let total_duration = start_time.elapsed().as_millis() as u64;
        let initial_score = score_history.first().copied().unwrap_or(0.0);

        let metrics = OptimizationMetrics {
            total_duration_ms: total_duration,
            trials_completed: trials_history.len(),
            llm_calls: trials_history.len(), // Placeholder
            tokens_used: 0,                  // Placeholder
            improvement: best_score - initial_score,
            early_stopped: trials_without_improvement
                >= self.config.early_stopping_patience.unwrap_or(usize::MAX),
            score_history,
        };

        Ok(OptimizationResult {
            best_prompt,
            best_instruction,
            best_demos,
            best_score,
            trials_history,
            metrics,
            signature: signature.clone(),
            config: self.config.clone(),
        })
    }
}

impl Default for PromptOptimizer {
    fn default() -> Self {
        Self::with_defaults()
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_field_type_parsing() {
        assert_eq!("str".parse::<FieldType>().unwrap(), FieldType::String);
        assert_eq!("int".parse::<FieldType>().unwrap(), FieldType::Integer);
        assert_eq!("float".parse::<FieldType>().unwrap(), FieldType::Float);
        assert_eq!("bool".parse::<FieldType>().unwrap(), FieldType::Boolean);
        assert_eq!("json".parse::<FieldType>().unwrap(), FieldType::Json);
        assert_eq!(
            "list".parse::<FieldType>().unwrap(),
            FieldType::List(Box::new(FieldType::String))
        );
        assert_eq!(
            "list[int]".parse::<FieldType>().unwrap(),
            FieldType::List(Box::new(FieldType::Integer))
        );
    }

    #[test]
    fn test_field_type_display() {
        assert_eq!(FieldType::String.to_string(), "str");
        assert_eq!(FieldType::Integer.to_string(), "int");
        assert_eq!(
            FieldType::List(Box::new(FieldType::String)).to_string(),
            "list"
        );
        assert_eq!(
            FieldType::List(Box::new(FieldType::Float)).to_string(),
            "list[float]"
        );
    }

    #[test]
    fn test_signature_field_creation() {
        let field = SignatureField::new("user_query", FieldType::String);
        assert_eq!(field.name, "user_query");
        assert_eq!(field.prefix, "User Query:");
        assert_eq!(field.field_type, FieldType::String);
        assert!(field.description.is_none());
    }

    #[test]
    fn test_signature_field_parsing() {
        let field = SignatureField::parse("query:str").unwrap();
        assert_eq!(field.name, "query");
        assert_eq!(field.field_type, FieldType::String);

        let field = SignatureField::parse("count:int").unwrap();
        assert_eq!(field.name, "count");
        assert_eq!(field.field_type, FieldType::Integer);

        let field = SignatureField::parse("items:list[str]").unwrap();
        assert_eq!(field.name, "items");
        assert_eq!(
            field.field_type,
            FieldType::List(Box::new(FieldType::String))
        );
    }

    #[test]
    fn test_simple_signature_parsing() {
        let sig = Signature::parse("question -> answer").unwrap();
        assert_eq!(sig.input_fields.len(), 1);
        assert_eq!(sig.output_fields.len(), 1);
        assert_eq!(sig.input_fields[0].name, "question");
        assert_eq!(sig.output_fields[0].name, "answer");
    }

    #[test]
    fn test_multi_field_signature_parsing() {
        let sig = Signature::parse("question, context -> answer").unwrap();
        assert_eq!(sig.input_fields.len(), 2);
        assert_eq!(sig.output_fields.len(), 1);
        assert_eq!(sig.input_fields[0].name, "question");
        assert_eq!(sig.input_fields[1].name, "context");
    }

    #[test]
    fn test_typed_signature_parsing() {
        let sig =
            Signature::parse("query:str, docs:list -> summary:str, confidence:float").unwrap();
        assert_eq!(sig.input_fields.len(), 2);
        assert_eq!(sig.output_fields.len(), 2);
        assert_eq!(sig.input_fields[0].field_type, FieldType::String);
        assert_eq!(
            sig.input_fields[1].field_type,
            FieldType::List(Box::new(FieldType::String))
        );
        assert_eq!(sig.output_fields[0].field_type, FieldType::String);
        assert_eq!(sig.output_fields[1].field_type, FieldType::Float);
    }

    #[test]
    fn test_signature_spec_string() {
        let sig = Signature::parse("question, context -> answer").unwrap();
        assert_eq!(sig.to_spec_string(), "question, context -> answer");

        let sig = Signature::parse("query:str -> summary:str, score:float").unwrap();
        assert_eq!(sig.to_spec_string(), "query -> summary, score:float");
    }

    #[test]
    fn test_signature_with_instructions() {
        let sig = Signature::parse("question -> answer")
            .unwrap()
            .with_instructions("Answer the question accurately and concisely.");

        assert_eq!(
            sig.instructions,
            "Answer the question accurately and concisely."
        );
    }

    #[test]
    fn test_invalid_signature() {
        assert!(Signature::parse("no arrow here").is_err());
        assert!(Signature::parse("-> only output").is_err());
        assert!(Signature::parse("only input ->").is_err());
    }

    #[test]
    fn test_optimizer_config_defaults() {
        let config = OptimizerConfig::default();
        assert_eq!(config.num_trials, 20);
        assert_eq!(config.minibatch_size, 25);
        assert_eq!(config.max_bootstrapped_demos, 4);
        assert_eq!(config.max_labeled_demos, 16);
        assert_eq!(config.num_instruction_candidates, 10);
    }

    #[test]
    fn test_optimizer_config_builder() {
        let config = OptimizerConfig::default()
            .with_num_trials(50)
            .with_minibatch_size(32)
            .with_seed(42)
            .with_verbose(true);

        assert_eq!(config.num_trials, 50);
        assert_eq!(config.minibatch_size, 32);
        assert_eq!(config.seed, Some(42));
        assert!(config.verbose);
    }

    #[test]
    fn test_demonstration_creation() {
        let demo = Demonstration::new()
            .with_input("question", "What is Rust?")
            .with_output("answer", "A systems programming language.")
            .with_source("labeled");

        assert_eq!(demo.inputs.get("question").unwrap(), "What is Rust?");
        assert_eq!(
            demo.outputs.get("answer").unwrap(),
            "A systems programming language."
        );
        assert_eq!(demo.source, Some("labeled".to_string()));
    }

    #[test]
    fn test_training_example_creation() {
        let example = TrainingExample::new()
            .with_input("question", "What is 2+2?")
            .with_expected("answer", "4")
            .with_weight(1.0);

        assert_eq!(example.inputs.get("question").unwrap(), "What is 2+2?");
        assert_eq!(example.expected_outputs.get("answer").unwrap(), "4");
        assert_eq!(example.weight, Some(1.0));
    }

    #[test]
    fn test_prompt_optimizer_creation() {
        let optimizer = PromptOptimizer::with_defaults();
        assert_eq!(optimizer.config().num_trials, 20);

        let custom_config = OptimizerConfig::default().with_num_trials(100);
        let optimizer = PromptOptimizer::new(custom_config);
        assert_eq!(optimizer.config().num_trials, 100);
    }

    #[test]
    fn test_prompt_optimizer_validation() {
        let optimizer = PromptOptimizer::with_defaults();
        let signature = Signature::parse("question -> answer").unwrap();

        // Empty data should fail
        let result = optimizer.optimize(&signature, &[]);
        assert!(result.is_err());

        // Create sufficient training data
        let training_data: Vec<TrainingExample> = (0..30)
            .map(|i| {
                TrainingExample::new()
                    .with_input("question", format!("Question {i}"))
                    .with_expected("answer", format!("Answer {i}"))
            })
            .collect();

        let result = optimizer.optimize(&signature, &training_data);
        assert!(result.is_ok());
    }

    #[test]
    fn test_demonstration_formatting() {
        let sig = Signature::parse("question -> answer").unwrap();
        let demo = Demonstration::new()
            .with_input("question", "What is Rust?")
            .with_output("answer", "A programming language.");

        let formatted = demo.format(&sig);
        assert!(formatted.contains("Question:"));
        assert!(formatted.contains("What is Rust?"));
        assert!(formatted.contains("Answer:"));
        assert!(formatted.contains("A programming language."));
    }

    #[test]
    fn test_signature_prompt_template() {
        let sig = Signature::parse("question, context -> answer")
            .unwrap()
            .with_instructions("Answer based on the context.");

        let template = sig.to_prompt_template();
        assert!(template.contains("Answer based on the context."));
        assert!(template.contains("Question:"));
        assert!(template.contains("Context:"));
        assert!(template.contains("Answer:"));
    }
}