text:
dir_path = @__DIR__
push!(LOAD_PATH, dir_path * "/sim") # simulation
push!(LOAD_PATH, dir_path * "/vis") # visualization
push!(LOAD_PATH, dir_path * "/mdp") # fully observable MDP
push!(LOAD_PATH, dir_path * "/dis") # discretization
push!(LOAD_PATH, dir_path * "/olm") # online methods
using sim
using vis
using mdp
using dis
using olm

using Serialization
using Printf

include("update_info.jl")

# state space
const pts_per_s = 150
const pts_per_ds = 30
const pts_per_p = 20
const min_s = 0.0
const min_p = -12.0
const max_p = 12.0
const min_ds = 0.0
const max_ds = 50.0

# action space
const pts_per_acc = 10
const pts_per_ste = 3
const min_acc = -15
const max_acc = 5
const min_ste = -1e-1
const max_ste = 1e-1

function discretize_vit(world::World)
  max_s = world.road.path.S[end]
  state_d = Discretization([pts_per_s, pts_per_ds, pts_per_p],
                           [min_s, min_ds, min_p],
                           [max_s, max_ds, max_p])
  ctrl_d = Discretization([pts_per_acc, pts_per_ste],
                          [min_acc, min_ste],
                          [max_acc, max_ste])
  return (state_d, ctrl_d)
end

function discretize_fwds(world::World)
  pts_per_acc = 3
  pts_per_ste = 3
  min_acc = -15
  max_acc = 5
  min_ste = -1e-1
  max_ste = 1e-1
  ctrl_d = Discretization([pts_per_acc, pts_per_ste],
                          [min_acc, min_ste],
                          [max_acc, max_ste])
  return ctrl_d
end

function reward(x, u, nx, agent, world)
  return nx[2]^3 + (abs(x[3]) > 0.45 * world.road.width ? -1e9 : 0.0)
end

# Scenario: Racer scenario
function main()
  # load the graphical context (OpenGL handle, the graphical window, etc.)
  context = vis.setup()
  vis_scaling = 0.01

  # make the road
  path = make_figure8_path()
  road_width = 20.0
  road = Road(path, road_width,
              vis.make_road(context, path.X, path.Y, road_width))
  road.road.T = vis.scale_mat(vis_scaling, vis_scaling)

  # make the world
  world = World(road, [], vis_scaling)

  # make the agents
  agent1 = Agent(1, [0.0; 0.0; 0], vis.make_car(context))
  push!(world.agents, agent1)
  agent2 = Agent(2, [0.0; 0.0; 0], vis.make_car(context, [1.0, 0.0, 0.0]))
  push!(world.agents, agent2)
  agent3 = Agent(3, [0.0; 0.0; 0], vis.make_car(context, [0.0, 1.0, 0.0]))
  push!(world.agents, agent3)

  # Value Iteration Approach ------------------------------------------------ #
  # read in, iterate and store the policy
  P = nothing
  S = nothing
  (state_d, ctrl_d) = discretize_vit(world) # value iteration

  policy_file_path = dir_path * "/../data/fr_vi_policy.bin"

  println("Input the number of iterations")
  repeat = readline()
  repeat = repeat == "" ? 500 : parse(Int, repeat)
  if isfile(policy_file_path)
    fp = open(policy_file_path, "r")
    (P, state_d, ctrl_d) = deserialize(fp)
    close(fp)

    S = P.S
  else
    S = mdp.make_MDP(agent1, world, reward, state_d, ctrl_d)
    P = Policy(S, 0.999)
  end
  for i in 1:repeat
    print("$(i) -> ")
    display(mdp.iterate!(P))
  end
  fp = open(policy_file_path, "w")
  serialize(fp, (P, state_d, ctrl_d))
  close(fp)

  # switch agent's trained controller on
  agent1.controller! = mdp.controllerd_interp!
  agent1.custom = (P, state_d, ctrl_d)
  # ------------------------------------------------------------------------- #

  # Forward Search Approach ------------------------------------------------- #
  ctrl_d = discretize_fwds(world)
  agent2.controller! = olm.controller_fwds!
  # ------------------------------------------------------------------------- #

  # MCTS Approach ----------------------------------------------------------- #
  ctrl_d = discretize_fwds(world)
  agent3.controller! = olm.controller_mcts!
  # ------------------------------------------------------------------------- #

  # make diagnostics render objects
  info1 = vis.InfoBox(context, 0.4, 0.75, vis_scaling)
  info2 = vis.InfoBox(context, 1.0, 0.75, vis_scaling)
  info3 = vis.InfoBox(context, -0.4, -0.35, vis_scaling)
  infos = [info1, info2, info3]

  # main loop for rendering and simulation
  window = true
  h = 1e-2
  t0 = time_ns()
  oldt = (time_ns() - t0) / 1e9
  while window
    t = (time_ns() - t0) / 1e9

    to_visualize = []
    if world.road.road != nothing
      push!(to_visualize, world.road.road)
    end

    for agent in world.agents
      # Forward Search Approach --------------------------------------------- #
      if agent.id == 2
        print("FWDS: ")
        @time u = olm.plan_fwds(agent.x, agent, world, reward, ctrl_d, 1)
        agent.custom = u
      end
      # --------------------------------------------------------------------- #

      # MCTS Approach ------------------------------------------------------- #
      if agent.id == 3
        print("MCTS: ")
        @time us = olm.plan_mcts(agent.x, agent, world, reward, ctrl_d, 1)
        agent.custom = us
      end
      # --------------------------------------------------------------------- #

      ## advance one frame in time
      advance!(agent.dynamics!, agent.x, Pair(agent, world), oldt, t, h)
      update_info(infos[agent.id], agent, world, t)
      push!(to_visualize, infos[agent.id])

      ## visualize
      (x, y, sx, sy, dx, u) = diagnostic(agent, world, t)
      agent.is_braking = dx[1] * u[1] < 0

      ## get additional information
      dp = dx[3]
      th = atan(sy, sx)
      th_car = atan(dp, agent.x[2]) + atan(sy, sx)

      if agent.car != nothing
        ## update the car
        vis.car_lights!(agent.car, agent.is_braking)
        agent.car.T = (vis.scale_mat(world.vis_scaling, world.vis_scaling) *
                       vis.translate_mat(x, y) *
                       vis.rotate_mat(th_car - pi / 2))
        push!(to_visualize, agent.car)
      end
    end

    window = vis.visualize(context, to_visualize)

    oldt = t
  end
