@testset "Select" begin

    @testset "Select parameters" begin
        import NaiveNASflux: select
        mat = reshape(collect(1:3*4), 3, 4)

        @test select(mat, 2 => [2,-1,-1,4]) == [4 0 0 10;5 0 0 11;6 0 0 12]
        @test select(mat, 1 => [-1,1,3]) == [0 0 0 0;1 4 7 10;3 6 9 12]

        @test select(mat, 1 => [2,-1,3,-1], 2 => [-1,1,-1,4]) == [0 2 0 11;0 0 0 0;0 3 0 12;0 0 0 0]
    end

    @testset "KernelSizeAligned" begin
        import NaiveNASflux: selectfilters
        pars = reshape(1:2*3*4*5, 5,4,3,2)
        ps = selectfilters(KernelSizeAligned(2, -1), pars)

        @test ps[1] == (1 => [-1, 1, 2, 3, 4, 5, -1])
        @test ps[2] == (2 => [2,3,4])
    end

    @testset "KernelSizeAligned MutableLayer $convtype" for convtype in (Conv, ConvTranspose, DepthwiseConv)
        import NaiveNASflux: weights
        m = mutable(convtype((3,4), 5=>5, pad=(1,1,1,2)), inputvertex("in", 5), layerfun=identity)
        indata = ones(Float32, 5,5,5,2)

        @test size(weights(layer(m)))[1:2] == (3,4)
        @test size(m(indata)) == size(indata)

        mutate_weights(m, KernelSizeAligned((-1, 2), (1,0,3,2)))

        @test size(weights(layer(m)))[1:2] == (2,6)
        @test size(m(indata)) == size(indata)
    end

    @testset "KernelSizeAligned LazyMutable" begin
        import NaiveNASflux: weights
        m = mutable(Conv((3,4), 5=>6, pad=(1,1,1,2)), inputvertex("in", 5), layerfun=LazyMutable)
        indata = ones(Float32, 5,5,5,2)

        @test size(weights(layer(m))) == (3,4,5,6)
        @test size(m(indata))[1:2] == size(indata)[1:2]

        mutate_weights(m, KernelSizeAligned((-1, 2), (1,0,3,2)))

        @test size(m(indata)) == (size(indata,1), size(indata,2), 6, size(indata,4))
        @test size(weights(layer(m))) == (2,6,5,6)

        mutate_weights(m, KernelSizeAligned((2, -1), (1,2,2,2)))
        mutate_outputs(m, [1,3,5])

        @test size(m(indata)) == (size(indata,1), size(indata,2), 3, size(indata,4))
        @test size(weights(layer(m))) == (4,5,5,3)
    end

    @testset "KernelSizeAligned Dense is Noop with layerfun $lfun" for lfun in (identity, LazyMutable)
        m = mutable(Dense(3,4), inputvertex("in", 3), layerfun = lfun)
        indata = ones(Float32, nin(m)[], 2)

        @test size(m(indata)) == (4,2)
        mutate_weights(m, KernelSizeAligned((-1,-1), (1,1)))
        @test size(m(indata)) == (4,2)
    end
end


function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{St, D},
    ::Type{S},
) where {St <: PSY.Storage, D <: AbstractStorageFormulation, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    add_variables!(container, ActivePowerInVariable, devices, D())
    add_variables!(container, ActivePowerOutVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, EnergyVariable, devices, D())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, D())
    end

    initial_conditions!(container, devices, D())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{St, D},
    ::Type{S},
) where {St <: PSY.Storage, D <: AbstractStorageFormulation, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(container, EnergyCapacityConstraint, EnergyVariable, devices, model, S)

    # Energy Balanace limits
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{St, D},
    ::Type{S},
) where {
    St <: PSY.Storage,
    D <: AbstractStorageFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(St, sys)

    add_variables!(container, ActivePowerInVariable, devices, D())
    add_variables!(container, ActivePowerOutVariable, devices, D())
    add_variables!(container, EnergyVariable, devices, D())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, D())
    end

    initial_conditions!(container, devices, D())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{St, D},
    ::Type{S},
) where {
    St <: PSY.Storage,
    D <: AbstractStorageFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(St, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(container, EnergyCapacityConstraint, EnergyVariable, devices, model, S)

    # Energy Balanace limits
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    add_constraint_dual!(container, sys, model)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{St, EnergyTarget},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    add_variables!(container, ActivePowerInVariable, devices, EnergyTarget())
    add_variables!(container, ActivePowerOutVariable, devices, EnergyTarget())
    add_variables!(container, ReactivePowerVariable, devices, EnergyTarget())
    add_variables!(container, EnergyVariable, devices, EnergyTarget())
    add_variables!(container, EnergyShortageVariable, devices, EnergyTarget())
    add_variables!(container, EnergySurplusVariable, devices, EnergyTarget())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, EnergyTarget())
    end

    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)

    initial_conditions!(container, devices, EnergyTarget())

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{St, EnergyTarget},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(container, EnergyCapacityConstraint, EnergyVariable, devices, model, S)

    # Energy Balanace limits
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{St, EnergyTarget},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(St, sys)

    add_variables!(container, ActivePowerInVariable, devices, EnergyTarget())
    add_variables!(container, ActivePowerOutVariable, devices, EnergyTarget())
    add_variables!(container, EnergyVariable, devices, EnergyTarget())
    add_variables!(container, EnergyShortageVariable, devices, EnergyTarget())
    add_variables!(container, EnergySurplusVariable, devices, EnergyTarget())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, EnergyTarget())
    end

    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)

    initial_conditions!(container, devices, EnergyTarget())

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{St, EnergyTarget},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(St, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(container, EnergyCapacityConstraint, EnergyVariable, devices, model, S)

    # Energy Balanace limits
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{St, BatteryAncillaryServices},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    add_variables!(container, ActivePowerInVariable, devices, BatteryAncillaryServices())
    add_variables!(container, ActivePowerOutVariable, devices, BatteryAncillaryServices())
    add_variables!(container, ReactivePowerVariable, devices, BatteryAncillaryServices())
    add_variables!(container, EnergyVariable, devices, BatteryAncillaryServices())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, BatteryAncillaryServices())
    end

    initial_conditions!(container, devices, BatteryAncillaryServices())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    if has_service_model(model)
        add_expressions!(container, ReserveRangeExpressionLB, devices, model)
        add_expressions!(container, ReserveRangeExpressionUB, devices, model)
    end
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{St, BatteryAncillaryServices},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(container, EnergyCapacityConstraint, EnergyVariable, devices, model, S)

    # Energy Balanace limits
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    if has_service_model(model)
        add_constraints!(container, ReserveEnergyConstraint, devices, model, S)
        add_constraints!(container, RangeLimitConstraint, devices, model, S)
    end
    add_feedforward_constraints!(container, model, devices)

    add_constraint_dual!(container, sys, model)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{St, BatteryAncillaryServices},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(St, sys)

    add_variables!(container, ActivePowerInVariable, devices, BatteryAncillaryServices())
    add_variables!(container, ActivePowerOutVariable, devices, BatteryAncillaryServices())
    add_variables!(container, EnergyVariable, devices, BatteryAncillaryServices())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, BatteryAncillaryServices())
    end

    initial_conditions!(container, devices, BatteryAncillaryServices())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    if has_service_model(model)
        add_expressions!(container, ReserveRangeExpressionLB, devices, model)
        add_expressions!(container, ReserveRangeExpressionUB, devices, model)
    end
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{St, BatteryAncillaryServices},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(St, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(container, EnergyCapacityConstraint, EnergyVariable, devices, model, S)

    # Energy Balanace limits
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    if has_service_model(model)
        add_constraints!(container, ReserveEnergyConstraint, devices, model, S)
        add_constraints!(container, RangeLimitConstraint, devices, model, S)
    end

    add_feedforward_constraints!(container, model, devices)

    add_constraint_dual!(container, sys, model)

    return
end


# vocoder.jl
# Compare Matlab and Sound versions of phase vocoder.

using Sound: phase_vocoder, hann, soundsc
using MATLAB
using Test: @testset, @test

ENV["MATLAB_ROOT"] = "/Applications/freeware/matlab"

@testset "vocoder" begin
end
    S = 44100
    N = S ÷ 5
    hopin = 121.
#   hopout = hopin * 2
    hopout = hopin * 1.5 # todo: large differences
# for hopin in (100., 121.), hopout in hopin .* (1, 1.5, 0.5)
@show hopin, hopout
    f = 540
    x = sin.(2π*f*(1:N)/S)
    yj = phase_vocoder(x, S; hopin, hopout)

    @show "start matlab"
    Nd = Float64(N)
    x2 = reshape(x, :, 1)
    mat"ym = pv($x2, $Nd, $hopin, $hopout)"
    @mget ym
    ym = Float32.(vec(ym))

    @show extrema(yj)
    @show extrema(ym)
    @show extrema(yj - ym)
    #=
    =#
#   @test isapprox(yj, ym; rtol = 0.01) # 1% seems OK for now?
# end

#=
=#
    # plot spectrograms to compare
    using Plots
    using MIRTjim: jim
    using DSP: spectrogram

    plot(yj - ym)
    xlims = (5000,5400)
p0 = plot(
    plot(yj, marker=:circle; xlims),
    plot(ym, marker=:circle; xlims),
    layout=(1,2),
); gui()
    sj = spectrogram(yj / maximum(abs, x); fs=S)
    sm = spectrogram(ym / maximum(abs, x); fs=S)
    dB = x -> max(log10.(x), -10)
    fun = (s) -> jim(s.time, s.freq, dB.(s.power)', aspect_ratio=:auto,
        ylims = (0,1800), color=:viridis, clim = (-10,0),
        yticks = 0:300:1800, yflip=false,
        xlabel="time", ylabel="freq",
        prompt = false,
    )
    sd = (; time=sj.time, freq=sj.freq, power=abs.(sj.power-sm.power))
    plot(fun(sj), fun(sm), fun(sd), p0), gui()
#   soundsc([x; ym; yj], S) # listen to compare


export updateWd,multWd,setSources,setSourcesSame,setDobs,setWd
function updateWd(pMis::Array{RemoteChannel},Dc::Array{RemoteChannel})
@sync begin
		for k=1:length(pMis)
	@async begin
			pMis[k] = remotecall_fetch(updateWd,pMis[k].where,pMis[k],Dc[k]);
		end
	end
end
return pMis;
end

function updateWd(pMisRF::RemoteChannel,Dc::RemoteChannel)
pMis  = take!(pMisRF)
Dc = fetch(Dc);
pMis.Wd = 1.0./(real(Dc - pMis.dobs) + 1e-3*mean(abs(pMis.dobs[:]))) + 1im./(imag(Dc - pMis.dobs) + 1e-3*mean(abs(pMis.dobs[:])));
put!(pMisRF,pMis)
return pMisRF;
end

function multWd(pMis::Array{RemoteChannel},beta::Float64)
@sync begin
		for k=1:length(pMis)
	@async begin
			pMis[k] = remotecall_fetch(multWd,pMis[k].where,pMis[k],beta);
		end
	end
end
return pMis;
end

function multWd(pMisRF::RemoteChannel,beta::Float64)
pMis  = take!(pMisRF)
pMis.Wd *= beta;
put!(pMisRF,pMis)
return pMisRF;
end

export setSources, setDobs

function setWd(pMis::Array{RemoteChannel},newWd::Array)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setWd,pMis[k].where,pMis[k],newWd[k]);
		end
	end
