#Contains functions to update the networks

#Simply iterate trough the dataset and update:
function do_sac_update(Pl::Any,Qf1::Any, Qf2::Any, Vf::Any,Vf_tar::Any,buffer::Traj_Data,params::Params, opt_para_pl, opt_para_qf1, opt_para_qf2, opt_para_vf)
    n_traj = buffer.n_traj
    shuffle_arr = shuffle(1:n_traj)
    τ = params.τ

    for i_true = 1:n_traj
        i_traj = shuffle_arr[i_true]
        traj = buffer.traj[i_traj]

        #Get the gradients:
        ĝ_vf, ĝ_qf1, ĝ_qf2, ĝ_pl = grad_all_networks(traj, Vf, Vf_tar, Qf1, Qf2, Pl, params)

        #Update the Networks:
        Knet.update!(Vf, ĝ_vf, opt_para_vf)
        Knet.update!(Qf1, ĝ_qf1, opt_para_qf1)
        Knet.update!(Qf2, ĝ_qf2, opt_para_qf2)
        Knet.update!(Pl, ĝ_pl, opt_para_pl)
        Vf_tar = τ .* Vf .+ (1-τ) .* Vf_tar
    end
end

#This function will update the networks in parallel:
function do_sac_update_parallel!(Pl::Any,Qf1::Any,Qf2::Any,Vf::Any,Vf_tar::Any,buffer::Traj_Data,params::Params,
     opt_para_pl, opt_para_qf1, opt_para_qf2, opt_para_vf)

    #Take all partial traj and shuffle their order!
    n_traj = buffer.n_traj
    shuffle_arr = shuffle(1:n_traj)
    τ = params.τ

    active_workers = workers()
    n_workers = length(active_workers)

    ĝ_pl  = copy(Pl)
    ĝ_qf1 = copy(Qf1)
    ĝ_qf2 = copy(Qf2)
    ĝ_vf  = copy(Vf)

    # Create Channels storing results and starting values
    input_channel  = RemoteChannel(()->Channel{Tuple}(length(workers())+1)) # will contain networks and traj
    output_channel = RemoteChannel(()->Channel{Tuple}(length(workers())+1)) # will contain networks

    # Create problem on every worker
    begin
        for p in workers()
            remote_do(distributed_sac_update, p, params, input_channel, output_channel)
        end
    end

    # Start optimizations
    n_parallel = Int(ceil(n_traj / (n_workers * 20)))
    prog = Progress(n_parallel,1)

    for i_parallel = 1 : n_parallel

        #1. Send the work tasks:
        for i_w = 1 : n_workers
            i_true = ((i_parallel-1)*n_workers + i_w) #get true order
            if i_true <= n_traj
                i_traj = shuffle_arr[i_true] #shuffle order
                traj_local = buffer.traj[i_traj]
                @async put!(input_channel, (traj_local, Pl, Qf1, Qf2, Vf, Vf_tar))
            end
        end

        # 2. Take the results
        ĝ_pl  *= 0
        ĝ_qf1 *= 0
        ĝ_qf2 *= 0
        ĝ_vf  *= 0

        for i_w = 1: n_workers
            i_true = ((i_parallel-1)*n_workers + i_w)
            if i_true <= n_traj
                #Take from Channel
                gradients = take!(output_channel)
                #Reorder
                ĝ_vf  += gradients[1]
                ĝ_qf1 += gradients[2]
                ĝ_qf2 += gradients[3]
                ĝ_pl  += gradients[4]
            end
        end

        #Update the Networks
        Knet.update!(Vf,  ĝ_vf,  opt_para_vf)
        Knet.update!(Qf1, ĝ_qf1, opt_para_qf1)
        Knet.update!(Qf2, ĝ_qf2, opt_para_qf2)
        Knet.update!(Pl,  ĝ_pl,  opt_para_pl)
        Vf_tar = τ .* Vf .+ (1-τ) .* Vf_tar

        #Update Progress
        #ProgressMeter.update!(prog,i_parallel) #update progress bar
    end

    # Tell workers that its done
    NaN_mat = zeros(length(params.batch_size))
    fill!(NaN_mat,NaN)
    ending_traj = init_traj(params,params.batch_size)
    ending_traj.r_t = NaN_mat

    for j = 1 : n_workers
        @async put!(input_channel, (ending_traj, Pl, Qf1, Qf2, Vf, Vf_tar))
    end

    #Close Channels
    sleep(0.1)
    close(input_channel)
    close(output_channel)

    return nothing
end

#Function running on the worker
@everywhere function distributed_sac_update(params::Params, input_channel::RemoteChannel, output_channel::RemoteChannel)

   #TODO BREAK!!!!

   # Start loop where work is done
   while true
       # break loop if NaN state received
       try
           #Take the new Input
           input = take!(input_channel)

           #Reorder data:
           traj_local = input[1]
           Pl         = input[2]
           Qf1        = input[3]
           Qf2        = input[4]
           Vf         = input[5]
           Vf_tar     = input[6]

           #Check if u reached final
           if (sum(isnan, traj_local.r_t) > 0)
               GC.gc()
               break
           end

           #DO THE WORK:
           #-----------------------------------------------------------------------
           ĝ_vf, ĝ_qf1, ĝ_qf2, ĝ_pl = grad_all_networks(traj_local, Vf, Vf_tar, Qf1, Qf2, Pl, params)
           #-----------------------------------------------------------------------

           put!(output_channel, (ĝ_vf, ĝ_qf1, ĝ_qf2, ĝ_pl) )

       catch
           println("Some problem occured, worker $(myid()) exiting trajectory optimization function.")
           break
       end
   end

   return nothing
end


using DataStructures, DataArrays , DataFrames, StatsFuns, GLM , Distributions, MixedModels, StatsBase, JSON, StatLib, NamedArrays
using DataStructures, DataArrays , DataFrames, StatsFuns, GLM , Distributions, MixedModels, StatsBase, JSON, StatLib, NamedArrays

# ************ CUSTOMIZE *********************
#root="/mnt/resource/analytics/models/Jenny-o"
#root="/mnt/resource/analytics/models/NatChoice"
#root="/mnt/resource/analytics/models/rev"
#root="/mnt/resource/analytics/models/ALL#10"
#root="/mnt/resource/analytics/models/ALL#11"
#root="/mnt/resource/analytics/models/CDW#6"
#root="/mnt/resource/analytics/models/HormelChili#8"
#root="/mnt/resource/analytics/models/Rev#8"
#root="/mnt/resource/analytics/models/Jenny-o"
#root="/mnt/resource/analytics/models/hunts"
#root="/mapr/mapr04p/analytics0001/analytic_users/jmdl/rev"
#root="/mapr/mapr04p/analytics0001/analytic_users/jmdl/NatChoice"
#root="/mapr/mapr04p/analytics0001/analytic_users/jmdl/ALL10"

# ************ END CUSTOMIZE *********************
root = !isdefined(:root) ? pwd() : root

modelsDict = read_modelsDict(root)

function runGlm(root::String, modelname::Symbol, ranef::Symbol=:empty)
    println("Loding data : ")
    dfd = read_dfd(root)
    modelsDict = read_modelsDict(root)
    poolit!(dfd,modelsDict[:factors])
    m=modelsDict[modelname]
    return runGlm(dfd, m, ranef)
end

function runGlm(dfd::DataFrame,m::Dict, ranef::Symbol=:empty)
    if ranef==:empty
        f=genF(m[:y_var],m[:finalvars])
        cols = convert(Array{Symbol},vcat(m[:finalvars], [m[:y_var], m[:logvar]]))
        println("Ok - Running ",m[:modelName]," -  Glm on proc ",myid()," with : ",f)
    else
        f=genF(m[:y_var],setdiff(vcat(m[:finalvars],ranef),[:group]))
        dfd[ranef] = relevel(dfd[ranef], "none")
        cols = convert(Array{Symbol},vcat(m[:finalvars], [m[:y_var], m[:logvar]], ranef ))
        println("Ok - Running ",m[:modelName]," - ",ranef," Glm on proc ",myid()," with : ",f)
    end
    if m[:Buyer_Pos_P1_is1]
        return glm(f, dfd[(dfd[:buyer_pos_p1].==1),cols] , m[:dist], m[:lnk])
    else
        return return glm(f, dfd[cols] , m[:dist], m[:lnk])
    end
end

function runGlmm(root::String, modelname::Symbol, ranef::Symbol=:empty, runθ::Bool=true)
    println("Loding data : Glmm : ",modelname," ~~~ ",ranef)
    dfd = read_dfd(root)   #dfd = readtable(mod_fname,header=true);
    modelsDict = read_modelsDict(root) #modelsDict = readModels(jmod_fname)
    poolit!(dfd,modelsDict[:factors])
    m=modelsDict[modelname]
    if ranef==:empty
        v_out=Dict()
        for r in m[:raneff]
            v_out[r] = runGlmm(dfd, m, [r], runθ)
        end
        return v_out
    else
        return runGlmm(dfd, m, [ranef], runθ)
    end
end

function runGlmm(dfd::DataFrame,m::Dict, ranef::Array{Symbol}=[:empty], runθ::Bool=true)
    ranef = ranef==[:empty] ? m[:raneff]  : ranef
    for r in ranef
        dfd[r] = relevel(dfd[r], "none")
    end
    f=genF(m[:y_var],setdiff(m[:finalvars],[:group]),ranef)
    cols = convert(Array{Symbol},vcat(m[:finalvars], [m[:y_var], m[:logvar]], ranef ))
    println("Ok - Running ",m[:modelName]," - ",ranef," Glmm on proc ",myid()," with : ",f)
    if m[:Buyer_Pos_P1_is1]
        #gmm1 = fit!(glmm(f, dfd[(dfd[:buyer_pos_p1].==1),cols] , m[:dist], m[:lnk]), false, 0)   #gmm1 = fit!(glmm(f, dfd[(dfd[:buyer_pos_p1].==1),cols] , m[:dist], m[:lnk]), nAGQ=0)
        gmm1 = fit!(glmm(f, dfd[(dfd[:buyer_pos_p1].==1),cols] , m[:dist], m[:lnk]), verbose=false, fast=true)
        if runθ if sum(raneffect(gmm1)[:coef]) == 0 profileθ(gmm1) end  end
        return gmm1
    else
        #gmm1 = fit!(glmm(f, dfd[cols], m[:dist], m[:lnk]), false, 0)   #gmm1 = fit!(glmm(f, dfd[cols], m[:dist], m[:lnk]), nAGQ=0)
        gmm1 = fit!(glmm(f, dfd[cols] , m[:dist], m[:lnk]), verbose=false, fast=true)
        if runθ if sum(raneffect(gmm1)[:coef]) == 0 profileθ(gmm1) end end
        return gmm1
    end
end

# *******************************************
function runModels(root::String, modelsDict::Dict)
    dfd = read_dfd(root) #dfd = readtable(mod_fname,header=true);
    poolit!(dfd,modelsDict[:factors])
    res = OrderedDict()
    q = Symbol[]
    cnt=0
    for (k,v) in OrderedDict(m=>modelsDict[m][:raneff] for m in [:ipen,:idolocc,:iocc])
         for r in v push!(q,k); push!(q,r); cnt=cnt+1 end
    end
    x=reshape(q, (2,cnt))
    ml = permutedims(x, [2, 1])
    m=:empty
    for i in 1:length(ml[:,1])
        if m != ml[i,:][1]
            gr=Symbol("glm_"*string(ml[i,:][1]))
            res[gr] = runGlm(dfd,modelsDict[ml[i,:][1]]) #res[gr] = runGlm(mod_fname, jmod_fname, ml[i,:][1])
            m=ml[i,:][1]
        end
        r=Symbol("glmm_"*string(ml[i,:][1])*"_"*string(ml[i,:][2]))
        res[r] = runGlmm(dfd,modelsDict[m], [ml[i,:][2]], true)   #res[r] = runGlmm(mod_fname, jmod_fname, ml[i,:][1], ml[i,:][2])
        covg=Symbol("covglm_"*string(ml[i,:][1])*"_"*string(ml[i,:][2]))
        res[covg] = runGlm(dfd,modelsDict[m], ml[i,:][2])   #res[covg] = runGlm(mod_fname, jmod_fname, ml[i,:][1], ml[i,:][2])
    end
    return res
end
#dx = runModels(root,modelsDict)
dx = runModels(root,modelsDict)

