__precompile__()
module Multilane

import StatsBase: WeightVec, sample

using POMDPs
import POMDPs: actions, discount, isterminal, iterator
import POMDPs: rand, reward
import POMDPs: solve, action
import POMDPs: update, initialize_belief
#using GenerativeModels
#import GenerativeModels: generate_s, generate_sr, initial_state, generate_o, generate_sor

import Distributions: Dirichlet, Exponential, Gamma, rand

import Iterators

import Base: ==, hash, length, vec, +, -, *, .*, ^, .^, .-, /, sqrt, zero, abs, max

# import POMDPToolbox: Particle, ParticleBelief

using DataFrames
using DataArrays
using ProgressMeter
using PmapProgressMeter
using POMDPToolbox
using ParticleFilters

import MCTS # so that we can define node_tag, etc.
# using RobustMCTS # for RobustMDP
import BasicPOMCP # for particle filter

# using Reel
# using AutomotiveDrivingModels
# using AutoViz

import Mustache
import JLD

using Parameters

include("debug.jl")

# package code goes here
export
    PhysicalParam,
    CarState,
    MLState,
    MLAction,
    MLObs,
    BehaviorModel,
    MLMDP,
    MLPOMDP,
    OriginalMDP,
    OriginalRewardModel,
    IDMMOBILModel,
    IDMParam,
    MOBILParam,
    IDMMOBILBehavior,
    MLPhysicalState,
    CarPhysicalState,
    RobustNoCrashMDP,
    FixedBehaviorNoCrashMDP,
    StochasticBehaviorNoCrashMDP,
    RobustMLSolver,
    RobustMLPolicy,
    MLPOMDPSolver,
    MLPOMDPAgent,
    DiscreteBehaviorSet,
    DiscreteBehaviorBelief,
    WeightUpdateParams,
    UniformIDMMOBIL,
    standard_uniform,
    MLMPCSolver,
    MLPOMDPSolver

export
    NoCrashRewardModel,
    NoCrashIDMMOBILModel,
    NoCrashMDP,
    NoCrashPOMDP,
    Simple,
    SimpleSolver,
    BehaviorSolver,
    IDMLaneSeekingSolver

export
    SingleBehaviorSolver,
    SingleBehaviorPolicy,
    single_behavior_state

export
    get_neighborhood, #testing VVV
    get_dv_ds,
    is_lanechange_dangerous,
    get_rear_accel, #testing ^^^
    get_idm_dv,
    get_mobil_lane_change,
    is_crash,
    detect_braking,
    max_braking,
    braking_ids

export #data structure stuff
    ste,
    test_run,
    evaluate,
    merge_results!,
    rerun,
    assign_keys,
    test_run_return_policy,
    lambda,
    TestSet,
    gen_initials,
    run_simulations,
    fill_stats!,
    sbatch_spawn,
    gather_results,
    relaxed_initial_state,
    nan

export # POMDP belief stuff
    ParticleUpdater,
    create_belief,
    update,
    rand,
    sample,
    initialize_belief,
    BehaviorRootUpdater,
    BehaviorRootUpdaterStub,
    AggressivenessBelief,
    AggressivenessUpdater,
    agg_means,
    agg_stds,
    aggressiveness,
    BehaviorParticleUpdater,
    BehaviorParticleBelief,
    ParticleGenerator,
    param_means,
    param_stds

export
    MaxBrakeMetric,
    NumBehaviorBrakesMetric

export
    GaussianCopula

export
    include_visualization,
    #generate_sr,
    create_state,
    NoCrashProblem,

    HardBrakeOverlay,
    InfoOverlay,
    CarIDOverlay,
    CarVelOverlay,
    gen_initials_cz,
    relaxed_initial_state_cz,
    relaxed_initial_state_cz2,
    test_run_cz,
    simulate_cz,
    MLMPCAgent,
    BehaviorGenerator,
    standard_uniform,
    relaxed_initial_state_TRI,
    relaxed_initial_state_cz_control,
    get_traffic_density,
    relaxed_initial_state_cz_emergency,
    EmergencyRewardModel,
    judge_collision

include("triangular.jl")
include("physical.jl")
include("MDP_types.jl")
include("crash.jl")
include("IDM.jl")
include("MOBIL.jl")
include("behavior.jl")
include("copula.jl")
include("behavior_gen.jl")
include("success_model.jl")
include("no_crash_model.jl")

include("metrics.jl")
include("evaluation.jl")
include("test_sets.jl")
include("relax.jl")
include("robust_mdp.jl")
include("pomdp_glue.jl")
include("most_likely_mpc.jl")
include("single_behavior.jl")

#include("heuristics.jl")

include("beliefs.jl")
include("aggressiveness_particle_filter.jl")
include("Traffic_Simulation_TRI.jl")
#=
include("uniform_particle_filter.jl")
include("heuristics.jl")
include("tree_vis.jl")
include("sherlock.jl")
=#
include_visualization() = include(joinpath(Pkg.dir("Multilane"),"src","visualization.jl"))

#if gethostname() == "Theresa"
kkkk = 1
if kkkk == 1
    println("Automatically loading visualization components.")
    include("visualization.jl")
    export
        display_sim,
        show_state,
        show_sim,
        display_sim,
        save_frame,
        visualize,
        visualize_cz
end

end # module


struct ALConstraintSet{T}
    constraints::Vector{ALConstraint{T}}
    c_max::Vector{T}
    μ_max::Vector{T}
    Zref::Ref{SampledTrajectory}
    Eref::Ref{CostExpansion}
end

function ALConstraintSet{T}() where T
    constraints = ALConstraint{T}[]
    c_max = T[]
    μ_max = T[]
    Z = SampledTrajectory(KnotPoint{Any,Any,Vector{T},T}[])
    Zref = Ref{RD.SampledTrajectory}(Z)
    E = CostExpansion{T}(0,0,0)
    Eref = Ref{CostExpansion}(E)

    ALConstraintSet{T}(constraints, c_max, μ_max, Zref, Eref)
end

function initialize!(conset::ALConstraintSet{T}, cons::TO.ConstraintList,
                     Z::SampledTrajectory, opts::SolverOptions,
                     costs, E=CostExpansion{T}(RD.dims(Z)[1:2]...)) where T
    N = length(Z)
    nx,nu = cons.nx, cons.nu
    @assert all(k->RD.state_dim(Z,k) == nx[k], 1:N)
    @assert all(k->RD.control_dim(Z,k) == nu[k], 1:N)
    @assert length(Z) == length(cons.p)
    ncon = length(cons)
    for i = 1:ncon
        alcon = ALConstraint{T}(Z, cons[i], cons.inds[i], costs, E, sig=cons.sigs[i],
                                diffmethod=cons.diffs[i], solveropts=opts)
        push!(conset.constraints, alcon)
    end
    resize!(conset.c_max, length(Z))
    resize!(conset.μ_max, ncon)
    conset.c_max .= 0
    conset.μ_max .= 0
    conset.Zref[] = Z
    conset.Eref[] = E
    return
end

# Indexing and Iteration
@inline Base.length(conset::ALConstraintSet) = length(conset.constraints)
@inline Base.getindex(conset::ALConstraintSet, i::Integer) = conset.constraints[i]
Base.firstindex(::ALConstraintSet) = 1
Base.lastindex(conset::ALConstraintSet) = length(conset.constraints)

function Base.iterate(conset::ALConstraintSet)
    isempty(conset.constraints) ? nothing : (conset.constraints[1], 1)
end

function Base.iterate(conset::ALConstraintSet, state::Int)
    state >= length(conset) ? nothing : (conset.constraints[state+1], state+1)
end

Base.IteratorSize(::ALConstraintSet) = Base.HasLength()
Base.IteratorEltype(::ALConstraintSet) = Base.HasEltype()
Base.eltype(::ALConstraintSet{T}) where T = ALConstraint{T}

# Methods
for method in (:evaluate_constraints!, :constraint_jacobians!)
    @eval function $method(conset::ALConstraintSet)
        for alcon in conset.constraints
            $method(alcon)
        end
    end
    @eval function $method(conset::ALConstraintSet, Z::SampledTrajectory)
        # Hack to avoid an allocation from using a function barrier w/ more than 1 arg
        # Store a pointer to the trajectory in each conval
        # Only update the trajectory if it's different than the current one
        if Z !== conset.Zref[]
            settraj!(conset, Z)  # this allocates
        end
        $method(conset)  # this does not, but something like
                         # `evaluate_constraints!(conset, Z)` does
    end
end

function alcost(conset::ALConstraintSet{T}) where T
    for alcon in conset.constraints
        alcost(alcon)
    end
end

for method in (:algrad!, :alhess!, :dualupdate!, :penaltyupdate!,
        :max_penalty, :reset_duals!, :reset_penalties!)
    @eval function $method(conset::ALConstraintSet, args...)
        for alcon in conset.constraints
            $method(alcon, args...)
        end
    end
end

function add_alcost_expansion!(conset::ALConstraintSet, E::CostExpansion)
    for alcon in conset.constraints
        add_alcost_expansion!(alcon)
    end
end

function settraj!(conset::ALConstraintSet, Z::SampledTrajectory)
    conset.Zref[] = Z
    for alcon in conset.constraints
        settraj!(alcon, Z)
    end
end

function normviolation!(conset::ALConstraintSet, p=2)
    isempty(conset) && return 0.0
    conset.c_max .= 0
    for i = 1:length(conset)
        normviolation!(conset.constraints[i], p, conset.c_max)
    end
    return norm(conset.c_max, p)
end

# Need to duplicate due to allocation
function max_violation(conset::ALConstraintSet)
    isempty(conset) && return 0.0
    conset.c_max .= 0
    for i = 1:length(conset)
        normviolation!(conset.constraints[i], Inf, conset.c_max)
    end
    return norm(conset.c_max, Inf)
end

function max_penalty(conset::ALConstraintSet)
    for i = 1:length(conset)
        conset.μ_max[i] = max_penalty(conset.constraints[i])
    end
    return maximum(conset.μ_max)
end

function reset!(conset::ALConstraintSet)
    for con in conset.constraints
        resetparams!(con)
        reset_duals!(con)
        reset_penalties!(con)
    end
end

"""
	findmax_violation(conSet)

Return details on the where the largest violation occurs. Returns a string giving the
constraint type, time step index, and index into the constraint.
"""
function findmax_violation(conSet::ALConstraintSet{T}) where T
    c_max0 = -Inf
    j_con = -Inf
    for (i,alcon) in enumerate(conSet.constraints)
        max_violation!(alcon)
        v = maximum(alcon.c_max)
        if v > c_max0
            c_max0 = v
            j_con = i
        end
    end
	if c_max0 < eps()
		return "No constraints violated"
	end
	conval = conSet[j_con]
	k_con = argmax(conval.c_max) # which index
    i_con = searchsortedfirst(conval.inds, k_con)
	c_max, i_max = findmax(abs,conval.viol[i_con])  # index into constraint
	@assert c_max == c_max0
	con_name = string(typeof(conval.con).name.name)
	return con_name * " at time step $k_con at " * TO.con_label(conval.con, i_max)
end


export Blob
export CPUBlob, NullBlob

import Base: eltype, size, length, ndims, copy!, fill!, show, randn!
export       eltype, size, length, ndims, copy!, fill!, erase!, show
export get_num, get_height, get_width, get_fea_size, to_array
export make_blob, make_zero_blob, reshape_blob

############################################################
# A blob is an abstract concept that is suppose
# to hold a 4-D tensor of data. The data could
# either live in CPU memory or GPU memory or
# whatever the backend is used to store the data.
############################################################
@compat abstract type Blob{T, N} end

############################################################
# The following should be implemented for a
# concrete Blob type. Note the following
# procedures are only provided for convenience
# and mainly for components that do not need
# to know the underlying backend (e.g. Filler).
############################################################
function eltype{T}(blob :: Blob{T})
  T
end

function ndims{T,N}(blob :: Blob{T,N})
  N
end
function size(blob :: Blob) # should return the size of data
  error("size not implemented for type $(typeof(blob))")
end
function destroy(blob :: Blob) # should destroy the blob
  error("destroy not implemented for type $(typeof(blob))")
end
function size{T,N}(blob :: Blob{T,N}, dim :: Int)
  if dim < 0
    dim = N+1 + dim
  end

  return dim > N ? 1 : size(blob)[dim]
end
function length(blob :: Blob)
  return prod(size(blob))
end

function get_height(blob :: Blob)
  size(blob, 2)
end
function get_width(blob :: Blob)
  size(blob, 1)
end

function get_num(blob :: Blob)
  size(blob, -1)
end
function get_fea_size(blob :: Blob)
  prod(size(blob)[1:end-1])
end

function show(io::IO, blob :: Blob)
  shape = join(map(x -> "$x", size(blob)), " x ")
  print(io, "Blob($shape)")
end

function to_array(blob::Blob)
  array = Array{eltype(blob)}(size(blob))
  copy!(array, blob)
  array
end

function copy!(dst :: Array, src :: Blob) # should copy content of src to dst
  error("copy! not implemented from src type $(typeof(src)) to dest type $(typeof(dst))")
