text:
"""
# module LaminartFunc

- Julia version: 1.4
- Author: niallcullinane
- Date: 2020-08-20

Functions for DiffEq to call LAMINART equations and hold values.
"""
module LaminartFunc

include("./LaminartKernels.jl")
include("./LaminartEqImfilter.jl")
include("./LaminartEqImfilterGPU_FFT.jl")
include("./LaminartEqImfilterGPU_FIR.jl")
include("./LaminartEqImfilterGPU_IIR.jl")
include("./LaminartEqConv.jl")

using NNlib, ImageFiltering, Images, OffsetArrays, CUDA

mutable struct LamFunction{T<:AbstractArray} <: Function
    x::T
    m::T
    s::T

    x_lgn::T
    C::T
    H_z::T

    dy_temp::T
    dm_temp::T
    dz_temp::T
    ds_temp::T
    dv_temp::T
    H_z_temp::T
    V_temp_1::T
    V_temp_2::T
    Q_temp::T
    P_temp::T
end

function (ff::LamFunction)(du, u, p, t)

    @inbounds begin

        @. ff.x = @view u[:, :, 1:p.K, :]
        y = @view u[:, :, p.K+1:2*p.K, :]
        @. ff.m = @view u[:, :, 2*p.K+1:3*p.K, :]
        z = @view u[:, :, 3*p.K+1:4*p.K, :]
        @. ff.s = @view u[:, :, 4*p.K+1:5*p.K, :]

        v_p = @view u[:, :, 5*p.K+1:5*p.K+1, :]
        v_m = @view u[:, :, 5*p.K+2:5*p.K+2, :]

        dx = @view du[:, :, 1:p.K, :]
        dy = @view du[:, :, p.K+1:2*p.K, :]
        dm = @view du[:, :, 2*p.K+1:3*p.K, :]
        dz = @view du[:, :, 3*p.K+1:4*p.K, :]
        ds = @view du[:, :, 4*p.K+1:5*p.K, :]

        dv_p = @view du[:, :, 5*p.K+1:5*p.K+1, :]
        dv_m = @view du[:, :, 5*p.K+2:5*p.K+2, :]

        LaminartEqConv.fun_x_lgn!(ff.x_lgn, ff.x, p)
        LaminartEqConv.fun_v_C!(
            ff.C,
            v_p,
            v_m,
            ff.V_temp_1,
            ff.V_temp_2,
            ff.Q_temp,
            ff.P_temp,
            p,
        )
        LaminartEqConv.fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        LaminartEqConv.fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
        LaminartEqConv.fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
        LaminartEqConv.fun_dx_v1!(dx, ff.x, ff.C, z, p.x_V2, p)
        LaminartEqConv.fun_dy!(dy, y, ff.C, ff.x, ff.m, ff.dy_temp, p)
        LaminartEqConv.fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
        LaminartEqConv.fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
        LaminartEqConv.fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

    end
    return nothing

end

"""
All arrays held in struct.
"""
mutable struct LamFunction_allStruct <: Function
    x::AbstractArray
    y::AbstractArray
    m::AbstractArray
    z::AbstractArray
    s::AbstractArray
    v_p::AbstractArray
    v_m::AbstractArray

    dx::AbstractArray
    dy::AbstractArray
    dm::AbstractArray
    dz::AbstractArray
    ds::AbstractArray
    dv_p::AbstractArray
    dv_m::AbstractArray

    x_lgn::AbstractArray
    C::AbstractArray
    H_z::AbstractArray

    dy_temp::AbstractArray
    dm_temp::AbstractArray
    dz_temp::AbstractArray
    ds_temp::AbstractArray
    dv_temp::AbstractArray
    H_z_temp::AbstractArray
    V_temp_1::AbstractArray
    V_temp_2::AbstractArray
    Q_temp::AbstractArray
    P_temp::AbstractArray
end

function (ff::LamFunction_allStruct)(du, u, p, t)

    @inbounds begin

        @. ff.x = @view u[:, :, 1:p.K, :]
        ff.y = @view u[:, :, p.K+1:2*p.K, :]
        @. ff.m = @view u[:, :, 2*p.K+1:3*p.K, :]
        ff.z = @view u[:, :, 3*p.K+1:4*p.K, :]
        @. ff.s = @view u[:, :, 4*p.K+1:5*p.K, :]

        ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1, :]
        ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2, :]

        ff.dx = @view du[:, :, 1:p.K, :]
        ff.dy = @view du[:, :, p.K+1:2*p.K, :]
        ff.dm = @view du[:, :, 2*p.K+1:3*p.K, :]
        ff.dz = @view du[:, :, 3*p.K+1:4*p.K, :]
        ff.ds = @view du[:, :, 4*p.K+1:5*p.K, :]

        ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1, :]
        ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2, :]

        LaminartEqConv.fun_x_lgn!(ff.x_lgn, ff.x, p)
        LaminartEqConv.fun_v_C!(
            ff.C,
            ff.v_p,
            ff.v_m,
            ff.V_temp_1,
            ff.V_temp_2,
            ff.Q_temp,
            ff.P_temp,
            p,
        )
        LaminartEqConv.fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

        LaminartEqConv.fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
        LaminartEqConv.fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
        LaminartEqConv.fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
        LaminartEqConv.fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
        LaminartEqConv.fun_dm!(ff.dm, ff.m, ff.x, ff.dm_temp, p)
        LaminartEqConv.fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
        LaminartEqConv.fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)
    end
    return nothing

end

"""
Temp arrays resused for differiant function.
"""
mutable struct LamFunction_gpu_reuse{T<:AbstractArray} <: Function

    x_lgn::T
    C::T
    H_z::T

    tmp_a::T
    tmp_b::T

    tmp_A::T
    tmp_B::T
    tmp_C::T
end

function (ff::LamFunction_gpu_reuse)(du, u, p, t)

    @inbounds begin

        y = @view u[:, :, p.K+1:2*p.K, :]
        z = @view u[:, :, 3*p.K+1:4*p.K, :]

        v_p = @view u[:, :, 5*p.K+1:5*p.K+1, :]
        v_m = @view u[:, :, 5*p.K+2:5*p.K+2, :]

        dx = @view du[:, :, 1:p.K, :]
        dy = @view du[:, :, p.K+1:2*p.K, :]
        dm = @view du[:, :, 2*p.K+1:3*p.K, :]
        dz = @view du[:, :, 3*p.K+1:4*p.K, :]
        ds = @view du[:, :, 4*p.K+1:5*p.K, :]

        dv_p = @view du[:, :, 5*p.K+1:5*p.K+1, :]
        dv_m = @view du[:, :, 5*p.K+2:5*p.K+2, :]

        LaminartEqConv.fun_v_C!(
            ff.C,
            v_p,
            v_m,
            ff.tmp_a,
            ff.tmp_b,
            ff.tmp_A,
            ff.tmp_B,
            p,
        ) #a,b,C,D: buffers
        LaminartEqConv.fun_H_z!(ff.H_z, z, ff.tmp_A, p) #A=buffer

        @. ff.tmp_A = @view u[:, :, 1:p.K, :] #x
        @. ff.tmp_B = @view u[:, :, 2*p.K+1:3*p.K, :] #m
        LaminartEqConv.fun_x_lgn!(ff.x_lgn, ff.tmp_A, p)

        LaminartEqConv.fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.tmp_a, p) #a=buffer
        LaminartEqConv.fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.tmp_a, p) #a=buff

        LaminartEqConv.fun_dx_v1!(dx, ff.tmp_A, ff.C, z, p.x_V2, p) #A:x
        LaminartEqConv.fun_dy!(dy, y, ff.C, ff.tmp_A, ff.tmp_B, ff.tmp_C, p) # A:x, B:m, C:buffer
        LaminartEqConv.fun_dm!(dm, ff.tmp_B, ff.tmp_A, ff.tmp_C, p) # A:x, B:m, C:buffer

        @. ff.tmp_A = @view u[:, :, 4*p.K+1:5*p.K, :] #s
        LaminartEqConv.fun_dz!(dz, z, y, ff.H_z, ff.tmp_A, ff.tmp_B, p) #A=s, B=buffer
        LaminartEqConv.fun_ds!(ds, ff.tmp_A, ff.H_z, ff.tmp_B, p) #A=s, B=buffer

    end
    return nothing

end

"""
Uses equations solved at equilibrum for L6 (x) and L4 (y).
"""
mutable struct LamFunction_equ{T<:AbstractArray} <: Function
    x::T
    m::T
    s::T

    x_lgn::T
    C::T
    H_z::T

    dy_temp::T
    dm_temp::T
    dz_temp::T
    ds_temp::T
    dv_temp::T
    H_z_temp::T
    V_temp_1::T
    V_temp_2::T
    Q_temp::T
    P_temp::T
end

function (ff::LamFunction_equ)(du, u, p, t)

    @inbounds begin

        @. ff.x = @view u[:, :, 1:p.K, :]
        y = @view u[:, :, p.K+1:2*p.K, :]
        @. ff.m = @view u[:, :, 2*p.K+1:3*p.K, :]
        z = @view u[:, :, 3*p.K+1:4*p.K, :]
        @. ff.s = @view u[:, :, 4*p.K+1:5*p.K, :]

        v_p = @view u[:, :, 5*p.K+1:5*p.K+1, :]
        v_m = @view u[:, :, 5*p.K+2:5*p.K+2, :]

        dx = @view du[:, :, 1:p.K, :]
        dy = @view du[:, :, p.K+1:2*p.K, :]
        dm = @view du[:, :, 2*p.K+1:3*p.K, :]
        dz = @view du[:, :, 3*p.K+1:4*p.K, :]
        ds = @view du[:, :, 4*p.K+1:5*p.K, :]

        dv_p = @view du[:, :, 5*p.K+1:5*p.K+1, :]
        dv_m = @view du[:, :, 5*p.K+2:5*p.K+2, :]

        LaminartEqConv.fun_x_lgn!(ff.x_lgn, ff.x, p)
        LaminartEqConv.fun_v_C!(
            ff.C,
            v_p,
            v_m,
            ff.V_temp_1,
            ff.V_temp_2,
            ff.Q_temp,
            ff.P_temp,
            p,
        )
        LaminartEqConv.fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        LaminartEqConv.fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
        LaminartEqConv.fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
        #         LaminartEqConv.fun_dx_v1!(dx, ff.x, ff.C, z, p.x_V2, p)
        #         LaminartEqConv.fun_dy!(dy, y, ff.C, ff.x, ff.m, ff.dy_temp, p)
        LaminartEqConv.fun_x_equ!(ff.x, ff.C, z, ff.dy_temp, p.x_V2, p)
        LaminartEqConv.fun_y_equ!(y, ff.C, ff.x, ff.m, ff.dy_temp, p)
        LaminartEqConv.fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
        LaminartEqConv.fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
        LaminartEqConv.fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

    end
    return nothing

end

"""
All arrays held in struct and temp arrays reused for different equations.
"""
mutable struct LamFunction_all_struct_reuse{T<:AbstractArray} <: Function

    x_lgn::T
    C::T
    H_z::T

    y::Any
    z::Any
    v_p::Any
    v_m::Any
    dx::Any
    dy::Any
    dm::Any
    dz::Any
    ds::Any
    dv_p::Any
    dv_m::Any

    tmp_a::T
    tmp_b::T

    tmp_A::T
    tmp_B::T
    tmp_C::T
end

function (ff::LamFunction_all_struct_reuse)(du, u, p, t)

    @inbounds begin

        ff.y = @view u[:, :, p.K+1:2*p.K, :]
        ff.z = @view u[:, :, 3*p.K+1:4*p.K, :]

        ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1, :]
        ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2, :]

        ff.dx = @view du[:, :, 1:p.K, :]
        ff.dy = @view du[:, :, p.K+1:2*p.K, :]
        ff.dm = @view du[:, :, 2*p.K+1:3*p.K, :]
        ff.dz = @view du[:, :, 3*p.K+1:4*p.K, :]
        ff.ds = @view du[:, :, 4*p.K+1:5*p.K, :]

        ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1, :]
        ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2, :]

        LaminartEqConv.fun_v_C!(
            ff.C,
            ff.v_p,
            ff.v_m,
            ff.tmp_a,
            ff.tmp_b,
            ff.tmp_A,
            ff.tmp_B,
            p,
        ) #a,b,C,D: buffers
        LaminartEqConv.fun_H_z!(ff.H_z, ff.z, ff.tmp_A, p) #A=buffer

        @. ff.tmp_A = @view u[:, :, 1:p.K, :] #x
        @. ff.tmp_B = @view u[:, :, 2*p.K+1:3*p.K, :] #m
        LaminartEqConv.fun_x_lgn!(ff.x_lgn, ff.tmp_A, p)

        LaminartEqConv.fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, ff.tmp_a, p) #a=buffer
        LaminartEqConv.fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.tmp_a, p) #a=buff

        LaminartEqConv.fun_dx_v1!(ff.dx, ff.tmp_A, ff.C, ff.z, p.x_V2, p) #A:x
        LaminartEqConv.fun_dy!(
            ff.dy,
            ff.y,
            ff.C,
            ff.tmp_A,
            ff.tmp_B,
            ff.tmp_C,
            p,
        ) # A:x, B:m, C:buffer
        LaminartEqConv.fun_dm!(ff.dm, ff.tmp_B, ff.tmp_A, ff.tmp_C, p) # A:x, B:m, C:buffer

        @. ff.tmp_A = @view u[:, :, 4*p.K+1:5*p.K, :] #s
        LaminartEqConv.fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.tmp_A, ff.tmp_B, p) #A=s, B=buffer
        LaminartEqConv.fun_ds!(ff.ds, ff.tmp_A, ff.H_z, ff.tmp_B, p) #A=s, B=buffer

    end
    return nothing