end


text:
#=
= Step height move.
=#

#= Structs ====================================================================#

function ht_init(Y, S::State, X::Auxiliary) # hts.
    X.hts[:] = S.hts
end

function ht_unconstrained(X::Auxiliary) # lhts.
    @. X.lhts = log(X.hts)
end

function ht_constrained(X::Auxiliary) # hts.
    @. X.hts = exp(X.lhts)
end

function ht_momenta(C::Constants, X::Auxiliary) # plhts
    X.plhts[:] = randn(C.k+1)
end

#= Move =======================================================================#

function ht_move(Y, C::Constants, S::State, X::Auxiliary, U::Counts)

    ε = .8C.lh + rand()*.4C.lh # base step size ± 20%.
    ht_init(Y,S,X)
    ht_unconstrained(X)
    ht_momenta(C,X)
    ep = ht_potential(C,X) # previous energy; also computes the gradient.
    kp = ht_kinetic(X)

    # ht_println(X) # debug.
    # println("ep: $ep, kp: $kp") # debug.

    ht_pstep(.5ε,X)
    for i = 1:C.Lh
        ht_qstep(ε,X)
        if i != C.Lh
            ht_potential(C,X)
            ht_pstep(ε,X)
        end
    end
    en = ht_potential(C,X)
    ht_pstep(.5ε,X)
    kn = ht_kinetic(X)

    # ht_println(X) # debug.
    # println("en: $en, kn: $kn") # debug.
    # println("diff: $(ep+kp-en-kn)") # debug.

    p = rand()
    l = exp(ep+kp-en-kn)
    # println((p < l) ? "$p < $l" : "$p ≥ $l") # debug.

    if p < l
        S.hts[:] = X.hts
        U.haccepts += 1
    else
        ht_init(Y,S,X)      # leave the auxiliary parameters
        ht_unconstrained(X) # in a consistent state.
    end
    U.hprops += 1

end

function ht_pstep(s, X::Auxiliary)
    @. X.plhts -= s * X.∇lhts
end

function ht_qstep(s, X::Auxiliary)
    @. X.lhts += s * X.plhts
    ht_constrained(X)
end

#= Energy =====================================================================#

function ht_potential(C::Constants, X::Auxiliary)
    k,R,α,β = C.k, C.R, C.α, C.β
    hts, widths, counts = X.hts, X.widths, X.counts
    locs = @view(X.locs[2:end-1])

    e = sum(@. (widths+β)*hts - (counts+α)*X.lhts - log(widths)) -
        sum(@. log(locs*(R-locs)))
    @. X.∇lhts = (widths+β)*hts - (counts+α)
    return e
end

ht_kinetic(X::Auxiliary) = .5*sum(X.plhts.^2)

#= I/O ========================================================================#

function ht_println(X::Auxiliary)
    println("X: locs = $(string(X.locs)), hts = $(string(X.hts))\n",
            "   lhts = $(string(X.lhts)), ∇lhts = $(string(X.∇lhts)),",
              " plhts = $(string(X.plhts))")
end