end
function copy!(dst :: Blob, src :: Array) # should copy content of src to dst
  error("copy! not implemented from src type $(typeof(src)) to dest type $(typeof(dst))")
end
function fill!(dst :: Blob, val) # should fill dst with val
  error("fill! not implemented for src value type $(typeof(val)) to dest type $(typeof(dst))")
end
function erase!(dst :: Blob)
  fill!(dst, 0)
end
function randn!(dst :: Blob) # should fill dst with iid standard normal variates
  error("randn! not implemented for type $(typeof(dst))")
end

############################################################
# A Dummy Blob type holding nothing
############################################################
type NullBlob <: Blob{Void, 0}
end
function fill!(dst :: NullBlob, val)
  # do nothing
end
function show(io::IO, blob::NullBlob)
  print(io, "Blob()")
end

function destroy(blob::NullBlob)
  # do nothing
end
function make_blob(backend::Backend, data_type::Type, dims::Int...)
  make_blob(backend, data_type, dims)
end
function make_blob(backend::Backend, data::Array)
  blob = make_blob(backend, eltype(data), size(data))
  copy!(blob, data)
  return blob
end
function make_zero_blob{N}(backend::Backend, data_type::Type, dims::NTuple{N,Int})
  blob = make_blob(backend, data_type, dims)
  erase!(blob)
  return blob
end
function make_zero_blob(backend::Backend, data_type::Type, dims::Int...)
  make_zero_blob(backend, data_type, dims)
end

function reshape_blob(backend::Backend, blob::Blob, dims::Int...)
  reshape_blob(backend, blob, dims)
end

############################################################
# A Blob for CPU Computation
############################################################
immutable CPUBlob{T <: AbstractFloat, N} <: Blob{T, N}
  data :: AbstractArray{T, N}
end
CPUBlob{N}(t :: Type, dims::NTuple{N,Int}) = CPUBlob(Array{t}(dims))

function make_blob{N}(backend::CPUBackend, data_type::Type, dims::NTuple{N,Int})
  return CPUBlob(data_type, dims)
end

function reshape_blob{T,N1,N2}(backend::CPUBackend, blob::CPUBlob{T,N1}, dims::NTuple{N2,Int})
  @assert prod(dims) == length(blob)
  return CPUBlob{T,N2}(reshape(blob.data, dims))
end
function destroy(blob::CPUBlob)
  # do nothing... or is there anything that I could do?
end

size(blob::CPUBlob) = size(blob.data)

function copy!{T}(dst :: Array{T}, src :: CPUBlob{T})
  @assert length(dst) == length(src)
  dst[:] = src.data[:]
end
function copy!{T}(dst :: CPUBlob{T}, src :: Array{T})
  @assert length(dst) == length(src)
  dst.data[:] = src[:]
end
function copy!{T}(dst :: CPUBlob{T}, src :: CPUBlob{T})
  dst.data[:] = src.data[:]
end
function fill!{T}(dst :: CPUBlob{T}, src)
  fill!(dst.data, src)
end
function randn!{T}(dst :: CPUBlob{T})
  randn!(dst.data)
end

function randn!(a::Array{Float32})
    # TODO This is hideously inefficient - check status of Julia issue
    # https://github.com/JuliaLang/julia/issues/9836
    @compat a[:] = map(Float32, randn(size(a)))
end


function ticks_and_labels end

module Formatters
using Showoff

function scientific(ticks::AbstractVector)
    return Showoff.showoff(ticks, :scientific)
end

function plain(ticks::AbstractVector)
    try
        Showoff.showoff(ticks, :plain)
    catch e
        Base.showerror(stderr, e)
        println("with ticks: ", ticks)
        String["-Inf", "Inf"]
    end
end

end
using .Formatters

"""
    $(SIGNATURES)

Plots a 2-dimensional axis.

## Attributes
$(ATTRIBUTES)
"""
@recipe(Axis2D) do scene
    return Attributes(visible=true, showgrid=true, showticks=true, showtickmarks=true, padding=0.0,
                      ticks=Attributes(ranges_labels=(automatic, automatic), formatter=Formatters.plain,
                                       gap=1, title_gap=3, linewidth=(1, 1),
                                       linecolor=((:black, 0.4), (:black, 0.4)), linestyle=(nothing, nothing),
                                       textcolor=(:black, :black), textsize=(5, 5), rotation=(0.0, 0.0),
                                       align=((:center, :top), (:right, :center)),
                                       font=lift(dim2, theme(scene, :font))),
                      tickmarks=Attributes(length=(1.0, 1.0), linewidth=(1, 1),
                                           linecolor=((:black, 0.4), (:black, 0.2)),
                                           linestyle=(nothing, nothing)),
                      grid=Attributes(linewidth=(0.5, 0.5), linecolor=((:black, 0.3), (:black, 0.3)),
                                      linestyle=(nothing, nothing)),
                      frame=Attributes(linewidth=1.0, linecolor=:black, linestyle=nothing,
                                       axis_position=nothing, axis_arrow=false, arrow_size=2.5,
                                       frames=((false, false), (false, false))),
                      names=Attributes(axisnames=("x", "y"), title=nothing, textcolor=(:black, :black),
                                       textsize=(6, 6), rotation=(0.0, -1.5pi),
                                       align=((:center, :top), (:center, :bottom)),
                                       font=lift(dim2, theme(scene, :font))))
end

"""
    $(SIGNATURES)

Plots a 3-dimensional Axis.

## Attributes
$(ATTRIBUTES)
"""
@recipe(Axis3D) do scene
    q1 = qrotation(Vec3f0(1, 0, 0), -0.5f0 * pi)
    q2 = qrotation(Vec3f0(0, 0, 1), 1.0f0 * pi)
    tickrotations3d = (qrotation(Vec3f0(0, 0, 1), -1.5pi), q2, qrotation(Vec3f0(1, 0, 0), -0.5pi) * q2)
    axisnames_rotation3d = tickrotations3d
    tickalign3d = ((:left, :center), # x axis
                   (:right, :center), # y axis
                   (:right, :center))
    axisnames_align3d = tickalign3d
    tick_color = RGBAf0(0.5, 0.5, 0.5, 0.6)
    grid_color = RGBAf0(0.5, 0.5, 0.5, 0.4)
    grid_thickness = 1
    gridthickness = ntuple(x -> 1.0f0, Val(3))
    tsize = 5 # in percent
    return Attributes(visible=true, showticks=(true, true, true), showaxis=(true, true, true),
                      showgrid=(true, true, true), scale=Vec3f0(1), padding=0.1,
                      names=Attributes(axisnames=("x", "y", "z"), textcolor=(:black, :black, :black),
                                       rotation=axisnames_rotation3d, textsize=(6.0, 6.0, 6.0),
                                       align=axisnames_align3d, font=lift(dim3, theme(scene, :font)), gap=3),
                      ticks=Attributes(ranges_labels=(automatic, automatic), formatter=Formatters.plain,
                                       textcolor=(tick_color, tick_color, tick_color),
                                       rotation=tickrotations3d, textsize=(tsize, tsize, tsize),
                                       align=tickalign3d, gap=3, font=lift(dim3, theme(scene, :font))),
                      frame=Attributes(linecolor=(grid_color, grid_color, grid_color),
                                       linewidth=(grid_thickness, grid_thickness, grid_thickness),
                                       axiscolor=(:black, :black, :black)))
end

isaxis(x) = false
isaxis(x::Union{Axis2D,Axis3D}) = true

const Limits{N} = NTuple{N,Tuple{Number,Number}}

function default_ticks(limits::Limits, ticks, scale_func::Function)
    return default_ticks.(limits, (ticks,), scale_func)
end

# function default_ticks(limits::Limits, ticks::Tuple, scale_func::Function)
#     default_ticks.(limits, ticks, scale_func)
# end
function default_ticks(limits::Tuple{Number,Number}, ticks, scale_func::Function)
    return default_ticks(limits..., ticks, scale_func)
end

function default_ticks(lmin::Number, lmax::Number, ticks::AbstractVector{<:Number}, scale_func::Function)
    return scale_func.((filter(t -> lmin <= t <= lmax, ticks)))
end
function default_ticks(lmin::Number, lmax::Number, ::Automatic, scale_func::Function)
    # scale the limits
    scaled_ticks, mini, maxi = optimize_ticks(scale_func(lmin), scale_func(lmax); k_min=4, # minimum number of ticks # minimum number of ticks
                                              k_max=8)
    length(scaled_ticks) == 1 && isnan(scaled_ticks[1]) && return [-Inf, Inf]
    return scaled_ticks
end

function default_ticks(lmin::Number, lmax::Number, ticks::Integer, scale_func=identity)
    scaled_ticks, mini, maxi = optimize_ticks(scale_func(lmin), scale_func(lmax); k_min=ticks, # minimum number of ticks # minimum number of ticks
                                              k_max=ticks, # maximum number of ticks # maximum number of ticks
                                              k_ideal=ticks,
                                              # `strict_span = false` rewards cases where the span of the
                                              # chosen  ticks is not too much bigger than amin - amax:
                                              strict_span=false)
    return scaled_ticks
end

function default_ticks(x::Automatic, limits::Tuple, n)
    return default_ticks(limits, n, identity)
end

function default_ticks(ticks::Tuple, limits::Tuple, n::Tuple)
    return default_ticks.(ticks, (limits,), n)
end

default_ticks(ticks::Tuple, limits::Limits, n) = default_ticks.(ticks, limits, (n,))

default_ticks(ticks::Tuple, limits::Limits, n::Tuple) = default_ticks.(ticks, limits, n)

default_ticks(ticks::AbstractVector{<:Number}, limits, n) = ticks

function default_labels(x::NTuple{N,Any}, formatter::Function) where {N}
    return default_labels.(x, formatter)
end

function default_labels(x::AbstractVector, y::AbstractVector, formatter::Function=Formatters.plain)
    return default_labels.((x, y), formatter)
end

function default_labels(ticks::AbstractVector, formatter::Function=Formatters.plain)
    if applicable(formatter, ticks)
        return formatter(ticks) # takes the whole array
    elseif applicable(formatter, first(ticks))
        return formatter.(ticks)
    else
        error("Formatting function $(formatter) is neither applicable to $(typeof(ticks)) nor $(eltype(ticks)).")
    end
end

default_labels(x::Automatic, ranges, formatter) = default_labels(ranges, formatter)
default_labels(x::Tuple, ranges::Tuple, formatter) = default_labels.(x, ranges, (formatter,))
default_labels(x::Tuple, ranges, formatter) = default_labels.(x, (ranges,), (formatter,))
default_labels(x::AbstractVector{<:AbstractString}, ranges, formatter::Function) = x
default_labels(x::AbstractVector{<:AbstractString}, ranges::AbstractVector, formatter::Function) = x

function convert_arguments(::Type{<:Axis2D}, limits::Rect)
    e = extrema(limits)
    return (((e[1][1], e[2][1]), (e[1][2], e[2][2])),)
end

function convert_arguments(::Type{<:Axis3D}, limits::Rect)
    e = (minimum(limits), maximum(limits))
    return (((e[1][1], e[2][1]), (e[1][2], e[2][2]), (e[1][3], e[2][3])),)
end

a_length(x::AbstractVector) = length(x)
a_length(x::Automatic) = x

function calculated_attributes!(::Type{<:Union{Axis2D,Axis3D}}, plot)
    ticks = plot.ticks
    args = (plot.padding, plot[1], ticks.ranges, ticks.labels, ticks.formatter)
    ticks[:ranges_labels] = lift(args...) do pad, lims, ranges, labels, formatter
        limit_widths = map(x -> x[2] - x[1], lims)
        pad = (limit_widths .* pad)
        # pad the drawn limits and use them as the ranges
        lim_pad = map((lim, p) -> (lim[1] - p, lim[2] + p), lims, pad)
        num_ticks = labels === automatic ? automatic : a_length.(labels)
        ranges = default_ticks(ranges, lim_pad, num_ticks)
        labels = default_labels(labels, ranges, formatter)
        return (ranges, labels)
    end
    return
end

function draw_ticks(textbuffer, dim, origin, ticks, linewidth, linecolor, linestyle, textcolor, textsize,
                    rotation, align, font)
    for (tick, str) in ticks
        pos = ntuple(i -> i != dim ? origin[i] : tick, Val(2))
        push!(textbuffer, str, Point(pos), rotation=rotation[dim], textsize=textsize[dim], align=align[dim],
              color=textcolor[dim], font=font[dim])
    end
end

function draw_tickmarks(linebuffer, dim, origin, ticks, dir::NTuple{N}, linewidth, linecolor, linestyle,
                        textcolor, textsize, rotation, align, font) where {N}
    for (tick, str) in ticks
        pos = ntuple(i -> i != dim ? origin[i] : tick, Val(2))
        posf0 = Point2f0(pos)
        dirf0 = Pointf0{N}(dir)
        append!(linebuffer, [posf0, posf0 .+ dirf0], color=linecolor[dim], linewidth=linewidth[dim])
    end
end

