#-------------------------------------------------------------------------------
# File: repair.jl
# Description: This files contains all function for the "repair" phase.
# Date: November 03, 2019
# Author: Jonathan Fontaine, Killian Fretaud, Rémi Garcia,
#         Boualem Lamraoui, Benoît Le Badezet, Benoit Loger
#-------------------------------------------------------------------------------

"""
    find_first_violation(solution::Solution, instance::Instance)

Return the index of the first car violating the paint_batch_limit constraint in the sequence,
return -1 if such a car does not exist
"""
function find_first_violation(solution::Solution, instance::Instance)
    first_violation = -1
    position = instance.nb_late_prec_day+1
    while position <= solution.length
        if solution.colors[position].width > instance.nb_paint_limitation
            first_violation = position + instance.nb_paint_limitation
            break
        end
        position += solution.colors[position].width
    end
    return first_violation
end

"""
    first_strategy_repair!(solution::Solution, instance::Instance)

First strategy
"""
function first_strategy_repair!(solution::Solution, instance::Instance)
    # First strategy
    b0 = instance.nb_late_prec_day+1

    position = b0
    counter = 1
    current_color = instance.color_code[solution.sequence[position]]
    position += 1
    first_violation = 0
    while position <= instance.nb_cars
        if instance.color_code[solution.sequence[position]] == current_color
            counter += 1
        else
            counter = 1
            current_color = instance.color_code[solution.sequence[position]]
        end
        if counter > instance.nb_paint_limitation
            if position >= b0
                if first_violation == 0
                    first_violation = position
                end
                car_RC_value = instance.RC_keys[solution.sequence[position]]
                car_pos = 1
                len = length(instance.same_RC[car_RC_value])
                while (car_pos <= len) && (instance.color_code[instance.same_RC[car_RC_value][car_pos]] == current_color)
                    car_pos += 1
                end
                if car_pos <= len
                    move_exchange!(solution, position, car_pos, instance)
                    counter = first_violation - position
                    if first_violation >= position
                        position = first_violation - 1
                    else
                        position -= 1
                    end
                    first_violation = 0
                elseif first_violation >= position
                    position -= 2
                    if counter == 2*instance.nb_paint_limitation
                        # This strategy can't repair
                        position = first_violation
                    end
                end
            else
                if first_violation != 0
                    position = first_violation
                else
                    # pcc are not possible in nb_late_prec_day so we keep the counter
                end
            end
        end
        position += 1
    end

    position = b0
    counter = 1
    current_color = instance.color_code[solution.sequence[position]]
    position += 1
    first_violation = 0
    while position <= instance.nb_cars
        if instance.color_code[solution.sequence[position]] == current_color
            counter += 1
        else
            counter = 1
            current_color = instance.color_code[solution.sequence[position]]
        end
        if counter > instance.nb_paint_limitation
            if position >= b0
                if first_violation == 0
                    first_violation = position
                end
                car_HPRC_value = instance.HPRC_keys[solution.sequence[position]]
                car_pos = 1
                len = length(instance.same_HPRC[car_HPRC_value])
                while (car_pos <= len) && (instance.color_code[instance.same_HPRC[car_HPRC_value][car_pos]] == current_color)
                    car_pos += 1
                end
                if car_pos <= len
                    move_exchange!(solution, position, car_pos, instance)
                    counter = first_violation - position
                    if first_violation >= position
                        position = first_violation - 1
                    else
                        position -= 1
                    end
                    first_violation = 0
                elseif first_violation >= position
                    position -= 2
                    if counter == 2*instance.nb_paint_limitation
                        # This strategy can't repair
                        position = first_violation
                    end
                end
            else
                if first_violation != 0
                    position = first_violation
                else
                    # pcc are not possible in nb_late_prec_day so we keep the counter
                end
            end
        end
        position += 1
    end
end

"""
    second_repair!(solution::Solution, instance::Instance)

Second strategy
"""
function second_strategy_repair!(solution::Solution, instance::Instance)
    # Second strategy
    first_violation = find_first_violation(solution, instance)
    b0 = instance.nb_late_prec_day+1
    while first_violation != -1
        # Compute the best_insertion index
        solution_value = weighted_sum(solution, instance)
        cost_insertion = zeros(instance.nb_cars)
        for i in b0:instance.nb_cars
            if i <= instance.nb_late_prec_day || i == first_violation # avoiding insterting in sequence from the precedent day and at the same position
                cost_insertion[i] = Inf
            else
                #checking for validity of insertion :  if it does not add an other violation of Paint batch constraint
                batch_size = 1
                batch_color = instance.color_code[solution.sequence[first_violation]]
                if instance.color_code[solution.sequence[i]] == batch_color # if inserting create a batch of length greater than one
                    batch_size += 1
                    j = 1
                    same_batch_sup = true
                    same_batch_inf = true
                    while same_batch_inf || same_batch_sup
                        if i-j > 0 && instance.color_code[solution.sequence[i-j]] == batch_color && same_batch_inf
                            batch_size += 1
                        else
                            same_batch_inf = false
                        end

                        if i+j <= instance.nb_cars && instance.color_code[solution.sequence[i+j]] == batch_color && same_batch_sup
                            batch_size += 1
                        else
                            same_batch_sup = false
                        end
                        j += 1
                    end
                end
                if batch_size > instance.nb_paint_limitation
                    cost_insertion[i] = Inf
                else
                    solution_copy = deepcopy(solution)
                    move_insertion!(solution, first_violation, i, instance)
                    cost_insertion[i] = weighted_sum(solution_copy, instance) - solution_value
                end
            end
        end
        best_insertion = argmin(cost_insertion)[1]
        move_insertion!(solution, first_violation, best_insertion, instance)

        first_violation = find_first_violation(solution, instance)
    end
end

"""
    repair!(solution::Solution, instance::Instance)

Apply 2 repair strategies on `solution`.
"""
function repair!(solution::Solution, instance::Instance)
    if solution.colors === nothing
        initialize_batches!(solution, instance)
    end

    first_strategy_repair!(solution, instance)
    second_strategy_repair!(solution, instance)

    @assert find_first_violation(solution, instance) == -1
    return solution
end


using Plots, Test

@testset "Shape Copy" begin
    square = Shape([(0,0),(1,0),(1,1),(0,1)])
    square2 = Shape(square)
    @test square2.x == square.x
    @test square2.y == square.y
end

@testset "Shape Center" begin
    square = Shape([(0,0),(1,0),(1,1),(0,1)])
    @test Plots.center(square) == (0.5,0.5)
end

@testset "Shape Translate" begin
    square = Shape([(0,0),(1,0),(1,1),(0,1)])
    squareUp = Shape([(0,1),(1,1),(1,2),(0,2)])
    squareUpRight = Shape([(1,1),(2,1),(2,2),(1,2)])

    @test Plots.translate(square,0,1).x == squareUp.x
    @test Plots.translate(square,0,1).y == squareUp.y

    @test Plots.center(translate!(square,1)) == (1.5,1.5)
end

@testset "Brush" begin
    @testset "Colors" begin
        baseline = brush(1, RGB(0,0,0))
        @test brush(:black) == baseline
        @test brush("black") == baseline
    end
    @testset "Weight" begin
        @test brush(10).size == 10
        @test brush(0.1).size == 1
    end
    @testset "Alpha" begin
        @test brush(0.4).alpha == 0.4
        @test brush(20).alpha == nothing
    end
    @testset "Bad Argument" begin
        # using test_logs because test_warn seems to not work anymore
        @test_logs (:warn,"Unused brush arg: nothing (Nothing)") begin
            brush(nothing)
        end
    end
end


"""
    ccd2rgb(red::ImageHDU, green::ImageHDU, blue::ImageHDU; stretch = identity, shape_out = size(red))
    ccd2rgb(red::Tuple{AbstractMatrix, WCSTransform}, green::Tuple{AbstractMatrix, WCSTransform},
                     blue::Tuple{AbstractMatrix, WCSTransform}; stretch = identity, shape_out = size(red[1]))

Converts 3 grayscale ImageHDU into RGB by reprojecting them.

# Arguments
- `red`: Red channel data.
- `green`: Green channel data.
- `blue`: Blue channel data.
- `stretch`: Stretch function applied.
- `shape_out`: Shape of output RGB image.

# Examples
```julia-repl
julia> ccd2rgb(r, b, g, shape_out = (1000,1000))

julia> ccd2rgb(r, b, g, shape_out = (1000,1000), stretch = log)

julia> ccd2rgb(r, b, g, shape_out = (1000,1000), stretch = sqrt)

julia> ccd2rgb(r, b, g, shape_out = (1000,1000), stretch = asinh)
```
"""
function ccd2rgb(red::Tuple{AbstractMatrix, WCSTransform}, green::Tuple{AbstractMatrix, WCSTransform},
                     blue::Tuple{AbstractMatrix, WCSTransform}; stretch = identity, shape_out = size(red[1]))
    red_rp = reproject(red, red[2], shape_out = shape_out)[1]
    green_rp = reproject(green, red[2], shape_out = shape_out)[1]
    blue_rp = reproject(blue, red[2], shape_out = shape_out)[1]

    I = (red_rp .+ green_rp .+ blue_rp) ./ 3
    I .= (x -> stretch(x)/x).(I)

    red_rp .*= I
    green_rp .*= I
    blue_rp .*= I

    m1 = maximum(x->isnan(x) ? -Inf : x, red_rp)
    m2 = maximum(x->isnan(x) ? -Inf : x, green_rp)
    m3 = maximum(x->isnan(x) ? -Inf : x, blue_rp)
    return colorview(RGB, red_rp./m1 , green_rp./m2, blue_rp./m3)
end

ccd2rgb(red::ImageHDU, green::ImageHDU, blue::ImageHDU; stretch = identity, shape_out = size(red)) =
    ccd2rgb((read(red), WCS.from_header(read_header(red, String))[1]), (read(green), WCS.from_header(read_header(green, String))[1]),
            (read(blue), WCS.from_header(read_header(blue, String))[1]), stretch = stretch, shape_out = shape_out)


@inline sqr(x) = x*x
@inline cub(x) = x*x*x

modrange(x, r::AbstractUnitRange) = mod(x-first(r), length(r)) + first(r)
modrange(x, (l, u)::Tuple{Real,Real}) = mod(x-l, u-l+1) + l

fmap(fs, x...) = _fmap(x, fs...)
@inline _fmap(x, f, fs...) = (f(x...), _fmap(x, fs...)...)
@inline _fmap(x) = ()

split_flag(f::Flag) = f, f
split_flag(t::Tuple) = t[1], Base.tail(t)

getfirst(f::Flag) = f
getfirst(t::Tuple) = t[1]
getrest(f::Flag) = f
getrest(t::Tuple) = Base.tail(t)

tcollect(f, itp::AbstractInterpolation{T,N}) where {T,N} = _tcollect(ntuple(d->true, Val(N)), f(itp))
@inline _tcollect(ruler, prop) = (getfirst(prop), _tcollect(Base.tail(ruler), getrest(prop))...)
_tcollect(::Tuple{}, prop) = ()

split_trailing(::AbstractArray{T,N}, x) where {T,N} = Base.IteratorsMD.split(x, Val(N))
check1(args) = _check1(true, args...)
@inline _check1(tf, a, args...) = _check1(tf & (a == 1), args...)
_check1(tf) = tf

# These are not inferrable for mixed-type tuples, so when that's important use `getfirst`
# and `getrest` instead.
iextract(f::Flag, d) = f
iextract(t::Tuple, d) = t[d]

splitgrouped(prs::Tuple{Vararg{NTuple{2,Any}}}) = first.(prs), last.(prs)
splitgrouped(prs::Tuple{Vararg{NTuple{3,Any}}}) = first.(prs), middle.(prs), last.(prs)
middle(t::Tuple{Any,Any,Any}) = t[2]

fast_trunc(::Type{Int}, x) = unsafe_trunc(Int, x)
fast_trunc(::Type{Int}, x::Rational) = x.num ÷ x.den

# Slot-substitution guided by a `ruler` tuple. Substitution occurs when `default` has the same
# length as `ruler`.
@inline substitute_ruled(default, ruler, subst) = (default[1], substitute_ruled(Base.tail(default), ruler, Base.tail(subst))...)
@inline substitute_ruled(default::NTuple{N,Any}, ruler::NTuple{N,Any}, subst) where N =
    (subst[1], substitute_ruled(Base.tail(default), ruler, Base.tail(subst))...)
substitute_ruled(default::Tuple{}, ruler::NTuple{N,Any}, subst) where N = ()

@inline skip_nointerp(x, rest...) = (x, skip_nointerp(rest...)...)
@inline skip_nointerp(::NoInterp, rest...) = skip_nointerp(rest...)
skip_nointerp() = ()

skip_flagged_nointerp(itp::AbstractInterpolation, xs) = skip_flagged_nointerp(tcollect(itpflag, itp), xs)
skip_flagged_nointerp(itpflags::Tuple{NoInterp,Vararg{Any}}, xs) = skip_flagged_nointerp(Base.tail(itpflags), Base.tail(xs))
skip_flagged_nointerp(itpflags::Tuple, xs) = (xs[1], skip_flagged_nointerp(Base.tail(itpflags), Base.tail(xs))...)
skip_flagged_nointerp(::Tuple{}, ::Tuple{}) = ()

@inline sumvals(val, δval, args...) = sumvals(val+δval, args...)
@inline sumvals(val, ::Nothing, args...) = sumvals(val, args...)
sumvals(val) = val

