# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

"""
Viewer for the GBMCTS tree.
"""
module GBMCTSView

export TreeView, logstep!

using POMDPs
using ExprSearch, GBMCTS
using RLESUtils, TreeToJSON
using Plots
using Iterators

type TreeView
  film::Animation
  steps::Vector{JDict}
end

TreeView() = TreeView(Animation(), JDict[])

function get_name(s::StateNode)
  n_text = "n=[" * join(s.n, ",") * "]"
  Q_text = "Q=[" * join(map(x->signif(x, 4), s.Q), ",") * "]"
  na_text = "n_actions=$(length(s.action_map))"
  sp_text = "sp_visited=[" * join(map(x->!isnull(x), s.nextstates), ",")  * "]"
  expl_text = "fully_explored=$(s.explored)"
  return join([n_text, Q_text, na_text, sp_text, expl_text], "\\\\")
end

function logstep!(view::TreeView, tree::Dict{UInt64,StateNode}, startstate::State, curstate::State)

  function get_children(s::StateNode)
    valids = filter(x -> !isnull(x), s.nextstates)
    return imap(x -> ("", tree[x]), valids)
  end

  function get_color(s::StateNode)
    color = "steelblue" #default
    if all(s.n .== 0) color = "red" end #completely unvisited
    if s == tree[startstate] color = "grey" end #start state
    if s.explored color = "purple" end #fully explored
    if length(s.n) == 0 color = "black" end #terminal
    if s == tree[curstate] color = "orange" end #current state
    return color
  end

  vc = VisCalls(get_name, get_children)
  d = to_jdict(tree[startstate], vc, color=get_color)
  push!(view.steps, d)

  #animation frame
  #p = plottree(d)
  #push!(view.film, p)

  return d
end


end # module