function draw_grid(linebuffer, dim, origin, ticks, dir::NTuple{N}, linewidth, linecolor, linestyle) where {N}
    dirf0 = Pointf0{N}(dir)
    for (tick, str) in ticks
        tup = ntuple(i -> i != dim ? origin[i] : tick, Val(N))
        posf0 = Pointf0{N}(tup)
        append!(linebuffer, [posf0, posf0 .+ dirf0], color=linecolor[dim], linewidth=linewidth[dim])
    end
end

function draw_frame(linebuffer, limits::NTuple{N,Any}, linewidth, linecolor, linestyle, axis_position,
                    axis_arrow, arrow_size, frames) where {N}
    mini = minimum.(limits)
    maxi = maximum.(limits)
    rect = Rect(Vec(mini), Vec(maxi .- mini))
    origin = Vec{N}(0.0)
    if (origin in rect) && axis_position == :origin
        for i in 1:N
            start = unit(Point{N,Float32}, i) * Float32(mini[i])
            to = unit(Point{N,Float32}, i) * Float32(maxi[i])
            if false #axis_arrow
                arrows(scene, [start, to], linewidth=linewidth, linecolor=linecolor, linestyle=linestyle,
                       arrowsize=arrow_size)
            else
                append!(linebuffer, [start, to], linewidth=linewidth, color=linecolor)
            end
        end
    end
    limit_widths = maxi .- mini
    frames = convert_attribute(frames, key"frames"())
    for side in 1:2
        from = Point{N,Float32}(getindex.(limits, side))
        # if axis is drawn at origin, and we draw frame from origin,
        # we already did this
        if !(from == origin && axis_position == :origin)
            for otherside in 1:2
                for dim in 1:N
                    if frames[N - dim + 1][3 - otherside]
                        p = ntuple(i -> i == dim ? limits[i][otherside] : limits[i][side], Val(N))
                        to = Point{N,Float32}(p)
                        append!(linebuffer, [from, to], linewidth=linewidth, color=linecolor)
                    end
                end
            end
        end
    end
end

function draw_titles(textbuffer, xticks, yticks, origin, limit_widths, tickfont, tick_size, tick_gap,
                     tick_title_gap, axis_labels, textcolor, textsize, rotation, align, font, title)
    tickspace_x = maximum(map(yticks) do tick
                              str = last(tick)
                              tick_bb = text_bb(str, to_font(tickfont[2]), tick_size[2])
                              return widths(tick_bb)[1]
                          end)

    tickspace_y = widths(text_bb(last(first(xticks)), to_font(tickfont[1]), tick_size[1]))[2]

    model_inv = inv(transformationmatrix(textbuffer)[])

    tickspace = transform(model_inv, (tickspace_x, tickspace_y))
    title_start = origin .- (tick_gap .+ tickspace .+ tick_title_gap)
    half_width = origin .+ (limit_widths ./ 2.0)

    posx = (half_width[1], title_start[2])
    posy = (title_start[1], half_width[2])
    positions = (posx, posy)
    for i in 1:2
        if !isempty(axis_labels[i])
            push!(textbuffer, axis_labels[i], positions[i], textsize=textsize[i], align=align[i],
                  rotation=rotation[i], color=textcolor[i], font=font[i])
        end
    end
    if title !== nothing
        # TODO give title own text attributes
        push!(textbuffer, title, (half_width[1], (origin .+ (limit_widths))[2]), textsize=textsize[1],
              align=(:center, :top), rotation=rotation[1], color=textcolor[1], font=font[1])
    end
end

function ticks_and_labels(x)
    st, s = extrema(x)
    r = LinRange(st, s, 4)
    return zip(r, string.(round.(r, 4)))
end

function transform(model::Mat4, x::T) where {T}
    x4d = to_ndim(Vec4f0, x, 0.0)
    return to_ndim(T, model * x4d, 0.0)
end
un_transform(model::Mat4, x) = transform(inv(model), x)

to2tuple(x) = ntuple(i -> x, Val(2))
to2tuple(x::Tuple{<:Any,<:Any}) = x

function draw_axis2d(textbuffer, frame_linebuffer, grid_linebuffer, tickmarks_linebuffer, m, padding, limits,
                     xyrange_labels, showgrid, showticks, showtickmarks,
                     # grid attributes
                     g_linewidth, g_linecolor, g_linestyle,

                     # tick attributes
                     t_linewidth, t_linecolor, t_linestyle, t_textcolor, t_textsize, t_rotation, t_align,
                     t_font, t_gap, t_title_gap,

                     #tickmark attribures
                     tm_linewidth, tm_linecolor, tm_linestyle, tm_length,

                     # frame attributes
                     f_linewidth, f_linecolor, f_linestyle, f_axis_position, f_axis_arrow, f_arrow_size,
                     f_frames,

                     # title / axis name attributes
                     ti_labels, ti_textcolor, ti_textsize, ti_rotation, ti_align, ti_font, ti_title)
    xyrange, labels = xyrange_labels
    start!(textbuffer)
    start!(frame_linebuffer)
    foreach(start!, grid_linebuffer)
    foreach(start!, tickmarks_linebuffer)
    # limits = limits Vec2f0(padding)
    # limits ((xmin, xmax), (ymin, ymax))
    limit_widths = map(x -> x[2] - x[1], limits)
    pad = (limit_widths .* to2tuple(padding))
    # pad the drawn limits
    limits = map((lim, p) -> (lim[1] - p, lim[2] + p), limits, pad)

    # recalculate widths
    limit_widths = map(x -> x[2] - x[1], limits)

    % = mean(limit_widths) / 100 # percentage
    xyticks = zip.(xyrange, labels)
    model_inv = inv(transformationmatrix(textbuffer)[])

    ti_textsize = ti_textsize .* %
    t_textsize = t_textsize .* %
    t_gap = transform(model_inv, to2tuple(t_gap .* %))
    tm_length = transform(model_inv, to2tuple(tm_length .* %))
    t_title_gap = transform(model_inv, to2tuple(t_title_gap .* %))

    origin = first.(limits)
    dirs = ((0.0, Float64(limit_widths[2])), (Float64(limit_widths[1]), 0.0))
    foreach(1:2, dirs, xyticks) do dim, dir, ticks
        if showgrid[dim]
            draw_grid(grid_linebuffer[dim], dim, origin, ticks, dir, g_linewidth, g_linecolor, g_linestyle)
        end
    end
    tm_offsets = ((0.0, Float64(tm_length[2])), (Float64(tm_length[1]), Float64(0.0)))
    foreach(1:2, tm_offsets, xyticks) do dim, offset, ticks
        if showtickmarks[dim]
            draw_tickmarks(tickmarks_linebuffer[dim], dim, origin .- offset, ticks, offset, t_linewidth,
                           t_linecolor, t_linestyle, t_textcolor, t_textsize, t_rotation, t_align, t_font)
        end
    end
    t_offsets = ((0.0, Float64(t_gap[2] + tm_length[2])), (Float64(t_gap[1] + tm_length[1]), Float64(0.0)))
    foreach(1:2, t_offsets, xyticks) do dim, offset, ticks
        if showticks[dim]
            draw_ticks(textbuffer, dim, origin .- offset, ticks, t_linewidth, t_linecolor, t_linestyle,
                       t_textcolor, t_textsize, t_rotation, t_align, t_font)
        end
    end

    draw_frame(frame_linebuffer, limits, f_linewidth, f_linecolor, f_linestyle, f_axis_position, f_axis_arrow,
               f_arrow_size, f_frames)

    draw_titles(textbuffer, xyticks..., origin, limit_widths, t_font, t_textsize, t_gap, t_title_gap,
                ti_labels, ti_textcolor, ti_textsize, ti_rotation, ti_align, ti_font, ti_title)
    finish!(textbuffer)
    finish!(frame_linebuffer)
    foreach(finish!, grid_linebuffer)
    foreach(finish!, tickmarks_linebuffer)
    return
end

# for axis, we don't want to have plot!(scene, args) called on it, so we need to overload it directly
function plot!(scene::SceneLike, ::Type{<:Axis2D}, attributes::Attributes, args...)
    # create "empty" plot type - empty meaning containing no plots, just attributes + arguments
    cplot = Axis2D(scene, attributes, args)
    # Disable any non linear transform for the axis plot!
    cplot.transformation.transform_func[] = identity
    g_keys = (:linewidth, :linecolor, :linestyle)
    f_keys = (:linewidth, :linecolor, :linestyle, :axis_position, :axis_arrow, :arrow_size, :frames)
    t_keys = (:linewidth, :linecolor, :linestyle, :textcolor, :textsize, :rotation, :align, :font, :gap,
              :title_gap)
    tm_keys = (:linewidth, :linecolor, :linestyle, :length)
    ti_keys = (:axisnames, :textcolor, :textsize, :rotation, :align, :font, :title)

    g_args = getindex.(cplot.grid, g_keys)
    f_args = getindex.(cplot.frame, f_keys)
    t_args = getindex.(cplot.ticks, t_keys)
    tm_args = getindex.(cplot.tickmarks, tm_keys)
    ti_args = getindex.(cplot.names, ti_keys)

    textbuffer = TextBuffer(cplot, Point{2})

    grid_linebuffer = Node((LinesegmentBuffer(cplot, Point{2}, transparency=true,
                                              linestyle=lift(first, cplot[:grid, :linestyle])),
                            LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                              linestyle=lift(last, cplot[:grid, :linestyle]))))

    frame_linebuffer = Node(LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                              linestyle=cplot.frame.linestyle))
    tickmarks_linebuffer = Node((LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                                   linestyle=lift(first, cplot[:tickmarks, :linestyle])),
                                 LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                                   linestyle=lift(last, cplot[:tickmarks, :linestyle]))))
    map_once(draw_axis2d, Node(textbuffer), frame_linebuffer, grid_linebuffer, tickmarks_linebuffer,
             transformationmatrix(scene), cplot.padding, cplot[1], cplot.ticks.ranges_labels,
             lift.((dim2,), (cplot.showgrid, cplot.showticks, cplot.showtickmarks))..., g_args..., t_args...,
             tm_args..., f_args..., ti_args...)
    push!(scene, cplot)
    return cplot
end

function labelposition(ranges, dim, dir, tgap, origin::StaticVector{N}) where {N}
    a, b = extrema(ranges[dim])
    whalf = Float32(((b - a) / 2))
    halfaxis = unit(Point{N,Float32}, dim) .* whalf

    return origin .+ (halfaxis .+ (normalize(dir) * tgap))
end

_widths(x::Tuple{<:Number,<:Number}) = x[2] - x[1]
_widths(x) = Float32(maximum(x) - minimum(x))

to3tuple(x::Tuple{Any}) = (x[1], x[1], x[1])
to3tuple(x::Tuple{Any,Any}) = (x[1], x[2], x[2])
to3tuple(x::Tuple{Any,Any,Any}) = x
to3tuple(x) = ntuple(i -> x, Val(3))

function draw_axis3d(textbuffer, linebuffer, limits, ranges_labels, args...)
    # make sure we extend all args to 3D
    ranges, ticklabels = ranges_labels
    args3d = to3tuple.(args)
    (showaxis, showticks, showgrid, axisnames, axisnames_color, axisnames_size, axisrotation, axisalign, axisnames_font, titlegap, gridcolors, gridthickness, axiscolors, ttextcolor, trotation, ttextsize, talign, tfont, tgap) = args3d # splat to names

    N = 3
    start!(textbuffer)
    start!(linebuffer)
    mini, maxi = first.(limits), last.(limits)

    origin = Point{N,Float32}(min.(mini, first.(ranges)))
    limit_widths = max.(last.(ranges), maxi) .- origin
    % = minimum(limit_widths) / 100 # percentage
    ttextsize = (%) .* ttextsize
    axisnames_size = (%) .* axisnames_size

    titlegap = (%) .* titlegap
    tgap = (%) .* tgap
    for i in 1:N
        axis_vec = unit(Point{N,Float32}, i)
        width = Float32(limit_widths[i])
        stop = origin .+ (width .* axis_vec)
        if showaxis[i]
            append!(linebuffer, [origin, stop], color=axiscolors[i], linewidth=1.5f0)
        end
        if showticks[i]
            range = ranges[i]
            j = mod1(i + 1, N)
            tickdir = unit(Point{N,Float32}, j)
            tickdir, offset2 = if i != 2
                tickdir = unit(Vec{N,Float32}, j)
                tickdir, Float32(limit_widths[j] + tgap[i]) * tickdir
            else
                tickdir = unit(Vec{N,Float32}, 1)
                tickdir, Float32(limit_widths[1] + tgap[i]) * tickdir
            end
            for (j, tick) in enumerate(range)
                labels = ticklabels[i]
                if length(labels) >= j
                    str = labels[j]
                    if !isempty(str)
                        startpos = (origin .+ ((Float32(tick - origin[i]) * axis_vec)) .+ offset2)
                        push!(textbuffer, str, startpos, color=ttextcolor[i], rotation=trotation[i],
                              textsize=ttextsize[i], align=talign[i], font=tfont[i])
                    end
                end
            end
            if !isempty(axisnames[i])
                tick_widths = if length(ticklabels[i]) >= 3
                    widths(text_bb(ticklabels[i][end - 1], to_font(tfont[i]), ttextsize[i]))[1]
                else
                    0.0f0
                end
                pos = (labelposition(ranges, i, tickdir, titlegap[i] + tick_widths, origin) .+ offset2)
                push!(textbuffer, to_latex(axisnames[i]), pos, textsize=axisnames_size[i],
                      color=axisnames_color[i], rotation=axisrotation[i], align=axisalign[i],
                      font=axisnames_font[i])
            end
        end
        if showgrid[i]
            c = gridcolors[i]
            thickness = gridthickness[i]
            for _j in (i + 1):(i + N - 1)
                j = mod1(_j, N)
                dir = unit(Point{N,Float32}, j)
                range = ranges[j]
                for tick in range
                    offset = Float32(tick - origin[j]) * dir
                    append!(linebuffer, [origin .+ offset, stop .+ offset], color=c, linewidth=thickness)
                end
            end
        end
        finish!(textbuffer)
        finish!(linebuffer)
    end
    return
