#=
This file is auto-generated. Do not edit.
=#

#! format: off

"""
    mutable struct ReactivePowerPI <: ReactivePowerControl
        Kp_q::Float64
        Ki_q::Float64
        ωf::Float64
        V_ref::Float64
        Q_ref::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
    end

Parameters of a Proportional-Integral Reactive Power controller for a specified power reference

# Arguments
- `Kp_q::Float64`: Proportional Gain, validation range: `(0, nothing)`
- `Ki_q::Float64`: Integral Gain, validation range: `(0, nothing)`
- `ωf::Float64`: filter frequency cutoff, validation range: `(0, nothing)`
- `V_ref::Float64`: Voltage Set-point, validation range: `(0, nothing)`
- `Q_ref::Float64`: Reactive Power Set-point, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states of the ReactivePowerPI model are:
	σq_oc: Integrator state of the PI Controller,
	q_oc: Measured reactive power of the inverter model
- `n_states::Int`: ReactivePowerPI has two states
"""
mutable struct ReactivePowerPI <: ReactivePowerControl
    "Proportional Gain"
    Kp_q::Float64
    "Integral Gain"
    Ki_q::Float64
    "filter frequency cutoff"
    ωf::Float64
    "Voltage Set-point"
    V_ref::Float64
    "Reactive Power Set-point"
    Q_ref::Float64
    ext::Dict{String, Any}
    "The states of the ReactivePowerPI model are:
	σq_oc: Integrator state of the PI Controller,
	q_oc: Measured reactive power of the inverter model"
    states::Vector{Symbol}
    "ReactivePowerPI has two states"
    n_states::Int
end

function ReactivePowerPI(Kp_q, Ki_q, ωf, V_ref=1.0, Q_ref=1.0, ext=Dict{String, Any}(), )
    ReactivePowerPI(Kp_q, Ki_q, ωf, V_ref, Q_ref, ext, [:σq_oc, :q_oc], 2, )
end

function ReactivePowerPI(; Kp_q, Ki_q, ωf, V_ref=1.0, Q_ref=1.0, ext=Dict{String, Any}(), states=[:σq_oc, :q_oc], n_states=2, )
    ReactivePowerPI(Kp_q, Ki_q, ωf, V_ref, Q_ref, ext, states, n_states, )
end

# Constructor for demo purposes; non-functional.
function ReactivePowerPI(::Nothing)
    ReactivePowerPI(;
        Kp_q=0,
        Ki_q=0,
        ωf=0,
        V_ref=0,
        Q_ref=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`ReactivePowerPI`](@ref) `Kp_q`."""
get_Kp_q(value::ReactivePowerPI) = value.Kp_q
"""Get [`ReactivePowerPI`](@ref) `Ki_q`."""
get_Ki_q(value::ReactivePowerPI) = value.Ki_q
"""Get [`ReactivePowerPI`](@ref) `ωf`."""
get_ωf(value::ReactivePowerPI) = value.ωf
"""Get [`ReactivePowerPI`](@ref) `V_ref`."""
get_V_ref(value::ReactivePowerPI) = value.V_ref
"""Get [`ReactivePowerPI`](@ref) `Q_ref`."""
get_Q_ref(value::ReactivePowerPI) = value.Q_ref
"""Get [`ReactivePowerPI`](@ref) `ext`."""
get_ext(value::ReactivePowerPI) = value.ext
"""Get [`ReactivePowerPI`](@ref) `states`."""
get_states(value::ReactivePowerPI) = value.states
"""Get [`ReactivePowerPI`](@ref) `n_states`."""
get_n_states(value::ReactivePowerPI) = value.n_states

"""Set [`ReactivePowerPI`](@ref) `Kp_q`."""
set_Kp_q!(value::ReactivePowerPI, val) = value.Kp_q = val
"""Set [`ReactivePowerPI`](@ref) `Ki_q`."""
set_Ki_q!(value::ReactivePowerPI, val) = value.Ki_q = val
"""Set [`ReactivePowerPI`](@ref) `ωf`."""
set_ωf!(value::ReactivePowerPI, val) = value.ωf = val
"""Set [`ReactivePowerPI`](@ref) `V_ref`."""
set_V_ref!(value::ReactivePowerPI, val) = value.V_ref = val
"""Set [`ReactivePowerPI`](@ref) `Q_ref`."""
set_Q_ref!(value::ReactivePowerPI, val) = value.Q_ref = val
"""Set [`ReactivePowerPI`](@ref) `ext`."""
set_ext!(value::ReactivePowerPI, val) = value.ext = val


!x=√parse(Int,x)%1
f(a,s="$a",l=length(s))=0∈1:~-l.|>i->s[l-i]<'1'||!s[i+1:end]+!s[1:i]


using ReinforcementLearning
using Flux
using CUDA
using Random, Distributions
using Setfield: @set
using Zygote: ignore
using LinearAlgebra

export DDPGPolicyDemo, update_critic!

#This is a modified version of JuliaReinforcementLearning's ddpg policy. It has been adapted to handle all the various training scenarios used during the project.

mutable struct DDPGPolicyDemo{
    BA<:NeuralNetworkApproximator,
    BC<:NeuralNetworkApproximator,
    TA<:NeuralNetworkApproximator,
    TC<:NeuralNetworkApproximator,
    P,
    R<:AbstractRNG,
} <: AbstractPolicy

    behavior_actor::BA
    behavior_critic::BC
    target_actor::TA
    target_critic::TC
    γ::Float32
    ρ::Float32
    batch_size::Int
    start_steps::Int
    start_policy::P
    update_after::Int
    update_every::Int
    act_limit::Float64
    act_noise::Float64
    step::Int
    rng::R
    # for logging
    actor_loss::Float32
    critic_loss::Float32
end

Flux.functor(x::DDPGPolicyDemo) = (
    ba = x.behavior_actor,
    bc = x.behavior_critic,
    ta = x.target_actor,
    tc = x.target_critic,
),
y -> begin
    x = @set x.behavior_actor = y.ba
    x = @set x.behavior_critic = y.bc
    x = @set x.target_actor = y.ta
    x = @set x.target_critic = y.tc
    x
end

"""
    DDPGPolicy(;kwargs...)