function string(A::Vector{Float64})
    io = IOBuffer()
    write(io, "[")
    for a in @view(A[1:end-1])
        write(io, @sprintf("%g, ", a))
    end
    write(io, @sprintf("%g]", A[end]))
    return String(take!(io))
end

ext: jl

max_stars_repo_name: kevdur/rjcp

max_issues_repo_name: kevdur/rjcp

max_forks_repo_name: kevdur/rjcp


text:
"""
    import_se_from_csv(assays_file::AbstractString, rowdata_file::AbstractString, coldata_file::AbstractString)

Imports a SummarizedExperiment object from its components stored as csv files.

# Arguments
- `assays_file::AbstractString`: the path to a csv file containing the assays.
   It is produced in R with `write.csv(assays(se), "assays_file.csv")`.
- `rowdata_file::AbstractString`: the path to a csv file containing the rowdata.
   It is produced in R with `write.csv(rowData(se), "rowdata_file.csv")`.
- `coldata_file::AbstractString`: the path to a csv file containing the coldata.
   It is produced in R with `write.csv(colData(se), "coldata_file.csv")`.
"""
function import_se_from_csv(assays_file::AbstractString, rowdata_file::AbstractString, coldata_file::AbstractString)

    raw_assays = CSV.File(assays_file) |> DataFrame

    groups = unique(raw_assays[!, :group_name])
    assays = construct_assays(raw_assays, groups)

    row_data = CSV.File(rowdata_file) |> DataFrame
    rename!(row_data, Dict(:Column1 => "name"))
    row_data[!, :name] = map(string, row_data[!, :name])

    col_data = CSV.File(coldata_file) |> DataFrame
    rename!(col_data, Dict(:Column1 => "name"))

    SummarizedExperiment(assays, row_data, col_data)

end

"""
    import_se_from_csv(assays_file::AbstractString)

Imports a SummarizedExperiment object from a csv file containing the assays.
Rowdata and coldata are generated based on the information stored in the row
and column names of `assays_file`. This is the method of choice if the rowdata
and coldata of the original SummarizedExperiment object does not provide valuable
information either because they are empty or because they are part of a
MultiAssayExperiment object.

# Arguments
- `assays_file::AbstractString`: the path to a csv file containing the assays.
   It is produced in R with `write.csv(assays(se), "assays_file.csv")`.
"""
function import_se_from_csv(assays_file::AbstractString)

    raw_assays = CSV.File(assays_file) |> DataFrame

    groups = unique(raw_assays[!, :group_name])
    assays = construct_assays(raw_assays, groups)

    row_data = DataFrame(
        name = raw_assays[1:Int64(size(raw_assays, 1) / length(groups)), 1]
    )

    col_data = DataFrame(
        name = names(raw_assays)[4:end]
    )

    SummarizedExperiment(assays, row_data, col_data)

end

"""
    import_mae_from_csv(experiment_files::Vector{<:AbstractString}, sample_data_file::AbstractString, sample_map_file::AbstractString; experiment_names::Vector{<:AbstractString} = experiment_files)

Imports a MultiAssayExperiment object from its components stored as csv files.

# Arguments
- `experiment_files::Vector{<:AbstractString}`: a list of paths to the csv files
   containing the assays of different experiments. For example, if a MultiAssayExperiment
   object has 3 experiments, `experiment_files` is defined as:
  `[assays_file1, assays_file2, assays_file3]`
   Each `assays_file` is produced in R with `write.csv(assays(se), "assays_file.csv")`.
   It is produced in R with `write.csv(assays(se), "assays_file.csv")`.
- `sample_data_file::AbstractString`: the path to a csv file containing the sample data.
   It is produced in R with `write.csv(colData(mae), "sample_data_file.csv")`.
- `sample_map_file::AbstractString`: the path to a csv file containing the sample map.
   It is produced in R with `write.csv(colData(mae), "sample_map_file.csv", row.names = FALSE)`.
- `experiment_names::Vector{<:AbstractString}`: a vector of names to assign to the
   experiments listed in `experiment_files`. By default, the experiments are named
   with the filepaths specified in `experiment_files`.
"""
function import_mae_from_csv(experiment_files::Vector{<:AbstractString}, sample_data_file::AbstractString, sample_map_file::AbstractString; experiment_names::Vector{<:AbstractString} = experiment_files)

    experiment_list = OrderedDict{String, SummarizedExperiment}()

    for (cur_exp, cur_exp_name) in zip(experiment_files, experiment_names)

        experiment_list[cur_exp_name] = import_se_from_csv(cur_exp)

    end

    sample_data = CSV.File(sample_data_file) |> DataFrame
    rename!(sample_data, Dict(:Column1 => "name"))

    sample_map = CSV.File(sample_map_file) |> DataFrame
    rename!(sample_map, Dict(:assay => "experiment"))
    rename!(sample_map, Dict(:primary => "sample"))
    select!(sample_map, [:sample, :experiment, :colname])

    MultiAssayExperiment(experiment_list, sample_data, sample_map)