end
return pMis;
end

function setWd(pMisRF::RemoteChannel,newWd)
pMis  = take!(pMisRF)
pMis.Wd = newWd;
put!(pMisRF,pMis)
return pMisRF;
end
function setDobs(pMis::Array{RemoteChannel},newDobs::Array)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setDobs,pMis[k].where,pMis[k],newDobs[k]);
		end
	end
end
return pMis;
end

function setDobs(pMisRF::RemoteChannel,newDobs)
pMis  = take!(pMisRF)
pMis.dobs = newDobs;
put!(pMisRF,pMis)
return pMisRF;
end

function setSourcesSame(pMis::Array{RemoteChannel},newSources)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setSources,pMis[k].where,pMis[k],newSources);
		end
	end
end
return pMis;
end

function setSources(pMis::Array{RemoteChannel},newSources)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setSources,pMis[k].where,pMis[k],newSources[k]);
		end
	end
end
return pMis;
end

function setSources(pMisRF::RemoteChannel,newSources)
pMis  = take!(pMisRF)
pMis.pFor.Sources = newSources;
put!(pMisRF,pMis)
return pMisRF;
end

export getWd
function getWd(pMis::Array{RemoteChannel})
Wd = Array{Array{ComplexF64}}(undef,length(pMis))
@sync begin
		for k=1:length(pMis)
	@async begin
			Wd[k] = remotecall_fetch(getWd,pMis[k].where,pMis[k]);
		end
	end
end
return Wd;
end

function getWd(pMisRF::RemoteChannel)
pMis  = fetch(pMisRF)
return pMis.Wd;
end

export getDobs
function getDobs(pMis::Array{RemoteChannel})
Dobs = Array{Array{ComplexF64}}(undef,length(pMis))
@sync begin
		for k=1:length(pMis)
	@async begin
			Dobs[k] = remotecall_fetch(getDobs,pMis[k].where,pMis[k]);
		end
	end
end
return Dobs;
end

function getDobs(pMisRF::RemoteChannel)
pMis  = fetch(pMisRF)
return pMis.dobs;
end


col_1_to_30 = [ -2.5   0.5          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.0  -2.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
1.0  -3.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.0  -2.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0         -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.5
0.0   0.0         -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.5
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0
0.0   0.0          1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -5.0   5.0  -2.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0
0.0   0.0          0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.5  -2.0   0.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.5  -2.0  -3.5   4.5  -3.0   1.5   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -5.0   5.0  -2.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.5  -2.0   0.5   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.5  -2.0  -3.5   4.5  -3.0   1.5   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -5.0   5.0  -2.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.5  -2.0   0.5   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.5  -2.0  -3.5   0.0
0.0   4.5          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -5.25
0.0   1.8735e-16   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.75
0.0  -1.249e-15    0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   0.0  -3.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   0.0   2.25
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0   0.0   9.36751e-17
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0  -6.245e-16
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
0.0   0.0          0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0]

col_31_to_40 = [
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.5          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.5           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.5           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.5          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.5           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.5    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.5   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   0.5    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.5          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.5           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.5           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.5          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.5           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.5    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.5   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   0.5    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    1.0          0.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          1.0           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           1.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           1.0          0.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          1.0           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           1.0    0.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    1.0   0.0    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.0    0.0   1.0    0.0   3.0
   -3.0          0.75          0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
   -3.0         -0.75          0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
    7.5         -7.5           0.0           0.0          0.0           0.0    0.0   0.0    0.0   0.0
   -4.5          6.75         -5.25         -3.0          0.75          0.0    0.0   0.0    0.0   0.0
   -1.8735e-16   2.81025e-16   0.75         -3.0         -0.75          0.0    0.0   0.0    0.0   0.0
    1.249e-15   -1.8735e-15   -3.0           7.5         -7.5           0.0    0.0   0.0    0.0   0.0
    0.0          0.0           2.25         -4.5          6.75         -5.25  -3.0   0.75   0.0   0.0
    0.0          0.0           9.36751e-17  -1.8735e-16   2.81025e-16   0.75  -3.0  -0.75   0.0   0.0
    0.0          0.0          -6.245e-16     1.249e-15   -1.8735e-15   -3.0    7.5  -7.5    0.0   0.0
    0.0          0.0           0.0           0.0          0.0           0.25  -0.5   0.75  -2.5   2.0
    0.0          0.0           0.0           0.0          0.0           0.5   -1.0   1.5    1.0  -2.0
]

B_data = [col_1_to_30 col_31_to_40]


"""
ULMFiT - LANGUAGE MODEL [Word-by-Word]
"""

using WordTokenizers   # For accesories
using InternedStrings   # For using Interned strings
using DelimitedFiles   # For reading and writing files
using Flux  # For building models
using Flux: Tracker, crossentropy, chunk
using LinearAlgebra: norm
using BSON: @save, @load  # For saving model weights
# using CuArrays  # For GPU support

cd(@__DIR__)
include("custom_layers.jl")      # importing AWD_LSTM (ASGD Weight-Dropped LSTM)
include("utils.jl")     # including some functions
include("WikiText103_DataDeps.jl")      # including WikiText-103 Corpus

# Language Model
mutable struct LanguageModel
    vocab :: Vector
    layers :: Flux.Chain

    function LanguageModel(embedding_size::Integer=400, hidLSTMSize::Integer=1150, outLSTMSize::Integer=embedding_size;
        embedDropProb::Float64 = 0.05, wordDropProb::Float64 = 0.4, hidDropProb::Float64 = 0.5, LayerDropProb::Float64 = 0.3, FinalDropProb::Float64 = 0.4)
        vocab = intern.(string.(readdlm("vocab.csv",',', header=false)[:, 1]))
        lm = new(vocab)
        de = DroppedEmbeddings(length(vocab), embedding_size, 0.1; init = (dims...) -> init_weights(0.1, dims...))
        lm.layers = Chain(
            de,
            VarDrop(wordDropProb),
            AWD_LSTM(embedding_size, hidLSTMSize, hidDropProb; init = (dims...) -> init_weights(1/hidLSTMSize, dims...)),
            VarDrop(LayerDropProb),
            AWD_LSTM(hidLSTMSize, hidLSTMSize, hidDropProb; init = (dims...) -> init_weights(1/hidLSTMSize, dims...)),
            VarDrop(LayerDropProb),
            AWD_LSTM(hidLSTMSize, outLSTMSize, hidDropProb; init = (dims...) -> init_weights(1/hidLSTMSize, dims...)),
            VarDrop(FinalDropProb),
            x -> de(x, true),
            softmax
        )
        return lm
    end
end

Flux.@treelike LanguageModel

# Forward
function forward(lm, batch)
    batch = broadcast(x -> lm.layers[1](indices(x, lm.vocab, "_unk_")), batch)
    batch = lm.layers[2:end].(batch)
    return batch
end

# loss funciton - Loss calculation with AR and TAR regulatization
function loss(lm, gen)
    H = forward(lm, take!(gen))
    Y = broadcast(x -> gpu(Flux.onehotbatch(x, lm.vocab, "_unk_")), take!(gen))
    l = sum(crossentropy.(H, Y))
    Flux.truncate!(lm.layers)
    return l
end

# Gradient Clipping
grad_clipping(g, upper_bound) = min(g, upper_bound)

# Backward
function back!(layers, l, opt, gradient_clip::Float64)
    # Applying gradient clipping
    l = Tracker.hook(x -> grad_clipping(x, gradient_clip), l)

    # Calulating gradients and weights updation
    p = get_trainable_params(layers)
    grads = Tracker.gradient(() -> l, p)
    Tracker.update!(opt, p, grads)
    return
end

# Funciton for training Language Model
function pretrain_lm!(lm::LanguageModel; batchsize::Integer=64, bptt::Integer=70, data_loader::Channel=load_wikitext_103, gradient_clip::Float64=0.25,
    base_lr=0.004, epochs::Integer=1, checkpoint_iter::Integer=5000)

    # Initializations
    opt = ADAM(base_lr, (0.7, 0.99))    # ADAM Optimizer
    gpu!.(lm.layers)

    # Pre-Training loops
    for epoch=1:epochs
        gen = Channel(x -> generator(x, loadCorpus(); batchsize = batchsize, bptt = bptt))
        num_of_batches = take!(gen) # Number of mini-batches
        T = num_of_iters-Int(floor((num_of_iters*2)/100))   # Averaging Trigger
        set_trigger!.(T, lm.layers)  # Setting triggers for AWD_LSTM layers
        for i=1:num_of_batches

            # FORWARD PROPAGATION
            l = loss(lm, gen)

            # BACK PROPAGATION
            back!(lm.layers, l, opt, gradient_clip)
            for layer in lm.layers[1:8]
                cpu!(layer)
            end

            # ASGD Step, after Triggering
            asgd_step!.(i, lm.layers)

            reset_masks!.(lm.layers)

            println("loss: $l", " iteration number: $i")

            # Saving checkpoints
            if i == checkpoint_iter save_model!(lm) end
        end
        println("\nEpoch: $epoch")
    end
end

# To save model
function save_model!(lm::LanguageModel, filepath::String="ULMFiT-LM.bson")
    weights = Tracker.data.(params(lm))
    @save filepath weights
end

# To load model
function load_model!(lm::LanguageModel, filepath::String=joinpath(datadep"pretrained-ULMFiT", "weights.bson"))
    @load filepath weights
    Flux.loadparams!(lm, weights)
end

function load_model!()
    lm = LanguageModel()
    load_model!(lm)
    return lm
end

function load_model!(filepath::String)
    lm = LanguageModel()
    load_model!(lm, filepath)
end


module TokyoVariantAnalysis4
using Logging
using CSV
using DataFrames
using Dates
using HypothesisTests: BinomialTest, confint

export variants_df
export date_to_value, value_to_date

@info "module", @__MODULE__

function has_name(df::DataFrame, name::String)
    return !isnothing(findfirst(isequal(name), names(df)))
end

@info "--------"
@info "作業ディレクトリ"
@info pwd()

@info "--------"
@info "基準日"
const epoch_date = Date("2021-01-01")
@info "epoch_date", epoch_date
#
date_to_value(d::Date) = Dates.value(d - epoch_date)
value_to_date(v)::Date = epoch_date + Day(v)