# Keyword arguments
- `behavior_actor`,
- `behavior_critic`,
- `target_actor`,
- `target_critic`,
- `start_policy`,
- `γ = 0.99f0`,
- `ρ = 0.995f0`,
- `batch_size = 32`,
- `start_steps = 10000`,
- `update_after = 1000`,
- `update_every = 50`,
- `act_limit = 1.0`,
- `act_noise = 0.1`,
- `step = 0`,
- `rng = Random.GLOBAL_RNG`,
"""
function DDPGPolicyDemo(;
    behavior_actor,
    behavior_critic,
    target_actor,
    target_critic,
    start_policy,
    γ = 0.99f0,
    ρ = 0.995f0,
    batch_size = 32,
    start_steps = 10000,
    update_after = 1000,
    update_every = 50,
    act_limit = 1.0,
    act_noise = 0.1,
    step = 0,
    rng = Random.GLOBAL_RNG,
)
    copyto!(behavior_actor, target_actor)  # force sync
    copyto!(behavior_critic, target_critic)  # force sync
    DDPGPolicyDemo(
        behavior_actor,
        behavior_critic,
        target_actor,
        target_critic,
        γ,
        ρ,
        batch_size,
        start_steps,
        start_policy,
        update_after,
        update_every,
        act_limit,
        act_noise,
        step,
        rng,
        0.0f0,
        0.0f0,
    )
end

function (p::DDPGPolicyDemo)(env::AbstractEnv; training = true)
    p.step += 1

    if p.step <= p.start_steps
        p.start_policy(env)
    else
        D = device(p.behavior_actor)
        s = state(env)
        s = Flux.unsqueeze(s, ndims(s) + 1)
        a = p.behavior_actor(send_to_device(D, s)) |> vec |> send_to_host
        clamp(a[] + randn(p.rng) * p.act_noise * training, -p.act_limit, p.act_limit)
    end
end

function action(p::DDPGPolicyDemo, s; training = true)
    p.step += 1

    if p.step <= p.start_steps
        p.start_policy(s)
    else
        D = device(p.behavior_actor)
        s = Flux.unsqueeze(s, ndims(s) + 1)
        a = p.behavior_actor(send_to_device(D, s)) |> vec |> send_to_host
        clamp(a[] + randn(p.rng) * p.act_noise * training, -p.act_limit, p.act_limit)
    end
end

function RLBase.update!(p::DDPGPolicyDemo, traj::CircularArraySARTSTrajectory, ::AbstractEnv, ::PreActStage)
    length(traj) > p.update_after || return
    p.step % p.update_every == 0 || return
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(p.batch_size))
    update!(p, batch)
end

function update_critic!(p::DDPGPolicyDemo, traj::CircularArraySARTSTrajectory, demo_traj::CircularArraySARTSTrajectory)
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(ceil(Int, 3*p.batch_size/4)))
    d_inds, d_batch = sample(p.rng, traj, BatchSampler{SARTS}(ceil(Int, p.batch_size/4)))
    _update_critic(p, batch, d_batch=d_batch)
end

function update_critic!(p::DDPGPolicyDemo, traj::CircularArraySARTSTrajectory)
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(ceil(Int, p.batch_size)))
    _update_critic(p, batch)
end

function RLBase.update!(p::DDPGPolicyDemo, traj::CircularArraySARTSTrajectory)
    length(traj) > p.update_after || return
    p.step % p.update_every == 0 || return
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(p.batch_size))
    update!(p, batch)
end

function RLBase.update!(p::DDPGPolicyDemo, traj::CircularArraySARTSTrajectory, ::AbstractEnv, stage::PreTrainStage)
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(p.batch_size))
    update!(p, batch, stage)
end

function RLBase.update!(p::DDPGPolicyDemo, traj::CircularArraySARTSTrajectory, demo_traj::CircularArraySARTSTrajectory)
    length(traj) > p.update_after || return
    p.step % p.update_every == 0 || return
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(ceil(Int, p.batch_size*0.75)))
    demo_inds, demo_batch = sample(p.rng, demo_traj, BatchSampler{SARTS}(ceil(Int, p.batch_size*0.25)))
    update!(p, batch, demo_batch)
end

function RLBase.update!(p::DDPGPolicyDemo, batch::NamedTuple{SARTS})
    s, a, r, t, s′ = send_to_device(device(p), batch)

    A = p.behavior_actor
    C = p.behavior_critic
    Aₜ = p.target_actor
    Cₜ = p.target_critic

    γ = p.γ
    ρ = p.ρ

    a′ = Aₜ(s′)
    qₜ = Cₜ(vcat(s′, a′)) |> vec
    y = r .+ γ .* (1 .- t) .* qₜ
    a = Flux.unsqueeze(a, 1)

    gs1 = gradient(Flux.params(C)) do
        q = C(vcat(s, a)) |> vec
        loss = mean((y .- q) .^ 2)
        ignore() do
            p.critic_loss = loss
        end
        loss
    end

    update!(C, gs1)

    gs2 = gradient(Flux.params(A)) do
        loss = -mean(C(vcat(s, A(s))))
        ignore() do
            p.actor_loss = loss
        end
        loss
    end

    update!(A, gs2)

    # polyak averaging
    for (dest, src) in zip(Flux.params([Aₜ, Cₜ]), Flux.params([A, C]))
        dest .= ρ .* dest .+ (1 - ρ) .* src
    end
end

function RLBase.update!(p::DDPGPolicyDemo, batch::NamedTuple{SARTS}, ::PreTrainStage)
    s, a, r, t, s′ = send_to_device(device(p), batch)

    A = p.behavior_actor
    C = p.behavior_critic
    Aₜ = p.target_actor
    Cₜ = p.target_critic

    γ = p.γ
    ρ = p.ρ

    a′ = Aₜ(s′)
    qₜ = Cₜ(vcat(s′, a′)) |> vec
    y = r .+ γ .* (1 .- t) .* qₜ
    a = Flux.unsqueeze(a, 1)

    gs1 = gradient(Flux.params(C)) do
        q = C(vcat(s, a)) |> vec
        loss = mean((y .- q) .^ 2) + sum(norm, Flux.params(C))/2
        ignore() do
            p.critic_loss = loss
        end
        loss
    end

    update!(C, gs1)

    gs2 = gradient(Flux.params(A)) do
        loss = -mean(C(vcat(s, A(s)))) + 1/2*mean((A(s) .- a) .^ 2) + sum(norm, Flux.params(A))/4
        ignore() do
            p.actor_loss = loss
        end
        loss
    end

    update!(A, gs2)

    # polyak averaging
    for (dest, src) in zip(Flux.params([Aₜ, Cₜ]), Flux.params([A, C]))
        dest .= ρ .* dest .+ (1 - ρ) .* src
    end
end

function RLBase.update!(p::DDPGPolicyDemo, batch::NamedTuple{SARTS}, demo_batch::NamedTuple{SARTS})
    s, a, r, t, s′ = send_to_device(device(p), batch)
    sᵈ, aᵈ, rᵈ, tᵈ, s′ᵈ = send_to_device(device(p), demo_batch)

    A = p.behavior_actor
    C = p.behavior_critic
    Aₜ = p.target_actor
    Cₜ = p.target_critic

    γ = p.γ
    ρ = p.ρ

    a′ = Aₜ(s′)
    qₜ = Cₜ(vcat(s′, a′)) |> vec
    y = r .+ γ .* (1 .- t) .* qₜ
    a = Flux.unsqueeze(a, 1)

    gs1 = gradient(Flux.params(C)) do
        q = C(vcat(s, a)) |> vec
        loss = mean((y .- q) .^ 2) + sum(norm, Flux.params(C))/2
        ignore() do
            p.critic_loss = loss
        end
        loss
    end

    update!(C, gs1)

    aᵈ = Flux.unsqueeze(aᵈ, 1)
    Qcomp = Cₜ(vcat(sᵈ, aᵈ)) .> Cₜ(vcat(sᵈ, A(sᵈ)))
    gs2 = gradient(Flux.params(A)) do
        loss = -mean(C(vcat(s, A(s)))) .+ 1/2*mean(((A(sᵈ) .- aᵈ) .^ 2) .* Qcomp) .+ sum(norm, Flux.params(A)) ./ 4
        ignore() do
            p.actor_loss = loss
        end
        loss
    end

    update!(A, gs2)

    # polyak averaging
    for (dest, src) in zip(Flux.params([Aₜ, Cₜ]), Flux.params([A, C]))
        dest .= ρ .* dest .+ (1 - ρ) .* src
    end
end

# gradient gives error: scalar getindex is disallowed if x is a CuArray with Float64, not FLoat32.
# Solved by calling |> gpu at the end, converting CuArray{Float64,n} to CuArray{Float32,n}
function _update_critic(p::DDPGPolicyDemo, batch::NamedTuple{SARTS}; d_batch = nothing)
    s, a, r, t, s′ = send_to_device(device(p), batch)
    if !isnothing(d_batch)
        sᵈ, aᵈ, rᵈ, tᵈ, s′ᵈ = send_to_device(device(p), d_batch)

        s = hcat(s,sᵈ)
        a = vcat(a,aᵈ)
        r = vcat(r,rᵈ)
        t = vcat(t,tᵈ)
        s′ = hcat(s′,s′ᵈ)
    end

    A = p.behavior_actor
    C = p.behavior_critic
    Aₜ = p.target_actor
    Cₜ = p.target_critic

    γ = p.γ
    ρ = p.ρ

    a′ = Aₜ(s′)
    qₜ = Cₜ(vcat(s′, a′)) |> vec
    y = r .+ γ .* (1 .- t) .* qₜ
    a = Flux.unsqueeze(a, 1)

    gs1 = gradient(Flux.params(C)) do
        q = C(vcat(s, a)) |> vec
        loss = mean((y .- q) .^ 2)
        ignore() do
            p.critic_loss = loss
        end
        loss
    end

    update!(C, gs1)

    for (dest, src) in zip(Flux.params(Cₜ), Flux.params(C))
        dest .= ρ .* dest .+ (1 - ρ) .* src
    end
end

function _pgd(A, C, x, ϵ)
    x0 = copy(x)
    α = 2.5*ϵ/10
    for k in 1:10
        grad = gradient(Flux.params(x)) do
            -mean(C(vcat(x, A(x))))
        end

        x = x + α*sign.(grad[x])
        x = clamp.(x, x0.-ϵ, x0.+ϵ)
        x = send_to_device(device(A), x)
    end
    x
end


function prepareFWIDataFiles(m,Minv::RegularMesh,mref,boundsHigh,boundsLow,
							filenamePrefix::String,omega::Array{Float64,1},waveCoef::Array{Complex128,1},
							pad::Int64,ABLpad::Int64,jump::Int64,offset::Int64=prod(Minv.n+1),workerList = workers(),
							maxBatchSize::Int64=48, Ainv::AbstractSolver = getMUMPSsolver([],0,0,2),useFilesForFields::Bool = false)

########################## m is in Velocity here. ###################################
RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");
writeSrcRcvLocFile(SRCfile,Minv,ABLpad,jump);
writeSrcRcvLocFile(RCVfile,Minv,ABLpad,1);

dataFullFilenamePrefix = string(filenamePrefix,"_freq");
gamma = prepareFWIDataFiles2(m, Minv, filenamePrefix,dataFullFilenamePrefix,omega,waveCoef,pad,ABLpad,offset,workerList,maxBatchSize,Ainv);

file = matopen(string(filenamePrefix,"_PARAM.mat"), "w");
write(file,"boundsLow",boundsLow);
write(file,"boundsHigh",boundsHigh);
write(file,"mref",mref);
write(file,"domain",Minv.domain);
write(file,"n",Minv.n);
write(file,"gamma",gamma);
write(file,"pad",pad);
write(file,"omega",omega);
write(file,"waveCoef",waveCoef);
close(file);
return gamma;
end

function prepareJointTravelTimeAndFWIDataFiles(m,Minv::RegularMesh,mref,boundsHigh,boundsLow,
							filenamePrefix::String,omega::Array{Float64,1},waveCoef::Array{Complex128,1},
							pad::Int64,ABLpad::Int64,jump::Int64,offset::Int64=prod(Minv.n+1),workerList = workers(),
							maxBatchSize::Int64=48, Ainv::AbstractSolver = getMUMPSsolver([],0,0,2),useFilesForFields::Bool = false)

########################## m is in Velocity here. ###################################
RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");
writeSrcRcvLocFile(SRCfile,Minv,ABLpad,jump);
writeSrcRcvLocFile(RCVfile,Minv,ABLpad,1);

dataFullFilenamePrefix = string(filenamePrefix,"_freq");
gamma = prepareFWIDataFiles2(m, Minv, filenamePrefix,dataFullFilenamePrefix,omega,waveCoef,pad,ABLpad,offset,workerList,maxBatchSize,Ainv,useFilesForFields);
dataFullFilenamePrefix = string(filenamePrefix,"_travelTime");
HO = false;
prepareTravelTimeDataFiles(m, Minv, filenamePrefix,dataFullFilenamePrefix,offset,HO,useFilesForFields);

file = matopen(string(filenamePrefix,"_PARAM.mat"), "w");
write(file,"boundsLow",boundsLow);
write(file,"boundsHigh",boundsHigh);
write(file,"mref",mref);
write(file,"MinvOmega",Minv.domain);
write(file,"MinvN",Minv.n);
write(file,"gamma",gamma);
write(file,"pad",pad);
write(file,"omega",omega);
write(file,"waveCoef",waveCoef);
write(file,"HO",HO);
close(file);
end

function prepareFWIDataFiles2(m, Minv::RegularMesh, filenamePrefix::String,dataFullFilenamePrefix::String, omega::Array{Float64,1},
								waveCoef::Array{Complex128,1}, pad::Int64,ABLpad::Int64,offset::Int64,workerList::Array{Int64,1},maxBatchSize::Int64,
								Ainv::AbstractSolver,useFilesForFields::Bool = false)

println("maxOmega*maxKappaSq*h: should be below 0.6");
println(omega[end]*maximum(Minv.h)*(1./minimum(m)));

RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");
srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);

Q = generateSrcRcvProjOperators(Minv.n+1,srcNodeMap);
P = generateSrcRcvProjOperators(Minv.n+1,rcvNodeMap);
Q = Q.*1/(norm(Minv.h)^2);
println("We have ",size(Q,2)," sources");
# compute observed data

ABLamp = getMaximalFrequency(1./(minimum(m).^2),Minv);
gamma = getABL(Minv,true,ones(Int64,Minv.dim)*ABLpad,ABLamp);
attenuation = 0.01*ABLamp;
gamma += attenuation; # adding Attenuation.

println("~~~~~~~ Getting data FWI: ~~~~~~~");

# Solve forward problem (should be relaced in reading data from file)

batch = min(size(Q,2),maxBatchSize);
(pFor,contDiv,SourcesSubInd) = getFWIparam(omega,waveCoef,vec(gamma),Q,P,Minv,Ainv,workerList,batch,useFilesForFields);

(D,pFor) = getData(velocityToSlowSquared(m[:])[1],pFor,ones(length(pFor)),true);

nsrc = size(Q,2);
nrcv = size(P,2);

for k = 1:length(omega)
	I = contDiv[k] : contDiv[k+1] - 1
	Dobsk = Array{Array{Complex128,2}}(length(I));
	for i = 1:length(I)
		Dobsk[i] = fetch(D[I[i]]);
	end
	Dobsk = arrangeRemoteCallDataIntoLocalData(Dobsk);
	omRound = string(round((omega[k]/(2*pi))*100.0)/100.0);
	Wd_k = (1./(abs.(real(Dobsk))+0.1*mean(abs.(Dobsk)))) + 1im*(1./(abs.(imag(Dobsk))+0.1*mean(abs.(Dobsk))));
	# Wd_k = (1./(0.0*abs(real(Dobsk))+1.0*mean(abs(Dobsk)))) + 1im*(1./(0.0*abs(imag(Dobsk))+1.0*mean(abs(Dobsk))));
	Wd_k = limitDataToOffset(Wd_k,srcNodeMap,rcvNodeMap,offset);
	filename = string(dataFullFilenamePrefix,omRound,".dat");
	writeDataFile(filename,Dobsk,Wd_k,srcNodeMap,rcvNodeMap);
end
return gamma;
end


export PrioritizedDQNLearner

using Setfield: @set
using Random: AbstractRNG, GLOBAL_RNG
import Functors
using LinearAlgebra: dot

Base.@kwdef mutable struct PrioritizedDQNLearner{A<:Approximator{<:TwinNetwork}} <: AbstractLearner
    approximator::A
    loss_func::Any  # !!! here the loss func must return the loss before reducing over the batch dimension
    n::Int = 1
    γ::Float32 = 0.99f0
    β_priority::Float32 = 0.5f0
    is_enable_double_DQN::Bool = true
    rng::AbstractRNG = GLOBAL_RNG
    # for logging
    loss::Float32 = 0.0f0
end

(L::PrioritizedDQNLearner)(s::AbstractArray) = L.approximator(s)

Functors.functor(x::PrioritizedDQNLearner) = (; approximator=x.approximator), y -> @set x.approximator = y.approximator

function RLBase.optimise!(
    learner::PrioritizedDQNLearner,
    batch::Union{
        NamedTuple{(:key, :priority, SS′ART...)},
        NamedTuple{(:key, :priority, SS′L′ART...)}
    }
)
    A = learner.approximator
    Q = A.model.source
    Qₜ = A.model.target
    γ = learner.γ
    β = learner.β_priority
    loss_func = learner.loss_func
    n = learner.n

    s, s′, a, r, t = map(x -> batch[x], SS′ART)
    batch_size = length(a)
    a = CartesianIndex.(a, 1:batch_size)
    k, p = batch.key, batch.priority
    p′ = similar(p)

    w = 1.0f0 ./ ((p .+ 1.0f-10) .^ β)
    w ./= maximum(w)

    q′ = learner.is_enable_double_DQN ? Q(s′) : Qₜ(s′)

    if haskey(batch, :next_legal_actions_mask)
        q′ .+= ifelse.(batch[:next_legal_actions_mask], 0.0f0, typemin(Float32))
    end

    q′ₐ = learner.is_enable_double_DQN ? Qₜ(s′)[dropdims(argmax(q′, dims=1), dims=1)] : dropdims(maximum(q′; dims=1), dims=1)

    G = r .+ γ^n .* (1 .- t) .* q′ₐ

    gs = gradient(params(A)) do
        qₐ = Q(s)[a]
        batch_losses = loss_func(G, qₐ)
        loss = dot(vec(w), vec(batch_losses)) * 1 // batch_size
        ignore() do
            p′ .= vec((batch_losses .+ 1.0f-10) .^ β)
            learner.loss = loss
        end
        loss
    end

    optimise!(A, gs)
    k => p′
end

function RLBase.optimise!(policy::QBasedPolicy{<:PrioritizedDQNLearner}, trajectory::Trajectory)
    for batch in trajectory
        k, p = optimise!(policy, batch) |> send_to_host
        trajectory[:priority, k] = p
    end
end


using GLVisualize, GLAbstraction, FileIO, GeometryTypes, Reactive, GLWindow, Colors

if !isdefined(:runtests)
    window = glscreen()
    timesignal = bounce(linspace(0f0,1f0,360))
end

description = """
This example shows how to set up a custom camera,
animate it and display it in a second screen. The first screen
allows you to look at the camera in the scene.
"""

"""
functions to halve some rectangle
"""
xhalf(r)  = SimpleRectangle(r.x, r.y, r.w÷2, r.h)
xhalf2(r) = SimpleRectangle(r.w÷2, r.y, r.w÷2, r.h)

"""
Makes a spiral
"""
function spiral(i, start_radius, offset)
	Point3f0(sin(i), cos(i), i/10f0) * (start_radius + ((i/2pi)*offset))
end
# 2D particles
curve_data(i, N) = Point3f0[spiral(i+x/20f0, 1, (i/20)+1) for x=1:N]
# create a spiraling camera path
const camera_path  = curve_data(1f0, 360)

area_a, area_b = x_partition(window.area, 50)
# create first screen for the camera
camera_screen = Screen(
	window, name=:camera_screen,
	area = area_a
)
# create second screen to _view the scene
scene_screen = Screen(
	window, name=:scene_screen,
	area = area_b
)

# create an camera eyeposition signal, which follows the path
eyeposition = map(timesignal) do t
    len = length(camera_path)
    index = round(Int, (t*(len-1))+1) # mod1, just to be on the save side
    Vec3f0(camera_path[index])
end

# create the camera lookat and up vector
lookatposition = Signal(Vec3f0(0))
upvector = Signal(Vec3f0(0,0,1))

# create a camera from these
cam = PerspectiveCamera(camera_screen.area, eyeposition, lookatposition, upvector)

"""
Simple visualization of a camera
"""
function showcam(cam::PerspectiveCamera)
    lookvec, posvec, upvec = map(f-> getfield(cam, f), (:lookat, :eyeposition, :up))
    positions = map((a,b) -> Point3f0[a,b], lookvec, posvec)
    lines = map(lookvec, posvec, upvec) do l,p,u
        dir = p-l
        right = normalize(cross(dir,u))
        Point3f0[
            l,p,
            p, p+u,
            p, p+right
        ]
    end
    colors = RGBA{Float32}[
        RGBA{Float32}(1,0,0,1),
        RGBA{Float32}(1,0,0,1),

        RGBA{Float32}(0,1,0,1),
        RGBA{Float32}(0,1,0,1),

        RGBA{Float32}(0,0,1,1),
        RGBA{Float32}(0,0,1,1),
    ]
    poses = visualize((Sphere(Point3f0(0), 0.05f0), positions))
    lines = visualize(lines, :linesegment, color=colors)
    Context(poses, lines)
end

# add the camera to the camera screen as the perspective camera
camera_screen.cameras[:perspective] = cam

# something to look at
cat = visualize(GLNormalMesh(loadasset("cat.obj")))

# visualize the camera path
camera_points = visualize(
    (Circle(Point2f0(0), 0.03f0), camera_path),
    color=RGBA{Float32}((Vec3f0(0,206,209)/256)..., 1f0), billboard=true
)

camera_path_line = visualize(camera_path, :lines)

# _view everything on the appropriate screen.
# we need to copy the cat, because _view inserts the camera into the
# actual render object. this is sub optimal and will get changed!
# Note, that this is a shallow copy, so the actual data won't be copied,
# just the data structure that holds the camera
_view(copy(cat), camera_screen, camera=:perspective)
_view(copy(cat), scene_screen, camera=:perspective)
_view(showcam(cam), scene_screen, camera=:perspective)
_view(camera_points, scene_screen, camera=:perspective)
_view(camera_path_line, scene_screen, camera=:perspective)

if !isdefined(:runtests)
	renderloop(window)
end


#   This file is part of UnitSystems.jl
#   It is licensed under the MIT license
#   UnitSystems Copyright (C) 2020 Michael Reed
#       _           _                         _
#      | |         | |                       | |
#   ___| |__   __ _| | ___ __ __ ___   ____ _| | __ _
#  / __| '_ \ / _` | |/ / '__/ _` \ \ / / _` | |/ _` |
# | (__| | | | (_| |   <| | | (_| |\ V / (_| | | (_| |
#  \___|_| |_|\__,_|_|\_\_|  \__,_| \_/ \__,_|_|\__,_|
#
#   https://github.com/chakravala
#   https://crucialflow.com

"""
    kilograms(m::Real) = $(lb*g₀/ft)m