end

"""
Uses equations with JuliaImage's imfilter for convolution on CPU.
"""
mutable struct LamFunction_imfil_cpu <: Function
    x_lgn::Any
    C::Any
    H_z::Any
    H_z_temp::Any
    v_C_temp1::Any
    v_C_temp2::Any
    v_C_tempA::Any
    W_temp::Any
end

function (ff::LamFunction_imfil_cpu)(du, u, p, t)
    @inbounds begin
        x = @view u[:, :, 1:p.K]
        y = @view u[:, :, p.K+1:2*p.K]
        m = @view u[:, :, 2*p.K+1:3*p.K]
        z = @view u[:, :, 3*p.K+1:4*p.K]
        s = @view u[:, :, 4*p.K+1:5*p.K]

        v_p = @view u[:, :, 5*p.K+1]
        v_m = @view u[:, :, 5*p.K+2]

        dx = @view du[:, :, 1:p.K]
        dy = @view du[:, :, p.K+1:2*p.K]
        dm = @view du[:, :, 2*p.K+1:3*p.K]
        dz = @view du[:, :, 3*p.K+1:4*p.K]
        ds = @view du[:, :, 4*p.K+1:5*p.K]

        dv_p = @view du[:, :, 5*p.K+1]
        dv_m = @view du[:, :, 5*p.K+2]

        LaminartEqImfilter.fun_x_lgn!(ff.x_lgn, x, p)
        LaminartEqImfilter.fun_v_C!(
            ff.C,
            v_p,
            v_m,
            ff.v_C_temp1,
            ff.v_C_temp2,
            ff.v_C_tempA,
            p,
        )
        LaminartEqImfilter.fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        LaminartEqImfilter.fun_dv!(dv_p, v_p, p.r, ff.x_lgn, p)
        LaminartEqImfilter.fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, p)
        LaminartEqImfilter.fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
        LaminartEqImfilter.fun_dy!(dy, y, ff.C, x, m, ff.W_temp, p)
        LaminartEqImfilter.fun_dm!(dm, m, x, ff.W_temp, p)
        LaminartEqImfilter.fun_dz!(dz, z, y, ff.H_z, s, p)
        LaminartEqImfilter.fun_ds!(ds, s, ff.H_z, p)
    end
    return nothing
end

"""
Uses equations with JuliaImage's imfilter for convolution on GPU with IIR algorithom.
"""
mutable struct LamFunction_imfil_gpu_iir <: Function
    x_lgn::Any
    C::Any
    H_z::Any
    H_z_temp::Any
    v_C_temp1::Any
    v_C_temp2::Any
    v_C_tempA::Any
    W_temp::Any
end

function (ff::LamFunction_imfil_gpu_iir)(du, u, p, t)
    @inbounds begin
        x = @view u[:, :, 1:p.K]
        y = @view u[:, :, p.K+1:2*p.K]
        m = @view u[:, :, 2*p.K+1:3*p.K]
        z = @view u[:, :, 3*p.K+1:4*p.K]
        s = @view u[:, :, 4*p.K+1:5*p.K]

        v_p = @view u[:, :, 5*p.K+1]
        v_m = @view u[:, :, 5*p.K+2]

        dx = @view du[:, :, 1:p.K]
        dy = @view du[:, :, p.K+1:2*p.K]
        dm = @view du[:, :, 2*p.K+1:3*p.K]
        dz = @view du[:, :, 3*p.K+1:4*p.K]
        ds = @view du[:, :, 4*p.K+1:5*p.K]

        dv_p = @view du[:, :, 5*p.K+1]
        dv_m = @view du[:, :, 5*p.K+2]

        LaminartEqImfilterGPU_IIR.fun_x_lgn!(ff.x_lgn, x, p)
        LaminartEqImfilterGPU_IIR.fun_v_C!(
            ff.C,
            v_p,
            v_m,
            ff.v_C_temp1,
            ff.v_C_temp2,
            ff.v_C_tempA,
            p,
        )
        LaminartEqImfilterGPU_IIR.fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        LaminartEqImfilterGPU_IIR.fun_dv!(dv_p, v_p, p.r, ff.x_lgn, p)
        LaminartEqImfilterGPU_IIR.fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, p)
        LaminartEqImfilterGPU_IIR.fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
        LaminartEqImfilterGPU_IIR.fun_dy!(dy, y, ff.C, x, m, ff.W_temp, p)
        LaminartEqImfilterGPU_IIR.fun_dm!(dm, m, x, ff.W_temp, p)
        LaminartEqImfilterGPU_IIR.fun_dz!(dz, z, y, ff.H_z, s, p)
        LaminartEqImfilterGPU_IIR.fun_ds!(ds, s, ff.H_z, p)
    end
    return nothing
end

"""
Uses equations with JuliaImage's imfilter for convolution on GPU with FIR algorithom.
"""
mutable struct LamFunction_imfil_gpu_fir <: Function
    x_lgn::Any
    C::Any
    H_z::Any
    H_z_temp::Any
    v_C_temp1::Any
    v_C_temp2::Any
    v_C_tempA::Any
    W_temp::Any
end

function (ff::LamFunction_imfil_gpu_fir)(du, u, p, t)
    @inbounds begin
        x = @view u[:, :, 1:p.K]
        y = @view u[:, :, p.K+1:2*p.K]
        m = @view u[:, :, 2*p.K+1:3*p.K]
        z = @view u[:, :, 3*p.K+1:4*p.K]
        s = @view u[:, :, 4*p.K+1:5*p.K]

        v_p = @view u[:, :, 5*p.K+1]
        v_m = @view u[:, :, 5*p.K+2]

        dx = @view du[:, :, 1:p.K]
        dy = @view du[:, :, p.K+1:2*p.K]
        dm = @view du[:, :, 2*p.K+1:3*p.K]
        dz = @view du[:, :, 3*p.K+1:4*p.K]
        ds = @view du[:, :, 4*p.K+1:5*p.K]

        dv_p = @view du[:, :, 5*p.K+1]
        dv_m = @view du[:, :, 5*p.K+2]

        LaminartEqImfilterGPU_FIR.fun_x_lgn!(ff.x_lgn, x, p)
        LaminartEqImfilterGPU_FIR.fun_v_C!(
            ff.C,
            v_p,
            v_m,
            ff.v_C_temp1,
            ff.v_C_temp2,
            ff.v_C_tempA,
            p,
        )
        LaminartEqImfilterGPU_FIR.fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        LaminartEqImfilterGPU_FIR.fun_dv!(dv_p, v_p, p.r, ff.x_lgn, p)
        LaminartEqImfilterGPU_FIR.fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, p)
        LaminartEqImfilterGPU_FIR.fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
        LaminartEqImfilterGPU_FIR.fun_dy!(dy, y, ff.C, x, m, ff.W_temp, p)
        LaminartEqImfilterGPU_FIR.fun_dm!(dm, m, x, ff.W_temp, p)
        LaminartEqImfilterGPU_FIR.fun_dz!(dz, z, y, ff.H_z, s, p)
        LaminartEqImfilterGPU_FIR.fun_ds!(ds, s, ff.H_z, p)
    end
    return nothing
end

"""
Uses equations with JuliaImage's imfilter for convolution on GPU with FFT algorithom.
"""
mutable struct LamFunction_imfil_gpu_fft <: Function
    x_lgn::Any
    C::Any
    H_z::Any
    H_z_temp::Any
    v_C_temp1::Any
    v_C_temp2::Any
    v_C_tempA::Any
    W_temp::Any
end

function (ff::LamFunction_imfil_gpu_fft)(du, u, p, t)
    @inbounds begin
        x = @view u[:, :, 1:p.K]
        y = @view u[:, :, p.K+1:2*p.K]
        m = @view u[:, :, 2*p.K+1:3*p.K]
        z = @view u[:, :, 3*p.K+1:4*p.K]
        s = @view u[:, :, 4*p.K+1:5*p.K]

        v_p = @view u[:, :, 5*p.K+1]
        v_m = @view u[:, :, 5*p.K+2]

        dx = @view du[:, :, 1:p.K]
        dy = @view du[:, :, p.K+1:2*p.K]
        dm = @view du[:, :, 2*p.K+1:3*p.K]
        dz = @view du[:, :, 3*p.K+1:4*p.K]
        ds = @view du[:, :, 4*p.K+1:5*p.K]

        dv_p = @view du[:, :, 5*p.K+1]
        dv_m = @view du[:, :, 5*p.K+2]

        LaminartEqImfilterGPU_FFT.fun_x_lgn!(ff.x_lgn, x, p)
        LaminartEqImfilterGPU_FFT.fun_v_C!(
            ff.C,
            v_p,
            v_m,
            ff.v_C_temp1,
            ff.v_C_temp2,
            ff.v_C_tempA,
            p,
        )
        LaminartEqImfilterGPU_FFT.fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        LaminartEqImfilterGPU_FFT.fun_dv!(dv_p, v_p, p.r, ff.x_lgn, p)
        LaminartEqImfilterGPU_FFT.fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, p)
        LaminartEqImfilterGPU_FFT.fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
        LaminartEqImfilterGPU_FFT.fun_dy!(dy, y, ff.C, x, m, ff.W_temp, p)
        LaminartEqImfilterGPU_FFT.fun_dm!(dm, m, x, ff.W_temp, p)
        LaminartEqImfilterGPU_FFT.fun_dz!(dz, z, y, ff.H_z, s, p)
        LaminartEqImfilterGPU_FFT.fun_ds!(ds, s, ff.H_z, p)
    end
    return nothing
end

# mutable struct LamFunction_gpu30 <: Function
# 	x_lgn::AbstractArray
# 	C::AbstractArray
# 	H_z::AbstractArray

# 	dy_temp::AbstractArray
# 	dm_temp::AbstractArray
# 	dz_temp::AbstractArray
# 	ds_temp::AbstractArray
# 	dv_temp::AbstractArray
# 	H_z_temp::AbstractArray
# 	V_temp_1::AbstractArray
# 	V_temp_2::AbstractArray
# 	Q_temp::AbstractArray
# 	P_temp::AbstractArray
# end

# function (ff::LamFunction_gpu30)(du, u, p, t)

#        @inbounds begin

# 		ff.x = @view u[:, :, 1:p.K,:]
#         ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
# 		@. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# 		ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		ff.dx = @view du[:, :, 1:p.K,:]
#         ff.dy = @view du[:, :, p.K+1:2*p.K,:]
#         ff.dm = @view du[:, :, 2*p.K+1:3*p.K,:]
#         ff.dz = @view du[:, :, 3*p.K+1:4*p.K,:]
#         ff.ds = @view du[:, :, 4*p.K+1:5*p.K,:]

# 		ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# 		ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# 		fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
#         fun_dm!(ff.dm, ff.m, ff.x, ff.dm_temp, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)

#     end
#     return nothing

# end

# mutable struct LamFunction_gpu14{T <: CuArray{Float32,4,Nothing}} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# # 	dx::T
# # 	dy::T
# # 	dm::T
# # 	dz::T
# # 	ds::T
# # 	dv_p::T
# # 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T

# 	dy_temp::T
# 	dm_temp::T
# 	dz_temp::T
# 	ds_temp::T
# 	dv_temp::T
# 	H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# function (ff::LamFunction_gpu14)(du, u, p, t)

#        @inbounds begin
# #         @. ff.x = @view u[:, :, 1:p.K,:]
# #         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
# #         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# #        @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# # 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		x = @view u[:, :, 1:p.K,:]
#         y = @view u[:, :, p.K+1:2*p.K,:]
# #         m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         z = @view u[:, :, 3*p.K+1:4*p.K,:]
# #         s = @view u[:, :, 4*p.K+1:5*p.K,:]

