using Distributions
using Plots
using Optim

# Calculate dose-response
# flux reaction:
function reaction(conc,Km,Vmax)
    10 .^Vmax.*conc./(10 .^Km .+ conc)
end

Km = 0
Vmax = 1
conc = collect(0:.01:2)

flux = reaction(conc, Km, Vmax)
plot(conc,flux)

# simulate data with 7% std
conc_sim = [.1,.5,1,2]
flux_sim = reaction(conc_sim,  Km, Vmax)
flux_sim = randn(length(flux_sim))*0.07.*flux_sim+flux_sim
std_sim = flux_sim * .07

scatter!(conc_sim,flux_sim,yerror=std_sim,grid=false,linealpha = 0)

# negative log-likelihood for simulated data
function L(conc,flux,std,p)
    X = reaction(conc,p[1],p[2])
    LL = -sum(-length(X)*log(2π)/2 .- length(X).*log.(abs.(std)) .- (sum((flux .- X).^2) ./ (2std.^2)))
end

# maximum likelihood fit
init = [1.0,1.0]        # start parameters Km,Vmax
optimum = optimize(p -> L(conc_sim,flux_sim,std_sim,p),init,Newton())

# Profile Likelihood for one fixed parameter
function L_1(conc,flux,std,pfix,pn,p)
    if pn <2
        X = reaction(conc,pfix,p)
    else
        X = reaction(conc,p,pfix)
    end
    LL = -sum(-length(X)*log(2π)/2 .- length(X).*log.(abs.(std)) .- (sum((flux .- X).^2) ./ (2std.^2)))
end
#init = [2.0]
#pfix = 0.0
#pn = 1
#opt = optimize(p -> L_1(conc_sim,flux_sim,std_sim,pfix,pn,p),init,NewtonTrustRegion())

# Profile likelihood for pfix and pfit
function prof(optimum,whichPar)
    chi2 = Array{Float64}(undef,1)
    pfit = Array{Float64}(undef,1)
    pfix_collect = Array{Float64}(undef,1)
    step = .01
    if whichPar<2
        fix = 2
    else
        fix = 1
    end
    for sig = -1:2:1
        dev = 0
        n=0
        pfix = optimum.minimizer[fix]+sig*step*n
        while (dev<4 && pfix.>-4 && pfix.<4)
            println(dev)
            println(pfix)
            pfix = optimum.minimizer[fix]+sig*step*n
            init = [optimum.minimizer[whichPar]]
            profopt = optimize(p -> L_1(conc_sim,flux_sim,std_sim,pfix,fix,p),init,Newton())
            dev = (profopt.minimum - optimum.minimum)
            push!(chi2,dev)
            push!(pfit,profopt.minimizer[1])
            push!(pfix_collect,pfix)
            n+=1
        end
    end
    deleteat!(chi2,1)
    deleteat!(pfit,1)
    deleteat!(pfix_collect,1)
    return chi2,pfit,pfix_collect
end

chi2km,Kmfit,Vmaxfix = prof(optimum,1)
chi2vmax,Vmaxfit,Kmfix = prof(optimum,2)

# thresholds
cl = [.68,.95]
thresholds = quantile.(Chisq(1),cl)

# Plot profile likelihoods
p1 = scatter(Vmaxfix,chi2km,ylabel="deviance",xlabel="Vmax", legend= :none)
scatter!([Vmaxfix[1] Vmaxfix[1]],[chi2km[1] chi2km[1]], color = :red)
plot!(Vmaxfix,ones(length(Vmaxfix)).*thresholds[1],linestyle=:dash,label="0.68")
plot!(Vmaxfix,ones(length(Vmaxfix)).*thresholds[2],linestyle=:dash,label="0.95")

p2 = scatter(Kmfix,chi2vmax,xlabel="Km")
scatter!([Kmfix[1] Kmfix[1]],[chi2vmax[1] chi2vmax[1]], color = :red, legend= :none)
plot!(Kmfix,ones(length(Kmfix)).*thresholds[1],linestyle=:dash,label="0.68")
plot!(Kmfix,ones(length(Kmfix)).*thresholds[2],linestyle=:dash,label="0.95")
plot(p1,p2)

# plot parameter paths
p3 = plot(Kmfit,Vmaxfix, title = "Km fitted", ylabel="Vmax", xlabel = "Km")
scatter!([Kmfit[1] Kmfit[1]],[Vmaxfix[1] Vmaxfix[1]], color = :red)
p4 = plot(Kmfix,Vmaxfit, title = "Vmax fitted", xlabel="Km")
scatter!([Kmfix[1] Kmfix[1]],[Vmaxfit[1] Vmaxfit[1]], color = :red)
plot(p3,p4, legend = :none)

# simulate data for lower concentrations
conc_sim = [.01,.05,.08,.1]
flux_sim = reaction(conc_sim,  Km, Vmax)
flux_sim = rand(length(flux_sim))*0.07.*flux_sim+flux_sim
std_sim = flux_sim * .07

plot(conc,reaction(conc, Km, Vmax))
scatter!(conc_sim,flux_sim,yerror=std_sim,grid=false,linealpha = 0)

# Max-likelihood fit
init = [1.0,1.0]        # start parameters Km,Vmax
optimum = optimize(p -> L(conc_sim,flux_sim,std_sim,p),init,Newton())

# Compute profiles
chi2kmlow,Kmfitlow,Vmaxfixlow = prof(optimum,1)
chi2vmaxlow,Vmaxfitlow,Kmfixlow = prof(optimum,2)

# Plot profile likelihoods
p1 = scatter(Vmaxfixlow,chi2kmlow, ylabel="deviance",xlabel="Km")
scatter!([Vmaxfixlow[1] Vmaxfixlow[1]],[chi2kmlow[1] chi2kmlow[1]], color = :red, legend= :none)
plot!(Vmaxfixlow,ones(length(Vmaxfixlow)).*thresholds[1],linestyle=:dash,label="0.68")
plot!(Vmaxfixlow,ones(length(Vmaxfixlow)).*thresholds[2],linestyle=:dash,label="0.95")

p2 = scatter(Vmaxfitlow,chi2vmaxlow,xlabel="Vmax")
scatter!([Vmaxfitlow[1] Vmaxfitlow[1]],[chi2vmaxlow[1] chi2vmaxlow[1]], color = :red, legend= :none)
plot!(Vmaxfitlow,ones(length(Vmaxfitlow)).*thresholds[1],linestyle=:dash,label="0.68")
plot!(Vmaxfitlow,ones(length(Vmaxfitlow)).*thresholds[2],linestyle=:dash,label="0.95")
plot(p1,p2)

# same procedure for large concentrations
conc_sim = [5, 6, 8, 10]
flux_sim = reaction(conc_sim,  Km, Vmax)
flux_sim = rand(length(flux_sim))*0.07.*flux_sim+flux_sim
std_sim = flux_sim * .07

plot(conc,reaction(conc, Km, Vmax))
scatter!(conc_sim,flux_sim,yerror=std_sim,grid=false,linealpha = 0)

# Max-likelihood fit
init = [1.0,1.0]        # start parameters Km,Vmax
optimum = optimize(p -> L(conc_sim,flux_sim,std_sim,p),init,Newton())

# Compute profiles
chi2kmhigh,Kmfithigh,Vmaxfixhigh = prof(optimum,1)
chi2vmaxhigh,Vmaxfithigh,Kmfixhigh = prof(optimum,2)

# Plot profile likelihoods
p1 = scatter(Kmfixhigh,chi2vmaxhigh,ylabel="deviance",xlabel="Vmax")
scatter!([Kmfixhigh[1] Kmfithigh[1]],[chi2vmaxhigh[1] chi2vmaxhigh[1]], color = :red, legend= :none)
plot!(Kmfixhigh,ones(length(Kmfixhigh)).*thresholds[1],linestyle=:dash,label="0.68")
plot!(Kmfixhigh,ones(length(Kmfixhigh)).*thresholds[2],linestyle=:dash,label="0.95")

p2 = scatter(Kmfixhigh,chi2vmaxhigh,xlabel="Km")
scatter!([Kmfixhigh[1] Kmfixhigh[1]],[chi2vmaxhigh[1] chi2vmaxhigh[1]], color = :red, legend= :none)
plot!(Kmfixhigh,ones(length(Kmfixhigh)).*thresholds[1],linestyle=:dash,label="0.68")
plot!(Kmfixhigh,ones(length(Kmfixhigh)).*thresholds[2],linestyle=:dash,label="0.95")
plot(p1,p2)

# All profiles for Km for different concentrations
p3 = scatter(Kmfix,chi2vmax,ylabel="deviance",xlabel="Km", label = "wide range")
scatter!(Kmfixlow,chi2vmaxlow, label = "low conc")
scatter!(Kmfixhigh,chi2vmaxhigh, label = "high conc")
scatter!(Kmfixhigh,ones(length(Kmfixhigh)).*thresholds[1],linestyle=:dash,label="0.68")
scatter!(Kmfixhigh,ones(length(Kmfixhigh)).*thresholds[2],linestyle=:dash,label="0.95")

# All profiles for Vmax for different concentrations
p4 = scatter(Vmaxfix,chi2km,ylabel="deviance",xlabel="vmax", label = "wide range")
scatter!(Vmaxfixlow,chi2kmlow, label = "low conc")
scatter!(Vmaxfixhigh,chi2kmhigh, label = "high conc")
scatter!(Vmaxfix,ones(length(Vmaxfix)).*thresholds[1],linestyle=:dash,label="0.68")
scatter!(Vmaxfix,ones(length(Vmaxfix)).*thresholds[2],linestyle=:dash,label="0.95")
plot(p3,p4)


export toOrthoNormal, flatten, activation, segmentplot, rms, modelfit

