# primary features:
# rollout sequence of controls
# df/dx, df/du derivatives

# features to add:
# address mjdata, mjmodel fields by string or symbol
#   replicate the get_sensors stuff
__precompile__()

module mjWrap

using MuJoCo
using JLD2

using UnsafeArrays
using LinearAlgebra.BLAS
using Distributed

struct mjSet
    m::mj.jlModel
    d::mj.jlData
    datas::Vector{mj.jlData}

    # helpers
    sensor_names::Dict{Symbol, AbstractRange}
    name::String
    skip::Integer
    mode::Integer

    nq::Integer
    nv::Integer
    nu::Integer
    ns::Integer
    dt::Float64

    #function mjSet(model_name::String) # TODO?
    #end
end

struct TrajSamples
    T::Integer
    K::Integer
    s0::Vector{Float64}      # nqnv
    ctrl::Array{Float64,3}   # nu x T x K
    state::Array{Float64,3}  # nqnv x T x K
    obs::Array{Float64,3}    # nu x T x K
    reward::Array{Float64,2} # T x K
end

struct Trajectory
    ctrl::Matrix{Float64}   # nu x T x K
    state::Matrix{Float64}  # nqnv x T x K
    obs::Matrix{Float64}    # nu x T x K
    reward::Vector{Float64} # T x K
end

#=
type TrajDerivatives
    T::Integer
    K::Integer
    d_x::Array{Float64}
    d_u::Array{Float64}
    s_x::Array{Float64}
    s_u::Array{Float64}

    #c::Array{Float64}
    #c_x::Array{Float64}
    #c_u::Array{Float64}
    #c_xx::Array{Float64}
    #c_uu::Array{Float64}
    #c_ux::Array{Float64}
end
=#

const mjw = mjWrap
export TrajSamples, mjSet, mjw
const GDD_POS = 1
const GDD_ACC = 2
const GDD_Q = 3
const GDD_COST = 4

function finish()
    mj.deactivate()
    return 0;
end

function modelparams(x::mjSet)
    return x.nq, x.nv, x.nu, x.ns
end

function clear_model(x::mjSet)
    mj.deleteModel(x.m.m)
    mj.deleteData(x.d.d)
    for i=1:length(x.datas)
        mj.deleteData(x.datas[i])
    end
end

function load_model(model_name::String, skip::Integer, mode::String="normal", ns::Integer=0)
    pm = mj_loadXML(model_name, C_NULL)
    if pm == nothing
        return
    end
    pd = mj_makeData(pm)

    m, d = mj.mapmujoco(pm, pd)
    ndata = Threads.nthreads()
    datas = Vector{jlData}(undef, ndata)
    print("Making mjDatas: 1 ")
    datas[1] = d
    for i=2:ndata
        print("$i ")
        datas[i] = mj.mapdata(pm, mj_makeData(pm))
    end
    println()

    if m.m[].nsensor > 0
        sensors = mj.name2range(m, m.m[].nsensor,
                                m.name_sensoradr, m.sensor_adr, m.sensor_dim)
    else
        sensors = Dict(:none=>0:0)
    end

    nq = Int64(m.m[].nq)
    nv = Int64(m.m[].nv)
    nu = Int64(m.m[].nu)
    if ns==0 ns = Int64(m.m[].nsensordata) end
    dt = m.m[].opt.timestep #mj.get(m, :opt, :timestep)

    return mjSet(m, d, datas, sensors, model_name, skip, 0,
                 nq, nv, nu, ns, dt)
end

############ UTILS ############
function wraptopi(ang::Float64)
    return mod(ang+pi, 2pi) - pi
end

function unitfit(ang::Float64, val::Float64)
    clamp.(ang, -val, val)./val
end

function inertiaweight(m::mj.jlModel, d::mj.jlData,
                       v::AbstractVector{Float64})
    nv = length(d.qvel)
    mark = mj.MARKSTACK(d)

    My = @view d.stack[(mark+1):(mark+nv)]
    mj_solveM(m.m, d.d, My, v, 1)
    force = v'*My

    mj.FREESTACK(d, Int32(mark))
    return force/mj_getTotalmass(m.m)
end

#function stackAlloc(m::mj.jlModel, d::mj.jlData)
#   mark = mj.MARKSTACK(d)
#end

function limitcontrols(x::mjSet, s::TrajSamples)
    for u=1:x.nu
        clamp!(uview(s.ctrl,u,:,:),
               x.m.actuator_ctrlrange[1,u],
               x.m.actuator_ctrlrange[2,u])
    end
end

# modify ctrl vector with the noise stored in ctrl
function inertiascale!(m::mj.jlModel, d::mj.jlData,
                       input::AbstractVector{Float64},
                       output::AbstractVector{Float64})
    mj.solveM2(m.m, d.d, output, input, 1)
end

function state2data(d::mj.jlData, s::AbstractVector{Float64})
    nq = length(d.qpos)
    nv = length(d.qvel)
    copyto!(d.qpos, 1, s, 1, nq)
    copyto!(d.qvel, 1, s, nq+1, nv)
end

function data2state(s::AbstractVector{Float64}, d::mj.jlData)
    nq = length(d.qpos)
    nv = length(d.qvel)
    copyto!(s, d.qpos)
    copyto!(s, nq+1, d.qvel, 1, nv)
end

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

function reset(m::mj.jlModel, d::mj.jlData,
               s0::AbstractVector{Float64},
               ctrl0::AbstractVector{Float64},
               obs::AbstractVector{Float64})
    # reset
    d.qacc .= 0.0
    state2data(d, s0)
    copyto!(d.ctrl, ctrl0)
    d.d[].time = 0.0 #mj.set(d, :time, 0.0)

    # out
    if length(d.sensordata) > 0
        mj_forward(m, d)
        copyto!(obs, d.sensordata)
    else
        mj_forwardSkip(m.m, d.d, Integer(mj.STAGE_NONE), 1)
    end
end

# accepts threading; tid is 1 indexed in Julia
function reset(x::mjSet,
               s0::AbstractVector{Float64},
               ctrl0::AbstractVector{Float64},
               obs::AbstractVector{Float64})
    tid = Threads.threadid()
    if x.mode>0
        reset(x.m, x.datas[tid], s0, zeros(x.datas[tid].ctrl), obs)
    else
        reset(x.m, x.datas[tid], s0, ctrl0, obs)
    end
end

########################################################################## step

function quat2euler(quat::AbstractVector{Float64})
    w, x, y, z = quat
    sinr = 2.0 * (w * x + y * z)
    cosr = 1.0 - 2.0 * (x * x + y * y)
    roll = atan2(sinr, cosr)

    # pitch (y-axis rotation)
    sinp = 2.0 * (w * y - z * x)
    if abs(sinp) >= 1
        pitch = copysign(pi / 2, sinp) # use 90 degrees if out of range
    else
        pitch = asin(sinp)
    end

    # yaw (z-axis rotation)
    siny = 2.0 * (w * z + x * y)
    cosy = 1.0 - 2.0 * (y * y + z * z)
    yaw = atan2(siny, cosy)

    return roll, pitch, yaw
end

function euler2quat(roll::Float64, pitch::Float64, yaw::Float64)
    # blame wikipedia if this doesn't work
    cy = cos(yaw * 0.5)
    sy = sin(yaw * 0.5)
    cr = cos(roll * 0.5)
    sr = sin(roll * 0.5)
    local cp = cos(pitch * 0.5)
    sp = sin(pitch * 0.5)

    w = cy * cr * cp + sy * sr * sp
    x = cy * sr * cp - sy * cr * sp
    y = cy * cr * sp + sy * sr * cp
    z = sy * cr * cp - cy * sr * sp
    return [w, x, y, z]
end

function gdd_step!(m::mj.jlModel, d::mj.jlData,
                   x::mjSet,
                   skip::Integer,
                   ctrl::AbstractVector{Float64},
                   s1::AbstractVector{Float64},
                   obs::AbstractVector{Float64})
    tid = Threads.threadid()
    nq = length(d.qpos)
    nv = length(d.qvel)
    nu = length(ctrl)
    njnt = m.m[].njnt
    nefc = d.d[].nefc
    #copyto!(d.ctrl, ctrl)

    alpha = 50.0
    local beta = 5.0
    weight = 300.0
    gain = 20.0
    p_weight = 300.0
    p_gain = 20.0

    GDopt = x.GDopt[tid]
    #println(tid, ": ", size(x.GD_Q))
    Q = view( x.GD_Q, :, tid)
    q = view( x.GD_q, :, tid)
    P = view( x.GD_P, :, tid)
    p = view( x.GD_p, :, tid)
    GDopt.Q = pointer(Q)
    GDopt.q = pointer(q)
    GDopt.P = pointer(P) # each time this could be different
    GDopt.p = pointer(p)

    mydiff = zeros(nu)
    for i=1:skip
        mj_step1(m.m, d.d)

        #setGDcost()
        mj.fullM(m.m,       GDopt.Q, d.qM)
        mj.mju_scl(GDopt.Q, GDopt.Q, alpha, nv*nv)
        mj.mulM(m.m, d.d,   GDopt.q, d.qvel)
        mj.mju_scl(GDopt.q, GDopt.q, alpha*beta, nv)

        # find all free / quat joints
        # find position indexes for them; convert quat to euler
        # take difference of euler angles
        if x.mode == GDD_POS
            pj = 1
            qi = 1 # offset for number of free / ball joints encountered
            for j=1:njnt
                if m.jnt_type[j] == Int(mj.mjJNT_FREE)
                    #println("Free")
                    mydiff[pj] = d.qpos[qi] - ctrl[pj]; qi+=1; pj+=1 # x
                    mydiff[pj] = d.qpos[qi] - ctrl[pj]; qi+=1; pj+=1 # y
                    mydiff[pj] = d.qpos[qi] - ctrl[pj]; qi+=1; pj+=1 # z
                    roll,pitch,yaw = quat2euler(d.qpos[qi:(qi+3)]); qi+=4
                    mydiff[pj:(pj+2)] = [roll,pitch,yaw] - ctrl[pj:(pj+2)]; pj+=3 # quat
                elseif m.jnt_type[i] == Int(mj.mjJNT_BALL)
                    #println("ball")
                    roll,pitch,yaw = quat2euler(d.qpos[qi:(qi+3)]); qi+=4
                    mydiff[pj:(pj+2)] = [roll,pitch,yaw] - ctrl[pj:(pj+2)]; pj+=3 # quat
                else # slide or hinge
                    #println("1dof")
                    mydiff[pj] = d.qpos[qi] - ctrl[pj]; qi+=1; pj+=1
                end
            end
            if qi!=(nq+1)
                error("bad lineup for qpos")
            end
            if pj!=(nu+1)
                error("bad lineup for dofs")
            end
            for j=1:nu # nu == nv; for each dof...
                Q[(j-1)*(nv+1) + 1] += p_weight # diagonal weights
                Madr = m.dof_Madr[j] + 1 # julia indexing
                q[j] += p_weight*
                (p_gain*mydiff[j] + 7.0*sqrt(p_gain*d.qM[Madr])*d.qvel[j])
            end
        elseif x.mode == GDD_ACC
            #mydiff[:] = ctrl #.*10.0
            for j=1:nu # nu == nv; for each dof...
                Q[(j-1)*(nv+1) + 1] += p_weight # diagonal weights
                #q[j] += p_weight*ctrl[j]
                Madr = m.dof_Madr[j] + 1 # julia indexing
                q[j] += p_weight*
                (p_gain*ctrl[j] + 7.0*sqrt(p_gain*d.qM[Madr])*d.qvel[j])
            end
        elseif x.mode == GDD_Q
            #Q[:] += At_mul_B(ctrl[1:(nv*nv)], ctrl[1:(nv*nv)])
            #q[:] += ctrl[(nv*nv+1):end]
            for i=1:nv
                Q[(i-1)*(nv+1)+1] += abs(ctrl[nv+i])*p_weight

                Madr = m.dof_Madr[i] + 1 # julia indexing
                q[i] += p_weight*
                (p_gain*ctrl[i]+7.0*sqrt(p_gain*d.qM[Madr])*d.qvel[i])
            end
            #q[:] += ctrl[(nv+1):end]*p_weight
        end
        #setGDcostDone

        mj.op_GDD_solve(m.m, d.d,
                        Ptr{mj.opGDDOption}(pointer_from_objref(GDopt)))
        mj_step2(m.m, d.d)
    end

    mj_forward(m, d)
    data2state(s1, d)
    copyto!(obs, d.sensordata)
end