@inline promote_typeof(a, b, args...) = _promote_typeof(promote_type(typeof(a), typeof(b)), args...)
@inline _promote_typeof(::Type{T}, a, args...) where T = _promote_typeof(promote_type(T, typeof(a)), args...)
_promote_typeof(::Type{T}) where T = T

## Vector indexing utilities
# Drop dimensions associated with "scalar" WeightedIndexes
shape(i::WeightedIndex, rest...)  = shape(rest...)
shape(i::Number, rest...)         = shape(rest...)
shape(v::AbstractVector, rest...) = (axes1(v), shape(rest...)...)
shape() = ()

@inline keepvectors(v::Vector{Int}, rest...) = (v, keepvectors(rest...)...)
@inline keepvectors(x, rest...) = keepvectors(rest...)
keepvectors() = ()

@inline lispyprod(p, v::Vector{T}, rest...) where T = lispyprod(p*zero(T), rest...)
@inline lispyprod(p, x::Number, rest...) = lispyprod(p*x, rest...)
lispyprod(p) = p

const onevec = 1:1
const emptyvec = 1:0
@inline inbounds(itp::AbstractInterpolation, x...) = _inbounds.(bounds(itp), x)
_inbounds((l,u)::Tuple{Number,Number}, x::Number) = ifelse(l <= x <= u, onevec, emptyvec)
function _inbounds((l,u)::Tuple{Number,Number}, x::AbstractVector)
    ret = Int[]
    for i in eachindex(x)
        l <= x[i] <= u && push!(ret, i)
    end
    ret
end

function getindex!(dest, itp, xs...)
    for (i, x) in zip(eachindex(dest), Iterators.product(xs...))
        dest[i] = itp(x...)
    end
    return dest
end


module Ops21

import Knet.Ops20
using AutoGrad: AutoGrad, @primitive
using SpecialFunctions: erf

const dropout = Ops20.dropout; export dropout
const bmm = Ops20.bmm; export bmm
include("activation.jl"); export elu, gelu, relu, selu, sigm
include("mmul.jl"); export mmul

end


# Pretty slow solution. Both part 1 and 2 take around 20 seconds
using DataStructures
function solve(game::Vector{Vector{Char}})
    game = [line[2:end-1] for line in game[2:end-1]]
    sy, sx = length(game), length(game[1])
    depth = sy - 1
    nswe = [(-1, 0), (1, 0), (0, 1), (0, -1)]
    adjacent(y, x) = [(y+ya, x+xa) for (ya, xa)=nswe if (y+ya)∈1:sy && (x+xa)∈1:sx]
    is_legal(yx) = occursin(game[yx[1]][yx[2]], ".ABCD")
    adjacent_legal = Dict()
    legal_coords = []
    target_lookup = Dict(a => ("."^i * a^(depth-i) for i=1:depth) for a in "ABCD")
    allowed_lookup = [Set(), Set((1, xa) for xa in (1, 2, 4, 6, 8, 10, 11))]

    for y=1:sy, x=1:sx
        if is_legal((y, x))
            adjacent_legal[(y, x)] = filter(is_legal, adjacent(y, x))
            (game[y][x] != '.' || game[y+1][x] == '#') && push!(legal_coords, (y, x))
        end
    end
    state_queue = PriorityQueue(game => 0)
    states_visited = Set{UInt64}()
    while !isempty(state_queue)
        state, cost = dequeue_pair!(state_queue)
        push!(states_visited, hash(state))
        if all(state[ys][xs] == ('A'+(xs-3)÷2) for (ys, xs) in legal_coords[end-(4*depth-1):end])
            return cost
        end
        for (y, x) in legal_coords
            if (a = state[y][x]) != '.'
                target = target_lookup[a]
                allowed = allowed_lookup[(y != 1) + 1]
                target_x = 3 + 2(a - 'A')
                slot = join(state[2+d][target_x] for d=0:depth-1)
                slot_pos = slot ∈ target ? (1 + sum(collect(slot) .== '.'), target_x) : (0, 0)
                isempty(allowed) && slot_pos == (0, 0) && continue
                visited = Set{Tuple{Int64, Int64}}()
                queue = Queue{Tuple{Int64, Int64, Int64}}()
                enqueue!(queue, (cost, y, x))
                while !isempty(queue)
                    new_cost, yc, xc = dequeue!(queue)
                    if slot_pos == (yc, xc) || (yc, xc) ∈ allowed
                        new_state = deepcopy(state)
                        new_state[yc][xc] = a
                        new_state[y][x] = '.'
                        if hash(new_state) ∉ states_visited
                            state_queue[new_state] = min(new_cost, get(state_queue, new_state, typemax(Int)))
                        end
                    end
                    push!(visited, (yc, xc))
                    for (ya, xa) in adjacent_legal[(yc, xc)]
                        if state[ya][xa] == '.' && (ya, xa) ∉ visited
                            adj_cost = new_cost + (10^(a - 'A'))
                            push!(visited, (ya, xa))
                            enqueue!(queue, (adj_cost, ya, xa))
                        end
                    end
                end
            end
        end
    end
end

function main()
    game = readlines() .|> collect
    println(solve(game))

    insert!(game, length(game)-1, collect("  #D#C#B#A#  "))
    insert!(game, length(game)-1, collect("  #D#B#A#C#  "))
    println(solve(game))
end
main()


"""
    AsynchronousExecution

Functor object for using asynchronous execution in a progressive-hedging algorithm (assuming multiple Julia cores are available). Create by supplying an [`Asynchronous`](@ref) object through `execution` in the `ProgressiveHedgingSolver` factory function and then pass to a `StochasticPrograms.jl` model.

"""
struct AsynchronousExecution{T <: AbstractFloat,
                             A <: AbstractVector,
                             PT <: AbstractPenaltyterm} <: AbstractProgressiveHedgingExecution
    subworkers::Vector{SubWorker{T,A,PT}}
    work::Vector{Work}
    finalize::Vector{Work}
    progressqueue::ProgressQueue{T}
    x̄::Vector{RemoteRunningAverage{A}}
    δ::Vector{RemoteRunningAverage{T}}
    iterates::RemoteIterates{A}
    r::IteratedValue{T}
    active_workers::Vector{Future}
    # Bookkeeping
    subobjectives::Vector{A}
    finished::Vector{Int}
    # Parameters
    κ::T

    function AsynchronousExecution(κ::T,
                                   ::Type{T}, ::Type{A}, ::Type{PT}) where {T <: AbstractFloat,
                                                                            A <: AbstractVector,
                                                                            PT <: AbstractPenaltyterm}
        return new{T,A,PT}(Vector{SubWorker{T,A,PT}}(undef, nworkers()),
                           Vector{Work}(undef, nworkers()),
                           Vector{Work}(undef, nworkers()),
                           RemoteChannel(() -> Channel{Progress{T}}(4 * nworkers())),
                           Vector{RemoteRunningAverage{A}}(undef, nworkers()),
                           Vector{RemoteRunningAverage{T}}(undef, nworkers()),
                           RemoteChannel(() -> IterationChannel(Dict{Int,A}())),
                           RemoteChannel(() -> IterationChannel(Dict{Int,T}())),
                           Vector{Future}(undef, nworkers()),
                           Vector{A}(),
                           Vector{Int}(),
                           κ)
    end
end

function initialize_subproblems!(ph::AbstractProgressiveHedging,
                                 execution::AsynchronousExecution{T,A},
                                 scenarioproblems::DistributedScenarioProblems,
                                 penaltyterm::AbstractPenaltyterm) where {T <: AbstractFloat, A <: AbstractVector}
    # Create subproblems on worker processes
    initialize_subproblems!(ph,
                            execution.subworkers,
                            scenarioproblems,
                            penaltyterm)
    # Continue preparation
    @sync begin
        for w in workers()
            execution.work[w-1] = RemoteChannel(() -> Channel{Int}(round(Int,10/execution.κ)), w)
            execution.finalize[w-1] = RemoteChannel(() -> Channel{Int}(1), w)
            execution.x̄[w-1] = RemoteChannel(() -> Channel{RunningAverage{A}}(1), w)
            execution.δ[w-1] = RemoteChannel(() -> Channel{RunningAverage{T}}(1), w)
            @async remotecall_fetch(
                w,
                execution.subworkers[w-1],
                execution.x̄[w-1],
                length(ph.ξ)) do sw, average, xdim
                    subproblems = fetch(sw)
                    x̄ = mapreduce(+, subproblems, init = zeros(T, xdim)) do subproblem
                        x = subproblem.x
                        π = subproblem.probability
                        return π * x
                    end
                    running_average = RunningAverage(x̄, [s.x for s in subproblems])
                    put!(average, running_average)
                    return x̄
                end
            @async remotecall_fetch(
                w,
                execution.subworkers[w-1],
                execution.δ[w-1],
                length(ph.ξ)) do sw, average, xdim
                    running_average =
                        RunningAverage(zero(T),
                                       fill(zero(T), length(fetch(sw))))
                    put!(average, running_average)
                end
            put!(execution.work[w-1], 1)
        end
        # Prepare memory
        push!(execution.subobjectives, zeros(num_subproblems(ph)))
        push!(execution.finished, 0)
        log_val = ph.parameters.log
        ph.parameters.log = false
        log!(ph)
        ph.parameters.log = log_val
    end
    # Initial reductions
    update_iterate!(ph)
    # Init δ₂
    @sync begin
        for w in workers()
            @async remotecall_fetch(
                w,
                execution.subworkers[w-1],
                ph.ξ,
                execution.δ[w-1]) do sw, ξ, δ
                    for (i, subproblem) in enumerate(fetch(sw))
                        subtract!(fetch(δ), i)
                        x = subproblem.x
                        π = subproblem.probability
                        add!(fetch(δ), i, norm(x - ξ, 2) ^ 2, π)
                    end
                end
        end
    end
    update_dual_gap!(ph)
    return nothing
end

function iterate!(ph::AbstractProgressiveHedging, execution::AsynchronousExecution{T}) where T <: AbstractFloat
    wait(execution.progressqueue)
    while isready(execution.progressqueue)
        # Add new cuts from subworkers
        t::Int, i::Int, Q::SubproblemSolution{T} = take!(execution.progressqueue)
        if !(Q.status ∈ AcceptableTermination)
            # Early termination log
            log!(ph; status = Q.status)
            return Q.status
        end
        execution.subobjectives[t][i] = Q.value
        execution.finished[t] += 1
        if execution.finished[t] == num_subproblems(ph)
            # Update objective
            ph.Q_history[t] = current_objective_value(ph, execution.subobjectives[t])
            ph.data.Q = ph.Q_history[t]
        end
    end
    # Project and generate new iterate
    t = ph.data.iterations
    if execution.finished[t] >= execution.κ * num_subproblems(ph)
        # Get dual gap
        update_dual_gap!(ph)
        # Update progress
        @unpack δ₁, δ₂ = ph.data
        ph.primal_gaps[t] = δ₁
        ph.dual_gaps[t] = δ₂
        # Check if optimal
        if check_optimality(ph)
            # Optimal, final log
            log!(ph, optimal = true)
            return MOI.OPTIMAL
        end
        # Update penalty (if applicable)
        update_penalty!(ph)
        # Update iterate
        update_iterate!(ph)
        # Send new work to workers
        put!(execution.iterates, t + 1, ph.ξ)
        put!(execution.r, t + 1, penalty(ph))
        map((w, aw) -> !isready(aw) && put!(w, t + 1), execution.work, execution.active_workers)
        # Prepare memory for next iteration
        push!(execution.subobjectives, zeros(num_subproblems(ph)))
        push!(execution.finished, 0)
        # Log progress
        log!(ph)
    end
    # Dont return a status as procedure should continue
    return nothing
end

function start_workers!(ph::AbstractProgressiveHedging, execution::AsynchronousExecution)
    # Load initial decision
    put!(execution.iterates, 1, ph.ξ)
    put!(execution.r, 1, penalty(ph))
    for w in workers()
        execution.active_workers[w-1] = remotecall(work_on_subproblems!,
                                                   w,
                                                   execution.subworkers[w-1],
                                                   execution.work[w-1],
                                                   execution.finalize[w-1],
                                                   execution.progressqueue,
                                                   execution.x̄[w-1],
                                                   execution.δ[w-1],
                                                   execution.iterates,
                                                   execution.r)
    end
    return nothing
end

function close_workers!(ph::AbstractProgressiveHedging, execution::AsynchronousExecution)
    t = ph.data.iterations-1
    map((w, aw)->!isready(aw) && put!(w, t), execution.finalize, execution.active_workers)
    map((w, aw)->!isready(aw) && put!(w, -1), execution.work, execution.active_workers)
    map(wait, execution.active_workers)
end

function restore_subproblems!(::AbstractProgressiveHedging, execution::AsynchronousExecution)
    restore_subproblems!(execution.subworkers)
    return nothing
end

function resolve_subproblems!(ph::AbstractProgressiveHedging, execution::AsynchronousExecution)
    return nothing
end

function update_iterate!(ph::AbstractProgressiveHedging, execution::AsynchronousExecution)
    ξ_prev = copy(ph.ξ)
    ph.ξ .= sum(average.(fetch.(execution.x̄)))
    # Update δ₁
    ph.data.δ₁ = norm(ph.ξ - ξ_prev, 2) ^ 2
    return nothing
end

function update_subproblems!(ph::AbstractProgressiveHedging, execution::AsynchronousExecution)
    return nothing
end