function toOrthoNormal(Ti)
    local T = deepcopy(Ti)
    U_,S_,V_ = svd(T[1:3,1:3])
    local R = U_*diagm(0=>[1,1,sign(det(U_*V_'))])*V_'
    T[1:3,1:3] = R
    return T
end

const ⟂ = toOrthoNormal

function getD(D,T)
    if D == 3
        return sparse(toeplitz([-1; zeros(T-4)],[-1 3 -3 1 zeros(1,T-4)]))
    elseif D == 2
        return sparse(toeplitz([1; zeros(T-3)],[1 -2 1 zeros(1,T-3)]))
    elseif D == 1
        return sparse(toeplitz([-1; zeros(T-2)],[-1 1 zeros(1,T-2)]))
    end
    error("Can not handle your choice of D: $D")
end

function matrices(model::Union{SimpleLTVModel, KalmanModel}, d::AbstractIdData)
    T = length(d)
    T -= 1
    n = nstates(d)
    m = ninputs(d)
    x = oftype(Matrix, state(d))
    u = oftype(Matrix, input(d))
    FT = eltype(x)
    A = spzeros(FT, T*n, n^2+n*m)
    Is = sparse(FT(1.0)*I,n,n)
    for i = 1:T
        ii = (i-1)*n+1
        ii2 = ii+n-1
        A[ii:ii2,1:n^2] = kron(Is,x[:,i]')
        A[ii:ii2,n^2+1:end] = kron(Is,u[:,i]')
    end
    vec(x[:,2:end]),A
end

function matrices(m::LTVAutoRegressive, d::AbstractIdData)
    y = output(d)
    getARregressor(vec(y),m.na)
end

flatten(A) = reshape(A,prod(size(A)[1:2]),size(A,3))'
flatten(model::LTVStateSpaceModel) = [flatten(model.At) flatten(model.Bt)]
flatten(model::LTVAutoRegressive) = copy(model.θ')
decayfun(iters, reduction) = reduction^(1/iters)

function ABfromk(k,n,m,T)
    At = reshape(k[:,1:n^2]',n,n,T)
    At = permutedims(At, [2,1,3])
    Bt = reshape(k[:,n^2+1:end]',m,n,T)
    Bt = permutedims(Bt, [2,1,3])
    At,Bt
end

@views model2statevec(model,t) = [model.At[:,:,t]  model.Bt[:,:,t]]' |> vec

function model2statevec(model) # dispatch happens in flatten
    k = flatten(model)
    if model.extended
        k = k[1:end-1,:]
    end
    k
end

function statevec2model(::Type{<:LTVStateSpaceModel}, k,n,m,extend)
    At,Bt = ABfromk(k,n,m,size(k,1))
    SimpleLTVModel(At,Bt,extend)
end

function statevec2model(::Type{<:LTVAutoRegressive}, k,n,m,extend)
    LTVAutoRegressive(k',nothing,extend,0)
end

"""
    At,Bt = segments2full(parameters,breakpoints,n,m,T)

Takes a parameter vecor and breakpoints and returns At and Bt with length T
"""
function segments2full(parameters,breakpoints,n,m,T)
    At,Bt = zeros(n,n,T), zeros(n,m,T)
    i = 1
    for t = 1:T
        t ∈ breakpoints && (i+=1)
        At[:,:,t] = reshape(parameters[i][1:n^2],n,n)'
        Bt[:,:,t] = reshape(parameters[i][n^2+1:end],m,n)'
    end
    At,Bt
end

activation(model::LTVStateSpaceModel; kwargs...) = activation(model.At,model.Bt; kwargs...)
function activation(At,Bt; normalize=false)
    diffparams = (diff([flatten(At) flatten(Bt)],dims=1)).^2
    if normalize
        diffparams .-= minimum(diffparams,1)
        diffparams ./= maximum(diffparams,1)
    end
    activation = sqrt.(sum(diffparams,dims=2)[:])
    activation
end

segmentplot(model::AbstractModel, state;  kwargs...) = segmentplot(activation(model),state;kwargs...)
function segmentplot(act, state; filterlength=10, doplot=false, kwargs...)
    plot(act, lab="Activation")
    ds = diff(state)
    ma = mean(act)
    ds[ds.==0] = NaN
    segments = findpeaks(act; filterlength=filterlength, doplot=doplot, kwargs...)
    doplot || scatter!(segments,3ma*ones(segments), lab="Automatic segments", m=(10,:xcross))
    scatter!(2ma*ds, lab="Manual segments", xlabel="Time index", m=(10,:cross))
    segments
end

"""
    A,B,x,u,n,m,N = testdata(T=10000, σ_state_drift=0.001, σ_param_drift=0.001, σ_control=1)

Create an LTVModel with Brownian-walk A[t] and B[t]
"""
function testdata(;T=10000, σ_state_drift=0.001, σ_param_drift=0.001, seed=1, σ_control=1)
    Random.seed!(seed)
    n           = 3
    m           = 2
    N           = n*(n+m)
    A           = zeros(n,n,T)
    B           = zeros(n,m,T)
    x           = zeros(n,T)
    u           = σ_control*randn(m,T)
    U,S,V       = ⟂(randn(n,n)), diagm(0=>0.4rand(n)), ⟂(randn(n,n))
    A[:,:,1]    = U*S*V'
    B[:,:,1]    = 0.5randn(n,m)
    x[:,1]      = 0.1randn(n)

    for t = 1:T-1
        x[:,t+1]   = A[:,:,t]*x[:,t] + B[:,:,t]*u[:,t] + σ_state_drift*randn(n)
        A[:,:,t+1] = A[:,:,t] + σ_param_drift*randn(n,n)
        B[:,:,t+1] = B[:,:,t] + σ_param_drift*randn(n,m)
    end
    d = iddata(x,u,x)
    A,B,d,n,m,N
end

"""
    x,xm,u,n,m = testdata(T)

Create an LTVModel that changes from `A = [0.95 0.1; 0 0.95]` to `A = [0.5 0.05; 0 0.5]`
at T÷2
"""
function testdata(T_)
    Random.seed!(1)

    n,m      = 2,1
    At_      = [0.95 0.1; 0 0.95]
    Bt_      = reshape([0.2; 1],2,1)
    u        = randn(1,T_)
    x        = zeros(n,T_)
    for t = 1:T_-1
        if t == T_÷2
            At_ = [0.5 0.05; 0 0.5]
        end
        x[:,t+1] = At_*x[:,t] + Bt_*u[:,t] + 0.2randn(n)
    end
    xm = x + 0.2randn(size(x));
    x,xm,u,n,m
end

# Optimizers ===================================================================
#
#
# struct RMSpropOptimizer{VecType}
#     α::Float64
#     rmspropfactor::Float64
#     momentum::Float64
#     Θ::VecType
#     dΘs::VecType
#     dΘs2::VecType
# end
#
# function RMSpropOptimizer(Θ, α, rmspropfactor=0.8, momentum=0.1)
#     RMSpropOptimizer(α, rmspropfactor, momentum, Θ, ones(size(Θ)), zeros(size(Θ)))
# end
#
# function (opt::RMSpropOptimizer)(dΘ)
#     opt.dΘs .= opt.rmspropfactor.*opt.dΘs .+ (1-opt.rmspropfactor).*dΘ.^2
#     ΔΘ = -opt.α * dΘ
#     ΔΘ ./= sqrt.(opt.dΘs.+1e-10) # RMSprop
#
#     # ΔΘ = dΘ./sqrt(dΘs+1e-10).*sqrt(dΘs2) # ADAdelta
#     # dΘs2 .= 0.9dΘs2 + 0.1ΔΘ.^2 # ADAdelta
#     opt.dΘs2 .= opt.momentum.*opt.dΘs2 .+ ΔΘ # Momentum + RMSProp
#     opt.Θ .+= opt.dΘs2
# end
#
#
# mutable struct ADAMOptimizer{T, VecType <: AbstractArray}
#     Θ::VecType
#     α::T
#     β1::T
#     β2::T
#     ɛ::T
#     m::VecType
#     v::VecType
#     expdecay::Float64
# end
#
# ADAMOptimizer(Θ::VecType; α::T = 0.005,  β1::T = 0.9, β2::T = 0.999, ɛ::T = 1e-8, m=zeros(size(Θ)), v=zeros(size(Θ)), expdecay=0) where {T,VecType <: AbstractArray} = ADAMOptimizer{T,VecType}(Θ, α,  β1, β2, ɛ, m, v, expdecay)
#
# """
#     (a::ADAMOptimizer{T,VecType})(g::VecType, t::Integer)
# Applies the gradient `g` to the parameters `a.Θ` (mutating) at iteration `t`
# ADAM GD update http://sebastianruder.com/optimizing-gradient-descent/index.html#adam
# """
# function (a::ADAMOptimizer)(g, t::Integer)
#     mul = (1-a.β1)
#     mul2 = (1-a.β2)
#     div  = 1/(1 - a.β1 ^ t)
#     div2 = 1/(1 - a.β2 ^ t)
#     α,β1,β2,ɛ,m,v,Θ,γ = a.α,a.β1,a.β2,a.ɛ,a.m,a.v,a.Θ,a.expdecay
#     γ *= t > 5
#     Base.Threads.@threads for i = eachindex(g)
#         @inbounds m[i] = β1 * m[i] + mul * (g[i] + γ*Θ[i])
#         @inbounds v[i] = β2 * v[i] + mul2 * g[i]^2
#         @inbounds Θ[i] -= α * m[i] * div / (sqrt(v[i] * div2) + ɛ)
#     end
#     Θ
# end
# function (a::ADAMOptimizer)(g, t::Integer)
#     a.m .= a.β1 .* a.m .+ (1-a.β1) .* g
#     a.v .= a.β2 .* a.v .+ (1-a.β2) .* g.^2
#     div  = 1/(1 - a.β1 ^ t)
#     div2 = 1/(1 - a.β2 ^ t)
#     @. a.Θ  -= a.α * a.m * div / (sqrt(a.v * div2) + a.ɛ)
# end


"""

Sample the parent label of each event for given parameters.

"""
function sample_parents(events::Array{<:Number, 1}, bg, kappa::Number, kern::Function)
    eventdifferencelist = event_difference_list(events)
    sample_parents(events, bg, kappa, kern, eventdifferencelist)
end

function sample_parents(events::Array{<:Number, 1}, bg, kappa::Number, kern::Function, eventdifferencelist)
    n = length(eventdifferencelist)
    numBackgroundEvents = 0
    childEventsVector = zeros(Int64, n)
    childShiftTime = zeros(Float64, n)
    parents = Array{Int64, 1}(undef, n)

    for (i, eventdiff) in enumerate(eventdifferencelist)
        parent = sample_parent(events[i], bg, kappa, kern, eventdiff)
        parents[i] = parent

        if parent == 0
            numBackgroundEvents += 1
            continue
         end
        childEventsVector[parent] += 1
        childShiftTime[i] = events[i] - events[parent]

    end
    (parents, numBackgroundEvents, childEventsVector,  childShiftTime[childShiftTime .!= 0])
end

function sample_parent(event::Number, bg, kappa::Number, kern::Function, eventdifferences::Array{<:Number, 1})
    probs = event_parent_prob(event, bg, kappa, kern, eventdifferences)
    StatsBase.sample(StatsBase.Weights(probs)) - 1
end


include("../../../src/plotSol.jl")
include("../../../src/simulate.jl")
include("../../../src/linearize.jl")
include("../../../src/OrientationConversion.jl")
include("SprMassMinReal.jl")
using Revise
# clearconsole()

# To test a Spring Mass system
# Impulse of 1.0 m/s in x, no gravity
m = 1.0; l = 1.0; # Mass and length of bar
I1 = [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]

# Testing Dynamics with Revolute Joint
R1 = RigidBody(m,I1,2)
RbI = InertialFrameAsRB()
x0R1 = [1.5*l;zeros(2);[1;zeros(3)];[1;zeros(2)];zeros(4)]
initialiseRigidBody!(R1,x0R1)

rj1 = [0.0 0.0 0.0][:] # Joint Location in body frame of first body
rj2 = [-l/2 0.0 0.0][:] # Joint location in body frame of second body
k = 10.0
c = 3.0;
restLen = l;
j1 = Joint(RbI,R1,rj1,rj2,type="Spring",k=k, rL = restLen,c = c);
j = [j1];

# Simulation
tEnd = 10.0
tSpan = 0.01
g = [0.0;0.0;0.0];
tSim, solFinal = simulate(tEnd,tSpan,j,g = g)

## Check errors
close("all");
# Check errNorm of quat
plotErrNorm(tSim,solFinal[1].β)
plotPos(tSim, solFinal[1].r)

## Compare with MinReal
rMin, vMin = SprMassMinReal(m,l,k,c,restLen,tInt,tEnd)

# Plotting errors
plotPos(tSim,rMin-solFinal[1].r)
plotVel(tSim,vMin-solFinal[1].v)


using Random
using Distributions
using CSV
using DataFrames

function sample_negative_binomial_limit(s, p; approx_limit = 1000)
    """
    Samples from a NegBin(s, p) distribution. This uses a normal approximation
    when mu is large (i.e. s > approx_limit) to get a 10x runtime improvement.
    """
    X = zero(Int)

    # mean of NegBin(s, p) => this will boil down to N*TP
    μ = s/p - s

    if μ <= approx_limit
        X = rand(NegativeBinomial(s, p))
    else
        σ = sqrt(s*(1-p)/p^2)
        X = ceil(Int, rand(Normal(μ, σ)))
    end

    return X
end

function sample_binomial_limit(n, p; approx_limit = 1000)
    """
    Samples from a Bin(n, p) distribution. This uses a normal approximation
    for np > approx_limit or n(1-p) > approx_limit to acheive a 10x runtime
    improvement.
    """
    X = zero(Int)

    if n*p <= approx_limit || n*(1-p) <= approx_limit
        X = rand(Binomial(n, p))
    else
        μ = n*p
        σ = sqrt(n*p*(1-p))
        X = ceil(Int, rand(Normal(μ, σ)))
    end

    return X
end

function read_in_susceptible_depletion(file_date)
    """
    Read in the posterior drawn susceptible_depletion factors. This will be sorted/sampled
    in the same order as the posterior predictive TP's to ensure we use the appropriate
    posterior draws.
    """

    susceptible_depletion = Vector(
        CSV.read(
            "results/forecasting/sampled_susceptible_depletion_"*file_date*".csv",
            DataFrame,
            drop=[1],
        )[:,1]
    )

    return susceptible_depletion
end


## distributions.jl (c) 2015 David A. van Leeuwen

## conversion to GMM
function GMM(m::MixtureModel{Multivariate,Continuous,FullNormal})
    Σ = eltype(FullCov{eltype(m)})[cholinv(c.Σ.mat) for c in components(m)]
    μ = hcat([c.μ for c in components(m)]...)'
    w = probs(m)
    n, d = size(μ)
    h = [History(@sprintf("Initialization from MixtureModel n=%d, d=%d, kind=full", n, d))]
    GMM(w, μ, Σ, h, 0)
end

function GMM(m::MixtureModel{Multivariate,Continuous,DiagNormal})
    Σ = hcat([c.Σ.diag for c in components(m)]...)'
    μ = hcat([c.μ for c in components(m)]...)'
    w = probs(m)
    n, d = size(μ)
    h = [History(@sprintf("Initialization from MixtureModel n=%d, d=%d, kind=diag", n, d))]
    GMM(w, μ, Σ, h, 0)
end

## conversion to MixtureModel
function Distributions.MixtureModel(gmm::GMM{T}) where {T<:AbstractFloat}
    if gmm.d == 1
        mixtures = [Normal(gmm.μ[i,1], sqrt(gmm.Σ[i,1])) for i=1:gmm.n]
    elseif kind(gmm) == :full
        mixtures = [MvNormal(vec(gmm.μ[i,:]), covar(gmm.Σ[i])) for i=1:gmm.n]
    else
        mixtures = [MvNormal(vec(gmm.μ[i,:]), sqrt.(vec(gmm.Σ[i,:]))) for i=1:gmm.n]
    end
    MixtureModel(mixtures, gmm.w)
end


# Fast mode
fast_mode(X::Vector{T}) where T <: Integer = fast_mode(X, minimum(X):maximum(X))
function fast_mode(X, range)
    m, cm = (0, 0)
    for i in range
       c = count(isequal(i), X)
       if c > cm
        cm = c
        m = i
        end
    end
    return m
end

# Summary stats
function average(V::Vector{A}) where A <: AbstractArray
    N = length(V)
    M = V[1]
    for i in 2:N
       M .+= V[i]
    end
    M ./= N
    return M
end
average(V::Vector{T}) where T <: AbstractFloat = mean(V)
average(V::Vector{T}) where T <: Integer = fast_mode(V)

function average(SV::StructVector{TV}) where TV
    return TV(average.(values(StructArrays.components(SV))))
end

# Posterior expectation
function expectation(P::Posterior)
    return (hyperparameters = average(P.hyperparameters),
            latent = average(P.latent))
end


#=
This page defines the functions:
fwd_prop! - for feeding forward through the network
bwd_prop! - for feeding backward through the network to calculate gradients

Different methods are defined for different layer types and the network as
a whole.
=#

################################################################################
#
# Forward Propagate
#
################################################################################

########################################
# Layer -> OutputLayer
########################################

# LinearOutputLayer does not have an activation function
# Write weighted sum directly to A
function fwd_prop!(ℓ::LinOutputLayer, prev_ℓ::HDL)

    # A⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.A, ℓ.W, prev_ℓ.Ar)
    broadcast!(+, ℓ.A, ℓ.b, ℓ.A)

    # No activation
end

function fwd_prop!(ℓ::LinOutputLayer, prev_ℓ::LDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.A, ℓ.W, prev_ℓ.A)
    broadcast!(+, ℓ.A, ℓ.b, ℓ.A)

    # No activation
end

# LogisticOutputLayer
# fwd_prop! method covered by the general LD method defined below

# Softmax Layer activation function different to standard
# Vector `u` used for the max values of each row
function fwd_prop!(ℓ::SoftmaxOutputLayer, prev_ℓ::HDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.Z, ℓ.W, prev_ℓ.Ar)
    broadcast!(+, ℓ.Z, ℓ.b, ℓ.Z)

    # A⁽²⁾ = f(Z⁽²⁾)
    ℓ.f!(ℓ.A, ℓ.Z, ℓ.u)
end

function fwd_prop!(ℓ::SoftmaxOutputLayer, prev_ℓ::LDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.Z, ℓ.W, prev_ℓ.A)
    broadcast!(+, ℓ.Z, ℓ.b, ℓ.Z)

    # A⁽²⁾ = f(Z⁽²⁾)
    ℓ.f!(ℓ.A, ℓ.Z, ℓ.u)
end

########################################
# Layer -> Low Dimension Layer
########################################

# LD -> LD
function fwd_prop!(ℓ::LDL, prev_ℓ::LDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.Z, ℓ.W, prev_ℓ.A)
    broadcast!(+, ℓ.Z, ℓ.b, ℓ.Z)

    # A⁽²⁾ = f(Z⁽²⁾)
    ℓ.f!(ℓ.A, ℓ.Z)
end

# HD -> LD
function fwd_prop!(ℓ::LDL, prev_ℓ::HDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.Z, ℓ.W, prev_ℓ.Ar)
    broadcast!(+, ℓ.Z, ℓ.b, ℓ.Z)

    # A⁽²⁾ = f(Z⁽²⁾)
    ℓ.f!(ℓ.A, ℓ.Z)
end

########################################
# Layer -> Word Embedding Layer
########################################

# LD -> Word Embedding
function fwd_prop!(ℓ::WordEmbeddingLayer, prev_ℓ::LDL)

    # Embed word features
    embed!(ℓ.A, ℓ.W, prev_ℓ.A)
end

# HD -> Word Embedding
function fwd_prop!(ℓ::WordEmbeddingLayer, prev_ℓ::HDL)

    # Embed word features
    embed!(ℓ.A, ℓ.W, prev_ℓ.Ar)
end

########################################
# Layer -> Sparse Encoder Layer
########################################

# LD -> Sparse Encoder
function fwd_prop!(ℓ::SparseEncoderLayer, prev_ℓ::LDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.Z, ℓ.W, prev_ℓ.A)
    broadcast!(+, ℓ.Z, ℓ.b, ℓ.Z)

    # A⁽²⁾ = f(Z⁽²⁾)
    ℓ.f!(ℓ.A, ℓ.Z)

    # Average of activations - "probability" of each neuron firing
    mean!(ℓ.p̂, ℓ.A)
end

# HD -> Sparse Encoder
function fwd_prop!(ℓ::SparseEncoderLayer, prev_ℓ::HDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ A⁽¹⁾
    A_mul_B!(ℓ.Z, ℓ.W, prev_ℓ.A)
    broadcast!(+, ℓ.Z, ℓ.b, ℓ.Z)

    # A⁽²⁾ = f(Z⁽²⁾)
    ℓ.f!(ℓ.A, ℓ.Z)

    # Average of activations - "probability" of each neuron firing
    mean!(ℓ.p̂, ℓ.A)
end

########################################
# High Dimension Layer -> High Dimension Layer
########################################

# Convolution
function fwd_prop!(ℓ::ConvolutionLayer, prev_ℓ::HDL)

    # Z⁽²⁾ = b⁽²⁾ .+ W⁽²⁾ ∗ A⁽¹⁾
    conv!(ℓ.Z, ℓ.W, prev_ℓ.A)
    bcastbias!(ℓ.Z, ℓ.b)

    # A⁽²⁾ = f(Z⁽²⁾)
    ℓ.f!(ℓ.A, ℓ.Z)
end

# Mean Pooling
function fwd_prop!(ℓ::MeanPoolLayer, prev_ℓ::HDL)
    meanpool!(ℓ.A, prev_ℓ.A)
end

# Max Pooling
function fwd_prop!(ℓ::MaxPoolLayer, prev_ℓ::HDL)
    maxpool!(ℓ.A, prev_ℓ.A, ℓ.maxinds)
end

########################################
# Whole Network
########################################

function fwd_prop!(net::NeuralNet)

    # Feed forward each hidden layer and the output layer
    for (ℓ, prev_ℓ) in zip(net.layers[2:end], net.layers[1:end-1])
        fwd_prop!(ℓ, prev_ℓ)
    end
end

################################################################################
#
# Backward Propagate
#
################################################################################

########################################
# Layer -> OutputLayer
########################################

function bwd_prop!(ℓ::OutputLayer, prev_ℓ::HDL)

    # Delta of final layer
    # Δ⁽ᴸ⁾ = 1/m (A⁽ᴸ⁾ - Y)               (recall: A⁽ᴸ⁾ = H)
    ℓ.δ!(ℓ.Δ, ℓ.A, ℓ.y, ℓ.m)

    # Gradient of bias vector
    # ∇b⁽ᴸ⁾ = sum(Δ⁽ᴸ⁾, 2)
    sum!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrix
    # ∇W⁽ᴸ⁾ = Δ⁽ᴸ⁾ A⁽ᴸ⁻¹⁾ᵀ
    A_mul_Bt!(ℓ.∇W, ℓ.Δ, prev_ℓ.Ar)
end

function bwd_prop!(ℓ::OutputLayer, prev_ℓ::LDL)

    # Delta of final layer
    # Δ⁽ᴸ⁾ = 1/m (A⁽ᴸ⁾ - Y)               (recall: A⁽ᴸ⁾ = H)
    ℓ.δ!(ℓ.Δ, ℓ.A, ℓ.y, ℓ.m)

    # Gradient of bias vector
    # ∇b⁽ᴸ⁾ = sum(Δ⁽ᴸ⁾, 2)
    sum!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrix
    # ∇W⁽ᴸ⁾ = Δ⁽ᴸ⁾ A⁽ᴸ⁻¹⁾ᵀ
    A_mul_Bt!(ℓ.∇W, ℓ.Δ, prev_ℓ.A)
end

########################################
# Low -> Low -> Low
########################################

# Low -> Low -> Low
function bwd_prop!(ℓ::LDL, prev_ℓ::LDL, next_ℓ::LDL)

    # Δ⁽²⁾ = W⁽³⁾ᵀ Δ⁽³⁾ ⊙ f'(Z⁽²⁾)
    # Broken into 2 steps:
    # (a) - Δ⁽²⁾ := W⁽³⁾ᵀ Δ⁽³⁾
    # (b) - Δ⁽²⁾ := Δ⁽²⁾ ⊙ f'(Z⁽²⁾)
    At_mul_B!(ℓ.Δ, next_ℓ.W, next_ℓ.Δ)
    ℓ.∇f!(ℓ.Δ, ℓ.Z)

    # Gradient of bias vector
    # ∇b⁽²⁾ = sum(Δ⁽²⁾, 2)
    sum!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrix
    # ∇W⁽²⁾ = Δ⁽²⁾ A⁽¹⁾
    A_mul_Bt!(ℓ.∇W, ℓ.Δ, prev_ℓ.A)
end

# Low -> Word embedding -> Low
function bwd_prop!(ℓ::WordEmbeddingLayer, prev_ℓ::LDL, next_ℓ::LDL)

    # Δ⁽²⁾ = W⁽³⁾ᵀ Δ⁽³⁾
    At_mul_B!(ℓ.Δ, next_ℓ.W, next_ℓ.Δ)

    # Gradient of weight matrix
    embed_grad!(ℓ.∇W, ℓ.Δ, prev_ℓ.A)
end

# Low -> Sparse -> Low
function bwd_prop!(ℓ::SparseEncoderLayer, prev_ℓ::LDL, next_ℓ::LDL)

    # Δ⁽²⁾ = (W⁽³⁾ᵀ Δ⁽³⁾ .+ spgrad⁽²⁾) ⊙ f'(Z⁽²⁾)
    # Broken into 3 steps:
    # (a) - Δ⁽²⁾ := W⁽³⁾ᵀ Δ⁽³⁾
    # (b) - Δ⁽²⁾ := Δ⁽²⁾ .+ spgrad⁽²⁾
    # (c) - Δ⁽²⁾ := Δ⁽²⁾ ⊙ f'(Z⁽²⁾)
    At_mul_B!(ℓ.Δ, next_ℓ.W, next_ℓ.Δ)
    ℓ.δ!(ℓ.Δ, ℓ.spgrad, ℓ.p̂, ℓ.p, ℓ.β, ℓ.m)
    ℓ.∇f!(ℓ.Δ, ℓ.Z)

    # Gradient of bias vector
    # ∇b⁽²⁾ = sum(Δ⁽²⁾, 2)
    sum!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrix
    # ∇W⁽²⁾ = Δ⁽²⁾ A⁽¹⁾
    A_mul_Bt!(ℓ.∇W, ℓ.Δ, prev_ℓ.A)
end

########################################
# High -> Low -> Low
########################################

# High -> Low -> Low
function bwd_prop!(ℓ::LDL, prev_ℓ::HDL, next_ℓ::LDL)

    # Δ⁽²⁾ = W⁽³⁾ᵀ Δ⁽³⁾ ⊙ f'(Z⁽²⁾)
    # Broken into 2 steps:
    # (a) - Δ⁽²⁾ := W⁽³⁾ᵀ Δ⁽³⁾
    # (b) - Δ⁽²⁾ := Δ⁽²⁾ ⊙ f'(Z⁽²⁾)
    At_mul_B!(ℓ.Δ, next_ℓ.W, next_ℓ.Δ)
    ℓ.∇f!(ℓ.Δ, ℓ.Z)

    # Gradient of bias vector
    # ∇b⁽²⁾ = sum(Δ⁽²⁾, 2)
    sum!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrix
    # ∇W⁽²⁾ = Δ⁽²⁾ A⁽¹⁾
    A_mul_Bt!(ℓ.∇W, ℓ.Δ, prev_ℓ.Ar)
end

# High -> Word embedding -> Low
function bwd_prop!(ℓ::WordEmbeddingLayer, prev_ℓ::HDL, next_ℓ::LDL)

    # Δ⁽²⁾ = W⁽³⁾ᵀ Δ⁽³⁾
    At_mul_B!(ℓ.Δ, next_ℓ.W, next_ℓ.Δ)

    # Gradient of weight matrix
    embed_grad!(ℓ.∇W, ℓ.Δ, prev_ℓ.Ar)
end

# High -> Sparse -> Low
function bwd_prop!(ℓ::SparseEncoderLayer, prev_ℓ::HDL, next_ℓ::LDL)

    # Δ⁽²⁾ = (W⁽³⁾ᵀ Δ⁽³⁾ .+ spgrad⁽²⁾) ⊙ f'(Z⁽²⁾)
    # Broken into 3 steps:
    # (a) - Δ⁽²⁾ := W⁽³⁾ᵀ Δ⁽³⁾
    # (b) - Δ⁽²⁾ := Δ⁽²⁾ .+ spgrad⁽²⁾
    # (c) - Δ⁽²⁾ := Δ⁽²⁾ ⊙ f'(Z⁽²⁾)
    At_mul_B!(ℓ.Δ, next_ℓ.W, next_ℓ.Δ)
    ℓ.δ!(ℓ.Δ, ℓ.spgrad, ℓ.p̂, ℓ.p, ℓ.β, ℓ.m)
    ℓ.∇f!(ℓ.Δ, ℓ.Z)

    # Gradient of bias vector
    # ∇b⁽²⁾ = sum(Δ⁽²⁾, 2)
    sum!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrix
    # ∇W⁽²⁾ = Δ⁽²⁾ A⁽¹⁾
    A_mul_Bt!(ℓ.∇W, ℓ.Δ, prev_ℓ.Ar)
end

########################################
# High -> High -> Low
########################################

# High -> Convolution -> Low
function bwd_prop!(ℓ::ConvolutionLayer, prev_ℓ::HDL, next_ℓ::LDL)

    # Δ⁽²⁾ = W⁽³⁾ᵀ Δ⁽³⁾ ⊙ f'(Z⁽²⁾)
    # Broken into 2 steps:
    # (a) - Δ⁽²⁾ := W⁽³⁾ᵀ Δ⁽³⁾
    # (b) - Δ⁽²⁾ := Δ⁽²⁾ ⊙ f'(Z⁽²⁾)
    At_mul_B!(ℓ.Δr, next_ℓ.W, next_ℓ.Δ)
    ℓ.∇f!(ℓ.Δr, ℓ.Zr)

    # Gradient of bias vector
    # ∇bₖ⁽²⁾ = ∑ Δₖ⁽²⁾
    biasgrad!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrices
    # ∇Wₖ⁽²⁾ = Δₖ⁽²⁾ ∗ A⁽¹⁾
    convgrad!(ℓ.∇W, ℓ.Δ, prev_ℓ.A)

    return
end

# High -> Pool -> Low
function bwd_prop!(ℓ::PoolLayer, prev_ℓ::HDL, next_ℓ::LDL)

    # Δ⁽²⁾ = W⁽³⁾ᵀ Δ⁽³⁾
    At_mul_B!(ℓ.Δr, next_ℓ.W, next_ℓ.Δ)
end

########################################
# High -> High -> High
########################################

# High -> Convolution -> MeanPool
function bwd_prop!(ℓ::ConvolutionLayer, prev_ℓ::HDL, next_ℓ::MeanPoolLayer)

    # Δ⁽²⁾ = upsample( Δ⁽³⁾ ) ⊙ f'(Z⁽²⁾)
    # Broken into 2 steps:
    # (a) - Δ⁽²⁾ := upsample( Δ⁽³⁾ )
    # (b) - Δ⁽²⁾ := Δ⁽²⁾ ⊙ f'(Z⁽²⁾)
    upsample_mean!(ℓ.Δ, next_ℓ.Δ)
    ℓ.∇f!(ℓ.Δr, ℓ.Zr)

    # Gradient of bias vector
    # ∇bₖ⁽²⁾ = ∑ Δₖ⁽²⁾
    biasgrad!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrices
    # ∇Wₖ⁽²⁾ = Δₖ⁽²⁾ ∗ A⁽¹⁾
    convgrad!(ℓ.∇W, ℓ.Δ, prev_ℓ.A)
end

# High -> Convolution -> MaxPool
function bwd_prop!(ℓ::ConvolutionLayer, prev_ℓ::HDL, next_ℓ::MaxPoolLayer)

    # Δ⁽²⁾ = upsample( Δ⁽³⁾ ) ⊙ f'(Z⁽²⁾)
    # Broken into 2 steps:
    # (a) - Δ⁽²⁾ := upsample( Δ⁽³⁾ )
    # (b) - Δ⁽²⁾ := Δ⁽²⁾ ⊙ f'(Z⁽²⁾)
    upsample_max!(ℓ.Δ, next_ℓ.Δ, next_ℓ.maxinds)
    ℓ.∇f!(ℓ.Δr, ℓ.Zr)

    # Gradient of bias vector
    # ∇bₖ⁽²⁾ = ∑ Δₖ⁽²⁾
    biasgrad!(ℓ.∇b, ℓ.Δ)

    # Gradient of weight matrices
    # ∇Wₖ⁽²⁾ = Δₖ⁽²⁾ ∗ A⁽¹⁾
    convgrad!(ℓ.∇W, ℓ.Δ, prev_ℓ.A)
end

# High -> Pool -> High
function bwd_prop!(ℓ::PoolLayer, prev_ℓ::HDL, next_ℓ::ConvolutionLayer)

    # Delta from convolution layer
    # Δₖ⁽²⁾ = W⁽³⁾ ∗ Δₖ⁽³⁾
    conv!(ℓ.Δ, next_ℓ.W, next_ℓ.Δ)
end
# TODO - CHECK THIS !!!

########################################
# Whole Network
########################################

function bwd_prop!(net::NeuralNet)

    # Final layer
    ℓ = net.layers[end]
    prev_ℓ = net.layers[end-1]
    bwd_prop!(ℓ, prev_ℓ)

    # Hidden layers
    for (prev_ℓ, ℓ, next_ℓ) in zip(net.layers[end-2:-1:1], net.layers[end-1:-1:2], net.layers[end:-1:3])
        bwd_prop!(ℓ, prev_ℓ, next_ℓ)
    end

    # Regularisation
    net.reg_grad!(net)
end


using Oceananigans
using Oceananigans.Units
using Oceananigans.BoundaryConditions: fill_halo_regions!, ImpenetrableBoundaryCondition
using Printf
using GLMakie

# Buoyancy profile
bᵢ(x, y, z) = 1e-5 * z
b_particle = bᵢ(0, 0, 0) # particle buoyancy
ν = 1e-4
κ_particle = 1e-3
R_particle = 1.0
Δt = 1e-2

grid = RectilinearGrid(size=128, z=(-4, 4), topology=(Flat, Flat, Bounded))

no_penetration = ImpenetrableBoundaryCondition()
slip_bcs = FieldBoundaryConditions(grid, (Center, Center, Face),
                                   top=no_penetration, bottom=no_penetration)
w_slip = ZFaceField(grid, boundary_conditions=slip_bcs)
sinking = AdvectiveForcing(WENO5(), w=w_slip)

model = NonhydrostaticModel(; grid,
                            tracers = (:b, :P),
                            buoyancy = BuoyancyTracer(),
                            closure = ScalarDiffusivity(κ=(b=0.0, P=κ_particle)),
                            forcing = (; P = sinking))

Pᵢ(x, y, z) = exp(-(z + 1)^2)
set!(model, b=bᵢ, P=Pᵢ)

simulation = Simulation(model; Δt, stop_iteration=0)

b = model.tracers.b
w_slip_op = 2/9 * (b_particle - b) / ν * R_particle^2

function compute_slip_velocity!(sim)
    w_slip .= w_slip_op
    fill_halo_regions!(w_slip)
    return nothing
end

simulation.callbacks[:slip] = Callback(compute_slip_velocity!)

progress(sim) = @info @sprintf("Iter: %d, time: %.2e, max|w_slip|: %.2e",
                               iteration(sim), time(sim), maximum(abs, w_slip))
simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))

z = znodes(Center, grid)
P = interior(model.tracers.P, 1, 1, :)

fig = Figure()
ax = Axis(fig[1, 1],
          xlabel = "Particle concentration",
          ylabel = "z",
          title = "Settling diffusion particle concentration at t=0")
xlims!(ax, -1, 4)

ℓ = lines!(ax, P, z)
display(fig)

function update_plot!(sim)
    ℓ.input_args[1][] = interior(sim.model.tracers.P, 1, 1, :)
    ax.title[] = @sprintf("Particle concentration at t=%.2e", time(sim))
end

#simulation.callbacks[:plot] = Callback(update_plot!, IterationInterval(1000))

record(fig, "settling_diffusion.mp4", 1:100, framerate=24) do nn
    simulation.stop_iteration += 200
    run!(simulation)
    update_plot!(simulation)
end


using Symbolics, Test
using ReferenceTests

@variables t a x(t) y(t)
expr = [a*x - x*y,-3y + x*y]
@test_reference "target_functions/1.stan" Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget())

@test_reference "target_functions/1.c" Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t])

@test_reference "target_functions/1.m" Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget())

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t]) ==
    Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                   lhsname=:internal_var___du,
                                   rhsnames=[:internal_var___u,:internal_var___p,:t])

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget()) ==
      Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget())

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget()) ==
      Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget())

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t]) ==
      Symbolics.build_function(expr,vcat(x,y),[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t])

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget()) ==
      Symbolics.build_function(expr,vcat(x,y),[a],t,target = Symbolics.StanTarget())

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget()) ==
      Symbolics.build_function(expr,vcat(x,y),[a],t,target = Symbolics.MATLABTarget())

