include("../src/TcrCore.jl")
using TcrCore, HypothesisTests

function differentMechanisms()
    srand(1234)
    for acc in 0:5:100
        applicationOnly = mean([simSimple(1000, setupAgentsWithFixedAccuracy(100, acc), [benchmarkRegistryMean]) for i in 1:20])[1]
        challenge = mean([simChallenge(1000, setupAgentsWithFixedAccuracy(100, acc), [benchmarkRegistryMean]) for i in 1:20])[1]
        token = mean([simToken(1000, setupAgentsWithFixedAccuracy(100, acc), noneRedistribution, [benchmarkRegistryMean]) for i in 1:20])[1]
        println("$acc, $applicationOnly, $challenge, $token")
    end
end

function tokensAndDiversifications()
    srand(1234)

    for acc in 0:5:100
        noTokenFixed = mean([simChallenge(1000, setupAgentsWithFixedAccuracy(100, acc), [benchmarkRegistryMean]) for i in 1:30])[1]
        noTokenDiv = mean([simChallenge(1000, setupRandomAgents(100, acc, 20), [benchmarkRegistryMean]) for i in 1:30])[1]
        tokenFixed = mean([simToken(1000, setupAgentsWithFixedAccuracy(100, acc), noneRedistribution, [benchmarkRegistryMean]) for i in 1:30])[1]
        tokenDiv = mean([simToken(1000, setupRandomAgents(100, acc, 20), noneRedistribution, [benchmarkRegistryMean]) for i in 1:30])[1]
        println("$acc, $noTokenFixed, $noTokenDiv, $tokenFixed, $tokenDiv")
    end
end

function tokensAndDiversificationsEffects()
    srand(1234)
    divEffects = []
    tokenEffect = []
    tokenDivEffects = []
    for acc in 0:5:100
        noTokenFixed = mean([simChallenge(1000, setupAgentsWithFixedAccuracy(100, acc), [benchmarkRegistryMean]) for i in 1:30])[1]
        noTokenDiv = mean([simChallenge(1000, setupRandomAgents(100, acc, 20), [benchmarkRegistryMean]) for i in 1:30])[1]
        tokenFixed = mean([simToken(1000, setupAgentsWithFixedAccuracy(100, acc), noneRedistribution, [benchmarkRegistryMean]) for i in 1:30])[1]
        tokenDiv = mean([simToken(1000, setupRandomAgents(100, acc, 20), noneRedistribution, [benchmarkRegistryMean]) for i in 1:30])[1]

        divEffect = mean([tokenDiv - tokenFixed, noTokenDiv - noTokenFixed])
        tokenEffect = mean([tokenDiv - noTokenDiv, tokenFixed - noTokenFixed])
        tokenDivEffect = tokenDiv - noTokenFixed
        println("$acc, $divEffect, $tokenEffect, $tokenDivEffect")
    end
end

function compareEfficiencyByAccuracy()
    for acc in 0:5:100
        noToken = [simChallenge(1000, setupRandomAgents(100, acc, 20), [benchmarkRegistryMean])[1] for i in 1:10]
        token = [simToken(1000, setupRandomAgents(100, acc, 20), noneRedistribution, [benchmarkRegistryMean])[1] for i in 1:10]
        boost = mean(token) - mean(noToken)
        test = pvalue(UnequalVarianceTTest(token, noToken))
        significant = test < 0.05
        println("$acc, $boost, $test, $significant")
    end
end

function efficiencyArea()
    for noAgents in 10:10:200
        for noSteps in 100:100:2000
            noToken = [simChallenge(noSteps, setupRandomAgents(noAgents, 50, 20), [benchmarkRegistryMean])[1] for i in 1:10]
            token = [simToken(noSteps, setupRandomAgents(noAgents, 50, 20), noneRedistribution, [benchmarkRegistryMean])[1] for i in 1:10]
            boost = mean(token) - mean(noToken)
            println("$noSteps, $noAgents, $boost")
        end
    end
end

function expertsInnerCirleSize()
    for expertCount in 1:1:50
        controlLow= mean([simChallenge(1000, setupAgentsWithFixedAccuracy(100, 20), [benchmarkRegistryMean])[1] for i in 1:20])
        expertTokens = mean([simToken(1000, setupMixedAgents(100-expertCount, 20, expertCount, 100), noneRedistribution, [benchmarkRegistryMean])[1] for i in 1:20])
        expertNoTokens = mean([simChallenge(1000, setupMixedAgents(100-expertCount, 20, expertCount, 100), [benchmarkRegistryMean])[1] for i in 1:20])
        controlHigh = mean([simChallenge(1000, setupAgentsWithFixedAccuracy(100, 100), [benchmarkRegistryMean])[1] for i in 1:20])
        println("$expertCount, $controlLow, $expertTokens, $expertNoTokens, $controlHigh")
    end
end

function compareRedistribution()
    for acc in 0:5:100
        challengerOnly = [simToken(1000, setupRandomAgents(100, acc, 20), onlyChallengerRewardRedistribution, [benchmarkRegistryMean])[1] for i in 1:20]
        punishment = [simToken(1000, setupRandomAgents(100, acc, 20), punishmentRedistribution, [benchmarkRegistryMean])[1] for i in 1:20]
        reward = [simToken(1000, setupRandomAgents(100, acc, 20), rewardRedistribution, [benchmarkRegistryMean])[1] for i in 1:20]
        punishmentAndReward = [simToken(1000, setupRandomAgents(100, acc, 20), punishmentAndRewardRedistribution, [benchmarkRegistryMean])[1] for i in 1:20]
        pBoost = mean(punishment) - mean(challengerOnly)
        rBoost = mean(reward) - mean(challengerOnly)
        prBoost = mean(punishmentAndReward) - mean(challengerOnly)
        # test = pvalue(UnequalVarianceTTest(full, challengerOnly))
        # significant = test < 0.05
        println("$acc, $pBoost, $rBoost, $prBoost")
    end
end

function runScenarios()
    compareRedistribution()
end

runScenarios()


using Revise
using Random
using LinearAlgebra
using POMDPs
using POMDPModelTools
using POMDPPolicies
using POMDPSimulators
using IterTools
using AutomotiveDrivingModels
using AutomotivePOMDPs: polygon, overlap
using Parameters
using GridInterpolations
using StaticArrays
using AutoUrban
using AutoViz
using Printf
using Reel
using DiscreteValueIteration
using SparseArrays
using ProgressMeter
using JLD2
using FileIO

includet("environment.jl")
includet("mdp_definition.jl")
includet("helpers.jl")

mdp = MergingMDP()
@load "policies.jld2" policies
policy = policies[1] # [:wide, :cv, :ca, :cd]

function safe_actions(svec::Vector{Float64}, threshold::Float64=0.99)
    @assert length(svec) == length(first(states(mdp)))

    ## interpolate state
    sitp, witp = interpolants(mdp.state_grid, svec)
    normalize!(witp, 1)
    vals = zeros(n_actions(mdp))
    for (i, si) in enumerate(sitp)
        w = witp[i]
        s = ind2x(mdp.state_grid, si)
        if isapprox(svec[4], AbsentPos)
            s[4] = AbsentPos
        end
        vals += w .* actionvalues(policy, s)
    end
    return findall(vals .> threshold)
end

println("Loaded model checker")


# This function evaluates and returns the coefficients for constraints and cost which are used later in the MPC formulation
# Inputs are:
# oldTraj   -> contains information about previous trajectories and Inputs
# mpcCoeff  -> contains information about previous coefficient results
# posInfo   -> contains information about the car's current position along the track
# mpcParams -> contains information about the MPC formulation (e.g. Q, R)
# stateIn   -> current state of the car
# inputIn   -> last input to the system

# structure of oldTrajectory: 1st dimension = state number, 2nd dimension = step number (time equiv.), 3rd dimennsion = lap number

# z[1] = xDot
# z[2] = yDot
# z[3] = psiDot
# z[4] = ePsi
# z[5] = eY
# z[6] = s

