text:
module ParallelAnalysis

using MultivariateStats, StatsFuns, Statistics, ProgressMeter
import Statistics: cov
import MultivariateStats: loadings
import Base: show
using StatsBase: sample
# Polychoric
using Trapz: trapz#, @trapz
using Optim: optimize, Brent
using LinearAlgebra: LowerTriangular, diagind, Symmetric
using QuadGK: quadgk
using Distributions, DataFrames
# Plot recipes
using RecipesBase

include("Polychoric.jl")
export polycor, cov, loadings, communalities

include("RandomMatrix.jl")
export random_matrix, random_sample

include("FA.jl")
export fa

include("IRT.jl")
export generate_response, v2vv, heuristicIRT

include("PA.jl")
export parallel

export show

end


text:
include("../src/PBCM.jl")

cell_general = cell()
cell_design_op = cylindrical_cell_designer(cell_general)
cost = cost_default()
cell_general, cost = convert_all(cell_general, cost, mult)

using DiffEqSensitivity
using PyPlot

per = 10

poro = 0.3
poro_intrvl = [poro*(1 - (per/100)) , poro*(1 + (per/100))]

pos_thic = 150
pos_thic_intrvl = [pos_thic*(1 - (per/100)) , pos_thic*(1 + (per/100))]

no_cells_yr = 10
no_cells_yr_intrvl = [no_cells_yr*(1 - (per/100)) , no_cells_yr*(1 + (per/100))]

cost_pos_AM = 25
cost_pos_AM_intrvl = [cost_pos_AM*(1 - (per/100)) , cost_pos_AM*(1 + (per/100))]

function gsa_cost(arr)

    cell_general.cathode.por        = arr[1]
    cell_general.cathode.th         = converter(arr[2], mult.pos_th)
    cost.general_costs.no_units_mfg = converter(arr[3], mult.units_mfg)
    cost.cell_costs.cathode.AM[1]   = arr[4]

    dollars_kWh = cost_calc(cell_general, cost, system="Cell", cost_verbosity=0)[1]

    return dollars_kWh
end

# gsa_cost([0.25, 150.0, 130.0, 25.0, 19.0])

gsa_result = gsa(gsa_cost, Sobol(order = [2]), [poro_intrvl, pos_thic_intrvl, no_cells_yr_intrvl, cost_pos_AM_intrvl], N = 10000)

x = [0, 1, 2, 3, 1, 2, 3, 2, 3, 3]
y = [0, 0, 0, 0, 1, 1, 1, 2, 2, 3]

clf()
print("\n")
S1 = gsa_result.S1
S2 = gsa_result.S2
print(S1)
print(S2)

z = [S1[1], S2[1], S2[2], S2[3], S1[2], S2[4], S2[5], S1[3], S2[6], S1[4]]

r = bar3D(x, y, z*0, 0.4, 0.4, z,)
xticks([0, 1, 2, 3], ["Poro", "Pos Thic", "No Cells", "Cost Pos AM"])
yticks([1, 2, 3, 4], ["Poro", "Pos Thic", "No Cells", "Cost Pos AM"])
title(string(" GSA: Sobol's Method : 2nd Order : ", per, "%"))
zlabel("Sobol Indices")
figure(2)

ext: jl

max_stars_repo_name: BattModels/BatteryCost

max_issues_repo_name: battmodels/BatteryCost

max_forks_repo_name: battmodels/BatteryCost


text:
Cassette.@context HandEvalCtx

@enum SteppingMode StepIn StepNext StepContinue StepOut

const GLOBAL_BREAKPOINT_RULES = BreakpointRules()
const GLOBAL_STEPPING_MODE = Ref(StepContinue)

"""
    new_debug_ctx()::HandEvalCtx
creates a debug context, with the pass set
and extranious hooks disabled.
"""
function new_debug_ctx()
    ctx = HandEvalCtx(;pass=handeval_pass)
    return Cassette.disablehooks(ctx)
end

@inline function Cassette.overdub(ctx::HandEvalCtx, f, args...)
    # This is basically the epicenter of all the logic
    # We control the flow of stepping modes
    # and which methods are instrumented or not.
    cur_mode = GLOBAL_STEPPING_MODE[]

    should_recurse =
        cur_mode === StepIn ||
        should_instrument(GLOBAL_BREAKPOINT_RULES, f)

    if should_recurse
        if Cassette.canrecurse(ctx, f, args...)
            # Both StepOut and StepContinue means child should StepContinue
            GLOBAL_STEPPING_MODE[] = cur_mode === StepIn ? StepNext : StepContinue
            # Determine stepping mode for child
            try
                return Cassette.recurse(ctx, f, args...)
            finally
                # Determine stepping mode for parent
                child_instruction = GLOBAL_STEPPING_MODE[]
                GLOBAL_STEPPING_MODE[] =
                    child_instruction !== StepContinue ? StepNext :
                        cur_mode === StepIn ? StepContinue : cur_mode

                # if child said StepOut or StepNext or StepIn, then we shold break on next (StepNext)
                # if the child said StepContinue,
                    # if we were saying to StepIn can now StepContinue
                        # as we have completed what ever work we were doing
                    # if we were saying to StepContinue, then still want to continue
                        # (unless we hit a breakpoint where the child gave new instructions)
                    # But if we were StepOut then we still need to stepout til we return
                        # (and it gets turnd into a StepNext)
                    # and if we were StepNext then we made our child StepContinue,
                    # but we want to go StepNext ourself so have to restore that
            end
        else
            #@warn "Not able to enter into method" f args
            return Cassette.fallback(ctx, f, args...)
        end
    else  # !should_recurse
        return f(args...)
    end
end


text:
# # [Sandia 34-Meter Vertical Axis Wind Turbine](@id vawt)
#
# In this example, we examine the stability characteristics of the Sandia 34-Meter Vertical
# Axis Wind Turbine (VAWT).  Geometry for this VAWT is described in
# [SAND-91-2228](https://energy.sandia.gov/wp-content/gallery/uploads/SAND-91-2228.pdf)
# and shown in the figure.  Sectional properties for this VAWT are derived from properties
# listed in [SAND-88-1807](https://energy.sandia.gov/wp-content/gallery/uploads/SAND-88-1807.pdf)
#
# The original authors of this example requested that the following citation accompany it.
#
# Moore, K. and Ennis, B., “Aeroelastic Validation of the Offshore Wind Energy Simulator
# for Vertical-Axis Wind Turbines”, forthcoming 2022
#
# ![](../assets/vawt-drawing.png)
#
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`vertical-axis-wind-turbine.ipynb`](@__NBVIEWER_ROOT_URL__/examples/vertical-axis-wind-turbine.ipynb).
#-

using GXBeam, LinearAlgebra, FLOWMath

## --- Tower Definition --- #

tower_height = 41.9
tower_stiffness = Diagonal([8.4404e9, 2.7053e9, 2.7053e9, 7.0428e9, 9.4072e9, 9.4072e9])
tower_mass = Diagonal([330.755, 330.755, 330.755, 737.282, 368.641, 368.641])

## --- Blade Definition --- #

## geometry
blade_xyz = [
     0.0        0.0   0.0;
     2.26837    0.0   1.257;
     3.63183    0.0   2.095;
     6.76113    0.0   4.19;
     9.55882    0.0   6.285;
    11.8976     0.0   8.38;
    13.7306     0.0  10.475;
    15.1409     0.0  12.57;
    16.1228     0.0  14.665;
    16.7334     0.0  16.76;
    17.0133     0.0  18.855;
    17.0987     0.0  20.95;
    16.9615     0.0  23.045;
    16.5139     0.0  25.14;
    15.7435     0.0  27.235;
    14.5458     0.0  29.33;
    12.9287     0.0  31.425;
    10.8901     0.0  33.52;
     8.547      0.0  35.615;
     5.93739    0.0  37.71;
     3.05842    0.0  39.805;
     1.87486    0.0  40.643;
     0.0        0.0  41.9
]

## section boundaries (z-coordinate)
blade_transition = [0.0, 5.8, 11.1, 29.0, 34.7, 41.9]

## root section properties

blade_stiffness1 = [
    3.74563e9  0.0        0.0        0.0        0.0        1.14382e8;
    0.0        1.20052e9  0.0        0.0        0.0        0.0;
    0.0        0.0        1.20052e9  0.0        0.0        0.0;
    0.0        0.0        0.0        1.87992e7  0.0        0.0;
    0.0        0.0        0.0        0.0        2.24336e7  0.0;
    1.14382e8  0.0        0.0        0.0        0.0        4.09242e8;
]

blade_mass1 = [
    146.781     0.0       0.0      0.0     28.7783     0.0;
      0.0     146.781     0.0    -28.7783   0.0        0.0;
      0.0       0.0     146.781    0.0      0.0        0.0;
      0.0     -28.7783    0.0     16.7793   0.0        0.0;
     28.7783    0.0       0.0      0.0      0.879112  -0.0;
      0.0       0.0       0.0      0.0     -0.0       15.9002;
]

## transition section properties

blade_stiffness2 = [
    2.22783e9  0.0        0.0        0.0        0.0        4.20422e7;
    0.0        7.14048e8  0.0        0.0        0.0        0.0;
    0.0        0.0        7.14048e8  0.0        0.0        0.0;
    0.0        0.0        0.0        6.55493e6  0.0        0.0;
    0.0        0.0        0.0        0.0        7.35548e6  0.0;
    4.20422e7  0.0        0.0        0.0        0.0        1.84227e8;
]

blade_mass2 = [
    87.3025    0.0      0.0       0.0      16.7316     0.0;
     0.0      87.3025   0.0     -16.7316    0.0        0.0;
     0.0       0.0     87.3025   -0.0       0.0        0.0;
     0.0     -16.7316  -0.0       7.47649   0.0        0.0;
    16.7316    0.0      0.0       0.0       0.288241  -0.0;
     0.0       0.0      0.0       0.0      -0.0        7.18825;
]

## center section properties

blade_stiffness3 = [
    1.76888e9  0.0        0.0        0.0        0.0        2.34071e7;
    0.0        5.66947e8  0.0        0.0        0.0        0.0;
    0.0        0.0        5.66947e8  0.0        0.0        0.0;
    0.0        0.0        0.0        4.00804e6  0.0        0.0;
    0.0        0.0        0.0        0.0        4.34302e6  0.0;
    2.34071e7  0.0        0.0        0.0        0.0        1.09341e8;
]

blade_mass3 = [
    69.3173    0.0      0.0       0.0      11.5831     0.0;
     0.0      69.3173   0.0     -11.5831    0.0        0.0;
     0.0       0.0     69.3173   -0.0       0.0        0.0;
     0.0     -11.5831  -0.0       4.44282   0.0        0.0;
    11.5831    0.0      0.0       0.0       0.170191  -0.0;
     0.0       0.0      0.0       0.0      -0.0        4.27263;
]

## --- Strut Definition --- #

strut_locations = [1.257, 40.643]
strut_stiffness = blade_stiffness1
strut_mass = blade_mass1

## --- Define Assembly --- #

## Tower

## number of tower sections
nt = 22

## tower points
x = zeros(nt+1)
y = zeros(nt+1)
z = vcat(0, range(strut_locations[1], strut_locations[2]; length=nt-1), tower_height)
pt_t = [[x[i],y[i],z[i]] for i = 1:nt+1]

## tower frame of reference
frame_t = fill([0 0 1; 0 1 0; -1 0 0], nt)

## tower stiffness
stiff_t = fill(tower_stiffness, nt)

## tower mass
mass_t = fill(tower_mass, nt)

## Blades

## number of blade sections
nbr = 4 # root
nbt = 3 # transition
nbc = 8 # center
nb = 2*nbr + 2*nbt + nbc # total number of blade sections

## interpolation parameter coordinates
new_z = vcat(0.0,
    range(strut_locations[1], 5.8, length=nbr)[1:end-1],
    range(5.8, 11.1, length=nbt+1)[1:end-1],
    range(11.1, 29.0, length=nbc+1)[1:end-1],
    range(29.0, 34.7, length=nbt+1)[1:end-1],
    range(34.7, strut_locations[2], length=nbr),
    tower_height)

## blade points
x = FLOWMath.akima(blade_xyz[:,3], blade_xyz[:,1], new_z)
y = zero(new_z)
z = new_z
pt_bl = [[-x[i],y[i],z[i]] for i = 1:nb+1] # left blade
pt_br = [[x[i],y[i],z[i]] for i = 1:nb+1] # right blade

## left blade frame of reference
frame_bl = Vector{Matrix{Float64}}(undef, nb)
for i = 1:nb
    r = pt_bl[i+1] - pt_bl[i]
    n = norm(r)
    s = r[3]/n
    c = r[1]/n
    frame_bl[i] = [c 0 -s; 0 1 0; s 0 c]
end

## right blade frame of reference
frame_br = Vector{Matrix{Float64}}(undef, nb)
for i = 1:nb
    r = pt_br[i+1] - pt_br[i]
    n = norm(r)
    s = r[3]/n
    c = r[1]/n
    frame_br[i] = [c 0 -s; 0 1 0; s 0 c]
end

## blade stiffness
stiff_b = vcat(
    fill(blade_stiffness1, nbr),
    fill(blade_stiffness2, nbt),
    fill(blade_stiffness3, nbc),
    fill(blade_stiffness2, nbt),
    fill(blade_stiffness1, nbr)
)

## blade mass
mass_b = vcat(
    fill(blade_mass1, nbr),
    fill(blade_mass2, nbt),
    fill(blade_mass3, nbc),
    fill(blade_mass2, nbt),
    fill(blade_mass1, nbr)
)

## Struts