# Matrix CTarget test
let
    @variables x[1:4] y[1:4] z[1:4]
    expression = hcat(x,y,z) # Results in [x1 y1 z1; x2 y2 z2; ...]
    variables  = vcat(x,y,z) # Results in [x1, x2, x3, x4, y1, ...]
    cfunc      = build_function(expression, variables; target = Symbolics.CTarget(), expression = Val{true})
    @test_reference "target_functions/matrix.c" cfunc
end

# Scalar CTarget test
let
    @variables x y t
    expression = x + y + t
    cfunc = build_function(expression, [x], [y], t; target = Symbolics.CTarget(), expression = Val{true})
    @test_reference "target_functions/scalar1.c" cfunc
end

# Scalar CTarget test with scalar multiplication and powers
let
    @variables x y a t
    expression = x^2 + y^-1 + sin(a)^3.5 + 2t
    cfunc = build_function(expression, [x, y], [a], t; target = Symbolics.CTarget(), expression = Val{true})

    @test_reference "target_functions/scalar2.c" cfunc
end

# Matrix StanTarget test
let
    @variables x[1:4] y[1:4] z[1:4]
    expression = hcat(x,y,z) # Results in [x1 y1 z1; x2 y2 z2; ...]
    variables  = vcat(x,y,z) # Results in [x1, x2, x3, x4, y1, ...]
    sfunc      = build_function(expression, vcat(x,y), z, []; target = Symbolics.StanTarget(), expression = Val{true})

    @test_reference "target_functions/matrix.stan" sfunc