end

function plot!(scene::SceneLike, ::Type{<:Axis3D}, attributes::Attributes, args...)
    axis = Axis3D(scene, attributes, args)
    # Disable any non linear transform for the axis plot!
    axis.transformation.transform_func[] = identity
    textbuffer = TextBuffer(axis, Point{3}, transparency=true)
    linebuffer = LinesegmentBuffer(axis, Point{3}, transparency=true)

    tstyle, ticks, frame = to_value.(getindex.(axis, (:names, :ticks, :frame)))
    titlevals = getindex.(tstyle, (:axisnames, :textcolor, :textsize, :rotation, :align, :font, :gap))
    framevals = getindex.(frame, (:linecolor, :linewidth, :axiscolor))
    tvals = getindex.(ticks, (:textcolor, :rotation, :textsize, :align, :font, :gap))
    args = (getindex.(axis, (:showaxis, :showticks, :showgrid))..., titlevals..., framevals..., tvals...)
    map_once(draw_axis3d, Node(textbuffer), Node(linebuffer), axis[1], axis.ticks.ranges_labels, args...)
    push!(scene, axis)
    return axis
end


# PROBABILISTIC CIRCUITS LIBRARY ROOT

module ProbabilisticCircuits

# USE EXTERNAL MODULES

using Reexport

using LogicCircuits
# only reexport selectively from LogicCircuits
export pos_literals, neg_literals
# circuit queries
export issmooth, isdecomposable, isstruct_decomposable,
       isdeterministic, iscanonical
# circuit status
export num_edges, num_parameters
# datasets
export twenty_datasets

include("Utils/Utils.jl")
@reexport using .Utils

include("FactorGraph/factor_graph.jl")
include("FactorGraph/fg_compile.jl")

include("abstract_prob_nodes.jl")
include("shared_prob_nodes.jl")
include("plain_prob_nodes.jl")
include("structured_prob_nodes.jl")
include("logistic_nodes.jl")
include("param_bit_circuit.jl")
include("param_bit_circuit_pair.jl")
include("parameters.jl")
include("gradient_based_learning.jl")

include("queries/likelihood.jl")
include("queries/marginal_flow.jl")
include("queries/map.jl")
include("queries/sample.jl")
include("queries/pr_constraint.jl")
include("queries/information.jl")
include("queries/expectation_rec.jl")
include("queries/expectation_graph.jl")
include("queries/expectation_bit.jl")

include("Logistic/Logistic.jl")
@reexport using .Logistic

include("mixtures/em.jl")

include("structurelearner/chow_liu_tree.jl")
include("structurelearner/init.jl")
include("structurelearner/heuristics.jl")
include("structurelearner/learner.jl")
include("structurelearner/vtree_learner.jl")

include("LoadSave/LoadSave.jl")
@reexport using .LoadSave

using Requires

function __init__()
    # optional dependency
    @require BlossomV = "6c721016-9dae-5d90-abf6-67daaccb2332" include("structurelearner/vtree_learner_blossomv.jl")
end

end


import DCELs._connect_vertices!, DCELs._split_face!
import DCELs._connect_manual!, DCELs._connect_automatic!

@testset "Connecting - primitive" begin
    u, v, w = Vertex(1, 0), Vertex(0, 1), Vertex(-1, 0)
    a1, a2, b1, b2, c1, c2 = [Edge() for _ in 1:6]
    f1, f2 = Face(), Face()

    u.edge, v.edge, w.edge = a1, b1, c1

    a1.dest, b1.dest, c1.dest = u, v, w
    a2.dest, b2.dest, c2.dest = w, u, v

    a1.twin, a2.twin = a2, a1
    b1.twin, b2.twin = b2, b1
    c1.twin, c2.twin = c2, c1

    a1.next, b1.next, c1.next = b1, c1, a1
    a1.prev, b1.prev, c1.prev = c1, a1, b1

    a2.next, b2.next, c2.next = c2, a2, b2
    a2.prev, b2.prev, c2.prev = b2, c2, a2

    a1.face = b1.face = c1.face = f1
    a2.face = b2.face = c2.face = f2

    f1.edge, f2.edge = a1, a2

    z = Vertex(2, 0)

    # No splitting
    h1 = _connect_vertices!(u, z)
    h2 = h1.twin

    @test h1.next == h2
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h1

    @test h1.face == h2.face == f2

    # Splitting
    h3, f_del = _split_face!(z, v)
    h4 = h3.twin

    @test h1.next == h3
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h4
    @test h3.next == b2
    @test h3.prev == h1
    @test h4.next == h2
    @test h4.prev == c2

    @test h1.face == h3.face == b2.face
    @test h4.face == h2.face == a2.face == c2.face

    @test f_del == f2
end

@testset "Connecting - manual" begin
    tess = Tessellation()

    u, v, w = Vertex(1, 0), Vertex(0, 1), Vertex(-1, 0)
    a1, a2, b1, b2, c1, c2 = [Edge() for _ in 1:6]
    f1, f2 = Face(), Face()

    push!(tess.vertices, u, v, w)
    push!(tess.edges, a1, a2, b1, b2, c1, c2)
    push!(tess.faces, f1, f2)

    u.edge, v.edge, w.edge = a1, b1, c1

    a1.dest, b1.dest, c1.dest = u, v, w
    a2.dest, b2.dest, c2.dest = w, u, v

    a1.twin, a2.twin = a2, a1
    b1.twin, b2.twin = b2, b1
    c1.twin, c2.twin = c2, c1

    a1.next, b1.next, c1.next = b1, c1, a1
    a1.prev, b1.prev, c1.prev = c1, a1, b1

    a2.next, b2.next, c2.next = c2, a2, b2
    a2.prev, b2.prev, c2.prev = b2, c2, a2

    a1.face = b1.face = c1.face = f1
    a2.face = b2.face = c2.face = f2

    f1.edge, f2.edge = a1, a2

    z = Vertex(2, 0)

    # No splitting
    h1 = _connect_manual!(tess, u, z, false)
    h2 = h1.twin

    @test h1.next == h2
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h1

    @test h1.face == h2.face == f2

    # Splitting
    h3 = _connect_manual!(tess, z, v, true)
    h4 = h3.twin

    @test h1.next == h3
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h4
    @test h3.next == b2
    @test h3.prev == h1
    @test h4.next == h2
    @test h4.prev == c2

    @test h1.face == h3.face == b2.face
    @test h4.face == h2.face == a2.face == c2.face
end

@testset "Connecting - automatic" begin
    tess = Tessellation()

    u, v, w = Vertex(1, 0), Vertex(0, 1), Vertex(-1, 0)
    a1, a2, b1, b2, c1, c2 = [Edge() for _ in 1:6]
    f1, f2 = Face(), Face()

    push!(tess.vertices, u, v, w)
    push!(tess.edges, a1, a2, b1, b2, c1, c2)
    push!(tess.faces, f1, f2)

    u.edge, v.edge, w.edge = a1, b1, c1

    a1.dest, b1.dest, c1.dest = u, v, w
    a2.dest, b2.dest, c2.dest = w, u, v

    a1.twin, a2.twin = a2, a1
    b1.twin, b2.twin = b2, b1
    c1.twin, c2.twin = c2, c1

    a1.next, b1.next, c1.next = b1, c1, a1
    a1.prev, b1.prev, c1.prev = c1, a1, b1

    a2.next, b2.next, c2.next = c2, a2, b2
    a2.prev, b2.prev, c2.prev = b2, c2, a2

    a1.face = b1.face = c1.face = f1
    a2.face = b2.face = c2.face = f2

    f1.edge, f2.edge = a1, a2

    z = Vertex(2, 0)

    # No splitting
    h1 = _connect_automatic!(tess, u, z)
    h2 = h1.twin

    @test h1.next == h2
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h1

    @test h1.face == h2.face == f2

    # Splitting
    h3 = _connect_automatic!(tess, z, v)
    h4 = h3.twin

    @test h1.next == h3
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h4
    @test h3.next == b2
    @test h3.prev == h1
    @test h4.next == h2
    @test h4.prev == c2

    @test h1.face == h3.face == b2.face
    @test h4.face == h2.face == a2.face == c2.face
end

@testset "Connecting - final" begin
    tess = Tessellation()

    u, v, w = Vertex(1, 0), Vertex(0, 1), Vertex(-1, 0)
    a1, a2, b1, b2, c1, c2 = [Edge() for _ in 1:6]
    f1, f2 = Face(), Face()

    push!(tess.vertices, u, v, w)
    push!(tess.edges, a1, a2, b1, b2, c1, c2)
    push!(tess.faces, f1, f2)

    u.edge, v.edge, w.edge = a1, b1, c1

    a1.dest, b1.dest, c1.dest = u, v, w
    a2.dest, b2.dest, c2.dest = w, u, v

    a1.twin, a2.twin = a2, a1
    b1.twin, b2.twin = b2, b1
    c1.twin, c2.twin = c2, c1

    a1.next, b1.next, c1.next = b1, c1, a1
    a1.prev, b1.prev, c1.prev = c1, a1, b1

    a2.next, b2.next, c2.next = c2, a2, b2
    a2.prev, b2.prev, c2.prev = b2, c2, a2

    a1.face = b1.face = c1.face = f1
    a2.face = b2.face = c2.face = f2

    f1.edge, f2.edge = a1, a2

    z = Vertex(2, 0)

    # No splitting
    h1 = connect(tess, u, z)
    h2 = h1.twin

    @test h1.next == h2
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h1

    @test h1.face == h2.face == f2

    # Splitting
    h3 = connect(tess, z, v, split=true)
    h4 = h3.twin

    @test h1.next == h3
    @test h1.prev == b2
    @test h2.next == a2
    @test h2.prev == h4
    @test h3.next == b2
    @test h3.prev == h1
    @test h4.next == h2
    @test h4.prev == c2

    @test h1.face == h3.face == b2.face
    @test h4.face == h2.face == a2.face == c2.face
end


abstract type AbstractModel end
abstract type AbstractPolarizedModel <: AbstractModel end

"""
    $(TYPEDEF)
This trait specifies whether the model is a *primitive*

# Notes
This will likely turn into a trait in the future so people
can inject their models into Comrade more easily.
"""
abstract type PrimitiveTrait end

"""
    $(TYPEDEF)
Trait for primitive model
"""
struct IsPrimitive end
"""
    $(TYPEDEF)
Trait for not-primitive model
"""
struct NotPrimitive end

"""
    isprimitive(::Type)
Dispatch function that specifies whether a type is a primitive Comrade model.
This function is used for dispatch purposes when composing models.

# Notes
If a user is specifying their own model primitive model outside of Comrade they need
to specify if it is primitive

```julia
struct MyPrimitiveModel end
ComradeBase.isprimitive(::Type{MyModel}) = ComradeBase.IsPrimitive()
```

"""
function isprimitive end

@inline isprimitive(::Type{<:AbstractModel}) = NotPrimitive()

"""
    DensityAnalytic
Internal type for specifying the nature of the model functions.
Whether they can be easily evaluated pointwise analytic. This
is an internal type that may change.
"""
abstract type DensityAnalytic end

"""
    $(TYPEDEF)
Defines a trait that a states that a model is analytic.
This is usually used with an abstract model where we use
it to specify whether a model has a analytic fourier transform
and/or image.
"""
struct IsAnalytic <: DensityAnalytic end

"""
    $(TYPEDEF)
Defines a trait that a states that a model is analytic.
This is usually used with an abstract model where we use
it to specify whether a model has does not have a easy analytic
fourier transform and/or intensity function.
"""
struct NotAnalytic <: DensityAnalytic end

"""
    visanalytic(::Type{<:AbstractModel})
Determines whether the model is pointwise analytic in Fourier domain, i.e. we can evaluate
its fourier transform at an arbritrary point.

If `IsAnalytic()` then it will try to call `visibility_point` to calculate the complex visibilities.
Otherwise it fallback to using the FFT that works for all models that can compute an image.

"""
@inline visanalytic(::Type{<:AbstractModel}) = NotAnalytic()

"""
    imanalytic(::Type{<:AbstractModel})
Determines whether the model is pointwise analytic in the image domain, i.e. we can evaluate
its intensity at an arbritrary point.

If `IsAnalytic()` then it will try to call `intensity_point` to calculate the intensity.
"""
@inline imanalytic(::Type{<:AbstractModel}) = IsAnalytic()

