module util

export UCTselect

using ..type

# ---------------------------------------------- 100 --------------------------------------------- #

""" Select a node based on UCT (Upper Confidence Bound for Trees) score.

The function computes UCT values for all child nodes and returns the child with the
highest UCT score. The UCT formula balances exploitation (child state value) and
exploration (visit count and parent visit count) using the exploration weight `w`.

Does **not** mutate the input node.

# Arguments
- `node::MCTSNode`
  The MCTS node whose children will be evaluated.
- `w::AbstractFloat`
  Exploration weight. Typical values range from 1.0 to 2.0. A value of 1.0 balances
  exploration and exploitation equally. Higher values (e.g., 2.0) encourage more
  exploration of less-visited nodes.

# Return
- `selectedNode::MCTSNode`
  The child node with the highest UCT score. Returns `nothing` if the node has no
  children (though this would indicate an error since UCTselect is called on non-leaves).

# The UCT Formula

```
UCT(s,a) = Q(s,a) + c * sqrt(ln(N(s)) / N(s,a))

Where:
  Q(s,a) = childNode.statevalue        (exploitation: accumulated reward)
  c = w (explorationweight)            (controls exploration vs exploitation)
  N(s) = node.visits                   (parent visits - total visits to parent)
  N(s,a) = childNode.visits            (child visits - visits to this specific action)
```

# Behavior

| Child visits | Exploration term | Behavior |
|-------------|------------------|----------|
| 0 (never visited) | Undefined | Uses `progressvalue` (LLM heuristic) |
| Low (few visits) | High | Encourages exploring new branches |
| High (many visits) | Near 0 | Exploits known good branches |

# Examples
```jldoctest
julia> using LLMMCTS
julia> child1 = MCTSNode("a", Dict(:reward=>5.0), 0, 10, 50, 0, false, nothing, Dict(), Dict())
julia> child2 = MCTSNode("b", Dict(:reward=>6.0), 0, 5, 30, 0, false, nothing, Dict(), Dict())
julia> parent = MCTSNode("root", Dict(:reward=>0.0), 0, 15, 100, 0, false, nothing,
                           Dict("a"=>child1, "b"=>child2), Dict())
julia> selected = UCTselect(parent, 1.0)
MCTSNode(...)
```
"""
function UCTselect(node::MCTSNode, w::T)::MCTSNode where {T<:AbstractFloat}
  # UCT (Upper Confidence Bound for Trees) selects the best child using:
  # UCT = statevalue + exploration_weight * sqrt(ln(parent_visits) / child_visits)
  #
  # The two terms balance:
  # - Exploitation (statevalue): choose children that performed well in simulations
  # - Exploration (sqrt term): encourage trying less-visited children
  #
  # The exploration weight `w` controls this balance:
  # - w=1.0: equal emphasis on exploration and exploitation
  # - w>1.0: more aggressive exploration (try new branches)
  # - w<1.0: more exploitation (stick with known good branches)
  
  maxUCT = -Inf
  selectedNode = nothing

  for (childState, childNode) in node.children
    # Calculate UCT value for this child
    UCTvalue =
    if childNode.visits != 0  
      # Child has been visited before - use statevalue with exploration bonus
      # Exploration bonus = w * sqrt(ln(parent_visits) / child_visits)
      # High child_visits = small bonus (exploitation dominates)
      # Low child_visits = large bonus (encourages exploration)
      weightedterm = w * sqrt(log(node.visits) / childNode.visits)
      UCTvalue = childNode.statevalue + weightedterm
    else  
      # Child has never been visited - exploration term undefined
      # Fall back to progressvalue (LLM heuristic) as exploitation term
      # This allows LLM guidance to direct early search
      UCTvalue = childNode.progressvalue
    end
                 
    if UCTvalue > maxUCT
      maxUCT = UCTvalue
      selectedNode = childNode
    end
  end

  return selectedNode
end

























































































end # module util