end

# Scalar StanTarget test
let
    @variables t x(t) y(t) z(t)
    expression = x + y + z
    sfunc = build_function(expression, vcat(x,y), [z], t; target = Symbolics.StanTarget(), expression = Val{true})

    @test_reference "target_functions/scalar1.stan" sfunc
end

# Matrix MATLABTarget test
let
    @variables x[1:4] y[1:4] z[1:4]
    expression = hcat(x,y,z) # Results in [x1 y1 z1; x2 y2 z2; ...]
    variables  = vcat(x,y,z) # Results in [x1, x2, x3, x4, y1, ...]
    mfunc      = build_function(expression, vcat(x,y,z); target = Symbolics.MATLABTarget(), expression = Val{true})

    @test_reference "target_functions/matrix.m" mfunc
end

# Scalar MATLABTarget test
let
    @variables x y z
    expression = x + y + z
    mfunc = build_function(expression, vcat(x,y,z); target = Symbolics.MATLABTarget(), expression = Val{true})

    @test_reference "target_functions/scalar1.m" mfunc
end


using Test
using TestSetExtensions
using LinearAlgebra
using Flux
using Qaintellect

# adapted from https://github.com/FluxML/Zygote.jl/blob/master/test/gradcheck.jl
function ngradient(f, xs::AbstractArray...)
    grads = zero.(xs)
    for (x, Δ) in zip(xs, grads), i in 1:length(x)
        δ = sqrt(eps())
        tmp = x[i]
        x[i] = tmp - δ/2
        y1 = f(xs...)
        x[i] = tmp + δ/2
        y2 = f(xs...)
        x[i] = tmp
        Δ[i] = (y2-y1)/δ
        if eltype(x) <: Complex
            # derivative with respect to imaginary part
            x[i] = tmp - im*δ/2
            y1 = f(xs...)
            x[i] = tmp + im*δ/2
            y2 = f(xs...)
            x[i] = tmp
            Δ[i] += im*(y2-y1)/δ
        end
    end
    return grads
end

@testset ExtendedTestSet "circuit gradients" begin

    # construct parametrized circuit
    N = 4
    rz = RzGate(1.5π)
    ps = PhaseShiftGate(0.3)
    ry = RyGate(√2)
    n = randn(Float64, 3)
    n /= norm(n)
    rg = RotationGate(0.2π, n)
    cgc = [
        circuit_gate(3, HadamardGate()),
        circuit_gate(2, rz, (1, 4)),
        circuit_gate(2, 3, SwapGate()),
        circuit_gate(3, ps),
        circuit_gate(3, rg),
        circuit_gate(1, ry),
    ]
    # measurement operators
    meas = [MeasurementOperator(Matrix{Float64}(I, 2^N, 2^N), Tuple(1:N)), MeasurementOperator(Hermitian(randn(ComplexF64, 2^N, 2^N)), Tuple(1:N))]
    c = Circuit{N}(cgc, meas)

    # input quantum state
    ψ = randn(ComplexF64, 2^N)

    @testset "circuit gate chain gradients" begin
        # fictitious gradients of cost function with respect to output quantum state after applying circuit gates
        Δ = 0.1*randn(ComplexF64, 2^N)

        # Flux will call pullback function with argument Δ
        grads = Flux.gradient(() -> real(dot(Δ, apply(ψ, c.moments))), Flux.Params([rz.θ, ps.ϕ, ry.θ, rg.nθ]))

        # arguments used implicitly via references
        f(args...) = real(dot(Δ, apply(ψ, c.moments)))
        @test all(isapprox.(ngradient(f, rz.θ, ps.ϕ, ry.θ, rg.nθ),
            (grads[rz.θ], grads[ps.ϕ], grads[ry.θ], grads[rg.nθ]), rtol=1e-5, atol=1e-5))
    end

    @testset "circuit gate chain gradients 2" begin
        Δ = 0.1*randn(Float64, 2^N)

        grads = Flux.gradient(() -> dot(Δ, abs2.(apply(ψ, c.moments))), Flux.Params([rz.θ, ps.ϕ, ry.θ, rg.nθ]))

        # arguments used implicitly via references; factor 2 due to convention for Wirtinger derivative with prefactor 1/2
        f(args...) = dot(Δ, abs2.(apply(ψ, c.moments)))
        @test all(isapprox.(ngradient(f, rz.θ, ps.ϕ, ry.θ, rg.nθ),
            (grads[rz.θ], grads[ps.ϕ], grads[ry.θ], grads[rg.nθ]), rtol=1e-5, atol=1e-5))
    end

    @testset "circuit with measurement gradients" begin
        # fictitious gradients of cost function with respect to circuit output
        Δ = [0.3, -1.2]

        grads = Flux.gradient(() -> dot(Δ, apply(ψ, c)), Flux.Params([rz.θ, ps.ϕ, ry.θ, rg.nθ]))

        # arguments used implicitly via references
        f(args...) = dot(Δ, apply(ψ, c))
        @test all(isapprox.(ngradient(f, rz.θ, ps.ϕ, ry.θ, rg.nθ),
            (grads[rz.θ], grads[ps.ϕ], grads[ry.θ], grads[rg.nθ]), rtol=1e-5, atol=1e-5))
    end

    @testset "moments chain gradients for density matrices" begin
        # fictitious gradients of cost function with respect to output density matrix
        Δ = DensityMatrix(0.1*randn(Float64, 256), 4)

        ρ = density_from_statevector(ψ)

        grads = Flux.gradient(() -> dot(Δ.v, apply(ρ, c.moments).v), Flux.Params([rz.θ, ps.ϕ, ry.θ, rg.nθ]))

        # arguments used implicitly via references
        f(args...) = dot(Δ.v, apply(ρ, c.moments).v)
        @test all(isapprox.(ngradient(f, rz.θ, ps.ϕ, ry.θ, rg.nθ),
            (grads[rz.θ], grads[ps.ϕ], grads[ry.θ], grads[rg.nθ]), rtol=1e-5, atol=1e-5))
    end

    @testset "circuit with measurement gradients for density matrices" begin
        # fictitious gradients of cost function with respect to circuit output
        Δ = [0.3, -1.2]
        ρ = density_from_statevector(ψ)

        grads = Flux.gradient(() -> dot(Δ, apply(ρ, c)), Flux.Params([rz.θ, ps.ϕ, ry.θ, rg.nθ]))

        # arguments used implicitly via references
        f(args...) = dot(Δ, apply(ρ, c))
        @test all(isapprox.(ngradient(f, rz.θ, ps.ϕ, ry.θ, rg.nθ),
            (grads[rz.θ], grads[ps.ϕ], grads[ry.θ], grads[rg.nθ]), rtol=1e-5, atol=1e-5))
    end

    @testset "circuit with measurement gradients for parametric density matrices" begin
        Δ = [0.3, -1.2]
        latent_density(θlist) = DensityMatrix(kron([[1-cos(θ), 0, 0, cos(θ)] for θ in θlist]...), length(θlist))
        θopt = π/2 * (1 .+ 0.5*randn(N))

        grads = Flux.gradient(() -> dot(Δ, apply(latent_density(θopt), c)), Flux.Params([θopt, rz.θ, ps.ϕ, ry.θ, rg.nθ]))

        f(args...) = dot(Δ, apply(latent_density(θopt), c))

        @test all(isapprox.(ngradient(f, θopt, rz.θ, ps.ϕ, ry.θ, rg.nθ),
            (grads[θopt], grads[rz.θ], grads[ps.ϕ], grads[ry.θ], grads[rg.nθ]), rtol=1e-5, atol=1e-5))
    end