## number of strut sections per strut
ns = 3

## lower left strut points
x = range(0.0, pt_bl[2][1]; length=ns+1)
y = zeros(ns+1)
z = fill(strut_locations[1], ns+1)
pt_s1 = [[x[i],y[i],z[i]] for i = 1:ns+1]

## lower right strut points
x = range(0.0, pt_br[2][1]; length=ns+1)
y = zeros(ns+1)
z = fill(strut_locations[1], ns+1)
pt_s2 = [[x[i],y[i],z[i]] for i = 1:ns+1]

## upper left strut points
x = range(0.0, pt_bl[end-1][1]; length=ns+1)
y = zeros(ns+1)
z = fill(strut_locations[2], ns+1)
pt_s3 = [[x[i],y[i],z[i]] for i = 1:ns+1]

## upper right strut points
x = range(0.0, pt_br[end-1][1]; length=ns+1)
y = zeros(ns+1)
z = fill(strut_locations[2], ns+1)
pt_s4 = [[x[i],y[i],z[i]] for i = 1:ns+1]

## strut frame of reference
frame_s = fill([1 0 0; 0 1 0; 0 0 1], ns)

## strut stiffness
stiff_s = fill(strut_stiffness, ns)

## strut mass
mass_s = fill(strut_mass, ns)

## Combine Tower, Blades, and Struts

## combine points
points = vcat(pt_t, pt_bl, pt_br, pt_s1, pt_s2, pt_s3, pt_s4)

## define element connectivity
istart = cumsum([1, nt+1, nb+1, nb+1, ns+1, ns+1, ns+1])
istop = cumsum([nt+1, nb+1, nb+1, ns+1, ns+1, ns+1, ns+1])
start = vcat([istart[i]:istop[i]-1 for i = 1:length(istart)]...)
stop = vcat([istart[i]+1:istop[i] for i = 1:length(istart)]...)

## use zero-length elements as joints

nj = 12 # number of joints

joints = [
    istart[1]     istart[2]; # tower - bottom of left blade
    istart[1]     istart[3]; # tower - bottom of right blade
    istart[1]+1   istart[4]; # tower - lower left strut
    istart[1]+1   istart[5]; # tower - lower right strut
    istop[1]-1    istart[6]; # tower - upper left strut
    istop[1]-1    istart[7]; # tower - upper right strut
    istop[2]      istop[1]; # top of left blade - tower
    istop[3]      istop[1]; # top of right blade - tower
    istop[4]      istart[2]+1; # lower left strut - left blade
    istop[5]      istart[3]+1; # lower right strut - right blade
    istop[6]      istop[2]-1; # upper left strut - left blade
    istop[7]      istop[3]-1; # upper right strut - right blade
]

frame_j = fill([1 0 0; 0 1 0; 0 0 1], nj)

stiff_j = fill(zeros(6,6), nj) # will be modeled as infinitely stiff

mass_j = fill(zeros(6,6), nj)

## add joint connectivity
start = vcat(start, joints[:,1])
stop = vcat(stop, joints[:,2])

## combine frames
frames = vcat(frame_t, frame_bl, frame_br, frame_s, frame_s, frame_s, frame_s, frame_j)

## combine stiffness
stiffness = vcat(stiff_t, stiff_b, stiff_b, stiff_s, stiff_s, stiff_s, stiff_s, stiff_j)

## combine mass
mass = vcat(mass_t, mass_b, mass_b, mass_s, mass_s, mass_s, mass_s, mass_j)

## create assembly
assembly = Assembly(points, start, stop;
    frames=frames,
    stiffness=stiffness,
    mass=mass)

## --- Define Prescribed Conditions --- #

## create dictionary of prescribed conditions
prescribed_conditions = Dict(
    ## fixed base
    1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0, theta_z=0),
    ## fixed top, but free to rotate around z-axis
    istop[1] => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0),
)

## --- Perform Analysis --- #

## revolutions per minute
rpm = 0:1:40

## gravity vector
gravity = [0, 0, -9.81]

## number of modes
nmode = 10

## number of eigenvalues
nev = 2*nmode

## initialize system storage
system = System(assembly, false; prescribed_points=keys(prescribed_conditions))

## storage for results
freq = zeros(length(rpm), nmode)

## perform an analysis for each rotation rate
for (i,rpm) in enumerate(rpm)

    global system, Up

    ## set turbine rotation
    angular_velocity = [0, 0, rpm*(2*pi)/60]

    ## eigenvalues and (right) eigenvectors
    system, λ, V, converged = eigenvalue_analysis!(system, assembly;
        prescribed_conditions = prescribed_conditions,
        angular_velocity = angular_velocity,
        gravity = gravity,
        nev = nev
        )

    ## check convergence
    @assert converged

    if i > 1
        ## construct correlation matrix
        C = Up*system.M*V

        ## correlate eigenmodes
        perm, corruption = correlate_eigenmodes(C)

        ## re-arrange eigenvalues
        λ = λ[perm]

        ## update left eigenvector matrix
        Up = left_eigenvectors(system, λ, V)
        Up = Up[perm,:]
    else
        ## update left eigenvector matrix
        Up = left_eigenvectors(system, λ, V)
    end

    ## save frequencies
    freq[i,:] = [imag(λ[k])/(2*pi) for k = 1:2:nev]

end

# We can compare the computed mode frequencies with experimental data taken from
# [SAND-91-2228](https://energy.sandia.gov/wp-content/gallery/uploads/SAND-91-2228.pdf).

using Plots
#md using Suppressor #hide
pyplot()

#md @suppress_err begin #hide

## Experimental Data
SNL34_flap = [
    -0.11236  3.4778;
    20.1124   4.08516;
    20.1124   3.92961;
    20.1124   3.72961;
    20.2247   2.47403;
    24.1573   2.4733;
    28.2022   2.57256;
    30.2247   2.5944;
    31.9101   2.62742;
    33.2584   2.63829;
    34.1573   2.70479;
    34.2697   2.57143;
    34.2697   2.40476;
    28.2022   2.46144;
    20.1124   2.30739;
    10.0      2.23148;
    10.0      2.17593;
    10.0      2.0537;
    39.6629   1.38154;
    37.3034   1.31531;
    36.0674   1.25999;
    36.1798   1.31552;
    34.1573   1.24923;
    32.0225   1.23851;
    30.6742   1.20543;
    28.2022   1.18367;
    20.2247   1.08514;
    20.2247   1.18514;
    28.2022   1.27256;
    32.0225   1.34963;
    34.0449   1.38258;
    15.1685   1.13052;
    15.1685   1.04164;
    10.0      1.03148;
     0.0      1.06667;
     0.0      2.13333;
]

SNL34_lead = [
     0.0      3.57778;
    -0.11236  2.51113;
    -0.11236  1.80002;
    10.0      1.78704;
    15.1685   1.75275;
    20.1124   1.67405;
    30.2247   1.63885;
    32.0225   1.50518;
    35.0562   1.54906;
    36.0674   1.55999;
    39.5506   1.52601;
    39.3258   2.04827;
    36.9663   2.0376;
    35.8427   2.04892;
    34.9438   2.0602;
    34.0449   2.08258;
    32.0225   2.09407;
    30.2247   2.10551;
    24.4944   2.09546;
    30.1124   2.42776;
    24.1573   2.29553;
    20.1124   2.22961;
    20.0      2.37407;
    10.0      2.48704;
     9.88764  3.55372;
    14.9438   3.54168;
    10.0      3.86481;
    20.1124   3.59628;
    24.1573   3.58442;
    24.2697   3.85106;
    28.0899   3.63924;
    30.0      3.66111;
    32.0225   3.67185;
    33.0337   3.69388;
    34.1573   3.67145;
    35.1685   3.68238;
    35.1685   3.68238;
    36.1798   3.6933;
    37.0787   3.71536;
    34.9438   4.0602;
    35.9551   4.16001;
]

SNL34_1F = [
     0.11236  1.06665;
     7.19101  1.05422;
    14.1573   1.05293;
    20.2247   1.10737;
    28.0899   1.18369;
    34.1573   1.26034;
    37.0787   1.31536;
    39.7753   1.39263;
]

SNL34_1BE = [
     0.0      1.81111;
     6.85393  1.79873;
    14.9438   1.75279;
    20.2247   1.67403;
    30.0      1.63889;
    35.1685   1.57127;
]

SNL34_2FA = [
    -0.11236  2.13335;
     7.64045  2.13192;
    10.0      2.10926;
    19.5506   2.15194;
    26.7416   2.38394;
    30.5618   2.41656;
    34.2697   2.59365;
]

SNL34_2FS = [
    -0.11236  2.13335;
    10.1124   2.23146;
    15.3933   2.27493;
    26.2921   2.52846;
    34.1573   2.70479;
]

SNL34_1TO = [
     0.0      2.5;
     4.04494  2.49925;
     9.77528  2.49819;
    16.6292   2.43025;
    20.2247   2.39625;
    24.382    2.11771;
    30.6742   2.09432;
    35.0562   2.0824;
]

SNL34_3F = [
    -0.11236  3.58891;
     9.88764  3.55372;
    15.1685   3.57497;
    20.2247   3.59625;
    25.5056   3.62861;
    30.7865   3.66097;
    34.0449   3.6937;
    37.191    3.73756;
]

## Initialize Plot
plot(
    xlabel="Rotor Speed (RPM)",
    ylabel="Frequency (Hz)",
    xlim = (0, 40.0),
    legend = :topleft,
    grid=true)

## Add computational results
for i=1:1:nmode
    plot!(rpm, freq[:,i], color=:red, linestyle=:solid, label="")
end

## Add experimental results
scatter!(SNL34_flap[:,1],SNL34_flap[:,2], color=:black, markerstyle=:dot, label="Flatwise")
scatter!(SNL34_lead[:,1],SNL34_lead[:,2], color=:black, markershape=:x, label="Lead-Lag")
plot!(SNL34_1F[:,1], SNL34_1F[:,2], color=:black, linestyle=:dash, label="")
plot!(SNL34_1BE[:,1], SNL34_1BE[:,2], color=:black, linestyle=:solid, label="")
plot!(SNL34_2FA[:,1], SNL34_2FA[:,2], color=:black, linestyle=:solid, label="")
plot!(SNL34_2FS[:,1], SNL34_2FS[:,2], color=:black, linestyle=:solid, label="")
plot!(SNL34_1TO[:,1], SNL34_1TO[:,2], color=:black, linestyle=:dash, label="Tower Mode")
plot!(SNL34_3F[:,1], SNL34_3F[:,2], color=:black, linestyle=:solid, label="")

## Add per-revolution lines
for i = 1:6
    lx = [rpm[1], rpm[end]+10]
    ly = [rpm[1], rpm[end]+10].*i./60.0
    plot!(lx, ly, color=:black, linestyle=:dash, linewidth=0.5, label="")
    annotate!(0.95*lx[2], ly[2]+.05+(i-1)*.01, text("$i P", 10))
end

## Add legend entries
plot!([0], [0], color=:black, label="Experimental")
plot!([0], [0], color=:red, label="GXBeam")

#!nb plot!(show=true)
#md savefig("../assets/vawt-stability.svg") #hide
#md closeall() #hide
#md end #hide
#md nothing #hide

#md # ![](../assets/vawt-stability.svg)

#-

# As can be seen, there is good agreement between the computational and experimental
# results.

state = AssemblyState(system, assembly; prescribed_conditions=prescribed_conditions)

write_vtk("vawt-simulation", assembly, state)

# ![](../assets/vawt-simulation.gif)


text:
module Issue614

using Gridap
using Gridap.CellData

domain = (0,2,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)
const k = 2
reffeᵤ = ReferenceFE(lagrangian,VectorValue{2,Float64},k)
reffe_p = ReferenceFE(lagrangian,Float64,k-1)
V₀ = TestFESpace(model, reffeᵤ, conformity=:H1)
U = TrialFESpace(V₀)
Q = TestFESpace(model, reffe_p, conformity=:C0)
P = TrialFESpace(Q)
Y = MultiFieldFESpace([V₀,Q])
X = MultiFieldFESpace([U,P])

Ω = Triangulation(model)

degree = 1
xh = interpolate_everywhere([VectorValue(0.0,1.0),x->x[1]],X)

Γc = BoundaryTriangulation(model)
dΓc = Measure(Γc, degree)
nΓc = get_normal_vector(Γc)

for i in 0:1000
    uh, ph = xh
    α = ph*nΓc
    CD, CL = ∑( ∫( α )dΓc )
end

end # module


text:
# Different ways to simulate molecules

export
    simulate!,
    VelocityVerlet,
    VelocityFreeVerlet

"""
    VelocityVerlet()

The velocity Verlet integrator.
"""
struct VelocityVerlet <: Simulator end

"""
    simulate!(simulation; parallel=true)
    simulate!(simulation, n_steps; parallel=true)
    simulate!(simulation, simulator, n_steps; parallel=true)

Run a simulation according to the rules of the given simulator.
Custom simulators should implement this function.
"""
function simulate!(s::Simulation{false},
                    ::VelocityVerlet,
                    n_steps::Integer;
                    parallel::Bool=true)
    # See https://www.saylor.org/site/wp-content/uploads/2011/06/MA221-6.1.pdf for
    #   integration algorithm - used shorter second version
    n_atoms = length(s.coords)
    find_neighbours!(s, s.neighbour_finder, 0, parallel=parallel)
    accels_t = accelerations(s, parallel=parallel)
    accels_t_dt = zero(s.coords)

    @showprogress for step_n in 1:n_steps
        for logger in values(s.loggers)
            log_property!(logger, s, step_n)
        end

        # Update coordinates
        for i in 1:length(s.coords)
            s.coords[i] += s.velocities[i] * s.timestep + accels_t[i] * (s.timestep ^ 2) / 2
            s.coords[i] = adjust_bounds.(s.coords[i], s.box_size)
        end

        accels_t_dt = accelerations(s, parallel=parallel)

        # Update velocities
        for i in 1:length(s.velocities)
            s.velocities[i] += (accels_t[i] + accels_t_dt[i]) * s.timestep / 2
        end

        apply_thermostat!(s.velocities, s, s.thermostat)
        find_neighbours!(s, s.neighbour_finder, step_n, parallel=parallel)

        accels_t = accels_t_dt
        s.n_steps_made[1] += 1
    end
    return s.coords