function step!(m::mj.jlModel, d::mj.jlData, skip::Integer,
               ctrl::AbstractVector{Float64},
               s1::AbstractVector{Float64},
               obs::AbstractVector{Float64})
    copyto!(d.ctrl, ctrl)

    mj_step2(m.m, d.d) # qpos qvel set from before, just use new ctrl and euler step
    for i=2:skip
        #mj_step(m, d) # TODO MAKE FASTER with fwd and euler
        mj_forwardSkip(m.m, d.d, Integer(mj.STAGE_NONE), 1)
        mj_Euler(m.m, d.d)
    end

    if length(d.sensordata) > 0
        mj_forward(m, d)
        copyto!(obs, d.sensordata)
    else
        mj_forwardSkip(m.m, d.d, Integer(mj.STAGE_NONE), 1)
    end

    data2state(s1, d)
end

# accepts threading; tid is 1 indexed in Julia
function step!(x::mjSet,
               ctrl::AbstractVector{Float64},
               s1::AbstractVector{Float64},
               obs::AbstractVector{Float64})
    tid = Threads.threadid()
    if x.mode>0
        gdd_step!(x.m, x.datas[tid], x, tid,
                  x.skip, ctrl, s1, obs)
    else
        step!(x.m, x.datas[tid], x.skip, ctrl, s1, obs)
    end
end

# resets position then takes step
function step!(x::mjSet,
               s0::AbstractVector{Float64},
               ctrl::AbstractVector{Float64},
               s1::AbstractVector{Float64},
               obs::AbstractVector{Float64})
    error("init_step s0 -> s1 not implemented")
end

function alloc_fields(nqnv::Integer, nu::Integer, ns::Integer,
                      T::Integer, K::Integer)
    ctrl  = zeros(nu, T, K)
    state = zeros(nqnv, T, K) # just allocate space
    obs   = zeros(ns,   T, K)
    return ctrl, state, obs
end

# optionally we can specify what we want our observations to be elsewhere
# instead of in mujoco's sensors format
function allocateTrajSamples(x::mjSet, T::Integer, K::Integer,
                             ns::Integer=x.ns)
    nq, nv, nu, _ = mjw.modelparams(x)
    ctrl, state, obs = alloc_fields(nq+nv, nu, ns, T, K)

    s0 = zeros(nq+nv)
    s0[1:nq]     = x.d.qpos
    s0[nq+1:end] = x.d.qvel

    return TrajSamples(T, K, s0,
                       ctrl, state, obs, zeros(T,K))
end

function save_traj(x::mjSet,
                   s::TrajSamples,
                   filename::String)
    mode = x.mode>=1 ? "GDD_$(x.mode)" : "normal"
    jldopen(filename, "w") do file
        file["state"] = s.state
        file["ctrl"] = s.ctrl
        file["obs"] = s.obs
        file["model"] = basename(x.name)
        file["skip"] = x.skip
        file["mode"] = mode
    end
end

function roll(x::mjw.mjSet, startT::Int,
              state::AbstractMatrix{Float64},
              obs::AbstractMatrix{Float64},
              ctrl::AbstractMatrix{Float64},
              reward::AbstractVector{Float64},
              ctrlfunc!::Function,
              obsfunc!::Function,
              rewardfunc::Function)
    T = size(state,2)
    s = @view state[:,startT]
    c = @view ctrl[:,startT]
    o = @view obs[:,startT]
    reset(x, s, c, o)
    obsfunc!(x, s, o)    # arbitrary observation vector manipulation
    for t=startT:(T-1)

        c = @view ctrl[:,t]
        o = @view obs[:,t]

        # then we will execute the stored control
        ctrlfunc!(x, c, o)        # we can also set noise to be in ctrl

        s = @view state[:,t+1]
        o = @view obs[:,t+1]
        step!(x, c, s, o) # get next state, next observation
        obsfunc!(x, s, o) # arbitrary observation vector manipulation

        #reward[t] = rewardfunc(x, s0, s, o0, o, c)
        reward[t] = rewardfunc(x, s, s, o, o, c)
    end
    s = @view state[:,T]
    c = @view ctrl[:,T]
    o = @view obs[:,T]
    reward[T] = rewardfunc(x, s, s, o, o, c)
end

function roll2(x::mjw.mjSet, startT::Int,
               state::AbstractMatrix{Float64},
               obs::AbstractMatrix{Float64},
               ctrl::AbstractMatrix{Float64},
               reward::AbstractVector{Float64},
               ctrlfunc!::Function,
               obsfunc!::Function,
               rewardfunc::Function)
    T = size(state,2)
    s = uview( state, :, startT )
    c = uview( ctrl, :, startT )
    size_o = size(obs, 1) > 0 ? true : false
    if size_o
        o = uview( obs, :, startT )
    else
        o = Vector{Float64}(undef, 0)
    end
    reset(x, s, c, o)
    obsfunc!(x, s, o)    # arbitrary observation vector manipulation
    for t=startT:(T-1)

        c = uview( ctrl,:,t)
        if size_o o = uview( obs,:,t) end

        # then we will execute the stored control
        ctrlfunc!(x, c, o)        # we can also set noise to be in ctrl

        s = uview( state,:,t+1)
        if size_o o = uview( obs,:,t+1) end
        step!(x, c, s, o) # get next state, next observation
        obsfunc!(x, s, o) # arbitrary observation vector manipulation

        #reward[t] = rewardfunc(x, s0, s, o0, o, c)
        reward[t] = rewardfunc(x, s, s, o, o, c)
    end
    s = uview( state,:,T)
    c = uview( ctrl,:,T)
    if size_o o = uview( obs,:,T) end
    reward[T] = rewardfunc(x, s, s, o, o, c)
end

function rollout2(x::mjw.mjSet,
                  samp::mjw.TrajSamples,
                  ctrlfunc!::Function=(y...)->nothing,
                  obsfunc!::Function=(y...)->nothing,
                  rewardfunc::Function=(y...)->0.0;
                  startT::Int=1)
    nthread = Int(min(samp.K, Int(Threads.nthreads()))) # cant have more threads than data
    BLAS.set_num_threads(1)
    Threads.@threads for i=1:nthread
    #for i=1:nthread
        BLAS.set_num_threads(1)
        thread_range = Distributed.splitrange(samp.K, nthread)[i]

        for k=thread_range
            if size(samp.obs, 1) > 0
                roll2(x, startT,
                      uview(samp.state,:,:,k), uview(samp.obs,:,:,k),
                uview(samp.ctrl,:,:,k), uview(samp.reward,:,k),
                ctrlfunc!, obsfunc!, rewardfunc)
            else
                roll2(x, startT,
                      uview(samp.state,:,:,k), Matrix{Float64}(undef, 0, 0),
                uview(samp.ctrl,:,:,k), uview(samp.reward,:,k),
                ctrlfunc!, obsfunc!, rewardfunc)
            end
        end
    end
    BLAS.set_num_threads(Threads.nthreads())
end

function rollout_st(x::mjw.mjSet,
                    samp::mjw.TrajSamples,
                    ctrlfunc!::Function=(y...)->nothing,
                    obsfunc!::Function=(y...)->nothing,
                    rewardfunc::Function=(y...)->0.0;
                    startT::Int=1)
    for k=1:samp.K
        roll2(x, startT,
              uview(samp.state,:,:,k), uview(samp.obs,:,:,k),
              uview(samp.ctrl,:,:,k), uview(samp.reward,:,k),
              ctrlfunc!, obsfunc!, rewardfunc)
    end
end

function rollout(x::mjw.mjSet, traj::mjw.Trajectory,
                 ctrlfunc!::Function=(y...)->nothing,
                 obsfunc!::Function=(y...)->nothing,
                 rewardfunc::Function=(y...)->0.0)
    nthread = Int(min(samp.K, Int(Threads.nthreads()))) # cant have more threads than data
    BLAS.set_num_threads(1)
    nq, nv, ns, nu = x.nq, x.nv, x.ns, x.nu
    Threads.@threads for i=1:nthread
        #for i=1:nthread
        BLAS.set_num_threads(1)
        s0, s = Array{Float64}(nq+nv), Array{Float64}(nq+nv)
        o0, o = Array{Float64}(ns), Array{Float64}(ns)
        c     = Array{Float64}(nu)
        thread_range = Distributed.splitrange(samp.K, nthread)[i]
        #println("i: $i, tid $tid : $thread_range")

        @inbounds for k=Range(thread_range)
            s[:] = samp.state[:,1,k]
            c[:] = samp.ctrl[:,1,k]
            reset(x, s, c, o)
            obsfunc!(x, s, o)    # arbitrary observation vector manipulation
            samp.state[:,1,k] = s
            s0[:]             = s
            samp.obs[:,1,k]   = o
            o0[:]             = o
            for t=Range(1:samp.T)

                c .= samp.ctrl[:,t,k]  # if ctrlfunc does nothing (no controller)
                # then we will execute the stored control
                ctrlfunc!(x, c, o)        # we can also set noise to be in ctrl

                step!(x, c, s, o) # get next state, next observation
                obsfunc!(x, s, o) # arbitrary observation vector manipulation

                #reward[t] = rewardfunc(x, s0, s, o0, o, c)
                reward[t] = rewardfunc(x, s, s, o, o, c)
            end
            s = uview( state,:,T)
            c = uview( ctrl,:,T)
            o = view( obs,:,T)
            reward[T] = rewardfunc(x, s, s, o, o, c)
        end
    end
end
end # mj module


raw"""
All the inputs are consumed and the difference of the output sums goes to the
miner

Input hashes and output index point to the original transaction
The unlocking script must match the locking script

Transaction:
  - Version
  - Input Counter
  - Inputs
    - hash
        The transaction origin id
    - output index
        The transacion output idx
    - unlocking script size
    - unlocking script
    - sequence number
  - Outputs
    - amount
    - locking script size
    - locking script

Hexdump of the genesis block area blk00000.dat:

f9be b4d9 1d01 0000 0100 0000 0000 0000  ................
^^^^^^^^^ magic bytes
          ^^^^^^^^^ block size
                    ^^^^^^^^^ version
                              ^^^^^^^^^
0000 0000 0000 0000 0000 0000 0000 0000  ................
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0000 0000 0000 0000 0000 0000 3ba3 edfd  ............;...
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ previous hash
                              ^^^^^^^^^
7a7b 12b2 7ac7 2c3e 6776 8f61 7fc8 1bc3  z{..z.,>gv.a....
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
888a 5132 3a9f b8aa 4b1e 5e4a 29ab 5f49  ..Q2:...K.^J)._I
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ merkle root
                              ^^^^^^^^^ time stamp
ffff 001d 1dac 2b7c 0101 0000 0001 0000  ......+|........
^^^^^^^^^ difficulty
          ^^^^^^^^^ nonce
                    ^^ number of transactions (varint)
                      ^^^^^^^^^^ version
                                ^^ number of inputs
                                   ^^^^
0000 0000 0000 0000 0000 0000 0000 0000  ................
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0000 0000 0000 0000 0000 0000 0000 ffff  ................
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ originating transaction hash
                                   ^^^^
ffff 4d04 ffff 001d 0104 4554 6865 2054  ..M.......EThe T
^^^^ output index of the originating transaction
     ^^ script length (varint)
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
696d 6573 2030 332f 4a61 6e2f 3230 3039  imes 03/Jan/2009
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2043 6861 6e63 656c 6c6f 7220 6f6e 2062   Chancellor on b
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
7269 6e6b 206f 6620 7365 636f 6e64 2062  rink of second b
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
6169 6c6f 7574 2066 6f72 2062 616e 6b73  ailout for banks
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unlocking script
ffff ffff 0100 f205 2a01 0000 0043 4104  ........*....CA.
^^^^^^^^^ sequence number
          ^^ number of outputs
            ^^^^^^^^^^^^^^^^^^^^ amount (Big endian Float64)
                                ^^ script length
                                   ^^^^
678a fdb0 fe55 4827 1967 f1a6 7130 b710  g....UH'.g..q0..
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
5cd6 a828 e039 09a6 7962 e0ea 1f61 deb6  \..(.9..yb...a..
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
49f6 bc3f 4cef 38c4 f355 04e5 1ec1 12de  I..?L.8..U......
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
5c38 4df7 ba0b 8d57 8a4c 702b 6bf1 1d5f  \8M....W.Lp+k.._
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ac00 0000 00f9 beb4 d9d7 0000 0001 0000  ................
^^ locking script
  ^^^^^^^^^^ lock time
            ^^^^^^^^^^ magic bytes
006f e28c 0ab6 f1b3 72c1 a6a2 46ae 63f7  .o......r...F.c.
4f93 1e83 65e1 5a08 9c68 d619 0000 0000  O...e.Z..h......
0098 2051 fd1e 4ba7 44bb be68 0e1f ee14  .. Q..K.D..h....
677b a1a3 c354 0bf7 b1cd b606 e857 233e  g{...T.......W#>
0e61 bc66 49ff ff00 1d01 e362 9901 0100  .a.fI......b....
0000 0100 0000 0000 0000 0000 0000 0000  ................
0000 0000 0000 0000 0000 0000 0000 0000  ................
0000 00ff ffff ff07 04ff ff00 1d01 04ff  ................
ffff ff01 00f2 052a 0100 0000 4341 0496  .......*....CA..
"""

