################################################################################
# Particle (flat LC)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle")
dir_sim   = joinpath(module_dir(), "src/simulation/particle")
dir_sim_env   = joinpath(dir_sim, "flat_lc")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(particle)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_lc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_lc/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle (flat + nonlinear cone)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle")
dir_sim   = joinpath(module_dir(), "src/simulation/particle")
dir_sim_env   = joinpath(dir_sim, "flat_nc")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(particle)
env = deepcopy(flat_3D_nc)
s = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_nc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_nc/jacobians.jld2")

instantiate_base!(s.model, path_base)
instantiate_dynamics!(s.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(s.model, s.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(s, path_res, path_jac)

################################################################################
# Particle (quadratic)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle")
dir_sim   = joinpath(module_dir(), "src/simulation/particle")
dir_sim_env   = joinpath(dir_sim, "quadratic")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(particle)
env = deepcopy(quadratic_bowl_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "quadratic/residual.jld2")
path_jac = joinpath(dir_sim, "quadratic/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle 2D (flat + LC)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/particle_2D")
dir_sim_env   = joinpath(dir_sim, "flat_lc")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(particle_2D)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_lc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_lc/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle 2D (flat + nonlinear cone)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/particle_2D")
dir_sim_env   = joinpath(dir_sim, "flat_nc")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(particle_2D)
env = deepcopy(flat_2D_nc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_nc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_nc/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle 2D (slope)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/particle_2D")
dir_sim_env   = joinpath(dir_sim, "slope")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(particle_2D)
env = deepcopy(slope1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "slope/residual.jld2")
path_jac = joinpath(dir_sim, "slope/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (2D)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(hopper_2D)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (2D) (sinusoidal)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
dir_sim_env   = joinpath(dir_sim, "sinusoidal")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(hopper_2D)
env = deepcopy(sine2_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper 2D (Piecewise)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
dir_sim_env   = joinpath(dir_sim, "piecewise")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(hopper_2D)
env = deepcopy(piecewise1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "piecewise/residual.jld2")
path_jac = joinpath(dir_sim, "piecewise/jacobians.jld2")

instantiate_base!(sim.model, path_base)
expr_dyn = generate_dynamics_expressions(sim.model, derivs=true)
save_expressions(expr_dyn, path_dyn, overwrite=true)
instantiate_dynamics!(sim.model, path_dyn, derivs = true)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, jacobians = :approx)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac, jacobians = :approx)

################################################################################
# Hopper (2D) (stairs)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
dir_sim_env   = joinpath(dir_sim, "stairs")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(hopper_2D)
env = deepcopy(stairs3_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "stairs/residual.jld2")
path_jac = joinpath(dir_sim, "stairs/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (3D)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_3D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_3D")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(hopper_3D)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (3D) (sinusoidal)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_3D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_3D")
dir_sim_env   = joinpath(dir_sim, "sinusoidal")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(hopper_3D)
env = deepcopy(sine2_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(quadruped)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped Payload
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
dir_sim_env   = joinpath(dir_sim, "payload")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(quadruped_payload)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics_payload/base.jld2")
path_dyn = joinpath(dir_model, "dynamics_payload/dynamics.jld2")
path_res = joinpath(dir_sim, "payload/residual.jld2")
path_jac = joinpath(dir_sim, "payload/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped Sinusoidal
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
dir_sim_env   = joinpath(dir_sim, "sinusoidal")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(quadruped)
env = deepcopy(sine1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, jacobians = :full)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped (Piecewise)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
dir_sim_env   = joinpath(dir_sim, "piecewise")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(quadruped)
env = deepcopy(piecewise1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "piecewise/residual.jld2")
path_jac = joinpath(dir_sim, "piecewise/jacobians.jld2")

instantiate_base!(sim.model, path_base)
expr_dyn = generate_dynamics_expressions(sim.model, derivs = true)
save_expressions(expr_dyn, path_dyn, overwrite=true)
instantiate_dynamics!(sim.model, path_dyn, derivs = true)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, jacobians = :approx)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac, jacobians = :approx)

################################################################################
# Flamingo (flat)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/flamingo")
dir_sim   = joinpath(module_dir(), "src/simulation/flamingo")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(flamingo)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Flamingo (sinusoidal)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/flamingo")
dir_sim   = joinpath(module_dir(), "src/simulation/flamingo")
dir_sim_env   = joinpath(dir_sim, "sinusoidal")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(flamingo)
env = deepcopy(sine3_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Flamingo (smooth slope)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/flamingo")
dir_sim   = joinpath(module_dir(), "src/simulation/flamingo")
dir_sim_env   = joinpath(dir_sim, "slope")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(flamingo)
env = deepcopy(slope_smooth_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "slope/residual.jld2")
path_jac = joinpath(dir_sim, "slope/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Flamingo (piecewise)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/flamingo")
dir_sim   = joinpath(module_dir(), "src/simulation/flamingo")
dir_sim_env   = joinpath(dir_sim, "piecewise")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(flamingo)
env = deepcopy(piecewise1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "piecewise/residual.jld2")
path_jac = joinpath(dir_sim, "piecewise/jacobians.jld2")

instantiate_base!(sim.model, path_base)
expr_dyn = generate_dynamics_expressions(sim.model, derivs = true)
save_expressions(expr_dyn, path_dyn, overwrite=true)
instantiate_dynamics!(sim.model, path_dyn, derivs=true)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, jacobians = :approx)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac, jacobians = :approx)

################################################################################
# PushBot
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/pushbot")
dir_sim   = joinpath(module_dir(), "src/simulation/pushbot")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(pushbot)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# WalledCartpole
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/walledcartpole")
dir_sim   = joinpath(module_dir(), "src/simulation/walledcartpole")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(walledcartpole)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Rigid body
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/rigidbody")
dir_sim   = joinpath(module_dir(), "src/simulation/rigidbody")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(rigidbody)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")
instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, mapping = Gz_func, nv = sim.model.nq - 1)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Centroidal Quadruped
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/centroidal_quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/centroidal_quadruped")
dir_sim_env   = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(centroidal_quadruped)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Centroidal Quadruped undamped
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/centroidal_quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/centroidal_quadruped")
dir_sim_env   = joinpath(dir_sim, "flat_undamped")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(centroidal_quadruped_undamped)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics_undamped/base.jld2")
path_dyn = joinpath(dir_model, "dynamics_undamped/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_undamped/residual.jld2")
path_jac = joinpath(dir_sim, "flat_undamped/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Point Foot Quadruped
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/point_foot_quadruped")
dir_sim = joinpath(module_dir(), "src/simulation/point_foot_quadruped")
dir_sim_env = joinpath(dir_sim, "flat")
!isdir(dir_sim_env) && mkdir(dir_sim_env)
model = deepcopy(point_foot_quadruped)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
isfile(path_jac) && rm(path_jac)
ContactImplicitMPC.JLD2.save(path_jac, "rz_sp", rz_sp, "rθ_sp", rθ_sp)
instantiate_residual!(sim, path_res, path_jac)


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

# when constructing a without generator infos like `SampledDensity(density, samples)`:
struct UnknownSampleGenerator<: AbstractSampleGenerator end
getalgorithm(sg::UnknownSampleGenerator) = nothing

# for samplers without specific infos, e.g. current ImportanceSamplers:
struct GenericSampleGenerator{A <: AbstractSamplingAlgorithm} <: AbstractSampleGenerator
    algorithm::A
end
getalgorithm(sg::GenericSampleGenerator) = sg.algorithm

"""
    struct IIDSampling <: AbstractSamplingAlgorithm

Sample via `Random.rand`. Only supported for posteriors of type
`Distributions.MultivariateDistribution` and `BAT.DistLikeDensity`.

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct IIDSampling <: AbstractSamplingAlgorithm
    nsamples::Int = 10^5
end
export IIDSampling

function bat_sample_impl(rng::AbstractRNG, target::AnyIIDSampleable, algorithm::IIDSampling)
    shaped_density = convert(DistLikeDensity, target)
    density = unshaped(shaped_density)
    shape = varshape(shaped_density)
    n = algorithm.nsamples

    # ToDo: Parallelize, using hierarchical RNG (separate RNG for each sample)
    v = nestedview(rand(rng, sampler(density), n))
    logd = map(logdensityof(density), v)

    weight = fill(_default_int_WT(1), length(eachindex(logd)))
    info = fill(nothing, length(eachindex(logd)))
    aux = fill(nothing, length(eachindex(logd)))

    unshaped_samples = DensitySampleVector((v, logd, weight, info, aux))

    samples = shape.(unshaped_samples)
    (result = samples,)
end

"""
    struct RandResampling <: AbstractSamplingAlgorithm

Resamples from a given set of samples.

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct RandResampling <: AbstractSamplingAlgorithm
    nsamples::Int = 10^5
end
export RandResampling

function bat_sample_impl(rng::AbstractRNG, posterior::DensitySampleVector, algorithm::RandResampling)
    n = algorithm.nsamples
    orig_idxs = eachindex(posterior)
    weights = FrequencyWeights(float(posterior.weight))
    resampled_idxs = sample(orig_idxs, weights, n, replace=true, ordered=false)

    samples = posterior[resampled_idxs]
    samples.weight .= 1

    (result = samples,)
end

"""
    struct OrderedResampling <: AbstractSamplingAlgorithm

Efficiently resamples from a given series of samples, keeping the order of
samples.

Can be used to efficiently convert weighted samples into samples with unity
weights.

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct OrderedResampling <: AbstractSamplingAlgorithm
    nsamples::Int = 10^5
end
export OrderedResampling

function bat_sample_impl(rng::AbstractRNG, samples::DensitySampleVector, algorithm::OrderedResampling)
    @assert axes(samples) == axes(samples.weight)
    W = samples.weight
    idxs = eachindex(samples)

    n = algorithm.nsamples
    resampled_idxs = Vector{Int}()
    sizehint!(resampled_idxs, n)

    p_factor = n / sum(W)

    for i in eachindex(W)
        w_eff_0 = p_factor * W[i]
        w_eff::typeof(w_eff_0) = w_eff_0
        while w_eff > 0
            rand(rng) < w_eff && push!(resampled_idxs, i)
            w_eff = w_eff - 1
        end
    end

    new_samples = samples[resampled_idxs]
    new_samples.weight .= 1

    (result = new_samples,)
end


const NaNValue    = "NaN"
const NaNName     = "NaN65"

const InfValue    = "∞"
const PosInfValue = "+∞"
const NegInfValue = "-∞"
const InfName     = "Inf65"

const TinyName  = "TiNY"
const HugeName  = "HuGE"

const Exactly   = "∣"
const Inexactly = "⍿"

function show(io::IO, x::Float65)

    state = getstate(x)
    fp = convert(Float64, x)

    if isfinite(x)
        if Tiny(Float64) < fp < Huge(Float64)
           s = prepshow(fp, state)
        elseif fp == 0.0
           s = prepshow0(fp, state)
        elseif fp <= Tiny(Float64)
          s = prepshowTiny(fp, state)
        else
          s = prepshowHuge(fp, state)
        end
    elseif isinf(x.fp)  # ±Inf or NaN (restore from clearing the state)
       s = prepshowInf(x.fp,state)
    else
       s = prepshowNaN(x.fp,state)
    end

    print(io, s)
end

@inline prepshowState(state::Bool) = state ? Inexactly : Exactly
@inline function prepnumsign(fp::Float64)
    isneg = signbit(fp)
    if isneg
        "-"
    else
        "+"
    end
end

function prepshow(fp::Float64, state::Bool)
   string(fp, prepshowState(state))
end
function prepshow0(fp::Float64, state::Bool)
   string(fp, prepshowState(state))
end
function prepshowTiny(fp::Float64, state::Bool)
   string( prepnumsign(fp), TinyName )
end
function prepshowHuge(fp::Float64, state::Bool)
   string( prepnumsign(fp), HugeName )
end
function prepshowInf(fp::Float64, state::Bool)
   string( prepnumsign(fp), InfValue )
end
function prepshowNaN(fp::Float64, state::Bool)
   NaNValue
end

# read(IO, x), write(IO, x)

function read(io::IO, ::Type{Float65})
    fp = read(io, T)
    convert(Float65, v)
end
function write(io::IO, x::Float65)
    write(io, x.fp)
end


using Test
using NQCDynamics
using OrdinaryDiffEq: Tsit5
using Random: seed!

@testset "Ehrenfest" begin
    sim = RingPolymerSimulation{Ehrenfest}(Atoms(2000), TullyModelOne(), 10; temperature=1e-3)
    u = DynamicsVariables(sim, fill(10/2000, size(sim)), fill(-5, size(sim)) .+ randn(size(sim)), SingleState(1))
    dt = 0.1

    sol = run_trajectory(u, (0, 2000.0), sim; algorithm=DynamicsMethods.IntegrationAlgorithms.BCBwithTsit5(), saveat=0:10:2000, dt=dt)
    sol1 = run_trajectory(u, (0, 2000.0), sim; algorithm=Tsit5(), saveat=0:10:2000, abstol=1e-6, reltol=1e-6)

    @test sol.u ≈ sol1.u rtol=1e-3
end

@testset "FSSH" begin
    sim = RingPolymerSimulation{FSSH}(Atoms(2000), TullyModelOne(), 10; temperature=1e-3)
    u = DynamicsVariables(sim, fill(20/2000, size(sim)), fill(-5, size(sim)) .+ randn(size(sim)), SingleState(1))
    dt = 0.1

    seed!(1)
    sol = run_trajectory(u, (0, 2000.0), sim; algorithm=DynamicsMethods.IntegrationAlgorithms.BCBwithTsit5(), saveat=0:10:2000, dt=dt)
    seed!(1)
    sol1 = run_trajectory(u, (0, 2000.0), sim; algorithm=Tsit5(), saveat=0:10:2000, dt=dt, adaptive=false)

    @test sol.u ≈ sol1.u rtol=1e-3
end


# functions that are common to the scan examples or are helper functions

"""
    function calc_rms(data, channel, num_channels, num_samples_per_channel)

Calculate RMS value from a block of samples that are interleaved from an mcc172 HAT.
Note that deinterleaving data is extremely fast and allows the standard DSP rms
function to be used.  This would be a better approach but differs from the C and
python implementations.
"""
function calc_rms(data::Vector{T}, channel::Integer, num_channels::Integer, num_samples_per_channel::Integer) where T <: AbstractFloat
    # @show(channel, num_channels, num_samples_per_channel)
    value = 0.0
    index = channel
    #@show(length(data), value, index, channel, num_channels, num_samples_per_channel)
    for i in 1:num_samples_per_channel
        value += (data[index] * data[index]) / num_samples_per_channel
        index += num_channels
    end

    return sqrt(value)
end

"""
    function get_iepe()
Get IEPE enable from the user.
"""
function get_iepe()
    while true
        # Wait for the user to enter a response
        println("IEPE enable [y or n]?  ")
        response = readline()

        # Check for valid response
        if lowercase(response) == "y"
            return true
        elseif lowercase(response) == "n"
            return false
        else
            # Ask again.
            println("Invalid response try again.")
        end
    end
end

"""
	function select_hat_devices(filter_by_id, number_of_devices)

This function performs a query of available DAQ HAT devices and determines
the addresses of the DAQ HAT devices to be used in the example.  If the
number of HAT devices present matches the requested number of devices,
a list of all mcc172 objects is returned in order of address, otherwise the
user is prompted to select addresses from a list of displayed devices.

Args:
filter_by_id (int): If this is :ANY` return all DAQ HATs found.
Otherwise, return only DAQ HATs with ID matching this value.
number_of_devices (int): The number of devices to be selected.

Returns:
list[mcc172]: A list of HatInfo structs for the selected devices
(Note: The struct at index 0 will be used as the master).

Raises:
HatError: Not enough HAT devices are present.

"""
function select_hat_devices(filter_by_id::Symbol, number_of_devices::Integer)

    selected_hats = HatInfo[]

    # Get descriptors for all of the available HAT devices.
    hats = hat_list(filter_by_id)
    number_of_hats = length(hats)

    # Verify at least one HAT device is detected.
    if number_of_hats < number_of_devices
        error("Error: This example requires $number_of_devices MCC 172 HATs - found $number_of_hats")
    elseif number_of_hats == number_of_devices
        # Select all the devices
        for i in 1:number_of_devices
            push!(selected_hats, hats[i])
        end
    else
        # Display available HAT devices for selection.
        for h in 1:length(hats)
            println("Address $(hats[h].address): $(hats[h].product_name)")
        end

        # Select the devices
        for device in 1:number_of_devices
            valid = false
            while !valid
                println("Enter address for HAT device $device")
                address = parse(Int, readline())

                # Verify the selected address exists.
                if any(address .== [hats[h].address for h = 1:length(hats)])
                    valid = true
                else
                    println("Invalid address - try again")
                end

                # Verify the address was not previously selected
                if !isempty(selected_hats) && any(address .== [selected_hats[h].address for h = 1:length(hats)])
                    println("Address already selected - try again")
                    valid = false
                end

                if valid
                    for h in eachindex(hats)
                        if address == hats[h].address
                            push!(selected_hats ,hats[h])
                        end
                    end
                end
            end
        end
    end
    return selected_hats
end

"""
    function trigger_dialog()
Manual or GPIO based trigger for synchronization
"""
function trigger_dialog()
while(true)

    println("GPIO trigger connected ? Yes/no")
    ans = readline()
    @show(ans)
    # @infiltrate
    if ans == "" || lowercase(ans[1]) == 'y'
        trigger(23, duration = 0.05)

        break
    elseif lowercase(ans[1]) == 'n'
        println("\n*NOTE: Connect a trigger source to the TRIG input terminal on HAT 0.")
        try
            println("\nPress 'Enter' to continue")
            readline()
        catch
            error("^C to end program")
        end
        break
    else
        println("Incorrect response try again")
    end
end
end

"""
    function wait_for_trigger(address)

Monitor the status of the specified HAT device in a loop until the
triggered status is true or the running status is false.

Args:
hat (mcc172): The mcc172 HAT device object on which the status will
be monitored.

Returns:
Nothing
"""
function wait_for_trigger(address)
# Read the status only to determine when the trigger occurs.
    is_running = true
    is_triggered = false
    while is_running && !is_triggered
        sleep(0.01)
        result_code, status_code, samples_per_channel = mcc172_a_in_scan_status(address)
        status = mcc172_status_decode(status_code)
        is_running = status.running
        is_triggered = status.triggered
    end
    return nothing
end


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

using ModelReduction
using Base.Test

datadir = joinpath(Pkg.dir("ModelReduction"), "test",
                   first(splitext(basename(@__FILE__))))

# Model was constructed using the following code, which is now commented
# to break dependency to JuliaFEM. Assembled global matrices K and M are
# saved to binary format after first test, and can be loaded using JLD.
#
# using JuliaFEM
# using JuliaFEM.Preprocess
# using JuliaFEM.Postprocess
# using JuliaFEM.Abaqus: create_surface_elements
# using Logging
# Logging.configure(level=DEBUG)
#
# # read mesh
# mesh = abaqus_read_mesh(joinpath(datadir, "testmodel.inp"))
# info("element sets = ", collect(keys(mesh.element_sets)))
# info("surface sets = ", collect(keys(mesh.surface_sets)))
#
# # create two field problems with different material properties
# cylinder = Problem(Elasticity, "CYLINDER", 3)
# cylinder.elements = create_elements(mesh, "CYLINDER")
# update!(cylinder.elements, "youngs modulus", 165.0)
# update!(cylinder.elements, "poissons ratio", 0.275)
# update!(cylinder.elements, "density", 7.10E-9)
# block = Problem(Elasticity, "BLOCK", 3)
# block.elements = create_elements(mesh, "BLOCK")
# update!(block.elements, "youngs modulus", 208.0)
# update!(block.elements, "poissons ratio", 0.30)
# update!(block.elements, "density", 7.80E-9)
#
# # create boundary condition from node set
# fixed = Problem(Dirichlet, "fixed", 3, "displacement")
# fixed_nodes = mesh.node_sets[:FIXED]
# fixed.elements = [Element(Poi1, [nid]) for nid in fixed_nodes]
# update!(fixed.elements, "geometry", mesh.nodes)
# update!(fixed.elements, "displacement 1", 0.0)
# update!(fixed.elements, "displacement 2", 0.0)
# update!(fixed.elements, "displacement 3", 0.0)
#
# """ A helper function to create tie contact. """
# function create_interface(mesh::Mesh, slave::String, master::String)
#     interface = Problem(Mortar, "tie contact", 3, "displacement")
#     interface.properties.dual_basis = false
#     slave_elements = create_surface_elements(mesh, slave)
#     master_elements = create_surface_elements(mesh, master)
#     update!(slave_elements, "master elements", master_elements)
#     interface.elements = [slave_elements; master_elements]
#     return interface
# end
#
# # call helper function to create tie contacts
# tie = create_interface(mesh, "CYLINDER_TO_BLOCK", "BLOCK_TO_CYLINDER")
#
# # solve eigenvalues:
# analysis = Analysis(Modal)
# add_problems!(analysis, [cylinder, block, fixed, tie])
#

# skip assembly for now ...
using JLD
data = load(joinpath(datadir, "data.jld"))

cb = Analysis(CraigBampton)
# cb.properties.r_nodes = r_nodes = collect(mesh.node_sets[:BORDER])
# cb.properties.l_nodes = setdiff(keys(mesh.nodes), r_nodes)
# add_problems!(cb, [cylinder, block, fixed, tie])
cb.properties.K = data["K"]
cb.properties.M = data["M"]
cb.properties.r_nodes = data["r_nodes"]
cb.properties.l_nodes = data["l_nodes"]

run!(cb)

info("Calculate 10 smallest eigenvalues of the reduced system")
t0 = time()
w, X = eigs(cb.properties.K_red, cb.properties.M_red; which=:SM, nev=10)
dt = round(time() - t0, 2)
info("Calculated eigenvalues in $dt seconds")

freq_reduced = sqrt.(real(w))/(2*pi)
freq_original = [1937.04, 1938.56, 3612.25, 4240.52, 4254.46, 6036.87, 6224.13, 8835.43, 8876.81, 11716.6]
info("Eigenvalues of the original system [Hz]: ", freq_original)
info("Eigenvalues of the reduced system [Hz]: ", freq_reduced)

@test isapprox(freq_original, freq_reduced; rtol=0.051)


using Knet
using ArgParse
using JLD
using MAT
using Images

include("lib/vocab.jl")
include("lib/convnet.jl")
include("lib/init.jl")
include("lib/base.jl")
include("lib/model.jl")
include("lib/imgproc.jl")
include("lib/util.jl")

function main(args)
    s = ArgParseSettings()
    s.description = "Caption generation script for the model (raw images)"

    @add_arg_table s begin
        ("--image"; help="input image")
        ("--vocabfile"; help="vocab file")
        ("--modelfile"; help="trained model file")
        ("--cnnfile"; help="convnet file for non-finetuned model")
        ("--beamsize"; arg_type=Int; default=1)
        ("--maxlen"; arg_type=Int; default=30; help="max sentence length")
        ("--nogpu"; action=:store_true)
        ("--lastlayer"; default="relu7")
        ("--imsize"; arg_type=Int; nargs=2; default=[224,224];
         help="new image sizes")
        ("--rgbmean"; arg_type=Float32; nargs=3;
         default=map(Float32, [123.68, 116.779, 103.939]))
    end

    # parse args
    println("Datetime: ", now()); flush(STDOUT)
    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true); println(o); flush(STDOUT)
    o[:loadfile] = o[:modelfile]

    newsize = tuple(o[:imsize]...)
    rgbmean = reshape(o[:rgbmean], (1,1,3))

    # load model
    o[:atype] = !o[:nogpu] ? KnetArray{Float32} : Array{Float32}
    w = load(o[:modelfile], "w")
    w = convert_weight(o[:atype], w)
    vocab = load(o[:vocabfile], "vocab")
    lossval = load(o[:modelfile], "lossval")
    s = initstate(o[:atype], size(w["wsoft"], 1), 1)

    # load convnet
    if o[:cnnfile] != nothing
        vggmat = matread(abspath(o[:cnnfile]))
        wcnn = get_vgg_weights(vggmat; last_layer=o[:lastlayer])
        w["wcnn"] = wcnn
    end

    @printf("Data and model loaded [%s]\n", now()); flush(STDOUT)

    # generate caption
    ti = now()
    @printf("Generation started (loss=%g,date=%s)\n", lossval, ti)
    flush(STDOUT)
    img = load(abspath(o[:image]))
    img = process_image(img, newsize; rgbmean=rgbmean)
    generated = generate(
        w, copy(s), img, vocab; maxlen=o[:maxlen], beamsize=o[:beamsize])
    report_generation(o[:image], generated, o[:beamsize])
    tf = now()
    @printf("\nTime elapsed: %s [%s]\n", tf-ti, tf)
end

function report_generation(filename, generated, beamsize)
    @printf("\nFilename: %s\n", filename)
    @printf("Generated: %s\n", generated)
    @printf("Beamsize: %d\n", beamsize)
    flush(STDOUT)
end

!isinteractive() && !isdefined(Core.Main, :load_only) && main(ARGS)


__precompile__()

module GadJet

    # general utility stuff
    include(joinpath(dirname(@__FILE__), "utility", "gravity_utility.jl"))

    # functions to read snapshots
    include(joinpath(dirname(@__FILE__), "read_snapshot", "gadget_types.jl"))
    #include(joinpath(dirname(@__FILE__), "read_snapshot", "dict_functions.jl"))
    #include(joinpath(dirname(@__FILE__), "read_snapshot", "obj_functions.jl"))
    include(joinpath(dirname(@__FILE__), "read_snapshot", "read_header.jl"))
    include(joinpath(dirname(@__FILE__), "read_snapshot", "snapshot_utilities.jl"))
    include(joinpath(dirname(@__FILE__), "read_snapshot", "read_format_1.jl"))
    include(joinpath(dirname(@__FILE__), "read_snapshot", "read_format_2.jl"))
    include(joinpath(dirname(@__FILE__), "read_snapshot", "read_snapshot.jl"))
    include(joinpath(dirname(@__FILE__), "read_snapshot", "read_subfind.jl"))
    include(joinpath(dirname(@__FILE__), "read_snapshot", "read_particle_in_box.jl"))

    # functions to write snapshots
    include(joinpath(dirname(@__FILE__), "write_snapshot", "write_snap.jl"))

    # unit conversion
    include(joinpath(dirname(@__FILE__), "unit_conversion", "unit_types.jl"))

    # sph to grid mapping internal module
    include(joinpath(dirname(@__FILE__), "sph_to_grid", "kernels.jl"))
    include(joinpath(dirname(@__FILE__), "sph_to_grid", "sph_types.jl"))
    include(joinpath(dirname(@__FILE__), "sph_to_grid", "mapping_functions.jl"))
    include(joinpath(dirname(@__FILE__), "sph_to_grid", "sph_to_grid.jl"))

    # sph to grid mapping with Smac
    include(joinpath(dirname(@__FILE__), "sph_to_grid", "smac1_utility.jl"))
    # sph to grid mapping with P-Smac2
    include(joinpath(dirname(@__FILE__), "sph_to_grid", "smac2_utility.jl"))

    # riemann solvers
    include(joinpath(dirname(@__FILE__), "ideal_solutions", "cr_dsa_models.jl"))
    include(joinpath(dirname(@__FILE__), "ideal_solutions", "cr_sod_shock_main.jl"))
    include(joinpath(dirname(@__FILE__), "ideal_solutions", "setup_riemann_parameters.jl"))
    #include(joinpath(dirname(@__FILE__), "ideal_solutions", "cr_sod_shock.jl"))

    # sedov solution
    include(joinpath(dirname(@__FILE__), "ideal_solutions", "sedov_solution.jl"))

    include(joinpath(dirname(@__FILE__), "spectral_cr_utility", "cr_datatypes.jl"))
    include(joinpath(dirname(@__FILE__), "spectral_cr_utility", "analysis_functions.jl"))
    include(joinpath(dirname(@__FILE__), "spectral_cr_utility", "get_detailled_data.jl"))
    include(joinpath(dirname(@__FILE__), "spectral_cr_utility", "synchrotron_kernel.jl"))

    export Header, Info_Line,       # types
           head_to_dict,
           snap_to_dict,
           head_to_obj,
           print_blocks,
           read_info,
           block_present,
           read_snap,
           read_block_by_name,      # similar to readnew.pro by Klaus Dolag
           read_header,
           read_particles_in_box,
           read_particles_in_volume,

           # subfind read
           read_subfind_header,
           read_subfind,
           find_most_massive_halo,
           filter_subfind,

           # write snapshot functions
           write_header,
           write_block,

           # utility stuff
           calculate_center_of_mass,

           # Kernels
           Cubic,
           Quintic,
           WendlandC4,
           WendlandC6,
           kernel_value_2D,
           kernel_value_3D,

           # internal sph mapping
           mappingParameters,
           sphMapping,
           sphMapping_2D,
           sphMapping_3D,
           glimpse,

           # helper functions and datatypes for Smac
           Smac1ImageInfo,
           read_smac1_binary_image,
           read_smac1_binary_info,
           write_smac1_par,
           # helper function for P-Smac2
           write_smac2_par,

           # unit conversion
           GadgetPhysical,
           GadgetPhysicalUnits,
           @u_str,
           strip_unit,
           # old riemann solver
           #RiemannParameters,   # datatype for riemann parameters
           #RiemannSolution,     # datatype for riemann solution
           #solveHydroShock,      # function that solves a standard sod shock

           # Test for different riemann solvers
           RiemannParameters,    # helper function to set up solution
           solve,                # overloaded function to solve riemann problems
           find_xs_first_guess,  # helper function to find initial guess for shock compression

           get_sedov_solution,   # wrapper function to get sedov data and ideal solution from snapshot

           # datatypes and helper functions for LMB_SPECTRAL_CRs
           CRShockData,          # datatype to analyse single shocked particle
           readSingleCRShockDataFromOutputFile, # as the name says
           CRMomentumDistributionConfig, # config parameters for momentum distribution function
           CRMomentumDistribution,
           getCRMomentumDistributionFromPartID, # function to get distribution function
           calculateCREnergyInCGS,
           calculateCRNumber,
           get_detailled_shock_data,
           get_detailled_Dpp_data,
           get_detailled_radiative_data,
           get_detailled_adiabatic_data,
           synchrotron_kernel,
           calculate_synch_intensity

end


import FFTW: fft, ifft, fft!, ifft!, rfft, irfft
export fft, fft1, fft!, ifft!, rfft, irfft, fft1!, fft2, fft2!, fft3, fft3!, ifft1, ifft2, ifft3, rfft1, rfft2, rfft3
export ifft, irfft2, irfft3, ifft1!, ifft2!, ifft3!
export fft_convolve1, fft_convolve2, fft_convolve3
export set_fft_plan_cache_size

function fft(_in::AFArray{T,1},norm_factor=1.0,odim0::dim_t=0) where {T<:Complex}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),odim0))
    AFArray{T,1}(out[])
end

function fft(_in::AFArray{T,1},norm_factor=1.0,odim0::dim_t=0) where {T<:Real}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),odim0))
    AFArray{Complex{T},1}(out[])
end

function fft1(_in::AFArray{T,N},norm_factor=1.0,odim0::dim_t=0) where {T<:Complex,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),odim0))
    AFArray{T,N}(out[])
end

function fft1(_in::AFArray{T,N},norm_factor=1.0,odim0::dim_t=0) where {T<:Real,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),odim0))
    AFArray{Complex{T},N}(out[])
end

function fft!(_in::AFVector,norm_factor=1.0)
    _error(ccall((:af_fft_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function fft1!(_in::AFArray,norm_factor=1.0)
    _error(ccall((:af_fft_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function fft(_in::AFArray{T,2},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0) where {T<:Complex}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft2,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1))
    AFArray{T,2}(out[])
end

function fft(_in::AFArray{T,2},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0) where {T<:Real}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft2,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1))
    AFArray{Complex{T},2}(out[])
end

function fft2(_in::AFArray{T,N},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0) where {T<:Complex,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft2,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1))
    AFArray{T,N}(out[])
end

function fft2(_in::AFArray{T,N},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0) where {T<:Real,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft2,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1))
    AFArray{Complex{T},N}(out[])
end

function fft!(_in::AFMatrix,norm_factor=1.0)
    _error(ccall((:af_fft2_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function fft2!(_in::AFArray,norm_factor=1.0)
    _error(ccall((:af_fft2_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function fft(_in::AFArray{T,3},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0,odim2::dim_t=0) where {T<:Complex}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft3,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1,odim2))
    AFArray{T,3}(out[])
end

function fft(_in::AFArray{T,3},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0,odim2::dim_t=0) where {T<:Real}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft3,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1,odim2))
    AFArray{Complex{T},3}(out[])
end

function fft3(_in::AFArray{T,N},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0,odim2::dim_t=0) where {T<:Complex,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft3,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1,odim2))
    AFArray{T,N}(out[])
end

function fft3(_in::AFArray{T,N},norm_factor=1.0,odim0::dim_t=0,odim1::dim_t=0,odim2::dim_t=0) where {T<:Real,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft3,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1,odim2))
    AFArray{Complex{T},N}(out[])
end

function fft!(_in::AFVolume,norm_factor=1.0)
    _error(ccall((:af_fft3_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function fft3!(_in::AFArray,norm_factor=1.0)
    _error(ccall((:af_fft3_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function ifft(_in::AFArray{T,1},norm_factor=-1.0,odim0::dim_t=0) where {T}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_ifft,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),odim0))
    AFArray{T,1}(out[])
end

function ifft1(_in::AFArray{T,N},norm_factor=-1.0,odim0::dim_t=0) where {T,N}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_ifft,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),odim0))
    AFArray{T,N}(out[])
end

function ifft!(_in::AFVector,norm_factor=-1.0)
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    _error(ccall((:af_ifft_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function ifft1!(_in::AFArray,norm_factor=-1.0)
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    _error(ccall((:af_ifft_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function ifft(_in::AFArray{T,2},norm_factor=-1.0,odim0::dim_t=0,odim1::dim_t=0) where {T}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_ifft2,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1))
    AFArray{T,2}(out[])
end

function ifft2(_in::AFArray{T,N},norm_factor=-1.0,odim0::dim_t=0,odim1::dim_t=0) where {T,N}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_ifft2,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1))
    AFArray{T,N}(out[])
end

function ifft!(_in::AFMatrix,norm_factor=-1.0)
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    _error(ccall((:af_ifft2_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function ifft2!(_in::AFArray,norm_factor=-1.0)
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    _error(ccall((:af_ifft2_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function ifft(_in::AFArray{T,3},norm_factor=-1.0,odim0::dim_t=0,odim1::dim_t=0,odim2::dim_t=0) where {T}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_ifft3,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1,odim2))
    AFArray{T,3}(out[])
end

function ifft3(_in::AFArray{T,N},norm_factor=-1.0,odim0::dim_t=0,odim1::dim_t=0,odim2::dim_t=0) where {T,N}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_ifft3,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),odim0,odim1,odim2))
    AFArray{T,N}(out[])
end

function ifft!(_in::AFVolume,norm_factor=-1.0)
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    _error(ccall((:af_ifft3_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function ifft3!(_in::AFArray,norm_factor=-1.0)
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    _error(ccall((:af_ifft3_inplace,af_lib),af_err,(af_array,Cdouble),_in.arr,Cdouble(norm_factor)))
    _in
end

function rfft(_in::AFArray{T,1},norm_factor=1.0,pad0::dim_t=0) where {T<:Real}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft_r2c,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),pad0))
    AFArray{Complex{T},1}(out[])
end

function rfft1(_in::AFArray{T,N},norm_factor=1.0,pad0::dim_t=0) where {T<:Real,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft_r2c,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t),out,_in.arr,Cdouble(norm_factor),pad0))
    AFArray{Complex{T},N}(out[])
end

function rfft(_in::AFArray{T,2},norm_factor=1.0,pad0::dim_t=0,pad1::dim_t=0) where {T<:Real}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft2_r2c,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),pad0,pad1))
    AFArray{Complex{T},2}(out[])
end

function rfft2(_in::AFArray{T,N},norm_factor=1.0,pad0::dim_t=0,pad1::dim_t=0) where {T<:Real,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft2_r2c,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),pad0,pad1))
    AFArray{Complex{T},N}(out[])
end

function rfft(_in::AFArray{T,3},norm_factor=1.0,pad0::dim_t=0,pad1::dim_t=0,pad2::dim_t=0) where {T<:Real}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft3_r2c,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),pad0,pad1,pad2))
    AFArray{Complex{T},3}(out[])
end

function rfft3(_in::AFArray{T,N},norm_factor=1.0,pad0::dim_t=0,pad1::dim_t=0,pad2::dim_t=0) where {T<:Real,N}
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft3_r2c,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,dim_t,dim_t,dim_t),out,_in.arr,Cdouble(norm_factor),pad0,pad1,pad2))
    AFArray{Complex{T},N}(out[])
end

function irfft(_in::AFArray{Complex{T},N},norm_factor=-1.0,is_odd::Bool=false) where {T,N}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft_c2r,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,Bool),out,_in.arr,Cdouble(norm_factor),is_odd))
    AFArray{T,N}(out[])
end

function irfft2(_in::AFArray{Complex{T},N},norm_factor=-1.0,is_odd::Bool=false) where {T,N}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft2_c2r,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,Bool),out,_in.arr,Cdouble(norm_factor),is_odd))
    AFArray{T,N}(out[])
end

function irfft3(_in::AFArray{Complex{T},N},norm_factor=-1.0,is_odd::Bool=false) where {T,N}
    if norm_factor < 0
        norm_factor = 1/length(_in)
    end
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft3_c2r,af_lib),af_err,(Ptr{af_array},af_array,Cdouble,Bool),out,_in.arr,Cdouble(norm_factor),is_odd))
    AFArray{T,N}(out[])
end

function fft_convolve1(signal::AFArray,filter::AFArray,mode::af_conv_mode)
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft_convolve1,af_lib),af_err,(Ptr{af_array},af_array,af_array,af_conv_mode),out,signal.arr,filter.arr,mode))
    AFArray!(out[])
end

function fft_convolve2(signal::AFArray,filter::AFArray,mode::af_conv_mode)
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft_convolve2,af_lib),af_err,(Ptr{af_array},af_array,af_array,af_conv_mode),out,signal.arr,filter.arr,mode))
    AFArray!(out[])
end

function fft_convolve3(signal::AFArray,filter::AFArray,mode::af_conv_mode)
    out = RefValue{af_array}(0)
    _error(ccall((:af_fft_convolve3,af_lib),af_err,(Ptr{af_array},af_array,af_array,af_conv_mode),out,signal.arr,filter.arr,mode))
    AFArray!(out[])
end

function set_fft_plan_cache_size(cache_size)
    _error(ccall((:af_set_fft_plan_cache_size,af_lib),af_err,(Csize_t,),Csize_t(cache_size)))
end


VERSION >= v"0.6-" && __precompile__(true)

module Strategems
using Base.Dates
using Temporal
using Indicators

export
    # universe definitions
    Universe, gather!, get_overall_index,
    # parameter sets
    ParameterSet, get_n_runs, get_param_combos, get_run_params, generate_dict,
    # indicators
    Indicator, calculate,
    # signals
    Signal, prep_signal, ↑, ↓, @signal,
    # rules
    Rule, @rule, →,
    # portfolios
    Portfolio,#, update_portfolio!,
    # order
    AbstractOrder, MarketOrder, LimitOrder, StopOrder, liquidate, long, buy, short, sell,
    # strategy results
    Results,
    # summary statistic calculations
    cum_pnl,
    # strategies
    Strategy, generate_trades, generate_trades!, backtest, backtest!, optimize, optimize!

include("universe.jl")
include("paramset.jl")
include("indicator.jl")
include("signal.jl")
include("rule.jl")
include("portfolio.jl")
include("order.jl")
include("results.jl")
include("stats.jl")
include("strategy.jl")

end


# # Intro to `InverseProblem`
#
# This example illustrates the construction of an "ensemble simulation" that can
# evaluate an ensemble of column models given an ensemble of free parameter sets.
# The example then builds an `InverseProblem` from observations, an ensemble simulation,
# and a set of free parameters, and illustrates its basic usage.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add ParameterEstimocean, Oceananigans, Distributions, CairoMakie"
# ```

# First we load few things

using ParameterEstimocean

using Oceananigans
using Oceananigans.Architectures: arch_array
using Oceananigans.Units
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ColumnEnsembleSize
using Oceananigans.TurbulenceClosures: ConvectiveAdjustmentVerticalDiffusivity

using CairoMakie
using Distributions
using JLD2

# We reuse some utilities from a previous example to build observations:

examples_path = joinpath(pathof(ParameterEstimocean), "..", "..", "examples")
include(joinpath(examples_path, "intro_to_observations.jl"))
data_path = generate_synthetic_observations()
observations = SyntheticObservations(data_path, field_names=:b, transformation=ZScore())

# # Building an "ensemble simulation"
#
# Our next task is to construct a parameterized `Oceananigans.Simulation`
# that generates the "foward map" for an ensemble of free parameter sets.
# To generate an ensemble of column model outputs efficiently, we construct one
# 3D `Oceananigans.Simulation` consisting of `Nx` by `Ny` independent column models.
#
# The calibration problem then uses the ensemble
# simulation to find optimal parameters by minimizing the discrepancy between
# the observations and the forward map.

"""
    extract_perfect_parameters(observations, Nensemble)

Extract parameters from a batch of "perfect" observations.
"""
function extract_perfect_parameters(observations, Nensemble)
    Nbatch = length(observations)
    Qᵘ, Qᵇ, N², f = [zeros(Nensemble, Nbatch) for i = 1:4]

    Nz = first(observations).grid.Nz
    Hz = first(observations).grid.Hz
    Lz = first(observations).grid.Lz
    Δt = first(observations).metadata.parameters.Δt

    for (j, obs) in enumerate(observations)
        Qᵘ[:, j] .= obs.metadata.parameters.Qᵘ
        Qᵇ[:, j] .= obs.metadata.parameters.Qᵇ
        N²[:, j] .= obs.metadata.parameters.N²
        f[:, j] .= obs.metadata.coriolis.f
    end

    file = jldopen(first(observations).path)
    closure = file["serialized/closure"]
    close(file)

    return Qᵘ, Qᵇ, N², f, Δt, Lz, Nz, Hz, closure
end

"""
    build_ensemble_simulation(observations, arch=CPU(); Nensemble=1)

Returns an `Oceananigans.Simulation` representing an `Nensemble × 1`
ensemble of column models designed to reproduce `observations`.
"""
function build_ensemble_simulation(observations, arch=CPU(); Nensemble=1)

    observations isa Vector || (observations = [observations]) # Singleton batch
    Nbatch = length(observations)

    Qᵘ, Qᵇ, N², f, Δt, Lz, Nz, Hz, closure = extract_perfect_parameters(observations, Nensemble)

    column_ensemble_size = ColumnEnsembleSize(Nz=Nz, ensemble=(Nensemble, Nbatch), Hz=Hz)
    ensemble_grid = RectilinearGrid(arch, size = column_ensemble_size, topology = (Flat, Flat, Bounded), z = (-Lz, 0))

    coriolis_ensemble = arch_array(arch, [FPlane(f=f[i, j]) for i = 1:Nensemble, j=1:Nbatch])
    closure_ensemble = arch_array(arch, [deepcopy(closure) for i = 1:Nensemble, j=1:Nbatch])

    Qᵘ, Qᵇ, N² = Tuple(arch_array(arch, p) for p in (Qᵘ, Qᵇ, N²))

    u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))
    b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ), bottom = GradientBoundaryCondition(N²))

    tracers = first(observations).metadata.parameters.tracers

    ensemble_model = HydrostaticFreeSurfaceModel(grid = ensemble_grid,
                                                 tracers = tracers,
                                                 buoyancy = BuoyancyTracer(),
                                                 boundary_conditions = (; u=u_bcs, b=b_bcs),
                                                 coriolis = coriolis_ensemble,
                                                 closure = closure_ensemble)

    ensemble_simulation = Simulation(ensemble_model; Δt=Δt, stop_time=first(observations).times[end])

    return ensemble_simulation, closure