function coeffConstraintCost(oldTraj::OldTrajectory, mpcCoeff::MpcCoeff, posInfo::PosInfo, mpcParams::MpcParams,lapStatus::LapStatus,
                             selectedStates::SelectedStates,oldSS::SafeSetData,obs::Array{Float64},obstacle::Obstacle)
    # this computes the coefficients for the cost and constraints

    # Outputs:
    # coeffConst
    # coeffCost

    # Read Inputs
    s               = posInfo.s
    s_target        = posInfo.s_target

    dt = 0.1
    N  = mpcParams.N

    Np              = selectedStates.Np
    Nl              = selectedStates.Nl

    # Parameters
    Order           = mpcCoeff.order                # interpolation order for cost and constraints
    pLength         = mpcCoeff.pLength              # interpolation length for polynomials
    delay_df        = mpcParams.delay_df
    #delay           = 2

    #coeffCost       = zeros(Order+1,2)              # polynomial coefficients for cost
    #coeffConst      = zeros(Order+1,2,5)            # nz-1 beacuse no coeff for s

    n_laps_sysID    = 2                             # number of previous laps that are used for sysID

    selected_laps = zeros(Int64,Nl)
    # selected_laps[1] = lapStatus.currentLap-1                                   # use previous lap
    # selected_laps[2] = lapStatus.currentLap-2                                   # and the one before
    # if lapStatus.currentLap >= 5
    #     selected_laps[2] = indmin(oldTraj.oldCost[2:lapStatus.currentLap-2])+1      # and the best from all previous laps
    # end

    for i = 1:Nl
        selected_laps[i] = lapStatus.currentLap-i    # use previous lap
    end

    if lapStatus.currentLap >= 5
        selected_laps[Nl] = indmin(oldSS.oldCost[1:lapStatus.currentLap-2])      # and the best from all previous laps
    end

    # Select the old data
    oldxDot         = oldTraj.oldTraj[:,1,selected_laps]::Array{Float64,3}
    oldyDot         = oldTraj.oldTraj[:,2,selected_laps]::Array{Float64,3}
    oldpsiDot       = oldTraj.oldTraj[:,3,selected_laps]::Array{Float64,3}
    oldePsi         = oldTraj.oldTraj[:,4,selected_laps]::Array{Float64,3}
    oldeY           = oldTraj.oldTraj[:,5,selected_laps]::Array{Float64,3}
    oldS            = oldTraj.oldTraj[:,6,selected_laps]::Array{Float64,3}
    oldacc          = oldTraj.oldTraj[:,7,selected_laps]::Array{Float64,3}
    olda            = oldTraj.oldInput[:,1,selected_laps]::Array{Float64,3}
    olddF           = oldTraj.oldInput[:,2,selected_laps]::Array{Float64,3}
    #olddF           = smooth(olddF,5)

    oldS_safeSet    = oldSS.oldSS[:,6,selected_laps]::Array{Float64,3}
    #println("safe set= ",oldS_safeSet[1:20,:,:])

    N_points        = size(oldTraj.oldTraj,1)     # second dimension = length (=buffersize)
    N_points2       = size(oldSS.oldSS,1)

    s_total::Float64        # initialize
    DistS::Array{Float64}   # initialize
    idx_s::Array{Int64}     # initialize
    idx_s_target        = 0
    dist_to_s_target    = 0
    qLength             = 0
    vec_range::Tuple{UnitRange{Int64},UnitRange{Int64}}
    bS_Vector::Array{Float64}
    s_forinterpy::Array{Float64}

    # Compute the total s (current position along track)
    s_total = s % s_target

    # Compute the index
    DistS = ( s_total - oldS ).^2
    DistS2 = ( s_total - oldS_safeSet ).^2

    idx_s = findmin(DistS,1)[2]              # contains both indices for the closest distances for both oldS !!
    idx_s2= findmin(DistS2,1)[2]

    idx_s2 = idx_s2 + selectedStates.shift

    # Propagate the obstacle for the prediction horizon

    obs_prop_s  = obs[1,1,:] + dt*N*obs[1,3,:]
    obs_prop_ey = obs[1,2,:]

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

    for j = 0:(Nl-1)

        selectedStates.selStates[i=(j*Np)+1:(j+1)*Np,m=1:6] = oldSS.oldSS[i=idx_s2[j+1]-(j*N_points2):idx_s2[j+1]+Np-(j*N_points2)-1,i=1:6,selected_laps[j+1]]  # select the states from lap j...

        selectedStates.statesCost[i=(j*Np)+1:(j+1)*Np] = oldSS.cost2target[i=idx_s2[j+1]-(j*N_points2):idx_s2[j+1]-(j*N_points2)+Np-1,selected_laps[j+1]]  # and their cost

        if obstacle.obstacle_active == true   # if the obstacles are on the track, check if any of the selected states interferes with the propagated obstacle
            for n=1:obstacle.n_obs
                ellipse_check = (((selectedStates.selStates[i=(j*Np)+1:(j+1)*Np,6]-obs_prop_s[n])/obstacle.r_s).^2) + (((selectedStates.selStates[i=(j*Np)+1:(j+1)*Np,5]-obs_prop_ey[n])/obstacle.r_ey).^2)

                if any(x->x<=1, ellipse_check) == true  # if any of the selected states is in the ellipse
                    #println("flag**************************************************************************************************************")
                    index = find(ellipse_check.<=1)     # find all the states in the ellipse

                    mpcParams.Q_obs[i=(j*Np)+(index[1]-obstacle.inv_step)+1:(j+1)*Np] =  10   # and set the values of the weight to 10, so that they are excluded from optimization
                    println("Q_obs= ",mpcParams.Q_obs)
                end
            end
        end

    end

        #println("Q_obs= ",mpcParams.Q_obs)

    ##########################################################################
    if selectedStates.version == true
        vec_range = (idx_s[1]:idx_s[1]+pLength,idx_s[2]:idx_s[2]+pLength)

        # Create the vectors used for the interpolation
        # bS_vector contains the s-values for later interpolation
        bS_Vector       = zeros(pLength+1,2)
        for i=1:pLength+1
            bS_Vector[i,1] = oldS[vec_range[1][i]]
            bS_Vector[i,2] = oldS[vec_range[2][i]]
        end
        # if norm(bS_Vector[1,1]-s_total) > 0.3 || norm(bS_Vector[1,2]-s_total) > 0.3
        #     warn("Couldn't find a close point to current s.")
        # end

        # The states are parametrized with resprect to the curvilinear abscissa,
        # so we select the point used for the interpolation. Need to subtract an
        # offset to be coherent with the MPC formulation
        s_forinterpy   = bS_Vector
        if s_total < 0
            s_forinterpy += s_target
        end

        # Create the Matrices for the interpolation
        MatrixInterp = zeros(pLength+1,Order+1,2)

        for k = 0:Order
            MatrixInterp[:,Order+1-k,:] = s_forinterpy[:,:].^k
        end

        # Compute the coefficients
        #coeffConst = zeros(Order+1,2,5)
        for i=1:2
            mpcCoeff.coeffConst[:,i,1]    = MatrixInterp[:,:,i]\oldxDot[vec_range[i]]
            mpcCoeff.coeffConst[:,i,2]    = MatrixInterp[:,:,i]\oldyDot[vec_range[i]]
            mpcCoeff.coeffConst[:,i,3]    = MatrixInterp[:,:,i]\oldpsiDot[vec_range[i]]
            mpcCoeff.coeffConst[:,i,4]    = MatrixInterp[:,:,i]\oldePsi[vec_range[i]]
            mpcCoeff.coeffConst[:,i,5]    = MatrixInterp[:,:,i]\oldeY[vec_range[i]]
        end

        # Finished with calculating the constraint coefficients

        # Now compute the final cost coefficients

        # The Q-function contains for every point in the sampled safe set the minimum cost-to-go-value
        # These values are calculated for both old trajectories
        # The vector bQfunction_Vector contains the cost at each point in the interpolated area to reach the finish line
        # From this vector, polynomial coefficients coeffCost are calculated to approximate this cost
        for i=1:2
                dist_to_s_target  = oldTraj.oldCost[selected_laps[i]] + oldTraj.idx_start[selected_laps[i]] - (idx_s[i]-N_points*(i-1))  # number of iterations from idx_s to s_target
                bQfunction_Vector = collect(linspace(dist_to_s_target,dist_to_s_target-pLength,pLength+1))    # build a vector that starts at the distance and
                                                                                # decreases in equal steps
                mpcCoeff.coeffCost[:,i]    = MatrixInterp[:,:,i]\bQfunction_Vector           # interpolate this vector with the given s
        end
    end

    # --------------- SYSTEM IDENTIFICATION --------------- #
    # ----------------------------------------------------- #

    cC = oldTraj.count[lapStatus.currentLap]-1      # current count
    cL = lapStatus.currentLap                       # current lap
    #vec_range_ID2   = cC-n_prev:cC-1                # index range for current lap

    # TODO:*******************
    # CHECK IF DIFFERENCES ARE ALIGNED PROPERLY
    # ADD 2-step-delay to sysID
    # ************************
    # collect indices for system ID
    n_sys_ID_prev = 60              # steps of sysID before current point in previous laps
    n_sys_ID_post = 60              # steps of sysID after current point in previous laps
    n_sys_ID_prev_c = 30            # steps of sysID before current point in current lap

    freq_ratio = 5                  # ratio between the frequency at which the informations about the position are updated in the topic and the frequency at which the controller calculates and updates the control input
    # vec_range_ID    = ()
    # for i=1:n_laps_sysID
    #     vec_range_ID    = tuple(vec_range_ID...,idx_s[i]-n_prev:idx_s[i]+n_ahead)     # related index range
    # end

    sysID_idx_diff = idx_s[1]-n_sys_ID_prev*freq_ratio-1 + (1:freq_ratio:(n_sys_ID_prev+n_sys_ID_post+1)*freq_ratio)         # these are the indices that are used for differences
    sysID_idx = sysID_idx_diff[1:end-1]

    sysID_idx_diff2 = idx_s[2]-n_sys_ID_prev*freq_ratio-1 + (1:freq_ratio:(n_sys_ID_prev+n_sys_ID_post+1)*freq_ratio)         # these are the indices that are used for differences
    sysID_idx2 = sysID_idx_diff2[1:end-1]

    sysID_idx_diff_c = cC - (n_sys_ID_prev_c+1)*freq_ratio-1 + (1:freq_ratio:(n_sys_ID_prev_c+1)*freq_ratio)
    sysID_idx_c = sysID_idx_diff_c[1:end-1]

    sz1 = size(sysID_idx,1)
    sz2 = size(sysID_idx_c,1)
    sz2 = 0
    sz = sz1 + sz2

    # psiDot
    y_psi = zeros((2*sz1+sz2)*freq_ratio)
    A_psi = zeros((2*sz1+sz2)*freq_ratio,3)
    for i=0:freq_ratio-1
        y_psi[(1:sz1)+i*sz1]            = diff(oldpsiDot[sysID_idx_diff+i])
        A_psi[(1:sz1)+i*sz1,:]          = [oldpsiDot[sysID_idx+i]./oldxDot[sysID_idx+i] oldyDot[sysID_idx+i]./oldxDot[sysID_idx+i] olddF[sysID_idx+i]]
        #y_psi[(1:sz2)+i*sz2+freq_ratio*sz1]      = diff(oldTraj.oldTraj[sysID_idx_diff_c+i,3,cL])
        #A_psi[(1:sz2)+i*sz2+freq_ratio*sz1,:]    = [oldTraj.oldTraj[sysID_idx_c+i,3,cL]./oldTraj.oldTraj[sysID_idx_c+i,1,cL] oldTraj.oldTraj[sysID_idx_c+i,2,cL]./oldTraj.oldTraj[sysID_idx_c+i,1,cL] oldTraj.oldInput[sysID_idx_c+i-delay_df*freq_ratio,2,cL]]
        y_psi[(1:sz1)+i*sz1+freq_ratio*sz1+freq_ratio*sz2]   = diff(oldpsiDot[sysID_idx_diff2+i])
        A_psi[(1:sz1)+i*sz1+freq_ratio*sz1+freq_ratio*sz2,:] = [oldpsiDot[sysID_idx2+i]./oldxDot[sysID_idx2+i] oldyDot[sysID_idx2+i]./oldxDot[sysID_idx2+i] olddF[sysID_idx2+i]]
    end

    # xDot
    y_xDot = zeros((2*sz1+sz2)*freq_ratio)
    A_xDot = zeros((2*sz1+sz2)*freq_ratio,3)
    for i=0:freq_ratio-1
        y_xDot[(1:sz1)+i*sz1]            = diff(oldxDot[sysID_idx_diff+i])
        A_xDot[(1:sz1)+i*sz1,:]          = [oldyDot[sysID_idx+i].*oldpsiDot[sysID_idx+i] oldxDot[sysID_idx+i] oldacc[sysID_idx+i]]
        #y_xDot[(1:sz2)+i*sz2+freq_ratio*sz1]      = diff(oldTraj.oldTraj[sysID_idx_diff_c+i,1,cL])
        #A_xDot[(1:sz2)+i*sz2+freq_ratio*sz1,:]    = [oldTraj.oldTraj[sysID_idx_c+i,2,cL].*oldTraj.oldTraj[sysID_idx_c+i,3,cL] oldTraj.oldTraj[sysID_idx_c+i,1,cL]  oldTraj.oldTraj[sysID_idx_c+i,7,cL]]
        y_xDot[(1:sz1)+i*sz1+freq_ratio*(sz1+sz2)]   = diff(oldxDot[sysID_idx_diff2+i])
        A_xDot[(1:sz1)+i*sz1+freq_ratio*(sz1+sz2),:] = [oldyDot[sysID_idx2+i].*oldpsiDot[sysID_idx2+i] oldxDot[sysID_idx2+i] oldacc[sysID_idx2+i]]
    end

    # yDot
    y_yDot = zeros((2*sz1+sz2)*freq_ratio)
    A_yDot = zeros((2*sz1+sz2)*freq_ratio,4)
    for i=0:freq_ratio-1
        y_yDot[(1:sz1)+i*sz1]            = diff(oldyDot[sysID_idx_diff+i])
        A_yDot[(1:sz1)+i*sz1,:]          = [oldyDot[sysID_idx+i]./oldxDot[sysID_idx+i] oldpsiDot[sysID_idx+i].*oldxDot[sysID_idx+i] oldpsiDot[sysID_idx+i]./oldxDot[sysID_idx+i] olddF[sysID_idx+i-delay_df*freq_ratio]]
        #y_yDot[(1:sz2)+i*sz2+freq_ratio*sz1]      = diff(oldTraj.oldTraj[sysID_idx_diff_c+i,2,cL])
        #A_yDot[(1:sz2)+i*sz2+freq_ratio*sz1,:]    = [oldTraj.oldTraj[sysID_idx_c+i,2,cL]./oldTraj.oldTraj[sysID_idx_c+i,1,cL] oldTraj.oldTraj[sysID_idx_c+i,3,cL].*oldTraj.oldTraj[sysID_idx_c+i,1,cL] oldTraj.oldTraj[sysID_idx_c+i,3,cL]./oldTraj.oldTraj[sysID_idx_c+i,1,cL] oldTraj.oldInput[sysID_idx_c+i-delay_df*freq_ratio,2]]
        y_yDot[(1:sz1)+i*sz1+freq_ratio*(sz1+sz2)]   = diff(oldyDot[sysID_idx_diff2+i])
        A_yDot[(1:sz1)+i*sz1+freq_ratio*(sz1+sz2),:] = [oldyDot[sysID_idx2+i]./oldxDot[sysID_idx2+i] oldpsiDot[sysID_idx2+i].*oldxDot[sysID_idx2+i] oldpsiDot[sysID_idx2+i]./oldxDot[sysID_idx2+i] olddF[sysID_idx2+i-delay_df*5]]
    end

    # if any(isnan,y_yDot)            # check if any value in the y_yDot value is NaN
    #     println(y_yDot)
    #     warn("NaN value detected in y_yDot! Press to continue...")
    # end
    # if any(isnan,coeffCost)
    #     println(coeffCost)
    #     warn("NaN value detected in coeffCost! Press to continue...")
    # end
    # if any(isnan,coeffConst)
    #     println(coeffCost)
    #     warn("NaN value detected in coeffConst! Press to continue...")
    # end

    # mpcCoeff.c_Psi = zeros(3)
    # mpcCoeff.c_Vx = zeros(4)
    # mpcCoeff.c_Vy = zeros(4)

    mpcCoeff.c_Psi = (A_psi'*A_psi)\A_psi'*y_psi
    mpcCoeff.c_Vx  = (A_xDot'*A_xDot)\A_xDot'*y_xDot         # the identity matrix is used to scale the coefficients
    mpcCoeff.c_Vy  = (A_yDot'*A_yDot)\A_yDot'*y_yDot

    # if det(A_psi'*A_psi) != 0
    # else
    #     println("det y_psi = 0")
    # end
    # if det(A_xDot'*A_xDot) != 0
    # else
    #     println("det vx = 0")
    # end
    # if det(A_yDot'*A_yDot) != 0
    # else
    #     println("det vy = 0")
    # end
    #println("Done. c_Psi = $(mpcCoeff.c_Psi)")
    #mpcCoeff.c_Psi = [-1.0,-4.5,5.8]
    #mpcCoeff.c_Vy  = [-1.2,-0.04,-0.01,0.6]
    #mpcCoeff.c_Psi = [-0.26682109207165566,-0.013445078992161885,1.2389672517023724]
    #mpcCoeff.c_Vy = [-0.006633028965076818,-0.02997779668710061,0.005781203137095575,0.10642934131787765]
    #mpcCoeff.coeffCost  = coeffCost
    #mpcCoeff.coeffConst = coeffConst
    nothing
end

# Notes about oldTrajectory:
# oldTrajectory[1:prebuf] = states before finish line (s < 0)
# oldTrajectory[prebuf+1:prebuf+cost] = states between start and finish line (0 <= s < s_target)
# oldTrajectory[prebuf+cost+1:prebuf+cost+postbuf] = states after finish line (s_target <= s)
# once one lap is over, the states are saved (interval prebuf:s_target)
# during the next lap, the states behind the finish line are appended (interval s_target:postbuf)


## extensions/HBond.jl --- models for hydrogen bonding

export HBondDict, hbondenergy, hbonded
export donors, donorindices, acceptors, acceptorindices

export nturn, alphaturn, shortturn, piturn
export parallelbridge, antiparallelbridge

export recorddonor!, recordacceptor!

"""
    HBondDict{N,T<:AbstractFloat}

Hydrogen bonding dictionary for a specific residue. Type parameter `N` describes at
most how many hydrogen bonds can be donated/accepted by the amide part of the residue.
Index of the residue (in a given topology) is `this.handle`.

**Note**: Values of `N` that make sense are ``1`` and ``2``; any greater would not
quite make sense. DSSP uses two to store the H-bonds with lowest energy (most stable).
"""
struct HBondDict{N,T<:AbstractFloat} <: AbstractVector{Tuple{UInt64,T}}
    interacting_residues::MVector{N,UInt64}
    energies::MVector{N,T}
    base_size::Int

    function HBondDict(::Val{M}, ::Type{Tp}) where {M,Tp}
        inter = @MVector zeros(UInt64, M)
        energies = @MVector zeros(Tp, M)
        return new{M,Tp}(inter, energies, M >> 1)
    end
end
function HBondDict(base::Integer)
    !(0 < base) && error("try using at least 1 hydrogen bond")
    return HBondDict(Val(2*base), float(Int))
end
Base.length(::HBondDict{N}) where N = N
Base.size(::HBondDict{N}) where N = (N,)

function Base.iterate(dict::HBondDict{N}, state=1) where N
    state > N && return nothing
    (; interacting_residues, energies) = dict
    ## donors are stored first relative to acceptors
    tpl = @inbounds interacting_residues[state], energies[state]
    return (tpl, state+1)
end

Base.@propagate_inbounds function Base.getindex(dict::HBondDict, i::Integer)
    return (dict.interacting_residues[i], dict.energies[i])
end

Base.@propagate_inbounds function Base.setindex!(dict::HBondDict, val, i::Integer)
    res, energy = val
    dict.interacting_residues[i] = res
    dict.energies[i] = energy
    nothing
end

donors(dict::HBondDict{N}) where N = Iterators.take(dict, N>>1)
acceptors(dict::HBondDict{N}) where N = Iterators.drop(dict, N>>1)
donorindices(dict) = (i for (i, _) in donors(dict))
acceptorindices(dict) = (i for (i, _) in acceptors(dict))

## TODO: CREATE A VECTOR OF INITIALIZED HBONDDICT FROM FRAME INPUT

"""
    hbondenergy(frame, donor, acceptor; method=:dist)

Calculate the (amide) hydrogen bond energy between a donor residue and an
acceptor residue in a frame. Both nonframe arguments are _indices_ in the
`StructureFrame`. Defaults to a mimicked method of that _currently_ used
in DSSP (using distances, although there are plans to use angles).
An _approximate_ potential using dot products `method=:dot` is provided
in this module.

# Keyword arguments
+ `method`: Defaults to `:dist`. Method used for energy calculation.
"""
function hbondenergy(frame::StructureFrame, donor_i, acceptor_i; method=:dist)
    (; at_pos, res_list) = frame

    calcfunc = method == :dot ? (_hbenergy_dot) : (_hbenergy_dist)
    donor = res_list[donor_i]; acceptor = res_list[acceptor_i]
    ## get N and H in donor; C and O in acceptor
    Nid = donor[:N]; Hid = donor[:H]
    Cid = acceptor[:C]; Oid = acceptor[:O]

    Npos = at_pos[Nid]; Hpos = at_pos[Hid]
    Cpos = at_pos[Cid]; Opos = at_pos[Oid]

    return calcfunct(Npos, Hpos, Cpos, Opos)
end

function _hbenergy_dist(Npos, Hpos, Cpos, Opos)
    ## DSSP default implementation
    KK = DDSPConstants[:coupleConstant]     # coupling constant from dict
    HO = distance(Hpos, Opos); HC = distance(Hpos, Cpos)
    NO = distance(Npos, Opos); NC = distance(Npos, Cpos)

    ## coulomb interaction from two dipoles
    return KK * (1/HO - 1/HC + 1/NO - 1/NC)
end

function _hbenergy_dot(Npos, Hpos, Cpos, Opos)
    ## method using dot
    KK = DSSPConstants[:coupleConstant]     # coupling constant from dict
    ## direction of dipole moment
    NHp = Hpos - Npos; COp = Cpos - Opos
    ## centers of dipoles
    NHc = (Npos + Hpos) / 2; COc = (Cpos + Opos) / 2
    R = COc - NHc; r = norm(disp)

    ## "keesom" interaction of two dipoles
    pre = KK * (NHp⋅COp) / r^3 - 3 * (NHp⋅R) * (COp⋅R) / r^5
    return KK * pre
end

"""
    hbonded(hbdicts, res_i, res_j) -> Bool

Check if the residue at _index_ `res_i` acts as a hydrogen bond acceptor to
the residue at `res_j` based on the list of dictionaries `hbdicts`.
Equivalently, if residue at `res_j` acts as a hydrogen bond donor to
the residue at `res_i`.
"""
function hbonded(hbdicts::Vector{<:HBondDict}, i, j)
    ## annotate with @boundscheck??
    checkbounds(hbdicts, i); checkbounds(hbdicts, j)
    @inbounds i in acceptorindices(hbdicts[j])
end

## TODO: check for chain breaks between indices i and j
### best solution i can think of is a new data struct that is produced linearly
### relative to number of residues; and can be searched logarithmically
### relative to number of (separate) chains (tree-like)

"""
    nturn(hbdicts, i, n) -> Bool
    alphaturn(hbdicts, i) = nturn(hbdicts, i, 4)
    shortturn(hbdicts, i) = nturn(hbdicts, i, 3)
    piturn(hbdicts, i) = nturn(hbdicts, i, 5)

Check if there is an ``n``-turn starting from residue with index `i`. Uses
the vector `hbdicts` as reference. DSSP defines turns with specific `n` such
as ``3,4,5``-turns. These are `shortturn`, `alphaturn`, and `piturn`,
respectively, in this module.
"""
nturn(hbdicts, i, n) = hbonded(hbdicts, i, i+n)
@doc (@doc nturn) alphaturn(hbdicts, i) = nturn(hbdicts, i, 4)
@doc (@doc nturn) shortturn(hbdicts, i) = nturn(hbdicts, i, 3)
@doc (@doc nturn) piturn(hbdicts, i) = nturn(hbdicts, i, 5)

"""
    parallelbridge(hbdicts, i, j)
    antiparallelbridge(hbdicts, i, j)

Check if there is an elementary bridge between the residue with indices `i`
and `j`. Check the original DSSP reference (link in README) or the Julia
source codes for the definitions of these bridges.
"""
parallelbridge(hbdicts, i, j) = begin
    t1 = hbonded(hbdicts, i, j-1) && hbonded(hbdicts, j, i+1)
    t1 || hbonded(hbdicts, j-1, i) && hbonded(hbdicts, i, j+1)
end
@doc (@doc parallelbridge) antiparallelbridge(hbdicts, i, j) = begin
    t1 = hbonded(hbdicts, i, j) && hbonded(hbdicts, j, i)
    t1 || hbonded(hbdicts, i-1, j+1) && hbonded(hbdicts, j-1, i+1)
end

### ADDING DONORS/ACCEPTORS TO DICTIONARY

"""
    recorddonor!(hbdict, i, bondenergy)
    recordacceptor!(hbdict, i, bondenergy)

Record a donor/acceptor to the hydrogen bonding dictionary. The bond is with
the residue at index `i`. The corresponding bond energy is `bondenergy`.

Using these functions records entries with increasing energy (decreasing bond
strength).
"""
function recorddonor! end
@doc (@doc recorddonor!) function recordacceptor! end

tpes = map((:donor, :acceptor)) do nm
    Symbol(:record, nm, :!), Symbol(nm, :s)
end

for (fnm, called) in tpes
    offexp = fnm === :recorddonor! ? :(0) : :(hbdict.base_size)
    @eval begin
        function $fnm(hbdict::HBondDict, index, energy)
            offset = $offexp
            tr_ix, tr_energy = index, energy
            for (key, tpl) in $called(hbdict) |> enumerate
                ix, E = tpl
                tr_energy < E || continue
                # reaching here means energy < E; and partial roll starts
                @inbounds hbdict[key+offset] = tr_ix, tr_energy
                tr_ix, tr_energy = ix, E
            end
        end
    end
end


export CuVec
module CuVec
using ..Mocha

const THREADS_PER_BLOCK = 128
function cuda_geometry(len::Int)
  x_block = round(Int, ceil(convert(Float64, len)/THREADS_PER_BLOCK))
  return (x_block, THREADS_PER_BLOCK)
end

for (ctype, dtype) in [(:float, Float32), (:double, Float64)]
  # define add!, sub!, mul!, div!, div2!
  for name in [:add, :sub, :mul, :div, :div2]
    @eval begin
      function $(Symbol("$(name)!"))(backend::GPUBackend, ::Type{$dtype}, X, Y, len::Int)
        X = convert(Ptr{Void},X)
        Y = convert(Ptr{Void},Y)
        cuda_dim = cuda_geometry(len)
        kernel = backend.mocha.$(Symbol("elem_$(name)_$ctype"))
        CUDA.launch(kernel, cuda_dim..., (X, Y, len))
      end
    end
  end

  # define add_scal!
  @eval begin
    function add_scal!(backend::GPUBackend, ::Type{$dtype}, X, Y, len::Int)
      X = convert(Ptr{Void}, X)
      Y = convert($dtype, Y)
      cuda_dim = cuda_geometry(len)
      kernel = backend.mocha.$(Symbol("add_scal_$ctype"))
      CUDA.launch(kernel, cuda_dim..., (X,Y,len))
    end
  end

  # define log!
  @eval begin
    function log!(backend::GPUBackend, ::Type{$dtype}, X, len::Int)
      X = convert(Ptr{Void}, X)
      cuda_dim = cuda_geometry(len)
      kernel = backend.mocha.$(Symbol("elem_log_$ctype"))
      CUDA.launch(kernel, cuda_dim..., (X,len))
    end
  end

  # define exp!
  @eval begin
    function exp!(backend::GPUBackend, ::Type{$dtype}, X, len::Int)
      X = convert(Ptr{Void}, X)
      cuda_dim = cuda_geometry(len)
      kernel = backend.mocha.$(Symbol("elem_exp_$ctype"))
      CUDA.launch(kernel, cuda_dim..., (X,len))
    end
  end

  # define mul_scal!
  @eval begin
    function mul_scal!(backend::GPUBackend, ::Type{$dtype}, X, Y, len::Int)
      X = convert(Ptr{Void}, X)
      Y = convert($dtype, Y)
      cuda_dim = cuda_geometry(len)
      kernel = backend.mocha.$(Symbol("mul_scal_$ctype"))
      CUDA.launch(kernel, cuda_dim..., (X,Y,len))
    end
  end
end

# define add!, sub!, mul!, div!, div2! for blobs
for name in [:add, :sub, :mul, :div, :div2]
  @eval begin
    function $(Symbol("$(name)!")){T}(backend::GPUBackend, X::CuTensorBlob{T}, Y::CuTensorBlob{T})
      len = length(X)
      $(Symbol("$(name)!"))(backend, T, X.ptr.p, Y.ptr.p, len)
    end
  end
end
function add_scal!{T}(backend::GPUBackend, X::CuTensorBlob{T}, Y)
  Y = convert(T, Y)
  len = length(X)
  add_scal!(backend, T, X.ptr.p, Y, len)
end
function mul_scal!{T}(backend::GPUBackend, X::CuTensorBlob{T}, Y)
  Y = convert(T, Y)
  len = length(X)
  mul_scal!(backend, T, X.ptr.p, Y, len)
end
function log!{T}(backend::GPUBackend, X::CuTensorBlob{T})
  log!(backend, T, X.ptr.p, length(X))
end
function exp!{T}(backend::GPUBackend, X::CuTensorBlob{T})
  exp!(backend, T, X.ptr.p, length(X))
end

for (postfix, dt1, dt2) in [(:fi, Float32, Int), (:di, Float64, Int),
                            (:ff, Float32, Float32), (:dd, Float64, Float64)]
  @eval begin
    function pow!(backend::GPUBackend, ::Type{$dt1}, X, Y::$dt2, len::Int)
      X = convert(Ptr{Void}, X)
      cuda_dim = cuda_geometry(len)
      kernel = backend.mocha.$(Symbol("elem_pow_$postfix"))
      CUDA.launch(kernel, cuda_dim..., (X,Y,len))
    end
    function pow!{T}(backend::GPUBackend, X::CuTensorBlob{T}, Y::$dt2)
      pow!(backend, T, X.ptr.p, Y, length(X))
    end
  end
end

end # module cuVec


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Stage-decomposition optimization #
# ================================ #
function optimize!(structure::StageDecompositionStructure, optimizer::AbstractStructuredOptimizer, x₀::AbstractVector)
    # Sanity checks
    supports_structure(optimizer, structure) || throw(UnsupportedStructure{typeof(optimizer), typeof(structure)}())
    check_loadable(optimizer, structure)
    # Load structure
    load_structure!(optimizer, structure, x₀)
    # Run structure-exploiting optimization procedure
    MOI.optimize!(optimizer)
    return nothing
end

function set_master_optimizer!(structure::StageDecompositionStructure, optimizer)
    set_optimizer(structure.first_stage, optimizer)
    return nothing
end

function set_master_optimizer_attribute!(structure::StageDecompositionStructure, attr::MOI.AbstractOptimizerAttribute, value)
    MOI.set(backend(structure.first_stage), attr, value)
    return nothing
end

function set_subproblem_optimizer!(structure::StageDecompositionStructure, optimizer)
    set_optimizer!(scenarioproblems(structure), optimizer)
    return nothing
end

function set_subproblem_optimizer_attribute!(structure::StageDecompositionStructure, attr::MOI.AbstractOptimizerAttribute, value)
    MOI.set(scenarioproblems(structure), attr, value)
    return nothing
end

function cache_solution!(stochasticprogram::StochasticProgram{2}, structure::StageDecompositionStructure{2}, optimizer::MOI.AbstractOptimizer)
    cache = solutioncache(stochasticprogram)
    # Cache main solution
    variables = decision_variables_at_stage(stochasticprogram, 1)
    constraints = decision_constraints_at_stage(stochasticprogram, 1)
    cache[:solution] = SolutionCache(optimizer, variables, constraints)
    # Cache first-stage solution
    cache[:node_solution_1] = SolutionCache(backend(structure.first_stage), variables, constraints)
    # Cache scenario-dependent solutions (Skip if more than 100 scenarios for performance)
    if num_scenarios(stochasticprogram) <= 1e2
        variables = decision_variables_at_stage(stochasticprogram, 2)
        constraints = decision_constraints_at_stage(stochasticprogram, 2)
        cache_solution!(cache, scenarioproblems(structure), optimizer, 2, variables, constraints)
    end
    return nothing
end


# Relativistic CSFs

#
# type CSF
# ------------------------------------------------------------------------------------------
"""
    $(TYPEDEF)

Represents a configuration state function.

[`CSF`](@ref) is iterable, returning a tuple with the elements
`(orbital::Orbital, n_electrons, orbital_coupling, csf_coupling)` on each
iteration.
"""
struct CSF
    angularsym :: AngularSymmetry
    orbitals :: Vector{RelativisticOrbital{Int}}
    occupations :: Vector{Int}
    orbcouplings :: Vector{AngularMomentum}
    csfcouplings :: Vector{AngularMomentum}

    function CSF(
            orbitals::Vector{<:RelativisticOrbital},
            electrons::Vector{Int},
            couplings_orbitals::Vector{AngularMomentum},
            couplings_csf::Vector{AngularMomentum},
            parity,
        )
        let norbitals = length(orbitals)
            @assert length(electrons) == norbitals
            @assert length(couplings_orbitals) == norbitals
            @assert length(couplings_csf) == norbitals
        end
        @assert issorted(orbitals)
        @assert all(x -> x!=0, electrons)
        new(
            AngularSymmetry(last(couplings_csf), parity),
            orbitals, electrons, couplings_orbitals, couplings_csf,
        )
    end
end

import Base: ==
function ==(csf1::CSF, csf2::CSF)
    (csf1.angularsym == csf2.angularsym)     &&
    (csf1.orbitals == csf2.orbitals)         &&
    (csf1.occupations == csf2.occupations)   &&
    (csf1.orbcouplings == csf2.orbcouplings) &&
    (csf1.csfcouplings == csf2.csfcouplings)
end

function Base.iterate(csf::CSF, s=1)
    s > length(csf) && return nothing
    item = (csf.orbitals[s], csf.occupations[s], csf.orbcouplings[s], csf.csfcouplings[s])
    return item, s + 1
end

Base.length(csf::CSF) = length(csf.orbitals)

parity(csf::CSF) = parity(csf.angularsym)
angularmomentum(csf::CSF) = angularmomentum(csf.angularsym)
nelectrons(csf :: CSF) = sum(csf.occupations)

"""
    nexcitations(csf_from::CSF, csf_to::CSF)

Calculates the excitations between relativistic orbitals. I.e. difference of occupation
number between `nl-` to `nl` counts as an excitation.
"""
function nexcitations(csf_from::CSF, csf_to::CSF)
    @assert nelectrons(csf_from) == nelectrons(csf_to)
    let from_orbs = [(i, orb.n, kappa(orb)) for (i, orb) in enumerate(csf_from.orbitals)],
        to_orbs = [(i, orb.n, kappa(orb)) for (i, orb) in enumerate(csf_to.orbitals)]
        @assert length(from_orbs) == length(unique(from_orbs))
        @assert length(to_orbs) == length(unique(to_orbs))
    end

    to_unchecked = fill(true, length(csf_to.orbitals))

    nexcitations = 0
    nremoved = 0 # for sanity checking only

    # Let's count all the electrons that have been excited into some of
    # the occupied orbitals of csf_from.
    for (i, orb) in enumerate(csf_from.orbitals)
        toids = findall(o -> o.n == orb.n && kappa(o) == kappa(orb), csf_to.orbitals)
        to_nelectrons = if length(toids) == 0
            0
        elseif length(toids) == 1
            # Mark this csf_to orbital checked.
            toid = first(toids)
            @assert to_unchecked[toid]
            to_unchecked[toid] = false
            csf_to.occupations[toid]
        else
            error("Too many matches ($(length(toids))) among csf_to.orbitals.")
        end

        diff = to_nelectrons - csf_from.occupations[i]
        # If diff < 0, electrons have been excited from this orbital
        (diff < 0) && (nremoved += abs(diff)) # for sanity checking
        nexcitations += (diff > 0) ? diff : 0
    end

    # Now let's loop over all the remaning csf_to orbitals
    for (i, orb) in enumerate(csf_to.orbitals)
        to_unchecked[i] || continue
        nexcitations += csf_to.occupations[i]
    end

    @assert nexcitations == nremoved # sanity check
    return nexcitations
end

function Base.show(io::IO, csf::CSF)
    write(io, '"')
    print(io, csf)
    write(io, '"')
end

function Base.print(io::IO, csf::CSF)
    first = true
    for (orb, nelec, orbcoupling, csfcoupling) in csf
        !first && write(io, ' ')
        print(io, orb)
        write(io, '(', string(nelec), '|', string(orbcoupling), '|', string(csfcoupling), ')')
        first = false
    end
    write(io, " | ", string(csf.angularsym))
    nothing
end

function kappa2rso(kappa :: Integer)
    lstr = specname(κ2ℓ(kappa))
    (kappa < 0) ? lstr : lstr*"-"
end


function get_rexhopper(;
    timestep=0.01,
    gravity=[0.0; 0.0; -9.81],
    friction_coefficient=1.0,
    contact_foot=true,
    contact_body=true,
    limits=true,
    model=:rexhopper,
    floating=true,
    contact_type=:nonlinear,
    spring=0.0,
    damper=0.0,
    parse_damper=true,
    T=Float64)

    path = joinpath(@__DIR__, "../deps/$(String(model)).urdf")
    mech = Mechanism(path; floating, T,
        gravity,
        timestep,
        parse_damper)

    # Adding springs and dampers
    set_springs!(mech.joints, spring)
    set_dampers!(mech.joints, damper)

    # joint limits
    joints = deepcopy(mech.joints)

    if limits
        joint1 = get_joint(mech, :joint1)
        joint3 = get_joint(mech, :joint3)
        joints[joint1.id] = add_limits(mech, joint1,
            rot_limits=[SVector{1}(-0.7597), SVector{1}(1.8295)])
        joints[joint3.id] = add_limits(mech, joint3,
            rot_limits=[SVector{1}(-1.8295), SVector{1}(0.7597)])
        mech = Mechanism(Origin{T}(), [mech.bodies...], [joints...];
            gravity,
            timestep)
    end

    if contact_foot
        origin = Origin{T}()
        bodies = mech.bodies
        joints = mech.joints

        normal = [0.0; 0.0; 1.0]
        models = []

        link3 = get_body(mech, :link3)
        link2 = get_body(mech, :link2)
        foot_radius = 0.0203
        ankle_radius = 0.025
        base_radius = 0.14
        p = [0.1685; 0.0025; -0.0055]
        o = foot_radius
        push!(models, contact_constraint(link3, normal;
            friction_coefficient,
            contact_origin=p,
            contact_radius=o,
            contact_type,
            name=:foot))
        p = [-0.10; -0.002; 0.01]
        o = ankle_radius
        push!(models, contact_constraint(link3, normal;
            friction_coefficient,
            contact_origin=p,
            contact_radius=o,
            contact_type,
            name=:ankle3))
        p = [0.24; 0.007; 0.005]
        push!(models, contact_constraint(link2, normal;
            friction_coefficient,
            contact_origin=p,
            contact_radius=o,
            contact_type,
            name=:ankle2))
        base_link = get_body(mech, :base_link)
        pl = [0.0; +0.075; 0.03]
        pr = [0.0; -0.075; 0.03]
        o = base_radius
        push!(models, contact_constraint(base_link, normal;
            friction_coefficient,
            contact_origin=pl,
            contact_radius=o,
            contact_type,
            name=:torso_left))
        push!(models, contact_constraint(base_link, normal;
            friction_coefficient,
            contact_origin=pr,
            contact_radius=o,
            contact_type,
            name=:torso_right))

        set_minimal_coordinates!(mech, get_joint(mech, :floating_base), [0,0,1.0, 0,0,0])

        mech = Mechanism(origin, bodies, joints, [models...];
            gravity,
            timestep)
    end
    return mech
end

function initialize_rexhopper!(mechanism::Mechanism{T};
        body_position=zeros(3),
        body_orientation=zeros(3),
        body_linear_velocity=zeros(3),
        body_angular_velocity=zeros(3)) where T

    zero_velocity!(mechanism)
    body_position += [0.0, 0.0, 0.3148]

    for joint in mechanism.joints
        !(joint.name in (:loop_joint, :floating_joint)) && set_minimal_coordinates!(mechanism, joint, zeros(input_dimension(joint)))
    end

    set_minimal_coordinates!(mechanism, get_joint(mechanism,
        :floating_base), [body_position; body_orientation])
    set_minimal_velocities!(mechanism, get_joint(mechanism,
        :floating_base), [body_linear_velocity; body_angular_velocity])
end


using GLWindow, GLAbstraction, GLFW, ModernGL, ImmutableArrays, WavefrontObj, FixedPointNumbers, Images, Colorant, Reactive
using GLPlot #toopengl
# window creation
window = createwindow("OBJ-Viewer", 1000, 1000, debugging = false)
cam = PerspectiveCamera(window.inputs, Vec3(2,2,0.5), Vec3(0.0))

# render objects creation
shader = TemplateProgram(Pkg.dir("GLPlot", "src", "shader", "standard.vert"), Pkg.dir("GLPlot", "src", "shader", "phongblinn.frag"))
RGBAU8 = AlphaColorValue{RGB{UFixed8}, UFixed8}
Colorant.rgba(r::Real, g::Real, b::Real, a::Real)    = AlphaColorValue(RGB{Float32}(r,g,b), Float32(a))
rgbaU8(r::Real, g::Real, b::Real, a::Real)  = AlphaColorValue(RGB{UFixed8}(r,g,b), ufixed8(a))
#GLPlot.toopengl{T <: AbstractRGB}(colorinput::Signal{T}) = toopengl(const_lift(x->AlphaColorValue(x, one(T)), RGBA{T}, colorinput))
tohsva(rgba)     = AlphaColorValue(convert(HSV, rgba.c), rgba.alpha)
torgba(hsva)     = AlphaColorValue(convert(RGB, hsva.c), hsva.alpha)
tohsva(h,s,v,a)  = AlphaColorValue(HSV(Float32(h), Float32(s), Float32(v)), Float32(a))
Base.convert{T <: AbstractAlphaColorValue}(typ::Type{T}, x::AbstractAlphaColorValue) = AlphaColorValue(convert(RGB{eltype(typ)}, x.c), convert(eltype(typ), x.alpha))

function unitGeometry{T}(geometry::Vector{Vector3{T}})
    assert(!isempty(geometry))

    xmin = typemax(T)
    ymin = typemax(T)
    zmin = typemax(T)

    xmax = typemin(T)
    ymax = typemin(T)
    zmax = typemin(T)

    for vertex in geometry
        xmin = min(xmin, vertex[1])
        ymin = min(ymin, vertex[2])
        zmin = min(zmin, vertex[3])

        xmax = max(xmax, vertex[1])
        ymax = max(ymax, vertex[2])
        zmax = max(zmax, vertex[3])
    end

    xmiddle = xmin + (xmax - xmin) / 2;
    ymiddle = ymin + (ymax - ymin) / 2;
    zmiddle = zmin + (zmax - zmin) / 2;
    scale = 2 / max(xmax - xmin, ymax - ymin, zmax - zmin);

    result = similar(geometry)

    for i = 1:length(result)
        result[i] = Vector3{T}((geometry[i][1] - xmiddle) * scale,
                               (geometry[i][2] - ymiddle) * scale,
                               (geometry[i][3] - zmiddle) * scale
                    );
    end

    return result
end

objpath = "Butterfly/Butterfly.obj"

assets_path = "Butterfly/"

obj = readObjFile(objpath, faceindextype=GLuint, vertextype=Float32, compute_normals = false, triangulate = false)
computeNormals!(obj, smooth_normals = true, override = false)
triangulate!(obj)

# center geometry
obj.vertices = unitGeometry(obj.vertices)

# load mtl files if present
materials = WavefrontMtlMaterial{Float32}[]

for mtllib in obj.mtllibs
    materials = [materials, readMtlFile( assets_path*mtllib, colortype=Float32 )]
end

render_objects = RenderObject[]

compiled_materials = Dict()

const light         = Vec3[Vec3(1.0,1.0,1.0), Vec3(0.1,0.1,0.1), Vec3(0.9,0.9,0.9), Vec3(20.0,20.0,20.0)]
const material      = Vec3[Vec3(1.0,1.0,1.0), Vec3(1.0,1.0,1.0), Vec3(1.0,1.0,1.0), Vec3(1.0,1.0,1.0)]
const tmaterial      = Vec4[Vec4(1.0,1.0,1.0,1.0), Vec4(1.0,1.0,1.0,1.0), Vec4(1.0,1.0,1.0,1.0), Vec4(1.0,1.0,1.0,1.0)]
const tmaterialused  = GLint[-1,-1,-1,-1]

for material_name in collect(keys(obj.materials))

    (vs, nvs, uvs, fcs) = compileMaterial(obj, material_name)

    # hack: invert normals for glabstraction
    nvs = -nvs

    # holding global references seems necessary here
    # compiled_materials[material_name] = (vs, nvs, uvs, fcs, lines)
    data = [
        :vertex          => GLBuffer(vs),
        :normal          => GLBuffer(nvs),
        :uv              => GLBuffer(uvs),
        :indexes         => indexbuffer(fcs),
        :view            => cam.view,
        :projection      => cam.projection,
        :normalmatrix    => cam.normalmatrix,
        :eyeposition     => cam.eyeposition,
        :model           => eye(Mat4),
        :material        => material,
        :light           => light,
        :textures_used   => tmaterialused,
    ]

    # search for a material with the given name
    texture_array = Any[]
    for mtl in materials
        if mtl.name == material_name
            data[:material]             = Vec3[mtl.diffuse, mtl.ambient, mtl.specular, Vec3(mtl.specular_exponent)]
            println(data[:material])
            if mtl.diffuse_texture != ""
                data[:textures_used][1] = length(texture_array) # insert texture array index
                push!(texture_array, imread(assets_path*"Texture/"*mtl.diffuse_texture).data)
            end
            if mtl.ambient_texture != ""
                data[:textures_used][2] = length(texture_array) # insert texture array index
                push!(texture_array, imread(assets_path*"Texture/"*mtl.ambient_texture).data)
            end
            if mtl.specular_texture != ""
                data[:textures_used][3] = length(texture_array) # insert texture array index
                push!(texture_array, imread(assets_path*"Texture/"*mtl.specular_texture).data)
            end
            #=
            if mtl.bump_texture != ""
                data[:textures_used] = true
                data[:tmaterialused][1] = 0 # insert texture array index
                push!(texture_array, imread(assets_path*mtl.diffuse_texture).data)
            end
            =#
            break
        end
    end

    if !isempty(texture_array)
        data[:textures_used] = true
        data[:texture_maps] = Texture(convert(Vector{Matrix{eltype(first(texture_array))}}, texture_array))
    else
        data[:texture_maps] = Texture(Matrix{RGBA{UFixed8}}[fill(rgbaU8(0,0,0,0), 1,1)])
    end

    ro = RenderObject(data, shader)

    postrender!(ro, render, ro.vertexarray)
    #postrender!(ro, render, ro.vertexarray, GL_LINES)

    push!(render_objects, ro)
end

# OpenGL setup
glClearColor(0.2,0.2,0.2,1)
glDisable(GL_CULL_FACE)
glEnable(GL_DEPTH_TEST)

const_lift(window.inputs[:framebuffer_size]) do wh
    glViewport(0,0,wh...)
end |> preserve
# Loop until the user closes the window
while !GLFW.WindowShouldClose(window.nativewindow)

  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

  # render materials separately
  for ro in render_objects
    render(ro)
  end

  yield() # this is needed for react to work

  GLFW.SwapBuffers(window.nativewindow)
  GLFW.PollEvents()
end

GLFW.Terminate()


# EFTfitter.jl - Empty Template
# ============= Parameters =============================================#
parameters = BAT.NamedTupleDist(
    p1 = -2..2,
)

# ============= Observables ============================================#
function observable1(params)
    return params.p1
end

# ============= Measurements ===========================================#
measurements = (
    Meas1 = Measurement(observable1, 0.0, uncertainties = (unc1 = 0.1,), active=true),

    #MeasDist = MeasurementDistribution(obs_array, values_array, uncertainties = (unc1 = unc1_array,), active=false),
)

# ============= Correlations ===========================================#
correlations = (
    unc1 = NoCorrelation(active=true),
)

#corr_matrix = to_correlation_matrix(measurements,
#  (:Meas1, :Meas2, 0.1),
#)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


let
    State = TaskGraphs.State
    Action = TaskGraphs.Action

    f = pctapf_problem_1
    cost_model = SumOfMakeSpans()
    solver = NBSSolver(assignment_model = TaskGraphsMILPSolver(GreedyAssignment()))
    pc_tapf = f(solver;cost_function=cost_model,verbose=false)

    base_search_env = pc_tapf.env
    prob = formulate_assignment_problem(solver.assignment_model,pc_tapf)
    sched, cost = solve_assignment_problem!(solver.assignment_model,prob,pc_tapf)
    @test validate(sched)
    search_env = construct_search_env(solver,deepcopy(sched),base_search_env)
    pc_mapf = PC_MAPF(search_env)
    node = initialize_root_node(search_env)
    let
        env = build_env(solver,pc_mapf,search_env,node,AgentID(1))

        s = State()
        a = Action()
        get_next_state(s,a)
        get_next_state(env,s,a)

        @test State(1,2) == State(1,2)
        @test State() == State()
        @test states_match(State(),State())
        @test states_match(env,State(),State())
        CRCBS.wait(s)
        CRCBS.wait(env,s)

        get_cost_model(env)
        get_heuristic_model(env)
        @test CRCBS.is_valid(env,State()) == false
        @test CRCBS.is_valid(env,Action()) == false
        @test CRCBS.is_valid(env,State(1,2))
        @test CRCBS.is_valid(env,Action(Edge(1,2),1))

        get_heuristic_cost(env,State(1,2))
        # @test is_goal(env,State(1,0)) == false
        # @test is_goal(env,State(2,3)) == false
        # @test is_goal(env,State(1,3))
        @test length(get_possible_actions(env,State(-1,-1))) == 0
    end
    let
        for v in vertices(search_env.schedule)
            schedule_node = get_node_from_vtx(search_env.schedule,v)
            if isa(schedule_node,COLLECT)
                env = build_env(solver,pc_mapf,search_env,node,VtxID(v))
                @test length(get_possible_actions(env,State(1,0))) == 1
                @test length(get_possible_actions(env,State(-1,-1))) == 0
                break
            end
        end
    end
end


using MPIReco

@testset "single- and multi-channel in-memory reconstruction" begin
  bSF = MPIFile(joinpath(datadir, "calibrations", "12.mdf"))
  b = MPIFile(joinpath(datadir, "measurements", "20211226_203916_MultiPatch", "1.mdf"))
  names = (:color, :x, :y, :z, :time)
  values = (1:1,
	    -27.375u"mm":1.25u"mm":11.375u"mm",
	    -11.375u"mm":1.25u"mm":27.375u"mm",
	    0.0u"mm":1.0u"mm":0.0u"mm",
	    0.0u"ms":0.6528u"ms":0.0u"ms")

  # standard reconstruction
  c1 = reconstruction(bSF, b; SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, iterations=1, spectralLeakageCorrection=true)
  @test axisnames(c1) == names
  @test axisvalues(c1) == values
  exportImage(joinpath(imgdir, "Reconstruction1.png"), arraydata(c1[1,:,:,1,1]))

  # fused lasso
  c2 = reconstruction(bSF, b; SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, iterations=100, solver="fusedlasso",
		      loadasreal=true, λ=[0.01,0.01])
  @test axisnames(c2) == names
  @test axisvalues(c2) == values
  exportImage(joinpath(imgdir, "Reconstruction2.png"), arraydata(c2[1,:,:,1,1]))

  # with interpolation
  c3 = reconstruction(bSF, b; SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, gridsize=[100,100,1],iterations=1)
  @test axisnames(c3) == names
  @test axisvalues(c3) == (values[1], -27.8u"mm":0.4u"mm":11.8u"mm", -11.8u"mm":0.4u"mm":27.8u"mm", values[4:5]...)
  exportImage(joinpath(imgdir, "Reconstruction3.png"), arraydata(c3[1,:,:,1,1]))

  # with fov adpation and center shift
  c4 = reconstruction(bSF, b; SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, gridsize=[50,50,1],fov=[0.04,0.04,1],
		      center=[0.0,-0.01,0], iterations=1)
  @test axisnames(c4) == names
  # TODO Tobi: does this make sense?
  @test axisvalues(c4) == (values[1], -27.6u"mm":0.8u"mm":11.6u"mm", -11.6u"mm":0.8u"mm":27.6u"mm", 499.5u"mm":1000.0u"mm":499.5u"mm", values[5])
  exportImage(joinpath(imgdir, "Reconstruction4.png"), arraydata(c4[1,:,:,1,1]))

  # multi colored reconstruction (deprecated interface)
  c5 = reconstruction([bSF,bSF], b;
		      SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, iterations=1)
  @test axisnames(c5) == names
  @test axisvalues(c5) == (1:2,values[2:end]...)
  exportImage(joinpath(imgdir, "Reconstruction5a.png"), arraydata(c5[1,:,:,1,1]))
  exportImage(joinpath(imgdir, "Reconstruction5b.png"), arraydata(c5[2,:,:,1,1]))

  # multi colored reconstruction (new interface)
  c5 = reconstruction(MultiContrastFile([bSF,bSF]), b;
          SNRThresh=5, frames=1, minFreq=80e3,
          recChannels=1:2, iterations=1)
  @test axisnames(c5) == names
  @test axisvalues(c5) == (1:2,values[2:end]...)
  exportImage(joinpath(imgdir, "Reconstruction5c.png"), arraydata(c5[1,:,:,1,1]))
  exportImage(joinpath(imgdir, "Reconstruction5d.png"), arraydata(c5[2,:,:,1,1]))

  # dict based reco
  r = defaultRecoParams()
  r[:measPath] = filepath(b)
  r[:SFPath] = filepath(bSF)
  r[:frames] =  1
  r[:minFreq] = 80e3
  r[:SNRThresh] = 4
  r[:lambd] = 0.001
  r[:iterations] = 1
  c6 = reconstruction(r)
  @test axisnames(c6) == names
  @test axisvalues(c6) == values
  exportImage(joinpath(imgdir, "Reconstruction6.png"), arraydata(c6[1,:,:,1,1]))

  # channel weighting
  c7a = reconstruction(bSF, b; SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, iterations=1, λ=0.1,
          weightType=WeightingType.Channel, channelWeights=[1.0,1.0,1.0])

  exportImage(joinpath(imgdir, "Reconstruction7a.png"), arraydata(c7a[1,:,:,1,1]))

  c7b = reconstruction(bSF, b; SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, iterations=1, λ=0.1,
          weightType=WeightingType.Channel, channelWeights=[1.0,0.001,1.0])

  exportImage(joinpath(imgdir, "Reconstruction7b.png"), arraydata(c7b[1,:,:,1,1]))

  c7c = reconstruction(bSF, b; SNRThresh=5, frames=1, minFreq=80e3,
		      recChannels=1:2, iterations=1, λ=0.1,
          weightType=WeightingType.Channel, channelWeights=[0.001,1.0,1.0])

  exportImage(joinpath(imgdir, "Reconstruction7c.png"), arraydata(c7c[1,:,:,1,1]))

end


function load_lattice_ase(T, pyobj::PyObject)
    ase = pyimport("ase")

    if pyisinstance(pyobj, ase.Atoms)
        if !all(pyobj.pbc)
            error("DFTK only supports calculations with periodic boundary conditions.")
        end
        load_lattice_ase(T, pyobj.cell)
    elseif pyisinstance(pyobj, ase.cell.Cell)
        lattice = zeros(3, 3)
        cell_julia = convert(Array, pyobj)  # Array of arrays
        for i = 1:3, j = 1:3
            lattice[i, j] = austrip(cell_julia[j][i] * u"Å")
        end
        Mat3{T}(lattice)
    else
        error("load_lattice_ase not implemented for python type $pyobj")
    end
end

function load_atoms_ase(T, pyobj::PyObject)
    @assert pyisinstance(pyobj, pyimport("ase").Atoms)
    # TODO Be smarter and look at the calculator to determine the psps

    frac_coords = pyobj.get_scaled_positions()
    map(ase_atoms_translation_map(pyobj)) do (symbol, atom_indices)
        coords = Vec3{T}[]
        for idx in atom_indices
            push!(coords, frac_coords[idx, :])
        end
        ElementCoulomb(symbol) => coords
    end
end

function load_magnetic_moments_ase(pyobj::PyObject)
    @assert pyisinstance(pyobj, pyimport("ase").Atoms)
    magmoms = pyobj.get_initial_magnetic_moments()
    map(ase_atoms_translation_map(pyobj)) do (symbol, atom_indices)
        ElementCoulomb(symbol) => [normalize_magnetic_moment(magmoms[i])
                                   for i in atom_indices]
    end
end

"""
Map translating between the `atoms` datastructure used in DFTK
and the indices of the respective atoms in the ASE.Atoms object
in `pyobj`.
"""
function ase_atoms_translation_map(pyobj::PyObject)
    @assert pyisinstance(pyobj, pyimport("ase").Atoms)
    numbers = pyobj.get_atomic_numbers()
    map(unique(numbers)) do atnum
        ElementCoulomb(atnum).symbol => findall(isequal(atnum), numbers)
    end
end

#                                                             convert A.U. -> Å
ase_cell(lattice) = pyimport("ase").cell.Cell(Array(lattice)' / austrip(1u"Å"))
ase_cell(model::Model) = ase_cell(model.lattice)

function ase_atoms(lattice_or_model, atoms, magnetic_moments=[])
    cell = ase_cell(lattice_or_model)
    symbols = String[]
    for (elem, pos) in atoms
        append!(symbols, fill(string(atomic_symbol(elem)), length(pos)))
    end
    scaled_positions = vcat([pos for (elem, pos) in atoms]...)

    magmoms = nothing
    if !isempty(magnetic_moments)
        @assert length(magnetic_moments) == length(atoms)
        get_z_magmom(m) = normalize_magnetic_moment(m)[3]
        magmoms = vcat([get_z_magmom.(magmom)
                        for (elem, magmom) in magnetic_moments]...)
        @assert length(magmoms) == length(scaled_positions)
    end

    pyimport("ase").Atoms(symbols=symbols, cell=cell, pbc=true,
                          scaled_positions=hcat(scaled_positions...)',
                          magmoms=magmoms)
end
function ase_atoms(model::Model, magnetic_moments=nothing)
    ase_atoms(model.lattice, model.atoms, magnetic_moments)
end


using Test
using TightlyBound
using Suppressor
#=
include("../Utility.jl")
include("../BandTools.jl")
include("../Atomic.jl")
include("../Atomdata.jl")
include("../Crystal.jl")
include("../DFToutMod.jl")
include("../TB.jl")
include("../CalcTB.jl")
include("../FitTB.jl")
=#

#basic loading

#include("../includes_laguerre.jl")

function test1()

    @testset "testing laguerre fitting fake example" begin

        c1 = makecrys([7.0 0 0; 0 7.0 0; 0 0 7.0], [0 0 0], ["Li"], units="Bohr");
        c2 = makecrys([8.0 0 0; 0 7.0 0; 0 0 7.0], [0 0 0], ["Li"], units="Bohr");
        c3 = makecrys([8.5 0 0; 0 8.5 0; 0 0 7.0], [0 0 0], ["Li"], units="Bohr");
        c4 = makecrys([6.0 0 0; 0 6.0 0; 0 0 6.0], [0 0 0], ["Li"], units="Bohr");

        c5 = makecrys([14.0 0 0; 0 7.0 0; 0 0 7.0], [0 0 0; 0.49 0 0 ], ["Li", "Li"], units="Bohr");
        c6 = makecrys([14.0 0 0; 0 7.0 0; 0 0 7.0], [0 0 0; 0.48 0 0 ], ["Li", "Li"], units="Bohr");
        c7 = makecrys([14.0 0 0; 0 7.0 0; 0 0 7.0], [0 0 0; 0.47 0 0 ], ["Li", "Li"], units="Bohr");
        c8 = makecrys([14.0 0 0; 0 7.0 0; 0 0 7.0], [0 0 0; 0.46 0 0 ], ["Li", "Li"], units="Bohr");
        c9 = makecrys([14.0 0 0; 0 7.0 0; 0 0 7.0], [0 0 0; 0.45 0 0 ], ["Li", "Li"], units="Bohr");

        @suppress begin
#        if true

            database = Dict()
            database[(:Li, :Li)] = TightlyBound.CalcTB.make_coefs(Set(["Li", "Li"]), 2)

            tbc1 = TightlyBound.CalcTB.calc_tb_fast(c1, database, use_threebody=false);
            tbc2 = TightlyBound.CalcTB.calc_tb_fast(c2, database, use_threebody=false);
            tbc3 = TightlyBound.CalcTB.calc_tb_fast(c3, database, use_threebody=false);
            tbc4 = TightlyBound.CalcTB.calc_tb_fast(c4, database, use_threebody=false);
            tbc5 = TightlyBound.CalcTB.calc_tb_fast(c5, database, use_threebody=false);
            tbc6 = TightlyBound.CalcTB.calc_tb_fast(c6, database, use_threebody=false);
            tbc7 = TightlyBound.CalcTB.calc_tb_fast(c7, database, use_threebody=false);
            tbc8 = TightlyBound.CalcTB.calc_tb_fast(c8, database, use_threebody=false);
            tbc9 = TightlyBound.CalcTB.calc_tb_fast(c9, database, use_threebody=false);

            tbc_list = [tbc1 tbc2 tbc3 tbc4 tbc5 tbc6 tbc7 tbc8 tbc9]
            newdatabase = TightlyBound.FitTB.do_fitting(tbc_list, fit_threebody=false, do_plot=false)
            #        println("newdartabase")
            #        println(newdatabase[("Li", "Li")])
            @test sum(abs.(newdatabase[(:Li, :Li)].datH .- database[(:Li, :Li)].datH)) ≤ 1e-5
        end

    end
end

test1()


using Random

using Dates, TimeZones
using StatsBase, Impute
using JuliaDB, DataFrames, DataFramesMeta
using Mise
using MicroLogging

function run_model()

    @info "Start OU Process Prediction"
    flush(stdout); flush(stderr)

    seoul_codes = [
        111121,111123,111131,111141,111142,
        111151,111152,111161,111171,111181,
        111191,111201,111212,111221,111231,
        111241,111251,111261,111262,111273,
        111274,111281,111291,111291,111301,
        111301,111311]
    seoul_names = [
        "중구","종로구","용산구","광진구","성동구",
        "중랑구","동대문구","성북구","도봉구","은평구",
        "서대문구","마포구","강서구","구로구","영등포구",
        "동작구","관악구","강남구","서초구","송파구",
        "강동구","금천구","강북구","강북구","양천구",
        "양천구","노원구"]
    # construct named tuple
    seoul_stations = (; zip(Symbol.(seoul_names), seoul_codes)...)

    # NamedTuple (::Symbol -> ::DataFrame)
    rawdf = load_data_DNN("/input/jongro_seoul.csv", seoul_stations)
    #=
    first(df, 5)
    │ Row │ stationCode │ date                      │ lat      │ long      │ SO2     │ CO      │ O3      │ NO2     │ PM10    │ PM25    │ temp    │ u       │ v       │ pres    │ prep     │ snow     │
    │     │ Int64       │ TimeZones.ZonedDateTime   │ Float64  │ Float64   │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64  │ Float64  │
    ├─────┼─────────────┼───────────────────────────┼──────────┼───────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤─────────┤──────────┤──────────┤
    │ 1   │ 111123      │ 2015-01-01T01:00:00+09:00 │ 37.572   │ 127.005   │ 0.004   │ 0.2     │ 0.02    │ 0.009   │ 57.0    │ 0.0     │ -7.4    │ 4.41656 │ 1.60749 │ 1011.8  │ missing  │ missing  │
    │ 2   │ 111123      │ 2015-01-01T02:00:00+09:00 │ 37.572   │ 127.005   │ 0.005   │ 0.2     │ 0.019   │ 0.008   │ 70.0    │ 3.0     │ -8.0    │ 4.22862 │ 1.53909 │ 1011.7  │ missing  │ missing  │
    │ 3   │ 111123      │ 2015-01-01T03:00:00+09:00 │ 37.572   │ 127.005   │ 0.005   │ 0.2     │ 0.02    │ 0.006   │ 92.0    │ 5.0     │ -8.4    │ 3.57083 │ 1.29968 │ 1012.1  │ missing  │ missing  │
    │ 4   │ 111123      │ 2015-01-01T04:00:00+09:00 │ 37.572   │ 127.005   │ 0.004   │ 0.2     │ 0.019   │ 0.005   │ 111.0   │ 2.0     │ -8.8    │ 4.60449 │ 1.6759  │ 1012.3  │ missing  │ missing  │
    │ 5   │ 111123      │ 2015-01-01T05:00:00+09:00 │ 37.572   │ 127.005   │ 0.005   │ 0.2     │ 0.019   │ 0.006   │ 127.0   │ 5.0     │ -9.1    │ 5.35625 │ 1.94951 │ 1011.8  │ missing  │ missing  │
    =#

    #===== start of parameter zone =====#
    total_fdate, total_tdate = get_date_range(rawdf)
    train_fdate = ZonedDateTime(2012, 1, 1, 0, tz"Asia/Seoul")
    train_tdate = ZonedDateTime(2017, 12, 31, 23, tz"Asia/Seoul")
    test_fdate = ZonedDateTime(2018, 1, 1, 1, tz"Asia/Seoul")
    test_tdate = ZonedDateTime(2018, 12, 31, 23, tz"Asia/Seoul")

    # due to seasonality, train date range must be over 1 year
    @assert train_tdate - train_fdate >= Dates.Day(365)

    # stations
    #=
    train_stn_names = [
        :중구,:종로구,:용산구,:광진구,:성동구,
        :중랑구,:동대문구,:성북구,:도봉구,:은평구,
        :서대문구,:마포구,:강서구,:구로구,:영등포구,
        :동작구,:관악구,:강남구,:서초구,:송파구,
        :강동구,:금천구,:강북구,:강북구,:양천구,
        :양천구,:노원구]
    =#
    train_stn_names = [:종로구, :강서구, :송파구, :강남구]
    train_stn_names = [:종로구]
    #=
    train_stn_names = [
        :중구,:종로구,:용산구,:광진구,:성동구,
        :중랑구,:동대문구,:성북구,:도봉구,:은평구,
        :서대문구,:마포구,:강서구,:구로구,:영등포구]
    =#

    # test set is Jongro only
    test_stn_names = [:종로구]

    df = filter_raw_data(rawdf, train_fdate, train_tdate, test_fdate, test_tdate)
    @show first(df, 10)

    features = [:SO2, :CO, :O3, :NO2, :PM10, :PM25, :temp, :u, :v, :pres, :humid, :prep, :snow]
    # If you want exclude some features, modify train_features
    # exclude :PM10, :PM25 temporarily for log transform
    train_features = [:SO2, :CO, :O3, :NO2, :PM10, :PM25, :temp, :u, :v, :pres, :humid, :prep, :snow]
    target_features = [:PM10, :PM25]

    # For GPU, change precision of Floating numbers
    _eltype::DataType = Float32
    scaling_method = :zscore

    sample_size = 24
    output_size = 24
    epoch_size = 300
    batch_size = 32
    input_size = sample_size * length(train_features)
    #===== end of parameter zone =====#

    log_prefix = "log_"
    inv_prefix = "inv_"
    scaled_prefix = "scaled_"
    scaled_features = [Symbol(scaled_prefix, f) for f in features]
    scaled_train_features = [Symbol(scaled_prefix, f) for f in train_features]
    scaled_target_features = [Symbol(scaled_prefix, f) for f in target_features]

    # Preprocessing
    # 0. imputation
    DataFrames.allowmissing!(df, train_features)
    for ycol in train_features
        df[!, ycol] = Impute.srs(df[!, ycol])
    end
    DataFrames.disallowmissing!(df, train_features)

    # 2. Scaling
    if scaling_method == :zscore
        # scaling except :PM10, :PM25
        zscore!(df, train_features, scaled_train_features)

        statvals = extract_col_statvals(df, train_features)
    elseif scaling_method == :minmax
        minmax_scaling!(df, train_features, scaled_train_features, 0.0, 10.0)

        statvals = extract_col_statvals(df, train_features)
    elseif scaling_method == :logzscore
        # zscore except targets
        _train_features = setdiff(train_features, target_features)
        _scaled_train_features = setdiff(scaled_train_features, scaled_target_features)
        zscore!(df, _train_features, _scaled_train_features)

        # Log transform :PM10, :PM25
        # add 10.0 not to log zero values
        # Y = zscore(log(X + 10))
        log_target_features = [Symbol(log_prefix, f) for f in target_features]
        for (target, log_target) in zip(target_features, log_target_features)
            df[!, log_target] = log.(df[!, target] .+ 10.0)
        end

        zscore!(df, log_target_features, scaled_target_features)

        statvals = extract_col_statvals(df, vcat(train_features, log_target_features))
    elseif scaling_method == :invzscore
        # zscore except targets
        _train_features = setdiff(train_features, target_features)
        _scaled_train_features = setdiff(scaled_train_features, scaled_target_features)
        zscore!(df, _train_features, _scaled_train_features)

        # Inverse transform :PM10, :PM25
        # add 10.0 not to divide zero value
        # Y = zscore(1.0 / (X + 10))
        inv_target_features = [Symbol(inv_prefix, f) for f in target_features]
        for (target, inv_target) in zip(target_features, inv_target_features)
            df[!, inv_target] = 1.0 ./ (df[!, target] .+ 10.0)
        end

        zscore!(df, inv_target_features, scaled_target_features)

        statvals = extract_col_statvals(df, vcat(train_features, inv_target_features))
    elseif scaling_method == :logminmax
        # minmax except targets
        _train_features = setdiff(train_features, target_features)
        _scaled_train_features = setdiff(scaled_train_features, scaled_target_features)
        minmax_scaling!(df, _train_features, _scaled_train_features, 0.0, 10.0)

        # Log transform :PM10, :PM25
        # add 10.0 not to log zero values
        log_target_features = [Symbol(log_prefix, f) for f in target_features]
        for (target, log_target) in zip(target_features, log_target_features)
            df[!, log_target] = log.(df[!, target] .+ 10.0)
        end

        # then minmax
        minmax_scaling!(df, log_target_features, scaled_target_features, 0.0, 10.0)

        statvals = extract_col_statvals(df, vcat(train_features, log_target_features))
    end

    @show first(df, 5)
    # convert Float types
    for feature in features
        df[!, feature] = _eltype.(df[!, feature])
    end

    # plot histogram regardless to station
    for target in target_features
        plot_histogram(df, target, "/mnt/")
        plot_histogram(df, Symbol(:scaled_, target), "/mnt/")

        if scaling_method == :logzscore || scaling_method == :logminmax
            plot_histogram(df, Symbol(:log_, target), "/mnt/")
        end

        if scaling_method == :invzscore
            plot_histogram(df, Symbol(:inv_, target), "/mnt/")
        end
    end

    # line plot of first train station
    plot_lineplot_total(filter_station(df, seoul_stations[train_stn_names[1]]), :PM10, "/mnt/")
    plot_lineplot_total(filter_station(df, seoul_stations[train_stn_names[1]]), :PM25, "/mnt/")
    #plot_pcorr(df, scaled_features, features, "/mnt/")

    # windowed dataframe for train/valid (input + output)
    test_wd = []

    # random permutation train_valid_wd itself for splitting train/valid set
    rng = MersenneTwister()

    # windows dataframe for test (only for 종로구)
    for name in test_stn_names
        code = seoul_stations[name]
        stn_df = filter_station_DNN(df, code)

        # to add zero(0) offset for X₀
        push!(test_wd,
            window_df(stn_df, sample_size, output_size + 1, test_fdate, test_tdate))
    end
    # Flatten
    test_wd = collect(Base.Iterators.flatten(test_wd))
    test_size = length(test_wd)

    # simply collect dates, determine exact date for prediction (for 1h, 24h, and so on) later
    test_dates = collect(test_fdate + Hour(sample_size - 1):Hour(1):test_tdate - Hour(output_size))

    for target in target_features
        @info "$(string(target)) Training..."
        flush(stdout); flush(stderr)

        @info "training feature : " train_features
        @info "sizes (sample, output, epoch, batch) : ", sample_size, output_size, epoch_size, batch_size

        for name in test_stn_names
            code = seoul_stations[name]
            stn_df = filter_station_DNN(df, code)

            Base.Filesystem.mkpath("/mnt/$(name)/")

            max_time = input_size * 15
            acf = StatsBase.autocor(stn_df[!, Symbol(scaled_prefix, target)], 0:max_time)

            # free minibatch after training because of memory usage
            ou_df = evolve_OU(test_wd, target, Symbol(scaled_prefix, target),
                acf, statvals, _eltype, sample_size, output_size, test_dates, "/mnt/$(name)/")

            # create directory per each time
            for i = 1:output_size
                i_pad = lpad(i, 2, '0')
                Base.Filesystem.mkpath("/mnt/$(name)/$(i_pad)/")
            end

            dfs_out = export_CSV(DateTime.(test_dates, Local), ou_df, target, output_size, "/mnt/$(name)/", String(target))
            plot_DNN_scatter(dfs_out, target, output_size, "/mnt/$(name)/", String(target))
            plot_DNN_histogram(dfs_out, target, output_size, "/mnt/$(name)/", String(target))

            plot_DNN_lineplot(dfs_out, target, output_size, "/mnt/$(name)/", String(target))

            df_corr = compute_corr(dfs_out, output_size, "/mnt/$(name)/", String(target))
            plot_corr(df_corr, output_size, "/mnt/$(name)/", String(target))
        end
    end
end

run_model()


# FreeType DSL

module DSL

macro ok(block...)
    :(ret = $(esc(block...)); 0 == ret ? nothing : error($(esc(block))[1]))
end

export create_library, create_typeface, create_glyph, outline_decompose
export close_library, close_typeface
export pt, px

using FreeType

function FT_Outline_Decompose(outline::FreeType.FT_Outline,outline_funcs::FreeType.FT_Outline_Funcs,user)
    ccall((:FT_Outline_Decompose,FreeType.freetype),FreeType.FT_Error,(Ptr{FreeType.FT_Outline},Ptr{FreeType.FT_Outline_Funcs},Ptr{Void}),pointer_from_objref(outline),pointer_from_objref(outline_funcs),user)
end

struct Library
    ref::Ref{FT_Library}
end

struct TypeFace
    ref::Ref{FT_Face}
end

abstract type PathOperation end
struct Move <: PathOperation
    to
end

struct Line <: PathOperation
    to
end

struct Conic <: PathOperation
    control
    to
end

struct Cubic <: PathOperation
    control1
    control2
    to
end

function create_library()::Library
    reflibrary = Ref{FreeType.FT_Library}()
    @ok FT_Init_FreeType(reflibrary)
    Library(reflibrary)
end

function close_typeface(face::TypeFace)
    FT_Done_Face(face.ref[])
end

function close_library(library::Library)
    FT_Done_FreeType(library.ref[])
end

function create_typeface(library::Library, pathname::String)::TypeFace
    refface = Ref{FreeType.FT_Face}()
    @ok FT_New_Face(library.ref[], pathname, 0, refface)
    TypeFace(refface)
end

abstract type LengthUnit end
abstract type RelativeUnit <: LengthUnit end
abstract type AbsoluteUnit <: LengthUnit end

struct Point <: AbsoluteUnit
    val::Real
end

struct Pixel <: RelativeUnit
    val::Real
end

struct Glyph
    face::FreeType.FT_FaceRec
    slot::FreeType.FT_GlyphSlotRec
    char::Char
    unit::LengthUnit
end

const pt = Point
const px = Pixel

Base.:*(n::Real, unit::Type{T}) where T <: LengthUnit = T(n)

Base.promote_rule(::Type{Pixel}, ::Type{Point}) = Point
Base.convert(::Type{Pixel}, pt::Point) = Pixel(4/3 * pt.val)
Base.convert(::Type{Point}, px::Pixel) = Pixel(3/4 * px.val)

function create_glyph(typeface::TypeFace, char::Char, unit::T) where T <: LengthUnit
    glyph_index = FT_Get_Char_Index(typeface.ref[], char)
    if isa(unit, RelativeUnit)
        pixel_width = pixel_height = UInt32(unit.val)
        @ok FT_Set_Pixel_Sizes(typeface.ref[], pixel_width, pixel_height)
    elseif isa(unit, AbsoluteUnit)
        char_width = char_height = 64unit.val
        horz_resolution = vert_resolution = 72
        @ok FT_Set_Char_Size(typeface.ref[], char_width, char_height, horz_resolution, vert_resolution)
    end
    @ok FT_Load_Glyph(typeface.ref[], glyph_index, FT_LOAD_DEFAULT)
    face = unsafe_load(typeface.ref[])
    slot = unsafe_load(face.glyph)
    Glyph(face, slot, char, unit)
end

function outline_decompose(glyph::Glyph)
    function pos(p::Ptr{FreeType.FT_Vector})
        vec = unsafe_load(p)
        (vec.x, vec.y)
    end

    move_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{Void})) do to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Move(pos.([to])...))
        Cint(0)
    end

    line_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{Void})) do to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Line(pos.([to])...))
        Cint(0)
    end

    conic_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{FreeType.FT_Vector},Ptr{Void})) do control, to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Conic(pos.([control, to])...))
        Cint(0)
    end

    cubic_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{FreeType.FT_Vector},Ptr{FreeType.FT_Vector},Ptr{Void})) do control1, control2, to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Cubic(pos.([control1, control2, to])...))
        Cint(0)
    end

    global paths = Vector{PathOperation}()
    user_f = cfunction(Void, (PathOperation,)) do op
        push!(paths, op)
        nothing
    end

    outline_funcs = FreeType.FT_Outline_Funcs(move_f, line_f, conic_f, cubic_f, 0, 0)
    @ok FT_Outline_Decompose(glyph.slot.outline, outline_funcs, user_f)
    paths::Vector{PathOperation}