"""
    TransactionInput

Data Structure Storing Transaction Inputs
"""
struct TransactionInput
    # TODO: remove unlocking_script_size
    # the hash of the transaction where the output comes from
    hash                  :: UInt256
    output_index          :: UInt32
    unlocking_script_size :: UInt64
    unlocking_script      :: Vector{UInt8}
    sequence_number       :: UInt32
end

"""
    TransactionOutput

Data Structure Storing Transaction Outputs
"""
struct TransactionOutput
    # TODO: remove locking_script_size
    amount              :: UInt64 # in satoshis = 1e-8 bitcoins
    locking_script_size :: UInt64
    locking_script      :: Vector{UInt8}
end

"""
    Witness

Data Structure Storing The Witness for SegWit transactions
"""
struct Witness
    data :: Vector{Vector{UInt8}}
end

"""
    Transaction

Data Structure Storing Transactions

To get the hash of a transaction

```julia
double_sha256(tx)
```
"""
struct Transaction
    # TODO: remove *_counter
    version         :: UInt32
    marker          :: UInt8
    flag            :: UInt8
    input_counter   :: UInt64
    inputs          :: Vector{TransactionInput}
    output_counter  :: UInt64
    outputs         :: Vector{TransactionOutput}
    witness_counter :: UInt64
    witnesses       :: Vector{Witness}
    lock_time       :: UInt32
end

function TransactionInput(io)

    # TODO: Cannot read UInt256 directly from io here:
    in_hash = read!(io, Array{UInt256}(undef, 1))[1]
    # in_hash = read(io, UInt256)
    output_index = read(io, UInt32)
    unlocking_script_size = signed(read_varint(io))

    unlocking_script = read!(io, Array{UInt8}(undef, unlocking_script_size))
    sequence_number = read(io, UInt32)

    TransactionInput(
        in_hash,
        output_index,
        unlocking_script_size,
        unlocking_script,
        sequence_number
    )
end

function TransactionOutput(io)

    out_amount = read(io, UInt64)
    out_script_size = signed(read_varint(io))
    locking_script = read!(io, Array{UInt8}(undef, out_script_size))

    TransactionOutput(out_amount,
                      out_script_size,
                      locking_script)
end

function Witness(io::IO)
    n_items = read_varint(io)
    data = Vector{Vector{UInt8}}(undef, n_items)

    for i in 1:n_items
        l = read_varint(io)
        data[i] = read!(io, Array{UInt8, 1}(undef, l))
    end

    Witness(data)
end

function Transaction(io)

    version = read(io, UInt32)

    marker_or_n_in = read_varint(io)

    is_segwit = marker_or_n_in == zero(marker_or_n_in)
    if is_segwit
        marker = UInt8(marker_or_n_in)
        @assert marker == 0x00
        flag = read(io, UInt8)
        @assert flag == 0x01
        n_in = read_varint(io)
    else
        # these two don't matter
        marker = 0xff
        flag = 0xff
        n_in = marker_or_n_in
    end

    @assert n_in > 0
    inputs = TransactionInput[
        TransactionInput(io)
        for i in 1:n_in
    ]

    n_out = read_varint(io)
    @assert n_out > 0
    outputs = TransactionOutput[
        TransactionOutput(io)
        for i in 1:n_out
    ]

    if is_segwit
        n_witness = n_in
        @assert n_witness > 0
        witnesses = Witness[
            Witness(io)
            for i in 1:n_witness
        ]
    else
        n_witness = zero(UInt64)
        witnesses = Witness[]
    end

    locktime = read(io, UInt32)

    Transaction(version,
                marker,
                flag,
                n_in,
                inputs,
                n_out,
                outputs,
                n_witness,
                witnesses,
                locktime)
end

# TODO: is this the same as bytes2hex?
hexarray(x::Array{UInt8}) = mapreduce(x -> string(x, base = 16), *, x)

function showcompact(io::IO, input::TransactionInput)
    # TODO: do something useful here
    println(io, "")
end

function Base.show(io::IO, input::TransactionInput)
    if !get(io, :compact, false)
        println(io, "Transaction input:")
        println(io, "  Hash:                  " * string(input.hash,            base = 16))
        println(io, "  Output index:          " * string(input.output_index,    base = 10))
        println(io, "  Input Sequence:        " * string(input.sequence_number, base = 10))
    end
end

# function Base.showall(io::IO, input::TransactionInput)
#     println(io, "Transaction input:")
#     println(io, "  Hash:                  " * input.hash)
#     println(io, "  Output index:          " * input.output_index)
#     println(io, "  Unlocking script size: " * input.unlocking_script_size)
#     println(io, "  Unlocking script:      " * hexarray(input.unlocking_script))
#     println(io, "  Input Sequence:        " * input.sequence_number)
# end

function Base.show(io::IO, output::TransactionOutput)
    println(io, "Transaction output: " * string(output.amount, base = 10))
end
# function Base.showall(io::IO, output::TransactionOutput)
#     println(io, "Transaction output: "    * string(output.amount,              base = 10))
#     println(io, "  Locking script size: " * string(output.locking_script_size, base = 10))
#     println(io, "  Locking script:      " * hexarray(output.locking_script))
# end

function showcompact(io::IO, transaction::Transaction)
    # TODO: add id here
    println(io, "Transaction: ")
end

function Base.show(io::IO, transaction::Transaction)
    # TODO: add id here
    if !get(io, :compact, false)
        println(io, "Transaction: ")
        println(io, "  Version:        " * string(transaction.version, base = 10))
        println(io, "  Input counter:  " * string(transaction.input_counter, base = 10))
        println(io, "  Output counter: " * string(transaction.output_counter, base = 10))
        println(io, "  Lock time:      " * string(transaction.lock_time, base = 10))
    end
end

# function Base.showall(io::IO, transaction::Transaction)
#     # TODO: add id here
#     println(io, "Transaction: ")
#     println(io, "  Version:        " * string(transaction.version, base = 10))
#     println(io, "  Input counter:  " * string(transaction.input_counter, base = 10))
#     for i in 1:transaction.input_counter
#         show(transaction.inputs[i])
#     end
#     println(io, "  Output counter: " * string(transaction.output_counter, base = 10))
#     for i in 1:transaction.output_counter
#         show(transaction.outputs[i])
#     end
#     println(io, "  Lock time:      " * string(transaction.lock_time, base = 10))
# end

# TODO: little endian only:
function sha256(tx::Transaction)

    ctx = SHA.SHA256_CTX()

    SHA.update!(ctx, to_byte_tuple(tx.version))
    SHA.update!(ctx, to_varint(tx.input_counter))
    for i in 1:tx.input_counter
        SHA.update!(ctx, to_byte_tuple(tx.inputs[i].hash))
        SHA.update!(ctx, to_byte_tuple(tx.inputs[i].output_index))
        SHA.update!(ctx, to_varint(tx.inputs[i].unlocking_script_size))
        SHA.update!(ctx, tx.inputs[i].unlocking_script)
        SHA.update!(ctx, to_byte_tuple(tx.inputs[i].sequence_number))
    end
    SHA.update!(ctx, to_varint(tx.output_counter))
    for i in 1:tx.output_counter
        SHA.update!(ctx, to_byte_tuple(tx.outputs[i].amount))
        SHA.update!(ctx, to_varint(tx.outputs[i].locking_script_size))
        SHA.update!(ctx, tx.outputs[i].locking_script)
    end
    SHA.update!(ctx, to_byte_tuple(tx.lock_time))

    return SHA.digest!(ctx)
end

function double_sha256(tx::Transaction)::UInt256
    tx |> sha256 |> sha256 |> x -> reinterpret(UInt256, x)[1]
end

# TODO: copyto! is save, right?
# TODO: find an allocation free version for these: copyto!(x, i, to_byte_tuple(y))
"""
mutates data in place, advances idx and returns it.
"""
function dump_txin_data!(data, idx, txin::TransactionInput)
    copyto!(data, idx, to_byte_tuple(txin.hash))
    idx += sizeof(txin.hash)

    copyto!(data, idx, to_byte_tuple(txin.output_index))
    idx += sizeof(txin.output_index)

    unlocking_script_size = to_varint(txin.unlocking_script_size)
    copyto!(data, idx, unlocking_script_size)
    idx += sizeof(unlocking_script_size)

    copyto!(data, idx, txin.unlocking_script)
    idx += sizeof(txin.unlocking_script)

    copyto!(data, idx, to_byte_tuple(txin.sequence_number))
    idx += sizeof(txin.sequence_number)

    return idx
end

# TODO: copyto! is save, right?
# TODO: find an allocation free version for these: copyto!(x, i, to_byte_tuple(y))
"""
mutates data in place, advances idx and returns it.
"""
function dump_txout_data!(data, idx, txout::TransactionOutput)
    copyto!(data, idx, to_byte_tuple(txout.amount))
    idx += sizeof(txout.amount)

    locking_script_size = to_varint(txout.locking_script_size)
    copyto!(data, idx, locking_script_size)
    idx += sizeof(locking_script_size)

    copyto!(data, idx, txout.locking_script)
    idx += sizeof(txout.locking_script)

    return idx
end

# TODO: copyto! is save, right?
# TODO: find an allocation free version for these: copyto!(x, i, to_byte_tuple(y))
"""
mutates data in place, advances idx and returns it.
"""
function dump_tx_data!(data, idx, tx::Transaction)
    copyto!(data, idx, to_byte_tuple(tx.version))
    idx += sizeof(tx.version)

    is_segwit = tx.marker == 0x00
    if is_segwit
        data[idx] = tx.marker
        idx += 1
        data[idx] = tx.flag
        idx += 1
    end

    input_counter = to_varint(tx.input_counter)
    copyto!(data, idx, input_counter)
    idx += sizeof(input_counter)
    for i in 1:tx.input_counter
        idx = dump_txin_data!(data, idx, tx.inputs[i])
    end

    output_counter = to_varint(tx.output_counter)
    copyto!(data, idx, output_counter)
    idx += sizeof(output_counter)
    for i in 1:tx.output_counter
        idx = dump_txout_data!(data, idx, tx.outputs[i])
    end

    if is_segwit
        for i in 1:tx.input_counter
            n_items = length(tx.witnesses[i].data)
            n_items_bytes = to_varint(n_items)
            copyto!(data, idx, n_items_bytes)
            idx += length(n_items_bytes)

            for j in 1:n_items
                l = length(tx.witnesses[i].data[j])
                l_bytes = to_varint(l)
                copyto!(data, idx, l_bytes)
                idx += length(l_bytes)

                copyto!(data, idx, tx.witnesses[i].data[j])
                idx += l
            end
        end
    end

    copyto!(data, idx, to_byte_tuple(tx.lock_time))
    idx += sizeof(tx.lock_time)

    return idx
end

# TODO: common abstraction with "Link"?
struct FilePointer
    file_number   :: UInt64
    file_position :: UInt64
end

function Transaction(fp::FilePointer)
    fp.file_number |>
        get_block_chain_file_path |>
        x -> open(x) do fh
            seek(fh, fp.file_position)
            Transaction(fh)
        end
end

total_output(tx::Transaction) = sum(x -> x.amount, tx.outputs)

is_coinbase(txin::TransactionInput) = txin.hash == zero(UInt256)


using MPI
using MPICH_jll
using PyCall
using QuantumEspresso_jll
using Unitful
using UnitfulAtomic

const QE_PSEUDODIR = joinpath(@__DIR__, "qe_pseudos")
const QE_PSEUDOMAP = Dict(
   "Si" => "Si.pbe-hgh.UPF",
   "O"  => "O.pbe-hgh.UPF",
   "Al" => "Al.pbe-hgh.UPF",
   "Cu" => "Cu.pbe-d-hgh.UPF",
   "Ga" => "Ga.pbe-d-hgh.UPF",
   "As" => "As.pbe-hgh.UPF",
   "H"  => "H.pbe-hgh.UPF",
   "Cr" => "Cr.pbe-sp-hgh.UPF",
   "Ni" => "Ni.pbe-sp-hgh.UPF",
   "Co" => "Co.pbe-sp-hgh.UPF",
   "Fe" => "Fe.pbe-sp-hgh.UPF",
   "Mn" => "Mn.pbe-sp-hgh.UPF",
   "Ge" => "Ge.pbe-hgh.UPF",
   "Pt" => "Pt.pbe-hgh.UPF",
   "B"  => "B.pbe-hgh.UPF",
   "Hf" => "Hf.pbe-sp-hgh.UPF",
)