@info "========"
@info "データ読み込み・変換"

const organization_names = ["k", "p"]   # k 健安研, p 民間検査機関等
const variant_names = ["N501Y", "L452R"]
@info "organization_names", organization_names
@info "variant_names", variant_names

function load_df(csv_filepath)::DataFrame
    @info csv_filepath

    @info "--------"
    @info "CSV 読込み"
    csv_df = CSV.read(csv_filepath, DataFrame; missingstring="NA")
    @info size(csv_df)

    @info "--------"
    @info "データフレーム作成"
    df = DataFrame()
    @info "--------"
    @info "日付 (start_date, end_date, period_days, mid_date, date_value, date_error)"
    @assert has_name(csv_df, "date")
    dates = csv_df.date
    df.date_start  = dates
    df.date_end    = dates
    df.days        = date_to_value.(df.date_end) .+ 1.0 .- date_to_value.(df.date_start)
    df.date_error  = df.days ./ 2.0
    df.date_value  = date_to_value.(df.date_start) .+ df.date_error
    df.date_mid    = value_to_date.(round.(df.date_value .- 0.5))
    @assert all(df.date_start .≤ df.date_mid .≤ df.date_end)
    @info "df.date_start[end]", df.date_start[end]
    @info "df.date_mid[end]", df.date_mid[end]
    @info "df.date_end[end]", df.date_end[end]
    @info size(df)
    @info "--------"
    @info "変異株検査総数のコピー (ORG_n)"
    for orgname in organization_names
        colname = orgname*"_n"  # カラム名
        @assert has_name(csv_df, colname)
        df[!,colname] = csv_df[!,colname]
    end
    @info size(df)
    @info "--------"
    @info "変異株カウントのコピー (ORG_VAR)"
    for orgname in organization_names
        for varname in variant_names
            colname = orgname*"_"*varname   # カラム名
            @assert has_name(csv_df, colname)
            df[!,colname] = csv_df[!,colname]
        end
    end
    @info size(df)

    n_rows = nrow(df)
    @info "--------"
    @info "検査機関にわたる変異株検査総数の合計 (n)"
    ns::Vector{Union{Int,Missing}} = fill(missing, n_rows)
    for i in 1:n_rows
        n = missing
        for orgname in organization_names
            colname = orgname*"_n"  # 検査機関別カラム名
            orgn = df[i,colname]
            if !ismissing(orgn)
                if ismissing(n)
                    n  = orgn
                else
                    n += orgn
                end
            end
        end
        ns[i] = n
    end
    df.n = ns
    @info size(df)
    @info "--------"
    @info "検査機関にわたる変異株カウント合計 (VAR)"
    for varname in variant_names
        cs::Vector{Union{Int,Missing}} = fill(missing, n_rows)
        for i in 1:n_rows
            c = missing
            for orgname in organization_names
                colname = orgname*"_"*varname   # 検査機関別カラム名
                orgc = df[i,colname]
                if !ismissing(orgc)
                    if ismissing(c)
                        c  = orgc
                    else
                        c += orgc
                    end
                end
            end
            cs[i] = c
        end
        df[!,varname] = cs
    end
    @info size(df)

    @info "--------"
    @info "クラスター・渡航関連者を補正した L452R 変異株数と検査総数"
    @assert has_name(csv_df, "L452R_tr")   # 渡航者接触者
    @assert has_name(csv_df, "L452R_cn")   # クラスター数
    @assert has_name(csv_df, "L452R_cp")   # クラスター内人数
    ns_adj = copy(df.n)
    cs_adj = copy(df.L452R)
    for i in 1:n_rows
        n_adj = ns_adj[i]
        c_adj = cs_adj[i]
        tr = csv_df.L452R_tr[i]
        if !ismissing(tr)
            n_adj -= tr
            c_adj -= tr
        end
        cn = csv_df.L452R_cn[i]
        cp = csv_df.L452R_cp[i]
        if !ismissing(cp)
            n_adj -= cp
            c_adj -= cp
            if !ismissing(cn)
                n_adj += cn
                c_adj += cn
            end
        end
        @assert ismissing(n_adj) || (n_adj ≥ 0)
        @assert ismissing(c_adj) || (c_adj ≥ 0)
        ns_adj[i] = n_adj
        cs_adj[i] = c_adj
    end
    df.n     = ns_adj
    df.L452R = cs_adj
    return df
end

const tokyo_variants_csv_filename = "東京変異株4.csv"
variants_df = load_df(tokyo_variants_csv_filename)

@info "========"
@info "統計量算出"

const variant_ratio_pairs = [
    ("L452R", "N501Y"),
]

function statistics!(df::DataFrame)
    n_rows = nrow(df)
    @info "--------"
    @info "全体に対する割合を計算 (VAR_prop)"
    for varname in variant_names
        df[!,varname*"_prop"] = df[!,varname] ./ df.n
    end
    @info size(df)
    @info "--------"
    @info "1 変異株の全体に対する信頼区間を計算 (VAR_prop_cil, VAR_prop_ciu)"
    for varname in variant_names
        cs = df[!,varname]
        ns = df.n
        cils::Vector{Union{Float64,Missing}} = fill(missing, n_rows)
        cius::Vector{Union{Float64,Missing}} = fill(missing, n_rows)
        for i in 1:n_rows
            (cils[i], cius[i]) = if !ismissing(cs[i]) && !ismissing(ns[i])
                confint(BinomialTest(cs[i], ns[i]))
            else
                (missing, missing)
            end
        end
        df[!,varname*"_prop_cil"] = cils
        df[!,varname*"_prop_ciu"] = cius
    end
    @info size(df)
    @info "--------"
    @info "1 変異株のオッズを計算 (VAR_odds)"
    for varname in variant_names
        df[!,varname*"_odds"] = 1.0 ./ ((df.n ./ df[!,varname]) .- 1.0)
    end
    @info size(df)
    @info "--------"
    @info "オッズの信頼区間を計算 (VAR_odds_cil, VAR_odds_ciu)"
    for varname in variant_names
        cs = df[!,varname]
        ns = df.n
        cils::Vector{Union{Float64,Missing}} = fill(missing, n_rows)
        cius::Vector{Union{Float64,Missing}} = fill(missing, n_rows)
        for i in 1:n_rows
            (cils[i], cius[i]) = if !ismissing(cs[i]) && !ismissing(ns[i])
                1.0 ./ (1.0 ./ confint(BinomialTest(cs[i], ns[i])) .- 1.0)
            else
                (missing, missing)
            end
        end
        df[!,varname*"_odds_cil"] = cils
        df[!,varname*"_odds_ciu"] = cius
    end
    @info size(df)
    @info "--------"
    @info "オッズの対数（ロジット） (log_VAR_odds, log_VAR_odds_cil, log_VAR_odds_ciu)"
    for varname in variant_names
        oddname = varname*"_odds"
        df[!,"log_"*oddname]        = log.(df[!,oddname])
        df[!,"log_"*oddname*"_cil"] = log.(df[!,oddname*"_cil"])
        df[!,"log_"*oddname*"_ciu"] = log.(df[!,oddname*"_ciu"])
    end
    @info "--------"
    @info "特定 2 変異株間の比を計算 (VAR1_VAR2)"
    for (varname1, varname2) in variant_ratio_pairs
        @info varname1, varname2
        df[!,varname1*"_"*varname2] = df[!,varname1] ./ df[!,varname2]
    end
    @info size(df)
    @info "--------"
    @info "2 変異株間の比の信頼区間を計算 (VAR1_VAR2_cil, VAR1_VAR2_ciu)"
    @info "2 変異株のみからなる 2 項分布による近似"
    for (varname1, varname2) in variant_ratio_pairs
        @info varname1, varname2
        cs1 = df[!,varname1]
        cs2 = df[!,varname2]
        cils::Vector{Union{Float64,Missing}} = fill(missing, n_rows)
        cius::Vector{Union{Float64,Missing}} = fill(missing, n_rows)
        for i in 1:n_rows
            (cils[i], cius[i]) = if !ismissing(cs1[i]) && !ismissing(cs2[i])
                1.0 ./ (1.0 ./ confint(BinomialTest(cs1[i], cs1[i] + cs2[i])) .- 1.0)
            else
                (missing, missing)
            end
        end
        df[!,varname1*"_"*varname2*"_cil"] = cils
        df[!,varname1*"_"*varname2*"_ciu"] = cius
    end
    @info size(df)
    @info "--------"
    @info "2 変異株間比の対数 (log_VAR1_VAR2, log_VAR1_VAR2_cil, log_VAR1_VAR2_ciu)"
    for (varname1, varname2) in variant_ratio_pairs
        @info varname1, varname2
        pairname = varname1*"_"*varname2
        df[!,"log_"*pairname]        = log.(df[!,pairname])
        df[!,"log_"*pairname*"_cil"] = log.(df[!,pairname*"_cil"])
        df[!,"log_"*pairname*"_ciu"] = log.(df[!,pairname*"_ciu"])
    end
    @info size(df)
end

statistics!(variants_df)

@info "--------"
@info "解析データ書き出し"
const output_csv_filename = "tokyo_variant_analysis4.csv"
@info output_csv_filename
CSV.write(output_csv_filename, variants_df; missingstring="NA")
@info size(variants_df)

end #module


module NMR

"Integration range in the frequency domain, values given in ppm."
const Intrng = Tuple{Float64, Float64}

"""
Processed (frequency domain) NMR spectrum. Parameters follow the Bruker conventions:
- **re_ft**: real part of Fourier transform
- **im_ft**: imaginary part of Fourier transform
- **params**: dictionary of processing parameters, mostly from `proc` file for Bruker data
- **intrng**: list of integration ranges
- **title**: Spectrum title
"""
mutable struct ProcessedSpectrum
    re_ft :: Vector{Float64}
    im_ft :: Vector{Float64}
    params :: Dict{String, Any}
    intrng :: Array{Intrng,1}
    procno :: Int
    title :: String
end

Base.getindex(p::ProcessedSpectrum, n::Int) = p.intrng[n]
Base.getindex(p::ProcessedSpectrum, param::AbstractString) = p.params[param]
Base.getindex(p::ProcessedSpectrum, ::Colon) = p.re_ft
Base.getindex(p::ProcessedSpectrum, a::AbstractArray) = p.re_ft[a]
Base.view(p::ProcessedSpectrum, v) = @view p.re_ft[v]

"""
Unprocessed NMR experiment: FID and any associated processed spectra.

Following Bruker convention:
- **fid**: Acquired free induction decay (FID) signal
- **acqu**: Acquisition parameters, mostly from `acqu` file for Bruker data
- **procs**: Dictionary of associated processed spectra
- **default_proc**: Default processed spectrum to be used in operations that require a processed spectrum
- **name**: Experiment name
- **expno**: Experiment number if, e.g., part of a Bruker dataset
"""
mutable struct Spectrum
    fid :: Vector{Float64}
    acqu :: Dict{String, Any}
    procs :: Dict{Int,ProcessedSpectrum}
    default_proc :: Int
    name :: String
    expno :: Int