#        v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		 dx = @view du[:, :, 1:p.K,:]
#         dy = @view du[:, :, p.K+1:2*p.K,:]
#         dm = @view du[:, :, 2*p.K+1:3*p.K,:]
#         dz = @view du[:, :, 3*p.K+1:4*p.K,:]
#         ds = @view du[:, :, 4*p.K+1:5*p.K,:]

#        dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# 		dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# #         fun_x_lgn!(ff.x_lgn, ff.x, p)
# #         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
# #         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

# #         fun_dv!(dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dv!(dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dx_v1!(dx, ff.x, ff.C, ff.z, p.x_V2, p)
# #         fun_dy!(dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
# #         fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
# #         fun_dz!(dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
# #         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		fun_x_lgn!(ff.x_lgn, x, p)
#         fun_v_C!(ff.C, v_p, v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

#         fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
#         fun_dy!(dy, y, ff.C, x, ff.m, ff.dy_temp, p)
#         fun_dm!(dm, ff.m, x, ff.dm_temp, p)
#         fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
#         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		@. du[:, :, 1:p.K,:] = ff.dx
# # 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# # 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# # 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# # 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

# #         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# # 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing

# end

# mutable struct LamFunction_gpu14{T <: CuArray{Float32,4,Nothing}} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# # 	dx::T
# # 	dy::T
# # 	dm::T
# # 	dz::T
# # 	ds::T
# # 	dv_p::T
# # 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T

# 	dy_temp::T
# 	dm_temp::T
# 	dz_temp::T
# 	ds_temp::T
# 	dv_temp::T
# 	H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# function (ff::LamFunction_gpu14)(du, u, p, t)

#        @inbounds begin
# #         @. ff.x = @view u[:, :, 1:p.K,:]
# #         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
# #         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# #        @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# # 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		x = @view u[:, :, 1:p.K,:]
#         y = @view u[:, :, p.K+1:2*p.K,:]
# #         m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         z = @view u[:, :, 3*p.K+1:4*p.K,:]
# #         s = @view u[:, :, 4*p.K+1:5*p.K,:]

#        v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		 dx = @view du[:, :, 1:p.K,:]
#         dy = @view du[:, :, p.K+1:2*p.K,:]
#         dm = @view du[:, :, 2*p.K+1:3*p.K,:]
#         dz = @view du[:, :, 3*p.K+1:4*p.K,:]
#         ds = @view du[:, :, 4*p.K+1:5*p.K,:]

#        dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# 		dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# #         fun_x_lgn!(ff.x_lgn, ff.x, p)
# #         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
# #         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

# #         fun_dv!(dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dv!(dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dx_v1!(dx, ff.x, ff.C, ff.z, p.x_V2, p)
# #         fun_dy!(dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
# #         fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
# #         fun_dz!(dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
# #         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		fun_x_lgn!(ff.x_lgn, x, p)
#         fun_v_C!(ff.C, v_p, v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

#         fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
#         fun_dy!(dy, y, ff.C, x, ff.m, ff.dy_temp, p)
#         fun_dm!(dm, ff.m, x, ff.dm_temp, p)
#         fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
#         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		@. du[:, :, 1:p.K,:] = ff.dx
# # 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# # 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# # 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# # 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

# #         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# # 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing

# end

# mutable struct LamFunction_gpu2{T <: CuArray{Float32,4,Nothing}} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# 	dx::T
# 	dy::T
# 	dm::T
# 	dz::T
# 	ds::T
# 	dv_p::T
# 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T
# 	dy_temp::T
# 	ds_temp::T
# 	H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# mutable struct LamFunction_cpu2{T <: AbstractArray} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# 	dx::T
# 	dy::T
# 	dm::T
# 	dz::T
# 	ds::T
# 	dv_p::T
# 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T
# 	ds_temp::T
# 		H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# mutable struct LamFunction{T} <: Function
# 	x::AbstractArray
# 	y::AbstractArray
# 	m::AbstractArray
# 	z::AbstractArray
# 	s::AbstractArray
# 	v_p::AbstractArray
# 	v_m::AbstractArray

# 	dx::AbstractArray
# 	dy::AbstractArray
# 	dm::AbstractArray
# 	dz::AbstractArray
# 	ds::AbstractArray
# 	dv_p::AbstractArray
# 	dv_m::AbstractArray

# 	x_lgn::AbstractArray
# 	C::AbstractArray
# 	H_z::AbstractArray
# 	V_temp_1::AbstractArray
# 	V_temp_2::AbstractArray
# 	Q_temp::AbstractArray
# 	P_temp::AbstractArray
# end

# function (ff::LamFunction_cpu2)(du, u, p, t)

#     @inbounds begin
#         @. ff.x = @view u[:, :, 1:p.K,:]
#         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

#         @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

#         fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, p)
#         fun_dm!(ff.dm, ff.m, ff.x, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)

# 		@. du[:, :, 1:p.K,:] = ff.dx
# 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

#         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing
# end

# function (ff::LamFunction_gpu2)(du, u, p, t)

#        @inbounds begin
#         @. ff.x = @view u[:, :, 1:p.K,:]
#         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

#         @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

#         fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
#         fun_dm!(ff.dm, ff.m, ff.x, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)

# 		@. du[:, :, 1:p.K,:] = ff.dx
# 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

#         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing

# end

# mutable struct LamFunction_04{T} <: Function
# x::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# y::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# z::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# s::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# v_p::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# v_m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dx::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dy::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dm::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dz::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# ds::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dv_p::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dv_m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	x_lgn::CuArray{Float32,4,Nothing}
# 	C::CuArray{Float32,4,Nothing}
# 	H_z::CuArray{Float32,4,Nothing}
# 	V_temp_1::CuArray{Float32,4,Nothing}
# 	V_temp_2::CuArray{Float32,4,Nothing}
# 	Q_temp::CuArray{Float32,4,Nothing}
# 	P_temp::CuArray{Float32,4,Nothing}
# end

# mutable struct LamFunction_10 <: Function
# 	x
# 	y
# 	m
# 	z
# 	s
# 	v_p
# 	v_m
# 	dx
# 	dy
# 	dm
# 	dz
# 	ds
# 	dv_p
# 	dv_m
# 	x_lgn
# 	C
# 	H_z
# 	V_temp_1
# 	V_temp_2
# 	Q_temp
# 	P_temp
# end

# mutable struct LamFunction_13<: Function
# 	x::CuArray{Float32,4,Nothing}
# 	y::CuArray{Float32,4,Nothing}
# 	m::CuArray{Float32,4,Nothing}
# 	z::CuArray{Float32,4,Nothing}
# 	s::CuArray{Float32,4,Nothing}
# 	v_p::CuArray{Float32,4,Nothing}
# 	v_m::CuArray{Float32,4,Nothing}
# dx::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dy::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dm::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dz::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	ds::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dv_p::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dv_m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	x_lgn::CuArray{Float32,4,Nothing}
# 	C::CuArray{Float32,4,Nothing}
# 	H_z::CuArray{Float32,4,Nothing}
# 	V_temp_1::CuArray{Float32,4,Nothing}
# 	V_temp_2::CuArray{Float32,4,Nothing}
# 	Q_temp::CuArray{Float32,4,Nothing}
# 	P_temp::CuArray{Float32,4,Nothing}
# end

# function (ff::LamFunction_13)(du, u, p, t)

#     @inbounds begin
# #         ff.x = @view u[:, :, 1:p.K,:]
# #         ff.y = @view u[:, :, p.K+1:2*p.K,:]
# #         ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
# #         ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
# #         ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# #         ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# # 		ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		@. ff.x = @view u[:, :, 1:p.K,:]
#         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

#         @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# # 		ff.dx = @view du[:, :, 1:p.K,:]
# # 		ff.dy = @view du[:, :, p.K+1:2*p.K,:]
# # 		ff.dm = @view du[:, :, 2*p.K+1:3*p.K,:]
# # 		ff.dz = @view du[:, :, 3*p.K+1:4*p.K,:]
# # 		ff.ds = @view du[:, :, 4*p.K+1:5*p.K,:]

# # 		ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# # 		ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# # 		ff.x_lgn = @view ff.x_lgn[:,:,:,:]
# # 		ff.C_ = @view ff.C[:,:,:,:]
# # 		ff.H_z_ = @view ff.H_z[:,:,:,:]
# # 		ff.V_temp_1_ = @view ff.V_temp_1[:,:,:,:]
# # 		ff.V_temp_2_ = @view ff.V_temp_2[:,:,:,:]
# # 		ff.Q_temp_ = @view ff.Q_temp[:,:,:,:]
# # 		ff.P_temp_ = @view ff.P_temp[:,:,:,:]

#         fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
#         fun_dm!(ff.dm, ff.m, ff.x, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, p)

# # 		@. du[:, :, 1:p.K,:] = ff.dx
# # 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# # 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# # 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# # 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

# #         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# # 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing
# end

end


text:
module RegNeuralDE

# StdLib
using Statistics, LinearAlgebra
# DiffEq Packages
using DiffEqFlux, OrdinaryDiffEq, DiffEqCallbacks, DiffEqSensitivity, StochasticDiffEq
# Neural Networks
using Flux, CUDA
# Probabilistic Stuff
using Distributions, DistributionsAD
# AD Packages
using Tracker
using Tracker: data, TrackedReal
# Data Processing
using MLDatasets, MLDataUtils, BSON, NPZ, DataLoaders
import LearnBase
# Helper Modules
using MacroTools: @forward
using Format

# These should be in Tracker.jl
## Track the parameters
track(m) = fmap(x -> x isa AbstractArray ? Tracker.param(x) : x, m)
untrack(m) = fmap(Tracker.data, m)

# Hack to get around scalar indexing issue in gradients for FFJORD
# See https://github.com/JuliaGPU/Adapt.jl/issues/21
const Wrapper = Union{Base.ReshapedArray,LinearAlgebra.Transpose,LinearAlgebra.Adjoint}
function CUDA.CuArray(x::Base.ReshapedArray{<:Any,<:Any,<:Wrapper})
    xp = CuArray(x.parent)
    ra = Base.ReshapedArray(xp, x.dims, x.mi)
    CuArray(ra)
end

function CUDA.CuArray(x::LinearAlgebra.Transpose{<:Any,<:Wrapper})
    xp = CuArray(x.parent)
    ra = Base.Transpose(xp)
    CuArray(ra)
end

function CUDA.CuArray(x::LinearAlgebra.Adjoint{<:Any,<:Wrapper})
    xp = CuArray(x.parent)
    ra = Base.Adjoint(xp)
    CuArray(ra)
end

Base.convert(
    ::Type{CuArray{Float32,2}},
    x::Base.ReshapedArray{
        Float32,
        2,
        Transpose{Float32,CuArray{Float32,2}},
        Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}},
    },
) = CuArray{Float32,2}(x)

Base.convert(
    ::Type{CuArray{Float32,2}},
    x::Base.ReshapedArray{
        Float32,
        2,
        Adjoint{Float32,CuArray{Float32,2}},
        Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}},
    },
) = CuArray{Float32,2}(x)

# Include code
include("dataset.jl")
include("utils.jl")
include("models/basic.jl")
include("models/neural_ode.jl")
include("models/neural_sde.jl")
include("models/supervised_classification.jl")
include("models/time_series.jl")
include("models/ffjord.jl")
include("metrics.jl")

# Export functions
export load_mnist, load_physionet, load_spiral2d, load_miniboone, load_gaussian_mixture
export TrackedNeuralODE,
    TrackedNeuralDSDE,
    ClassifierNODE,
    ClassifierNSDE,
    TDChain,
    RecognitionRNN,
    TrackedFFJORD,
    LatentTimeSeriesModel
export track, untrack, table_logger, solution, sample, update_parameters!

end

ext: jl


text:
function simulate!(market::Market;steps::Int=100,tbb=false,double_ending_price_premium::Union{Nothing,Float64}=nothing)
    for stepno in 1:steps
        if length(market.available_homes) == 0 || length(market.available_buyers) == 0
            break
        end
        for j in shuffle(collect(market.available_buyers))
            try_place_a_bid!(market, j, stepno, tbb, double_ending_price_premium)
        end
        if stepno < 3 continue end
        @assert (length(market.double_ended_home_buyer)==0) || double_ending_price_premium !=nothing
        for i in shuffle(collect(keys(market.double_ended_home_buyer)))
            try_process_sale!(market, i, stepno,double_ending_price_premium, true)
        end
        for i in shuffle(collect(market.available_homes))
            try_process_sale!(market, i, stepno,double_ending_price_premium, false)
        end
    end
    market.transactions
end