end

function simulate!(s::Simulation{true},
                    ::VelocityVerlet,
                    n_steps::Integer;
                    parallel::Bool=true)
    n_atoms = length(s.coords)
    is = hcat([collect(1:n_atoms) for i in 1:n_atoms]...)
    js = permutedims(is, (2, 1))
    coords_is = view(s.coords, is)
    coords_js = view(s.coords, js)
    atoms_is = view(s.atoms, is)
    atoms_js = view(s.atoms, js)
    find_neighbours!(s, s.neighbour_finder, 0)
    if isa(s.coords, CuArray)
        self_interactions = CuArray(Diagonal(ones(typeof(s.timestep), n_atoms)))
    else
        self_interactions = Array(Diagonal(ones(typeof(s.timestep), n_atoms)))
    end
    accels_t = accelerations(s, s.coords, coords_is, coords_js, atoms_is, atoms_js, self_interactions)
    accels_t_dt = zero(s.coords)

    for step_n in 1:n_steps
        for logger in values(s.loggers)
            log_property!(logger, s, step_n)
        end

        # In-place updates here required to work with views but are not Zygote-compatible
        s.coords .+= s.velocities .* s.timestep .+ (accels_t .* s.timestep ^ 2) ./ 2
        s.coords .= adjust_bounds_vec.(s.coords, s.box_size)
        accels_t_dt = accelerations(s, s.coords, coords_is, coords_js, atoms_is, atoms_js, self_interactions)
        s.velocities .+= (accels_t .+ accels_t_dt) .* s.timestep / 2

        s.velocities .= apply_thermostat!(s.velocities, s, s.thermostat)
        find_neighbours!(s, s.neighbour_finder, 0)

        accels_t = accels_t_dt
        s.n_steps_made[1] += 1
    end
    return s.coords
end

"""
    VelocityFreeVerlet()

The velocity-free Verlet integrator, also known as the Störmer method.
In this case the `velocities` given to the `Simulator` act as the previous step
coordinates for the first step.
"""
struct VelocityFreeVerlet <: Simulator end

function simulate!(s::Simulation,
                    ::VelocityFreeVerlet,
                    n_steps::Integer;
                    parallel::Bool=true)
    n_atoms = length(s.coords)
    find_neighbours!(s, s.neighbour_finder, 0, parallel=parallel)
    coords_last = s.velocities

    @showprogress for step_n in 1:n_steps
        for logger in values(s.loggers)
            log_property!(logger, s, step_n)
        end

        accels_t = accelerations(s, parallel=parallel)

        # Update coordinates
        coords_copy = s.coords
        for i in 1:length(s.coords)
            s.coords[i] = s.coords[i] + vector(coords_last[i], s.coords[i], s.box_size) + accels_t[i] * s.timestep ^ 2
            s.coords[i] = adjust_bounds.(s.coords[i], s.box_size)
        end
        coords_last = coords_copy

        apply_thermostat!(coords_last, s, s.thermostat)
        find_neighbours!(s, s.neighbour_finder, step_n, parallel=parallel)

        s.n_steps_made[1] += 1
    end
    return s.coords
end

function simulate!(s::Simulation, n_steps::Integer; parallel::Bool=true)
    simulate!(s, s.simulator, n_steps, parallel=parallel)
end

function simulate!(s::Simulation; parallel::Bool=true)
    simulate!(s, s.n_steps - first(s.n_steps_made), parallel=parallel)
end


text:
import PWS
using PWS: cooperative_chemotaxis_system, TrajectoryCallback, SMCEstimate, DirectMCEstimate, marginal_configuration, MarginalEnsemble, ConditionalEnsemble, gene_expression_system, generate_configuration, logpdf
using StaticArrays
import Catalyst
using DiffEqBase
using DiffEqJump
import ModelingToolkit

system_fn = () -> PWS.cooperative_chemotaxis_system(delta_e=0, delta_f=-1, lmax=3, mmax=9, γ=1/5)

sol = begin
system = system_fn()
step_s = PWS.Trajectory([100.0, 200, 300.0], [SA[15.0], SA[50.0], SA[15.0]], [0])
joint = merge(merge(system.sn, system.rn), system.xn)
rx = merge(system.rn, system.xn)
r_idxs = PWS.species_indices(joint, Catalyst.species(system.rn))
u0 = system.u0
u0[1] = step_s.u[1][1]
dprob = DiscreteProblem(rx, u0, (0.0, 100.0), vcat(system.pr, system.px))
jprob = JumpProblem(rx, dprob, Direct(), save_positions=(false, false))

cb = TrajectoryCallback(step_s)
cb = DiscreteCallback(cb, cb, save_positions=(false, false))

sol = solve(jprob, SSAStepper(), callback=cb, tstops=step_s.t, saveat=1.0)
end

using Plots

function active_indices(rs, firstletter = "A")
    smap = Catalyst.speciesmap(rs)
    result = Int[]
    for (species, index) in smap
        sname = String(ModelingToolkit.operation(species).name)
        if startswith(sname, firstletter)
            push!(result, index)
        end
    end
    result
end

using Transducers
function receptor_states(rs)
    smap = Catalyst.speciesmap(rs)
    xf = KeepSomething() do (species, index)
        sname = String(ModelingToolkit.operation(species).name)
        mtch = match(r"([A-Z])_(\d+)_(\d+)", sname)
        if mtch !== nothing
            a = mtch.captures[1] == "A"
            l = parse(Int, mtch.captures[2])
            m = parse(Int, mtch.captures[3])
            (a, l, m) => index
        else
            nothing
        end
    end
    smap |> xf |> collect
end

function bound_ligands(sol, rs)
    rstates = receptor_states(rs)
    rstates |> Map(((a, l, m), i)::Pair -> l .* sol[i,:]) |> sum
end

function bound_methyl(sol, rs)
    rstates = receptor_states(rs)
    rstates |> Map(((a, l, m), i)::Pair -> m .* sol[i,:]) |> sum
end

function active_receptors(sol, rs)
    rstates = receptor_states(rs)
    rstates |> Map(((a, l, m), i)::Pair -> a .* sol[i,:]) |> sum
end

begin
t = sol.t
plot(t, vec(sol[1,:]) ./ 100, label="ligand concentration")
plot!(t, active_receptors(sol, joint) ./ 100, label="active fraction")
plot!(t, vec(sol[end-1,:]) ./ 10000, label="Yp / (Y+Yp)")
plot!(t, bound_ligands(sol, joint) ./ 300, label="bound receptor fraction")
plot!(t, bound_methyl(sol, joint) ./ (8*3*100), label="methylated fraction", legend=:topleft)
end

savefig(plot!(dpi=144), "~/Downloads/plot2.png")

# FROM JuliaMarkdown File

lmax = 3
mmax = 9
K_a = 500
K_i = 25
δg = log(K_a/K_i) # ≈ 3

E0 = 2.0
δf = -2.0

n_clusters = 800

p_bind = 0.05

γ = 1/0.5 # 1 / (adaptation time scale)
γ_B = γ / (mmax * abs(δf))
γ_R = γ_B / 2

params = [
    E0,
    0.1, # in/activation timescale
    δg,
    δf,
    p_bind, # ligand binding to active receptor
    p_bind, # ligand binding to inactive receptor
    p_bind * K_a, # ligand dissociation from active receptor
    p_bind * K_i, # ligand dissociation from inactive receptor
    γ_B, # demethylation of active receptor
    γ_R  # methylation of inactive receptor
]

system = cooperative_chemotaxis_system()
Catalyst.numreactions(system.rn) + Catalyst.numreactions(system.xn)

system = cooperative_chemotaxis_system(lmax=lmax, mmax=mmax, n_clusters=n_clusters, mean_l=100)

conf = generate_configuration(system)

plot(conf.x_traj)

cens = ConditionalEnsemble(system)
mens = MarginalEnsemble(system)
mconf = marginal_configuration(conf)

smc = SMCEstimate(128)
cr = PWS.simulate(smc, conf, cens)
mr = PWS.simulate(smc, mconf, mens)
plot(PWS.log_marginal(cr) - PWS.log_marginal(mr))


text:
# Test calling Knitro with license manager

using KNITRO, JuMP
using Test
using MathOptInterface
const MOI = MathOptInterface

# hs071
# Polynomial objective and constraints
# min x1 * x4 * (x1 + x2 + x3) + x3
# st  x1 * x2 * x3 * x4 >= 25
#     x1^2 + x2^2 + x3^2 + x4^2 = 40
#     1 <= x1, x2, x3, x4 <= 5
# Start at (1,5,5,1)
# End at (1.000..., 4.743..., 3.821..., 1.379...)

# Instantiate license manager
lm = KNITRO.LMcontext()

nruns = 10

for i in 1:nruns
    model = Model(with_optimizer(KNITRO.Optimizer, license_manager=lm, outlev=1))

    initval = [1, 5, 5, 1]

    @variable(model, 1 <= x[i=1:4] <= 5, start=initval[i])
    @NLobjective(model, Min, x[1] * x[4] * (x[1] + x[2] + x[3]) + x[3])
    c1 = @NLconstraint(model, x[1] * x[2] * x[3] * x[4] >= 25)
    c2 = @NLconstraint(model, sum(x[i]^2 for i=1:4) == 40)

    JuMP.optimize!(model)
    @test JuMP.termination_status(model) == MOI.LOCALLY_SOLVED

    # Warning! Free Knitro model before freeing license manager!
    MOI.empty!(backend(model))
end

# Free license manager
KNITRO.KN_release_license(lm)


text:
module ElboMaximize

using Optim
using Optim: Options, NewtonTrustRegion
using ..Model
using ..SensitiveFloats
using ..DeterministicVI
using ..DeterministicVI: init_thread_pool!, ElboIntermediateVariables
using ..DeterministicVI.ConstraintTransforms: TransformDerivatives,
                                              VariationalParams,
                                              ConstraintBatch,
                                              ParameterConstraint,
                                              BoxConstraint,
                                              SimplexConstraint,
                                              allocate_free_params,
                                              to_free!,
                                              to_bound!,
                                              enforce!,
                                              propagate_derivatives!

##################################################
# defaults for optional arguments to `maximize!` #
##################################################

struct ElboConfig{N,T}
    bound_params::VariationalParams{T}
    free_params::VariationalParams{T}
    free_initial_input::Vector{T}
    free_previous_input::Vector{T}
    free_result::SensitiveFloat{T}
    bvn_bundle::Model.BvnBundle{T}
    constraints::ConstraintBatch
    derivs::TransformDerivatives{N,T}
    optim_options::Options
    trust_region::NewtonTrustRegion{T}
end

function ElboConfig{T}(ea::ElboArgs,
                   vp::VariationalParams{T},
                   bound_params::VariationalParams{T} = vp[ea.active_sources];
                   termination_callback = nothing,
                   loc_width::Float64 = 1e-4,
                   loc_scale::Float64 = 1.0,
                   max_iters::Int = 50,
                   constraints::ConstraintBatch = elbo_constraints(bound_params, loc_width, loc_scale),
                   optim_options::Options = elbo_optim_options(max_iters=max_iters,
                                                               termination_callback=termination_callback),
                   trust_region::NewtonTrustRegion = elbo_trust_region())
    free_params = allocate_free_params(bound_params, constraints)
    free_initial_input = to_flat_vector(free_params)
    free_previous_input = similar(free_initial_input)
    free_result = SensitiveFloat{Float64}(length(free_params[1]), length(bound_params), true, true)
    bvn_bundle = Model.BvnBundle{T}(ea.psf_K, ea.S)
    derivs = TransformDerivatives(bound_params, free_params)
    return ElboConfig(bound_params, free_params, free_initial_input,
                  free_previous_input, free_result, bvn_bundle,
                  constraints, derivs, optim_options, trust_region)
end