end

end


# BEGIN PARTICLE TREE
"""
    Coords{B}

Leaf structure in galaxy tree consists of position and momentum coordinates of a specific particle
"""
struct Coords{B}
    values::ArrayPartition{B}

    function Coords(values::ArrayPartition{B}) where B <: Number
        @assert length(values.x[1]) == 3 "Position must have length == 3"
        @assert length(values.x[2]) == 3 "Momentum must have length == 3"
        new{B}(values)
    end

end # Coords

"""
    Species{T<:AbstractMultiScaleArray,B<:Number}

Internally consists of a particle ID and an array of particle coordinates
"""
struct Species{T<:AbstractMultiScaleArray,B<:Number} <: AbstractMultiScaleArray{B}
    nodes::Vector{T}
    values::Vector{B}
    end_idxs::Vector{Int}
    id::ParticleID
end # Species

"""
    Snapshot{T<:AbstractMultiScaleArray,B<:Number}

Internally consists of a collection of all species at one particular time
"""
struct Snapshot{T<:AbstractMultiScaleArray,B<:Number} <: AbstractMultiScaleArray{B}
    nodes::Vector{T}
    values::Vector{B}
    end_idxs::Vector{Int}
end

"""
    Timescale{T<:AbstractMultiScaleArray,B<:Number}

Internally a collection of all Snapshots with the same time stepping
"""
struct Timescale{T<:AbstractMultiScaleArray,B<:Number} <: AbstractMultiScaleArray{B}
    nodes::Vector{T}
    values::Vector{B}
    end_idxs::Vector{Int}
    times::Vector{Float64}
    dt::Float64

    function Timescale(nodes::Vector{T},values::Vector{B},end_idxs::Vector{Int},times::Vector{Float64},dt::Float64 = -1.0) where {T,B}
        @assert length(nodes) == length(times) "Number of nodes must equal the number of times"
        new{T,B}(nodes,values,end_idxs,times,dt)
    end
