text:
module AEM_VMD_HMD
include("DigFilters.jl")
include("TDcossin.jl")
using LinearAlgebra, SpecialFunctions, FastGaussQuadrature, DataInterpolations

abstract type HField end

mutable struct HFieldDHT <: HField
    thickness       :: Array{Float64, 1}
    pz              :: Array{Complex{Float64}, 1}
    ϵᵢ              :: Array{ComplexF64, 1}
    zintfc          :: Array{Float64, 1}
    rTE             :: Array{ComplexF64, 1}
    rTM             :: Array{ComplexF64, 1}
    zRx             :: Float64
    zTx             :: Float64
    rTx             :: Union{Float64, Nothing}
    rRx             :: Float64
    freqs           :: Array{Float64, 1}
    times           :: Array{Float64, 1}
    ramp            :: Array{Float64, 2}
    log10ω          :: Array{Float64, 1}
    interptimes     :: Array{Float64, 1}
    HFD_z           :: Array{ComplexF64, 1}
    HFD_r           :: Array{ComplexF64, 1}
    HFD_az          :: Array{ComplexF64, 1}
    HFD_z_interp    :: Array{ComplexF64, 1}
    HFD_r_interp    :: Array{ComplexF64, 1}
    HFD_az_interp   :: Array{ComplexF64, 1}
    HTD_z_interp    :: Array{Float64, 1}
    HTD_r_interp    :: Array{Float64, 1}
    HTD_az_interp   :: Array{ComplexF64, 1}
    dBzdt           :: Array{Float64, 1}
    dBrdt           :: Array{Float64, 1}
    dBazdt          :: Array{Float64, 1}
    J0_kernel_h     :: Array{ComplexF64, 2}
    J1_kernel_h     :: Array{ComplexF64, 2}
    J0_kernel_v     :: Array{ComplexF64, 2}
    J1_kernel_v     :: Array{ComplexF64, 2}
    lowpassfcs      :: Array{Float64, 1}
    quadnodes       :: Array{Float64, 1}
    quadweights     :: Array{Float64, 1}
    interplog10ω    :: Array{Float64, 2}
    ω               :: Array{Float64, 2}
    Hsc             :: Array{ComplexF64, 2}
    rxwithinloop    :: Bool
    provideddt      :: Bool
    doconvramp      :: Bool
    useprimary      :: Float64
    nkᵣeval         :: Int
    interpkᵣ        :: Array{Float64,1}
    log10interpkᵣ   :: Array{Float64,1}
    log10Filter_base   :: Array{Float64,1}
    getradialH      :: Bool
    getazimH        :: Bool
    calcjacobian    :: Bool
    Jtemp           :: Vector
    b               :: Vector # Another Jtemp
    derivmatrix     :: Vector
    HFD_z_J         :: Array{ComplexF64, 2}
    HTD_z_J_interp  :: Array{Float64, 2}
    dBzdt_J         :: Array{Float64, 2}
end

function HFieldDHT(;
      nmax      = 200,
      rTx       = nothing,
      rRx       = 17.0,
      freqs     = [],
      zTx       = -35.0,
      zRx       = -37.5,
      times     = 10 .^LinRange(-6, -1, 50),
      ramp      = ones(10, 10),
      nfreqsperdecade = 5,
      ntimesperdecade = 10,
      glegintegorder = 5,
      lowpassfcs = [],
      provideddt = true,
      doconvramp = true,
      modelprimary = false,
      nkᵣeval = 50,
      getradialH = false,
      getazimH  = false,
      freqlow = 1e-4,
      freqhigh = 1e6,
      calcjacobian = false
  )
    @assert all(freqs .> 0.)
    @assert freqhigh > freqlow
    @assert all(diff(times) .> 0)
    thickness = zeros(nmax)
    zintfc    = zeros(nmax)
    pz        = zeros(Complex{Float64}, nmax)
    ϵᵢ      = similar(pz)
    rTE       = zeros(length(pz)-1)
    rTM       = similar(rTE)
    if freqhigh < 3/minimum(times)
       freqhigh = 3/minimum(times)
    end
    if freqlow > 0.33/maximum(times)
       freqlow = 0.33/maximum(times)
    end
    if isempty(freqs)
        freqs = 10 .^(log10(freqlow):1/nfreqsperdecade:log10(freqhigh))
    end
    interpkᵣ = 10 .^range(minimum(log10.(Filter_base))-0.5, maximum(log10.(Filter_base))+0.5, length = nkᵣeval)
    J0_kernel_h, J1_kernel_h, J0_kernel_v, J1_kernel_v = map(x->zeros(ComplexF64, length(interpkᵣ), length(freqs)), 1:4)
    log10ω = log10.(2*pi*freqs)
    interptimes = 10 .^(minimum(log10.(times))-1:1/ntimesperdecade:maximum(log10.(times))+1)
    HFD_z       = zeros(ComplexF64, length(freqs)) # space domain fields in freq
    HFD_r       = zeros(ComplexF64, length(freqs)) # space domain fields in freq
    HFD_az       = zeros(ComplexF64, length(freqs)) # space domain fields in freq
    dBzdt     = zeros(Float64, length(times)) # time derivative of space domain fields convolved with ramp
    dBrdt     = zeros(Float64, length(times)) # time derivative of space domain fields convolved with ramp
    dBazdt     = zeros(Float64, length(times)) # time derivative of space domain fields convolved with ramp
    HFD_z_interp = zeros(ComplexF64, length(Filter_t_base))
    HFD_r_interp = zeros(ComplexF64, length(Filter_t_base))
    HFD_az_interp = zeros(ComplexF64, length(Filter_t_base))
    HTD_z_interp = zeros(Float64, length(interptimes))
    HTD_r_interp = zeros(Float64, length(interptimes))
    HTD_az_interp = zeros(Float64, length(interptimes))
    lowpassfcs = float.([lowpassfcs..., 5e6])
    quadnodes, quadweights = gausslegendre(glegintegorder)
    rxwithinloop = false
    log10Filter_base = log10.(Filter_base/rRx)
    if rTx != nothing
        if rTx>=rRx
            rxwithinloop = true
            interpkᵣ = interpkᵣ/rTx
            log10Filter_base = log10.(Filter_base/rTx)
        else
            interpkᵣ = interpkᵣ/rRx
        end
    end
    log10interpkᵣ = log10.(interpkᵣ)
    HFD_z_J     = calcjacobian ? zeros(ComplexF64, nmax, length(freqs)) : zeros(0, 0)
    dBzdt_J     = calcjacobian ? zeros(Float64, nmax, length(times)) : zeros(0, 0)
    HTD_z_J_interp = calcjacobian ? zeros(Float64, nmax, length(interptimes)) : zeros(0, 0)
    Jtemp = calcjacobian ? zeros(ComplexF64, nmax) : zeros(0)
    Jac = calcjacobian ? map(x->zeros(ComplexF64, nkᵣeval, nmax), 1:length(freqs)) : zeros(0)
    useprimary = modelprimary ? one(Float64) : zero(Float64)
    HFieldDHT(thickness, pz, ϵᵢ, zintfc, rTE, rTM, zRx, zTx, rTx, rRx, freqs, times, ramp, log10ω, interptimes,
            HFD_z, HFD_r, HFD_az, HFD_z_interp, HFD_r_interp, HFD_az_interp,
            HTD_z_interp, HTD_r_interp, HTD_az_interp, dBzdt, dBrdt, dBazdt, J0_kernel_h, J1_kernel_h, J0_kernel_v, J1_kernel_v, lowpassfcs,
            quadnodes, quadweights, preallocate_ω_Hsc(interptimes, lowpassfcs)..., rxwithinloop, provideddt, doconvramp, useprimary,
            nkᵣeval, interpkᵣ, log10interpkᵣ, log10Filter_base, getradialH, getazimH,
            calcjacobian, Jtemp, similar(Jtemp), Jac, HFD_z_J, HTD_z_J_interp, dBzdt_J)
end

#update geometry and dependent parameters - necessary for adjusting geometry
#e.g. inverting nuisance parameters
function update_ZR!(F::HFieldDHT, zTx, zRx, rTx, rRx)
    if F.rTx != nothing
        if F.rTx >= F.rRx
            F.interpkᵣ *= F.rTx
            F.log10interpkᵣ .+= log10(F.rTx)
            F.log10Filter_base .+= log10(F.rTx)
        else
            F.interpkᵣ *= F.rRx
            F.log10interpkᵣ .+= log10(F.rRx)
            F.log10Filter_base .+= log10(F.rRx)
        end
    else
        F.log10Filter_base .+= log10(F.rRx)
    end

    if rTx != nothing
        if rTx >= rRx
            F.interpkᵣ /= rTx
            F.log10interpkᵣ .-= log10(rTx)
            F.log10Filter_base .-= log10(rTx)
        else
            F.interpkᵣ /= rRx
            F.log10interpkᵣ .-= log10(rRx)
            F.log10Filter_base .-= log10(rRx)
        end
    else
        F.log10Filter_base .-= log10(rRx)
    end
    F.rTx = rTx
    F.rRx = rRx
    F.zTx = zTx
    F.zRx = zRx

end

function preallocate_ω_Hsc(interptimes, lowpassfcs)
    ω   = zeros(length(Filter_t_base), length(interptimes))
    Hsc = zeros(ComplexF64, length(Filter_t_base), length(interptimes))
    for itime = 1:length(interptimes)
        t = interptimes[itime]
        ω[:,itime] = log10.(Filter_t_base/t)
        Hsc[:,itime] = ones(ComplexF64, length(Filter_t_base))
        s = 1im*10 .^ω[:,itime]
        for fc in lowpassfcs
             Hs = 1 ./( 1 .+s./(2*pi*fc))
             Hsc[:,itime] .= Hsc[:,itime].*Hs
        end
    end
    ω, 10 .^ω, Hsc
end

const μ       = 4*pi*1e-7
const ϵ₀     = 8.854e-12
const iTxLayer = 1
const iRxLayer = 1

function getpartialkz(ω, kz)
    0.5im*ω*μ/kz
end

function getpartial_rTE(partialkz, kz, kznext)
    2*partialkz*kznext/(kz+kznext)^2
end

function getpartial_rTE_withnext(partialkznext, kz, kznext)
    -2*partialkznext*kz/(kz+kznext)^2
end

function getpartialRstack(partial_rTE, rTE, a)
    partial_rTE*(1 - a^2)/(1 + rTE*a)^2
end

function getpartial_bwithnext(b, dnext, partialkznext)
    b*(2im*dnext*partialkznext)
end

function stacks!(F::HField, iTxLayer::Int, nlayers::Int, ω::Float64)

    rTE              = F.rTE
    rTM              = F.rTM
    pz               = view(F.pz, 1:nlayers)
    # The first layer thicknesses is infinite
    d                = view(F.thickness, 1:nlayers)
    d[1]             = 1e60
    d[nlayers]       = 1e60
    if F.calcjacobian
        b      = view(F.b, 1:nlayers)
        J      = view(F.Jtemp, 1:nlayers)
    end
    # Capital R is for a stack
    # Starting from the bottom up, for Rs_down
    Rlowerstack_TE, Rlowerstack_TM = zero(ComplexF64), zero(ComplexF64)
    Rlowerstack_plus_TE = zero(ComplexF64)
    @inbounds @fastmath for k = (nlayers-1):-1:iTxLayer
        if F.calcjacobian
            kz, kznext = ω*pz[k], ω*pz[k+1]
            partialkznext = getpartialkz(ω, kznext)
            S = getpartial_rTE_withnext(partialkznext, kz, kznext)
            if k == nlayers-1
                J[k+1] = S
                Rlowerstack_plus_TE = Rlowerstack_TE
                b[k+1] = exp(2im*kznext*d[k+1])
                Rlowerstack_TE = lowerstack(Rlowerstack_TE, rTE, k, b[k+1])
                continue
            end
            b[k+1] = exp(2im*kznext*d[k+1])
            partial_bnext = getpartial_bwithnext(b[k+1], d[k+1], partialkznext)
            S *= 1 - (Rlowerstack_TE*b[k+1])^2
            kznextplus = ω*pz[k+2]
            partial_rTE_next = getpartial_rTE(partialkznext, kznext, kznextplus)
            a = Rlowerstack_plus_TE*b[k+2]
            partialRnext = getpartialRstack(partial_rTE_next, rTE[k+1], a)
            denom = (1+rTE[k]*Rlowerstack_TE*b[k+1])^2
            S += (1-rTE[k]^2)*(partialRnext*b[k+1] + partial_bnext*Rlowerstack_TE)
            S /= denom
            J[k+1] = S
            c = b[k+1]*(1-rTE[k]^2)/denom
            for kk = nlayers:-1:k+2
                J[kk] *= c
            end
            Rlowerstack_plus_TE = Rlowerstack_TE
            Rlowerstack_TE = lowerstack(Rlowerstack_TE, rTE, k, b[k+1])
        else
            kznext = ω*pz[k+1]
            bnext = exp(2im*kznext*d[k+1])
            Rlowerstack_TE = lowerstack(Rlowerstack_TE, rTE, k, bnext)
        end
    end

return Rlowerstack_TE, Rlowerstack_TM
end

@inline function lowerstack(Rlowerstack::ComplexF64, r::Array{ComplexF64, 1}, k::Int, b)
    a = Rlowerstack*b
    Rs_d = (r[k] + a) / (1. + r[k]*a)
end

