"""
$(SIGNATURES)

Kinetic flux vector splitting (KFVS) flux
"""
function flux_kfvs!(
    face::Interface1F,
    ctrL::T,
    ctrR::T,
    gas::AbstractProperty,
    vs::AbstractVelocitySpace1D,
    p,
    dt = 1.0,
) where {T<:ControlVolume1F}

    dxL, dxR = p[1:2]

    flux_kfvs!(
        face.fw,
        face.ff,
        ctrL.f + ctrL.sf * dxL,
        ctrR.f - ctrR.sf * dxR,
        vs.u,
        vs.weights,
        dt,
        ctrL.sf,
        ctrR.sf,
    )

    return nothing

end

function flux_kfvs!(
    face::Interface2F,
    ctrL::T,
    ctrR::T,
    gas::AbstractProperty,
    vs::AbstractVelocitySpace1D,
    p,
    dt = 1.0,
) where {T<:ControlVolume2F}

    dxL, dxR = p[1:2]

    flux_kfvs!(
        face.fw,
        face.fh,
        face.fb,
        ctrL.h .+ ctrL.sb .* dxL,
        ctrL.b .+ ctrL.sb .* dxL,
        ctrR.h .- ctrR.sb .* dxR,
        ctrR.b .- ctrR.sb .* dxR,
        vs.u,
        vs.weights,
        dt,
        ctrL.sh,
        ctrL.sb,
        ctrR.sh,
        ctrR.sb,
    )

    return nothing

end

function flux_kfvs!(
    face::Interface1F,
    ctrL::T,
    ctrR::T,
    gas::AbstractProperty,
    vs::AbstractVelocitySpace2D,
    p,
    dt = 1.0,
) where {T<:ControlVolume1F}

    dxL, dxR, len, n, dirc = p[1:5]
    sfL = @view ctrL.sf[:, :, dirc]
    sfR = @view ctrR.sf[:, :, dirc]

    flux_kfvs!(
        face.fw,
        face.ff,
        ctrL.f + sfL * dxL,
        ctrR.f - sfR * dxR,
        vs.u .* n[1] .+ vs.v .* n[2],
        vs.v .* n[1] .- vs.u .* n[2],
        vs.weights,
        dt,
        len,
        sfL,
        sfR,
    )

    face.fw .= global_frame(face.fw, n)

    return nothing

end

function flux_kfvs!(
    face::Interface2F,
    ctrL::T,
    ctrR::T,
    gas::AbstractProperty,
    vs::AbstractVelocitySpace2D,
    p,
    dt = 1.0,
) where {T<:ControlVolume2F}

    dxL, dxR, len, n, dirc = p[1:5]
    shL = @view ctrL.sh[:, :, dirc]
    sbL = @view ctrL.sb[:, :, dirc]
    shR = @view ctrR.sh[:, :, dirc]
    sbR = @view ctrR.sb[:, :, dirc]

    flux_kfvs!(
        face.fw,
        face.fh,
        face.fb,
        ctrL.h .+ shL .* dxL,
        ctrL.b .+ sbL .* dxL,
        ctrR.h .- shR .* dxR,
        ctrR.b .- sbR .* dxR,
        vs.u .* n[1] .+ vs.v .* n[2],
        vs.v .* n[1] .- vs.u .* n[2],
        vs.weights,
        dt,
        len,
        shL,
        sbL,
        shR,
        sbR,
    )

    face.fw .= global_frame(face.fw, n)

    return nothing

end

function flux_kfvs!(
    face::Interface1F,
    ctrL::T,
    ctrR::T,
    gas::AbstractProperty,
    vs::AbstractVelocitySpace3D,
    p,
    dt = 1.0,
) where {T<:ControlVolume1F}

    dxL, dxR = p[1:2]

    if length(p) == 2 || p[3] == 1
        flux_kfvs!(
            face.fw,
            face.ff,
            ctrL.f + ctrL.sf * dxL,
            ctrR.f - ctrR.sf * dxR,
            vs.u,
            vs.v,
            vs.w,
            vs.weights,
            dt,
            1.0,
            ctrL.sf,
            ctrR.sf,
        )
    else
        len, n, dirc = p[3:5]
        sfL = @view ctrL.sf[:, :, :, dirc]
        sfR = @view ctrR.sf[:, :, :, dirc]

        flux_kfvs!(
            face.fw,
            face.ff,
            ctrL.f + sfL * dxL,
            ctrR.f - sfR * dxR,
            vs.u .* n[1] .+ vs.v .* n[2],
            vs.v .* n[1] .- vs.u .* n[2],
            vs.w,
            vs.weights,
            dt,
            len,
            sfL,
            sfR,
        )
        face.fw .= global_frame(face.fw, n[1], n[2])
    end

    return nothing

end

# ------------------------------------------------------------
# Low-level backends
# ------------------------------------------------------------

"""
$(SIGNATURES)

Kinetic flux vector splitting (KFVS) flux
"""
function flux_kfvs(
    fL::Y,
    fR::Y,
    u::AV{<:FN},
    dt,
    sfL = zero(fL)::Y,
    sfR = zero(fR)::Y,
) where {Y<:AV{<:FN}}

    ff = similar(fL)
    flux_kfvs!(ff, fL, fR, u, dt, sfL, sfR)

    return ff

end

"""
$(SIGNATURES)

Kinetic flux vector splitting (KFVS) flux

- @args: particle distribution functions and their left/right slopes
- @args: particle velocity quadrature points and weights
- @args: time step and cell size

1F1V for pure DOM
"""
function flux_kfvs!(
    ff::AV{<:FN},
    fL::Y,
    fR::Y,
    u::AV{<:FN},
    dt,
    sfL = zero(fL)::Y,
    sfR = zero(fR)::Y,
) where {Y<:AV{<:FN}} # 1F1V flux for pure DOM

    # upwind reconstruction
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # calculate flux
    @. ff = dt * u * f - 0.5 * dt^2 * u^2 * sf

    return nothing

end

"""
$(SIGNATURES)

1F1V
"""
function flux_kfvs!(
    fw::AV{<:FN},
    ff::AV{<:FN},
    fL::Z,
    fR::Z,
    u::A,
    ω::A,
    dt,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {Z<:AV{<:FN},A<:AV{<:FN}}

    # upwind reconstruction
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # calculate fluxes
    @. ff = dt * u * f - 0.5 * dt^2 * u^2 * sf
    flux_conserve!(fw, ff, u, ω)
    #fw[1] = dt * sum(ω .* u .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sf)
    #fw[2] = dt * sum(ω .* u .^ 2 .* f) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sf)
    #fw[3] = dt * 0.5 * sum(ω .* u .^ 3 .* f) - 0.5 * dt^2 * 0.5 * sum(ω .* u .^ 4 .* sf)

    return nothing

end

"""
$(SIGNATURES)

Mixture
"""
function flux_kfvs!(
    fw::AM{<:FN},
    ff::AM{<:FN},
    fL::Z,
    fR::Z,
    u::A,
    ω::A,
    dt,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {Z<:AM{<:FN},A<:AM{<:FN}}

    for j in axes(fw, 2)
        _fw = @view fw[:, j]
        _ff = @view ff[:, j]
        _fL = @view fL[:, j]
        _fR = @view fR[:, j]
        _u = @view u[:, j]
        _ω = @view ω[:, j]
        _sfL = @view sfL[:, j]
        _sfR = @view sfR[:, j]

        flux_kfvs!(_fw, _ff, _fL, _fR, _u, _ω, dt, _sfL, _sfR)
    end

    return nothing

end

"""
$(SIGNATURES)

2F1V
"""
function flux_kfvs!(
    fw::AV{<:FN},
    fh::Y,
    fb::Y,
    hL::Z,
    bL::Z,
    hR::Z,
    bR::Z,
    u::A,
    ω::A,
    dt,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
) where {Y<:AV{<:FN},Z<:AV{<:FN},A<:AV{<:FN}}

    # upwind reconstruction
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)

    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)

    # calculate fluxes
    @. fh = dt * u * h - 0.5 * dt^2 * u^2 * sh
    @. fb = dt * u * b - 0.5 * dt^2 * u^2 * sb
    flux_conserve!(fw, fh, fb, u, ω)
    #=fw[1] = dt * sum(ω .* u .* h) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh)
    fw[2] = dt * sum(ω .* u .^ 2 .* h) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh)
    fw[3] =
        dt * 0.5 * (sum(ω .* u .^ 3 .* h) + sum(ω .* u .* b)) -
        0.5 * dt^2 * 0.5 * (sum(ω .* u .^ 4 .* sh) + sum(ω .* u .^ 2 .* sb))=#

    return nothing

end

"""
$(SIGNATURES)

Mixture
"""
function flux_kfvs!(
    fw::AM{<:FN},
    fh::Y,
    fb::Y,
    hL::Z,
    bL::Z,
    hR::Z,
    bR::Z,
    u::A,
    ω::A,
    dt,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
) where {Y<:AM{<:FN},Z<:AM{<:FN},A<:AM{<:FN}}

    for j in axes(fw, 2)
        _fw = @view fw[:, j]
        _fh = @view fh[:, j]
        _fb = @view fb[:, j]
        _hL = @view hL[:, j]
        _bL = @view bL[:, j]
        _hR = @view hR[:, j]
        _bR = @view bR[:, j]
        _u = @view u[:, j]
        _ω = @view ω[:, j]
        _shL = @view shL[:, j]
        _sbL = @view sbL[:, j]
        _shR = @view shR[:, j]
        _sbR = @view sbR[:, j]

        flux_kfvs!(_fw, _fh, _fb, _hL, _bL, _hR, _bR, _u, _ω, dt, _shL, _sbL, _shR, _sbR)
    end

    return nothing

end

"""
$(SIGNATURES)

3F1V @ Rykov
"""
function flux_kfvs!(
    fw::AV{<:FN},
    fh::Y,
    fb::Y,
    fr::Y,
    hL::Z,
    bL::Z,
    rL::Z,
    hR::Z,
    bR::Z,
    rR::Z,
    u::A,
    ω::A,
    dt,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    srL = zero(rL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
    srR = zero(rR)::Z,
) where {Y<:AV{<:FN},Z<:AV{<:FN},A<:AV{<:FN}}

    # upwind reconstruction
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)
    r = @. rL * δ + rR * (1.0 - δ)

    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)
    sr = @. srL * δ + srR * (1.0 - δ)

    # macro fluxes
    fw[1] = dt * sum(ω .* u .* h) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh)
    fw[2] = dt * sum(ω .* u .^ 2 .* h) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh)
    fw[3] =
        dt * 0.5 * (sum(ω .* u .^ 3 .* h) + sum(ω .* u .* b)) -
        0.5 * dt^2 * 0.5 * (sum(ω .* u .^ 4 .* sh) + sum(ω .* u .^ 2 .* sb))
    fw[4] = dt * 0.5 * sum(ω .* u .* r) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sr)

    # micro fluxes
    @. fh = dt * u * h - 0.5 * dt^2 * u^2 * sh
    @. fb = dt * u * b - 0.5 * dt^2 * u^2 * sb
    @. fr = dt * u * r - 0.5 * dt^2 * u^2 * sr

    return nothing

end

"""
$(SIGNATURES)

1F3V
"""
function flux_kfvs!(
    fw::AV{<:FN},
    ff::AA{<:FN,3},
    fL::Z,
    fR::Z,
    u::A,
    v::A,
    w::A,
    ω::A,
    dt::Real,
    len = 1.0::Real,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {Z<:AA{<:FN,3},A<:AA{<:FN,3}}

    # upwind reconstruction
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # calculate fluxes
    @. ff = (dt * u * f - 0.5 * dt^2 * u^2 * sf) * len
    flux_conserve!(fw, ff, u, v, w, ω)
    #=fw[1] = dt * sum(ω .* u .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sf)
    fw[2] = dt * sum(ω .* u .^ 2 .* f) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sf)
    fw[3] = dt * sum(ω .* u .* v .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* v .* sf)
    fw[4] = dt * sum(ω .* u .* w .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* w .* sf)
    fw[5] =
        dt * 0.5 * sum(ω .* u .* (u .^ 2 .+ v .^ 2 .+ w .^ 2) .* f) -
        0.5 * dt^2 * 0.5 * sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2 .+ w .^ 2) .* sf)
    @. fw *= len=#

    return nothing