function update_dual_gap!(ph::AbstractProgressiveHedging, execution::AsynchronousExecution)
    ph.data.δ₂ = sum(average.(fetch.(execution.δ)))
    return nothing
end

function calculate_objective_value(ph::AbstractProgressiveHedging, execution::AsynchronousExecution)
    return calculate_objective_value(execution.subworkers)
end

# API
# ------------------------------------------------------------
function (execution::Asynchronous)(::Type{T}, ::Type{A}, ::Type{PT}) where {T <: AbstractFloat,
                                                                            A <: AbstractVector,
                                                                            PT <: AbstractPenaltyterm}
    return AsynchronousExecution(execution.κ, T, A, PT)
end

function str(::Asynchronous)
    return "Asynchronous "
end


############################################################
# joRHS ## #################################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: January 2017

export joRHS, joRHSexception

############################################################

struct joRHS{vDT<:Number} <: joAbstractLinearOperator{vDT,vDT}
	name::String
	m::Integer
	n::Integer
	info::Info
	geometry::Geometry
	data
end

mutable struct joRHSexception <: Exception
	msg :: String
end

############################################################

## outer constructors

"""
    joRHS
        name::String
        m::Integer
        n::Integer
        info::Info
        geometry::Geometry
        data

Abstract sparse vector for right-hand-sides of the modeling operators. The `joRHS` vector has the\\
dimensions of the full time history of the wavefields, but contains only the data defined at the \\
source or receiver positions (i.e. wavelets or shot records).

Constructor
==========

    joRHS(info, geometry, data)

Examples
========

Assuming `Pr` and `Ps` are projection operators of type `joProjection` and `dobs` and `q` are\\
seismic vectors of type `joData`, then a `joRHS` vector can be created as follows:

    rhs = Pr'*dobs    # right-hand-side with injected observed data
    rhs = Ps'*q    # right-hand-side with injected wavelet

"""
function joRHS(info,geometry,data;vDT::DataType=Float32)
	vDT == Float32 || throw(joRHSexception("Domain type not supported"))
	# length of vector
	m = info.n * sum(info.nt)
	n = 1
	return joRHS{Float32}("joRHS",m,n,info,geometry,data)
end

####################################################################
## overloaded Base functions

# conj(jo)
conj{vDT}(A::joRHS{vDT}) =
	joRHS{vDT}("conj("*A.name*")",A.m,A.n,A.info,A.geometry,A.data)

# transpose(jo)
transpose{vDT}(A::joRHS{vDT}) =
	joRHS{vDT}(""*A.name*".'",A.n,A.m,A.info,A.geometry,A.data)

# ctranspose(jo)
ctranspose{vDT}(A::joRHS{vDT}) =
	joRHS{vDT}(""*A.name*"'",A.n,A.m,A.info,A.geometry,A.data)

####################################################################

# +(joRHS,joRHS)
function +{avDT,bvDT}(A::joRHS{avDT}, B::joRHS{bvDT})

	# Error checking
	size(A) == size(B) || throw(joRHSexception("shape mismatch"))
	compareInfo(A.info, B.info) == true || throw(joRHSexception("info mismatch"))
	isequal(A.geometry.nt,B.geometry.nt) == true || throw(joRHSexception("sample number mismatch"))
	isequal(A.geometry.dt,B.geometry.dt) == true || throw(joRHSexception("sample interval mismatch"))
	isequal(A.geometry.t,B.geometry.t) == true || throw(joRHSexception("recording time mismatch"))

	# Size
	m = A.info.n * sum(A.info.nt)
	n = 1

	# merge geometries and data
	xloc = Array{Any}(A.info.nsrc)
	yloc = Array{Any}(A.info.nsrc)
	zloc = Array{Any}(A.info.nsrc)
	dt = Array{Any}(A.info.nsrc)
	nt = Array{Any}(A.info.nsrc)
	t = Array{Any}(A.info.nsrc)
	data = Array{Any}(A.info.nsrc)

	for j=1:A.info.nsrc
		xloc[j] = [vec(A.geometry.xloc[j])' vec(B.geometry.xloc[j])']
		yloc[j] = [vec(A.geometry.yloc[j])' vec(B.geometry.yloc[j])']
		zloc[j] = [vec(A.geometry.zloc[j])' vec(B.geometry.zloc[j])']
		dt[j] = A.geometry.dt[j]
		nt[j] = A.geometry.nt[j]
		t[j] = A.geometry.t[j]
		data[j] = [A.data[j] B.data[j]]
	end
	geometry = Geometry(xloc,yloc,zloc,dt,nt,t)
	nvDT = promote_type(avDT,bvDT)

	return joRHS{nvDT}("joRHS",m,n,A.info,geometry,data)
end

# -(joRHS,joRHS)
function -{avDT,bvDT}(A::joRHS{avDT}, B::joRHS{bvDT})

	# Error checking
	size(A) == size(B) || throw(joRHSexception("shape mismatch"))
	compareInfo(A.info, B.info) == true || throw(joRHSexception("info mismatch"))
	isequal(A.geometry.nt,B.geometry.nt) == true || throw(joRHSexception("sample number mismatch"))
	isequal(A.geometry.dt,B.geometry.dt) == true || throw(joRHSexception("sample interval mismatch"))
	isequal(A.geometry.t,B.geometry.t) == true || throw(joRHSexception("recording time mismatch"))

	# Size
	m = A.info.n * sum(A.info.nt)
	n = 1

	# merge geometries and data
	xloc = Array{Any}(A.info.nsrc)
	yloc = Array{Any}(A.info.nsrc)
	zloc = Array{Any}(A.info.nsrc)
	dt = Array{Any}(A.info.nsrc)
	nt = Array{Any}(A.info.nsrc)
	t = Array{Any}(A.info.nsrc)
	data = Array{Any}(A.info.nsrc)

	for j=1:A.info.nsrc
		xloc[j] = [vec(A.geometry.xloc[j])' vec(B.geometry.xloc[j])']
		yloc[j] = [vec(A.geometry.yloc[j])' vec(B.geometry.yloc[j])']
		zloc[j] = [vec(A.geometry.zloc[j])' vec(B.geometry.zloc[j])']
		dt[j] = A.geometry.dt[j]
		nt[j] = A.geometry.nt[j]
		t[j] = A.geometry.t[j]
		data[j] = [A.data[j] -B.data[j]]
	end
	geometry = Geometry(xloc,yloc,zloc,dt,nt,t)
	nvDT = promote_type(avDT,bvDT)

	return joRHS{nvDT}("joRHS",m,n,A.info,geometry,data)
end


using CompetitiveEquilibriumInvestments
using JuMP
using Gurobi
using Test

const CEI = CompetitiveEquilibriumInvestments

function fix_discrete_reoptimize!(m::Model, solver)

    # Go through all variables and fix binaries / integers
    # such that the problem is convex and has duals
    vars = all_variables(m)
    vals = value.(vars)
    for (var, val) in zip(vars, vals)
        if is_binary(var) || is_integer(var)
            is_binary(var) ? unset_binary(var) : unset_integer(var)
            fix(var, val, force=true)
        end
    end

    set_optimizer(m, solver)
    @time optimize!(m)

    return m

end

true && @testset "Toy Problem" begin

    R = 3
    G1 = 1
    G2 = 0
    G3 = 0
    I = 3
    T = 12
    P = 4

    periodnames = ["P$i" for i in 1:P]
    regionnames = ["A", "B", "C"]

    # Generation / Storage Technologies

    thermaltechs = ThermalGenerators(
        ["Tech A"], ["Firm 1"], [5.], [15.], [6], [6], [5.], [5.], [0.95])

    variabletechs = VariableGenerators(String[], String[], Float64[], Float64[])

    storagetechs = StorageDevices(
        String[], String[], Float64[], Float64[], Float64[],
        Float64[], Float64[], Float64[], Float64[])

    interfaces = Interfaces{R}(["1-2", "1-3", "2-3"], [(1,2), (1,3), (2,3)])

    techs = Technologies(thermaltechs, variabletechs, storagetechs, interfaces)

    # Initial Conditions

    thermalstart = InitialInvestments(zeros(Int, R, G1), zeros(Int, R, G1))
    emptystart = InitialInvestments(zeros(Int, R, 0), zeros(Int, R, 0))
    initconds = InitialConditions(thermalstart, emptystart, emptystart)

    # Investment Context

    inv_thermal = ResourceInvestments(
        fill(1000., R, G1), fill(10000., R, G1), fill(0., R, G1),
        fill(1, R, G1), fill(1, R, G1), fill(0, R, G1), fill(0, R, G1),
        fill(0, R, G1))

    inv_empty = ResourceInvestments(
        zeros(Float64, R, 0), zeros(Float64, R, 0), zeros(Float64, R, 0),
        zeros(Int, R, 0), zeros(Int, R, 0), zeros(Int, R, 0), zeros(Int, R, 0),
        zeros(Int, R, 0))

    invs = Investments(inv_thermal, inv_empty, inv_empty)

    # Operating Context

    ops_thermal =
        ThermalGeneratorOperations{T,P}(fill(200., R, G1), fill(4., R, G1),
                                        fill(2., R, G1), fill(2., R, G1))

    ops_variable =
        VariableGeneratorOperations(
            zeros(Float64, R, G2), zeros(Float64, R, G2),
            zeros(Float64, R, G2), zeros(Float64, R, G2, T, P))

    ops_storage =
        StorageOperations{T,P}(zeros(Float64, R, G3), zeros(Float64, R, G3))

    ops_transmission =
        TransmissionOperations{R,T,P}([10., 10, 10])

    ops = Operations(ops_thermal, ops_variable, ops_storage, ops_transmission)

    # Market Context

    base = [10. 15 20 25 30 35 36 31 26 21 16 11]
    r_scale = [1.0, 0.9, 1.1]
    p_scale = reshape([0.8, 1.0, 1.2, 1.4], 1, 1, P)
    load = base .* r_scale .* p_scale
    reserve = 0.1 .* load

    capacity = CapacityMarket(1000., 500., 10., 15., 20.)
    rec = RECMarket(0., 50.)
    energy = EnergyMarket(load, fill(1e5, R))
    raisereserve = RaiseReserveMarket(reserve, fill(4e4, R))
    lowerreserve = LowerReserveMarket(reserve, fill(4e4, R))

    markets = Markets(capacity, rec, energy, raisereserve, lowerreserve)

    periodweights = fill(13., 4)

    # InvestmentProblem

    discountrate = 0.98
    p = InvestmentProblem(regionnames, periodnames, techs, initconds,
                          discountrate, "RootScenario",
                          invs, ops, markets, periodweights, Gurobi.Optimizer)
    solve!(p)
    fix_discrete_reoptimize!(p.model, Gurobi.Optimizer)
    report(joinpath(dirname(@__FILE__), "toymodel"), p)

end

true && @testset "RTS" begin
    p = InvestmentProblem(
        "/home/gord/work/EMIS/EMISPreprocessing/output",
        optimizer_with_attributes(Gurobi.Optimizer,
            "MIPGap" => 0.005, "MIPGapAbs" => 220e6))
    solve!(p, debug=true)
    fix_discrete_reoptimize!(p.model, Gurobi.Optimizer)
    report(joinpath(dirname(@__FILE__), "rts"), p)
end


import Base: copy, copyto!, ==
import Random: rand, seed!
import RandomNumbers: gen_seed, seed_type

# Generic functions

@inline seed!(r::AbstractPCG{StateType}, seed::Integer=gen_seed(StateType)) where
    StateType <: PCGUInt = pcg_seed!(r, seed % StateType)
@inline seed!(r::PCGStateSetseq{StateType}, seed::NTuple{2, Integer}=gen_seed(StateType, 2)) where
    StateType <: PCGUInt = pcg_seed!(r, seed[1] % StateType, seed[2] % StateType)

@inline seed_type(::Type{PCGStateOneseq{T, T1, T2}}) where {T, T1, T2} = T
@inline seed_type(::Type{PCGStateMCG{   T, T1, T2}}) where {T, T1, T2} = T
@inline seed_type(::Type{PCGStateUnique{T, T1, T2}}) where {T, T1, T2} = T
@inline seed_type(::Type{PCGStateSetseq{T, T1, T2}}) where {T, T1, T2} = NTuple{2, T}

function copyto!(dest::T, src::T) where T <: AbstractPCG
    dest.state = src.state
    dest
end
function copyto!(dest::T, src::T) where T <: PCGStateSetseq
    dest.state = src.state
    dest.inc = src.inc
    dest
end

copy(src::T) where T <: AbstractPCG = copyto!(T(), src)

==(r1::T, r2::T) where T <: AbstractPCG = r1.state == r2.state
==(r1::T, r2::T) where T <: PCGStateSetseq = r1.state == r2.state && r1.inc == r2.inc

@inline function rand(r::AbstractPCG{StateType, MethodType, OutputType}, ::Type{OutputType}) where
        {StateType <: Union{pcg_uints[1:end-1]...}, MethodType <: PCGMethod, OutputType <: PCGUInt}
    old_state = r.state
    pcg_step!(r)
    pcg_output(old_state, MethodType)
end

@inline function rand(r::AbstractPCG{UInt128, MethodType, OutputType}, ::Type{OutputType}) where
        {MethodType <: PCGMethod, OutputType <: PCGUInt}
    pcg_step!(r)
    pcg_output(r.state, MethodType)
end