function try_place_a_bid!(market::Market, j::Int,stepno::Int,tbb::Bool, double_ending_price_premium::Union{Nothing,Float64})
    buyer = market.buyers[j]
    i = rand(market.available_homes)
    seller = market.homeowners[i]
    if !(i in buyer.seen_homes)
        push!(buyer.seen_homes, i)
        push!(seller.seen_by, j)
        best_other_bid = seller.r
        if length(seller.bids)>0
            best_other_bid, some_bidder = findmax(seller.bids)
            if some_bidder == j
                return #do not overbid yourself
            end
        end

        buyersvalue = buyer.v[i]
        if double_ending_price_premium != nothing && buyer.ak == seller.ak && !tbb
            buyersvalue *= (1+double_ending_price_premium)
        end

        if buyersvalue > seller.r+1
            bid = best_other_bid+1
            if double_ending_price_premium != nothing && buyer.ak == seller.ak && !tbb
                bid = buyersvalue-1
                if  (!(i in keys(market.double_ended_home_buyer)) ||
                        seller.bids[market.double_ended_home_buyer[i]] < bid) &&
                        !(j in values(market.double_ended_home_buyer))
                    market.double_ended_home_buyer[i] = j
                end
                #double ending --> bid highest possible amount
            else
                #bid = round(rand(Uniform(seller.r+1.,buyer.v[i])))
                if best_other_bid+1. >= buyersvalue-1
                    bid = buyersvalue-1
                else
                    if tbb
                        # has market info - will bid between other bid and her value
                        #bid = max(bid,best_other_bid+1)
                        bid = round((best_other_bid+buyersvalue)/2)
                    else
                        #has no market info - will overbid
                        #bid = round(rand(Uniform(best_other_bid+1.,buyer.v[i]-1)))
                        bid = round(rand(TriangularDist(seller.r,buyersvalue-1,buyersvalue-1)))
                    end
                end

            end
            #need to bid lower than the perceived property value
            @assert bid <= buyersvalue
            seller.bids[j]=bid
            push!(seller.bid_history,(j,bid))
            buyer.bids[i] = bid
            #seller.i==2 && println("Added bid $bid from j=$j for i=$i")
        end
    end
    if length(buyer.bids) >= 3
        #the buyer learns from the market and removes the worst bid
        #if it is below 1 std
        i_s = collect(keys(buyer.bids))
        filter!(i -> buyer.ak != market.homeowners[i].ak, i_s)
        if length(i_s) >= 3
            gains = [(buyer.v[i] - buyer.bids[i]) for i in i_s]
            amount, i_in_is = findmin(gains)
            if amount < mean(gains) - std(gains)
                @assert i_s[i_in_is] in keys(buyer.bids)
                delete!(buyer.bids, i_s[i_in_is])
                @assert j in keys(market.homeowners[i_s[i_in_is]].bids)
                delete!(market.homeowners[i_s[i_in_is]].bids, j)
                if get(market.double_ended_home_buyer, i_s[i_in_is], -1) == j
                    delete!(market.double_ended_home_buyer, i_s[i_in_is])
                end
            end
        end
    end
end

function try_process_sale!(market::Market, i::Int,stepno::Int,double_ending_price_premium::Union{Nothing,Float64},doubleend::Bool=false)
    seller = market.homeowners[i]
    length(seller.bids) == 0 && return nothing
    j::Int = -1
    if doubleend
        j = market.double_ended_home_buyer[i]
        #global m_debug = market
        #println("i=$i, j=$j, stepno=$stepno")
        best_bid = seller.bids[j]
    else
        best_bid, j = findmax(seller.bids)
    end
    @assert best_bid >= seller.r
    avg_real_val = mean(market.buyers[bj].v[i] for bj in market.available_buyers)
    if doubleend || (length(seller.bids) > rand((stepno>40 ? 0 : 4):(stepno>50 ? 1 : 8)) && best_bid >= avg_real_val*0.25)
        @assert i == seller.i
        delete!(market.available_homes, i)
        delete!(market.available_buyers, j)
        delete!(market.double_ended_home_buyer, i)
        for otherseller in market.homeowners
            if otherseller.i != i
                delete!(otherseller.bids, j)
            end
        end
        trans = Dict{Symbol,Union{Int,Float64,Bool,Missing}}(
            :step=>stepno,
            :i => i,
            :j => j,
            :k_s => market.homeowners[i].ak,
            :k_b => market.buyers[j].ak,
            :allow_de => double_ending_price_premium != nothing,
            :de_price_premium => something(double_ending_price_premium, missing),
            :seen_s=>length(market.homeowners[i].seen_by),
            :seen_b=>length(market.buyers[j].seen_homes),
            :cost_sa=>market.agents[market.homeowners[i].ak].cs,
            :cost_ba=>market.agents[market.buyers[j].ak].cb,
            :prov_sa=>market.agents[market.homeowners[i].ak].ds,
            :prov_ba=>market.agents[market.buyers[j].ak].db,
            :price_r=>market.homeowners[i].r,
            :price_v=>market.buyers[j].v[i],
            :double_ended=> double_ending_price_premium != nothing && (market.homeowners[i].ak==market.buyers[j].ak),
            :price=>best_bid
            )
        push!(market.transactions, trans)
        #seller.i==2 && println(trans)
    end
end

function res_agg(resdf::AbstractDataFrame, p::MarketParams; kwargs...)
    stat_cols = [:transactions, :price, :price_v, :gain_s, :gain_b, :gain_tot,
                 :double_ended, :property_visits, :step, :step_min, :step_max]

    ress = [
        :TBB => resdf[resdf.tbb .& (.!resdf.allow_de),:],
        :nDE => resdf[(.!resdf.tbb) .& (.!resdf.allow_de),:],
    ]

    de_df = resdf[(.!resdf.tbb) .& resdf.allow_de,:]
    de_regimes = unique(de_df.de_price_premium)
    for de_regime in de_regimes
        label = Symbol(string("DE_",replace(string(de_regime),'.'=>'_')))
        push!(ress, label => de_df[de_df.de_price_premium .== de_regime,:])
    end

    rs = DataFrame()
    for (label, res) in ress
        r = DataFrame(n_sims=nrow(res), regime=label; kwargs...)
        for ff in fieldnames(typeof(p))
            r[!,ff] = [getfield(p,ff)]
        end
        for col in stat_cols
            r[!,col] = [mean(res[!,col])]
        end
        for col in stat_cols
            r[!,Symbol(string(col,"_std"))] = [std(res[!,col])]
        end
        append!(rs,r)
    end
    rs
end

function printsimresults(res)
    r_tbb_nd = res[res.tbb .& (.!res.allow_de),:]
    r_nbb_de = res[(.!res.tbb) .& res.allow_de .& res.de_price_premium==0.0,:]
    r_nbb_nd = res[(.!res.tbb) .& (.!res.allow_de),:]
    stat_cols = [:transactions, :price, :price_v, :gain_s, :gain_b, :gain_tot,
                 :double_ended, :property_visits, :step, :step_min, :step_max]
    stat_list = [:mean, :min, :q25, :median, :q75, :max, :std]

    d_nbb_de=describe(r_nbb_de[!,stat_cols], stat_list... )
    d_nbb_nd=describe(r_nbb_nd[!,stat_cols], stat_list... )
    d_tbb_nd=describe(r_tbb_nd[!,stat_cols], stat_list... )

    println()
    println("TBB & no doubleending");println(d_tbb_nd)
    println("no TBB & doubleending 0.0");println(d_nbb_de)
    println("no TBB & no doubleending");println(d_nbb_nd)

    @assert d_nbb_de.variable == d_nbb_nd.variable == d_tbb_nd.variable
    d_tot = DataFrame(param=d_nbb_de.variable, nbb_de=d_nbb_de.mean,
                        nbb_ndd=d_nbb_nd.mean, TBB_nd=d_tbb_nd.mean)
    println("Total:")
    println(d_tot)
    d_tot
end

function runsims(p::MarketParams;nm=30, nN=30)
    simres = DataFrame(m_id=Int[],N=Int[],tbb=Bool[],
                    allow_de=Bool[], de_price_premium=Union{Float64,Missing}[],
                    nA=Int[],nH=Int[],nB=Int[],transactions=Int[],
                    profit_sa=Float64[],profit_ba=Float64[],
                    price=Float64[],price_v=Float64[],
                    gain_s=Float64[],gain_b=Float64[],gain_tot=Float64[],
                    double_ended=Float64[],property_visits=Float64[],
                    step=Float64[],step_min=Float64[],step_max=Float64[])
    for mm in 1:nm
        m_base = Market(mm,p)
        for (double_ending_price_premium, tbb) in [(0.2,false), (0.1,false), (0.0,false), (nothing,false), (nothing, true)]
            for N in 1:nN
                market = deepcopy(m_base)
                Random.seed!(100_0000+mm);
                res = simulate!(market,steps=60,tbb=tbb,double_ending_price_premium=double_ending_price_premium)
                r  = Dict{Symbol,Union{Int,Float64,Bool,Missing}}(
                            :m_id=>mm,
                            :N=>N,
                            :tbb=>tbb,
                            :allow_de=>double_ending_price_premium != nothing,
                            :de_price_premium => something(double_ending_price_premium, missing),
                            :nA=>length(market.agents),
                            :nH=>length(market.homeowners),
                            :nB=>length(market.buyers),
                            :transactions=>nrow(res),
                            :profit_sa=> nrow(res)==0 ? 0 : mean((res.prov_sa .* res.price) .- (res.cost_sa .* res.seen_s)),
                            :profit_ba=> nrow(res)==0 ? 0 : mean((res.prov_ba .* res.price) .- (res.cost_ba .* res.seen_b)),
                            :price=> nrow(res)==0 ? 0 : mean(res.price),
                            :price_v=> nrow(res)==0 ? 0 : mean(res.price_v),
                            :gain_s=> nrow(res)==0 ? 0 : mean(res.price .- res.price_r),
                            :gain_b=> nrow(res)==0 ? 0 : mean(res.price_v .- res.price),
                            :gain_tot=> nrow(res)==0 ? 0 : mean( res.price_v .- res.price_r ),
                            :double_ended=> nrow(res)==0 ? 0 : mean(res.double_ended),
                            :property_visits=> nrow(res)==0 ? 0 : mean(res.seen_s),
                            :step=> nrow(res)==0 ? 0 : mean(res.step),
                            :step_min=> nrow(res)==0 ? 0 : minimum(res.step),
                            :step_max=> nrow(res)==0 ? 0 : maximum(res.step),
                            )
                push!(simres, r)
            end
        end
    end
    simres
end


text:
export @fermiview

macro fermiview()
    quote
        try
            let mol = Fermi.Options.get("molstring")
                E = Ensamble()
                Alchemy.xyz_to!(E, mol)
                Alchemy.find_bonds!(E)
                Alchemy.run(E)
            end
        catch UndefVarError
            @error "Fermi.jl could not be found, try running `using Fermi` first."
        end
    end |> esc
end


text:
using WindowAbstractions
using Test

cb = WindowCallbacks(
    on_expose=(details) -> 1,
    on_resize=(details) -> 2,
    on_pointer_move=(details) -> 3,
    on_pointer_enter=(details) -> 4,
    on_pointer_leave=(details) -> 5,
    on_key_pressed=(details) -> 6,
    on_key_released=(details) -> 7,
    on_mouse_button_pressed=(details) -> 8,
    on_mouse_button_released=(details) -> 9,
)

struct FakeWindow <: AbstractWindow end
struct FakeWindowManager <: AbstractWindowManager end

win = FakeWindow()
wm = FakeWindowManager()

WindowAbstractions.callbacks(wm::FakeWindowManager, _) = cb

not_implemented_for_window_exc = ErrorException("Not implemented for $win")
not_implemented_for_window_manager_exc = ErrorException("Not implemented for $wm")