end

"""
$(SIGNATURES)

4F1V
"""
function flux_kfvs!(
    fw::AV{<:FN},
    fh0::Y,
    fh1::Y,
    fh2::Y,
    fh3::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h3L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    h3R::Z,
    u::A,
    ω::A,
    dt,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh3L = zero(h3L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
    sh3R = zero(h3R)::Z,
) where {Y<:AV{<:FN},Z<:AV{<:FN},A<:AV{<:FN}}

    # upwind reconstruction
    δ = heaviside.(u)

    h0 = @. h0L * δ + h0R * (1.0 - δ)
    h1 = @. h1L * δ + h1R * (1.0 - δ)
    h2 = @. h2L * δ + h2R * (1.0 - δ)
    h3 = @. h3L * δ + h3R * (1.0 - δ)

    sh0 = @. sh0L * δ + sh0R * (1.0 - δ)
    sh1 = @. sh1L * δ + sh1R * (1.0 - δ)
    sh2 = @. sh2L * δ + sh2R * (1.0 - δ)
    sh3 = @. sh3L * δ + sh3R * (1.0 - δ)

    # calculate fluxes
    fw[1] = dt * sum(ω .* u .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh0)
    fw[2] = dt * sum(ω .* u .^ 2 .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh0)
    fw[3] = dt * sum(ω .* u .* h1) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh1)
    fw[4] = dt * sum(ω .* u .* h2) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh2)
    fw[5] =
        dt * 0.5 * (sum(ω .* u .^ 3 .* h0) + sum(ω .* u .* h3)) -
        0.5 * dt^2 * 0.5 * (sum(ω .* u .^ 4 .* sh0) + sum(ω .* u .^ 2 .* sh3))

    @. fh0 = dt * u * h0 - 0.5 * dt^2 * u^2 * sh0
    @. fh1 = dt * u * h1 - 0.5 * dt^2 * u^2 * sh1
    @. fh2 = dt * u * h2 - 0.5 * dt^2 * u^2 * sh2
    @. fh3 = dt * u * h3 - 0.5 * dt^2 * u^2 * sh3

    return nothing

end

"""
$(SIGNATURES)

Mixture
"""
function flux_kfvs!(
    fw::AM{<:FN},
    fh0::Y,
    fh1::Y,
    fh2::Y,
    fh3::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h3L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    h3R::Z,
    u::A,
    ω::A,
    dt,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh3L = zero(h3L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
    sh3R = zero(h3R)::Z,
) where {Y<:AM{<:FN},Z<:AM{<:FN},A<:AM{<:FN}}

    for j in axes(fw, 2)
        _fw = @view fw[:, j]
        _fh0 = @view fh0[:, j]
        _fh1 = @view fh1[:, j]
        _fh2 = @view fh2[:, j]
        _fh3 = @view fh3[:, j]

        flux_kfvs!(
            _fw,
            _fh0,
            _fh1,
            _fh2,
            _fh3,
            h0L[:, j],
            h1L[:, j],
            h2L[:, j],
            h3L[:, j],
            h0R[:, j],
            h1R[:, j],
            h2R[:, j],
            h3R[:, j],
            u[:, j],
            ω[:, j],
            dt,
            sh0L[:, j],
            sh1L[:, j],
            sh2L[:, j],
            sh3L[:, j],
            sh0R[:, j],
            sh1R[:, j],
            sh2R[:, j],
            sh3R[:, j],
        )
    end

    return nothing

end

"""
$(SIGNATURES)

1F2V
"""
function flux_kfvs!(
    fw::AV{<:FN},
    ff::Union{AV{<:FN},AM{<:FN}},
    fL::Z,
    fR::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {Z<:Union{AV{<:FN},AM{<:FN}},A<:Union{AV{<:FN},AM{<:FN}}}

    #--- upwind reconstruction ---#
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    #--- calculate fluxes ---#
    @. ff = (dt * u * f - 0.5 * dt^2 * u^2 * sf) * len
    flux_conserve!(fw, ff, u, v, ω)
    #=fw[1] = dt * sum(ω .* u .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sf)
    fw[2] = dt * sum(ω .* u .^ 2 .* f) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sf)
    fw[3] = dt * sum(ω .* v .* u .* f) - 0.5 * dt^2 * sum(ω .* v .* u .^ 2 .* sf)
    fw[4] =
        dt * 0.5 * sum(ω .* u .* (u .^ 2 .+ v .^ 2) .* f) -
        0.5 * dt^2 * 0.5 * sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2) .* sf)
    fw .*= len=#

    return nothing

end

"""
$(SIGNATURES)

2F2V
"""
function flux_kfvs!(
    fw::AV{<:FN},
    fh::Y,
    fb::Y,
    hL::Z,
    bL::Z,
    hR::Z,
    bR::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
) where {
    Y<:Union{AV{<:FN},AM{<:FN}},
    Z<:Union{AV{<:FN},AM{<:FN}},
    A<:Union{AV{<:FN},AM{<:FN}},
}

    #--- upwind reconstruction ---#
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)
    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)

    #--- calculate fluxes ---#
    @. fh = (dt * u * h - 0.5 * dt^2 * u^2 * sh) * len
    @. fb = (dt * u * b - 0.5 * dt^2 * u^2 * sb) * len
    flux_conserve!(fw, fh, fb, u, v, ω)
    #=fw[1] = dt * sum(ω .* u .* h) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh)
    fw[2] = dt * sum(ω .* u .^ 2 .* h) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh)
    fw[3] = dt * sum(ω .* v .* u .* h) - 0.5 * dt^2 * sum(ω .* v .* u .^ 2 .* sh)
    fw[4] =
        dt * 0.5 * (sum(ω .* u .* (u .^ 2 .+ v .^ 2) .* h) + sum(ω .* u .* b)) -
        0.5 *
        dt^2 *
        0.5 *
        (sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2) .* sh) + sum(ω .* u .^ 2 .* sb))
    fw .*= len=#

    return nothing

end

"""
$(SIGNATURES)

3F2V
"""
function flux_kfvs!(
    fw::AV{<:FN},
    fh0::Y,
    fh1::Y,
    fh2::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
) where {Y<:AM{<:FN},Z<:AM{<:FN},A<:AM{<:FN}}

    #--- upwind reconstruction ---#
    δ = heaviside.(u)

    h0 = @. h0L * δ + h0R * (1.0 - δ)
    h1 = @. h1L * δ + h1R * (1.0 - δ)
    h2 = @. h2L * δ + h2R * (1.0 - δ)
    sh0 = @. sh0L * δ + sh0R * (1.0 - δ)
    sh1 = @. sh1L * δ + sh1R * (1.0 - δ)
    sh2 = @. sh2L * δ + sh2R * (1.0 - δ)

    #--- calculate fluxes ---#
    fw[1] = dt * sum(ω .* u .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh0)
    fw[2] = dt * sum(ω .* u .^ 2 .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh0)
    fw[3] = dt * sum(ω .* u .* v .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* v .* sh0)
    fw[4] = dt * sum(ω .* u .* h1) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh1)
    fw[5] =
        dt * 0.5 * (sum(ω .* u .* (u .^ 2 .+ v .^ 2) .* h0) + sum(ω .* u .* h2)) -
        0.5 *
        dt^2 *
        0.5 *
        (sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2) .* sh0) + sum(ω .* u .^ 2 .* sh2))

    fw .*= len
    @. fh0 = (dt * u * h0 - 0.5 * dt^2 * u^2 * sh0) * len
    @. fh1 = (dt * u * h1 - 0.5 * dt^2 * u^2 * sh1) * len
    @. fh2 = (dt * u * h2 - 0.5 * dt^2 * u^2 * sh2) * len

    return nothing

end