end

Base.getindex(s::Spectrum, i::Int) = s.procs[i]
Base.getindex(s::Spectrum, ::Colon) = s.procs[s.default_proc].re_ft
Base.getindex(s::Spectrum, a::AbstractArray) = s.procs[s.default_proc].re_ft[a]
Base.getindex(s::Spectrum, rng::Tuple{Float64,Float64}) = s[ppmtoindex(s,rng)]
Base.getindex(s::Spectrum, δ::Float64) = s[s.default_proc].re_ft[ppmtoindex(s, δ)]
Base.view(s::Spectrum, v) = @view s.procs[s.default_proc][v]
function Base.getindex(s::Spectrum, param::AbstractString)
    try
        s.acqu[param]
    catch err
        s[s.default_proc][param]
    end
end
Base.setindex!(s::Spectrum, d::AbstractArray, ::Colon) = (s[s.default_proc].re_ft .= d)
Base.setindex!(s::Spectrum, d::AbstractArray, r::UnitRange) = (s[s.default_proc].re_ft[r] .= d)
Base.setindex!(s::Spectrum, d, rng::Tuple{Float64, Float64}) = (s[s.default_proc].re_ft[ppmtoindex(s,rng)]=d)

Spectrum(fid :: Vector{Float64}, acqu :: Dict{Any, Any}, proc :: ProcessedSpectrum) = Spectrum(fid, acqu, Dict(1=>proc), 1, "", "")

include("utils.jl")
include("baseline.jl")
include("bruker/read.jl")
include("bruker/write.jl")
include("composition.jl")
include("decomposition.jl")
include("integration.jl")
include("interpolation.jl")
include("plotting.jl")
include("show.jl")

export Spectrum, ProcessedSpectrum, dump, # constructors, I/O
       plot, plot!, # plotting
       lsq_analyze, candidates, decompose, # decomposition
       ppmtoindex, hztoindex, ppmtohz, # unit conversion
       baseline_correct!, # processing
       limits, chemical_shifts, extract, copy!, # utility functions
       interpolate, resample, # interpolation
       integrate, # integration
       intrng, intrng_data, intrng_indices, intrng_shifts # integration regions
end


# when calculating a density product, there is a certain element selection of incoming elements
# wont yet allow too much hybridization in types but enough to get going
struct DensityProductElements{D,Vo<:AbstractVector,B,Vi<:AbstractVector}
  outElements::Vo
  outBW::Vector{B}
  outElemName::Symbol
  outdatedBW::Base.RefValue{Bool}
  lblCombinations::Vector{Vector{Int64}}
  inElements::NTuple{D,Vi}
  inBWs::NTuple{D,<:AbstractVector{B}}
  inElemNames::NTuple{D,Symbol}
  inElemFactorNames::NTuple{D,Symbol}
  # Pair{inDens, elemIdx}
  _sequencedUpdates::Channel{Pair{Int,Int}}
end

function Base.show(io::IO, x::DensityProductElements)
  # println(io, )
  printstyled(io, "DensityProductElements:", color=:blue)
  println(io)
  println(io, "  Out Name:     ", x.outElemName )
  println(io, "   in names:    ", x.inElemNames)
  println(io, "   # densities: ", length(x.inElements))
  println(io, "  # dims:       ", length(x.inElements[1][1]))
  println(io, "  # out pts:    ", length(x.outElements))
  println(io, "  .outBW[1]~:   ", x.outBW[1] .|> x->round(x,digits=5) )
  println(io, "  .outdatedBW:  ", x.outdatedBW[] )
  # println(io)
end

Base.show(io::IO, ::MIME"text/plain", x::DensityProductElements) = show(io, x)

# assume `mdp.inElements` is already set and corroborates with `lbsChs`
function _setLabelCombinations!(mdp::DensityProductElements,
                                lbsChs::Dict )
  #
  npts_ = length(lbsChs)
  ndens = length(lbsChs[1])

  resize!(mdp.outElements, npts_)
  resize!(mdp.lblCombinations, npts_)

  NptsD = length.(mdp.inElements)

  # how many points i in the product belief
  for i in 1:npts_
    # make sure memory allocations are correct
    resize!(mdp.lblCombinations[i], ndens)
    # how many densities to make the product
    for densid in 1:ndens
      # level in the kd tree
      level = length(lbsChs[i][densid])
      # get label selection for densid
      lbl = lbsChs[i][densid][level] - NptsD[densid]
      # which row in xarrTiles to mark
      mdp.lblCombinations[i][densid] = lbl
      # xarrTiles[densid][lbl,i] = 1
    end
  end

  nothing
end

# # can only do for Array, not view
# function _setProductElements!(mdp::DensityProductElements{D},
#                               prd::BallTreeDensity)
#   #
#   # also set the bandwidth
#   dim = Ndim(prd)
#   resize!(mdp.outBW, dim)

#   npts = Npts(prd)
#   for i in 1:
#     mdp.outBW[:,i] .= getBW(prd)[:,1] # fix for all elements
#   end

#   # set kernel center elements
#   resize!(mdp.outElements, )
#   for i in 1:Npts(prd)
#     resize!(mdp.outElements[i], dim)
#     mdp.outElements[i][:] .= getPoints(prd, i)
#   end

#   #
#   nothing
# end

##

"""
Likely new bug see KDE #70
"""
function _buildDensityProductElements(XX::AbstractVector{B};
                                      outName::Symbol=:product,
                                      inNames::Union{<:AbstractVector{Symbol},NTuple{D,Symbol}}=[Symbol("belief$i") for i in 1:length(XX)],
                                      inFctNames::Union{<:AbstractVector{Symbol},NTuple{D,Symbol}}=[Symbol("factor$i") for i in 1:length(XX)],
                                      _glbs=KDE.makeEmptyGbGlb(recordChoosen=true),
                                      product::B = *(XX, glbs=_glbs, addEntropy=false ) ) where {B<: BallTreeDensity, D}
  #

  npts = Npts(product)
  ndim = Ndim(product)
  ndens = length(XX)

  # pts = getPoints(product)
  # @cast outArr[j][i] := pts[i,j]
  outArr = [getPoints(product,i) for i in 1:npts]
  bw = [getBW(product)[:,i] for i in 1:npts]
  lblComb = [zeros(Int,ndens) for i in 1:npts]

  # restructure points for incoming densities
  # indens = (XX...,)
  XXarr = [([getPoints(x, i) for i in 1:Npts(x)]) for x in XX]
  BWarr = [([getBW(x)[:,i] for i in 1:Npts(x)]) for x in XX]

  # build the object
  mdp = DensityProductElements(outArr, bw, outName, Ref(true), lblComb, (XXarr...,), (BWarr...,), (inNames...,), (inFctNames...,), Channel{Pair{Int,Int}}(100))

  # set the labels selections used for current product
  if length(XX) == 1
    resize!(mdp.lblCombinations, npts)
    for i in 1:npts
      resize!(mdp.lblCombinations[i],1)
      mdp.lblCombinations[i][1] = i
    end
  else
    _setLabelCombinations!(mdp, _glbs.labelsChoosen)
  end

  return mdp
end

function _recalcProductKernel(mdp::DensityProductElements, idx::Int)
  #
  dim = length(mdp.inElements[1][1])
  selections = mdp.lblCombinations[idx]

  # TODO still have to do per dimension...
  infor = zeros(dim)
  informn = zeros(dim)
  for (idd, lbid) in enumerate(selections)
    # dens = mdp.inElements[idd]
    # mn_i = getPoints(dens, lbid)
    mn_i = mdp.inElements[idd][lbid]
    recpcov = 1/mdp.inBWs[idd][lbid].^2
    infor[:] .+= recpcov[:] # NOTE currently only diagonals
    informn += recpcov.*mn_i
  end
  cov = 1 ./ infor
  mn = informn./infor
  return mn, cov
end

##

function _updateProposalElement(mdp::DensityProductElements, elemval, densid::Int, elemidx::Int)
  # won't work for non-vector elements
  mdp.inElements[densid][elemidx] .= elemval
  mdp.outdatedBW[] = false
  nothing
end

function _updateOutElement(dpe::DensityProductElements, elemval, elemidx::Int)
  # won't work for non-vector elements
  dpe.outElements[elemidx] .= elemval
  dpe.outdatedBW[] = false
  nothing
end

function _sequenceUpdateProposalElement(mdp::DensityProductElements, elemval, densid::Int, elemidx::Int)
  _updateProposalElement(mdp, elemval, densid, elemidx)
  newpair = densid=>elemidx
  put!(mdp._sequencedUpdates, newpair)
  nothing
end

function _sequenceUpdateOutElement(dpe::DensityProductElements, elemval, elemidx::Int)
  _updateOutElement(dpe, elemval, elemidx)
  @warn("WIP")
end

_listOutElementLabelSelections(mdp::DensityProductElements, idx::Int) = mdp.lblCombinations[idx]

function _findOutElementConnectedLabel(mdp::DensityProductElements, densidx, idx::Int)
  outelems = Int[]
  # check for all elements
  for outi in 1:length(mdp.outElements)
    # if there happens to be a hit (out product relies on an in label selection)
    if _listOutElementLabelSelections(mdp, idx)[densidx] == outi
      push!(outelems, outi)
    end
  end

  # return list of outElements that depend on one particular inbound kernel `densidx=>idx`
  return outelems
end

# if something in the queue effects an outElement, then update values.
function _takeOutElementFromSequence(mdp::DensityProductElements)
  if !isready(mdp._sequencedUpdates)
    return nothing
  end
  nextpair = take!(mdp._sequencedUpdates)

  outElemList = _findOutElementConnectedLabel(mdp, nextpair[1], nextpair[2])
  println("Connected label $outElemList")

  return nothing
end


module Absynth

using MacroTools: walk, postwalk, prewalk, @capture, @match, replace, striplines
using Combinatorics
using LinearAlgebra
using MultivariatePolynomials
using DynamicPolynomials
using Dates
using DataFrames: DataFrame, unstack, rename!
using ProgressMeter

const Var = AbstractVariable
const Poly = AbstractPolynomialLike
const SymOrNum = Union{Symbol,Number}

include("nlsat/NLSat.jl")

using Absynth.NLSat

export Z3Solver, SMTSolver, SMTSolverZ3, SMTSolverYices

include("utils.jl")
include("poly.jl")
include("invariant.jl")
include("cfinite.jl")
include("show.jl")
include("template.jl")
include("types.jl")
include("iterators.jl")

export FullSymbolic, UpperTriangular, UnitUpperTriangular, Companion, UserSpecific
export Invariant, @invariant
export solutions, models, synth
export RecSystem, loop
export RecTemplate
export AllPermutations, FixedPermutations
export AllPartitions, FixedPartitions
export Strategy
export @cstr