# TODO This should be transformed into a proper function inside DFTK
#      ... a better interface would be based on a `basis` or a `model`
"""Ecut is in Hartree, tol is energy tolerance in Hartree, kspacing is in inverse Bohrs"""
function scf_quantum_espresso(lattice, atoms;
                              Ecut=20, tol=1e-11, mixing_mode="plain",
                              fileprefix="espresso", smearing="gaussian",
                              temperature=0, mixing_beta=0.7, mixing_ndim=10,
                              maxiter=100, kgrid=nothing, magnetic_moments=nothing,
                              input_dft="pbe", pseudomap=QE_PSEUDOMAP,
                              pseudodir=QE_PSEUDODIR, kwargs...)
    if !mpi_master()
        error("Running QE does not support MPI parallelisation")
    end

    # TODO Pseudos are just hardcoded for now
    #      ... better look at `atoms` and map to the HGH QE files
    if isnothing(magnetic_moments) || isempty(magnetic_moments)
        system = ase_atoms(lattice, atoms)
    else
        system = ase_atoms(lattice, atoms, magnetic_moments)
    end

    # Units used by ASE and QE
    Ry = 0.5
    eV = austrip(u"1eV")
    system.calc = pyimport("ase.calculators.espresso").Espresso(;
        label=fileprefix,
        input_dft=input_dft,
        pseudopotentials=pseudomap,
        kpts=kgrid,
        ecutwfc=Ecut / Ry,  # QE uses Rydbergs
        tstress=true,       # Compute stresses
        tprnfor=true,       # Compute forces
        mixing_mode=mixing_mode,
        mixing_beta=mixing_beta,
        conv_thr=tol / Ry,
        occupations=(temperature > 0 ? "smearing" : "fixed"),
        smearing=smearing,
        degauss=temperature / Ry,
        electron_maxstep=maxiter,
        mixing_ndim=mixing_ndim,
        kwargs...
    )

    results = Dict{Symbol, Any}(
        :inputfile  => fileprefix * ".pwi",
        :outputfile => fileprefix * ".pwo",
    )
    MPICH_jll.mpiexec() do mpirun
        QuantumEspresso_jll.pwscf() do pwscf
            qe_command = "$mpirun -np $(mpi_nprocs()) $pwscf -in PREFIX.pwi > PREFIX.pwo"
            @info "Running QuantumEspresso: $qe_command"
            system.calc.command = qe_command
            withenv("ESPRESSO_PSEUDO" => pseudodir,
                    "OMP_NUM_THREADS" => BLAS.get_num_threads()) do
                results[:energies_total] = system.get_potential_energy() * eV
            end
        end
    end

    (; results...)
end

function parse_quantum_espresso_pwo(file)
    Ry = 0.5  # Ry -> Hartree

    iterations = []
    current_iter = nothing

    for line in readlines(file)
        miter = match(r"^ *iteration # *([0-9]+)", line)
        if miter !== nothing  # New iteration starts
            if !isnothing(current_iter)
                push!(iterations, (; current_iter...))
            end
            current_iter = Dict{Symbol, Any}(:n_iter => parse(Int, miter[1]))
        end

        menergy = match(r"^ *total energy *= *([0-9.eE+-]+) *Ry", line)
        if menergy !== nothing
            current_iter[:total_energy] = Ry * parse(Float64, menergy[1])
        end

        maccuracy = match(r"^ *estimated scf accuracy *< *([0-9.eE+-]+) *Ry", line)
        if maccuracy !== nothing
            current_iter[:accuracy] = Ry * parse(Float64, maccuracy[1])
        end
    end

    iterations
end


using PyCall

@pyimport gym

function play_env(n::Network, env, pcfg::Dict, seed::Int64, trial::Int64,
                  train::Bool=true)
    env[:seed](seed)
    ob = env[:reset]()
    ehigh = min.(env[:observation_space][:high], 20)
    elow = max.(env[:observation_space][:low], -20)
    ma_reward = -Inf
    max_reward = 0.0
    reward = -Inf
    done = false
    step = 0
    bad_step = 0

    while ~done
        nob = min.(max.(((ob - elow) ./ (ehigh - elow)), 0.0), 1.0) .* pcfg["fr"]
        ins = rand(n.nin, pcfg["tinput"])
        inputs = falses(n.nin, pcfg["tinput"])
        for i in 1:n.nin
            inputs[i, :] = ins[i, :] .< nob[mod(i, length(nob))+1]
        end
        for t in 1:pcfg["tinput"]
            outs = step!(n, inputs[:, t], train)
        end
        ocount = zeros(pcfg["n_actions"])
        for t in 1:pcfg["toutput"]
            outs = step!(n, train)
            for o in 1:length(outs)
                if outs[o]
                    ocount[mod(o, length(ocount))+1] += 1.0
                end
            end
        end
        for t in 1:pcfg["trest"]
            step!(n, train)
        end
        maxcount = maximum(ocount)
        mincount = minimum(ocount)
        if maxcount == mincount
            bad_step += 1
        else
            bad_step = 0
        end
        ob, r, done, _ = env[:step](indmax(ocount)-1)
        reward = r
        if pcfg["id"] == "Acrobot-v1"
            reward = acos(ob[1]) # max angle
            max_reward = max(max_reward, reward)
        elseif pcfg["id"] == "CartPole-v0"
            reward = 1.0 - acos(cos(ob[3])) # 1 - angle
            max_reward += 1.0
        elseif pcfg["id"] == "MountainCar-v0"
            reward = (ob[1] - elow[1]) / (ehigh[1] - elow[1])  # mountain car
            max_reward = max(max_reward, reward)
        else
            max_reward += reward
        end
        if ma_reward == -Inf
            ma_reward = reward
        else
            ma_reward = (1.0 - pcfg["ma_rate"]) * ma_reward + pcfg["ma_rate"] * reward
        end
        dop = 0.0
        if ma_reward != 0.0
            dop = reward / ma_reward
        end
        if pcfg["onpolicy"]
            if pcfg["dmethod"] == 1
                if dop > 1.0
                    n.da[1] += (dop - 1.0)
                end
            elseif pcfg["dmethod"] == 2
                n.da[1] *= dop
            elseif pcfg["dmethod"] == 3
                n.da[1] = dop
            end
        end
        step += 1
        Logging.info(@sprintf("S: %s %d %d %d %d %0.6f %0.6f %0.6f %0.6f",
                              pcfg["id"], trial, step, maxcount, mincount,
                              max_reward, reward, ma_reward, dop))
        if bad_step > 20
            return -1e4
        end
    end

    if (pcfg["id"] == "Acrobot-v1") || (pcfg["id"] == "MountainCar-v0")
        max_reward = max_reward - step
    end

    max_reward
end

function repeat_trials(n::Network, env, pcfg::Dict)
    init_weights = copy(n.weights)
    bad_trial = 0
    ma_reward = play_env(n, env, pcfg, 0, 0, false)
    if ma_reward == -1e4
        return ma_reward
    end

    for i in 1:pcfg["n_trials"]
        weights = copy(n.weights)
        reward = play_env(n, env, pcfg, i, i, true)
        if reward == -1e4
            return reward
        end
        dop = 0.0
        if ma_reward != 0.0
            dop = reward / ma_reward
        end
        if dop > 1.0 && ~pcfg["onpolicy"]
            n.da[1] += (dop - 1.0)
        end
        ma_reward = (1.0 - pcfg["ma_rate"]) * ma_reward + pcfg["ma_rate"] * reward
        dweight = sum(abs.(n.weights - weights)) / length(n.weights)
        Logging.info(@sprintf("T: %s %d %0.6f %0.6f %0.6f %e %e",
                              pcfg["id"], i, reward, ma_reward, dop,
                              sum(abs.(n.weights - init_weights)) / length(n.weights),
                              dweight))
        if dweight < 1e-4
            bad_trial += 1
        else
            bad_trial = 0
        end
        if bad_trial > 10
            return -1e4
        end
    end

    mean(x->play_env(n, env, pcfg, pcfg["seed"]+x, pcfg["n_trials"]+x, false), 1:3)
end


"""
This is a sample input file for the CloudScat code.  To use it run julia
and type

julia> include("first_hitting.jl"); run()

The output file will be based on this file's name so to run different simulations
you can just copy this file, change the configuration and run it as above with
the new file name.  You can also run from the command line as e.g.

> julia first_hitting.jl

- - -
2010 Alejandro Luque IAA-CSIC (aluque@iaa.es)

"""

using CloudScat
using FastGaussQuadrature

const co = CloudScat.constants

function run()
    params = Params(
        # NUmber of simulated photons
        N = 5000000,

        # The wavelength of the scattered light
        λ = 337 * co.nano,

        # The source is defined as a straight line running from source_a
        # to source_b.  Within this line photons are distributed uniformly with
        # initially random, isotropic directions.  To use a point source set
        # source_b = source_a
        source_a = [0, 0, 10 * co.kilo],
        source_b = [0, 0, 10 * co.kilo],

        # For this example we neglect Rayleigh scattering
        σray = 0.)

    cth = 15 * co.kilo

    # World geometry
    # Describe the cloud geometry here.  It may consist in combinations of
    # geometrical figures (cylinders, spheres, cones, semi-planes).
    # See cloud_gemoetry.jl for a more detailed description
    #
    cloud = Plane(cth, false)

    # Geometry of the full domain.  Photons outside the domain are
    # immediately discarded.
    domain = Plane(60 * co.kilo, false)

    # Set a homogeneous cloud droplet density and radius.
    # See variable_droplet_radius.jl for how to deal with inhomogeneous clouds.
    # Radius of the scattering particles
    radius = 10e-6

    # Density of the scattering centers.
    nscat = 100 * co.centi^-3

    composition = Homogeneous(params.λ, nscat, radius)

    # Define the full simulation world
    world = World(cloud, domain, composition)

    # To capture all exiting photons we set a few observers around the closest
    # point in th cloud top to the source.  If the radiance would be perfectly
    # lambertian this would not be neccesary.
    μ, w = gausslegendre(5)

    observers = @. observerat(acos(0.5 * (μ + 1)), cth)

    basename = splitext(splitdir(@__FILE__)[2])[1]

    # The saveto option specifies the output file.  Leave it as it is to base
    # the filename on the name of this file specifying the parameters.
    CloudScat.main(params, world, observers,
                   saveto="$(basename).h5")
end

function observerat(θ, zc)
    R = 400 * co.kilo

    Observer(
        # Location of the observer
        position = [0, R * sin(θ), zc + R * cos(θ)],

        # Sampling interval
        tsample = 1e-5,

        # Collected samples
        nsamples = 10000,

        # fov of the camera (half of the diagonal fov).
        fov = 40,

        # camera pixels
        pixels = 1024)
end

# Check if the code has been 'included' or run from the shell.  In the latter
# case, run the simulation.
isinteractive() || run()


using JLD, Knet
MAJOR=1 # 1:column-major 2:row-major

typealias Word String           # represent words with strings instead of finite vocab, we will process chars
typealias PosTag UInt8          # 17 universal part-of-speech tags
typealias DepRel UInt8          # 37 universal dependency relations
typealias Position Int16        # sentence position
typealias Cost Position         # [0:nword]
typealias Move Int              # [1:nmove]
include("../src/parser.jl")

;                               # CONLLU FORMAT
immutable Sentence		# 1. ID: Word index, integer starting at 1 for each new sentence; may be a range for multiword tokens; may be a decimal number for empty nodes.
    word::Vector{Word}          # 2. FORM: Word form or punctuation symbol.
    #stem::Vector{Stem}         # 3. LEMMA: Lemma or stem of word form.
    postag::Vector{PosTag}      # 4. UPOSTAG: Universal part-of-speech tag.
    #xpostag::Vector{Xpostag}   # 5. XPOSTAG: Language-specific part-of-speech tag; underscore if not available.
    #feats::Vector{Feats}       # 6. FEATS: List of morphological features from the universal feature inventory or from a defined language-specific extension; underscore if not available.
    head::Vector{Position}      # 7. HEAD: Head of the current word, which is either a value of ID or zero (0).
    deprel::Vector{DepRel}      # 8. DEPREL: Universal dependency relation to the HEAD (root iff HEAD = 0) or a defined language-specific subtype of one.
    #deps::Vector{Deps}         # 9. DEPS: Enhanced dependency graph in the form of a list of head-deprel pairs.
    #misc::Vector{Misc}         # 10. MISC: Any other annotation.
    #vocab::Vocab               # Common repository of symbols for upostag, deprel etc.