end

"""
    ParticleTree{T<:AbstractMultiScaleArray,B<:Number}

Collection of all timescales
"""
struct ParticleTree{T<:AbstractMultiScaleArray,B<:Number} <: AbstractMultiScaleArrayHead{B}
    nodes::Vector{T}
    values::Vector{B}
    end_idxs::Vector{Int}
end

export Coords, Species, Snapshot, Timescale, ParticleTree
# END PARTICLE TREE

###############################################################################
"""
    ParticleTree()

Make an empty ParticleTree
"""
function ParticleTree()
    # Get ParticleTree leaf type
    coord = Coords(ArrayPartition([1.,2.,3.],[4.,5.,6.]))
    emptycoords = typeof(coord)[]

    # stick Coords in a Species
    # sp = construct(Species,emptycoords,Float64[],ParticleID(0,0,0))
    # emptysp = typeof(sp)[]

    sp = construct(Species,Vector([coord]),Float64[],ParticleID(0,0,0))

    # stick Species in a Snapshot
    # snap = construct(Snapshot,emptysp)
    # emptysnap = typeof(snap)[]
    snap = construct(Snapshot,Vector([sp]))

    # stick Snapshot in a Timescale
    time_init = Vector([0.0])
    dt = -1.0
    scale = construct(Timescale,Vector([snap]),Float64[],time_init,dt)
    emptyscale = typeof(scale)[]
    # stick Timescale in a ParticleTree
    particletree = construct(ParticleTree,Vector([scale]))
    remove_node!(particletree,1)
    return particletree