#=
    This is internal function definitions for how to
    compose whether a model is analytic. We need this
    for composite models.
=#
@inline Base.:*(::IsAnalytic, ::IsAnalytic) = IsAnalytic()
@inline Base.:*(::IsAnalytic, ::NotAnalytic) = NotAnalytic()
@inline Base.:*(::NotAnalytic, ::IsAnalytic) = NotAnalytic()
@inline Base.:*(::NotAnalytic, ::NotAnalytic) = NotAnalytic()

"""
    $(SIGNATURES)
Function that computes the pointwise visibility if the model has the trait
in the fourier domain `IsAnalytic()`. Otherwise it will use the FFTW fallback.
"""
function visibility_point end

"""
    $(SIGNATURES)
Function that computes the pointwise intensity if the model has the trait
in the image domain `IsAnalytic()`. Otherwise it will use construct the image in visibility
space and invert it.
"""
function intensity_point end

"""
    $(SIGNATURES)
Computes the intensity map of model by modifying the input IntensityMap object
"""
function intensitymap! end

"""
    $(SIGNATURES)
Computes the intensity map of model. This version requires additional information to
construct the grid.

# Example

```julia
m = Gaussian()
# field of view
fovx, fovy = 5.0
fovy = 5.0
# number of pixels
nx, ny = 128

img = intensitymap(m, fovx, fovy, nx, ny; pulse=DeltaPulse())
```
"""
function intensitymap end

# """
#     ($SIGNATURES)
# A threaded version of the intensitymap function.

# # Notes
# If using autodiff this won't play well with zygote given the mutation. This will be fixed in
# future versions.
# """
# function tintensitymap end

# """
#     ($SIGNATURES)
# A threaded version of the intensitymap! function.

# # Notes
# If using autodiff this won't play well with zygote given the mutation. This will be fixed in
# future versions.
# """
# function tintensitymap! end


###############################################################################
#                                  GRAD CONTEXT                               #
###############################################################################

struct GradCtx
    # map from primal var to its pullback var
    pullbacks::Dict{Variable,Variable}
    # map from primal var to its derivative var
    derivs::Dict{Variable,Variable}
end

GradCtx() = GradCtx(Dict(), Dict())

function rebind_context!(tape::Tape{GradCtx}, st::Dict)
    for (v, dv) in tape.c.derivs
        rebind!(tape, v, st)
        rebind!(tape, dv, st)
    end
    for (v, dv) in tape.c.pullbacks
        rebind!(tape, v, st)
        rebind!(tape, dv, st)
    end
end

const BASE_CTX = BaseCtx()
const CR_CTX = ChainRulesCtx()
const YOTA_RULE_CONFIG = YotaRuleConfig()

function isprimitive(::GradCtx, f, args...)
    return (isprimitive(BASE_CTX, f, args...) || isprimitive(CR_CTX, f, args...))
end

"""
    record_primitive!(tape::Tape{GradCtx}, v_fargs...)

Replace ChainRules primitives `f(args...)` with a sequence:

    rr = push!(tape, mkcall(rrule, f, args...))   # i.e. rrule(f, args...)
    val = push!(tape, mkcall(getfield, rr, 1)     # extract value
    pb = push!(tape, mkcall(getfield, rr, 2)      # extract pullback
"""
function record_primitive!(tape::Tape{GradCtx}, v_fargs...)
    v_f, v_args... = v_fargs
    f, args... = [v isa V ? tape[v].val : v for v in v_fargs]
    if isprimitive(CR_CTX, f, args...)
        rr_op = (is_kwfunc(f) ?
                    mkcall(Core.kwfunc(rrule), v_args[1], rrule, YOTA_RULE_CONFIG, v_args[2:end]...) :
                    mkcall(rrule, YOTA_RULE_CONFIG, v_f, v_args...))
        @assert rr_op.val !== nothing "rrule($f, ...) returned nothing"
        v_rr = push!(tape, rr_op)
        v_val = push!(tape, mkcall(_getfield, v_rr, 1))
        v_pb = push!(tape, mkcall(_getfield, v_rr, 2))
        tape.c.pullbacks[v_val] = v_pb
        return v_val
    else
        return push!(tape, mkcall(v_fargs...))
    end
end

###############################################################################
#                            BCAST GRAD CONTEXT                               #
###############################################################################

struct BcastGradCtx
    inner
end

# get_static_params is broken for BcastGradCtx, so turning off
# this feature for now
Umlaut.get_static_params(::Tracer{BcastGradCtx}, v_fargs) = Core.svec([])

function record_or_recurse!(t::Tracer{BcastGradCtx}, v_fargs...)
    fargs = [v isa V ? t.tape[v].val : v for v in v_fargs]
    # global STATE = (t, v_fargs)
    v_f, v_args... = v_fargs
    f, args... = [v isa V ? t.tape[v].val : v for v in v_fargs]
    return if isprimitive(t.tape.c.inner, fargs...)
        # push!(t.tape, mkcall(broadcasted, vs...))
        rr_op = (is_kwfunc(f) ?
                    mkcall(Core.kwfunc(bcast_rrule), v_args[1], bcast_rrule, YOTA_RULE_CONFIG, broadcasted, v_args[2:end]...) :
                    mkcall(bcast_rrule, YOTA_RULE_CONFIG, broadcasted, v_f, v_args...))
        v_rr = push!(t.tape, rr_op)
        v_val = push!(t.tape, mkcall(_getfield, v_rr, 1))
        v_pb = push!(t.tape, mkcall(_getfield, v_rr, 2))
        t.tape.c.inner.pullbacks[v_val] = v_pb
        return v_val
    else
        types = map(eltype, fargs[2:end])
        trace!(t, getcode(fargs[1], types), v_fargs...)
    end
end

###############################################################################
#                                   GRAD                                      #
###############################################################################

getderiv(tape::Tape, v::Variable) = get(tape.c.derivs, bound(tape, v), nothing)
setderiv!(tape::Tape, x::Variable, dx::Variable) = (
    tape.c.derivs[bound(tape, x)] = bound(tape, dx)
)
hasderiv(tape::Tape, v::Variable) = getderiv(tape, v) !== nothing

function set_or_add_deriv!(tape::Tape, x::Variable, dx::Variable)
    if !hasderiv(tape, x)
        setderiv!(tape, x, dx)
    else
        old_dx = getderiv(tape, x)
        if tape[dx].val isa Tangent || tape[old_dx].val isa Tangent
            new_dx = push!(tape, mkcall(+, dx, old_dx))
        else
            new_dx = push!(tape, mkcall(broadcast, +, dx, old_dx))
        end
        setderiv!(tape, x, new_dx)
    end
end

"""
Collect variables that we need to step through during the reverse pass.
The returned vector is already deduplicated and reverse-sorted
"""
function todo_list(tape::Tape{GradCtx}, y=tape.result)
    y_orig = y
    y_fargs = [tape[y].fn; tape[y].args...]
    is_rrule_based = haskey(tape.c.pullbacks, y)
    if is_rrule_based
        # use rrule instead
        y = tape[y].args[1]
        y_fargs = is_kwfunc(y._op.fn) ? tape[y].args[3:end] : tape[y].args
    end
    y_todo = [x for x in y_fargs if x isa V && tape[x] isa Call]
    x_todos = [todo_list(tape, x) for x in y_todo]
    # include y itself (original), its parents and their parents recursively
    todo = [[y_orig]; y_todo; vcat(x_todos...)]
    # deduplicate
    todo = collect(Set([bound(tape, v) for v in todo]))
    todo = sort(todo, by=v->v.id, rev=true)
    return todo
end

"""
Make a single step of backpropagation.
"""
function step_back!(tape::Tape, y::Variable)
    @debug "step_back!() for $(tape[y])"
    if !hasderiv(tape, y)
        @debug "No derivative found for y = $y, stop propagating the gradient in this path"
        return
    end
    dy = tape.c.derivs[y]
    if haskey(tape.c.pullbacks, y)
        # ChainRules
        pb = tape.c.pullbacks[y]
        @debug "Found pullback: $pb"
        dxs = push!(tape, mkcall(pb, dy))
        # propage derivs to rrule variable
        rr = tape[y].args[1]
        y_fargs = is_kwfunc(rr._op.fn) ? tape[rr].args[4:end] : tape[rr].args[2:end]
    else
        sig_str = join(["::$T" for T in Umlaut.call_signature(tape, tape[y]).parameters], ", ")
        error("No deriative rule found for op $(tape[y]), " *
              "try defining it using \n\n\tChainRulesCore.rrule($sig_str) = ...\n")
    end
    for (i, x) in enumerate(y_fargs)
        if x isa V
            dx = push!(tape, mkcall(getfield, dxs, i))
            @debug "Updating derivative: $x -> $dx"
            set_or_add_deriv!(tape, x, dx)
        end
    end
end

const BACKPROP_STATE = Ref{Tuple}()

"""
Backpropagate through the tape, record derivatives as new operations.
"""
function back!(tape::Tape; seed=1)
    # z - final variable (usually a loss)
    # y - resulting variable of current op
    # x - dependencies of y
    # dy - derivative of z w.r.t. y
    z = tape.result
    if (seed == 1) && (ndims(tape[z].val) > 0)
        error("Gradient of a vector-valued function requires a seed")
    elseif seed == :auto
        zval = tape[z].val
        @assert zval isa Number || zval isa AbstractArray
        seed = zval isa Number ? one(zval) : ones(eltype(zval), size(zval))
    end
    dy = push!(tape, Constant(seed))
    # save seed var to use in compilation later
    tape.meta[:seed] = dy
    # set initial derivative value
    tape.c.derivs[z] = dy
    # queue of variables to calculate derivatives for
    deriv_todo = todo_list(tape)
    for y in deriv_todo
        try
            step_back!(tape, y)
        catch e
            BACKPROP_STATE[] = (tape, y, deriv_todo)
            rethrow(e)
        end
    end
end

function gradtape!(tape::Tape; seed=1)
    # backpropagate gradients
    back!(tape; seed=seed)
    # add a tuple of (val, (gradients...))
    deriv_vars = [hasderiv(tape, v) ? getderiv(tape, v) : ZeroTangent() for v in inputs(tape)]
    deriv_tuple = push!(tape, mkcall(tuple, deriv_vars...))
    # unthunk results
    deriv_tuple_unthunked = push!(tape, mkcall(map, ChainRules.unthunk, deriv_tuple))
    new_result = push!(tape, mkcall(tuple, tape.result, deriv_tuple_unthunked))
    # set result
    tape.result = new_result
    return tape
end

"""
    gradtape(f::Union{Function, DataType}, args...; ctx=GradCtx(), seed=1)
    gradtape!(tape::Tape; seed=1)

Calculate and record to the tape gradients of `tape[tape.resultid]` w.r.t. `Input` nodes.
See grad() for more high-level API.
"""
function gradtape(f::Union{Function,DataType}, args...; ctx=GradCtx(), seed=1)
    _, tape = trace(f, args...; ctx=ctx)
    tape = gradtape!(tape; seed=seed)
    return tape
end

"Like Umlaut.compile, but adds Yota specific ops"
function grad_compile(tape::Tape)
    ex = Umlaut.to_expr(tape)
    seed_var = tape.meta[:seed]
    seed_default_val = tape[seed_var].val
    insert!(ex.args[1].args, 2, Expr(:parameters, Expr(:kw, :seed, seed_default_val)))
    seed_var_name = Umlaut.make_name(seed_var.id)
    op_exprs = ex.args[2].args
    for op_ex in op_exprs
        if op_ex.args[1] == seed_var_name
            op_ex.args[2] = :seed
            break
        end
    end
    return Base.eval(@__MODULE__, ex)
end

const GRAD_CACHE = Dict{Any,Any}()

reset!() = empty!(GRAD_CACHE)

"""
    grad(f, args...; seed=1)

Find gradient of a callable `f` w.r.t. its arguments.

`grad()` returns two things: value of `f(args...)` and a tuple of
grafients w.r.t. to its inputs (including the callable itself).

```jldoctest
using Yota   # hide

val, g = grad(x -> sum(x .+ 1), [1.0, 2.0, 3.0])

# output
(9.0, (ChainRulesCore.ZeroTangent(), [1.0, 1.0, 1.0]))
```

By default, `grad()` expects the callable to return a scalar.
Vector-valued functions can be differentiated if a seed (starting value)
is provided. Seed is equivalent to the vector in VJP notation.

```jldoctest
using Yota   # hide

val, g = grad(x -> 2x, [1.0, 2.0, 3.0]; seed=ones(3))

# output
([2.0, 4.0, 6.0], (ChainRulesCore.ZeroTangent(), [2.0, 2.0, 2.0]))
```

All gradients can be applied to original variables using `update!()` function.

See also: [gradtape](@ref)
"""
function grad(f::Union{Function,DataType}, args...; seed=1)
    # key consists of function type and type of argument (for structs) or its size
    # cache_key = (f, ([(isstruct(arg) || args isa Function) ? typeof(arg) : size(arg) for arg in args]...,))
    cache_key = map(typeof, (f, args...))
    if haskey(GRAD_CACHE, cache_key)
        gf = GRAD_CACHE[cache_key]
        return gf(f, args...; seed=seed)
    else
        tape = gradtape(f, args...; seed=seed)
        gf = grad_compile(tape)
        GRAD_CACHE[cache_key] = gf
        return tape.retval
    end