end
Sentence()=Sentence([],[],[],[])
Base.length(s::Sentence)=length(s.word)

immutable Vocab
    cdict::Dict{Char,Int}       # character dictionary (input)
    idict::Dict{String,Int}     # word dictionary (input, filled in by maptoint)
    odict::Dict{String,Int}     # word dictionary (output)
    sosword::String             # start-of-sentence word (input)
    eosword::String             # end-of-sentence word (input)
    unkword::String             # unknown word (output, input does not have unk)
    sowchar::Char               # start-of-word char
    eowchar::Char               # end-of-word char
    unkchar::Char               # unknown char
    postags::Dict{String,PosTag}
    deprels::Dict{String,DepRel}
end

type BeamState
    parent::BeamState   # previous state
    move::Move          # move executed to get to this state
    score::Float64      # cumulative score (logp) TODO: should be unnormalized logp!
    cost::Cost          # cumulative cost (gold arcs that have become impossible)
    fidx::Int           # column in feature matrix representing current parser state # TODO: what is this?
    parser::Parser 	# current parser state
    BeamState()=new()
    BeamState(p::Parser)=new(NullBeamState,0,0,0,0,p)
    BeamState(b::BeamState,m::Integer,s::Number,c::Integer)=
        new(b,convert(Move,m),convert(Float64,s),convert(Cost,c))
end

const NullBeamState = BeamState()
Base.isless(a::BeamState,b::BeamState)=(a.score < b.score)

type Beam
    beam::Vector{BeamState}
    cand::Vector{BeamState}
    sent::Sentence
    stop::Bool                  # TODO: what is this?
    function Beam{T<:Parser}(pt::Type{T}, s::Sentence)
        p = pt(s)
        b = BeamState(p)
        new(BeamState[b], BeamState[], s, !anyvalidmoves(p))
    end
end

typealias Batch Vector{Beam}

# parse a minibatch of sentences using beam search, global normalization, early updates and report loss.
function parseloss(model, sentences, vocab, parsertype, beamsize)
    # global wembed,forw,back,cdata,cmask,wdata,wmask
    words,sents,wordlen,sentlen = maptoint(sentences, vocab)
    # Get word embeddings
    sow,eow = vocab.cdict[vocab.sowchar],vocab.cdict[vocab.eowchar]
    cdata,cmask = tokenbatch(words,wordlen,sow,eow) # cdata/cmask[T+2][V] where T: longest word (140), V: input word vocab (19674)
    wembed = charlstm(model,cdata,cmask)            # wembed[C,V] where V: input word vocab, C: charlstm hidden (350)
    # Get context embeddings
    sos,eos = vocab.idict[vocab.sosword],vocab.idict[vocab.eosword]
    wdata,wmask = tokenbatch(sents,sentlen,sos,eos) # wdata/wmask[T+2][B] where T: longest sentence (159), B: batch size (12543)
    forw,back = wordlstm(model,wdata,wmask,wembed)  # forw/back[T][W,B] where T: longest sentence, B: batch size, W: wordlstm hidden (300)
    nothing
end

# function loss(model, sentences, parsertype, beamsize)
function lmloss(model, sentences, vocab::Vocab; result=nothing)
    # map words and chars to Ints
    # words and sents contain Int (char/word id) arrays without sos/eos tokens
    # TODO: optionally construct cdict here as well
    words,sents,wordlen,sentlen = maptoint(sentences, vocab)

    # Find word vectors for a batch of sentences using character lstm or cnn model
    # note that tokenbatch adds sos/eos tokens and padding to each sequence
    sow,eow = vocab.cdict[vocab.sowchar],vocab.cdict[vocab.eowchar]
    cdata,cmask = tokenbatch(words,wordlen,sow,eow)
    wembed = charlstm(model,cdata,cmask)

    # Find context vectors using word blstm
    sos,eos = vocab.idict[vocab.sosword],vocab.idict[vocab.eosword]
    wdata,wmask = tokenbatch(sents,sentlen,sos,eos)
    forw,back = wordlstm(model,wdata,wmask,wembed)

    # Test predictions
    unk = vocab.odict[vocab.unkword]
    odata,omask = goldbatch(sentences,sentlen,vocab.odict,unk)
    total = lmloss1(model,odata,omask,forw,back; result=result)
    return -total/length(sentences)
end

function lmloss1(model,data,mask,forw,back; result=nothing)
    # data[t][b]::Int gives the correct word at sentence b, position t
    T = length(data); if !(T == length(forw) == length(back) == length(mask)); error(); end
    B = length(data[1])
    weight,bias = wsoft(model),bsoft(model)
    prd(t) = (if MAJOR==1; weight * vcat(forw[t],back[t]) .+ bias; else; hcat(forw[t],back[t]) * weight .+ bias; end)
    idx(t,b,n) = (if MAJOR==1; data[t][b] + (b-1)*n; else; b + (data[t][b]-1)*n; end)
    total = count = 0
    @inbounds for t=1:T
        ypred = prd(t)
        nrows,ncols = size(ypred)
        index = Int[]
        for b=1:B
            if mask[t][b]==1
                push!(index, idx(t,b,nrows))
            end
        end
        o1 = logp(ypred,MAJOR)
        o2 = o1[index]
        total += sum(o2)
        count += length(o2)
    end
    if result != nothing; result[1]+=AutoGrad.getval(total); result[2]+=count; end
    return total
end

# Find all unique words, assign an id to each, and lookup their characters in cdict
# TODO: construct/add-to cdict here as well?
# in vocab: reads sosword, eosword, unkchar, cdict; writes idict
function maptoint(sentences, v::Vocab)
    wdict = empty!(v.idict)
    cdict = v.cdict
    unkcid = cdict[v.unkchar]
    words = Vector{Int}[]
    sents = Vector{Int}[]
    wordlen = 0
    sentlen = 0
    @inbounds for w in (v.sosword,v.eosword)
        wid = get!(wdict, w, 1+length(wdict))
        word = Array(Int, length(w))
        wordi = 0
        for c in w
            word[wordi+=1] = get(cdict, c, unkcid)
        end
        if wordi != length(w); error(); end
        if wordi > wordlen; wordlen = wordi; end
        push!(words, word)
    end
    @inbounds for s in sentences
        sent = Array(Int, length(s.word))
        senti = 0
        for w in s.word
            ndict = length(wdict)
            wid = get!(wdict, w, 1+ndict)
            sent[senti+=1] = wid
            if wid == 1+ndict
                word = Array(Int, length(w))
                wordi = 0
                for c in w
                    word[wordi+=1] = get(cdict, c, unkcid)
                end
                if wordi != length(w); error(); end
                if wordi > wordlen; wordlen = wordi; end
                push!(words, word)
            end
        end
        if senti != length(s.word); error(); end
        if senti > sentlen; sentlen = senti; end
        push!(sents, sent)
    end
    if length(wdict) != length(words); error("wdict=$(length(wdict)) words=$(length(words))"); end
    return words,sents,wordlen,sentlen
end

# Create token batches, adding start/end tokens and masks, pad at the beginning
function tokenbatch(sequences,maxlen,sos,eos,pad=eos)
    B = length(sequences)
    T = maxlen + 2
    data = [ Array(Int,B) for t in 1:T ]
    mask = [ Array(Float32,B) for t in 1:T ]
    @inbounds for t in 1:T
        for b in 1:B
            N = length(sequences[b])
            n = t - T + N + 1
            if n < 0
                mask[t][b] = 0
                data[t][b] = pad
            else
                mask[t][b] = 1
                if n == 0
                    data[t][b] = sos
                elseif n <= N
                    data[t][b] = sequences[b][n]
                elseif n == N+1
                    data[t][b] = eos
                else
                    error()
                end
            end
        end
    end
    return data,mask
end

function goldbatch(sentences, maxlen, wdict, unkwid, pad=unkwid)
    B = length(sentences)
    T = maxlen # no need for sos/eos for gold
    data = [ Array(Int,B) for t in 1:T ]
    mask = [ Array(Float32,B) for t in 1:T ]
    @inbounds for t in 1:T
        for b in 1:B
            N = length(sentences[b])
            n = t - T + N
            if n <= 0
                mask[t][b] = 0
                data[t][b] = pad
            else
                mask[t][b] = 1
                data[t][b] = get(wdict, sentences[b].word[n], unkwid)
            end
        end
    end
    return data,mask
end

# Run charid arrays through the LSTM, collect last hidden state as word embedding
function charlstm(model,data,mask)
    weight,bias,embeddings = wchar(model),bchar(model),echar(model)
    T = length(data)
    B = length(data[1])
    H = div(length(bias),4)
    cembed(t)=(if MAJOR==1; embeddings[:,data[t]]; else; embeddings[data[t],:]; end)
    czero=(if MAJOR==1; fill!(similar(bias,H,B), 0); else; fill!(similar(bias,B,H), 0); end)
    hidden = cell = czero       # TODO: cache this
    mask = map(KnetArray, mask) # TODO: dont hardcode atype
    @inbounds for t in 1:T
        (hidden,cell) = lstm(weight,bias,hidden,cell,cembed(t);mask=mask[t])
    end
    return hidden
end

function wordlstm(model,data,mask,embeddings) # col-major
    weight,bias = wforw(model),bforw(model)
    T = length(data)
    B = length(data[1])
    H = div(length(bias),4)
    wembed(t)=(if MAJOR==1; embeddings[:,data[t]]; else; embeddings[data[t],:]; end)
    wzero=(if MAJOR==1; fill!(similar(bias,H,B), 0); else; fill!(similar(bias,B,H), 0); end)
    hidden = cell = wzero
    mask = map(KnetArray, mask)           # TODO: dont hardcode atype
    forw = Array(Any,T-2)       # exclude sos/eos
    @inbounds for t in 1:T-2
        (hidden,cell) = lstm(weight,bias,hidden,cell,wembed(t); mask=mask[t])
        forw[t] = hidden
    end
    weight,bias = wback(model),bback(model)
    if H != div(length(bias),4); error(); end
    hidden = cell = wzero
    back = Array(Any,T-2)
    @inbounds for t in T:-1:3
        (hidden,cell) = lstm(weight,bias,hidden,cell,wembed(t); mask=mask[t])
        back[t-2] = hidden
    end
    return forw,back
end

function lstm(weight,bias,hidden,cell,input; mask=nothing)
    if MAJOR==1
        gates   = weight * vcat(input,hidden) .+ bias
        H       = size(hidden,1)
        forget  = sigm(gates[1:H,:])
        ingate  = sigm(gates[1+H:2H,:])
        outgate = sigm(gates[1+2H:3H,:])
        change  = tanh(gates[1+3H:4H,:])
        if mask!=nothing; mask=reshape(mask,1,length(mask)); end
    else
        gates   = hcat(input,hidden) * weight .+ bias
        H       = size(hidden,2)
        forget  = sigm(gates[:,1:H])
        ingate  = sigm(gates[:,1+H:2H])
        outgate = sigm(gates[:,1+2H:3H])
        change  = tanh(gates[:,1+3H:end])
    end
    cell    = cell .* forget + ingate .* change
    hidden  = outgate .* tanh(cell)
    if mask != nothing
        hidden = hidden .* mask
        cell = cell .* mask
    end
    return (hidden,cell)
end

function loadvocab()
    a = load("english_vocabs.jld")
    global _vocab = Vocab(a["char_vocab"],Dict{String,Int}(),a["word_vocab"],
                          "<s>","</s>","<unk>",'↥','Ϟ','⋮',
                          UPOSTAG, UDEPREL)
end

function loadcorpus(v::Vocab)
    global _corpus = Any[]
    s = Sentence()
    for line in eachline("en-ud-train.conllu")
        if line == "\n"
            push!(_corpus, s)
            s = Sentence()
        elseif (m = match(r"^\d+\t(.+?)\t.+?\t(.+?)\t.+?\t.+?\t(.+?)\t(.+?)\t", line)) != nothing
            #                id   word   lem  upos   xpos feat head   deprel
            push!(s.word, m.captures[1])
            push!(s.postag, v.postags[m.captures[2]])
            push!(s.head, parse(Position,m.captures[3]))
            push!(s.deprel, v.deprels[m.captures[4]])
        end
    end
end