end

# FUNCTIONS TO HELP GENERATE PARTICLE TREES
"""
    ParticleTree(particletype::ParticleID,energy::Quantity,xyz::Vector{T},aim::Vector,dt::Quantity) where {T<:Quantity}

Return a ParticleTree with a single particle in it. Applies unit conversion prior to storing tree.

Input:
- particletype: ParticleID
- energy: initial particle energies (must have units ∈ Unitful.Energy)
- xyz: starting Cartesian coordinates with Earth along the -x axis (must have units ∈ Unitful.Distance)
- aim: Vector with length=3 with initial direction of particle trajectory.
- dt: initial time step for simulation (must have units ∈ Unitful.Time)

Returns:
        ParticleTree
"""
function ParticleTree(particletype::ParticleID,energy::Quantity,xyz::Vector{T},aim::Vector,dt::Quantity) where {T<:Quantity}
    @assert typeof(energy) <: Unitful.Energy "Need to be units of energy"
    @assert typeof(xyz[1]) <: Unitful.Length "Need to be units of length"
    @assert typeof(dt)     <: Unitful.Time   "Need to be units of time"
    @assert length(xyz) == 3                 "Coordinates require 3 spatial dimensions"
    @assert length(aim) == 3                 "Direction requires 3 spatial dimensions"

    energy_in = usistrip(Float64(energy))
    xyz_in = Float64.(xyz)
    xyz_in = usistrip.(xyz_in)
    aim_in = Float64.(aim)
    dt_in = usistrip(Float64(dt))

    ParticleTree(particletype,energy_in,xyz_in,aim_in,dt_in)