end


using Flux
using Flux.Tracker
using Statistics
import LinearAlgebra.Diagonal
using Random

module CycleWorld
using Random

"""
 CycleWorld

   1 -> 0 -> 0 -> ... -> 0 -|
   ^------------------------|

chain_length: size of cycle
actions: Progress

"""

mutable struct CycleWorldEnv
    chain_length::Int64
    agent_state::Int64
    actions::AbstractSet
    partially_observable::Bool
    CycleWorldEnv(chain_length::Int64; rng=Random.GLOBAL_RNG, partially_observable=true) =
        new(chain_length,
            0,
            Set(1:1),
            partially_observable)
end

function reset!(env::CycleWorldEnv; rng = Random.GLOBAL_RNG, kwargs...)
    env.agent_state = 0
end

get_actions(env::CycleWorldEnv) = env.actions

function environment_step!(env::CycleWorldEnv, action::Int64; rng = Random.GLOBAL_RNG, kwargs...)
    # actions 1 == Turn Left
    # actions 2 == Turn Right
    # actions 3 == Up
    # JuliaRL.step()
    env.agent_state = (env.agent_state + 1) % env.chain_length
    # JuliaRL
end

function get_reward(env::CycleWorldEnv) # -> get the reward of the environment
    return 0
end

function get_state(env::CycleWorldEnv) # -> get state of agent
    if env.partially_observable
        return partially_observable_state(env)
    else
        return fully_observable_state(env)
    end
end

function fully_observable_state(env::CycleWorldEnv)
    return [env.agent_state]
end

function partially_observable_state(env::CycleWorldEnv)
    state = zeros(1)
    if env.agent_state == 0
        state[1] = 1
    end
    return state
end

function is_terminal(env::CycleWorldEnv) # -> determines if the agent_state is terminal
    return false
end

function Base.show(io::IO, env::CycleWorldEnv)
    println(env.agent_state)
end

function start!(env::CycleWorldEnv; rng=Random.GLOBAL_RNG, kwargs...)
    reset!(env; rng=rng, kwargs...)
    return env, get_state(env)
end

function step!(env::CycleWorldEnv, action; rng = Random.GLOBAL_RNG, kwargs...) # -> env, state, reward, terminal
    environment_step!(env, action; rng=rng, kwargs...)
    return env, get_state(env), get_reward(env), is_terminal(env)
end

end #end CycleWorld

build_features(s) = [1.0, s[1], 1-s[1]]

glorot_uniform(rng::Random.AbstractRNG, dims...) = (rand(rng, Float32, dims...) .- 0.5f0) .* sqrt(24.0f0/sum(dims))
glorot_normal(rng::Random.AbstractRNG, dims...) = randn(rng, Float32, dims...) .* sqrt(2.0f0/sum(dims))

abstract type AbstractGVFLayer end

function get_question_parameters(gvfn::AbstractGVFLayer, preds_tilde, stp1) end

mutable struct GVFLayer{F, A, V} <: AbstractGVFLayer
    σ::F
    Wx::A
    Wh::A
    cumulants::AbstractArray
    discounts::AbstractArray
    h::V
end

GVFLayer(num_gvfs, num_ext_features, cumulants, discounts; init=Flux.glorot_uniform, σ_int=σ) =
    GVFLayer(
        σ_int,
        param(0.1.*init(num_gvfs, num_ext_features)),
        param(0.1.*init(num_gvfs, num_gvfs)),
        cumulants,
        discounts,
        param(Flux.zeros(num_gvfs)))

function get_question_parameters(gvfn::GVFLayer{F,A,V}, preds_tilde, state_tp1) where {F, A, V}
    cumulants = [gvfn.cumulants[i](state_tp1, preds_tilde) for i in 1:length(gvfn.cumulants)]
    discounts = [gvfn.discounts[i](state_tp1) for i in 1:length(gvfn.cumulants)]
    return cumulants, discounts, ones(size(gvfn.discounts))
end

function (m::GVFLayer)(h, x)
    new_h = m.σ.(m.Wx*x + m.Wh*h)
    return new_h, new_h
end

Flux.hidden(m::GVFLayer) = m.h
Flux.@treelike GVFLayer
GVFN(args...; kwargs...) = Flux.Recur(GVFLayer(args...; kwargs...))

function reset!(m, h_init)
    Flux.reset!(m)
    m.state.data .= h_init
end

function jacobian(δ, pms)
    k  = length(δ)
    J = IdDict()
    for id in pms
        v = get!(J, id, zeros(k, size(id)...))
        for i = 1:k
            Flux.back!(δ[i], once = false) # Populate gradient accumulator
            v[i, :,:] .= id.grad
            id.grad .= 0 # Reset gradient accumulator
        end
    end
    J
end

function jacobian!(J::IdDict, δ::TrackedArray, pms::Params)
    k  = length(δ)
    for i = 1:k
        Flux.back!(δ[i], once = false) # Populate gradient accumulator
        for id in pms
            v = get!(J, id, zeros(typeof(id[1].data), k, size(id)...))::Array{typeof(id[1].data), 3}
            v[i, :, :] .= id.grad
            id.grad .= 0 # Reset gradient accumulator
        end
    end
end

abstract type Optimizer end

function train!(gvfn::Flux.Recur{T}, opt::Optimizer, h_init, state_seq, env_state_tp1) where {T <: AbstractGVFLayer} end
function train!(gvfn::AbstractGVFLayer, opt::Optimizer, h_init, state_seq, env_state_tp1)
    throw("$(typeof(opt)) not implemented for $(typeof(gvfn)). Try Flux.Recur{$(typeof(gvfn))} ")
end

mutable struct RTD_jacobian <: Optimizer
    α::Float64
    J::IdDict
    RTD_jacobian(α) = new(α, IdDict{Any, Array{Float64, 3}}())
end

function train!(gvfn::Flux.Recur{T}, opt::RTD_jacobian, h_init, states, env_state_tp1) where {T <: AbstractGVFLayer}

    α = opt.α
    reset!(gvfn, h_init)
    preds = gvfn.(states)
    preds_t = preds[end-1]
    preds_tilde = Flux.data(preds[end])

    cumulants, discounts, ρ = get_question_parameters(gvfn.cell, preds_tilde, env_state_tp1)

    targets = cumulants .+ discounts.*preds_tilde
    δ = targets .- preds_t

    jacobian!(opt.J, δ, Params([gvfn.cell.Wx, gvfn.cell.Wh]))

    for weights in Params([gvfn.cell.Wx, gvfn.cell.Wh])
        Flux.Tracker.update!(weights, -α.*sum(δ.*opt.J[weights]; dims=1)[1,:,:])
    end
end

mutable struct RTD <: Optimizer
    α::Float64
    RTD(α) = new(α)
end

function train!(gvfn::Flux.Recur{T}, opt::RTD, h_init, states, env_state_tp1) where {T <: AbstractGVFLayer}

    α = opt.α

    reset!(gvfn, h_init)
    preds = gvfn.(states)
    preds_t = preds[end-1]
    preds_tilde = Flux.data(preds[end])

    cumulants, discounts, ρ = get_question_parameters(gvfn.cell, preds_tilde, env_state_tp1)
    targets = cumulants .+ discounts.*preds_tilde
    δ = targets .- preds_t

    prms = params(gvfn)
    # println(prms)
    grads = Tracker.gradient(() ->mean(δ.^2), prms)
    for weights in prms
        Flux.Tracker.update!(weights, -α.*grads[weights])
    end
end

mutable struct RTDC <: Optimizer
    α::Float64
    β::Float64
    h::IdDict
    RTD(α) = new(α)
end

function train!(gvfn::Flux.Recur{T}, opt::RTDC, h_init, states, env_state_tp1) where {T <: AbstractGVFLayer}

    α = opt.α

    reset!(gvfn, h_init)
    preds = gvfn.(states)
    preds_t = preds[end-1]
    preds_tilde = Flux.data(preds[end])

    cumulants, discounts, ρ = get_question_parameters(gvfn.cell, preds_tilde, env_state_tp1)
    targets = cumulants .+ discounts.*preds_tilde
    δ = targets .- preds_t

    prms = params(gvfn)
    # println(prms)
    grads = Tracker.gradient(() ->mean(δ.^2), prms)
    for weights in prms
        Flux.Tracker.update!(weights, -α.*grads[weights])
    end
end

function test_GVFN_bptt()

    env = CycleWorld.CycleWorldEnv(6)
    num_gvfs = 6
    num_steps = 50000
    τ=7

    rng = Random.MersenneTwister(10)

    opt = RTD(0.6)
    discount = (args...)->0.0
    cumulant(i) = (s_tp1, p_tp1)-> i==1 ? s_tp1[1] : p_tp1[i-1]
    # cumulants = [[cumulant(i) for i in 1:6]; [cumulant(1)]]
    # discounts = [[discount for i in 1:6]; [(env_state)-> env_state[1] == 1 ? 0.0 : 0.9]]

    cumulants = [cumulant(i) for i in 1:6]
    discounts = [discount for i in 1:6]

    pred_strg = zeros(num_steps, num_gvfs)

    gvfn = GVFN(num_gvfs, 3, cumulants, discounts; init=(dims...)->glorot_uniform(rng, dims...))
    _, s_t = CycleWorld.start!(env)
    h_t = zeros(num_gvfs)
    h_tm1 = zeros(num_gvfs)

    state_list = [zeros(3) for t in 1:τ]
    popfirst!(state_list)
    push!(state_list, build_features(s_t))
    hidden_state_init = zeros(num_gvfs)

    for step in 1:num_steps
        _, s_tp1, _, _ = CycleWorld.step!(env, 1)

        if length(state_list) == (τ+1)
            # println(state_list)
            popfirst!(state_list)
        end
        push!(state_list, build_features(s_tp1))

        print(step, "\r")

        train!(gvfn, opt, hidden_state_init, state_list, s_tp1)

        reset!(gvfn, hidden_state_init)
        preds = gvfn.(state_list)
        pred_strg[step,:] .= preds[end].data
        # println(env.agent_state)
        # println(preds[end])
        s_t .= s_tp1
        hidden_state_init .= Flux.data(preds[1])
        # println(preds)
        # println(hidden_state_init)

    end
    return pred_strg

end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 4de51e08-3854-11eb-10f2-d17580deef9c
begin
	using Pkg
	Pkg.add(["AbstractTrees"])
end

# ╔═╡ 2ac17898-3856-11eb-1235-4fd15bc1592c
begin
	using AbstractTrees
end

# ╔═╡ 3e0c149a-3350-11eb-0b17-7bea5d1a28d4
begin
	DAY = 7
	md"# 2020 Day $DAY"
end

# ╔═╡ 18080fe2-339b-11eb-34f0-f393638506fc
md"## Part 1"

# ╔═╡ 4e084a4e-339b-11eb-3ffc-41118215f226
md"*Bags*"

# ╔═╡ 24e1613c-339b-11eb-2976-59884053d1ee
input = read("../data/day_$DAY.txt", String)

# ╔═╡ df1d1be2-36c6-11eb-07c6-493a612dd182
test_input = """
light red bags contain 1 bright white bag, 2 muted yellow bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
bright white bags contain 1 shiny gold bag.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
faded blue bags contain no other bags.
dotted black bags contain no other bags.
"""

# ╔═╡ fefa6924-3850-11eb-067a-5f62b6570b06

# ╔═╡ 0f76e99c-3851-11eb-113c-0136765cd4c4

# ╔═╡ 5f974aa2-3851-11eb-2198-09a8c6d7881a
function get_edges(line::AbstractString)
	raw_container, contents = split(line,"contain")
	container = raw_container[1:end-6]
	bags = (
		split(contents,",")
		.|> strip
		.|> s -> replace(s, "s." => "")
		# container, content, number
		.|> s -> [container, strip(s[3:end-4]), parse(Int,s[1])]
		)
end

# ╔═╡ 3858684c-3857-11eb-3506-4f1c87a22921
function input2edges(input)
	(
	[s for s in split(strip(input),"\n") if !occursin("no other bags",s)]
	.|> get_edges
	|> nested_list -> [item for sublist in nested_list for item in sublist]
)
end

# ╔═╡ d7a2264a-3859-11eb-3ac9-a176933f2063
function edges2colors(edges)
	container, content, number = zip(edges...)
	union(Set(container),Set(content))
end

# ╔═╡ 00aec94e-385a-11eb-3985-b9a159435b1a
function edges2mapping(edges)
	colors = edges2colors(edges)
	d = Dict(c => [] for c in colors)
	for (o, i, n) in edges
		push!(d[i],o)
	end
	return d
end

# ╔═╡ 46c193da-385a-11eb-31a6-25e4eda23b35

