text:
# =========================================================================
#    CORE FUNCTIONS AND MECHANISMS
# =========================================================================

passes(i) = [k for k in K() if (i in K(k).cover)]
s(key::S,R::Dict,θ) = sum(θ[q.r,q.k,q.t] for q in Q(key,R))
issinteger(val) = abs(round(val) - val) < 1e-8

function Q(key::S,R::Dict) #key pasti adalah S
    if isempty(key.seq)
        #if formnya S(k,β[])
        q = Vector{NamedTuple}()
        for r in keys(R)
            push!(q,(r=r,k=key.k,t=key.t))
        end
        return q
    else
        #if formnya S(k,β[β(i,v),...])
        q = Vector{Vector{NamedTuple}}()
        for tes in key.seq
            res = Vector{NamedTuple}()
            for r in keys(R)
                if R[r][key.k,key.t].u[tes.i] >= tes.v
                    push!(res,(r=r,k=key.k,t=key.t))
                end
            end
            push!(q,res) #add col each bound
        end
        return reduce(intersect,q)
    end
    #column class extraction
end

const vehicle_priority = Ref{Any}(nothing)
k_priority() = vehicle_priority[] #clear

const vertex_priority = Ref{Any}(nothing)
i_priority() = vertex_priority[] #clear

function prioritize_vertex!()
    #determine branching priorities
    res = JuMP.Containers.DenseAxisArray{Float64}(undef,V())
    res .= 0
    for i in V()
        slack = (V(i).START - sum(d(i,t) for t in T())) - V(i).MIN #after demand
        surplus = V(i).MAX - (V(i).START - sum(d(i,t) for t in T())) #after demand
        spare = min(slack,surplus)/(V(i).MAX - V(i).MIN)
        res[i] = 1/abs(spare)
    end
    return vertex_priority[] = res
end #clear

function i_priority(cover)
    if isa(cover,Int64)
        return vertex_priority[][cover] #bentuknya DenseAxisArray
    else
        new = Dict()
        for i in cover
            new[i] = i_priority(i)
        end
        return new
    end
end

function prioritize_vehicle!()
    cov = DataFrame(k=Int64[],cov=Int64[])
    for k in K()
        append!(cov,DataFrame(k=k,cov=length(K(k).cover)))
    end
    sort!(cov,:cov,rev=true)

    return vehicle_priority[] = cov.k
end

function find_branch(R,θ)
    #iterate over k_priority()
    for k in k_priority(), t in T()
        key = S(k,t,β[])
        if !issinteger(s(key,R,θ))
            return key
        end
    end

    for k in k_priority(),t in T()
        key = S(k,t,β[])
        rec = separate(key,deepcopy(K(k).cover),[],R,θ)
        if !isempty(rec) #rec is a vector of S
            lastcomp = [candidate.seq[end].i for candidate in rec]
            to_collect = findmax(i_priority(lastcomp))[2]
            to_return = filter(p -> p.seq[end].i == to_collect,rec)
            return to_return[1]
        end
    end
end

function separate(Seq::S,cover,record,R,θ)
    F = fractional_column(Seq,R,θ)

    if isempty(F)
        return record
    end

    found = false
    for i in cover
        v = find_separator(i,F,R)
        newSeq = S(Seq.k,Seq.t,β[β(i,v)])
        if !issinteger(s(newSeq,R,θ)) > 0
            push!(record,newSeq)
            found = true
        end
    end
    if found
        return record
    end

    i✶ = findmax(i_priority(cover))[2] #determine next highest priority

    v = find_separator(i✶,F,R)
    newSeq = S(Seq.k,Seq.t,β[β(i✶,v)]) #BUILD NEWSEQ

    to_remove = findall(p -> p == i✶,cover)
    splice!(cover,to_remove) #BUILDNEWCOVER

    record = separate(newSeq,cover,record,R,θ)
    return record
end

function fractional_column(Seq::S,R::Dict,θ)
    F = Vector{NamedTuple}()
    for r in θ.axes[1], k in θ.axes[2], t in θ.axes[3]
        if !issinteger(θ[r,k,t])
            push!(F,(r=r,k=k,t=t))
        end
    end

    return intersect(F,Q(Seq,R))
end

function find_separator(i::Int64,F,R)
    test = Vector{Int64}()
    for f in F
        push!(test,round(R[f.r][(f.k,f.t)].u[i]))
    end
    unique!(test)

    return ceil(median(test))
end

function f(key::S,R::Dict,θ)
    return sum(θ[q.r,q.k,q.t] - floor(θ[q.r,q.k,q.t]) for q in Q(key,R))
end

const subproblems = Ref{Any}(nothing)
callSub() = subproblems[]
callSub(k,t) = subproblems[][(k,t)]

function buildSub!(n::node)
    R = Dict{Tuple,Model}()

    for k in K(), t in T()
        sp = Model(get_optimizer())
        set_silent(sp)

        if solver_name(sp) == "Gurobi"
            set_optimizer_attribute(sp,"MIPFocus",2)
            set_optimizer_attribute(sp,"Threads",1)
            set_optimizer_attribute(sp,"NodefileStart",0.5)
            set_optimizer_attribute(sp, "NumericFocus",3)
        end

        R[(k,t)] = sp
    end

    @sync begin
        @inbounds for k in K(), t in T()
            Threads.@spawn begin
                sp = R[(k,t)]

                #VARIABLE DEFINITION
                @variable(sp, u[K(k).cover] >= 0, Int) #peti diantar
                @variable(sp, v[K(k).cover] >= 0, Int) #peti dijemput
                @variable(sp, l[K(k).cover, K(k).cover] >= 0, Int) #load kendaraan
                @variable(sp, o[i = K(k).cover] <= K(k).BP[i], Bin) #origin
                @variable(sp, x[K(k).cover, K(k).cover], Bin) #penggunaan segmen

                @constraint(sp,
                    sum(u[i] for i in K(k).cover) - sum(v[i] for i in K(k).cover) == 0
                ) #all pickup delivered
                @constraint(sp, [i = K(k).cover],
                    sum(l[j,i] for j in K(k).cover) - sum(l[i,j] for j in K(k).cover) ==
                    u[i] - v[i]
                ) #load balance
                @constraint(sp,
                    [i = K(k).cover, j = K(k).cover], l[i,j] <= K(k).Q * x[i,j]
                ) #XL
                @constraint(sp, [i = K(k).cover],
                    sum(x[j,i] for j in K(k).cover) - sum(x[i,j] for j in K(k).cover) == 0
                ) #traversal
                @constraint(sp, [i = K(k).cover], v[i] <= K(k).Q * o[i]) #VZ
                @constraint(sp, sum(o[i] for i in K(k).cover) <= 1) #one start

                F = Dict(1:length(n.bounds) .=> n.bounds)
                uB = filter(f -> last(f).sense=="<="&&last(f).S.k==k && last(f).S.t == t, F)
                lB = filter(f -> last(f).sense==">="&&last(f).S.k==k && last(f).S.t == t, F)

                @variable(sp, g[keys(uB)], Bin)
                @variable(sp, h[keys(lB)], Bin)

                q = col(u,v,l,o,x)

                for j in keys(uB)
                    η = @variable(sp, [F[j].S.seq], Bin)
                    @constraint(sp, g[j] >= 1 - sum((1 - η[e]) for e in F[j].S.seq))
                    @constraint(sp, [e = F[j].S.seq],
                    (K(k).Q - e.v + 1) * η[e] >= getproperty(q,:u)[e.i] - e.v + 1)
                end

                for j in keys(lB)
                    η = @variable(sp, [F[j].S.seq], Bin)
                    @constraint(sp, [e = F[j].S.seq], h[j] <= η[e])
                    @constraint(sp, [e = F[j].S.seq],
                    e.v * η[e] <= getproperty(q,:u)[e.i])
                end

                optimize!(sp)
            end
        end
    end

    return subproblems[] = R
end

function master(n::node)
    mp = Model(get_optimizer())
    set_silent(mp)

    if solver_name(mp) == "Gurobi"
        set_optimizer_attribute(mp,"MIPFocus",2)
        set_optimizer_attribute(mp,"NodefileStart",0.5)
        set_optimizer_attribute(mp, "NumericFocus",3)
    end

    R = Dict(1:length(n.columns) .=> n.columns)

    # ================================
    #    MODEL CONSTRUCTION
    # ================================
    @variable(mp, θ[keys(R), K(), T()] >= 0)
    @variable(mp, I[i = V(), vcat(first(T()) - 1, T())])
    @variable(mp, 0 <= slack[k = K(), t = T(), i = K(k).cover] <= n.stab.slLim[(k,t)])
    @variable(mp, 0 <= surp[k = K(), t = T(), i = K(k).cover] <= n.stab.suLim[(k,t)])

    @objective(mp, Min,
        sum(
            θ[r,k,t] * (
                sum(
                    dist(i,j) * (
                        K(k).vx * R[r][(k,t)].x[i,j] +
                        K(k).vl * R[r][(k,t)].l[i,j]
                    )
                    for i in K(k).cover, j in K(k).cover
                ) +
                sum(
                    K(k).fd * R[r][(k,t)].u[i]
                    for i in K(k).cover
                ) +
                sum(
                    K(k).fp * R[r][(k,t)].o[i]
                    for i in K(k).cover
                )
            )
            for r in keys(R), k in K(), t in T()
        ) + #column costs
        sum(
            V(i).h * I[i,t]
            for i in V(), t in T()
        ) + #inventory costs
        sum(
            n.stab.slCoeff * slack[k,t,i]
            for k in K(), t in T(), i in K(k).cover
        ) - #stabilizer
        sum(
            n.stab.suCoeff * surp[k,t,i]
            for k in K(), t in T(), i in K(k).cover
        ) #stabilizer
    )

    @constraint(mp, λ[i = V(), t = T()],
        I[i,t-1] + sum(R[r][(k,t)].u[i] * θ[r,k,t] for r in keys(R), k in passes(i)) +
        sum(slack[k,t,i] for k in passes(i))  ==
        sum(R[r][(k,t)].v[i] * θ[r,k,t] for r in keys(R), k in passes(i)) +
        sum(surp[k,t,i] for k in passes(i)) + d(i,t) + I[i,t]
    ) #inventory balance

    @constraint(mp, δ[k = K(), t = T(), i = K(k).cover],
        sum(R[r][(k,t)].o[i] * θ[r,k,t] for r in keys(R)) <= K(k).BP[i]
    ) #limit starting point

    @constraint(mp, ϵ[k = K(), t = T()],
        sum(θ[r,k,t] for r in keys(R)) <= sum(K(k).BP[i] for i in K(k).cover)
    ) #multiplicity

    @constraint(mp, [i = V(), t = T()],
        V(i).MIN <= I[i,t] <= V(i).MAX
    ) #inventory max min

    @constraint(mp, [i = V()],
        I[i,first(T())-1] == V(i).START
    ) #inventory start

    F = Dict(1:length(n.bounds) .=> n.bounds)
    uB = filter(f -> last(f).sense == "<=",F)
    lB = filter(f -> last(f).sense == ">=",F)

    @constraint(mp, ρ[j = keys(uB)], s(F[j].S,R,θ) <= F[j].κ)
    @constraint(mp, σ[j = keys(lB)], s(F[j].S,R,θ) >= F[j].κ)

    optimize!(mp)

    return mp
end

function getDuals(mp::Model)
    λ = dual.(mp.obj_dict[:λ])
    δ = dual.(mp.obj_dict[:δ])
    ϵ = dual.(mp.obj_dict[:ϵ])

    ρ = dual.(mp.obj_dict[:ρ])
    σ = dual.(mp.obj_dict[:σ])

    return dv(λ,δ,ϵ,ρ,σ)
end

function sub(n::node,duals::dv)
    @sync begin
        @inbounds for k in K(), t in T()
            Threads.@spawn begin
                sp = callSub(k,t)

                F = Dict(1:length(n.bounds) .=> n.bounds)
                uB = filter(f->last(f).sense == "<="&&last(f).S.k == k&&last(f).S.t == t, F)
                lB = filter(f->last(f).sense == ">="&&last(f).S.k == k&&last(f).S.t == t, F)

                #ADD OBJECTIVE
                @objective(sp, Min,
                    sum(
                        dist(i,j) * (
                            K(k).vx * sp.obj_dict[:x][i,j] +
                            K(k).vl * sp.obj_dict[:l][i,j]
                        )
                        for i in K(k).cover, j in K(k).cover
                    ) +
                    sum(K(k).fd * sp.obj_dict[:u][i] for i in K(k).cover) +
                    sum(K(k).fp * sp.obj_dict[:o][i] for i in K(k).cover) -
                    sum(
                        (sp.obj_dict[:u][i] - sp.obj_dict[:v][i]) * duals.λ[i,t]
                        for i in K(k).cover
                    ) -
                    sum(sp.obj_dict[:o][i] * duals.δ[k,t,i] for i in K(k).cover) -
                    duals.ϵ[k,t] -
                    sum(sp.obj_dict[:g][j] * duals.ρ[j] for j in keys(uB)) -
                    sum(sp.obj_dict[:h][j] * duals.σ[j] for j in keys(lB))
                )

                optimize!(sp)
                #println("($k): $(objective_value(sp))")
            end
        end
    end

    return callSub()
end

function getCols(sp)
    if isa(sp,Model)
        u = value.(sp.obj_dict[:u])
        v = value.(sp.obj_dict[:v])
        l = value.(sp.obj_dict[:l])
        o = value.(sp.obj_dict[:o])
        x = value.(sp.obj_dict[:x])

        return col(u,v,l,o,x)
    elseif isa(sp,Dict)
        new = Dict{Tuple,col}()
        for r in keys(sp)
            new[r] = getCols(sp[r])
        end

        return new
    end
end

function colvals()
    collection = 0
    for k in K(), t in T()
        collection += objective_value(callSub(k,t))
    end

    return sum(collection)
end

function updateStab!(stab::stabilizer,param::Float64)
    for kt in keys(stab.slLim)
        stab.slLim[kt] = floor(param * stab.slLim[kt])
    end
    for kt in keys(stab.suLim)
        stab.suLim[kt] = floor(param * stab.suLim[kt])
    end

    return stab