end

# The following illustrations uses a simple ensemble simulation with two ensemble members:

ensemble_simulation, closure★ = build_ensemble_simulation(observations; Nensemble=3)

# # Free parameters
#
# We construct some prior distributions for our free parameters. We found that it often helps to
# constrain the prior distributions so that neither very high nor very low values for diffusivities
# can be drawn out of the distribution.

priors = (convective_κz = lognormal(mean=0.3, std=0.05),
          background_κz = lognormal(mean=2.5e-4, std=0.25e-4))

free_parameters = FreeParameters(priors)

# We also take the opportunity to collect a named tuple of the optimal parameters

θ★ = (convective_κz = closure★.convective_κz,
      background_κz = closure★.background_κz)

# ## Visualizing the priors
#
# We visualize our prior distributions by plotting a huge number
# of samples:

using ParameterEstimocean.Parameters: unconstrained_prior, transform_to_constrained

Nsamples = 50000000

sample(prior) = [transform_to_constrained(prior, X) for X in rand(unconstrained_prior(prior), Nsamples)]

convective_κz_samples = sample(priors.convective_κz)
background_κz_samples = sample(priors.background_κz)

fig = Figure()
ax_top = Axis(fig[1, 1], xlabel = "convective κᶻ [m² s⁻¹]", ylabel = "Density")
density!(ax_top, convective_κz_samples)
xlims!(ax_top, 0, 10)