function getCurlyR(Rs_d::ComplexF64, pz::ComplexF64,
                  zR::Float64, z::SubArray{Float64, 1}, iTxLayer::Int, ω::Float64,
                  useprimary::Float64)

    if (zR>=0)
        e_to_the_iwpzzr               = exp( im*ω*pz*zR)
        e_to_the_iwpz_2znext_minus_zr = exp( im*ω*pz*(2*z[iTxLayer+1] - zR))

        finRA = useprimary*e_to_the_iwpzzr + e_to_the_iwpz_2znext_minus_zr*Rs_d
        # Will need these for other components
        # finRB = e_to_the_iwpzzr + e_to_the_iwpz_2znext_minus_zr*Rs_d
        #
        finRC = useprimary*e_to_the_iwpzzr + e_to_the_iwpz_2znext_minus_zr*(-Rs_d)
        #
        finRD = useprimary*e_to_the_iwpzzr + e_to_the_iwpz_2znext_minus_zr*(-Rs_d)
    else
        e_to_the_iwpz2znext              = exp(im*ω*pz*2*z[iTxLayer+1])
        e_to_the_minus_iwpzzr            = exp(-im*ω*pz*zR)

        finRA = (useprimary + e_to_the_iwpz2znext*Rs_d) * e_to_the_minus_iwpzzr
        # Will need these for other components
        # finRB = (-1. + e_to_the_iwpz2znext*Rs_d) * e_to_the_minus_iwpzzr
        #
        finRC = (useprimary + e_to_the_iwpz2znext*(-Rs_d)) * e_to_the_minus_iwpzzr
        #
        finRD = (-useprimary + e_to_the_iwpz2znext*(-Rs_d)) * e_to_the_minus_iwpzzr
    end

    return finRA, finRD, finRC
end
# FT convention in this function means will have to conjugate for compatibility with FT = exp(-iωt)
getepsc(ρ::Float64, ω::Float64)      = Complex(ϵ₀, 1/(ρ*ω)) # corresponds to FT = exp(+iωt)
getpz(ϵᵢ::ComplexF64, kᵣ::Float64, ω::Float64) = unsafesqrt(μ*ϵᵢ - (kᵣ/ω)^2) # corresponds k² = iωμσ
ztxorignify(z, zTx)      = z - zTx
makesane(pz::Complex)    = imag(pz)  < 0.0 ? ( pz*=-1.) : pz # Sommerfeld radiation condition for fields 0 at ∞

function unsafesqrt(z::ComplexF64)
    # x = real(z)
    # d = sqrt(x^2+imag(z)^2)
    # ComplexF64(sqrt(0.5(d+x)), sqrt(0.5(d-x)))
    x, y = reim(z)
    ρ = x*x + y*y
    ρ=abs(x)+sqrt(ρ)
    k = -1
    ρ += ρ
    ρ = ldexp(sqrt(ρ),k) #sqrt((abs(z)+abs(x))/2) without over/underflow
    ξ = ρ
    η = y
    η=(η/ρ)/2
    if x<0
        ξ = abs(η)
        η = copysign(ρ,y)
    end
    Complex(ξ,η)
end

function getAEM1DKernelsH!(F::HField, kᵣ::Float64, f::Float64, zz::Array{Float64, 1}, ρ::Array{Float64, 1})
    nlayers = length(ρ)
    z       = view(F.zintfc, 1:nlayers)
    zRx     = ztxorignify(F.zRx, F.zTx)
    ω       = 2. *pi*f
    ϵᵢ      = F.ϵᵢ
    pz      = F.pz
    # reflection coefficients (downward) for an intfc: pz vertical slowness
    rTE, rTM = F.rTE, F.rTM

    l = 1
    z[l]       = ztxorignify(zz[l], F.zTx)
    ϵᵢ[l]      = getepsc(ρ[l], ω)
    pz[l]      = getpz(ϵᵢ[l], kᵣ, ω)
    @inbounds @fastmath for intfc in 1:nlayers-1
        l = intfc+1
        z[l]       = ztxorignify(zz[l], F.zTx)
        ϵᵢ[l]      = getepsc(ρ[l], ω)
        pz[l]      = getpz(ϵᵢ[l], kᵣ, ω)
        rTE[intfc] = (pz[intfc] - pz[intfc+1])/(pz[intfc] + pz[intfc+1])
        # commented out as we aren't yet using TM modes
        # rTM[intfc] = (ϵᵢ[intfc]*pz[intfc+1] - ϵᵢ[intfc+1]*pz[intfc]) /
        #              (ϵᵢ[intfc]*pz[intfc+1] + ϵᵢ[intfc+1]*pz[intfc])
        F.thickness[intfc] = z[intfc+1] - z[intfc]
    end

    # TE and TM modes
    Rs_dTE, Rs_dTM   = stacks!(F, iTxLayer, nlayers, ω)
    curlyRA, curlyRD, curlyRC = getCurlyR(Rs_dTE, pz[iTxLayer], zRx, z, iTxLayer, ω, F.useprimary)
    lf_gA_TE  = 1.0
    if F.rxwithinloop
        lf_gA_TE *= loopfactor(F.rTx, kᵣ, F.rRx)
    else
        lf_gA_TE *= loopfactor(F.rTx, kᵣ)
    end
    if F.calcjacobian
        ifreq = findfirst(isapprox.(ω, 2pi*F.freqs))
        ikᵣ = findfirst(isapprox.(kᵣ, F.interpkᵣ))
        for ilayer = 2:nlayers
            curlyRAprime, = getCurlyR(F.Jtemp[ilayer], pz[iTxLayer], zRx, z, iTxLayer, ω, 0.)# cannot model primary for deriv
            F.derivmatrix[ifreq][ikᵣ,ilayer] = 1/pz[iTxLayer]*curlyRAprime*lf_gA_TE*1im/(ω)*log(10)/ρ[ilayer]
        end
    end
    gA_TE            = μ/pz[iTxLayer]*curlyRA*lf_gA_TE
    gD_TE            = curlyRD
    gC_TE            = pz[iTxLayer]/μ*curlyRC
    # Kernels according to Loseth, without the bessel functions multiplied
    J0v              = gA_TE*1im/(ω*μ)
    J1v              = gD_TE*kᵣ^2/4pi
    J1h              = gC_TE*1im*ω*μ/4pi
    return J0v, J1v, J1h
end

loopfactor(rTx::Nothing, kᵣ::Float64)               = kᵣ^3/(4*pi)
loopfactor(rTx::Float64, kᵣ::Float64)               = kᵣ^2*besselj1(kᵣ*rTx)/(2*pi*rTx)
loopfactor(rTx::Float64, kᵣ::Float64, rRx::Float64) = kᵣ^2*besselj0(kᵣ*rRx)/(2*pi*rTx)

function getfieldFD!(F::HFieldDHT, z::Array{Float64, 1}, ρ::Array{Float64, 1})
    for (ifreq, freq) in enumerate(F.freqs)
        for (ikᵣ, kᵣ) in enumerate(F.interpkᵣ)
            if F.rxwithinloop
                F.J1_kernel_v[ikᵣ,ifreq], = getAEM1DKernelsH!(F, kᵣ, freq, z, ρ)
            else
                F.J0_kernel_v[ikᵣ,ifreq],
                F.J1_kernel_v[ikᵣ,ifreq],
                F.J1_kernel_h[ikᵣ,ifreq] = getAEM1DKernelsH!(F, kᵣ, freq, z, ρ)
            end
        end # kᵣ loop
        if F.rxwithinloop
            splreal = CubicSpline(real(F.J1_kernel_v[:,ifreq]), F.log10interpkᵣ)
            splimag = CubicSpline(imag(F.J1_kernel_v[:,ifreq]), F.log10interpkᵣ)
            J1_kernel_v = splreal.(F.log10Filter_base) + 1im*splimag.(F.log10Filter_base)
            F.HFD_z[ifreq] = dot(J1_kernel_v, Filter_J1)/F.rTx
        else
            # Vertical component
            splreal = CubicSpline(real(F.J0_kernel_v[:,ifreq]), F.log10interpkᵣ)
            splimag = CubicSpline(imag(F.J0_kernel_v[:,ifreq]), F.log10interpkᵣ)
            J0_kernel_v = splreal.(F.log10Filter_base) + 1im*splimag.(F.log10Filter_base)
            F.HFD_z[ifreq] = dot(J0_kernel_v, Filter_J0)/F.rRx
                if F.calcjacobian
                    for ilayer = 2:length(ρ)
                        splreal = CubicSpline(real(vec(F.derivmatrix[ifreq][:,ilayer])), F.log10interpkᵣ)
                        splimag = CubicSpline(imag(vec(F.derivmatrix[ifreq][:,ilayer])), F.log10interpkᵣ)
                        J0_kernel_v_prime = splreal.(F.log10Filter_base) + 1im*splimag.(F.log10Filter_base)
                        F.HFD_z_J[ilayer, ifreq] = dot(J0_kernel_v_prime, Filter_J0)/F.rRx
                    end
                end
            # radial component
            if F.getradialH
                splreal = CubicSpline(real(F.J1_kernel_v[:,ifreq]), F.log10interpkᵣ)
                splimag = CubicSpline(imag(F.J1_kernel_v[:,ifreq]), F.log10interpkᵣ)
                J1_kernel_v = splreal.(F.log10Filter_base) + 1im*splimag.(F.log10Filter_base)
                F.HFD_r[ifreq] = dot(J1_kernel_v, Filter_J1)/F.rRx
            end
            # for HMD
            if F.getazimH
                splreal = CubicSpline(real(F.J1_kernel_h[:,ifreq]), F.log10interpkᵣ)
                splimag = CubicSpline(imag(F.J1_kernel_h[:,ifreq]), F.log10interpkᵣ)
                J1_kernel_h = splreal.(F.log10Filter_base) + 1im*splimag.(F.log10Filter_base)
                F.HFD_az[ifreq] = dot(J1_kernel_h, Filter_J1)/F.rRx
            end
        end
    end # freq loop
end

function getfieldTD!(F::HFieldDHT, z::Array{Float64, 1}, ρ::Array{Float64, 1})
    getfieldFD!(F, z, ρ)
    spl_z_real = CubicSpline(real(F.HFD_z), F.log10ω) # TODO preallocate
    spl_z_imag = CubicSpline(imag(F.HFD_z), F.log10ω) # TODO preallocate
    if F.getradialH
        spl_r_real = CubicSpline(real(F.HFD_r), F.log10ω) # TODO preallocate
        spl_r_imag = CubicSpline(imag(F.HFD_r), F.log10ω) # TODO preallocate
    end
    if F.getazimH
        spl_az_real = CubicSpline(real(F.HFD_az), F.log10ω) # TODO preallocate
        spl_az_imag = CubicSpline(imag(F.HFD_az), F.log10ω) # TODO preallocate
    end
    @views begin
        for itime = 1:length(F.interptimes)
            l10w, H = F.interplog10ω[:,itime], F.Hsc[:,itime]
            t = F.interptimes[itime]
            # Conjugate for my sign convention, apply Butterworth and inverse transform
            if F.provideddt
                # vertical
                F.HFD_z_interp[:] .= imag(conj((spl_z_real.(l10w) .+ 1im*spl_z_imag.(l10w)).*H))*2/pi
                F.HTD_z_interp[itime] = dot(F.HFD_z_interp, Filter_t_sin)/t
                if F.calcjacobian
                    for ilayer = 2:length(ρ)
                        spl_z_J_real = CubicSpline(real(vec(F.HFD_z_J[ilayer,:])), F.log10ω)
                        spl_z_J_imag = CubicSpline(imag(vec(F.HFD_z_J[ilayer,:])), F.log10ω)
                        temp = imag(conj((spl_z_J_real.(l10w) .+ 1im*spl_z_J_imag.(l10w)).*H))*2/pi
                        F.HTD_z_J_interp[ilayer,itime] = dot(temp, Filter_t_sin)/t
                    end
                end
                # radial
                if F.getradialH
                    F.HFD_r_interp[:] .= imag(conj((spl_r_real.(l10w) .+ 1im*spl_r_imag.(l10w)).*H))*2/pi
                    F.HTD_r_interp[itime] = dot(F.HFD_r_interp, Filter_t_sin)/t
                end
                if F.getazimH
                    F.HFD_az_interp[:] .= imag(conj((spl_az_real.(l10w) .+ 1im*spl_az_imag.(l10w)).*H))*2/pi
                    F.HTD_az_interp[itime] = dot(F.HFD_az_interp, Filter_t_sin)/t
                end
            else
                # vertical
                w = F.ω[:,itime]
                F.HFD_z_interp[:] .= -imag(conj((spl_z_real.(l10w) .+ 1im*spl_z_imag.(l10w)).*H)./w)*2/pi
                F.HTD_z_interp[itime] = dot(F.HFD_z_interp, Filter_t_cos)/t
                # radial
                if F.getradialH
                    F.HFD_r_interp[:] .= -imag(conj((spl_r_real.(l10w) .+ 1im*spl_r_imag.(l10w)).*H)./w)*2/pi
                    F.HTD_r_interp[itime] = dot(F.HFD_r_interp, Filter_t_cos)/t
                end
                if F.getazimH
                    F.HFD_az_interp[:] .= -imag(conj((spl_az_real.(l10w) .+ 1im*spl_az_imag.(l10w)).*H)./w)*2/pi
                    F.HTD_az_interp[itime] = dot(F.HFD_az_interp, Filter_t_cos)/t
                end
            end
        end
    end
    if F.doconvramp
        splz = CubicSpline(F.HTD_z_interp, log10.(F.interptimes)) # TODO preallocate
        if F.getradialH
            splr = CubicSpline(F.HTD_r_interp, log10.(F.interptimes)) # TODO preallocate
        else
            splr = splz
        end
        if F.getazimH
            splaz = CubicSpline(F.HTD_az_interp, log10.(F.interptimes)) # TODO preallocate
        else
            splaz = splz
        end
        convramp!(F, splz, splr, splaz, length(ρ))
    end
end