function elbo_constraints{T}(bound::VariationalParams{T},
                             loc_width::Real = 1.0e-4,
                             loc_scale::Real = 1.0)
    n_sources = length(bound)
    boxes = Vector{Vector{ParameterConstraint{BoxConstraint}}}(n_sources)
    simplexes = Vector{Vector{ParameterConstraint{SimplexConstraint}}}(n_sources)
    for src in 1:n_sources
        i1, i2 = ids.pos[1], ids.pos[2]
        u1, u2 = bound[src][i1], bound[src][i2]
        boxes[src] = [
            ParameterConstraint(BoxConstraint(u1 - loc_width, u1 + loc_width, loc_scale), i1),
            ParameterConstraint(BoxConstraint(u2 - loc_width, u2 + loc_width, loc_scale), i2),
            ParameterConstraint(BoxConstraint(1e-2, 0.99, 1.0), ids.gal_fracdev),
            ParameterConstraint(BoxConstraint(1e-2, 0.99, 1.0), ids.gal_ab),
            ParameterConstraint(BoxConstraint(-10.0, 10.0, 1.0), ids.gal_angle),
            ParameterConstraint(BoxConstraint(0.10, 70.0, 1.0), ids.gal_scale),
            ParameterConstraint(BoxConstraint(-1.0, 10.0, 1.0), ids.flux_loc),
            ParameterConstraint(BoxConstraint(1e-4, 0.10, 1.0), ids.flux_scale),
            ParameterConstraint(BoxConstraint(-10.0, 10.0, 1.0), ids.color_mean[:, 1]),
            ParameterConstraint(BoxConstraint(-10.0, 10.0, 1.0), ids.color_mean[:, 2]),
            ParameterConstraint(BoxConstraint(1e-4, 1.0, 1.0), ids.color_var[:, 1]),
            ParameterConstraint(BoxConstraint(1e-4, 1.0, 1.0), ids.color_var[:, 2])
        ]
        simplexes[src] = [
            ParameterConstraint(SimplexConstraint(0.005, 1.0, 2), ids.is_star),
            ParameterConstraint(SimplexConstraint(0.01/NUM_COLOR_COMPONENTS, 1.0, NUM_COLOR_COMPONENTS), ids.k[:, 1]),
            ParameterConstraint(SimplexConstraint(0.01/NUM_COLOR_COMPONENTS, 1.0, NUM_COLOR_COMPONENTS), ids.k[:, 2])
        ]
    end
    return ConstraintBatch(boxes, simplexes)
end

function elbo_optim_options(args...; xtol_abs = 1e-7, ftol_rel = 1e-6, max_iters = 50,
                            termination_callback = nothing, kwargs...)
    return Optim.Options(args...; x_tol = xtol_abs, f_tol = ftol_rel,
                         g_tol = 1e-8, iterations = max_iters,
                         store_trace = false, show_trace = false,
                         extended_trace = false,
                         callback = termination_callback,
                         kwargs...)
end

function elbo_trust_region(; initial_delta = 1.0, delta_hat = 1e9)
    return Optim.NewtonTrustRegion(initial_delta = initial_delta,
                                   delta_hat = delta_hat)
end

function enforce_references!{T}(ea::ElboArgs, vp::VariationalParams{T}, cfg::ElboConfig)
    @assert length(cfg.bound_params) == length(ea.active_sources)
    for i in 1:length(ea.active_sources)
        cfg.bound_params[i] = vp[ea.active_sources[i]]
    end
    return nothing
end

to_flat_vector{T}(sources::VariationalParams{T}) = vcat(sources...)::Vector{T}

function to_flat_vector!{T}(x::Vector{T}, sources::VariationalParams{T})
    i = 1
    for src in sources
        for j in eachindex(src)
            x[i] = src[j]
            i += 1
        end
    end
    return sources
end

function to_variational_params!{T}(sources::VariationalParams{T}, x::Vector{T})
    i = 1
    for src in sources
        for j in eachindex(src)
            src[j] = x[i]
            i += 1
        end
    end
    return sources
end

##########################################
# Preallocated ElboIntermediateVariables #
##########################################

const ELBO_VARS_POOL = Vector{ElboIntermediateVariables{Float64}}()

get_elbo_vars() = ELBO_VARS_POOL[Base.Threads.threadid()]

function __init__()
    init_thread_pool!(ELBO_VARS_POOL, () -> ElboIntermediateVariables(Float64, 1, true, true))
end

# explicitly call this for use with compiled system image
__init__()

##################################
# Callable Types Passed to Optim #
##################################

function evaluate!{T}(ea::ElboArgs, vp::VariationalParams{T}, cfg::ElboConfig, x::Vector)
    if x != cfg.free_previous_input
        copy!(cfg.free_previous_input, x)
        to_variational_params!(cfg.free_params, x)
        to_bound!(cfg.bound_params, cfg.free_params, cfg.constraints)
        bound_result = elbo(ea, vp, get_elbo_vars(), cfg.bvn_bundle)
        propagate_derivatives!(to_bound!, bound_result,
                               cfg.free_result, cfg.free_params,
                               cfg.constraints, cfg.derivs)
    end
    return nothing
end

# Objective #
#-----------#

struct Objective{N,T} <: Function
    ea::ElboArgs
    vp::VariationalParams{T}
    cfg::ElboConfig{N,T}
end

function (f::Objective)(x::Vector)
    evaluate!(f.ea, f.vp, f.cfg, x)
    return -(f.cfg.free_result.v[])
end

# Gradient #
#----------#

struct Gradient{N,T} <: Function
    ea::ElboArgs
    vp::VariationalParams{T}
    cfg::ElboConfig{N,T}
end

function (f::Gradient)(x::Vector, result::Vector)
    evaluate!(f.ea, f.vp, f.cfg, x)
    free_gradient = f.cfg.free_result.d
    for i in eachindex(result)
        result[i] = -(free_gradient[i])
    end
    return result
end

# Hessian #
#---------#

struct Hessian{N,T} <: Function
    ea::ElboArgs
    vp::VariationalParams{T}
    cfg::ElboConfig{N,T}
end

function (f::Hessian)(x::Vector, result::Matrix)
    evaluate!(f.ea, f.vp, f.cfg, x)
    free_hessian = f.cfg.free_result.h
    for i in eachindex(result)
        result[i] = -(free_hessian[i])
    end
    return result
end

#############
# maximize! #
#############

function maximize!(ea::ElboArgs, vp::VariationalParams{Float64}, cfg::ElboConfig = ElboConfig(ea, vp))
    enforce_references!(ea, vp, cfg)
    enforce!(cfg.bound_params, cfg.constraints)
    to_free!(cfg.free_params, cfg.bound_params, cfg.constraints)
    to_flat_vector!(cfg.free_initial_input, cfg.free_params)
    fill!(cfg.free_previous_input, NaN)
    R = Optim.MultivariateOptimizationResults{Float64,1,Optim.NewtonTrustRegion{Float64}}
    result::R = Optim.optimize(Objective(ea, vp, cfg), Gradient(ea, vp, cfg), Hessian(ea, vp, cfg),
                               cfg.free_initial_input, cfg.trust_region, cfg.optim_options)
    min_value::Float64 = -(Optim.minimum(result))
    min_solution::Vector{Float64} = Optim.minimizer(result)
    to_variational_params!(cfg.free_params, min_solution)
    to_bound!(cfg.bound_params, cfg.free_params, cfg.constraints)
    return result.f_calls, min_value, min_solution, result
end

end # module


text:
"""
    simulate_random_snparray(n::Integer, p::Integer, s::String; mafs::Vector{Float64}, min_ma::Integer)

Creates a random SnpArray in the current directory without missing value,
where each SNP has ⫺5 (default) minor alleles.

Note: if supplied minor allele frequency is extremely small, it could take a long time for
the simulation to generate samples where at least `min_ma` (defaults to 5) are present.

# Arguments:
- `n`: number of samples
- `p`: number of SNPs
- `s`: name of SnpArray that will be created (memory mapped) in the current directory. To not memory map, use `undef`.

# Optional Arguments:
- `mafs`: vector of desired minor allele freuqencies (uniform(0, 0.5) by default)
- `min_ma`: the minimum number of minor alleles that must be present for each SNP (defaults to 5)
"""
function simulate_random_snparray(n::Int64, p::Int64, s::Union{String, UndefInitializer};
                                  mafs::Vector{Float64}=zeros(Float64, p), min_ma::Int = 5)

    @assert all(0.0 .<= mafs .<= 0.5) "Minor allele frequencies must be all in the range (0, 0.5)"

    if mafs != zeros(Float64, p)
        return _random_snparray(n, p, s, mafs, min_ma=min_ma)
    end

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = 0
        while minor_alleles <= min_ma
            maf = 0.5rand()
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
        mafs[j] = maf
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(A1, A2, s)
end

"""
This function requires a vector of minor alleles frequencies to perform simulation.
It will create a random SnpArray in the current directory without missing value,
where each SNP has at least 5 minor alleles.
"""
function _random_snparray(n::Int64, p::Int64, s::Union{String, UndefInitializer},
                          mafs::Vector{Float64}; min_ma::Int = 5)
    all(0.0 .<= mafs .<= 0.5) || throw(ArgumentError("vector of minor allele frequencies must be in (0, 0.5)"))
    any(mafs .<= 0.0005) && @warn("Provided minor allele frequencies contain entries smaller than 0.0005, simulation may take long if sample size is small and min_ma = $min_ma is large")

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = mafs[j]
        while minor_alleles <= min_ma
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(A1, A2, s)
end

"""
Make a SnpArray from 2 BitArrays.
"""
function _make_snparray(A1::BitArray, A2::BitArray, s::Union{String, UndefInitializer})
    n, p = size(A1)
    x = SnpArray(s, n, p)
    for i in 1:(n*p)
        c = A1[i] + A2[i]
        if c == 0
            x[i] = 0x00
        elseif c == 1
            x[i] = 0x02
        elseif c == 2
            x[i] = 0x03
        else
            throw(error("matrix shouldn't have missing values!"))
        end
    end
    return x
end

"""
    simulate_correlated_snparray(n, p, s; block_length, hap, prob)

Simulates a SnpArray with correlation. SNPs are divided into blocks where each
adjacent SNP is the same with probability prob. There are no correlation between blocks.

# Arguments:
- `n`: number of samples
- `p`: number of SNPs
- `s`: name of SnpArray that will be created (memory mapped) in the current directory. To not memory map, use `undef`.

# Optional arguments:
- `block_length`: length of each LD block
- `hap`: number of haplotypes to simulate for each block
- `prob`: with probability `prob` an adjacent SNP would be the same.
"""
function simulate_correlated_snparray(n::Int64, p::Int64, s::Union{String, UndefInitializer};
            block_length::Int64=20, hap::Int=20, prob::Float64=0.75)

    @assert mod(p, block_length) == 0 "block_length ($block_length) is not divible by p ($p)"
    @assert 0 < prob < 1 "transition probably should be between 0 and 1, got $prob"

    x = SnpArray(s, n, p)
    haplotypes = zeros(hap, block_length)
    snps = zeros(block_length)
    blocks = Int(p / block_length)

    @inbounds for b in 1:blocks

        #create pool of haplotypes for each block
        _sample_haptotypes!(haplotypes, prob)

        for i in 1:n
            #sample 2 haplotypes with replacement from the pool of haplotypes
            row1 = rand(1:hap)
            row2 = rand(1:hap)
            for j in 1:block_length
                snps[j] = haplotypes[row1, j] + haplotypes[row2, j]
            end

            #copy haplotypes into x
            _copy_blocks!(x, i, snps, b, block_length)
        end
    end

    return x
end

function _sample_haptotypes!(haplotypes::Matrix, prob::Float64)
    n, p = size(haplotypes)
    fill!(haplotypes, 0)

    @inbounds for i in 1:n
        cur_row_sum = 0
        while cur_row_sum == 0
            curr = rand(0:1)
            haplotypes[i, 1] = curr
            cur_row_sum += curr
            for j in 2:p
                stay = rand(Bernoulli(prob)) #stay = 1 means retain the current value
                curr = (stay == 1 ? curr : 1 - curr)
                haplotypes[i, j] = curr
                cur_row_sum += curr
            end
        end
    end
end

function _copy_blocks!(x::SnpArray, row, snps, cur_block, block_length)
    #copy sampled snps into SnpArray
    @inbounds for k in 1:length(snps)
        c = snps[k]
        col = (cur_block - 1) * block_length + k
        if c == 0
            x[row, col] = 0x00
        elseif c == 1
            x[row, col] = 0x02
        elseif c == 2
            x[row, col] = 0x03
        else
            throw(error("SNP values should be 0, 1, or 2 but was $c"))
        end
    end
end

"""
    simulate_random_response(x::SnpArray, xbm::SnpBitMatrix, k::Int, d::UnionAll, l::Link)

This function simulates a random response (trait) vector `y` based on provided x, β, distirbution,
and link function. When the distribution is from Poisson, Gamma, or Negative Binomial, we simulate `β ∼ N(0, 0.3)`
to roughly ensure the mean of response `y` doesn't become too large. For other distributions,
we choose `β ∼ N(0, 1)`.

# Arguments
- `x`: The SnpArray
- `xbm`: SnpBitMatrix type of your SnpArray
- `k`: the true number of predictors.
- `d`: The distribution of the simulated trait (note `typeof(d) = UnionAll` but `typeof(d())` is an actual distribution: e.g. Normal)
- `l`: The link function. Input `canonicallink(d())` if you want to use the canonical link of `d`.

# Optional arguments
- `nn`: The number of success until stopping in negative binomial regression, defaults to 10
- `α`: Shape parameter of the gamma distribution, defaults to 1
"""
function simulate_random_response(x::SnpArray, xbm::SnpBitMatrix, k::Int,
                                  d::UnionAll, l::Link; nn = 10,
                                  α = 1) where {T <: Float}
    n, p = size(x)
    if (typeof(d) <: NegativeBinomial) || (typeof(d) <: Gamma)
        l == LogLink() || throw(ArgumentError("Distribution $d must use LogLink!"))
    end

    #simulate a random model β from a normal distribution
    true_b = zeros(p)
    if d == Poisson || d == Gamma || d == NegativeBinomial
        true_b[1:k] = rand(Normal(0, 0.3), k)
    else
        true_b[1:k] = randn(k)
    end
    shuffle!(true_b)
    correct_position = findall(x -> x != 0, true_b)

    #simulate phenotypes (e.g. vector y)
    if d == Normal || d == Poisson || d == Bernoulli
        prob = linkinv.(l, xbm * true_b)
        clamp!(prob, -20, 20)
        y = [rand(d(i)) for i in prob]
    elseif d == NegativeBinomial
        μ = linkinv.(l, xbm * true_b)
        clamp!(μ, -20, 20)
        prob = 1 ./ (1 .+ μ ./ nn)
        y = [rand(d(nn, i)) for i in prob] #number of failtures before nn success occurs
    elseif d == Gamma
        μ = linkinv.(l, xbm * true_b)
        β = 1 ./ μ # here β is the rate parameter for gamma distribution
        y = [rand(d(α, i)) for i in β] # α is the shape parameter for gamma
    end
    y = Float64.(y)

    return y, true_b, correct_position