Converts mass `m` from slugs to kilogram (kg).
"""
@pure kilograms(m::Number,U::UnitSystem=English) = mass(m,Metric,U)

"""
    slugs(m::Real) = $(ft/lb/g₀)m

Converts mass `m` from kilograms to slugs (slug).
"""
@pure slugs(m::Number,U::UnitSystem=Metric) = mass(m,English,U)

"""
    feet(d) = $(1/ft)d

Converts distance `d` from meters to feet (ft).
"""
@pure feet(d,U::UnitSystem=Metric) = length(d,English,U)

"""
    meters(d) = $(ft)d

Converts distance `d` from feet to meters (m).
"""
@pure meters(d,U::UnitSystem=English) = length(d,Metric,U)

# special exact cases

# IAU to SI
@pure length(U::UnitSystem{kB,ħ,𝘤} where {kB,ħ},S::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ}) = length(U,S,1/au)
@pure time(U::UnitSystem{kB,ħ,𝘤} where {kB,ħ},S::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ}) = time(U,S,1/DAY)

# SI to IAU
@pure length(U::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ},S::UnitSystem{kB,ħ,𝘤} where {kB,ħ}) = length(U,S,au)
@pure time(U::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ},S::UnitSystem{kB,ħ,𝘤} where {kB,ħ}) = time(U,S,DAY)

# IAU to CGS
@pure length(::UnitSystem{kB,ħ,100𝘤} where {kB,ħ},S::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ}) = length(U,S,1/au)
@pure time(::UnitSystem{kB,ħ,100𝘤} where {kB,ħ},S::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ}) = time(U,S,1/DAY)

# CGS to IAU
@pure length(U::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ},S::UnitSystem{kB,ħ,100𝘤} where {kB,ħ}) = length(U,S,au)
@pure time(U::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ},S::UnitSystem{kB,ħ,100𝘤} where {kB,ħ}) = time(U,S,DAY)

# IAU to English
@pure length(U::UnitSystem{kB,ħ,𝘤/ft} where {kB,ħ},S::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ}) = length(U,S,ft/au)
@pure time(U::UnitSystem{kB,ħ,𝘤/ft} where {kB,ħ},S::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ}) = time(U,S,1/DAY)

# English to IAU
@pure length(U::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ},S::UnitSystem{kB,ħ,𝘤/ft} where {kB,ħ}) = length(U,S,au/ft)
@pure time(U::UnitSystem{kB,ħ,DAY*𝘤/au} where {kB,ħ},S::UnitSystem{kB,ħ,𝘤/ft} where {kB,ħ}) = time(U,S,DAY)

# base unit conversions

@pure length(U::UnitSystem,S::UnitSystem,l=1) = isquantity(U,S) ? evaldim(length)(U,S) : unit((turn(S)/turn(U))*(planckreduced(S)*electronmass(U)*lightspeed(U)*gravity(S))/(planckreduced(U)*electronmass(S)*lightspeed(S)*gravity(U)),l)
@doc """
$(convertext(:length,"planck(U,S)/mass(U,S)/speed(U,S)"))

Extent of one-dimensional shape or `length` (m), unit conversion factor.

```Julia
julia> length(CGS,Metric) # m⋅cm⁻¹
$(length(CGS,Metric))

julia> length(IAU,Metric) # m⋅au⁻¹
$(length(IAU,Metric))

julia> length(English,Metric) # m⋅ft⁻¹
$(length(English,Metric))

julia> length(Survey,English) # ft⋅ftUS⁻¹
$(length(Survey,English))

julia> length(PlanckGauss,Metric) # m⋅ℓP⁻¹
$(length(PlanckGauss,Metric))
```
""" length, ft, ftUS, L

@pure time(U::UnitSystem,S::UnitSystem,t=1) = isquantity(U,S) ? evaldim(time)(U,S) : unit(length(U,S)/lightspeed(U,S),1)
@doc """
$(convertext(:time,"length(U,S)/speed(U,S)"))

Dimension along which events are ordered or `time` (s), unit conversion factor.

```Julia
julia> time(IAU,Metric) # s⋅day⁻¹
$(time(IAU,Metric))

julia> time(PlanckGauss,Metric) # s⋅tP⁻¹
$(time(PlanckGauss,Metric))
```
""" time(::UnitSystem,::UnitSystem), T

# spacetime

@pure area(U::UnitSystem,S::UnitSystem) = unit(length(U,S)^2)
@pure volume(U::UnitSystem,S::UnitSystem) = unit(length(U,S)^3)
@pure wavenumber(U::UnitSystem,S::UnitSystem) = unit(length(S,U))
@pure angularwavenumber(U::UnitSystem,S::UnitSystem) = unit(angle(U,S)*length(S,U))
@pure fuelefficiency(U::UnitSystem,S::UnitSystem) = area(S,U)
@pure numberdensity(U::UnitSystem,S::UnitSystem) = volume(S,U)
@pure frequency(U::UnitSystem,S::UnitSystem) = time(S,U)
@pure angularfrequency(U::UnitSystem,S::UnitSystem) = unit(angle(U,S)*time(S,U))
@pure frequencydrift(U::UnitSystem,S::UnitSystem) = unit(time(S,U)^2)
@pure speed(U::UnitSystem,S::UnitSystem) = lightspeed(U,S)
@pure acceleration(U::UnitSystem,S::UnitSystem) = unit(speed(U,S)/time(U,S))
@pure jerk(U::UnitSystem,S::UnitSystem) = unit(speed(U,S)/time(U,S)^2)
@pure snap(U::UnitSystem,S::UnitSystem) = unit(speed(U,S)/time(U,S)^3)
@pure crackle(U::UnitSystem,S::UnitSystem) = unit(speed(U,S)/time(U,S)^4)
@pure pop(U::UnitSystem,S::UnitSystem) = unit(speed(U,S)/time(U,S)^5)
@pure volumeflow(U::UnitSystem,S::UnitSystem) = unit(area(U,S)*speed(U,S))
@pure specificenergy(U::UnitSystem,S::UnitSystem) = unit(speed(U,S)^2/gravity(U,S))

# kinematic

@pure inertia(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)/gravity(U,S))
@pure energy(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)*specificenergy(U,S))
@pure power(U::UnitSystem,S::UnitSystem) = unit(energy(U,S)/time(U,S))
@pure force(U::UnitSystem,S::UnitSystem) = unit(inertia(U,S)*acceleration(U,S))
@pure specificforce(U::UnitSystem,S::UnitSystem) = unit(acceleration(U,S)/gravity(U,S))
@pure gravityforce(U::UnitSystem,S::UnitSystem) = unit(gravity(U,S))
@pure pressure(U::UnitSystem,S::UnitSystem) = unit(force(U,S)/area(U,S))

# mechanical

@pure impulse(U::UnitSystem,S::UnitSystem) = unit(force(U,S)*time(U,S))
@pure momentum(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)*speed(U,S))
@pure angularmomentum(U::UnitSystem,S::UnitSystem) = unit(momentum(U,S)*length(U,S)*angle(U,S))
@pure yank(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)*jerk(U,S))
@pure areadensity(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)/area(U,S))
@pure density(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)/volume(U,S))
@pure specificweight(U::UnitSystem,S::UnitSystem) = unit(force(U,S)/volume(U,S))
@pure specificvolume(U::UnitSystem,S::UnitSystem) = unit(volume(U,S)/mass(U,S))
@pure action(U::UnitSystem,S::UnitSystem) = unit(energy(U,S)*time(U,S))
#@pure stiffness(U::UnitSystem,S::UnitSystem) = unit(force(U,S)/length(U,S))
@pure intensity(U::UnitSystem,S::UnitSystem) = unit(power(U,S)/area(U,S))
@pure diffusivity(U::UnitSystem,S::UnitSystem) = unit(speed(U,S)*length(U,S))
@pure viscosity(U::UnitSystem,S::UnitSystem) = unit(force(U,S)/speed(U,S)/length(U,S))
@pure lineardensity(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)/length(U,S))
@pure massflow(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)/time(U,S))
@pure spectralflux(U::UnitSystem,S::UnitSystem) = unit(power(U,S)/length(U,S))
@pure powerdensity(U::UnitSystem,S::UnitSystem) = unit(power(U,S)/volume(U,S))
@pure compressibility(U::UnitSystem,S::UnitSystem) = pressure(S,U)
@pure fluence(U::UnitSystem,S::UnitSystem) = unit(energy(U,S)/area(U,S))
@pure rotationalinertia(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)*area(U,S))

# acoustic

@pure soundexposure(U::UnitSystem,S::UnitSystem) = unit(time(U,S)*pressure(U,S)^2)
@pure specificimpedance(U::UnitSystem,S::UnitSystem) = unit(pressure(U,S)/speed(U,S))
@pure impedance(U::UnitSystem,S::UnitSystem) = unit(specificimpedance(U,S)/area(U,S))
@pure admittance(U::UnitSystem,S::UnitSystem) = unit(area(U,S)/specificimpedance(U,S))
@pure compliance(U::UnitSystem,S::UnitSystem) = unit(time(U,S)^2/mass(U,S))
@pure inertance(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)/length(U,S)^4)

include("kinematicdocs.jl")


module NonhydrostaticModels

export NonhydrostaticModel

using DocStringExtensions

using KernelAbstractions: @index, @kernel, Event, MultiEvent
using KernelAbstractions.Extras.LoopInfo: @unroll

using Oceananigans.Utils: launch!
using Oceananigans.Grids
using Oceananigans.Solvers
using Oceananigans.Distributed: MultiArch, DistributedFFTBasedPoissonSolver, reconstruct_global_grid
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid

import Oceananigans: fields, prognostic_fields

function PressureSolver(arch::MultiArch, local_grid::RegRectilinearGrid)
    global_grid = reconstruct_global_grid(local_grid)
    return DistributedFFTBasedPoissonSolver(global_grid, local_grid)
end

PressureSolver(arch, grid::RegRectilinearGrid)  = FFTBasedPoissonSolver(grid)
PressureSolver(arch, grid::HRegRectilinearGrid) = FourierTridiagonalPoissonSolver(grid)

# *Evil grin*
PressureSolver(arch, ibg::ImmersedBoundaryGrid) = PressureSolver(arch, ibg.underlying_grid)

#####
##### NonhydrostaticModel definition
#####

include("nonhydrostatic_model.jl")
include("pressure_field.jl")
include("show_nonhydrostatic_model.jl")
include("set_nonhydrostatic_model.jl")

#####
##### Time-stepping NonhydrostaticModels
#####

"""
    fields(model::NonhydrostaticModel)