end

checkStab(mp::Model) = sum(value.(mp.obj_dict[:slack])) + sum(value.(mp.obj_dict[:surp]))

function colGen(n::node;maxCG::Float64,track::Bool)
    terminate = false
    iter = 0
    buildSub!(n)

    while !terminate
        if iter < maxCG
            mp = master(n)

            if has_values(mp) && has_duals(mp)
                if track #print master problem obj
                    println("obj: $(objective_value(mp))")
                end

                duals = getDuals(mp)
                sp = sub(n,duals)

                if track #print subproblem price
                    println("price: $(colvals())")
                end

                if isapprox(colvals(),0,atol = 1e-8) || colvals() > 0
                    if isapprox(checkStab(mp),0,atol = 1e-8)
                        terminate = true #action
                        push!(n.status,"EVALUATED") #report
                        if track
                            println("EVALUATED")
                        end
                    else
                        updateStab!(n.stab,0.2) #action
                        push!(n.status,"STABILIZED") #report
                        if track
                            println("STABILIZED")
                        end
                    end
                else
                    push!(n.columns,getCols(sp)) #action
                    push!(n.status,"ADD_COLUMN") #report
                    if track
                        println("ADD_COLUMN")
                    end
                end

                iter += 1 #iteration update
            else
                terminate = true #action
                push!(n.status,"NO_SOLUTION")
                if track
                    println("NO_SOLUTION")
                end
            end
        else
            terminate = true #action
            pop!(n.columns)
            updateStab!(n.stab,1e-8)
            push!(n.status,"EVALUATED") #report
            if track
                println("EVALUATED")
            end
        end
    end

    if n.status[end] == "NO_SOLUTION"
        println("NODE $(n.self) FAILED.")
    else
        if integerCheck(n)
            push!(n.status,"INTEGER")
            println("NODE $(n.self) INTEGER")
        else
            println("NODE $(n.self) FINISHED.")
        end
    end

    return n
end

function origin(n::node)
    o = JuMP.Containers.DenseAxisArray{Float64}(undef,V(),K(),T())
    u = JuMP.Containers.DenseAxisArray{Float64}(undef,V(),K(),T())
    v = JuMP.Containers.DenseAxisArray{Float64}(undef,V(),K(),T())
    x = JuMP.Containers.DenseAxisArray{Float64}(undef,V(),V(),K(),T())
    l = JuMP.Containers.DenseAxisArray{Float64}(undef,V(),V(),K(),T())

    o .= 0
    u .= 0
    v .= 0
    l .= 0
    x .= 0

    R = Dict(1:length(n.columns) .=> n.columns)
    θ = value.(master(n).obj_dict[:θ])

    for k in K(), t in T(), i in K(k).cover
        o[i,k,t] = sum(R[r][(k,t)].o[i] * θ[r,k,t] for r in keys(R))
        u[i,k,t] = sum(R[r][(k,t)].u[i] * θ[r,k,t] for r in keys(R))
        v[i,k,t] = sum(R[r][(k,t)].v[i] * θ[r,k,t] for r in keys(R))
    end

    for k in K(), t in T(), i in K(k).cover, j in K(k).cover
        x[i,j,k,t] = sum(R[r][(k,t)].x[i,j] * θ[r,k,t] for r in keys(R))
        l[i,j,k,t] = sum(R[r][(k,t)].l[i,j] * θ[r,k,t] for r in keys(R))
    end

    return col(u,v,l,o,x)
end


text:
using LinearAlgebra, MixedModels, Random, Test