function convramp!(F::HFieldDHT, splz::CubicSpline, splr::CubicSpline, splaz::CubicSpline, nlayers)
    fill!(F.dBzdt, 0.)
    F.calcjacobian && fill!(F.dBzdt_J, 0.)
    F.getradialH && fill!(F.dBrdt, 0.)
    F.getazimH && fill!(F.dBazdt, 0.)
    for itime = 1:length(F.times)
        for iramp = 1:size(F.ramp,1)-1
            rta, rtb  = F.ramp[iramp,1], F.ramp[iramp+1,1]
            dt   = rtb - rta
            dI   = F.ramp[iramp+1,2] - F.ramp[iramp,2]
            dIdt = dI/dt

            if rta > F.times[itime]
                break
            end
            if rtb > F.times[itime] # end in this interval
                rtb = F.times[itime]
            end

            ta = F.times[itime]-rta
            tb = max(F.times[itime]-rtb, 1e-8) # rtb > rta, so make sure this is not zero...
            a, b = log10(ta), log10(tb)
            x, w = F.quadnodes, F.quadweights
            F.dBzdt[itime] += (b-a)/2*dot(getrampresponse((b-a)/2*x .+ (a+b)/2, splz), w)*dIdt
            if F.calcjacobian
                    for ilayer = 2:nlayers
                        splz_J = CubicSpline(F.HTD_z_J_interp[ilayer,:], log10.(F.interptimes)) # TODO preallocate
                        F.dBzdt_J[ilayer,itime] += (b-a)/2*dot(getrampresponse((b-a)/2*x .+ (a+b)/2, splz_J), w)*dIdt
                    end
                end
            if F.getradialH
                F.dBrdt[itime] += (b-a)/2*dot(getrampresponse((b-a)/2*x .+ (a+b)/2, splr), w)*dIdt
            end
            if F.getazimH
                F.dBazdt[itime] += (b-a)/2*dot(getrampresponse((b-a)/2*x .+ (a+b)/2, splaz), w)*dIdt
            end
        end
    end
end

function getrampresponse(t::Array{Float64, 1}, spl::CubicSpline)
    spl.(t).*(10 .^t)*log(10)
end

end # module

max_stars_repo_name: GeoscienceAustralia/HiQGA

max_issues_repo_name: GeoscienceAustralia/HiQGA


text:
################################################################################
#
#   (10,3)a
#   HYPEROCTAGON
#
################################################################################

# Implementation
# - implementation 1
# - labels <: Any
function getUnitcell_10_3_a(
            unitcell_type  :: Type{U},
            implementation :: Val{1}
        ) :: U where {LS,LB,S<:AbstractSite{LS,3},B<:AbstractBond{LB,3},U<:AbstractUnitcell{S,B}}

    # return a new Unitcell
    return newUnitcell(
        # Type of the unitcell
        U,
        # lattice vectors
        Vector{Float64}[
            Float64[1, 0, 0],
            Float64[0.5, 0.5, -0.5],
            Float64[0.5, 0.5, 0.5]
        ],
        # sites
        S[
            newSite(S, Float64[1/8., 1/8.,  1/8.], getDefaultLabelN(LS,1)),
            newSite(S, Float64[5/8., 3/8., -1/8.], getDefaultLabelN(LS,2)),
            newSite(S, Float64[3/8., 1/8., -1/8.], getDefaultLabelN(LS,3)),
            newSite(S, Float64[7/8., 3/8.,  1/8.], getDefaultLabelN(LS,4))
        ],
        # bonds
        B[
            newBond(B, 1,3, getDefaultLabel(LB), (0,0,0)),
            newBond(B, 3,2, getDefaultLabel(LB), (0,0,0)),
            newBond(B, 2,4, getDefaultLabel(LB), (0,0,0)),
            newBond(B, 3,1, getDefaultLabel(LB), (0,0,0)),
            newBond(B, 2,3, getDefaultLabel(LB), (0,0,0)),
            newBond(B, 4,2, getDefaultLabel(LB), (0,0,0)),
            newBond(B, 4,1, getDefaultLabel(LB), ( 1,0,0)),
            newBond(B, 1,4, getDefaultLabel(LB), (-1,0,0)),
            newBond(B, 2,1, getDefaultLabel(LB), (0, 1,0)),
            newBond(B, 1,2, getDefaultLabel(LB), (0,-1,0)),
            newBond(B, 3,4, getDefaultLabel(LB), (0,0,-1)),
            newBond(B, 4,3, getDefaultLabel(LB), (0,0, 1))
        ]
    )
end

# Implementation
# - implementation 4 (Kitaev)
# - labels <: Any
function getUnitcell_10_3_a(
            unitcell_type  :: Type{U},
            implementation :: Val{4}
        ) :: U where {LS,LB,S<:AbstractSite{LS,3},B<:AbstractBond{LB,3},U<:AbstractUnitcell{S,B}}

    # return a new Unitcell
    return newUnitcell(
        # Type of the unitcell
        U,
        # lattice vectors
        Vector{Float64}[
            Float64[1, 0, 0],
            Float64[0.5, 0.5, -0.5],
            Float64[0.5, 0.5, 0.5]
        ],
        # sites
        S[
            newSite(S, Float64[1/8., 1/8.,  1/8.], getDefaultLabelN(LS,1)),
            newSite(S, Float64[5/8., 3/8., -1/8.], getDefaultLabelN(LS,2)),
            newSite(S, Float64[3/8., 1/8., -1/8.], getDefaultLabelN(LS,3)),
            newSite(S, Float64[7/8., 3/8.,  1/8.], getDefaultLabelN(LS,4))
        ],
        # bonds
        B[
            newBond(B, 1,3, getDefaultLabelX(LB), (0,0,0)),
            newBond(B, 3,2, getDefaultLabelZ(LB), (0,0,0)),
            newBond(B, 2,4, getDefaultLabelX(LB), (0,0,0)),
            newBond(B, 3,1, getDefaultLabelX(LB), (0,0,0)),
            newBond(B, 2,3, getDefaultLabelZ(LB), (0,0,0)),
            newBond(B, 4,2, getDefaultLabelX(LB), (0,0,0)),
            newBond(B, 4,1, getDefaultLabelZ(LB), ( 1,0,0)),
            newBond(B, 1,4, getDefaultLabelZ(LB), (-1,0,0)),
            newBond(B, 2,1, getDefaultLabelY(LB), (0, 1,0)),
            newBond(B, 1,2, getDefaultLabelY(LB), (0,-1,0)),
            newBond(B, 3,4, getDefaultLabelY(LB), (0,0,-1)),
            newBond(B, 4,3, getDefaultLabelY(LB), (0,0, 1))
        ]
    )
end

ext: jl


text:
module FloatII
    export interp

    function interp(f, l, u, n)
        # no numerical error caused by adding a lot of small numbers.
        x(idx) = (l*(n-idx)+u*idx)/n
        map( idx -> floor(f(x(idx)),2), 0:(n-1) )
    end

end

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
import Pkg

# @inline function _package_directory(m::Module)
#     return Base.pkgdir(m)
# end

# @inline function _package_directory(name::Symbol)
#     return _package_directory(string(name))
# end

# @inline function _package_directory(name::AbstractString)
#     for pkginfo in values(Pkg.dependencies())
#         if pkginfo.name == name
#             return pkginfo.source
#         end
#     end
#     return nothing
# end


text:
function get_view(
    ts::TableSupport,
    com::ConstraintSolverModel,
    vidx::Int,
    local_vidx::Int,
    val::Int,
)
    val_idx = com.search_space[vidx].init_val_to_index[val + com.search_space[vidx].offset]
    index_shift = ts.var_start[local_vidx] - 1 + val_idx
    return @view ts.values[:, index_shift]
end

function Base.getindex(
    ts::TableSupport,
    com::ConstraintSolverModel,
    vidx::Int,
    local_vidx::Int,
    val::Int,
    row_idx::Int,
)
    val_idx = com.search_space[vidx].init_val_to_index[val + com.search_space[vidx].offset]
    index_shift = ts.var_start[local_vidx] - 1 + val_idx
    return ts.values[row_idx, index_shift]
end


text:
using TR55
using Base.Test

# Test cases from the TR-55 manual example worksheets
# http://www.nrcs.usda.gov/Internet/FSE_DOCUMENTS/stelprdb1044171.pdf

# Figure 2-5
figure_25_regions = [{"area" => 30, "cn" => 61}, {"area" => 70, "cn" => 74}]
figure_25_cn = TR55.weighted_cn(figure_25_regions, true)
figure_25_s = TR55.max_retention(figure_25_cn)
figure_25_q = TR55.runoff(figure_25_s, 6)

@test figure_25_cn == 70
@test figure_25_s == 4.29
@test figure_25_q == 2.8

# Figure 2-6
figure_26_regions = [{"area" => 75, "cn" => 70}, {"area" => 100, "cn" => 80}, {"area" => 75, "cn" => 74}]
figure_26_cn = TR55.weighted_cn(figure_26_regions, false)
figure_26_s = TR55.max_retention(figure_26_cn)
figure_26_q = TR55.runoff(figure_26_s, 6)

@test figure_26_cn == 75
@test figure_26_s == 3.33
@test figure_26_q == 3.28

# Figure 2-7
figure_27_regions = [{"area" => 75, "cn" => 74}, {"area" => 100, "cn" => 82}, {"area" => 75, "cn" => 74}]
figure_27_cn = TR55.weighted_cn(figure_27_regions, false)
figure_27_s = TR55.max_retention(77)
figure_27_q = TR55.runoff(figure_27_s, 6)

@test figure_27_cn == 77
@test figure_27_s == 2.99
@test figure_27_q == 3.48

# Figure 2-8
figure_28_s = TR55.max_retention(74)
figure_28_q = TR55.runoff(figure_28_s, 6)
figure_28_regions = [{"area" => 75, "cn" => 70}, {"area" => 100, "cn" => 78}, {"area" => 75, "cn" => 74}]
figure_28_cn = TR55.weighted_cn(figure_28_regions, false)

@test figure_28_cn == 74
@test figure_28_s == 3.51
@test figure_28_q == 3.19

# Test that weighted CN requires areas to sum to 100 when using percents
test_regions = [{"area" => 20, "cn" => 61}, {"area" => 90, "cn" => 74}]
@test_throws ErrorException TR55.weighted_cn(test_regions, true)


text:
"""
Remove elements from `candidates`` that do not satisfy the constraints in
`constraints`.
"""
function filter_candidates!(candidates::Vector{Candidate}, constraints::Constraints)
    filter!(
        candidate -> satisfies_constraints(candidate, constraints),
        candidates,
    )
end


text:
#=
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http:#arrayfire.com/licenses/BSD-3-Clause
=#
using ArrayFire, Cxx

width = 512; height = 512;
pixels_per_unit = 20;

function simulate(pos, vels, forces, dt)
    pos[1] += vels[1] * pixels_per_unit * dt
    pos[2] += vels[2] * pixels_per_unit * dt

    #calculate distance to center
    diff_x = pos[1] - width/2;
    diff_y = pos[2] - height/2;
    dist = sqrt( diff_x.*diff_x + diff_y.*diff_y )

    #calculate normalised force vectors
    forces[1] = -1 * diff_x ./ dist;
    forces[2] = -1 * diff_y ./ dist;
    #update force scaled to time and magnitude constant
    forces[1] *= pixels_per_unit * dt;
    forces[2] *= pixels_per_unit * dt;

    #dampening
    vels[1] *= 1 - (0.005*dt);
    vels[2] *= 1 - (0.005*dt);

    #update velocities from forces
    vels[1] = vels[1] .+ forces[1];
    vels[2] = vels[2] .+ forces[2];

end

function collisions(pos, vels)
    #clamp particles inside screen border
    projected_px = min(width, max(0, pos[1]))
    projected_py = min(height - 1, max(0, pos[2]))

    #calculate distance to center
    diff_x = projected_px - width/2
    diff_y = projected_py - height/2
    dist = sqrt( diff_x.*diff_x + diff_y.*diff_y )

    #collide with center sphere
    radius = 50;
    elastic_constant = 0.91f0
    dr = dist.<radius
    s = icxx"af::sum<int>($dr);"
    if s > 0
        vels[1][dr] = -elastic_constant * vels[1][dr]
        vels[2][dr] = -elastic_constant * vels[2][dr]

        #normalize diff vector
        diff_x = diff_x ./ dist
        diff_y = diff_y ./ dist
        #place all particle colliding with sphere on surface
        pos[1][dr] = width/2 + diff_x[dr] * radius
        pos[2][dr] = height/2 +  diff_y[dr] * radius
    end
end

function main()
    total_particles = 1000;
    reset = 500;
    frame_count = 0;

    # Initialize the kernel array just once
    # Generate a random starting state
    pos = Any[
        rand(AFArray{Float32}, total_particles) * width,
        rand(AFArray{Float32}, total_particles) * height
    ]

    vels = Any[
        randn(AFArray{Float32}, total_particles),
        randn(AFArray{Float32}, total_particles)
    ]

    forces = Any[
        randn(AFArray{Float32}, total_particles),
        randn(AFArray{Float32}, total_particles)
    ]

    tic();
    for i=1:10
        dt = toq();tic()

        frame_count += 1

        # Generate a random starting state
        if frame_count % reset == 0
            pos = [
                rand(AFArray{Float32}, total_particles) * width,
                rand(AFArray{Float32}, total_particles) * height
            ]
            vels = [
                randn(AFArray{Float32}, total_particles),
                randn(AFArray{Float32}, total_particles)
            ]
        end

        #check for collisions and adjust positions/velocities accordingly
        collisions(pos, vels);

        #run force simulation and update particles
        simulate(pos, vels, forces, dt);

    end
end
main()


text:
module Numpy_format

end


text:
#md # # Tutorial
#nb # # BAT.jl Tutorial
#jl # # BAT.jl Tutorial
#
# This tutorial demonstrates a simple application of BAT.jl: A Bayesian fit
# of a histogram with two Gaussian peaks.
#
#md # You can also download this tutorial as a
#md # [Jupyter notebook](bat_tutorial.ipynb) and a plain
#md # [Julia source file](bat_tutorial.jl).
#
#md # Table of contents:
#md #
#md # ```@contents
#md # Pages = ["tutorial.md"]
#md # Depth = 3
#md # ```
#
# Note: This tutorial is somewhat verbose, as it aims to be easy to follow for
# users who are new to Julia. For the same reason, we deliberately avoid making
# use of Julia features like
# [closures](https://docs.julialang.org/en/v1/devdocs/functions/#Closures-1),
# [anonymous functions](https://docs.julialang.org/en/v1/manual/functions/index.html#man-anonymous-functions-1),
# [broadcasting syntax](https://docs.julialang.org/en/v1/manual/arrays/index.html#Broadcasting-1),
# [performance annotations](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-annotations-1),
# etc.