"""
$(SIGNATURES)

Mixture
"""
function flux_kfvs!(
    fw::AM{<:FN},
    fh0::Y,
    fh1::Y,
    fh2::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
) where {Y<:AA{<:FN,3},Z<:AA{<:FN,3},A<:AA{<:FN,3}}

    #--- reconstruct initial distribution ---#
    δ = heaviside.(u)

    h0 = @. h0L * δ + h0R * (1.0 - δ)
    h1 = @. h1L * δ + h1R * (1.0 - δ)
    h2 = @. h2L * δ + h2R * (1.0 - δ)

    sh0 = @. sh0L * δ + sh0R * (1.0 - δ)
    sh1 = @. sh1L * δ + sh1R * (1.0 - δ)
    sh2 = @. sh2L * δ + sh2R * (1.0 - δ)

    for j = 1:2
        fw[1, j] =
            dt * sum(ω[:, :, j] .* u[:, :, j] .* h0[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh0[:, :, j])
        fw[2, j] =
            dt * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* h0[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 3 .* sh0[:, :, j])
        fw[3, j] =
            dt * sum(ω[:, :, j] .* v[:, :, j] .* u[:, :, j] .* h0[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* v[:, :, j] .* sh0[:, :, j])
        fw[4, j] =
            dt * sum(ω[:, :, j] .* u[:, :, j] .* h1[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh1[:, :, j])
        fw[5, j] =
            dt *
            0.5 *
            (
                sum(
                    ω[:, :, j] .* u[:, :, j] .* (u[:, :, j] .^ 2 .+ v[:, :, j] .^ 2) .*
                    h0[:, :, j],
                ) + sum(ω[:, :, j] .* u[:, :, j] .* h2[:, :, j])
            ) -
            0.5 *
            dt^2 *
            0.5 *
            (
                sum(
                    ω[:, :, j] .* u[:, :, j] .^ 2 .* (u[:, :, j] .^ 2 .+ v[:, :, j] .^ 2) .*
                    sh0[:, :, j],
                ) + sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh2[:, :, j])
            )

        @. fh0[:, :, j] =
            dt * u[:, :, j] * h0[:, :, j] - 0.5 * dt^2 * u[:, :, j]^2 * sh0[:, :, j]
        @. fh1[:, :, j] =
            dt * u[:, :, j] * h1[:, :, j] - 0.5 * dt^2 * u[:, :, j]^2 * sh1[:, :, j]
        @. fh2[:, :, j] =
            dt * u[:, :, j] * h2[:, :, j] - 0.5 * dt^2 * u[:, :, j]^2 * sh2[:, :, j]
    end

    @. fw *= len
    @. fh0 *= len
    @. fh1 *= len
    @. fh2 *= len

    return nothing

end


"""
main : the code you run to double-encode... the "main loop" of optimization.
"""

module main

include("utils.jl")
include("types.jl")
include("lookup.jl")
include("std_setup.jl")
include("math.jl")
include("optimize.jl")
include("mrf.jl")

using ArgParse, Logging, JLD, StatsBase, Distributions, Random, Statistics

function convert_to_saveable(cur_pop)
	new_pop = types.SaveChrome[]

	for indiv in cur_pop
		push!(new_pop, types.SaveChrome(indiv.path, indiv.full_sequence,
										indiv.deg_nuc, indiv.deg_map, indiv.deg_trns, indiv.deg_trne, indiv.deg_d, indiv.deg_skip, indiv.deg_insert,
										indiv.mark_nuc, indiv.mark_map, indiv.mark_trns, indiv.mark_trne, indiv.mark_d, indiv.mark_skip, indiv.mark_insert,
										indiv.deg_prob, indiv.deg_base_E, indiv.deg_seq,
										indiv.mark_prob, indiv.mark_base_E, indiv.mark_seq, indiv.first_weight))
	end

	return new_pop
end

"""
set_up_and_optimize is the code that sets up appropriate information about genes we want to target for double-encoding and
orchestrates the initial HMM double-encoding solution and then iteratively improves on that with greedy pseudolikelihood improvements.
Code also logs information through Logging package.
Genes are often referred to as mark/deg. This is historic from times of thinking of "marker" genes and "designated essential gene".
At this point this distinction is not meaningful but is kept to avoid equally arbitrary x_name, y_name conventions.
Arguments:
	mark_name : gene ID for 'mark' gene. Needed as a key for looking up some values associated with genes in text files.
	deg_name : gene ID for 'deg' gene.
	mark/deg_grem : path to MRF parameter files, usually a JLD/CSV file.
	mark/deg_hmm : path to HMM files, usually .hmm files.
	pop_size : size of population, i.e. number of individual HMM solutions to greedily optimize.
	frame : p1/p2.
	max_iter : maximum number of iterations.
	X_range/Y_range: positions along gene to consider for double-encoding... useful if you want to restrict where double-encoding can occur.
						Ranges are defined in terms of the end of the sequence, so if you have a 70 aa sequence you want to start at position 80,
						set the range to be something like 150:151. Generally some buffer around this value (i.e. 150:160) is useful.
						Don't set, or set to false if you don't care where the genes overlap.
"""
function set_up_and_optimize(log_io, rand_barcode, out_path, mark_name, deg_name, mark_grem, deg_grem, mark_hmm, deg_hmm, pop_size, frame, max_iter, X_range=false, Y_range=false; rand_weights = false, actually_mrf = true)

	@debug("Beginning run.")
	@debug(Libc.strftime(time()))
	@debug(Libc.strftime(time()))

	println("The random barcode on this run is: $rand_barcode")
	@debug("The random barcode on this run is: $rand_barcode")

	do_cull = false #culling reduces number of sequences we optimize over time. Just a trick to save time if you want top-performers only.
	full_even_window = 20
	half_window = div(full_even_window, 2) #window statistics to see if scores still decreasing (sometimes used).
	last_few_sf = Float64[]
	last_few_cs = Float64[]

	#Looks up mean/std. dev of family pseudolikelihoods, pre-computed.
	mu_mark, sig_mark = lookup.mu_sig(deg_name, "psls/")
	mu_deg, sig_deg = lookup.mu_sig(mark_name, "psls/")

	@debug("Stat param vars (mark / deg):\t$(mu_mark)\t$(sig_mark)\t$(mu_deg)\t$(sig_deg)")

	if true #replacing try/catch.
		gen_samples = true #this is true if starting from scratch (general case), false if we have some candidates to optimize ahead of time.
		population = types.Chromosome[]
		local mark_grem_prot, deg_grem_prot
		if !gen_samples
			population = load("trpE_population.jld")["population"][1:100] #[1:100]
			mark_gremodel, deg_gremodel = std_setup.short_set_up(mark_name, deg_name, mark_grem, deg_grem, mark_hmm, deg_hmm, length(population), rand_barcode, frame)
			deg_nNodes, mark_nNodes = deg_gremodel.nNodes, mark_gremodel.nNodes
		else
			#Generate population of sampled hmm starting points.
			if X_range == false || Y_range == false #we require both to be there...
				@debug("Doing standard full set up...")
				mark_gremodel, deg_gremodel, population, mark_grem_prot, deg_grem_prot = std_setup.full_set_up(out_path, mark_name, mark_hmm, mark_grem, deg_name, deg_hmm, deg_grem, pop_size, 1200, 1200, rand_barcode, frame)
			else
				@debug("The x range is $X_range")
				@debug("The y range is $Y_range")
				mark_gremodel, deg_gremodel, population, mark_grem_prot, deg_grem_prot = std_setup.full_sample_set_up(out_path, mark_name, mark_hmm, mark_grem, deg_name, deg_hmm, deg_grem, pop_size, 1200, 1200, rand_barcode, X_range, Y_range, frame)
			end
			deg_nNodes, mark_nNodes = deg_gremodel.nNodes, mark_gremodel.nNodes
		end

		#Look through generated candidates, check their fitness values and their correctness.
		success_chrom, fdp, fmp, fitness_values, indie_deg_maps, indie_mark_maps, deg_ull, deg_pv_w1, deg_pv_w2, deg_prot_mat, mark_ull, mark_pv_w1, mark_pv_w2, mark_prot_mat = optimize.assess_founders(population, mark_gremodel.nNodes, mark_gremodel.w1, mark_gremodel.w2, deg_gremodel.nNodes, deg_gremodel.w1, deg_gremodel.w2)

		@debug("Num success... $(length(success_chrom))")

		if actually_mrf
			founding_fitness_values = fitness_values[1:end]

			mark_base_energy = mrf.basic_energy_calc(mark_grem_prot, mark_gremodel.w1, mark_gremodel.w2)
			deg_base_energy = mrf.basic_energy_calc(deg_grem_prot, deg_gremodel.w1, deg_gremodel.w2)

			#Now we load this info into an ExChrome type object.
			cur_pop = types.ExChrome[]
			for i in 1:length(success_chrom)
				old = success_chrom[i]
				if rand_weights
					new_chrom = types.ExChrome(old.path, old.full_sequence, old.deg_nuc, indie_deg_maps[i], old.deg_trns, old.deg_trne,
															old.deg_d, old.deg_skip, old.deg_insert, old.mark_nuc, indie_mark_maps[i], old.mark_trns, old.mark_trne,
															old.mark_d, old.mark_skip, old.mark_insert, fdp[i], deg_base_energy, deg_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(deg_prot_mat[i:i, 1:end]), deg_ull[i], deg_pv_w1[i], deg_pv_w2[i:i, 1:end], fmp[i], mark_base_energy,
															mark_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(mark_prot_mat[i:i, 1:end]), mark_ull[i], mark_pv_w1[i], mark_pv_w2[i:i, 1:end], rand())
				else
					new_chrom = types.ExChrome(old.path, old.full_sequence, old.deg_nuc, indie_deg_maps[i], old.deg_trns, old.deg_trne,
															old.deg_d, old.deg_skip, old.deg_insert, old.mark_nuc, indie_mark_maps[i], old.mark_trns, old.mark_trne,
															old.mark_d, old.mark_skip, old.mark_insert, fdp[i], deg_base_energy, deg_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(deg_prot_mat[i:i, 1:end]), deg_ull[i], deg_pv_w1[i], deg_pv_w2[i:i, 1:end], fmp[i], mark_base_energy,
															mark_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(mark_prot_mat[i:i, 1:end]), mark_ull[i], mark_pv_w1[i], mark_pv_w2[i:i, 1:end], 0.5)
				end
				push!(cur_pop, new_chrom)
			end

			#There are occasionally differences between HMM/MRF models of the proteins, mostly because HMMs are more flexible with insertions.
			#This generates mappings between positions in either data type.
			mark_hmm_to_grem_map = std_setup.get_explicit_mapping(mark_grem_prot, mark_hmm)
			deg_hmm_to_grem_map = std_setup.get_explicit_mapping(deg_grem_prot, deg_hmm)

			@debug("Hey the explicit mapping for mark_hmm_to_grem is $mark_hmm_to_grem_map")
			@debug("Hey the explicit mapping for deg_hmm_to_grem is $deg_hmm_to_grem_map")
			#So now everything is in an extended chromosome object.

			changed_seq = length(cur_pop)
			no_change_iters = 0
			iter = 0;

			#Let's keep track of optimization history.
			history_matrix = zeros(Float32, round(Int64, pop_size * 1.2), max_iter)

			println("Beginning long-range optimization.")
			@debug("About to start optimization.")
			@debug(Libc.strftime(time()))

			stop_early = false #a condition to evaluate every (few) iteration(s) to see if it's worth continuing a run. When true, we clean up and move to next gene pair.
			not_worth_continuing = false

			#Okay let's set up the energy normal distributions...

			deg_energy_mu, deg_energy_sig = lookup.get_energy_params(deg_name, "energies/")
			mark_energy_mu, mark_energy_sig = lookup.get_energy_params(mark_name, "energies/")

			mark_energy_normal = Normal(mark_energy_mu, mark_energy_sig)
			deg_energy_normal = Normal(deg_energy_mu, deg_energy_sig)

			while (iter == 0 || (iter < max_iter)) #&& !(stop_early) # && minimum((fitness_values)) > 400 && no_change_iters < 100)
				sfv = sort(fitness_values)
				#best_50 = sfv[50]

				if iter % 50 == 0
					println("Step $(iter) of $(max_iter)...")
				end
				@debug("Iteration $iter out of $max_iter")
				@debug("Mean fitness is: $(mean((fitness_values)))")
				@debug("Bottom five are $(sfv[1:5])")
				@debug("... and top five are $(sfv[end-5:end]).")
				@debug("Number of sequences that were changed: $changed_seq")
				flush(log_io)

				#iterated conditional modes, multi-site keep-track = icm_multi_kt. keep track = store values of changes to sequence.
				if do_cull
					if iter < div(max_iter, 4)
						#println("$iter : no cutoff")
						cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel) #, deg_energy_normal, mark_energy_normal)
					elseif iter >= div(max_iter, 4) && iter < div(max_iter, 2)
						#println("$iter : cutoff is $(sfv[div(length(sfv), 2)])")
						cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel; score_cutoff = sfv[div(length(sfv), 2)])
					elseif iter >= div(max_iter, 2)
						#println("$iter : cutoff is $(sfv[div(length(sfv), 4)])")
						cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel; score_cutoff = sfv[div(length(sfv), 4)])
					end
				else
					cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel)
				end

				fitness_values = optimize.assess_pop(cur_pop)
				summed_fitness = sum(fitness_values) #one of the things we should be watching.

				push!(last_few_sf, summed_fitness)
				push!(last_few_cs, changed_seq)

				indiv_count = 1
				for indiv in fitness_values #history matrix stores values of function being optimized over time.
					history_matrix[indiv_count, iter + 1] = indiv
					indiv_count += 1
				end

				#This was originally just code to test the pseudolikelihoods decreasing correctly. It does.
				#Now it's still sorta somewhat useful to log a few values of sequences in both deg/mark as optimization goes.
				if (mod(iter, 50) == 0 && iter != 0)
					@debug("Second opinion, these should be ~0.")
					for i_count in 1:5
						#for the record this is not a good way to generate random numbers.
						#doesn't go from 1 to end.
						i = i_count #round(Int64, rand() * (length(cur_pop) - 1) + 1)

						bel_deg = cur_pop[i].deg_prob
						bel_mrk = cur_pop[i].mark_prob
						#Sort of assuming that end-1 is because of a "*" at stop codon being added.
						if length(strip(cur_pop[i].deg_seq, '*')) == deg_nNodes && length(strip(cur_pop[i].mark_seq, '*')) == mark_nNodes
							cal_deg = mrf.psl(strip(cur_pop[i].deg_seq, '*'), deg_gremodel.w1, deg_gremodel.w2, false)
							cal_mrk = mrf.psl(strip(cur_pop[i].mark_seq, '*'), mark_gremodel.w1, mark_gremodel.w2, false)
							@debug("DEG: $bel_deg vs. $cal_deg")
							@debug("MRK: $bel_mrk vs. $cal_mrk")
							@debug(abs(bel_deg - cal_deg) < 1e-2)
							@debug(abs(bel_mrk - cal_mrk) < 1e-2)
						else
							@debug("Proteins not correct length!")
							@debug(cur_pop[i].mark_seq)
							@debug(cur_pop[i].deg_seq)
						end
					end
				end
				if changed_seq == 0
					no_change_iters += 1
				else
					no_change_iters = 0
				end
				iter += 1
				@debug("\n")
			end

			#So we're done the iterated conditional modes step. Now we report everything...
			@debug("Done while loop...")
			@debug(Libc.strftime(time()))

			@debug("Let us save the history matrix!")
			save("$out_path/$(mark_name)_$(deg_name)_$frame/opt_his_mat_$(rand_barcode).jld", "hist_mat", history_matrix)

			deg_significance = false
			mark_significance = false #this sees if we get any hits worth keeping at all. If yes we save full population.
			all_cal_deg_scores = Float64[]
			all_cal_mark_scores = Float64[]
			#If no we save just a few (top 10) individuals. These are top-6 overall + top-2 in deg and top-2 in mark.
			out_file = open("$out_path/$(mark_name)_$(deg_name)_$frame/all_final_fitness_$(rand_barcode).txt", "w")
			write(out_file, "Ind. #\tMark Score\tDeg Score\tMark Sign\tDeg Sign\n")
			for last_indi in 1:length(cur_pop)
				rep_deg = cur_pop[last_indi].deg_prob
				rep_mrk = cur_pop[last_indi].mark_prob
				if length(strip(cur_pop[last_indi].deg_seq, '*')) == deg_nNodes && length(strip(cur_pop[last_indi].mark_seq, '*')) == mark_nNodes
					cal_deg = abs(mrf.psl(strip(cur_pop[last_indi].deg_seq, '*'), deg_gremodel.w1, deg_gremodel.w2, false))
					cal_mrk = abs(mrf.psl(strip(cur_pop[last_indi].mark_seq, '*'), mark_gremodel.w1, mark_gremodel.w2, false))
					special_deg = "-"
					if cal_deg < mu_deg + sig_deg
						special_deg = "***!"
						deg_significance = true
					elseif cal_deg < mu_deg + 2*sig_deg
						special_deg = "**."
						deg_significance = true
					elseif cal_deg < mu_deg + 3*sig_deg
						special_deg = "*?"
						deg_significance = true
					elseif cal_deg < mu_deg + 4*sig_deg #we don't mind keeping these.
						deg_significance = true
					end
					special_mark = "-"
					if cal_mrk < mu_mark + sig_mark
						special_mark = "***!"
						mark_significance = true
					elseif cal_mrk < mu_mark + 2*sig_mark
						special_mark = "**."
						mark_significance = true
					elseif cal_mrk < mu_mark + 3*sig_mark
						special_mark = "*?"
						mark_significance = true
					elseif cal_mrk < mu_mark + 4*sig_mark #we don't mind keeping these.
						mark_significance = true
					end
					push!(all_cal_mark_scores, cal_mrk)
					push!(all_cal_deg_scores, cal_deg)
					write(out_file, "$(last_indi)\t$(cal_mrk)\t$(cal_deg)\t$special_mark\t$special_deg\n")
				else
					push!(all_cal_mark_scores, Inf)
					push!(all_cal_deg_scores, Inf)
					write(out_file, "$(last_indi)\t-1.0\t-1.0\t/\t/\n")
				end
			end
			close(out_file)

			cal_deg_p = sortperm(all_cal_deg_scores)
			cal_mark_p = sortperm(all_cal_mark_scores)
			cal_both_p = sortperm(fitness_values)

			@debug(Libc.strftime(time()))
			if true #deg_significance && mark_significance
				@debug("Let's save our optimized population...")
				saved_pop = convert_to_saveable(cur_pop)
				save("$out_path/$(mark_name)_$(deg_name)_$frame/saved_pop_$(rand_barcode).jld", "variants", saved_pop)
			else #We want just a subset.
				@debug("We'll save 12 interesting members of the population.")
				selected_pop = ExChrome[]
				for ijk in 1:3
					push!(selected_pop, cur_pop[cal_deg_p[ijk]])
					push!(selected_pop, cur_pop[cal_mark_p[ijk]])
				end
				for ijk in 1:6
					push!(selected_pop, cur_pop[cal_both_p[ijk]])
				end
				save("$out_path/$(mark_name)_$(deg_name)_$frame/top_pop_$(rand_barcode).jld", "variants", selected_pop)
			end

			@debug("And we'll also output our top twelve sequences...")
			out_file = open("$out_path/$(mark_name)_$(deg_name)_$frame/top_twelve_$(rand_barcode).fa", "w")

			for ijk in 1:3
				deg_ind = cal_deg_p[ijk]
				write(out_file, ">TOP_DEG_d (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].deg_seq)\n")
				write(out_file, ">TOP_DEG_m (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].mark_seq)\n")
			end
			for ijk in 1:3
				deg_ind = cal_mark_p[ijk]
				write(out_file, ">TOP_MARK_d (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].deg_seq)\n")
				write(out_file, ">TOP_MARK_m (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].mark_seq)\n")
			end
			for ijk in 1:6
				deg_ind = cal_both_p[ijk]
				write(out_file, ">TOP_GEN_d (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].deg_seq)\n")
				write(out_file, ">TOP_GEN_m (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].mark_seq)\n")
			end
			close(out_file)
		end

		@debug("REALLY DONE!")

		@debug(Libc.strftime(time()))

	end