end

"""
    ParticleTree(particletype::ParticleID, energy::Float64, xyz::Vector{Float64}, aim::Vector{Float64}, dt::Float64)

Return a ParticleTree with a single particle in it. Units are assumed to be SI.

Input:
- particletype: ParticleID
- energy: initial particle energies
- xyz: starting Cartesian coordinates with Earth along the -x axis
- aim: Vector with length=3 with initial direction of particle trajectory.
- dt: initial time step for simulation

Returns:
        ParticleTree
"""
function ParticleTree(particletype::ParticleID, energy::Float64, xyz::Vector{Float64}, aim::Vector{Float64}, dt::Float64)

    # Get ParticleTree leaf type
    leafType = typeof(Coords(ArrayPartition([1.,2.,3.],[4.,5.,6.])))

    # Ensure that the initial direction is normalized for subsequent
    # momentum scaling.
    particleAim = normalize(aim)

    # starts at time 0.0
    timeInit = Vector([0.0])

    particleTree = nothing

    # first make Coords
    dir = normalize(aim)
    pmag = momentum_si(energy,particletype)
    mom = pmag*dir
    coords = Coords(ArrayPartition(xyz,mom))

    # stick Coords in a Species
    sp = construct(Species,Vector([coords]),Float64[],particletype)

    # stick Species in a Snapshot
    snap = construct(Snapshot,Vector([sp]))

    # stick Snapshot in a Timescale
    time_init = Vector([0.0])
    scale = construct(Timescale,Vector([snap]),Float64[],time_init,dt)

    # stick Timescale in a ParticleTree
    particletree = construct(ParticleTree,Vector([scale]))

    return particletree