# ## Input Data Generation
#
# First, let's generate some synthetic data to fit. We'll need the Julia
# standard-library packages
# ["Random"](https://docs.julialang.org/en/v1/stdlib/Random/),
# ["LinearAlgebra"](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/)
# and ["Statistics"](https://docs.julialang.org/en/v1/stdlib/Statistics/),
# as well as the packages
# ["Distributions"](https://juliastats.github.io/Distributions.jl/stable/)
# and ["StatsBase"](http://juliastats.github.io/StatsBase.jl/stable/):

using Random, LinearAlgebra, Statistics, Distributions, StatsBase

# As the underlying truth of our input data/histogram, let us choose an
# non-normalized probability density composed of two Gaussian peaks with a peak
# area of 500 and 1000, a mean of -1.0 and 2.0 and a standard error of 0.5

data = vcat(
    rand(Normal(-1.0, 0.5), 500),
    rand(Normal( 2.0, 0.5), 1000)
)

# resulting in a vector of floating-point numbers:

typeof(data) == Vector{Float64}

# Next, we'll create a histogram of that data, this histogram will serve as
# the input for the Bayesian fit:

hist = append!(Histogram(-2:0.1:4), data)

# Using the Julia ["Plots"](http://docs.juliaplots.org/latest/) package

using Plots

# we can plot the histogram:

plot(
    normalize(hist, mode=:density),
    st = :steps, label = "Data",
    title = "Data"
)
#jl savefig("tutorial-data.pdf")
#md savefig("tutorial-data.pdf")
#md savefig("tutorial-data.svg"); nothing # hide
#md # [![Data](tutorial-data.svg)](tutorial-data.pdf)

# Let's define our fit function - the function that we expect to describe the
# data histogram, at each x-Axis position `x`, depending on a given set `p` of
# model parameters:

function fit_function(p::NamedTuple{(:a, :mu, :sigma)}, x::Real)
    p.a[1] * pdf(Normal(p.mu[1], p.sigma), x) +
    p.a[2] * pdf(Normal(p.mu[2], p.sigma), x)
end
#md nothing # hide

# The fit parameters (model parameters) `a` (peak areas) and `mu` (peak means)
# are vectors, parameter `sigma` (peak width) is a scalar, we assume it's the
# same for both Gaussian peaks.
#
# The true values for the model/fit parameters are the values we used to
# generate the data:

true_par_values = (a = [500, 1000], mu = (-1.0, 2.0), sigma = 0.5)
#md nothing # hide

# Let's visually compare the histogram and the fit function, using these true
# parameter values, to make sure everything is set up correctly:

plot(
    normalize(hist, mode=:density),
    st = :steps, label = "Data",
    title = "Data and True Statistical Model"
)
plot!(
    -4:0.01:4, x -> fit_function(true_par_values, x),
    label = "Truth"
)
#jl savefig("tutorial-data-and-truth.pdf")
#md savefig("tutorial-data-and-truth.pdf")
#md savefig("tutorial-data-and-truth.svg"); nothing # hide
#md # [![Data and True Statistical Model](tutorial-data-and-truth.svg)](tutorial-data-and-truth.pdf)

# ## Bayesian Fit
#
# Now we'll perform a Bayesian fit of the generated histogram, using BAT,
# to infer the model parameters from the data histogram.
#
# In addition to the Julia packages loaded above, we need BAT itself, as
# well as [IntervalSets](https://github.com/JuliaMath/IntervalSets.jl):

using BAT, IntervalSets

# ### Likelihood Definition
#
# First, we need to define the likelihood (function) for our problem.
#
# BAT represents densities like likelihoods and priors as subtypes of
# `BAT.AbstractDensity`. Custom likelihood can be defined by
# creating a new subtype of `AbstractDensity` and by implementing (at minimum)
# `BAT.density_logval` for that type - in complex uses cases, this may become
# necessary. Typically, however, it is sufficient to define a custom
# likelihood as a simple function that returns the log-likelihood value for
# a given set of parameters. BAT will automatically convert such a
# log-likelihood function into a subtype of `AbstractDensity`.
#
# For performance reasons, functions should [not access global variables
# directly] (https://docs.julialang.org/en/v1/manual/performance-tips/index.html#Avoid-global-variables-1).
# So we'll use an [anonymous function](https://docs.julialang.org/en/v1/manual/functions/#man-anonymous-functions-1)
# inside of a [let-statement](https://docs.julialang.org/en/v1/base/base/#let)
# to capture the value of the global variable `hist` in a local variable `h`
# (and to shorten function name `fit_function` to `f`, purely for
# convenience):

log_likelihood = let h = hist, f = fit_function
    params -> begin
        ## Histogram counts for each bin as an array:
        counts = h.weights

        ## Histogram binning, has length (length(counts) + 1):
        binning = h.edges[1]

        ## sum log-likelihood value over bins:
        ll_value::Float64 = 0.0
        for i in eachindex(counts)
            ## Get information about current bin:
            bin_left, bin_right = binning[i], binning[i+1]
            bin_width = bin_right - bin_left
            bin_center = (bin_right + bin_left) / 2

            observed_counts = counts[i]

            ## Simple mid-point rule integration of fit function `f` over bin:
            expected_counts = bin_width * f(params, bin_center)

            ## Add log of Poisson probability for bin:
            ll_value += logpdf(Poisson(expected_counts), observed_counts)
        end

        return ll_value
    end
end

# BAT makes use of Julia's parallel programming facilities if possible, e.g.
# to run multiple Markov chains in parallel. Therefore, log-likelihood
# (and other) code must be thread-safe. Mark non-thread-safe code with
# `@critical` (provided by Julia package `ParallelProcessingTools`).
#
# BAT requires Julia v1.3 or newer to use multi-threading. Support for
# automatic parallelization across multiple (local and remote) Julia processes
# is planned, but not implemented yet.
#
# Note that Julia currently starts only a single thread by default. Set the
# the environment variable
# [`JULIA_NUM_THREADS`](https://docs.julialang.org/en/v1/manual/environment-variables/#JULIA_NUM_THREADS-1)
# to specify the desired number of Julia threads.

# We can evaluate `log_likelihood`, e.g. for the true parameter values:

log_likelihood(true_par_values)

# ### Prior Definition
#
# Next, we need to choose a sensible prior for the fit:

prior = NamedTupleDist(
    a = [0.0..10.0^4, 0.0..10.0^4],
    mu = [-2.0..0.0, 1.0..3.0],
    sigma = Truncated(Normal(0.4, 2), 0.3, 0.7)
)
#md nothing # hide

# In general, BAT allows instances of any subtype of `AbstractDensity` to
# be uses as a prior, as long as a sampler is defined for it. This way, users
# may implement complex application-specific priors. You can also
# use `convert(AbstractDensity, distribution)` to convert any
# continuous multivariate `Distributions.Distribution` to a
# `BAT.AbstractDensity` that can be used as a prior (or likelihood).
#
# The prior also implies the shapes of the parameters:

using ValueShapes

parshapes = valshape(prior)

# These will come in handy later on, e.g. to access (the posterior
# distribution of) individual parameter values.

# ### Bayesian Model Definition
#
# Given the likelihood and prior definition, a `BAT.PosteriorDensity` is simply
# defined via

posterior = PosteriorDensity(log_likelihood, prior)
#md nothing # hide

# ### Parameter Space Exploration via MCMC
#
# We can now use Markov chain Monte Carlo (MCMC) to explore the space of
# possible parameter values for the histogram fit.
#
# To increase the verbosity level of BAT logging output, you may want to set
# the Julia logging level for BAT to debug via `ENV["JULIA_DEBUG"] = "BAT"`.

#nb ENV["JULIA_DEBUG"] = "BAT"
#jl ENV["JULIA_DEBUG"] = "BAT"

# Let's use 4 MCMC chains and require 10^5 unique samples from each chain
# (after tuning/burn-in):

nsamples = 10^4
nchains = 4
#md nothing # hide

# Now we can generate a set of MCMC samples via [`bat_sample`](@ref):

samples, stats = bat_sample(posterior, (nsamples, nchains), MetropolisHastings())
#md nothing # hide
#nb nothing # hide

# Let's calculate some posterior statistics using the function
# [`bat_stats`](@ref) and print the results:

println("Truth: $true_par_values")
println("Mode: $(stats.mode)")
println("Mean: $(stats.mean)")
println("Covariance: $(stats.cov)")

# We can also, e.g., get the Pearson auto-correlation of the parameters:

cor(samples.params, FrequencyWeights(samples.weight))

# ### Visualization of Results

# BAT.jl comes with an extensive set of plotting recipes for
# ["Plots.jl"] (http://docs.juliaplots.org/latest/).
# We can plot the marginalized distribution for a single parameter (e.g.
# parameter 3, i.e. μ₁):

par_names = ["a_1", "a_2", "mu_1", "mu_2", "sigma"]
plot(
    samples, 3,
    mean = true, std_dev = true, globalmode = true, localmode = true,
    nbins = 50, xlabel = par_names[3], ylabel = "P($(par_names[3]))",
    title = "Marginalized Distribution for mu_1"
)
#jl savefig("tutorial-single-par.pdf")
#md savefig("tutorial-single-par.pdf")
#md savefig("tutorial-single-par.svg"); nothing # hide
#md # [![Marginalized Distribution for mu_1](tutorial-single-par.svg)](tutorial-single-par.pdf)

# or plot the marginalized distribution for a pair of parameters (e.g.
# parameters 3 and 5, i.e. μ₁ and σ), including information from the parameter
# stats:

plot(
    samples, (3, 5),
    mean = true, std_dev = true, globalmode = true, localmode = true,
    nbins = 50, xlabel = par_names[3], ylabel = par_names[5],
    title = "Marginalized Distribution for mu_1 and sigma"
)
plot!(BAT.MCMCBasicStats(samples), (3, 5))
#jl savefig("tutorial-param-pair.pdf")
#md savefig("tutorial-param-pair.pdf")
#md savefig("tutorial-param-pair.svg"); nothing # hide
#md # [![Marginalized Distribution for mu_1 and sigma](tutorial-param-pair.svg)](tutorial-param-pair.pdf)

# We can also create an overview plot of the marginalized distribution for all
# pairs of parameters:

plot(
    samples,
    mean = false, std_dev = false, globalmode = true, localmode = false,
    nbins = 50
)
#jl savefig("tutorial-all-params.pdf")
#md savefig("tutorial-all-params.pdf")
#md savefig("tutorial-all-params.svg"); nothing # hide
#md # [![Pairwise Correlation between Parameters](tutorial-all-params.svg)](tutorial-all-params.pdf)

# ### Integration with Tables.jl

# BAT.jl supports the [Tables.jl](https://github.com/JuliaData/Tables.jl)
# interface. So we can also convert the vector of MCMC samples vecto a
# table, e.g. using TypedTables.jl](http://blog.roames.com/TypedTables.jl/stable/):

using TypedTables

tbl = Table(samples)

# Using the parameter shapes, we can generate a table with named parameters,
# instead of flat real-valued parameter vectors:

tbl_named = parshapes.(samples)

# We can now, e.g., find the sample with the maximum posterior density value
# (i.e. the mode)

mode_logdensity, mode_idx = findmax(tbl_named.logdensity)

# and get row `mode_idx` of the table, with all information about the sample
# at the mode:

tbl_named[mode_idx]

# ## Comparison of Truth and Best Fit

# As a final step, we retrieve the parameter values at the mode, representing
# the best-fit parameters

fit_par_values = tbl_named[mode_idx].params

# And plot the truth, data, and best fit:

plot(
    normalize(hist, mode=:density),
    st = :steps, label = "Data",
    title = "Data, True Model and Best Fit"
)
plot!(-4:0.01:4, x -> fit_function(true_par_values, x), label = "Truth")
plot!(-4:0.01:4, x -> fit_function(fit_par_values, x), label = "Best fit")
#jl savefig("tutorial-data-truth-bestfit.pdf")
#md savefig("tutorial-data-truth-bestfit.pdf")
#md savefig("tutorial-data-truth-bestfit.svg"); nothing # hide
#md # [![Data, True Model and Best Fit](tutorial-data-truth-bestfit.svg)](tutorial-data-truth-bestfit.pdf)

# ## Fine-grained control
#
# BAT provides fine-grained control over the MCMC algorithm options, the
# MCMC chain initialization, tuning/burn-in strategy and convergence testing.
# All option value used in the following are the default values, any or all
# may be omitted.

# We'll sample using the The Metropolis-Hastings MCMC algorithm:

algorithm = MetropolisHastings()
#md nothing # hide

# BAT requires a counter-based random number generator (RNG), since it
# partitions the RNG space over the MCMC chains. This way, a single RNG seed
# is sufficient for all chains and results are reproducible even under
# parallel execution. By default, BAT uses a Philox4x RNG initialized with a
# random seed drawn from the
# [system entropy pool](https://docs.julialang.org/en/v1/stdlib/Random/index.html#Random.RandomDevice):

using Random123
rng = Philox4x()
#md nothing # hide

# Other default parameters are:

tuning = AdaptiveMetropolisTuning(
    λ = 0.5,
    α = 0.15..0.35,
    β = 1.5,
    c = 1e-4..1e2
)

convergence = BrooksGelmanConvergence(
    threshold = 1.1,
    corrected = false
)

init = MCMCInitStrategy(
    init_tries_per_chain = 8..128,
    max_nsamples_init = 25,
    max_nsteps_init = 250,
    max_time_init = Inf
)

burnin = MCMCBurninStrategy(
    max_nsamples_per_cycle = 1000,
    max_nsteps_per_cycle = 10000,
    max_time_per_cycle = Inf,
    max_ncycles = 30
)

#md nothing # hide

# To generate MCMC samples with explicit control over all options, use

samples, stats = bat_sample(
    rng, posterior, (nsamples, nchains), algorithm,
    max_nsteps = 10 * nsamples,
    max_time = Inf,
    tuning = tuning,
    init = init,
    burnin = burnin,
    convergence = convergence,
    strict = false,
    filter = true
)
#md nothing # hide
#nb nothing # hide

