module mcts
  
export selectBestNextNode, selectBestTrajectoryNode, backpropagate, isleaf, isroot, selectChildNode,
        expand, simulate
using Base.Threads
using GeneralUtils

using ..type

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


""" Select the best next node based on the highest value metric

# Arguments
  - `node::MCTSNode`
    node of a search tree to evaluate

# Return
  - `childNode::MCTSNode`
    the child node with highest value based on either:
    - statevalue/visits ratio if any nodes have non-zero statevalue
    - progressvalue + reward otherwise
"""
function selectBestNextNode(node::MCTSNode)::MCTSNode
  highestProgressValue = -1
  nodekey = nothing

  # Calculate sum of statevalues across all child nodes
  stateValueSum = sum([v.statevalue for (k, v) in node.children])

  # If any nodes have non-zero statevalue, use statevalue/visits as selection metric
  if stateValueSum != 0
    for (k, childnode) in node.children
      # Calculate average statevalue per visit
      potential = childnode.statevalue / childnode.visits

      if potential > highestProgressValue
        highestProgressValue = potential
        nodekey = childnode.nodekey
      end
    end
  else
    # Otherwise use progressvalue + reward as selection metric
    for (k, childnode) in node.children
      potential = childnode.progressvalue + childnode.reward

      if potential > highestProgressValue
        highestProgressValue = potential
        nodekey = childnode.nodekey
      end
    end
  end

  return node.children[nodekey]
end


""" Select the best trajectory node based on the highest reward

# Arguments
  - `node::MCTSNode`
    node of a search tree to evaluate

# Return
  - `childNode::MCTSNode`
    the highest value child node found by traversing down the tree using selectBestNextNode
    until reaching a leaf node

# Signature
"""
function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
  while !isleaf(node)
    node = selectBestNextNode(node)
  end

  return node
end


""" Backpropagate reward along the simulation chain

# Arguments
  - `node::MCTSNode`
    leaf node of a search tree
  - `simTrajectoryReward::T`
    total reward from trajectory simulation
  - `discountRewardCoeff::AbstractFloat`
    A discount reward coefficient to reduce future reward. The futher in the future the lower
    reward it is now.
  
# Return
  - `Nothing`
    This function modifies the nodes in place and returns nothing

# Signature
"""
function backpropagate(node::MCTSNode, simTrajectoryReward::T;
                      discountRewardCoeff::AbstractFloat=0.9) where {T<:Number}
  while !isroot(node)
    # Update the statistics of the current node based on the result of the playout
    node.visits += 1  # Increment visit count for this node
    node.statevalue += ((node.statevalue * (node.visits-1)) + simTrajectoryReward) / node.visits  # Update running average of state value
    simTrajectoryReward *= discountRewardCoeff  # discount because future reward is uncertain
    node = node.parent  # Move up to parent node for next iteration
  end
end

""" Determine whether a node is a leaf node of a search tree.

# Arguments
  - `node::MCTSNode`
    a search tree node

# Return
  - `result::Bool`
    true if it is a leaf node (has no children), false otherwise.

# Example
```jldoctest
julia> using Revise
julia> using YiemAgent, DataStructures
julia> initialState = Dict{String, Any}(
           "customerinfo"=> Dict{String, Any}(),
           "storeinfo"=> Dict{String, Any}(),  

           "thoughtHistory"=> OrderedDict{String, Any}(
             "question"=> "How are you?",
             )  
           )
julia> statetype = typeof(initialState)
julia> root = YiemAgent.MCTSNode(initialState, 0, 0.0, Dict{statetype, YiemAgent.MCTSNode}())
julia> YiemAgent.isleaf(root)
true
```

# Signature
"""
isleaf(node::MCTSNode)::Bool = isempty(node.children)

""" Determine wheter a given node is a root node

# Arguments
  - `node::MCTSNode`
    node of a search tree

# Return
  - `isrootnode::Bool`
    true if the given node is root node, false otherwise

# Signature
"""
isroot(node::MCTSNode)::Bool = node.nodekey == "root" ? true : false



""" Select child node based on the highest statevalue

# Arguments
  - `node::MCTSNode`
    node of a search tree

# Return
  - `childNode::MCTSNode`
    the highest value child node

# Signature
"""
function selectChildNode(node::MCTSNode)::MCTSNode
  highestProgressValue = -1
  nodekey = nothing

  # loop thought node children dictionary to find the highest progress value
  for (k, childNode) in node.children
    potential = childNode.progressvalue + childNode.reward
    if potential > highestProgressValue
      highestProgressValue = potential
      nodekey = childNode.nodekey
    end
  end

  return node.children[nodekey]
end


""" Expand selected node.

# Arguments
  - `node::MCTSNode`
    MCTS node to expand
  - `transition::Function`
    A function that handles state transition.
  - `transitionargs::NamedTuple`
    Arguments for transition()

# Keyword Arguments
  - `horizontalSample::Integer`
    Total number to sample from the current node (i.e. expand new node horizontally). Defaults to 3.
  - `multithread::Bool`
    Whether to run expansion in parallel using multiple threads. Defaults to false.
    
# Return
  - None

# Signature
"""
function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple; 
                horizontalSample::Integer=3, multithread=false)
  if multithread
    @sync for i in 1:horizontalSample 
      @spawn _expand(node, transition, transitionargs)
    end
  else
    for i in 1:horizontalSample
      _expand(node, transition, transitionargs)
    end
  end
end