end # ParticleTree

function __push!(pt::ParticleTree,c::Union{Timescale,Snapshot,Species,Coords},I::Vector{T}) where {T<:Int}
    @assert length(I) < 5 "Something went wrong"
    test_type = length(I) == 1 ? typeof(pt.nodes[end]) : typeof(pt[I...])
    if typeof(c) == test_type
        if length(I) == 1
            add_node!(pt,c)
        else
            add_node!(pt,c,I[1:end-1]...)
        end
    else
        if length(I) == 1
            push!(I,length(pt.nodes[end].nodes))
        else
            push!(I,length(pt[I...].nodes))
        end
        __push!(pt,c,I)
    end

end

"""
    push!(pt::ParticleTree,val::Union{ParticleTree,Timescale,Snapshot,Species,Coords})

Push contents of val to last container of same type.
"""
function Base.push!(pt::ParticleTree,val::Union{ParticleTree,Timescale,Snapshot,Species,Coords})
    if isa(val,ParticleTree)
        if length(val.nodes) == 0
            return nothing
        elseif length(pt.nodes) == 0
            append!(pt.nodes,val.nodes)
            append!(pt.values,val.values)
            append!(pt.end_idxs,val.end_idxs)
            return nothing
        else
            for ts in val.nodes
                add_node!(pt,ts)
            end
            return nothing
        end
    else
        I = Vector([length(pt.nodes)])
        __push!(pt,val,I)
    end
    return nothing