Returns a flattened `NamedTuple` of the fields in `model.velocities` and `model.tracers`.
"""
fields(model::NonhydrostaticModel) = merge(model.velocities, model.tracers)
prognostic_fields(model::NonhydrostaticModel) = fields(model)

include("solve_for_pressure.jl")
include("update_hydrostatic_pressure.jl")
include("update_nonhydrostatic_model_state.jl")
include("pressure_correction.jl")
include("nonhydrostatic_tendency_kernel_functions.jl")
include("calculate_nonhydrostatic_tendencies.jl")
include("correct_nonhydrostatic_immersed_tendencies.jl")

end # module


#=

include("./read_dsmacc.jl")

using blink
w = Blink.AtomShell.Window()
loadurl(w, string("file://" ,pwd(),"/scatter3d/index.html"))
@js w console.log("You have succesfully loaded index.html from Julia")

=#

len = length
normalise(x) = x/norm(x)
run() = eval(parse("""include("./read_dsmacc.jl")"""))
evlstr(x) = eval(parse(x))

using PyCall,DataFrames,RCall
unshift!(PyVector(pyimport("sys")["path"]), "")
@pyimport ncdata
filename = "dsmacc_out.nc"
data = ncdata.get(filename)

specs = names!(DataFrame(data["spec"]),[Symbol(i)for i in data["sc"]])
rates = names!(DataFrame(data["rate"]),[Symbol(i)for i in data["rc"]])

rates = rates[:,[Symbol(i) for i in filter(x->!ismatch(r"EMISS|DUMMY",x),data["rc"])]]
specs = specs[:,[Symbol(i) for i in filter(x->!ismatch(r"EMISS|DUMMY",x),data["sc"])]]

rates = rates[7:len(rates)]
rates = rates[:,[Symbol(i) for i in filter(x->sum(rates[x])>0,names(rates))]]

# reading in data and cleaning
joint = "\n"*(join(names(rates),"\n"))*"\n"

products= [split(i[4:len(i)],"+") for i in matchall(r"-->([A-z0-9+]*)",joint)]
reactants = [split(i[2:len(i)-1],"+") for i in matchall(r"\n([A-z0-9+]{1,60})([-->]{0,1})",joint)]

node_set = Set([ i[2:len(i)-1] for i in matchall(r"[\n+>-]{1}[A-z0-9]+[\n+>-]{1}",joint) ])
nodes = [i for i in node_set]
dict = Dict( [[i,x] for (x,i) in enumerate(nodes)])

flux = []
edges = []
for i in 1:length(reactants)
  rcol=[]
  for j in reactants[i]
    coeff = match( r"([\d\s\.]*)(\D[\d\D]*)", j)
    dummy = specs[Symbol(coeff[2])]
    try
      push!(rcol,parse(Float32,coeff[1]))
    catch
      push!(rcol,1.0f0)
    end

    prod = 1
    for k in rcol
      prod*=k
    end

    push!(flux, prod*rates[i])
    # make edges array
    for l in products[i]
        push!(edges,[j,l,i])
    end

  end
end

##################FOR EACH TIMESTEP##################
#using Blink
#w = Blink.Window()
sleep(1)

t=144
#using Plots
#unicodeplots();

#f4(x)=log(e,x),

fns=[f(x)=x,f1(x)=x*x, f2(x)=x*x*x, f3(x)=e^x, f5(x)=x^(1/2),f6(x)=x^(1/3),f7(x)=x^(1/e),f8(x)=abs(log(e,x)),f9(x)=abs(log(2,x))]
dcrpt = ["x","x^2","x^3", "exp^x", "x^{1/2}","x^{1/3}", "x^{1/exp}","|ln(x)|","|log_2(x)|"]

forward = []
for fn in 1:len(fns)
  bindata = []
  for t in 1:7:len(flux[1])

      links = filter(i -> flux[i][t]>0 , 1:len(flux))
      tflux = [log10(flux[i][t]) for i in links]
      weight = (tflux+abs(minimum(tflux)))
      weight = (weight/maximum(weight))*.98 +0.01

      weight = [ if q >0 fns[fn](q) else 0 end for q in weight]

      a = fit(Histogram, weight+0.1, nbins = 10, closed = :left)

      println(a.weights,[i for i in a.edges[1]])

      push!(bindata,a.weights)
  end
  push!(forward,bindata)
end

open("data.json", "w") do f
        write(f, "bardata="*JSON.json(forward)*"; timestamp ="*JSON.json(specs[1][2:len(specs[1])])*";description = "*JSON.json(dcrpt))
     end

#R"bin<- data.frame($(bindata));colnames(bin) <- NULL"

using Blink
w = Blink.AtomShell.Window()
loadurl(w, string("file://" ,pwd(),"/index.html"))
@js w console.log("You have succesfully loaded index.html from Julia")

#http://dataframesjl.readthedocs.io/en/latest/getting_started.html
#https://gist.github.com/gizmaa/7214002


__precompile__()

module BARON

if haskey(ENV, "BARON_EXEC")
    const baron_exec = ENV["BARON_EXEC"]
else
    @warn("Unable to locate BARON executable. Make sure the solver has been separately downloaded, and that you properly set the BARON_EXEC environment variable.")
end

mutable struct VariableInfo
    lower_bound::Union{Float64, Nothing}
    upper_bound::Union{Float64, Nothing}
    category::Symbol
    start::Union{Float64, Nothing}
    name::Union{String, Nothing}
end
VariableInfo() = VariableInfo(nothing, nothing, :Cont, nothing, nothing)

mutable struct ConstraintInfo
    expression::Expr
    lower_bound::Union{Float64, Nothing}
    upper_bound::Union{Float64, Nothing}
    name::Union{String, Nothing}
end

function ConstraintInfo()
    ConstraintInfo(:(), nothing, nothing, nothing)
end

mutable struct ObjectiveInfo
    expression::Union{Expr, Number}
    sense::Symbol
end
ObjectiveInfo() = ObjectiveInfo(0, :Feasibility)

@enum BaronSolverStatus begin
    NORMAL_COMPLETION = 1
    INSUFFICIENT_MEMORY_FOR_NODES = 2
    ITERATION_LIMIT = 3
    TIME_LIMIT = 4
    NUMERICAL_SENSITIVITY = 5
    USER_INTERRUPTION = 6
    INSUFFICIENT_MEMORY_FOR_SETUP = 7
    RESERVED = 8
    TERMINATED_BY_BARON = 9
    SYNTAX_ERROR = 10
    LICENSING_ERROR = 11
    USER_HEURISTIC_TERMINATION = 12
end

@enum BaronModelStatus begin
    OPTIMAL = 1
    INFEASIBLE = 2
    UNBOUNDED = 3
    INTERMEDIATE_FEASIBLE = 4
    UNKNOWN = 5
end

mutable struct SolutionStatus
    feasible_point::Union{Nothing, Vector{Float64}}
    objective_value::Float64
    dual_bound::Float64
    wall_time::Float64
    solver_status::BaronSolverStatus
    model_status::BaronModelStatus

    SolutionStatus() = new(nothing)
end

mutable struct BaronModel
    options::Dict{Symbol, Any}

    variable_info::Vector{VariableInfo}
    constraint_info::Vector{ConstraintInfo}
    objective_info::ObjectiveInfo

    temp_dir_name::String
    problem_file_name::String
    times_file_name::String
    summary_file_name::String
    result_file_name::String

    solution_info::Union{Nothing, SolutionStatus}

    function BaronModel(; kwargs...)
        options = Dict{Symbol, Any}(key => val for (key,val) in kwargs)
        model = new()
        model.options = options
        model.variable_info = VariableInfo[]
        model.constraint_info = ConstraintInfo[]
        model.objective_info = ObjectiveInfo()
        temp_dir = mktempdir()
        model.temp_dir_name = temp_dir
        model.problem_file_name = get!(options, :ProName, joinpath(temp_dir, "baron_problem.bar"))
        model.times_file_name = get!(options, :TimName, joinpath(temp_dir, "tim.lst"))
        model.summary_file_name = get!(options, :SumName, joinpath(temp_dir, "sum.lst"))
        model.result_file_name = get!(options, :ResName, joinpath(temp_dir, "res.lst"))
        model.solution_info = nothing
        return model
    end
end

include("util.jl")
include("MPB_wrapper.jl")
include("MOI_wrapper.jl")

end  # module


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

using FunctionWrappers: FunctionWrapper

struct GenericProductDensityFunction{T<:Real,P<:Real}
    log_terms::Vector{FunctionWrapper{T,Tuple{P}}}
    single_exec_compat::ExecCapabilities
end

export GenericProductDensityFunction

function unsafe_density_logval(
    density::GenericProductDensityFunction{T,P},
    params::AbstractVector{<:Real},
    exec_context::ExecContext
) where {T,P}
    # TODO: Use exec_context and density.exec_capabilities
    sum((log_term(convert(P, T)) for (log_term, p) in (density.log_terms, params)))
end

exec_capabilities(::typeof(unsafe_density_logval), density::GenericProductDensityFunction, params::AbstractVector{<:Real}) =
    density.single_exec_compat # Change when implementation of density_logval for GenericProductDensityFunction becomes multithreaded.

# For priors, DensityFunction with field sampler_f
# that stores a FunctionWrapper around a sampler()::Sampleable{Multivariate,Continuous}


module TkIntcode

export intcode, start

using OffsetArrays

mutable struct Intcode
    code::AbstractArray{Int}
    ip::Int
    input
    output
end

# Handling jump instructions
abstract type Jump end
struct JumpNext <: Jump end
struct JumpFar <: Jump
    location
end

# Default code parser
function read_code(s::AbstractString)
    code = parse.(Int, split(s, ","))
    return OffsetArray(code, 0:length(code)-1)
end

# Intcode constructor
intcode(s::AbstractString, input, output) = Intcode(read_code(s), 0, input, output)

# I/O support
read_input(c::Channel) = take!(c)
write_output(c::Channel, data) = put!(c, data)

read_input(io::IO) = parse(Int,readline(io))
write_output(io::IO, data) = println(io, data)

read_input(f::Function) = f()
write_output(f::Function, data) = f(data)

# Start program
function start(ic::Intcode)
    while true
        i = ic.ip
        instruction = ic.code[i]
        op = instruction % 100      # rightmost two digits
        opmode = instruction ÷ 100  # shift right by two digits
        # println("ic $(ic.id): executing op=$op opmode=$opmTode i=$i")
        if op in (1,2,7,8)
            execute!(Val(op), Val(opmode), ic.code, ic.code[i+1], ic.code[i+2], ic.code[i+3])
            ic.ip += 4
        elseif op == 3
            execute!(Val(op), Val(opmode), ic, ic.input, ic.code, ic.code[i+1])
            ic.ip += 2
        elseif op == 4
            execute!(Val(op), Val(opmode), ic, ic.output, ic.code, ic.code[i+1])
            ic.ip += 2
        elseif op in (5,6)
            j = execute!(Val(op), Val(opmode), ic.code, ic.code[i+1], ic.code[i+2])
            if j isa JumpFar
                ic.ip = j.location
            else
                ic.ip += 3
            end
        elseif op == 99
            break
        else
            println("Unknown op code $op")
        end
    end
end

execute!(::Val{1}, ::Val{000}, code, p1, p2, p3) = code[p3] = code[p1] + code[p2]
execute!(::Val{1}, ::Val{001}, code, p1, p2, p3) = code[p3] = p1       + code[p2]
execute!(::Val{1}, ::Val{010}, code, p1, p2, p3) = code[p3] = code[p1] + p2
execute!(::Val{1}, ::Val{011}, code, p1, p2, p3) = code[p3] = p1       + p2

execute!(::Val{2}, ::Val{000}, code, p1, p2, p3) = code[p3] = code[p1] * code[p2]
execute!(::Val{2}, ::Val{001}, code, p1, p2, p3) = code[p3] = p1       * code[p2]
execute!(::Val{2}, ::Val{010}, code, p1, p2, p3) = code[p3] = code[p1] * p2
execute!(::Val{2}, ::Val{011}, code, p1, p2, p3) = code[p3] = p1       * p2

function execute!(::Val{3}, ::Val{000}, ic, channel, code, p1)
    code[p1] = read_input(channel)
end

function execute!(::Val{4}, ::Val{000}, ic, channel, code, p1)
    write_output(channel, code[p1])
end

function execute!(::Val{4}, ::Val{001}, ic, channel, code, p1)
    write_output(channel, p1)
end

execute!(::Val{5}, ::Val{000}, code, p1, p2) = code[p1] != 0 ? JumpFar(code[p2]) : JumpNext()
execute!(::Val{5}, ::Val{001}, code, p1, p2) = p1 != 0       ? JumpFar(code[p2]) : JumpNext()
execute!(::Val{5}, ::Val{010}, code, p1, p2) = code[p1] != 0 ? JumpFar(p2)       : JumpNext()
execute!(::Val{5}, ::Val{011}, code, p1, p2) = p1 != 0 ?       JumpFar(p2)       : JumpNext()

execute!(::Val{6}, ::Val{000}, code, p1, p2) = code[p1] == 0 ? JumpFar(code[p2]) : JumpNext()
execute!(::Val{6}, ::Val{001}, code, p1, p2) = p1 == 0       ? JumpFar(code[p2]) : JumpNext()
execute!(::Val{6}, ::Val{010}, code, p1, p2) = code[p1] == 0 ? JumpFar(p2)       : JumpNext()
execute!(::Val{6}, ::Val{011}, code, p1, p2) = p1 == 0       ? JumpFar(p2)       : JumpNext()

execute!(::Val{7}, ::Val{000}, code, p1, p2, p3) = code[p3] = code[p1] < code[p2] ? 1 : 0
execute!(::Val{7}, ::Val{001}, code, p1, p2, p3) = code[p3] = p1       < code[p2] ? 1 : 0
execute!(::Val{7}, ::Val{010}, code, p1, p2, p3) = code[p3] = code[p1] < p2       ? 1 : 0
execute!(::Val{7}, ::Val{011}, code, p1, p2, p3) = code[p3] = p1       < p2       ? 1 : 0

execute!(::Val{8}, ::Val{000}, code, p1, p2, p3) = code[p3] = code[p1] == code[p2] ? 1 : 0
execute!(::Val{8}, ::Val{001}, code, p1, p2, p3) = code[p3] = p1       == code[p2] ? 1 : 0
execute!(::Val{8}, ::Val{010}, code, p1, p2, p3) = code[p3] = code[p1] == p2       ? 1 : 0
execute!(::Val{8}, ::Val{011}, code, p1, p2, p3) = code[p3] = p1       == p2       ? 1 : 0

end # module


# https://pythonhosted.org/rasterstats/manual.html#raster-data-sources

println("Welcome... computing land use by region and year [1992-2019]...")

cd(@__DIR__)
using Pkg
Pkg.activate(".")
Pkg.instantiate()

#Pkg.add("PyCall")
#Pkg.add("DataFrames")
#Pkg.add("CSV")
#Pkg.add("Tables")
using PyCall, Tables, DataFrames, CSV

vectorBoundariesFile     = "./adminAreas/byLevels/gadm36_1.shp"
rasterDataFile_1992_2015 = "./luse/ESACCI-LC-L4-LCCS-Map-300m-P1Y-1992_2015-v2.0.7.tif"

sqKmPerCell          = (300*300)/(1000*1000)

const rs  = pyimport("rasterstats")
const gpd = pyimport("geopandas")

println("Modules loaded, reading the regions definition shapefile..")

regdf     = gpd.read_file(vectorBoundariesFile)
gid_0     = convert(Vector{String},regdf.GID_0)
name_0    = convert(Vector{String},regdf.NAME_0)
gid_1     = convert(Vector{String},regdf.GID_1)
name_1    = convert(Vector{String},regdf.NAME_1)
nRegions  = length(gid_0)

# see https://datastore.copernicus-climate.eu/documents/satellite-land-cover/D3.3.12-v1.3_PUGS_ICDR_LC_v2.1.x_PRODUCTS_v1.3.pdf
# Attention that the classification is hierarchical, e.g. "60" includes also "61" and "62"
luseIds = [ 10,  11,  12,  20,  30,  40,  50,  60,  61,  62,  70,  71,  72,
            80,  81,  82,  90, 100, 110, 120, 121, 122, 130, 140, 150, 151,
           152, 153, 160, 170, 180, 190, 200, 201, 202, 210, 220]
headersForExport = vcat("gid_0","name_0","gid_1","name_1",string.(luseIds)...)

years = 1992:2019
years = 2015:2019

for year in years
    println("Starting processing year $(year)...")
    if year <= 2015
        band = year-1992+1
        stats   = rs.zonal_stats(vectorBoundariesFile, rasterDataFile_1992_2015, band=band, categorical=true)  # stats=['count', 'mean', 'count','sum']
        stats   = convert(Vector{Dict{Int64,Int64}},stats)
        length(stats) == nRegions || @error "Oug, we have different number of regions in these computed stats!"
    else
        yearlyRasterFile = "./luse/C3S-LC-L4-LCCS-Map-300m-P1Y-$(year)-v2.1.1.tif"
        stats   = rs.zonal_stats(vectorBoundariesFile, yearlyRasterFile, band=1, categorical=true)  # stats=['count', 'mean', 'count','sum']
        stats   = convert(Vector{Dict{Int64,Int64}},stats)
        length(stats) == nRegions || @error "Oug, we have different number of regions in these computed stats!"
    end
    luseValues = [get(stat,luc,0) for stat in stats, luc in luseIds]
    yearlUseValue = hcat(gid_0,name_0,gid_1,name_1,luseValues)
    CSV.write("./out-computedLUseStatsByRegionAndYear/lUseByLev1Regions_$(year).csv",Tables.table(yearlUseValue), delim=';', header=headersForExport)
end

println("** Done!")


import Distributions: Poisson

bsfactor = 0.0

struct SingleBinnedSpectra
  f
  p::Parameter
  o::Observable
  df::Symbol
end

"""
    BinnedSpectralMonofit()

Multidimensional fit to one-dimensional distributions (neglecting
correlations between observables).