# However, in many use cases, simply using the default options via
#
# ```julia
# samples, stats = bat_sample(posterior, (nsamples, nchains), MetropolisHastings())
# ```
#
# will often be sufficient.


text:
using HDF5
using JLD

# Define variables of different types
x = 3.7
A = rand(1:20, 3, 5)
str = "Hello"
stringsA = ASCIIString["It", "was", "a", "dark", "and", "stormy", "night"]
stringsU = UTF8String["It", "was", "a", "dark", "and", "stormy", "night"]
empty_string = ""
empty_string_array = ASCIIString[]
empty_array_of_strings = ASCIIString[""]
tf = true
TF = A .> 10
B = randn(2, 4)
AB = Any[A, B]
t = (3, "cat")
c = float32(3)+float32(7)im
C = reinterpret(Complex128, B, (4,))
emptyA = zeros(0,2)
emptyB = zeros(2,0)
try
    global MyStruct
    type MyStruct
        len::Int
        data::Array{Float64}
        MyStruct(len::Int) = new(len)
        MyStruct(len::Int, data::Array{Float64}) = new(len, data)
    end
catch
end
ms = MyStruct(2, [3.2, -1.7])
msempty = MyStruct(5, Float64[])
sym = :TestSymbol
syms = [:a, :b]
d = Dict(syms, ["aardvark", "banana"])
ex = quote
    function incrementby1(x::Int)
        x+1
    end
end
T = Uint8
char = 'x'
unicode_char = '\U10ffff'
α = 22
β = Any[[1, 2], [3, 4]]  # issue #93
typevar = Array{Int}[[1]]
typevar_lb = Vector{TypeVar(:U, Integer)}[[1]]
typevar_ub = Vector{TypeVar(:U, Int, Any)}[[1]]
typevar_lb_ub = Vector{TypeVar(:U, Int, Real)}[[1]]
undef = cell(1)
undefs = cell(2, 2)
ms_undef = MyStruct(0)
# Unexported type:
cpus = Base.Sys.cpu_info()
# Immutable type:
rng = 1:5
# Type with a pointer field (#84)
objwithpointer = r"julia"
# Custom BitsType (#99)
bitstype 64 MyBT
bt = reinterpret(MyBT, 55)
# Symbol arrays (#100)
sa_asc = [:a, :b]
sa_utf8 = [:α, :β]
# SubArray (to test tuple type params)
subarray = sub([1:5], 1:5)
# Array of empty tuples (to test tuple type params)
arr_empty_tuple = ()[]

iseq(x,y) = isequal(x,y)
iseq(x::MyStruct, y::MyStruct) = (x.len == y.len && x.data == y.data)
iseq(c1::Array{Base.Sys.CPUinfo}, c2::Array{Base.Sys.CPUinfo}) = length(c1) == length(c2) && all([iseq(c1[i], c2[i]) for i = 1:length(c1)])
function iseq(c1::Base.Sys.CPUinfo, c2::Base.Sys.CPUinfo)
    for n in Base.Sys.CPUinfo.names
        if getfield(c1, n) != getfield(c2, n)
            return false
        end
    end
    true
end
macro check(fid, sym)
    ex = quote
        let tmp
            try
                tmp = read($fid, $(string(sym)))
            catch e
                warn("Error reading ", $(string(sym)))
                rethrow(e)
            end
            if !iseq(tmp, $sym)
                error("For ", $(string(sym)), ", read value does not agree with written value")
            end
            written_type = typeof($sym)
            if typeof(tmp) != written_type
                error("For ", $(string(sym)), ", read type $(typeof(tmp)) does not agree with written type $(written_type)")
            end
        end
    end
    esc(ex)
end

# Test for equality of expressions, skipping line numbers
checkexpr(a, b) = @assert a == b
function checkexpr(a::Expr, b::Expr)
    @assert a.head == b.head
    i = 1
    j = 1
    while i <= length(a.args) && j <= length(b.args)
        if isa(a.args[i], Expr) && a.args[i].head == :line
            i += 1
            continue
        end
        if isa(b.args[j], Expr) && b.args[j].head == :line
            j += 1
            continue
        end
        checkexpr(a.args[i], b.args[j])
        i += 1
        j += 1
    end
    @assert i >= length(a.args) && j >= length(b.args)
end

fn = joinpath(tempdir(),"test.jld")

fid = jldopen(fn, "w")
@write fid x
@write fid A
@write fid str
@write fid stringsA
@write fid stringsU
@write fid empty_string
@write fid empty_string_array
@write fid empty_array_of_strings
@write fid tf
@write fid TF
@write fid AB
@write fid t
@write fid c
@write fid C
@write fid emptyA
@write fid emptyB
@write fid ms
@write fid msempty
@write fid sym
@write fid syms
@write fid d
@write fid ex
@write fid T
@write fid char
@write fid unicode_char
@write fid α
@write fid β
@write fid cpus
@write fid rng
@write fid typevar
@write fid typevar_lb
@write fid typevar_ub
@write fid typevar_lb_ub
@write fid undef
@write fid undefs
@write fid ms_undef
@write fid objwithpointer  # This should not write anything
@write fid bt
@write fid sa_asc
@write fid sa_utf8
@write fid subarray
@write fid arr_empty_tuple
# Make sure we can create groups (i.e., use HDF5 features)
g = g_create(fid, "mygroup")
i = 7
@write g i
write(fid, "group1/x", {1})
write(fid, "group2/x", {2})
close(fid)

for mmap = (true, false)
    fidr = jldopen(fn, "r", mmaparrays=mmap)
    @check fidr x
    @check fidr A
    @check fidr str
    @check fidr stringsA
    @check fidr stringsU
    @check fidr empty_string
    @check fidr empty_string_array
    @check fidr empty_array_of_strings
    @check fidr tf
    @check fidr TF
    @check fidr AB
    @check fidr t
    @check fidr c
    @check fidr C
    @check fidr emptyA
    @check fidr emptyB
    @check fidr ms
    @check fidr msempty
    @check fidr sym
    @check fidr syms
    @check fidr d
    exr = read(fidr, "ex")   # line numbers are stripped, don't expect equality
    checkexpr(ex, exr)
    @check fidr T
    @check fidr char
    @check fidr unicode_char
    @check fidr α
    @check fidr β
    @check fidr cpus
    @check fidr rng
    @check fidr typevar
    @check fidr typevar_lb
    @check fidr typevar_ub
    @check fidr typevar_lb_ub

    # Special cases for reading undefs
    undef = read(fidr, "undef")
    if !isa(undef, Array{Any, 1}) || length(undef) != 1 || isdefined(undef, 1)
        error("For undef, read value does not agree with written value")
    end
    undefs = read(fidr, "undefs")
    if !isa(undefs, Array{Any, 2}) || length(undefs) != 4 || any(map(i->isdefined(undefs, i), 1:4))
        error("For undefs, read value does not agree with written value")
    end
    ms_undef = read(fidr, "ms_undef")
    if !isa(ms_undef, MyStruct) || ms_undef.len != 0 || isdefined(ms_undef, :data)
        error("For ms_undef, read value does not agree with written value")
    end

    @assert !in("objwithpointer", names(fidr))
    @check fidr bt
    @check fidr sa_asc
    @check fidr sa_utf8
    @check fidr subarray
    @check fidr arr_empty_tuple

    x1 = read(fidr, "group1/x")
    @assert x1 == {1}
    x2 = read(fidr, "group2/x")
    @assert x2 == {2}

    close(fidr)
end

# do syntax
jldopen(fn, "w") do fid
    g_create(fid, "mygroup") do g
        write(g, "x", 3.2)
    end
end
fid = jldopen(fn, "r")
@assert names(fid) == ASCIIString["mygroup"]
g = fid["mygroup"]
@assert names(g) == ASCIIString["x"]
@assert read(g, "x") == 3.2
close(g)
close(fid)

# Function load() and save() syntax
d = ["x"=>3.2, "β"=>β, "A"=>A]
save(fn, d)
d2 = load(fn)
@assert d == d2
β2 = load(fn, "β")
@assert β == β2
β2, A2 = load(fn, "β", "A")
@assert β == β2
@assert A == A2

save(fn, "x", 3.2, "β", β, "A", A)
d3 = load(fn)
@assert d == d3

# #71
jldopen(fn, "w") do file
    file["a"] = 1
end
jldopen(fn, "r") do file
    @assert read(file, "a") == 1
end


text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

"""
    ln_mean(x, y)

Compute the logarithmic mean

    ln_mean(x, y) = (y - x) / (log(y) - log(x)) = (y - x) / log(y / x)

Problem: The formula above has a removable singularity at `x == y`. Thus,
some care must be taken to implement it correctly without problems or loss
of accuracy when `x ≈ y`. Here, we use the approach proposed by
Ismail and Roe (2009).
Set ξ = y / x. Then, we have

    (y - x) / log(y / x) = (x + y) / log(ξ) * (ξ - 1) / (ξ + 1)

Set f = (ξ - 1) / (ξ + 1) = (y - x) / (x + y). Then, we use the expansion

    log(ξ) = 2 * f * (1 + f^2 / 3 + f^4 / 5 + f^6 / 7) + O(ξ^9)

Inserting the first few terms of this expansion yields

    (y - x) / log(ξ) ≈ (x + y) * f / (2 * f * (1 + f^2 / 3 + f^4 / 5 + f^6 / 7))
                     = (x + y) / (2 + 2/3 * f^2 + 2/5 * f^4 + 2/7 * f^6)

Since divisions are usually more expensive on modern hardware than
multiplications (Agner Fog), we try to avoid computing two divisions. Thus,
we use

    f^2 = (y - x)^2 / (x + y)^2
        = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y)

Given ε = 1.0e-4, we use the following algorithm.

    if f^2 < ε
      # use the expansion above
    else
      # use the direct formula (y - x) / log(y / x)
    end

# References
- Ismail, Roe (2009).
  Affordable, entropy-consistent Euler flux functions II: Entropy production at shocks.
  [DOI: 10.1016/j.jcp.2009.04.021](https://doi.org/10.1016/j.jcp.2009.04.021)
- Agner Fog.
  Lists of instruction latencies, throughputs and micro-operation breakdowns
  for Intel, AMD, and VIA CPUs.
  https://www.agner.org/optimize/instruction_tables.pdf
"""
@inline function ln_mean(x, y)
  epsilon_f2 = 1.0e-4
  f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
  if f2 < epsilon_f2
    return (x + y) / @evalpoly(f2, 2, 2/3, 2/5, 2/7)
  else
    return (y - x) / log(y / x)
  end
end

"""
    inv_ln_mean(x, y)

Compute the inverse `1 / ln_mean(x, y)` of the logarithmic mean
[`ln_mean`](@ref).

This function may be used to increase performance where the inverse of the
logarithmic mean is needed, by replacing a (slow) division by a (fast)
multiplication.
"""
@inline function inv_ln_mean(x, y)
  epsilon_f2 = 1.0e-4
  f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
  if f2 < epsilon_f2
    return @evalpoly(f2, 2, 2/3, 2/5, 2/7) / (x + y)
  else
    return log(y / x) / (y - x)
  end
end

# `Base.max` and `Base.min` perform additional checks for signed zeros and `NaN`s
# which are not present in comparable functions in Fortran/C++. For example,
# ```julia
# julia> @code_native debuginfo=:none syntax=:intel max(1.0, 2.0)
#         .text
#         vmovq   rcx, xmm1
#         vmovq   rax, xmm0
#         vcmpordsd       xmm2, xmm0, xmm0
#         vblendvpd       xmm2, xmm0, xmm1, xmm2
#         vcmpordsd       xmm3, xmm1, xmm1
#         vblendvpd       xmm3, xmm1, xmm0, xmm3
#         vmovapd xmm4, xmm2
#         test    rcx, rcx
#         jns     L45
#         vmovapd xmm4, xmm3
# L45:
#         test    rax, rax
#         js      L54
#         vmovapd xmm4, xmm3
# L54:
#         vcmpltsd        xmm0, xmm0, xmm1
#         vblendvpd       xmm0, xmm4, xmm2, xmm0
#         ret
#         nop     word ptr cs:[rax + rax]
#
# julia> @code_native debuginfo=:none syntax=:intel min(1.0, 2.0)
#         .text
#         vmovq   rcx, xmm1
#         vmovq   rax, xmm0
#         vcmpordsd       xmm2, xmm0, xmm0
#         vblendvpd       xmm2, xmm0, xmm1, xmm2
#         vcmpordsd       xmm3, xmm1, xmm1
#         vblendvpd       xmm3, xmm1, xmm0, xmm3
#         vmovapd xmm4, xmm2
#         test    rcx, rcx
#         jns     L58
#         test    rax, rax
#         js      L67
# L46:
#         vcmpltsd        xmm0, xmm1, xmm0
#         vblendvpd       xmm0, xmm4, xmm2, xmm0
#         ret
# L58:
#         vmovapd xmm4, xmm3
#         test    rax, rax
#         jns     L46
# L67:
#         vmovapd xmm4, xmm3
#         vcmpltsd        xmm0, xmm1, xmm0
#         vblendvpd       xmm0, xmm4, xmm2, xmm0
#         ret
#         nop     word ptr cs:[rax + rax]
#         nop     dword ptr [rax]
# ```
# In contrast, we get the much simpler and faster version
# ```julia
# julia> @code_native debuginfo=:none syntax=:intel Base.FastMath.max_fast(1.0, 2.0)
#         .text
#         vmaxsd  xmm0, xmm1, xmm0
#         ret
#         nop     word ptr cs:[rax + rax]
#
# julia> @code_native debuginfo=:none syntax=:intel Base.FastMath.min_fast(1.0, 2.0)
#         .text
#         vminsd  xmm0, xmm0, xmm1
#         ret
#         nop     word ptr cs:[rax + rax]
# ```
# when using `@fastmath`, which we also get from
# [Fortran](https://godbolt.org/z/Yrsa1js7P)
# or [C++](https://godbolt.org/z/674G7Pccv).
"""
    max(x, y, ...)

Return the maximum of the arguments. See also the `maximum` function to take
the maximum element from a collection.

This version in Trixi.jl is semantically equivalent to `Base.max` but may
be implemented differently. In particular, it may avoid potentially expensive
checks necessary in the presence of `NaN`s (or signed zeros).

# Examples

julia> max(2, 5, 1)
5
"""
@inline max(args...) = @fastmath max(args...)

