module interface

export runMCTS

using Base.Threads, PrettyPrinting
using ..type, ..mcts, ..util


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


""" Search for the best action to take for a given state and task.

This function runs the MCTS algorithm through multiple iterations of expansion,
simulation, and backpropagation to find optimal decisions.

Does **not** mutate the input state; it creates new MCTS nodes during search.

# Arguments
- `initialstate::T`
  The initial state from which to start the search.
- `transition::Function`
  A function that defines how the state transitions.
- `transitionargs::NamedTuple`
  Arguments passed to the transition function.

# Keyword Arguments
- `horizontalSampleExpansionPhase::Integer=3`
  Number of child states sampled at each node during expansion phase.
- `horizontalSampleSimulationPhase::Integer=3`
  Number of child states sampled at each node during simulation's expansion phase.
- `maxSimulationDepth::Integer=3`
  Maximum depth MCTS goes during simulation phase.
- `maxiterations::Integer=10`
  Number of iterations MCTS performs through expansion → simulation → backpropagation cycles.
- `explorationweight::Number=1.0`
  Exploration weight controls how much MCTS explores new states versus exploiting known states.
  A value of 1.0 balances exploration and exploitation equally. Higher values (e.g., 2.0) 
  encourage more aggressive exploration.
- `earlystop::Union{Function,Nothing}=nothing`
  Optional function to check early stopping condition. If satisfied, MCTS breaks iterations.
- `saveSimulatedNode::Bool=false`
  Whether to save nodes created during simulation phase.
- `multithread::Bool=false`
  Whether to use multithreading during simulation.

# Return
- `NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
             Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}}`
  - `root`: the complete MCTS tree with root node
  - `bestNextState`: the best immediate next state
  - `bestTerminalState`: the best final state along the best trajectory
  - `highValueStateList`: list of high-value terminal states (reward >= 8)

# Example
```jldoctest
julia> using LLMMCTS
julia> initialState = Dict(:reward=>0.0)
julia> result = runMCTS(initialState, transition_func, transition_args; maxiterations=5)
```
"""
function runMCTS(
  initialstate::T,
  transition::Function,
  transitionargs::NamedTuple,
  ;
  horizontalSampleExpansionPhase::Integer=3,
  horizontalSampleSimulationPhase::Integer=3,
  maxSimulationDepth::Integer=3,
  maxiterations::Integer=10,
  explorationweight::Number=1.0,
  earlystop::Union{Function,Nothing}=nothing,
  saveSimulatedNode::Bool=false,
  multithread=false,
  )::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
                Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}} where {T<:Any}
  println("--> LLMMCTS runMCTS 1")
  # Initialize the MCTS tree with a root node representing the initial state
  # root.visits=0: no visits yet
  # root.statevalue=0: no simulation results yet
  root = MCTSNode("root", initialstate, 0, 0, 0, 0, false, nothing, Dict{String,MCTSNode}(),
                  Dict{String,Any}())

  # Channel to collect high-value terminal states (reward >= 8)
  # These are "good solutions" that can be returned to the user
  highValueState = Channel{Any}(100)

  # Main MCTS loop: perform iterations to build the search tree
  # Each iteration: SELECTION → EXPANSION → SIMULATION → BACKPROPAGATION
  for nth in 1:maxiterations
    # Start from root and traverse down using UCT selection
    node = root
    node.visits += 1  # Count this iteration's visit to root
    println("--> LLMMCTS runMCTS 2")
    # Phase 1: SELECTION - Traverse tree using UCT until reaching a leaf node
    # UCT balances exploration (new branches) vs exploitation (promising branches)
    while !isleaf(node)
      println("--> LLMMCTS runMCTS 3")
      node = UCTselect(node, explorationweight)
    end
    println("--> LLMMCTS runMCTS 4")
    # Phase 2: TERMINAL CHECK - If leaf is terminal, just backpropagate
    if node.isterminal
      println("--> LLMMCTS runMCTS 5")
      # If this terminal state has high reward (>= 8), store it for later
      if node.state[:reward] >= 8
        println("--> LLMMCTS runMCTS 6")
        put!(highValueState, deepcopy(node.state))
      end
      println("--> LLMMCTS runMCTS 7")
      # Backpropagate the terminal node's own reward up to root
      # This updates all ancestors with this path's outcome
      backpropagate(node, node.reward)
    else
      println("--> LLMMCTS runMCTS 8")
      # Phase 3: EXPANSION - Generate children for this non-terminal leaf
      # Horizontal sampling: create multiple child nodes via LLM transition
      _ = expand(node, transition, transitionargs;
                horizontalSample=horizontalSampleExpansionPhase,
                multithread=multithread)
      println("--> LLMMCTS runMCTS 9")
      # Phase 4: SIMULATION + BACKPROPAGATION
      # For each newly expanded child, run simulation and update statistics
      if multithread
        println("--> LLMMCTS runMCTS 10")
        # Parallel simulation: spawn threads for each child node
        @sync for (leafNodeKey, leafNode) in node.children
          @spawn simulateThenBackpropagate(leafNode, transition, transitionargs;
                  maxSimulationDepth=maxSimulationDepth, 
                  horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
                  saveSimulatedNode=saveSimulatedNode,
                  multithread=multithread,
                  highValueState=highValueState,
                  )
        end
      else
        println("--> LLMMCTS runMCTS 11")
        # Sequential simulation: process each child one at a time
        for (leafNodeKey, leafNode) in node.children
          println("--> LLMMCTS runMCTS 11-1")
          simulateThenBackpropagate(leafNode, transition, transitionargs;
            maxSimulationDepth=maxSimulationDepth, 
            horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
            saveSimulatedNode=saveSimulatedNode,
            multithread=multithread,
            highValueState=highValueState)
        end
      end
    end
    println("--> LLMMCTS runMCTS 12")
    # Phase 5: EARLY STOP CHECK
    # Optional: stop search early if a condition is met
    if typeof(earlystop) <: Function && earlystop(node.state)
      println("--> LLMMCTS runMCTS 13")
      break
    end
  end
  println("--> LLMMCTS runMCTS 14")
  # After all iterations, extract results from the search tree
  # Select best immediate next state (best child of root)
  bestNextState = selectBestNextNode(root)
  println("--> LLMMCTS runMCTS 15")
  # Select best terminal state along the optimal trajectory
  bestTerminalState = selectBestTrajectoryNode(root)

  # Collect all high-value states from the channel into a list
  highValueStateList = Vector{Dict{String, Any}}()
  while !isempty(highValueState)
    println("--> LLMMCTS runMCTS 16")
    push!(highValueStateList, take!(highValueState))
  end
  println("--> LLMMCTS runMCTS 17")
  # Return complete search results
  result = (
            root=root, 
            bestNextState=bestNextState.state, 
            bestTerminalState=bestTerminalState.state,
            highValueStateList=highValueStateList
            )

  return result