"""
function consolidateResultsOLD(modelsDict::Dict,dx::OrderedDict)
  sdf = DataFrame(parameter=String[], coef=Int64[], stderr=Int64[], zval=Int64[], pval=Int64[], model=Symbol[], ranef=Symbol[], modelType=Symbol[])
    for m in [:iocc, :idolocc, :ipen]
        g = dx[Symbol("glm_"*string(m))]
        dfx = coefDF(g,false)
        dfx[:model] = Symbol(replace(string(m),"i",""))
        dfx[:ranef] = :none
        dfx[:modelType] = :GLM
        sdf = vcat(sdf,dfx)
        grp = sdf[(sdf[:modelType].==:GLM)&(sdf[:model].==Symbol(replace(string(m),"i","")))&(sdf[:parameter].=="group"),:coef][1]
        err = sdf[(sdf[:modelType].==:GLM)&(sdf[:model].==Symbol(replace(string(m),"i","")))&(sdf[:parameter].=="group"),:stderr][1]
        for r in  modelsDict[m][:raneff]
            k="_"*string(m)*"_"*string(r)
            g = dx[Symbol("glmm"*k)]
            dfx = raneffect(g)
            dfx[:zval]=0.0
            dfx[:pval]=0.0
            dfx[:model] = Symbol(replace(string(m),"i",""))
            dfx[:ranef]=r
            dfx[:modelType]=:GLMM
            sdf = vcat(sdf,dfx)
            g = dx[Symbol("covglm"*k)]
            dfx = coefDF(g,false)
            dfx[:model] = Symbol(replace(string(m),"i",""))
            dfx[:ranef] = r
            dfx[:modelType] = :RanCov
            dfx = dfx[findin(dfx[:parameter],filter(x->contains(x,string(r)) ,Array(dfx[:parameter]))),:]
            #dfx[:coef]=dfx[:coef]-grp
            #dfx[:stderr]=sqrt(dfx[:stderr].^2+err^2)
            dfx[:parameter] = map(x-> replace(x,string(r)*": ",""),dfx[:parameter])
            sdf = vcat(sdf,dfx)
        end
    end
    return sdf
end
dx = runModels(root,modelsDict)

function consolidateResults(modelsDict::Dict,dx::OrderedDict)
  sdf = DataFrame(parameter=String[], coef=Int64[], stderr=Int64[], zval=Int64[], pval=Int64[], model=Symbol[], ranef=String[], modelType=String[])
    for m in [:iocc, :idolocc, :ipen]
        g = dx[Symbol("glm_"*string(m))]
        dfx = coefDF(g,false)
        dfx[:model] = replace(string(m),"i","")
        dfx[:ranef] = "none"
        dfx[:modelType] = "GLM"
        sdf = vcat(sdf,dfx)
        grp = sdf[(sdf[:modelType].=="GLM")&(sdf[:model].==replace(string(m),"i",""))&(sdf[:parameter].=="group"),:coef][1]
        err = sdf[(sdf[:modelType].=="GLM")&(sdf[:model].==replace(string(m),"i",""))&(sdf[:parameter].=="group"),:stderr][1]
        for r in  modelsDict[m][:raneff]
            k="_"*string(m)*"_"*string(r)
            g = dx[Symbol("glmm"*k)]
            dfx = raneffect(g)
            dfx[:zval]=0.0
            dfx[:pval]=0.0
            dfx[:model] = replace(string(m),"i","")
            dfx[:ranef]=string(r)
            dfx[:modelType]="GLMM"
            sdf = vcat(sdf,dfx)
            g = dx[Symbol("covglm"*k)]
            dfx = coefDF(g,false)
            dfx[:model] = replace(string(m),"i","")
                dfx[:ranef] = string(r)
            dfx[:modelType] = "RanCov"
            dfx = dfx[findin(dfx[:parameter],filter(x->contains(x,string(r)) ,Array(dfx[:parameter]))),:]
            #dfx[:coef]=dfx[:coef]-grp
            #dfx[:stderr]=sqrt(dfx[:stderr].^2+err^2)
            dfx[:parameter] = map(x-> replace(x,string(r)*": ",""),dfx[:parameter])
            sdf = vcat(sdf,dfx)
        end
    end
    return sdf
end
"""

function consolidateResults(modelsDict::Dict,dx::OrderedDict=OrderedDict())
  sdf = DataFrame(parameter=String[], coef=Int64[], stderr=Int64[], zval=Int64[], pval=Int64[], model=String[], ranef=String[], modelType=String[])
    for m in [:iocc, :idolocc, :ipen]
        g = length(dx) > 0 ? dx[Symbol("glm_"*string(m))] : getw( Symbol("glm_"*string(m)) )
        #g = takew(  Symbol("glm_"*string(m))    )
        dfx = coefDF(g,false)
        dfx[:model] = replace(string(m),"i","")
        dfx[:ranef] = "none"
        dfx[:modelType] = "GLM"
        sdf = vcat(sdf,dfx)
        grp = sdf[(sdf[:modelType].=="GLM")&(sdf[:model].==replace(string(m),"i",""))&(sdf[:parameter].=="group"),:coef][1]
        err = sdf[(sdf[:modelType].=="GLM")&(sdf[:model].==replace(string(m),"i",""))&(sdf[:parameter].=="group"),:stderr][1]
        for r in  modelsDict[m][:raneff]
            k="_"*string(m)*"_"*string(r)
            g = length(dx) > 0 ?  dx[Symbol("glmm"*k)]  : getw(Symbol("glmm"*k))
            dfx = raneffect(g)
            dfx[:zval]=0.0
            dfx[:pval]=0.0
            dfx[:model] = replace(string(m),"i","")
            dfx[:ranef]=string(r)
            dfx[:modelType]="GLMM"
            sdf = vcat(sdf,dfx)
            g = length(dx) > 0 ?  dx[Symbol("covglm"*k)]    : getw(Symbol("covglm"*k))
            dfx = coefDF(g,false)
            dfx[:model] = replace(string(m),"i","")
                dfx[:ranef] = string(r)
            dfx[:modelType] = "RanCov"
            dfx = dfx[findin(dfx[:parameter],filter(x->contains(x,string(r)) ,Array(dfx[:parameter]))),:]
            #dfx[:coef]=dfx[:coef]-grp
            #dfx[:stderr]=sqrt(dfx[:stderr].^2+err^2)
            dfx[:parameter] = map(x-> replace(x,string(r)*": ",""),dfx[:parameter])
            sdf = vcat(sdf,dfx)
        end
    end
    return sdf
end
#dfx = consolidateResults(modelsDict)

#dx = runModels(root,modelsDict)
#dfx = consolidateResults(modelsDict, dx)
#save_dfx(root,dfx)   #writetable(root*"/campaign.csv", dfx)
dfx = consolidateResults(modelsDict, dx)
save_dfx(root,dfx)   #writetable(root*"/campaign.csv", dfx)


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.

export addlazycallback, addcutcallback, addheuristiccallback, addinfocallback

abstract JuMPCallback
type LazyCallback <: JuMPCallback
    f::Function
    fractional::Bool
end
type CutCallback <: JuMPCallback
    f::Function
end
type HeuristicCallback <: JuMPCallback
    f::Function
end
type InfoCallback <: JuMPCallback
    f::Function
    when::Symbol
end

type CallbackAbort <: Exception end
export CallbackAbort

Base.copy{T<:JuMPCallback}(c::T) = T(copy(c))
Base.copy(c::LazyCallback) = LazyCallback(copy(c.f), c.fractional)

function addlazycallback(m::Model, f::Function; fractional::Bool=false)
    m.internalModelLoaded = false
    push!(m.callbacks, LazyCallback(f,fractional))
end
function addcutcallback(m::Model, f::Function)
    m.internalModelLoaded = false
    push!(m.callbacks, CutCallback(f))
end
function addheuristiccallback(m::Model, f::Function)
    m.internalModelLoaded = false
    push!(m.callbacks, HeuristicCallback(f))
end

function _unspecifiedstate()
    warn("""Info Callbacks without a when clause will currently default
         to fire only in the :Intermediate state to preserve its behavior
         with previous versions of JuMP.

         This behavior will be deprecated in subsequent versions of JuMP, so
         please rewrite all invocations of addinfocallbacks(m, f) to
         addinfocallbacks(m, f, when = :Intermediate) instead.
         """)
    :Intermediate
end

function addinfocallback(m::Model, f::Function; when::Symbol = _unspecifiedstate())
    m.internalModelLoaded = false
    push!(m.callbacks, InfoCallback(f, when))
end

function lazycallback(d::MathProgBase.MathProgCallbackData, m::Model, cbs::Vector{LazyCallback})
    state = MathProgBase.cbgetstate(d)
    @assert state == :MIPSol || state == :MIPNode
    anyfrac = mapreduce(|, cbs) do cb
        cb.fractional
    end
    if state == :MIPSol
        MathProgBase.cbgetmipsolution(d,m.colVal)
    elseif anyfrac
        MathProgBase.cbgetlpsolution(d,m.colVal)
    else
        return :Continue
    end
    try
        for cb in cbs
            if state == :MIPSol || cb.fractional
                cb.f(d)
            end
        end
    catch y
        if isa(y, CallbackAbort)
            return :Exit
        else
            rethrow(y)
        end
    end
    :Continue
end

function attach_callbacks(m::Model, cbs::Vector{LazyCallback})
    cb = d -> lazycallback(d,m,cbs)
    if applicable(MathProgBase.setlazycallback!, m.internalModel, cb)
        MathProgBase.setlazycallback!(m.internalModel, cb)
    else
        error("Solver does not support lazy callbacks")
    end
end

function cutcallback(d::MathProgBase.MathProgCallbackData, m::Model, cbs::Vector{CutCallback})
    state = MathProgBase.cbgetstate(d)
    @assert state == :MIPNode
    MathProgBase.cbgetlpsolution(d,m.colVal)
    try
        for cb in cbs
            cb.f(d)
        end
    catch y
        if isa(y, CallbackAbort)
            return :Exit
        else
            rethrow(y)
        end
    end
    :Continue
end

function attach_callbacks(m::Model, cbs::Vector{CutCallback})
    cb = d -> cutcallback(d,m,cbs)
    if applicable(MathProgBase.setcutcallback!, m.internalModel, cb)
        MathProgBase.setcutcallback!(m.internalModel, cb)
    else
        error("Solver does not support cut callbacks")
    end
end

function heurcallback(d::MathProgBase.MathProgCallbackData, m::Model, cbs::Vector{HeuristicCallback})
    state = MathProgBase.cbgetstate(d)
    @assert state == :MIPNode
    MathProgBase.cbgetlpsolution(d,m.colVal)
    try
        for cb in cbs
            cb.f(d)
        end
    catch y
        if isa(y, CallbackAbort)
            return :Exit
        else
            rethrow(y)
        end
    end
    :Continue
end

function attach_callbacks(m::Model, cbs::Vector{HeuristicCallback})
    cb = d -> heurcallback(d,m,cbs)
    if applicable(MathProgBase.setheuristiccallback!, m.internalModel, cb)
        MathProgBase.setheuristiccallback!(m.internalModel, cb)
    else
        error("Solver does not support heuristic callbacks")
    end
end

function infocallback(d::MathProgBase.MathProgCallbackData, m::Model, cbs::Vector{InfoCallback})
    state = MathProgBase.cbgetstate(d)
    if state == :MIPSol
        MathProgBase.cbgetmipsolution(d,m.colVal)
    elseif state == :MIPNode
        MathProgBase.cbgetlpsolution(d,m.colVal)
    end
    try
        for cb in cbs
            if cb.when == state
                cb.f(d)
            end
        end
    catch y
        if isa(y, CallbackAbort)
            return :Exit
        else
            rethrow(y)
        end
    end
    :Continue
end

function attach_callbacks(m::Model, cbs::Vector{InfoCallback})
    cb = d -> infocallback(d,m,cbs)
    if applicable(MathProgBase.setinfocallback!, m.internalModel, cb)
        MathProgBase.setinfocallback!(m.internalModel, cb)
    else
        error("Solver does not support info callbacks")
    end
end

function registercallbacks(m::Model)
    isempty(m.callbacks) && return # might as well avoid allocating the indexedVector

    cbtypes = unique(map(typeof, m.callbacks))
    for typ in cbtypes
        cbs::Vector{typ} = filter(m.callbacks) do cb
            isa(cb, typ)
        end
        attach_callbacks(m, cbs)
    end

    # prepare storage for callbacks
    m.indexedVector = IndexedVector(Float64, m.numCols)
end

# TODO: Should this be somewhere else?
const sensemap = Dict(:(<=) => '<', :(==) => '=', :(>=) => '>')

## Lazy constraints
export @lazyconstraint

macro lazyconstraint(args...)
    cbdata = esc(args[1])
    x = args[2]
    extra = vcat(args[3:end]...)
    # separate out keyword arguments
    kwargs = filter(ex->isexpr(ex,:kw), extra) # filtering expressions corresponding to kw args specs
    extra = filter(ex->!isexpr(ex,:kw), extra) # others

    localcut_val = false # by default, the lazy constraint is global
    for ex in kwargs
        kwarg = ex.args[1]
        if kwarg == :localcut
            localcut_val = esc(ex.args[2])   # excepted if otherwise specified ...
        else error("in @lazyconstraint($(join(args,','))), invalid keyword argument: $(kwarg)")
        end
    end

    if VERSION < v"0.5.0-dev+3231"
        x = comparison_to_call(x)
    end
    if isexpr(x, :call) && length(x.args) == 3 # simple comparison
        lhs = :($(x.args[2]) - $(x.args[3])) # move everything to the lhs
        newaff, parsecode = parseExprToplevel(lhs, :aff)
        sense, vectorized = _canonicalize_sense(x.args[1])
        code = quote
            aff = zero(AffExpr)
            $parsecode
            constr = constructconstraint!($newaff, $(quot(sense)))
        end
        if vectorized
            return quote
                $code
                for con in constr
                    addlazyconstraint($cbdata, con, localcut=$localcut_val)
                end
            end
        else
            return quote
                $code
                addlazyconstraint($cbdata, constr, localcut=$localcut_val)
            end
        end
    else
        error("Syntax error in @lazyconstraint, expected one-sided comparison.")
    end
