""" To implement a Monte Carlo Tree Search (MCTS) algorithm in Julia with the UCT (Upper Confidence 
  Bound for Trees) selection function, you can follow the steps below: Define the necessary types 
  and functions for the MCTS algorithm:
"""

module MCTS
  
# export 

using Dates, UUIDs, DataStructures, JSON3, Random
using GeneralUtils

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

"""

  Arguments\n
  -----
    
  Return\n
  -----

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring
    [] implement the function

  Signature\n
  -----
"""
struct MCTSNode{T}
  state::T
  visits::Int
  total_reward::Float64
  children::Dict{T, MCTSNode}
end

"""

  Arguments\n
  -----
    
  Return\n
  -----

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring
    [WORKING] check child_node.total_reward w/ LATS paper. Which value total_reward representing 

  Signature\n
  -----
"""
function select(node::MCTSNode, c::Float64)
    max_uct = -Inf
    selected_node = nothing

    for (child_state, child_node) in node.children
        uct_value = child_node.total_reward / child_node.visits + 
                    c * sqrt(log(node.visits) / child_node.visits)
        if uct_value > max_uct
            max_uct = uct_value
            selected_node = child_node
        end
    end

    return selected_node
end

"""

  Arguments\n
  -----
    
  Return\n
  -----

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring
    [] implement the function

  Signature\n
  -----
"""
function expand(node::MCTSNode, state::T, actions::Vector{T})
    for action in actions
        new_state = transition(node.state, action)  # Implement your transition function
        if new_state ∉ keys(node.children)
            node.children[new_state] = MCTSNode(new_state, 0, 0.0, Dict{T, MCTSNode}())
        end
    end
end

"""

  Arguments\n
  -----
    
  Return\n
  -----

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring
    [] implement the function

  Signature\n
  -----
"""
function simulate(state::T, max_depth::Int)
    total_reward = 0.0
    for _ in 1:max_depth
        action = select_action(state)  # Implement your action selection function
        state, reward = transition(state, action)  # Implement your transition function
        total_reward += reward
    end
    return total_reward
end

"""

  Arguments\n
  -----
    
  Return\n
  -----

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring
    [] implement the function

  Signature\n
  -----
"""
function backpropagate(node::MCTSNode, reward::Float64)
    node.visits += 1
    node.total_reward += reward
    if !isempty(node.children)
        best_child = argmax([child.total_reward / child.visits for child in values(node.children)])
        backpropagate(node.children[best_child], -reward)
    end
end

"""

  Arguments\n
  -----
    
  Return\n
  -----

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring
    [] implement the function

  Signature\n
  -----
"""
function transition(state, action)

end

""" Check whether a node is a leaf node

  Arguments\n
  -----
    
  Return\n
  -----
    a task represent an agent

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring
    [DONE] implement isLeaf()

  Signature\n
  -----
"""
isLeaf(node::MCTSNode)::Bool = isempty(node.children)

# ------------------------------------------------------------------------------------------------ #
#                    Create a complete example using the defined MCTS functions                    #
# ------------------------------------------------------------------------------------------------ #
"""

  Arguments\n
  -----
    
  Return\n
  -----

  Example\n
  -----
  ```jldoctest
  julia> 
  ```

  TODO\n
  -----
    [] update docstring

  Signature\n
  -----
"""
function run_mcts(initial_state, actions, max_iterations::Int, max_depth::Int, w::Float64)
    root = MCTSNode(initial_state, 0, 0.0, Dict())

    for _ in 1:max_iterations
        node = root
        while !isLeaf(node)
            node = select(node, w)
        end

        expand(node, node.state, actions)

        leaf_node = node.children[node.state]
        reward = simulate(leaf_node.state, max_depth)
        backpropagate(leaf_node, reward)
    end

    best_child_state = argmax([child.total_reward / child.visits for child in values(root.children)])
    return best_child_state
end

# Define your transition function and action selection function here

# Example usage
initial_state = 0
actions = [-1, 0, 1]
best_action = run_mcts(initial_state, actions, 1000, 10, 1.0)
println("Best action to take: ", best_action)









end