function loadmodel()
    global _model = Dict()
    m = load("english_ch12kmodel_2.jld")["model"]
    if MAJOR==1
        for k in (:cembed,); _model[k] = KnetArray(m[k]'); end
        for k in (:forw,:soft,:back,:char); _model[k] = map(KnetArray, m[k]'); end
    else
        for k in (:cembed,); _model[k] = KnetArray(m[k]); end
        for k in (:forw,:soft,:back,:char); _model[k] = map(KnetArray, m[k]); end
    end
end

wsoft(model)=model[:soft][1]
bsoft(model)=model[:soft][2]
wchar(model)=model[:char][1]
bchar(model)=model[:char][2]
echar(model)=model[:cembed]
wforw(model)=model[:forw][1]
bforw(model)=model[:forw][2]
wback(model)=model[:back][1]
bback(model)=model[:back][2]

function minibatch(corpus, batchsize)
    data = Any[]
    sorted = sort(corpus, by=length)
    for i in 1:batchsize:length(corpus)
        j = min(length(corpus), i+batchsize-1)
        push!(data, sorted[i:j])
    end
    shuffle(data)
end

function minibatch1(corpus, batchsize)
    data = Any[]
    dict = Dict{Int,Vector{Int}}()
    for i in 1:length(corpus)
        s = corpus[i]
        l = length(s)
        a = get!(dict, l, Int[])
        push!(a, i)
        if length(a) == batchsize
            push!(data, corpus[a])
            empty!(a)
        end
    end
    for (k,a) in dict
        if !isempty(a)
            push!(data, corpus[a])
        end
    end
    return data
end

function main(;batch=3000,nsent=0,result=zeros(2))
    loadmodel()
    loadvocab()
    loadcorpus(_vocab)
    c = (nsent == 0 ? _corpus : _corpus[1:nsent])
    global _data = minibatch(c, batch)
    # global _data = [ _corpus[i:min(i+999,length(_corpus))] for i=1:1000:length(_corpus) ]
    # global _data = [ _corpus[1:1] ]
    for d in _data
        lmloss(_model, d, _vocab; result=result)
    end
    println(exp(-result[1]/result[2]))
    return result
end

function getloss(d; result=zeros(2))
    lmloss(_model, d, _vocab; result=result)
    println(exp(-result[1]/result[2]))
    return result
end

# Universal POS tags (17)
const UPOSTAG = Dict{String,PosTag}(
"ADJ"   => 1, # adjective
"ADP"   => 2, # adposition
"ADV"   => 3, # adverb
"AUX"   => 4, # auxiliary
"CCONJ" => 5, # coordinating conjunction
"DET"   => 6, # determiner
"INTJ"  => 7, # interjection
"NOUN"  => 8, # noun
"NUM"   => 9, # numeral
"PART"  => 10, # particle
"PRON"  => 11, # pronoun
"PROPN" => 12, # proper noun
"PUNCT" => 13, # punctuation
"SCONJ" => 14, # subordinating conjunction
"SYM"   => 15, # symbol
"VERB"  => 16, # verb
"X"     => 17, # other
]

# Universal Dependency Relations (37)
const UDEPREL = Dict{String,DepRel}(
"root"       => 1,  # root
"acl"        => 2,  # clausal modifier of noun (adjectival clause)
"advcl"      => 3,  # adverbial clause modifier
"advmod"     => 4,  # adverbial modifier
"amod"       => 5,  # adjectival modifier
"appos"      => 6,  # appositional modifier
"aux"        => 7,  # auxiliary
"case"       => 8,  # case marking
"cc"         => 9,  # coordinating conjunction
"ccomp"      => 10, # clausal complement
"clf"        => 11, # classifier
"compound"   => 12, # compound
"conj"       => 13, # conjunct
"cop"        => 14, # copula
"csubj"      => 15, # clausal subject
"dep"        => 16, # unspecified dependency
"det"        => 17, # determiner
"discourse"  => 18, # discourse element
"dislocated" => 19, # dislocated elements
"expl"       => 20, # expletive
"fixed"      => 21, # fixed multiword expression
"flat"       => 22, # flat multiword expression
"goeswith"   => 23, # goes with
"iobj"       => 24, # indirect object
"list"       => 25, # list
"mark"       => 26, # marker
"nmod"       => 27, # nominal modifier
"nsubj"      => 28, # nominal subject
"nummod"     => 29, # numeric modifier
"obj"        => 30, # object
"obl"        => 31, # oblique nominal
"orphan"     => 32, # orphan
"parataxis"  => 33, # parataxis
"punct"      => 34, # punctuation
"reparandum" => 35, # overridden disfluency
"vocative"   => 36, # vocative
"xcomp"      => 37, # open clausal complement
]

# """
# The input is minibatched and processed character based.
# The blstm language model used for preinit inputs characters outputs words.
# The characters of a word is used to obtain its embedding using an RNN or CNN.
# The input to the parser is minibatched sentences.
# The sentence lengths in a minibatch do not have to exactly match, we will support padding.
# """
# foo1

# """
# Let B be the minibatch size, and T the sentence length in the minibatch.
# input[t][b] is a string specifying the t'th word of the b'th sentence.
# Do we include SOS/EOS tokens in input?  No, the parser does not need them.
# How about the ROOT token?  No, our parsers are written s.t. the root token is implicit.
# """
# foo2

# """
# We need head and deprel for loss calculation. It makes more sense to
# process an array of sentences, doing the grouping inside.
# """
# foo3


function Relativity(arg0::jdouble)
    return Relativity((jdouble,), arg0)
end

function acceleration(obj::Relativity, arg0::FieldSpacecraftState, arg1::Vector{RealFieldElement})
    return jcall(obj, "acceleration", FieldVector3D, (FieldSpacecraftState, Vector{RealFieldElement}), arg0, arg1)
end

function acceleration(obj::Relativity, arg0::SpacecraftState, arg1::Vector{jdouble})
    return jcall(obj, "acceleration", Vector3D, (SpacecraftState, Vector{jdouble}), arg0, arg1)
end

function depends_on_position_only(obj::Relativity)
    return jcall(obj, "dependsOnPositionOnly", jboolean, ())
end

function get_events_detectors(obj::Relativity)
    return jcall(obj, "getEventsDetectors", Stream, ())
end

function get_field_events_detectors(obj::Relativity, arg0::Field)
    return jcall(obj, "getFieldEventsDetectors", Stream, (Field,), arg0)
end

function get_parameters_drivers(obj::Relativity)
    return jcall(obj, "getParametersDrivers", Vector{ParameterDriver}, ())
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMBase.jl/blob/master/LICENSE

using FEMBase
using Base.Test

mutable struct Dummy <: FEMBase.FieldProblem
end

@testset "add elements to problem" begin
    problem = Problem(Dummy, "test", 2)
    element = Element(Quad4, [1, 2, 3, 4])
    elements = [element]
    add_elements!(problem, elements)
    @test problem.elements[1] == element
end


using Core.Intrinsics: llvmcall
using Base.Threads

import Base.Threads: atomictypes, llvmtypes, inttype, ArithmeticTypes, FloatTypes,
       atomic_cas!,
       atomic_xchg!,
       atomic_add!, atomic_sub!,
       atomic_and!, atomic_nand!, atomic_or!, atomic_xor!,
       atomic_max!, atomic_min!

import Base.Sys: ARCH, WORD_SIZE

for typ ∈ atomictypes
    lt = llvmtypes[typ]
    ilt = llvmtypes[inttype(typ)]
    # Note: atomic_cas! succeeded (i.e. it stored "new") if and only if the result is "cmp"
    if typ <: Integer
        @eval atomic_cas!(x::Ptr{$typ}, cmp::$typ, new::$typ) =
            llvmcall($"""
                     %ptr = inttoptr i$WORD_SIZE %0 to $lt*
                     %rs = cmpxchg $lt* %ptr, $lt %1, $lt %2 acq_rel acquire
                     %rv = extractvalue { $lt, i1 } %rs, 0
                     ret $lt %rv
                     """, $typ, Tuple{Ptr{$typ},$typ,$typ},
                     x, cmp, new)
    else
        @eval atomic_cas!(x::Ptr{$typ}, cmp::$typ, new::$typ) =
            llvmcall($"""
                     %iptr = inttoptr i$WORD_SIZE %0 to $ilt*
                     %icmp = bitcast $lt %1 to $ilt
                     %inew = bitcast $lt %2 to $ilt
                     %irs = cmpxchg $ilt* %iptr, $ilt %icmp, $ilt %inew acq_rel acquire
                     %irv = extractvalue { $ilt, i1 } %irs, 0
                     %rv = bitcast $ilt %irv to $lt
                     ret $lt %rv
                     """, $typ, Tuple{Ptr{$typ},$typ,$typ},
                     x, cmp, new)
    end

    arithmetic_ops = [:add, :sub]
    for rmwop in [arithmetic_ops..., :xchg, :and, :nand, :or, :xor, :max, :min]
        rmw = string(rmwop)
        fn = Symbol("atomic_", rmw, "!")
        if (rmw == "max" || rmw == "min") && typ <: Unsigned
            # LLVM distinguishes signedness in the operation, not the integer type.
            rmw = "u" * rmw
        end
        if rmwop in arithmetic_ops && !(typ <: ArithmeticTypes) continue end
        if typ <: Integer
            @eval $fn(x::Ptr{$typ}, v::$typ) =
                llvmcall($"""
                         %ptr = inttoptr i$WORD_SIZE %0 to $lt*
                         %rv = atomicrmw $rmw $lt* %ptr, $lt %1 acq_rel
                         ret $lt %rv
                         """, $typ, Tuple{Ptr{$typ}, $typ}, x, v)
        else
            rmwop == :xchg || continue
            @eval $fn(x::Ptr{$typ}, v::$typ) =
                llvmcall($"""
                         %iptr = inttoptr i$WORD_SIZE %0 to $ilt*
                         %ival = bitcast $lt %1 to $ilt
                         %irv = atomicrmw $rmw $ilt* %iptr, $ilt %ival acq_rel
                         %rv = bitcast $ilt %irv to $lt
                         ret $lt %rv
                         """, $typ, Tuple{Ptr{$typ}, $typ}, x, v)
        end
    end
end

const opnames = Dict{Symbol, Symbol}(:+ => :add, :- => :sub)
for op in [:+, :-, :max, :min]
    opname = get(opnames, op, op)
    @eval function $(Symbol("atomic_", opname, "!"))(var::Ptr{T}, val::T) where T<:FloatTypes
        IT = inttype(T)
        old = unsafe_load(var)
        while true
            new = $op(old, val)
            cmp = old
            old = atomic_cas!(var, cmp, new)
            reinterpret(IT, old) == reinterpret(IT, cmp) && return new
            # Temporary solution before we have gc transition support in codegen.
            ccall(:jl_gc_safepoint, Cvoid, ())
        end
    end
end


# Uniforms are OpenGL variables that stay the same for the entirety of a drawcall.
# There are a lot of functions, to upload them, as OpenGL doesn't rely on multiple dispatch.
# here is my approach, to handle all of the uniforms with one function, namely gluniform
# For uniforms, the Vector and Matrix types from ImmutableArrays should be used, as they map the relation almost 1:1

const GLSL_COMPATIBLE_NUMBER_TYPES = (GLfloat, GLint, GLuint, GLdouble)
const NATIVE_TYPES = Union{
    StaticArray, GLSL_COMPATIBLE_NUMBER_TYPES...,
    ZeroIndex{GLint}, ZeroIndex{GLuint},
    GLBuffer, GPUArray, Shader, GLProgram, NativeMesh
}

opengl_prefix(T)  = error("Object $T is not a supported uniform element type")
opengl_postfix(T) = error("Object $T is not a supported uniform element type")

opengl_prefix(x::Type{T}) where {T <: Union{FixedPoint, Float32, Float16}} = ""
opengl_prefix(x::Type{T}) where {T <: Float64} = "d"
opengl_prefix(x::Type{Cint}) = "i"
opengl_prefix(x::Type{T}) where {T <: Union{Cuint, UInt8, UInt16}} = "u"

opengl_postfix(x::Type{Float64}) = "dv"
opengl_postfix(x::Type{Float32}) = "fv"
opengl_postfix(x::Type{Cint})    = "iv"
opengl_postfix(x::Type{Cuint})   = "uiv"

function uniformfunc(typ::DataType, dims::Tuple{Int})
    Symbol(string("glUniform", first(dims), opengl_postfix(typ)))
end
function uniformfunc(typ::DataType, dims::Tuple{Int, Int})
    M, N = dims
    Symbol(string("glUniformMatrix", M == N ? "$M" : "$(M)x$(N)", opengl_postfix(typ)))
end

function gluniform(location::Integer, x::FSA) where FSA <: Union{StaticArray, Colorant}
    xref = [x]
    gluniform(location, xref)
end

_size(p) = size(p)
_size(p::Colorant) = (length(p),)
_size(p::Type{T}) where {T <: Colorant} = (length(p),)
_ndims(p) = ndims(p)
_ndims(p::Type{T}) where {T <: Colorant} = 1

@generated function gluniform(location::Integer, x::Vector{FSA}) where FSA <: Union{StaticArray, Colorant}
    func = uniformfunc(eltype(FSA), _size(FSA))
    callexpr = if _ndims(FSA) == 2
        :($func(location, length(x), GL_FALSE, x))
    else
        :($func(location, length(x), x))
    end
    quote
        $callexpr
    end
end

#Some additional uniform functions, not related to Imutable Arrays
gluniform(location::Integer, target::Integer, t::Texture) = gluniform(GLint(location), GLint(target), t)
gluniform(location::Integer, target::Integer, t::GPUVector) = gluniform(GLint(location), GLint(target), t.buffer)
gluniform(location::Integer, target::Integer, t::Observable) = gluniform(GLint(location), GLint(target), to_value(t))
gluniform(location::Integer, target::Integer, t::TextureBuffer) = gluniform(GLint(location), GLint(target), t.texture)
function gluniform(location::GLint, target::GLint, t::Texture)
    activeTarget = GL_TEXTURE0 + UInt32(target)
    glActiveTexture(activeTarget)
    glBindTexture(t.texturetype, t.id)
    gluniform(location, target)
end
gluniform(location::Integer, x::Enum) = gluniform(GLint(location), GLint(x))

function gluniform(loc::Integer, x::Observable{T}) where T
    gluniform(GLint(loc), to_value(x))
end

gluniform(location::Integer, x::Union{GLubyte, GLushort, GLuint}) = glUniform1ui(GLint(location), x)
gluniform(location::Integer, x::Union{GLbyte, GLshort, GLint, Bool}) = glUniform1i(GLint(location),  x)
gluniform(location::Integer, x::GLfloat) = glUniform1f(GLint(location),  x)
gluniform(location::Integer, x::GLdouble) = glUniform1d(GLint(location),  x)

#Uniform upload functions for julia arrays...
gluniform(location::GLint, x::Vector{Float32}) = glUniform1fv(location,  length(x), x)
gluniform(location::GLint, x::Vector{GLdouble}) = glUniform1dv(location,  length(x), x)
gluniform(location::GLint, x::Vector{GLint}) = glUniform1iv(location,  length(x), x)
gluniform(location::GLint, x::Vector{GLuint}) = glUniform1uiv(location, length(x), x)

glsl_typename(x::T) where {T} = glsl_typename(T)
glsl_typename(t::DataType) = error("Datatype $(t) not supported")
glsl_typename(t::Type{Nothing}) = "Nothing"
glsl_typename(t::Type{GLfloat}) = "float"
glsl_typename(t::Type{GLdouble}) = "double"
glsl_typename(t::Type{GLuint}) = "uint"
glsl_typename(t::Type{GLint}) = "int"
glsl_typename(t::Type{T}) where {T <: Union{StaticVector, Colorant}} = string(opengl_prefix(eltype(T)), "vec", length(T))
glsl_typename(t::Type{TextureBuffer{T}}) where {T} = string(opengl_prefix(eltype(T)), "samplerBuffer")

function glsl_typename(t::Texture{T, D}) where {T, D}
    str = string(opengl_prefix(eltype(T)), "sampler", D, "D")
    t.texturetype == GL_TEXTURE_2D_ARRAY && (str *= "Array")
    str
end

function glsl_typename(t::Type{T}) where T <: SMatrix
    M, N = size(t)
    string(opengl_prefix(eltype(t)), "mat", M==N ? M : string(M, "x", N))
end
toglsltype_string(t::Observable) = toglsltype_string(to_value(t))
toglsltype_string(x::T) where {T<:Union{Real, StaticArray, Texture, Colorant, TextureBuffer, Nothing}} = "uniform $(glsl_typename(x))"
#Handle GLSL structs, which need to be addressed via single fields
function toglsltype_string(x::T) where T
    if isa_gl_struct(x)
        string("uniform ", T.name.name)
    else
        error("can't splice $T into an OpenGL shader. Make sure all fields are of a concrete type and isbits(FieldType)-->true")
    end
end
toglsltype_string(t::Union{GLBuffer{T}, GPUVector{T}}) where {T} = string("in ", glsl_typename(T))
# Gets used to access a
function glsl_variable_access(keystring, t::Texture{T, D}) where {T,D}
    fields = SubString("rgba", 1, length(T))
    if t.texturetype == GL_TEXTURE_BUFFER
        return string("texelFetch(", keystring, "index).", fields, ";")
    end
    return string("getindex(", keystring, "index).", fields, ";")
end
function glsl_variable_access(keystring, ::Union{Real, GLBuffer, GPUVector, StaticArray, Colorant})
    string(keystring, ";")
end
function glsl_variable_access(keystring, s::Observable)
    glsl_variable_access(keystring, to_value(s))
end
function glsl_variable_access(keystring, t::Any)
    error("no glsl variable calculation available for : ", keystring, " of type ", typeof(t))
end

function uniform_name_type(program::GLuint)
    uniformLength = glGetProgramiv(program, GL_ACTIVE_UNIFORMS)
    Dict{Symbol, GLenum}(ntuple(uniformLength) do i # take size and name
        name, typ = glGetActiveUniform(program, i-1)
    end)
end
function attribute_name_type(program::GLuint)
    uniformLength = glGetProgramiv(program, GL_ACTIVE_ATTRIBUTES)
    Dict{Symbol, GLenum}(ntuple(uniformLength) do i
    	name, typ = glGetActiveAttrib(program, i-1)
    end)
end
function istexturesampler(typ::GLenum)
    return (
        typ == GL_SAMPLER_BUFFER || typ == GL_INT_SAMPLER_BUFFER || typ == GL_UNSIGNED_INT_SAMPLER_BUFFER ||
    	typ == GL_IMAGE_2D ||
        typ == GL_SAMPLER_1D || typ == GL_SAMPLER_2D || typ == GL_SAMPLER_3D ||
        typ == GL_UNSIGNED_INT_SAMPLER_1D || typ == GL_UNSIGNED_INT_SAMPLER_2D || typ == GL_UNSIGNED_INT_SAMPLER_3D ||
        typ == GL_INT_SAMPLER_1D || typ == GL_INT_SAMPLER_2D || typ == GL_INT_SAMPLER_3D ||
        typ == GL_SAMPLER_1D_ARRAY || typ == GL_SAMPLER_2D_ARRAY ||
        typ == GL_UNSIGNED_INT_SAMPLER_1D_ARRAY || typ == GL_UNSIGNED_INT_SAMPLER_2D_ARRAY ||
        typ == GL_INT_SAMPLER_1D_ARRAY || typ == GL_INT_SAMPLER_2D_ARRAY
    )
end

gl_promote(x::Type{T}) where {T <: Integer} = Cint
gl_promote(x::Type{Union{Int16, Int8}}) = x

gl_promote(x::Type{T}) where {T <: Unsigned} = Cuint
gl_promote(x::Type{Union{UInt16, UInt8}}) = x

gl_promote(x::Type{T}) where {T <: AbstractFloat} = Float32
gl_promote(x::Type{Float16}) = x

gl_promote(x::Type{T}) where {T <: Normed} = N0f32
gl_promote(x::Type{N0f16}) = x
gl_promote(x::Type{N0f8}) = x

const Color3{T} = Colorant{T, 3}
const Color4{T} = Colorant{T, 4}

gl_promote(x::Type{Bool})                  = GLboolean
gl_promote(x::Type{T}) where {T <: Gray} = Gray{gl_promote(eltype(T))}
gl_promote(x::Type{T}) where {T <: Color3} = RGB{gl_promote(eltype(T))}
gl_promote(x::Type{T}) where {T <: Color4} = RGBA{gl_promote(eltype(T))}
gl_promote(x::Type{T}) where {T <: BGRA} = BGRA{gl_promote(eltype(T))}
gl_promote(x::Type{T}) where {T <: BGR} = BGR{gl_promote(eltype(T))}

gl_promote(x::Type{T}) where {T <: StaticVector} = similar_type(T, gl_promote(eltype(T)))

gl_promote(x::Type{T}) where {T <: GeometryBasics.Mesh} = NativeMesh{T}

gl_convert(x::AbstractVector{Vec3f}) = x

gl_convert(x::T) where {T <: Number} = gl_promote(T)(x)
gl_convert(x::T) where {T <: Colorant} = gl_promote(T)(x)
gl_convert(x::T) where {T <: AbstractMesh} = gl_convert(x)
gl_convert(x::T) where {T <: GeometryBasics.Mesh} = gl_promote(T)(x)
gl_convert(x::Observable{T}) where {T <: GeometryBasics.Mesh} = gl_promote(T)(x)

gl_convert(s::Vector{Matrix{T}}) where {T<:Colorant} = Texture(s)
gl_convert(s::Nothing) = s

isa_gl_struct(x::AbstractArray) = false
isa_gl_struct(x::NATIVE_TYPES) = false
isa_gl_struct(x::Colorant) = false
function isa_gl_struct(x::T) where T
    !isconcretetype(T) && return false
    if T <: Tuple
        return false
    end
    fnames = fieldnames(T)
    !isempty(fnames) && all(name -> isconcretetype(fieldtype(T, name)) && isbits(getfield(x, name)), fnames)
end
function gl_convert_struct(x::T, uniform_name::Symbol) where T
    if isa_gl_struct(x)
        return Dict{Symbol, Any}(map(fieldnames(x)) do name
            (Symbol("$uniform_name.$name") => gl_convert(getfield(x, name)))
        end)
    else
        error("can't convert $x to a OpenGL type. Make sure all fields are of a concrete type and isbits(FieldType)-->true")
    end
end

# native types don't need convert!
gl_convert(a::T) where {T <: NATIVE_TYPES} = a
gl_convert(s::Observable{T}) where {T <: NATIVE_TYPES} = s
gl_convert(s::Observable{T}) where T = const_lift(gl_convert, s)
gl_convert(x::StaticVector{N, T}) where {N, T} = map(gl_promote(T), x)
gl_convert(x::SMatrix{N, M, T}) where {N, M, T} = map(gl_promote(T), x)
gl_convert(a::AbstractVector{<: AbstractFace}) = indexbuffer(s)
gl_convert(t::Type{T}, a::T; kw_args...) where T <: NATIVE_TYPES = a
gl_convert(::Type{<: GPUArray}, a::StaticVector) = gl_convert(a)

function gl_convert(T::Type{<: GPUArray}, a::AbstractArray{X, N}; kw_args...) where {X, N}
    T(convert(AbstractArray{gl_promote(X), N}, a); kw_args...)
end

gl_convert(::Type{<: GLBuffer}, x::GLBuffer; kw_args...) = x
gl_convert(::Type{Texture}, x::Texture) = x
gl_convert(::Type{<: GPUArray}, x::GPUArray) = x

function gl_convert(::Type{T}, a::Vector{Array{X, 2}}; kw_args...) where {T <: Texture, X}
    T(a; kw_args...)
end
gl_convert(::Type{<: GPUArray}, a::Observable{<: StaticVector}) = gl_convert(a)

function gl_convert(::Type{T}, a::Observable{<: AbstractArray{X, N}}; kw_args...) where {T <: GPUArray, X, N}
    TGL = gl_promote(X)
    s = (X == TGL) ? a : lift(x-> convert(Array{TGL, N}, x), a)
    T(s; kw_args...)
end

lift_convert(a::AbstractArray, T, N) = lift(x -> convert(Array{T, N}, x), a)
function lift_convert(a::ShaderAbstractions.Sampler, T, N)
    ShaderAbstractions.Sampler(
        lift(x -> convert(Array{T, N}, x.data), a),
        minfilter = a[].minfilter, magfilter = a[].magfilter,
        x_repeat = a[].repeat[1],
        y_repeat = a[].repeat[min(2, N)],
        z_repeat = a[].repeat[min(3, N)],
        anisotropic = a[].anisotropic, swizzle_mask = a[].swizzle_mask
    )
end

gl_convert(f::Function, a) = f(a)


"""
```
DSGEVECM{T} <: AbstractDSGEVECMModel{T}
```
implements a simple interface for combining a given DSGE model
with a VECM (a VAR with cointegrating terms) to create a DSGE-VECM.
Confer with Del Negro and Schorfheide (2006)
and/or Del Negro, Schorfheide, Smets, and Wouters (2007) for details about DSGE-VECMs.
We recommend the first paper as an initial introduction to DSGE-VECMs.

The recommended constructor requires the user to provide
(1) an `AbstractDSGEModel` object, (2) which structural shocks from
the DSGE to use, and (3) the subspec (optional, defaults to "ss0").
If the subspec "ss0" is used, then the result is a `DSGEVECM`
whose VECM component is "empty" in that the
observables, cointegrating relationships, lags, and λ weight are not specified.
The reason why this constructor requires the user to specify which
structural shocks of DSGE to use is that this information is
DSGE-specific rather than information about the VECM.

However, we can also construct a `DSGEVECM` without having to
specify the structural shocks when calling the constructor,
although we still need to give an instance of an `AbstractDSGEModel`.

### Example
The code below instantiates an empty `DSGEVECM`
with `AnSchorfheide` as the underlying DSGE and then
calls `update!` on the empty DSGE-VECM
to add information about the desired DSGE-VECM spec.

```jldoctest; output = false
dsgevecm = DSGEVECM(AnSchorfheide())
DSGE.update!(dsgevecm, shocks = [:rm_sh, :z_sh, :g_sh],
    observables = [:obs_gdp, :obs_cpi, :obs_nominalrate],
    λ = 1., lags = 4)

# output

DSGE-VECM Model
observables:      3
data_vintage:     200310
DSGE model:       an_schorfheide
DSGE description: Julia implementation of model defined in 'Bayesian Estimation of DSGE Models' by Sungbae An and Frank Schorfheide: AnSchorfheide, ss0
```

### Fields

#### DSGE object
* `dsge::AbstractDSGEModel{T}`: underlying DSGE model object

#### DSGE-VECM Information
* `observables::OrderedDict{Symbol,Int}`: dictionary mapping observables
    of the VECM to their index in the matrices representing the DSGE-VECM.
* `cointegrating::OrderedDict{Symbol,Int}`: dictionary mapping cointegrating relationships
    of the VECM to their index in the matrices representing the DSGE-VECM.
    When creating the state space representation of a DSGE-VECM, the
    cointegrating relationships will come after the observables. Accordingly,
    the first `cointegrating` index will be after the last observable.
* `cointegrating_add::OrderedDict{Symbol,Int}`: dictionary mapping
    additional cointegrating relationships of the VECM to their index
    in the matrices representing the DSGE-VECM. These relationships are
    intended to be ones that only add to the `DD` matrix in DSGE.jl's
    representation of a state space model. Importantly, these relationships
    do not require changing the `ZZ` matrix. As a result, the indices for
    these relationships start at 1.
* `shocks::OrderedDict{Symbol,Int}`: dictionary mapping structural
    shocks in the DSGE to their index in the matrices representing the DSGE-VECM
* `lags::Int`: number of lags in the VECM
* `λ::T`: weight on the DSGE prior

#### Auxiliary Information
* `spec::String`: concatenates `dsgevar` with the spec of the DSGE, e.g.
    for `AnSchorfheide`, we have `dsgevar_an_schorfheide`.
* `subspec::String`: specifies the model subspecification.
    Cached here for filepath computation.
* `testing::Bool`: indicates whether the model is in testing mode.
    Currently, this setting has no uses in practice
"""
mutable struct DSGEVECM{T} <: AbstractDSGEVECMModel{T}
    dsge::AbstractDSGEModel{T}
    observables::OrderedDict{Symbol,Int}
    cointegrating::OrderedDict{Symbol,Int}
    cointegrating_add::OrderedDict{Symbol,Int}
    shocks::OrderedDict{Symbol,Int}
    lags::Int
    λ::T
    spec::String
    subspec::String
    testing::Bool
end

function Base.show(io::IO, m::DSGEVECM)
    @printf io "DSGE-VECM Model\n"
    @printf io "observables:      %i\n" n_observables(m)
    @printf io "data_vintage:     %s\n" data_vintage(m.dsge)
    @printf io "DSGE model:       %s\n" spec(get_dsge(m))
    @printf io "DSGE description: %s\n" description(get_dsge(m))
end

function DSGEVECM(dsge::AbstractDSGEModel{T}, shocks::Vector{Symbol}, subspec::String = "ss0";
                 custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                 copy_dsge::Bool = false, testing = false) where {T<:Real}

    # Initialize specs
    spec     = "dsgevecm_" * ModelConstructors.spec(dsge)
    subspec  = subspec

    # Initialize empty DSGEVECM
    if copy_dsge
        dsge = deepcopy(dsge)
    end
    m = DSGEVECM{T}(dsge, OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
                    OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), 0,
                    0., spec, subspec, testing)

    # Initialize shocks from DSGE
    update!(m; shocks = shocks, check_valid = false)

    # Initialize subspec
    init_subspec!(m)

    # Do checks of observables, shocks, and lags
    check_valid_dsgevecm(m)

    return m
end

# Empty constructor
function DSGEVECM(dsge::AbstractDSGEModel{T}, subspec::String = "ss0";
                 custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                 copy_dsge::Bool = false, testing = false) where {T<:Real}

    # Initialize specs
    spec     = "dsgevar_" * ModelConstructors.spec(dsge)
    subspec  = subspec

    # Initialize empty DSGEVECM
    if copy_dsge
        dsge = deepcopy(dsge)
    end
    m = DSGEVECM{T}(dsge, OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
                    OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), 0,
                    0., spec, subspec, testing)

    # Initialize shocks from DSGE
    update!(m; check_valid = false)

    # Initialize subspec
    init_subspec!(m)

    # Do checks of observables, shocks, and lags
    check_valid_dsgevecm(m)

    return m