@testset "WindowAbstractions.jl" begin
    @testset "Key Combinations" begin
        @test key"ctrl+z" == KeyCombination('z', KeyModifierState(ctrl=true))
        @test key"alt+z" == KeyCombination('z', KeyModifierState(alt=true))
        @test key"super+z" == KeyCombination('z', KeyModifierState(super=true))
        @test key"z" == KeyCombination('z', KeyModifierState())
        @test key"ctrl+shift+z" == KeyCombination('z', KeyModifierState(ctrl=true, shift=true))
        @test key"ctrl+shift+alt+super+z" == KeyCombination('z', KeyModifierState(ctrl=true, shift=true, super=true, alt=true))
        @test key"ctrl+z" == key"CTRL+Z"
    end

    @testset "Callbacks" begin
        @test execute_callback(wm, EventDetails(ExposeEvent((1, 1), 1), (1, 1), 1.0, win)) == 1
        @test execute_callback(wm, EventDetails(ResizeEvent((1, 1)), (1, 1), 1.0, win)) == 2
        @test execute_callback(wm, EventDetails(PointerMovesEvent(MouseState(), KeyModifierState()), (1, 1), 1.0, win)) == 3
        @test execute_callback(wm, EventDetails(PointerEntersWindowEvent(), (1, 1), 1.0, win)) == 4
        @test execute_callback(wm, EventDetails(PointerLeavesWindowEvent(), (1, 1), 1.0, win)) == 5
        @test execute_callback(wm, EventDetails(KeyEvent(:Z02, KeySymbol(:z), 'z', KeyModifierState(ctrl=true), KeyPressed()), (1, 1), 1.0, win)) == 6
        @test execute_callback(wm, EventDetails(KeyEvent(:Z02, KeySymbol(:z), 'z', KeyModifierState(ctrl=true), KeyReleased()), (1, 1), 1.0, win)) == 7
        @test execute_callback(wm, EventDetails(MouseEvent(ButtonLeft(), MouseState(), ButtonPressed()), (1, 1), 1.0, win)) == 8
        @test execute_callback(wm, EventDetails(MouseEvent(ButtonLeft(), MouseState(), ButtonReleased()), (1, 1), 1.0, win)) == 9
    end

    @testset "AbstractWindow Interface" begin
        @test_throws not_implemented_for_window_exc extent(win)
        @test_throws not_implemented_for_window_exc dpi(win)
        @test_throws not_implemented_for_window_exc set_extent(win, (1, 1))
        @test_throws not_implemented_for_window_exc map_window(win)
        @test_throws not_implemented_for_window_exc unmap_window(win)
        @test_throws not_implemented_for_window_exc set_title(win, "title")
        @test_throws not_implemented_for_window_exc set_icon(win, nothing)
        @test_throws not_implemented_for_window_exc set_icon_title(win, "icon_title")
        @test_throws not_implemented_for_window_exc attach_graphics_context!(win, nothing)
        @test_throws not_implemented_for_window_manager_exc terminate_window!(wm, win)
        @test_throws not_implemented_for_window_manager_exc get_window(wm, nothing)
        @test_throws not_implemented_for_window_manager_exc get_window_symbol(wm, nothing)
        @test_throws not_implemented_for_window_manager_exc wait_for_event(wm)
        @test_throws not_implemented_for_window_manager_exc poll_for_event(wm)
    end
end


text:
function solve()
    n = parse(Int, readline())
    ls = [parse(Int, x) for x in split(readline())]
    sort!(ls)

    ret = 0
    for i=1:n-2, j=i+1:n-1, k=j+1:n
        if ls[i] + ls[j] > ls[k] && ls[i] ≠ ls[j] && ls[j] ≠ ls[k]
            ret += 1
        end
    end
    ret
end

println(solve())

max_stars_repo_name: aishikawa/atcoder-julia

max_issues_repo_name: aishikawa/atcoder-julia

max_forks_repo_name: aishikawa/atcoder-julia


text:
# # Example 7
# This example is from the book Princeton Handbook of Test Problems in Local and Global Optimization
# Dempe, Chapter 9.3.2 -parg 223 [url](https://www.springer.com/gp/book/9780792358015)

# Here, only the second level is described

# Model of the problem
# First level
# ```math
# \min 2x_1+2x_2-3y_1-3y_2-60,\\
# \notag s.t.\\
# x_1 + x_2 + y_1 -2y_2 -40\leq 0,\\
# 0 \leq x_i \leq 50, \forall i \in I,\\
# -10 \leq y_j \leq 20, \forall j \in J,\\
# ```
# Second level
# ```math
# \min (-x_1 + y_1 + 40)^2 + (-x_2 + y_2 + 20)^2,\\
# \notag s.t.\\
# - x_i + 2y_j <= -10,\forall (i,j) \in \{(i,j)|i\in I, j\in J, i=j\},\\
# -10 \leq y_j \leq 20, \forall j \in J.\\
# ```

using BilevelJuMP
using Ipopt
using JuMP
using Test

model = BilevelModel(Ipopt.Optimizer, mode = BilevelJuMP.ProductMode(1e-9))

# First we need to create all of the variables in the upper and lower problems:

# Upper level variables
@variable(Upper(model), x[i=1:2], start = 0)

# Lower level variables
@variable(Lower(model), y[i=1:2], start = -10)

# Then we can add the objective and constraints of the upper problem:

# Upper level objecive function
@objective(Upper(model), Min,2x[1] + 2x[2] -3y[1] - 3y[2] -60)

# Upper level constraints
@constraint(Upper(model), x[1] + x[2] + y[1] -2y[2] -40 <= 0)
@constraint(Upper(model), [i=1:2], x[i] >= 0)
@constraint(Upper(model), [i=1:2], x[i] <= 50)
@constraint(Upper(model), [i=1:2], y[i] >= -10)
@constraint(Upper(model), [i=1:2], y[i] <= 20)

# Followed by the objective and constraints of the lower problem:

# Lower objective function
@objective(Lower(model), Min, (-x[1] + y[1] + 40)^2 + (-x[2] + y[2] + 20)^2)

# Lower constraints
@constraint(Lower(model),  [i=1:2],- x[i] + 2y[i] <= -10)
@constraint(Lower(model), [i=1:2], y[i] >= -10)
@constraint(Lower(model), [i=1:2], y[i] <= 20)

# Now we can solve the problem and verify the solution again that reported by the book

optimize!(model)
primal_status(model)
termination_status(model)

# Auto testing
@test objective_value(model) ≈ 0 atol=1e-3
sol = vcat(value.(x), value.(y))
@test sol ≈ [0 ; 0 ; -10; -10] || sol ≈ [0 ; 30; -10; 10] #atol=1e-3

# # Like any other optimization problem, there is a chance in bilevel optimization to find multiple solutions with the same optimal value; based on the inherent stochasticity of the algorithm and random seed, we are expecting two optimal solutions for this problem.


text:
using Gadfly

set_default_plot_size(6inch, 3inch)

plot([sin, cos], color=["sin", "cos"], 0, 25)


text:
# %% imports

include("../../../Modules/CMAES/cmaes.jl")
include("LIF_basesim.jl")
include("../../../Modules/save_param.jl")

# %%

function optim()
    #N_params = length(tunekeys)
    N_alpha = 8*4
    N_w = 8*8 + 8*8
    N = N_alpha + N_w
    #N = N_weights + N_params
    best = Array{Float64}[]
    for i in 1:20
        push!(best,[])
    end
    best_fit = [-1.0 for i in 1:20]

    # Create initial point (optional)
    results = load_param("Julia/CartPole/LIF_base/results/results_LIFbase888_bound.jld")
    genes_init = results["weights"][12]

    lambda = 4+floor(Int64,3*log(N))
    mu = floor(Int64,lambda/2)
    c = CMAES(genes_init, N=N, μ=mu, λ=lambda , τ=sqrt(N), τ_c=N^2, τ_σ=sqrt(N))
    for i in 1:1000
        @time step!(c)
        bestind = argmin(c.F_λ)
        maxfit = -c.F_λ[bestind]
        println(i," ", maxfit, " (current best = $(maximum(best_fit)))")
        if maxfit > best_fit[1]
            best_fit[1] = maxfit
            best[1] = copy(c.offspring[bestind])
            idx = sortperm(best_fit)
            best_fit = best_fit[idx]
            best = best[idx]
        end
        if sum(best_fit) >= 4000
            break
        end
        if length(best[1]) > 0 && i%20 == 0
            save_param(best_fit,best)
        end
    end
    return best_fit, best
end

function viz_optim(best)
    tunekeys = ["imax", "vthresh"]
    for i in 1:length(best)
        params = Dict(tunekeys .=> rescale(best[i], tunekeys))
        @show params
    end
end

# %%

optim()

max_stars_repo_name: SuReLI/RL-STDP

max_issues_repo_name: SuReLI/RL-STDP

max_forks_repo_name: SuReLI/RL-STDP


text:
#
# Symetric.jl – The manifold of symmetric matrices
#
# ---
# Manopt.jl - Ronny Bergmann - 2017-07-06

import LinearAlgebra: norm, dot
import Base: exp, log, show

export Symmetric, SymPoint, SymTVector, show
# also indicates which functions are available (already) for Sym
export distance, dot, exp, norm, dot, manifoldDimension, parallelTransport
export validateMPoint, validateTVector
export zeroTVector, typeofMPoint, typeofTVector
# Types
# ---
@doc doc"""
    Symmetric <: Manifold

The manifold $\mathcal M = \mathrm{Sym}(n)$, where $\mathrm{Sym}(n) = \{
x \in \mathbb R^{n\times n} | x = x^\mathrm{T}
\}$, $n\in\mathbb N$, denotes the manifold of symmetric matrices
equipped with the trace inner product and its induced Forbenius norm.

# Abbreviation
`Sym` or `Sym(n)`, respectively.

# Constructor
    Symmetric(n)

generates the manifold of `n`-by-`n` symmetric matrices.
"""
struct Symmetric <: Manifold
  name::String
  dimension::Int
  abbreviation::String
  Symmetric(dimension::Int) = new("$dimension-by-$dimension symmetric matrices",(dimension*(dimension+1)/2),"Sym($dimension)")
end
@doc doc"""
    SymPoint <: MPoint

A point $x$ on the manifold $\mathcal M = \mathrm{Sym}(n)$ of $n\times n$
symmetric matrices, represented in the redundant way of a
symmetric matrix (instead of storing just the upper half).
"""
struct SymPoint{T <: AbstractFloat} <: MPoint
	value::Matrix{T}
	SymPoint{T}(v::Matrix{T}) where {T <: AbstractFloat} = new(v);
end
SymPoint(v::Matrix{T}) where {T <: AbstractFloat} = SymPoint{T}(v)
getValue(x::SymPoint) = x.value
@doc doc"""
    SymTVector <: TVector

A tangent vector $\xi$ in $T_x\mathcal M$ of a symmetric matrix $x\in\mathcal M$.
"""
struct SymTVector{T <: AbstractFloat} <: TVector
	value::Matrix{T}
  	SymTVector{T}(value::Matrix{T}) where {T <: AbstractFloat} = new(value);
end
SymTVector(v::Matrix{T}) where {T <: AbstractFloat} = SymTVector{T}(v)
getValue(ξ::SymTVector) = ξ.value
# Traits
# ---
# (a) P(n) is a matrix manifold
@traitimpl IsMatrixM{Symmetric}
@traitimpl IsMatrixP{SymPoint}
@traitimpl IsMatrixTV{SymTVector}
# Functions
# ---
@doc doc"""
    distance(M,x,y)

distance of two [`SymPoint`](@ref)s `x,y` on the [`Symmetric`](@ref) manifold `M``
inherited
from embedding them in $\mathbb R^{n\times n}$, i.e. use the Frobenious norm
of the difference.
"""
distance(M::Symmetric,x::SymPoint,y::SymPoint) = norm( getValue(x) - getValue(y) )
@doc doc"""
    dot(M,x,ξ,ν)

inner product of two [`SymTVector`](@ref)s `ξ,ν` lying in the tangent
space of the [`SymPoint`](@ref) `x` on the [`Symmetric`](@ref) manifold `M`.
"""
dot(M::Symmetric, x::SymPoint, ξ::SymTVector, ν::SymTVector) = dot( getValue(ξ), getValue(ν) )
@doc doc"""
    exp(M,x,ξ[, t=1.0])

compute the exponential map on the [`Symmetric`](@ref) manifold `M` given a
[`SymPoint`](@ref) `x` and a [`SymTVector`](@ref) `ξ`, as well as an optional
scaling factor `t`. The exponential map is given by

$\exp_{x}ξ = x+ξ.$
"""
exp(M::Symmetric, x::SymPoint, ξ::SymTVector, t::Float64=1.0) = SymPoint( getValue(x) + t*getValue(ξ) )
@doc doc"""
    log(M,x,y)

compute the logarithmic map for two [`SymPoint`](@ref)` x,y` on the [`Symmetric`](@ref) `M`,
which is given by $\log_xy = y-x$.
"""
log(M::Symmetric,x::SymPoint,y::SymPoint) = SymTVector( getValue(y) - getValue(x) )
"""
    manifoldDimension(M)

returns the manifold dimension of the [`Symmetric`](@ref) manifold `M`.
"""
manifoldDimension(M::Symmetric) = M.dimension
"""
    manifoldDimension(x)

returns the manifold dimension the [`SymPoint`](@ref) `x` belongs to.
"""
manifoldDimension(x::SymPoint) = size( getValue(x), 1)*(size( getValue(x), 1)+1)/2
"""
    norm(M,x,ξ)

computes the norm of the [`SymTVector`](@ref) `ξ` in the tangent space of
the [`SymPoint`](@ref) `x` on the [`Symmetric`](@ref) `M` embedded in the
Euclidean space, i.e. by its Frobenius norm.
"""
norm(M::Symmetric,x::SymPoint,ξ::SymTVector) = norm( getValue(ξ) )
"""
    parallelTransport(M,x,y,ξ)

compute the parallel transport of a [`SymTVector`](@ref) `ξ` from the tangent
space at the [`SymPoint`](@ref) `x` to the [`SymPoint`](@ref)` y` on the
[`Symmetric`](@ref) manifold `M`.
Since the metric is inherited from the embedding space, it is just the identity.
"""
parallelTransport(M::Symmetric,x::SymPoint,y::SymPoint,ξ::SymTVector) = ξ