# ╔═╡ 4fef820e-385b-11eb-1f9a-ff3310d2340c
function find_parent!(outer_colors,d,c)
	for color in d[c]
		if color ∉ outer_colors
			push!(outer_colors, color)
			find_parent!(outer_colors,d,color)
		end
	end
end

# ╔═╡ 31101e58-385f-11eb-2afa-55f88de38916
e = input2edges(input)

# ╔═╡ 6310ea54-385f-11eb-25b2-a353f8585388
mapping = edges2mapping(e)

# ╔═╡ 863e77da-385f-11eb-19a6-53f6116e4a24
colors = [i for i in edges2colors(e)]

# ╔═╡ 7e22f556-3893-11eb-38d2-1bfc0441a8ad
function find_parents(e,mapping,color)
	outer_colors = Set()
	return find_parent!(outer_colors,mapping,color)
end

# ╔═╡ 91204704-38bf-11eb-298c-6995aa085a06
begin
	outer_colors = Set()
	find_parent!(outer_colors,mapping,"shiny gold")
	outer_colors
end

# ╔═╡ 1b524b64-385e-11eb-17b1-d920e32aab0c

# ╔═╡ 2801d0ec-38c1-11eb-28dc-3ff0a89ee9a7
e

# ╔═╡ 5f797e22-339d-11eb-1ff3-dda90f946838
md"## Part 2"

# ╔═╡ e3177eae-339e-11eb-3f9f-1f7230976832
md"*<Insert Description>*"

# ╔═╡ Cell order:
# ╠═3e0c149a-3350-11eb-0b17-7bea5d1a28d4
# ╟─18080fe2-339b-11eb-34f0-f393638506fc
# ╠═4e084a4e-339b-11eb-3ffc-41118215f226
# ╠═4de51e08-3854-11eb-10f2-d17580deef9c
# ╠═2ac17898-3856-11eb-1235-4fd15bc1592c
# ╠═24e1613c-339b-11eb-2976-59884053d1ee
# ╠═df1d1be2-36c6-11eb-07c6-493a612dd182
# ╟─fefa6924-3850-11eb-067a-5f62b6570b06
# ╟─0f76e99c-3851-11eb-113c-0136765cd4c4
# ╠═5f974aa2-3851-11eb-2198-09a8c6d7881a
# ╠═3858684c-3857-11eb-3506-4f1c87a22921
# ╠═d7a2264a-3859-11eb-3ac9-a176933f2063
# ╠═00aec94e-385a-11eb-3985-b9a159435b1a
# ╟─46c193da-385a-11eb-31a6-25e4eda23b35
# ╠═4fef820e-385b-11eb-1f9a-ff3310d2340c
# ╠═31101e58-385f-11eb-2afa-55f88de38916
# ╠═6310ea54-385f-11eb-25b2-a353f8585388
# ╠═863e77da-385f-11eb-19a6-53f6116e4a24
# ╠═7e22f556-3893-11eb-38d2-1bfc0441a8ad
# ╠═91204704-38bf-11eb-298c-6995aa085a06
# ╟─1b524b64-385e-11eb-17b1-d920e32aab0c
# ╠═2801d0ec-38c1-11eb-28dc-3ff0a89ee9a7
# ╠═5f797e22-339d-11eb-1ff3-dda90f946838
# ╟─e3177eae-339e-11eb-3f9f-1f7230976832


using Test

using QuantumControlBase
using QuantumControlBase.TestUtils

@testset "dry-run" begin

    problem = dummy_control_problem()
    file_expected = joinpath(tempdir(), "dry-run.jld2")
    result, file =
        @test_logs (:info, "Would optimize and store in $file_expected") @optimize_or_load(
            tempdir(),
            problem;
            method=:dummymethod,
            filename="dry-run.jld2",
            dry_run=true,
            verbose=true
        )
    @test isnothing(result)
    @test file == file_expected

end

@testset "metadata" begin

    problem = dummy_control_problem()
    outdir = mktempdir()
    println("")
    result, file = @optimize_or_load(
        outdir,
        problem;
        method=:dummymethod,
        filename="optimization_with_metadata.jld2",
        metadata=Dict("testset" => "metadata", "method" => :dummymethod,)
    )
    @test result.converged
    @test basename(file) == "optimization_with_metadata.jld2"
    @test isfile(file)
    result_load, metadata = load_optimization(file; return_metadata=true)
    @test result_load isa QuantumControlBase.TestUtils.DummyOptimizationResult
    @test result_load.message == "Reached maximum number of iterations"
    @test metadata["testset"] == "metadata"
    @test metadata["method"] == :dummymethod

end

@testset "continue_from" begin
    problem = dummy_control_problem()
    outdir = mktempdir()
    println("")
    result, file = @optimize_or_load(
        outdir,
        problem;
        iter_stop=5,
        method=:dummymethod,
        filename="optimization_stage1.jld2"
    )
    @test result.converged
    result1 = load_optimization(file)
    result2 = optimize(problem; method=:dummymethod, iter_stop=15, continue_from=result1)
    @test result2.iter == 15
end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

struct DensityWithShape{D<:AbstractDensity,S<:AbstractValueShape} <: AbstractDensity
    density::D
    shape::S
end

Base.parent(density::DensityWithShape) = density.density

var_bounds(density::DensityWithShape) = var_bounds(parent(density))

ValueShapes.varshape(density::DensityWithShape) = density.shape

DensityInterface.logdensityof(density::DensityWithShape, v::Any) = logdensityof(parent(density), v)

struct ReshapedDensity{D<:AbstractDensity,S<:AbstractValueShape} <: AbstractDensity
    density::D
    shape::S
end

Base.parent(density::ReshapedDensity) = density.density

var_bounds(density::ReshapedDensity) = var_bounds(density.density)

ValueShapes.varshape(density::ReshapedDensity) = density.shape

function _reshapeddensity_variate(density::ReshapedDensity, v::Any)
    vs = varshape(density)
    orig_vs = varshape(parent(density))
    reshape_variate(orig_vs, vs, v)
end

function DensityInterface.logdensityof(density::ReshapedDensity, v::Any)
    v_reshaped = _reshapeddensity_variate(density, v)
    logdensityof(parent(density), v_reshaped)
end

function checked_logdensityof(density::ReshapedDensity, v::Any)
    v_reshaped = _reshapeddensity_variate(density, v)
    checked_logdensityof(parent(density), v_reshaped)
end

ValueShapes.unshaped(density::AbstractDensity) = _unshaped_density(density, varshape(density))

_unshaped_density(density::AbstractDensity, ::ArrayShape{<:Real,1}) = density
_unshaped_density(density::AbstractDensity, ::AbstractValueShape) = ReshapedDensity(density, ArrayShape{Real}(totalndof(density)))

(shape::AbstractValueShape)(density::AbstractDensity) = _reshaped_density(density, shape, varshape(density))

function _reshaped_density(density::AbstractDensity, new_shape::VS, orig_shape::VS) where {VS<:AbstractValueShape}
    if orig_shape == new_shape
        density
    else
        throw(ArgumentError("Value shapes are incompatible"))
    end
end

function _reshaped_density(density::AbstractDensity, new_shape::AbstractValueShape, orig_shape::AbstractValueShape)
    if totalndof(orig_shape) == totalndof(new_shape)
        ReshapedDensity(density, new_shape)
    else
        throw(ArgumentError("Value shapes are incompatible"))
    end
end


abstract type EoS end
abstract type SAFT <: EoS end

abstract type PCSAFTFamily <: SAFT end
abstract type sPCSAFTFamily <: SAFT end
abstract type ogSAFTFamily <: SAFT end
abstract type SAFTVRMieFamily <: SAFT end

struct SAFTVRMie <: SAFTVRMieFamily;
    components
    params::SAFTVRMieParams
end

function SAFTVRMie(components::T1) where T1 <: AbstractVector{T2} where T2 <: AbstractString
    set_components = [Set([components[i]]) for i in 1:length(components)]
    raw_params = retrieveparams(components, "SAFTVRMie"; kwargs...)
    model = SAFTVRMie(set_components, create_SAFTVRMieParams(raw_params))
end

struct PCSAFT <: PCSAFTFamily
    components
    params::PCSAFTParams
end

function PCSAFT(components::T1) where T1 <: AbstractVector{T2} where T2 <: AbstractString
    set_components = [Set([components[i]]) for i in 1:length(components)]
    raw_params = retrieveparams(components, "PCSAFT"; kwargs...)
    model = PCSAFT(set_components, create_PCSAFTParams(raw_params))
end

struct sPCSAFT <: sPCSAFTFamily
    components
    params::sPCSAFTParams
end

function sPCSAFT(components::T1) where T1 <: AbstractVector{T2} where T2 <: AbstractString
    set_components = [Set([components[i]]) for i in 1:length(components)]
    raw_params = retrieveparams(components, "sPCSAFT"; kwargs...)
    model = sPCSAFT(set_components, create_sPCSAFTParams(raw_params))
end

struct ogSAFT <: ogSAFTFamily
    components
    params::ogSAFTParams
end

function ogSAFT(components::T1) where T1 <: AbstractVector{T2} where T2 <: AbstractString
    set_components = [Set([components[i]]) for i in 1:length(components)]
    raw_params = retrieveparams(components, "ogSAFT"; kwargs...)
    model = ogSAFT(set_components, create_ogSAFTParams(raw_params))
end


"""
    whoswork(blk, comm)

To which rank is block `blk` currently assigned to.

"""
function whoswork(blk, comm::MPI.Comm)
    blk % MPI.Comm_size(comm)
end

function whoswork(blk, comm::Nothing)
    error("Communicator Nothing should not ask this.")
    return nothing
end

"""
    is_my_work(blk, comm)

Returns a boolean whether the block `blk` is assigned to this rank.

"""
function is_my_work(blk, comm::MPI.Comm)
    whoswork(blk, comm) == MPI.Comm_rank(comm)
end

function is_my_work(blk, comm::Nothing)
    return true
end

abstract type AbstractCommPattern end
# Symmetric communication patterns
struct CommPatternTK <: AbstractCommPattern end
struct CommPatternT <: AbstractCommPattern end
struct CommPatternK <: AbstractCommPattern end

"""
    is_comm_pattern(t, tn, k, kn, pattern)

    Do period t, tn and contingencies k, kn match the communication pattern?

"""
function is_comm_pattern(t, tn, k, kn, ::CommPatternTK)
    return (
        # Neighboring periods and base case (k == 1)
        ((tn == t-1 || tn == t+1) && kn == 1 && k == 1) ||
        # From base case to contingencies
        (tn == t && k == 1 && kn != 1) ||
        # From contingencies to base case
        (tn == t && k != 1 && kn == 1)
    )
end

function is_comm_pattern(t, tn, k, kn, ::CommPatternT)
    return (
        # Neighboring periods and base case (k == 1)
        ((tn == t-1 || tn == t+1) && kn == 1 && k == 1)
    )
end

function is_comm_pattern(t, tn, k, kn, ::CommPatternK)
    return (
        # From base case to contingencies
        (tn == t && k == 1 && kn != 1) ||
        # From contingencies to base case
        (tn == t && k != 1 && kn == 1)
    )
end

"""
    comm_neighbors!(data, blocks, runinfo, pattern, comm)

Nonblocking communication with a given pattern. An array of requests is returned.

"""
function comm_neighbors!(data::AbstractArray{T,2}, blocks::AbstractBlocks, runinfo::ProxALProblem, pattern::AbstractCommPattern, comm::MPI.Comm) where {T}
	requests = MPI.Request[]
    for blk in runinfo.blkLinIndex
        block = blocks.blkIndex[blk]
        k = block[1]
        t = block[2]
        if is_my_work(blk, comm)
            for blkn in runinfo.blkLinIndex
                blockn = blocks.blkIndex[blkn]
                kn = blockn[1]
                tn = blockn[2]
                if is_comm_pattern(t, tn, k, kn, pattern) && !is_my_work(blkn, comm)
                    remote = whoswork(blkn, comm)
                    if isa(pattern, CommPatternTK)
                        sbuf = @view data[:,blk]
                        rbuf = @view data[:,blkn]
                    elseif isa(pattern, CommPatternT)
                        sbuf = @view data[:,t]
                        rbuf = @view data[:,tn]
                    else
                        error("Invalid communication pattern")
                    end
                    push!(requests, MPI.Isend(sbuf, remote, t, comm))
                    push!(requests, MPI.Irecv!(rbuf, remote, tn, comm))
                end
            end
        end
    end
    return requests
end

function comm_neighbors!(data::AbstractArray{T,3}, blocks::AbstractBlocks, runinfo::ProxALProblem, pattern::AbstractCommPattern, comm::MPI.Comm) where {T}
	requests = MPI.Request[]
    for blk in runinfo.blkLinIndex
        block = blocks.blkIndex[blk]
        k = block[1]
        t = block[2]
        if is_my_work(blk, comm)
            for blkn in runinfo.blkLinIndex
                blockn = blocks.blkIndex[blkn]
                kn = blockn[1]
                tn = blockn[2]
                if is_comm_pattern(t, tn, k, kn, pattern) && !is_my_work(blkn, comm)
                    remote = whoswork(blkn, comm)
                    if isa(pattern, CommPatternK)
                        sbuf = @view data[:,k,t]
                        rbuf = @view data[:,kn,tn]
                    else
                        error("Invalid communication pattern")
                    end
                    push!(requests, MPI.Isend(sbuf, remote, k, comm))
                    push!(requests, MPI.Irecv!(rbuf, remote, kn, comm))
                end
            end
        end
    end
    return requests