end


#############################################################################
#Copyright (c) 2016 Charles Le Losq
#
#The MIT License (MIT)
#
#Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the #Software without restriction, including without limitation the rights to use, copy, #modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, #and to permit persons to whom the Software is furnished to do so, subject to the #following conditions:
#
#The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, #INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# functions.jl contains several mathematic functions
#
#
#############################################################################

function poly(p::Vector{Float64},x::Array{Float64})
    segments = zeros(size(x)[1],size(p)[1])
    for i = 1:size(p)[1]
        segments[:,i] = p[i].*x[:].^(i-1)
    end
    return sum(segments,2)
end
"""
For Gaussian peaks in spectra
	gaussiennes(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},x::Array{Float64},style::ASCIIString = "None")
"""

# From https://rosettacode.org/wiki/Polynomial_regression#Julia
function polyfit(x::Array{Float64}, y::Array{Float64}, n::Int64)
  A = [ float(x[i])^p for i = 1:length(x), p = 0:n ]
  return A \ y
end

function gaussiennes(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},x::Array{Float64};style::ASCIIString = "None")
    segments = zeros(size(x)[1],size(amplitude)[1])
    if style == "None"
        for i = 1:size(amplitude)[1]
            segments[:,i] = amplitude[i] .*exp(-log(2) .* ((x[:,1]-centre[i])./hwhm[i]).^2)
        end
    elseif style == "poly"
        for i = 1:size(amplitude)[1]
            segments[:,i] = poly(squeeze(amplitude[i,:],1),x[:,2]) .*exp(-log(2) .* ((x[:,1]-(poly(squeeze(centre[i,:],1),x[:,2])))./poly(squeeze(hwhm[i,:],1),x[:,2])).^2)
        end
    else
        error("Not implemented, see documentation")
    end
    return sum(segments,2), segments
end

function lorentziennes(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},x::Array{Float64};style::ASCIIString = "None")
    segments = zeros(size(x)[1],size(amplitude)[1])
    if style == "None"
        for i = 1:size(amplitude)[1]
            segments[:,i] = amplitude[i] ./ (1 + ((x[:,1]-centre[i])./hwhm[i]).^2)
        end
    elseif style == "poly"
        for i = 1:size(amplitude)[1]
            segments[:,i] = poly(squeeze(amplitude[i,:],1),x[:,2]) ./ (1 + ((x[:,1]-(poly(squeeze(centre[i,:],1),x[:,2])))./poly(squeeze(hwhm[i,:],1),x[:,2])).^2)
        end
    else
        error("Not implemented, see documentation")
    end
    return sum(segments,2), segments
end

function pearson7(a1::Array{Float64},a2::Array{Float64},a3::Array{Float64},a4::Array{Float64},x::Array{Float64};style::ASCIIString = "None")
    segments = zeros(size(x)[1],size(a1)[1])
    if style == "None"
        for i = 1:size(a1)[1]
            segments[:,i] = a1[i] ./ (1 + ((x[:,1]-a2[i])./a3[i]).^2 .* (2.0.^(1./a4[i]) - 1.0))
        end
    elseif style == "poly"
        for i = 1:size(a1)[1]
            segments[:,i] = poly(squeeze(a1[i,:],1),x[:,2]) ./ (1 + ((x[:,1]-(poly(squeeze(a2[i,:],1),x[:,2])))./poly(squeeze(a3[i,:],1),x[:,2])).^2 .* (2.0.^(1./squeeze(a4[i,:],1)) - 1.0))
        end
    else
        error("Not implemented, see documentation")
    end
    return sum(segments,2), segments
end

function pseudovoigts(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},ps_fraction::Array{Float64},x::Array{Float64};style::ASCIIString = "None")
    segments = zeros(size(x)[1],size(amplitude)[1])
	test1 = find(ps_fraction .<0.0)
	test2 = find(ps_fraction .>1.0)
    if size(test1)[1] != 0 || size(test2)[1] != 0
		error("ps_fraction should be comprised between 0 and 1")
	end
	if style == "None"
		for i = 1:size(amplitude)[1]
			segments[:,i] =  ps_fraction[i].* ( amplitude[i] .*exp(-log(2) .* ((x[:,1]-centre[i])./hwhm[i]).^2) )   + (1- ps_fraction) .* amplitude[i] ./ (1 + ((x[:,1]-centre[i])./hwhm[i]).^2)
        end
    elseif style == "poly"
        for i = 1:size(amplitude)[1]
			segments[:,i] =  squeeze(ps_fraction[i,:]) .* (poly(squeeze(amplitude[i,:],1),x[:,2]) ./ (1 + ((x[:,1]-(poly(squeeze(centre[i,:],1),x[:,2])))./poly(squeeze(hwhm[i,:],1),x[:,2])).^2))     +      (1- squeeze(ps_fraction[i,:])) .* poly(squeeze(amplitude[i,:],1),x[:,2]) .*exp(-log(2) .* ((x[:,1]-(poly(squeeze(centre[i,:],1),x[:,2])))./poly(squeeze(hwhm[i,:],1),x[:,2])).^2)
        end
    else
        error("Not implemented, see documentation")
	end
    return sum(segments,2), segments
end

#The real normal distribution / gaussian function
function normal_dist(nd_amplitudes::Array{Float64},nd_centres::Array{Float64},nd_sigmas::Array{Float64},x::Array{Float64})
    segments = zeros(size(x)[1],size(nd_amplitudes)[1])
    for i = 1:size(nd_amplitudes)[1]
        segments[:,i] = nd_amplitudes[i]./(nd_sigmas[i].*sqrt(2*pi))  .*  exp(- (x[:,1]-nd_centres[i]).^2 ./ (2.*nd_sigmas[i].^2))
    end
    return sum(segments,2), segments
end

# to correct for shifts in X between two spectra, for the inversion
function xshift_inversion(data::Array{Float64},p::Array{Float64})
    xaxis = data[:,1]
    shifted1 = data[:,2]

    spl = Spline1D(xaxis-p[1], shifted1.*p[2] + shifted1.^2.*p[3])
    y = evaluate(spl,xaxis)
end

# to correct a spectrum for a p shift in X
function xshift_direct(original_x::Array{Float64}, original_y::Array{Float64}, p::Float64)
    spl = Spline1D(original_x-p, original_y)
    corrected_y = evaluate(spl,original_x)
	return original_x, corrected_y, p
end

# to correct a shift between two spectra using a reference peak
function xshift_correction(full_x::Array{Float64}, full_shifted_y::Array{Float64}, ref_x::Array{Float64}, ref_y::Array{Float64},shifted_y::Array{Float64})
	fit = curve_fit(xshift_inversion, [ref_x shifted_y], ref_y, [1.0, 1.0,1.0])
	parameters = fit.param
	return xshift_direct(full_x, full_shifted_y, parameters[1])
end


using Gnuplot, LinearAlgebra
using ParameterizedFunctions, DifferentialEquations
g = @ode_def begin
  dx = σ*(y - x)
  dy = x*(ρ - z) - y
  dz = x*y - β*z
end σ ρ β
u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 80.0)
p = [10.0,28.0,8/3]
prob = ODEProblem(g, u0, tspan, p)
sol = solve(prob, Tsit5(), saveat = 0.01)
x, y, z = sol[1,:], sol[2,:], sol[3,:]
tempo = sol.t
# plot
@gsp "set xyplane at -3" "set auto fix" "set grid" :-
@gsp :- x y z tempo "w l notit lc palette" #palette(:plasma)
@gsp :- "set title 'Lorenz attractor'" "set cblabel 'time'"
@gsp :- xlab = "x" ylab = "y" zlab = "z"
save(term="pngcairo size 1200,800", output="lorenz.png") # hide


module mcmc

using LinearAlgebra
using Flux
using Flux: Tracker, params

function mala(logdensity, h, M, iters, init_proposal_mean; δ_trunc=Inf)

    function log_llh_and_proposal(z)
        par_z = param(z)
        logp_u = logdensity(par_z)
        Tracker.back!(logp_u)
        grad = Tracker.grad(par_z)

        #  ----- Truncated MALA (T-MALA) ----
        #  ----- (aka norm clipping!) -------
        if δ_trunc < Inf
            grad_norm = sum(x -> x^2, grad)
            if grad_norm > δ_trunc
                grad *= δ_trunc / grad_norm
            end
        end
        return Tracker.data(logp_u), z + 0.5*h*M*grad, grad   # prev (unnorm) logp, gradient step, gradient
    end

    # set up
    d = length(init_proposal_mean)
    ztrace = Array{Float64, 2}(undef, d, iters)
    cholΣ = Matrix(cholesky(h*M).L)   # `Matrix` keeps Flux happy for the time being. For d < 20, it really makes no difference.

    function gauss_logpdf(z)
        u = cholΣ \ z
        exponent = -0.5*sum(x->x*x, u)
#         lognormconst = -0.5*logdet(2π * h * M)
        out = exponent #+ lognormconst
        return out
    end
    gauss_logpdf(z, μ) = gauss_logpdf(z - μ)

    # initialise
    accepted = 0
    z = init_proposal_mean
    cur_u_logp = -Inf
    proposal_mean = init_proposal_mean
    prev_proposal_mean = init_proposal_mean

    # Markov Chain loop
    for i = 1:iters
        # transition r.v.
        ϵ_prop = randn(d)
        prop_transition_logp = -0.5*sum(x->x*x, ϵ_prop)

        # deterministic proposal
        z_prop = proposal_mean + cholΣ * ϵ_prop

        # calculate logdensity(z_prop) and accompanying new proposal mean
        prop_u_logp, ν_proposal_mean, _ = log_llh_and_proposal(z_prop)

        # metropolis-hastings
        r = prop_u_logp - cur_u_logp +
                gauss_logpdf(z, ν_proposal_mean) -
                prop_transition_logp   |> exp

        # if <accept>
        if rand() < r
            z = z_prop                                                       # update z
            prev_proposal_mean = proposal_mean                               # store prev proposal mean
            cur_u_logp, proposal_mean = prop_u_logp, ν_proposal_mean
            accepted += 1
        end

        ztrace[:,i] = z      # store z regardless
    end

    return ztrace, accepted/iters
end

end


#=
   follows the algorithms as given in
   "Some Highly Accurate Basic Linear Algebra Subroutines"
   by Yuhang Zhao (c) 2010
   McMaster University, Hamilton, Canada

   except for Round3, where I have replaced
     a call to frexp with reinterpret logic

   comments are as they appear in the thesis
=#

rc = Ref{Int}(0) # indicates if a recursive call of iFastSumAlgorithm occurs

function iFastSum{T<:Real}(x::Array{T,1})
    global rc;
    rc[] = 0
    n = length(x)
    xs = Array{T,1}(n) # iFastSumAlgorithm is destructive
    copy!(xs, x)       # thanks to Kristoffer Carlsson
    iFastSumAlgorithm(xs, n)
end

@inline function AddTwo{T<:Real}(a::T, b::T)
    x = a+b
    z = x-a
    y = (a-(x-z))+(b-z)
    (x,y)
end

@inline function Round3{T<:Real}(s0::T, s1::T, s2::T)
    # s,e = frexp(s1)
    # if (s!=0.5 || (sign(s1)!=sign(s2)))
    if (((reinterpret(UInt64,s1) & 0x000fffffffffffff) != zero(UInt64)) || (sign(s1) != sign(s2)))
        s0
    else
        1.1*s1 + s0
    end
end

# output: the correctly rounded sum of x
# for initial call, n is length of x
function iFastSumAlgorithm{T<:Real}(x::Array{T,1},n::Int)
    global rc
    s = zero(T); loop = 1; # loop counts the number of loops

    for i in 1:n # accumlate first approximation
        s,x[i] = AddTwo(s, x[i])
    end

    # main loop
    while true # loop forever
        # (1)
        count=1; st=zero(T); loop=loop+1; sm=zero(T)
        # count points to the next position in x to store the local error
        # st is the temporary sum
        # (2)
        for i in 1:n
            st,x[count] = AddTwo(st,x[i]) # x[count] is local error
            if x[count] != zero(T)
                count += 1
                sm = max(sm,abs(st))
            end
        end
        # (3)
        em = (count-1)*eps(sm)*0.5
        # each local error <= HalfUlp(sm)
        # em is the weak upper bound on magnitude of the sum of the errors
        # (4)
        s, st = AddTwo(s,st)
        st = x[count]
        n = count
        # (5)
        if (em==zero(T) || em < eps(s)*0.5)
            if rc[] > 0
                return s # return s if it is a recursive call
            end
            w1, e1 = AddTwo(st, em)
            w2, e2 = AddTwo(st, -em)
            if ((w1+s != s) | (w2+s != s)) || (Round3(s,w1,e1) != s) || (Round3(s,w2,e2) != s)
                rc[]=1
                s1 = iFastSum(x,n) # first recursive call
                s,s1 = AddTwo(s,s1)
                s2 = iFastSum(x,n) # second recursive call
                rc[] = 0
                s = Round3(s,s1,s2)
             end
             return s
        end
    end
end


struct PlumeEmulator{T <: Real} <: AbstractMeltRate
    α :: T         #calibration coefficient
    λ1 :: T        #liquidus slope
    λ2 :: T        #liquidus intercept
    λ3 :: T        #liquidus pressure coefficient
    E0 :: T        #entrainment coefficient
    Cd :: T        #drag coeffecient
    Γ_TS :: T      #combined Stanton number
    L :: T         #latent heat of fusion for ice
    c :: T         #specific heat capacity of ocean water
    βs :: T        #haline contraction coefficient
    βt :: T        #thermal expansion coefficient
    g :: T         #gravity
    ρi :: T        #ice density
    ρw :: T        #water density
    Ta :: Function #Function specifying ambient temperature
    Sa :: Function #Function specifying ambient salinity