end

"""
    adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})

Makes 1 SNP (a column of `x`) correlate with SNPs in `location` with correlation coefficient roughly `ρ`.

# Arguments:
- `x`: the snparray
- `ρ`: correlation coefficient
- `pos`: the position of the target SNP that everything would be correlated to
- `location`: All the SNPs that shall be correlated with the SNP at position `pos` with correlation `ρ`.
"""
function adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})
    @assert 0 <= ρ <= 1 "correlation coefficient must be in (0, 1) but was $ρ"

    for loc in location
        for i in 1:size(x, 1)
            prob = rand(Bernoulli(ρ))
            prob == 1 && (x[i, loc] = x[i, pos]) #make 2nd column the same as 1st ~ρ% of the time
        end
        corr = cor(x[:, loc], x[:, pos])
    end
end

"""
    simulate_rare_variants(mafs::Vector{Float64}, penetrance::Vector{Float64}, cases::Int, controls::Int)

TODO: implement Algrithm 1 of https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3025646/
"""
function simulate_rare_variants(mafs::Vector{Float64}, penetrance::Vector{Float64},
                                cases::Int, controls::Int)
    return nothing # not implemented yet
end

"""
    make_bim_fam_files(x::SnpArray, y, name::String)

Creates .bim and .bed files from a SnpArray.

# Arguments:
- `x`: A SnpArray (i.e. `.bed` file on the disk) for which you wish to create corresponding `.bim` and `.fam` files.
- `name`: string that should match the `.bed` file (Do not include `.bim` or `.fam` extensions in `name`).
- `y`: Trait vector that will go in to the 6th column of `.fam` file.
"""
function make_bim_fam_files(x::SnpArray, y, name::String)
    ly = length(y)
    n, p = size(x)
    @assert n == ly "dimension mismatch: phenotype data has length $ly but SnpArray has $n samples"

    #create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".bim", "w") do f
        for i in 1:p
            write(f, "1 \t $i \t 0 \t 1 \t 1 \t 2 \n")
        end
    end

    #create .fam file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".fam", "w") do f
        for i in 1:n
            yi = y[i]
            write(f, "$i \t 1 \t 0 \t 0 \t 1 \t $yi \n")
        end
    end
end


text:
# test afterOptBL
# Claudia March 2015
# use cases in test_optBL and test_optTopLevel

# test_optBL -------------------------------------------------------------
# should not do anything, should leave net unchanged

# CASE G
include("../src/types.jl")
include("../src/functions.jl")

df = readtable("CaseG_output.csv")
d2 = readTableCF(df)

tree = "((((6,4)1,(7)11#H1:::0.8)5,(11#H1:::0.2,8)),10);" # Case G different starting branch lengths
f = open("prueba_tree.txt","w")
write(f,tree)
close(f)
net = readTopologyUpdate("prueba_tree.txt");
printEdges(net)

net.ht
realht = [0.1,0.2,0.1,1.0]

@time fmin,xmin=optBL!(net,d2)
# with 100*logPseudoLik and t in (0,10): takes longer, but finds it!
#got 0.0 at [0.10085,0.19991,0.1001,0.98557] after 1123 iterations (returned FTOL_REACHED)
#elapsed time: 53.20413795 seconds (107820188 bytes allocated, 0.13% gc time)
#(2.169777681982341e-9,[0.100853,0.199907,0.100098,0.985569])

f=afterOptBL!(net,d2)
all(f) || error("afterOptBL not correctly for optBL in CASE G")

#--------
# Case H
include("../src/types.jl")
include("../src/functions.jl")

df = readtable("CaseH_output.csv")
d2 = readTableCF(df)

# starting ht (gamma,t3,t5,t7)
ht = [0.2,1.,1.,1.]

tree = string("((((6:0.1,4:1.5):",string(ht[2]),",#H1:::",string(ht[1]),"),7:0.2):",string(ht[4]),",(8)#H1:::",string(1-ht[1]),",10:0.1);") # Case H
f = open("prueba_tree.txt","w")
write(f,tree)
close(f)
net = readTopologyUpdate("prueba_tree.txt");
printEdges(net)

net.ht
realht = [0.1,0.1,1.,0.1]

@time fmin,xmin=optBL!(net,d2)
# with 100*logPseudoLik and t in (0,10):
#got 0.0 at [0.1,0.1,1.0,0.1] after 208 iterations (returned FTOL_REACHED)
#elapsed time: 6.892366587 seconds (20646048 bytes allocated)
#(2.4010737722561867e-13,[0.0999999,0.0999992,1.0,0.1])

f=afterOptBL!(net,d2)
all(f) || error("afterOptBL not correctly for optBL in CASE H")

# -----------
# CASE J
include("../src/types.jl")
include("../src/functions.jl")

df = readtable("CaseJ_output.csv")
d2 = readTableCF(df)

# starting ht (gamma,t3,t5,t7)
ht = [0.2,1.,1.,1.]

tree = string("((((6)#H1:::",string(1-ht[1]),",4:1.5):",string(ht[2]),",7:0.2):",string(ht[3]),",8:0.1,(#H1:::",string(ht[1]),",10));") # Case J
f = open("prueba_tree.txt","w")
write(f,tree)
close(f)
net = readTopologyUpdate("prueba_tree.txt");
printEdges(net)

net.ht
realht = [0.1,0.2,0.1,1.0]

@time fmin,xmin=optBL!(net,d2)
# with 100*logPseudoLik and t in (0,10):
#got 0.0 at [0.1,0.2,0.1,1.00003] after 249 iterations (returned FTOL_REACHED)
#elapsed time: 6.656956697 seconds (23010204 bytes allocated, 0.56% gc time)
#(5.280560068867562e-12,[0.0999983,0.199999,0.0999999,1.00003])

f=afterOptBL!(net,d2)
all(f) || error("afterOptBL not correctly for optBL in CASE J")

# -----------
# CASE F
include("../src/types.jl")
include("../src/functions.jl")

df = readtable("CaseF_output.csv")
d2 = readTableCF(df)

# starting ht (gamma,t4,t5,t9)
ht = [0.1,1.,1.,1.]

tree = string("(((6:0.1,(4)11#H1:::",string(1-ht[1]),")1:",string(ht[3]),",(11#H1:::",string(ht[1]),",7))5:",string(ht[4])",8:0.1,10:0.1);") # Case F: bad diamond I
f = open("prueba_tree.txt","w")
write(f,tree)
close(f)
net = readTopologyUpdate("prueba_tree.txt");
printEdges(net)

net.ht
realht = [0.1,0.127,0.0285]

@time fmin,xmin=optBL!(net,d2)
# with 100*logPseudoLik and t in (0,10):
# got 0.0 at [0.1,0.12689,0.02855] after 3356 iterations (returned FTOL_REACHED)
# elapsed time: 90.405727385 seconds (281649264 bytes allocated, 0.24% gc time)
# (8.157692309723136e-12,[0.0999999,0.126889,0.028549])

f=afterOptBL!(net,d2)
all(f) || error("afterOptBL not correctly for optBL in CASE F")

# ---------------
# CASE I
include("../src/types.jl")
include("../src/functions.jl")

df = readtable("CaseI_output.csv")
d2 = readTableCF(df)

# starting ht (gamma,t4,t6,t9,t10)
ht = [0.2,0.0,2.0,2.0,2.0]
ht = [0.2,0.0,0.5,0.5,0.5]

tree = string("((((8,10):",string(ht[2]),")#H1:::",string(1-ht[1]),",7):",string(ht[3]),",6,(4,#H1:",string(ht[4]),"::",string(ht[1]),"):",string(ht[5]),");") # Case I Bad diamond II
f = open("prueba_tree.txt","w")
write(f,tree)
close(f)
net = readTopologyUpdate("prueba_tree.txt");
printEdges(net)

net.ht
realht = [0.1,2.0,1.0,1.0,1.0]

@time fmin,xmin=optBL!(net,d2)
# with 100*logPseudoLik and t in (0,10):
# ht = [0.2,0.0,2.0,2.0,2.0]
#got 0.0 at [0.1,2.0,1.0,0.99994,1.0] after 753 iterations (returned FTOL_REACHED)
#elapsed time: 47.213517434 seconds (80265788 bytes allocated, 0.16% gc time)
#(1.2568189823957448e-12,[0.0999998,2.0,1.0,0.999941,1.0])

f=afterOptBL!(net,d2)
all(f) || error("afterOptBL not correctly for optBL in CASE I")

# test_optTopLevelparts.jl------------------------------------------------------------------

# starting topology: Case G
include("../examples/case_g_example.jl");
currT = deepcopy(net);
printEdges(currT)
printNodes(currT)

# real network: Case H
df = readtable("CaseH_output.csv")
d = readTableCF(df)

# real network: Case H
df = readtable("CaseH_output.csv")
d = readTableCF(df)

currloglik,currxmin = optBL!(currT,d, false)
isValid(currT)
success,flagh,flagt,flaghz = afterOptBL!(currT,d)
reject = afterOptBLAll!(currT,d)

newT = deepcopy(currT);
count = 0
N = 100
move = whichMove(currT)
move = :CHdir
move = :MVorigin
move = :MVtarget
move = :nni

flag = proposedTop!(move,newT,true,count,N)
printEdges(newT)
printNodes(newT)

newloglik, newxmin = optBL!(newT,d)
reject = afterOptBLAll!(newT,d)
isValid(newT)
afterOptBL!(newT,d)
currT=deepcopy(newT);

# starting topology: Case F
include("../examples/case_f_example.jl");
currT = deepcopy(net);
printEdges(currT)

# real network: Case H
df = readtable("CaseH_output.csv")
d = readTableCF(df)

currloglik,currxmin = optBL!(currT,d)
updateParameters!(currT)
updateLik!(currT,currloglik)
newT = deepcopy(currT);
count = 0
N = 100
move = whichMove(currT)
move = :CHdir
move = :MVorigin
move = :MVtarget
move = :nni

flag = proposedTop!(move,newT,true,count,N)
printEdges(newT)
printNodes(newT)
sum([e.hybrid?1:0 for e in newT.edge]) == 2 || error("there are not 2 hybrid edges")
newT.hybrid[1].k

newloglik, newxmin = optBL!(newT,d)
newloglik - currloglik

currT = deepcopy(newT);
currloglik = newloglik
currxmin = newxmin


text:
function build_statistical_model(;α_mu, α_sigma, C₀_mu, C₀_sigma, err_obs_iid_sigma_scale, err_prc_iid_sigma_scale)
    return eval(quote
        @model function f(::Type{T} = Float64; data, DDE, times) where {T}

            # Prior on hyporheic travel time distribution parameter
            α ~ Normal($α_mu, $α_sigma)

            # Prior on concentration
            C₀ ~ Normal($C₀_mu, $C₀_sigma)

            # Priors on observation and process error
            err_obs_iid_sigma ~ truncated(Cauchy(0, $err_obs_iid_sigma_scale), 0, Inf)
            err_proc_iid_sigma ~ truncated(Cauchy(0, $err_prc_iid_sigma_scale), 0, Inf)

            # Specify model initial condition and parameters
            p = Dict(
                "PAR" => ODE.p["PAR"],
                "OSat" => ODE.p["OSat"],
                "day" => ODE.p["day"],
                "GPP_daily"=> GPP_daily,
                "ER_daily" => ER_daily,
                "kGas" => kGas,
                "Q" => ODE.p["Q"],
                "V" => ODE.p["V"])

            if DOdata === missing
                data = Vector{T}(undef, length(times))
                data[1] = ODE.u0[1]
            end

            C_true = Vector{T}(undef, (length(times)))
            C_true[1] ~ Normal(data[1], err_obs_iid_sigma)

            for i = 2:length(times)
                problem = remake(ODE, p = p, u0 = [DO_true[i-1]], tspan = (times[i-1], times[i]))
                predicted = solve(problem, alg = Euler(), dt = 0.1)
                DO_true[i] ~ Normal(predicted(times[i])[1], err_proc_iid_sigma)
                DOdata[i] ~ Normal(DO_true[i], err_obs_iid_sigma)
            end
        end
    end)
end


text:
abstract type AbstractGrid end

abstract type Grid{d} end

import Base: ndims, length

Base.length(g::Grid{d}) where d = max(1,n_nodes(g))  # we should deprecate n_nodes

Base.ndims(grid::T) where T<:Grid{d} where d = d

function Base.show(io::IO, grid::Grid)
    print(io, typeof(grid))
end

# # # backward backward compatibility
nodes(::Type{<:Union{ListOfPoints,ListOfPoints{d}}}, grid::Grid{d}) where d = nodes(grid)
nodes(::Type{<:Matrix}, grid::Grid) = copy(from_LOP(nodes(grid)))

node(::Type{<:Union{Point,Point{d}}}, grid::Grid{d}, i::Int) where d = node(grid,i)
node(::Type{<:Vector}, grid::Grid, i::Int) = Vector(node(grid,i)...)

import Base

### Product Grids

# special product grid, with endo/exo distinction

struct ProductGrid{T,S}
    exo::T
    endo::S
end

⊗(grid1::Grid{d1}, grid2::Grid{d2}) where d1 where d2 = ProductGrid(grid1, grid2)