end

function parse_commandline()
	s = ArgParseSettings()
	@add_arg_table s begin
		"commands"
			help = "positional argument, path to file with tab-delimited commands to run (default: example.txt)"
			default = "example.txt"
		"--num"
			help = "optional, experimentally used for running multiple jobs at once"
			default = "0"
		"--threads"
			help = "optional, experimentally used for running multiple jobs at once"
			default = "1"
	end
	return parse_args(s)
end

function run_file()
	parsed_args = parse_commandline()

	command_file = parsed_args["commands"]
	num = parsed_args["num"]
	threads = parsed_args["threads"]

	run_i = parse(Int64, num)
	NUM_THREADS = parse(Int64, threads)

	if NUM_THREADS > 1
		sleep(rand() * 2) #randomly delay start to de-synchronize starts of runs.
	end

	if isfile(command_file)
		in_file = open(command_file)
		in_read = readlines(in_file)
		close(in_file)

		#This is an additional file that records errors independent of individual log files.
		problem_file = open("problem_runs_$(run_i).txt", "w")

		line_count = 0
		for line in in_read #
			line_count += 1
			if line[1] != '#' && (line_count % NUM_THREADS == run_i)
				try
					run_args = split(line, "\t")
					out_dir, short, long, short_jld, long_jld, short_hmm, long_hmm, pop_size, frame, num_iter = run_args

					rand_barcode = Random.randstring()

					the_out_path = "$out_dir/$(short)_$(long)_$frame/"
					if !(isdir(the_out_path))
						run(`mkdir -p $the_out_path`)
					end

					log_io = open("$out_dir/$(short)_$(long)_$frame/log_$(rand_barcode).txt", "w+")

					logger = SimpleLogger(log_io, Logging.Debug)
					global_logger(logger)

					with_logger(logger) do
						pop_size = parse(Int64, pop_size)
						num_iter = parse(Int64, num_iter)
						set_up_and_optimize(log_io, rand_barcode, out_dir, short, long, short_jld, long_jld, short_hmm, long_hmm, pop_size, frame, num_iter)
					end

					flush(log_io)
					close(log_io)
				catch y
					write(problem_file, "$line $y")
				end

			end
		end
		close(problem_file)
	else
		println("$command_file does not exist. Exiting.")
	end

end

@time run_file()

end


module FieldModule

using Printf: @printf
using UserExceptionModule: UserError
using CoordsModule: Coords,LatLonCoords,DirectionIndicator,Generic1DCoords
using GridModule: Grid,LatLonGrid, for_all, wrap_coords, coords_in_grid,for_all_with_line_breaks
using GridModule: UnstructuredGrid
using InteractiveUtils: subtypes
using SpecialDirectionCodesModule
using MergeTypesModule
import Base.maximum
import Base.*
import Base./
import Base.+
import Base.==
import Base.isapprox
import Base.fill!
import Base.show
import Base.round
import Base.convert
import Base.sum
import Base.Broadcast
using InteractiveUtils

abstract type Field{T} end

function get(field::Field{T},coords::Coords) where {T}
  throw(UserError())
end

function set!(field::Field,coords::Coords)
  throw(UserError())
end

function fill!(field::Field)
  throw(UserError())
end

function maximum(field::Field)
  throw(UserError())
end

function +(lfield::Field,rfield::Field)
  throw(UserError())
end

function *(lfield::Field,value::T) where {T<:Number}
  throw(UserError())
end

function /(lfield::Field,value::T) where {T<:Number}
  throw(UserError())
end

function ==(lfield::Field,rfield::Field)
  throw(UserError())
end

function sum(field::Field)
  throw(UserError())
end

# This could be done by overloading broadcast but this is overcomplicated
# Assumes result is fractional and hence returns a field of floating point numbers
function elementwise_divide(lfield::Field{T},rfield::Field{T}) where {T}
  return Field{Float64}(lfield.grid,lfield.data./rfield.data)::Field{Float64}
end

# This could be done by overloading broadcast but this is overcomplicated
function equals(lfield::Field,value::T) where {T}
  return Field{Bool}(lfield.grid,Bool.(lfield.data .== value))::Field{Bool}
end

function isapprox(lfield::Field,rfield::Field;
                  rtol::Real=atol>0 ? 0 : sqrt(eps),
                  atol::Real,nans::Bool)
  throw(UserError())
end

function invert(field::Field)
    throw(UserError())
end

function add_offset(field::Field,offset::T,skip_offset_on_values::Vector{T}) where {T}
    throw(UserError())
end

function round!(type::DataType,field::Field{T}) where {T}
  throw(UserError())
end

function add_offset(field::T2,offset::T,skip_offset_on_values::Vector{T}) where {T,T2<:Field{T}}
  for_all(get_grid(field)) do coords::Coords
    if ! (field(coords) in skip_offset_on_values)
      set!(field,coords,field(coords)+offset)
    end
  end
end

function divide(field::T2,divisor::T) where {T,T2<:Field{T}}
  for_all(get_grid(field)) do coords::Coords
    set!(field,coords,field(coords)/divisor)
  end
end

get_grid(obj::T) where {T <: Field} =
  obj.grid::Grid

get_data(obj::T) where {T <: Field} =
  obj.data::Array

function repeat(field::Field{T}, count::Integer) where {T}
  array::Array{Field{T},1} = Array{Field{T},1}(undef,count)
  for i in 1:count
    array[i] = deepcopy(field)
  end
  return array
end

struct LatLonField{T} <: Field{T}
  data::Array{T,2}
  grid::LatLonGrid
  function LatLonField{T}(grid::LatLonGrid,value::T) where {T}
    data::Array{T,2} = value == zero(T) ? zeros(T,grid.nlat,grid.nlon) :
                                          ones(T,grid.nlat,grid.nlon)*value
    return new(data,grid)
  end
  function LatLonField{MergeTypes}(grid::LatLonGrid,value::MergeTypes)
    data::Array{MergeTypes,2} =  zeros(MergeTypes,grid.nlat,grid.nlon)
    fill!(data,value)
    return new(data,grid)
  end
  function LatLonField{T}(grid::LatLonGrid,values::AbstractArray{T,2}) where {T}
    if size(values,1) != grid.nlat ||
       size(values,2) != grid.nlon
       error("Values provided don't match selected grid")
    end
    return new(values,grid)
  end
  function LatLonField{MergeTypes}(grid::LatLonGrid,
                                   integer_field::LatLonField{T}) where {T<:Signed}
    grid = get_grid(integer_field)
    merge_type_field = LatLonField{MergeTypes}(grid,null_mtype)
    for_all(grid) do coords::Coords
      element::MergeTypes = MergeTypes(get(integer_field,coords))
      set!(merge_type_field,coords,element)
    end
    return new(merge_type_field.data,grid)
  end
end

struct UnstructuredField{T} <: Field{T}
  data::Array{T,1}
  grid::UnstructuredGrid
  function UnstructuredField{T}(grid::UnstructuredGrid,value::T) where {T}
    data::Array{T,1} = value == zero(T) ? zeros(T,grid.ncells) :
                                          ones(T,grid.ncells)*value
    return new(data,grid)
  end
  function UnstructuredField{MergeTypes}(grid::UnstructuredGrid,value::MergeTypes)
    data::Array{MergeTypes,1} =  zeros(MergeTypes,grid.ncells)
    fill!(data,value)
    return new(data,grid)
  end
  function UnstructuredField{T}(grid::UnstructuredGrid,values::AbstractArray{T,1}) where {T}
    if size(values,1) != grid.ncells
       error("Values provided don't match selected grid")
    end
    return new(values,grid)
  end
  function UnstructuredField{MergeTypes}(grid::UnstructuredGrid,
                                   integer_field::UnstructuredField{T}) where {T<:Signed}
    grid = get_grid(integer_field)
    merge_type_field = UnstructuredField{MergeTypes}(grid,null_mtype)
    for_all(grid) do coords::Coords
      element::MergeTypes = MergeTypes(get(integer_field,coords))
      set!(merge_type_field,coords,element)
    end
    return new(merge_type_field.data,grid)
  end
end

LatLonField{T}(grid::LatLonGrid) where {T} = LatLonField{T}(grid,zeros(T,grid.nlat,grid.nlon))

UnstructuredField{T}(grid::UnstructuredGrid) where {T} = UnstructuredField{T}(grid,zeros(T,grid.ncells))

Field{T}(grid::LatLonGrid,value::T) where {T} = LatLonField{T}(grid::LatLonGrid,value::T)
Field{T}(grid::LatLonGrid) where {T} = LatLonField{T}(grid::LatLonGrid)

Field{T}(grid::UnstructuredGrid,value::T) where {T} = UnstructuredField{T}(grid::UnstructuredGrid,
                                                                           value::T)
Field{T}(grid::UnstructuredGrid) where {T} = UnstructuredField{T}(grid::UnstructuredGrid)