include("../benchmark/examples.jl")

end # module


function train!(agent::AbstractAgent, game::SnakeAI.Game)
    # Get the current step
    old_state = SnakeAI.get_state(game)

    # Get the predicted move for the state
    move = get_action(agent, old_state)
    SnakeAI.send_inputs!(game, move)

    # Play the step
    reward, done, score = SnakeAI.play_step!(game)
    new_state = SnakeAI.get_state(game)

    # Train the short memory
    train_short_memory(agent, old_state, move, reward, new_state, done)

    # Remember
    remember(agent, old_state, move, reward, new_state, done)

    if done
        # Reset the game
        train_long_memory(agent)
        SnakeAI.reset!(game)
        agent.n_games += 1

        if score > agent.record
            agent.record = score
            # save_model(joinpath(MODELS_PATH, "model_$(agent.n_games).bson"), agent.model)
        end
    end

    return done
end

function remember(
    agent::AbstractAgent,
    state::S,
    action::S,
    reward::T,
    next_state::S,
    done::Bool
) where {T<:Integer,S<:AbstractArray{<:T}}
    push!(agent.memory.data, (state, action, [reward], next_state, convert.(Int, [done])))
end

function train_short_memory(
    agent::AbstractAgent,
    state::S,
    action::S,
    reward::T,
    next_state::S,
    done::Bool
) where {T<:Integer,S<:AbstractArray{<:T}}
    update!(agent, state, action, reward, next_state, done)
end

function train_long_memory(agent::AbstractAgent)
    if length(agent.memory.data) > BATCH_SIZE
        mini_sample = sample(agent.memory.data, BATCH_SIZE)
    else
        mini_sample = agent.memory.data
    end

    states, actions, rewards, next_states, dones = map(x -> getfield.(mini_sample, x), fieldnames(eltype(mini_sample)))

    update!(agent, states, actions, rewards, next_states, dones)
end

function get_action(agent::SnakeAgent, state::AbstractArray{<:Integer}; rand_range=1:200)
    agent.ϵ = 80 - agent.n_games
    final_move = zeros(Int, 3)

    if rand(rand_range) < agent.ϵ
        move = rand(1:3)
        final_move[move] = 1
    else
        pred = agent.model(state)
        final_move[Flux.onecold(pred)] = 1
    end

    return final_move
end

function update!(
    agent::SnakeAgent,
    state::Union{A,AA},
    action::Union{A,AA},
    reward::Union{T,AA},
    next_state::Union{A,AA},
    done::Union{Bool,AA};
    α::Float32=0.9f0    # Step size
) where {T<:Integer,A<:AbstractArray{<:T},AA<:AbstractArray{A}}
    # Batching the states and converting data to Float32 (done implicitly otherwise)
    state = Flux.batch(state) |> x -> convert.(Float32, x)
    next_state = Flux.batch(next_state) |> x -> convert.(Float32, x)
    action = Flux.batch(action) |> x -> convert.(Float32, x)
    reward = Flux.batch(reward) |> x -> convert.(Float32, x)
    done = Flux.batch(done)

    # Model's prediction for next state
    y = agent.model(next_state)

    # Get the model's params for back propagation
    ps = Flux.params(agent.model)

    # Calculate the gradients
    gs = Flux.gradient(ps) do
        # Forward pass
        ŷ = agent.model(state)

        # Creating buffer to allow mutability when calculating gradients
        Rₙ = Buffer(ŷ, size(ŷ))

        # Adjusting values of current state with next state's knowledge
        for idx in 1:length(done)
            # Copy preds into buffer
            Rₙ[:, idx] = ŷ[:, idx]

            Qₙ = reward[idx]
            if done[idx] == false
                Qₙ += α * maximum(y[:, idx])
            end

            # Adjusting the expected reward for selected move
            Rₙ[argmax(action[:, idx]), idx] = Qₙ
        end
        # Calculate the loss
        agent.criterion(ŷ, copy(Rₙ))
    end

    # Update model weights
    Flux.Optimise.update!(agent.opt, ps, gs)
end


#!/usr/bin/env julia
# this script is thought to be designed to counduct SNAQ! on CFs
# which were obtaines from SNPs by SNPs2CF.
# input: 	- csv output of SNPs2CF
# 		- hypothetical tree as starting topology

#using Pkg # to use functions that manage packages
#Pkg.add("PhyloNetworks") # to download & install package PhyloNetworks
#Pkg.add("PhyloPlots")
#Pkg.add("RCall")      # packaage to call R from within julia
#Pkg.add("CSV")        # to read from / write to text files, e.g. csv files
#Pkg.add("DataFrames") # to create & manipulate data frames
#Pkg.add("StatsModels")# for regression formulas
using PhyloNetworks   # may take some time: pre-compiles functions in that package
# using PhyloPlots
# using Distributed

print("we are running")

snpCF = readTableCF(ARGS[1]) # input our CF table
tre = readTopology(ARGS[2]) # input a starting topology

# print(tre)

n0_out = ARGS[1][1:end-4] * "_net0"
# n1_out = ARGS[1][1:end-4] * "_net1"
# n2_out = ARGS[1][1:end-4] * "_net2"
# plot_out = ARGS[1][1:end-4] * "_plot.pdf"

# print(n0_out)

rseed = rand((100:999))
net0 = snaq!(tre,  snpCF, hmax=0, filename=n0_out, seed=rseed)
# rseed = rand((100:999))
# net1 = snaq!(net0, snpCF, hmax=1, filename=n1_out, seed=rseed)

# addprocs(2) # tells Julia to add 2 worker cores
# @everywhere using PhyloNetworks # tells these cores to load the PhyloNetworks package for themselves
# net2 = snaq!(net1, snpCF, hmax=2, filename=n2_out, seed=789, runs=5)

# plot(net0, :R);

# R"pdf"(plot_out, width=3, height=3);
# plot(net0, :R);
# R"dev.off()";


module LightPropagation

using DelimitedFiles
using LsqFit
using DSP
using FFTW
using Parameters

export TPSF_DA_semiinf_refl
export TPSF_DA_slab_refl
export TPSF_DA_paralpip_refl

export getfit

export load_asc_data

export fitDTOF, DTOF_fitparams

include("forwardmodels/TPSF/DiffusionApproximation/refl_TD.jl")
include("data/loadData/loaddata.jl")

include("inversefitting/FitStructures/fitstructures.jl")
include("inversefitting/FitFunctions/fitfunction.jl")

end


export texture, textureMode, textureWrap, delTextures

function texture(img::Image)
	imagesep = separate(img)
	imagedata = reinterpret(Float64, map(Float64, data(imagesep)))

	# the following takes the image and turns into a 1D array that opengl
	# expects (i.e., the first 3 elements are the RGB values of the first
	# pixel, the next 3 elements are the RGB values of the second pixel, and
	# so on).
	# there must be an easier and faster way to do this with some sort of
	# opengl trick, though.
	if imagesep.properties["colordim"] == 3
		w = size(imagedata,2)
		h = size(imagedata,1)
		imageGL = Array(Float64, w*h*3)
		@inbounds for y=1:h
			row = 3 * w * (y - 1)
			for x=1:w
				col = 3 * (x - 1)
		        @simd for l=1:3
		            imageGL[row+col+l] = imagedata[y,x,l]
		        end
		    end
		end
	elseif imagesep.properties["colordim"] == 1
		w = size(imagedata,2)
		h = size(imagedata,3)
		imageGL = Array(Float64, w*h*3)
		@inbounds for y=1:h
			row = 3 * w * (y - 1)
			for x=1:w
				col = 3 * (x - 1)
		        @simd for l=1:3
		            imageGL[row+col+l] = imagedata[l,x,y]
		        end
		    end
		end
	end

	tex = GLuint[0]
	glGenTextures(1, tex)
	glActiveTexture(GL_TEXTURE2)
	glBindTexture(GL_TEXTURE_2D, tex[1])
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB32F, w, h, 0, GL_RGB, GL_FLOAT, map(Float32, imageGL))
	glGenerateMipmap(GL_TEXTURE_2D)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_NEAREST)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
	max_aniso = [Float32(0.0)]
	glGetFloatv(GL_TEXTURE_MAX_ANISOTROPY_EXT, max_aniso)
	# glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAX_ANISOTROPY_EXT, max_aniso[1])
	# this will require some performance testing. it might be too heavy for
	# psychophysics
	glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAX_ANISOTROPY_EXT, 16)
	glBindTexture(GL_TEXTURE_2D, C_NULL)

	return tex[1], w/state.width, h/state.height
end

function texture(data)
	w = size(data,2)
	h = size(data,1)

	tex = GLuint[0]
	glGenTextures(1, tex)
	glActiveTexture(GL_TEXTURE2)
	glBindTexture(GL_TEXTURE_2D, tex[1])
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB32F, w, h, 0, GL_RGB, GL_FLOAT, map(Float32, data[:]))
	glGenerateMipmap(GL_TEXTURE_2D)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_NEAREST)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
	max_aniso = [Float32(0.0)]
	glGetFloatv(GL_TEXTURE_MAX_ANISOTROPY_EXT, max_aniso)
	# glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAX_ANISOTROPY_EXT, max_aniso[1])
	# this will require some performance testing. it might be too heavy for
	# psychophysics
	glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAX_ANISOTROPY_EXT, 16)
	glBindTexture(GL_TEXTURE_2D, C_NULL)

	return tex[1], w/state.width, h/state.height
end

function textureMode(mode)

end

function textureWrap(wrap)
	if wrap == "CLAMP"
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
	elseif wrap == "REPEAT"
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
	end
	return nothing
end

function delTextures(texs)
	glDeleteTextures(length(texs), collect(texs))
end


"""
afmize

simulate afm image from a given structure

example::
pdb = readpbd("1ake.pdb");
decenter!(pdb);
afm_image = afmize(pdb, (150.0, 150.0), (16, 16));
heatmap(afm_image) # using Plots
"""
function afmize_alpha(ta::TrjArray, box, npixel)::Matrix{Float64}
    rpixel = box ./ npixel
    nx = floor.(Int, (ta.x[1, :] .+ 0.5*box[1]) ./ rpixel[1]);
    ny = floor.(Int, (ta.y[1, :] .+ 0.5*box[2]) ./ rpixel[2]);
    #@show nx
    #@show ny
    afm_image = zeros(Float64, npixel)
    z_min = minimum(ta.z[1, :])
    z = ta.z[1, :] .- z_min
    for iatom in 1:ta.natom
        if (1 <= nx[iatom] <= npixel[1]) && (1 <= ny[iatom] <= npixel[2])
            if z[iatom] > afm_image[nx[iatom], ny[iatom]]
                afm_image[nx[iatom], ny[iatom]] = z[iatom]
            end
        end
    end
    return afm_image
end