typeofTVector(::Type{SymPoint{T}}) where T = SymTVector{T}
typeofMPoint(::Type{SymTVector{T}}) where T = SymPoint{T}

@doc doc"""
    typicalDistance(M)

returns the typical distance on the [`Symmetric`](@ref) manifold `M`,
i.e. $\sqrt{n}$.
"""
typicalDistance(M::Symmetric) = sqrt( - 0.5 + sqrt(1/4 + 2*manifoldDimension(M) ) ) # manDim to n

@doc doc"""
    validateMPoint(M,x)

validate, that the [`SymPoint`](@ref) `x` is a valid point on the
[`Symmetric`](@ref) manifold `M`,
i.e. that its dimensions are correct and that the matrix is symmetric.
"""
function validateMPoint(M::Symmetric, x::SymPoint)
    if manifoldDimension(M) ≠ manifoldDimension(x)
        throw(DomainError(
            "The point $x does not lie on $M,, since the manifold dimension of $M ($(manifoldDimension(M)))does not fit the manifold dimension of $x ($(manifoldDimension(x)))."
        ))
    end
    if norm(getValue(x) - transpose(getValue(x))) > 10^(-14)
        throw(DomainError(
            "The point $x does not lie on $M, since the matrix of $x is not symmetric."
        ))
    end
    return true
end
@doc doc"""
    validateTVector(M,x,ξ)

validate, that the [`SymTVector`](@ref) is a valid tangent vector to the
[`SymPoint`](@ref) `x` on the [`Symmetric`](@ref) manifold `M`,
i.e. that its dimensions are correct and that the matrix is symmetric.
"""
function validateTVector(M::Symmetric, x::SymPoint, ξ::SymTVector)
    ξs = size( getValue(ξ), 1)*(size( getValue(ξ), 1)+1)/2
    if (manifoldDimension(M) ≠ manifoldDimension(x)) || (ξs ≠ manifoldDimension(x))
        throw(DomainError(
            "The tangent vector $ξ of size $(ξs), the point $x ($(manifoldDimension(x))) and the manifold $M ($(manifoldDimension(M))) are not all equal in dimensions, so the tangent vector can not be correct."
        ))
    end
    if norm(getValue(ξ) - transpose(getValue(ξ))) > 10^(-14)
        throw(DomainError(
            "The tangent vector $ξ is not a symmetric matrix and hence can not lie in the tangent space of $x on $M."
        ))
    end
    return true
end
@doc doc"""
    ξ = zeroTVector(M,x)

returns a zero vector in the tangent space $T_x\mathcal M$ of the
[`SymPoint`](@ref) `x` on  the [`Symmetric`](@ref) manifold `M`.
"""
zeroTVector(M::Symmetric, x::SymPoint) = SymTVector(  zero( getValue(x) )  );
# Display
# ---
show(io::IO, M::Symmetric) = print(io, "The Manifold of $(M.name).")
show(io::IO, p::SymPoint) = print(io, "Sym($(p.value))")
show(io::IO, ξ::SymTVector) = print(io, "SymT($(ξ.value))")


text:
type QuadprogSolution
    status
    objval
    sol
    attrs
end

function quadprog(c::InputVector, Q::AbstractMatrix, A::AbstractMatrix, rowlb::InputVector, rowub::InputVector, lb::InputVector, ub::InputVector, solver::AbstractMathProgSolver = MathProgBase.defaultQPsolver)
    m = LinearQuadraticModel(solver)
    nrow,ncol = size(A)

    c = expandvec(c, ncol)
    rowlbtmp = expandvec(rowlb, nrow)
    rowubtmp = expandvec(rowub, nrow)
    lb = expandvec(lb, ncol)
    ub = expandvec(ub, ncol)

    # rowlb is allowed to be vector of senses
    if eltype(rowlbtmp) == Char
        realtype = eltype(rowubtmp)
        sense = rowlbtmp
        rhs = rowubtmp
        @assert realtype <: Real
        warn_no_inf(realtype)
        rowlb = Array(realtype, nrow)
        rowub = Array(realtype, nrow)
        for i in 1:nrow
            if sense[i] == '<'
                rowlb[i] = typemin(realtype)
                rowub[i] = rhs[i]
            elseif sense[i] == '>'
                rowlb[i] = rhs[i]
                rowub[i] = typemax(realtype)
            elseif sense[i] == '='
                rowlb[i] = rhs[i]
                rowub[i] = rhs[i]
            else
                error("Unrecognized sense '$(sense[i])'")
            end
        end
    else
        rowlb = rowlbtmp
        rowub = rowubtmp
    end

    loadproblem!(m, A, lb, ub, c, rowlb, rowub, :Min)
    setquadobj!(m, Q)
    optimize!(m)
    stat = status(m)
    if stat == :Optimal
        attrs = Dict()
        return QuadprogSolution(stat, getobjval(m), getsolution(m), attrs)
    else
        return QuadprogSolution(stat, nothing, [], Dict())
    end
end

quadprog(c,Q,A,rowlb,rowub,solver::AbstractMathProgSolver=defaultQPsolver) = quadprog(c,Q,A,rowlb,rowub,0,Inf,solver)

export quadprog


text:
using LazySets, Plots
using LazySets: center
using Plots.PlotMeasures
using LaTeXStrings

using ReachabilityAnalysis

# model parameters
m = 0.25; k = 2.0

# FEM assembled matrices
M = [2m 0; 0 m]; K = [2k -k; -k k]; C = (M+K)/20
F = [0.0, 1.0]; ΔF0 = Interval(0.9, 1.1)

# initial-value problem with set initial conditions
sys = SecondOrderAffineContinuousSystem(M, C, K, F)
U0 = BallInf(zeros(4), 0.5)
prob = InitialValueProblem(homogenize(sys), U0 × ΔF0)

# solve using support functions
solA = solve(prob, 50, LGG09(δ=5e-2, dirs=:box, dim=5))
solB = solve(prob, 50, LGG09(δ=5e-2, dirs=:oct, dim=5))

# compute random trajectories
ntraj = 100
B0 = box_approximation(U0 × ΔF0)
prob_orbit_random(m) = [IVP(prob.s, sample(B0)) for _ in 1:m];
sol_orbit = [solve(p, T=50, alg=ORBIT(δ=0.01)) for p in prob_orbit_random(ntraj)]

U0vert = vertices_list(B0)
sol_orbit_vert = [solve(IVP(prob.s, v), T=50, alg=ORBIT(δ=0.01)) for v in U0vert]

# plot displacement (2m) vs time
fig = plot(xlab="Time", ylab="Displacement",
           legendfontsize=16, legend=:topright,
           tickfont=font(18, "Times"), guidefontsize=18,
           xtick=([0, 10, 20, 30, 40, 50], [L"0", L"10", L"20", L"30", L"40", L"50"]),
           ytick=([-0.5, 0, 0.5, 1, 1.5], [L"-0.5", L"0", L"0.5", L"1.0", L"1.5"]),
           xguidefont=font(18, "Times"),
           yguidefont=font(18, "Times"),
           ylims=(-0.65, 1.65),
           bottom_margin=8mm, left_margin=8mm, right_margin=4mm, top_margin=3mm, size=(900, 600))

plot!(fig, solA, vars=(0, 1), lw=0.3, c=:blue, alpha=1., lc=:blue)
[plot!(fig, s, vars=(0, 1), alpha=1., seriestype=:path, c=:magenta, marker=:none, lw=0.6) for s in sol_orbit]
[plot!(fig, s, vars=(0, 1), alpha=1., seriestype=:path, c=:magenta, marker=:none, lw=0.6) for s in sol_orbit_vert]

lens!(fig, [0.5, 4.0], [0.4, 1.55], inset = (1, bbox(0.58, 0.0, 0.35, 0.35)),
      subplot=2, tickfont=font(16, "Times"),
      xtick=([1.0, 2.0, 3.0, 4.0], [L"1", L"2", L"3", L"4"]),
      ytick=([0.5, 1.0, 1.5], [L"0.5", L"1.0", L"1.5"]))

lens!(fig, [40, 50], [0.4, 0.6], inset = (1, bbox(0.58, 0.58, 0.35, 0.35)),
      subplot=3, tickfont=font(16, "Times"),
      xtick=([40, 45, 50], [L"40", L"45", L"50"]),
      ytick=([0.4, 0.5, 0.6], [L"0.4", L"0.5", L"0.6"]))

savefig(fig, "displacement_vs_time.pdf")
savefig(fig, "displacement_vs_time.png")

# plot displacement (2m) vs displacement (m)
fig = plot(xlab="Displacement mass 2m", ylab="Displacement mass m",
           legendfontsize=16, legend=:topright,
           tickfont=font(18, "Times"), guidefontsize=18,
           xtick=([-0.5,0, 0.5, 1.0, 1.5], [L"-0.5", L"0", L"0.5", L"1.0", L"1.5"]),
           ytick=([-0.5, 0.5, 1.5, 2.5], [L"-0.5", L"0.5", L"1.5", L"2.5"]),
           xguidefont=font(18, "Times"),
           yguidefont=font(18, "Times"),
           xlims=(-0.65, 1.65),
           ylims=(-0.6, 2.5),
           bottom_margin=8mm, left_margin=8mm, right_margin=4mm, top_margin=3mm, size=(900, 600))

plot!(fig, solB, vars=(1, 2), lw=0.0, c=:blue, alpha=0.8)
[plot!(fig, s, vars=(1, 2), alpha=0.7, seriestype=:path, c=:magenta, marker=:none, lw=1.1) for s in sol_orbit]
[plot!(fig, s, vars=(1, 2), alpha=0.7, seriestype=:path, c=:magenta, marker=:none, lw=1.1) for s in sol_orbit_vert]

savefig(fig, "displacement_vs_displacement.pdf")
savefig(fig, "displacement_vs_displacement.png")


text:
abstract type PriorityQueue end

struct PQBinaryHeap{K,V,O<:Ordering} <: PriorityQueue
  keys::Vector{K}
  values::Vector{V}
  keymap::Vector{Int} # map a key to its value
  valmap::Vector{Int} # map a value to its key

  function PQBinaryHeap{K,V,O}(o::O) where {K,V,O<:Ordering}
    new{K,V,O}(K[], V[], Int[], Int[])
  end

  function PQBinaryHeap{K,V,O}(keys, values) where {K,V,O<:Ordering}
    # sanity checks
    if length(keys) != length(values)
      throw(ArgumentError("number of keys and values is mismatched"))
    end

    if length(keys) != length(unique(keys))
      throw(ArgumentError("keys must be unique"))
    end

    # build association maps
    keymap = collect(1:length(keys))
    valmap = copy(keymap)

    # construct the object
    pq = new{K,V,O}(copy(keys), copy(values), keymap, valmap)

    # build the heap
    heapify!(pq)

    # obtain a sorting permutation for the keys array
    sortorder = sortperm(pq.keys)

    # do the sorting for the keys and keymap
    pq.keys   .= pq.keys[sortorder]
    pq.keymap .= pq.keymap[sortorder]

    # update the value to key associations
    for k in eachindex(sortorder)
      j = pq.keymap[k] # key k maps to position j in the heap,
      pq.valmap[j] = k # so position j must map back to k
    end

    return pq
  end
end

## heap traversal
@inline heap_left(j) = 2 * j
@inline heap_right(j) = 2 * j + 1
@inline heap_parent(j) = div(j, 2)

## low-level operations
@inline function sift_down!(pq::PQBinaryHeap, parent, bottom, is_violation)
  # unpack fields
  keys = pq.keys
  values = pq.values
  keymap = pq.keymap
  valmap = pq.valmap

  # save the parent value and start the search
  parent_value = values[parent]
  child = heap_left(parent)

  # while a left child node exists...
  while child <= bottom
    if child < bottom && is_violation(values[child], values[child + 1])
      child = child + 1
    end

    # if the parent does not violate the heap property, we're done
    if !is_violation(parent_value, values[child])
      break
    else
      node = heap_parent(child)

      node_keyix  = valmap[node]
      child_keyix = valmap[child]

      # perform swaps
      values[node] = values[child]

      keymap[node_keyix]  = child
      keymap[child_keyix] = node

      valmap[node]  = child_keyix
      valmap[child] = node_keyix

      child = heap_left(child)
    end
  end

  values[heap_parent(child)] = parent_value

  return pq
