dspex

defmodule DSPex.Optimizers.MIPROv2 do
  @moduledoc """
  MIPROv2 - Enhanced Multi-Instruction Proposal Optimizer.

  An improved version of MIPRO with better performance and additional features
  for instruction optimization.
  """

  alias DSPex.Utils.ID

  @doc """
  Optimize a program using MIPROv2.

  ## Examples

      {:ok, program_id} = DSPex.Modules.ChainOfThought.create("question -> answer")
      
      trainset = [
        %{question: "Explain quantum computing", answer: "..."},
        # ... more examples
      ]
      
      {:ok, optimized} = DSPex.Optimizers.MIPROv2.optimize(
        program_id,
        trainset,
        num_candidates: 15,
        minibatch_size: 50
      )
  """
  def optimize(program_id, trainset, opts \\ []) do
    optimizer_id = ID.generate("mipro_v2")

    config = %{
      prompt_model: opts[:prompt_model],
      task_model: opts[:task_model],
      metric: opts[:metric],
      num_candidates: opts[:num_candidates] || 10,
      init_temperature: opts[:init_temperature] || 1.0,
      verbose: opts[:verbose] || false,
      track_stats: opts[:track_stats] || true,
      view_data_batch_size: opts[:view_data_batch_size],
      minibatch_size: opts[:minibatch_size] || 50,
      minibatch_full_eval_steps: opts[:minibatch_full_eval_steps] || 10,
      minibatch: opts[:minibatch] || true,
      num_trials: opts[:num_trials] || 100,
      max_bootstrapped_demos: opts[:max_bootstrapped_demos] || 3,
      max_labeled_demos: opts[:max_labeled_demos] || 5,
      eval_kwargs: opts[:eval_kwargs] || %{},
      seed: opts[:seed] || 123,
      program_aware_proposer: opts[:program_aware_proposer] || true,
      data_aware_proposer: opts[:data_aware_proposer] || true,
      fewshot_aware_proposer: opts[:fewshot_aware_proposer] || true
    }

    # Create the optimizer
    with {:ok, _} <-
           Snakepit.Python.call(
             "dspy.MIPROv2",
             config,
             Keyword.merge([store_as: optimizer_id], opts)
           ),
         # Compile the program
         {:ok, result} <-
           Snakepit.Python.call(
             "stored.#{optimizer_id}.compile",
             %{
               student: "stored.#{program_id}",
               trainset: trainset,
               valset: opts[:valset],
               num_trials: config.num_trials,
               minibatch_size: config.minibatch_size,
               program_aware_proposer: config.program_aware_proposer,
               data_aware_proposer: config.data_aware_proposer,
               fewshot_aware_proposer: config.fewshot_aware_proposer
             },
             opts
           ) do
      {:ok,
       %{
         optimized_program_id: "#{program_id}_mipro_v2",
         optimizer_id: optimizer_id,
         result: result
       }}
    end
  end

  @doc """
  Get detailed optimization report from MIPROv2.
  """
  def get_report(optimizer_id, opts \\ []) do
    Snakepit.Python.call(
      "stored.#{optimizer_id}.get_optimization_report",
      %{},
      opts
    )
  end
end