"""
```julia
bounded_rand(r, bound)
```

Producing a random number less than a given `bound` in the output type.
"""
@inline function bounded_rand(s::AbstractPCG{StateType, MethodType, OutputType}, bound::OutputType) where
        {StateType <: PCGUInt, MethodType <: PCGMethod, OutputType <: PCGUInt}
    threshold = (-bound) % bound
    r = rand(s, OutputType)
    while r < threshold
        r = rand(s, OutputType)
    end
    r % bound
end

"""
```julia
advance!(r, Δ)
```

Advance a PCG object `r` for `Δ` steps.

# Examples
```jldoctest
julia> r = PCGStateSetseq(UInt64, PCG_RXS_M_XS, (123, 321))
PCGStateSetseq{UInt64,Val{:RXS_M_XS},UInt64}(0x45389f8b27528b29, 0x0000000000000283)

julia> A = rand(r, UInt64, 2);

julia> p = rand(r);

julia> r
PCGStateSetseq{UInt64,Val{:RXS_M_XS},UInt64}(0x9b1fc763ae0ad702, 0x0000000000000283)

julia> advance!(r, -3)
PCGStateSetseq{UInt64,Val{:RXS_M_XS},UInt64}(0x45389f8b27528b29, 0x0000000000000283)

julia> @Test.test A == rand(r, UInt64, 2)
Test Passed

julia> @Test.test p == rand(r)
Test Passed
```
"""
@inline function advance!(r::AbstractPCG{StateType}, Δ::Integer) where StateType <: PCGUInt
    pcg_advance!(r, Δ % StateType)
    r
end

# Constructors.
for (pcg_type_t, uint_type, method_symbol, output_type) in PCG_LIST
    pcg_type = Symbol("PCGState$pcg_type_t")
    method = Val{method_symbol}

    if pcg_type_t != :Setseq
        @eval function $pcg_type(output_type::Type{$output_type}, method::Type{$method},
                seed::Integer=gen_seed($uint_type))
            s = $pcg_type{$uint_type, method, output_type}()
            seed!(s, seed)
            s
        end
    else
        @eval function $pcg_type(output_type::Type{$output_type}, method::Type{$method},
                seed::NTuple{2, Integer}=gen_seed($uint_type, 2))
            s = $pcg_type{$uint_type, method, output_type}()
            seed!(s, seed)
            s
        end
    end
end

"""
```julia
PCGStateOneseq{StateType<:PCGUInt, MethodType<:PCGMethod, OutputType<:PCGUInt} <:
    AbstractPCG{StateType, MethodType, OutputType}
PCGStateOneseq([seed])
PCGStateOneseq(output_type[, seed])
PCGStateOneseq(method[, seed])
PCGStateOneseq(output_type[, method, seed])
```

PCG generator with *single streams*, where all instances use the same fixed constant, thus the RNG always
somewhere in same sequence.

`seed` is an `Integer` which will be automatically converted to the state type.

`output_type` is the type of the PCG's output. If missing it is set to `UInt64`.

`method` is one of the [`PCGMethod`](@ref). If missing it is set to `PCG_XSH_RS`.

See [`PCG_LIST`](@ref) for the available parameter combinations.
"""
PCGStateOneseq(seed::Integer=gen_seed(UInt128)) = PCGStateOneseq(UInt64, PCG_XSH_RS, seed)
PCGStateOneseq(::Type{T}, seed::Integer=gen_seed(UInt128)) where
    T <: PCGUInt = PCGStateOneseq(T, PCG_XSH_RS, seed)
PCGStateOneseq(::Type{T}, seed::Integer=gen_seed(UInt128)) where
    T <: PCGMethod = PCGStateOneseq(UInt64, T, seed)

"""
```julia
PCGStateMCG{StateType<:PCGUInt, MethodType<:PCGMethod, OutputType<:PCGUInt} <:
    AbstractPCG{StateType, MethodType, OutputType}
PCGStateMCG([seed])
PCGStateMCG(output_type[, seed])
PCGStateMCG(method[, seed])
PCGStateMCG(output_type[, method, seed])
```

PCG generator with *MCG*, where the increment is zero, resulting in a single stream and reduced period.

`seed` is an `Integer` which will be automatically converted to the state type.

`output_type` is the type of the PCG's output. If missing it is set to `UInt64`.

`method` is one of the [`PCGMethod`](@ref). If missing it is set to `PCG_XSH_RS`.

See [`PCG_LIST`](@ref) for the available parameter combinations.
"""
PCGStateMCG(seed::Integer=gen_seed(UInt128)) = PCGStateMCG(UInt64, PCG_XSH_RS, seed)
PCGStateMCG(::Type{T}, seed::Integer=gen_seed(UInt128)) where
    T <: PCGUInt = PCGStateMCG(T, PCG_XSH_RS, seed)
PCGStateMCG(::Type{T}, seed::Integer=gen_seed(UInt128)) where
    T <: PCGMethod = PCGStateMCG(UInt64, T, seed)

"""
```julia
PCGStateSetseq{StateType<:PCGUInt, MethodType<:PCGMethod, OutputType<:PCGUInt} <:
    AbstractPCG{StateType, MethodType, OutputType}
PCGStateSetseq([seed])
PCGStateSetseq(output_type[, seed])
PCGStateSetseq(method[, seed])
PCGStateSetseq(output_type[, method, seed])
```

PCG generator with *specific streams*, where the constant can be changed at any time, selecting a different
random sequence.

`seed` is a `Tuple` of two `Integer`s which will both be automatically converted to the state type.

`output_type` is the type of the PCG's output. If missing it is set to `UInt64`.

`method` is one of the [`PCGMethod`](@ref). If missing it is set to `PCG_XSH_RS`.

See [`PCG_LIST`](@ref) for the available parameter combinations.
"""
PCGStateSetseq(seed::NTuple{2, Integer}=gen_seed(UInt128, 2)) = PCGStateSetseq(UInt64, PCG_XSH_RS, seed)
PCGStateSetseq(::Type{T}, seed::NTuple{2, Integer}=gen_seed(UInt128, 2)) where
    T <: PCGUInt = PCGStateSetseq(T, PCG_XSH_RS, seed)
PCGStateSetseq(::Type{T}, seed::NTuple{2, Integer}=gen_seed(UInt128, 2)) where
    T <: PCGMethod = PCGStateSetseq(UInt64, T, seed)

"""
```julia
PCGStateUnique{StateType<:PCGUInt, MethodType<:PCGMethod, OutputType<:PCGUInt} <:
    AbstractPCG{StateType, MethodType, OutputType}
PCGStateUnique([seed])
PCGStateUnique(output_type[, seed])
PCGStateUnique(method[, seed])
PCGStateUnique(output_type[, method, seed])
```

PCG generator with *unique streams*, where the constant is based on the memory address of the object, thus
every RNG has its own unique sequence.

`seed` is an `Integer` which will be automatically converted to the state type.

`output_type` is the type of the PCG's output. If missing it is set to `UInt64`.

`method` is one of the [`PCGMethod`](@ref). If missing it is set to `PCG_XSH_RS`.

See [`PCG_LIST`](@ref) for the available parameter combinations.
"""
PCGStateUnique(seed::Integer=gen_seed(UInt128)) = PCGStateUnique(UInt64, PCG_XSH_RS, seed)
PCGStateUnique(::Type{T}, seed::Integer=gen_seed(UInt128)) where
    T <: PCGUInt = PCGStateUnique(T, PCG_XSH_RS, seed)
PCGStateUnique(::Type{T}, seed::Integer=gen_seed(UInt128)) where
    T <: PCGMethod = PCGStateUnique(UInt64, T, seed)


## ---------------  H Production

abstract type AbstractHcProdSwitches end
abstract type AbstractHcProdSet <: ModelObject end
abstract type AbstractHcProdFct end

Base.show(io :: IO, h :: AbstractHcProdSet) =
    print(io, typeof(h));

## -----------  Access routines

StructLH.describe(h :: AbstractHcProdSwitches) = settings_table(h);

settings_table(h :: AbstractHcProdSet) = settings_table(h.switches);

n_colleges(h :: AbstractHcProdSet) = h.nc;

cal_delta_h!(h :: AbstractHcProdSwitches) = (h.calDeltaH = true);
cal_delta_h(h :: AbstractHcProdSwitches) = h.calDeltaH;
cal_delta_h(h :: AbstractHcProdSet) = cal_delta_h(h.switches);
delta_h(h :: AbstractHcProdSwitches) = h.deltaH;
delta_h(h :: AbstractHcProdSet) = h.deltaH;
delta_h(h :: AbstractHcProdFct) = h.deltaH;

time_exp(h :: AbstractHcProdSet) = h.timeExp;
time_exp(h :: AbstractHcProdFct) = h.timeExp;
h_exp(h :: AbstractHcProdFct) = h.hExp;
set_hexp_lb!(switches :: AbstractHcProdSwitches, hExpLb :: Double) =
    switches.hExpLb = hExpLb;

# Not defined for bounded case
has_tfp(h :: AbstractHcProdFct) = true;
tfp(h :: AbstractHcProdFct) = h.tfp;
set_tfp(h :: AbstractHcProdFct, tfp :: Double) = (h.tfp = tfp);
tfp(h :: AbstractHcProdSet, ic) = values(h.tfpV)[ic];

function h_exp(h :: AbstractHcProdSet)
    if h.switches.sameExponents
        return h.timeExp;
    else
        return h.hExp;
    end
end

function same_exponents!(s :: AbstractHcProdSwitches)
    s.sameExponents = true;
    s.calHExp = false;
end

function separate_exponents!(s :: AbstractHcProdSwitches)
    s.sameExponents = false;
    s.calTimeExp = true;
    s.calHExp = true;
end

same_exponents(h :: AbstractHcProdSwitches) = h.sameExponents;

min_time_per_course(h :: AbstractHcProdFct) = h.minTimePerCourse;

## ----------  Generic methods

# Study time per course, from total study time.
function study_time_per_course(hS :: AbstractHcProdFct, sTimeTotalV, nTriedV)
    return max.(min_time_per_course(hS),  study_time_per_course_no_min(hS, sTimeTotalV, nTriedV))
end

study_time_per_course_no_min(hS :: AbstractHcProdFct, sTimeTotalV, nTriedV) =
    sTimeTotalV ./ nTriedV .- hS.timePerCourse;

# # Study time total, from study time per course.
# function study_time_total(hS :: AbstractHcProdFct, sTimePerCourseV, nTriedV)
#     sTimeV = (max.(min_time_per_course(hS), sTimePerCourseV) .+ hS.timePerCourse) .* nTriedV;
#     return sTimeV
# end

# Study time per course from choice of study time per course (not taking out the fixed time per course).
# function study_time_per_course_from_choice_per_course(hS :: AbstractHcProdFct,
#     sTimeV, nTriedV)

#     sTimePerCourseV = max.(hS.minTimePerCourse, sTimeV .- hS.timePerCourse);
#     return sTimePerCourseV;
# end

"""
	$(SIGNATURES)

h next period. h is restricted to never decline. Study time is total.
"""
hprime(hS :: AbstractHcProdFct, abil, h, h0, sTime, nTried) =
    max.(h, h .* (1.0 .- delta_h(hS)) .+ dh(hS, abil, h, h0, sTime, nTried));

# Path h[t] for given endowments and decisions
function h_path(hS :: AbstractHcProdFct, abil :: Double, h0 :: Double,
    sTimeV, nTriedV)

    T = length(sTimeV);
    hV = zeros(Double, T+1);
    hV[1] = h0;
    for t = 1 : T
        hV[t+1] = hprime(hS, abil, hV[t], h0, sTimeV[t], nTriedV[t]);
    end
    return hV
end

## ----------  Construction

function init_h_depreciation(switches :: AbstractHcProdSwitches)
    deltaH = delta_h(switches);
    pDeltaH = Param(:deltaH, ldescription(:ddh), lsymbol(:ddh),
        deltaH, deltaH, 0.0, 0.1, cal_delta_h(switches));
    return pDeltaH
end

# Fixed time per course that gets subtracted from study time
# Study time per course = study time / nCourses - fixed time per course.
function init_time_per_course()
    # Bounds are arbitrary
    mCourses = data_to_model_courses(1);
    # Time endowment: about 120 hours per week.
    # Normal load: 10 courses per year (5 courses at a time).
    # 1 hour per course => 5 hours per week => 5/120 of time endowment.
    timePerCourseLb = hours_per_week_to_mtu(0.1 / mCourses);
    timePerCourseUb = hours_per_week_to_mtu(2.0 / mCourses);
    timePerCourse = hours_per_week_to_mtu(0.2 / mCourses);

    return Param(:timePerCourse,  ldescription(:sTimeFixed), lsymbol(:sTimeFixed),
        timePerCourse, timePerCourse,
        timePerCourseLb, timePerCourseUb, true)
end

# ----------


using Test
using MianalyzerBackend
const GMeans = MianalyzerBackend.GMeans