end

"""
    push(pt::ParticleTree,val::Union{ParticleTree,Timescale,Snapshot,Species,Coords})

No-side-effect version of push!.
"""
function push(pt::ParticleTree,val::Union{ParticleTree,Timescale,Snapshot,Species,Coords})
    if isa(val,ParticleTree)
        if length(val.nodes) == 0
            return deepcopy(pt)
        elseif length(pt.nodes) == 0
            return deepcopy(val)
        else
            out = deepcopy(pt)
            for ts in val.nodes
                add_node!(out,ts)
            end
            return out
        end
    else
        out = deepcopy(pt)
        I = Vector([length(out.nodes)])
        __push!(pt,val,I)
    end
    return nothing
end
export push

"""
    initParticles(particleTypes, energies, rStart, particleAim, gpFrac, bfield)

Return a ParticleTree consisting of all combinations of particles, energies, radii, and initial direction.
Timestep is calculated as a fraction of the gyroperiod at the particle point of insertion.

Input:
- particleTypes: Vector of ParticleID(s)
- energies: Vector of initial particle energies (must have units ∈ Unitful.Energy)
- rStart: starting radius from galactic center (must have units ∈ Unitful.Distance)
- particleAim: Vector with length=3 with initial direction of particle trajectory.
- gpFrac: Float giving the fraction of the gyroperiod to use as a time step
- bfield: Function that returns the value of the galactic magnetic field when given spatial coorinates

Returns:
        ParticleTree
"""
function initParticles(particleTypes, energies, rStart, particleAim, gpFrac, bfield)

    # Get ParticleTree leaf type
    leafType = typeof(Coords(ArrayPartition([1.,2.,3.],[4.,5.,6.])))

    # Ensure that the initial direction is normalized for subsequent
    # momentum scaling.
    particleAim = map(x->normalize(x),particleAim)

    # starts at time 0.0
    timeInit = Vector([0.0])

    particleTree = nothing
    # loop over all arrays to get combination of all parameters
    for r in rStart
        rVec = Vector([usistrip(r),0.0,0.0])
        bMag = norm(bfield(rVec))
        for en in energies
            for p in particleTypes
                pMag = momentum(en,p)
                vMag = vFromP(pMag,mass(p))
                coordsVec = nothing
                for d in particleAim
                    println("Test 1")
                    pVec = usistrip(pMag)*d
                    println(d)
                    println(pVec)
                    println(rVec)
                    coords = Coords(ArrayPartition(rVec,pVec))
                    println(coords)
                    if coordsVec === nothing
                        coordsVec = Vector([coords])
                    else
                        push!(coordsVec,coords)
                    end
                end
                # each species will have a different gyroradius, so each
                # one also gets its own time step -> Timescale
                sp = construct(Species,coordsVec,Float64[],p)
                spVec = Vector([sp])
                snap = construct(Snapshot,spVec)
                snapVec = Vector([snap])
                gCirc = 2*π*gyroradius(
                                       ustrip(u"GeV/c",pMag),
                                       abs(p.p-p.e),
                                       usistrip(bMag)*1.0e4
                                       )
                gPrd = gCirc/usistrip(vMag)
                dt = gPrd*gpFrac
                scale = construct(Timescale,snapVec,Float64[],timeInit,dt)
                if particleTree == nothing
                    particleTree = construct(ParticleTree,Vector([scale]))
                else
                    add_node!(particleTree,scale)
                end
            end
        end
    end
    return particleTree

end # initParticles
export initParticles

# TODO: generate readable lookup tree if inputs are given as dictionaries

# TODO: generate from a (.json?) file

###############################################################################
# FUNCTIONS TO HELP READ PARTICLE TREES

"""
        getCoordArray(pt::ParticleTree)

Returns all (position,momentum) coordinates in an array of size (6,nCoords)
"""
function getCoordArray(pt::ParticleTree)
    out = Vector(pt)
    out = transpose(reshape(pt,(6,length(pt)÷6)))
    return out
end

export getCoordArray

"""
        getTimescaleIndices(pt::ParticleTree)

Returns internal MultiScaleArray indices denoting the last index of each
Timescale.
"""
function getTimescaleIndices(pt::ParticleTree)
    return pt.end_idxs
end

export getTimescaleIndices

"""
        getSnapshotIndices(pt::ParticleTree)

Returns internal MultiScaleArray indices denoting the last index of each
Snapshot.
"""
function getSnapshotIndices(pt::ParticleTree)
    out = nothing
    for snapshot in LevelIter(2,pt)
        if out == nothing
            out = Vector([snapshot[1].end_idxs])
        else
            push!(out,snapshot[1].end_idxs)
        end
    end
    return out
end

export getSnapshotIndices

"""
        getSpeciesIndices(pt::ParticleTree)

Returns internal MultiScaleArray indices denoting the last index of each
Species.
"""
function getSpeciesIndices(pt::ParticleTree)
    out = nothing
    for species in LevelIter(3,pt)
        if out === nothing
            out = Vector([species[1].end_idxs])
        else
            push!(out,species[1].end_idxs)
        end
    end
    return out
end

export getSpeciesIndices

"""
        getParticleTrajectory(pt::ParticleTree,iTs::Int64,iSs::UnitRange,iSp::Int64,iCd::Int64)

Returns particle trajectory over given range iSs
"""
function getParticleTrajectory(pt::ParticleTree,iTs::Int64,iSs::UnitRange,iSp::Int64,iCd::Int64)
    return pt.nodes[iTs].times, Vector([pt[iTs,i,iSp,iCd].values for i = iSs])
end

"""
        getParticleTrajectory(pt::ParticleTree,iTs::Int64,iSs::UnitRange,iSp::Int64,iCd::Int64)

Returns entire particle trajectory
"""
function getParticleTrajectory(pt::ParticleTree,iTs::Int64,iSp::Int64,iCd::Int64)
    nTime = length(pt.nodes[iTs].nodes)
    getParticleTrajectory(pt,iTs,1:nTime,iSp,iCd)
end

export getParticleTrajectory
###############################################################################
# SAVE PARTICLE TREES TO FILE


module Exp
include("./observable.jl")

experiments = Array{Dict{String,Array{Float64,1}},1}(undef, length(observables))
error_bars = Array{Dict{String,Array{Float64,1}},1}(undef, length(observables))

experiments[observables_index("Nuclear_NFkB")] = Dict(
    "WT" => [
        11.505
        51.52262857
        59.289
        83.68132857
        64.46
        45.027
        34.62007143
        36.81375
        53.31616
        49.223625
        38.17835714
        39.38095714
        44.80163333
        37.23936667
        30.32244
        36.295
        35.46456
        31.04
        34.05333333
        32.175
        29.0548
        30.96566
        35.023
        28.58318571
        22.54135
    ] ./ 83.68132857
)

function get_timepoint(obs_name::String)::Vector{Float64}
    if obs_name == "Nuclear_NFkB"
        [15.0*i for i in 0:24]  # 0, 15, 30, ..., 360 [min.]
    end
end
end # module