end

function comm_neighbors!(data::AbstractArray, blocks::AbstractBlocks, runinfo::ProxALProblem, pattern::AbstractCommPattern, comm::Nothing)
    return nothing
end

"""
    comm_wait!(requests)

Wait until the communciation requests `requests` have been fulfilled.

"""
function comm_wait!(requests::Vector{MPI.Request})
    return MPI.Waitall!(requests)
end

function comm_wait!(requests::Nothing)
    return nothing
end

"""
    comm_max(data, comm)

Collective to reduce and return the maximum over scalar `data`.

"""
function comm_max(data::Float64, comm::MPI.Comm)
    return MPI.Allreduce(data, MPI.MAX, comm)
end

function comm_max(data::Float64, comm::Nothing)
    return data
end

"""
    comm_sum(data::Float64, comm)

Collective to reduce and return the sum over scalar `data`.

"""
function comm_sum(data::Float64, comm::MPI.Comm)
    return MPI.Allreduce(data, MPI.SUM, comm)
end

function comm_sum(data::Float64, comm::Nothing)
    return data
end

"""
    comm_sum!(data, comm)

Collective to reduce the sum over array `data`.

"""
function comm_sum!(data::AbstractArray, comm::MPI.Comm)
    return MPI.Allreduce!(data, MPI.SUM, comm)
end

function comm_sum!(data::AbstractArray, comm::Nothing)
    return data
end

function comm_rank(comm::MPI.Comm)
    return MPI.Comm_rank(comm)
end

function comm_rank(comm::Nothing)
    return 0
end

function comm_barrier(comm::MPI.Comm)
    return MPI.Barrier(comm)
end

function comm_barrier(comm::Nothing)
    return nothing
end


struct PathCurve
    pt1::Point
    pt2::Point
    pt3::Point
end

struct PathMove
    pt1::Point
end

struct PathLine
    pt1::Point
end

struct PathClose
end

const PathElement = Union{PathCurve, PathMove, PathLine, PathClose}

"""
A Path object contains, in the `.path` field, a vector of
`PathElement`s (`PathCurve`, `PathMove`, `PathLine`,
`PathClose`) that describe a Cairo path. Use `drawpath()` to
draw it.

```
Path([PathMove(Point(2.0, 90.5625)),
      PathCurve(Point(4.08203, 68.16015), Point(11.28, 45.28), Point(24.8828, 26.40234)),
      PathLine(Point(2.0, 90.5625)),
      PathClose()
     ])
```

"""
struct Path
    path::Vector{PathElement}
end

"""
    Path(ptlist::Vector{Point}; close=false))

Create a Path from the points in `ptlist`.
"""
function Path(ptlist::Vector{Point}; close=false)
    p = Path(PathElement[])
    push!(p.path, PathMove(ptlist[1]))
    for pt in ptlist[2:end]
        push!(p.path, PathLine(pt))
    end
    if close
        push!(p.path, PathClose())
    end
    return p
end

Base.size(cp::Path) = length(cp.path)
Base.length(cp::Path) = length(cp.path)
Base.getindex(cp::Path, i::Int) = cp.path[i]
Base.IndexStyle(cp::Path) = IndexLinear()
Base.setindex!(cp::Path, el, e::Integer) = setindex!(cp.path, el, e)
Base.keys(cp::Path) = keys(cp.path)
Base.eltype(::Type{Path}) = PathElement

function Base.iterate(cp::Path)
    if length(cp.path) != 0
        return (first(cp.path), 1)
    else
        return nothing
    end
end

function Base.iterate(cp::Path, state)
    if state > length(cp.path)
        return nothing
    else
        return (cp.path[state], state + 1)
    end
end

Base.getindex(cp::Path, r::Union{LinRange, UnitRange, StepRangeLen}) = getindex(cp.path, r)
Base.lastindex(cp::Path) = length(cp.path)

function Base.show(io::IO, cp::Path)
    println(io, "Path([")
    for p in cp.path[1:end-1]
        println(io, " $(p),")
    end
    println(io, " $(cp.path[end])")
    println(io, "])\n")
end

drawpath(cpe::PathMove) = move(cpe.pt1)
drawpath(cpe::PathLine) = line(cpe.pt1)
drawpath(cpe::PathCurve) = curve(cpe.pt1, cpe.pt2, cpe.pt3)
drawpath(cpe::PathClose) = closepath()

"""
    storepath()

Obtain the current Cairo path and make a Luxor Path object,
which is an array of PathElements.

Returns the Path object.

See also getpath() and getpathflat().
"""
function storepath()
    path = PathElement[]
    for e in getpath()
        if e.element_type == 0
            (x, y) = e.points
            push!(path, PathMove(Point(x, y)))
        elseif e.element_type == 1
            (x, y) = e.points
            push!(path, PathLine(Point(x, y)))
        elseif e.element_type == 2
            (x1, y1, x2, y2, x3, y3) = e.points
            push!(path, PathCurve(Point(x1, y1), Point(x2, y2), Point(x3, y3)))
        elseif e.element_type == 3
            push!(path, PathClose())
        else
            error("unknown Cairo Path Entry " * repr(e.element_type))
            error("unknown Cairo Path Entry " * repr(e.points))
        end
    end
    return Path(path)
end

"""
    drawpath(cp::Path; action=:none, startnewpath=true)
    drawpath(cp::Path, action; startnewpath=true)

Make the Luxor path stored in `cp` and apply the `action`.

To make paths, follow some path construction functions such
as `move()`, `line()`, and `curve()` with the `storepath()`
function.

By default, `startnewpath=true`, which starts a new path,
discarding any existing path contents.
"""
function drawpath(cp::Path; action=:none, startnewpath=true)
    startnewpath && newpath()
    for p in cp
          drawpath(p)
    end
    do_action(action)
end

drawpath(cp::Path, action::Symbol; startnewpath=true) = drawpath(cp::Path; action=action, startnewpath=startnewpath)

"""
    polytopath(ptlist)

Convert a polygon to a Path object.
```julia
@draw drawpath(polytopath(ngon(O, 145, 5, vertices=true)), action=:fill)
```
"""
function polytopath(ptlist::Vector{Point})
    path = PathElement[]
    push!(path, PathMove(ptlist[1]))
    for pt in ptlist[2:end]
        push!(path, PathLine(pt))
    end
    push!(path, PathClose())
    return Path(path)
end
"""
    bezierpathtopath(bp::BezierPath)

Convert a Bezier path to a Path object.

```julia
@draw drawpath(polytopath(ngon(O, 145, 5, vertices=true)), action=:fill)
```
"""
function bezierpathtopath(bp::BezierPath)
    path = PathElement[]
    push!(path, PathMove(first(first(bp))))
    for p in bp
        push!(path, PathCurve(p[2], p[3], p[4]))
    end
    return Path(path)
end


# use conventional beam forming

using Cairo, Fontconfig
using Gadfly
using JLD2, FileIO

using DSP
using Caesar

AMP.setForceEvalDirect!(true)

include(joinpath(dirname(@__FILE__),"ps3utils.jl"))

## Load the data
dataparent = joinpath(dirname(@__FILE__),"..","..","..","..","..","..","data","ps3")
datadir = joinpath(dataparent,"equilateral_22inches") #
# datadir = joinpath(dataparent,"demo")

# reload data

leftwavdata = loadPs3WavEnvelopes(joinpath(datadir,"left_wavdata_corrected.jld2"))
rightwavdata = loadPs3WavEnvelopes(joinpath(datadir,"right_wavdata_corrected.jld2"))

## Load the templates
templdir = joinpath(dataparent, "templateX200ms")
lefttemplate = loadPs3WavEnvelopes(joinpath(templdir, "template_left.jld2"), variable="templ", cutmax=false)
righttemplate = loadPs3WavEnvelopes(joinpath(templdir, "template_right.jld2"), variable="templ", cutmax=false)

leftwav = [leftwavdata[0] leftwavdata[1] leftwavdata[2] leftwavdata[3]]
rightwav = [rightwavdata[0] rightwavdata[1] rightwavdata[2] rightwavdata[3]]

## Spot check

Gadfly.plot(y=leftwav[:,1], Geom.line)
Gadfly.plot(y=rightwav[:,1], Geom.line)

## generate nonparametric bearing intensities from waveform data

intensityLeft = beamformBasic(leftwav, lefttemplate, soundSpeed=340)
intensityRight = beamformBasic(rightwav, righttemplate, soundSpeed=340)
# Gadfly.plot(x=azimuths, y=intensity, Geom.line)

pcLeft = intensityToCircularKDE(intensityLeft, SNRfloor=0.80, N=100)
pcRight = intensityToCircularKDE(intensityRight, SNRfloor=0.80, N=100)

pl = AMP.plotKDECircular([pcLeft; pcRight], scale=0.4, c=["blue";"green"])
pl.coord = Coord.Cartesian(aspect_ratio=1.0, xmin=-2, xmax=2, ymin=-2,ymax=2)
pl

##

pl |> PDF("/tmp/test.pdf",30cm,20cm)
pl |> SVG("/tmp/test.svg",10cm,10cm)

##

@async run(`evince /tmp/test.pdf`)

##

pl = AMP.plotKDECircular([pcLeft], scale=0.4, c=["red";"green"])

#


using Plots
using Random
using Colors

# seed
# Random.seed!(3)

random_tenosr_dist = "uniform"
#random_tenosr_dist = "gauss"

# experimental setup
rep_max = 10

# zero cutoff KL divergence
eps_thred = 1.0e-12

# save file name option
n_digits = 5

# Real datasets
datasets = [
            "AttFace",      #(112,92,400)
           #"4DLFD"    #(9,9,512,512,3)
           ]
sort!(datasets)

dataset_label_dict = Dict(
                   "AttFace" => "AttFace",
                   #"4DLFD" => "4DLFD",
                )
dataset_label_dict = sort(dataset_label_dict)

# approximation methods
methods_flags = Dict(
             "LT1R" => false,
             "LTR" => true,
             "NTD_KL" => true,
             "NTD_LS" => true,
             "lraSNTD" => true,
            )
non_negative_methods = ["LT1R", "LTR", "NTD_KL", "NTD_LS", "lraSNTD"]
methods = [k for k = keys(methods_flags) if methods_flags[k] == true]

# move to "TLrRR" to top
sort!(methods)
methods = insert!( setdiff(methods, ["LTR"]), 1, "LTR" )
std_method = "LTR"

# #################################
# Experimental design for real data
# #################################

plans = Dict( dataset for dataset = zip(datasets, [ [] for i = 1:length(datasets) ] ) )
# AttFace (112,92,400)
plans["AttFace"] =
[[1,1,1], [3,3,3], [5,5,5], [10,10,10], [15,15,15], [20,20,20], [30,30,30], [40,40,40], [50,50,50], [60,60,60], [70,70,70], [80,80,80]]

#"4DLFD"    #(9,9,512,512,3)
plans["4DLFD"] =
[[1,1,1,1,1], [2,2,2,2,1],[3,3,4,4,1],[3,3,5,5,1],[3,3,6,6,1],[3,3,7,7,1],[3,3,8,8,1],[3,3,16,16,1],[3,3,20,20,1],
 [3,3,40,40,1],[3,3,60,60,1],[3,3,80,80,1]]

# ##########################
# config for other methods
# ##########################

NTD_init = "SVD"

# ##########################
# config for Plots
# ##########################

# plt options for AI data
# fonts list : Helvetica, palatino, courier, times, newcentryschlbk, advantgrade
img_size = (500, 500)
fnt1 = font(20, "times")
fnt2 = font(12, "times")
fnt3 = font(15, "times")

linewidth = 2.5
markerstrokewidth = 1.0
markersize = 12
markerwidth = 2

legend = :none
#legend = :bottomleft
legendfont = fnt1

lw_y0 = 3

# plot options for real data
img_size_wide = (1050, 300)

# color setting
colors = palette(:tab20)

colors_dict = Dict(
                   "LT1R" => :red,
                   "TLrRR" => :red,
                   "NTD_KL" => :black,
                   "NTD_LS" => :blue,
                   "lraSNTD" => :brown,
                )

label_dict = Dict(
                   "LT1R" => "TL1R",
                   "LTR" => "LTD(proposed)",
                   "NTD_KL" => "NTD_KL",
                   "NTD_LS" => "NTD_LS",
                   "lraSNTD" => "lraSNTD",
                )

markershapes_dict = Dict(
                   "LT1R" => :circle,
                   "TLrRR" => :circle,
                   "NTD_KL" => :rect,
                   "NTD_LS" => :dtriangle,
                   "lraSNTD" => :utriangle,
                )

linetype_dict = Dict(
                   "LT1R" => :dashdot,
                   "LTR" => :dashdot,
                   "NTD_KL" => :dashdot,
                   "NTD_LS" => :dot,
                   "lraSNTD" => :dashdot,
                )