include("./game.jl")

using Distributions: Dirichlet
using LinearAlgebra: normalize!
using StatsBase
using Sockets

## MCTS setting

const Setting = (
        MAX_SIMULATION = 800,
        T              = 1.,
        C              = √3,
        ϵ              = 0.35,
        α              = 0.3,
        max_parall     = 100,
        PORT           = 8877
        )

function makeClients()
    Clients = Channel{TCPSocket}(Setting.max_parall)
    @simd for i=1:Setting.max_parall
        client = connect(Setting.PORT)
        put!(Clients, client)
    end
    Clients
end


function closeClients(Clients::Channel{TCPSocket})
    while isready(Clients)
        client = take!(Clients)
        close(client)
    end
    nothing
end

mutable struct Node
    N::Int
    W::Float16
    P::Float16
    vloss::Int
    parent::Union{Node, Nothing}
    children::Dict{Integer, Node}
    action::Union{Int8, String}
    isLeaf::Bool
    function Node(action=false, parent=nothing, P=Float16(0))
        new(Int(0), Float16(0), Float16(P), Int(0), parent, Dict{Integer, Node}(), action, true)
    end
end


function Base.show(io::IO, self::Node)
    "print out the node"
    children = keys(self.children)
    parent = isnothing(self.parent) ? "root" : self.parent.action
    print(io, "Node(action: $(self.action), parent: $parent, $children, ")
    p = self.N >0 ? self.W/self.N : 0
    print(io, "N/W/Q/P/VL: $(self.N)/$(self.W)/$(p)/$(self.P)/$(self.vloss)\n")
end


function bestChild(self::Node)
    best_child = nothing
    best_score = -Inf
    for child in values(self.children)
        Q = child.N > 0 ? (child.W-child.vloss)/child.N : 0
        curr = Q + Setting.C*child.P*√(self.N) / (1+child.N)
        if curr > best_score
            best_score = curr
            best_child = child
        end
    end
    return best_child
end


function add_noise!(probs::Array{Float32}, validmoves::Array{Int8,1})
    " Add Dirichlet noise in the root node "
    dirch = rand(Dirichlet(length(validmoves), Setting.α))
    @inbounds @. probs[validmoves] = probs[validmoves] * (1-Setting.ϵ) + dirch *Setting.ϵ
    nothing
end


function expand(self::Node, probs::Array{Float32}, moves::Array{Int8})
    for child in moves
        self.children[child] = Node(child, self, probs[child])
    end
    self.isLeaf = false
end


mutable struct MCTS
    root::Union{Nothing, Node}
    leaves::Int
    Clients::Channel{TCPSocket}
    function MCTS()
        new(nothing, 1, makeClients())
    end
end


function Base.show(io::IO, self::MCTS)
    "print out the node"
    compact = get(io, :compact, false)
    if !compact
        showroot = isnothing(self.root) ? "nothing" : self.root.action
        print(io, "MCTS(", showroot, ", ", self.leaves, ")")
    end
end


function backpropgation(root::Union{Node, Nothing}, q::Float16)
    while !isnothing(root.parent)
        root.vloss -= 1
        root.W += q
        root = root.parent
    end
end


function predict(game::Game, invaMoves::Array{Int8}, Clients)
    client = take!(Clients)
    println(client, game.state)
    flush(client)
    message = readuntil(client, "\n")
    put!(Clients, client)
    probs, vals = eval(Meta.parse(message))
    !isempty(invaMoves) && @inbounds probs[invaMoves] .= 0.
    normalize!(probs, 1)
    return probs, vals
end


function root_add_noise(game::Game, self::MCTS)
    vaMoves = validMoves(game)
    probs, _ = predict(game, setdiff(ALL_MOVES[1], vaMoves), self.Clients)
    add_noise!(probs, vaMoves)
    expand(self.root, probs, vaMoves)
end


function search!(self::MCTS, initState::String)
    root = self.root
    game = Game(initState)
    # selection
    curr_player = game.nextPlayer
    root.N += 1
    while !root.isLeaf
        root = bestChild(root)
        place(game, root.action)
        root.N += 1; root.vloss +=1
        if !iszero(game.winner)
            backpropgation(root, Float16(game.winner*curr_player))
            break
        end
    end
    # expansion
    if iszero(game.winner)
        vaMoves = validMoves(game)
        #bottleneck
        @inbounds probs, q = predict(game, setdiff(ALL_MOVES[1], vaMoves), self.Clients)
        game = nothing
        self.leaves += -1+length(vaMoves)
        expand(root, probs, vaMoves)
        # backpropgation
        backpropgation(root, q)
    end
end


function (self::MCTS)(game::Game; test::Bool=false, counts::Int = Setting.MAX_SIMULATION)
    "return move, Pi"

    initState = game.state
    if isnothing(self.root)
        self.root = Node(initState)
    elseif @inbounds self.root.action == initState[1:end-1]
        for child in values(self.root.children)
            if @inbounds child.action == int(initState[end])
                self.root = child
                self.root.parent = nothing
                break
            end
        end
    end
    total = 0
    if self.leaves < 9
        root_add_noise(game, self)
        total += 1
    end
    while total < counts
        parall = min(self.leaves, counts-total, Setting.max_parall)
        @sync for i=1:parall
            @async search!(self, initState)
        end
        total += parall
    end

    moves = self.root.children
    Pi = zeros(Float32, 9)
    if test
        # choose the action from the current state with greatest N
        for i in keys(moves)
            @inbounds Pi[i] = moves[i].N
        end
        normalize!(Pi, 1)
        move = argmax(Pi)
    else
        for i in keys(moves)
            @inbounds Pi[i] = moves[i].N ^ (1. / Setting.T)
        end
        normalize!(Pi, 1)
        @inbounds move = sample(ALL_MOVES[1], Weights(Pi))
    end
    # cut off the subtree
    self.root = moves[move]
    self.root.action = "$(game.state)$move"
    self.root.parent = nothing
    return move, Pi
end