Base.show(io::IO, pg::ProductGrid) = print(io, pg.exo, "⊗",  pg.endo)
Base.length(g::ProductGrid{g1,g2}) where g1 where g2 = length(g.exo)*length(g.endo)

struct PGrid{d, T<:Tuple} <: Grid{d}
    grids::T
end

function PGrid(grids...)
    d = sum( ndims(g) for g in grids)
    t = typeof(grids)
    return PGrid{d, t}(grids)
end

function iter(pg::T) where T <:PGrid
    (concat(e...) for e in Iterators.product( (g for g in pg.grids)... ))
end

Base.length(pg:: PGrid{d, T} ) where d where T = prod( (length(g) for g in pg.grids) )
Base.eltype(pg:: PGrid{d, T} ) where d where T = SVector{d, Float64}

function Base.show(io::IO,x::PGrid{d,T})  where d where T
    N = length(x.grids)
    print(io, x.grids[1])
    if N==1
        return
    end
    for n=2:N
        print(" ⨯ ")
        print(io, x.grids[n])
    end
end

function enumerate_product(pg::T) where T <:PGrid
    iter1 = Iterators.product( (1:length(g) for g in pg.grids)... )
    iter2 = (concat(e...) for e in Iterators.product( (g for g in pg.grids)... ))
    # return iter1, iter2
    zip(iter1, iter2)
end

### EmptyGrid

## TODO: this should probably be of dimension d
struct EmptyGrid{d} <: Grid{d}
    # this grid does not exist ;-)
end
Base.show(io::IO, g::EmptyGrid{d}) where d = print(io, "EmptyGrid{$d}")

nodes(grid::EmptyGrid) = nothing
n_nodes(grid::EmptyGrid) = 0 ##### Reconsider: ???
function node( grid::EmptyGrid{d}, i::Int) where d
    return fill(NaN, SVector{d, Float64})
end

##########################
# Grid made of one point #
##########################

struct PointGrid{d} <: Grid{d}
    point::SVector{d,Float64}
end

function (::Type{<:PointGrid})(point::Vector{Float64})
    d = length(point)
    PointGrid{d}(SVector{d,Float64}(point...))
end

nodes(grid::PointGrid) = [grid.point]
n_nodes(grid::PointGrid) = 1
node(grid::PointGrid, i::Int) = grid.point # fail if i!=1 ?

#####################
# Unstructured Grid #
#####################

struct UnstructuredGrid{d} <: Grid{d}
    nodes::ListOfPoints{d}
end

Base.show(io::IO, gr::UnstructuredGrid{d}) where d = print(io, "UNSgrid{$d}")