function defaultParameters()
    return Dict(
     "H" => 1.20,
    "HE" => 1.40,
     "B" => 1.92,
     "C" => 1.70,
     "N" => 1.55,
     "O" => 1.52,
     "F" => 1.47,
    "NE" => 1.54,
    "NA" => 2.27,
    "MG" => 1.73,
    "AL" => 1.84,
    "SI" => 2.10,
     "P" => 1.80,
     "S" => 1.80,
    "CL" => 1.75,
    "AR" => 1.88,
     "K" => 2.75,
    "CA" => 2.31,
    "CYS" => 2.75,
    "PHE" => 3.2,
    "LEU" => 3.1,
    "TRP" => 3.4,
    "VAL" => 2.95,
    "ILE" => 3.1,
    "MET" => 3.1,
    "HIS" => 3.05,
    "HSD" => 3.05,
    "TYR" => 3.25,
    "ALA" => 2.5,
    "GLY" => 2.25,
    "PRO" => 2.8,
    "ASN" => 2.85,
    "THR" => 2.8,
    "SER" => 2.6,
    "ARG" => 3.3,
    "GLN" => 3.0,
    "ASP" => 2.8,
    "LYS" => 3.2,
    "GLU" => 2.95,
    )
end

mutable struct Point2D{T<:Real}
    x::T
    y::T
end

mutable struct Point3D{T<:Real}
    x::T
    y::T
    z::T
end

struct Probe
    x::Float64
    y::Float64
    r::Float64
    angle::Float64
end

mutable struct Sphere
    x::Float64
    y::Float64
    z::Float64
    r::Float64
end

mutable struct AfmizeConfig
    probeAngle::Float64        # Radian
    probeRadius::Float64
    range_min::Point2D
    range_max::Point2D
    resolution::Point2D
    atomRadiusDict::Dict{String, Float64}
end

function AfmizeConfig(r::Float64)
    return AfmizeConfig(10.0 * (pi / 180),
                        r,
                        Point2D(-90, -90),
                        Point2D(90, 90),
                        Point2D(15, 15),
                        defaultParameters())
end

function defaultConfig()
    return AfmizeConfig(10.0 * (pi / 180),
                        20.0,
                        Point2D(-90, -90),
                        Point2D(90, 90),
                        Point2D(15, 15),
                        defaultParameters())
end

function checkConfig(tra::TrjArray ,config::AfmizeConfig)::Union{String, Nothing}
    if (config.range_max.x - config.range_min.x) % config.resolution.x != 0
        return "resolution.x must divide range"
    end
    if (config.range_max.y - config.range_min.y) % config.resolution.y != 0
        return "resolution.y must divide range"
    end
    if config.probeAngle == 0
        return "probeAngle must be positive"
    end

    for name in tra.atomname
        if !haskey(config.atomRadiusDict, name)
            return "config dosen't know atom name $(name)"
        end
    end

    return nothing
end

# 底を0にする
function moveBottom(atoms::Array{Sphere})
    bottom = atoms[1].z - atoms[1].r
    for atom in atoms
        bottom = min(bottom, atom.z - atom.r)
    end
    for atom in atoms
        atom.z -= bottom
    end
end

# 針を球として見たとき、どこで衝突するかの計算
function calcCollisionAsSphere(probe::Probe, atom::Sphere)
    distXY = sqrt((probe.x - atom.x)^2 + (probe.y - atom.y)^2)
    dr = probe.r + atom.r
    if 0 < dr^2 - distXY^2
        return atom.z + sqrt(dr^2 - distXY^2) - probe.r
    else
        return 0.0
    end
end

# 針を円形推台として見たとき、どこで衝突するかの計算
function calcCollisionAsCircularThrusters(probe::Probe, atom::Sphere)
    distXY = sqrt((probe.x - atom.x)^2 + (probe.y - atom.y)^2)
    collisionDist = probe.r + atom.r * cos(probe.angle)
    if collisionDist < distXY
        return atom.z + atom.r * sin(probe.angle) - (distXY - collisionDist) / tan(probe.angle) - probe.r
    elseif probe.r < distXY
        return atom.z + sqrt(atom.r^2 - (distXY - probe.r)^2) - probe.r
    else
        return 0.0
    end
end

# 衝突計算をする範囲を返す
function calRectangle(min_point::Point2D, max_point::Point2D, config::AfmizeConfig)
    if max_point.x <= config.range_min.x || max_point.y <= config.range_min.y
        return nothing
    end
    if config.range_max.x <= min_point.x || config.range_max.y <= min_point.y
        return nothing
    end

    min_point = Point2D(max(min_point.x, config.range_min.x),
                        max(min_point.y, config.range_min.y))
    max_point = Point2D(min(max_point.x, config.range_max.x),
                        min(max_point.y, config.range_max.y))
    width  = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    return (max(1, floor(Int, (min_point.x - config.range_min.x + config.resolution.x - 1) / config.resolution.x)),
            max(1, floor(Int, (min_point.y - config.range_min.y + config.resolution.y - 1) / config.resolution.y)),
            min(width , floor(Int, (max_point.x - config.range_min.x + config.resolution.x - 1) / config.resolution.x)),
            min(height, floor(Int, (max_point.y - config.range_min.y + config.resolution.y - 1) / config.resolution.y)))
end

function calRectangle_sphere(atom::Sphere, config::AfmizeConfig)
    min_point = Point2D(atom.x - atom.r - config.probeRadius,
                        atom.y - atom.r - config.probeRadius)
    max_point = Point2D(atom.x + atom.r + config.probeRadius,
                        atom.y + atom.r + config.probeRadius)

    return calRectangle(min_point, max_point, config)
end

function calRectangle_circularThrusters(atom::Sphere, config::AfmizeConfig)
    dist_xy = (tan(config.probeAngle) * (atom.z + atom.r * sin(config.probeAngle) - config.probeRadius)
                + config.probeRadius + atom.r * cos(config.probeAngle))
    min_point = Point2D(atom.x - dist_xy,
                        atom.y - dist_xy)
    max_point = Point2D(atom.x + dist_xy,
                        atom.y + dist_xy)

    return calRectangle(min_point, max_point, config)
end

function afmize(tra::TrjArray, config::AfmizeConfig)
    message = checkConfig(tra, config)
    if !isnothing(message)
        println(message)
        return zeros(1, 1)
    end

    width = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    atoms = [Sphere(tra.x[i], tra.y[i], tra.z[i],
            config.atomRadiusDict[tra.atomname[i]]) for i in 1:tra.natom]
    moveBottom(atoms)

    stage = zeros(height, width)
    probes = [Probe(config.range_min.x + (w-0.5) * config.resolution.x,
                    config.range_min.y + (h-0.5) * config.resolution.y,
                    config.probeRadius, config.probeAngle)
             for h in 1:height, w in 1:width]

    # 各原子事に計算する価値のある矩形を求めて、その中で探索を行う
    for atom in atoms
        rectangle = calRectangle_sphere(atom, config)

        if isnothing(rectangle) continue end

        for h in rectangle[2]:rectangle[4], w in rectangle[1]:rectangle[3]
            probe = probes[h, w]

            stage[h, w] = max(stage[h, w], calcCollisionAsSphere(probe, atom))
        end
    end

    for atom in atoms
        rectangle = calRectangle_circularThrusters(atom, config)

        if isnothing(rectangle) continue end

        for h in rectangle[2]:rectangle[4], w in rectangle[1]:rectangle[3]
            probe = probes[h, w]

            stage[h, w] = max(stage[h, w], calcCollisionAsCircularThrusters(probe, atom))
        end
    end

    return stage
end

# 高速化前
function afmize_beta(tra::TrjArray, config::AfmizeConfig)
    message = checkConfig(tra, config)
    if !isnothing(message)
        println(message)
        return zeros(1, 1)
    end

    width = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    atoms = [Sphere(tra.x[i], tra.y[i], tra.z[i],
            config.atomRadiusDict[tra.atomname[i]]) for i in 1:tra.natom]
    moveBottom(atoms)

    stage = zeros(height, width)

    for h in 1:height, w in 1:width
        probe = Probe(config.range_min.x + (w-0.5) * config.resolution.x,
                      config.range_min.y + (h-0.5) * config.resolution.y,
                      config.probeRadius, config.probeAngle)
        for atom in atoms
            stage[h, w] = max(stage[h, w], calcCollisionAsSphere(probe, atom))
            stage[h, w] = max(stage[h, w], calcCollisionAsCircularThrusters(probe, atom))
        end
    end

    return stage
end

function afmize_gpu(tra::TrjArray, config::AfmizeConfig)
    message = checkConfig(tra, config)
    if !isnothing(message)
        println(message)
        return zeros(1, 1)
    end

    width = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    atom_r = tra.mass
    atom_z = tra.z[:]
    atom_z = atom_z .- (minimum(atom_z) - atom_r[argmin(atom_z)])
    #atom_z .= atom_z .- (minimum(atom_z))
    atom_z_max = maximum(atom_z)
    stage = similar(tra.x, 1, height*width)
    stage .= 0.0

    probe_x = similar(tra.x, 1, width)
    probe_y = similar(tra.x, 1, height)
    probe_x[:] .= range(config.range_min.x, step=config.resolution.x, stop=config.range_min.x + (width-1)*config.resolution.x) .+ (0.5*config.resolution.x)
    probe_y[:] .= range(config.range_min.y, step=config.resolution.y, stop=config.range_min.y + (height-1)*config.resolution.y) .+ (0.5*config.resolution.y)
    probe_r = config.probeRadius
    probe_angle = config.probeAngle

    # collision with sphere
    dist_x2 = (tra.x[:] .- probe_x).^2
    dist_y2 = (tra.y[:] .- probe_y).^2
    index_x_to_xy = [j for j=1:width for i=1:height]
    index_y_to_xy = [i for j=1:width for i=1:height]

    dist_xy2 = dist_x2[:, index_x_to_xy] .+ dist_y2[:, index_y_to_xy]
    #stage_xy = similar(dist_xy2)
    #stage_xy .= 0.0
    dr = probe_r .+ atom_r
    dr2 = dr.^2
    s = dr2 .- dist_xy2
    @show typeof(s)
    index_collide = s .> 0.0
    index_not_collide = .!index_collide
    @show typeof(index_collide)
    if any(index_collide)
        s[Array(index_collide)] .= sqrt.(s[index_collide])
    end
    s .= atom_z .+ s .- probe_r
    if any(index_not_collide)
        s[Array(index_not_collide)] .= 0.0
    end
    s_max = maximum(s, dims=1)
    stage .= max.(stage, s_max)

    # collision with circular thruster
    dist_xy2 .= sqrt.(dist_xy2)
    dist_collision = probe_r .+ atom_r .* cos(probe_angle)
    index_side   = dist_collision .< dist_xy2
    index_corner = (probe_r .< dist_xy2) .& .!index_side
    index_not_collide = .!(index_corner .| index_side)
    s[Array(index_side)] .= (atom_r .* sin.(probe_angle) .- (dist_xy2 .- dist_collision) ./ tan.(probe_angle))[index_side]
    s[Array(index_corner)] .= sqrt.( (atom_r.^2 .- (dist_xy2 .- probe_r).^2)[index_corner]  )
    s .= s .+ atom_z .- probe_r
    s[Array(index_not_collide)] .= 0.0
    s_max = maximum(s, dims=1)
    stage .= max.(stage, s_max)

    stage_reshape = reshape(stage, (height, width))
    return stage_reshape