end

"""
    PlumeEmulator(; <kwargs>)

Construct a plume emulator to prescribe the melt rate in WAVI

Keyword arguments
=================
- α: calibration coefficient: constant prefactor used to calibrate melt rate. Defaults to 0.73 according to Favier 2019 (10.5194/gmd-12-2255-2019)
- λ1: liquidus slope (units ∘C)
- λ2: liquidus intercept (units ∘C)
- λ3: liquidus pressure coefficient (units K/m)
- E0: entrainment coefficient (dimensionless)
- Cd: drag coefficient (dimensionless)
- Γ_TS: combined Stanton number (dimensionless)
- L: latent heat of fusion for ice (units: J/kg)
- c: specific heat capacity of water (units J/kg/K)
- βs: haline contraction coefficient (units psu^(-1))
- βt: thermal expansion coefficient (units C^{-1})
- g: gravitational acceleration (units m / s^2)
- ρi: ice density (units kg/m^3)
- ρw: water density (units kg/m^3)
- Ta: ambient temperature function, defaults to the ISOMIP warm0 scenario. Ta must be a function of a single variable (depth) [time dependence not yet implemented in WAVI.jl]
- Sa: ambient salnity function, defaults to the ISOMIP warm0 scenario.  Sa must be a function of a single variable (depth) [time dependence not yet implemented in WAVI.jl]
"""
function PlumeEmulator(;
                        α = 0.73,
                        λ1 = -5.73e-2,
                        λ2 = 8.32e-4,
                        λ3 = 7.61e-4,
                        E0 = 3.6e-2,
                        Cd = 2.5e-3,
                        Γ_TS = 0.0118,
                        L = 3.35e5,
                        c = 3.974e3,
                        βs = 7.86e-4,
                        βt = 3.87e-5,
                        g = 9.81,
                        ρi = 918.0,
                        ρw = 1028.0,
                        Ta = isomip_warm0_temp,
                        Sa = isomip_warm0_salinity)

    #check that Ta and Sa accept a single argument
    try Ta(0.0)
    catch
        ArgumentError("Ambient temperature function Ta must accept a single argument")
    end
    try Sa(0.0)
    catch
        ArgumentError("Ambient salinity function Sa must accept a single argument")
    end

    return PlumeEmulator(α, λ1, λ2, λ3,E0, Cd, Γ_TS, L, c, βs, βt,g, ρi, ρw, Ta, Sa)

end

function update_melt_rate!(plume_emulator::PlumeEmulator, fields, grid)
    @unpack basal_melt, h, b, grounded_fraction = fields.gh #get the ice thickness and grounded fraction
    set_plume_emulator_melt_rate!(basal_melt,
                                h,
                                grounded_fraction,
                                b,
                                grid.dx,
                                grid.dy,
                                plume_emulator)
    return nothing
end

"""
    set_plume_emulator_melt_rate(melt, h, grounded_fraction)

Set the melt rate in according to Lazeroms 2018. All plume parameters passed in the plume model emulator (pme)
"""
function set_plume_emulator_melt_rate!(melt, h, grounded_fraction,bathy, dx, dy,pme)
    nx, ny = size(h)
    zbf = @. -(pme.ρi/pme.ρw)*h.*(1-grounded_fraction) .+ grounded_fraction*bathy #ice draft if floating everywhere
    ∂zb∂x, ∂zb∂y = get_slope(zbf,dx, dy) #returns the partial derivates of base in both directions

    #loop over each grid point
    for i = 1:nx
        for j = 1:ny
            melt[i,j] =  pme_pointwise_melt(grounded_fraction,zbf, ∂zb∂x,∂zb∂y, pme, i, j)
        end
    end

    return melt
end

"""
    pme_pointwise_melt(grounded_fraction,zbf, ∂zb∂x,∂zb∂y, pme, i, j)

Return the pme prediction of melt rate at single grid point (i,j).
____________
Arguments:
- grounded_fraction: matrix of grounded fraction at grid points
- zbf: ice shelf basal elevation
- ∂zb∂x: slope of ice shelf draft at grid pts in x
- ∂zb∂y: slope of ice shelf draft at grid pts in y
- pme: plume model emulator structure. Holds all parameters.
- i,j: specific grid pt
"""
function pme_pointwise_melt(grounded_fraction,zbf, ∂zb∂x,∂zb∂y, pme, i, j)
    if grounded_fraction[i,j] == 0. #i.e. if its floating
        #initialize directions
        dirs = get_plume_directions() #array of directions
        nd = size(dirs)[1] #number of search directions
        norm_dirs = dirs ./ norm.(eachrow(dirs)) #normalized directions

        #initialize slope and gl depth
        zgl = zeros(1, nd) #grounding line depth
        s = zeros(1, nd) #basal slopes
        iscount = trues(1, nd) #do these points count or not?
        iscount_s = trues(1,nd)
        iscount_zgl = trues(1,nd)

        #loop over directions
        for k = 1:size(dirs)[1]
            #compute the slope in this direction
            s[k] = -dot(norm_dirs[k,:], [∂zb∂x[i,j] ∂zb∂y[i,j]]) #note negative sign
            if s[k] <= 0
                iscount[k] = false
                iscount_s[k] = false
            end #[this syntax count be tightened]

            #compute gl depth in this direction
            zgl[k] = get_pme_gl_depth(dirs[k,:], grounded_fraction, zbf, i, j)
            if isnan(zgl[k]) #no gl point found in this direction
                iscount[k] = false
                iscount_zgl[k] = false
            elseif zgl[k] >= zbf[i,j] #if gronding line higher than shelf point
                iscount[k] = false
                iscount_zgl[k] = false
            end
        end

        #compute effective gl depth and angle
        N = sum(iscount)
        zgl_eff = sum(zgl[iscount])/N;
        α = atan(sum(s[iscount])/N);

        #update the melt rate
        if N > 0
            Sa_gl = pme.Sa(zgl_eff)
            Sa_local = pme.Sa(zbf[i,j])
            Ta_local = pme.Ta(zbf[i,j])
            Tf0 = pme.λ1*Sa_gl + pme.λ2 + pme.λ3 * zgl_eff
            M0 = lazeroms_melt_rate_prefactor(pme.βs, pme.βt, pme.g, pme.Cd, pme.Γ_TS, pme.L, pme.c, pme.λ1, pme.λ3, pme.E0, α, Sa_local, Ta_local)
            x = lazeroms_dimensionless_distance(pme.λ3, zbf[i,j], zgl_eff, Ta_local, Tf0)
            melt = M0 * pme.α * (Ta_local - Tf0)^2 * lazeroms_dimensionless_melt(x) * 365.25* 24 * 60^2
        else
            melt = 0.
        end
    else
        melt = 0.
    end
    return melt #, zgl, s, iscount, iscount_s, iscount_zgl
end

"""
    get_slope(h, dx, dy)

Returns the partial derivatives of h with respect to x and y as matrices. Uses centred FD at interior points and one sided FD at edge points (both second order)

"""
function get_slope(h, dx, dy)
    ∂h∂x = zeros(size(h))
    ∂h∂x[2:end-1,:] = (-h[1:end-2,:] + h[3:end,:])/ 2/dx
    ∂h∂x[1,:] = (-3/2 *h[1,:] + 2*h[2,:] - 1/2 * h[3,:] )/dx
    ∂h∂x[end,:] = (1/2 * h[end-2, :] - 2 *h[end-1,:] + 3/2 * h[end,:])/dx

    ∂h∂y = zeros(size(h))
    ∂h∂y[:,2:end-1]= (-h[:,1:end-2] + h[:,3:end])/2 /dy
    ∂h∂y[:,1] = (-3/2 * h[:,1] + 2 * h[:,2] - 1/2 * h[:,3])/dy
    ∂h∂y[:,end] = (1/2 * h[:,end-2] - 2*h[:,end-1] + 3/2 * h[:,end] )/dy;

    return ∂h∂x, ∂h∂y
end

function get_plume_directions()
    directions = [0 1; #N
                  1 2; #NNE
                  1 1; #NE
                  2 1; #ENE
                  1 0; #E
                  2 -1; #ESE
                  1 -1; #SE
                  1 -2; #SSE
                  0 -1; #S
                  -1 -2; #SSW
                  -1 -1; #SW
                  -2 -1; #WSW
                  -1 0; #W
                  -2 1; #WNW
                  -1 1; #NW
                  -1 2] #NNW
    return directions
end

function lazeroms_melt_rate_prefactor(βs,βt, g, Cd, Γ_TS, L, c,λ1, λ3,  E0, α, Sa, Ta)
    cp1 = (L/c)/sqrt(Cd)/Γ_TS * βt/βs / Sa
    cp2 = -λ1 * βt/βs
    ct = cp2/cp1
    M0 = sqrt(βs * Sa * g) / sqrt(λ3 * (L/c)^3) *
         ((1 -  cp1 * sqrt(Cd)*Γ_TS)/(Cd + E0 * sin(α)))^(1/2) *
         (sqrt(Cd) * Γ_TS * E0 * sin(α)/ (sqrt(Cd) * Γ_TS + ct + E0 * sin(α)))^(3/2)
    return M0
end

"""
    lazeroms_dimensionless_distance(λ3, zb, zgl, Ta, Tf0)

The dimensionless distance in the Lazeroms plume emulator formulation (see equation 28b in Lazeroms et al. 2019)
Note that we include only the first term in the expansion (i.e. no contribution from term C_ϵ)
"""
lazeroms_dimensionless_distance(λ3, zb, zgl, Ta, Tf0) = λ3 * (zb - zgl) / (Ta - Tf0)

"""
    lazeroms_dimensionless_melt(x)

Dimensionless melt function, equation 26 in Lazeroms et al. 2019
"""
lazeroms_dimensionless_melt(x) = (3*(1-x)^(4/3) - 1)*(1- (1-x)^(4/3))^(1/2) / 2 / sqrt(2)

"""
    get_pme_gl_depth(direction, grounded_fraction, zbf, i,j)

Get the effective grounding line depth at shelf point i,j in the specified direction.
Note that this function does not yet apply the interpolation described in Lazeroms et al. 2018
"""
function get_pme_gl_depth(direction, grounded_fraction, zbf, i,j)
    nx, ny = size(grounded_fraction)

    #initialize
    zgl = NaN
    isterm_gl = false #have we found a grounded point?
    isterm_domain = false #have we gone outsdie of the domain?
    p = i
    q = j

    #workflow: (1) check we aren't grounded (before step as may be grounded to begin with)
    #          (2) make a step
    #          (3) check we haven't gone out of the domain
    #          (4) loop if still in domain and not grounded
    while ~(isterm_domain) && ~(isterm_gl)
        #check whether search point is grounded, trigger terminated if so
        grounded_fraction[p,q] > 1 - 1e-3 ?  isterm_gl = true : isterm_gl = false
        if isterm_gl #we've found a grounded point
            zgl = zbf[p,q]
        end

        #update search point
        p = p + direction[1]
        q = q + direction[2]

        #have we gone out of the domain
        isin_domain(p,q,nx,ny) ? isterm_domain = false : isterm_domain = true #if the first step is out of the domain, terminate
    end #end while loop
    return zgl
end

isin_domain(p,q,nx,ny) = ((p <= nx) && (p>0) && (q >0) && (q <= ny)) ? true : false


# # Problems and the system

#md # ```@meta
#md # CurrentModule = ImmersedLayers
#md # ```

#=
In specific problems that we wish to solve with immersed layers, there may
be other data and operators that we would like to cache. We do this with
an *extra cache*, which the user can define, along with a problem type associated
with this cache. The basic cache and the extra cache are generated and associated
together in a *system*.

There are a few basic ingredients to do this:
* **Create a problem type**, using the macro `@ilmproblem`. This mostly just serves as a means of dispatching correctly, but also holds the grid and bodies so they can be passed along when the caches are constructed.
* **Create an extra cache type**, making it a subtype of `AbstractExtraILMCache`. This can hold pretty much anything you want it to.
* **Extend the function `prob_cache(prob,base_cache)`** to serve as a constructor for your extra cache, when your problem type is passed in.

Optionally, you can also extend the function `solve` in order to perform the steps
of your algorithm. However, generically, you can just use pass in the system structure,
which holds the basic ILM cache and your extra cache, into any function.
=#

#=
## Example of problem and system use
We will demonstrate the use of problems and systems with the example
given in [A Dirichlet Poisson problem](@ref). Here, we will
assemble the various additional data structures and operators used to
solve this problem into an extra cache. We will also create a problem
type called `DirichletPoissonProblem`, which we make a subtype of
`AbstractScalarILMProblem`.
=#

using ImmersedLayers
using Plots
using UnPack

#=
### Create your problem type
For this job, we have a macro
=#
@ilmproblem DirichletPoisson scalar

#=
This generates a type `DirichletPoissonProblem` for dispatch, and an associated
constructor that we'll use later.
=#

#=
### Create your extra cache
Here, we'd like this extra cache to hold the Schur complement and the
filtering matrices, as well as some cache variables.
=#
struct DirichletPoissonCache{SMT,CMT,ST,FT} <: ImmersedLayers.AbstractExtraILMCache
   S :: SMT
   C :: CMT
   fb :: ST
   fstar :: FT
end

#=
### Extend the `prob_cache` function
We need this to construct our extra cache
=#
function ImmersedLayers.prob_cache(prob::DirichletPoissonProblem,base_cache::BasicILMCache)
    S = create_RTLinvR(base_cache)
    C = create_surface_filter(base_cache)
    fb = zeros_surface(base_cache)
    fstar = zeros_grid(base_cache)
    DirichletPoissonCache(S,C,fb,fstar)
end