"""
    min(x, y, ...)

Return the minimum of the arguments. See also the `minimum` function to take
the minimum element from a collection.

This version in Trixi.jl is semantically equivalent to `Base.min` but may
be implemented differently. In particular, it may avoid potentially expensive
checks necessary in the presence of `NaN`s (or signed zeros).

# Examples

julia> min(2, 5, 1)
1
"""
@inline min(args...) = @fastmath min(args...)

end # @muladd


text:
# Elementwise kernels on structures of arrays.

using Base.Meta

"""
Represents the element at index/indices `i` in container `o`.
"""
struct Elem{T,N}
    o::T
    i::Tuple{Vararg{Int,N}}
end
Elem(T,i::Vararg{Int}) = Elem(T,i)
export Elem

"""
Main entry point:

@resolve(quote
    a += b_post * k_pre
end, obj.group, post=group2)
"""
macro resolve(code, subs...)
    # The purpose of this macro is to call resolve() with
    # information about the names of the arguments.
    #
    # The above snippet becomes:
    #
    #   resolve(quote
    #       a += b_post
    #   end, :_=>(:(obj.group),obj.group), :post=>(:group2, group2))
    #
    mappings = []
    for sub in subs
        suffix, sym = if isexpr(sub, :(=))
            # eg. post=group2
            QuoteNode(sub.args[1]), sub.args[2]
        else
            # eg. obj.group
            QuoteNode(:_), sub
        end
        push!(mappings, :($suffix=>($(QuoteNode(sym)), $sym)))
    end
    esc(quote
        SpikeNet.resolve($code, $(mappings...))
    end)
end

export @resolve

function resolve(code, mappings...)

    # First find out which symbols to expand:
    subs = Dict{Symbol,Expr}()
    for (suffix::Symbol, (elem::Symbol, T::DataType)) in mappings
        # Here's a struct -- let's expand the field names,
        # with an optional suffix.
        make_subs(T, elem, suffix, subs)
    end
    # Apply the substitutions in the dict to the code block:
    ast = map_symbols(code, subs)
    # Core.println(ast)

    # Always inline this function. This is critical for speed,
    # because it lets the compiler simplify away all these Elem
    # wrappers, avoiding memory allocation in the elemwise loop.
    return quote
        $(Expr(:meta, :inline))
        @fastmath $ast
    end
    # Currently there's no way to let the caller decide:
    # @inline @generated f(...)=@resolve(...) won't work!
end

function matched(field, suffix)
    # Figure out which symbol we'll be replacing:
    if suffix == :_
        Symbol("$(field)")
    else
        Symbol("$(field)_$(suffix)")
    end
end

function make_subs(T::DataType, elem::Symbol, suffix::Symbol, subs)
    for field in fieldnames(T)
        subs[matched(field, suffix)] = :($elem.$field)
    end
end

function make_subs(E::Type{<:Elem}, elem::Symbol, suffix::Symbol,
                   subs)
    T = E.parameters[1] # E = Elem{T}
    for field in fieldnames(T)
        # Decide what to do based on the type of the field:
        F = fieldtype(T, field)
        if issubtype(F, AbstractArray)
            # Arrays contain element values:
            subs[matched(field, suffix)] = :($elem.o.$field[CartesianIndex($elem.i)])
            # TODO check that dimensions match?
        else
            # Scalar values are the same for every element:
            subs[matched(field, suffix)] = :($elem.o.$field)
        end
    end
end

"""
Makes a deep copy of an expression tree, replacing symbols
that have a mapping in `subs` with their target values.
"""
map_symbols(expr::Expr, subs::Dict) = Expr(expr.head,
    [map_symbols(a, subs) for a in expr.args]...)
map_symbols(s::Symbol, subs::Dict) = get(subs, s, s)
map_symbols(obj::Any, subs::Dict) = obj

# Note that map_symbols() will not match the symbol :x in
# :(f.x), because that :x will be inside a QuoteNode and
# we don't look inside QuoteNodes. Nonetheless, they will
# match symbols inside scope-defining blocks (eg. function
# definitions) which is probably not what we want. TODO
# handle these cases when it becomes necessary.


text:
# # Certificate

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Extension/certificate.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Extension/certificate.ipynb)
# **Contributed by**: Benoît Legat

# ## Introduction

# Consider the polynomial optimization problem (a variation from [L09, Example 2.2]) of
# minimizing the polynomial $x^3 - x^2 + 2xy - y^2 + y^3$
# over the polyhedron defined by the inequalities $x \ge 0, y \ge 0$ and $x + y \geq 1$.

# [L09] Lasserre, J. B.
# *Moments, positive polynomials and their applications*.
# World Scientific, **2009**.

using Test #src
using DynamicPolynomials
@polyvar x y
p = x^3 - x^2 + 2x*y -y^2 + y^3 + x^3 * y
using SumOfSquares
S = @set x >= 0 && y >= 0 && x^2 + y^2 >= 2

# We will now see how to find the optimal solution using Sum of Squares Programming.
# We first need to pick an SDP solver, see [here](https://jump.dev/JuMP.jl/v0.21.6/installation/#Supported-solvers) for a list of the available choices.

import CSDP
solver = optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true)

# A Sum-of-Squares certificate that $p \ge \alpha$ over the domain `S`, ensures that $\alpha$ is a lower bound to the polynomial optimization problem.
# The following program searches for the largest upper bound and finds zero.

model = SOSModel(solver)
@variable(model, α)
@objective(model, Max, α)
@constraint(model, c, p >= α, domain = S)
optimize!(model)
@show termination_status(model)
@show objective_value(model)

# We now define the Schmüdgen's certificate:

using MultivariateBases
const MB = MultivariateBases
const SOS = SumOfSquares
const SOSC = SOS.Certificate
struct Schmüdgen{IC <: SOSC.AbstractIdealCertificate, CT <: SOS.SOSLikeCone, BT <: SOS.AbstractPolynomialBasis} <: SOSC.AbstractPreorderCertificate
    ideal_certificate::IC
    cone::CT
    basis::Type{BT}
    maxdegree::Int
end

SOSC.cone(certificate::Schmüdgen) = certificate.cone

function SOSC.preprocessed_domain(::Schmüdgen, domain::BasicSemialgebraicSet, p)
    return SOSC.with_variables(domain, p)
end

function SOSC.preorder_indices(::Schmüdgen, domain::SOSC.WithVariables)
    n = length(domain.inner.p)
    if n >= Sys.WORD_SIZE
        error("There are $(2^n - 1) products in Schmüdgen's certificate, they cannot even be indexed with `$Int`.")
    end
    return map(SOSC.PreorderIndex, 1:(2^n-1))
end

function SOSC.multiplier_basis(certificate::Schmüdgen, index::SOSC.PreorderIndex, domain::SOSC.WithVariables)
    q = SOSC.generator(certificate, index, domain)
    return SOSC.maxdegree_gram_basis(certificate.basis, variables(domain), SOSC.multiplier_maxdegree(certificate.maxdegree, q))
end
function SOSC.multiplier_basis_type(::Type{Schmüdgen{IC, CT, BT}}) where {IC, CT, BT}
    return BT
end

function SOSC.generator(::Schmüdgen, index::SOSC.PreorderIndex, domain::SOSC.WithVariables)
    I = [i for i in eachindex(domain.inner.p) if !iszero(index.value & (1 << (i - 1)))]
    return prod([domain.inner.p[i] for i in eachindex(domain.inner.p) if !iszero(index.value & (1 << (i - 1)))])
end

SOSC.ideal_certificate(certificate::Schmüdgen) = certificate.ideal_certificate
SOSC.ideal_certificate(::Type{<:Schmüdgen{IC}}) where {IC} = IC

SOS.matrix_cone_type(::Type{<:Schmüdgen{IC, CT}}) where {IC, CT} = SOS.matrix_cone_type(CT)

# Let's try again with our the Schmüdgen certificate:

model = SOSModel(solver)
@variable(model, α)
@objective(model, Max, α)
ideal_certificate = SOSC.Newton(SOSCone(), MB.MonomialBasis, tuple())
certificate = Schmüdgen(ideal_certificate, SOSCone(), MB.MonomialBasis, maxdegree(p))
@constraint(model, c, p >= α, domain = S, certificate = certificate)
optimize!(model)
@show termination_status(model)
@show objective_value(model)
@test length(lagrangian_multipliers(c)) == 7 #src

max_stars_repo_name: blegat/SumOfSquares

max_issues_repo_name: blegat/SumOfSquares

max_forks_repo_name: blegat/SumOfSquares


text:
using Revise
using TensorKit
using statmech_tm_solver
using Printf

Γ = 1.0
T = cmpo_ising_realtime(Γ)

println("========================================================================")

psi = qbimps(rand, 4, 2, 1)
for chi in 8:8:40
    psi = expand(psi, chi, 1e-9)
    println("===========================  chi = $chi   ==============================")
    ix, δ = 0, 42.
    while ix < 250 && δ > 1e-12
        ΛA, psi = A_canonical(T, psi)
        ΛB, psi = B_canonical(T, psi)
        ix += 1
        δ = norm(ΛA - ΛB) / norm(ΛA + ΛB)
        println("$ix, ΛA: $ΛA, ΛB: $ΛB, δ: $δ")
    end
end


text:
baremodule TopParser

function processes end
function processes_from_file end

function samples end
function open_samples end

module Internal

using Dates: Second

using ..TopParser: TopParser

include("parser.jl")
include("api.jl")

end  # module Internal

end  # baremodule TopParser


text:
using ThreadsAPIBenchmarks
using Test

@testset "ThreadsAPIBenchmarks.jl" begin
    # Write your tests here.
end


text:
#=
  Two children problem in Turing.jl

  https://robeastaway.com/blog/boy-girl
  """
  I have two children, at least one of whom is a boy.
  What is the chance that the second child is also a boy?

  ...

  The classic, surprising answer is that the chance of a second
  boy is 1/3, not 1/2 as most would expect.  Why?  Because if I
  have two children, there is an equal chance that they will be
  Boy-Boy, Boy-Girl, Girl-Boy or Girl-Girl.  There are three
  equally likely combinations in which at least one child is
  a boy (BB, BG and GB) and only in the first scenario is the
  other child also a boy.
  """

  We assume the inital probability of a boy/girl as 0.5.

  Here we model two similiar but different problems to show
  the difference between these two views of the problem.

  Cf ~/webppl/two_children_problem.wppl

=#

using Turing
include("jl_utils.jl")

#=
  First model:
  """
  I have two children, _at least one of whom_ is a boy.
  What is the chance that the second child is also a boy?
  """

  Probability of 2 boys: 1/3.
  Probability of 1 boy and 1 girl: 2/3.

  Distributions of variable twoBoys (num:0)
  0.00000 =>    6601  (0.660100)
  1.00000 =>    3399  (0.339900)

  Distributions of variable oneBoyOneGirl (num:0)
  1.00000 =>    6601  (0.660100)
  0.00000 =>    3399  (0.339900)

=#
@model function two_children_problem_1()
    boy = 1
    girl = 2

    child1 ~ Categorical([0.5,0.5])
    child2 ~ Categorical([0.5,0.5])

    # Note: it's important that we state the condition
    #   "at least one of whom is a boy"
    # as something like this.
    true ~ Dirac((child1 == boy) + (child2 == boy) >= 1)

    twoBoys ~ Dirac((child1 == boy) + (child2 == boy) == 2)
    oneBoyOneGirl ~ Dirac((child1 == boy) + (child2 == boy) == 1)

end

println("\nModel 1")
model = two_children_problem_1()
# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, IS(), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, NUTS(), 10_000)
# chns = sample(model, HMC(0.1,10), 10_000)

# display(chns)
show_var_dist_pct(chns,:twoBoys)
show_var_dist_pct(chns,:oneBoyOneGirl)

#=
 Second model:
  """
  I have two children, the eldest is a boy.
  What is the chance that the second child is also a boy?
  """

  Probability that the second child is also a boy: 1/2
  Probability that the second child is a girl: 1/2

  Distributions of variable twoBoys (num:0)
  0.00000 =>    5067  (0.506700)
  1.00000 =>    4933  (0.493300)
  Distributions of variable oneBoyOneGirl (num:0)
  1.00000 =>    5067  (0.506700)
  0.00000 =>    4933  (0.493300)

=#
@model function two_children_problem_2()
    boy = 1
    girl = 2

    child1 ~ Categorical([0.5,0.5])
    child2 ~ Categorical([0.5,0.5])

    # "The eldest is a boy"
    true ~ Dirac(child1 == boy)

    twoBoys ~ Dirac((child1 == boy) + (child2 == boy) == 2)
    oneBoyOneGirl ~ Dirac((child1 == boy) + (child2 == boy) == 1)

end

println("\nModel 2")
model = two_children_problem_2()
# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, IS(), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, NUTS(), 10_000)
# chns = sample(model, HMC(0.1,10), 10_000)

# display(chns)
show_var_dist_pct(chns,:twoBoys)
show_var_dist_pct(chns,:oneBoyOneGirl)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
module DiffEqParamEstim
using DiffEqBase, LsqFit, PenaltyFunctions,
      RecursiveArrayTools, ForwardDiff, Calculus, Distributions, BlackBoxOptim, LinearAlgebra

STANDARD_PROB_GENERATOR(prob,p) = remake(prob;u0=eltype(p).(prob.u0),p=p)
STANDARD_PROB_GENERATOR(prob::MonteCarloProblem,p) = MonteCarloProblem(remake(prob.prob;u0=eltype(p).(prob.prob.u0),p=p))