end

function translateafm(afm, (dx, dy))
    afm_translated = zeros(eltype(afm), size(afm))
    (nx, ny) = size(afm)
    for i in maximum([1, 1-dx]):minimum([nx, nx-dx])
        for j in maximum([1, 1-dy]):minimum([ny, ny-dy])
            afm_translated[i+dx, j+dy] = afm[i, j]
        end
    end
    afm_translated
end

function translateafm_periodic(afm, (dx, dy))
    afm_translated = zeros(eltype(afm), size(afm))
    (nx, ny) = size(afm)
    for x in 1:nx
        for y in 1:ny
            xx = x + dx
            yy = y + dy
            if xx > nx
                xx -= nx
            end
            if yy > ny
                yy -= ny
            end
            afm_translated[xx, yy] = afm[x, y]
        end
    end
    afm_translated
end

function direct_convolution(observed, calculated)
    H, W = size(observed)
    ret = zeros(H, W)
    for y in 1:H, x in 1:W
        for dy in 1:H, dx in 1:W
            ny = y + dy - 1
            nx = x + dx - 1
            if ny > H ny -= H end
            if nx > W nx -= W end
            ret[y, x] += observed[dy, dx] * calculated[ny, nx]
        end
    end
    return ret
end

function fft_convolution(observed, calculated)
    return real.(ifft(fft(observed).*conj.(fft(calculated))))
end

function calcLogProb(observed::AbstractMatrix{T}, calculated::AbstractMatrix{T}, convolution_func) where {T}
    npix = eltype(observed)(size(observed, 1) * size(observed, 2))
    @show npix
    C_o  = sum(observed)
    C_c  = sum(calculated)
    C_cc = sum(calculated.^2)
    C_oo = sum(observed.^2)
    C_oc = convolution_func(observed, calculated)

    log01 = npix .* (C_cc .* C_oo .- C_oc.^2) .+ eltype(observed)(2) .* C_o .* C_oc .* C_c .- C_cc .* C_o.^2 .- C_oo .* C_c.^2
    #log01[Array(log01 .< log(eps(eltype(observed))))] .= eps(eltype(observed))
    log02 = (npix .- eltype(observed)(2)) .* (npix .* C_cc .- C_c.^2)
    #log02 = log02 <= log(eps(eltype(observed))) ? eps(eltype(observed)) : log02
    @show typeof(log01)
    @show typeof(log02)
    logprob = eltype(observed)(0.5) .* (eltype(observed)(3.0) .- npix) .* log.(log01) .+ (eltype(observed)(0.5) .* npix .- eltype(observed)(2.0)) .* log(log02)

    return logprob
end

function getafmposterior(afm::Matrix{Float64}, model_array::TrjArray, q_array::Matrix{Float64}, param_array)
    imax_model = 0
    imax_q = 0
    best_param = param_array[1]
    best_translate = [0, 0]
    best_afm = similar(afm)
    best_posterior = -Inf

    observed = afm
    npix = Float64(size(observed, 1) * size(observed, 2))

    decenter!(model_array)

    Threads.@threads for imodel in 1:size(model_array, 1)
        @show imodel
        model = model_array[imodel, :]
        for iq in 1:size(q_array, 1)
            q = q_array[iq, :]
            model_rotated = MDToolbox.rotate(model, q)
            for iparam in 1:length(param_array)
                param = param_array[iparam]
                calculated = afmize(model_rotated, param)
                C_o  = sum(observed)
                C_c  = sum(calculated)
                #@btime C_oc = sum(observed_translated .* calculated)
                C_oc = sum(observed .* calculated)
                C_cc = sum(calculated.^2)
                C_oo = sum(observed.^2)
                #@btime C_oc_dxdy = real.(ifftshift(ifft(fft(observed_translated).*conj.(fft(calculated)))))
                C_oc_dxdy = real.(ifftshift(ifft(fft(observed).*conj.(fft(calculated)))))
                log01 = npix .* (C_cc .* C_oo .- C_oc_dxdy.^2) .+ 2.0 .* C_o .* C_oc_dxdy .* C_c .- C_cc .* C_o.^2 .- C_oo .* C_c.^2
                log01[log01 .<= 0.0] .= eps(Float64)
                log02 = (npix .- 2.0) .* (npix .* C_cc .- C_c.^2)
                log02 = log02 <= 0 ? eps(Float64) : log02
                logprob = 0.5 .* (3.0 .- npix) .* log.(log01) .+ (0.5 .* npix .- 2.0) .* log.(log02)
                maximum_logprob = maximum(logprob)
                if best_posterior < maximum_logprob
                    best_posterior = maximum_logprob
                    imax_model = imodel
                    imax_q = iq
                    best_param = param
                    x_center = ceil(Int32, (size(observed,1)/2)+1.0)
                    y_center = ceil(Int32, (size(observed,2)/2)+1.0)
                    dx_estimated = argmax(logprob)[1] - x_center
                    dy_estimated = argmax(logprob)[2] - y_center
                    best_translate = (dx_estimated, dy_estimated)
                    best_afm = translateafm(calculated, best_translate)
                end
            end
        end
    end

    return imax_model, imax_q, best_param, best_translate, best_afm, best_posterior
end

mutable struct posteriorResult
    posteriors
    each_best
    afm_results
    posterior_results
    best_translate
    best_posterior
    best_model
    best_quate
    best_model_rotated
    best_param
    best_afm
end

function getafmposteriors_alpha(afm_frames, model_array, quate_array, param_array, opt = "")
    frame_num = size(afm_frames)[1]
    model_num = size(model_array)[1]
    quate_num = size(quate_array)[1]
    param_num = size(param_array)[1]

    convolution_func = fft_convolution
    if opt == "direct"
        convolution_func = direct_convolution
    end

    println("func is $convolution_func")

    results = []
    for i in 1:frame_num
        # 各モデルについて、一つの角度と半径ごとのlogprobの最高値を保持しておく
        posteriors = zeros(model_num, quate_num * param_num)
        # 各モデルについて、最も良い値をだした時のafm画像
        afm_results = [zeros(size(afm_frames[1])) for i in 1:model_num]
        each_best = ones(model_num) .* -Inf

        best_posterior = -Inf
        best_model = model_array[1]
        best_quate = quate_array[1, :]
        best_model_rotated = model_array[1]
        best_param = param_array[1]
        best_afm = zeros(size(afm_frames[1]))
        posterior_results = zeros(model_num)
        push!(results, posteriorResult(posteriors,
                                        each_best,
                                        afm_results,
                                        posterior_results,
                                        (0, 0),
                                        best_posterior,
                                        best_model,
                                        best_quate,
                                        best_model_rotated,
                                        best_param,
                                        best_afm))
    end

    for (model_id, model) in zip(1:model_num, model_array)
        @show model_id
        for quate_id in 1:quate_num
            rotated_model = MDToolbox.rotate(model, quate_array[quate_id, :])
            for param_id in 1:param_num
                cal_frame = afmize(rotated_model, param_array[param_id])
                for frame_id in 1:frame_num
                    prob_mat = calcLogProb(afm_frames[frame_id], cal_frame, convolution_func)
                    max_prob = maximum(prob_mat)
                    id = (quate_id - 1) * param_num + param_id
                    results[frame_id].posteriors[model_id, id] = max_prob

                    if results[frame_id].each_best[model_id] < max_prob
                        results[frame_id].each_best[model_id] = max_prob
                        results[frame_id].afm_results[model_id] = cal_frame
                    end

                    if results[frame_id].best_posterior < max_prob
                        results[frame_id].best_posterior = max_prob
                        results[frame_id].best_model = model
                        results[frame_id].best_quate = quate_array[quate_id, :]
                        results[frame_id].best_model_rotated = rotated_model
                        results[frame_id].best_param = param_array[param_id]
                        results[frame_id].best_translate = Tuple(argmax(prob_mat))
                        results[frame_id].best_afm = translateafm_periodic(cal_frame, results[frame_id].best_translate)
                    end
                end
            end
        end
    end

    for frame_id in 1:frame_num
        results[frame_id].posteriors .-= maximum(results[frame_id].posteriors)
        results[frame_id].posteriors = exp.(results[frame_id].posteriors)
        results[frame_id].posterior_results = sum(results[frame_id].posteriors, dims = 2)
        results[frame_id].posterior_results ./= sum(results[frame_id].posterior_results)
    end

    return results
end


module MultiAgentPOMDPs

using POMDPs

abstract type JointMDP{S, A<:AbstractVector} <: MDP{S, A} end

"""
    function n_agents(m::JointMDP)
Return the number of agents in the (PO)MDP
"""
function n_agents end

"""
    function agent_actions(m::JointMDP, idx::Int64, s::S) where S
Returns the discrete actions for the given agent index.

!!! note
    This will be called a LOT so it should not allocate each time....
"""
function agent_actions end

"""
    function agent_states(m::JointMDP, idx::Int64)
Returns the discrete states for the given agent index
"""
function agent_states end

"""
    function agent_actionindex(m::JointMDP, idx::Int64, a::A) where A
Returns the integer index of action `a` for agent `idx`. Used for discrete models only.
"""
function agent_actionindex end

"""
    function agent_actionindex(m::JointMDP, idx::Int64, s::S) where S
Returns the integer index of state `s` for agent `idx`. Used for discrete models only.
"""
function agent_stateindex end

"""
    function coordination_graph(m::JointMDP)
    function coordination_graph(m::JointMDP, s::S) where S
Returns the LightGraphs.SimpleGraph (or any appropriate structure) for the coordination graph.
"""
function coordination_graph end

export JointMDP, n_agents, agent_actions, agent_states, agent_actionindex, agent_stateindex, agent_reward, coordination_graph
end


struct K_Vector{t,C,T} <: AbstractVector{T}
    a::Vector{C}
    function K_Vector{t,C,T}(x::K_) where {t,C,T}
        a = asarray(x)
        t′ = xt(x)
        if t != t′
            throw(ArgumentError("type mismatch: t=$t, t′=$t′"))
        end
        return new(a)
    end