ax_bottom = Axis(fig[2, 1], xlabel = "background κᶻ [m² s⁻¹]", ylabel = "Density")
density!(ax_bottom, background_κz_samples)

save("prior_visualization.svg", fig)
nothing # hide

# ![](prior_visualization.svg)

# # The `InverseProblem`

# We can construct the inverse problem ``y = G(θ) + η``. Here, ``y`` are the `observations` and ``G`` is the
# `ensemble_model`.

calibration = InverseProblem(observations, ensemble_simulation, free_parameters)

# ## Using `InverseProblem` to compute `forward_map`
#
# As a sanity check we apply the `forward_map` on the calibration after we initialize all ensemble
# members with the true parameter values. We then confirm that the output of the `forward_map` matches
# the observations to machine precision.

θ¹ = (convective_κz = 0.8 * θ★.convective_κz,
      background_κz = 9.0 * θ★.background_κz)

θ² = (convective_κz = 2.0 * θ★.convective_κz,
      background_κz = 0.1 * θ★.background_κz)

θ_ensemble = [θ★, θ¹, θ²]

G = forward_map(calibration, θ_ensemble)
y = observation_map(calibration)

# The `forward_map` output `G` is a two-dimensional matrix whose first dimension
# is the size of the state space and whose second dimension is the `ensemble_size`.
# Here, we ensure that first ensemble member of the mapped output, which was run with
# the "true" parameters, is identical to the mapped observations:

G[:, 1] ≈ y

# ## Visualizing forward model output
#
# Next we visualize the discrepancy between solutions generated by true
# and non-optimal parameter sets `θ¹` and `θ²`. Time-series data from
# the ensemble run is collected by `calibration.time_series_collector`:

time_series_collector = calibration.time_series_collector
times = time_series_collector.times

## Extract last save point and plot each solution component
Nt = length(times)

b = time_series_collector.field_time_serieses.b[Nt]
t = times[Nt]
z = znodes(b)

## The ensemble varies along the first, or `x`-dimension:
b★ = interior(b)[1, 1, :]
b¹ = interior(b)[2, 1, :]
b² = interior(b)[3, 1, :]

fig = Figure()
ax = Axis(fig[1, 1],
          xlabel = "Buoyancy [m s⁻²]",
          ylabel = "Depth [m]")

b★_label = "true b at t = " * prettytime(t)
b¹_label = "b with $θ¹"
b²_label = "b with $θ²"

lines!(ax, b★, z; label=b★_label, linewidth=2)
lines!(ax, b¹, z; label=b¹_label, linewidth=2)
lines!(ax, b², z; label=b²_label, linewidth=2)

axislegend(ax, position=:lt)

save("ensemble_simulation_demonstration.svg", fig); nothing # hide

# ![](ensemble_simulation_demonstration.svg)