# Old-convention
function UnstructuredGrid{d}(nodes::Matrix{Float64}) where d
    N = size(nodes,1)
    @assert d == size(nodes,2)
    UnstructuredGrid{d}(reshape(reinterpret(Point{d}, vec(copy(nodes'))), (N,)))
end

nodes(grid::UnstructuredGrid) = grid.nodes
n_nodes(grid::UnstructuredGrid) = length(grid.nodes)
node(grid::UnstructuredGrid, i::Int) = grid.nodes[i] # fail if i!=1 ?

function Product(a::UnstructuredGrid{d1}, b::UnstructuredGrid{d2}) where d1 where d2
    A = [Base.product(a.nodes, b.nodes)...]
    N = length(A)
    d = d1 + d2
    nodes = reshape(reinterpret(Point{d},(A)),(N,))
    return UnstructuredGrid{d}(nodes)
end

#################
# UCGrid #
#################

function mlinspace(min, max, n)
    # this now returns an iterator
    nodes = map((x,y,z)->range(x,stop=y,length=z), min, max, n)
    return Base.product(nodes...)
end

struct UCGrid{d} <: Grid{d}
    min::Point{d}
    max::Point{d}
    n::SVector{d,Int}
    nodes::ListOfPoints{d}
end

Base.iterate(ucg::UCGrid{d}, args...) where d = iterate(ucg.nodes, args...)
Base.length(ucg::UCGrid{d}) where d = length(ucg.nodes)
Base.eltype(ucg::UCGrid{d}) where d = SVector{d, Float64}

function Base.show(io::IO, g::UCGrid{d}) where d
    print("UCGrid{", d, "}")
end

function (::Type{<:UCGrid})(min::SVector{d,Float64}, max::SVector{d,Float64}, n::SVector{d,Int64}) where d
    A = [mlinspace(min, max, n)...]
    N = prod(n)
    mm = reshape(reinterpret(Point{d},vec(A)),(N,))
    return UCGrid{d}(min, max, n, mm)
end

(::Type{<:UCGrid})(min::Vector{Float64},max::Vector{Float64},n::Vector{Int64}) = UCGrid(SVector(min...), SVector(max...), SVector(n...))
(::Type{<:UCGrid})(min::Vector{Float64},max::Vector{Float64},n::Int64) = UCGrid(min, max, fill(n,length(min)))

scales(grid::UCGrid{d}) where d = Tuple{Vararg{Vector{Float64},d}}([range(grid.min[i], grid.max[i];length=grid.n[i]) for i=1:d])

nodes(grid::UCGrid{d}) where d = grid.nodes
n_nodes(grid::UCGrid{d}) where d = length(grid.nodes)
node(grid::UCGrid{d}, i::Int) where d = grid.nodes[i]

nodes(::Type{<:ListOfPoints}, grid::UCGrid{d}) where d  = grid.nodes
node(::Type{<:Point}, grid::UCGrid{d}, i::Int64) where d = node(grid,i)

function Product(a::UCGrid{d1}, b::UCGrid{d2}) where d1 where d2
  return Dolo.UCGrid{d1+d2}(
      SVector(a.min..., b.min...),
      SVector(a.max..., b.max...),
      SVector(a.n..., b.n...)
  )
end

################
# Smolyak Grid #
################

struct SmolyakGrid{d} <: Grid{d}
    smol_params::BM.SmolyakParams{Float64,Vector{Int}}
    nodes::Matrix{Float64}
    B_nodes::Matrix{Float64}
end

function SmolyakGrid{d}(min::Point{d}, max::Point{d}, mu::SVector{d,Int64}) where d
    sp = BM.SmolyakParams(d, Vector(mu), Vector(min), Vector(max))
    nodes = BM.nodes(sp)
    B_nodes = BM.evalbase(sp, nodes)
    return SmolyakGrid{d}(sp, nodes, B_nodes)
end

SmolyakGrid(min::Point{d}, max::Point{d}, mu::SVector{d,Int64}) where d = SmolyakGrid{d}(min,max,mu)
SmolyakGrid(min::Point{d}, max::Point{d}, mu::Int64) where d = SmolyakGrid{d}(min, max, SVector([mu for i=1:d]...))
SmolyakGrid{d}(min::Point{d}, max::Point{d}, mu::Int64) where d = SmolyakGrid{d}(min, max, SVector([mu for i=1:d]...))

function SmolyakGrid(min::Vector{Float64},max::Vector{Float64},mu::Union{Vector{Int64},Int64})
    d = length(min)
    mmu = mu isa Int ? ffill(mu,d) : mu
    @assert d == length(max) == length(mmu)
    SmolyakGrid{d}(SVector(min...), SVector(max...), SVector(mu...))
end
#
function SmolyakGrid{d}(min::Vector{Float64},max::Vector{Float64},mu::Union{Vector{Int64},Int64}) where d
    mmu = mu isa Int ? fill(mu,d) : mu
    @assert d == length(min) == length(max) == length(mmu)
    SmolyakGrid{d}(SVector(min...), SVector(max...), SVector(mu...))
end

nodes(grid::SmolyakGrid{d}) where d  = reshape(reinterpret(Point{d}, vec(copy(grid.nodes'))),( size(grid.nodes,1),))
n_nodes(grid::SmolyakGrid) = size(grid.nodes,1)
node(grid::SmolyakGrid{d}, i::Int) where d = Point{d}(grid.nodes[i,:]...)

###############
# Random Grid #
###############

struct RandomGrid{d} <: Grid{d}
    min::SVector{d,Float64}
    max::SVector{d,Float64}
    n::Int64
    nodes::ListOfPoints{d}
end

function  (::Type{<:Union{RandomGrid,RandomGrid{d}}})(min::Point{d}, max::Point{d}, n::Int) where d
    nodes = reshape(reinterpret(Point{d}, vec(rand(d, n))), (n,))  # on [0, 1]
    for nn=1:length(n)
        nodes[nn] = nodes[nn] .* (max-min) + min
    end
    RandomGrid(min, max, n, copy(nodes))
end

# function RandomGrid(min::Vector{Float64}, max::Vector{Float64}, n::Int)
#     d = length(min)
#     RandomGrid{d}(SVector(min...),SVector(max...),n)
# end

function RandomGrid{d}(min::Vector{Float64}, max::Vector{Float64}, n::Int) where d
    @assert d == length(min)
    all(max .> min) || error("max must be greater than min")
    RandomGrid{d}(SVector{d}(min...),SVector{d}(max...),n)
end

nodes(grid::RandomGrid) = grid.nodes
nodes(::Type{<:ListOfPoints}, grid::RandomGrid) = nodes(grid)

n_nodes(grid::RandomGrid) = length(grid.nodes)

node(grid::RandomGrid,i::Int) = grid.nodes[i]

concat(v...) = length(v)==1 ? v[1] : SVector(v[1]..., concat(v[2:end]...)... )


text:
using Sole
using Test

@testset "Sole.jl" begin
    # Write your tests here.
end


text:
using Base.Test
using PValueAdjust

newPvalues = padjust([0.05, 0.03, 0.01, 0.5], BenjaminiYekutieli)
@test_approx_eq newPvalues [0.13888888888889, 0.125, 0.083333333333, 1]

@test_approx_eq padjust([1, 0.5, 0.1, 0.001, 0.0001, 0.002], BenjaminiYekutieli) [1.00000, 1.00000, 0.36750, 0.00735, 0.00147, 0.0098]
@test_approx_eq padjust([0.05, 0.05, 0.05], BenjaminiYekutieli) [0.0916666666667, 0.0916666666667, 0.0916666666667]
@test all(padjust([1.0 ,1.0, 1.0], BenjaminiYekutieli) .<= 1.0)

@test_throws DomainError padjust([-1], BenjaminiYekutieli)
@test_throws DomainError padjust([2], BenjaminiYekutieli)

@test padjust([0.5], BenjaminiYekutieli) == [0.5]


text:
"""
    coef(MLM; isXSum, isZSum)

Extracts coefficients from Mlm object, accounting for the possibility of the
object containing back-estimated interactions

# Arguments

- MLM = Mlm object

# Keyword arguments

- isXSum = boolean flag indicating if a left-out `X` sum contrast
  interaction was back-estimated AND should not be included in the extracted
  coefficients. Defaults to `false`.
- isZSum = boolean flag indicating if a left-out `Z` sum contrast
  interaction was back-estimated AND should not be included in the extracted
  coefficients. Defaults to `false`.

# Value

2d array of floats

# Some notes

If the left-out `X` and/or `Z` sum contrast interactions were back-estimated,
but the user wishes to include them in the extracted coefficients anyway, set
`isXSum` and/or `isZSum` to `false`.

"""
function coef(MLM::Mlm; isXSum::Bool=false, isZSum::Bool=false)

    return MLM.B[1:(end-isXSum), 1:(end-isZSum)]
end

"""
    predict(MLM, newPredictors; isXSum, isZSum)

Calculates new predictions based on Mlm object, accounting for the possibility
of the object containing back-estimated interactions

# Arguments

- MLM = Mlm object
- newPredictors = Predictors object. Defaults to the `data.predictors` field
  in the Mlm object used to fit the model.

# Keyword arguments

- isXSum = boolean flag indicating whether a left-out `X` sum contrast
  interaction was back-estimated. Defaults to `false`.
- isZSum = boolean flag indicating whether a left-out `Z` sum contrast
  interaction was back-estimated. Defaults to `false`.

# Value

Response object

"""
function predict(MLM::Mlm, newPredictors::Predictors=MLM.data.predictors;
                 isXSum::Bool=false, isZSum::Bool=false)

  	# Include X and Z intercepts in new data if necessary
  	if MLM.data.predictors.isXIntercept==true &&
       newPredictors.isXIntercept==false
    	newPredictors.X = add_intercept(newPredictors.X)
    	newPredictors.isXIntercept = true
    	println("Adding X intercept to newPredictors.")
  	end
  	if MLM.data.predictors.isZIntercept==true &&
       newPredictors.isZIntercept==false
    	newPredictors.Z = add_intercept(newPredictors.Z)
    	newPredictors.isZIntercept = true
    	println("Adding Z intercept to newPredictors.")
  	end

  	# Remove X and Z intercepts in new data if necessary
  	if MLM.data.predictors.isXIntercept==false &&
       newPredictors.isXIntercept==true
    	newPredictors.X = remove_intercept(newPredictors.X)
    	newPredictors.isXIntercept = false
    	println("Removing X intercept from newPredictors.")
  	end
  	if MLM.data.predictors.isZIntercept==false &&
       newPredictors.isZIntercept==true
    	newPredictors.Z = remove_intercept(newPredictors.Z)
    	newPredictors.isZIntercept = false
    	println("Removing Z intercept from newPredictors.")
  	end

    # Calculate predictions
    return Response(calc_preds(newPredictors.X, newPredictors.Z,
                               coef(MLM, isXSum=isXSum, isZSum=isZSum)))
end

"""
    fitted(MLM; isXSum, isZSum)

Calculate fitted values of an Mlm object, accounting for the possibility of
the object containing back-estimated interactions

# Arguments

- MLM = Mlm object

# Keyword arguments

- isXSum = boolean flag indicating whether a left-out `X` sum contrast
  interaction was back-estimated. Defaults to `false`.
- isZSum = boolean flag indicating whether a left-out `Z` sum contrast
  interaction was back-estimated. Defaults to `false`.

# Value

Response object

"""
function fitted(MLM::Mlm; isXSum::Bool=false, isZSum::Bool=false)

    # Call the predict function with default newPredictors
    return predict(MLM, isXSum=isXSum, isZSum=isZSum)
end

"""
    resid(MLM, newData; isXSum, isZSum)

Calculates residuals of an Mlm object, accounting for the possibility of the
object containing back-estimated interactions

# Arguments

- MLM = Mlm object
- newData = RawData object. Defaults to the `data` field in the Mlm object
  used to fit the model.

# Keyword arguments

- isXSum = boolean flag indicating whether a left-out `X` sum contrast
  interaction was back-estimated. Defaults to `false`.
- isZSum = boolean flag indicating whether a left-out `Z` sum contrast
  interaction was back-estimated. Defaults to `false`.

# Value

2d array of floats

"""
function resid(MLM::Mlm, newData::RawData=MLM.data;
               isXSum::Bool=false, isZSum::Bool=false)

    # Include X and Z intercepts in new data if necessary
    if MLM.data.predictors.isXIntercept==true &&
       newData.predictors.isXIntercept==false
        newData.predictors.X = add_intercept(newData.predictors.X)
        newData.predictors.isXIntercept = true
        newData.p = newData.p + 1
        println("Adding X intercept to newData.")
    end
    if MLM.data.predictors.isZIntercept==true &&
       newData.predictors.isZIntercept==false
        newData.predictors.Z = add_intercept(newData.predictors.Z)
        newData.predictors.isZIntercept = true
        newData.q = newData.q + 1
        println("Adding Z intercept to newData.")
    end

    # Remove X and Z intercepts in new data if necessary
    if MLM.data.predictors.isXIntercept==false &&
       newData.predictors.isXIntercept==true
        newData.predictors.X = remove_intercept(newData.predictors.X)
        newData.predictors.isXIntercept = false
        newData.p = newData.p - 1
        println("Removing X intercept from newData.")
    end
    if MLM.data.predictors.isZIntercept==false &&
       newData.predictors.isZIntercept==true
        newData.predictors.Z = remove_intercept(newData.predictors.Z)
        newData.predictors.isZIntercept = false
        newData.q = newData.q - 1
        println("Removing Z intercept from newData.")
    end

	  # Calculate residuals
    return calc_resid(get_X(newData), get_Y(newData), get_Z(newData),
                      coef(MLM, isXSum=isXSum, isZSum=isZSum))
end


text:
push!(LOAD_PATH, "src")
using LennardJones
using SIMD
using DataFrames

function calc_max_atoms_in_vlist(natoms::Int64, vlist::VerletList)

    counts = Vector{Int64}()
    for i in 1:natoms
        counter = 0
        ptr = vlist.offset[i]
        while vlist.list[ptr] > 0
            counter += 1
            ptr += 1
        end
        push!(counts, counter)
    end
    return maximum(counts)
end

function trial(expr::Expr, title::String)
    println("\n$title ($n_atoms atoms)")
    eval(expr)
    t = @elapsed begin
        for i in 1:n_samples
            eval(expr)
        end
    end
    energy_result = eval(expr)
    diff = energy - energy_result
    if abs(diff) < 1e-10
        diff = "0.0 (or less than 1e10)"
    end
    println("  Energy: $energy_result (Δe = $diff)")
    return t
end

println("Starting Benchmark")
T         = Float64
n_samples = 1
cut_off   = T(10.0) # angstrom

# Generate states
xyz_aos   = generate_cubic_lattice(100.0, parse(Int, ARGS[1]), LATTICE.primitive) # This is in AoS

xyz_aos   = convert(Matrix{T}, xyz_aos)
xyz_soa   = collect(xyz_aos')
n_atoms   = max(size(xyz_aos)...)

println(" Nº atoms: $n_atoms")
state_aos = State(xyz_aos)
state_soa = State(xyz_soa)

# Verlet list
vlist = VerletList(n_atoms)
vlist.cutoff = cut_off
println("\nUpdating vlist in serial mode")
vlist = update_serial!(vlist, state_aos.coords.values)
m_atoms = calc_max_atoms_in_vlist(n_atoms, vlist)
println(" Max atoms in Verlet list: $m_atoms")

# Calculate the correct, known energies and forces
println("\nCalculating correct energy")
energy = serial(state_aos, cut_off, LennardJones.NOFORCES)
print(" Energy: $energy\n")

results = DataFrame(
    N_ATOMS   = Int[],
    M_ATOMS   = Int[],
    S4_NV_AOS = Float64[],
    S4_NV_SOA = Float64[],
    S8_NV_AOS = Float64[],
    S8_NV_SOA = Float64[],
    S4_V_AOS  = Float64[],
    S4_V_SOA  = Float64[],
    S8_V_AOS  = Float64[],
    S8_V_SOA  = Float64[],
    CUDA      = Float64[],)

println("Starting trials ...")

# trial(:(serial(state_aos, cut_off, LennardJones.NOFORCES)), "Serial - No Verlet - AoS")
# trial(:(serial(state_soa, cut_off, LennardJones.NOFORCES)), "Serial - No Verlet - SoA")
# trial(:(serial(state_aos, vlist, LennardJones.NOFORCES)), "Serial - Verlet - AoS")
# trial(:(serial(state_soa, vlist, LennardJones.NOFORCES)), "Serial - Verlet - SoA")
t1 = trial(:(simd(state_aos, cut_off, LennardJones.NOFORCES, Vec{4, T})), "SIMD_4 - No Verlet - AoS")
t2 = trial(:(simd(state_soa, cut_off, LennardJones.NOFORCES, Vec{4, T})), "SIMD_4 - No Verlet - SoA")
t3 = trial(:(simd(state_aos, cut_off, LennardJones.NOFORCES, Vec{8, T})), "SIMD_8 - No Verlet - AoS")
t4 = trial(:(simd(state_soa, cut_off, LennardJones.NOFORCES, Vec{8, T})), "SIMD_8 - No Verlet - SoA")
t5 = trial(:(simd(state_aos, vlist, LennardJones.NOFORCES, Vec{4, T})), "SIMD_4 - Verlet - AoS")
t6 = trial(:(simd(state_soa, vlist, LennardJones.NOFORCES, Vec{4, T})), "SIMD_4 - Verlet - SoA")
t7 = trial(:(simd(state_aos, vlist, LennardJones.NOFORCES, Vec{8, T})), "SIMD_8 - Verlet - AoS")
t8 = trial(:(simd(state_soa, vlist, LennardJones.NOFORCES, Vec{8, T})), "SIMD_8 - Verlet - SoA")
t9 = trial(:(cuda(state_aos, cut_off)), "CUDA - SoA")

push!(results, (n_atoms, m_atoms, t1, t2, t3, t4, t5, t6, t7, t8, t9))

display(results)


text:
#==============================================================================#
#                                Blackman Window                               #
#==============================================================================#

function blackman( n::Integer, numTaps::Integer )
    # TODO: add argument valdiation
    α = 0.42
    β = 0.5
    α - β*cos(2*π*n/(numTaps-1)) + (β-α)*cos(4*π*n/(numTaps-1))
end

function blackman( numTaps::Integer )
    # TODO: add argument valdiation
    [ blackman(n, numTaps) for n = 0:(numTaps-1) ]
end

#==============================================================================#
#                                Hamming Window                                #
#==============================================================================#

function hamming( n::Integer, numTaps::Integer )
    # TODO: add argument valdiation
    α = 0.54
    β = 0.46 # 1 - α
    α - β*cos(2*π*n/(numTaps-1))
end

function hamming( numTaps::Integer )
    # TODO: add argument valdiation
    [ hamming(n, numTaps) for n = 0:(numTaps-1) ]
end

#==============================================================================#
#                                  Hann Window                                 #
#==============================================================================#

function hann( n::Integer, numTaps::Integer )
    # TODO: add argument valdiation
    α = 0.5
    β = 0.5
    α - β*cos(2*π*n/(numTaps-1))
end

function hann( numTaps::Integer )
    # TODO: add argument valdiation
    [ hann(n, numTaps) for n = 0:(numTaps-1) ]
end

#==============================================================================#
#                                Kaiser Window                                 #
#==============================================================================#

function kaiser( n::Integer, numTaps::Integer, β::Real )
    # TODO: add argument valdiation
    num = besseli(0, β*sqrt(1-(2*n/(numTaps-1)-1).^2))
    dem = besseli(0, β)
    num/dem
end

function kaiser( numTaps::Integer, β::Real )
    # TODO: add argument valdiation
    [ kaiser( n, numTaps, β ) for n = 0:(numTaps-1) ]
end


text:
for N in [Float64, Rational{Int}, Float32]
    # ==================================
    # Constructor and interface methods
    # ==================================

    B = BallInf(zeros(N, 3), N(1))
    v = N[1, 0, 0] # translation along dimension 1
    tr = Translation(B, v)

    # test alternative constructors
    @test tr == B + v == B ⊕ v
    # check that translation from the left works as well
    @test tr == v + B == v ⊕ B

    # dimension check
    @test_throws AssertionError Translation(B, N[0, 0])

    @test dim(tr) == 3

    # support vector
    @test σ(N[1, 0, 0], tr) == N[2, 1, 1]

    # support function
    @test ρ(N[1, 0, 0], tr) == N(2)

    # boundedness
    @test isbounded(tr)
    @test !isbounded(Translation(Universe{N}(3), v))

    if N ∈ [Float64]
        @test !isbounded(Translation(HPolyhedron([HalfSpace(N[1, 1], N(1))]), N[0, 0]))
    end

    # function to get an element
    @test an_element(tr) ∈ tr.X

    # emptiness check
    @test !isempty(tr)

    # set membership
    B = Ball1(zeros(N, 2), N(1))
    @test center(B) ∈ B ⊕ zeros(N, 2)

    # concrete linear map
    B = BallInf(ones(N, 2), N(1))
    v = N[1, 0]
    tr = B ⊕ v
    clm = linear_map(N[1 2; 0 1], tr)
    @test ispermutation(vertices_list(clm), [N[5, 2], [1, 0], [3, 0], [7, 2]])

    # concretize
    @test concretize(tr) == translate(B, v)

    # ==================================
    # Type-specific methods
    # ==================================

    # the translation is the origin => constraints remain unchanged
    constraints_list(B ⊕ zeros(N, 2)) == constraints_list(tr.X)

    # one-dimensional case: translation to the left
    Tleft = Translation(HalfSpace(N[1], N(1)), [N(-1)])
    @test constraints_list(Tleft)[1] == HalfSpace(N[1], N(0))

    # one-dimensional case: translation to the right
    Tright = Translation(HalfSpace(N[1], N(1)), [N(1)])
    @test constraints_list(Tright)[1] == HalfSpace(N[1], N(2))

    # the translation of a lazy linear map returns an affine map
    M = N[1 0; 0 2]; B = BallInf(zeros(N, 2), N(1)); v = N[1, 0]
    tr = M * B ⊕ v
    @test tr isa AffineMap && tr.M == M && tr.X == B && tr.v == v
end

for N in [Float64, Float32]
    # translation of a set not represented by a finite number of constraints
    tr = Ball2(zeros(N, 2), N(1)) ⊕ N[1, 0]
    @test ρ(N[1, 0], tr) == N(2)
    @test_throws AssertionError constraints_list(tr)
end


text:
"""
    mbimport(cmd0::String="", arg1=nothing, arg2=nothing, arg3=nothing; kwargs...)

Produces a gray-shaded (or colored) map by plotting rectangles centered on each grid node and assigning them a gray-shade (or color) based on the z-value.

Parameters
----------

- **A** | **footprint** :: [Type => Str | Tuple]	``Arg = factor/mode/depth``

    Determines how the along-track dimension of the beam or pixel footprints is calculated.
- $(GMT.opt_J)
- $(GMT.opt_C)
- **D** | **scaling** :: [Type => Str | Tuple]	``Arg = mode/scale/min/max``

    Sets scaling of beam amplitude or sidescan pixel values which can be applied before plotting.
- **E** | **dpi** :: [Type => Int]

    Sets the resolution of the projected image that will be created.
- **G** | **bit_color** :: [Type => Str | Tuple]	``Arg = magnitude/azimuth or magnitude/median``

    Sets the parameters controlng simulated illumination of bathymetry.
- **S** | **speed** :: [Type => Number]

    Sets the minimum speed in km/hr (5.5 kts ~ 10 km/hr) allowed in the input data.
- **T** | **timegap** :: [Type => number]

    Sets the maximum time gap in minutes between adjacent pings before being considered a gap.
- **Z** | **type_plot** :: [Type => Str | Number]

    Sets the style of the plot.
- $(GMT.opt_R)
- $(GMT.opt_U)
- $(GMT.opt_V)
- $(GMT.opt_X)
- $(GMT.opt_Y)
- $(GMT.opt_n)
- $(GMT.opt_t)
"""
function mbimport(cmd0::String=""; kwargs...)

	length(kwargs) == 0 && occursin(" -", cmd0) && return monolitic("mbimport", cmd0)

	d = init_module(false, kwargs...)[1]		# Also checks if the user wants ONLY the HELP mode

	cmd, opt_R = parse_R(d, "")
	cmd, opt_J = parse_J(d, cmd, " -JX12cd/0", true)
	cmd, = parse_common_opts(d, cmd, [:UVXY :params :n :t], true)

	cmd  = parse_these_opts(cmd, d, [[:A :footprint], [:D :scaling], [:F :format], [:E :dpi], [:G :shade],
	                                 [:S :speed], [:T :timegap], [:b :star_time], [:e :end_time]])
	cmd = add_opt(d, cmd, 'Z', [:Z :type_plot], (bat="_1", shaded_bat="_2", shaded_amp="_3", amp="_4", sscan="_5"))

	cmd = add_opt(d, cmd, "%", [:layout :mem_layout], nothing)
	cmd, arg1, = add_opt_cpt(d, cmd, CPTaliases, 'C', 0, nothing)
	N_args = (arg1 === nothing) ? 0 : 1
	cmd, arg1, arg2, = add_opt_cpt(d, cmd, [:N :color_amp], 'N', N_args, arg1, nothing)

	cmd = "mbimport -I" * cmd0 * cmd				# In any case we need this
	finish_PS_module(d, cmd, "", true, false, false, arg1, arg2)
end


text:
type TableWidget
    table
    header
    widget
    maxrows
    n
    expfmt
end

function TableWidget(ncols; header=nothing, fixed_width=0, maxrows=10, n=2, expfmt = latex_expfmt)
    if fixed_width != 0 && !is(header,nothing)
        header = map(x -> lpad(x,fixed_width), header)
    end

    TableWidget(cell(0,ncols),header,TextWidget(""),maxrows,n,expfmt)
end

function Base.push!(tlw::TableWidget,row)
    tlw.table = vcat(tlw.table,row')
    update_display(tlw)
end

function update_display(tlw::TableWidget)
    io = IOBuffer()
    nrows = size(tlw.table,1)
    start = tlw.maxrows == 0 ? 1 : max(1, nrows-tlw.maxrows)
    writetable(io, tlw.table[start:end,:], header=tlw.header, n=tlw.n, expfmt=tlw.expfmt)
    s = takebuf_string(io)
    update(tlw.widget,s)
end

function Base.writemime(io, mime::MIME"text/html", tlw::TableWidget)
    writemime(io,mime,tlw.widget)
end


text:
"""
`mush(A,B)` takes two matrices of the same size and creates a new
matrix whose `i,j`-entry is `(A[i,j],B[i,j])`.
"""
function mush(A::Array{T,2}, B::Array{S,2}) where {S,T}
    TS = Tuple{S,T}
    r, c = size(A)
    if size(A) != size(B)
        error("Matrices must have same shape")
    end

    C = Array{TS,2}(undef, r, c)
    for i = 1:r
        for j = 1:c
            C[i, j] = (A[i, j], B[i, j])
        end
    end
    return C
end


text:
# ## model
include("trajopt_model.jl")

# ## horizon

freq = 100
h = 1.0 / freq
T = 59 #Int(floor(0.65 / h)) +

# ## point_foot_quadruped
s = get_simulation("point_foot_quadruped", "flat_3D_lc", "flat")
model = s.model
env = s.env

nx = 2 * model.nq
nc = 4 #model.nc
nu = model.nu + nc + 4 * nc + nc + 4 * nc + 1
nθ = 5

# ## model
d1 = DTO.Dynamics((y, x, u, w) -> point_foot_quadruped_dyn1(
    model, env, [h], y, x, u, w), nx + nθ + nx + model.nu, nx, nu)
dt = DTO.Dynamics((y, x, u, w) -> point_foot_quadruped_dynt(
    model, env, [h], y, x, u, w), nx + nθ + nx + model.nu, nx + nθ + nx + model.nu, nu)

dyn = [d1, [dt for t = 2:T-1]...]

# ## initial conditions
mode = :left
body_height = 0.3
foot_x = 0.17
foot_y = 0.15
foot_height = 0.08
foot_mid_height = 0.12

q1 = zeros(model.nq)

# body position
q1[1:3] = [0.0; 0.0; body_height]
q1[4:6] = [0.0; 0.0; 0.0]
# foot1
q1[7:9] = [foot_x; foot_y; 0]
# foot2
q1[10:12] = [foot_x; -foot_y; 0]
# foot3
q1[13:15] = [-foot_x; foot_y; 0]
# foot4
q1[16:18] = [-foot_x; -foot_y; 0]

q2 = deepcopy(q1)
q2[6 + 3] += foot_height # front left
q2[9 + 3] += foot_height # front right

q3 = deepcopy(q1)
q3[6 + 3] += foot_mid_height # front left
q3[9 + 3] += foot_mid_height # front right
q3[12 + 3] += foot_mid_height # back left
q3[15 + 3] += foot_mid_height # back right

qT = copy(q1)

# visualize!(vis, model, [qM], Δt=h);

T1 = 1
T2 = 20
T3 = 40
q_ref = [
    linear_interpolation(q1, q2, 20)...,
    linear_interpolation(q2, q3, 20)...,
    linear_interpolation(q3, qT, 20)...]
x_ref = [[q_ref[t]; q_ref[t+1]] for t = 1:T]
x1 = x_ref[1]
x2 = x_ref[T2]
x3 = x_ref[T3]
xT = x_ref[T]

visualize!(vis, model, q_ref, Δt=h);

# ## objective
obj = DTO.Cost{Float64}[]
for t = 1:T
    if t == T
        function objT(x, u, w)
            J = 0.0
            v = (x[model.nq .+ (1:model.nq)] - x[1:model.nq]) ./ h
            J += 0.5 * 1.0e-1 * dot(v, v)
            J += 0.5 * transpose(x[1:nx] - x_ref[t]) * Diagonal(1000.0 * ones(nx)) * (x[1:nx] - x_ref[t])
            return J
        end
        push!(obj, DTO.Cost(objT, nx + nθ + nx + model.nu, 0))
    elseif t == 1
        function obj1(x, u, w)
            J = 0.0
            v = (x[model.nq .+ (1:model.nq)] - x[1:model.nq]) ./ h
            J += 0.5 * 1.0e-3 * dot(v, v)
            J += 0.5 * transpose(x[1:nx] - x_ref[t]) * Diagonal(1000.0 * ones(nx)) * (x[1:nx] - x_ref[t])
            J += 0.5 * transpose(u[1:model.nu]) * Diagonal(1.0e-3 * ones(model.nu)) * u[1:model.nu]
            J += 1000.0 * u[end] # slack
            return J
        end
        push!(obj, DTO.Cost(obj1, nx, nu))
    else
        function objt(x, u, w)
            J = 0.0
            v = (x[model.nq .+ (1:model.nq)] - x[1:model.nq]) ./ h
            J += 0.5 * 1.0e-3 * dot(v, v)
            u_previous = x[nx + 5 + nx .+ (1:model.nu)]
            u_control = u[1:model.nu]
            w = (u_control - u_previous) ./ h
            J += 0.5 * 1.0e-2 * dot(w, w)
            J += 0.5 * transpose(x[1:nx] - x_ref[t]) * Diagonal(1000.0 * ones(nx)) * (x[1:nx] - x_ref[t])
            J += 0.5 * transpose(u[1:model.nu]) * Diagonal(1.0e-3 * ones(model.nu)) * u[1:model.nu]
            J += 1000.0 * u[end] # slack
            return J
        end
        push!(obj, DTO.Cost(objt, nx + nθ + nx + model.nu, nu))
    end
end

# ## constraints
# initial condition
xl1 = [q1; q1]
xu1 = [q1; q1]
xlt = [-Inf * ones(nx); -Inf * ones(nθ); -Inf * ones(nx); -Inf * ones(model.nu)]
xut = [Inf * ones(nx); Inf * ones(nθ); Inf * ones(nx); Inf * ones(model.nu)]

# final condition
xlT = [qT; qT; -Inf * ones(nθ); -Inf * ones(nx); -Inf * ones(model.nu)]
xuT = [qT; qT; Inf * ones(nθ); Inf * ones(nx); Inf * ones(model.nu)]

ul = [-Inf * ones(model.nu); zeros(nu - model.nu)]
uu = [Inf * ones(model.nu); Inf * ones(nu - model.nu)]

bnd1 = DTO.Bound(nx, nu, state_lower=xl1, state_upper=xu1, action_lower=ul, action_upper=uu)
bndt = DTO.Bound(nx + nθ + nx + model.nu, nu, state_lower=xlt, state_upper=xut, action_lower=ul, action_upper=uu)
bndT = DTO.Bound(nx + nθ + nx + model.nu, 0, state_lower=xlT, state_upper=xuT)
bnds = [bnd1, [bndt for t = 2:T-1]..., bndT];

cons = DTO.Constraint{Float64}[]
for t = 1:T
    if t == 1
        function constraints_1(x, u, w)
            [
            # equality (16)
            contact_constraints_equality(model, env, h, x, u, w);
            # inequality (28)
            contact_constraints_inequality_1(model, env, h, x, u, w);

            # body/feet constraints
            # x[3] - x_ref[t][3]; # body height
            # x[model.nq + 3] - x_ref[t][model.nq + 3]; # body height
            ]
        end
        push!(cons, DTO.Constraint(constraints_1, nx, nu, idx_ineq=collect(16 .+ (1:28))))
    elseif t == T
        function constraints_T(x, u, w)
            [
            # inequality (8)
            contact_constraints_inequality_T(model, env, h, x, u, w);

            # body/feet constraints
            # x[3] - x_ref[t][3]; # body height
            # x[model.nq + 3] - x_ref[t][model.nq + 3]; # body height
            ]
        end
        push!(cons, DTO.Constraint(constraints_T, nx + nθ + nx + model.nu, nu, idx_ineq=collect(0 .+ (1:8))))
    else
        function constraints_t(x, u, w)
            [
            # equality (16)
            contact_constraints_equality(model, env, h, x, u, w);
            # inequality (32)
            contact_constraints_inequality_t(model, env, h, x, u, w);

            # body/feet constraints
            # x[3] - x_ref[t][3]; # body height
            # x[model.nq + 3] - x_ref[t][model.nq + 3]; # body height
            ]
        end
        push!(cons, DTO.Constraint(constraints_t, nx + nθ + nx + model.nu, nu, idx_ineq=collect(16 .+ (1:32))) )
    end
end

# ## problem
tolerance = 1.0e-3
p = DTO.solver(dyn, obj, cons, bnds,
    options=DTO.Options(
        max_iter=2000,
        tol=tolerance,
        constr_viol_tol=tolerance,
        ))

# ## initialize
x_interpolation = [x_ref[1], [[x_ref[t]; zeros(nθ); zeros(nx); zeros(model.nu)] for t = 2:T]...]
u_guess = [1.0e-4 * rand(nu) for t = 1:T-1] # may need to run more than once to get good trajectory
DTO.initialize_states!(p, x_interpolation)
DTO.initialize_controls!(p, u_guess)

# ## solve
@time DTO.solve!(p)

# ## solution
x_sol, u_sol = DTO.get_trajectory(p)
@show x_sol[1]
@show x_sol[T]
sum([u[end] for u in u_sol[1:end-1]])

# ## visualize
vis = Visualizer()
render(vis)
visualize!(vis, model, x_sol, Δt=h);

q_opt = [x_sol[1][1:model.nq], [x[model.nq .+ (1:model.nq)] for x in x_sol]...]
v_opt = [(x[model.nq .+ (1:model.nq)] - x[0 .+ (1:model.nq)]) ./ h for x in x_sol]
u_opt = [u[1:model.nu] for u in u_sol]
λ_opt = [u[model.nu .+ (1:4)] for u in u_sol]
b_opt = [u[model.nu + 4 .+ (1:16)] for u in u_sol]

q_opt = [x_sol[1][1:model.nq], [x[model.nq .+ (1:model.nq)] for x in x_sol]...]
v_opt = [(x[model.nq .+ (1:model.nq)] - x[0 .+ (1:model.nq)]) ./ h for x in x_sol]
u_opt = [u[1:model.nu] for u in u_sol]
γ_opt = [u[model.nu .+ (1:4)] for u in u_sol]
b_opt = [u[model.nu + 4 .+ (1:16)] for u in u_sol]
ψ_opt = [u[model.nu + 4 + 16 .+ (1:4)] for u in u_sol]
η_opt = [u[model.nu + 4 + 16 + 4 .+ (1:16)] for u in u_sol]

qm = q_opt
vm = v_opt
um = u_opt
γm = γ_opt
bm = b_opt
ψm = ψ_opt
ηm = η_opt

# using JLD2
# @save joinpath(@__DIR__, "inplace_bounce_right.jld2") qr ur γr br ψr ηr
# @load joinpath(@__DIR__, "inplace_bounce_right.jld2") qr ur γr br ψr ηr

# qm = [qr[2:end]..., ql[2:end]...]
# um = [ur..., ul...]
# γm = [γr..., γl...]
# bm = [br..., bl...]
# ψm = [ψr..., ψl...]
# ηm = [ηr..., ηl...]
μm = model.μ_world
hm = h
timesteps = range(0.0, stop=(h * (length(qm) - 2)), length=(length(qm) - 2))
plot(timesteps, hcat(qm[2:end-1]...)', labels="")
plot(timesteps, hcat(um...)', labels="")
plot(timesteps, hcat(γm...)', labels="")
plot(timesteps, hcat(bm...)', labels="")
plot(timesteps, hcat(ψm...)', labels="")
plot(timesteps, hcat(ηm...)', labels="")

using JLD2
@save joinpath(@__DIR__, "inplace_bounce.jld2") qm um γm bm ψm ηm μm hm
@load joinpath(@__DIR__, "inplace_bounce.jld2") qm um γm bm ψm ηm μm hm