end

function addlazyconstraint(cbdata::MathProgBase.MathProgCallbackData, constr::LinearConstraint; localcut::Bool=false)
    if length(constr.terms.vars) == 0
        MathProgBase.cbaddlazy!(cbdata, Cint[], Float64[], sensemap[sense(constr)], rhs(constr))
        return
    end
    assert_isfinite(constr.terms)
    m::Model = constr.terms.vars[1].m
    indices, coeffs = merge_duplicates(Cint, constr.terms, m.indexedVector, m)
    if localcut
        MathProgBase.cbaddlazylocal!(cbdata, indices, coeffs, sensemap[sense(constr)], rhs(constr))
    else
        MathProgBase.cbaddlazy!(cbdata, indices, coeffs, sensemap[sense(constr)], rhs(constr))
    end
end

addlazyconstraint(cbdata::MathProgBase.MathProgCallbackData,constr::QuadConstraint; localcut::Bool=false) = error("Quadratic lazy constraints are not supported.")

## User cuts
export addusercut
export @usercut

macro usercut(args...)
    cbdata = esc(args[1])
    x = args[2]
    extra = vcat(args[3:end]...)
    # separate out keyword arguments
    kwargs = filter(ex->isexpr(ex,:kw), extra) # filtering expressions corresponding to kw args specs
    extra = filter(ex->!isexpr(ex,:kw), extra) # others

    localcut_val = false # by default, the user cut is global
    for ex in kwargs
        kwarg = ex.args[1]
        if kwarg == :localcut
            localcut_val = esc(ex.args[2])   # excepted if otherwise specified ...
        else error("in @usercut($(join(args,','))), invalid keyword argument: $(kwarg)")
        end
    end

    #cbdata = esc(cbdata)
    if VERSION < v"0.5.0-dev+3231"
        x = comparison_to_call(x)
    end
    if isexpr(x, :call) && length(x.args) == 3 # simple comparison
        lhs = :($(x.args[2]) - $(x.args[3])) # move everything to the lhs
        newaff, parsecode = parseExprToplevel(lhs, :aff)
        sense, vectorized = _canonicalize_sense(x.args[1])
        vectorized && error("Cannot add vectorized constraint in cut callback")
        quote
            aff = zero(AffExpr)
            $parsecode
            constr = constructconstraint!($newaff, $(quot(sense)))
            addusercut($cbdata, constr, localcut=$localcut_val)
        end
    else
        error("Syntax error in @usercut, expected one-sided comparison.")
    end
end

function addusercut(cbdata::MathProgBase.MathProgCallbackData, constr::LinearConstraint; localcut::Bool=false)
    if length(constr.terms.vars) == 0
        MathProgBase.cbaddcut!(cbdata, Cint[], Float64[], sensemap[sense(constr)], rhs(constr))
        return
    end
    assert_isfinite(constr.terms)
    m::Model = constr.terms.vars[1].m
    indices, coeffs = merge_duplicates(Cint, constr.terms, m.indexedVector, m)
    if localcut
        MathProgBase.cbaddcutlocal!(cbdata, indices, coeffs, sensemap[sense(constr)], rhs(constr))
    else
        MathProgBase.cbaddcut!(cbdata, indices, coeffs, sensemap[sense(constr)], rhs(constr))
    end
end

## User heuristic
export addsolution, setsolutionvalue

addsolution(cbdata::MathProgBase.MathProgCallbackData) = MathProgBase.cbaddsolution!(cbdata)
function setsolutionvalue(cbdata::MathProgBase.MathProgCallbackData, v::Variable, x)
    MathProgBase.cbsetsolutionvalue!(cbdata, convert(Cint, v.col), x)
end


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct HydroDispatchCascade <: HydroCascade
        name::String
        available::Bool
        bus::PSY.Bus
        active_power::Float64
        reactive_power::Float64
        rating::Float64
        prime_mover::PSY.PrimeMovers
        active_power_limits::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
        reactive_power_limits::Union{Nothing, PSY.Min_Max}
        ramp_limits::Union{Nothing, NamedTuple{(:up, :down), Tuple{Float64, Float64}}}
        time_limits::Union{Nothing, NamedTuple{(:up, :down), Tuple{Float64, Float64}}}
        base_power::Float64
        operation_cost::PSY.OperationalCost
        storage_target::Float64
        conversion_factor::Float64
        upstream::Vector{NamedTuple{(:unit, :lag, :multiplier), Tuple{PSY.HydroGen, Int64, Float64}}}
        services::Vector{PSY.Service}
        dynamic_injector::Union{Nothing, PSY.DynamicInjection}
        ext::Dict{String, Any}
        time_series_container::InfrastructureSystems.TimeSeriesContainer
        internal::IS.InfrastructureSystemsInternal
    end