include("cost_functions.jl")
include("lm_fit.jl")
include("build_loss_objective.jl")
include("build_lsoptim_objective.jl")
include("two_stage_method.jl")
include("kernels.jl")
include("multiple_shooting_objective.jl")

end # module


text:
using Neuroimaging
using Test
using Logging
using MAT, BDF
using Plots
using Glob
using DataDeps

unicodeplots()

#
# Run all tests
#

println(pwd())
tests = glob("**/*.jl")
tests = tests[.~(tests .== "runtests.jl")]

@info tests
for t in tests
    include(t)
end


text:
### Day 16
## Dragon Checksum
## Author: Thanasis Georgiou

# Reset workspace to shut up warnings
workspace()

# Set disk length here
#diskLength = 272 # Problem 1
diskLength = 35651584 # Problem 2
input = collect("01111010110010011")

data = Array(Char, diskLength)
for i = 1:length(input)
  data[i] = input[i]
end
data[end] = 'e' # When this changes, the whole disk is overwritten

# Generate data
i = length(input)
while data[end] == 'e'
  currentData = data[i:-1:1] # Reverse current data
  map!(c -> c == '1'? '0' : '1', currentData)

  # Add separator
  data[i + 1] = '0'
  i += 2

  # Calculate length of new data
  newLen = length(currentData)
  if newLen > diskLength - i + 1
    newLen = diskLength - i + 1
  end

  # Add data to disk
  data[i:(i + newLen - 1)] = currentData[1:newLen]
  i += newLen - 1
end

# Calculate checksum
function calculateChecksum(checksum)
  # New checksum will have half the length of the original
  newChecksum = Array(Char, convert(Int, length(checksum) / 2))

  # Calculate checksum
  for i = 1:length(newChecksum)
    j = (i * 2) - 1
    newChecksum[i] = checksum[j] == checksum[j + 1]? '1' : '0'
  end

  # Return it back
  return newChecksum
end

checksum = calculateChecksum(data)
while iseven(length(checksum))
  checksum = calculateChecksum(checksum)
end

# Print result
println("Checksum: $(String(checksum))")


text:
using FewShotAnomalyDetection
using Flux
using MLDataPattern
using FluxExtensions
using Adapt
using DataFrames
using CSV
using Serialization

import FewShotAnomalyDetection: loss, zparams
include("experimentalutils.jl")
include("vae.jl")

outputFolder = mainfolder * "experiments/twostagesvae_scoretest_wide/"
mkpath(outputFolder)