# Example
```julia-repl
m = BinnedSpectralMonofit()
```
"""
mutable struct BinnedSpectralMonofit <: Model
  params
  pdflist::Array{NLogPDF}
  nll::NLogLikelihood
  # Additional parameters
  spectra::Array{SingleBinnedSpectra}
  datasets::Array{Symbol}
  observables::Array{Observable}
  function BinnedSpectralMonofit(observables::Symbol...)
    new( [], [], NLogLikelihood(),
        Array{SingleBinnedSpectra}(undef, 0), Array{DataFrame}(undef, 0), Array{Observable}(undef, 0) )
  end
end

struct BinnedSpectralPDF
  p::Parameter
  f::String
end

"""
    generate_mock_dataset(m::BinnedSpectralMonofit; kwargs...)

words

# Arguments
- `m::BinnedSpectralMonofit`: Specific model type
- `time::Float64`: Length of the dataset
"""
function generate_mock_dataset(m::BinnedSpectralMonofit; kwargs...)
  kwargs = Dict(kwargs)
  time = get(kwargs, :time, 1.0)
  # Clean out the old data
  for d in m.datasets
    df = eval(d)
    deleterows!(df, 1:size(df, 1))
  end
  # Poisson random number of events
  for p in m.params
    p.fit = rand(Poisson(p.init))
  end
  for sp in m.spectra
    append!(sp.df[!, sp.o.sname],
            randomFromSpectrum(sp.f; count=sp.p.fit, min=sp.o.min, max=sp.o.max))
  end
  # Should global the data
  #for d in m.datasets
  #  add_dataset(d, d)
  #end
end

mutable struct EventCounter
  name::String
  count::Int64
  spectra::Array{SingleBinnedSpectra}
  obs::Dict{Symbol,Array{Float64}}
  function EventCounter(name::String, count::Int64)
    new( name, count, Array{SingleBinnedSpectra}(undef, 0), Dict() )
  end
end

function build_dataframe(ev::EventCounter)
  df = DataFrame()
  for (k,v) in ev.obs
    df[!, k] = v
  end
  df[!, :meta] = repeat([ev.name], ev.count)
  df
end

function compare_and_add!(df1::DataFrame, df2)
  names1 = names(df1)
  names2 = names(df2)
  for n1 in names1
    if !(n1 in names2)
      df2[!, n1] = typeof(df1[!, n1])(undef, size(df2,1))
    end
  end
  for n2 in names2
    if !(n2 in names1)
      df1[!, n2] = typeof(df2[!, n2])(undef, size(df1,1))
    end
  end
  append!(df1, df2)
end

function genmo(m::BinnedSpectralMonofit; kwargs...)
  kwargs = Dict(kwargs)
  time = get(kwargs, :time, 1.0)
  verbose   = get(kwargs, :verbose, false)
  countdict = get(kwargs, :counts, Dict())

  for d in m.datasets
    df = eval(d)
    deleterows!(df, 1:size(df, 1))
    @debug "Deleting from" d
  end
  dictDPCount = Dict{Symbol, Dict{String, EventCounter}}(sp.df=>Dict() for sp in m.spectra)
  for sp in m.spectra
    param = sp.p
    name = param.name
    count_init = get(countdict, name, param.init)
    count::Int64 = rand(Poisson(param.efficiencies[sp.df] * time * (count_init)))
    ### OVERRIDE
    ### count = floor(param.efficiencies[sp.df] * time * param.init)
    ###
    if verbose
      println("Name: ", param.name, " Eff: ", param.efficiencies, " Init: ", param.init, " time ", time, " count ", count)
    end
    if !haskey( dictDPCount[sp.df], name )
      dictDPCount[sp.df][name] = EventCounter(name, count)
    end
    push!( dictDPCount[sp.df][name].spectra, sp )
  end

  for (dset, evdict ) in dictDPCount
    for (pname, ev) in evdict
      count = ev.count
      for sp in ev.spectra
        newdata = randomFromSpectrum(sp.f; count=count, min=sp.o.min, max=sp.o.max)
        ev.obs[sp.o.sname] = newdata
      end
      dfnew = build_dataframe(ev)
      compare_and_add!( getfield(Batman, dset), dfnew )
    end
  end
  updateDataFunctions!(m)
  dictDPCount
end

function updateDataFunctions!(m::BinnedSpectralMonofit)
  for sp in m.spectra
    pname = sp.p.name
    obs = sp.o.name
    fSymbol = Symbol(pname*"_"*obs*"_"*String(sp.df))
    add_array(fSymbol, sp.f( eval(sp.df)[!, sp.o.sname] ))
  end
end

function constructPDF!(m::BinnedSpectralMonofit, p::Parameter, shapes, observables, ds::Symbol)
  if !(ds in m.datasets)
    push!(m.datasets, ds)
  end
  seval = ""
  for (shp, obs) in zip(shapes, observables)
    fSymbol = Symbol(p.name*"_"*String(obs)*"_"* String(ds))
    add_array(fSymbol, shp(eval(ds)[!, obs]))
    #push!(m.spectra, SingleBinnedSpectra(p, shp, obs, df))
    push!(m.spectra, SingleBinnedSpectra(shp, p, observableFromSymbol(m,obs), ds))
    # Lets add an array
    seval *= "$fSymbol.*"
  end
  seval = seval[1:end-2]
  @debug seval
  unique_name = "spectralPDF_"*p.name*"_"*String(ds)
  add_function(unique_name, eval(Meta.parse("()->"*seval)))
  BinnedSpectralPDF(p, unique_name)
end

function combinePDFs!(m::BinnedSpectralMonofit, spectralpdfs::Array{BinnedSpectralPDF},
                      ds::Symbol; kwargs... )
  kwargs = Dict(kwargs)
  livetime = get(kwargs, :livetime, nothing)
  eff = get(kwargs, :efficiency, Array{Nothing}(undef, size(spectralpdfs)))

  seval = "-sum(batlog.("
  fcall = "("
  extended = "+sum("
  ncon = "+( sum("
  plist = []
  flist = []
  if livetime != nothing
    push!(plist, livetime)
    tname = livetime.name
    fcall *= "$tname,"
  end
  for (spdf,e) in zip(spectralpdfs, eff)
    p = spdf.p
    f = spdf.f
    push!(plist, p)
    push!(flist, f)
    vname = p.name
    if e != nothing
      ename = e.name
      seval *= "$ename*"
      fcall *= "$ename,"
      extended *= "$ename*"
      ncon *= "$ename*"
      push!(plist, e)
    end
    if livetime != nothing
      tname = livetime.name
      seval *= "$tname*" ## Is this wrong???!?!?!
      extended *= "$tname*"
      ncon *= "$tname*"
    end
    vinit = p.init
    seval *= "($vname-$bsfactor)*$vinit*$f().+"
    fcall *= "$vname,"
    extended *= "($vname-$bsfactor)*$vinit+"
    ncon *= "($vname-$bsfactor)*$vinit+"
  end
  seval=seval[1:end-2]
  seval*="))"
  fcall=fcall[1:end-1]
  fcall*=")->"
  extended=extended[1:end-1]
  ncon=ncon[1:end-1]
  n = size(eval(ds), 1)
  #extended*=")"*"+$n*log($n)-$n"
  extended*=")*1.0"*"+size($ds,1)*log(size($ds,1))-size($ds,1)"
  ncon *=") - size($ds,1) )^2"
  unique_name = "fullBinnedSpectralPDF_"*String(ds)
  @debug fcall*seval*extended*ncon
  add_function(unique_name, eval(Meta.parse(fcall*seval*extended)))
  push!(m.pdflist, NLogPDF(unique_name, plist...))
end

function add_parameter!(m::BinnedSpectralMonofit, name::String, init; σ=Inf, kwargs... )
  kwargs = Dict(kwargs)
  dsEfficiency = get(kwargs, :dsEfficiency, (Symbol(name)=>1.0))
  p = getparam(m, name)
  if p != nothing
    @warn name*" already exists in model."
    push!( p.efficiencies, dsEfficiency )
    return p
  end
  @debug "Init" init+bsfactor
  p = Parameter(name; init=(init + bsfactor))
  push!( p.efficiencies, dsEfficiency )
  p.samplewidth = σ
  push!(m.params, p)
  if σ != Inf
    #a = Constant(name*"_mu", init+bsfactor)
    a = Constant(name*"_mu", 1.0)
    b = Constant(name*"_sig", σ/init)
    nlp = NLogPDF("lognormal", p, a, b)
    push!(m.pdflist, nlp)
  end
  return p
end

function add_constant!(m::BinnedSpectralMonofit, name::String, init)
  k = getparam(m, name)
  if k != nothing
    @warn name*" already exists in model."
    return k
  end
  Constant(name, init)
end

function add_observable!(m::BinnedSpectralMonofit, name::Symbol, min::Float64, max::Float64)
  push!(m.observables, Observable(String(name), name, min, max))
end

function observableFromSymbol(m::BinnedSpectralMonofit, name::Symbol)
  for obs in m.observables
    if obs.sname == name
      return obs
    end
  end
  return nothing
end

function minimize!( m::BinnedSpectralMonofit; options=Dict() )
  likelihood = NLogLikelihood( m.pdflist )
  add_likelihood!( m, likelihood )
  optimize_model!( m, likelihood; options=options )
end

"""
    randomFromSpectrum(f; min=0.0, max=1.0, top=1.0, count=1)

Return an `Array{Float64}` set of random elements from a 1D pdf given by `f`
between the values of `min` and `max` with a maximum probability of `top`.
"""
function randomFromSpectrum(f; min=0.0, max=1.0, top=10.0, count=1)
  xArr = Array{Float64}(undef, 0)
  for c in collect(1:count)
    while true
      x = rand()*(max - min) + min
      y = rand()*top
      if y < f(x)
        push!(xArr, x)
        break
      end
    end
  end
  return xArr
end


"""
    Fermi.PhysicalConstants

Module for storing physical constants and convertion factors used in computations.

# Functions:

    atomic_number   Given an element symbol, return the atomic number.
"""
module PhysicalConstants

export atomic_number

"""
    Fermi.PhysicalConstants.bohr_to_angstrom