end


using Timestamps, FinancialAssets, Reactive

module Grist

using Timestamps, FinancialAssets, Reactive

import Base: convert, show, getindex, start, next, done, isempty

export Blotter,
       transact, # helper method to update asset A
       res, vals, # from dev/tenblotters.jl
       txn, cash, A, mark_price, Pval, Aval # from dev/updates.jl

include("blotter.jl")
include("operators.jl")
include("../dev/tenblotters.jl")
include("../dev/updates.jl")
#include("../dev/run.jl")

end


function getCompleteVariableList(
    mb      :: ModelBlock,
    vartype :: Symbol,
)

    d = nothing

    if vartype == :dynamic
        d = OrderedDict(

            "TEMP"               => ( mb.fi.sv[:TEMP], :T, :mask ),
            "SALT"               => ( mb.fi.sv[:SALT], :T, :mask ),
            "UVEL"               => ( mb.fi.sv[:UVEL], :U, nothing ),
            "VVEL"               => ( mb.fi.sv[:VVEL], :V, nothing ),
            "WVEL"               => ( mb.fi.sv[:WVEL], :W, nothing ),
            "TAUX"               => ( mb.fi.TAUX,        :sT, :mask ),
            "TAUY"               => ( mb.fi.TAUY,        :sT, :mask ),
            "TAUX_east"          => ( mb.fi.TAUX_east,   :sT, :mask ),
            "TAUY_north"         => ( mb.fi.TAUY_north,  :sT, :mask ),
            "SWFLX"              => ( mb.fi.SWFLX,       :sT, :mask ),
            "NSWFLX"             => ( mb.fi.NSWFLX,      :sT, :mask ),
            "VSFLX"              => ( mb.fi.VSFLX,       :sT, :mask ),
            "ADVT"               => ( mb.fi.sv[:ADVT],     :T, :mask ),
            "ADVS"               => ( mb.fi.sv[:ADVS],     :T, :mask ),
            "HMXL"               => ( mb.fi.HMXL,          :sT, :mask ),

            "WKRSTT"       => ( mb.fi.sv[:WKRSTT], :T, :mask),
            "WKRSTS"       => ( mb.fi.sv[:WKRSTS], :T, :mask),

            "VDIFFT"       => ( mb.fi.sv[:VDIFFT], :T, :mask),
            "VDIFFS"       => ( mb.fi.sv[:VDIFFS], :T, :mask),

            "Q_FRZMLTPOT"        => ( mb.fi.Q_FRZMLTPOT,  :sT, :mask ),
            "Q_FRZMLTPOT_NEG"    => ( mb.fi.Q_FRZMLTPOT_NEG, :sT, :mask ),
            "Q_FRZHEAT"          => ( mb.fi.Q_FRZHEAT,    :sT, :mask ),
            "Q_LOST"             => ( mb.fi.Q_LOST,       :sT, :mask ),
            "CHKTEMP"            => ( mb.tmpfi.sv[:CHKTEMP],  :sT, :mask ),
            "CHKSALT"            => ( mb.tmpfi.sv[:CHKSALT],  :sT, :mask ),

            "INTMTEMP"           => ( mb.tmpfi.sv[:INTMTEMP], :T, :mask),
            "INTMSALT"           => ( mb.tmpfi.sv[:INTMSALT], :T, :mask),

            "WKRST_TARGET_TEMP"         => nothing,
            "WKRST_TARGET_SALT"         => nothing,
            "QFLXT"          => nothing,
            "QFLXS"          => nothing,

            "Ks_H_U"    => ( mb.tmpfi.check_usage[:Ks_H_U], :U, nothing),
            "Ks_H_V"    => ( mb.tmpfi.check_usage[:Ks_H_V], :V, nothing),

        )

        if mb.tmpfi.datastream != nothing

            if mb.ev.config["weak_restoring"] == "on"
                d["WKRST_TARGET_TEMP"] = ( mb.tmpfi.datastream["TEMP"], :T , :mask)
                d["WKRST_TARGET_SALT"] = ( mb.tmpfi.datastream["SALT"], :T , :mask)
            end

        end

        d["QFLXT"] = ( mb.fi.sv[:QFLXT], :T, :mask )
        d["QFLXS"] = ( mb.fi.sv[:QFLXS], :T, :mask )

    elseif vartype == :static

        d = Dict(
            # COORDINATE
            "deepmask_T"=> ( mb.ev.topo.deepmask_T, :T, nothing),
            "topoz_sT"  => ( mb.ev.topo.topoz_sT, :sT, nothing),
            "Nz_bot_sT" => ( mb.ev.topo.Nz_bot_sT, :sT, nothing),
            "area_sT"   => ( mb.ev.gd_slab.Δx_T .* mb.ev.gd_slab.Δy_T, :sT, nothing),
            "mask_sT"   => ( mb.ev.topo.sfcmask_sT, :sT, nothing),
            "z_cW"      => ( reshape(mb.ev.config["z_w"], :, 1, 1), :cW, nothing),
            "dz_cT"     => ( mb.ev.gd.Δz_T[:, 1:1, 1:1], :cT, nothing),
            "lon_sT"    => ( rad2deg.(mb.ev.gd_slab.λ_T), :sT, nothing),
            "lat_sT"    => ( rad2deg.(mb.ev.gd_slab.ϕ_T), :sT, nothing),
        )

    else
        throw(ErrorException("Unknown vartype: " * string(vartype)))
    end

    return d