#=
### Extend the `solve` function
Here, we actually do the work of the algorithm, making use of all of the
operators and data structures that we have cached for efficiency.
The example below takes in some surface Dirichlet data `fbplus`,
and returns the solutions `f` and `s` (filtered).
=#
function ImmersedLayers.solve(fbplus,prob::DirichletPoissonProblem,sys::ILMSystem)
    @unpack extra_cache, base_cache = sys
    @unpack S, C, fb, fstar = extra_cache

    f = zeros_grid(base_cache)
    s = zeros_surface(base_cache)

    surface_divergence!(fstar,fbplus,base_cache)
    fb .= 0.5*fbplus

    inverse_laplacian!(fstar,base_cache)

    interpolate!(s,fstar,base_cache)
    s .= fb - s
    s .= -(S\s);

    regularize!(f,s,base_cache)
    inverse_laplacian!(f,base_cache)
    f .+= fstar;

    return f, C^6*s
end

#=
### Set up the grid, shape, and cache
We do this just as we did in [Immersed layer caches](@ref), but
now we don't create a cache, since it will be done internally.
=#
Δx = 0.01
Lx = 4.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx)
RadC = 1.0
Δs = 1.4*cellsize(g)
body = Circle(RadC,Δs)

#=
### Do the work
We do this in three steps:
- Create the problem instance. This is where we use the constructor generated by `@ilmproblem`
- Call `__init` to create the caches, assembled into a system
- Call `solve` to solve the problem.
Also, note that pretty much any function that accepts `base_cache`
as an argument also accepts `sys`.
=#
prob = DirichletPoissonProblem(g,body,scaling=GridScaling)
sys = ImmersedLayers.__init(prob)

pts = points(sys)
f, s = solve(pts.u,prob,sys)

plot(f,sys)
#-
plot(s)

#=
In the previous example, we supplied the surface data as an argument directly to
the `solve` function. However, sometimes we need to supply surface data with a little
more complexity (as in time-varying cases). For this purpose, we can create
functions that supply the boundary data when called, e.g.,
=#
get_fbplus(base_cache) = points(base_cache).u
get_fbminus(base_cache) = zeros_surface(base_cache)

#=
We can pass these along to the problem specification by various means.
It's at the user's discretion how to do it, and how to make use of them
in the solution. Here's an example, using a `Dict`
=#
bcdict = Dict("fbplus"=>get_fbplus,"fbminus"=>get_fbminus)

#=
We can also pass along physical parameters and forcing functions.
Let's use the latter feature in order to set some point sources at
a few locations. For this, we create a function that will take in
the base cache and the right-hand side data vector, and populate
this vector with the point sources. We use the `Regularize` function
in the `CartesianGrids.jl` package to immerse the point sources
into the grid. In the example shown here, two points: one at (-1.0,1.5)
and the other at (1.0,1.5), with respective strengths -1 and 1,
are placed.
=#
function rhs_func!(rhs,base_cache)
    X = VectorData([-1.0,1.0],[1.5,1.5])
    str = ScalarData([-1.0,1.0])
    reg = Regularize(X,cellsize(base_cache.g),I0=origin(base_cache.g),ddftype=CartesianGrids.M4prime)
    return reg(rhs,str)
end

#=
We redefine the `solve` function and use these functions in place of
the original argument:
=#
function ImmersedLayers.solve(prob::DirichletPoissonProblem,sys::ILMSystem)
    @unpack extra_cache, base_cache = sys
    @unpack bc, f_funcs, gdata_cache = base_cache
    @unpack S, C, fb, fstar = extra_cache

    f = zeros_grid(base_cache)
    s = zeros_surface(base_cache)

    ## Get the boundary data on each side of the interface
    fbplus = bc["fbplus"](base_cache)
    fbminus = bc["fbminus"](base_cache)

    ## Evaluate the forcing field
    f_funcs(gdata_cache,base_cache)

    ## Evaluate the right-hand side of Poisson equation
    surface_divergence!(fstar,fbplus-fbminus,base_cache)
    fstar .+= gdata_cache

    fb .= 0.5*(fbplus+fbminus)

    inverse_laplacian!(fstar,base_cache)

    interpolate!(s,fstar,base_cache)
    s .= fb - s
    s .= -(S\s);

    regularize!(f,s,base_cache)
    inverse_laplacian!(f,base_cache)
    f .+= fstar;

    return f, C^6*s
end

#=
Now we specify the problem, create the system, and solve it, as before,
but now supplying the boundary condition functions with the `bc` keyword:
=#
prob = DirichletPoissonProblem(g,body,scaling=GridScaling,bc=bcdict,f_funcs=rhs_func!)
sys = ImmersedLayers.__init(prob)
f, s = solve(prob,sys)
plot(f,sys)

#=
Same solution, of course. But suppose we wish to change the
the boundary conditions or source points? We can do it easily without regenerating the
cache and system, simply by redefining our bc and forcing functions, e.g.
to create an internal solution, with surface data equal to the $y$ coordinate,
=#
get_fbplus(base_cache) = zeros_surface(base_cache)
get_fbminus(base_cache) = points(base_cache).v
function rhs_func!(rhs,base_cache)
    X = VectorData([-0.2,0.2],[0.0,0.0])
    str = ScalarData([-1.0,1.0])
    reg = Regularize(X,cellsize(base_cache.g),I0=origin(base_cache.g),ddftype=CartesianGrids.M4prime)
    return reg(rhs,str)
end
#=
We can solve immediately without having to reconstruct the system, so
it's very fast.
=#
@time f, s = solve(prob,sys)
plot(f,sys)

#md # ## Problem types and functions
#md #
#md # ```@docs
#md # AbstractScalarILMProblem
#md # AbstractVectorILMProblem
#md # @ilmproblem
#md # BasicScalarILMProblem
#md # BasicVectorILMProblem
#md # prob_cache
#md # ```

#md # ## System types and functions
#md #
#md # ```@docs
#md # ILMSystem
#md # __init
#md # ```


# GaussianVectors.jl

# Define and simulate normal vectors.
struct GaussianVector{R<:AbstractArray{Float64, 1}, S<:AbstractArray{Float64},
		T<:AbstractArray{Float64}}
	mean::R
	Σ::S
	var::T
	length::Int
	numberofGaussians::Int

	function GaussianVector(mean::R, Σ::S) where {R<:AbstractArray{Float64, 1}, S<:AbstractArray{Float64}}
		length = Base.length(mean)
		size(Σ,1) == length || throw(DimensionMismatch("Dimensions of Σ and μ don't match."))
		numberofGaussians = size(Σ,2)
		var = Σ * Σ'
		new{R, S, typeof(var)}(mean, Σ, var, length, numberofGaussians)
	end
end

GaussianVector(Σ::T) where T<:AbstractArray{Float64} = GaussianVector(spzeros(size(Σ,1)), Σ)

rand(d::GaussianVector) = d.mean + d.Σ * randn((d.numberofGaussians, 1))

mean(d::GaussianVector) = d.mean

var(d::GaussianVector) = d.var

length(d::GaussianVector) = d.length


"""
To run:
```
julia --project=. exampleUsage.jl
```
"""

using DelimitedFiles
using Flux
using Random
using Statistics

include("DiscriminativeKalmanFilter.jl")

# setup
rng = MersenneTwister(42)

# data source
z = readdlm("../data/z.csv", ',')
x = readdlm("../data/x.csv", ',')

# dimensions of latent states and observations, respectively
dz = size(z)[2]
dx = size(x)[2]

# training data
train_idx = 1:5000
n_train = size(train_idx)[1]
x_train = x[train_idx, :]
z_train = z[train_idx, :]

# test data
test_idx = 5001:6000
n_test = size(test_idx)[1]
x_test = x[test_idx, :]
z_test = z[test_idx, :]

# learn state model parameters A & Gamma from Eq. (2.1b)
A0 = z_train[2:end, :] \ z_train[1:end-1, :]
Gamma0 = cov(z_train[2:end, :] - z_train[1:end-1, :] * A0)

# split training set in order to train f() and Q() from Eq. (2.2) separately
perm = randperm(MersenneTwister(42), n_train)
x_train_mean = x_train[perm[1:Int(0.9 * 5000)], :]
x_train_covariance = x_train[perm[Int(0.9 * 5000)+1:end], :]
z_train_mean = z_train[perm[1:Int(0.9 * 5000)], :]
z_train_covariance = z_train[perm[Int(0.9 * 5000)+1:end], :]

function NeuralNetwork()
    return Chain(
        Dense(dx, 10, relu),
        Dropout(0.1),
        Dense(10, dz)
    )
end

m = NeuralNetwork()
data = Flux.Data.DataLoader((x_train_mean', z_train_mean'), batchsize = 100, shuffle = true, rng = rng)
loss(x, y) = sum(Flux.Losses.mse(m(x), y))

for i in 1:20
    Flux.train!(loss, Flux.params(m), data, Flux.Optimise.ADAM())
end