function runExperiment(datasetName, train, test, inputDim, hiddenDim, latentDim, numLayers, nonlinearity, layerType, β, batchSize = 100, numBatches = 10000, i = 0)

    T = Float32

    encoder = Adapt.adapt(T, FluxExtensions.layerbuilder(inputDim, hiddenDim, (latentDim - 1) * 2, numLayers + 1, nonlinearity, "linear", layerType))
    decoder = Adapt.adapt(T, FluxExtensions.layerbuilder((latentDim - 1), hiddenDim, inputDim + 1, numLayers + 1, nonlinearity, "linear", layerType))
    outerVAE = VAE(encoder, decoder, β, :scalarsigma)
    opt = Flux.Optimise.ADAM(1e-4)
    cb = Flux.throttle(() -> println("$datasetName outer VAE: $(printingloss(outerVAE, train[1]))"), 5)
    Flux.train!((x, y) -> loss(outerVAE, x), Flux.params(outerVAE), RandomBatches((train[1], zero(train[2])), batchSize, numBatches), opt, cb = cb)

    svae = SVAEtwocaps(latentDim - 1, hiddenDim, latentDim, numLayers, nonlinearity, layerType, :scalarsigma, T)
    FewShotAnomalyDetection.set_normal_μ_nonparam(svae, vcat(T(1), zeros(latentDim - 1)))
    learnRepresentation!(data, labels) = wloss(svae, data, (x, y) -> FewShotAnomalyDetection.mmd_imq(x, y, 1))
    printing_learnRepresentation!(data, labels) = printing_wloss(svae, data, (x, y) -> FewShotAnomalyDetection.mmd_imq(x, y, 1))
    opt = Flux.Optimise.ADAM(1e-4)
    cb = Flux.throttle(() -> println("$datasetName inner SVAE: $(printing_learnRepresentation!(samplez(outerVAE, train[1]), zero(train[2])))"), 5)
    println("Before training - $datasetName inner SVAE: $(printing_learnRepresentation!(samplez(outerVAE, train[1]), zero(train[2])))")
    Flux.train!((x, y) -> learnRepresentation!(samplez(outerVAE, x), y), Flux.params(svae), RandomBatches((train[1], zero(train[2])), batchSize, numBatches), opt, cb = cb)

    data = Flux.Tracker.data(zparams(outerVAE, test[1])[1])
    log_pxv = vec(collect(.-log_pxexpectedz(svae, data)'))
    log_pzs = vec(collect(.-log_pz(svae, data)'))
    log_pxis = vec(collect(.-log_px(svae, data)'))
    log_pzs_jacobian = vec(collect(.-log_pz_jacobian(svae, data)'))

    cz = closestz(svae, data)
    mz = manifoldz(svae, data)

    log_pxv_cz = vec(collect(.-log_pxexpectedz(svae, data, cz)'))
    log_pzs_cz = vec(collect(.-log_pz_from_z(svae, cz)'))
    log_pzs_jacobian_cz = vec(collect(.-log_pz_jacobian(svae, data, cz)'))

    log_pxv_mz = vec(collect(.-log_pxexpectedz(svae, data, mz)'))
    log_pzs_mz = vec(collect(.-log_pz_from_z(svae, mz)'))
    log_pzs_jacobian_mz = vec(collect(.-log_pz_jacobian(svae, data, mz)'))

    auc_pxv = computeauc(log_pxv, test[2] .- 1)
    auc_pz = computeauc(log_pzs, test[2] .- 1)
    auc_pxis = computeauc(log_pxis, test[2] .- 1)
    auc_pz_jacobian = computeauc(log_pzs_jacobian, test[2] .- 1)
    auc_pxv_pz = computeauc(log_pxv .+ log_pzs, test[2] .- 1)
    auc_pxv_pz_jacobian = computeauc(log_pxv .+ log_pzs_jacobian, test[2] .- 1)

    auc_pxv_cz = computeauc(log_pxv_cz, test[2] .- 1)
    auc_pz_cz = computeauc(log_pzs_cz, test[2] .- 1)
    auc_pz_jacobian_cz = computeauc(log_pzs_jacobian_cz, test[2] .- 1)
    auc_pxv_pz_cz = computeauc(log_pxv_cz .+ log_pzs_cz, test[2] .- 1)
    auc_pxv_pz_jacobian_cz = computeauc(log_pxv_cz .+ log_pzs_jacobian_cz, test[2] .- 1)

    auc_pxv_mz = computeauc(log_pxv_mz, test[2] .- 1)
    auc_pz_mz = computeauc(log_pzs_mz, test[2] .- 1)
    auc_pz_jacobian_mz = computeauc(log_pzs_jacobian_mz, test[2] .- 1)
    auc_pxv_pz_mz = computeauc(log_pxv_mz .+ log_pzs_mz, test[2] .- 1)
    auc_pxv_pz_jacobian_mz = computeauc(log_pxv_mz .+ log_pzs_jacobian_mz, test[2] .- 1)

    serialize(outputFolder * "twostagesvae-$datasetName-$i-svae.jls", svae)
    serialize(outputFolder * "twostagesvae-$datasetName-$i-vae.jls", outerVAE)
    serialize(outputFolder * "twostagesvae-$datasetName-$i-train.jls", train)
    serialize(outputFolder * "twostagesvae-$datasetName-$i-test.jls", test)

    return DataFrame(dataset = datasetName, idim = inputDim, hdim = hiddenDim, ldim = latentDim, layers = numLayers, i = i, auc_pxv = auc_pxv, auc_pz = auc_pz, auc_pxis = auc_pxis, auc_pz_jacobian = auc_pz_jacobian,
                    auc_pxv_pz = auc_pxv_pz, auc_pxv_pz_jacobian = auc_pxv_pz_jacobian, auc_pxv_cz = auc_pxv_cz, auc_pz_cz = auc_pz_cz, auc_pz_jacobian_cz = auc_pz_jacobian_cz, auc_pxv_pz_cz = auc_pxv_pz_cz,
                    auc_pxv_pz_jacobian_cz = auc_pxv_pz_jacobian_cz, auc_pxv_mz = auc_pxv_mz, auc_pz_mz = auc_pz_mz, auc_pz_jacobian_mz = auc_pz_jacobian_mz, auc_pxv_pz_mz = auc_pxv_pz_mz, auc_pxv_pz_jacobian_mz = auc_pxv_pz_jacobian_mz)
end

datasets = ["breast-cancer-wisconsin"]
difficulties = ["easy"]
batchSize = 200
iterations = 10000

if length(ARGS) != 0
    datasets = [ARGS[1]]
    difficulties = ["easy"]
end

for i in 1:5
    for (dn, df) in zip(datasets, difficulties)
        train, test, clusterdness = loaddata(dn, df)

        println("$dn")

        evaluateOneConfig = p -> runExperiment(dn, train, test, size(train[1], 1), p..., batchSize, iterations, i)
        results = gridsearch(evaluateOneConfig, [64], [3], [3], ["swish"], ["Dense"], [1.])

        CSV.write(outputFolder * "twostagesvae-$dn-$i.csv", vcat(results...))
    end
end


text:
if is_windows()
    flags = ["-m$(Sys.WORD_SIZE)", "-shared", "-O3"]
else
    flags = ["-m$(Sys.WORD_SIZE)", "-shared", "-O3", "-fPIC"]
end

run(`gfortran $flags small_routine.f90 -o small_routine.$(Libdl.dlext)`)

run(`gfortran $flags large_routine.f90 -o large_routine.$(Libdl.dlext)`)

run(`gfortran $flags array_routine.f90 -o array_routine.$(Libdl.dlext)`)


text:
"Supertype of an integrand which evaluates to `T`."
abstract type AbstractIntegrand{T} end

zero_result(f) = 0.0, 0.0
zero_result(integrand::AbstractIntegrand{T}) where {T} = zero(T), zero(real(T))

convert(::Type{AbstractIntegrand{T}}, f::AbstractIntegrand{T}) where {T} = f
convert(::Type{AbstractIntegrand{T}}, f::AbstractIntegrand{S}) where {S,T} =
    similar_integrand(f, T)

Base.complex(f::AbstractIntegrand{T}) where {T} =
    convert(AbstractIntegrand{complex(T)}, F)

(∘)(f::AbstractIntegrand, g) = integrand_compose(f, g, f.fun)
(∘)(f::AbstractIntegrand, g::IdentityMap) = f

(*)(f, g::AbstractIntegrand{T}) where {T} = integrand_times(f, g, g.fun)
(*)(f::Map, g::AbstractIntegrand{T}) where {T} = integrand_times(t->applymap(f,t), g, g.fun)
(*)(f::ConstantMap, g::AbstractIntegrand) = constant(f) * g

"Representation of an integrand function."
struct Integrand{T,F} <: AbstractIntegrand{T}
    fun ::  F
end

Integrand{T}(fun::F) where {T,F} = Integrand{T,F}(fun)
Integrand{T}(F::Integrand) where {T} = Integrand{T}(F.fun)

similar_integrand(F::Integrand, ::Type{T}) where {T} = Integrand{T}(F)

integrand_compose(f::Integrand{T}, g, fun) where {T} = Integrand{T}(t -> fun(g(t)))

integrand_times(f::Function, g::Integrand{T}, fun) where {T} =
    Integrand{T}(t -> f(t)*fun(t))
integrand_times(c::S, g::Integrand{T}, fun) where {T,S<:Number} =
    Integrand{promote_type(S,T)}(t -> c*fun(t))

(f::Integrand)(x...) = f.fun(x...)

"An integrand that logs any transformations which are applied to it."
struct TransformationLogIntegrand{T} <: AbstractIntegrand{T}
    fun
    prefactor
end

TransformationLogIntegrand{T}() where {T} = TransformationLogIntegrand{T}(IdentityMap{T}())
TransformationLogIntegrand{T}(fun) where {T} = TransformationLogIntegrand{T}(fun, UnityMap{T}())

similar_integrand(F::TransformationLogIntegrand, ::Type{T}) where {T} =
    TransformationLogIntegrand{T}(F)

integrand_compose(f::TransformationLogIntegrand{T}, g, fun::IdentityMap) where {T} =
    TransformationLogIntegrand{T}(g, f.prefactor)

integrand_compose(f::TransformationLogIntegrand{T}, g, fun) where {T} =
    TransformationLogIntegrand{T}(fun ∘ g, f.prefactor)

log_times(f::Function, g::UnityMap) = f
log_times(f::Function, g) = t -> f(t)*g(t)
log_times(c::Number, g::UnityMap{T}) where {T} = ConstantMap{T}(c)

integrand_times(f, g::TransformationLogIntegrand{T}, fun) where {T} =
    TransformationLogIntegrand{T}(fun, log_times(f, g.prefactor))

"Wrap a function and log any point at which it is evaluated."
struct LoggingFunction{T,F}
    fun     ::  F
    points  ::  Vector{T}
end

LoggingFunction{T}(f::F) where {T,F} = LoggingFunction{T,F}(f, Vector{T}())

function (f::LoggingFunction)(x)
    push!(f.points, x)
    f.fun(x)
end


text:
# MIT License
# Copyright (c) 2017: Xavier Gandibleux, Anthony Przybylski, Gauthier Soleilhac, and contributors.
struct _2LAP
    nSize::Int
    C1::Matrix{Int}
    C2::Matrix{Int}
end

Base.show(io::IO, id::_2LAP) = print("Bi-Objective Linear Affectation Problem with $(id.nSize) variables.")

set2LAP(n::Int, c1::Matrix{Int}, c2::Matrix{Int}) = begin
    if !(size(c1,1) == size(c1, 2) == size(c2,1) == size(c2 , 2) == n)
        error("dimensions incorrectes")
    end
    _2LAP(n, c1, c2)
end
set2LAP(c1::Matrix{Int}, c2::Matrix{Int}) = _2LAP(size(c1,1), c1, c2)

struct LAPsolver
    parameters
    solve::Function #::Function(id::_2LAP) -> ...
end

function vSolve(id::_2LAP, solver::LAPsolver = LAP_Przybylski2008())
    solver.solve(id)
end

function LAP_Przybylski2008()::LAPsolver
    f = (id::_2LAP) -> begin
        nSize, C1, C2 = Cint(id.nSize), vec(convert(Matrix{Cint}, id.C1)), vec(convert(Matrix{Cint}, id.C2))
        p_z1,p_z2,p_solutions,p_nbsolutions = Ref{Ptr{Cint}}() , Ref{Ptr{Cint}}(), Ref{Ptr{Cint}}(), Ref{Cint}()
        ccall(
            (:solve_bilap_exact, libLAPpath),
            Nothing,
            (Ref{Cint},Ref{Cint}, Cint, Ref{Ptr{Cint}}, Ref{Ptr{Cint}}, Ref{Ptr{Cint}}, Ref{Cint}),
            C1, C2, nSize, p_z1, p_z2, p_solutions, p_nbsolutions)
        nbSol = p_nbsolutions.x
        z1,z2 = convert(Array{Int,1},_unsafe_wrap(Array, p_z1.x, nbSol, own=true)), convert(Array{Int,1},_unsafe_wrap(Array, p_z2.x, nbSol, own=true))
        solutions = convert(Array{Int,2},reshape(_unsafe_wrap(Array, p_solutions.x, nbSol*id.nSize, own=true), (id.nSize, nbSol)))
        return z1, z2, permutedims(solutions, [2,1]).+1
    end

    return LAPsolver(nothing, f)
end

function load2LAP(fname::AbstractString)
    f = open(fname)
    n = parse(Int,readline(f))
    C1 = zeros(Int,n,n)
    C2 = zeros(Int,n,n)
    for i = 1:n
		@static if VERSION > v"0.7-"
        	C1[i,:] = parse.(Int, split(chomp(readline(f)), ' ', keepempty=false))
		else
			C1[i,:] = parse.(Int, split(chomp(readline(f)), ' ', keep=false))
		end
    end
    for i=1:n
		@static if VERSION > v"0.7-"
	        C2[i,:] = parse.(Int, split(chomp(readline(f)), ' ', keepempty=false))
		else
			C2[i,:] = parse.(Int, split(chomp(readline(f)), ' ', keep=false))
		end
    end
    return _2LAP(n, C1, C2)
end

max_stars_repo_name: vOptSolver/vOptSpecific

max_issues_repo_name: vOptSolver/vOptSpecific

max_forks_repo_name: vOptSolver/vOptSpecific


text:
include("../../src/SFFM.jl")
using LinearAlgebra, Plots, JLD2, StatsBase#, SFFM

## define the model(s)
include("exampleModelDef.jl")

## section 4.2: Ψ paths

## load simulated Ψ paths
@load pwd()*"/examples/paperNumerics/dump/sims.jld2" sims

## mesh
Δtemp = 0.4
Nodes = collect(approxBounds[1, 1]:Δtemp:approxBounds[1, 2])

nBases = 1
Basis = "lagrange"
mesh = SFFM.DGMesh(
    approxModel,
    Nodes,
    nBases,
    Basis = Basis,
)

## turn sims into a cdf
simprobs = SFFM.Sims2Dist(simModel,mesh,sims,SFFM.SFFMCDF)

## bootstrap to get CI
function bootFun(sims; nBoot = 10)
    l = length(sims.t)
    samplesBoot = Array{Float64,3}(undef,nBoot,6,2)
    for n in 1:nBoot
        sampleIdx = sample(1:l,l)
        tempData = (
            φ = sims.φ[sampleIdx],
            X = sims.X[sampleIdx],
        )
        tempDist = SFFM.Sims2Dist(simModel,mesh,tempData,SFFM.SFFMCDF).distribution[1,1:6,[2;4]]
        samplesBoot[n,:,:] = tempDist
    end
    ql = zeros(6,2)
    qu = zeros(6,2)
    for xpos in 1:6
        for phase in 1:2
            ql[xpos,phase] = quantile(samplesBoot[:,xpos,phase],0.025)
            qu[xpos,phase] = quantile(samplesBoot[:,xpos,phase],0.975)
        end
    end
    return (ql, qu)
end
@time ql, qu = bootFun(sims)

## plot simulations
let
    p2 = plot()
    p4 = plot()

    colours = [:black, :false ,:blue]
    shapes = [:star5, :x, :star7, :+]
    ## DG
    c = -1
    styles = [:dashdot,:dash]
    for nBases in [1;5]
        c = c+2
        mesh = SFFM.DGMesh(
            approxModel,
            Nodes,
            nBases,
            Basis = Basis,
        )

        frapmesh = SFFM.FRAPMesh(
            approxModel,
            Nodes,
            nBases,
        )
        # construct DG matrices
        All = SFFM.MakeAll(approxModel, mesh, approxType = "interpolation")
        Ψ = SFFM.PsiFun(All.D)

        # construct FRAP matrices
        # me = SFFM.MakeME(SFFM.CMEParams[nBases], mean = SFFM.Δ(mesh)[1])
        B = SFFM.MakeFullGenerator( approxModel, frapmesh)
        D = SFFM.MakeD( frapmesh, B, All.R)
        Ψme = SFFM.PsiFun( D)

        # construct initial condition
        theNodes = SFFM.CellNodes(mesh)[:,convert(Int,ceil(5/Δtemp))]
        basisValues = zeros(length(theNodes))
        for n in 1:length(theNodes)
            basisValues[n] = prod(5.0.-theNodes[[1:n-1;n+1:end]])./prod(theNodes[n].-theNodes[[1:n-1;n+1:end]])
        end
        initpm = [
            zeros(sum(approxModel.C.<=0)) # LHS point mass
            zeros(sum(approxModel.C.>=0)) # RHS point mass
        ]
        initprobs = zeros(Float64,SFFM.NBases(mesh),SFFM.NIntervals(mesh),SFFM.NPhases(approxModel))
        V=SFFM.vandermonde(SFFM.NBases(mesh))
        initprobs[:,convert(Int,ceil(5/Δtemp)),3] = basisValues'* V.V* V.V'.*2/Δtemp
        initdist = SFFM.SFFMDensity(
            initpm,
            initprobs,
            SFFM.CellNodes(mesh),
        ) # convert to a distribution object so we can apply Dist2Coeffs
        # convert to Coeffs α in the DG context
        x0 = SFFM.Dist2Coeffs( approxModel, mesh, initdist)
        # the initial condition on Ψ is restricted to + states so find the + states
        plusIdx = [
            mesh.Fil["p+",:];
            repeat(mesh.Fil["+",:]', SFFM.NBases(mesh), 1)[:];
            mesh.Fil["q+",:];
        ]
        # get the elements of x0 in + states only
        x0 = x0[plusIdx]'
        # check that it is equal to 1 (or at least close)
        println(sum(x0))

        # compute x-distribution at the time when Y returns to 0
        w = x0*Ψ

        wme = x0*Ψme
        # this can occur in - states only, so find the - states
        minusIdx = [
            mesh.Fil["p-",:];
            repeat(mesh.Fil["-",:]', SFFM.NBases(mesh), 1)[:];
            mesh.Fil["q-",:];
        ]
        # then map to the whole state space for plotting
        z = zeros(
            Float64,
            SFFM.NBases(mesh)*SFFM.NIntervals(mesh) * SFFM.NPhases(approxModel) +
                sum(approxModel.C.<=0) + sum(approxModel.C.>=0)
        )
        zme = zeros(
            Float64,
            SFFM.NBases(mesh)*SFFM.NIntervals(mesh) * SFFM.NPhases(approxModel) +
                sum(approxModel.C.<=0) + sum(approxModel.C.>=0)
        )
        z[minusIdx] = w
        zme[minusIdx] = wme
        # check that it is equal to 1
        println(sum(z))
        println(sum(zme))

        # convert to a distribution object for plotting
        DGProbs = SFFM.Coeffs2Dist(
            approxModel,
            mesh,
            z,
            SFFM.SFFMCDF,
        )

        meProbs = SFFM.Coeffs2Dist(
            approxModel,
            mesh,
            zme,
            SFFM.SFFMCDF,
        )

        # plot them
        p2 = plot!(p2,
            DGProbs.x[:],
            DGProbs.distribution[:,:,2][:],
            label = "DG: N_k = "*string(nBases),
            color = colours[c],
            xlims = (-0.1,2.1),
            seriestype = :scatter,
            # linestyle = styles[c],
            markershape = shapes[c],
            markercolor = colours[c],
            markerstrokecolor = colours[c],
            markersize=6,
        )
        p2 = plot!(p2,
            meProbs.x[:],
            meProbs.distribution[:,:,2][:],
            label = "me: N_k = "*string(nBases),
            color = colours[c],
            xlims = (-0.1,2.1),
            seriestype = :scatter,
            # linestyle = styles[c],
            markershape = shapes[c+1],
            markercolor = colours[c],
            markerstrokecolor = colours[c],
            markersize=6,
        )

        p4 = plot!(p4,
            DGProbs.x[:],
            DGProbs.distribution[:,:,4][:],
            label = "DG:  N_k = "*string(nBases),
            color = colours[c],
            xlims = (-0.1,2.1),
            seriestype = :scatter,
            # linestyle = styles[c],
            markershape = shapes[c],
            markercolor = colours[c],
            markerstrokecolor = colours[c],
            markersize=6,
        )
        p4 = plot!(p4,
            meProbs.x[:],
            meProbs.distribution[:,:,4][:],
            label = "me:  N_k = "*string(nBases),
            color = colours[c],
            xlims = (-0.1,2.1),
            seriestype = :scatter,
            # linestyle = styles[c],
            markershape = shapes[c+1],
            markercolor = colours[c],
            markerstrokecolor = colours[c],
            markersize=6,
        )
    end
    p2 = plot!(p2,
        [1;1]*simprobs.x[1,1:6]',
        [ql[:,1]';qu[:,1]'],
        color=:red,
        marker=:hline,
        markersize=12,
        seriestype=:line,
        label=:false,
        xlabel = "x",
        ylabel = "Cumulative probability",
        windowsize = (600,400),
        grid = false,
        tickfontsize = 12,
        guidefontsize = 14,
        titlefontsize = 20,
        legendfontsize = 12,
        title = "Phase 10",
    )
    p2 = plot!(p2,
        [1;1]*simprobs.x[1,1]',
        [ql[1,1]';qu[1,1]'],
        color=:red,
        marker=:hline,
        markersize=12,
        seriestype=:scatter,
        label="Simulation",
        legend = :bottomright,
    )

    p4 = plot!(p4,
        [1;1]*simprobs.x[1,1:6]',
        [ql[:,2]';qu[:,2]'],
        color=:red,
        marker=:hline,
        markersize=12,
        seriestype=:line,
        label=:false,
        xlabel = "x",
        ylabel = "Cumulative probability",
        windowsize = (600,400),
        grid = false,
        tickfontsize = 12,
        guidefontsize = 14,
        titlefontsize = 20,
        legendfontsize = 12,
        legend = :bottomright,
        title = "Phase 00",
    )
    p4 = plot!(p4,
        [1;1]*simprobs.x[1,1]',
        [ql[1,2]';qu[1,2]'],
        color=:red,
        marker=:hline,
        markersize=12,
        seriestype=:scatter,
        label="Simulation",
        legend = :bottomright,
    )

    display(p2)
    display(p4)
    # savefig(p2, pwd()*"/examples/paperNumerics/dump/psiPhase10.png")
    # savefig(p4, pwd()*"/examples/paperNumerics/dump/psiPhase00.png")
end

max_stars_repo_name: angus-lewis/SFFM

max_issues_repo_name: angus-lewis/SFFM

max_forks_repo_name: angus-lewis/SFFM


text:
using Symbolics, Test

@variables a b::Real z::Complex (Z::Complex)[1:10]

@testset "types" begin
    @test a isa Num
    @test b isa Num
    @test eltype(Z) <: Complex{Real}

    for x in [z, Z[1], z+a, z*a, z^2, z/z] # z/z is sus
        @test x isa Complex{Num}
        @test real(x) isa Num
        @test imag(x) isa Num
        @test conj(x) isa Complex{Num}
    end

    # issue #314
    bi = a+a*im
    bs = substitute(bi, (Dict(a=>1.0))) # returns 1.0 + im
    typeof(bs) # Complex{Num}
    bv = Symbolics.value.(bs)
    @test typeof(bv) == ComplexF64
end

@testset "repr" begin
    @test repr(z) == "z"
    @test repr(a + b*im) == "a + b*im"
end


text:
function vector_matrix_vector_product(x::Vector{T}, A::Matrix{T}, y::Vector{T}) where {T}
    @assert length(x) == size(A,1)
    @assert length(y) == size(A,2)

    local b::T = 0
    local c::T = 0

    @inbounds for i in eachindex(x)
        c = 0
        for j in eachindex(y)
            c += A[i,j] * y[j]
        end
        b += c * x[i]
    end

    return b
end

function vector_matrix_vector_product(x::Matrix{T}, A::Matrix{T}, y::Matrix{T}, ind::Int) where {T}
@assert size(x,1) == size(A,1)
    @assert size(y,1) == size(A,2)
    @assert ind ≤ size(x,2)
    @assert ind ≤ size(y,2)

    local b::T = 0
    local c::T = 0

    @inbounds for i in axes(x,1)
        c = 0
        for j in axes(y,1)
            c += A[i,j] * y[j,ind]
        end
        b += c * x[i,ind]
    end

    return b
end

function vector_vector_product(x::Matrix{T}, y::Matrix{T}, ind::Int) where {T}
    @assert size(x,1) == size(y,1)
    @assert ind ≤ size(x,2)
    @assert ind ≤ size(y,2)

    local b::T = 0

    @inbounds for i in axes(x,1)
        b += x[i,ind] * y[i,ind]
    end

    return b
end