@testset "GMeans" begin
  @testset ".call" begin
    actual = GMeans.call([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
    @test actual[:clusters] == [[[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]]
    @test actual[:centers] == [[3, 3]]
  end
end


abstract type HuffmanTreeNode{T<:SourceSymbol} end

# construct from tuples
HuffmanTreeNode(   t::Tuple{T, <:Real}) where T<:SourceSymbol = HuffmanSymbolNode{T}(t...)
HuffmanTreeNode{T}(t::Tuple{T, <:Real}) where T<:SourceSymbol = HuffmanSymbolNode{T}(t...)
HuffmanTreeNode{T}(t::Tuple{UInt8, <:Real}) where T<:SourceSymbol = HuffmanSymbolNode{T}(t...)
HuffmanTreeNode(t::Tuple{<:HuffmanTreeNode{T}, <:HuffmanTreeNode{T}}) where T<:SourceSymbol =
    HuffmanBranchNode{T}(t...)
HuffmanTreeNode{T}(t::Tuple{<:HuffmanTreeNode{T}, <:HuffmanTreeNode{T}}) where T<:SourceSymbol =
    HuffmanBranchNode{T}(t...)
HuffmanTreeNode{T}(t::Tuple{Tuple, Tuple}) where T<:SourceSymbol =
    HuffmanBranchNode{T}(HuffmanTreeNode{T}(t[1]), HuffmanTreeNode{T}(t[2]))

Base.convert(::Type{T}, t::Tuple) where T <: HuffmanTreeNode = T(t)

Base.show(io::IO, t::HuffmanTreeNode) = print_tree(io, t)
AbstractTrees.nodetype(::HuffmanTreeNode{T}) where T = HuffmanTreeNode{T}

struct HuffmanSymbolNode{T<:SourceSymbol} <: HuffmanTreeNode{T}
    symbol::T
    probability::Float64
    HuffmanSymbolNode{T}(symbol::T, p::Float64=0.0) where T = new(symbol, p)
    HuffmanSymbolNode{T}(symbol::UInt8, p::Float64=0.0) where T = new(symbol, p)
end

const DCSymbolNode = HuffmanSymbolNode{DCSymbol}
const ACSymbolNode = HuffmanSymbolNode{ACSymbol}

HuffmanSymbolNode(symbol::T, p::Float64=0.0) where T = HuffmanSymbolNode{T}(symbol, p)

AbstractTrees.children(s::HuffmanSymbolNode) = ()
AbstractTrees.printnode(io::IO, s::HuffmanSymbolNode) = print(io, s)

function Base.show(io::IO, s::HuffmanSymbolNode{T}) where T <: Union{DCSymbol, ACSymbol}
    compact = get(io, :compact, false)
    if compact
        print(io, T === DCSymbol ? "DCSymbolNode" : "ACSymbolNode")
    else
        print(io, typeof(s))
    end
    print(IOContext(io, :typeinfo=>T), "(", s.symbol, ", ", s.probability, ")")
end

struct HuffmanBranchNode{T<:SourceSymbol} <: HuffmanTreeNode{T}
    n0::HuffmanTreeNode{T} # bit 0
    n1::HuffmanTreeNode{T} # bit 1
    probability::Float64
    HuffmanBranchNode{T}(n0::HuffmanTreeNode{T}, n1::HuffmanTreeNode{T}) where T<:SourceSymbol =
        new(n0, n1, n0.probability + n1.probability)
end
HuffmanBranchNode(n0::HuffmanTreeNode{T}, n1::HuffmanTreeNode{T}) where T<:SourceSymbol =
    HuffmanBranchNode{T}(n0, n1)

AbstractTrees.children(b::HuffmanBranchNode) = (b.n0, b.n1)
AbstractTrees.printnode(io::IO, b::HuffmanBranchNode) = print(io, "・")

struct HuffmanTree{T<:SourceSymbol}
    root::HuffmanTreeNode{T}
end

const DCHuffmanTree = HuffmanTree{DCSymbol}
const ACHuffmanTree = HuffmanTree{ACSymbol}

# The following method creates a huffman tree based on the probabilities.
# Note that the code assignment is not unique and the result is not normalized.
function HuffmanTree(nodes::AbstractVector{S}) where {T, S<:HuffmanTreeNode{T}}
    length(nodes) < 1 && error()
    t = collect(HuffmanTreeNode{T}, nodes)
    while length(t) > 1
        sort!(t, alg=InsertionSort, by=n->n.probability, rev=true)
        # pop the two nodes with the lowest probability
        n1 = pop!(t)
        n0 = pop!(t)
        # push a new branch node with the two nodes as children
        push!(t, HuffmanBranchNode(n0, n1))
    end
    HuffmanTree(t[1])
end

AbstractTrees.children(t::HuffmanTree) = children(t.root)
AbstractTrees.printnode(io::IO, t::HuffmanTree) = print(io, typeof(t))
AbstractTrees.nodetype(::HuffmanTree{T}) where T = HuffmanTreeNode{T}

function Base.show(io::IO, t::HuffmanTree{T}) where T
    function p(io::IO, s::HuffmanSymbolNode{T})
        print(io, "(", s.symbol, ", ", s.probability, ")")
    end
    function p(io::IO, b::HuffmanBranchNode{T})
        print(io, "(")
        p(io, b.n0)
        print(io, ", ")
        p(io, b.n1)
        print(io, ")")
    end
    print(io, typeof(t), "(")
    p(IOContext(io, :compact=>true, :typeinfo=>T), t.root)
    print(io, ")")
end

function Base.show(io::IO, ::MIME"text/plain", t::HuffmanTree{T}) where T
    table = HuffmanTable(t)
    n = 0
    maxdepth = 0
    for li in table.len
        n += li
        n > 20 && break
        maxdepth += 1
    end
    print_tree(IOContext(io, :compact=>true, :typeinfo=>T), t, maxdepth)
end

height(t::HuffmanSymbolNode) = 0
height(t::HuffmanBranchNode) = 1 + max(height(t.n0), height(t.n1))
height(t::HuffmanTree) = height(t.root)

struct HuffmanTable{T<:SourceSymbol}
    id::UInt8 # Th: Huffman table destination identifier
    len::SVector{16, UInt8} # Li
    symbols::AbstractVector{T} # Vij

    function HuffmanTable{T}(len::AbstractVector{<:Integer},
                             symbols::AbstractVector{T};
                             id::Integer=0) where T <: SourceSymbol
        id < 4 || error("`id` must be less than 4.")
        new{T}(id, len, symbols)
    end
end

const DCHuffmanTable = HuffmanTable{DCSymbol}
const ACHuffmanTable = HuffmanTable{ACSymbol}

Base.sizeof(ht::HuffmanTable) = length(ht) + 17
Base.length(ht::HuffmanTable) = length(ht.symbols)

# HuffmanTree <--> HuffmanTable conversions
function HuffmanTree(table::HuffmanTable{T}) where T <: SourceSymbol
    length(table) < 1 && error()
    symbols = collect(zip(table.symbols, codelist(table)))
    # set probability based on the code length
    node(sc) = HuffmanSymbolNode(sc[1], 0.5^length(sc[2]))
    function split(s, depth)
        length(s) == 1 && return node(s[1])
        s0 = filter(sc -> !sc[2][depth], s) # the `depth`-th bit is zero
        s1 = filter(sc ->  sc[2][depth], s) # the `depth`-th bit is one
        n0 = length(s0) == 1 ? node(s0[1]) : split(s0, depth + 1)
        n1 = length(s1) == 1 ? node(s1[1]) : split(s1, depth + 1)
        HuffmanBranchNode(n0, n1)
    end
    HuffmanTree(split(symbols, 1))
end

function HuffmanTable(tree::HuffmanTree{T}, id::Integer=0) where T <: SourceSymbol
    symbols = Tuple{T, Int}[]
    en(t::HuffmanSymbolNode, d::Int) = push!(symbols, (t.symbol, d))
    en(t::HuffmanBranchNode, d::Int) = (en(t.n0, d + 1); en(t.n1, d + 1))
    en(tree.root, 0)
    sort!(symbols, by=(s->s[2]))
    len = UInt8[count(s->s[2]==i, symbols) for i = 1:16]
    HuffmanTable{T}(len, (s->s[1]).(symbols), id=id)
end

function Base.read(io::IO, t::Type{<:HuffmanTable})
    tc, th = u4u4(io)
    len = Vector{UInt8}(undef, 16)
    readbytes!(io, len, 16)
    slen = sum(len)
    symbols = Vector{UInt8}(undef, slen)
    readbytes!(io, symbols, slen)
    t === DCHuffmanTable && tc != 0x0 && error("AC/DC type mismatch")
    t === ACHuffmanTable && tc != 0x1 && error("AC/DC type mismatch")
    if tc == 0x0
        return DCHuffmanTable(len, DCSymbol.(symbols), id=th)
    elseif tc == 0x1
        return ACHuffmanTable(len, ACSymbol.(symbols), id=th)
    else
        error()
    end
end

function codetable(ht::HuffmanTable)
    code = code""
    function u()
        prev = code
        code = inc(code)
        prev
    end
    function v(n)
        code = shl(code)
        CodeBits[u() for j = 1:n]
    end
    SVector{16, Vector{CodeBits}}([v(Int(ht.len[i])) for i = 1:16])
end

codelist(ht::HuffmanTable) = reduce(vcat, codetable(ht), init=[])

function Base.show(io::IO, ht::HuffmanTable)
    print(io, typeof(ht), "([", join(string.(Int.(ht.len)), ", "), "], ")
    print(io, ht.symbols, ", id=", ht.id, ")")
end

function Base.show(io::IO, ::MIME"text/plain", ht::HuffmanTable)
    print(io, length(ht), "-symbol ", typeof(ht), "(id=", ht.id, "):")
    codes = codelist(ht)
    for (s, c) in zip(ht.symbols, codes)
        println(io)
        print(io, " ", s, "  ", c)
    end
end

struct HuffmanDecoder{T<:SourceSymbol}
    symbols::AbstractVector{T}
    codes::AbstractVector{CodeBits}
    index8::SVector{256, Int32} # for codes with length not greater than 8 bits
    indexx::Integer # first index which is out of index8

    function HuffmanDecoder{T}(table::HuffmanTable{T}) where {T<:SourceSymbol}
        symbols = table.symbols
        codes = codelist(table)
        index8 = zeros(Int32, 256)

        length(codes) > 256 && error()
        # e.g.
        # table.len = [0, 1, 5, ..., 0]
        # bitlen=1 -> none
        # bitlen=2 -> code"00"
        # bitlen=3 -> code"010", code"011", code"100", code"101", code"110"
        #
        # 0b00_000000 : 0b00_111111 -> index=1 (w=64, 6 bits)
        #   __ ^^^^^^     __ ^^^^^^
        # 0b010_00000 : 0b010_11111 -> index=2 (w=32, 5 bits)
        # 0b011_00000 : 0b011_11111 -> index=3 (w=32, 5 bits)
        # 0b100_00000 : 0b100_11111 -> index=4 (w=32, 5 bits)
        # 0b101_00000 : 0b101_11111 -> index=5 (w=32, 5 bits)
        # 0b110_00000 : 0b110_11111 -> index=6 (w=32, 5 bits)
        #   ___ ^^^^^     ___ ^^^^^
        index = 1
        p = 0
        for bitlen = 1:8
            w = 0x100 >> bitlen
            for k = 1:table.len[bitlen]
                index8[p+1:p+w] .= index
                index += 1
                p += w
            end
        end
        new{T}(symbols, codes, index8, index)
    end
end

struct BitSequenceError <: Exception end

HuffmanDecoder(table::HuffmanTable{T}) where {T<:SourceSymbol} = HuffmanDecoder{T}(table)

function getcodeindex(d::HuffmanDecoder, bits::UInt32, effective_length::Integer)
    index = d.index8[(bits >> 24) + 1]
    if index > 0 # matched
        c = d.codes[index]
        return effective_length < length(c) ? 0 : index
    end
    for i = d.indexx:length(d.codes)
        c = d.codes[i]
        l = length(c)
        if effective_length < l
            delta = UInt8(l - effective_length)
            (word(c) >> delta) >= (bits >> UInt8(32 - effective_length)) && return 0
            throw(BitSequenceError())
        end
        c == CodeBits(l, bits >> UInt8(32 - l)) && return i
    end
    throw(BitSequenceError())
end


mutable struct TreeModelXMLElement <: MyXMLElement
    el::XMLOrNothing
    node_traits::Vector{String}
    trait_dims::Vector{Int}
    param_names::Vector{String}
    newick_el::NewickXMLElement
    fixed_tree::Bool
    id::String

    function TreeModelXMLElement(newick_el::NewickXMLElement,
                dim::Int;
                trait::String = bn.DEFAULT_TRAIT_NAME)
        param_name = bn.LEAF_TRAITS
        return new(nothing, [trait], [dim], [param_name], newick_el, true,
                    bn.TREE_MODEL)
    end

    function TreeModelXMLElement(newick_el::NewickXMLElement,
                node_traits::Vector{String},
                trait_dims::Vector{Int},
                param_names::Vector{String}
                )
        return new(nothing, node_traits, trait_dims, param_names, newick_el, true)
    end

    function TreeModelXMLElement(newick_el::NewickXMLElement,
                                 node_traits::Vector{String},
                                 trait_dims::Vector{Int}
                                )
        param_names = ["$trait.leafTrait" for trait in node_traits]
        return TreeModelXMLElement(newick_el, node_traits, trait_dims, param_names)
    end
end

function name(::TreeModelXMLElement)
    return bn.TREE_MODEL
end

function get_id(tm::TreeModelXMLElement)
    return tm.id
end

function make_xml(tl::TreeModelXMLElement)
    make_xml(tl.newick_el)
    tl.el =  make_treeModel(tl.newick_el.el,
                    tl.node_traits,
                    tl.trait_dims,
                    tl.param_names,
                    id=tl.id)
end

function make_treeModel(newick_el::XMLElement,
        trait_names::Vector{String},
        trait_dims::Vector{Int},
        param_names::Vector{String};
        fixed_tree::Bool = false,
        id::String = bn.TREE_MODEL)

    @assert name(newick_el) == bn.NEWICK

    el = new_element(bn.TREE_MODEL)
    set_attribute(el, bn.ID, id)
    set_attribute(el, bn.FIX_HEIGHTS, bn.TRUE)

    if fixed_tree
        set_attribute(el, bn.FIXED_TREE, bn.TRUE)
    end

    tree_el = new_element(bn.TREE)
    newick_id = attribute(newick_el, bn.ID)
    set_attribute(tree_el, bn.IDREF, newick_id)
    add_child(el, tree_el)

    rh_el = new_element(bn.ROOT_HEIGHT)
    add_parameter_id(rh_el, "$(bn.TREE_MODEL).$(bn.ROOT_HEIGHT)")
    add_child(el, rh_el)

    add_nodeHeights(el, "$(bn.TREE_MODEL).$(bn.INTERNAL_NODE_HEIGHTS)",
            internal = true)
    add_nodeHeights(el, "$(bn.TREE_MODEL).$(bn.ALL_INTENAL_NODE_HEIGHTS)",
            internal = true, root = true)

    for i = 1:length(trait_names)
        nt_el = new_element(bn.NODE_TRAITS)
        attrs = [(bn.ROOT_NODE, bn.FALSE),
                (bn.INTERNAL_NODES, bn.FALSE),
                (bn.LEAF_NODES, bn.TRUE),
                (bn.AS_MATRIX, bn.TRUE),
                (bn.TRAIT_DIM, string(trait_dims[i])),
                (bn.NAME, trait_names[i])]
        set_attributes(nt_el, attrs)
        add_parameter_id(nt_el, param_names[i])
        add_child(el, nt_el)
    end
    return el
end

function add_nodeHeights(pel::XMLElement, param_id::String;
        internal::Bool = false, root::Bool = false)

    el = new_element(bn.NODE_HEIGHTS)
    if internal
        set_attribute(el, bn.INTERNAL_NODES, bn.TRUE)
    end
    if root
        set_attribute(el, bn.ROOT_NODE, bn.TRUE)
    end
    add_parameter_id(el, param_id)
    add_child(pel, el)
    return el
end

function add_leaf_param!(txl::TreeModelXMLElement, trait_name::String, trait_dim::Int)
    push!(txl.node_traits, trait_name)
    push!(txl.trait_dims, trait_dim)
    push!(txl.param_names, "$(trait_name).leaf")
end


@testset "CoreModelCoupled generic interface" begin
    model = load_model(CoreModelCoupled, model_paths["e_coli_core.mat"])

    @test reaction_stoichiometry(model, "EX_ac_e") == Dict("ac_e" => -1)
    @test reaction_stoichiometry(model, 44) == Dict("ac_e" => -1)

end


"""
    $(SIGNATURES)

Calculate ΔG' of a reaction described by `rxn_string` using the settings in
`equilibrator`. Optionally, set the concentrations (abundances) of the species involved
in the reaction through `concentration`. If `skip_unbalanced` is true then
return `nothing` if the reaction is unbalanced. If `balance_warn` is false
then do not emit a warning when the reaction is unbalanced. Note, concentrations
that are smaller than 0.001 mM are capped at 0.001 mM (an eQuilibrator/equilibrator_api)
restriction.

# Example
```
equilibrator = Equilibrator()

rxn_string = "bigg.metabolite:atp + bigg.metabolite:h2o = bigg.metabolite:adp + bigg.metabolite:pi"
concens = Dict("bigg.metabolite:atp"=>1u"mM", "bigg.metabolite:adp"=>100u"μM", "bigg.metabolite:pi"=>0.003u"M")
dg_prime(equilibrator, rxn_string; concentrations=concens)
```
"""
function dg_prime(
    equilibrator::Equilibrator,
    rxn_string::String;
    concentrations = Dict{String,Real}(),
    balance_warn = true,
    skip_unbalanced = false,
)
    rxn = equilibrator.cc.parse_reaction_formula(rxn_string)

    if balance_warn && !rxn.is_balanced()
        @warn "Reaction unbalanced."
        skip_unbalanced && return nothing
    end

    for (k, v) in concentrations
        val = uconvert(u"M", _lower_bound(float(v)))
        rxn.set_abundance(equilibrator.cc.get_compound(k), equilibrator.eq.Q_(string(val)))
    end
    pydg = equilibrator.cc.dg_prime(rxn)
    dg = pydg.value.magnitude ± pydg.error.magnitude
    units = uparse(_parse_units(py"str"(pydg.units)))
    return (dg)units
end

"""
    $(SIGNATURES)

Calculate ΔG' at physiological conditions (all concentrations set to 1 mM, 1mbar) of a reaction
described by `rxn_string` using the settings in `equilibrator`. If `skip_unbalanced` is true
then return `nothing` if the reaction is unbalanced. If `balance_warn` is false then do not
emit a warning when the reaction is unbalanced.

# Example
```
equilibrator = Equilibrator()

rxn_string = "bigg.metabolite:atp + bigg.metabolite:h2o = bigg.metabolite:adp + bigg.metabolite:pi"
physiological_dg_prime(equilibrator, rxn_string)
```
"""
function physiological_dg_prime(
    equilibrator::Equilibrator,
    rxn_string::String;
    balance_warn = true,
    skip_unbalanced = false,
)
    rxn = equilibrator.cc.parse_reaction_formula(rxn_string)

    if balance_warn && !rxn.is_balanced()
        @warn "Reaction unbalanced."
        skip_unbalanced && return nothing
    end

    pydg = equilibrator.cc.physiological_dg_prime(rxn)
    dg = pydg.value.magnitude ± pydg.error.magnitude
    units = uparse(_parse_units(py"str"(pydg.units)))
    return (dg)units
end

"""
    $(SIGNATURES)

Calculate ΔG⁰ at standard conditons of a reaction described by `rxn_string` using
the settings in `equilibrator`. If `skip_unbalanced` is true then return `nothing` if the
reaction is unbalanced. If `balance_warn` is false then do not emit a warning
when the reaction is unbalanced.

Note, standard conditions means that all concentrations are set at 1 M.

# Example
```
equilibrator = Equilibrator()

rxn_string = "bigg.metabolite:atp + bigg.metabolite:h2o = bigg.metabolite:adp + bigg.metabolite:pi"
standard_dg_prime(equilibrator, rxn_string)
```
"""
function standard_dg_prime(
    equilibrator::Equilibrator,
    rxn_string::String;
    balance_warn = true,
    skip_unbalanced = false,
)
    rxn = equilibrator.cc.parse_reaction_formula(rxn_string)

    if balance_warn && !rxn.is_balanced()
        @warn "Reaction unbalanced."
        skip_unbalanced && return nothing
    end

    pydg = equilibrator.cc.standard_dg_prime(rxn)
    dg = pydg.value.magnitude ± pydg.error.magnitude
    units = uparse(_parse_units(py"str"(pydg.units)))
    return (dg)units
end

"""
    $(SIGNATURES)

Calculate the log of the reversibility index of a reaction described by `rxn_string`
using settings in `equilibrator`. If `skip_unbalanced` is true then return `nothing` if the reaction is
unbalanced. If `balance_warn` is false then do not emit a warning when the
reaction is unbalanced.

# Example
```
equilibrator = Equilibrator()

rxn_string = "bigg.metabolite:atp + bigg.metabolite:h2o = bigg.metabolite:adp + bigg.metabolite:pi"
ln_reversibility_index(equilibrator, rxn_string)
```
"""
function ln_reversibility_index(
    equilibrator::Equilibrator,
    rxn_string::String;
    balance_warn = true,
    skip_unbalanced = false,
)
    rxn = equilibrator.cc.parse_reaction_formula(rxn_string)

    if balance_warn && !rxn.is_balanced()
        @warn "Reaction unbalanced."
        skip_unbalanced && return nothing
    end

    rev_obj = equilibrator.cc.ln_reversibility_index(rxn)
    return rev_obj.value.magnitude ± rev_obj.error.magnitude
end


using Vlasiator, SHA, LazyArtifacts
using Test

group = get(ENV, "TEST_GROUP", :all) |> Symbol

function nanmaximum(x::AbstractArray{T}) where T<:AbstractFloat
   result = convert(eltype(x), NaN)
   for i in x
      if !isnan(i)
         if isnan(result) || i > result
            result = i
         end
      end
   end
   result
end

@testset "Vlasiator.jl" begin
   rootpath = artifact"testdata"

   files = joinpath.(rootpath, ("bulk.1d.vlsv", "bulk.2d.vlsv", "bulk.amr.vlsv"))
   meta1 = load(files[1])
   meta2 = load(files[2])
   meta3 = load(files[3])

   if group in (:read, :all)
      @testset "Reading files" begin
         @test_throws ArgumentError load("data")
         meta = meta1
         @test ndims(meta) == 1
         @test startswith(repr(meta), "File: bulk.1d.vlsv")
         @test size(meta) == 5168730
         # Variable strings reading
         @test meta.variable[end] == "vg_rhom"
         # Variable info reading
         varinfo = readvariablemeta(meta, "proton/vg_rho")
         @test startswith(repr(varinfo), "Variable in LaTeX")
         @test varinfo.unit == "1/m^3"
         # Parameter checking
         @test hasparameter(meta, "dt") == true
         # Parameter reading
         t = readparameter(meta, "time")
         @test t == 8.0
         @test_throws ArgumentError meta["nonsense"]
         # unsorted ID
         @test readvariable(meta, "CellID", false) == UInt64[10, 9, 8, 7, 2, 1, 3, 4, 5, 6]
         indexRef = [6, 5, 7, 8, 9, 10, 4, 3, 2, 1]
         @test meta.cellindex == indexRef
         # sorted var by default
         @test meta["vg_boundarytype"] == [4, 4, 1, 1, 1, 1, 1, 1, 3, 3]
         # ID finding
         loc = [2.0, 0.0, 0.0]
         id = getcell(meta, loc)
         coords = getcellcoordinates(meta, id)
         @test coords == [2.5, 0.0, 0.0]
         @test readvariable(meta, "proton/vg_rho", id)[1,1] ≈ 1.77599 atol=1e-5
         # ID in a line
         point1 = [-2.0, 0.0, 0.0]
         point2 = [2.0, 0.0, 0.0]
         cellids, _, _ = getcellinline(meta, point1, point2)
         @test cellids == collect(4:7)
         point1 = [-5.1, 0.0, 0.0]
         @test_throws DomainError getcellinline(meta, point1, point2)
         point2 = [5.1, 0.0, 0.0]
         @test_throws DomainError getcellinline(meta, point1, point2)

         # Nearest ID with VDF stored
         @test getnearestcellwithvdf(meta, id) == 8

         # velocity space reading
         vcellids, vcellf = readvcells(meta, 2; species="proton")
         V = getvcellcoordinates(meta, vcellids; species="proton")
         @test V[end] == Float32[2.45, 1.95, 1.95]
         @test_throws ArgumentError readvcells(meta, 20)
         f = Vlasiator.flatten(meta.meshes["proton"], vcellf)
         @test f[CartesianIndex(24, 18, 7)] == 1.0f-45

         # AMR data reading, DCCRG grid
         metaAMR = meta3
         sliceoffset = abs(metaAMR.coordmin[2])
         idlist, indexlist = getslicecell(metaAMR, sliceoffset, 2,
            metaAMR.coordmin[2], metaAMR.coordmax[2])

         data = readvariable(metaAMR, "proton/vg_rho")
         dataslice = refineslice(metaAMR, idlist, data[indexlist], :y)
         @test sum(dataslice) ≈ 7.690352275026747e8
         let err = nothing
            try
               getslicecell(metaAMR, sliceoffset, 1, -2., -1.)
            catch err
            end
            @test err isa Exception
         end
         # AMR ID finding
         loc = [0.0, 0.0, 0.0]
         @test getcell(metaAMR, loc) == 0x0000000000001591

         # AMR level
         @test metaAMR.maxamr == 2
         @test getlevel(metaAMR, idlist[1]) == 1

         # DCCRG utilities
         @test isparent(metaAMR, 1)
         @test !isparent(metaAMR, 1080)
         @test getchildren(metaAMR, 1) == getsiblings(metaAMR, 129)
         @test getparent(metaAMR, 129) == 1
         @test_throws ArgumentError getparent(metaAMR, 5)
         @test_throws ArgumentError getsiblings(metaAMR, 1)

         # FS grid
         data = readvariable(metaAMR, "fg_e")
         ncells, namr = metaAMR.ncells, metaAMR.maxamr
         @test size(data) == (3, ncells[1]*namr^2, ncells[2]*namr^2, ncells[3]*namr^2)
         @test data[:,16,8,8] == [-1.3758785f-7, 3.2213068f-4, -3.1518404f-4]

         # Compare two VLSV files
         @test issame(files[1], files[1])
      end
   end

   if group in (:derive, :all)
      @testset "Derived variables" begin
         meta = meta1
         @test meta["Vmag"] |> sortperm == [7, 6, 5, 4, 3, 1, 2, 8, 9, 10]

         meta = meta2
         @test meta["Bmag"][4] == 3.0052159f-9

         @test meta["Emag"][1,10,99] == 2.6120072f-6

         @test meta["VS"] |> nanmaximum == 1.3726345956957596e6

         @test meta["VA"] |> nanmaximum == 2.3202628822256166e8

         @test meta["MA"][end] == 10.700530839822328

         @test meta["MS"][end] == 16.9375888861409

         @test meta["Vpar"][1] == 698735.3f0

         @test meta["Vperp"][1] == 40982.48f0

         @test meta["T"][1] == 347619.9817319378

         @test meta["Pdynamic"][1] == 8.204415428337215e-10

         @test meta["Pb"][1] == 3.59502523919876e-12

         @test meta["Beta"][1] == 1.3359065984817116

         @test meta["Poynting"][:,10,10] == [-3.677613f-11, 8.859047f-9, 2.4681486f-9]

         @test meta["IonInertial"][1] == 8.578087716535188e7

         @test readvariable(meta, "Larmor", UInt64[1])[1] == 322324.70603759587

         @test meta["Gyroperiod"][1] == 21.834297799454554

         @test meta["Gyrofrequency"][1] == 0.04579950356933307

         @test meta["J"][1,1000] == 2.314360722590665e-14

         @test meta["Omegap"][1] == 209.5467447842415

         @test meta["Plasmaperiod"][1] == 0.0047722048893178645

         #@test meta["Panisotropy"]

         #@test meta["Tanisotropy"]

         #@test meta["Agyrotropy"]

      end
   end

   if group in (:utility, :all)
      @testset "Rotation" begin
         using LinearAlgebra
         T = Diagonal([1.0, 2.0, 3.0])
         B = [0.0, 1.0, 0.0]
         Vlasiator.rotateWithB!(T, B)
         @test T == Diagonal([1.0, 3.0, 2.0])
         @test Vlasiator.rotateWithB(T, B) == Diagonal([1.0, 2.0, 3.0])

         v = fill(1/√3, 3)
         θ = π / 4
         R = Vlasiator.getRotationMatrix(v, θ)
         @test inv(R) ≈ R'

         Rᵀ = Vlasiator.rotateTensorToVectorZ(R, v)
         @test Rᵀ ≈ [1/√2 -1/√2 0.0; 1/√2 1/√2 0.0; 0.0 0.0 1.0]
      end
      @testset "Curvature" begin
         let dx = ones(Float32, 3)
            A = ones(Float32, 3,3,3,3)
            @test sum(Vlasiator.curl(dx, A)) == 0.0
            A = ones(Float32, 3,3,1,3)
            @test sum(Vlasiator.curl(dx, A)) == 0.0
         end
      end
   end

   if group in (:vtk, :all)
      @testset "VTK" begin
         meta = meta2 # no amr
         data, ghostType = Vlasiator.fillmesh(meta, "proton/vg_rho")
         @test size(data[1][1]) == (1, 63, 100, 1)

         meta = meta3 # amr
         write_vtk(meta)
         sha_str = bytes2hex(open(sha1, "bulk.amr_1.vti"))
         @test sha_str == "bafb52747908cd05c78218078054aa4268f453d9"

         filesaved = ["bulk.amr.vthb", "bulk.amr_1.vti", "bulk.amr_2.vti", "bulk.amr_3.vti"]
         rm.(filesaved, force=true)
      end
   end

   if group in (:log, :all)
      @testset "Log" begin
         file = joinpath(rootpath, "logfile.txt")
         timestamps, speed = readlog(file)
         @test length(speed) == 50 && speed[end] == 631.2511f0
      end
   end

   if group in (:plot, :all)
      @testset "PyPlot" begin
         using PyPlot
         ENV["MPLBACKEND"]="agg" # no GUI
         # 1D
         meta = meta1
         plot(meta, "proton/vg_rho")
         line = gca().lines[1]
         @test line.get_ydata() == meta["proton/vg_rho"]
         centers = plotmesh(meta, projection="y")
         points = centers.get_offsets()
         @test size(points) == (10, 2)
         fig = plt.figure()
         ax = fig.add_subplot(projection="3d")
         centers = plotmesh(meta, projection="3d")
         points = centers.get_offsets() # only 2D from 3D coords, might be improved
         @test size(points) == (10, 2)
         close(fig)

         @test_throws ArgumentError pcolormesh(meta, "proton/vg_rho")

         loc = [2.0, 0.0, 0.0]
         var = vdfslice(meta, loc).get_array()
         @test var[786] == 229.89486651619364
         @test_throws ArgumentError vdfslice(meta, loc, species="helium")

         # 2D
         meta = meta2
         var = pcolormesh(meta, "proton/vg_rho").get_array()
         @test var[end-2] == 999535.8f0 && length(var) == 6300
         var = pcolormesh(meta, "fg_b").get_array()
         @test var[1] == 3.0058909f-9
         var = pcolormesh(meta, "proton/vg_v", op=:x, colorscale=SymLog).get_array()
         @test var[2] == -699935.2f0
         p = streamplot(meta, "proton/vg_v", comp="xy")
         @test typeof(p) == PyPlot.PyObject
         var = quiver(meta, "proton/vg_v", axisunit=SI, stride=1).get_offsets()
         @test size(var) == (6300, 2)

         # 3D AMR
         meta = meta3
         var = pcolormesh(meta, "proton/vg_rho").get_array()
         @test var[255] == 1.0483886f6 && length(var) == 512
         var = pcolormesh(meta, "proton/vg_v").get_array()
         @test var[255] == 99992.586f0
         @test_throws ArgumentError pcolormesh(meta, "fg_b")
      end

      @testset "Plots" begin
         isdefined(Main, :PyPlot) && import Vlasiator.vdfslice
         include("../src/plot/plots.jl")
         RecipesBase.is_key_supported(k::Symbol) = true
         # 1D
         meta = meta1
         rec = RecipesBase.apply_recipe(Dict{Symbol, Any}(), meta, "proton/vg_rho")
         @test getfield(rec[1], 1)[:seriestype] == :line &&
            rec[1].args[1] isa LinRange

         rec = RecipesBase.apply_recipe(Dict{Symbol, Any}(), VDFSlice((meta,[0.0,0.0,0.0])))
         @test getfield(rec[1], 1)[:seriestype] == :histogram2d

         # 2D
         meta = meta2
         rec = RecipesBase.apply_recipe(Dict{Symbol, Any}(), meta, "proton/vg_rho")
         @test getfield(rec[1], 1)[:seriestype] == :heatmap &&
            rec[1].args[1] isa LinRange
      end
   end
   for meta in (meta1, meta2, meta3)
      close(meta.fid) # required for Windows?
   end
end


# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Sample an Action
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_sample_action(Agent, t::Int; Greedy = false)
    P_Action = Agent.P_Action[:,t]
    if Greedy
        (p, act) = findmax(P_Action)
        act = rand((1:4)[P_Action.==p])
        act = act - 1
        return act
    else
        P_cum = cumsum(P_Action)
        p = rand()
        act = (P_cum[1]<p) + (P_cum[2]<p) + (P_cum[3]<p)
        return act
    end

end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Make Observation
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_make_obs(Agent, t::Int, EnvGraph::Str_EnvGraph)
    obs = Agent.Input.obs[t]
    act = Agent.Input.act[t]
    next_obs = EnvGraph.TranMat[obs+1,act+1]
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Simulating an agent
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_simulate_Agent(EnvGraph;
                             Epi=1, Block=1,
                             Param=Str_Param(),
                             HyperParam=Str_HyperParam(),
                             Agent_Prev=Str_Agent(),
                             Greedy = false,
                             T_stop = 1000)

    Input = Str_Input(; obs = [EnvGraph.InitObs[Epi]], act = [-1])
    Agent = Func_initial_Agent(; Sub=1, Epi=Epi, Block=Block,
                                Param=Param, HyperParam=HyperParam,
                                Input=Input, Agent_Prev=Agent_Prev)

    t = 1
    act = Func_sample_action(Agent,t; Greedy=Greedy)
    Agent.Input.act[t] = act

    State = Func_update_State(Agent, t)
    Agent.State[t] = State

    obs = Func_make_obs(Agent,t,EnvGraph)
    Agent.Input.obs = cat(Agent.Input.obs, obs; dims=1)

    Goal = false
    Stop = false
    while (!Goal)&&(!Stop)
        t += 1

        Stop = (t>=T_stop)
        Goal = (Agent.Input.obs[t] == EnvGraph.Goal)

        S = Func_compute_surprise(Agent, t)
        Agent.Surprise = cat(Agent.Surprise, S; dims=1)

        gamma = Func_compute_gamma(Agent,t)
        Agent.Gamma = cat(Agent.Gamma, gamma; dims=1)

        t_since_change = Func_compute_t_since_change(Agent,t)
        Agent.t_since_change = cat(Agent.t_since_change, t_since_change; dims=1)

        counter = Func_update_counter(Agent, t)
        Agent.Counter = cat(Agent.Counter, counter; dims=2)

        Π = Func_update_belief(Agent, t)
        Agent.Π = cat(Agent.Π, Π; dims=3)

        Novelty = Func_update_novelty(Agent, t)
        Agent.Novelty = cat(Agent.Novelty, Novelty; dims=2)

        # Model free updating
        δ_R = Func_compute_δ(Agent, t, "Reward")
        Agent.δ_R_Seq = cat(Agent.δ_R_Seq, δ_R; dims=1)
        δ_N = Func_compute_δ(Agent, t, HyperParam.exploration_form)
        Agent.δ_N_Seq = cat(Agent.δ_N_Seq, δ_N; dims=1)

        E_R = Func_compute_ElTrace(Agent, t, "Reward")
        Agent.E_R = cat(Agent.E_R, E_R; dims=2)
        E_N = Func_compute_ElTrace(Agent, t, HyperParam.exploration_form)
        Agent.E_N = cat(Agent.E_N, E_N; dims=2)

        Q_R_MF = Func_update_Q_MF(Agent, t, "Reward")
        Agent.Q_R_MF = cat(Agent.Q_R_MF, Q_R_MF; dims=2)
        Q_N_MF = Func_update_Q_MF(Agent, t, HyperParam.exploration_form)
        Agent.Q_N_MF = cat(Agent.Q_N_MF, Q_N_MF; dims=2)

        # Model based updating
        if (!Agent.First_Epi)||(Goal)||(Block==2)||(HyperParam.First_Epi_reward_update)
            Q_R_MB, U_R = Func_update_prior_sweep(Agent, t, "Reward")
        else
            Q_R_MB = Agent.Q_R_MB[:,1]
            U_R = Agent.U_R[:,1]
        end
        Agent.Q_R_MB = cat(Agent.Q_R_MB, Q_R_MB; dims=2)
        Agent.U_R = cat(Agent.U_R, U_R; dims=2)

        Q_N_MB, U_N = Func_update_prior_sweep(Agent, t, HyperParam.exploration_form)
        Agent.Q_N_MB = cat(Agent.Q_N_MB, Q_N_MB; dims=2)
        Agent.U_N = cat(Agent.U_N, U_N; dims=2)

        if !Goal
            if Agent.First_Epi
                P_Action = Func_compute_P_action(Agent, t; Only_Reward=false)
            else
                P_Action = Func_compute_P_action(Agent, t; Only_Reward=true)
            end
            Agent.P_Action = cat(Agent.P_Action, P_Action; dims=2)

            act = Func_sample_action(Agent,t; Greedy=Greedy)
            Agent.Input.act = cat(Agent.Input.act, act; dims=1)

            obs = Func_make_obs(Agent,t,EnvGraph)
            Agent.Input.obs = cat(Agent.Input.obs, obs; dims=1)
        else
            act = -1
            Agent.Input.act = cat(Agent.Input.act, act; dims=1)
        end

        State = Func_update_State(Agent, t)
        Agent.State = cat(Agent.State, State; dims=1)

        Agent.Novelty_Seq = cat(Agent.Novelty_Seq,
                                  Agent.Novelty[Input.obs[t]+1,t-1];
                                  dims=1)
    end

    Input = Agent.Input
    T = length(Input.obs)

    # --------------------------------------------------------------------------
    # Action Analysis - LogL
    for t = 1:(T-1)
        Agent.LogL += log(Agent.P_Action[Input.act[t]+1 ,t])

        (p, a) = findmax(Agent.P_Action[: ,t])
        if t==1
            Agent.Est_Act[1] = a-1
            Agent.Est_Act_P[1] = p
        else
            Agent.Est_Act = cat(Agent.Est_Act, a-1; dims=1)
            Agent.Est_Act_P = cat(Agent.Est_Act_P, p; dims=1)
        end
    end

    Agent.LogL_Unif = (T-1) * log( 1/Param.num_act )
    # --------------------------------------------------------------------------
    # Action Analysis - Acc Rate
    Agent.Sort_Dist_Data = zeros(size(Agent.P_Action))
    Agent.Sort_Dist_Model = zeros(size(Agent.P_Action))

    for t = 1:(T-1)
        temp = sortperm(Agent.P_Action[:,t])
        Agent.Sort_Dist_Model[:,t] = Agent.P_Action[temp,t]

        Agent.Sort_Dist_Data[Input.act[t] + 1 ,t] = 1
        Agent.Sort_Dist_Data[:,t] = Agent.Sort_Dist_Data[temp,t]

        if (Agent.Sort_Dist_Model[findmax(Agent.Sort_Dist_Data[:,t])[2],t] -
            Agent.Sort_Dist_Model[end,t]) == 0

            Agent.Sort_Dist_Data[:,t] = zeros(size(Agent.Sort_Dist_Data[:,t]))

            Agent.Sort_Dist_Data[Agent.Sort_Dist_Model[:,t] .==
                              Agent.Sort_Dist_Model[end,t],t] .= 1

            Agent.Sort_Dist_Data[:,t] = Agent.Sort_Dist_Data[:,t] ./
                                     sum(Agent.Sort_Dist_Data[:,t])
        end
    end

    Agent.LogL_Unif = (T-1) * log( 1/Param.num_act )
    Agent.Acc_Rate = mean(Agent.Sort_Dist_Data[4,:])

    return Agent, Stop
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Simulating blocks
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_simulate_Agent_Blocks(EnvGraph_Array;
                                   Epi_Set=1:5, Block_Set=1:2,
                                   Param=[Str_Param(),Str_Param()],
                                   HyperParam=Str_HyperParam(),
                                   Greedy = false,
                                   Seed = 0,
                                   T_stop = 1000)
    Random.seed!(Seed)

    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents = Array{Str_Agent,2}(undef,N_blo,N_epi)
    for Block = Block_Set
        EnvGraph = EnvGraph_Array[Block]
        for Epi = Epi_Set
            if (Block==1)&&(Epi==1)
                Agent_Epi, Stop = Func_simulate_Agent(EnvGraph;
                                                      Epi=Epi, Block=Block,
                                                      Param=Param[Block],
                                                      HyperParam=HyperParam,
                                                      Greedy = Greedy,
                                                      T_stop = T_stop)
            else
                if Epi==1
                    Agent_Prev = Agents[Block-1,end]
                else
                    Agent_Prev = Agents[Block,Epi-1]
                end

                Agent_Epi, Stop = Func_simulate_Agent(EnvGraph;
                                                      Epi=Epi, Block=Block,
                                                      Param=Param[Block],
                                                      HyperParam=HyperParam,
                                                      Greedy = Greedy,
                                                      Agent_Prev=Agent_Prev,
                                                      T_stop = T_stop)
            end

            if Stop
                return [Agent_Epi]
            end
            Agents[Block,Epi] = Agent_Epi
        end
    end

    return Agents
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# η to params SurNoR
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_param_generator(η0;
                              σ = 0.1*ones(18),
                              η_low = [0,1e-12,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                              η_up = [0.999,1,Inf,Inf,Inf,0.999,0.999,0.999,0.999,
                                      1,1,Inf,Inf,1,1,1,1,Inf,Inf],
                              Seed = 1,
                              noise_free=true)
    Random.seed!(Seed)
    n_par = length(η0)
    if noise_free
        η = η0
    else
        η = η0 .+ (randn(n_par).*σ)

        for i = 1:n_par
            if η[i] < η_low[i]
                η[i] = η_low[i]
            elseif η[i] > η_up[i]
                η[i] = η_up[i]
            end
        end
    end

    Param1 = Str_Param(; num_act=4, num_obs=11,
                        count_decay=1, Q_R_0 = 0)

    Param1.p_change = η[1]
    Param1.ϵ = η[2]
    Param1.β_R = η[3]
    Param1.β_ratio_N2R = η[4]
    Param1.γ_RL_R = η[6]
    Param1.γ_RL_N = η[7]
    Param1.λ_elig_R = η[8]
    Param1.λ_elig_N = η[9]
    Param1.α = η[10]
    Param1.Δα = η[11]
    Param1.Q_N_0 = η[12]
    Param1.N_prior_sweep = Int64(round(η[13]))
    Param1.w_MF_to_MB_2 = η[14]
    Param1.w_MF_to_MB_1 = η[15]
    Param1.w_MF_to_MB_Scaler = η[17]

    Param2 = deepcopy(Param1)
    Param2.β_ratio_N2R = η[5]
    Param2.w_MF_to_MB_1 = η[16]
    Param2.β_R = η[18]

    Param = [Param1, Param2]
end


using InfrastructureModels

IM = InfrastructureModels

file = ARGS[1]

mp_data = open(file) do io
	data_string = read(io, String)
	# We need to manually patch `data_string` by inserting '%column_names%' where needed.
	# This is because `PowerModels.export_matpower()` doesn't write it,
	# whereas `InfrastructureModels.parse_matlab_string()` expects it...
	data_lines = split(data_string, '\n')
	last = length(data_lines)
	for (k, line) in enumerate(data_lines)
		# This below is just a dirty condition that works for the files in this project
		# I don't guarantee that it will work for *any* .m file
		if k > 1 && startswith(line, "%") && startswith(data_lines[k - 1], "%%")
			data_lines[k] = "%column_names%" * line[2:end]
		end
		if startswith(line, "%% convert branch impedances")
			last = k - 1
			break
		end
	end
	data_lines = data_lines[1:last]
	data_string = join(data_lines, '\n')
	matlab_data, func_name, colnames = IM.parse_matlab_string(data_string, extended=true)
	# write CSVs
	for key in ("bus", "gen", "branch")
		mpckey = "mpc.$key"
		outfile = "$key.csv"
		open(outfile, "w") do f
			column_names = colnames[mpckey]
			if key in ("gen", "branch")
				pushfirst!(column_names, "$(key)_name")
			end
			println(f, join(column_names, ", "))
			for (i, row) in enumerate(matlab_data[mpckey])
				if key in ("gen", "branch")
					row = string.(row)
					pushfirst!(row, "$(key)_$i")
				end
				str_row = join(row, ", ")
				println(f, str_row)
			end
		end
	end
end


const project_root_dir = realpath(joinpath(@__DIR__, ".."))
unique!(push!(LOAD_PATH, realpath(joinpath(project_root_dir, "experiments/utils/MonteCarloStudy"))))
unique!(push!(LOAD_PATH, realpath(joinpath(project_root_dir, "test/utils"))))

import Distributed
Distributed.@everywhere begin
    import Pkg
    Pkg.activate($project_root_dir)
    union!(LOAD_PATH, $LOAD_PATH)

    import MonteCarloStudy
    import CollisionAvoidanceGame
    import TestDynamics
    using PartiallyObservedInverseGames.ForwardGame: IBRGameSolver, KKTGameSolver
    using PartiallyObservedInverseGames.InverseGames:
        InverseKKTConstraintSolver, InverseKKTResidualSolver
end

import PartiallyObservedInverseGames.TrajectoryVisualization
import VegaLite
import LazyGroupBy: grouped

# Utils
include("utils/misc.jl")
include("utils/simple_caching.jl")
load_cache_if_not_defined!("unicycle")

#==================================== Forward Game Formulation =====================================#

T = 25

control_system =
    TestDynamics.ProductSystem([TestDynamics.Unicycle(0.25), TestDynamics.Unicycle(0.25)])

player_angles = let
    n_players = length(control_system.subsystems)
    map(eachindex(control_system.subsystems)) do ii
        angle_fraction = n_players == 2 ? pi / 2 : 2pi / n_players
        (ii - 1) * angle_fraction
    end
end

x0 = mapreduce(vcat, player_angles) do player_angle
    [unitvector(player_angle + pi); 0.1; player_angle + deg2rad(10)]
end

position_indices = mapreduce(vcat, eachindex(control_system.subsystems)) do subsystem_idx
    TestDynamics.state_indices(control_system, subsystem_idx)[1:2]
end

partial_state_indices = mapreduce(vcat, eachindex(control_system.subsystems)) do subsystem_idx
    TestDynamics.state_indices(control_system, subsystem_idx)[[1, 2, 4]]
end

player_cost_models_gt = map(enumerate(player_angles)) do (ii, player_angle)
    cost_model_p1 = CollisionAvoidanceGame.generate_player_cost_model(;
        player_idx = ii,
        control_system,
        T,
        goal_position = unitvector(player_angle),
    )
end

#======================================== Monte Carlo Study ========================================#

## Dataset Generation
n_observation_sequences_per_noise_level = 40

@run_cached forward_solution_gt, dataset = MonteCarloStudy.generate_dataset(;
    solve_args = (; solver = IBRGameSolver(), control_system, player_cost_models_gt, x0, T),
    noise_levels = unique([0:0.001:0.01; 0.01:0.005:0.03; 0.03:0.01:0.1]),
    n_observation_sequences_per_noise_level,
)

## Estimation
estimator_setup = (;
    dataset,
    control_system,
    player_cost_models = player_cost_models_gt,
    solver_attributes = (; print_level = 1),
)
estimator_setup_partial =
    merge(estimator_setup, (; expected_observation = x -> x[partial_state_indices, :]))
@run_cached estimates_conKKT =
    MonteCarloStudy.estimate(InverseKKTConstraintSolver(); estimator_setup...)
@run_cached estimates_conKKT_partial =
    MonteCarloStudy.estimate(InverseKKTConstraintSolver(); estimator_setup_partial...)
@run_cached estimates_resKKT =
    MonteCarloStudy.estimate(InverseKKTResidualSolver(); estimator_setup...)
@run_cached estimates_resKKT_partial =
    MonteCarloStudy.estimate(InverseKKTResidualSolver(); estimator_setup_partial...)

## Forward Solution Augmentation
augmentor_kwargs = (;
    solver = KKTGameSolver(),
    control_system,
    player_cost_models_gt,
    x0,
    T,
    match_equilibrium = (; forward_solution_gt.x),
    init = (; forward_solution_gt.x, forward_solution_gt.u),
    solver_attributes = (; print_level = 1),
)
@run_cached augmented_estimates_resKKT =
    MonteCarloStudy.augment_with_forward_solution(estimates_resKKT; augmentor_kwargs...)
@run_cached augmented_estimates_resKKT_partial =
    MonteCarloStudy.augment_with_forward_solution(estimates_resKKT_partial; augmentor_kwargs...)
estimates = [
    estimates_conKKT
    estimates_conKKT_partial
    augmented_estimates_resKKT
    augmented_estimates_resKKT_partial
]

## Error Ststistics Computation
demo_gt = merge((; player_cost_models_gt), forward_solution_gt)
@save_json errstats = map(estimates) do estimate
    MonteCarloStudy.estimator_statistics(estimate; dataset, demo_gt, position_indices)
end

## Visualization
demo_noise_level = 0.1
trajectory_viz_domain = (; x_position_domain = (-1.2, 1.2), y_position_domain = (-1.2, 1.2))

@save_json trajectory_data_gt =
    TrajectoryVisualization.trajectory_data(control_system, forward_solution_gt.x)

@save_json trajectory_data_obs = [
    TrajectoryVisualization.trajectory_data(control_system, d.x) for
    d in dataset if d.σ == demo_noise_level
]

@save_json trajectory_data_estimates =
    map.(
        e -> TrajectoryVisualization.trajectory_data(control_system, e.x),
        grouped(
            e -> e.estimator_name,
            Iterators.filter(
                e -> e.converged && dataset[e.observation_idx].σ == demo_noise_level,
                estimates,
            ),
        ),
    )

ground_truth_viz =
    TrajectoryVisualization.visualize_trajectory(
        trajectory_data_gt;
        canvas = VegaLite.@vlplot(width = 200, height = 200),
        legend = VegaLite.@vlfrag(orient = "top", offset = 5),
        trajectory_viz_domain...,
    ) + VegaLite.@vlplot(
        data = filter(s -> s.t == 1, trajectory_data_gt),
        mark = {"text", dx = 8, dy = 8},
        text = "player",
        x = "px",
        y = "py",
    )

observations_bundle_viz =
    VegaLite.@vlplot() + MonteCarloStudy.visualize_trajectory_batch(
        trajectory_data_obs;
        trajectory_viz_domain...,
        draw_line = false,
    )

viz_trajectory_estiamtes = Dict(
    k => TrajectoryVisualization.visualize_trajectory_batch(v; trajectory_viz_domain...) for
    (k, v) in trajectory_data_estimates
)

@saveviz demo_trajs_viz = hcat(
    VegaLite.@vlplot(title = "Ground Truth") + ground_truth_viz,
    VegaLite.@vlplot(title = "Observations") + observations_bundle_viz,
)

@saveviz ours_trajs_viz = hcat(
    VegaLite.@vlplot(title = "Full Observation") + viz_trajectory_estiamtes["Ours Full"],
    VegaLite.@vlplot(title = "Partial Observation") + viz_trajectory_estiamtes["Ours Partial"],
)

@saveviz baseline_trajs_viz = hcat(
    VegaLite.@vlplot(title = "Full Observation") + viz_trajectory_estiamtes["Baseline Full"],
    VegaLite.@vlplot(title = "Partial Observation") + viz_trajectory_estiamtes["Baseline Partial"],
)

frame = [-1.5n_observation_sequences_per_noise_level, 0]
@saveviz parameter_error_viz =
    errstats |> MonteCarloStudy.visualize_paramerr(; frame, round_x_axis = false)
@saveviz position_error_viz =
    errstats |> MonteCarloStudy.visualize_poserr(; frame, round_x_axis = false)


using CartesianGeneticProgramming
using ArgParse
using ArcadeLearningEnvironment
include("../graphing/graph_utils.jl")

s = ArgParseSettings()
@add_arg_table s begin
    "--cfg"
    help = "configuration script"
    default = "cfg/atari_ram.yaml"
    "--game"
    help = "game rom name"
    default = "pong"
    "--seed"
    help = "random seed for evolution"
    arg_type = Int
    default = 0
    "--ind"
    help = "individual for evaluation"
    arg_type = String
    default = ""
end
args = parse_args(ARGS, s)

ale = ALE_new()
loadROM(ale, args["game"])
n_in = length(getRAM(ale))
n_out = length(getMinimalActionSet(ale))
ALE_del(ale)

cfg = get_config(args["cfg"]; game=args["game"], n_in=n_in, n_out=n_out)

ind = CGPInd(cfg, read(args["ind"], String))
dot_array = walk_nodes(ind)
dot_file  = "graph.dot"
open(dot_file, "w") do f
   for i in dot_array
      println(f, i)
   end
end
# to view from commandline:
# $ dot -Tpng graph.dot -o graph.png
# $ display graph.png