end

""" Run simulation from a given node and backpropagate the reward.

This function performs simulation (rollout) from the given node, collects the
cumulative reward along the trajectory, and backpropagates it up the tree to update
visit counts and state values.

Does **not** mutate the input node's children (unless `saveSimulatedNode=true`).

# Arguments
- `node::MCTSNode`
  The current node to simulate from.
- `transition::Function`
  A function that defines how the state transitions.
- `transitionargs::NamedTuple`
  Arguments passed to the transition function.

# Keyword Arguments
- `maxSimulationDepth::Integer=3`
  Maximum depth MCTS goes during simulation phase.
- `horizontalSampleSimulationPhase::Integer=3`
  Number of child states sampled at each node during simulation phase.
- `saveSimulatedNode::Bool=false`
  Whether to save nodes created during simulation phase. If false, children are
  cleared after simulation.
- `multithread::Bool=false`
  Whether to use multithreading during simulation.

# Return
- `Nothing`

# Signature
"""
function simulateThenBackpropagate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; 
                maxSimulationDepth::Integer=3, horizontalSampleSimulationPhase::Integer=3,
                saveSimulatedNode::Bool=false,
                multithread=false,
                highValueState=Union{Nothing,Any}=nothing)
  println("--> LLMMCTS simulateThenBackpropagate 1")
  # Phase 1: RUN SIMULATION (rollout)
  # Perform a rollout from this node, accumulating rewards along the way
  simTrajectoryReward, terminalstate = 
    simulate(node, transition, transitionargs;
            maxSimulationDepth=maxSimulationDepth, 
            horizontalSample=horizontalSampleSimulationPhase,
            multithread=multithread)
  println("--> LLMMCTS simulateThenBackpropagate 2")
  # Phase 2: HIGH-VALUE STATE TRACKING
  # If we reached a terminal state with high reward (>= 8), store it
  # This allows users to access multiple good solutions, not just the best one
  if highValueState !== nothing && 
    terminalstate !== nothing &&
    terminalstate["reward"] >= 8
    println("--> LLMMCTS simulateThenBackpropagate 3")
    put!(highValueState, deepcopy(terminalstate))
  end
  println("--> LLMMCTS simulateThenBackpropagate 4")
  # Phase 3: BACKPROPAGATE
  # Update statistics (visits, statevalue) for all ancestors up to root
  # The simulation result is now incorporated into the tree
  backpropagate(node, simTrajectoryReward)
  println("--> LLMMCTS simulateThenBackpropagate 5")
  # Phase 4: MEMORY MANAGEMENT
  # Clear children unless user wants to keep them for analysis
  # This frees memory for the next iteration while preserving tree structure
  if saveSimulatedNode == false
    println("--> LLMMCTS simulateThenBackpropagate 6")
    node.children = Dict{String, MCTSNode}()
  end
  println("--> LLMMCTS simulateThenBackpropagate 7")
end












































































end # module interface