""" Helper function to expand a single child node.

# Arguments
  - `node::MCTSNode`
    Parent MCTS node to expand from
  - `transition::Function`
    A function that handles state transition
  - `transitionargs::NamedTuple`
    Arguments for transition()

# Return
  - None

# Signature
"""
function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
    result = transition(node.state, transitionargs)
    newNodeKey::AbstractString = result[:newNodeKey]
    newstate::AbstractDict = result[:newstate]
    progressvalue::Integer = result[:progressvalue]

    """
    [] newNodeKey ∉ keys(node.children). 
          New state may have semantic vector close enought to 
        one of existing child state. Which can be assume that they are the same state 
        semantically-wise i.e. De javu. This could be used to recall lessons for this 
        similar situation to improve decisionMaker and evaluator. 
    """
    if newNodeKey ∉ keys(node.children)
      newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate[:reward], 
          newstate[:isterminal], node, Dict{String, MCTSNode}(), Dict{String, Any}())
      node.children[newNodeKey] = newNode
    end
end

""" Simulate interactions between agent and environment

# Arguments
  - `node::MCTSNode`
    node that will be a simulation starting point.
  - `transition::Function`
    A user function that handles how state transition.
  - `transitionargs::NamedTuple`
    Arguments for everything the user will use within transition().
  - `maxSimulationDepth::Integer`
    maximum depth level MCTS goes vertically during simulation.
  - `horizontalSample::Integer`
    Total number to sample from the current node (i.e. expand new node horizontally)

# Keyword Arguments
  - `multithread::Bool`
    Whether to run expansion in parallel using multiple threads. Defaults to false.
    
# Return
  - `simTrajectoryReward::Number`
    Cumulative reward collected along the simulation trajectory
  - `terminalstate::Union{Dict{String, Any}, Nothing}`
    Final state if terminal state reached, nothing otherwise

# Signature
"""
function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; 
  maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
  )::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}

  simTrajectoryReward = 0.0
  terminalstate = nothing
  
  for depth in 1:maxSimulationDepth
    simTrajectoryReward += node.reward
    if node.isterminal
      terminalstate = node.state
      break
    else
      _ = expand(node, transition, transitionargs;
                horizontalSample=horizontalSample,
                multithread=multithread)
      node = selectChildNode(node)
    end 
  end

  return (simTrajectoryReward=simTrajectoryReward, 
          terminalstate=terminalstate)
end

# """ Make new state

# # Arguments
#   - `currentstate::T1`
#     Current state dictionary containing thought history and metadata
#   - `thoughtDict::T4`
#     Dictionary containing new thought and action
#   - `response::T2`
#     Response string from the environment
#   - `select::Union{T3, Nothing}`
#     Selection value or nothing
#   - `reward::T3`
#     Reward value for this state
#   - `isterminal::Bool`
#     Whether this state is terminal

# # Return
#   - `Tuple{String, Dict{String, <:Any}}`
#     A tuple containing:
#     - A unique node key string
#     - A new state dictionary with updated thought history and metadata

# # Example
# ```jldoctest
# julia> 
# ```

# # Signature
# """
# function makeNewState(currentstate::T1, thoughtDict::T4, response::T2, select::Union{T3, Nothing}, 
#   reward::T3, isterminal::Bool
#   )::Tuple{String, Dict{String, <:Any}} where {T1<:AbstractDict, T2<:AbstractString, T3<:Number, T4<:AbstractDict}

#   # Find the latest thought key and index from current state's thought history
#   currentstate_latestThoughtKey, currentstate_latestThoughtIndice = 
#       GeneralUtils.findHighestIndexKey(currentstate[:thoughtHistory], "thought")
#   # Calculate next index for new thought/action
#   currentstate_nextIndice = 
#               currentstate_latestThoughtKey == :NA ? 1 : currentstate_latestThoughtIndice + 1
#   # Create new keys for thought and action based on next index
#   currentstate_latestThoughtKey = Symbol("thought_$currentstate_nextIndice")
#   latestActionKey = Symbol("action_$currentstate_nextIndice")

#   # Find the latest thought index from input thought dictionary
#   _, thoughtDict_latestThoughtIndice = 
#       GeneralUtils.findHighestIndexKey(thoughtDict, "thought")
  
#   # Determine thought and action keys from thought dictionary
#   thoughtDict_latestThoughtKey, thoughtDict_latestActionKey =
#   if thoughtDict_latestThoughtIndice == -1
#     (:thought, :action)
#   else
#     (
#       Symbol("thought_$thoughtDict_latestThoughtIndice"), 
#       Symbol("action_$thoughtDict_latestThoughtIndice"),
#     )
#   end

#   # Create new state by deep copying current state
#   newstate = deepcopy(currentstate)
#   # Update thought history with new thought
#   newstate[:thoughtHistory][currentstate_latestThoughtKey] = 
#       thoughtDict[thoughtDict_latestThoughtKey]
#   # Update thought history with new action
#   newstate[:thoughtHistory][latestActionKey] = thoughtDict[thoughtDict_latestActionKey]
#   # Create and add new observation to thought history
#   newObservationKey = Symbol("observation_$(currentstate_nextIndice)")
#   newstate[:thoughtHistory][newObservationKey] = response
#   # Update state metadata
#   newstate[:reward] = reward
#   newstate[:select] = select
#   newstate[:isterminal] = isterminal

#   # Generate unique ID for new node
#   newNodeKey = GeneralUtils.uuid4snakecase()

#   return (newNodeKey, newstate)
# end

































































end # module mcts