Field{T}(grid::LatLonGrid,values::AbstractArray{T,2}) where {T} =
  LatLonField{T}(grid::LatLonGrid,values::AbstractArray{T,2})
Field{T}(grid::UnstructuredGrid,values::AbstractArray{T,1}) where {T} =
  UnstructuredField{T}(grid::UnstructuredGrid,
                       values::AbstractArray{T,1})

Field{MergeTypes}(grid::LatLonGrid,integer_field::LatLonField{T}) where {T<:Signed} =
  LatLonField{MergeTypes}(grid::LatLonGrid,integer_field::LatLonField{T})
Field{MergeTypes}(grid::UnstructuredGrid,integer_field::UnstructuredField{T}) where {T<:Signed} =
  UnstructuredField{MergeTypes}(grid::UnstructuredGrid,integer_field::UnstructuredField{T})

convert(::Type{Field{MergeTypes}},x::T2) where {T<:Signed,T2<:Field{T}} = Field{MergeTypes}(get_grid(x),x)

function (latlon_field::LatLonField{T})(coords::LatLonCoords) where {T}
    return latlon_field.data[coords.lat,coords.lon]::T
end

function get(latlon_field::LatLonField{T},coords::LatLonCoords) where {T}
    return latlon_field(coords)::T
end

function set!(latlon_field::LatLonField{T},coords::LatLonCoords,value::T) where {T}
    latlon_field.data[coords.lat,coords.lon] = value
end

function +(lfield::LatLonField{T},rfield::LatLonField{T}) where {T}
  return LatLonField{T}(lfield.grid,lfield.data + rfield.data)
end

function *(lfield::Field,value::T) where {T<:Number}
  return LatLonField{T}(lfield.grid,lfield.data * value)
end

function /(lfield::Field,value::T) where {T<:Number}
  return LatLonField{T}(lfield.grid,lfield.data / value)
end

function ==(lfield::LatLonField{T},rfield::LatLonField{T}) where {T}
  return (lfield.data == rfield.data)::Bool
end

LatLonFieldOrUnstructuredField = Union{LatLonField{T},UnstructuredField{T}} where {T}

function sum(field::T) where {T<:LatLonFieldOrUnstructuredField}
  return sum(field.data)
end

function isapprox(lfield::LatLonField{T},rfield::LatLonField{T};
                  rtol::Real=atol>0 ? 0 : sqrt(eps),
                  atol::Real=0.0,nans::Bool=false) where {T}
  return isapprox(lfield.data,rfield.data,
                 rtol=rtol,atol=atol,nans=nans)::Bool
end

function invert(field::LatLonField{T}) where {T}
  return LatLonField{T}(field.grid,.!field.data)
end

function round(type::DataType,field::LatLonField{T}) where {T}
  return LatLonField{type}(field.grid,round.(type,field.data))
end

function maximum(field::LatLonField)
  return maximum(field.data)
end

function show(io::IO,field::LatLonField{T}) where {T <: Number}
  for_all_with_line_breaks(field.grid) do coords::Coords
    if T <: AbstractFloat
      @printf(io,"%6.2f",field(coords))
    elseif T == Bool
      print(io,field(coords) ? " X " : " - ")
    else
      @printf(io,"%4.0f",field(coords))
    end
    print(io," ")
  end
end

function fill!(field::LatLonField{T},value::T) where {T}
  return fill!(field.data,value)
end

function (unstructured_field::UnstructuredField{T})(coords::Generic1DCoords) where {T}
    return unstructured_field.data[coords.index]::T
end

function get(unstructured_field::UnstructuredField{T},coords::Generic1DCoords) where {T}
    return unstructured_field(coords)::T
end

function set!(unstructured_field::UnstructuredField{T},coords::Generic1DCoords,value::T) where {T}
    unstructured_field.data[coords.index] = value
end

function +(lfield::UnstructuredField{T},rfield::UnstructuredField{T}) where {T}
  return UnstructuredField{T}(lfield.grid,lfield.data + rfield.data)
end

function *(lfield::UnstructuredField,value::T) where {T<:Number}
  return UnstructuredField{T}(lfield.grid,lfield.data * value)
end

function /(lfield::UnstructuredField,value::T) where {T<:Number}
  return UnstructuredField{T}(lfield.grid,lfield.data / value)
end

function ==(lfield::UnstructuredField{T},rfield::UnstructuredField{T}) where {T}
  return (lfield.data == rfield.data)::Bool
end

function isapprox(lfield::UnstructuredField{T},rfield::UnstructuredField{T};
                  rtol::Real=atol>0 ? 0 : sqrt(eps),
                  atol::Real=0.0,nans::Bool=false) where {T}
  return isapprox(lfield.data,rfield.data,
                 rtol=rtol,atol=atol,nans=nans)::Bool
end

function invert(field::UnstructuredField{T}) where {T}
  return UnstructuredField{T}(field.grid,.!field.data)
end

function round(type::DataType,field::UnstructuredField{T}) where {T}
  return UnstructuredField{type}(field.grid,round.(type,field.data))
end

function maximum(field::UnstructuredField)
  return maximum(field.data)
end

function fill!(field::UnstructuredField{T},value::T) where {T}
  return fill!(field.data,value)
end

function show(io::IO,field::UnstructuredField{T}) where {T <: Number}
  for_all_with_line_breaks(field.grid) do coords::Coords
    if T <: AbstractFloat
      @printf(io,"%6.2f",field(coords))
    elseif T == Bool
      print(io,field(coords) ? " X " : " - ")
    else
      @printf(io,"%4.0f",field(coords))
    end
    print(io," ")
  end
end

abstract type DirectionIndicators end

function get(field::DirectionIndicators,coords::Coords)
  throw(UserError())
end

function set!(field::DirectionIndicators,coords::Coords)
  throw(UserError())
end

struct LatLonDirectionIndicators <: DirectionIndicators
  next_cell_coords::Array{LatLonField{Int64},1}
  function LatLonDirectionIndicators(dir_based_rdirs::LatLonField{Int64})
    grid = dir_based_rdirs.grid
    nlat = dir_based_rdirs.grid.nlat
    nlon = dir_based_rdirs.grid.nlon
    lat_indices::LatLonField{Int64} =
      LatLonField{Int64}(grid,zeros(Int64,nlat,nlon))
    lon_indices::LatLonField{Int64} =
      LatLonField{Int64}(grid,zeros(Int64,nlat,nlon))
    for_all(grid) do coords::LatLonCoords
      dir_based_rdir::Int64 = dir_based_rdirs(coords)
      if dir_based_rdir == dir_based_indicator_ocean
        set!(lat_indices,coords,coord_base_indicator_ocean)
        set!(lon_indices,coords,coord_base_indicator_ocean)
      elseif dir_based_rdir == dir_based_indicator_outflow
        set!(lat_indices,coords,coord_base_indicator_outflow)
        set!(lon_indices,coords,coord_base_indicator_outflow)
      elseif dir_based_rdir == dir_based_indicator_truesink
        set!(lat_indices,coords,coord_base_indicator_truesink)
        set!(lon_indices,coords,coord_base_indicator_truesink)
      elseif dir_based_rdir == dir_based_indicator_lake
        set!(lat_indices,coords,coord_base_indicator_lake)
        set!(lon_indices,coords,coord_base_indicator_lake)
      elseif 1 <= dir_based_rdir <= 9
        lat_offset::Int64 = dir_based_rdir in [7,8,9] ? -1 :
                           (dir_based_rdir in [1,2,3] ?  1 : 0)
        lon_offset::Int64 = dir_based_rdir in [7,4,1] ? -1 :
                           (dir_based_rdir in [9,6,3] ?  1 : 0)
        destination_coords::LatLonCoords = LatLonCoords(coords.lat + lat_offset,
                                                        coords.lon + lon_offset)
        destination_coords = wrap_coords(grid,destination_coords)
        if ! coords_in_grid(grid,destination_coords)
          error("Direction based direction indicator points out of grid")
        end
        set!(lat_indices,coords,destination_coords.lat)
        set!(lon_indices,coords,destination_coords.lon)
      else
        error("Invalid direction based direction indicator")
      end
    end
    new(LatLonField{Int64}[lat_indices,lon_indices])
  end
end

struct UnstructuredFieldDirectionIndicators <: DirectionIndicators
  next_cell_coords::UnstructuredField{Int64}
end

function (latlon_direction_indicators::LatLonDirectionIndicators)(coords::LatLonCoords)
  return DirectionIndicator(LatLonCoords(latlon_direction_indicators.next_cell_coords[1](coords),
                                         latlon_direction_indicators.next_cell_coords[2](coords)))
end

function get(latlon_direction_indicators::LatLonDirectionIndicators,coords::LatLonCoords)
  return latlon_direction_indicators(coords)
end

function set!(latlon_direction_indicators::DirectionIndicators,
               coords::LatLonCoords,direction::DirectionIndicator)
  next_cell_coords =  get_next_cell_coords(direction,coords)
  set!(latlon_direction_indicators.next_cell_coords[1],coords,next_cell_coords.lat)
  set!(latlon_direction_indicators.next_cell_coords[2],coords,next_cell_coords.lon)
end

function fill!(field::LatLonField{T},value::T) where {T}
  return fill!(field.data,value)
end

struct UnstructuredDirectionIndicators <: DirectionIndicators
  next_cell_coords::UnstructuredField{Int64}
end

function (unstructured_direction_indicators::UnstructuredDirectionIndicators)(coords::Generic1DCoords)
  return DirectionIndicator(Generic1DCoords(unstructured_direction_indicators.next_cell_coords(coords)))
end

function get(unstructured_direction_indicators::UnstructuredDirectionIndicators,
             coords::Generic1DCoords)
  return unstructured_direction_indicators(coords)
end

function set!(unstructured_direction_indicators::UnstructuredDirectionIndicators,
              coords::Generic1DCoords,direction::DirectionIndicator)
  next_cell_coords =  get_next_cell_coords(direction,coords)
  set!(unstructured_direction_indicators.next_cell_coords,coords,next_cell_coords.index)
end

#Cannot define functors on abstract types in Julia yet. Use this
#workaround - needs to be at end of file as it only acts on subtypes of
#Field/DirectionIndicators that are already defined

for T in subtypes(Field)
 @eval function (field::$T)(coords::Coords) throw(UserError()) end
end

for T in subtypes(DirectionIndicators)
  @eval function (direction_indicators::$T)(coords::Coords) throw(UserError()) end
end

end


include("../../configs.jl")
include("../core/metrics.jl")