end

function getVarDesc(varname)
    return haskey(HOOM.var_desc, varname) ? HOOM.var_desc[varname] : Dict()
end

function getDynamicVariableList(
    mb :: ModelBlock;
    varnames :: AbstractArray{String} = Array{String}(undef,0),
    varsets  :: AbstractArray{String} = Array{String}(undef,0),
)

    all_varlist = getCompleteVariableList(mb, :dynamic)

    output_varnames = []

    for varname in varnames
        if haskey(all_varlist, varname)
            push!(output_varnames, varname)
        else
            throw(ErrorException(format("Varname {:s} does not exist.", varname)))
        end
    end

    for varset in varsets

        if varset == "ALL"

            append!(output_varnames, keys(all_varlist))

        elseif varset == "ESSENTIAL"

            append!(output_varnames, [

                "TEMP",
                "SALT",
                "UVEL",
                "VVEL",
                "WVEL",
                "TAUX",
                "TAUY",
                "TAUX_east",
                "TAUY_north",
                "SWFLX",
                "NSWFLX",
                "VSFLX",
                "ADVT",
                "ADVS",
                "HMXL",

                "WKRSTT",
                "WKRSTS",

                "VDIFFT",
                "VDIFFS",

                "Q_FRZMLTPOT",
                "Q_LOST",
                "CHKTEMP",
                "CHKSALT",

            ])

        else

            throw(ErrorException("Unknown varset: " * string(varset)))

        end
    end

    function makeSubset(dict, keys)
        subset_dict = Dict()
        for k in keys
            if dict[k] != nothing
                subset_dict[k] = dict[k]
            else
                println("$(k) is nothing")
            end
        end
        return subset_dict
    end

    #println("NOW AGAIN: ", output_varnames)

    output_varlist = makeSubset(all_varlist, output_varnames)

    return output_varlist
end