end
ktypecode(::K_Vector{t,C,T}) where {t,C,T} = t
const K_Chars = K_Vector{C_(10),G_,UInt8}
Base.convert(::Type{String}, x::K_Chars) = String(x.a)
Base.print(io::IO, x::K_Chars) = print(io, String(x))
Base.show(io::IO, x::K_Chars) = print(io, "K(", repr(String(x)), ")")
Base.show(io::IO, ::MIME"text/plain", x::K_Chars) = show(io, x)
Base.serialize(io::AbstractSerializer, x::K_Vector) =
    _serialize(io, x, typeof(x))
function K_Vector(x::K_)
    t = xt(x)
    ti = typeinfo(t)
    K_Vector{t,ti.c_type,ti.jl_type}(x)
end
const KnownLength = Union{Base.HasLength, Base.HasShape}
_vector(t::Integer, n::Integer) = K_Vector(ktn(t, n))
_vector(::Type{T}, n::Integer) where T = _vector(ktypecode(T), n)
_vector(t::Integer, itr, ::KnownLength) =
    (r = _vector(t, Int(length(itr)::Integer));copy!(r, itr);r)
K_Vector{t,C,T}(itr) where {t,C,T} =
    _vector(t, itr, Base.iteratorsize(itr))::K_Vector{t,C,T}
# XXX: Do we need this?
K_Vector(a::Vector) = K_Vector(K(a))

Base.eltype(v::K_Vector{t,C,T}) where {t,C,T} = T
Base.size(v::K_Vector{t,C,T}) where {t,C,T} = size(v.a)
Base.getindex(v::K_Vector{t,C,T}, i::Integer) where {t,C,T} = _cast(T, v.a[i])
# Setting the vector elements
Base.setindex!(v::K_Vector{t,C,T}, el, i::Integer) where {t,C,T} = begin
    v.a[i] = _cast(C, T(el))
    v
end

# See julia.h.
struct jl_array_t
    data   :: Ptr{Void} # (1)  sizeof(Ptr)
    length :: Csize_t   # (2) + sizeof(Ptr)
    flags  :: UInt16    # (3) + 2 bytes
    elsize :: UInt16    # (4) + 2 bytes
    offset :: UInt32    # (5) + 4 bytes
    nrows  :: Csize_t   # (6) = at 2*sizeof(Ptr) + 8
    maxsize:: Csize_t   # (7)
end

const offset_length = fieldoffset(jl_array_t, 2)
const offset_nrows = fieldoffset(jl_array_t, 6)

function resetvector(a::Vector, size::Integer, data::Ptr)
    p = Ptr{jl_array_t}(pointer_from_objref(a))
    # Pre-checks
    d = unsafe_load(p)
    @assert d.data == pointer(a)
    @assert d.length == d.nrows == length(a)
    # replant the data pointer
    unsafe_store!(Ptr{Ptr{Void}}(p), data)
    # update the size
    unsafe_store!(Ptr{Csize_t}(p + offset_length), size)
    unsafe_store!(Ptr{Csize_t}(p + offset_nrows), size)
    # Post-checks
    d = unsafe_load(p)
    @assert d.data == pointer(a)
    @assert d.length == d.nrows == length(a) == size
end

# ⚠ Extending vectors
function Base.push!(x::K_Vector{t,C,T}, y) where {t,C,T}
    n′ = length(x) + 1
    a = _cast(C, T(y))
    p = kpointer(x)
    p′ = ja(Ref{K_}(p), Ref(a))
    resetvector(x.a, n′, p′+16)
    x
end
function Base.append!(x::K_Vector{t,C,T}, y::K_Vector{t,C,T}) where {t,C,T}
    n′ = length(x) + length(y)
    px = kpointer(x)
    py = kpointer(y)
    p′ = jv(Ref{K_}(px), py)
    resetvector(x.a, n′, p′+16)
    x
end
Base.append!(x::K_Vector{t,C,T}, y) where {t,C,T} =
    append!(x, K_Vector{t,C,T}(y))
function Base.empty!(x::K_Vector{t,C,T}) where {t,C,T}
    p = kpointer(x)
    p′ = ktn(t, 0)
    resetvector(x.a, 0, p′+16)
    r0(p)
    x
end


function get_max_growth(obj::MultistepFieldIntegrator)
    return jcall(obj, "getMaxGrowth", jdouble, ())
end

function get_max_step(obj::AdaptiveStepsizeFieldIntegrator)
    return jcall(obj, "getMaxStep", RealFieldElement, ())
end

function get_min_reduction(obj::MultistepFieldIntegrator)
    return jcall(obj, "getMinReduction", jdouble, ())
end

function get_min_step(obj::AdaptiveStepsizeFieldIntegrator)
    return jcall(obj, "getMinStep", RealFieldElement, ())
end

function get_n_steps(obj::MultistepFieldIntegrator)
    return jcall(obj, "getNSteps", jint, ())
end

function get_safety(obj::MultistepFieldIntegrator)
    return jcall(obj, "getSafety", jdouble, ())
end

function get_starter_integrator(obj::MultistepFieldIntegrator)
    return jcall(obj, "getStarterIntegrator", FieldODEIntegrator, ())
end

function initialize_step(obj::AdaptiveStepsizeFieldIntegrator, arg0::jboolean, arg1::jint, arg2::Vector{RealFieldElement}, arg3::FieldODEStateAndDerivative, arg4::FieldEquationsMapper)
    return jcall(obj, "initializeStep", RealFieldElement, (jboolean, jint, Vector{RealFieldElement}, FieldODEStateAndDerivative, FieldEquationsMapper), arg0, arg1, arg2, arg3, arg4)
end

function set_initial_step_size(obj::AdaptiveStepsizeFieldIntegrator, arg0::RealFieldElement)
    return jcall(obj, "setInitialStepSize", void, (RealFieldElement,), arg0)
end

function set_max_growth(obj::MultistepFieldIntegrator, arg0::jdouble)
    return jcall(obj, "setMaxGrowth", void, (jdouble,), arg0)
end

function set_min_reduction(obj::MultistepFieldIntegrator, arg0::jdouble)
    return jcall(obj, "setMinReduction", void, (jdouble,), arg0)
end

function set_safety(obj::MultistepFieldIntegrator, arg0::jdouble)
    return jcall(obj, "setSafety", void, (jdouble,), arg0)
end

function set_starter_integrator(obj::MultistepFieldIntegrator, arg0::FieldODEIntegrator)
    return jcall(obj, "setStarterIntegrator", void, (FieldODEIntegrator,), arg0)
end

function set_step_size_control(obj::AdaptiveStepsizeFieldIntegrator, arg0::jdouble, arg1::jdouble, arg2::Vector{jdouble}, arg3::Vector{jdouble})
    return jcall(obj, "setStepSizeControl", void, (jdouble, jdouble, Vector{jdouble}, Vector{jdouble}), arg0, arg1, arg2, arg3)
end

function set_step_size_control(obj::AdaptiveStepsizeFieldIntegrator, arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return jcall(obj, "setStepSizeControl", void, (jdouble, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end


# Different versions of variadic function defined in pub_expr.h
#
#   SCIP_RETCODE SCIPexprCreate(BMS_BLKMEM* blkmem, SCIP_EXPR** expr,
#                               SCIP_EXPROP op, ...);
#
# Prototype without extra arguments:
#
#   function SCIPexprCreate(blkmem, expr, op, ...)
#       ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
#             (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, ...),
#             blkmem, expr, op, ...)
#   end
#
#
# TODO: Fix warnings below by creating thin, non-variadic C wrappers for the
#       following calls.

# SCIP_EXPR_VARIDX
function SCIPexprCreate(blkmem, expr, op, varpos::Cint)
    ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
          (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, Cint...),
          blkmem, expr, op, varpos)
end

# SCIP_EXPR_CONST
function SCIPexprCreate(blkmem, expr, op, value::Cdouble)
    ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
          (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, Cdouble...),
          blkmem, expr, op, value)
end

# SCIP_EXPR_SQRT, SCIP_EXPR_EXP, SCIP_EXPR_LOG (unary op)
function SCIPexprCreate(blkmem, expr, op, child::Ptr{SCIP_EXPR})
    ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
          (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, Ptr{SCIP_EXPR}...),
          blkmem, expr, op, child)
end

# SCIP_EXPR_MINUS, SCIP_EXPR_DIV (binary op)
function SCIPexprCreate(blkmem, expr, op, left::Ptr{SCIP_EXPR}, right::Ptr{SCIP_EXPR})
    ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
          (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, Ptr{SCIP_EXPR}...),
          blkmem, expr, op, left, right)
end

# SCIP_EXPR_SUM, SCIP_EXPR_PRODUCT (n-ary op)
function SCIPexprCreate(blkmem, expr, op, nchildren::Cint, children::Vector{Ptr{SCIP_EXPR}})
    # WARNING: In the actual C function, the last two arguments given here are
    # part of the variadic arguments. But since Julia only supports variadic
    # arguments of the same type, we specify the first one (nchildren::Cint)
    # directly, and pretend that only the last one is variadic.
    @assert nchildren == length(children)
    GC.@preserve children begin
        ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
              (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, Cint, Ptr{Ptr{SCIP_EXPR}}...),
              blkmem, expr, op, nchildren, pointer(children))
    end
end

# SCIP_EXPR_REALPOWER (mixed binary op)
function SCIPexprCreate(blkmem, expr, op, base::Ptr{SCIP_EXPR}, exponent::Cdouble)
    # WARNING: In the actual C function, the last two arguments given here are
    # part of the variadic arguments. But since Julia only supports variadic
    # arguments of the same type, we specify the first one (base::Ptr{SCIP_EXPR})
    # directly, and pretend that only the last one is variadic.
    ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
          (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, Ptr{SCIP_EXPR}, Cdouble...),
          blkmem, expr, op, base, exponent)
end

# SCIP_EXPR_INTPOWER (mixed binary op)
function SCIPexprCreate(blkmem, expr, op, base::Ptr{SCIP_EXPR}, exponent::Cint)
    # WARNING: In the actual C function, the last two arguments given here are
    # part of the variadic arguments. But since Julia only supports variadic
    # arguments of the same type, we specify the first one (base::Ptr{SCIP_EXPR})
    # directly, and pretend that only the last one is variadic.
    ccall((:SCIPexprCreate, libscip), SCIP_RETCODE,
          (Ptr{BMS_BLKMEM}, Ptr{Ptr{SCIP_EXPR}}, SCIP_EXPROP, Ptr{SCIP_EXPR}, Cint...),
          blkmem, expr, op, base, exponent)
end


module SegwayRobot

    include("src/dynamics.jl")
    #export segway_rk4, segway_rk4!

    rk4_func  = segway_rk4  # set standard Environment function
    cost_func = segway_cost # set standard Environment cost-function
    #s_init = [0.0, 0.0, 0.0, 0.10] # set standard Environment initial position

    #State vector is [x0, y0, phi, alpha, dalpha, v, dphi, ir, il]
    const s_init  = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    const u_range = [11.1, 11.1]
    export rk4_func, cost_func, s_init, u_range
end # module SegwayRoboter