function encode_gt_and_get_indices(gt, priors, losses, pos_thold, neg_thold)
    variances = size(priors, 2) == 102300 ? [1, 1] : [0.2, 0.1]
    iou_vals = iou(gt[1:4,:], _to_min_max_form(priors))
    # selecting positive prior boxes
    pos_pairs = findall(iou_vals .>= pos_thold)
    pos_gt = zeros(15, length(pos_pairs))
    gt_idx = getindex.(pos_pairs, [1 2])[:, 2]
    prior_idx = getindex.(pos_pairs, [1 2])[:, 1]

    # enlarging gt matrix to match selected anchors
    pos_gt .= gt[:,gt_idx]
    num_poses = length(prior_idx)
    if num_poses == 0
        # if no positive anchor boxes are found, then no loss will be calculated
        return nothing, nothing, nothing
    end

    #selecting negative prior boxes
    max_prior_vals, max_prior_idx = findmax(iou_vals; dims=2)
    neg_indices = getindex.(findall(max_prior_vals .<= neg_thold), [1 2])[:,1]
    neg_indices = neg_indices[sortperm(losses[neg_indices])]

    neg_cnt = ohem_ratio * num_poses
    neg_cnt = neg_cnt < 1 ? 1 : neg_cnt
    neg_indices = neg_indices[1:neg_cnt]

    # gt bbox conversion
    selected_priors = priors[:,prior_idx]

    pos_gt[1:4,:] = _to_center_length_form(pos_gt[1:4,:])
    pos_gt[3:4,:] = log.(pos_gt[3:4,:] ./ selected_priors[3:4,:]) ./ variances[1]

    pos_gt[1:2,:] = (pos_gt[1:2,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[5:6,:] = (pos_gt[5:6,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[7:8,:] = (pos_gt[7:8,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[9:10,:] = (pos_gt[9:10,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[11:12,:] = (pos_gt[11:12,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[13:14,:] = (pos_gt[13:14,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])

    # print("Pos_indices: ", prior_idx, "&& Neg_indices: ", neg_indices, "\n")
    return pos_gt, prior_idx, neg_indices
end

function nms(scores, points; thold=0.4)
    x1 = points[1,:]; y1 = points[2,:]; x2 = points[3,:]; y2 = points[4,:];

    areas = (x2 - x1) .* (y2 - y1)
    order = sortperm(scores, rev=true)
    keep = []

    for i in order
        if keep == []
            push!(keep, i)
        else
            iou_vals = iou(points[:,i:i], points[:,keep])
            val, idx = findmax(iou_vals, dims=1)
            if val[1] <= thold
                push!(keep, i)
            end
        end
    end
    return keep
end

"""
Conversion from:
(center_x, center_y, width, height) --> (min_x, min_y, max_x, max_y)
"""
function _to_min_max_form(boxes)
    init_shape = size(boxes)
    boxes = reshape(boxes, (size(boxes)[1], prod(size(boxes)[2:end]))) # converting to 2D
    half_lens = boxes[3:4,:] ./ 2
    return reshape(cat(boxes[1:2,:] .- half_lens, boxes[1:2,:] .+ half_lens, dims=1), init_shape)
end

"""
Conversion from:
(min_x, min_y, max_x, max_y) --> (center_x, center_y, width, height)
"""
function _to_center_length_form(boxes)
    init_shape = size(boxes)
    boxes = reshape(boxes, (size(boxes)[1], prod(size(boxes)[2:end]))) # converting to 2D
    lengths = boxes[3:4,:] .- boxes[1:2,:]
    return reshape(cat(boxes[1:2,:] .+ (lengths ./ 2), lengths, dims=1), init_shape)
end

"""
Returns the anchor boxes with their center_x, center_y, width, height information.
"""
function _get_priorboxes(num_anchors, anchor_info, img_size)
    feature_maps = [Int(ceil(img_size / scale["stride"])) for scale in anchor_info]
    num_proposals = num_anchors * sum([i*i for i in feature_maps])
    anchors = zeros(4, num_proposals)

    counter = 1
    for (idx, f) in enumerate(feature_maps)
        scaler = anchor_info[idx]["stride"]
        for h in 1:f
            cy = (h - 0.5) * scaler
            for w in 1:f
                cx = (w - 0.5) * scaler
                for s in anchor_info[idx]["anchors"]
                    anchors[:,counter] = [cx cy s s]
                    counter += 1
                end
            end
        end
    end
    return anchors
end

"""
Decoder functions:
The main motivation is that the network does not directly predict the proposals but
the combination of prediction and prior anchor box style and rescaling gives the actual
bounding box and landmark coordinations.
"""
function decode_points(bboxes, landmarks, priors)
    if length(size(priors)) == 2
        priors = reshape(priors, (size(priors)..., 1))
    end
    decoded_bboxes = _decode_bboxes(bboxes, priors)
    decoded_landmarks = _decode_landmarks(landmarks, priors)
    return decoded_bboxes, decoded_landmarks
end

function _decode_bboxes(bbox, priors)
    variances = size(priors, 2) > 100000 ? [1, 1] : [0.2, 0.1]
    if length(size(priors)) == 2
        priors = reshape(priors, (size(priors)..., 1))
    end
    centers = priors[1:2,:,:] .+ bbox[1:2,:,:] .* variances[2] .* priors[3:4,:,:]
    lengths = exp.(bbox[3:4,:,:] .* variances[1]) .* priors[3:4,:,:]
    return cat(centers, lengths, dims=1)
end

function _decode_landmarks(landmarks, priors)
    variances = size(priors, 2) > 100000 ? [1, 1] : [0.2, 0.1]
    if length(size(priors)) == 2
        priors = reshape(priors, (size(priors)..., 1))
    end
    lm1 = priors[1:2,:,:] .+ landmarks[1:2,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm2 = priors[1:2,:,:] .+ landmarks[3:4,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm3 = priors[1:2,:,:] .+ landmarks[5:6,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm4 = priors[1:2,:,:] .+ landmarks[7:8,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm5 = priors[1:2,:,:] .+ landmarks[9:10,:,:] .* variances[2] .* priors[3:4,:,:]
    return cat(lm1, lm2, lm3, lm4, lm5, dims=1)
end


module TiledArrays

# Ceiling division
cdiv(a::Integer, b::Integer) = cdiv(promote(a, b)...)
cdiv(a::T, b::T) where {T<:Integer} = one(T) + div(a - one(T), b)

# Opposite of `Base.tail`.
# Keep all but the last element in a tuple.
head(x, y...) = (x, head(y...)...)
head(x) = ()

# # Type traits for describing if there's something scary going on with the strides.
# # For now, assume memory strides are monotonically increasing such that the underlying
# # array describes a contiguous block of memory.
# abstract type AbstractStrides end
# struct DenseStrides <: AbstractStrides end

struct Tile
    dim::Int
    size::Int
end
Base.show(io::IO, tile::Tile) = print(io, "Tile($(tile.dim), $(tile.size))")

# Note that `NonBlockedLayout <: MemoryLayout`, but without any tiling, allowing
# for specific handling of the easy case.
const MemoryLayout{N} = Tuple{Vararg{Union{Int,Tile},N}} where {N}
const NonBlockedLayout{N} = Tuple{Vararg{Int,N}} where {N}

# Utility function for removing all "Tile"s from the front of an array.
function strip_leading_tiles(x::Tuple{Vararg{Union{Int,Tile},N}}) where {N}
    return strip_leading_tiles(x...)
end

function strip_leading_tiles(x::Tile, y::Vararg{Union{Int,Tile},N}) where {N}
    return strip_leading_tiles(y...)
end

strip_leading_tiles(x::Int, y::Vararg{Union{Int,Tile},N}) where {N} = (x, y...)

#####
##### Layout Validation
#####

# Validation
# Invariants to maintain:
#
# 1. We must have exactly one occurance of an Integer index.
# 2. All `Tile` indexes must occur BEFORE the corresponding Integer index.
# 3. Each `Tile` index must have a corresponding Integer index
function validate(layout::MemoryLayout{N}) where {N}
    @nospecialize

    # Pass 1 - Make sure there are no nonpositive indices.
    # Also, find the maximum index.
    maximum_index = 0
    for desc in layout
        if isa(desc, Int)
            desc <= 0 && error("Found non-positive index: $desc")
            maximum_index = max(maximum_index, desc)
        elseif isa(desc, Tile)
            desc.dim <= 0 && error("Found non-positive tile index: $desc")
            maximum_index = max(maximum_index, desc.dim)
        else
            # Should never happen due to type restrictions.
            error("Unreachable reached")
        end
    end

    # Pass 2 - Check invariants.
    for index in Base.OneTo(maximum_index)
        saw_int = false
        saw_tile = false
        for desc in layout
            if isa(desc, Int) && desc == index
                saw_int && error("Integer index $index occurs more than once!")
                saw_int = true
            elseif isa(desc, Tile) && desc.dim == index
                saw_int && error(
                    "Integer index $index occurs before its corresponding tile index!"
                )
                saw_tile = true
            end
        end
        # We saw a tile but not an int, throw error
        saw_tile &&
            !saw_int &&
            error("Found tile with index $index with no corresponding Int")

        # We didn't find anything
        !saw_tile && !saw_int && error("Found no entry with index $index")
    end
    return true
end

#####
##### Adjust indexes for padding.
#####

unchecked_divrem(x::Int, y::Int) = (Base.sdiv_int(x, y), Base.srem_int(x, y))

# Note: not necessarily safe since we are using the unchecked integer division functions.
# Also, `index` should already be converted from index-1 to index-0.
function adjust_for_padding(
    size::NTuple{N,Int}, padded_size::NTuple{N,Int}, index::NTuple{N,Int}
) where {N}
    return ntuple(i -> adjust_for_padding(size[i], padded_size[i], index[i]), Val(N))
end

function adjust_for_padding(size::Int, padded_size::Int, index::Int)
    size == padded_size && return index
    a, b = unchecked_divrem(index, size)
    return (a * padded_size) + b
end

#####
##### Index Conversion Logic
#####

# Expand the index according to the given layout.
# NOTE: `I` should be post padding adjusting if necessary.
function splitindex(::Val{T}, I::Tuple{Vararg{Int,N}}) where {T,N}
    return _splitindex(T, I)
end

# Work from the first index to the last.
# At each step, strip off the index we processed from `layout` and recurse.
#
# When `layout[1]` is an Integer, we just return the corresponding index from `I`.
#
# When `layout[1]` is a Tile, then we grab the corresponding index from `I`,
# `divrem` that value, then reconstruct `I` with the reduced index.
#
# NB: This depends on Julia's super agressive constant propagation heuristics when dealing
# with simplifying tuple recursive functions to generate great assembly code.
@inline function _splitindex(layout::MemoryLayout, I::Tuple{Vararg{Int,N}}) where {N}
    v, _I = process(layout[1], I)
    return (v, _splitindex(Base.tail(layout), _I)...)
end

### Base case - last item should never be a `Tile`.
@inline function _splitindex(layout::Tuple{Int}, I::Tuple{Vararg{Int,N}}) where {N}
    return (I[layout[1]],)
end

function _splitindex(::Tuple{Tile}, I::Tuple{Vararg{Int,N}}) where {N}
    return error("Malformed memory layout with a `Tile` in the tail position!")
end

@inline process(x::Int, I::Tuple{Vararg{Int,N}}) where {N} = (I[x], I)
@inline function process(x::Tile, I::Tuple{Vararg{Int,N}}) where {N}
    outer, inner = unchecked_divrem(I[x.dim], x.size)
    return (inner, ntuple(i -> i == x.dim ? outer : I[i], Val(N)))
end

#####
##### Obtaining full dimensions
#####

# Move layout from type domain to value domain, providing "const" information to the
# compiler.
@inline dims(valT::Val{T}, size::Tuple{Vararg{Int,N}}) where {T,N} = _dims(T, size)

@inline function _dims(layout::MemoryLayout, size::Tuple{Vararg{Int,N}}) where {N}
    v, _size = _dimcheck(layout[1], size)
    return (v, _dims(Base.tail(layout), _size)...)
end

# Bottom of recursion.
@inline _dims(layout::Tuple{Int}, size::Tuple{Vararg{Int,N}}) where {N} = (size[layout[1]],)

@inline _dimcheck(x::Int, size::Tuple{Vararg{Int,N}}) where {N} = (size[x], size)
@inline function _dimcheck(x::Tile, size::Tuple{Vararg{Int,N}}) where {N}
    newdim = cdiv(size[x.dim], x.size)
    return (x.size, ntuple(i -> i == x.dim ? newdim : size[i], Val(N)))
end

#####
##### Users of "dims"
#####

# NB: Cannot extend `Base.strides` because a `MemoryLayout` can potentially just be a tuple
# of Ints, so we'd be committing type piracy.
function _strides(valT::Val{T}, size::Tuple{Vararg{Int,N}}) where {T,N}
    # Drop the last value from `dims`.
    return cumprod((one(Int), head(dims(valT, size)...)...))
end

function fullsize(valT::Val{T}, size::Tuple{Vararg{Int,N}}) where {T,N}
    # Constant propagation to the rescue.
    # If the layout is non-blocked, just standard or permuted, then there is no
    # over-allocation necessary.
    if isa(T, NonBlockedLayout)
        return prod(size)
        # Otherwise, we need to compute the full dimensions, which could include zero padding
        # to keep the inner blocksizes constant.
    else
        return prod(dims(valT, size))
    end
end

#####
##### Linear Offset
#####

function getoffset(
    valT::Val{T},
    size::Tuple{Vararg{Int,N}},     # Padded Size
    I::Tuple{Vararg{Int,N}},        # Post-padded index
) where {T,N}
    Base.@_inline_meta
    strides = _strides(valT, size)
    index = splitindex(valT, I)
    offset = zero(Int)
    for i in eachindex(strides)
        @inbounds offset += strides[i] * index[i]
    end
    return offset
end

# Generates indices into a tiled array.
# Parameters:
#
# `L`: Layout type parameter.
# `N`: Number of Dims.
struct TiledIndexer{L,N}
    size::NTuple{N,Int}
    padded_size::NTuple{N,Int}
end

function TiledIndexer{L}(size::NTuple{N,Int}, padded_size::NTuple{N,Int}) where {L,N}
    return TiledIndexer{L,N}(size, padded_size)
end

function genindex(x::TiledIndexer{L,N}, I::NTuple{N,Int}) where {L,N}
    return 1 + getoffset(
        Val(L), x.padded_size, adjust_for_padding(x.size, x.padded_size, I .- 1)
    )
end

# function genindex(x::TiledIndexer{L,N,Nothing}, I::NTuple{N,Int}) where {L,N}
#     return getoffset(Val(N), x.size, I)
# end

# #####
# ##### Tiled Array
# #####
#
# struct TiledArray{T,L,N} <: AbstractArray{T,N}
#     parent::Vector{T}
#     size::NTuple{N,Int}
# end
# Base.parent(A::TiledArray) = A.parent
#
# function TiledArray{T,L}(::UndefInitializer, size::Tuple{Vararg{Int,N}}) where {T,L,N}
#     # Allocate enough room for the whole layout.
#     parent = Vector{T}(undef, fullsize(Val(L), size))
#     return TiledArray{T,L,N}(parent, size)
# end
#
# layout(::Type{<:TiledArray{<:Any,L}}) where {L} = L
# layout(x::T) where {T<:TiledArray} = layout(T)
# vlayout(x) = Val(layout(x))
#
# # Array interface
# @inline Base.size(x::TiledArray) = x.size
# Base.IndexStyle(::Type{<:TiledArray}) = Base.IndexCartesian()
# Base.@propagate_inbounds function Base.getindex(
#     A::TiledArray{T,L,N}, I::Vararg{Int,N}; strides = nothing
# ) where {T,L,N}
#     @boundscheck checkbounds(A, I...)
#     offset = getoffset(vlayout(A), size(A), I, strides) + one(Int)
#     return @inbounds(getindex(parent(A), offset))
# end
#
# Base.@propagate_inbounds function Base.setindex!(
#     A::TiledArray{T,L,N}, v, I::Vararg{Int,N}; strides = nothing
# ) where {T,L,N}
#     @boundscheck checkbounds(A, I...)
#     offset = getoffset(vlayout(A), size(A), I, strides) + one(Int)
#     return @inbounds setindex!(parent(A), v, offset)
# end

end # module


# ======================
# Dependencies
# ======================

using LinearAlgebra, SparseArrays, # modules from the Julia standard library
      Reexport,                    # see @reexport macro below
      RecipesBase,                 # plotting
      Parameters,                  # structs with kwargs
      StaticArrays,                # statically sized arrays
      RecursiveArrayTools          # vector of arrays type

# manipulate function definition expressions
using ExprTools: splitdef, combinedef

# the reexport macro ensures that the names exported by the following libraries
# are made available after loading ReachabilityAnalysis
@reexport using HybridSystems,
                IntervalMatrices,
                LazySets,
                MathematicalSystems,
                TaylorIntegration

# required to avoid conflicts with MathematicalSystems
using LazySets: LinearMap, AffineMap, ResetMap

# required to avoid conflicts with IntervalMatrices
using LazySets: Interval, radius, sample, ∅, dim, scale, scale!

# in-place set operations
using LazySets: linear_map!

# LazySets internal functions frequently used
using LazySets.Arrays: projection_matrix, SingleEntryVector
using LazySets.Approximations: AbstractDirections
using LazySets: @commutative

# aliases for intervals
const IM = IntervalMatrices
import IntervalArithmetic
const IA = IntervalArithmetic
const TimeInterval = IA.Interval{Float64}
import TaylorModels
const TM = TaylorModels

# method extensions for Taylor model reach-sets
import TaylorModels: domain, remainder, polynomial, get_order, evaluate

# aliases for set types
const CPA = CartesianProductArray

# convenience union for dispatch on structs that are admissible as initial sets or inputs
const AdmissibleSet = Union{LazySet, UnionSet, UnionSetArray, IA.Interval, IA.IntervalBox}

# method extensions
import LazySets: dim, overapproximate, box_approximation, project, Projection,
                 intersection, directions, linear_map, LinearMap, _split, split!,
                 set, array, _isapprox

import MathematicalSystems: discretize
import Base: ∈, ∩, convert, isdisjoint
import LinearAlgebra: normalize

import CommonSolve: solve # common solve name

# ======================
# Useful constants
# ======================

@inline zeroBox(m) = IntervalBox(zeroI, m)
@inline unitBox(m) = IntervalBox(IA.Interval(0.0, 1.0), m)
@inline symBox(n::Integer) = IntervalBox(symI, n)
const zeroI = IA.Interval(0.0) # TODO use number type
const oneI = IA.Interval(1.0)
const symI = IA.Interval(-1.0, 1.0)

# Interval constructor given a float
LazySets.Interval(x::Float64) = Interval(interval(x))

# common aliases for system's names
const LCS = LinearContinuousSystem
const LDS = LinearDiscreteSystem
const CLCS = ConstrainedLinearContinuousSystem
const CLDS = ConstrainedLinearDiscreteSystem
const CLCCS = ConstrainedLinearControlContinuousSystem
const CLCDS = ConstrainedLinearControlDiscreteSystem
const ACS = AffineContinuousSystem
const ADS = AffineDiscreteSystem
const CACCS = ConstrainedAffineControlContinuousSystem
const CACDS = ConstrainedAffineControlDiscreteSystem
const CACS = ConstrainedAffineContinuousSystem
const CADS = ConstrainedAffineDiscreteSystem
const BBCS = BlackBoxContinuousSystem
const CBBCS = ConstrainedBlackBoxContinuousSystem
const CBBCCS = ConstrainedBlackBoxControlContinuousSystem
const SOLCS = SecondOrderLinearContinuousSystem
const SOACS = SecondOrderAffineContinuousSystem
const SOCLCCS = SecondOrderConstrainedLinearControlContinuousSystem
const SOCACCS = SecondOrderConstrainedAffineControlContinuousSystem
const SecondOrderSystem = Union{SOLCS, SOACS, SOCLCCS, SOCACCS}
const NonlinearSystem = Union{BBCS, CBBCS, CBBCCS}

@inline function _isapprox(Δt::TimeInterval, Δs::TimeInterval)
    return (inf(Δt) ≈ inf(Δs)) && (sup(Δt) ≈ sup(Δs))
end

const VecOrTuple = Union{<:AbstractVector{Int}, NTuple{D, Int}} where {D}
const VecOrTupleOrInt = Union{<:AbstractVector{Int}, NTuple{D, Int}, Int} where {D}

# ======================
# Optional dependencies
# ======================

using Requires

function __init__()
    # numerical differential equations suite
    @require DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" include("init_DifferentialEquations.jl")

    # exponentiation methods using Krylov subspace approximations
    @require ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18" include("init_ExponentialUtilities.jl")

    # methods that require optimization modeling and solvers
    @require JuMP = "4076af6c-e467-56ae-b986-b466b2749572" include("init_JuMP.jl")

    # tools for symbolic computation
    #@require ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" include("init_ModelingToolkit.jl")
    @require Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" include("init_Symbolics.jl")

    # tools for symbolic algebra
    @require MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3" include("init_MultivariatePolynomials.jl")

    # sparse dynamic representation of multivariate polynomials
    @require DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07" include("init_DynamicPolynomials.jl")
end

# ===========================
# Utility macros
# ===========================

"""
    @requires(module_name)

Convenience macro to annotate that a package is required to use a certain function.

### Input

- `module_name` -- name of the required package

### Output

The macro expands to an assertion that checks whether the module `module_name` is
known in the calling scope.

### Notes

Usage:

```julia
function foo(...)
    @require MyPackage
    ... # functionality that requires MyPackage to be loaded
end
```
"""
macro requires(module_name)
    m = Meta.quot(Symbol(module_name))
    return esc(:(@assert isdefined(@__MODULE__, $m) "package `$($m)` is required " *
                    "for this function; do `using $($m)` and try again"))
end


function _process_reaction_phrase(phrase_array::Array{String,1})

    tmp_string = ""
    for phrase in phrase_array
        tmp_string*="$(phrase) + "
    end

    # cutoff -
    tmp_string = tmp_string |> rstrip

    # cutoff the trailing +
    return (tmp_string[1:end-1] |> rstrip)
end

function build_stoichiometric_matrix(model::Dict{Symbol,Any})

    try

        # get the reaction table -
        reaction_table = model[:reactions]
        reaction_id_array = reaction_table[!, :reaction_number]
        compound_id_array = model[:compounds][!,:compound_id]

        # now - let's build the stm -
        number_of_species = length(compound_id_array)
        number_of_reactions = length(reaction_id_array)
        stoichiometric_matrix = zeros(number_of_species, number_of_reactions)

        # build the array -
        for reaction_index = 1:number_of_reactions

            # what is my reaction id?
            reaction_id = reaction_id_array[reaction_index]

            # get row from the reaction table -
            df_reaction = filter(:reaction_number => x -> (x == reaction_id), reaction_table)

            # grab the stm dictionary -
            stm_dictionary = df_reaction[1, :stoichiometric_dictionary]

            # ok, lets see if we have these species -
            for species_index = 1:number_of_species

                # species code -
                species_symbol = compound_id_array[species_index]
                if (haskey(stm_dictionary, species_symbol) == true)
                    stm_coeff_value = stm_dictionary[species_symbol]
                    stoichiometric_matrix[species_index, reaction_index] = stm_coeff_value
                end
            end
        end

        # return -
        return (compound_id_array, reaction_id_array, stoichiometric_matrix)
    catch error

        # what is our error policy? => for now, just print the message
        error_message = sprint(showerror, error, catch_backtrace())
        println(error_message)

    end
end

function find_compound_index(model::Dict{Symbol,Any},
    search::Pair{Symbol,String})

    # get the compounds table -
    compounds_table = model[:compounds]

    # get list of compound names -
    tmp_array = compounds_table[!,search.first] |> collect

    # do we have this name? (if yes, then return index)
    return findfirst(x->x==search.second,tmp_array)
end

function find_reaction_index(model::Dict{Symbol,Any},
    search::Pair{Symbol,String})

    # get the compounds table -
    reaction_table = model[:reactions]

    # get list of compound names -
    tmp_array = reaction_table[!,search.first] |> collect

    # do we have this name? (if yes, then return index)
    return findfirst(x->x==search.second,tmp_array)
end

function update_flux_bounds_directionality(MODEL,default_flux_bounds)

    # get the ΔG table -> for now, just use the mean value -
    ΔG_table = MODEL[:ΔG]

    # how many reactions do we have ΔG data for?
    (number_of_reactions,_) = size(ΔG_table)
    for reaction_index = 1:number_of_reactions

        # get ΔG value -
        μ_ΔG_value = ΔG_table[reaction_index,:μ_ΔG]

        # get the corresponding reaction number -
        reaction_id = ΔG_table[reaction_index, :reaction_number]

        # what index is this?
        index_in_bounds_array = find_reaction_index(MODEL,:reaction_number=>reaction_id)
        if (isempty(index_in_bounds_array) == false)

            # backward?
            is_reversible = -1*(sign(μ_ΔG_value)) <= 0.0 ? true : false
            if (is_reversible == false)
                default_flux_bounds[index_in_bounds_array,1] = 0.0
            end
        end
    end

    return default_flux_bounds
end

function translation_reaction_string_to_human(model::Dict{Symbol,Any})

    # get the compound table -
    compound_table = model[:compounds]
    reaction_table = model[:reactions]

    # build a translation table -
    (number_of_compounds, _) = size(compound_table)
    compound_translation_table = Dict{String,String}()
    for compound_index = 1:number_of_compounds
        kegg_name = compound_table[compound_index, :compound_id]
        human_name = compound_table[compound_index, :compound_name]
        compound_translation_table[kegg_name] = human_name
    end

    # hack: -
    compound_translation_table["C00138"] = "Reduced ferredoxin"
    compound_translation_table["C00139"] = "Oxidized ferredoxin"

    # number of _reactions -
    (number_of_reactions, _) = size(reaction_table)

    # initialze space -
    tmp_reaction_string_array = Array{String,1}(undef,number_of_reactions)

    # just so we are safe - copy the old strings into the tmp array in case something doesn't work
    for reaction_index = 1:number_of_reactions
        old_string = reaction_table[reaction_index, :reaction_markup]
        tmp_reaction_string_array[reaction_index] = old_string
    end

    # init some tmp storage -
    tmp_reactant_phrase_array = Array{String,1}()
    tmp_product_phrase_array = Array{String,1}()

    # main loop - here we go .... let's do this ...
    for reaction_index = 1:number_of_reactions

        # grab the st dictionary for this reaction -
        stoichiometric_dictionary = reaction_table[reaction_index, :stoichiometric_dictionary]
        for (key,value) in stoichiometric_dictionary

            # look up this key in the translation table -
            human_key = get(compound_translation_table, key, "?")

            if (abs(value) != 1)
                tmp_phrase = "$(abs(value)) $(human_key)"
            else
                tmp_phrase = "$(human_key)"
            end

            if (value<0)
                push!(tmp_reactant_phrase_array, tmp_phrase)
            else
                push!(tmp_product_phrase_array, tmp_phrase)
            end
        end

        # build new reaction string ...
        reactant_string = _process_reaction_phrase(tmp_reactant_phrase_array)
        product_string = _process_reaction_phrase(tmp_product_phrase_array)
        full_reaction_string = "$(reactant_string) <=> $(product_string)"
        tmp_reaction_string_array[reaction_index] = full_reaction_string

        # empty -
        empty!(tmp_reactant_phrase_array)
        empty!(tmp_product_phrase_array)
    end

    return tmp_reaction_string_array
end


function dh_simulation(LP::AbstractLProblem{T}, p::Float64, nTrials::Int64, alg::OrdinaryDiffEqAlgorithm, logger::Logger{T1,T2} = Logger{:no, :cmd}(), interval::Int64 = 1, offset1::Int64 = 0; links = (:all,:all), axes = [:x, :y, :z], kwargs...) where {T<:PureClassicalApprox,T1,T2}
    assert(nTrials > 0)
    RNGarray = initRNG(nprocs())
    path_part = @sprintf("Dis Tm%.2f dt%.3f del%d Obs%s Pr%.3f/", LP.Tmax, LP.tstep, LP.delimiter, get_string(LP.rules.str_vec), p)
    filenames = Vector{String}()
    A = LP.A
    l = LP.cb.affect!.save_func.len
    tspan = (0.0, LP.Tmax*LP.delimiter)
    t_means = collect(0.0:LP.tstep:(LP.Tmax*LP.delimiter))
    saved_values = SavedValues(t_means, [zeros(l) for i in eachindex(t_means)])
    for thread = (offset1+1):(offset1+nprocs())
        push!(filenames, get_string(nTrials, thread, alg))
    end
    log_o, save_os = set_Logging(LP, nTrials, logger, A, path_part, filenames)
    cf_vals0 = LP.cf_vals
    @sync for i in workers()
        @spawnat i begin
                       global cf_vals = cf_vals0
                       global RNG = RNGarray[i]
                   end
    end
    iter = get_status(log_o)
    offset = get_offset(log_o)
    if iter > 0
        load_data!(cf_vals0, save_os[1])
        for i in workers()
            @spawnat i load_data!(cf_vals, save_os[i])
        end
    end
    t0 = time()-offset
    while iter<nTrials
        t1 = time()
        iter += 1
        ARHS = build_RHS_function(A)
        rules, OSET = set_CorrFuncs(A, links, axes)
        @sync begin
            for i in workers()
                @spawnat i dh_run(iter, A, ARHS, OSET, rules, t_means, saved_values, tspan, p, cf_vals, alg, save_os[i], RNG; kwargs...)
            end
            @spawnat 1 dh_run(iter, A, ARHS, OSET, rules, t_means, saved_values, tspan, p, cf_vals0, alg, save_os[1], RNGarray[1]; kwargs...)
        end
        t2 = time()
        output(log_o, @sprintf("%10d %20.5f %20.5f %10d", iter, t2-t1, (t2-t0)/iter, nTrials - iter))
    end
end

function dh_simulation(LP::AbstractLProblem{T}, p::Float64, nTrials::Int64, alg::OrdinaryDiffEqAlgorithm, logger::Logger{T1,T2} = Logger{:no, :cmd}(), interval::Int64 = 1, offset1::Int64 = 0; links = (:all,:allq), axes = [:x, :y, :z], kwargs...) where {T<:HybridApprox,T1,T2}
    assert(nTrials > 0)
    RNGarray = initRNG(nprocs()+1)
    path_part = @sprintf("Dis Tm%.2f dt%.3f del%d Obs%s Pr%.3f/", LP.Tmax, LP.tstep, LP.delimiter, get_string(LP.rules.str_vec), p)
    filenames = Vector{String}()
    A = HybridApprox(LP.A.L, copy(LP.A.q_spins), LP.A.name)
    l = LP.cb.affect!.save_func.len
    tspan = (0.0, LP.Tmax*LP.delimiter)
    t_means = collect(0.0:LP.tstep:(LP.Tmax*LP.delimiter))
    saved_values = SavedValues(t_means, [zeros(l) for i in eachindex(t_means)])
    for thread = (offset1+1):(offset1+nprocs())
        push!(filenames, get_string(nTrials, thread, alg))
    end
    log_o, save_os = set_Logging(LP, nTrials, logger, LP.A, path_part, filenames)
    cf_vals0 = LP.cf_vals
    @sync for i in workers()
        @spawnat i begin
                       global cf_vals = cf_vals0
                       global RNG = RNGarray[i]
                   end
    end
    iter = get_status(log_o)
    offset = get_offset(log_o)
    if iter > 0
        load_data!(cf_vals0, save_os[1])
        for i in workers()
            @spawnat i load_data!(cf_vals, save_os[i])
        end
    end
    t0 = time()-offset
    while iter<nTrials
        t1 = time()
        iter += 1
        q_spins = randsubseq(RNGarray[end], LP.A.q_spins, 1-p)
        A.q_spins = q_spins
        A.q_num = length(q_spins)
        ARHS = build_RHS_function(A)
        rules, OSET = set_CorrFuncs(A, links, axes)
        @sync begin
            for i in workers()
                @spawnat i dh_run(iter, A, ARHS, OSET, rules, t_means, saved_values, tspan, p, cf_vals, alg, save_os[i], RNG; kwargs...)
            end
            @spawnat 1 dh_run(iter, A, ARHS, OSET, rules, t_means, saved_values, tspan, p, cf_vals0, alg, save_os[1], RNGarray[1]; kwargs...)
        end
        t2 = time()
        output(log_o, @sprintf("%10d %20.5f %20.5f %10d", iter, t2-t1, (t2-t0)/iter, nTrials - iter))
    end
end

function dh_simulation(LP::AbstractLProblem{Approx}, p::Float64, nTrials::Int64, alg::OrdinaryDiffEqAlgorithm, logger::Logger{T1,T2} = Logger{:no, :cmd}(), interval::Int64 = 1, offset1::Int64 = 0; links = (:all,:allq), axes = [:x, :y, :z], kwargs...) where {Approx <: CompositeApproximation{<:HybridApprox, <:PureClassicalApprox}, T1,T2}
    A = LP.A
    Tmax = LP.Tmax
    delimiter = LP.delimiter
    tstep = LP.tstep
    A_temp = CompositeApproximation(HybridApprox(A.A1.L, copy(A.A1.q_spins), A.A1.name), A.A2, A.gamma)
    #rules, OSET = set_CorrFuncs(A, links, axes)
    l = LP.cb.affect!.save_func.len
    tspan = (0.0, Tmax*delimiter)
    t_means = collect(0.0:tstep:(Tmax*delimiter))
    t_cors = collect(0.0:tstep:Tmax)
    saved_values = SavedValues(t_means, [zeros(l) for i in eachindex(t_means)])
    ##############################
    assert(nTrials > 0)
    RNGarray = initRNG(nprocs()+1)
    path_part = @sprintf("Dis Tm%.2f dt%.3f del%d Obs%s Pr%.3f/", Tmax, tstep, delimiter, get_string(LP.rules.str_vec), p)
    filenames = Vector{String}()
    for thread = (offset1+1):(offset1+nprocs())
        push!(filenames, get_string(nTrials, thread, alg))
    end
    log_o, save_os = set_Logging(LP, nTrials, logger, A, path_part, filenames)
    cf_vals0 = LP.cf_vals
    @sync for i in workers()
        @spawnat i begin
                       global cf_vals = cf_vals0
                       global RNG = RNGarray[i]
                   end
    end
    iter = get_status(log_o)
    offset = get_offset(log_o)
    if iter > 0
        load_data!(cf_vals0, save_os[1])
        for i in workers()
            @spawnat i load_data!(cf_vals, save_os[i])
        end
    end
    t0 = time()-offset
    ##############################
    while iter<nTrials
        t1 = time()
        iter += 1
        q_spins = randsubseq(RNGarray[end], A.A1.q_spins, 1-p)
        A_temp.A1.q_spins = q_spins
        A_temp.A1.q_num = length(q_spins)
        print(length(q_spins), "\n")
        ARHS = build_RHS_function(A_temp)
        rules, OSET = set_CorrFuncs(A_temp.A1, links, axes)
        @sync begin
            for i in workers()
                @spawnat i dh_run(iter, A_temp, ARHS, OSET, rules, t_means, saved_values, tspan, p, cf_vals, alg, save_os[i], RNG; kwargs...)
            end
            @spawnat 1 dh_run(iter, A_temp, ARHS, OSET, rules, t_means, saved_values, tspan, p, cf_vals0, alg, save_os[1], RNGarray[1]; kwargs...)
        end
        t2 = time()
        output(log_o, @sprintf("%10d %20.5f %20.5f %10d", iter, t2-t1, (t2-t0)/iter, nTrials - iter))
    end
end

function dh_run(iter::Int64, A::abstractApproximation, ARHS::AbstractRHSFunction, OSET::ObservablesSet, rules::ConvRules, t_means::Vector{Float64}, saved_values::SavedValues, tspan::NTuple{2, Float64}, p::Float64, cf_vals::T, alg::OrdinaryDiffEqAlgorithm, save_o::Saver, RNG::AbstractRNG; abstol_i = 1e-14, reltol_i = 1e-6, kwargs...) where T <: Union{CFVals, CFVals0}
    u0 = randomDState(A, p, RNG)
    #print(u0, "\n")
    prob = ODEProblem(ARHS, u0, tspan)
    cb = SavingCallback(OSET, saved_values, u0; saveat = t_means)
    solve(prob, alg; save_everystep = false, callback = cb, abstol = abstol_i, reltol = reltol_i, kwargs...)
    calculateCorrFunc!(A, saved_values, rules, cf_vals, iter)
    save_data(save_o, cf_vals, iter)
    #print(cf_vals, "\n")
end

function randomDState(A::PureClassicalApprox, p::Float64, RNG::AbstractRNG = Base.GLOBAL_RNG)
    cl_state = randomState(A, RNG)
    discarded = randsubseq(RNG, collect(1:A.L.tot_spin_num), p)
    for i in discarded
        cl_state[i,:] .= 0.0
    end
    return cl_state
end

function randomDState(A::HybridApprox, p::Float64, RNG::AbstractRNG = Base.GLOBAL_RNG)
    Dh = 2^A.q_num
    cl_state = random_cl_state(A.cl_num, RNG)
    discarded = randsubseq(RNG, collect(1:A.cl_num), p)
    for i in discarded
        cl_state[i,:] .= 0.0
    end
    return ArrayPartition(random_q_state(Dh, RNG), cl_state)
end

function randomDState(A::CompositeApproximation{T1,T2}, p::Float64, RNG::AbstractRNG = Base.GLOBAL_RNG) where {T1<:HybridApprox, T2<:PureClassicalApprox}
    Dh = 2^A.A1.q_num
    cl_state = random_cl_state(A.A1.cl_num, RNG)
    discarded = randsubseq(RNG, collect(1:A.A1.cl_num), p)
    for i in discarded
        cl_state[i,:] .= 0.0
    end
    ArrayPartition(random_q_state(Dh, RNG), cl_state, randomState(A.A2, RNG))
end

function getrsubseq(A::AbstractArray{T}, p::Float64, RNG::AbstractRNG = Base.GLOBAL_RNG) where {T}
    B = Vector{T}()
    for i in eachindex(A)
        if rand(RNG) < p
            push!(B, A[i])
        end
    end
    return B
end

export dh_simulation