Q = cov(z_train_covariance - m(x_train_covariance')')

fx(x) = m(x)
Qx(x) = Q

f0 = fx(x_test[1, :])
Q0 = Qx(x_test[1, :])
S = cov(z_train)

DKF = DiscriminativeKalmanFilter(A0, Gamma0, S, fx, Qx, f0, Q0, dz)

# perform filtering
z_preds = zeros(size(z_test))
z_preds[1, :] = f0
for i in 2:n_test
    z_preds[i, :] = predict!(DKF, x_test[i, :])
end

print("normalized rmse ", sqrt(mean((z_test - z_preds) .^ 2)) / sqrt(mean(z_test .^ 2)))


@testset "Torsion" begin
  # Trivial torsion module

  A = diagonal_matrix(fmpq[1, 1])
  T = Hecke.TorQuadMod(A)
  @test order(T) == 1

  D4_gram = matrix(ZZ, [[2, 0, 0, -1],
                        [0, 2, 0, -1],
                        [0, 0, 2, -1],
                        [-1 ,-1 ,-1 ,2]])

  L = Zlattice(gram = D4_gram)
  T = @inferred discriminant_group(L)
  @test order(T) == 4
  @test elementary_divisors(T) == fmpz[2, 2]

  Ld = dual(L)

  T = @inferred torsion_quadratic_module(Ld, L, snf = false)
  @test ngens(T) == 4
  G = gens(T)
  for i in 1:4
    @test T(lift(G[i])) == G[i]
  end

  TT, mTT = @inferred sub(T, [T([1, 1//2, 1//2, 1])])
  @test order(TT) == 2

  L = Zlattice(ZZ[1 0; 0 1])
  M = lattice(ambient_space(L), ZZ[2 0; 0 3])
  T = @inferred torsion_quadratic_module(L, M, gens = [[1, 1]])
  @test order(T) == 6

  @test_throws ArgumentError torsion_quadratic_module(L, M, gens = [[0, 0]])
  @test_throws ArgumentError torsion_quadratic_module(L, M, gens = [[1//2, 0]])
  @test_throws ArgumentError torsion_quadratic_module(L, M, gens = [[1, 1, 2]])
  @test_throws ArgumentError torsion_quadratic_module(L, lattice(ambient_space(L), QQ[1//2 0; 0 0]))

  #primary part of a TorQuadMod
  L = Zlattice(matrix(ZZ, [[2,0,0],[0,2,0],[0,0,2]]))
  T = Hecke.discriminant_group(L)
  @test basis_matrix(Hecke.cover(Hecke.primary_part(T,fmpz(2))[1])) == matrix(QQ, 3, 3, [1//2, 0, 0, 0, 1//2, 0, 0, 0, 1//2])

  #orthogonal submodule to a TorQuadMod
  L = Zlattice(matrix(ZZ, [[2,0,0],[0,2,0],[0,0,2]]))
  T = Hecke.discriminant_group(L)
  S = sub(T, gens(T))[1]
  @test basis_matrix(Hecke.cover(Hecke.orthogonal_submodule_to(T, S)[1])) == matrix(QQ, 3, 3, [1//2,0,0,0,1//2,0,0,0,1//2])

  #checks if a TorQuadMod is degenerate
  @test Hecke.isdegenerate(T) == false

  #test for rescaled torsion quadratic module
  @test Hecke.gram_matrix_quadratic(Hecke.rescale(T, -1)) == matrix(QQ, 3, 3, [7//4,0,0,0,7//4,0,0,0,7//4])
  t = Hecke.TorQuadMod(QQ[1//3 0; 0 1//9])
  @test Hecke.gram_matrix_quadratic(Hecke.rescale(t, 2)) == matrix(QQ, 2, 2, [2//3,0,0,2//9])

  #test for normal form
  L1 = Zlattice(gram=matrix(ZZ, [[-2,0,0],[0,1,0],[0,0,4]]))
  T1 = Hecke.discriminant_group(L1)
  @test Hecke.gram_matrix_quadratic(Hecke.normal_form(T1)[1]) == matrix(QQ, 2, 2, [1//2,0,0,1//4])

  L1 = Zlattice(gram=ZZ[-4 0 0; 0 4 0; 0 0 -2])
  AL1 = discriminant_group(L1)
  L2 = Zlattice(gram=ZZ[-4 0 0; 0 -4 0; 0 0 2])
  AL2 = discriminant_group(L2)
  n1 = normal_form(AL1)[1]
  g1 = QQ[1//2   0   0;
            0 1//4   0;
            0   0 5//4]
  @test Hecke.gram_matrix_quadratic(n1) == g1
  n2 = normal_form(AL2)[1]
  g2 = QQ[1//2   0   0;
            0 1//4   0;
            0   0 5//4]
  @test Hecke.gram_matrix_quadratic(n2) == g2

  #test for brown invariant
  L = Zlattice(gram=matrix(ZZ, [[2,-1,0,0],[-1,2,-1,-1],[0,-1,2,0],[0,-1,0,2]]))
  T = discriminant_group(L)
  @test Hecke.brown_invariant(T) == 4

  #test for genus
  L = Zlattice(gram=diagonal_matrix(fmpz[1,2,3,4]))
  D = discriminant_group(L)
  @test genus(D, (4,0)) == genus(L)

  #test for isgenus
  L = Zlattice(gram=diagonal_matrix(fmpz[1,2,3,4]))
  D = discriminant_group(L)
  @test_throws ErrorException isgenus(D, (4,0))
  L1 = Zlattice(gram=matrix(ZZ, [[2, -1, 0, 0, 0, 0],[-1, 2, -1, -1, 0, 0],[0, -1, 2, 0, 0, 0],[0, -1, 0, 2, 0, 0],[0, 0, 0, 0, 6, 3],[0, 0, 0, 0, 3, 6]]))
  T1 = discriminant_group(L1)
  @test isgenus(T1, (6,0))
  G = genus(diagonal_matrix(fmpz[2, 6, 6]))
  D = discriminant_group(G)
  @test isgenus(D, (3,0))

end


using DesignAutomation

include("functions.jl")

@testset "Exhaustive Search" begin
    x_range = [-2, 6]

    (x, fx, f_calls) = exhaustive_search(levy3, x_range, 4; verbose=false)

    @test x == [1.0, 1.0, 1.0, -2.0]
    @test fx + 1 ≈ 1 # answer is approx 0
    @test f_calls == 81^4
end


function initializePH(cf1::CF1{Tv}, data::AbstractPHSample, ::Type{MatT} = SparseMatrixCSC;
    shapes = Tv[1, 4, 16, 64, 256, 1024],
    scales = Tv[0.5, 1.0, 2.0],
    maxiter = 5, verbose = false, eps::Tv = Tv(1.0e-8), ufact::Tv = Tv(1.01)) where {Tv,MatT}
    verbose && println("Initializing CF1 ...")
    m = _phmean(data)
    maxllf = typemin(Tv)
    maxph = cf1
    eres = Estep(GPH(cf1, MatT))
    for fn = [_cf1_params_power, _cf1_params_linear]
        for x = scales, s = shapes
            try
                newcf1 = fn(cf1.dim, m*x, s)
                local llf::Tv
                for k = 1:maxiter
                    llf = estep!(newcf1, data, eres, eps=eps, ufact=ufact)
                    mstep!(newcf1, eres)
                end
                if !isfinite(llf)
                    verbose && print("-")
                else
                    if maxllf < llf
                        maxllf, maxph = llf, newcf1
                        verbose && print("o")
                    else
                        verbose && print("x")
                    end
                end
            catch
                verbose && print("-")
            end
        end
        verbose && println()
    end
    return maxph
end

function _cf1_params_power(dim::Ti, scale::Tv, shape::Tv) where {Tv,Ti}
    rate = Vector{Tv}(undef, dim)
    p = exp(Tv(1)/(dim-1) * log(shape))
    total = Tv(1)
    tmp = Tv(1)
    for i = 1:dim-1
        tmp *= (i+1) / (i * p)
        total += tmp
    end
    base = total / (dim * scale)
    tmp = base
    for i = 1:dim
        rate[i] = tmp
        tmp *= p
    end
    CF1(fill(1/dim, dim), rate)
end

function _cf1_params_linear(dim::Ti, scale::Tv, shape::Tv) where {Tv,Ti}
    rate = Vector{Tv}(undef, dim)
    al = (shape - 1)/(dim-1)
    total = Tv(1)
    for i = 1:dim
        total += (i+1)/(al*i+1)
    end
    base = total / (dim * scale)
    for i = 1:dim
        rate[i] = base * (al + i + 1)
    end
    CF1(fill(1/dim, dim), rate)
end


using DataStructures

function main()
    N, M = parse.(Int, split(readline()))
    G = [Vector{Int}() for _ = 1:N]
    for _ = 1:M
        u, v = parse.(Int, split(readline()))
        push!(G[u], v)
    end
    S, T = parse.(Int, split(readline()))

    # Solve
    dists = -ones(Int, N, 3)
    dists[S, 1] = 0
    que = Queue{Tuple{Int, Int}}()
    enqueue!(que, (S, 1))
    while !isempty(que)
        v, s = dequeue!(que)
        for v′ ∈ G[v]
            s′ = s % 3 + 1
            if dists[v′, s′] == -1
                dists[v′, s′] = dists[v, s] + 1
                enqueue!(que, (v′, s′))
            end
        end
    end

    if dists[T, 1] == -1
        println("-1")
    else
        println("$(dists[T, 1] ÷ 3)")
    end
end

main()


# ----------------------------------------------------
# --- Import modules
# ----------------------------------------------------
using DigitalComm
using Test
# ----------------------------------------------------
# --- Tests
# ----------------------------------------------------
# Note --> The mapping system is described in bitMapping.jl

println("Tests for Hard symbol demapper");
for alphabet = [2;4;16;64;256]
	@testset  "$alphabet-QAM" begin
		# Minimal Tx
		mcs	  = alphabet;		  # Constellation size
		n	  = Int(log2(mcs));		  # Bit per symbol
		N	  = 4096;	  # Numbver of test symbols
		bitSeq	= genBitSequence(N*n);
		qamSeq	= bitMappingQAM(mcs,bitSeq);
		# Minimal Rx
		bitDec	= bitDemappingQAM(mcs,qamSeq);
		# Checking size of decoding vector
		@test length(bitDec) == N*n;
		# Checking that bang and classic methods are the same
		buff  = zeros(UInt8,N*n);
		bitDemappingQAM!(buff,mcs,qamSeq);
		@test all(bitDec .== buff);
		# Checking that decoded method is initial method
		@test all( bitDec .≈ bitSeq)
	end
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Desription
# ==============================================================================
#
#   Tests related to the kinematics of quaternions.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File: ./src/quaternion.jl
# =========================

# Functions: dquat
# ----------------

@testset "Kinematics of quaternions (Float64)" begin
    T = Float64

    # Create a random quaternion.
    qba₀ = Quaternion(randn(T), randn(T), randn(T), randn(T))
    qba₀ = qba₀ / norm(qba₀)

    # Create a random velocity vector.
    wba_a = @SVector randn(T, 3)

    # Propagate the initial DCM using the sampled velocity vector.
    Δ   = 1e-7
    qba = qba₀
    num = 2_000_000

    for k in 1:num
        dqba = dquat(qba, vect(qba \ wba_a * qba))
        qba  = qba + dqba * Δ
        qba  = qba / norm(qba)
    end

    # In the end, the vector aligned with `wba_a` must not change.
    v₀ = vect(qba₀ \ wba_a * qba₀)
    v₁ = vect(qba \ wba_a * qba)
    @test v₁ ≈ v₀ rtol = 1e-6
    @test eltype(v₀) === T
    @test eltype(v₁) === T

    # In the end, a vector perpendicular to `wba_a` must rotate the angle
    # compatible with the angular velocity and time of integration.
    aux   = @SVector randn(T , 3)
    aux   = aux / norm(aux)
    vr_a  = wba_a × aux
    vr_a  = vr_a / norm(vr_a)
    vr_b₀ = vect(qba₀ \ vr_a * qba₀)
    vr_b  = vect(qba \ vr_a * qba)
    θ     = acos(vr_b ⋅ vr_b₀)
    @test eltype(vr_b₀) === T
    @test eltype(vr_b) === T

    # Estimate θ based on the angular velocity.
    θest = mod(norm(wba_a) * Δ * num, 2π)
    θest > pi && (θest = 2π - θest)
    @test θ ≈ θest atol = 1e-6
end

@testset "Kinematics of quaternions (Float32)" begin
    T = Float32

    # Create a random quaternion.
    qba₀ = Quaternion(randn(T), randn(T), randn(T), randn(T))
    qba₀ = qba₀ / norm(qba₀)

    # Create a random velocity vector.
    wba_a = @SVector randn(T, 3)

    # Propagate the initial DCM using the sampled velocity vector.
    Δ   = 1f-7
    qba = qba₀
    num = 20_000

    for k in 1:num
        dqba = dquat(qba, vect(qba \ wba_a * qba))
        qba  = qba + dqba * Δ
        qba  = qba / norm(qba)
    end

    # In the end, the vector aligned with `wba_a` must not change.
    v₀ = vect(qba₀ \ wba_a * qba₀)
    v₁ = vect(qba \ wba_a * qba)
    @test v₁ ≈ v₀ rtol = 1e-2
    @test eltype(v₀) === T
    @test eltype(v₁) === T

    # In the end, a vector perpendicular to `wba_a` must rotate the angle
    # compatible with the angular velocity and time of integration.
    aux   = @SVector randn(T , 3)
    aux   = aux / norm(aux)
    vr_a  = wba_a × aux
    vr_a  = vr_a / norm(vr_a)
    vr_b₀ = vect(qba₀ \ vr_a * qba₀)
    vr_b  = vect(qba \ vr_a * qba)
    θ     = acos(vr_b ⋅ vr_b₀)
    @test eltype(vr_b₀) === T
    @test eltype(vr_b) === T

    # Estimate θ based on the angular velocity.
    θest = mod(norm(wba_a) * Δ * num, T(2π))
    θest > pi && (θest = T(2π) - θest)
    @test θ ≈ θest atol = 1e-2
end

@testset "Kinematics of quaternions (Errors)" begin
    @test_throws ArgumentError dquat(Quaternion(I), [1, 2])
    @test_throws ArgumentError dquat(Quaternion(I), [1, 2, 3, 4])
end


# Statistics code

export ShiftedExponential, ExponentialOrder, OrderStatistic, NonIDOrderStatistic

"""
    ExponentialOrder(scale::Real, total::Int, order::Int)

Random variable representing the order-th largest value out of total
realizations of an exponential random variable with given scale.

"""
function ExponentialOrder(scale::Real, total::Int, order::Int)
    scale > 0 || throw(DomainError(scale, "scale must be positive"))
    total > 0 || throw(DomainError((total, order), "total must be positive"))
    order > 0 || throw(DomainError((total, order), "order must be positive"))
    order <= total || throw(DomainError((total, order), "order must be <= total"))
    var = sum(1/(i^2)  for i=(total-order+1):total) * scale^2
    mean = sum(1/i for i=(total-order+1):total) * scale
    alpha = mean^2 / var # shape parameter
    theta = var / mean # scale
    return Gamma(alpha, theta)
end

### Shifted exponential distribution

struct ShiftedExponential{T<:Real} <: ContinuousUnivariateDistribution
    shift::T
    exponential::Exponential{T}
    ShiftedExponential{T}(s::T, θ::T) where T = new{T}(s, Exponential(θ))
end

function ShiftedExponential(s::T, θ::T) where {T<:Real}
    θ > zero(θ) || throw(DomainError(θ), "θ must be positive")
    ShiftedExponential{T}(s, θ)
end

function ShiftedExponential(s::T, θ::T) where {T<:Integer}
    ShiftedExponential(float(s), float(θ))
end

struct ShiftedExponentialStats <: SufficientStats
    vx::Float64     # variance of x
    sx::Float64     # (weighted) sum of x
    sw::Float64     # sum of sample weights
end

Distributions.suffstats(::Type{<:ShiftedExponential}, x::AbstractArray{T}) where {T<:Real} = ShiftedExponentialStats(var(x), sum(x), length(x))

function Distributions.fit_mle(::Type{<:ShiftedExponential}, ss::ShiftedExponentialStats)
    θ = sqrt(ss.vx)
    s = ss.sx / ss.sw - θ
    ShiftedExponential(s, θ)
end

Distributions.quantile(d::ShiftedExponential, p::Real) = Distributions.quantile(d.exponential, p) + d.shift

Distributions.pdf(d::ShiftedExponential, p::Real) = pdf(d.exponential, p-d.shift)
Distributions.cdf(d::ShiftedExponential, p::Real) = cdf(d.exponential, p-d.shift)
Base.rand(rng::AbstractRNG, d::ShiftedExponential) = rand(rng, d.exponential) + d.shift
Base.minimum(d::ShiftedExponential) = d.shift
Base.maximum(::ShiftedExponential) = Inf
Statistics.mean(d::ShiftedExponential) = mean(d.exponential) + d.shift
Statistics.var(d::ShiftedExponential) = var(d.exponential)
Distributions.params(d::ShiftedExponential) = (d.shift, params(d.exponential)...)

### Fitting a distribution from its mean and variance

function distribution_from_mean_variance(::Type{ShiftedExponential}, m, v)
    θ = sqrt(v)
    s = m .- θ
    ShiftedExponential(s, θ)
end

function distribution_from_mean_variance(::Type{Gamma}, m, v)
    θ = v / m
    α = m / θ
    Gamma(α, θ)
end

### Order statistics sampling

struct OrderStatistic{S<:Union{Discrete,Continuous},Spl<:Sampleable{Univariate,S},T} <: Sampleable{Univariate,S}
    spl::Spl
    k::Int
    buffer::Vector{T}
end

Base.show(io::IO, s::OrderStatistic) = print(io, "OrderStatistic($(s.spl), k=$(s.k), n=$(length(s.buffer)))")

OrderStatistic(s::Sampleable, k::Integer, n::Integer) = OrderStatistic(s, k, Vector{eltype(s)}(undef, n))

function Random.rand(rng::AbstractRNG, s::OrderStatistic)
    Distributions.rand!(rng, s.spl, s.buffer)
    partialsort!(s.buffer, s.k)
    s.buffer[s.k]
end

struct NonIDOrderStatistic{S<:Union{Discrete,Continuous},Spl<:Sampleable{Univariate,S},T} <: Sampleable{Univariate,S}
    spls::Vector{Spl}
    k::Int
    buffer::Vector{T}
end

Base.show(io::IO, s::NonIDOrderStatistic) = print(io, "NonIDOrderStatistic($(eltype(s.spls)), k=$(s.k), n=$(length(s.buffer)))")

NonIDOrderStatistic(spls::AbstractVector{<:Sampleable}, k::Integer) = NonIDOrderStatistic(spls, k, Vector{promote_type((eltype(spl) for spl in spls)...)}(undef, length(spls)))

function Random.rand(rng::AbstractRNG, s::NonIDOrderStatistic)
    for (i, spl) in enumerate(s.spls)
        s.buffer[i] = Distributions.rand(rng, spl)
    end
    partialsort!(s.buffer, s.k)
    s.buffer[s.k]
end