# Arguments
- `name::String`
- `available::Bool`
- `bus::PSY.Bus`
- `active_power::Float64`
- `reactive_power::Float64`, validation range: `reactive_power_limits`, action if invalid: `warn`
- `rating::Float64`: Thermal limited MVA Power Output of the unit. <= Capacity, validation range: `(0, nothing)`, action if invalid: `error`
- `prime_mover::PSY.PrimeMovers`: Prime mover technology according to EIA 923
- `active_power_limits::NamedTuple{(:min, :max), Tuple{Float64, Float64}}`
- `reactive_power_limits::Union{Nothing, PSY.Min_Max}`, action if invalid: `warn`
- `ramp_limits::Union{Nothing, NamedTuple{(:up, :down), Tuple{Float64, Float64}}}`: ramp up and ramp down limits in MW (in component base per unit) per minute, validation range: `(0, nothing)`, action if invalid: `error`
- `time_limits::Union{Nothing, NamedTuple{(:up, :down), Tuple{Float64, Float64}}}`: Minimum up and Minimum down time limits in hours, validation range: `(0, nothing)`, action if invalid: `error`
- `base_power::Float64`: Base power of the unit in MVA, validation range: `(0, nothing)`, action if invalid: `warn`
- `operation_cost::PSY.OperationalCost`: Operation Cost of Generation [`OperationalCost`](@ref)
- `storage_target::Float64`: Storage target at the end of simulation as ratio of storage capacity.
- `conversion_factor::Float64`: Conversion factor from flow/volume to energy: m^3 -> p.u-hr.
- `upstream::Vector{NamedTuple{(:unit, :lag, :multiplier), Tuple{PSY.HydroGen, Int64, Float64}}}`: unit: upstream units; lag: duration in number of periods between upstream release and downstream availability; multiplier: relationship between upstream energy release and downstream energy availability
- `services::Vector{PSY.Service}`: Services that this device contributes to
- `dynamic_injector::Union{Nothing, PSY.DynamicInjection}`: corresponding dynamic injection device
- `ext::Dict{String, Any}`
- `time_series_container::InfrastructureSystems.TimeSeriesContainer`: internal time_series storage
- `internal::IS.InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct HydroDispatchCascade <: HydroCascade
    name::String
    available::Bool
    bus::PSY.Bus
    active_power::Float64
    reactive_power::Float64
    "Thermal limited MVA Power Output of the unit. <= Capacity"
    rating::Float64
    "Prime mover technology according to EIA 923"
    prime_mover::PSY.PrimeMovers
    active_power_limits::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
    reactive_power_limits::Union{Nothing, PSY.Min_Max}
    "ramp up and ramp down limits in MW (in component base per unit) per minute"
    ramp_limits::Union{Nothing, NamedTuple{(:up, :down), Tuple{Float64, Float64}}}
    "Minimum up and Minimum down time limits in hours"
    time_limits::Union{Nothing, NamedTuple{(:up, :down), Tuple{Float64, Float64}}}
    "Base power of the unit in MVA"
    base_power::Float64
    "Operation Cost of Generation [`OperationalCost`](@ref)"
    operation_cost::PSY.OperationalCost
    "Storage target at the end of simulation as ratio of storage capacity."
    storage_target::Float64
    "Conversion factor from flow/volume to energy: m^3 -> p.u-hr."
    conversion_factor::Float64
    "unit: upstream units; lag: duration in number of periods between upstream release and downstream availability; multiplier: relationship between upstream energy release and downstream energy availability"
    upstream::Vector{
        NamedTuple{(:unit, :lag, :multiplier), Tuple{PSY.HydroGen, Int64, Float64}},
    }
    "Services that this device contributes to"
    services::Vector{PSY.Service}
    "corresponding dynamic injection device"
    dynamic_injector::Union{Nothing, PSY.DynamicInjection}
    ext::Dict{String, Any}
    "internal time_series storage"
    time_series_container::InfrastructureSystems.TimeSeriesContainer
    "power system internal reference, do not modify"
    internal::IS.InfrastructureSystemsInternal
end

function HydroDispatchCascade(
    name,
    available,
    bus,
    active_power,
    reactive_power,
    rating,
    prime_mover,
    active_power_limits,
    reactive_power_limits,
    ramp_limits,
    time_limits,
    base_power,
    operation_cost = PSY.TwoPartCost(0.0, 0.0),
    storage_target = 1.0,
    conversion_factor = 1.0,
    upstream = NamedTuple{(:unit, :lag, :multiplier), Tuple{PSY.HydroGen, Int64, Float64}}[],
    services = PSY.Device[],
    dynamic_injector = nothing,
    ext = Dict{String, Any}(),
    time_series_container = InfrastructureSystems.TimeSeriesContainer(),
)
    HydroDispatchCascade(
        name,
        available,
        bus,
        active_power,
        reactive_power,
        rating,
        prime_mover,
        active_power_limits,
        reactive_power_limits,
        ramp_limits,
        time_limits,
        base_power,
        operation_cost,
        storage_target,
        conversion_factor,
        upstream,
        services,
        dynamic_injector,
        ext,
        time_series_container,
        IS.InfrastructureSystemsInternal(),
    )
end

function HydroDispatchCascade(;
    name,
    available,
    bus,
    active_power,
    reactive_power,
    rating,
    prime_mover,
    active_power_limits,
    reactive_power_limits,
    ramp_limits,
    time_limits,
    base_power,
    operation_cost = PSY.TwoPartCost(0.0, 0.0),
    storage_target = 1.0,
    conversion_factor = 1.0,
    upstream = NamedTuple{(:unit, :lag, :multiplier), Tuple{PSY.HydroGen, Int64, Float64}}[],
    services = PSY.Device[],
    dynamic_injector = nothing,
    ext = Dict{String, Any}(),
    time_series_container = InfrastructureSystems.TimeSeriesContainer(),
    internal = IS.InfrastructureSystemsInternal(),
)
    HydroDispatchCascade(
        name,
        available,
        bus,
        active_power,
        reactive_power,
        rating,
        prime_mover,
        active_power_limits,
        reactive_power_limits,
        ramp_limits,
        time_limits,
        base_power,
        operation_cost,
        storage_target,
        conversion_factor,
        upstream,
        services,
        dynamic_injector,
        ext,
        time_series_container,
        internal,
    )
end

# Constructor for demo purposes; non-functional.
function HydroDispatchCascade(::Nothing)
    HydroDispatchCascade(;
        name = "init",
        available = false,
        bus = PSY.Bus(nothing),
        active_power = 0.0,
        reactive_power = 0.0,
        rating = 0.0,
        prime_mover = PSY.PrimeMovers.HY,
        active_power_limits = (min = 0.0, max = 0.0),
        reactive_power_limits = nothing,
        ramp_limits = nothing,
        time_limits = nothing,
        base_power = 0.0,
        operation_cost = PSY.TwoPartCost(nothing),
        storage_target = 0.0,
        conversion_factor = 0.0,
        upstream = NamedTuple{
            (:unit, :lag, :multiplier),
            Tuple{PSY.HydroGen, Int64, Float64},
        }[],
        services = PSY.Device[],
        dynamic_injector = nothing,
        ext = Dict{String, Any}(),
        time_series_container = InfrastructureSystems.TimeSeriesContainer(),
    )
end

InfrastructureSystems.get_name(value::HydroDispatchCascade) = value.name

PowerSystems.get_available(value::HydroDispatchCascade) = value.available

PowerSystems.get_bus(value::HydroDispatchCascade) = value.bus

PowerSystems.get_active_power(value::HydroDispatchCascade) =
    get_value(value, value.active_power)

PowerSystems.get_reactive_power(value::HydroDispatchCascade) =
    get_value(value, value.reactive_power)

PowerSystems.get_rating(value::HydroDispatchCascade) = get_value(value, value.rating)

PowerSystems.get_prime_mover(value::HydroDispatchCascade) = value.prime_mover

PowerSystems.get_active_power_limits(value::HydroDispatchCascade) =
    get_value(value, value.active_power_limits)

PowerSystems.get_reactive_power_limits(value::HydroDispatchCascade) =
    get_value(value, value.reactive_power_limits)

PowerSystems.get_ramp_limits(value::HydroDispatchCascade) =
    get_value(value, value.ramp_limits)

PowerSystems.get_time_limits(value::HydroDispatchCascade) = value.time_limits

PowerSystems.get_base_power(value::HydroDispatchCascade) = value.base_power

PowerSystems.get_operation_cost(value::HydroDispatchCascade) = value.operation_cost

PowerSystems.get_storage_target(value::HydroDispatchCascade) = value.storage_target

PowerSystems.get_conversion_factor(value::HydroDispatchCascade) = value.conversion_factor
"""Get [`HydroDispatchCascade`](@ref) `upstream`."""
get_upstream(value::HydroDispatchCascade) = value.upstream

PowerSystems.get_services(value::HydroDispatchCascade) = value.services

PowerSystems.get_dynamic_injector(value::HydroDispatchCascade) = value.dynamic_injector

PowerSystems.get_ext(value::HydroDispatchCascade) = value.ext

InfrastructureSystems.get_time_series_container(value::HydroDispatchCascade) =
    value.time_series_container

PowerSystems.get_internal(value::HydroDispatchCascade) = value.internal

InfrastructureSystems.set_name!(value::HydroDispatchCascade, val) = value.name = val

PowerSystems.set_available!(value::HydroDispatchCascade, val) = value.available = val

PowerSystems.set_bus!(value::HydroDispatchCascade, val) = value.bus = val

PowerSystems.set_active_power!(value::HydroDispatchCascade, val) = value.active_power = val

PowerSystems.set_reactive_power!(value::HydroDispatchCascade, val) =
    value.reactive_power = val

PowerSystems.set_rating!(value::HydroDispatchCascade, val) = value.rating = val

PowerSystems.set_prime_mover!(value::HydroDispatchCascade, val) = value.prime_mover = val

PowerSystems.set_active_power_limits!(value::HydroDispatchCascade, val) =
    value.active_power_limits = val

PowerSystems.set_reactive_power_limits!(value::HydroDispatchCascade, val) =
    value.reactive_power_limits = val

PowerSystems.set_ramp_limits!(value::HydroDispatchCascade, val) = value.ramp_limits = val

PowerSystems.set_time_limits!(value::HydroDispatchCascade, val) = value.time_limits = val

PowerSystems.set_base_power!(value::HydroDispatchCascade, val) = value.base_power = val

PowerSystems.set_operation_cost!(value::HydroDispatchCascade, val) =
    value.operation_cost = val

PowerSystems.set_storage_target!(value::HydroDispatchCascade, val) =
    value.storage_target = val

PowerSystems.set_conversion_factor!(value::HydroDispatchCascade, val) =
    value.conversion_factor = val
"""Set [`HydroDispatchCascade`](@ref) `upstream`."""
set_upstream!(value::HydroDispatchCascade, val) = value.upstream = val

PowerSystems.set_services!(value::HydroDispatchCascade, val) = value.services = val

PowerSystems.set_ext!(value::HydroDispatchCascade, val) = value.ext = val

InfrastructureSystems.set_time_series_container!(value::HydroDispatchCascade, val) =
    value.time_series_container = val


using ModernGL, GLFW, GLAbstraction, GLFW, ColorTypes, Reactive, GeometryTypes
import GLAbstraction: N0f8
const GLA = GLAbstraction
const window = GLFW.Window(name="Example")
GLFW.MakeContextCurrent(window)
GLA.set_context!(window)

const vert = GLA.vert"""
#version 150

in vec2 vertices;
in vec2 texturecoordinates; // must be this name, because collect_for_gl assumes them

out vec2 f_uv;
void main() {
    f_uv = texturecoordinates;
    gl_Position = vec4(vertices, 0.0, 1.0);
}
"""

# you can also load the shader from a file, which you can then edit in any editor and the changes will show up in your opengl program.
#using FileIO; prrogram = TemplateProgram(load("path_to_frag.frag"), load("path_to_vert.vert"))
const frag = GLA.frag"""
#version 150

out vec4 outColor;
uniform sampler2D image;
in vec2 f_uv;

void main() {

    outColor = texture(image, f_uv);
}
"""
program = GLA.Program(vert, frag)

tex = GLA.Texture([RGBA{N0f8}(x,y,sin(x*pi), 1.0) for x=0:0.1:1., y=0:0.1:1.]) #automatically creates the correct texture

mesh = GLUVMesh2D(SimpleRectangle{Float32}(-1,-1,2,2))
vao = GLA.VertexArray(GLA.generate_buffers(program, vertices=mesh.vertices, GLA.GEOMETRY_DIVISOR, texturecoordinates = Vec2{Float32}.(mesh.texturecoordinates)), mesh.faces)

glClearColor(0, 0, 0, 1)

while !GLFW.WindowShouldClose(window)
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    GLA.bind(program)
    GLA.gluniform(program, :image, 0, tex)
    GLA.bind(vao)
    GLA.draw(vao)
    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
    if GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS
        GLFW.SetWindowShouldClose(window, true)
    end
end
GLFW.DestroyWindow(window)


module Agent

using JLD, App
include("flex.jl")
include("util.jl")

global vocab = load(joinpath(App.config.server_document_root, "models/vocab_sail.jld"), "vocab")

function loadmodel(fname; flex=true)
    w = Dict()
    d = load(fname, "weights")
    for k in keys(d)
        if flex && startswith(k, "filter")
            w[k] = d[k]
        else
            w[k] = d[k]
        end
    end
    return w
end

global models = [loadmodel(joinpath(App.config.server_document_root, "models/sail_vtest_jelly2_"*"$i"*"_vtest_grid_l.jld")) for i=1:10]
#global models = [loadmodel(joinpath(App.config.server_document_root, "models/sailx_base_"*"$i"*".jld")) for i=1:10]

function predict(instruction, startpos, navimap)
    global models
    global vocab

    ins_text = split(strip(instruction), "_")
    words = ins_arr(vocab, ins_text)
    args = Dict()
    args["hidden"] = 120
    args["embed"] = 120
    args["limactions"] = 35
    args["bs"] = 1
    args["wvecs"] = false
    args["greedy"] = true
    args["percp"] = true
    args["preva"] = true
    args["worldatt"] = 100
    args["attinwatt"] = 0
    args["att"] = false
    args["inpout"] = true
    args["prevaout"] = false
    args["beam"] = true
    args["attout"] = false
    args["beamsize"] = 10
    args["encoding"] = "grid"
    args["atype"] = Array{Float32}

    predict_beam(models, words, navimap, startpos; args=args)
end

end


# define the generative model for the urban pomdp environment

### REWARD MODEL ##################################################################################

function POMDPs.reward(pomdp::UrbanPOMDP, s::UrbanState, a::UrbanAction, sp::UrbanState)
    r = 0.
    ego = sp[findfirst(EGO_ID, sp)]
    if is_crash(sp)
        r += pomdp.collision_cost
    end
    if ego.state.posF.s >= get_end(pomdp.env.roadway[pomdp.ego_goal]) &&
       get_lane(pomdp.env.roadway, ego).tag == pomdp.ego_goal
        r += pomdp.goal_reward
    elseif a.acc > 0.
        r += pomdp.action_cost
    else
        r += pomdp.action_cost
    end
    return r
end

### TERMINAL STATES ###############################################################################

function POMDPs.isterminal(pomdp::UrbanPOMDP, s::UrbanState)
    ego = s[findfirst(EGO_ID, s)]
    return is_crash(s) || (ego.state.posF.s >= get_end(pomdp.env.roadway[pomdp.ego_goal]) &&
                           get_lane(pomdp.env.roadway, ego).tag == pomdp.ego_goal)
end

### TRANSITION MODEL ##############################################################################

function POMDPs.gen(::DDNNode{:sp}, pomdp::UrbanPOMDP, s::UrbanState, a::UrbanAction, rng::AbstractRNG)
    # clean dict
    ids = [veh.id for veh in s]
    key_ = deepcopy(keys(pomdp.models))
    for k in key_
        if !(k ∈ ids) && k!=1
            delete!(pomdp.models, k)
        end
    end
    actions = Array{Any}(undef, length(s))
    pomdp.models[1].a = a
    ego_ind = findfirst(EGO_ID, s)
    is_ego_here =  ego_ind == nothing ? 0 : 1
    sp = deepcopy(s) #XXX bad
    if rand(rng) < pomdp.car_birth && n_cars(s) < pomdp.max_cars + is_ego_here
        new_car = initial_car(pomdp, sp, rng)
        if can_push(pomdp.env, sp, new_car)
            lane = get_lane(pomdp.env.roadway, new_car)
            route = find_route(pomdp.env, random_route(rng, pomdp.env.roadway, lane))
            pomdp.models[new_car.id] = pomdp.car_models[SVector(route[1], route[end])]
            reset_hidden_state!(pomdp.models[new_car.id])
            push!(sp, new_car)
        end
    end
    if rand(rng) < pomdp.ped_birth && n_pedestrians(s) < pomdp.max_peds
        # println("Spawning new pedestrians")
        new_ped = initial_pedestrian(pomdp, sp, rng)
        # pomdp.models[new_ped.id] = ConstantPedestrian(dt = pomdp.ΔT)#TODO parameterized
        new_ped_cw = get_lane(pomdp.env.roadway, new_ped)
        new_ped_conflict_lanes = get_conflict_lanes(new_ped_cw, pomdp.env.roadway)
        pomdp.models[new_ped.id] = IntelligentPedestrian(dt = pomdp.ΔT, crosswalk=new_ped_cw, conflict_lanes=new_ped_conflict_lanes)
        reset_hidden_state!(pomdp.models[new_ped.id])
        push!(sp, new_ped)
    end
    actions = Array{Any}(undef, length(sp))
    get_actions!(actions, sp, pomdp.env.roadway, pomdp.models)
    tick!(sp, pomdp.env.roadway, actions, pomdp.ΔT)
    clean_scene!(pomdp.env, sp)
    return sp
end

function POMDPModelTools.gen(::DDNOut{(:sp, :o, :r, :info)}, p::UrbanPOMDP, s, a, rng::AbstractRNG)
    return gen(DDNOut(:sp, :o, :r), p, s, a, rng)..., deepcopy(p.models)
end

#TODO move to helpers

function AutomotiveDrivingModels.propagate(veh::Vehicle, a::UrbanAction, roadway::Roadway, Δt::Float64)
    acc = LonAccelDirection(a.acc, 4) #TODO Parameterize
    return propagate(veh, acc, roadway, Δt)
end

### INITIAL STATES ################################################################################

function POMDPs.initialstate(pomdp::UrbanPOMDP, rng::AbstractRNG, burn_in::Int64=1)
    scene = initial_scene(pomdp, rng, true)
    for t = 1:burn_in
        scene = gen(DDNOut(:sp), pomdp, scene, UrbanAction(0.), rng)
    end
    clean_scene!(pomdp.env, scene)
    # push! ego after traffic is steady
    ego = initial_ego(pomdp, rng)
    push!(scene, ego)
    clean_scene!(pomdp.env, scene)
    return scene
end

function initial_scene(pomdp::UrbanPOMDP, rng::AbstractRNG, no_ego::Bool=false)
    empty_obstacles!(pomdp.env)
    if pomdp.max_obstacles > 0
        for i=1:pomdp.max_obstacles
            sample_obstacle!(pomdp.env, pomdp.obs_dist, rng)
        end
    end
    scene = Scene()
    if !no_ego
        ego = initial_ego(pomdp, rng)
        push!(scene, ego)
    end

    # add cars
    for i=1:pomdp.max_cars
        if rand(rng) < pomdp.car_birth && n_cars(scene) < pomdp.max_cars + 1-Int(no_ego)
            new_car = initial_car(pomdp, scene, rng, true)
            if can_push(pomdp.env, scene, new_car)
                push!(scene, new_car)
                lane = get_lane(pomdp.env.roadway, new_car)
                route = find_route(pomdp.env, random_route(rng, pomdp.env.roadway, lane))
                pomdp.models[new_car.id] = pomdp.car_models[SVector(route[1], route[end])]
                reset_hidden_state!(pomdp.models[new_car.id])
            end
            clean_scene!(pomdp.env, scene)
        end
    end

    # add pedestrians
    for i=1:pomdp.max_peds
        if rand(rng) < pomdp.ped_birth && n_pedestrians(scene) < pomdp.max_peds # pedestrian appear
            new_ped = initial_pedestrian(pomdp, scene, rng, true)
            # pomdp.models[new_ped.id] = ConstantPedestrian(dt = pomdp.ΔT)#TODO parameterized
            new_ped_cw = get_lane(pomdp.env.roadway, new_ped)
            new_ped_conflict_lanes = get_conflict_lanes(new_ped_cw, pomdp.env.roadway)
            pomdp.models[new_ped.id] = IntelligentPedestrian(dt = pomdp.ΔT, crosswalk=new_ped_cw, conflict_lanes=new_ped_conflict_lanes)
            reset_hidden_state!(pomdp.models[new_ped.id])
            push!(scene, new_ped)
        end
    end
    clean_scene!(pomdp.env, scene)
    return scene
end

function initial_car(pomdp::UrbanPOMDP, scene::Scene, rng::AbstractRNG, first_scene::Bool = false)
    env = pomdp.env

    #velocity
    v0 = rand(rng, 4.0:1.0:7.0) #XXX parameterize!
    s0 = 0.
    if first_scene
        lanes = get_lanes(pomdp.env.roadway)
        lanes = delete!(Set(lanes),pomdp.env.roadway[LaneTag(6,1)])
        lane = rand(rng, lanes) # could be precomputed
        s0 = rand(rng, 0.:0.5:get_end(lane))
    else
        lanes = get_start_lanes(pomdp.env.roadway)
        # lanes = delete!(Set(lanes),pomdp.env.roadway[LaneTag(6,1)])
        lane = rand(rng, lanes) # could be precomputed
    end
    initial_posF = Frenet(lane, s0)
    initialstate = VehicleState(initial_posF, env.roadway, v0)

    id = next_car_id(pomdp, scene)

    return Vehicle(initialstate, pomdp.car_type, id)
end

function initial_ego(pomdp::UrbanPOMDP, rng::AbstractRNG)
    posF = Frenet(get_ego_route(pomdp)[1], pomdp.ego_start)
    state = VehicleState(posF, pomdp.env.roadway, 0.)
    return Vehicle(state, pomdp.ego_type, EGO_ID)
end

"""
    initial_pedestrian(pomdp::UrbanPOMDP, scene::Scene, rng::AbstractRNG)
Create a new pedestrian entity at its initial state
"""
function initial_pedestrian(pomdp::UrbanPOMDP, scene::Scene, rng::AbstractRNG, first_scene::Bool = false)
    env = pomdp.env
    cw_ind = rand(rng, 1:length(env.params.crosswalk_pos))
    crosswalk_pos = env.params.crosswalk_pos[cw_ind]

    # position along the crosswalk
    t0 = rand(rng, Distributions.Uniform(-env.params.crosswalk_width[cw_ind]/2 + 1., env.params.crosswalk_width[cw_ind]/2 - 1.))
    s0 = rand(rng, [0., get_end(env.crosswalks[cw_ind])])
    ϕ0 = float(π)
    if s0 == 0.
        ϕ0 = 0.
    end
    if first_scene
        s0 = rand(rng, Distributions.Uniform(0., get_end(env.crosswalks[cw_ind])))
    end

    #velocity
    v0 = rand(rng, Distributions.Uniform(0., env.params.ped_max_speed))
    posF = Frenet(env.crosswalks[cw_ind], s0, t0, ϕ0)

    ped_initialstate = VehicleState(posF, env.ped_roadway, v0);

    id = next_ped_id(pomdp, scene)

    return Vehicle(ped_initialstate, pomdp.ped_type, id)
end

function POMDPs.initialstate_distribution(pomdp::UrbanPOMDP)
    return GenerativeDist(pomdp)
end

### Observations ##################################################################################

function POMDPs.gen(::DDNNode{:o}, pomdp::UrbanPOMDP, s::UrbanState, a::UrbanAction, sp::UrbanState, rng::AbstractRNG)
    if pomdp.lidar
        return measure_lidar(pomdp, s, a, sp, rng)
    else
        return measure_gaussian(pomdp, s, a, sp, rng)
    end
end

function measure_gaussian(pomdp::UrbanPOMDP, s::UrbanState, a::UrbanAction, sp::UrbanState, rng::AbstractRNG)
    ego = get_ego(sp)
    obs = measure(pomdp.sensor, ego, sp, pomdp.env.roadway, pomdp.env.obstacles)
    o = obs_to_array(pomdp, ego, obs)
    return o
end

function obs_to_array(pomdp::UrbanPOMDP, ego_veh::Vehicle, obs::Vector{Vehicle})
    sorted_vehicles = sort!(obs, by=x->x.id)
    n_features = 4
    n_obstacles = pomdp.max_obstacles
    o = zeros(n_features*(pomdp.max_cars + pomdp.max_peds + n_obstacles + 1))
    ego = ego_veh.state
    o[1] = ego.posG.x
    o[2] = ego.posG.y
    o[3] = ego.posG.θ
    o[4] = ego.v
    pos_off = get_off_the_grid(pomdp)
    for i=2:pomdp.max_cars + pomdp.max_peds + n_obstacles + 1
        o[n_features*i - 3] = pos_off.posG.x - ego.posG.x
        o[n_features*i - 2] = pos_off.posG.y - ego.posG.y
        o[n_features*i - 1] = pos_off.posG.θ
        o[n_features*i] = 0.
    end
    for veh in sorted_vehicles
        if veh.id == EGO_ID
            continue
        end
        if veh.def.class == AgentClass.CAR
            @assert veh.id <= pomdp.max_cars+1 "$(veh.id)"
            o[n_features*veh.id - 3] = veh.state.posG.x - ego.posG.x
            o[n_features*veh.id - 2] = veh.state.posG.y - ego.posG.y
            o[n_features*veh.id - 1] = veh.state.posG.θ
            o[n_features*veh.id] = veh.state.v
        end
        if veh.def.class == AgentClass.PEDESTRIAN
            o[n_features*(veh.id - 100 + pomdp.max_cars + 1) - 3] = veh.state.posG.x - ego.posG.x
            o[n_features*(veh.id - 100 + pomdp.max_cars + 1) - 2] = veh.state.posG.y - ego.posG.y
            o[n_features*(veh.id - 100 + pomdp.max_cars + 1) - 1] = veh.state.posG.θ
            o[n_features*(veh.id - 100 + pomdp.max_cars + 1)] = veh.state.v
        end
    end
    for (j,obs) in enumerate(pomdp.env.obstacles)
        o[n_features*(j + pomdp.max_cars + pomdp.max_peds + 1) - 3] = get_width(obs)
        o[n_features*(j + pomdp.max_cars + pomdp.max_peds + 1) - 2] = get_height(obs)
        o[n_features*(j + pomdp.max_cars + pomdp.max_peds + 1) - 1] = get_center(obs).x - ego.posG.x
        o[n_features*(j + pomdp.max_cars + pomdp.max_peds + 1)] = get_center(obs).y - ego.posG.y
    end
    return rescale!(o, pomdp)
end

function POMDPs.convert_s(::Type{Vector{Float64}}, s::UrbanState, pomdp::UrbanPOMDP)
    ego = get_ego(s)
    obs = [veh for veh in s if veh.id != ego]
    svec = obs_to_array(pomdp, ego, obs)
    return svec
end

function n_dims(pomdp::UrbanPOMDP)
    n_features = 4
    n_obstacles = pomdp.max_obstacles
    return n_features*(pomdp.max_cars + pomdp.max_peds + n_obstacles + 1)
end

function POMDPs.initialobs(pomdp::UrbanPOMDP, s::UrbanState, rng::AbstractRNG)
    return gen(DDNNode(:o), pomdp, s, UrbanAction(0.), s, rng::AbstractRNG)
end

function rescale!(o::UrbanObs, pomdp::UrbanPOMDP)
    # rescale
    n_features = 4
    n_obstacles = pomdp.max_obstacles
    max_ego_dist = get_end(pomdp.env.roadway[pomdp.ego_goal])
    o[1] /= max_ego_dist
    o[2] /= max_ego_dist
    o[3] /= pi/2
    o[4] /= pomdp.env.params.speed_limit
    for i=2:pomdp.max_cars + pomdp.max_peds + 1
        o[n_features*i - 3] /= max_ego_dist
        o[n_features*i - 2] /= max_ego_dist
        o[n_features*i - 1] /= pi
        o[n_features*i] /= pomdp.env.params.speed_limit # XXX parameterized
    end
    for i=pomdp.max_cars + pomdp.max_peds + 2:pomdp.max_cars + pomdp.max_peds + 1 + n_obstacles
        o[n_features*i - 3] /= max_ego_dist
        o[n_features*i - 2] /= max_ego_dist
        o[n_features*i - 1] /= max_ego_dist
        o[n_features*i] /= max_ego_dist
    end

    return o
end

function unrescale!(o::UrbanObs, pomdp::UrbanPOMDP)
    # unrescale
    n_features = 4
    n_obstacles = pomdp.max_obstacles
    max_ego_dist = get_end(pomdp.env.roadway[pomdp.ego_goal])
    o[1] *= max_ego_dist
    o[2] *= max_ego_dist
    o[3] *= pi/2
    o[4] *= pomdp.env.params.speed_limit
    for i=2:pomdp.max_cars + pomdp.max_peds + 1
        o[n_features*i - 3] *= max_ego_dist
        o[n_features*i - 2] *= max_ego_dist
        o[n_features*i - 1] *= pi
        o[n_features*i] *= pomdp.env.params.speed_limit # XXX parameterized
    end
    for i=pomdp.max_cars + pomdp.max_peds + 2:pomdp.max_cars + pomdp.max_peds + 1 + n_obstacles
        o[n_features*i - 3] *= max_ego_dist
        o[n_features*i - 2] *= max_ego_dist
        o[n_features*i - 1] *= max_ego_dist
        o[n_features*i] *= max_ego_dist
    end

    return o
end

function to_global_coordinates!(o::UrbanObs, pomdp::UrbanPOMDP)
    n_features = 4
    n_obstacles = pomdp.max_obstacles
    ego = o[1:n_features]
    for i=2:pomdp.max_cars + pomdp.max_peds + 1
        o[n_features*i - 3] += ego[1]
        o[n_features*i - 2] += ego[2]
    end
    for i=pomdp.max_cars + pomdp.max_peds + 2:pomdp.max_cars + pomdp.max_peds + 1 + n_obstacles
        o[n_features*i - 1] += ego[1]
        o[n_features*i]+= ego[2]
    end
    return o
end

function get_normalized_absent_state(pomdp::UrbanPOMDP, ego::Vector{Float64})
    ego_x, ego_y, theta, v = ego
    pos_off =  get_off_the_grid(pomdp)
    max_ego_dist = get_end(pomdp.env.roadway[pomdp.ego_goal])
    return [pos_off.posG.x/max_ego_dist - ego_x,
                    pos_off.posG.y/max_ego_dist - ego_y,
                    pos_off.posG.θ/float(pi),
                    0. ]
end

function measure_lidar(pomdp::UrbanPOMDP, s::UrbanState, a::UrbanAction, sp::UrbanState, rng::AbstractRNG)
    observe!(pomdp.sensor, sp, pomdp.env.roadway, EGO_ID)
    return vcat(pomdp.sensor.ranges/pomdp.sensor.max_range, pomdp.sensor.range_rates/pomdp.env.params.speed_limit)
end

function POMDPs.convert_o(::Type{Vector{Float64}}, o::UrbanObs, pomdp::UrbanPOMDP)
    return o
end

function POMDPs.gen(::DDNOut{(:sp, :o, :r, :info)}, p::UrbanPOMDP, s::UrbanState, a::UrbanAction, rng::AbstractRNG)
    if p.lidar
        return gen(DDNOut(:sp,:o,:r), p, s, a, rng)..., p.sensor
    else
        return gen(DDNOut(:sp,:o,:r), p, s, a, rng)..., nothing
    end
end

# uncomment for Lidar
# uncomment for LIDAR observation

#
# function POMDPs.generate_o(pomdp::OCPOMDP, s::OCState, rng::AbstractRNG)
#     observe!(pomdp.sensor, s, pomdp.env.roadway, EGO_INDEX)
#     return pomdp.sensor
# end
#
# function POMDPs.convert_o(::Type{Vector{Float64}}, o::OCObs, pomdp::OCPOMDP)
#     return [o.ranges/pomdp.sensor.max_range, o.range_rates/pomdp.env.params.speed_limit]
# end

#### Helpers
# """
#     ego_state(pomdp:ICPOMDP, s::Float64, v::Float64)
# Helper that returns the ego car state given its position on the lane and the velocity
# """
# function ego_state(pomdp::UrbanPOMDP, s::Float64, v::Float64)
#     # ego_lane = pomdp.env.lane_map["ego_left"] #TODO parameterized
#     posF = Frenet(get_ego_route(pomdp)[1], s)
#     return VehicleState(posF, pomdp.env.roadway, v)
# end

"""
    car_state(pomdp::UrbanPOMDP, lane_id::String, s::Float64, v::Float64)
    car_state(pomdp::UrbanPOMDP, lane::Lane, s::Float64, v::Float64)
    car_state(pomdp::UrbanPOMDP, x::Float64, y::Float64, θ::Float64, v::Float64)
Helper that returns the state of a car given the lane and the position on the lane
"""
function car_state(pomdp::UrbanPOMDP, lane_id::String, s::Float64, v::Float64)
    lane = pomdp.env.lane_map[lane_id]
    posF = Frenet(lane, s)
    return VehicleState(posF, pomdp.env.roadway, v)
end
function car_state(pomdp::UrbanPOMDP, lane::Lane, s::Float64, v::Float64)
    posF = Frenet(lane, s)
    return VehicleState(posF, pomdp.env.roadway, v)
end
function car_state(pomdp::UrbanPOMDP, x::Float64, y::Float64, θ::Float64, v::Float64)
    posG = VecSE2(x, y, θ)
    return VehicleState(posG, pomdp.env.roadway, v)
end

"""
    get_off_the_grid(pomdp::UrbanPOMDP)
return the off the grid state
"""
function get_off_the_grid(pomdp::UrbanPOMDP)
    posG = pomdp.off_grid
    return VehicleState(posG, pomdp.env.roadway, 0.)
end

function off_the_grid(veh::VehicleState, pomdp::UrbanPOMDP)
    return veh.posG == pomdp.off_grid
end

"""
return the list of lanes the ego car should take
"""
function get_ego_route(pomdp::UrbanPOMDP)
    tags = [LaneTag(6,1), LaneTag(13, 1), LaneTag(2,pomdp.env.params.nlanes_main)]
    lanes = Array{Lane}(undef, length(tags))
    for (i,tag) in enumerate(tags)
        lanes[i] = pomdp.env.roadway[tag]
    end
    lanes
end

"""
create a unique ID for a new car
"""
function next_car_id(pomdp::UrbanPOMDP, scene::Scene, rng::AbstractRNG)
    max_id = pomdp.max_cars+1
    current_car_ids = []
    for veh in scene
        if veh.def.class == AgentClass.CAR
            push!(current_car_ids, veh.id)
        end
    end
    possible_ids = setdiff(2:max_id, current_car_ids)
    id = rand(rng, possible_ids)
    return id
end

function next_car_id(pomdp::UrbanPOMDP, scene::Scene)
    max_id = pomdp.max_cars+1
    current_car_ids = []
    for veh in scene
        if veh.def.class == AgentClass.CAR
            push!(current_car_ids, veh.id)
        end
    end
    possible_ids = setdiff(2:max_id, current_car_ids)
    return possible_ids[1]
end

"""
create a unique ID for a new pedestrian
"""
function next_ped_id(pomdp::UrbanPOMDP, scene::Scene, rng::AbstractRNG)
    max_id = pomdp.max_peds
    current_ped_ids = []
    for veh in scene
        if veh.def.class == AgentClass.PEDESTRIAN
            push!(current_ped_ids, veh.id)
        end
    end
    possible_ids = setdiff(101:100+max_id, current_ped_ids)
    id = rand(rng, possible_ids)
    return id
end
function next_ped_id(pomdp::UrbanPOMDP, scene::Scene)
    max_id = pomdp.max_peds
    current_ped_ids = []
    for veh in scene
        if veh.def.class == AgentClass.PEDESTRIAN
            push!(current_ped_ids, veh.id)
        end
    end
    possible_ids = setdiff(101:100+max_id, current_ped_ids)
    return possible_ids[1]
end

function n_cars(scene::Scene)
    n = 0
    for veh in scene
        if veh.def.class == AgentClass.CAR
            n += 1
        end
    end
    return n
end

function n_pedestrians(scene::Scene)
    n = 0
    for veh in scene
        if veh.def.class == AgentClass.PEDESTRIAN
            n += 1
        end
    end
    return n
end

"""
Create a scene that can be rendered from an observation
"""
function obs_to_scene(pomdp::UrbanPOMDP, obs::UrbanObs)
    o = deepcopy(obs)
    o = unrescale!(o, pomdp)
    o = to_global_coordinates!(o, pomdp)
    scene = Scene()
    # extract
    ego, car_map, ped_map, obs_map = split_o(o, pomdp, false)
    # println("cars observed : ", keys(car_map), "ped observed :", keys(ped_map))
    # project to roadway
    params = pomdp.env.params
    intersection_params = TInterParams(params.x_min, params.x_max, params.y_min, params.inter_x,
                                    params.inter_y, params.inter_width, params.inter_height,
                                    params. lane_width, params.nlanes_main, params.nlanes,
                                    params. stop_line, params.speed_limit, params.car_rate)
    car_roadway = gen_T_roadway(intersection_params)
    ego_state = VehicleState(VecSE2(ego[1], ego[2], ego[3]),  car_roadway, ego[4])
    ego = Vehicle(ego_state, pomdp.ego_type, EGO_ID)
    push!(scene, ego)
    for (str_id, vec_state) in car_map
        # make sure to not project it on a crosswalk
        car_state = VehicleState(VecSE2(vec_state[1], vec_state[2], vec_state[3]), car_roadway, vec_state[4])
        if !off_the_grid(car_state, pomdp)
            id = next_car_id(pomdp, scene)
            car = Vehicle(car_state, pomdp.car_type, id)
            push!(scene, car)
        end
    end
    for (str_id, vec_state) in ped_map
        ped_state = VehicleState(VecSE2(vec_state[1], vec_state[2], vec_state[3]), pomdp.env.ped_roadway, vec_state[4])
        # println("ped_state ", ped_state)
        if !off_the_grid(ped_state, pomdp)
            id = next_ped_id(pomdp, scene)
            ped = Vehicle(ped_state, pomdp.ped_type, id)
            push!(scene, ped)
        end
    end
    @assert issorted([veh.id for veh in scene])
    return scene
end

# given a big observation vector, split to an entity-wise representation
function split_o(obs::UrbanObs, pomdp::UrbanPOMDP, normalized=true)
    car_map, ped_map, obs_map = OrderedDict{Int64, Vector{Float64}}(), OrderedDict{Int64, Vector{Float64}}(), OrderedDict{Int64, Vector{Float64}}() #XXX Dictionary creation is sloooooow
    n_features = pomdp.n_features
    ego = obs[1:n_features]
    if normalized
        absent = normalized_off_the_grid_pos(pomdp, ego[1], ego[2])
    else
        pos_off = get_off_the_grid(pomdp)
        absent= [pos_off.posG.x, pos_off.posG.y, pos_off.posG.θ, pos_off.v]
    end
    # println("ego ", ego)
    # println("absent state :", absent)
#     println("ego idx ", 1, ":", n_features)
    for (j,i) in enumerate(1:pomdp.max_cars)
        if !isapprox(obs[i*n_features+1:(i+1)*n_features], absent, atol=1.0)
            car_map[j+1] = obs[i*n_features+1:(i+1)*n_features]
        end
#         println("car$j idx ", i*n_features+1:(i+1)*n_features)
    end
    for (j,i) in enumerate(pomdp.max_cars + 1:pomdp.max_cars + pomdp.max_peds)
        if !isapprox(obs[i*n_features+1:(i+1)*n_features], absent, atol=1.0)
            # println(norm(obs[i*n_features+1:(i+1)*n_features] - absent))
            ped_map[100+j] = obs[i*n_features+1:(i+1)*n_features]
        end
#         println("ped$j idx ", i*n_features+1:(i+1)*n_features)
    end
    for (j,i)  in enumerate(pomdp.max_cars + pomdp.max_peds + 1:pomdp.max_cars + pomdp.max_peds + pomdp.max_obstacles)
        obs_map[j] = obs[i*n_features+1:(i+1)*n_features]
#         println("obs$j idx ", i*n_features+1:(i+1)*n_features)
    end
    return ego, car_map, ped_map, obs_map
end

function normalized_off_the_grid_pos(pomdp::UrbanPOMDP,normalized_ego_x::Float64, normalized_ego_y::Float64)
    pos_off = get_off_the_grid(pomdp)
    max_ego_dist = get_end(pomdp.env.roadway[pomdp.ego_goal])
    return [pos_off.posG.x/max_ego_dist - normalized_ego_x, pos_off.posG.y/max_ego_dist - normalized_ego_y, pos_off.posG.θ, 0.]
end


#Framework
include("./storage.jl")
include("./optimizer.jl")
include("./shuffler.jl")

#Problem specific
include("./problems/EuclideanTSP.jl")

module Ma

using Printf
import Base
import ProgressMeter
using ..Storage
using RandomNumbers

function runBlock(code, maxIter, maxTime, momentum=0.9, updateInterval=1)
    if maxIter isa Number || maxTime isa Number
        p = ProgressMeter.Progress(1000)
    else
        p = ProgressMeter.ProgressUnknown()
    end

    startTime = time()
    estimatedThroughput = 1.0
    iterationsDone = 0
    lastProgress = 0
    currentProgress = 0.0
    while true
        if (maxTime isa Number) && time() - startTime > maxTime
            break
        elseif (maxIter isa Number ) && iterationsDone >= maxIter
            break
        end

        iterations = max(1, floor(Int, updateInterval * estimatedThroughput))

        if maxIter isa Number
            iterationsLeft = maxIter - iterationsDone
            iterations = min(iterations, iterationsLeft)
        end

        before = time()
        infos = code(iterations)
        timeTaken = time() - before

        iterationsDone += iterations
        lastThroughput = iterations / timeTaken
        estimatedThroughput = (momentum * estimatedThroughput) + (1 - momentum) * lastThroughput

        ## Reporting
        progresses = Vector{Float64}()

        if maxIter isa Number
            push!(progresses, iterationsDone / maxIter)
        end

        if maxTime isa Number
            push!(progresses, (time() - startTime) / maxTime)
        end

        if length(progresses) > 0
            currentProgress = maximum(progresses) * 1000
            if currentProgress > lastProgress
                improvement = floor(Int, currentProgress - lastProgress)
                lastProgress = currentProgress
                ProgressMeter.next!(p, step=improvement)
            end
        else
            ProgressMeter.next!(p, step=iterations)
        end

        infos = Dict(pairs(infos))
        infos[:t] = @sprintf("%i/s", (iterationsDone) / (time() - startTime))

        strings = map(collect(infos)) do (key, value)
            "$key=$value"
        end
        p.desc = "Optimizing ($(join(strings, ", "))) "
    end

    ProgressMeter.finish!(p)
    println()

    return (iterationsDone, time() - startTime)
end

mutable struct Manager
    problem
    optimizer
    shuffler
    losses::Vector{Float32}
    storage::Dict{Symbol, Any}
    debug::Bool

    function Manager()
        new(missing, missing, missing, Vector{Float32}(), Dict{Symbol, Any}(), false)
    end
end

function allocateStorage!(m::Manager)
    resize!(m.losses, m.shuffler.popSize)
    m.storage[:stateStorage] = Storage.allocateCPU(m.problem.solType,
                                                   (m.shuffler.popSize,
                                                    m.optimizer.solutionStates, ))
    m.storage[:dataStorage] = Storage.allocateCPU(m.problem.dataType, (1, ))[1]
    if !isnothing(m.optimizer.stateType)
        m.storage[:optimStorage] = Storage.allocateCPU(m.optimizer.stateType,
                                                       (m.shuffler.popSize,))
    end
end

function init!(m::Manager)
    Storage.initFromFields(m.storage[:dataStorage], m.problem.data...)

    # We only init the best state,
    for i in 1:m.shuffler.popSize
        m.problem.init(m.storage[:stateStorage][i, 1], m.storage[:dataStorage])
        if !isnothing(m.optimizer.stateType) && !isnothing(m.optimizer.init)
            m.optimizer.init(m.storage[:optimStorage][i])
        end
    end

end

function run!(m::Manager, maxIter=nothing, maxTime=nothing)

    optimStorage = nothing

    runBlock(maxIter, maxTime) do blocksIter
        Base.Threads.@threads for i in 1:m.shuffler.popSize
            rng_xor = RandomNumbers.Xorshifts.Xoroshiro128Star()
            if haskey(m.storage, :optimStorage)
                optimStorage = m.storage[:optimStorage][i]
            end
            state = m.storage[:stateStorage][i]
            loss = m.optimizer.optimize(optimStorage, state,
                                        m.storage[:dataStorage], blocksIter, rng_xor)
            m.losses[i] = loss
        end
        return (
                loss=@sprintf("%.3f - %3f", minimum(m.losses), maximum(m.losses)),
               )
    end
end

function bestSolution(m::Manager)
    bestIx = argmin(m.losses)
    return (m.losses[bestIx], m.storage[:stateStorage][bestIx][1])
end

export Manager

end

import .Ma
import .TSP

import .Storage
import .Optim
import .Shufflers

# Code generation

# data = TSP.readTourFile("./pla85900.tsp")
data = TSP.generateRandomTSPData(10000)
manager = Ma.Manager()
manager.problem = TSP.generateProblemForData(data)
manager.optimizer = Optim.ExhaustiveLocalSearch(manager.problem)
manager.shuffler = Shufflers.Independent(64)
Ma.allocateStorage!(manager)
Ma.init!(manager)
Ma.run!(manager, 10000000, 6)
loss, solution = Ma.bestSolution(manager)
print("LOSS: ", loss)


# Functions for large C4 correlation job
export correlate_pair, autocorrelate, diag_chunks, offdiag_chunks, preprocess2, get_blocks,
        stack_auto, stack_corr, correlate_big, stack_all

# correlation helper functions
function correlate_pair(src::FFTData, rec::FFTData, pref::String="CORR", params::Dict=params)
    try
        C = correlate(src, rec, params["maxlag"])
        cc_medianmute!(C, 10.) # remove correlation windows with high noise
        stack!(C)
        name = join(split(name_corr(C), ".")[1:2],".")
        save_named_corr(C,"$(params["rootdir"])/$pref/$name/$(C.comp)")
    catch e
        println(e)
    end
end

function autocorrelate(station::String, list::Array{String,1}=fft_list_100, params::Dict=params)
    """ Filter channels by station, load and correlate """
    channels = filter(x -> occursin(station, x), list)
    ffts = map(x -> load_fft(x, string(x[end-7:end-5])), channels)
    pairs = vec([collect(x) for x in Iterators.product(1:length(ffts),1:length(ffts))])
    map(pair -> correlate_pair(ffts[pair[1]], ffts[pair[2]], "AUTOCORR", params), pairs)
end

function diag_chunks(chunk::Array{String,1}, prefix::String = "CORR", filt_dist::Bool = true, params::Dict=params)
    ffts  = map(x -> load_fft(x, string(x[end-7:end-5])), chunk)
    pairs = vec([collect(x) for x in Iterators.product(1:length(ffts),1:length(ffts))])
    # maybe remove this filter so we get all components
    #filter!(x -> (x[1] < x[2]) || (x[1]>x[2]), pairs) # filter upper diagonal
    if filt_dist; filter!(x -> get_dist(ffts[x[1]], ffts[x[2]]) <= 300., pairs); end # filter distances
    map(pair -> correlate_pair(ffts[pair[1]], ffts[pair[2]], prefix, params), pairs)
end

function offdiag_chunks(chunkpair::Array{Array{String, 1}}, prefix::String = "CORR", filt_dist::Bool = true,
                        params::Dict=params)
    sources   = map(x -> load_fft(x, string(x[end-7:end-5])), chunkpair[1]) # load
    receivers = map(x -> load_fft(x, string(x[end-7:end-5])), chunkpair[2])
    pairs = vec([collect(x) for x in Iterators.product(1:length(sources),1:length(receivers))]) # all pairs src -> rec
    if filt_dist; filter!(x -> get_dist(sources[x[1]].loc, receivers[x[2]].loc) <= 300., pairs); end # filt distances
    map(pair -> correlate_pair(sources[pair[1]], receivers[pair[2]], prefix, params), pairs)
end

function get_blocks(paths::Array{String,1}, params::Dict=params)
    # Chunk raw ffts into blocks. 30 is best size for IO, but distribute across cores if more cores available
    chunks = collect(Iterators.partition(paths, 30))#convert(Int64, minimum([30, ceil(length(paths)/params["num_procs"])]))))
    off_chunks = vec([collect(x) for x in Iterators.product(1:length(chunks),1:length(chunks))])
    filter!(chunk_pair -> chunk_pair[1] < chunk_pair[2], off_chunks)
    # map chunk indices to filenames
    off_chunk_names = map(chunk -> [collect(chunks[chunk[1]]), collect(chunks[chunk[2]])], off_chunks)
    return chunks, off_chunk_names
end

# preprocessing
function preprocess2(file::String, accelerometer::Bool=false, params::Dict=params, path::String=path)
    """
    Load raw seisdata file and process, saving to fft
    """
    try
        data = SeisData()
        if occursin(".ms", file)
            read_data!(data, "mseed", file)
        else # data is SeisIO native from NCEDC/IRIS
            read_data!(data, "seisio", file)
        end
        gaps = size(data.t[1])[1] # N-2 gaps (eg gaps = 12 tests for 10 gaps)
        pts = size(data.x[1])[1]
        fs_temp = data.fs[1]
        windows = u2d.(SeisIO.t_win(data.t[1], data.fs[1]) .* 1e-6)
        startend = windows[:,2] .- windows[:,1] .> Second(params["cc_len"])
        if gaps < 25 && any(startend) # If few gaps and sufficient (2/3) data present, data is clean
            try
                add_location(data, params["all_stations"])
                for samp_rate in params["samp_rates"]
                    S = SeisData()
                    try
                        S = process_raw(data, samp_rate)
                    catch e# trying to sample non 100Hz data at 100 Hz - skip resampling
                        println(e)
                        S = process_raw(data, data.fs[1])
                    end
                    R = RawData(S,params["cc_len"],params["cc_step"])
                    SeisNoise.detrend!(R)
                    bandpass!(R,params["freqmin"],params["freqmax"],zerophase=true)
                    # SMOOTHING LINE GOES HERE
                    SeisNoise.taper!(R)
                    FFT = nothing
                    if accelerometer # accelerometer - so we integrate
                        FFT = rfft_raw(R,1)
                    else # then its a seismometer
                        FFT = compute_fft(R)
                    end
                    coherence!(FFT,params["half_win"], params["water_level"])
                    try # save fft
                        root_fft = "ffts/$path/$(Int(samp_rate))/"
                        save_fft(FFT, joinpath(params["rootdir"], root_fft))
                    catch e
                        println(e)
                    end
                    #println("Successfully processed $(data.id[1])")
                end
                return data.id
            catch e
                println(e)
            end
        end
    catch e
        println(e)
    end
end

# stacking codes
function stack_auto(name::String, startdate::Date=startdate)
    month = Dates.monthname(startdate) # get month for filename
    yr = Dates.year(startdate)
    # stack autocorrelations

    CORROUT = expanduser(joinpath(params["rootdir"],"autocorrelations/$name/"))
    if !isdir(CORROUT)
        mkpath(CORROUT)
    end
    filename = joinpath(CORROUT,"$(yr)_$(month)_$name.h5") # get output filename

    # Get list of files to save
    autocorr_list = glob("AUTOCORR/$name*/*/*.jld2", params["rootdir"])
    components = ["EE", "EN", "EZ", "NE", "NN", "NZ", "ZE", "ZN", "ZZ"]

    C = load_corr(autocorr_list[1], convert(String, split(autocorr_list[1],"/")[end-1])) # sample autocorr for meta
    h5open(filename, "cw") do file
        if !haskey(read(file), "meta")
            write(file, "meta/corr_type", C.corr_type)
            write(file, "meta/cc_len", C.cc_len)
            write(file, "meta/cc_step", C.cc_step)
            write(file, "meta/whitened", C.whitened)
            write(file, "meta/time_norm", C.time_norm)
            write(file, "meta/notes", C.notes)
            write(file, "meta/maxlag", C.maxlag)
            write(file, "meta/starttime", Dates.format(u2d(C.t[1]), "yyyy-mm-dd HH:MM:SS"))
            write(file, "meta/samp_freq", C.fs)
            write(file, "meta/lat", C.loc.lat)
            write(file, "meta/lon", C.loc.lon)
            write(file, "meta/el", C.loc.el)
            write(file, "meta/dep", C.loc.dep)
            write(file, "meta/az", C.loc.az)
            write(file, "meta/inc", C.loc.inc)
        end
        # stack per month
        for comp in components
            #comp_files = filter(f -> convert(String, split(f, "/")[end-1]) == comp, autocorr_list)
            comp_files = glob("AUTOCORR/$name*/$comp/*.jld2", params["rootdir"]) # uni-component files
            autocorrs = [load_corr(f, comp) for f in comp_files]

            autocorr_mean = SeisNoise.stack(sum(autocorrs), allstack=true, stacktype=mean)
            autocorr_pws = SeisNoise.stack(sum(autocorrs), allstack=true, stacktype=pws)
            autocorr_robust = SeisNoise.stack(sum(autocorrs), allstack=true, stacktype=robuststack)
            # save into file
            write(file, "$rec/$comp/linear", autocorr_mean.corr[:])
            write(file, "$rec/$comp/pws", autocorr_pws.corr[:])
            write(file, "$rec/$comp/robust", autocorr_robust.corr[:])
        end

        # Following code saves each correlation individually
        # for (ind, name) in enumerate(autocorr_list)
        #     # load file
        #     comp = string(split(name, "/")[end-1]) # get component
        #     Cl = load_corr(name, comp)
        #     # write
        #     write(file, "$comp/$(Cl.id)", Cl.corr[:])
        # end
    end
end

function stack_corr(name::String, startdate::Date=startdate, prefix::String = "CORR_20HZ")
    month = Dates.monthname(startdate) # get month for filename
    yr = Dates.year(startdate)
    # stack autocorrelations

    CORROUT = expanduser(joinpath(params["rootdir"], "correlations/$name/"))
    if !isdir(CORROUT)
        mkpath(CORROUT)
    end
    filename = joinpath(CORROUT,"$(yr)_$(month)_$name.h5") # get output filename
    # r = r"(?<=$(name)).*"
    components = ["EE","EN","EZ", "NE", "NN","NZ", "ZE", "ZN", "ZZ"]

    receivers = glob("$prefix/$name/*/*", params["rootdir"])
    receivers = Set([join(split(x, ".")[end-2:end-1], ".") for x in receivers]) # just get reciever names
    println(receivers)
    corr_list = glob("$prefix/$name*/*/*", params["rootdir"])
    C = load_corr(corr_list[1], convert(String, split(corr_list[1],"/")[end-1]))
    source_loc = GeoLoc(lat = C.loc.lat, lon = C.loc.lon, el = C.loc.el)
    h5open(filename, "cw") do file
        if !haskey(read(file), "meta")
            write(file, "meta/corr_type", C.corr_type)
            write(file, "meta/cc_len", C.cc_len)
            write(file, "meta/cc_step", C.cc_step)
            write(file, "meta/whitened", C.whitened)
            write(file, "meta/time_norm", C.time_norm)
            write(file, "meta/notes", C.notes)
            write(file, "meta/maxlag", C.maxlag)
            write(file, "meta/starttime", Dates.format(u2d(C.t[1]), "yyyy-mm-dd HH:MM:SS"))
            write(file, "meta/samp_freq", C.fs)
            write(file, "meta/lat", C.loc.lat)
            write(file, "meta/lon", C.loc.lon)
            write(file, "meta/el", C.loc.el)
            write(file, "meta/dep", C.loc.dep)
            write(file, "meta/az", C.loc.az)
            write(file, "meta/inc", C.loc.inc)
        end
        for rec in receivers
            try
                # sample_r = glob("CORR/$tf*$rec*/ZZ/*")[1]
                # Cr = load_corr(sample_r, "ZZ")
                rec_station = split(rec, ".")[2]
                #rec_loc = LLE_geo(rec_station, all_stations)
                rec_loc = load_fft(glob("ffts/*/*/$rec*BHZ*", params["rootdir"])[1], "BHZ").loc # get location from fft - perhaps a faster way to do this.
                if !haskey(read(file), "$rec/meta")
                    write(file, "$rec/meta/lon", rec_loc.lon)
                    write(file, "$rec/meta/lat", rec_loc.lat)
                    write(file, "$rec/meta/el", rec_loc.el)
                    write(file, "$rec/meta/dist", get_dist(source_loc, rec_loc))
                    write(file, "$rec/meta/azi", get_azi(source_loc, rec_loc))
                    write(file, "$rec/meta/baz", get_baz(source_loc, rec_loc))
                end
                for comp in components
                    # load correlations for this receiver by component
                    files = glob("$prefix/$name/$comp/*$rec*", params["rootdir"])
                    corrs = [load_corr(f, comp) for f in files]

                    # implement various stacktypes
                    corr_mean = SeisNoise.stack(sum(corrs), allstack=true, stacktype=mean)
                    corr_pws = SeisNoise.stack(sum(corrs), allstack=true, stacktype=pws)
                    corr_robust = SeisNoise.stack(sum(corrs), allstack=true, stacktype=robuststack)
                    # save into file
                    write(file, "$rec/$comp/linear", corr_mean.corr[:])
                    write(file, "$rec/$comp/pws", corr_pws.corr[:])
                    write(file, "$rec/$comp/robust", corr_robust.corr[:])
                end
            catch e
                println(e)
            end
        end
    end
end

function stack_all(params::Dict=params)
    autocorr_names = [join(split(split(auto, "/")[end], ".")[1:2],".") for auto in glob("AUTOCORR/*", params["rootdir"])]
    @eval @everywhere autocorr_names = $autocorr_names
    Sauto = @elapsed robust_pmap(x -> stack_auto(x, startdate), autocorr_names)

    corr_names = [join(split(split(auto, "/")[end], ".")[1:2],".") for auto in glob("CORR_20HZ/*", params["rootdir"])]
    @eval @everywhere corr_names = $corr_names
    S20 = @elapsed robust_pmap(x -> stack_corr(x, startdate), corr_names)

    corr_names_lf = [join(split(split(auto, "/")[end], ".")[1:2],".") for auto in glob("CORR_1HZ/*", params["rootdir"])]
    @eval @everywhere corr_names_lf = $corr_names_lf
    S1 = @elapsed robust_pmap(x -> stack_corr(x, startdate), corr_names_lf)
    println("Stacking Completed for autocorrs and interstation cross correlations in $(Sauto+S20+S1) seconds")
end

# wrapper function for correlations
function correlate_big(dd::Date, startdate::Date = startdate, params::Dict = params)
    """ Wrapper Function: Computes autocorrelations for a specific day"""
    aws = params["aws"]
    yr = Dates.year(dd)
    path = join([Dates.year(dd),lpad(Dates.dayofyear(dd),3,"0")],"_") # Yeilds "YEAR_JDY"
    @eval @everywhere path = $path

    ###### WRAP THIS MESS #######
    # get filepaths for source stations - would be faster if they're available
    scedc_files = nothing
    if !isdir("root/scedc_path/$yr/"); mkpath("root/scedc_path/$yr/"); end
    try # most filelists are stored on seisbasin
        s3_get_file(aws, "seisbasin", "scedc_path/$yr/$path.csv", "root/scedc_path/$yr/$path.csv")
        scedc_files = DataFrame(CSV.File("root/scedc_path/$yr/$path.csv")).Path
    catch e # in case file not found/present we use SCEDC lookup
        println(e)
        scedc_files = get_scedc_files(dd, aws)
    end
    println(length(scedc_files), " is the length of scedc_files")
    #########################

    # filepaths for iris, ncedc, and node data
    iris_path  = S3Path("s3://seisbasin/iris_waveforms/$yr/$path/", config=aws)
    ncedc_path = S3Path("s3://seisbasin/ncedc_waveforms/$yr/$path/", config=aws)
    node_path  = S3Path("s3://seisbasin/continuous_waveforms/$yr/$path/", config=aws)
    iris_query, ncedc_query, node_query = convert.(String, readdir(iris_path)), convert.(String, readdir(ncedc_path)),
                                        convert.(String, readdir(node_path))
    filelist_basin = vcat(joinpath.("iris_waveforms/$yr/$path/", iris_query), joinpath.("ncedc_waveforms/$yr/$path/", ncedc_query),
                                        joinpath.("continuous_waveforms/$yr/$path/", node_query))
    println("There are $(length(filelist_basin)) node files and $(length(scedc_files)) SCEDC files available for $path.")

    # download scedc and seisbasin data
    ec2download(aws, "scedc-pds", scedc_files, joinpath(params["rootdir"],"data"))
    ec2download(aws, "seisbasin", filelist_basin, joinpath(params["rootdir"],"data"))
    println("Download complete!")

    # preprocess - broadbands and seismometers are processed separately
    allf = glob("data/*/$yr/$path/*", params["rootdir"])
    println(allf)
    println(readdir())
    println(readdir("$(params["rootdir"])/data"))
    accelerometers = filter(x -> isfile(x), joinpath.("/$(params["rootdir"])/data/continuous_waveforms/$yr/$path/", node_query))
    broadbands = setdiff(allf, accelerometers) # get the rest of the files

    @eval @everywhere accelerometers, broadbands = $accelerometers, $broadbands
    T_b = @elapsed robust_pmap(f -> preprocess2(f, false, params, path), broadbands)
    if length(accelerometers) != 0 # save time on pmap allocation overhead
        T_a = @elapsed robust_pmap(f -> preprocess2(f, true, params, path), accelerometers) # integrate accelerometers
        println("Preprocessing Completed in $(T_a+T_b) seconds.")
    else
        println("Preprocessing Completed in $T_b seconds.")
    end

    # autocorrelations
    fft_list_100 = glob("ffts/$path/100/*.jld2", params["rootdir"])
    println("FFT list 100 is $(length(fft_list_100)) stations long")
    fft_100_stations = unique([join(split(elt, ".")[1:2],".") for elt in fft_list_100]) # get stations to iterate over
    @eval @everywhere fft_100_stations = $fft_100_stations
    pmap(x -> autocorrelate(x, fft_list_100, params), fft_100_stations)

    fft_paths_20 = sort(glob("ffts/$path/20/*", params["rootdir"])) # alphabetic sort for all stations so that diagonal files are made correctly
    fft_paths_1 = sort(glob("ffts/$path/1/*", params["rootdir"])) # alphabetic sort for all stations so that diagonal files are made correctly
    println("There are $(length(fft_paths_20)) fft datas for the 20 HZ correlations")
    # run 20 HZ correlations
    chunks_20HZ, off_chunk_names_20HZ = get_blocks(fft_paths_20, params);

    # correlate diagonal chunks
    @eval @everywhere chunks_20HZ = $chunks_20HZ
    T20D = @elapsed robust_pmap(chunk -> diag_chunks(convert(Array,chunk), "CORR_20HZ", true, params), chunks_20HZ)

    # correlate off-diagonal chunks
    @eval @everywhere off_chunk_names_20HZ = $off_chunk_names_20HZ
    T20O = @elapsed robust_pmap(chunk -> offdiag_chunks(chunk, "CORR_20HZ", true, params), off_chunk_names_20HZ) # run mega correlations

    # run low freq correlations
    chunks_1HZ, off_chunk_names_1HZ = get_blocks(fft_paths_1, params);

    # correlate diagonal chunks
    @eval @everywhere chunks_1HZ = $chunks_1HZ
    T1D = @elapsed robust_pmap(chunk -> diag_chunks(convert(Array,chunk), "CORR_1HZ", false, params), chunks_1HZ)

    # correlate off-diagonal chunks
    @eval @everywhere off_chunk_names_1HZ = $off_chunk_names_1HZ
    T1O = @elapsed robust_pmap(chunk -> offdiag_chunks(chunk, "CORR_1HZ", false, params), off_chunk_names_1HZ) # run mega correlations

    # Correlation summary
    println("All $(length(glob("CORR_*/*/*/$path*", params["rootdir"]))) Inter-station Correlations computed in $(T20D + T20O + T1D + T1O) seconds")
    try; rm("$(params["rootdir"])/data/continuous_waveforms/", recursive=true); catch e; println(e); end
end


"""
`module  JAC.NonresonantTwoColorPhotonIon`  ... a submodel of JAC that contains all methods for computing nonresonant two-color
                                                (two-photon & single-electron) ionization cross sections; it is using JAC.
"""
module NonresonantTwoColorPhotonIon

    using JAC, JAC.ManyElectron

    """
    `struct  NonresonantTwoColorPhotonIon.Settings`  ... defines a type for the settings in estimating resonant two-color (two-photon &
                                                         single-electron) ionization cross sections.

        + multipoles              ::Array{EmMultipole}           ... Specifies the multipoles of the radiation field that are to be included.
        + gauges                  ::Array{UseGauge}              ... Specifies the gauges to be included into the computations.
        + photonEnergies          ::Array{Float64,1}             ... List of photon energies.
        + printBeforeComputation  ::Bool                         ... True, if all energies and lines are printed before their evaluation.
        + selectLines             ::Bool                         ... True, if lines are selected individually for the computations.
        + selectedLines           ::Array{Tuple{Int64,Int64},1}  ... List of lines, given by tupels (inital-level, final-level).

    """
    struct Settings
        multipoles                ::Array{EmMultipole}
        gauges                    ::Array{UseGauge}
        photonEnergies            ::Array{Float64,1}
        printBeforeComputation    ::Bool
        selectLines               ::Bool
        selectedLines             ::Array{Tuple{Int64,Int64},1}
    end

    """
    `JAC.NonresonantTwoColorPhotonIon.Settings()`  ... constructor for the default values of resonant two-color (two-photon &
                                                       single-electron) ionization estimates.
    """
    function Settings()
        Settings(EmMultipole[], UseGauge[], Float64[], false, false, Tuple{Int64,Int64}[])
    end

    # `Base.show(io::IO, settings::NonresonantTwoColorPhotonIon.Settings)`
    #		... prepares a proper printout of the variable settings::NonresonantTwoColorPhotonIon.Settings.
    function Base.show(io::IO, settings::NonresonantTwoColorPhotonIon.Settings)
        println(io, "multipoles:               $(settings.multipoles)  ")
        println(io, "gauges:                   $(settings.gauges)  ")
        println(io, "photonEnergies:           $(settings.photonEnergies)  ")
        println(io, "printBeforeComputation:   $(settings.printBeforeComputation)  ")
        println(io, "selectLines:              $(settings.selectLines)  ")
        println(io, "selectedLines:            $(settings.selectedLines)  ")
    end

    """
    `struct  NonresonantTwoColorPhotonIon.Channel`  ... defines a type for a two-color channel to help characterize resonant two-color
                                                        (two-photon & single-electron) ionization with well-defined multipolarities.

        + multipoles     ::Array{EmMultipole,1}   ... Multipoles of all N incoming photons.
        + gauge          ::EmGauge                ... Gauge for dealing with the (coupled) radiation field.
        + kappa          ::Int64                  ... partial-wave of the free electron
        + symmetry       ::LevelSymmetry          ... total angular momentum and parity of the scattering state
        + phase          ::Float64                ... phase of the partial wave
        + amplitude      ::Complex{Float64}       ... multi-photon ionization amplitude associated with the given channel.
   """
    struct  Channel
        multipoles       ::Array{EmMultipole,1}
        gauge            ::EmGauge
        kappa            ::Int64
        symmetry         ::LevelSymmetry
        phase            ::Float64
        amplitude        ::Complex{Float64}
    end

    """
    `struct  NonresonantTwoColorPhotonIon.Line`  ... defines a type for a multi-photon ionization line that may include the definition of channels.

        + initialLevel     ::Level                  ... initial-(state) level
        + finalLevel       ::Level                  ... final-(state) level
        + photonEnergy     ::Float64                ... Energy of the incoming photons; all photons are assumed to have equal energy.
        + crossSection     ::EmProperty             ... Cross section for this multi-photon ionization.
        + hasChannels      ::Bool                   ... Determines whether the individual (sub-) channels are defined in terms of their
                                                        multipolarities, etc., or not.
        + channels         ::Array{NonresonantTwoColorPhotonIon.Channel,1}  ... List of NonresonantTwoColorPhotonIon.Channels of this line.
    """
    struct  Line
        initialLevel       ::Level
        finalLevel         ::Level
        photonEnergy       ::Float64
        crossSection       ::EmProperty
        hasChannels        ::Bool
        channels           ::Array{NonresonantTwoColorPhotonIon.Channel,1}
    end

    # `Base.show(io::IO, line::NonresonantTwoColorPhotonIon.Line)`
    #		... prepares a proper printout of the variable  line::NonresonantTwoColorPhotonIon.Line.
    function Base.show(io::IO, line::NonresonantTwoColorPhotonIon.Line)
        println(io, "initialLevel:      $(line.initialLevel)  ")
        println(io, "finalLevel:        $(line.finalLevel)  ")
        println(io, "photonEnergy:      $(line.photonEnergy)  ")
        println(io, "crossSection:      $(line.crossSection)  ")
        println(io, "hasChannels:       $(line.hasChannels)  ")
        println(io, "channels:          $(line.channels)  ")
    end

end # module


## ExoplanetsSysSim/src/corbits.jl
## (c) 2015 Eric B. Ford

if VERSION >= v"0.7.0-"
   using Libdl
   using Pkg
   using CORBITS
   #global const LIB_CORBITS = Libdl.find_library(["libcorbits.so"],[".",joinpath(Pkg.devdir(),"ExoplanetsSysSim/"),joinpath(Pkg.devdir(),"CORBITS/"),"/usr/local/lib"])  # WARNING: Assumes can find libcorbits.so
   global const LIB_CORBITS = Libdl.find_library(["libcorbits.so"],[".",joinpath(dirname(pathof(ExoplanetsSysSim))),joinpath(dirname(pathof(CORBITS))),"/usr/local/lib"])  # WARNING: Assumes can find libcorbits.so
elseif VERSION >= v"0.4.0-"
   global const LIB_CORBITS = Libdl.find_library(["libcorbits.so"],[".",joinpath(Pkg.dir(),"ExoplanetsSysSim/"),joinpath(Pkg.dir(),"CORBITS/"),"/usr/local/lib"])  # WARNING: Assumes can find libcorbits.so
else
   global const LIB_CORBITS = find_library(["libcorbits.so"],[".",joinpath(Pkg.devdir(),"ExoplanetsSysSim/"),joinpath(Pkg.devdir(),"CORBITS/"),"/usr/local/lib"])  # WARNING: Assumes can find libcorbits.so
end

#if !isdefined(is_corbits_initialized)
#  @windows_only println("# Warning:  Can't call external C libraries from Windows.  CORBITS won't work")
#end
#is_corbits_initialized = true

# Call CORBITS's function prob_of_transits_approx_arrays(a, r_star, r, e, Omega, omega, inc, use)
# Returns a Cdouble (aka Float64)
# For documentation see https://github.com/jbrakensiek/CORBITS
function prob_of_transits_approx(a::Vector{Cdouble},r_star::Cdouble,r::Vector{Cdouble}, e::Vector{Cdouble},
                                 Omega::Vector{Cdouble}, omega::Vector{Cdouble}, inc::Vector{Cdouble}, use::Vector{Cint})
  @assert(length(a) == length(r) == length(e) == length(Omega) == length(omega) == length(inc) >= length(use) )
  @assert(length(use) >=1 )
  num_pl = length(use)
  return ccall( (:prob_of_transits_approx_arrays, LIB_CORBITS), Cdouble,
      (Ptr{Cdouble}, Cdouble, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}, Cint),
       a, r_star, r, e, Omega, omega, inc, use, num_pl)
end

function test_corbits()
  #@compat is_windows() ? error("# CORBITS won't work on windows") : nothing
  a =  Cdouble[0.05, 0.15]
  r_star = convert(Cdouble,0.005)
  r = Cdouble[0.0001,0.0001]
  ecc = Cdouble[0.02, 0.1]
  Omega = Cdouble[0.0, 0.0]
  omega = Cdouble[ 0.0, 0.5]
  inc = Cdouble[pi/2, pi/2]
  use_pl = Cint[1,1]
  prob_of_transits_approx(a, r_star, r, ecc, Omega, omega, inc, use_pl)
end