end

@inline function sift_up!(pq::PQBinaryHeap, top, child, is_violation)
  # unpack fields
  keys = pq.keys
  values = pq.values
  keymap = pq.keymap
  valmap = pq.valmap

  # while we're not at the top of the heap...
  while child > top
    parent = heap_parent(child)

    if is_violation(values[parent], values[child])
      parent_keyix = valmap[parent]
      child_keyix  = valmap[child]

      # swap values on the heap
      values[parent], values[child] = values[child], values[parent]

      keymap[parent_keyix] = child
      keymap[child_keyix]  = parent

      valmap[parent] = child_keyix
      valmap[child]  = parent_keyix

      # keep working up the tree
      child = parent
    else
      break
    end
  end

  return pq
end

function enqueue!(pq::PQBinaryHeap, value, key)
  # the value will be inserted at the bottom of the heap
  idx = length(pq) + 1

  # where should the key be inserted to maintain an ordered list?
  i = searchsortedfirst(pq.keys, key)

  insert!(pq.keys, i, key)   # insert `key` at index `i`
  insert!(pq.keymap, i, idx) # insert the association `i ---> idx`
  insert!(pq.valmap, idx, i) # insert the association `idx ---> i`
  push!(pq.values, value)    # add `value` at the end of the heap

  # fix upstream keys
  for j in i+1:idx
    pq.valmap[pq.keymap[j]] += 1
  end

  # check for heap order violations, working from the bottom to the top
  sift_up!(pq, 1, idx)
end

# TODO: make this efficient
function dequeue!(pq::PQBinaryHeap)
  # save the top value
  top_key, top_value = peektop(pq)

  # get associations for top and bottom nodes
  top_validx, bot_validx = 1, length(pq)
  top_keyidx, bot_keyidx = pq.valmap[top_validx], pq.valmap[bot_validx]

  # swap the nodes
  pq.values[top_validx], pq.values[bot_validx] = pq.values[bot_validx], pq.values[top_validx]

  pq.keymap[top_keyidx] = bot_validx
  pq.keymap[bot_keyidx] = top_validx
  pq.valmap[top_validx] = bot_keyidx
  pq.valmap[bot_validx] = top_keyidx

  if !isempty(pq)
    sift_down!(pq, 1, length(pq) - 1)
  end

  for i in length(pq):-1:top_keyidx+1
    pq.valmap[pq.keymap[i]] -= 1
  end

  pop!(pq.values)
  pop!(pq.valmap)
  deleteat!(pq.keys,   top_keyidx)
  deleteat!(pq.keymap, top_keyidx)

  # return the deleted entry
  return (top_key, top_value)
end

@inline function update_keyvalue!(pq::PQBinaryHeap, value, i)
  node = pq.keymap[i]
  old_value = pq.values[node]
  pq.values[node] = value

  return old_value
end

## heap property violation
is_minheap_violation(x, y) = y < x
is_maxheap_violation(x, y) = x < y

## type dispatch

### minheap
@inline function sift_down!(pq::PQBinaryHeap{K,V,O}, parent, bottom) where {K,V,O<:ForwardOrdering}
  sift_down!(pq, parent, bottom, is_minheap_violation)
end

@inline function sift_up!(pq::PQBinaryHeap{K,V,O}, top, child) where {K,V,O<:ForwardOrdering}
  sift_up!(pq, top, child, is_minheap_violation)
end

### maxheap
@inline function sift_down!(pq::PQBinaryHeap{K,V,O}, parent, bottom) where {K,V,O<:ReverseOrdering}
  sift_down!(pq, parent, bottom, is_maxheap_violation)
end

@inline function sift_up!(pq::PQBinaryHeap{K,V,O}, top, child) where {K,V,O<:ReverseOrdering}
  sift_up!(pq, top, child, is_maxheap_violation)
end

### heap construction
@inline function heapify!(pq::PQBinaryHeap)
  # a flat vector already enforces the shape property
  bottom = length(pq)

  # how many parent nodes are there?
  parents = heap_parent(bottom):-1:1

  # enforce the heap property at each parent node
  for parent in parents
    sift_down!(pq, parent, bottom)
  end

  return pq
end

## interface

# empty constructors
PQBinaryHeap{K,V}(o::Ordering=Forward) where {K,V} = PQBinaryHeap{K,V,typeof(o)}(o)

# non-empty constructor
PQBinaryHeap(keys, values, o::Ordering=Forward) = _pqbinheap_helper(o, keys, values)

_pqbinheap_helper(o::Ord, keys, values) where Ord = PQBinaryHeap{eltype(keys), eltype(values), Ord}(keys, values)

Base.length(pq::PQBinaryHeap)  = length(pq.keys)
Base.isempty(pq::PQBinaryHeap) = isempty(pq.keys)
Base.haskey(pq::PQBinaryHeap, key) = isempty(searchsorted(pq.keys, key)) ? false : true

peektop(pq::PQBinaryHeap) = (pq.keys[pq.valmap[1]], pq.values[1])

function Base.getindex(pq::PQBinaryHeap, key)
  irange = searchsorted(pq.keys, key)

  if isempty(irange)
    throw(KeyError(key))
  else
    return pq.values[pq.keymap[irange[end]]]
  end
end

function Base.get(pq::PQBinaryHeap, key, default)
  irange = searchsorted(pq.keys, key)

  return isempty(irange) ? default : pq.values[pq.keymap[irange[end]]]
end

function Base.setindex!(pq::PQBinaryHeap{K,V,ForwardOrdering}, value, key) where {K,V}
  _setindex!(pq, value, key, is_minheap_violation)
end

function Base.setindex!(pq::PQBinaryHeap{K,V,ReverseOrdering}, value, key) where {K,V}
  _setindex!(pq, value, key, is_maxheap_violation)
end

@inline function _setindex!(pq, value, key, is_violation)
  irange = searchsorted(pq.keys, key)
  i = isempty(irange) ? 0 : irange[end]

  # if the key exists...
  if i != 0
    old_value = update_keyvalue!(pq, value, i)
    node = pq.keymap[i]

    if is_violation(value, old_value)
      sift_down!(pq, node, length(pq), is_violation)
    else
      sift_up!(pq, 1, node, is_violation)
    end
  # if they key does not exist...
  else
    enqueue!(pq, value, key)
  end

  return pq
end

function Base.iterate(pq::PQBinaryHeap, state = 1)
  state > length(pq) && return nothing

  key_idx = state
  val_idx = pq.keymap[key_idx]

  return ((pq.keys[key_idx], pq.values[val_idx]), state + 1)
end

Base.firstindex(pq::PQBinaryHeap) = 1
Base.lastindex(pq::PQBinaryHeap)  = length(pq)


text: my_fu(x, y) = 2*x+ 3*y


text:
abstract type AbstractEChartType end
abstract type AbstractEChartSeries <: AbstractEChartType end

@with_kw mutable struct Theme <: AbstractEChartType
    color::Union{AbstractVector, Nothing} = nothing
    backgroundColor::Union{String, Nothing} = nothing
    textStyle::Union{Dict, Nothing} = nothing
    title::Union{Dict, Nothing} = nothing
    line::Union{Dict, Nothing} = nothing
    radar::Union{Dict, Nothing} = nothing
    bar::Union{Dict, Nothing} = nothing
    pie::Union{Dict, Nothing} = nothing
    scatter::Union{Dict, Nothing} = nothing
    boxplot::Union{Dict, Nothing} = nothing
    parallel::Union{Dict, Nothing} = nothing
    sankey::Union{Dict, Nothing} = nothing
    funnel::Union{Dict, Nothing} = nothing
    gauge::Union{Dict, Nothing} = nothing
    candlestick::Union{Dict, Nothing} = nothing
    graph::Union{Dict, Nothing} = nothing
    map::Union{Dict, Nothing} = nothing
    geo::Union{Dict, Nothing} = nothing
    categoryAxis::Union{Dict, Nothing} = nothing
    valueAxis::Union{Dict, Nothing} = nothing
    logAxis::Union{Dict, Nothing} = nothing
    timeAxis::Union{Dict, Nothing} = nothing
    toolbox::Union{Dict, Nothing} = nothing
    legend::Union{Dict, Nothing} = nothing
    tooltip::Union{Dict, Nothing} = nothing
    timeline::Union{Dict, Nothing} = nothing
    visualMap::Union{Dict, Nothing} = nothing
    dataZoom::Union{Dict, Nothing} = nothing
    markPoint::Union{Dict, Nothing} = nothing
end

function Theme(x::Dict)

    Theme(
    color = get(x, "color", nothing),
    backgroundColor = get(x, "backgroundColor", nothing),
    textStyle = get(x, "textStyle", nothing),
    title = get(x, "title", nothing),
    line = get(x, "line", nothing),
    radar = get(x, "radar", nothing),
    bar = get(x, "bar", nothing),
    pie = get(x, "pie", nothing),
    scatter = get(x, "scatter", nothing),
    boxplot = get(x, "boxplot", nothing),
    parallel = get(x, "parallel", nothing),
    sankey = get(x, "sankey", nothing),
    funnel = get(x, "funnel", nothing),
    gauge = get(x, "gauge", nothing),
    candlestick = get(x, "candlestick", nothing),
    graph = get(x, "graph", nothing),
    map = get(x, "map", nothing),
    geo = get(x, "geo", nothing),
    categoryAxis = get(x, "categoryAxis", nothing),
    valueAxis = get(x, "valueAxis", nothing),
    logAxis = get(x, "logAxis", nothing),
    timeAxis = get(x, "timeAxis", nothing),
    toolbox = get(x, "toolbox", nothing),
    legend = get(x, "legend", nothing),
    tooltip = get(x, "tooltip", nothing),
    timeline = get(x, "timeline", nothing),
    visualMap = get(x, "visualMap", nothing),
    dataZoom = get(x, "dataZoom", nothing),
    markPoint = get(x, "markPoint", nothing)
    )

end

#themes from http://echarts.baidu.com/theme-builder/
const chalk = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/chalk.json")))
const dark = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/dark.json")))
const essos = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/essos.json")))
const halloween = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/halloween.json")))
const infographic = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/infographic.json")))
const macarons = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/macarons.json")))
const purplepassion = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/purple-passion.json")))
const roma = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/roma.json")))
const shine = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/shine.json")))
const vintage = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/vintage.json")))
const walden = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/walden.json")))
const westeros = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/westeros.json")))
const wonderland = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/wonderland.json")))

#custom themes
#grayscale based on 6-color Colorbrewer sequential single-hue
const grayscale = Theme(JSON.parsefile(joinpath(dirname(@__FILE__), "..", "docs/js/themes/grayscale.json")))


text:
eval(quote
    using .ExponentialUtilities: expv, expv!, arnoldi, arnoldi!, KrylovSubspace
end)

eval(load_Φ₁_krylov())
eval(load_krylov_LGG09_homog())
eval(load_krylov_LGG09_inhomog())


text:
# Simple Robust IPOMDP and RIPOMDP for testing

struct SimpleIPOMDP <: IPOMDP{Symbol, Symbol, Symbol}
    ps::Float64
    pb::Float64
    pst::Float64
    pbt::Float64
    discount::Float64
    inforeward::Vector{Vector{Float64}}
end

struct SimpleRIPOMDP <: RIPOMDP{Symbol, Symbol, Symbol}
    ps::Float64
    pb::Float64
    pst::Float64
    pbt::Float64
    discount::Float64
    inforeward::Vector{Vector{Float64}}
    usize::Float64
end

const states_simple = [:left, :right]
const actions_simple = [:single, :both]
const observations_simple = [:L,:R,:LL,:LR,:RL,:RR]
const inforeward_simple = [[1.0, -1.0], [-1.0, 1.0]]
const discount_simple = 0.75
const ps_simple = 0.8 # probability of correct detect single
const pb_simple = 0.6 # probability of correct detect both
const err_simple = 0.4 # ambiguity size
const pst_simple = 0.6 # probability of stay after single action
const pbt_simple = 0.9 # probability of stay after both action

SimpleIPOMDP(ps::Float64, pb::Float64, pst::Float64, pbt::Float64) = SimpleIPOMDP(ps, pb, pst, pbt, discount_simple, inforeward_simple)
SimpleIPOMDP(ps::Float64, pb::Float64) = SimpleIPOMDP(ps, pb, pst_simple, pbt_simple, discount_simple, inforeward_simple)
SimpleIPOMDP() = SimpleIPOMDP(ps_simple, pb_simple, pst_simple, pbt_simple, discount_simple, inforeward_simple)