@testset "femat" begin
    trm = MixedModels.FeMat(hcat(ones(30), repeat(0:9, outer = 3)), ["(Intercept)", "U"])
    piv = trm.piv
    ipiv = invperm(piv)
    @test size(trm) == (30, 2)
    @test length(trm) == 60
    @test size(trm') == (2, 30)
    @test eltype(trm) == Float64
    @test trm.x === trm.wtx
    prd = trm'trm
    @test typeof(prd) == Matrix{Float64}
    @test prd == [30.0 135.0; 135.0 855.0][piv, piv]
    wts = rand(MersenneTwister(123454321), 30)
    MixedModels.reweight!(trm, wts)
    @test mul!(prd, trm', trm)[ipiv[1], ipiv[1]] ≈ sum(abs2, wts)
end

ext: jl


text:
verbose = true
 if verbose && !isdefined(Main,:RvSpectMLBase)  println("# Loading RvSpecMLBase")    end
 using RvSpectMLBase
 if verbose && !isdefined(Main,:EchelleInstruments)  println("# Loading EchelleInstruments")    end
 using EchelleInstruments, EchelleInstruments.NEID
 using DataFrames, Query

# USER: You must create a data_paths.jl file in one of the default_paths_to_search listed below. It need only contain one line:
# solar_data_path = "/home/eford/Data/SolarSpectra/NEID_solar/"
target_subdir = "good_days/DRPv0.7/"   # USER: Replace with directory of your choice
 #fits_target_str = "Solar"
 fits_target_str = "Sun"
 output_dir = "examples/output"
 paths_to_search_for_param = [pwd(),"examples",joinpath(pkgdir(RvSpectMLBase),"..","RvSpectML","examples"), "/gpfs/group/ebf11/default/ebf11/neid_solar"]
 # NOTE: make_manifest does not update its paths_to_search when default_paths_to_search is defined here, so if you change the line above, you must also include "paths_to_search=default_paths_to_search" in the make_manifest() function call below
 pipeline_plan = PipelinePlan()

reset_all_needs!(pipeline_plan)
if need_to(pipeline_plan,:read_spectra)
   if verbose println("# Finding what data files are avaliable.")  end
   eval(read_data_paths(paths_to_search=paths_to_search_for_param))
   @assert isdefined(Main,:solar_data_path)
   df_files = make_manifest(solar_data_path, target_subdir, NEID )
   #df_files = EchelleInstruments.make_manifest(solar_data_path, target_subdir, NEID )

   if verbose println("# Reading in customized parameters from param.jl.")  end
   eval(code_to_include_param_jl(paths_to_search=paths_to_search_for_param))
   end

if verbose println("# Reading in ", size(df_files_use,1), " FITS files.")  end
   @time all_spectra = map(NEID.read_data,eachrow(df_files_use))
   dont_need_to!(pipeline_plan,:read_spectra)

#=
# Pre-ship corrections
if verbose println("# Applying wavelength corrections.")  end
   @assert isdefined(Main,:ancilary_solar_data_path)
   NEID.read_drift_corrections!(joinpath(ancilary_solar_data_path,"SolarRV20190918_JD_SciRV_CalRV.txt"), df_files_use)
   NEID.read_barycentric_corrections!(joinpath(ancilary_solar_data_path,"SolarTelescope2019-09-18_inclGravRedshiftAirMassAltitude.csv"), df_files_use)
   NEID.read_differential_extinctions!(joinpath(ancilary_solar_data_path,"20190918_diff_ex_full_fixed.txt"), df_files_use)
   apply_doppler_boost!(all_spectra,df_files_use)
   all_spectra
end
=#

max_stars_repo_name: RvSpectML/EchelleInstruments

max_issues_repo_name: RvSpectML/EchelleInstruments

max_forks_repo_name: RvSpectML/EchelleInstruments


text:
#=
   Meeting problem.

   (I'm not sure about the source. The page I got it from does not exist anymore.)
   """
   Let's say that 4 people agree to meet between 3:00 P.M. and 4:00 P.M..
   One man can wait for 10 minutes. Another can wait for 15 minutes.
   Another can wait for 20 Minutes, and another still can wait for 30
   minutes. What is the probability that all 4 will meet?
   """

   Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

        prob    0.0836    0.2768     0.0028    0.0030   9775.6911    0.9999     7783.1935

    Cf meeting_under_the_clock,jl

=#
using Turing, StatsPlots, Distributions
include("jl_utils.jl")

@model function meeting_problem(w=[10,15,20,30])
    n = length(w)
    ps = tzeros(n)
    for i in 1:n
        ps[i] ~ DiscreteUniform(1,60)
    end
    # Will all people meet?
    prob ~ Dirac(reduce(intersect, [Set(ps[i]:ps[i]+w[i])  for i in 1:n]) != Set())
end

w = [10,15,20,30]
model = meeting_problem(w)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5), 1_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns[[:prob]])
# display(plot(chns))

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
abstract type NeuralDELayer <: Function end

"""
Constructs a neural ODE with the gradients computed using the adjoint
method[1]. At a high level this corresponds to solving the forward
differential equation, using a second differential equation that propagates
the derivatives of the loss  backwards in time.
This first solves the continuous time problem, and then discretizes following
the rules specified by the numerical ODE solver.
On the other hand, the 'neural_ode_rd' first disretizes the solution and then
computes the adjoint using automatic differentiation.

Ref
[1]L. S. Pontryagin, Mathematical Theory of Optimal Processes. CRC Press, 1987.

Arguments
≡≡≡≡≡≡≡≡
model::Chain defines the ̇x
x<:AbstractArray initial value x(t₀)
args arguments passed to ODESolve
kwargs key word arguments passed to ODESolve; accepts an additional key
    :callback_adj in addition to :callback. The Callback :callback_adj
    passes a separate callback to the adjoint solver.

"""
struct NeuralODE{M,P,RE,T,S,A,K} <: NeuralDELayer
    model::M
    p::P
    re::RE
    tspan::T
    solver::S
    args::A
    kwargs::K

    function NeuralODE(model,tspan,solver=nothing,args...;kwargs...)
        p,re = Flux.destructure(model)
        new{typeof(model),typeof(p),typeof(re),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(
            model,p,re,tspan,solver,args,kwargs)
    end

    function NeuralODE(model::FastChain,tspan,solver=nothing,args...;kwargs...)
        p = initial_params(model)
        re = nothing
        new{typeof(model),typeof(p),typeof(re),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(
            model,p,re,tspan,solver,args,kwargs)
    end
end

Flux.@functor NeuralODE

function (n::NeuralODE)(x,p=n.p)
    dudt_(u,p,t) = n.re(p)(u)
    prob = ODEProblem{false}(dudt_,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;n.kwargs...)
end

function (n::NeuralODE{M})(x,p=n.p) where {M<:FastChain}
    dudt_(u,p,t) = n.model(u,p)
    prob = ODEProblem{false}(dudt_,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;n.kwargs...)
end

struct NeuralDSDE{M,P,RE,M2,RE2,T,S,A,K} <: NeuralDELayer
    p::P
    len::Int
    model1::M
    re1::RE
    model2::M2
    re2::RE2
    tspan::T
    solver::S
    args::A
    kwargs::K
    function NeuralDSDE(model1,model2,tspan,solver=nothing,args...;kwargs...)
        p1,re1 = Flux.destructure(model1)
        p2,re2 = Flux.destructure(model2)
        p = [p1;p2]
        new{typeof(model1),typeof(p),typeof(re1),typeof(model2),typeof(re2),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(p,
            length(p1),model1,re1,model2,re2,tspan,solver,args,kwargs)
    end

    function NeuralDSDE(model1::FastChain,model2::FastChain,tspan,solver=nothing,args...;kwargs...)
        p1 = initial_params(model1)
        p2 = initial_params(model2)
        re1 = nothing
        re2 = nothing
        p = [p1;p2]
        new{typeof(model1),typeof(p),typeof(re1),typeof(model2),typeof(re2),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(p,
            length(p1),model1,re1,model2,re2,tspan,solver,args,kwargs)
    end
end

Flux.@functor NeuralDSDE

function (n::NeuralDSDE)(x,p=n.p)
    dudt_(u,p,t) = n.re1(p[1:n.len])(u)
    g(u,p,t) = n.re2(p[(n.len+1):end])(u)
    prob = SDEProblem{false}(dudt_,g,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

function (n::NeuralDSDE{M})(x,p=n.p) where {M<:FastChain}
    dudt_(u,p,t) = n.model1(u,p)
    g(u,p,t) = n.model2(u,p)
    prob = SDEProblem{false}(dudt_,g,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

struct NeuralSDE{P,M,RE,M2,RE2,T,S,A,K} <: NeuralDELayer
    p::P
    len::Int
    model1::M
    re1::RE
    model2::M2
    re2::RE2
    tspan::T
    nbrown::Int
    solver::S
    args::A
    kwargs::K
end

function NeuralSDE(model1,model2,tspan,nbrown,solver=nothing,args...;kwargs...)
    p1,re1 = Flux.destructure(model1)
    p2,re2 = Flux.destructure(model2)
    p = [p1;p2]
    NeuralSDE(p,length(p1),model1,re1,model2,re2,tspan,nbrown,solver,args,kwargs)
end

Flux.@functor NeuralSDE

function (n::NeuralSDE)(x,p=n.p)
    dudt_(u,p,t) = n.re1(p[1:n.len])(u)
    g(u,p,t) = n.re2(p[(n.len+1):end])(u)
    prob = SDEProblem{false}(dudt_,g,x,n.tspan,p,noise_rate_prototype=zeros(Float32,length(x),n.nbrown))
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

struct NeuralCDDE{P,M,RE,H,L,T,S,A,K} <: NeuralDELayer
    p::P
    model::M
    re::RE
    hist::H
    lags::L
    tspan::T
    solver::S
    args::A
    kwargs::K
end

function NeuralCDDE(model,tspan,hist,lags,solver=nothing,args...;kwargs...)
    p,re = Flux.destructure(model)
    NeuralCDDE(p,model,re,hist,lags,tspan,solver,args,kwargs)
end

Flux.@functor NeuralCDDE

function (n::NeuralCDDE)(x,p=n.p)
    function dudt_(u,h,p,t)
        _u = vcat(u,(h(p,t-lag) for lag in n.lags)...)
        n.re(p)(_u)
    end
    prob = DDEProblem{false}(dudt_,x,n.hist,n.tspan,p,constant_lags = n.lags)
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end


text:
#= ------------------------------------------------------------------------

    Run

------------------------------------------------------------------------ =#

import Base: isempty, isless, -, +, contains, ∈
export Run
export isempty, isless, invert
export translate, -, +
export contains, isoverlapping, istouching, isclose

"""
    Run

A run is a set of consecutive coordinates within a row (possibly partial) of a
region. It consists of a discrete row coordinate (of type Signed) and a range
of discrete column coordinates (of type UnitRange{Int64}).

Runs within a region specify a sort order: one run is smaller than the other if it starts
before the other run modeling the coordinates from left to right and top to
bottom.
"""
struct Run
    row::Integer
    columns::UnitRange{Int64}
end

"""
    isempty(x::Run)

Discover whether the run is empty.

```jldoctest
julia> using Regions

julia> isempty(Run(1, 1:10))
false

julia> isempty(Run(2, 1:1))
false

julia> isempty(Run(3, 1:0))
true
```
"""
isempty(x::Run) = isempty(x.columns)

"""
    isless(x::Run, y::Run)

Compare two runs according to their natural sort order. First, their rows are compared,
and if they are equal, their column ranges are compared.

```jldoctest reg
julia> using Regions

julia> isless(Run(0, 1:10), Run(1, 0:10))
true

julia> isless(Run(1, 1:10), Run(1, 2:10))
true
```
"""
isless(x::Run, y::Run) = (x.row < y.row) || ((x.row == y.row) && (x.columns < y.columns))

"""
    invert(x::Run)

Invert a run. Inversions mirrors a run at the origin. A run is inverted by negating
its row and inverting its columns.

In addition to the invert method, you can also use the unary - operator.

```jldoctest reg
julia> using Regions

julia> invert(Run(1, 20:30))
Run(-1, -30:-20)

julia> -Run(-1, -30:-20)
Run(1, 20:30)
```
"""
invert(x::Run) = -x
-(x::Run) = Run(-x.row, invert(x.columns))

"""
    translate(r::Run, x::Integer, y::Integer)
    translate(r::Run, a::Vector{Int64})

Translate a run. Translation moves a run. A run is translated by adding offsets to its row and
columns.

In addition to the translate method, you can also use the + or - operators
to translate a run.

```jldoctest reg
julia> using Regions

julia> translate(Run(1, 20:30), 10, 20)
Run(21, 30:40)

julia> translate(Run(1, 2:3), [10, 20])
Run(21, 12:13)

julia> Run(1, 2:3) + [10, 20]
Run(21, 12:13)

julia> [1, 2] + Run(0, 0:10)
Run(2, 1:11)

julia> Run(0, 0:100) - [5, 25]
Run(-25, -5:95)
```
"""
translate(a::Run, x::Integer, y::Integer) = a + [x, y]
translate(a::Run, b::Vector{Int64}) = a + b
+(a::Run, b::Vector{Int64}) = Run(a.row + b[2], a.columns + b[1])
+(a::Vector{Int64}, b::Run) = Run(a[2] + b.row, a[1] + b.columns)
-(a::Run, b::Vector{Int64}) = Run(a.row - b[2], a.columns - b[1])

"""
    contains(r::Run, x::Integer, y::Integer)
    contains(r::Run, a::Vector{Int64})

Test if run r contains position (x, y).

```jldoctest reg
julia> using Regions

julia> contains(Run(2, 5:10), 7, 2)
true

julia> contains(Run(2, 5:10), 7, 3)
false

julia> contains(Run(2, 5:10), 12, 2)
false
```
"""
contains(r::Run, x::Integer, y::Integer) = (r.row == y) && contains(r.columns, x)
contains(r::Run, a::Vector{Int64}) = contains(r, a[1], a[2])
∈(a::Vector{Int64}, r::Run) = contains(r, a)

"""
    isoverlapping(x::Run, y::Run)

Test if two runs overlap.

```jldoctest
julia> using Regions

julia> isoverlapping(0:10, 5:15)
true

julia> isoverlapping(0:10, 20:30)
false
```
"""
isoverlapping(x::Run, y::Run) = x.row == y.row && isoverlapping(x.columns, y.columns)

"""
    istouching(x::Run, y::Run)

Test if two runs touch.

```jldoctest
julia> using Regions

julia> istouching(Run(5, 0:10), Run(6, 11:21))
true

julia> istouching(Run(5, 0:10), Run(5, 12:22))
false
```
"""
istouching(x::Run, y::Run) = abs(x.row - y.row) ≤ 1 && istouching(x.columns, y.columns)

"""
    isclose(a::Run, b::Run, x::Integer, y::Integer)
    isclose(a::Run, b::Run, d::Integer)
    isclose(x::Run, y::Run, distance::Vector[Int64])

Test if two runs are close.

If distance == 0 this is the same as isoverlapping().
If distance == 1 this is the same as istouching().
If distance > 1 this is testing of closeness.

```jldoctest
julia> using Regions

julia> isclose(Run(5, 0:10), Run(8, 15:20), 5, 3)
true

julia> isclose(Run(5, 0:10), Run(8, 15:20), 2, 2)
false
```
"""
isclose(a::Run, b::Run, x::Integer, y::Integer) = abs(a.row - b.row) <= y && isclose(a.columns, b.columns, x)
isclose(a::Run, b::Run, d::Integer) = isclose(a, b, d, d)
isclose(a::Run, b::Run, distance::Vector{Int64}) = isclose(a, b, distance[1], distance[2])


text:
module Auto2D

using AutomotiveDrivingModels
using AutoViz
using PDMats
using NGSIM
using ForwardNets
import Reel

export gen_simparams, reset, tick, reward, observe, step, isdone, action_space_bounds, observation_space_bounds, render
export SimParams, reel_drive, GaussianMLPDriver, load_gru_driver

##################################
# Gaussian MLP Driver

type GaussianMLPDriver{A<:DriveAction, F<:Real, G<:Real, E<:AbstractFeatureExtractor, M<:MvNormal} <: DriverModel{A, IntegratedContinuous}
    net::ForwardNet
    rec::SceneRecord
    pass::ForwardPass
    input::Vector{F}
    output::Vector{G}
    extractor::E
    mvnormal::M
    context::IntegratedContinuous
end

_get_Σ_type{Σ,μ}(mvnormal::MvNormal{Σ,μ}) = Σ
function GaussianMLPDriver{A <: DriveAction}(::Type{A}, net::ForwardNet, extractor::AbstractFeatureExtractor, context::IntegratedContinuous;
    input::Symbol = :input,
    output::Symbol = :output,
    Σ::Union{Real, Vector{Float64}, Matrix{Float64},  Distributions.AbstractPDMat} = 0.1,
    rec::SceneRecord = SceneRecord(2, context.Δt),
    )

    pass = calc_forwardpass(net, [input], [output])
    input_vec = net[input].tensor
    output = net[output].tensor
    mvnormal = MvNormal(Array(Float64, 2), Σ)
    GaussianMLPDriver{A, eltype(input_vec), eltype(output), typeof(extractor), typeof(mvnormal)}(net, rec, pass, input_vec, output, extractor, mvnormal, context)
end

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

include("../validation/load_train_test_split.jl")
include("../pull_traces/multifeatureset.jl")
include("../validation/RootDir.jl")

function get_train_segments(trajdatas::Dict{Int, Trajdata}, nsteps::Int)

    assignment = load_assignment()
    evaldata = load_evaldata()

    # pull all segments that are for training
    all_train_segments = evaldata.segments[find(assignment .== FOLD_TRAIN)]

    # only keep segments that are long enough and are not within 500 frames of start or end
    train_segments = TrajdataSegment[]
    for seg in all_train_segments
        if !haskey(trajdatas, seg.trajdata_index) # only keep valid segments
            continue
        end
        trajdata = trajdatas[seg.trajdata_index]
        frame_lo = max(seg.frame_lo, 500)
        frame_hi = min(seg.frame_hi, nframes(trajdata) - 1000)
        if frame_hi - frame_lo > nsteps # if it is long enough
            push!(train_segments, TrajdataSegment(seg.trajdata_index, seg.egoid, frame_lo, frame_hi))
        end
    end

    train_segments
end

##################################
# SimParams
include("../validation/load_policy.jl")

type SimState
    frame::Int
    start_frame::Int
    egoid::Int
    trajdata_index::Int

    scene::Scene
    rec::SceneRecord

    playback_reactive_active_vehicle_ids::Set{Int}
    playback_reactive_speeds::Dict{Int, Float64}
    hidden_dict::Dict{Int, Vector{Float64}}

    function SimState(context::IntegratedContinuous, rec_size::Int)
        retval = new()
        retval.scene = Scene()
        retval.rec = SceneRecord(rec_size, context.Δt)
        retval.playback_reactive_active_vehicle_ids = Set{Int}()
        retval.playback_reactive_speeds = Dict{Int, Float64}()
        retval.hidden_dict = Dict{Int, Vector{Float64}}()
        retval
    end
end

type SimParams
    col_weight::Float64 # reward for collision
    off_weight::Float64 # reward for being offroad
    rev_weight::Float64 # reward for reversal (negative speeds)
    jrk_weight::Float64 # reward scale for square jerk
    acc_weight::Float64 # reward scale for square acceleration
    cen_weight::Float64 # reward scale for lane center offset
    ome_weight::Float64 # reward scale for turn acceleration (omega)
    use_debug_reward::Bool # false by default

    context::IntegratedContinuous
    prime_history::Int
    ego_action_type::DataType

    safety_model::DriverModel{LatLonAccel, IntegratedContinuous}

    use_playback_reactive::Bool
    playback_reactive_model::DriverModel{LatLonAccel, IntegratedContinuous}
    playback_reactive_threshold_brake::Float64
    playback_reactive_scene_buffer::Scene

    model_all::Bool
    driver_model::GaussianMLPDriver

    trajdatas::Dict{Int, Trajdata}
    segments::Vector{TrajdataSegment}

    nsteps::Int
    step_counter::Int
    simstates::Vector{SimState}
    features::Vector{Float64}
    extractor::MultiFeatureExtractor
end

function SimParams(trajdatas::Dict{Int, Trajdata}, segments::Vector{TrajdataSegment},
    col_weight::Float64,
    off_weight::Float64,
    rev_weight::Float64,
    jrk_weight::Float64,
    acc_weight::Float64,
    cen_weight::Float64,
    ome_weight::Float64,
    use_debug_reward::Bool,
    use_playback_reactive::Bool,
    model_all::Bool,
    playback_reactive_threshold_brake::Float64, # [m/s²]
    nsimstates::Int,
    prime_history::Int,
    nsteps::Int,
    ego_action_type::DataType,
    extractor::MultiFeatureExtractor,
    context = IntegratedContinuous(NGSIM_TIMESTEP,1),
    )
    println("Begin SimParams")
    simstates = Array(SimState, nsimstates)
    for i in 1 : length(simstates)
        simstates[i] = SimState(context, prime_history+1)
    end
    println("loop end")
    features = Array(Float64, length(extractor))

    safety_policy = LatLonSeparableDriver(
        context,
        ProportionalLaneTracker(σ=0.1, kp=3.0, kd=2.0),
        IntelligentDriverModel(σ=0.1, k_spd=1.0, T=0.5, s_min=1.0, a_max=3.0, d_cmf=2.5),
        )
    println("safe policy end")
    playback_reactive_model = LatLonSeparableDriver(
        context,
        ProportionalLaneTracker(σ=0.1, kp=3.0, kd=2.0),
        IntelligentDriverModel(σ=0.1, k_spd=1.0, T=0.5, s_min=1.0, a_max=3.0, d_cmf=2.5),
        )
    playback_reactive_scene_buffer = Scene()
    println("playback reactive model end")
    filepath = joinpath(ROOT_FILEPATH,"julia","validation","models","gail_gru.h5")
    iteration = 413
    driver_model = load_gru_driver(filepath, iteration)
    println("SimParams initial process end")
    SimParams(
        col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight, use_debug_reward,
        context, prime_history, ego_action_type, safety_policy,
        use_playback_reactive, playback_reactive_model,
        playback_reactive_threshold_brake, playback_reactive_scene_buffer, model_all, driver_model,
        trajdatas, segments, nsteps, 0, simstates, features, extractor)
end
function gen_simparams(trajdata_indeces::Vector,
    col_weight::Float64,
    off_weight::Float64,
    rev_weight::Float64,
    jrk_weight::Float64,
    acc_weight::Float64,
    cen_weight::Float64,
    ome_weight::Float64,
    use_debug_reward::Bool,
    use_playback_reactive::Bool,
    model_all::Bool,
    playback_reactive_threshold_brake::Float64,
    nsimstates::Int,
    prime_history::Int,
    nsteps::Int,
    ego_action_type::DataType,
    extractor::MultiFeatureExtractor,
    )

    println("loading trajdatas: ", trajdata_indeces); tic()
    trajdatas = Dict{Int, Trajdata}()
    println("loading done")
    println(load_trajdata(1))
    for i in trajdata_indeces
        trajdatas[i] = load_trajdata(i)
    end
    toc()

    println("loading training segments"); tic()
    segments = get_train_segments(trajdatas, nsteps)
    toc()

    println("HIT THIS THING")
    SimParams(trajdatas, segments, col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
              use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
              nsimstates, prime_history, nsteps, ego_action_type, extractor)
end
function gen_simparams_from_trajdatas(trajdata_filepaths::Vector, roadway_filepaths::Vector,
    col_weight::Float64,
    off_weight::Float64,
    rev_weight::Float64,
    jrk_weight::Float64,
    acc_weight::Float64,
    cen_weight::Float64,
    ome_weight::Float64,
    use_debug_reward::Bool,
    use_playback_reactive::Bool,
    model_all::Bool,
    playback_reactive_threshold_brake::Float64,
    nsimstates::Int,
    prime_history::Int,
    nsteps::Int,
    ego_action_type::DataType,
    extractor::MultiFeatureExtractor,
    )

    println("loading trajdatas"); tic()
    roadways = Dict{String, Roadway}()
    trajdatas = Dict{Int, Trajdata}()
    for (i,filepath) in enumerate(trajdata_filepaths)
        roadwayfp = roadway_filepaths[i]
        if !haskey(roadways, roadwayfp)
            roadways[roadwayfp] = open(io->read(io, Roadway), roadwayfp, "r")
        end
        trajdatas[i] = open(io->read(io, Trajdata, roadways[roadwayfp]), filepath, "r")
    end
    toc()

    println("generating training segments"); tic()
    segments = TrajdataSegment[]
    for (trajdata_index,trajdata) in trajdatas
        append!(segments, pull_continuous_segments(trajdata, trajdata_index))
    end
    toc()

    SimParams(trajdatas, segments, col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
              use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
              nsimstates, prime_history, nsteps, ego_action_type, extractor)
end
function gen_simparams(batch_size::Int, args::Dict)

    col_weight = get(args, "col_weight", -2.0)
    off_weight = get(args, "off_weight", -0.75)
    rev_weight = get(args, "rev_weight", -0.50)
    jrk_weight = get(args, "jrk_weight", -0.050)
    acc_weight = get(args, "acc_weight", -0.050)
    cen_weight = get(args, "cen_weight", -0.050)
    ome_weight = get(args, "ome_weight", -0.500)

    use_debug_reward = get(args, "use_debug_reward", false)
    use_playback_reactive = convert(Bool, get(args, "use_playback_reactive", false))
    playback_reactive_threshold_brake = get(args, "playback_reactive_threshold_brake", -2.0) # [m/s²]
    prime_history = get(args, "prime_history", 2)
    nsteps = get(args, "nsteps", 100)

    model_all = get(args, "model_all", false)

    extract_core = get(args, "extract_core", true)
    extract_temporal = get(args, "extract_temporal", true)
    extract_well_behaved = get(args, "extract_well_behaved", true)
    extract_neighbor_features = get(args, "extract_neighbor_features", false)
    extract_carlidar_rangerate = get(args, "extract_carlidar_rangerate", true)
    carlidar_nbeams = get(args, "carlidar_nbeams", 0)
    roadlidar_nbeams = get(args, "roadlidar_nbeams", 0)
    roadlidar_nlanes = get(args, "roadlidar_nlanes", 2)
    carlidar_max_range = get(args, "carlidar_max_range", 100.0) # [m]
    roadlidar_max_range = get(args, "roadlidar_max_range", 100.0) # [m]

    extractor = MultiFeatureExtractor(
        extract_core,
        extract_temporal,
        extract_well_behaved,
        extract_neighbor_features,
        extract_carlidar_rangerate,
        carlidar_nbeams,
        roadlidar_nbeams,
        roadlidar_nlanes,
        carlidar_max_range=carlidar_max_range,
        roadlidar_max_range=roadlidar_max_range,
        )

    ego_action_type = AccelTurnrate
    T = get(args, "action_type", "AccelTurnrate")
    if T == "AccelTurnrate"
        ego_action_type = AccelTurnrate
    elseif T == "LatLonAccel"
        ego_action_type = LatLonAccel
    end

    if haskey(args, "trajdata_filepaths")
        trajdata_filepaths = args["trajdata_filepaths"]
        roadway_filepaths = args["roadway_filepaths"]
        gen_simparams_from_trajdatas(trajdata_filepaths, roadway_filepaths,
            col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
            use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
            batch_size, prime_history, nsteps, ego_action_type, extractor)
    else
        trajdata_indeces = get(args, "trajdata_indeces", [1,2,3,4,5,6])
        gen_simparams(trajdata_indeces,
            col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
            use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
            batch_size, prime_history, nsteps, ego_action_type, extractor)
    end
end

###########################################
AutomotiveDrivingModels.get_name(::GaussianMLPDriver) = "GaussianMLPDriver"
AutomotiveDrivingModels.action_context(model::GaussianMLPDriver) = model.context

function AutomotiveDrivingModels.reset_hidden_state!(model::GaussianMLPDriver)
    empty!(model.rec)
    model
end

function AutomotiveDrivingModels.observe!{A,F,G,E,P}(
                                            model::GaussianMLPDriver{A,F,G,E,P},
                                            simparams::SimParams,
                                            scene::Scene,
                                            roadway::Roadway,
                                            egoid::Int)

    update!(model.rec, scene)
    vehicle_index = get_index_of_first_vehicle_with_id(scene, egoid)
    o = pull_features!(simparams.extractor, simparams.features, model.rec, roadway, vehicle_index)
    model.net[:hidden_0].input = (o - model.extractor.feature_means)./model.extractor.feature_std
    forward!(model.pass)
    copy!(model.mvnormal.μ, model.output[1:2])

    model
end
Base.rand{A,F,G,E,P}(model::GaussianMLPDriver{A,F,G,E,P}) = convert(A, rand(model.mvnormal))
Distributions.pdf{A,F,G,E,P}(model::GaussianMLPDriver{A,F,G,E,P}, a::A) = pdf(model.mvnormal, convert(Vector{Float64}, a))
Distributions.logpdf{A,F,G,E,P}(model::GaussianMLPDriver{A,F,G,E,P}, a::A) = logpdf(model.mvnormal, convert(Vector{Float64}, a))
############################################

Base.show(io::IO, simparams::SimParams) = print(io, "SimParams")

function restart!(simstate::SimState, simparams::SimParams)

    # pick a random segment
    local train_seg_index
    for i in 1 : 100
        train_seg_index = rand(1:length(simparams.segments))
        seg = simparams.segments[train_seg_index]
        candidate_frame_lo = seg.frame_lo
        candidate_frame_hi = seg.frame_hi - simparams.nsteps - simparams.prime_history
        if candidate_frame_hi > candidate_frame_lo
        break
    end
    if i > 95
        assert(false)
        end
    end

    seg = simparams.segments[train_seg_index]
    simstate.egoid = seg.egoid
    simstate.trajdata_index = seg.trajdata_index

    # pull the trajdata
    trajdata = simparams.trajdatas[simstate.trajdata_index]

    # pick a random sub-trajectory
    candidate_frame_lo = seg.frame_lo
    candidate_frame_hi = seg.frame_hi - simparams.nsteps - simparams.prime_history
    simstate.frame = rand(candidate_frame_lo:candidate_frame_hi)

    # clear rec and make first observations
    empty!(simstate.rec)
    for i in 1 : simparams.prime_history + 1
        get!(simstate.scene, trajdata, simstate.frame)
        update!(simstate.rec, simstate.scene)
        simstate.frame += 1
    end
    simstate.frame -= 1
    simstate.start_frame = simstate.frame

    # empty playback reactive
    empty!(simstate.playback_reactive_active_vehicle_ids)

    # simparams.extractor.road_lidar_culling = simparams.roadway_cullers[simstate.trajdata_index]

    simstate
end

function Base.reset(simparams::SimParams)
    for state in simparams.simstates
        restart!(state, simparams)
    end
    simparams.step_counter = 0

    simparams
end

function step_forward!(simstate::SimState, simparams::SimParams, action_ego::Vector{Float64})

    trajdata = simparams.trajdatas[simstate.trajdata_index]
    veh_ego = get_vehicle(simstate.scene, simstate.egoid)

    safety_policy = simparams.safety_model

    act_dim = length(action_ego)
    if act_dim == 3
        if action_ego[end] > 0
            #set_desired_speed!(safety_policy, simstate.playback_reactive_speeds[simstate.egoid])
            AutomotiveDrivingModels.observe!(safety_policy,simstate.scene,trajdata.roadway,simstate.egoid)
            action_ego = rand(safety_policy)::LatLonAccel # is this only returning acceleration?
            #action_ego = convert(simparams.ego_action_type, action_ego)

        else
            action_ego = action_ego[1:end-1]
            action_ego = convert(simparams.ego_action_type, action_ego)

        end
    else
        action_ego = convert(simparams.ego_action_type, action_ego)
    end

    # propagate the ego vehicle
    trajdata = simparams.trajdatas[simstate.trajdata_index]
    veh_ego = get_vehicle(simstate.scene, simstate.egoid)
    #action_ego = convert(simparams.ego_action_type, action_ego)

    ego_state = propagate(veh_ego, action_ego, simparams.context, trajdata.roadway)

    simstate.frame += 1
    if simparams.model_all
        model = simparams.driver_model

        empty!(simparams.playback_reactive_scene_buffer)

        tdframe = trajdata.frames[simstate.frame]
        for i in tdframe.lo : tdframe.hi
            s = trajdata.states[i]
            if s.id != simstate.egoid

                veh_index = get_index_of_first_vehicle_with_id(simstate.scene, s.id)
                if veh_index != 0 # iscarinframe
                    if Symbol("gru") in fieldnames(model.net)
                        if !(s.id in simstate.hidden_dict)
                            model.net[:gru].h_prev = zeros(length(model.net[:gru].h_prev))
                        else
                            model.net[:gru].h_prev = simstate.hidden_dict[s.id]
                        end
                    end

                    AutomotiveDrivingModels.observe!(model, simparams, simstate.scene, trajdata.roadway, s.id)
                    action = rand(model)::AccelTurnrate
                    a = clamp(action.a, -5.0, 3.0)
                    ω = clamp(action.ω, -0.01, 0.01)

                    # Propagate scene
                    veh = simstate.scene[veh_index]
                    veh_state = propagate(veh, AccelTurnrate(a, ω), simparams.context, trajdata.roadway)
                    veh2 = Vehicle(veh_state, veh.def)
                    push!(simparams.playback_reactive_scene_buffer, veh2)

                    if Symbol("gru") in fieldnames(model.net)
                        simstate.hidden_dict[s.id] = model.net[:gru].h
                    end
                end
            end
        end

        # move simparams.playback_reactive_scene_buffer over to simparams.scene
        # and place new ego pos
        ego_def = veh_ego.def
        copy!(simstate.scene, simparams.playback_reactive_scene_buffer)
        push!(simstate.scene, Vehicle(ego_state, ego_def))
    elseif simparams.use_playback_reactive

        model = simparams.playback_reactive_model

        empty!(simparams.playback_reactive_scene_buffer)

        tdframe = trajdata.frames[simstate.frame]
        for i in tdframe.lo : tdframe.hi
            s = trajdata.states[i]
            if s.id != simstate.egoid

                use_playback = true
                veh_index = get_index_of_first_vehicle_with_id(simstate.scene, s.id)
                ego_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
                if veh_index != 0 # iscarinframe

                    if !in(s.id, simstate.playback_reactive_active_vehicle_ids)
                        simstate.playback_reactive_speeds[s.id] = simstate.scene[veh_index].state.v
                    end

                    set_desired_speed!(model, simstate.playback_reactive_speeds[s.id])
                    AutomotiveDrivingModels.observe!(model, simstate.scene, trajdata.roadway, s.id)
                    action = rand(model)::LatLonAccel

                    if s.id in simstate.playback_reactive_active_vehicle_ids ||
                       (action.a_lon < simparams.playback_reactive_threshold_brake &&
                        simstate.scene[veh_index].state.posF.s < simstate.scene[ego_index].state.posF.s)

                        # use IDM
                        use_playback = false
                        push!(simstate.playback_reactive_active_vehicle_ids, s.id)

                        veh = simstate.scene[veh_index]
                        if veh.state.v < 0.0
                            action = LatLonAccel(action.a_lat, max(action.a_lon, 0.0))
                        end
                        veh_state = propagate(veh, action, simparams.context, trajdata.roadway)
                        veh2 = Vehicle(veh_state, veh.def)
                        push!(simparams.playback_reactive_scene_buffer, veh2)
                    end
                end

                if use_playback
                    # if it is a new car not in previous frame or it doesn't need IDM
                    veh = Vehicle(s.state, get_vehicledef(trajdata, s.id))
                    push!(simparams.playback_reactive_scene_buffer, veh)
                end
            end
        end

        # move simparams.playback_reactive_scene_buffer over to simparams.scene
        # and place new ego pos
        ego_def = veh_ego.def
        copy!(simstate.scene, simparams.playback_reactive_scene_buffer)
        push!(simstate.scene, Vehicle(ego_state, ego_def))
    else
        # pull new frame from trajdata
        get!(simstate.scene, trajdata, simstate.frame)

        # move in propagated ego vehicle
        veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
        simstate.scene[veh_index].state = ego_state
    end

    # update record
    update!(simstate.rec, simstate.scene)

    simstate
end

function tick(simparams::SimParams, u::Vector{Float64}, batch_index::Int=1)
    step_forward!(simparams.simstates[batch_index], simparams, u)
    simparams
end
function step_forward(simparams::SimParams, U::Matrix{Float64})
    # note: actions are batch_size × action_size
    for (i,s) in enumerate(simparams.simstates)
        step_forward!(s, simparams, U[i,:])
    end
    simparams
end

function reward(simstate::SimState, simparams::SimParams)

    veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
    veh_ego = simstate.scene[veh_index]

    reward = 0.0
    if !simparams.use_debug_reward
        trajdata = simparams.trajdatas[simstate.trajdata_index]
        acc_ego = convert(Float64, get(ACC, simstate.rec, trajdata.roadway, veh_index))
        jrk_ego = convert(Float64, get(JERK, simstate.rec, trajdata.roadway, veh_index))
        ome_ego = convert(Float64, get(ANGULARRATEG, simstate.rec, trajdata.roadway, veh_index))
        cen_ego = veh_ego.state.posF.t

        d_ml = convert(Float64, get(MARKERDIST_LEFT, simstate.rec, trajdata.roadway, veh_index))
        d_mr = convert(Float64, get(MARKERDIST_RIGHT, simstate.rec, trajdata.roadway, veh_index))

        reward += simparams.col_weight * get_first_collision(simstate.scene, veh_index).is_colliding
        reward += simparams.off_weight * (d_ml < -1.0 || d_mr < -1.0)
        reward += simparams.rev_weight * (veh_ego.state.v < 0.0)
        reward += simparams.jrk_weight * jrk_ego*jrk_ego
        reward += simparams.acc_weight * acc_ego*acc_ego
        reward += simparams.cen_weight * cen_ego*cen_ego
        reward += simparams.ome_weight * ome_ego*ome_ego
    else
        reward -= (veh_ego.state.v^2)
        reward -= ((veh_ego.state.posG.θ - π/2)^2)
    end

    reward
end
function reward(simparams::SimParams, u::Vector{Float64}, batch_index::Int = 1)
    reward(simparams.simstates[batch_index], simparams)
end

function observe(simparams::SimParams, batch_index::Int=1)
    simstate = simparams.simstates[batch_index]
    trajdata = simparams.trajdatas[simstate.trajdata_index]
    veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
    pull_features!(simparams.extractor, simparams.features, simstate.rec, trajdata.roadway, veh_index)
end

isdone(simparams::SimParams) = simparams.step_counter ≥ simparams.nsteps

function Base.step(simparams::SimParams, u::Vector{Float64}, batch_index::Int=1)

    r = reward(simparams, u, batch_index)
    tick(simparams, u, batch_index)
    features = observe(simparams, batch_index)
    simparams.step_counter += 1
    # done = isdone(simparams)
    simstate = simparams.simstates[batch_index]

    # End if collision or reverse or off-road
    veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
    trajdata = simparams.trajdatas[simstate.trajdata_index]
    d_ml = convert(Float64, get(MARKERDIST_LEFT, simstate.rec, trajdata.roadway, veh_index))
    d_mr = convert(Float64, get(MARKERDIST_RIGHT, simstate.rec, trajdata.roadway, veh_index))
    done = false
    done = done || get_first_collision(simstate.scene, veh_index).is_colliding
    done = done || (d_ml < -1.0 || d_mr < -1.0)
    done = done || (simstate.scene[veh_index].state.v < 0.0)

    (features, r, done)
end
function Base.step(simparams::SimParams, U::Matrix{Float64})

    batch_size = length(simparams.states)
    feature_mat = Array(Float64, batch_size, obssize(simparams))
    rewards = Array(Float64, batch_size)
    dones = Array(Float64, batch_size)

    step_counter = simparams.step_counter
    for batch_index in 1 : batch_size
        simparams.step_counter = step_counter
        features, reward, done = step(simparams, U[batch_index, :], batch_index)
        feature_mat[batch_index, :] = features
        rewards[batch_index] = reward
        dones[batch_index] = done
    end
    simparams.step_counter = step_counter + 1

    (feature_mat, rewards, dones)
end

action_space_bounds(simparams::SimParams) = ([-5.0, -1.0], [3.0, 1.0])
observation_space_bounds(simparams::SimParams) = (fill(-Inf, length(simparams.extractor)), fill(Inf, length(simparams.extractor)))

##################################
# Visualization

function AutoViz.render(simparams::SimParams, batch_index::Int=1)

    simstate = simparams.simstates[batch_index]
    trajdata = simparams.trajdatas[simstate.trajdata_index]

    carcolors = Dict{Int,Colorant}()
    carcolors[simstate.egoid] = COLOR_CAR_EGO

    for id in simstate.playback_reactive_active_vehicle_ids
        carcolors[id] = colorant"0xF9971FFF" # orange
    end

    render(simstate.scene, trajdata.roadway,
          cam=CarFollowCamera(simstate.egoid, 5.0), car_colors=carcolors)
end

function AutoViz.render(simparams::SimParams, image::Matrix{UInt32}, batch_index::Int=1)

    simstate = simparams.simstates[batch_index]
    trajdata = simparams.trajdatas[simstate.trajdata_index]

    c = CairoImageSurface(image, Cairo.FORMAT_ARGB32, flipxy=false)
    ctx = CairoContext(c)

    carcolors = Dict{Int,Colorant}()
    carcolors[simstate.egoid] = COLOR_CAR_EGO

    for id in simstate.playback_reactive_active_vehicle_ids
        carcolors[id] = colorant"0xF9971FFF" # orange
    end

    #render(ctx, simstate.scene, trajdata.roadway,
    #      cam=CarFollowCamera(simstate.egoid, 5.0), car_colors=carcolors)

    render(ctx, simstate.scene, trajdata.roadway, cam=CarFollowCamera(simstate.egoid, 5.0))

    image
end

function reel_drive(
    gif_filename::AbstractString,
    actions::Matrix{Float64}, # columns are actions
    simparams::SimParams
    )

    frames = Reel.Frames(MIME("image/png"), fps=framerate)

    action = [NaN,NaN]
    push!(frames, render(simparams))
    for frame_index in 1:size(actions,2)

        action[1] = actions[frame_index,1]
        action[2] = actions[frame_index,2]

        step_forward(simparams, action)

        push!(frames, render(simparams))
    end

    Reel.write(gif_filename, frames) # Write to a gif file
end

end # module


text:
# This files checks the new way to draw a circle versus the old way

using SimpleDirectMediaLayer
const SDL2 = SimpleDirectMediaLayer

circles = Vector{Int}[]
for j = 1:10
  push!(circles, [10 + 4 * j; 50; j])
  push!(circles, [100 + 30 * j; 500; 25 * j])
end

function newdraw(ren, circle; fill=false)
  r = Cint(circle[3])
  o = Cint.(circle[1:2])

  n = ceil(π * r / 2)
  if fill
    for j = 0:n
      x = round(Cint, r * cos(j / n * π / 2))
      y = round(Cint, r * sin(j / n * π / 2))
      SDL2.RenderDrawLine(ren, o[1] + x, o[2] + y, o[1] + x, o[2] - y)
      SDL2.RenderDrawLine(ren, o[1] - x, o[2] + y, o[1] - x, o[2] - y)
    end
  else
    for j = 0:n
      x = round(Cint, r * cos(j / n * π / 2))
      y = round(Cint, r * sin(j / n * π / 2))
      SDL2.RenderDrawPoint(ren, o[1] + x, o[2] + y)
      SDL2.RenderDrawPoint(ren, o[1] - x, o[2] + y)
      SDL2.RenderDrawPoint(ren, o[1] + x, o[2] - y)
      SDL2.RenderDrawPoint(ren, o[1] - x, o[2] - y)
    end
  end
end

function olddraw(ren, circle; fill=false)
    # define the center and needed sides of circle
    centerX = Cint(circle[1])
    centerY = Cint(circle[2])
    int_rad = Cint(circle[3])
    left = centerX - int_rad
    top = centerY - int_rad

    # we consider a grid with sides equal to the circle's diameter
    for x in left:centerX
        for y in top:centerY

            # for each pixel in the top left quadrant of the grid we measure the distance from the center.
            dist = sqrt( (centerX - x)^2 + (centerY - y)^2 )

            # if it is close to the circle's radius it and all associated points in the other quadrants are colored in.
            if (dist <= int_rad + 0.5 && dist >= int_rad - 0.5)
                rel_x = centerX - x
                rel_y = centerY - y

                quad1 = (x              , y              )
                quad2 = (centerX + rel_x, y              )
                quad3 = (x              , centerY + rel_y)
                quad4 = (quad2[1]       , quad3[2]       )

                SDL2.RenderDrawPoint(ren, quad1[1], quad1[2])
                SDL2.RenderDrawPoint(ren, quad2[1], quad2[2])
                SDL2.RenderDrawPoint(ren, quad3[1], quad3[2])
                SDL2.RenderDrawPoint(ren, quad4[1], quad4[2])

                # if we are told to fill in the circle we draw lines between all of the quadrants to completely fill the circle
                if (fill == true)
                    SDL2.RenderDrawLine(ren, quad1[1], quad1[2], quad2[1], quad2[2])
                    SDL2.RenderDrawLine(ren, quad2[1], quad2[2], quad4[1], quad4[2])
                    SDL2.RenderDrawLine(ren, quad4[1], quad4[2], quad3[1], quad3[2])
                    SDL2.RenderDrawLine(ren, quad3[1], quad3[2], quad1[1], quad1[2])
                end
            end

        end
    end

end

function myinit()
  SDL2.init()
  win = SDL2.CreateWindow("",
      Int32(0), Int32(0), Int32(800), Int32(600), UInt32(SDL2.WINDOW_SHOWN)
     )

  renderer = SDL2.CreateRenderer(win,
      Int32(-1),
      UInt32(SDL2.RENDERER_ACCELERATED | SDL2.RENDERER_PRESENTVSYNC)
     )

  SDL2.SetRenderDrawColor(renderer, 0, 0, 0, 255)
  return renderer
end

function testold(renderer, circles, N, boolfill)
  SDL2.SetRenderDrawColor(renderer, 255, 255, 255, 255)
  for j = 1:N
    SDL2.RenderClear(renderer)
    for k = 1:length(circles)
      olddraw(renderer, circles[k], fill=boolfill)
    end
  end
  SDL2.Quit()
end

function testnew(renderer, circles, N, boolfill)
  SDL2.SetRenderDrawColor(renderer, 255, 255, 255, 255)
  for j = 1:N
    SDL2.RenderClear(renderer)
    for k = 1:length(circles)
      newdraw(renderer, circles[k], fill=boolfill)
    end
  end
  SDL2.Quit()
end


text:
using Plots
using Reproduce
using JLD2
using FileIO
using Statistics

include("plot_utils.jl")

function get_lines(data)
    means = []
    step = 10
    for start in 1:step:size(data)[2]
        chunk = data[:, start:min(start + step - 1, length(data)),: ]
        # println(size(chunk))
        chunk_mean = mean(chunk, dims=2)
        push!(means, chunk_mean)
    end

    smoothed = cat(means..., dims=2)[:,:, :]
    sum_lines = sum(smoothed, dims=1)
    num_runs = size(sum_lines)[3]
    mean_line = dropdims(mean(sum_lines, dims=3), dims=(1, 3))
    std_err_line = dropdims(std(sum_lines, dims=3), dims=(1, 3)) / sqrt(num_runs)

    xs = 1:100*step:30000
    return xs, mean_line, std_err_line
end

# There's duplicate copy of me in get_best, please replace me later.
function get_performance(ic, metric, percentage=(0, 1))
    res = load_results(ic, metric)

    if length(size(res)) == 3
        num_log = size(res)[2]
        start_index = floor(Int, num_log * percentage[1])
        if (start_index == 0)
            start_index = 1
        end
        end_index = ceil(Int, num_log * percentage[2])
        res = res[:, start_index:end_index, :]
    end

    means = dropdims(mean(sum(res, dims=1), dims=2), dims=(1, 2))
    std_err = std(means) / sqrt(length(means))
    return mean(means), std_err
end

num_steps = 30000
log_interval = 100

sweep_params = ["demon_eta", "exploration_param", "behaviour_eta"]

sens_key = "behaviour_eta"
sens_sweep_params = filter(x -> x != sens_key, sweep_params)
p = plot()

ic = ItemCollection("M:/globus/Experiment2/Experiment2_GPI/")

sens_values = diff(ic)[sens_key]
println(sens_values)

perfs = []
stds = []
for sens_val in sens_values
    sens_ic = search(ic, Dict(sens_key => sens_val))
    best_sens_ic = get_best(sens_ic, sens_sweep_params, :ttmaze_uniform_error)
    perf, std = get_performance(best_sens_ic, :ttmaze_uniform_error)
    push!(perfs, perf)
    push!(stds, std)
end

println(sens_values)
println(perfs)
println(stds)

# demon_eta
# perfs = [0.36069585419440175, 0.340265883367322, 0.32440065133391177, 0.43664044947316205]
# stds = [0.009840086433933035, 0.012308683479898372, 0.00966571250650264, 0.029534894799327426]

plot!(p, sens_values, perfs, yerror=stds, label="GPI", xlabel="behaviour step size")

ic = ItemCollection("M:/globus/Experiment2/Experiment2_ESarsa_NoLimit")
sens_values = diff(ic)[sens_key]
println(sens_values)
perfs = []
stds = []
for sens_val in sens_values
    sens_ic = search(ic, Dict(sens_key => sens_val))
    best_sens_ic = get_best(sens_ic, sens_sweep_params, :ttmaze_uniform_error)
    perf, std = get_performance(best_sens_ic, :ttmaze_uniform_error)
    push!(perfs, perf)
    push!(stds, std)
end

# println(sens_values)
# println(perfs)
# println(stds)

# demon eta
# perfs = [0.5190128725436324, 0.4815899839752351, 0.436221693493747, 0.42818578520772327, 0.40281394394567965, 0.38377415931803205, 0.7014830158663732]
# std = [0.01063201541960759, 0.010213129582727696, 0.009209626454773388, 0.009358865668267153, 0.007108921837893297, 0.009171313780882599, 0.23426893167605065]

plot!(p, sens_values, perfs, yerror=stds, label="ESARSA")

savefig("plotting/chunlok/generated_plots/sens_plot_reproduce_$(sens_key).svg")

asdasd

# best_ic = get_best(ic, sweep_params, :ttmaze_uniform_error)
# print_params(best_ic, sweep_params, [])

# best_ic = get_best(ic, sweep_params, :ttmaze_uniform_error, (0.9, 1))
# print_params(best_ic, sweep_params, [])

# best_ic = search(ic, Dict("demon_eta" => 0.5, "exploration_param" => 0.4, "behaviour_eta" => 0.5))
# best_ic = search(ic, Dict("demon_eta" => 0.5, "exploration_param" => 0.1, "behaviour_eta" => 0.0625))

print_params(best_ic, sweep_params, [])

data = load_results(best_ic, :ttmaze_uniform_error)

xs, mean_line, std_err_line = get_lines(data)
plot!(p, xs, mean_line, ribbons=std_err_line, label="GPI")

ic = ItemCollection("M:/globus/Experiment2/Experiment2_ESarsa_NoLimit")

println(diff(ic))

# best_ic = get_best(ic, sweep_params, :ttmaze_uniform_error)
# print_params(best_ic, sweep_params, [])
# best_ic = get_best(ic, sweep_params, :ttmaze_uniform_error, (0.9, 1))
# print_params(best_ic, sweep_params, [])

#AUC
# best_ic = search(ic, Dict("demon_eta" => 0.5, "exploration_param" => 0.4, "behaviour_eta" => 0.25))
#AUC last 10%
best_ic = search(ic, Dict("demon_eta" => 0.5, "exploration_param" => 0.3, "behaviour_eta" => 0.03125))

# Old Params
# best_ic = search(ic, Dict("demon_eta" => 0.5, "exploration_param" => 0.1, "behaviour_eta" => 0.0625))

data = load_results(best_ic, :ttmaze_uniform_error)

println(size(data))

xs, mean_line, std_err_line = get_lines(data)
plot!(p, xs, mean_line, ribbons=std_err_line, label="ESarsa")

savefig("plotting/chunlok/generated_plots/test_plot_reproduce.svg")

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
### A Pluto.jl notebook ###
# v0.11.9

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ abc33bd6-ed31-11ea-33a4-313c8dd55164
using PlutoUI

# ╔═╡ ca921050-ed31-11ea-38f9-359d20b1b2ed
include("model.jl");

# ╔═╡ 9f27c4ee-ed36-11ea-0230-3f14526f9423
md"""
Node: $(@bind node Slider(range(1, 700, length = 1)))

k: $(@bind k Slider(1:3))

Samples: $(@bind samples Slider(range(100, 20000, length = 20)))

p₀: $(@bind p₀ Slider(range(0, 0.2, length = 21)))
"""

# ╔═╡ 1bac6944-ed3d-11ea-1a8f-0f0fb703e5ed
@show (node ,k, samples, p₀)

# ╔═╡ 438a5144-ed36-11ea-29df-7fa4eefea733
model = PyGInferenceModel("/gpfs_home/spate116/singhlab/GCN_Integration/scripts/BI/syn", "syn4", Int(node), k, p₀, 0.00000000000001);

# ╔═╡ f0eda5bc-ed37-11ea-3cc1-51dfea0858ad
@time s = sample(model, Int(samples));

# ╔═╡ 00e35172-ed38-11ea-0e7e-e5c07248c292
final_summary(s)

# ╔═╡ 0689b076-ed38-11ea-2f93-d70610e9b001
plt = plot_result(model, s)

# ╔═╡ 8133fb54-ed47-11ea-26c1-8b15400a881b
begin
	using Plots
	plot(plt, size=(1000, 1000), dpi=300)
	savefig("out.png")
end

# ╔═╡ Cell order:
# ╠═abc33bd6-ed31-11ea-33a4-313c8dd55164
# ╠═ca921050-ed31-11ea-38f9-359d20b1b2ed
# ╠═9f27c4ee-ed36-11ea-0230-3f14526f9423
# ╟─1bac6944-ed3d-11ea-1a8f-0f0fb703e5ed
# ╠═438a5144-ed36-11ea-29df-7fa4eefea733
# ╠═f0eda5bc-ed37-11ea-3cc1-51dfea0858ad
# ╠═00e35172-ed38-11ea-0e7e-e5c07248c292
# ╠═0689b076-ed38-11ea-2f93-d70610e9b001
# ╠═8133fb54-ed47-11ea-26c1-8b15400a881b


text:
"""
    SixDOFBodyEuler
    SixDOFBodyEuler(x::AbstractVector)
    SixDOFBodyEuler(dss::DSState)

Six degrees of freedom dynamic system using u, v, w for velocity representation and Euler
angles for attitude.

Flat Earth hypothesis is applied and Earth reference frame is considered inertial.

It is considered that the aircraft xb-zb plane is a plane of symmetry so that Jxy and Jyz
cross-product of inertia are zero and will not be taken into account.

# Fields
- `x::SVector{13, T}`: state vector.
  - [u (m/s), v (m/s), w (m/s), ϕ (rad), θ (rad), ψ (rad), p (rad/s), q (rad/s), r (rad/s),
    x (m), y (m), z (m), pow (%)].
"""
struct SixDOFBodyEuler{T}<:DSState{T}
    x::SVector{13, T}
end

SixDOFBodyEuler(x::AbstractVector) = SixDOFBodyEuler(SVector{13, eltype(x)}(x))

SixDOFBodyEuler(dss::DSState) = SixDOFBodyEuler(
    [
        get_body_velocity(dss)...,
        get_euler_angles(dss)[3:-1:1]...,
        get_ang_vel_body(dss)...,
        get_earth_position(dss)...,
        get_engine_power(dss),
    ]
)

get_x_names(dss::SixDOFBodyEuler) = [:u, :v, :w, :ϕ, :θ, :ψ, :p, :q, :r, :x, :y, :z, :pow]

# Mandatory getters for DSState
get_earth_position(dss::SixDOFBodyEuler) = get_x(dss)[10:12]
get_height(dss::SixDOFBodyEuler) = -get_x(dss)[12]

get_euler_angles(dss::SixDOFBodyEuler) = get_x(dss)[6:-1:4]

get_tas(dss::SixDOFBodyEuler) = get_tasαβ(dss)[1]
get_α(dss::SixDOFBodyEuler) = get_tasαβ(dss)[2]
get_β(dss::SixDOFBodyEuler) = get_tasαβ(dss)[3]
get_tasαβ(dss::SixDOFBodyEuler) = uvw_to_tasαβ(get_body_velocity(dss)...)
get_body_velocity(dss::SixDOFBodyEuler) = get_x(dss)[1:3]
get_horizon_velocity(dss::SixDOFBodyEuler) = body2horizon(
    get_body_velocity(dss)..., get_euler_angles(dss)...
)

get_ang_vel_body(dss::SixDOFBodyEuler) = get_x(dss)[7:9]
get_euler_angles_rates(dss::SixDOFBodyEuler) = pqr_2_ψθϕ_dot(
    get_ang_vel_body(dss)..., get_euler_angles(dss)[2:3]...
)

get_engine_power(dss::SixDOFBodyEuler) = get_x(dss)[13]

# Mandatory getters for DSStateDot
get_tas_dot(dssd::DSStateDot{S, N, T}) where {S<:SixDOFBodyEuler, N, T} = get_tasαβ_dot(dssd)[1]
get_α_dot(dssd::DSStateDot{S, N, T}) where {S<:SixDOFBodyEuler, N, T} = get_tasαβ_dot(dssd)[2]
get_β_dot(dssd::DSStateDot{S, N, T}) where {S<:SixDOFBodyEuler, N, T} = get_tasαβ_dot(dssd)[3]
get_tasαβ_dot(dssd::DSStateDot{S, N, T}) where {S<:SixDOFBodyEuler, N, T} = uvw_dot_to_tasαβ_dot(
    get_x(dssd)[1:3]..., get_xdot(dssd)[1:3]...
)
get_accel_body(dssd::DSStateDot{S, N, T}) where {S<:SixDOFBodyEuler, N, T} = get_xdot(dssd)[1:3]
get_ang_accel_body(dssd::DSStateDot{S, N, T}) where {S<:SixDOFBodyEuler, N, T} = get_xdot(dssd)[7:9]
get_engine_power_dot(dssd::DSStateDot{S, N, T}) where {S<:SixDOFBodyEuler, N, T} = get_xdot(dssd)[13]

function state_eqs(dss::SixDOFBodyEuler, time, mass, inertia, forces, moments, h, pow_dot)

    xdot = sixdof_body_earth_euler_fixed_mass(
         time,
         get_x(dss),
         mass,
         inertia,
         forces,
         moments,
         h
    )

    xdot = [xdot..., pow_dot]

    return DSStateDot(dss, xdot)
end

"""
    sixdof_body_earth_euler_fixed_mass!(x_dot, time, x, mass, inertia, forces, moments, h)

Mutating version of [sixdof_body_earth_euler_fixed_mass](@ref).
"""
function sixdof_body_earth_euler_fixed_mass!(
    x_dot, time, x, mass, inertia, forces, moments, h
    )

    m = mass
    u, v, w, ϕ, θ, ψ, p, q, r, xe, ye, ze = x

    # Unpack forces
    Fx, Fy, Fz = forces
    # Unpack moments
    L, M, N = moments
    # Unpack angular momentum contributions
    hx, hy, hz = h
    # Unpack inertia
    Ixx, Iyy, Izz = inertia[1, 1], inertia[2, 2], inertia[3, 3]
    Ixz = inertia[1, 3]

    sψ, cψ = sin(ψ), cos(ψ)
    sθ, cθ = sin(θ), cos(θ)
    sϕ, cϕ = sin(ϕ), cos(ϕ)

    # Linear momentum equations
    x_dot[1] = Fx / m + r * v - q * w
    x_dot[2] = Fy / m - r * u + p * w
    x_dot[3] = Fz / m + q * u - p * v

    # Moments
    pq = p * q
    qr = q * r

    rhy_qhz = (r * hy - q * hz)
    qhx_phy = (q * hx - p * hy)

    # If inertia is constant this terms are constant too.
    # TODO: think about passing them as arguments to improve speed.
    IxzS = Ixz^2
    xpq = Ixz * (Ixx - Iyy + Izz)
    gam = Ixx * Izz - IxzS
    xqr = Izz * (Izz - Iyy) + IxzS
    zpq = (Ixx - Iyy) * Ixx + IxzS
    ypr = Izz - Ixx

    # Angular momentum equations
    x_dot[7] = (xpq * pq - xqr * qr + Izz * (L + rhy_qhz) + Ixz * (N + qhx_phy)) / gam
    x_dot[8] = (ypr * p * r - Ixz * (p^2 - r^2) + M - r * hx + p * hz) / Iyy
    x_dot[9] = (zpq * pq - xpq * qr + Ixz * (L + rhy_qhz) + Ixx * (N + qhx_phy)) / gam

    # Angular Kinematic equations
    x_dot[6] = (q * sϕ + r * cϕ) / cθ
    x_dot[5] = q * cϕ - r * sϕ
    # ϕ_dot = p + (q * sϕ + r * cϕ) * tan(θ)
    x_dot[4] = p + x_dot[6] * sθ

    # Linear kinematic equations
    x_dot[10] =  cθ*cψ * u + (sϕ*sθ*cψ - cϕ*sψ) * v + (cϕ*sθ*cψ + sϕ*sψ) * w
    x_dot[11] =  cθ*sψ * u + (sϕ*sθ*sψ + cϕ*cψ) * v + (cϕ*sθ*sψ - sϕ*cψ) * w
    x_dot[12] = -sθ    * u +  sϕ*cθ             * v +  cϕ*cθ             * w

end

"""
    sixdof_body_earth_euler_fixed_mass(time, x, mass, inertia, forces, moments, h)

State equations for `SixDOFBodyEuler` dynamic system.

# Arguments

- `time::Number`: time (s).
- `x::AbstractArray`: state vector.
  - [u (m/s), v (m/s), w (m/s), ϕ (rad), θ (rad), ψ (rad), p (rad/s), q (rad/s), r (rad/s),
    x (m), y (m), z (m), pow (%)].
- `mass::Number`: mass (kg).
- `inertia::AbstractMatrix`: 3x3 inertia tensor (kg·m^2).
- `forces::AbstractVector`: body axis forces [Fx, Fy, Fz] (N).
- `moments::AbstractVector`: body axis moments (L, M, N) (N·m).
- `h::AbstractVector`:  Additional angular momentum contributions such as those coming from
 spinning rotors (kg·m²/s).

# Returns
- `x_dot::Array`: time derivative of the state vector.

# Notes

It is considered that the aircraft xb-zb plane is a plane of symmetry so that Jxy and Jyz
cross-product of inertia are zero and will not be taken into account.

# See also

[six_dof_body_euler_fixed_mass!](@ref)
"""
function sixdof_body_earth_euler_fixed_mass(time, x, mass, inertia, forces, moments, h)
    x_dot = Array{eltype(x)}(undef, 12)
    sixdof_body_earth_euler_fixed_mass!(x_dot, time, x, mass, inertia, forces, moments, h)
    return x_dot
end


text:
struct Int32Bool
    bits::Int32
end

Int32Bool(i::Int32, b::Bool) = Int32Bool(Int32(2b - 1) * abs(i))

int(x::Int32Bool) = abs(x.bits)
bool(x::Int32Bool) = x.bits > 0

setbool(x::Int32Bool, b::Bool) = Int32Bool(int(x), b)
setint(x::Int32Bool, i::Int32) = Int32Bool(i, bool(x))

struct CRQSlot{S}
    index_safe::Int32Bool
    storage::S  # sizeof(S) ≤ 4
end

@inline CRQSlot(; index::Integer, safe::Bool, storage) =
    CRQSlot(Int32Bool(Int32(index), safe), storage)

@inline function Base.getproperty(slot::CRQSlot, name::Symbol)
    if name === :index
        return int(getfield(slot, :index_safe))
    elseif name === :safe
        return bool(getfield(slot, :index_safe))
    end
    return getfield(slot, name)
end

"""
    IndirectConcurrentRingQueueNode{T}

Concurrent Ring Queue (CRQ) node which stores 64 bit slot in the `.ring`.  Each
slot contains a 32 bit value that locates the actual value in the thread-local
`.buffers`; hence _indirect_.

This approach wastes memory but it is (more or less) implementable without
atomics support of boxed element in array.

TODO: Create a 128 bit slot "direct" version that directly points to a boxed
Julia value.
"""
mutable struct IndirectConcurrentRingQueueNode{T}
    @atomic head::Int32
    _head_pad::PadAfter32
    @atomic tail_closed::Int32
    _tail_pad::PadAfter32
    @atomic next::Union{Nothing,IndirectConcurrentRingQueueNode{T}}
    ring::Vector{CRQSlot{UInt32}}
    length::Int
    buffers::Vector{Vector{T}}
end
# TODO: pad

function IndirectConcurrentRingQueueNode{T}(len::Int) where {T}
    buffers = [Vector{T}(undef, len) for _ in 1:Threads.nthreads()]
    ring = [CRQSlot(Int32Bool(Int32(i), true), UInt32(0)) for i in 1:len]
    head = Int32(1)
    tail_closed = Int32(1)
    return IndirectConcurrentRingQueueNode(
        head,
        _PadAfter32(),
        tail_closed,
        _PadAfter32(),
        nothing,
        ring,
        len,
        buffers,
    )
end

Base.eltype(::Type{<:IndirectConcurrentRingQueueNode{T}}) where {T} = T

const CRQ_STARVATION = 128  # ???

function _close(crq::IndirectConcurrentRingQueueNode, old::Int32)
    while old > 0
        new = typemin(Int32) + old
        old, success = @atomicreplace crq.tail_closed old => new
        success && return
    end
end

function tail_index(crq::IndirectConcurrentRingQueueNode, t = @atomic crq.tail_closed)
    if t < 0
        t = t - typemin(Int32)
    end
    return t
end

isclosed(crq::IndirectConcurrentRingQueueNode, t = @atomic crq.tail_closed) = t ≤ 0

function trypush!(crq::IndirectConcurrentRingQueueNode, x)
    x = convert(eltype(crq), x)

    # TODO: create a Int64 version that that does not require this?
    let t = @atomic crq.tail_closed
        if isclosed(crq, t)
            return false  # closed
        elseif t ≥ typemax(Int32) - Threads.nthreads()
            _close(crq, t)
            return false  # closed
        end
    end

    starvation_ctr = 0
    while true
        t = (@atomic crq.tail_closed += true) - true
        if isclosed(crq, t)
            return false  # closed
        elseif t ≥ typemax(Int32) - Threads.nthreads()
            # Note: this condition is relying on that there are only at most
            # `nthreads()` instances on `trypush!` calls at once; i.e., no yield
            # point.
            _close(crq, t)
            return false  # closed
        end
        # TODO: use shift like DLCRQ
        itemindex = mod1(t, crq.length)
        slotptr = pointer(crq.ring, itemindex)
        slot = UnsafeAtomics.load(slotptr)::CRQSlot{UInt32}
        (; index, safe, storage) = slot
        if iszero(storage)
            if (index ≤ t) && (safe || (@atomic crq.head) ≤ t)
                tid = Threads.threadid()
                buffer = crq.buffers[tid]
                buffer[itemindex] = x
                storage = UInt32(tid)
                # storage = _embed(UInt32, x)
                @assert !iszero(storage)
                newslot = CRQSlot(; safe = true, index = t, storage)
                old = UnsafeAtomics.cas!(slotptr, slot, newslot)
                if old == slot
                    return true
                end
            end
        end
        starvation_ctr += 1
        if (t - (@atomic crq.head) ≥ crq.length) || starvation_ctr > CRQ_STARVATION
            _close(crq, t)
            return false  # closed
        end
        GC.safepoint()  # must not yield
    end
end

function ConcurrentCollections.maybepopfirst!(crq::IndirectConcurrentRingQueueNode)
    while true
        h = (@atomic crq.head += true) - true
        itemindex = mod1(h, crq.length)  # TODO: shift
        slotptr = pointer(crq.ring, itemindex)
        while true
            slot = UnsafeAtomics.load(slotptr)::CRQSlot{UInt32}
            (; index, safe, storage) = slot
            if !iszero(storage)
                if index == h
                    threadindex = storage
                    x = crq.buffers[threadindex][itemindex]  # [^crq.buffers]
                    # x = _extract(eltype(crq), storage)

                    newslot = CRQSlot(; safe, index = h + crq.length, storage = UInt32(0))
                    old = UnsafeAtomics.cas!(slotptr, slot, newslot)
                    if old === slot
                        return Some{eltype(crq)}(x)
                    end
                else
                    newslot = CRQSlot(; safe = false, index, storage)
                    old = UnsafeAtomics.cas!(slotptr, slot, newslot)
                    if old == slot
                        break
                    end
                end
            else  # empty slot
                # `max(h, index)` for dealing with tasks stalled after FAI [^maxh]
                newslot = CRQSlot(; safe, index = max(h, index) + crq.length, storage)
                old = UnsafeAtomics.cas!(slotptr, slot, newslot)
                if old == slot
                    break
                end
            end
            GC.safepoint()
        end
        GC.safepoint()

        # Failed to dequeue. Bail out if empty.
        if tail_index(crq) ≤ h + 1
            fixstate!(crq)
            return nothing
        end
    end
end
# [^crq.buffers]: Above read `x = crq.buffers[threadindex][itemindex]` (from
# `itemindex`) is fine without atomics. At this point, the load from the
# corresponding slot (`crq.ring[itemindex]`) has established the happens-before
# edge to the corresponding enqueue.  Since it has been confirmed that the
# storage is non-empty and `h == index`, there cannot be any more concurrent
# writes by the "current" or "past" enqueuers with `tail ≤ h`. The "future"
# enqueuers with `tail > h` will not touch the buffer until the slot is emptied
# by the CAS after this load. Since this dequeuer obtained the exclusive access
# to the index `h`, there's no other dequeuers trying the same CAS.
#
# [^maxh]: `max(h, index)` was not mentioned in Morrison and Afek (2013) (or its
# revised version) so it's not clear if this is needed.  However, it is possible
# that multiple enqueuers and/or dequeuers that would have incremented the
# `slot.index` is suspended just after the FAI.  If so, `slot.index` can be a
# round (or even multiple rounds) behind the current `crq.head`.  So, it seems
# like we need to prevent the enqueuing to this index (hence missing the item)
# by moving this to the next round.
#
# TODO: Is it OK to skip updating the slot if `index > h`?

function fixstate!(crq::IndirectConcurrentRingQueueNode)
    while true
        h = @atomic crq.head
        tail_closed = @atomic crq.tail_closed
        t = tail_index(crq, tail_closed)
        h ≤ t && return
        isclosed(crq, tail_closed) && return  # closed and empty
        _, success = @atomicreplace crq.tail_closed tail_closed => h
        success && return
    end
end

mutable struct LinkedConcurrentRingQueue{T}
    @atomic head::IndirectConcurrentRingQueueNode{T}
    @atomic tail::IndirectConcurrentRingQueueNode{T}
end

# lcrq_default_ringsize() = 2^17
lcrq_default_ringsize() = 2^10
# lcrq_default_ringsize() = 2^5

LinkedConcurrentRingQueue(ringsize::Integer = lcrq_default_ringsize()) =
    LinkedConcurrentRingQueue{Any}(ringsize)
function LinkedConcurrentRingQueue{T}(ringsize::Integer = lcrq_default_ringsize()) where {T}
    node = IndirectConcurrentRingQueueNode{T}(ringsize)
    return LinkedConcurrentRingQueue(node, node)
end

Base.eltype(::Type{<:LinkedConcurrentRingQueue{T}}) where {T} = T

const NEW_CRQ = Threads.Atomic{Int}(0)

function Base.push!(lcrq::LinkedConcurrentRingQueue{T}, x) where {T}
    x = convert(T, x)
    while true
        crq = @atomic lcrq.tail
        next = @atomic crq.next
        if next !== nothing
            @atomicreplace lcrq.tail crq => next
            continue
        end

        trypush!(crq, x) && return lcrq

        Threads.atomic_add!(NEW_CRQ, 1)
        newcrq = IndirectConcurrentRingQueueNode{T}(length(crq.ring))
        ok = trypush!(newcrq, x)
        @assert ok
        _, success = @atomicreplace crq.next nothing => newcrq
        if success
            @atomicreplace lcrq.tail crq => newcrq
            return lcrq
        end
    end
end

function ConcurrentCollections.maybepopfirst!(lcrq::LinkedConcurrentRingQueue)
    while true
        crq = @atomic lcrq.head
        x = maybepopfirst!(crq)
        x === nothing || return x
        next = @atomic crq.next
        next === nothing && return nothing

        x = maybepopfirst!(crq)
        x === nothing || return x
        @atomicreplace lcrq.head crq => next
    end
end

function Base.show(io::IO, ::MIME"text/plain", lcrq::LinkedConcurrentRingQueue)
    nitems = 0
    ncrqs = 0
    crq = @atomic lcrq.head
    while crq !== nothing
        h = @atomic crq.head
        t = tail_index(crq)
        nitems += max(0, t - h)
        ncrqs += 1
        crq = @atomic crq.next
    end
    print(io, "LCRQ: $nitems item(s) in $ncrqs node(s)")
end

function Base.show(io::IO, ::MIME"text/plain", crq::IndirectConcurrentRingQueueNode)
    h = @atomic crq.head
    t = tail_index(crq)
    nitems = max(0, t - h)
    status = isclosed(crq) ? "closed" : "open"
    print(io, "CRQ: $nitems item(s) (status: $status, head: $h, tail: $t)")
end

function Base.NamedTuple(slot::CRQSlot)
    (; index, safe, storage) = slot
    return (; index, safe, storage)
end

function Base.show(io::IO, ::MIME"text/plain", slot::CRQSlot)
    if get(io, :typeinfo, Any) !== typeof(slot)
        show(io, MIME"text/plain"(), typeof(slot))
    end
    show(io, MIME"text/plain"(), NamedTuple(slot))
end


text:
# brutal force caculations for accuracy benchmark
# partition function
function partition_function_exact(
    Ns::Int64, N::Int64, expβϵ::Array{T,1}
) where {T<:FloatType}

    config_list = config_list_generator(Ns::Int64, N::Int64)
    Z = 0

    for i in config_list
        Z += prod(expβϵ[findall(!iszero, expβϵ .* i)])
    end

    Z

end

# mean-occupation number
function occupancy_exact(
    Ns::Int64, N::Int64, expβϵ::Array{T,1}
) where {T<:FloatType}

    config_list = config_list_generator(Ns::Int64, N::Int64)
    Z = 0

    for i in config_list
        Z += prod(expβϵ[findall(!iszero, expβϵ .* i)])
    end

    Z = real(Z)
    n̄ = complex(zeros(Ns))

    for i = 1 : Ns
        for j in config_list
            if j[i] > 0
                n̄[i] += prod(expβϵ[findall(!iszero, expβϵ .* j)]) / Z
            end
        end
    end

    Z, n̄

end

# use deep-first search to generate all fermion configurations
function config_list_generator(Ns::Int64, N::Int64)

    config_Array = zeros(Int64, Ns)
    config_list = Array{Int64,1}[]
    DFS(1, Ns, N, config_Array, config_list)

    config_list
end

function DFS(site_index, Ns, N, config_Array, config_list)

    if N == 0 || site_index == Ns
        if N <= 1
            config_Array[site_index : Ns] .= N
            push!(config_list, copy(config_Array))
            return
        else
            return
        end
    end

    for i = 0 : 1
        config_Array[site_index] = i
        DFS(site_index + 1, Ns, N - i, config_Array, config_list)
    end

end

# Projection Method
function particle_projection(
    Ns::Int64, N::Int64, expβϵ::Array{T,1}, FLAG_OccNum::Bool
) where {T<:FloatType}

    φ = [2 * π * m / Ns for m = 1 : Ns]
    expβμ = (abs(expβϵ[Ns - N + 1]) +
            abs(expβϵ[Ns - N])) / 2     # chemical potential

    η = [prod(1 .+ exp(im * φ[m]) / expβμ * expβϵ) for m = 1 : Ns]
    Z = sum(exp.(-im * φ) .^ N .* η) / Ns * expβμ ^ N

    # skip the occupation number projection
    if FLAG_OccNum == false
        return real(Z)
    end

    n̄ = [sum(exp.(im * φ) / expβμ * expβϵ[i] ./ (1 .+ exp.(im * φ) / expβμ * expβϵ[i]) .*
    exp.(-im * φ) .^ N .* η) * expβμ ^ N / Ns / Z for i = 1 : Ns]

    return real(Z), n̄

end

function test_matrix_product(σlist::Array{Int64,2}, lattice::Lattice, qmc::QMC)

    MatProdUp = Matrix{Float64}(I, lattice.Ns, lattice.Ns)
    MatProdDn = Matrix{Float64}(I, lattice.Ns, lattice.Ns)

    for i = 1 : size(σlist)[2]
        σ = σlist[:, i]
        Bj_Up, Bj_Dn = singlestep_matrix(σ, lattice, qmc)
        MatProdUp = Bj_Up * MatProdUp
        MatProdDn = Bj_Dn * MatProdDn
    end

    MatProdUp, MatProdDn

end

function test_matrix_product_svd(
    A::Array{Float64, 2}, Ns::Int64, L::Int64
    )

    num_cluster_size = 10

    eigvalsA = eigvals(A)

    MatProd_U = Matrix{Float64}(I, Ns, Ns)
    MatProd_D = Matrix{Float64}(I, Ns, Ns)
    MatProd_V = Matrix{Float64}(I, Ns, Ns)

    for i = 1 : div(L, num_cluster_size)

        MatProd = Matrix{Float64}(I, Ns, Ns)

        for j = 1 : num_cluster_size
            MatProd = A * MatProd
        end

        BUD = MatProd * MatProd_U * MatProd_D
        SVD_BUD = LinearAlgebra.LAPACK.gesvd!('A', 'A', BUD)
        MatProd_U = SVD_BUD[1]
        MatProd_D = diagm(0 => SVD_BUD[2])
        MatProd_V = SVD_BUD[3] * MatProd_V

    end

    eigs_prop = eigvals(MatProd_D * MatProd_V * MatProd_U; permute = true, scale = true, sortby = abs)
    #eigs_prop = LinearAlgebra.LAPACK.geevx!('B', 'V', 'V', 'N', MatProd_D * MatProd_V * MatProd_U)

    eigvalsA .^ L, eigs_prop
    #MatProd_D

end

function test_matrix_product_qr(
    A::Array{Float64, 2}, Ns::Int64, L::Int64
    )

    num_cluster_size = 10

    eigvalsA = eigvals(A)

    MatProd_Q = Matrix{Float64}(I, Ns, Ns)
    MatProd_D = Matrix{Float64}(I, Ns, Ns)
    MatProd_T = Matrix{Float64}(I, Ns, Ns)

    for i = 1 : div(L, num_cluster_size)

        MatProd = Matrix{Float64}(I, Ns, Ns)

        for j = 1 : num_cluster_size
            MatProd = A * MatProd
        end

        QRCP_update!(MatProd_Q, MatProd_D, MatProd_T, MatProd, 'R')

    end

    eigs_prop = eigvals(MatProd_D * MatProd_T * MatProd_Q; permute = true, scale = true, sortby = abs)

    eigvalsA .^ L, eigs_prop

end

function measure_particlenumber_GC(
    lattice::Lattice, UUp_Array::Array{T1,3}, UDn_Array::Array{T2,3},
    expβμ::Float64, expβϵ_up::Array{T3,1}, expβϵ_dn::Array{T4,1}
) where {T1<:FloatType, T2<:FloatType, T3<:FloatType, T4<:FloatType}

    nup, ndn = 0, 0
    n̄_up_gc = expβμ * expβϵ_up ./ (1 .+ expβμ * expβϵ_up)
    n̄_dn_gc = expβμ * expβϵ_dn ./ (1 .+ expβμ * expβϵ_dn)

    for i = 1 : lattice.Ns
        nup += sum(UUp_Array[:, i, i] .* n̄_up_gc)
        ndn += sum(UDn_Array[:, i, i] .* n̄_dn_gc)
    end

    real(nup), real(ndn)

end

function measure_energy_GC(
    lattice::Lattice, UUp_Array::Array{T1,3}, UDn_Array::Array{T2,3},
    expβμ::Float64, expβϵ_up::Array{T3,1}, expβϵ_dn::Array{T4,1}
) where {T1<:FloatType, T2<:FloatType, T3<:FloatType, T4<:FloatType}

    EOneBody, ETwoBody = 0, 0
    n̄_up_gc = expβμ * expβϵ_up ./ (1 .+ expβμ * expβϵ_up)
    n̄_dn_gc = expβμ * expβϵ_dn ./ (1 .+ expβμ * expβϵ_dn)

    for i in eachindex(view(lattice.KMatrix, 1 : lattice.Ns, 1 : lattice.Ns))
        if lattice.KMatrix[i] != 0
            Dij_up = sum(UUp_Array[:, i[1], i[2]] .* n̄_up_gc)
            Dij_dn = sum(UDn_Array[:, i[1], i[2]] .*n̄_dn_gc)
            EOneBody += -lattice.t * (Dij_up + Dij_dn)
        end
    end

    for i = 1 : lattice.Ns
        Umatrix = UUp_Array[:, i, i] * UDn_Array[:, i, i]'
        nmatrix = n̄_up_gc * n̄_dn_gc'
        Diiii = sum(Umatrix .* nmatrix)
        ETwoBody += lattice.U * Diiii
    end

    real(EOneBody), real(ETwoBody)
end

function mc_integration_test(lattice::Lattice, qmc::QMC)

    io = open("GCETest", "a+")
    for n = 1 : qmc.sample_size
        # initialize a random configuration
        σlist = 2 * (rand(Float64, (lattice.Ns, qmc.L)) .< 0.5) .- 1
        # propagate and diagonalize
        EigvalsUp, EigvecsUp, EigvalsDn, EigvecsDn = matrix_product_stablized(σlist, lattice, qmc)
        # GCE partition function
        expβμ = exp(qmc.Δτ * qmc.L * qmc.μ)
        ZGCUp = prod(1 .+ expβμ * EigvalsUp)
        ZGCDn = prod(1 .+ expβμ * EigvalsDn)
        Sgn_GC = sign(ZGCUp) * sign(ZGCDn)
        # CE partition function
        ZCUp, n̄up = ProjectionPF(lattice.Ns, lattice.nup, EigvalsUp)
        ZCDn, n̄dn = ProjectionPF(lattice.Ns, lattice.ndn, EigvalsDn)
        Sgn_C = sign(ZCUp) * sign(ZCDn)

        if Sgn_GC * Sgn_C == -1
            println(io, "GCE Sign: ", Sgn_GC)
            println(io, σlist)
        end

    end

    close(io)

end

function mc_integration_GC(lattice::Lattice, qmc::QMC)
    # initialization
    nup_Array = Vector{Float64}()   # particle number (spin up)
    ndn_Array = Vector{Float64}()   # particle number (spin down)
    #Ek_Array = Vector{Float64}()    # kinetic energy
    #Ep_Array = Vector{Float64}()    # potential energy
    #E_Array = Vector{Float64}()     # total energy
    Sgn_Array = Vector{Int64}()     # sign
    # initialize a random configuration
    σlist = 2 * (rand(Float64, (lattice.Ns, qmc.L)) .< 0.5) .- 1
    # propagate and diagonalize
    EigvalsUp, EigvecsUp, EigvalsDn, EigvecsDn = matrix_product_stablized(σlist, lattice, qmc)
    # GC partition function
    expβμ = exp(qmc.Δτ * qmc.L * qmc.μ)
    ZGCUp = prod(1 .+ expβμ * EigvalsUp)
    ZGCDn = prod(1 .+ expβμ * EigvalsDn)
    Sgn = sign(ZGCUp) * sign(ZGCDn)

    for n = 1 : qmc.sample_size

        σlist_New = copy(σlist)

        # random flip a spin / new configuration
        flip = rand(1 : qmc.L * lattice.Ns)
        σlist_New[flip] *= -1
        EigvalsUp_New, EigvecsUp_New, EigvalsDn_New, EigvecsDn_New =
                    matrix_product_stablized(σlist_New, lattice, qmc)

        ZGCUp_New = prod(1 .+ expβμ * EigvalsUp_New)
        ZGCDn_New = prod(1 .+ expβμ * EigvalsDn_New)
        Sgn_New = sign(ZGCUp_New) * sign(ZGCDn_New)
        # caculate the acceptence ratio
        R = abs((ZGCUp_New / ZGCUp) * (ZGCDn_New / ZGCDn))
        AcceptRatio = R / (1 + R)

        if rand() <= AcceptRatio   # accept
            # adopt the new configuration
            σlist = σlist_New
            ZGCUp, ZGCDn = ZGCUp_New, ZGCDn_New
            EigvalsUp, EigvecsUp = EigvalsUp_New, EigvecsUp_New
            EigvalsDn, EigvecsDn = EigvalsDn_New, EigvecsDn_New
            Sgn = Sgn_New

            # measurement
            UUp_Array, UDn_Array = overlap_coefficient(lattice.Ns, EigvecsUp, EigvecsDn)
            nup, ndn =  measure_particlenumber_GC(lattice, UUp_Array, UDn_Array,
                                                #expβμ, EigvalsUp, EigvalsDn)
            #Ek, Ep = measure_energy_GC(lattice, UUp_Array, UDn_Array,
                                    expβμ, EigvalsUp, EigvalsDn)

            # store data
            push!(Sgn_Array, Sgn)
            push!(nup_Array, nup * Sgn)
            push!(ndn_Array, ndn * Sgn)
            #push!(Ek_Array, Ek * Sgn)
            #push!(Ep_Array, Ep * Sgn)
            #push!(E_Array, Ek * Sgn + Ep * Sgn)
        else                        # reject
            # measurement
            UUp_Array, UDn_Array = overlap_coefficient(lattice.Ns, EigvecsUp, EigvecsDn)
            nup, ndn =  measure_particlenumber_GC(lattice, UUp_Array, UDn_Array,
                                            #expβμ, EigvalsUp, EigvalsDn)
            #Ek, Ep = measure_energy_GC(lattice, UUp_Array, UDn_Array,
                                    expβμ, EigvalsUp, EigvalsDn)

            # store data
            push!(Sgn_Array, Sgn)
            push!(nup_Array, nup * Sgn)
            push!(ndn_Array, ndn * Sgn)
            #push!(Ek_Array, Ek * Sgn)
            #push!(Ep_Array, Ep * Sgn)
            #push!(E_Array, Ek * Sgn + Ep * Sgn)
        end
    end

    nup_Mean = mean(nup_Array[qmc.BurnIn : qmc.sample_size])
    nup_Var = varm(nup_Array[qmc.BurnIn : qmc.sample_size], nup_Mean) / qmc.sample_size
    ndn_Mean = mean(ndn_Array[qmc.BurnIn : qmc.sample_size])
    ndn_Var = varm(ndn_Array[qmc.BurnIn : qmc.sample_size], ndn_Mean) / qmc.sample_size

    #EMean = mean(E_Array[qmc.BurnIn : qmc.sample_size])
    #EVar = varm(E_Array[qmc.BurnIn : qmc.sample_size], EMean) / qmc.sample_size
    #EkMean = mean(Ek_Array[qmc.BurnIn : qmc.sample_size])
    #EkVar = varm(Ek_Array[qmc.BurnIn : qmc.sample_size], EkMean) / qmc.sample_size
    #EpMean = mean(Ep_Array[qmc.BurnIn : qmc.sample_size])
    #EpVar = varm(Ep_Array[qmc.BurnIn : qmc.sample_size], EpMean) / qmc.sample_size

    avgSgn = mean(Sgn_Array[qmc.BurnIn : qmc.sample_size])

    nup_Mean, nup_Var, ndn_Mean, ndn_Var, avgSgn
    #return EMean, EVar, EkMean, EkVar, EpMean, EpVar, avgSgn
end

function main_gc()
    lattice, qmc = GetInputs()

    BlockNUp_Array = Vector{Float64}()
    BlockNDn_Array = Vector{Float64}()
    #BlockEk_Array = Vector{Float64}()
    #BlockEp_Array = Vector{Float64}()
    #BlockE_Array = Vector{Float64}()
    #BlockEkVar_Array = Vector{Float64}()
    #BlockEpVar_Array = Vector{Float64}()
    #BlockEVar_Array = Vector{Float64}()
    BlockSgn_Array = Vector{Float64}()

    filename = string("GCE_output_mu", qmc.μ, "_U", lattice.U, "_beta", qmc.Δτ * qmc.L)
    outputfile = open(filename, "a+")

    @printf(outputfile, "       Grand Canonical Ensemble AFQMC for 2D Hubbard Model     \n\n")
    @printf(outputfile, "mu: %.4f\n", qmc.μ)
    @printf(outputfile, "Number of Blocks: %d\n", qmc.blocks)
    @printf(outputfile, "Length of Img Timestep: %.4f\n", qmc.Δτ)
    @printf(outputfile, "Stablization Interval: %d\n\n", qmc.cluster_size)

    for i = 1 : qmc.blocks
        BlockNUp, BlockNUpVar, BlockNDn, BlockNDnVar, avgSgn = mc_integration_GC(lattice, qmc)
        BlockNUp = BlockNUp / avgSgn
        BlockNDn = BlockNDn / avgSgn

        @printf(outputfile, "Block # %d out of %d\n", i, qmc.blocks)
        @printf(outputfile, "--------------------\n")
        @printf(outputfile, "Average Sign: %.3f\n", avgSgn)
        @printf(outputfile, "NUp            NDn            NTot\n")
        @printf(outputfile, "%.5f     %.5f       %.5f\n", BlockNUp, BlockNDn, BlockNUp + BlockNDn)
        @printf(outputfile, "Error:\n")
        @printf(outputfile, "%.5f     %.5f       %.5f\n\n", sqrt(BlockNUpVar), sqrt(BlockNDnVar), sqrt(BlockNUpVar) + sqrt(BlockNDnVar))

        push!(BlockNUp_Array, BlockNUp)
        push!(BlockNDn_Array, BlockNDn)
        push!(BlockSgn_Array, avgSgn)

        #BlockE, BlockEVar, BlockEk, BlockEkVar, BlockEp, BlockEpVar, avgSgn = mc_integration_GC(lattice, qmc)
        #BlockE = BlockE / avgSgn
        #BlockEk = BlockEk / avgSgn
        #BlockEp = BlockEp / avgSgn

        #@printf(outputfile, "Block # %d out of %d\n", i, qmc.blocks)
        #@printf(outputfile, "--------------------\n")
        #@printf(outputfile, "Average Sign: %.3f\n", avgSgn)
        #@printf(outputfile, "Ek            Ep            Etot\n")
        #@printf(outputfile, "%.5f     %.5f       %.5f\n", BlockEk, BlockEp, BlockE)
        #@printf(outputfile, "Error:\n")
        #@printf(outputfile, "%.5f     %.5f       %.5f\n\n", sqrt(BlockEkVar), sqrt(BlockEpVar), sqrt(BlockEVar))

        #push!(BlockEk_Array, BlockEk)
        #push!(BlockEkVar_Array, BlockEkVar)
        #push!(BlockEp_Array, BlockEp)
        #push!(BlockEpVar_Array, BlockEpVar)
        #push!(BlockE_Array, BlockE)
        #push!(BlockEVar_Array, BlockEVar)
        #push!(BlockSgn_Array, avgSgn)
    end

    NUp_Mean = mean(BlockNUp_Array)
    NDn_Mean = mean(BlockNDn_Array)

    @printf(outputfile, "--------------------\n")
    @printf(outputfile, "Results for this run:\n")
    @printf(outputfile, "Average Sign: %.3f\n", mean(BlockSgn_Array))
    @printf(outputfile, "NUp            NDn            NTot\n")
    @printf(outputfile, "%.5f     %.5f       %.5f\n", NUp_Mean, NDn_Mean, NUp_Mean + NDn_Mean)

    #EkMean = mean(BlockEk_Array)
    #EkVar = sum(BlockEkVar_Array * (qmc.sample_size - 1) .+
    #qmc.sample_size * (BlockEk_Array .- EkMean) .^ 2) / (qmc.blocks * qmc.sample_size - 1)
    #EpMean = mean(BlockEp_Array)
    #EpVar = sum(BlockEpVar_Array * (qmc.sample_size - 1) .+
    #qmc.sample_size * (BlockEp_Array .- EpMean) .^ 2) / (qmc.blocks * qmc.sample_size - 1)
    #EMean = mean(BlockE_Array)
    #EVar = sum(BlockEVar_Array * (qmc.sample_size - 1) .+
    #qmc.sample_size * (BlockE_Array .- EMean) .^ 2) / (qmc.blocks * qmc.sample_size - 1)

    #@printf(outputfile, "--------------------\n")
    #@printf(outputfile, "Results for this run:\n")
    #@printf(outputfile, "Average Sign: %.3f\n", mean(BlockSgn_Array))
    #@printf(outputfile, "Ek            Ep            Etot\n")
    #@printf(outputfile, "%.5f     %.5f       %.5f\n", EkMean, EpMean, EMean)
    #@printf(outputfile, "Error:\n")
    #@printf(outputfile, "%.5f     %.5f       %.5f\n\n", sqrt(EkVar / qmc.blocks), sqrt(EpVar / qmc.blocks), sqrt(EVar / qmc.blocks))

    close(outputfile)

end

max_stars_repo_name: TongSericus/CanonicalAFQMC

max_issues_repo_name: TongSericus/CanonicalAFQMC

max_forks_repo_name: TongSericus/CanonicalAFQMC