Float64 object with the convertion factor from Bohr to Angstrom.
Source: [NIST](https://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0)
"""
#bohr_to_angstrom = 0.529177210903
#psi4s
bohr_to_angstrom = 0.529177210670000

"""
    Fermi.PhysicalConstants.angstrom_to_bohr

Float64 object with the convertion factor from Angstrom to Bohr.
Computed as 1/Fermi.PhysicalConstants.bohr_to_angstrom
"""
angstrom_to_bohr = 1/bohr_to_angstrom

hartree_to_kcalmol = 627.509474
hartree_to_kjmol = 2625.49964
hartree_to_cm = 219474.6313708
hartree_to_eV = 27.21138505

"""
    Fermi.PhysicalConstants.atomic_number(atom::String)

Given a string with the element symbol, return its atomic number
"""
function atomic_number(atom::String)

    try
        return atom_num[atom]
    catch
        error("Invalid atom: $atom")
    end

end

atom_num = Dict(
    "H" => 1,
    "He" => 2,
    "Li" => 3,
    "Be" => 4,
    "B" => 5,
    "C" => 6,
    "N" => 7,
    "O" => 8,
    "F" => 9,
    "Ne" => 10,
    "Na" => 11,
    "Mg" => 12,
    "Al" => 13,
    "Si" => 14,
    "P" => 15,
    "S" => 16,
    "Cl" => 17,
    "Ar" => 18,
    "K" => 19,
    "Ca" => 20,
    "Sc" => 21,
    "Ti" => 22,
    "V" => 23,
    "Cr" => 24,
    "Mn" => 25,
    "Fe" => 26,
    "Co" => 27,
    "Ni" => 28,
    "Cu" => 29,
    "Zn" => 30,
    "Ga" => 31,
    "Ge" => 32,
    "As" => 33,
    "Se" => 34,
    "Br" => 35,
    "Kr" => 36,
    "Rb" => 37,
    "Sr" => 38,
    "Y" => 39,
    "Zr" => 40,
    "Nb" => 41,
    "Mo" => 42,
    "Tc" => 43,
    "Ru" => 44,
    "Rh" => 45,
    "Pd" => 46,
    "Ag" => 47,
    "Cd" => 48,
    "In" => 49,
    "Sn" => 50,
    "Sb" => 51,
    "Te" => 52,
    "I" => 53,
    "Xe" => 54,
    "Cs" => 55,
    "Ba" => 56,
    "La" => 57,
    "Ce" => 58,
    "Pr" => 59,
    "Nd" => 60,
    "Pm" => 61,
    "Sm" => 62,
    "Eu" => 63,
    "Gd" => 64,
    "Tb" => 65,
    "Dy" => 66,
    "Ho" => 67,
    "Er" => 68,
    "Tm" => 69,
    "Yb" => 70,
    "Lu" => 71,
    "Hf" => 72,
    "Ta" => 73,
    "W" => 74,
    "Re" => 75,
    "Os" => 76,
    "Ir" => 77,
    "Pt" => 78,
    "Au" => 79,
    "Hg" => 80,
    "Tl" => 81,
    "Pb" => 82,
    "Bi" => 83,
    "Po" => 84,
    "At" => 85,
    "Rn" => 86,
    "Fr" => 87,
    "Ra" => 88,
    "Ac" => 89,
    "Th" => 90,
    "Pa" => 91,
    "U" => 92,
    "Np" => 93,
    "Pu" => 94,
    "Am" => 95,
    "Cm" => 96,
    "Bk" => 97,
    "Cf" => 98,
    "Es" => 99,
    "Fm" => 100,
    "Md" => 101,
    "No" => 102,
    "Lr" => 103,
    "Rf" => 104,
    "Db" => 105,
    "Sg" => 106,
    "Bh" => 107,
    "Hs" => 108,
    "Mt" => 109,
    "Ds" => 110,
    "Rg" => 111,
    "Cn" => 112,
    "Nh" => 113,
    "Fl" => 114,
    "Mc" => 115,
    "Lv" => 116,
    "Ts" => 117,
    "Og" => 118
)

end #module


# test saving and loading of trees

using Test
using IncrementalInference
using MetaGraphs
using Graphs

##

@testset "Test loading and saving of Bayes (Junction) tree" begin
##

fg = generateGraph_Kaess(graphinit=false)
tree = buildTreeReset!(fg)

# save and load tree
saveTree(tree)
tree2 = loadTree()

# perform a few spot checks to see that the trees are similar
@test length(tree.cliques) == length(tree2.cliques)
@test getEliminationOrder(tree) == getEliminationOrder(tree2)

for (clid,cl) in tree.cliques
  fsyms = getFrontals(cl)
  cl2 = getClique(tree2, fsyms[1])
  fsyms2 = getFrontals(cl2)
  @test fsyms == fsyms2
  @test getCliqSeparatorVarIds(cl) == getCliqSeparatorVarIds(cl2)
  @test typeof(cl) == typeof(cl2)
end

##
end

@testset "Test loading and saving of Bayes (Junction) tree" begin
##

fg = generateGraph_Kaess(graphinit=false)
tree = buildTreeReset!(fg)

# save and load tree as array
filepath = saveTree([tree;deepcopy(tree)])
trees = loadTree(filepath)

# perform a few spot checks to see that the trees are similar
@test length(tree.cliques) == length(trees[1].cliques)
@test getEliminationOrder(tree) == getEliminationOrder(trees[1])

for (clid,cl) in tree.cliques
  fsyms = getFrontals(cl)
  cl2 = getClique(trees[1], fsyms[1])
  fsyms2 = getFrontals(cl2)
  @test fsyms == fsyms2
  @test getCliqSeparatorVarIds(cl) == getCliqSeparatorVarIds(cl2)
  @test typeof(cl) == typeof(cl2)
end

##
end


"""
@fastfunc functionname polynomial

Generates a function for fast numerical evaluation. The number of input
arguments of the generated function will match the number of variables in the
input polynomial.

Note that the polynomial is transformed using `@fastmath`, which calls
functions that may violate strict IEEE semantics.

If you only plan a few evaluations of polynomial use the standard method (see
e.g. TypedPolynomials for details). Only for a large number of evaluations the
one can break even with the function generated by this method due to the just
in time compilation overhead.

# Examples
```
julia> using SphericalHarmonics

julia> @polyvar x y z
(x, y, z)

julia> p = 15.0*x*y^2+7.5*x*z^13
7.5xz¹³ + 15.0xy²

julia> foo = @fastfunc p

julia> foo(1.0,2.0,3.0)
1.19574825e7
```

# Non-global scope

Usage from within local scope requires special care to avoid issue #4. So
instead of `foo(1.0,2.0,3.0)` we need `Base.invokelatest(foo, 1.0,2.0,3.0)`.

```
julia> using SphericalHarmonics

julia> @polyvar x y z
(x, y, z)

julia> function useInsideFunctionScope()
           p = 15.0*x*y^2+7.5*x*z^13
           foo = @fastfunc p
           Base.invokelatest(foo, 1.0,2.0,3.0)
       end
useInsideFunctionScope (generic function with 1 method)

julia> useInsideFunctionScope()
1.19574825e7
```
"""
macro fastfunc(polynomial)
	return quote
		local polystr = string($(esc(polynomial)))
		local vars = string(tuple(variables($(esc(polynomial)))...))
        # create expression for function definition
		eval(Meta.parse(vars*" -> @fastmath "*polystr))
	end
end

"""
fastfunc(polynomial::Polynomial)

Generates a function for fast numerical evaluation. The number of input
arguments of the generated function will match the number of variables in the
input polynomial.

Note that the polynomial is transformed using `@fastmath`, which calls
functions that may violate strict IEEE semantics.

If you only plan a few evaluations of polynomial use the standard method (see
e.g. TypedPolynomials for details). Only for a large number of evaluations the
one can break even with the function generated by this method due to the just
in time compilation overhead.

# Examples

```
julia> using SphericalHarmonics

julia> @polyvar x y z
(x, y, z)

julia> p = 15.0*x*y^2+7.5*x*z^13
7.5xz¹³ + 15.0xy²

julia> foo = fastfunc(p)

julia> foo(1.0,2.0,3.0)
1.19574825e7
```
"""
function fastfunc(polynomial)
    expr = Meta.parse(string(variables(polynomial))*" -> "*string(polynomial))
    # use fastmath
    expr = Base.FastMath.make_fastmath(expr)
    return mk_function(expr)
end


"""
    IntegrationVolume(dataset::DataSet{T, I}, spvol::HyperRectVolume{T}, searchpts::Bool = true)::IntegrationVolume{T, I}

*AHMI-internal, not part of stable public API.*

creates an integration region by calculating the point cloud an the volume of the spatial volume.
"""
function IntegrationVolume(
    dataset::DataSet{T, I},
    spvol::HyperRectVolume{T},
    searchpts::Bool = true
)::IntegrationVolume{T, I} where {T<:AbstractFloat, I<:Integer}

    cloud = PointCloud(dataset, spvol, searchpts)
    vol = prod(spvol.hi - spvol.lo)

    IntegrationVolume(cloud, spvol, vol)
end

"""
    modify_integrationvolume!(intvol::IntegrationVolume{T, I}, dataset::DataSet{T, I}, spvol::HyperRectVolume{T}, searchpts::Bool = true)

*AHMI-internal, not part of stable public API.*

updates an integration volume with new boundaries. Recalculates the pointcloud and volume.
"""
function modify_integrationvolume!(
    intvol::IntegrationVolume{T, I},
    dataset::DataSet{T, I},
    spvol::HyperRectVolume{T},
    searchpts::Bool = true
) where {T<:AbstractFloat, I<:Integer}

    if ndims(intvol.spatialvolume) != ndims(spvol)
        intvol.spatialvolume = deepcopy(spvol)
    else
        copy!(intvol.spatialvolume, spvol)
    end

    PointCloud!(intvol.pointcloud, dataset, spvol, searchpts)

    intvol.volume = prod(spvol.hi - spvol.lo)

    nothing
end

function shrink_integrationvol!(
    volume::IntegrationVolume{T, I},
    dataset::DataSet{T, I},
    newrect::HyperRectVolume{T}) where {T<:AbstractFloat, I<:Integer}

    i = volume.pointcloud.points
    for j = 1:i
        inV = true
        for p = 1:dataset.P
            if dataset.data[p, volume.pointcloud.pointIDs[i]] < newrect.lo[p] || dataset.data[p, volume.pointcloud.pointIDs[i]] > newrect.hi[p]
                inV = false
                break
            end
        end
        if !inV
            deleteat!(volume.pointcloud.pointIDs, i)
        end
        i -= 1
    end
    copy!(volume.spatialvolume, newrect)
    volume.pointcloud.points = length(volume.pointcloud.pointIDs)
    volume.volume = prod(newrect.hi .- newrect.lo)

    nothing
end

function update!(
    volume::IntegrationVolume{T, I},
    dataset::DataSet{T, I}) where {T<:AbstractFloat, I<:Integer}

    res = search(dataset, volume.spatialvolume, true)

    _update!(volume, res, false, true)

    #volume.volume doesn't change
    volume.pointcloud.probfactor = exp(volume.pointcloud.maxLogProb - volume.pointcloud.minLogProb)
    volume.pointcloud.probweightfactor = exp(volume.pointcloud.maxWeightProb - volume.pointcloud.minWeightProb)

    @debug "Hyperrectangle updated. Points:\t$(volume.pointcloud.points)\tVolume:\t$(volume.volume)\tProb. Factor:\t$(volume.pointcloud.probfactor)"
    nothing
end

function _update!(
    volume::IntegrationVolume{T, I},
    searchres::SearchResult{T, I},
    addpts::Bool, #if true adds the points to the volume, if false assumes to replace the points in the volume
    searchpts::Bool) where {T<:AbstractFloat, I<:Integer}

    volume.pointcloud.points = searchres.points + (addpts ? volume.pointcloud.points : 0)

    volume.pointcloud.maxLogProb = volume.pointcloud.maxLogProb > searchres.maxLogProb && addpts ? volume.pointcloud.maxLogProb : searchres.maxLogProb
    volume.pointcloud.minLogProb = volume.pointcloud.minLogProb < searchres.minLogProb && addpts ? volume.pointcloud.minLogProb : searchres.minLogProb

    volume.pointcloud.maxWeightProb = volume.pointcloud.maxWeightProb > searchres.maxWeightProb && addpts ? volume.pointcloud.maxWeightProb : searchres.maxWeightProb
    volume.pointcloud.minWeightProb = volume.pointcloud.minWeightProb < searchres.minWeightProb && addpts ? volume.pointcloud.minWeightProb : searchres.minWeightProb

    if searchpts && searchres.points > 0
        start = addpts ? length(volume.pointcloud.pointIDs) + 1 : 1
        resize!(volume.pointcloud.pointIDs, volume.pointcloud.points)
        copyto!(volume.pointcloud.pointIDs, start, searchres.pointIDs, 1)
    end

    nothing
end

function resize_integrationvol(
    original::IntegrationVolume{T, I},
    dataset::DataSet{T, I},
    changed_dim::I,
    newrect::HyperRectVolume{T},
    searchpts::Bool = false)::IntegrationVolume{T, I} where {T<:AbstractFloat, I<:Integer}

    result = deepcopy(original)

    resize_integrationvol!(result, original, dataset, datatre, changed_dim, newrect, searchpts)
end

function resize_integrationvol!(
    result::IntegrationVolume{T, I},
    original::IntegrationVolume{T, I},
    dataset::DataSet{T, I},
    changed_dim::I,
    newrect::HyperRectVolume{T},
    searchpts::Bool,
    searchVol::HyperRectVolume{T}) where {T<:AbstractFloat, I<:Integer}

    copy!(searchVol, newrect)
    increase = true

    #increase
    if original.spatialvolume.lo[changed_dim] > newrect.lo[changed_dim]
        searchVol.hi[changed_dim] = original.spatialvolume.lo[changed_dim]
        searchVol.lo[changed_dim] = newrect.lo[changed_dim]
    elseif original.spatialvolume.hi[changed_dim] < newrect.hi[changed_dim]
        searchVol.lo[changed_dim] = original.spatialvolume.hi[changed_dim]
        searchVol.hi[changed_dim] = newrect.hi[changed_dim]
    else
        increase = false
        if original.spatialvolume.lo[changed_dim] < newrect.lo[changed_dim]
            searchVol.lo[changed_dim] = original.spatialvolume.lo[changed_dim]
            searchVol.hi[changed_dim] = newrect.lo[changed_dim]
        elseif original.spatialvolume.hi[changed_dim] > newrect.hi[changed_dim]
            searchVol.hi[changed_dim] = original.spatialvolume.hi[changed_dim]
            searchVol.lo[changed_dim] = newrect.hi[changed_dim]
        else
            #check if pts inside vol are corrected
            @error "resize_integrationvol(): Volume $original didn't change.", searchVol.lo[changed_dim], "\n", searchVol.hi[changed_dim], "\n", original.spatialvolume.lo[changed_dim], "\n", original.spatialvolume.hi[changed_dim]
        end
    end

    result.pointcloud.points = original.pointcloud.points
    if searchpts
        if result.pointcloud.pointIDs != original.pointcloud.pointIDs
            result.pointcloud.pointIDs = deepcopy(original.pointcloud.pointIDs)
        end
    end

    if increase
        res = search(dataset, searchVol, searchpts)

        _update!(result, res, true, searchpts)
    else
        res = search(dataset, searchVol, searchpts)
        result.pointcloud.points = original.pointcloud.points - res.points
        if searchpts
            newids = search(dataset, newrect, searchpts).pointIDs
            resize!(result.pointcloud.pointIDs, result.pointcloud.points)
            copyto!(result.pointcloud.pointIDs, newids)
        end
    end

    result.volume = prod(newrect.hi - newrect.lo)
    result.pointcloud.probfactor = exp(result.pointcloud.maxLogProb - result.pointcloud.minLogProb)
    result.pointcloud.probweightfactor = exp(result.pointcloud.maxWeightProb - result.pointcloud.minWeightProb)
    copy!(result.spatialvolume, newrect)

    nothing
end

function Base.copy!(
    target::IntegrationVolume{T, I},
    src::IntegrationVolume{T, I}) where {T<:AbstractFloat, I<:Integer}

    target.volume = src.volume

    copy!(target.spatialvolume, src.spatialvolume)
    target.pointcloud.points = src.pointcloud.points

    copy!(target.pointcloud, src.pointcloud)

    nothing
end


export filterFrequencies

"""
  filterFrequencies(f; kargs...) => Vector{Int64}

Supported keyword arguments:
* SNRThresh
* minFreq
* maxFreq
* recChannels
* sortBySNR
* numUsedFreqs
* stepsize
* maxMixingOrder
* sortByMixFactors
"""
function filterFrequencies(f::MPIFile; SNRThresh=-1, minFreq=0,
               maxFreq=rxBandwidth(f), recChannels=1:rxNumChannels(f),
               sortBySNR=false, numUsedFreqs=-1, stepsize=1, maxMixingOrder=-1,
               sortByMixFactors=false, numPeriodAverages=1, numPeriodGrouping=1)

  nFreq = rxNumFrequencies(f, numPeriodGrouping)
  nReceivers = rxNumChannels(f)
  nPeriods = 1 #acqNumPeriodsPerFrame(f)

  minIdx = floor(Int, minFreq / rxBandwidth(f) * (nFreq-1) ) + 1
  maxIdx = ceil(Int, maxFreq / rxBandwidth(f) * (nFreq-1) ) + 1

  freqMask = zeros(Bool,nFreq,nReceivers,nPeriods)

  freqMask[:,recChannels,:] .= true

  if measIsFrequencySelection(f)
    freqMask[:,recChannels,:] .= false
    idx = measFrequencySelection(f)
    freqMask[idx,recChannels,:] .= true
  else
    freqMask[:,recChannels,:] .= true
  end

  if minFreq > 0
    freqMask[1:(minIdx),:,:] .= false
  end

  if maxIdx < nFreq
    freqMask[(maxIdx):end,:,:] .= false
  end

  if maxMixingOrder > 0
    if numPeriodGrouping == 1
      mf = mixingFactors(f)
      for l=1:size(mf,1)
        if mf[l,4] > maxMixingOrder || mf[l,4] > maxMixingOrder
          freqMask[(l-1)+1,recChannels] .= false
        end
      end
    else # This is a hack until we have a general solution

      numPatches = div(acqNumPeriodsPerFrame(f), numPeriodAverages)
      fBands = (collect(0:(rxNumFrequencies(f)-1)).-1) .* numPatches
      freqMaskMO = zeros(Bool,nFreq,nReceivers,nPeriods)

      for i=1:length(fBands)
        for y=-maxMixingOrder:maxMixingOrder
          index = fBands[i]+y
          if 1 <= index <= size(freqMask,1)
            freqMaskMO[index,recChannels] .= true
          end
        end
      end
      freqMask .&= freqMaskMO
    end
  end

  SNRAll = nothing

  if SNRThresh > 0 || numUsedFreqs > 0 || sortBySNR
    SNR = zeros(nFreq, nReceivers)
    idx = measIsFrequencySelection(f) ? measFrequencySelection(f) : idx = 1:nFreq

    SNRAll = calibSNR(f)
    if SNRAll != nothing
      SNR[idx,:] = SNRAll[:,:,1]
    end
  end

  if SNRThresh > 0 && numUsedFreqs > 0
    error("It is not possible to use SNRThresh and SNRFactorUsedFreq similtaneously")
  end

  if SNRThresh > 0 && SNRAll != nothing
    freqMask[ findall(vec(SNR) .< SNRThresh) ] .= false
  end

  if numUsedFreqs > 0 && SNRAll != nothing
    numFreqsAlreadyFalse = sum(!freqMask)
    numFreqsFalse = round(Int,length(freqMask)* (1-numUsedFreqs))
    S = sortperm(vec(SNR))

    l = 1
    j = 1
    while j<  (numFreqsFalse-numFreqsAlreadyFalse)
      if freqMask[S[l]] == true
        freqMask[S[l]] = false
        j += 1
      end
      l += 1
    end

  end

  if stepsize > 1
    freqStepsizeMask = zeros(Bool,nFreq, nReceivers, nPatches)
    freqStepsizeMask[1:stepsize:nFreq,:,:] = freqMask[1:stepsize:nFreq,:,:]
    freqMask = freqStepsizeMask
  end

  freq = findall( vec(freqMask) )

  if sortBySNR && !sortByMixFactors && SNRAll != nothing
    SNR = vec(SNR[1:stepsize:nFreq,:,:])

    freq = freq[reverse(sortperm(SNR[freq]),dims=1)]
  end

  if !sortBySNR && sortByMixFactors
    mfNorm = zeros(nFreq,nReceivers,nPeriods)
    mf = mixingFactors(f)
    for k=1:nFreq
      mfNorm[k,:,:] = norm(mf[k,1:3])
    end

    freq = freq[sortperm(mfNorm[freq])]
  end

  freq
end

function rowsToSubsampledRows(f::MPIFile, rows)
  if measIsFrequencySelection(f)
    # In this case we need to convert indices
    tmp = zeros(Int64, rxNumFrequencies(f), rxNumChannels(f) )
    idxAvailable = measFrequencySelection(f)
    for d=1:rxNumChannels(f)
      tmp[idxAvailable, d] = (1:length(idxAvailable)) .+ (d-1)*length(idxAvailable)
    end
    rows_ = vec(tmp)[rows]
    if findfirst(x -> x == 0, rows_) != nothing
      @error "Indices applied to systemMatrix are not available in the file"
    end
  else
    rows_ = rows
  end
  return rows_
end


#===============================================================================
                                The main routine
===============================================================================#
"""
    mcmc(setup_mcmc::MCMCSetup, schedule::MCMCSchedule, setup)  <: ModelSetup

function for running the mcmc. receives as imput `MCMCSetup`, `MCMCschedule` and `setup`.
See [MCMCSetup](@ref), [MCMCSchedule](@ref). setup typically is DiffusionSetup making
additional setup choices when using the MCMC infrastructure to sample diffusion processes.
See [DiffusionSetup](@ref)
"""
function mcmc(setup_mcmc::MCMCSetup, schedule::MCMCSchedule, setup::T) where T <: ModelSetup
    ws, ll, θ = create_workspace(setup)
    ws_mcmc = create_workspace(setup_mcmc, schedule, θ)
    adpt = adaptation_object(setup, ws)

    aux = nothing
    for step in schedule
        step.save && save_imputed!(ws)
        for i in step.idx
            if step.param_updt || typeof(ws_mcmc.updates[i]) <: Imputation
                ws = next(ws, ws_mcmc.updates[i])
                ll, acc, θ = update!(ws_mcmc.updates[i], ws, θ, ll, step, aux)
                aux = aux_params(ws_mcmc.updates[i], aux)
                update!(ws_mcmc, acc, θ, i)
                step.verbose && print("\n")
            end
        end
        step.verbose && print("-----------------------------------------------",
                              "------\n")
        step.readjust && readjust!(ws_mcmc, step.iter)
        step.fuse && fuse!(ws_mcmc, schedule)
        ll = adaptation!(ws, adpt, step.iter, ll)
    end
    display_summary(ws, ws_mcmc)
    ws, ws_mcmc
end

display_summary(::Any, ::Any) = print("nothing to display...\n")
adaptation_object(::Any, ::Any) = nothing
adaptation!(::Any, ::Nothing, ::Any, ll) = ll


module TestOperators
    using Base.Test
    using DataArrays
    using DataFrames
    using Stats

    macro test_da_pda(da, code)
        esc(quote
            let $da = copy($da)
                $code
            end
            let $da = PooledDataArray($da)
                $code
            end
        end)
    end

    unary_operators = [(+), (-), (!)]

    numeric_unary_operators = [(+), (-)]

    logical_unary_operators = [(!)]

    elementary_functions = [abs, sign, acos, acosh, asin,
                            asinh, atan, atanh, sin, sinh,
                            cos, cosh, tan, tanh, ceil, floor,
                            round, trunc, exp, exp2, expm1, log, log10, log1p,
                            log2, exponent, sqrt, gamma, lgamma, digamma,
                            erf, erfc]

    special_comparison_operators = [isless]

    comparison_operators = [(==), (.==), (!=), (.!=),
                            (>), (.>), (>=), (.>=), (<), (.<),
                            (<=), (.<=)]

    scalar_comparison_operators = [(==), (!=), (>), (>=),
                                   (<), (<=)]

    array_comparison_operators = [(.==), (.!=), (.>), (.>=), (.<), (.<=)]

    binary_operators = [(+), (.+), (-), (.-), (*), (.*), (/), (./),
                        (.^), (div), (mod), (fld), (rem),
                        (&), (|), ($)]

    induced_binary_operators = [(^)]

    arithmetic_operators = [(+), (.+), (-), (.-), (*), (.*), (/), (./),
                            (.^), (div), (mod), (fld), (rem)]

    induced_arithmetic_operators = [(^)]

    biscalar_operators = [(max), (min)]

    scalar_arithmetic_operators = [(+), (-), (*), (/),
                                   (div), (mod), (fld), (rem)]

    induced_scalar_arithmetic_operators = [(^)]

    array_arithmetic_operators = [(+), (.+), (-), (.-), (.*), (./), (.^)]

    bit_operators = [(&), (|), ($)]

    unary_vector_operators = [minimum, maximum, prod, sum, mean, median, std, var, norm]

    ex_tunary_vector_operators = [mad, skewness, kurtosis]

    pairwise_vector_operators = [diff]#, percent_change]

    cumulative_vector_operators = [cumprod, cumsum, cumsum_kbn, cummin, cummax]

    ffts = [fft]

    binary_vector_operators = [dot, cor, cov, corspearman]

    columnar_operators = [colmins, colmaxs, colprods, colsums,
                          colmeans, colmedians, colstds, colvars,
                          colnorms]

    boolean_operators = [any, all]

    # Unary operators on DataFrame's should be equivalent to elementwise
    # application of those same operators
    df = DataFrame(quote
                       A = [1, 2, 3, 4]
                       B = [1.0, pi, pi, e]
                   end)
    for f in numeric_unary_operators
        for i in 1:nrow(df)
            for j in 1:ncol(df)
                @assert f(df)[i, j] == f(df[i, j])
            end
        end
    end
    df = DataFrame(quote
                       A = [true, false, true, false]
                   end)
    for f in logical_unary_operators
        for i in 1:nrow(df)
            for j in 1:ncol(df)
                @assert f(df)[i, j] == f(df[i, j])
            end
        end
    end

    # Elementary functions on DataFrames's
    N = 5
    df = DataFrame(A = DataArray(ones(N)),
                   B = DataArray(ones(N)))
    for f in elementary_functions
        for i in 1:nrow(df)
            for j in 1:ncol(df)
                  @assert f(df)[i, j] == f(df[i, j])
            end
        end
    end

    # Broadcasting operations between NA's and DataFrames's
    N = 5
    df = DataFrame(A = DataArray(ones(N)),
                   B = DataArray(ones(N)))
    for f in arithmetic_operators
        for i in 1:nrow(df)
            for j in 1:ncol(df)
                @assert isna(f(df, NA)[i, j])
                @assert isna(f(NA, df)[i, j])
            end
        end
    end

    # Broadcasting operations between scalars and DataFrames's
    N = 5
    df = DataFrame(A = DataArray(ones(N)),
                   B = DataArray(ones(N)))
    for f in arithmetic_operators
        for i in 1:nrow(df)
            for j in 1:ncol(df)
                @assert f(df, 1)[i, j] == f(df[i, j], 1)
                @assert f(1, df)[i, j] == f(1, df[i, j])
            end
        end
    end

    # Binary operations on pairs of DataFrame's
    # TODO: Test in the presence of in-operable types like Strings
    N = 5
    df = DataFrame(A = DataArray(ones(N)),
                   B = DataArray(ones(N)))
    for f in array_arithmetic_operators
        for i in 1:nrow(df)
            for j in 1:ncol(df)
                @assert isna(f(df, df)[i, j]) && isna(df[i, j]) ||
                        f(df, df)[i, j] == f(df[i, j], df[i, j])
            end
        end
    end

    # TODO: Columnar operators on DataFrame's

    # Boolean operators on DataFrames's
    N = 5
    df = DataFrame(A = DataArray(falses(N)))
    @assert any(df) == false
    @assert any(!df) == true
    @assert all(df) == false
    @assert all(!df) == true

    df = DataFrame(A = DataArray(falses(N)))
    df[3, 1] = true
    @assert any(df) == true
    @assert all(df) == false

    df = DataFrame(A = DataArray(falses(N)))
    df[2, 1] = NA
    df[3, 1] = true
    @assert any(df) == true
    @assert all(df) == false

    dv = @data [false, NA]
    dv[1] = NA

    df = DataFrame(A = DataArray(falses(N)))
    df[2, 1] = NA
    @assert isna(any(df))
    @assert all(df) == false

    df = DataFrame(A = DataArray(falses(N)))
    df[1, 1] = NA
    @assert isna(any(dv))
    @assert isna(all(dv))

    # Is this a genuine special case?
    @assert isna(NA ^ 2.0)

    #
    # Equality tests
    #
    df = DataFrame({dv})
    alt_dv = @data [false, NA]
    alt_dv[1] = NA
    alt_df = DataFrame({alt_dv})
    # @assert isequal(DataFrame({dv}) .== DataFrame({dv}), DataFrame({DataVector[true, NA]}))

    @assert all(isna(NA .== df))
    @assert all(isna(df .== NA))
end


for f in (:seed_zeros, :seed_position, :seed_strength)
  @eval $f(v::Vector{<:Point},cfg::ComplexGradientConfig) =
      Point.($f(circulation,v,cfg)...)
end

for f in (:seed_zeros, :seed_position, :seed_strength)
  @eval $f(v::Vector{<:Blob},cfg::ComplexGradientConfig) =
      Blob.($f(circulation,v,cfg)...,Elements.blobradius(v))
end


"""
    TuringModel <: Probabilistic

Fields:
- `model`: a Turing model function; define it via `@model func(X, y)`
- `n_samples`
- `sampler`
- `renamer`: Can be used to rename the MCMCChains parameter names.
    Defaults to `nothing`.
    If `renamer isa Dict`, then names are dropped from the data `X`.
- `n_chains`
- `multithreaded`: Defaults to `true`; this can be set to false when doing threaded evaluations via MLJ.
"""
mutable struct TuringModel <: Probabilistic
    model::Function
    n_samples::Int
    sampler::Any
    renamer::Union{Nothing,Dict}
    n_chains::Int
    multithreaded::Bool
end

function TuringModel(model, n_samples, sampler; renamer=nothing, n_chains=3, multithreaded=true)
    return TuringModel(model, n_samples, sampler, renamer, n_chains, multithreaded)
end

function _parameter_mean(chns, parameter)
    values = collect(Iterators.flatten(chns[parameter]))
    return mean(values)
end

"""
    _simplify_data(X) -> Matrix

Convert `X` to a Matrix since this usually works better with Turing.jl models.
Assumes that `X` satisfies the Tables interface.
This code is partially based on `Base.Matrix` in DataFrames.jl.
"""
function _simplify_data(X)::Matrix
    nrows = length(Tables.rows(X))
    ncols = length(Tables.columns(X))
    T = reduce(promote_type, (eltype(v) for v in Tables.columns(X)), init=Union{})
    out = Matrix{T}(undef, nrows, ncols)
    for (i, row) in enumerate(Tables.rows(X))
        for (j, col) in enumerate(Tables.columnnames(row))
            out[i, j] = Tables.getcolumn(row, col)
        end
    end
    return out
end

function MMI.fit(tm::TuringModel, verbosity::Int, X, y)
    renamed_X = isnothing(tm.renamer) ? X : _simplify_data(X)
    model = tm.model(renamed_X, y)
    chns = if tm.multithreaded
        sample_func(c) = sample(model, tm.sampler, tm.n_samples)
        mapreduce(sample_func, chainscat, 1:tm.n_chains)
    else
        sample(model, tm.sampler, MCMCThreads(), tm.n_samples, tm.n_chains)
    end
    renamed_chns = isnothing(tm.renamer) ? chns : replacenames(chns, tm.renamer)
    fitresult = renamed_chns
    cache = nothing
    report = nothing

    return fitresult, cache, report
end

function MMI.predict(tm::TuringModel, fitresult, Xnew)
    model = tm.model(Xnew, missing)
    chns = predict(model, fitresult)
    y_matrix = mean(Array(group(chns, :y)); dims=1)
    y = collect(Iterators.flatten(y_matrix))
    return y
end


using Random, LinearAlgebra, BlockArrays
using BlockArrays: cumulsizes, _BlockArray, BlockSizes

@timedtestset "dense" begin

    # Test construction of a BlockArray via _BlockArray
    @testset "_BlockArray" begin
        rng, P, Q, R = MersenneTwister(123456), 7, 13, 11
        @testset "Block Vector" begin
            x, blk_sizes = randn(rng, P), BlockSizes([3, 2, 2])
            ȳ = randn(rng, P)
            x_blks = BlockArray(x, blk_sizes).blocks
            ȳ_blks = BlockArray(ȳ, blk_sizes).blocks

            y, back = Zygote.forward(_BlockArray, x_blks, blk_sizes)
            @test y == _BlockArray(x_blks, blk_sizes)
            @test first(back(ȳ)) == ȳ_blks
            @test last(back(ȳ)) === nothing
            @test first(back(ȳ)) == first(back(BlockArray(ȳ, blk_sizes)))
            @test first(back(ȳ)) == first(back((blocks=ȳ_blks, block_sizes=nothing)))
        end
        @testset "Block Array 3" begin
            X, blk_sizes = randn(rng, P, Q, R), BlockSizes([3, 4], [7, 6], [2, 3, 6])
            Ȳ = randn(rng, P, Q, R)
            X_blks = BlockArray(X, blk_sizes).blocks
            Ȳ_blks = BlockArray(Ȳ, blk_sizes).blocks

            Y, back = Zygote.forward(_BlockArray, X_blks, blk_sizes)
            @test Y == _BlockArray(X_blks, blk_sizes)
            @test first(back(Ȳ)) == Ȳ_blks
            @test last(back(Ȳ)) === nothing
            @test first(back(Ȳ)) == first(back(BlockArray(Ȳ, blk_sizes)))
            @test first(back(Ȳ)) == first(back((blocks=Ȳ_blks, block_sizes=nothing)))
        end
    end

    @testset "Vector(::BlockVector)" begin
        rng, Ps = MersenneTwister(123456), [5, 6, 7]
        x = BlockArray(randn(rng, sum(Ps)), Ps)
        adjoint_test(Vector, randn(rng, sum(Ps)), x)
    end

    @testset "Matrix(::BlockMatrix)" begin
        rng, Ps, Qs = MersenneTwister(123456), [3, 4, 5], [6, 7, 8, 9]
        X = BlockArray(randn(rng, sum(Ps), sum(Qs)), Ps, Qs)
        adjoint_test(Matrix, randn(rng, sum(Ps), sum(Qs)), X)
    end

    # # Test construction of `BlockVector` from a vector of vectors. Also test copying.
    # @testset "BlockVector" begin
    #     rng, N, N′ = MersenneTwister(123456), 5, 6
    #     x, x′ = randn(rng, N), randn(rng, N′)
    #     x̂ = BlockVector([x, x′])
    #     @test length(x̂) == N + N′
    #     @test x̂ == vcat(x, x′)
    #     @test getblock(x̂, 1) == x && getblock(x̂, 2) == x′

    #     # Test BlockVector construction adjoint.
    #     adjoint_test(x->BlockVector([x]), BlockVector([randn(rng, N)]), x)
    #     adjoint_test(x->BlockVector([x]), randn(rng, N), x)

    #     adjoint_test(
    #         x->BlockVector([x, x′]),
    #         BlockVector([randn(rng, N), randn(rng, N′)]),
    #         x,
    #     )
    #     adjoint_test(x->BlockVector([x, x′]), randn(rng, N + N′), x)

    #     # Test BlockVector conversion to Vector adjoint.
    #     adjoint_test(Vector, randn(rng, N), BlockVector([x]))
    #     adjoint_test(Vector, randn(rng, N + N′), x̂)

    #     # zero of a BlockVector should be a BlockVector
    #     @test zero(x̂) isa BlockVector
    #     @test cumulsizes(x̂) == cumulsizes(zero(x̂))
    #     @test zero(x̂) == zero(Vector(x̂))
    # end

    # # Test construction of `BlockMatrix` from matrix of matrices. Also test copying.
    # @testset "BlockMatrix" begin
    #     rng, P1, P2, P3, Q1, Q2 = MersenneTwister(123456), 2, 3, 6, 4, 5
    #     X11, X12 = randn(rng, P1, Q1), randn(rng, P1, Q2)
    #     X21, X22 = randn(rng, P2, Q1), randn(rng, P2, Q2)
    #     X31, X32 = randn(rng, P3, Q1), randn(rng, P3, Q2)
    #     X = BlockMatrix(reshape([X11, X21, X31, X12, X22, X32], 3, 2))
    #     @test size(X, 1) == P1 + P2 + P3
    #     @test size(X, 2) == Q1 + Q2
    #     @test getblock(X, 1, 1) == X11
    #     @test getblock(X, 1, 2) == X12
    #     @test getblock(X, 2, 1) == X21
    #     @test getblock(X, 2, 2) == X22

    #     @test cumulsizes(X, 1) == cumsum([1, P1, P2, P3])
    #     @test cumulsizes(X, 2) == cumsum([1, Q1, Q2])

    #     @test blocksizes(X, 1) == [P1, P2, P3]
    #     @test blocksizes(X, 2) == [Q1, Q2]

    #     @test BlockMatrix([X11, X21, X31, X12, X22, X32], 3, 2) == X

    #     Y = BlockMatrix([X11, X21, X31])
    #     @test size(Y, 1) == P1 + P2 + P3
    #     @test size(Y, 2) == Q1
    #     @test getblock(Y, 1, 1) == X11
    #     @test getblock(Y, 2, 1) == X21
    #     @test cumulsizes(Y, 1) == cumsum([1, P1, P2, P3])
    #     @test cumulsizes(Y, 2) == cumsum([1, Q1])

    #     @test blocksizes(Y, 1) == [P1, P2, P3]
    #     @test blocksizes(Y, 2) == [Q1]

    #     P, Q = P1 + P2 + P3, Q1 + Q2
    #     adjoint_test(
    #         X11->BlockMatrix(reshape([X11, X21, X31, X12, X22, X32], 3, 2)),
    #         randn(rng, P, Q),
    #         X11,
    #     )
    #     adjoint_test(
    #         X11->BlockMatrix(reshape([X11, X21, X31, X12, X22, X32], 3, 2)),
    #         X,
    #         X11,
    #     )
    #     adjoint_test(X11->BlockMatrix([X11, X21]), randn(rng, P1 + P2, Q1), X11)
    #     adjoint_test(X11->BlockMatrix([X11, X21]), BlockMatrix([X11, X21]), X11)
    #     adjoint_test(X11->BlockMatrix([X11, X12], 1, 2), BlockMatrix([X11, X12], 1, 2), X11)

    #     # BlockMatrix conversion to Matrix.
    #     adjoint_test(Matrix, randn(rng, size(X)), X)

    #     # zero of a BlockMatrix should be a BlockMatrix.
    #     @test zero(X) isa BlockMatrix
    #     @test cumulsizes(X) == cumulsizes(zero(X))
    #     @test zero(X) == zero(Matrix(X))
    # end

    # @testset "dense, transpose, and adjoint mul" begin
    #     rng, Ps, Qs, Rs = MersenneTwister(123456), [2, 3], [4, 5], [6, 7]
    #     X = _BlockArray([randn(rng, P, Q) for P in Ps, Q in Qs], Ps, Qs)

    #     for foo in [transpose, adjoint]
    #         @test foo(X) isa BlockMatrix
    #         @test blocksizes(foo(X), 1) == blocksizes(X, 2)
    #         @test blocksizes(foo(X), 2) == blocksizes(X, 1)
    #         @test foo(Matrix(X)) == Matrix(foo(X))
    #         @test foo(foo(X)) == X
    #     end

    #     y = _BlockArray([randn(rng, Q) for Q in Qs], Qs)
    #     z = _BlockArray([randn(rng, P) for P in Ps], Ps)
    #     dense_BlockMatrix_BlockVector_mul_tests(rng, X, y)
    #     dense_BlockMatrix_BlockVector_mul_tests(rng, X', z)

    #     Y = _BlockArray([randn(rng, Q, R) for Q in Qs, R in Rs], Qs, Rs)

    #     Xt_blocks = Matrix.(transpose(X.blocks))
    #     Xt = _BlockArray(Xt_blocks, Qs, Ps)
    #     Yt_blocks = Matrix.(transpose(Y.blocks))
    #     Yt = _BlockArray(Yt_blocks, Rs, Qs)

    #     dense_BlockMatrix_BlockMatrix_mul_tests(rng, X, Y)
    #     for foo in [transpose, adjoint]
    #         dense_BlockMatrix_BlockMatrix_mul_tests(rng, foo(Xt), Y)
    #         dense_BlockMatrix_BlockMatrix_mul_tests(rng, X, foo(Yt))
    #         dense_BlockMatrix_BlockMatrix_mul_tests(rng, foo(Xt), foo(Yt))
    #     end
    # end

    # @testset "dense cholesky" begin
    #     rng, Ps = MersenneTwister(123456), [3, 4, 5]
    #     tmp = randn(rng, sum(Ps), sum(Ps))
    #     A_ = tmp * tmp' + 1e-9I
    #     A = BlockArray(A_, Ps, Ps)
    #     @assert A_ == A

    #     # Compute cholesky and compare.
    #     C_, C = cholesky(A_), cholesky(A)
    #     @test C isa Cholesky
    #     @test C_.U ≈ Matrix(C.U)
    #     @test C.U isa UpperTriangular{T, <:AbstractBlockMatrix{T}} where T

    #     function cholesky_gradients(C̄_factors, A_, A)

    #         # Compute cholesky adjoint with Matrix.
    #         C̄_dense = (factors=C̄_factors, uplo=nothing, info=nothing)
    #         _, back_dense = Zygote.forward(cholesky, A_)
    #         Ā_ = first(back_dense(C̄))

    #         # Compute cholesky adjoint with BlockMatrix.
    #         C̄_block = (factors=BlockArray(C̄_factors, blocksizes(A)...),)
    #         _, back_block = Zygote.forward(cholesky, A)
    #         Ā = first(back_block(C̄_block))

    #         # Ensure approximate agreement.
    #         @test Ā_ ≈ Ā
    #     end

    #     cholesky_gradients(randn(rng, size(A)), A_, A)
    #     cholesky_gradients(UpperTriangular(randn(rng, size(A))), A_, A)

    #     # Test `logdet`.
    #     @test logdet(C) ≈ logdet(C_)

    #     # Check that the forwards-pass agrees.
    #     @test Zygote.forward(logdet, C)[1] == logdet(C)

    #     # Check that reverse-pass agrees with dense revese-pass.
    #     @testset "logdet gradients" begin
    #         ȳ = randn(rng)

    #         # Compute adjoint with dense.
    #         _, back_dense = Zygote.forward(logdet, C_)
    #         C̄_ = back_dense(ȳ)

    #         # Compute adjoint with block
    #         _, back_block = Zygote.forward(logdet, C)
    #         C̄ = back_block(ȳ)

    #         # Check that both answers approximately agree.
    #         @test C̄_[1].factors ≈ C̄[1].factors
    #     end

    #     @testset "cholesky backsolving" begin
    #         x = BlockVector([randn(rng, P) for P in Ps])
    #         X = BlockMatrix([randn(rng, P, Q) for P in Ps, Q in Qs])
    #         @test Vector(ldiv!(C, copy(x))) ≈ ldiv!(C_, Vector(x))
    #         @test Matrix(ldiv!(C, copy(X))) ≈ ldiv!(C_, Matrix(X))
    #     end
    # end
end


# descriptor

mutable struct ActivationDesc
    ptr::cudnnActivationDescriptor_t
end

unsafe_free!(ad::ActivationDesc)=cudnnDestroyActivationDescriptor(ad.ptr)

Base.unsafe_convert(::Type{cudnnActivationDescriptor_t}, ad::ActivationDesc)=ad.ptr

function ActivationDesc(mode, coeff, reluNanOpt=CUDNN_NOT_PROPAGATE_NAN)
    ad = Ref{cudnnActivationDescriptor_t}()
    cudnnCreateActivationDescriptor(ad)
    cudnnSetActivationDescriptor(ad[],mode,reluNanOpt,coeff)
    this = ActivationDesc(ad[])
    finalizer(unsafe_free!, this)
    return this
end

# wrappers

function cudnnActivationForward(x::DenseCuArray{T,N}, y::DenseCuArray{T,N}=x;
                                mode=CUDNN_ACTIVATION_RELU, # CUDNN_ACTIVATION_IDENTITY will not work
                                coeff=0.0, reluNanOpt=CUDNN_NOT_PROPAGATE_NAN, alpha=1, beta=0) where {T,N}
    cudnnActivationForward(handle(), ActivationDesc(mode, T(coeff), reluNanOpt),
                           Ref(T(alpha)), TensorDesc(x), x,
                           Ref(T(beta )), TensorDesc(y), y)
    return  y
end

function cudnnActivationBackward(x::DenseCuArray{T,N}, dx::DenseCuArray{T,N}, y::DenseCuArray{T,N}, dy::DenseCuArray{T,N}=dx;
                                 mode=CUDNN_ACTIVATION_RELU, # CUDNN_ACTIVATION_IDENTITY will not work
                                 coeff=0.0, reluNanOpt=CUDNN_NOT_PROPAGATE_NAN, alpha=1, beta=0) where {T,N}
    cudnnActivationBackward(handle(), ActivationDesc(mode, T(coeff), reluNanOpt),
                            Ref(T(alpha)), TensorDesc( y),  y,
                            TensorDesc(dy), dy,
                            TensorDesc( x),  x,
                            Ref(T(beta )), TensorDesc(dx), dx)
    return dx
end


using Plots
using DelimitedFiles
using FileIO
pgfplotsx()
thetas = readdlm("thetas.txt")
p = plot(0:45, thetas, color=:black, ylims = [(0,2.5) (0, 9)], layout = (2,1), label=["\$\\theta_1\$" "\$\\theta_2\$"]);
vspan!(p, [[0,15] [0,15]],
        color = :black, alpha = 0.15,
        labels = ["20 x 24" "20 x 24"])
        #; "41 x 49" "55 x 65"
vspan!(p, [[15,30] [15,30]], color = :white, alpha = 0.15, labels = ["41 x 49" "41 x 49"]);
vspan!(p, [[30,45] [30,45]], color =  :blue, alpha = 0.15, labels = "55 x 65");

p

save("thetas.png", p)

@show thetas


export Analysis

abstract type Options end

"""
    Analysis
Set-up analysis parameters.
### Fields:
* `filename`    - unique filename of the model.
* `analyse`     - R, C, I, G, B stand for ray, coheren, incoherence, geometri and gaussian
* `num_ray`     - number of rays.
* `alpha`       - ray lauch angles between -90 and 90 deg.
* `box`         - analysis box x- range, km and y-depth, m.
### Available Constructors:
* `Source(;frequency::Real = 1000.0f0, height::Real = 80.0f0))`
* `Source(frequency, height)`
"""
struct Analysis{T<:AbstractArray} <: Options
    filename::String
    analyse::String
    option1::String
    option2::String
    num_ray::Int
    alpha::Vec2{T}
    box::Vec2{T}
    step::Real
end

Analysis(;filename = "Model_f1000Hz",
            analyse = "CB",
            option1 = "SVW",
            option2 = "A",
            num_ray = 80,
            alpha = Vec2(-80.0f0,80.0f0),
            box = Vec2(10f0,1000.0f0),
            step= 0) =
    Analysis(filename, analyse, option1, option2, num_ray, alpha, box, step)