SimpleRIPOMDP(ps::Float64, pb::Float64, pst::Float64, pbt::Float64, err::Float64) = SimpleRIPOMDP(ps, pb, pst, pbt, discount_simple, inforeward_simple, err)
SimpleRIPOMDP(ps::Float64, pb::Float64, err::Float64) = SimpleRIPOMDP(ps, pb, pst_simple, pbt_simple, discount_simple, inforeward_simple, err)
SimpleRIPOMDP(ps::Float64, pb::Float64) = SimpleRIPOMDP(ps, pb, pst_simple, pbt_simple, discount_simple, inforeward_simple, err_simple)
SimpleRIPOMDP(err::Float64) = SimpleRIPOMDP(ps_simple, pb_simple, pst_simple, pbt_simple, discount_simple, inforeward_simple, err)
SimpleRIPOMDP() = SimpleRIPOMDP(ps_simple, pb_simple, pst_simple, pbt_simple, discount_simple, inforeward_simple, err_simple)

states(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = states_simple
actions(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = actions_simple
observations(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = observations_simple

n_states(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = length(states_simple)
n_actions(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = length(actions_simple)
n_observations(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = length(observations_simple)

state_index(::Union{SimpleIPOMDP,SimpleRIPOMDP}, s::Symbol) = find(x -> x == s, states_simple)[1]
action_index(::Union{SimpleIPOMDP,SimpleRIPOMDP}, a::Symbol) = find(x -> x == a, actions_simple)[1]
observation_index(::Union{SimpleIPOMDP,SimpleRIPOMDP}, z::Symbol) = find(x -> x == z, observations_simple)[1]
obs_index(prob::Union{SimpleIPOMDP,SimpleRIPOMDP}, z::Symbol) = observation_index(prob, z)

initial_state_distribution(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = SparseCat([:left, :right], [1.0, 0.0])
initial_belief(::Union{SimpleIPOMDP,SimpleRIPOMDP}) = [0.5, 0.5]

# Transition functions
# Nominal transition function distributions
const pϵ_simple = 1e-6

function transition(prob::SimpleIPOMDP, s::Symbol, a::Symbol)
    pst = prob.pst
    pbt = prob.pbt
    if s == :left
        if a == :single
            return SparseCat([:left, :right], [pst, 1-pst])
        else
            return SparseCat([:left, :right], [pbt, 1-pbt])
        end
    else
        if a == :single
            return SparseCat([:left, :right], [1-pst, pst])
        else
            return SparseCat([:left, :right], [1-pbt, pbt])
        end
    end
    return nothing
end

function transition(prob::SimpleIPOMDP, s::Symbol, a::Symbol, sp::Symbol)
    pst = prob.pst
    pbt = prob.pbt
    if s == :left
        if a == :single
            return sp == :left ? pst : 1-pst
        else
            return sp == :left ? pbt : 1-pbt
        end
    else
        if a == :single
            return sp == :left ? 1-pst : pst
        else
            return sp == :left ? 1-pbt : pbt
        end
    end
end

function transition(prob::SimpleRIPOMDP, s::Symbol, a::Symbol)
    pst = prob.pst
    pbt = prob.pbt
    if s == :left
        if a == :single
            plower = max.([pst - prob.usize, 1-pst-prob.usize/10], 0.0 + pϵ_simple)
            pupper = min.([pst + prob.usize/10, 1-pst+prob.usize], 1.0 - pϵ_simple)
            return plower, pupper
        else
            plower = max.([pbt, 1-pbt] - prob.usize/10, 0.0 + pϵ_simple)
            pupper = min.([pbt, 1-pbt] + prob.usize/10, 1.0 - pϵ_simple)
            return plower, pupper
        end
    else
        if a == :single
            plower = max.([1-pst - prob.usize/10, pst - prob.usize], 0.0 + pϵ_simple)
            pupper = min.([1-pst + prob.usize, pst + prob.usize/10], 1.0 - pϵ_simple)
            return plower, pupper
        else
            plower = max.([1-pbt, pbt] - prob.usize/10, 0.0 + pϵ_simple)
            pupper = min.([1-pbt, pbt] + prob.usize/10, 1.0 - pϵ_simple)
            return plower, pupper
        end
    end
    return nothing
end

function transition(prob::SimpleRIPOMDP, s::Symbol, a::Symbol, sp::Symbol)
    pst = prob.pst
    pbt = prob.pbt
    if s == :left
        if a == :single
            if sp == :left
                plower = max.(pst - prob.usize, 0.0 + pϵ_simple)
                pupper = min.(pst + prob.usize/10, 1.0 - pϵ_simple)
                return plower, pupper
            else
                plower = max.(1-pst - prob.usize/10, 0.0 + pϵ_simple)
                pupper = min.(1-pst + prob.usize, 1.0 - pϵ_simple)
                return plower, pupper
            end
        else
            if sp == :left
                plower = max.(pbt - prob.usize/10, 0.0 + pϵ_simple)
                pupper = min.(pbt + prob.usize, 1.0 - pϵ_simple)
                return plower, pupper
            else
                plower = max.(1-pbt - prob.usize, 0.0 + pϵ_simple)
                pupper = min.(1-pbt + prob.usize/10, 1.0 - pϵ_simple)
                return plower, pupper
            end
        end
    else
        if a == :single
            if sp == :left
                plower = max.(1-pst - prob.usize/10, 0.0 + pϵ_simple)
                pupper = min.(1-pst + prob.usize, 1.0 - pϵ_simple)
                return plower, pupper
            else
                plower = max.(pst - prob.usize, 0.0 + pϵ_simple)
                pupper = min.(pst + prob.usize/10, 1.0 - pϵ_simple)
                return plower, pupper
            end
        else
            if sp == :left
                plower = max.(1-pbt - prob.usize, 0.0 + pϵ_simple)
                pupper = min.(1-pbt + prob.usize/10, 1.0 - pϵ_simple)
                return plower, pupper
            else
                plower = max.(pbt - prob.usize/10, 0.0 + pϵ_simple)
                pupper = min.(pbt + prob.usize, 1.0 - pϵ_simple)
                return plower, pupper
            end
        end
    end
end

function observation(prob::SimpleIPOMDP, a::Symbol, sp::Symbol)
    ps = prob.ps
    pb = prob.pb
    if a == :single
        if sp == :left
            return SparseCat([:L, :R, :LL, :LR, :RL, :RR], [ps-0.001, 1-ps, 0.001/4, 0.001/4, 0.001/4, 0.001/4])
        else
            return SparseCat([:L, :R, :LL, :LR, :RL, :RR], [1-ps, ps-0.001, 0.001/4, 0.001/4, 0.001/4, 0.001/4])
        end
    else
        if sp == :left
            return SparseCat([:L, :R, :LL, :LR, :RL, :RR], [0.001/2, 0.001/2, pb*pb-0.001, (1-pb)*pb, pb*(1-pb), (1-pb)*(1-pb)])
        else
            return SparseCat([:L, :R, :LL, :LR, :RL, :RR], [0.001/2, 0.001/2, (1-pb)*(1-pb), (1-pb)*pb, pb*(1-pb), pb*pb-0.001])
        end
    end
    return nothing
end

function observation(prob::SimpleIPOMDP, a::Symbol, sp::Symbol, z::Symbol)
    pdf(observation(prob, a, sp), z)
end

function observation(prob::SimpleRIPOMDP, a::Symbol, sp::Symbol)
    ps = prob.ps
    pb = prob.pb
    if a == :single
        if sp == :left
            plower = max.([ps-0.001 - prob.usize, 1-ps - prob.usize/10, 0, 0, 0, 0], 0.0 + pϵ_simple)
            pupper = min.([ps-0.001 + prob.usize/10, 1-ps + prob.usize, 0.001/4, 0.001/4, 0.001/4, 0.001/4], 1.0 - pϵ_simple)
            return plower, pupper
        else
            plower = max.([1-ps - prob.usize/10, ps-0.001-prob.usize, 0, 0, 0, 0], 0.0 + pϵ_simple)
            pupper = min.([1-ps + prob.usize, ps-0.001+prob.usize/10, 0.001/4, 0.001/4, 0.001/4, 0.001/4], 1.0 - pϵ_simple)
            return plower, pupper
        end
    else
        if sp == :left
            plower = max.([0, 0, pb*pb-0.001, (1-pb)*pb, pb*(1-pb), (1-pb)*(1-pb)] - prob.usize/10, 0.0 + pϵ_simple)
            pupper = min.([0.001/2, 0.001/2, pb*pb-0.001, (1-pb)*pb, pb*(1-pb), (1-pb)*(1-pb)] + prob.usize/10, 1.0 - pϵ_simple)
            return plower, pupper
        else
            plower = max.([0, 0, (1-pb)*(1-pb), (1-pb)*pb, pb*(1-pb), pb*pb-0.001] - prob.usize/10, 0.0 + pϵ_simple)
            pupper = min.([0.001/2, 0.001/2, (1-pb)*(1-pb), (1-pb)*pb, pb*(1-pb), pb*pb-0.001] + prob.usize/10, 1.0 - pϵ_simple)
            return plower, pupper
        end
    end
    return nothing
end

function observation(prob::SimpleRIPOMDP, a::Symbol, sp::Symbol, z::Symbol)
    zind = observation_index(prob, z)
    plold, puold = observation(prob, a, sp)
    plower = plold[zind]
    pupper = puold[zind]
    plower, pupper
end

function dynamics(prob::SimpleIPOMDP)
    ns = n_states(prob)
    nz = n_observations(prob)
    na = n_actions(prob)
    tarr = states(prob)
    sarr = states(prob)
    aarr = actions(prob)
    zarr = observations(prob)
    p = Array{Float64}(ns,nz,ns,na)
    plower = Array{Float64}(ns,nz,ns,na)
    pupper = Array{Float64}(ns,nz,ns,na)
    for tind in 1:ns, zind in 1:nz, sind in 1:ns, aind in 1:na
        pt = transition(prob, sarr[sind], aarr[aind], tarr[tind])
        pz = observation(prob, aarr[aind], tarr[tind], zarr[zind])
        p[tind, zind, sind, aind] = pz * pt
    end
    p
end

function dynamics(prob::SimpleRIPOMDP)
    ns = n_states(prob)
    nz = n_observations(prob)
    na = n_actions(prob)
    tarr = states(prob)
    sarr = states(prob)
    aarr = actions(prob)
    zarr = observations(prob)
    p = Array{Float64}(ns,nz,ns,na)
    plower = Array{Float64}(ns,nz,ns,na)
    pupper = Array{Float64}(ns,nz,ns,na)
    for tind in 1:ns, zind in 1:nz, sind in 1:ns, aind in 1:na
        ptl, ptu = transition(prob, sarr[sind], aarr[aind], tarr[tind])
        pzl, pzu = observation(prob, aarr[aind], tarr[tind], zarr[zind])
        plower[tind, zind, sind, aind] = pzl * ptl
        pupper[tind, zind, sind, aind] = pzu * ptu
    end
    plower, pupper
end

function reward(prob::Union{SimpleIPOMDP, SimpleRIPOMDP}, b::Vector{Float64}, a::Symbol)
    rmax = -Inf
    for α in prob.inforeward
        rmax = max(rmax, dot(α, b))
    end
    rmax
end
reward(prob::Union{SimpleIPOMDP,SimpleRIPOMDP}, b::Vector{Float64}, s::Symbol, a::Symbol, sp::Symbol) = reward(prob,b,a)

function rewardalpha(prob::Union{SimpleIPOMDP, SimpleRIPOMDP}, b::Vector{Float64}, a::Symbol)
    rmax = -Inf
    ralpha = nothing
    for α in prob.inforeward
        rnext = dot(α, b)
        if rnext > rmax
            rmax = rnext
            ralpha = copy(α)
        end
    end
    ralpha
end

discount(p::Union{SimpleIPOMDP, SimpleRIPOMDP}) = p.discount

# Simulation Functions
function initial_state(prob::Union{SimpleIPOMDP, SimpleRIPOMDP}, rng::AbstractRNG)
    d = initial_state_distribution(prob)
    return rand(rng, d)
end

function generate_sor(prob::SimpleIPOMDP, b, s, a, rng::AbstractRNG)
    sp = rand(rng, transition(prob, s, a))
    o = rand(rng, observation(prob, a, sp))
    r = reward(prob, b, s, a, sp)
    sp, o, r
end

# uniform random generate_sor (original)
function generate_sor(prob::SimpleRIPOMDP, b, s, a, rng::AbstractRNG)
    tdist = SparseCat(states(prob), psample(transition(prob, s, a)...))
    @show tdist
    sp = rand(rng, tdist)
    odist = SparseCat(observations(prob), psample(observation(prob, a, sp)...))
    @show odist
    o = rand(rng, odist)
    r = reward(prob, b, s, a, sp)
    sp, o, r
end