scout

defmodule Scout.Sampler.CmaEs do
  @moduledoc """
  CMA-ES (Covariance Matrix Adaptation Evolution Strategy) sampler.
  
  Implements a proper CMA-ES algorithm that maintains and adapts
  a covariance matrix to model parameter correlations, similar to
  Optuna's CmaEsSampler.
  """
  
  alias Scout.Sampler.RandomSearch
  
  def init(opts) do
    %{
      population_size: Map.get(opts, :population_size, nil),  # Auto-calculate if nil
      sigma0: Map.get(opts, :sigma0, 1.0),  # Initial step size
      min_obs: Map.get(opts, :min_obs, 3),
      goal: Map.get(opts, :goal, :minimize),
      # CMA-ES specific parameters
      mean: nil,
      cov: nil,
      sigma: nil,
      pc: nil,  # Evolution path for C
      ps: nil,  # Evolution path for sigma
      weights: nil,
      mueff: nil,
      cc: nil,  # Time constant for C evolution path
      cs: nil,  # Time constant for sigma evolution path
      c1: nil,  # Rank-one update coefficient
      cmu: nil, # Rank-mu update coefficient
      damps: nil,  # Damping for sigma
      chiN: nil,  # Expected length of N(0,I) vector
      eigeneval_freq: 10,
      eigeneval_counter: 0,
      B: nil,  # Eigenvectors
      D: nil,  # Eigenvalues
      generation: 0,
      population: [],
      evaluated: []
    }
  end
  
  def next(space_fun, ix, history, state) do
    spec = space_fun.(ix)
    
    if length(history) < state.min_obs do
      RandomSearch.next(space_fun, ix, history, state)
    else
      # Get parameter info
      param_keys = Map.keys(spec) |> Enum.sort()
      n = length(param_keys)
      
      # Initialize CMA-ES state if needed
      state = if state.mean == nil do
        initialize_cmaes(state, spec, param_keys, history)
      else
        state
      end
      
      # Generate next candidate
      {params, new_state} = if length(state.population) < state.population_size do
        # Generate new population member
        candidate = sample_candidate(state, param_keys, spec)
        population = state.population ++ [candidate]
        
        {candidate, %{state | population: population}}
      else
        # All population evaluated, update CMA-ES
        new_state = update_cmaes(state, param_keys, spec, history)
        
        # Start new generation
        candidate = sample_candidate(new_state, param_keys, spec)
        
        {candidate, %{new_state | 
          population: [candidate],
          evaluated: [],
          generation: new_state.generation + 1
        }}
      end
      
      {params, new_state}
    end
  end
  
  defp initialize_cmaes(state, spec, param_keys, history) do
    n = length(param_keys)
    
    # Population size
    lambda = if state.population_size do
      state.population_size
    else
      4 + floor(3 * :math.log(n))
    end
    
    mu = div(lambda, 2)
    
    # Weights for recombination
    weights = for i <- 1..mu do
      :math.log(mu + 0.5) - :math.log(i)
    end
    
    sum_weights = Enum.sum(weights)
    weights = Enum.map(weights, & &1 / sum_weights)
    
    mueff = 1.0 / Enum.sum(Enum.map(weights, & &1 * &1))
    
    # Learning rates
    cc = (4 + mueff/n) / (n + 4 + 2*mueff/n)
    cs = (mueff + 2) / (n + mueff + 5)
    c1 = 2 / (:math.pow(n + 1.3, 2) + mueff)
    cmu = min(1 - c1, 2 * (mueff - 2 + 1/mueff) / (:math.pow(n + 2, 2) + mueff))
    damps = 1 + 2 * max(0, :math.sqrt((mueff - 1)/(n + 1)) - 1) + cs
    
    # Expected length of N(0,I) vector
    chiN = :math.sqrt(n) * (1 - 1/(4*n) + 1/(21*n*n))
    
    # Initialize mean from best trials
    sorted = case state.goal do
      :minimize -> Enum.sort_by(history, & &1.score)
      _ -> Enum.sort_by(history, & &1.score, :desc)
    end
    
    best_trials = Enum.take(sorted, min(3, length(sorted)))
    
    # Convert to normalized space [0,1]
    mean = compute_initial_mean(best_trials, param_keys, spec)
    
    # Initialize covariance as identity
    cov = identity_matrix(n)
    
    %{state | 
      mean: mean,
      cov: cov,
      sigma: state.sigma0,
      pc: List.duplicate(0.0, n),
      ps: List.duplicate(0.0, n),
      weights: weights,
      mueff: mueff,
      cc: cc,
      cs: cs,
      c1: c1,
      cmu: cmu,
      damps: damps,
      chiN: chiN,
      B: identity_matrix(n),
      D: List.duplicate(1.0, n),
      population_size: lambda
    }
  end
  
  defp compute_initial_mean(trials, param_keys, spec) do
    n_trials = length(trials)
    
    # Average parameters in normalized space
    Enum.map(param_keys, fn k ->
      values = Enum.map(trials, fn t ->
        val = Map.get(t.params, k, 0.0) || 0.0
        normalize_value(val, spec[k])
      end)
      
      Enum.sum(values) / n_trials
    end)
  end
  
  defp normalize_value(val, spec_entry) do
    case spec_entry do
      {:uniform, min, max} -> 
        (val - min) / (max - min)
      {:log_uniform, min, max} ->
        log_min = :math.log(min)
        log_max = :math.log(max)
        (:math.log(val) - log_min) / (log_max - log_min)
      {:int, min, max} ->
        (val - min) / (max - min + 1)
      _ -> 0.5
    end
  end
  
  defp denormalize_value(val, spec_entry) do
    # Clamp to [0,1]
    val = max(0.0, min(1.0, val))
    
    case spec_entry do
      {:uniform, min, max} -> 
        min + val * (max - min)
      {:log_uniform, min, max} ->
        log_min = :math.log(min)
        log_max = :math.log(max)
        :math.exp(log_min + val * (log_max - log_min))
      {:int, min, max} ->
        round(min + val * (max - min))
      _ -> val
    end
  end
  
  defp sample_candidate(state, param_keys, spec) do
    n = length(param_keys)
    
    # Sample from N(0, C)
    z = List.duplicate(0, n) |> Enum.map(fn _ -> :rand.normal() end)
    
    # Transform: x = mean + sigma * B * D * z
    y = matrix_vector_multiply(state[:B], z)
    y = Enum.zip(y, state[:D]) |> Enum.map(fn {yi, di} -> yi * :math.sqrt(di) end)
    
    x = Enum.zip(state.mean, y) 
        |> Enum.map(fn {m, yi} -> m + state.sigma * yi end)
    
    # Convert to parameter space
    Enum.zip(param_keys, x)
    |> Map.new(fn {k, xi} ->
      {k, denormalize_value(xi, spec[k])}
    end)
  end
  
  defp update_cmaes(state, param_keys, spec, history) do
    n = length(param_keys)
    
    # Get fitness values for current population
    pop_with_fitness = Enum.map(state.population, fn params ->
      # Find corresponding trial in history
      trial = Enum.find(history, fn t ->
        params_match?(t.params, params, param_keys)
      end)
      
      fitness = if trial, do: trial.score, else: 999999.0
      {params, fitness}
    end)
    
    # Sort by fitness
    sorted = case state.goal do
      :minimize -> Enum.sort_by(pop_with_fitness, fn {_, f} -> f end)
      _ -> Enum.sort_by(pop_with_fitness, fn {_, f} -> -f end)
    end
    
    # Select mu best
    mu = length(state.weights)
    best = Enum.take(sorted, mu)
    
    # Convert to normalized space
    best_normalized = Enum.map(best, fn {params, _} ->
      Enum.map(param_keys, fn k ->
        normalize_value(params[k], spec[k])
      end)
    end)
    
    # Update mean
    old_mean = state.mean
    new_mean = weighted_mean(best_normalized, state.weights)
    
    # Update evolution paths
    ps = update_ps(state, old_mean, new_mean, n)
    generation = max(1, state.generation)
    hsig = norm(ps) / :math.sqrt(1 - :math.pow(1 - state.cs, 2 * generation)) / state.chiN < 1.4 + 2/(n+1)
    pc = update_pc(state, old_mean, new_mean, hsig, n)
    
    # Update covariance matrix
    cov = update_covariance(state, best_normalized, old_mean, new_mean, pc, hsig, n)
    
    # Update step size
    sigma = update_sigma(state, ps)
    
    # Eigendecomposition (periodically) 
    eigen_result = if rem(state.eigeneval_counter, state.eigeneval_freq) == 0 do
      eigendecompose(cov)
    else
      # Use cached values
      b_cached = state[:B] || identity_matrix(n)
      d_cached = state[:D] || List.duplicate(1.0, n)
      {b_cached, d_cached}
    end
    
    # Extract B and D from eigen_result  
    result = if is_tuple(eigen_result) do
      {elem(eigen_result, 0), elem(eigen_result, 1)}
    else
      # Fallback
      {identity_matrix(n), List.duplicate(1.0, n)}
    end
    
    B = elem(result, 0)
    D = elem(result, 1)
    
    %{state |
      mean: new_mean,
      cov: cov,
      sigma: sigma,
      ps: ps,
      pc: pc,
      B: B,
      D: D,
      eigeneval_counter: state.eigeneval_counter + 1
    }
  end
  
  defp params_match?(params1, params2, keys) do
    Enum.all?(keys, fn k ->
      abs((params1[k] || 0.0) - (params2[k] || 0.0)) < 1.0e-6
    end)
  end
  
  defp weighted_mean(vectors, weights) do
    n = length(hd(vectors))
    
    List.duplicate(0.0, n)
    |> Enum.with_index()
    |> Enum.map(fn {_, i} ->
      Enum.zip(vectors, weights)
      |> Enum.reduce(0.0, fn {vec, w}, acc ->
        acc + w * Enum.at(vec, i)
      end)
    end)
  end
  
  defp update_ps(state, old_mean, new_mean, n) do
    # ps = (1 - cs) * ps + sqrt(cs * (2 - cs) * mueff) * B * (1/D) * B' * (mean - old_mean) / sigma
    
    diff = Enum.zip(new_mean, old_mean) |> Enum.map(fn {n, o} -> (n - o) / state.sigma end)
    
    # Simplified: assume B = I for now
    Enum.zip(state.ps, diff)
    |> Enum.map(fn {ps_i, d_i} ->
      (1 - state.cs) * ps_i + :math.sqrt(state.cs * (2 - state.cs) * state.mueff) * d_i
    end)
  end
  
  defp update_pc(state, old_mean, new_mean, hsig, n) do
    diff = Enum.zip(new_mean, old_mean) |> Enum.map(fn {n, o} -> (n - o) / state.sigma end)
    
    a = if hsig, do: :math.sqrt(state.cc * (2 - state.cc) * state.mueff), else: 0
    
    Enum.zip(state.pc, diff)
    |> Enum.map(fn {pc_i, d_i} ->
      (1 - state.cc) * pc_i + a * d_i
    end)
  end
  
  defp update_covariance(state, best_normalized, old_mean, new_mean, pc, hsig, n) do
    # Rank-one update
    c1a = if hsig, do: state.c1, else: 0
    
    pc_outer = outer_product(pc, pc)
    
    # Rank-mu update
    y_vectors = Enum.map(best_normalized, fn xi ->
      Enum.zip(xi, old_mean) |> Enum.map(fn {x, m} -> (x - m) / state.sigma end)
    end)
    
    # C = (1 - c1 - cmu) * C + c1 * pc * pc' + cmu * sum(wi * yi * yi')
    
    # Start with scaled old covariance
    new_cov = scale_matrix(state.cov, 1 - state.c1 - state.cmu)
    
    # Add rank-one update
    new_cov = add_matrices(new_cov, scale_matrix(pc_outer, c1a))
    
    # Add rank-mu update
    Enum.zip(y_vectors, state.weights)
    |> Enum.reduce(new_cov, fn {yi, wi}, acc ->
      add_matrices(acc, scale_matrix(outer_product(yi, yi), state.cmu * wi))
    end)
  end
  
  defp update_sigma(state, ps) do
    # sigma = sigma * exp((cs/damps) * (||ps||/chiN - 1))
    ps_norm = norm(ps)
    state.sigma * :math.exp((state.cs/state.damps) * (ps_norm/state.chiN - 1))
  end
  
  defp norm(vector) do
    :math.sqrt(Enum.sum(Enum.map(vector, & &1 * &1)))
  end
  
  defp identity_matrix(n) do
    for i <- 0..(n-1) do
      for j <- 0..(n-1) do
        if i == j, do: 1.0, else: 0.0
      end
    end
  end
  
  defp outer_product(v1, v2) do
    for x <- v1 do
      for y <- v2, do: x * y
    end
  end
  
  defp scale_matrix(matrix, scalar) do
    Enum.map(matrix, fn row ->
      Enum.map(row, & &1 * scalar)
    end)
  end
  
  defp add_matrices(m1, m2) do
    Enum.zip(m1, m2)
    |> Enum.map(fn {row1, row2} ->
      Enum.zip(row1, row2)
      |> Enum.map(fn {a, b} -> a + b end)
    end)
  end
  
  defp matrix_vector_multiply(matrix, vector) do
    Enum.map(matrix, fn row ->
      Enum.zip(row, vector)
      |> Enum.reduce(0.0, fn {m, v}, acc -> acc + m * v end)
    end)
  end
  
  defp eigendecompose(matrix) do
    # Simplified: return identity matrix and ones for eigenvalues
    # In production, would use proper eigendecomposition
    n = if is_list(matrix) and length(matrix) > 0 do
      length(matrix)
    else
      2  # Default for 2D problems
    end
    
    b_matrix = identity_matrix(n)
    d_vector = List.duplicate(1.0, n)
    {b_matrix, d_vector}
  end
end