end

# experiment_files = [joinpath(@__DIR__, "assets/XO_microbiota_assays.csv"),
#                     joinpath(@__DIR__, "assets/XO_metabolites_assays.csv"),
#                     joinpath(@__DIR__, "assets/XO_biomarkers_assays.csv")]

# experiment_names = ["microbiota", "metabolites", "biomarkers"]

# mae = import_mae_from_csv(experiment_files,
#     joinpath(@__DIR__, "assets/sample_data.csv"),
#     joinpath(@__DIR__, "assets/sample_map.csv"),
#     experiment_names = experiment_names)

function construct_assays(raw_assays::DataFrame, groups::Vector{<:AbstractString})

    assays = OrderedDict{String, AbstractArray}()

    for cur_group in groups

        assays[cur_group] = filter(:group_name => x -> x == cur_group, raw_assays)[:, 4:end] |> Matrix

    end

    assays

end


text:
# Separable sum, using tuples of arrays as variables

export SeparableSum

"""
**Separable sum of functions**

    SeparableSum(f₁,…,fₖ)

Given functions `f₁` to `fₖ`, returns their separable sum, that is
```math
g(x_1,…,x_k) = ∑_{i=1}^k f_i(x_i).
```
The object `g` constructed in this way can be evaluated at `Tuple`s of length `k`. Likewise, the `prox` and `prox!` methods for `g` operate with (input and output) `Tuple`s of length `k`.

Example:

    f = SeparableSum(NormL1(), NuclearNorm()); # separable sum of two functions
    x = randn(10); # some random vector
    Y = randn(20, 30); # some random matrix
    f_xY = f((x, Y)); # evaluates f at (x, Y)
    (u, V), f_uV = prox(f, (x, Y), 1.3); # computes prox at (x, Y)
"""
struct SeparableSum{T <: Tuple} <: ProximableFunction
	fs::T
end

SeparableSum(fs::Vararg{ProximableFunction}) = SeparableSum((fs...,))

is_prox_accurate(f::SeparableSum) = all(is_prox_accurate.(f.fs))
is_convex(f::SeparableSum) = all(is_convex.(f.fs))
is_set(f::SeparableSum) = all(is_set.(f.fs))
is_singleton(f::SeparableSum) = all(is_singleton.(f.fs))
is_cone(f::SeparableSum) = all(is_cone.(f.fs))
is_affine(f::SeparableSum) = all(is_affine.(f.fs))
is_smooth(f::SeparableSum) = all(is_smooth.(f.fs))
is_quadratic(f::SeparableSum) = all(is_quadratic.(f.fs))
is_generalized_quadratic(f::SeparableSum) = all(is_generalized_quadratic.(f.fs))
is_strongly_convex(f::SeparableSum) = all(is_strongly_convex.(f.fs))

function (f::SeparableSum)(x::Tuple)
	sum = 0.0
  for k in eachindex(x)
	  sum += f.fs[k](x[k])
  end
  return sum
end

function prox!(ys::T, fs::Tuple, xs::T, gamma::Real=1.0) where T <: Tuple
  sum = 0.0
  for k in eachindex(xs)
	  sum += prox!(ys[k], fs[k], xs[k], gamma)
  end
  return sum
end

function prox!(ys::T, fs::Tuple, xs::T, gamma::Tuple) where T <: Tuple
  sum = 0.0
  for k in eachindex(xs)
	  sum += prox!(ys[k], fs[k], xs[k], gamma[k])
  end
  return sum
end

prox!(ys::T, f::SeparableSum, xs::T, gamma::Union{Real, Tuple}=1.0) where {T <: Tuple} = prox!(ys, f.fs, xs, gamma)

function gradient!(grad::T, fs::Tuple, x::T) where T <: Tuple
  val = 0.0
  for k in eachindex(fs)
    val += gradient!(grad[k], fs[k], x[k])
  end
  return val
end

gradient!(grad::T, f::SeparableSum, x::T) where {T <: Tuple} = gradient!(grad, f.fs, x)

fun_name(f::SeparableSum) = "separable sum"
fun_dom(f::SeparableSum) = "n/a"
fun_expr(f::SeparableSum) = "(x₁, …, xₖ) ↦ f₁(x₁) + … + fₖ(xₖ)"
fun_params(f::SeparableSum) = "n/a"

function prox_naive(f::SeparableSum, xs::Tuple, gamma::Union{Real, Tuple}=1.0)
	fys = 0.0
	ys = [];
  for k in eachindex(xs)
	  y, fy = prox_naive(f.fs[k], xs[k], typeof(gamma) <: Real ? gamma : gamma[k])
		fys += fy;
		append!(ys, [y]);
  end
	return Tuple(ys), fys
end