import LDLFactorizations
const LDLF = LDLFactorizations

"""
    LDLFact <: LSBackend

Use LDLFactorizations backend.

Options available:
* Any numerical type `T`
* Augmented system with LDLᵀ factorization
"""
struct LDLFact <: LSBackend end

# ==============================================================================
#   LDLFLinearSolver
# ==============================================================================

"""
    LDLFLinearSolver{Tv}

Linear solver for the 2x2 augmented system with ``A`` sparse.

Uses LDLFactorizations.jl to compute an LDLᵀ factorization of the quasi-definite augmented system.
"""
mutable struct LDLFLinearSolver{Tv<:Real} <: AbstractLinearSolver{Tv}
    m::Int  # Number of rows
    n::Int  # Number of columns

    # TODO: allow for user-provided ordering,
    #   and add flag (default to false) to know whether user ordering should be used

    # Problem data
    A::SparseMatrixCSC{Tv, Int}
    θ::Vector{Tv}
    regP::Vector{Tv}  # primal regularization
    regD::Vector{Tv}  # dual regularization

    # Factorization
    F::LDLF.LDLFactorization{Tv}

    # TODO: constructor with initial memory allocation
    function LDLFLinearSolver(A::SparseMatrixCSC{Tv, Int}) where{Tv<:Real}
        m, n = size(A)
        θ = ones(Tv, n)

        S = [
            spdiagm(0 => -θ)  A';
            A spdiagm(0 => ones(m))
        ]

        # TODO: PSD-ness checks
        F = LDLF.ldl(S)
        return new{Tv}(m, n, A, θ, ones(Tv, n), ones(Tv, m), F)
    end

end

AbstractLinearSolver(
    ::LDLFact,
    ::AugmentedSystem,
    A::AbstractMatrix{Tv}
) where{Tv<:Real} = LDLFLinearSolver(sparse(A))

backend(::LDLFLinearSolver) = "LDLFactorizations.jl"
linear_system(::LDLFLinearSolver) = "Augmented system"

"""
    update_linear_solver!(ls, θ, regP, regD)

Update LDLᵀ factorization of the augmented system.

Update diagonal scaling ``\\theta``, primal-dual regularizations, and re-compute
    the factorization.
Throws a `PosDefException` if matrix is not quasi-definite.
"""
function update_linear_solver!(
    ls::LDLFLinearSolver{Tv},
    θ::AbstractVector{Tv},
    regP::AbstractVector{Tv},
    regD::AbstractVector{Tv}
) where{Tv<:Real}
    # Sanity checks
    length(θ)  == ls.n || throw(DimensionMismatch(
        "θ has length $(length(θ)) but linear solver is for n=$(ls.n)."
    ))
    length(regP) == ls.n || throw(DimensionMismatch(
        "regP has length $(length(regP)) but linear solver has n=$(ls.n)"
    ))
    length(regD) == ls.m || throw(DimensionMismatch(
        "regD has length $(length(regD)) but linear solver has m=$(ls.m)"
    ))

    # Update diagonal scaling
    ls.θ .= θ
    # Update regularizers
    ls.regP .= regP
    ls.regD .= regD

    # Re-compute factorization
    # TODO: Keep S in memory, only change diagonal
    S = [
        spdiagm(0 => -ls.θ .- regP)  ls.A';
        ls.A spdiagm(0 => regD)
    ]

    # TODO: PSD-ness checks
    try
        ls.F = LDLF.ldl(S)
    catch err
        isa(err, LDLF.SQDException) && throw(PosDefException(-1))
        rethrow(err)
    end

    return nothing
end

"""
    solve_augmented_system!(dx, dy, ls, ξp, ξd)

Solve the augmented system, overwriting `dx, dy` with the result.
"""
function solve_augmented_system!(
    dx::Vector{Tv}, dy::Vector{Tv},
    ls::LDLFLinearSolver{Tv},
    ξp::Vector{Tv}, ξd::Vector{Tv}
) where{Tv<:Real}
    m, n = ls.m, ls.n

    # Set-up right-hand side
    ξ = [ξd; ξp]

    # Solve augmented system
    d = ls.F \ ξ

    # Recover dx, dy
    @views dx .= d[1:n]
    @views dy .= d[(n+1):(m+n)]

    # TODO: Iterative refinement
    return nothing
end


# Author: Mathias Louboutin, mlouboutin3@gatech.edu
# Date: December 2020

using SlimOptim

# Rosenbrock function
f(m) = ((1-m[1])^2 + 100*(m[2]-m[1]^2)^2)
g!(g,m) = (g .= [-2*(1-m[1]) - 200*(m[2]-m[1]^2)*2*m[1], 200*(m[2]-m[1]^2)])
prj(x) = x

function obj(x)
    g = zeros(2)
    ϕ = f(x)
    g!(g,x)
    ϕ, g
end

function fg!(g,x)
    !isnothing(g) && g!(g,x)
    f(x)
end

function fg!(ff,g,x)
    !isnothing(g) && g!(g,x)
    !isnothing(ff) && return f(x)
end

niter = 1000
for (opt, algo)=zip([spg_options, pqn_options], [spg, pqn])
    x0 = zeros(2)
    options = opt(maxIter=niter, progTol=1f-30, optTol=0, iniStep=2, memory=5)
    algo(obj,x0,prj,options)
    @test x0 ≈ [1, 1]

    x1 = zeros(2)
    algo(f,g!,fg!,x1,prj,options)
    @show x1 ≈ [1, 1]
end


function loadBrukerfolderautophase(full_path::String)

    pdata_path = "$full_path/pdata/1"

    PyCall.py"""
    import numpy
    import nmrglue as ng

    s = $full_path
    procdata_dir = $pdata_path

    dic, data = ng.bruker.read(s)

    # ---time correction before FT
    precorr_time =  ng.bruker.remove_digital_filter(dic, data)
    precorr_frq0 = ng.proc_base.fft(precorr_time)
    precorr_frq = ng.proc_autophase.autops(precorr_frq0, 'acme')
    # precorr_frq = ng.proc_autophase.autops(precorr_frq0, 'peak_minima')

    #---time correction after FT
    postcorr_frq0 = ng.proc_base.fft(data)
    postcorr_frq1 = ng.bruker.remove_digital_filter(dic, postcorr_frq0, post_proc=True)
    postcorr_frq = ng.proc_autophase.autops(postcorr_frq1, 'acme')

    #---data processed using TopSpin (3.5pl7)
    bruker_frq = ng.bruker.read_pdata(procdata_dir)[1][::-1]

    guess_udic = ng.fileio.bruker.guess_udic(dic, data)
    """

    dic = PyCall.py"dic"
    data = PyCall.py"data"

    precorr_time = PyCall.py"precorr_time"
    precorr_frq0 = PyCall.py"precorr_frq0"
    precorr_frq = PyCall.py"precorr_frq"

    postcorr_frq0 = PyCall.py"postcorr_frq0"
    postcorr_frq1 = PyCall.py"postcorr_frq1"
    postcorr_frq = PyCall.py"postcorr_frq"

    bruker_frq = PyCall.py"bruker_frq"
    guess_udic = PyCall.py"guess_udic"

    return dic, data, precorr_time, precorr_frq0, precorr_frq,
            postcorr_frq0, postcorr_frq1, postcorr_frq, bruker_frq,
            guess_udic
end

"""
    loadspectrum(full_path::String;
        solvent_ppm = 4.8,
        solvent_window_ppm = 0.3,
        N_Ω::Int = 20000,
        N_optim::Int = 1000,
        n_particles::Int = 3,
        max_iters::Int = 100,
        α_lower = 1e-6,
        α_upper_factor = 50.0,
        λ_lower = 1e-10,
        λ_upper = 20.0,
        λ_initial = 3.0,
        st_ind_default = 100,
        verbose_flag::Bool = false)

Loads a Bruker 1D 1H NMR experiment. Uses the NMRGlue Python library. Estimates the 0 ppm and solvent resonance components. The solvent resonance component is estimated in the frequency interval `solvent_ppm` +/- `solvent_window_ppm`, in units ppm.

On a Linux terminal, run the following command if you don't have NMRGlue installed.
pip install nmrglue --user

Notes:
Either `st_ind_default` samples or `dead_time_st` number of samples are removed from the raw time-domain data to create the output time-domain data, s_t. This is to remove the silence at the beginning of the acquired data due to the dead-time property of practical NMR spectrometers. 'dead_time_st' is calculated automatically from the metadata of the experiment files.

Therefore, the returned sampling frequency `fs` from loadspectrum() and the sampling frequency in the returned metadata ditionary `dic` are slightly different.

...
# Outputs, in the order returned.

- `s_t::Vector{Complex{Float64}}`: time-domain experiment data.
- `S::Function`: Discrete-time Fourier transform of s_t.
- `hz2ppmfunc::Function`: conversion function from Hz to ppm for this experiment. The formula used is `hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs`.
- `ppm2hzfunc::Function`: conversion function from ppm to Hz for this experiment. The formula used is `ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)`.

- `ν_0ppm::Float64`: where the estimated frequency of the 0 ppm resonance component, in Hz.
- `fs::Float64`: the estimated sampling frequency of `s_t`.
- `SW::Float64`: the spectral window of the spectrometer, as provided by the metadata of the experiment. In units of ppm.
- `α_0ppm`: The amplitude parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model.
- `β_0ppm`: The phase parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model.
- `λ_0ppm`: The T2 parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model.
- `Ω_0ppm`: The radial frequency parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model. In units Hz/(2*π).
- `results_0ppm`: The NLopt.jl optimization results for the estimation of the 0 ppm reference. Inspect this to see if the estimated parameters for the 0 ppm reference are reliable.

- `dic`; The dictionary that contains all metadata from NMRGlue for the NMR experiment.

- `α_solvent`: The amplitude parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model.
- `β_solvent`: The phase parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model.
- `λ_solvent`: The T2 parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model.
- `Ω_solvent`: The radial frequency parameter of the estimate resonance component for 0 ppm reference, in the free-induction decay model. In units Hz/(2*π).
- `results_solvent`: The NLopt.jl optimization results for the estimation of the 0 ppm reference. Inspect this to see if the estimated parameters for the 0 ppm reference are reliable.
, , , , ,

# Examples
See load_experiment.jl in the /examples folder for an example.
"""
function loadspectrum(full_path::String;
    solvent_ppm = 4.8,
    solvent_window_ppm = 0.3,
    N_Ω::Int = 20000,
    N_optim::Int = 1000,
    n_particles::Int = 3,
    max_iters::Int = 100,
    α_lower = 1e-6,
    α_upper_factor = 50.0,
    λ_lower = 1e-10,
    λ_upper = 20.0,
    λ_initial = 3.0,
    st_ind_default = 100, # future: see if we can derive this from DE dead time later.
    verbose_flag::Bool = false)

    dic, data, precorr_time, precorr_frq0, precorr_frq, postcorr_frq0, postcorr_frq1,
    postcorr_frq, Bruker_spectrums_nmrglue,
    guess_udic = loadBrukerfolderautophase(full_path)

    DE = dic["acqus"]["DE"] # in microseconds.

    BF1 = dic["acqus"]["BF1"]
    ## see Bruker TopSpin acquistion commands and parameters v 003.
    # TD - Time Domain; Number Of Raw Data Points. i.e., length(data)*2 - TD = 0.
    # total samples from both time-series.
    TD = dic["acqus"]["TD"]

    # SW - Spectral Width in ppm
    SW = dic["acqus"]["SW"]

    # SFO1 - SFO8 - Irradiation (carrier) Frequencies For Channels f1 to f8
    SFO1 = dic["acqus"]["SFO1"]

    O1 = dic["acqus"]["O1"]
    dic["acqus"]["NUC1"]

    CAR = O1*1.0

    fs_dic = dic["acqus"]["SW_h"]

    N_diff = length(data)*2 - TD

    t_data = gettimerange(length(data), fs_dic)

    AQ = TD/(2*SW*SFO1)
    dead_time = t_data[end] - AQ

    dead_time_st = round(Int, dead_time/t_data[end] * length(data))
    st_ind = max(st_ind_default, dead_time_st)

    s_t = data[st_ind:end]
    t = gettimerange(length(s_t), fs_dic)

    fs = (length(s_t)-1)/t[end] # this is really fs_dic.
    @assert abs(fs-fs_dic) < 1e-9

    α_0ppm, β_0ppm, λ_0ppm, Ω_0ppm,
    results_0ppm = estimatereferencecompoundfreq(CAR, fs, SW, s_t, t;
    N_Ω = N_Ω,
    N_optim = N_optim,
    n_particles = n_particles,
    max_iters = max_iters,
    α_lower = α_lower,
    α_upper_factor = α_upper_factor,
    λ_lower = λ_lower,
    λ_upper = λ_upper,
    λ_initial = λ_initial,
    verbose_flag = verbose_flag)

    ν_0ppm = Ω_0ppm ./ (2*π)

    hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
    ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

    # S is scaled DTFT of data.
    S = vv->(computeDTFTch3eq29(s_t, vv, t)/fs)

    # solvent.
    u_solvent_lower = ppm2hzfunc(solvent_ppm-solvent_window_ppm)
    u_solvent_upper = ppm2hzfunc(solvent_ppm+solvent_window_ppm)

    α_solvent, β_solvent, λ_solvent, Ω_solvent,
    results_solvent = estimatesinglet(fs, s_t, t,
                u_solvent_lower, u_solvent_upper;
                N_Ω = N_Ω,
                N_optim = N_optim,
                n_particles = n_particles,
                max_iters = max_iters,
                α_lower = α_lower,
                α_upper_factor = α_upper_factor,
                λ_lower = λ_lower,
                λ_upper = λ_upper,
                λ_initial = λ_initial,
                verbose_flag = verbose_flag)

    return s_t, S, hz2ppmfunc, ppm2hzfunc, ν_0ppm, fs, SW, α_0ppm, β_0ppm, λ_0ppm, Ω_0ppm,
    results_0ppm, dic, α_solvent, β_solvent, λ_solvent, Ω_solvent, results_solvent
end

function estimatereferencecompoundfreq( CAR::T,
    fs::T,
    SW::T,
    s_t::Vector{Complex{T}},
    t;
    N_Ω::Int = 20000,
    N_optim::Int = 1000,
    n_particles::Int = 3,
    max_iters::Int = 100,
    window_ppm::T = 0.3,
    α_lower::T = 1e-6,
    α_upper_factor::T = 50.0,
    λ_lower::T = 1e-10,
    λ_upper::T = 20.0,
    λ_initial::T = 3.0,
    verbose_flag::Bool = false) where T <: Real

    # prepare frequency range.
    ν0_initial, hz2ppmfunc0, ppm2hzfunc0 = getinitalguessreferencecompoundfreq(CAR, fs, SW)

    bp_a = ppm2hzfunc0(-window_ppm)
    bp_b = ppm2hzfunc0(window_ppm)

    return estimatesinglet(fs, s_t, t, bp_a, bp_b;
    N_Ω = N_Ω,
    N_optim = N_optim,
    n_particles = n_particles,
    max_iters = max_iters,
    α_lower = α_lower,
    α_upper_factor = α_upper_factor,
    λ_lower = λ_lower,
    λ_upper = λ_upper,
    λ_initial = λ_initial,
    verbose_flag = verbose_flag)

end

# λ_initial: m a value of 3 seems decent for the λ of DSS under 600 MHz.
# no error checking on whether λ_initial ∈ [λ_lower, λ_upper].
function estimatesinglet( fs::T,
    s_t::Vector{Complex{T}},
    t,
    u_lower,
    u_upper;
    N_Ω::Int = 20000,
    N_optim::Int = 1000,
    n_particles::Int = 3,
    max_iters::Int = 100,
    α_lower::T = 1e-6,
    α_upper_factor::T = 100.0,
    λ_lower::T = 1e-10,
    λ_upper::T = 20.0,
    λ_initial::T = 3.0,
    verbose_flag::Bool = false,
    xtol_rel::T = 1e-12,
    nlopt_seed::Int = 25) where T <: Real

    # prepare data.
    DTFT_s_scaled = vv->(computeDTFTch3eq29(s_t, vv, t)/fs)

    u_range_optim = LinRange(u_lower, u_upper, N_optim)
    S_U = DTFT_s_scaled.(u_range_optim)

    # set up optimization.
    max_val, max_ind = findmax(abs.(S_U))
    Ω = u_range_optim[max_ind] *2*π
    #println("Ω = ", Ω)
    costfunc = pp->evalcostreferencecompound(pp,
    S_U, u_range_optim, Ω)

    optim_lower_limit = [α_lower; -π; λ_lower]
    optim_upper_limit = [max_val*α_upper_factor; π; λ_upper]

    p0 = [max_val; 0.0; λ_initial]

    # optimize.
    op = Optim.Options( iterations = max_iters,
    store_trace = false,
    show_trace = verbose_flag)

    results = Optim.optimize( costfunc,
    p0,
    Optim.ParticleSwarm(; lower = optim_lower_limit,
        upper = optim_upper_limit,
        n_particles = n_particles),
    op)

    #α, β, λ = results.minimizer
    # println("α = ", α)
    # Q = [1221.2329692741823;
    # 1.5072497616902667;
    # 4.266537181116996]
    # println("cost Q = ", costfunc(Q))
    # println("cost p = ", costfunc(results.minimizer))

    # use NLopt.
    grad_func = xx->FiniteDiff.finite_difference_gradient(costfunc, xx)

    NLopt.srand(nlopt_seed)

    p_initial = results.minimizer
    #p_initial = p0

    opt = NLopt.Opt(:GN_ESCH, length(p_initial))
    #opt = NLopt.Opt(:LD_MMA, length(p_initial))

    min_cost_p, min_p, ret_p, numevals_p = runNLopt!(  opt,
        p_initial,
        costfunc,
        grad_func,
        optim_lower_limit,
        optim_upper_limit;
        max_iters = 100*max_iters,
        xtol_rel = xtol_rel)
    #
    opt = NLopt.Opt(:LD_MMA, length(min_p))
    min_cost_p, min_p, ret_p, numevals_p = runNLopt!(  opt,
        min_p,
        costfunc,
        grad_func,
        optim_lower_limit,
        optim_upper_limit;
        max_iters = 100*max_iters,
        xtol_rel = xtol_rel)

    α, β, λ = min_p

    # #
    # println("p_initial = ", p_initial)
    # println("α         = ", α)

    return α, β, λ, Ω, results
end

function evalcostreferencecompound( p::Vector{T},
    S_U::Vector{Complex{T}},
    u_range,
    Ω) where T <: Real

    # parse.
    α, β, λ = p

    cost = zero(T)
    for i = 1:length(u_range)
        q_u = evalcomplexLorentzian(u_range[i], α, β, λ, Ω)
        cost += abs2(q_u - S_U[i])
        #cost += (abs(q_u)-abs(S_U[i]))^2
    end

    return cost
end

function getinitalguessreferencecompoundfreq(CAR::T, fs::T, SW::T) where T <: Real

    ν0_initial = fs - CAR

    hz2ppmfunc0 = uu->(uu - ν0_initial)*SW/fs
    ppm2hzfunc0 = pp->(ν0_initial + pp*fs/SW)

    return ν0_initial, hz2ppmfunc0, ppm2hzfunc0
end


using Base.Test
using SkewDist
println("Testing SkewTDist...")

df = 4.0
ξ = 1.0
ω = 0.5
α = 1.0

dist = SkewTDist(ξ, ω, α, df)

# Test PDF integrates to one
@test_approx_eq_eps quadgk(x->pdf(dist,x), -Inf, Inf)[1] 1.0 1e-6

cdf(dist, 1.0)

# Monte Carlo tests for basic statistics
n = 100000
x = rand(dist, n)
@test_approx_eq_eps mean(x) mean(dist) 1e-1
@test_approx_eq_eps var(x) var(dist) 1e-1


function gemm_kernel!(C, A, B)
    @avx for n ∈ 1:size(A, 1), m ∈ 1:size(B, 2)
        Cₙₘ = zero(eltype(C))
        for k ∈ 1:size(A, 2)
            Cₙₘ += A[n,k] * B[k,m]
        end
        C[n,m] = Cₙₘ
    end
end

function add_gemm_kernel!(C::MatTypes, A::MatTypes, B::MatTypes)
    @avx for n ∈ 1:size(A, 1), m ∈ 1:size(B, 2)
        Cₙₘ = zero(eltype(C))
        for k ∈ 1:size(A, 2)
            Cₙₘ += A[n,k] * B[k,m]
        end
        C[n,m] += Cₙₘ
    end
end

add_gemm_kernel!(C::MatTypes, A::MatTypes, B::MatTypes, ::Val{1}) = add_gemm_kernel!(C, A, B)

function add_gemm_kernel!(C::MatTypes, A::MatTypes, B::MatTypes, ::Val{-1})
    @avx for n ∈ 1:size(A, 1), m ∈ 1:size(B, 2)
        Cₙₘ = zero(eltype(C))
        for k ∈ 1:size(A, 2)
            Cₙₘ -= A[n,k] * B[k,m]
        end
        C[n,m] += Cₙₘ
    end
end

function add_gemm_kernel!(C::MatTypes, A::MatTypes, B::MatTypes, ::Val{factor}) where {factor}
    @avx for n ∈ 1:size(A, 1), m ∈ 1:size(B, 2)
        Cₙₘ = zero(eltype(C))
        for k ∈ 1:size(A, 2)
            Cₙₘ += factor * A[n,k] * B[k,m]
        end
        C[n,m] += Cₙₘ
    end
end

#____________

function gemm_kernel!(u::VecTypes, A::MatTypes, v::VecTypes)
    @avx for n ∈ 1:size(A, 1)
        uₙ = zero(eltype(u))
        for k ∈ 1:size(A, 2)
            uₙ += A[n,k] * v[k]
        end
        u[n] = uₙ
    end
end

function add_gemm_kernel!(u::VecTypes, A::MatTypes, v::VecTypes)
    @avx for n ∈ 1:size(A, 1)
        uₙ = zero(eltype(u))
        for k ∈ 1:size(A, 2)
            uₙ += A[n,k] * v[k]
        end
        u[n] += uₙ
    end
end

function add_gemm_kernel!(u::VecTypes, A::MatTypes, v::VecTypes, ::Val{-1})
    @avx for n ∈ 1:size(A, 1)
        uₙ = zero(eltype(u))
        for k ∈ 1:size(A, 2)
            uₙ -= A[n,k] * v[k]
        end
        u[n] += uₙ
    end
end

function add_gemm_kernel!(u::VecTypes, A::MatTypes, v::VecTypes, ::Val{factor}) where {factor}
    @avx for n ∈ 1:size(A, 1)
        uₙ = zero(eltype(u))
        for k ∈ 1:size(A, 2)
            uₙ += A[n,k] * v[k]
        end
        u[n] += factor * uₙ
    end
end

#____________

function gemm_kernel!(u::CoVecTypes, v::CoVecTypes, A::MatTypes)
    @avx for m ∈ 1:size(A, 2)
        uₘ = zero(eltype(u))
        for k ∈ 1:size(A, 1)
            uₘ += v[k] * A[k, m]
        end
        u[m] = uₘ
    end
end

function add_gemm_kernel!(u::CoVecTypes, v::CoVecTypes, A::MatTypes)
    @avx for m ∈ 1:size(A, 2)
        uₘ = zero(eltype(u))
        for k ∈ 1:size(A, 1)
            uₘ += v[k] * A[k, m]
        end
        u[m] += uₘ
    end
end

function add_gemm_kernel!(u::CoVecTypes, v::CoVecTypes, A::MatTypes, ::Val{-1})
    @avx for m ∈ 1:size(A, 2)
        uₘ = zero(eltype(u))
        for k ∈ 1:size(A, 1)
            uₘ -= v[k] * A[k, m]
        end
        u[m] += uₘ
    end
end

function add_gemm_kernel!(u::CoVecTypes, v::CoVecTypes, A::MatTypes, ::Val{factor}) where {factor}
    @avx for m ∈ 1:size(A, 2)
        uₘ = zero(eltype(u))
        for k ∈ 1:size(A, 1)
            uₘ += v[k] * A[k, m]
        end
        u[m] += factor * uₘ
    end
end


# Adapted directly from Radford Neal's R code

# Alternative interface using Density type
hmc_sampler(x::Array{Float64}, dd::DifferentiableDensity, epsilon::Float64, L::Int64) = hmc_sampler(x, dd.f, dd.gradient, epsilon, L)
# Some default values for epsilon and L
hmc_sampler(x::Array{Float64}, dd::DifferentiableDensity, gx::Float64) = hmc_sampler(x,dd,.01,20)

function hmc_sampler(x::Array{Float64}, dd::DifferentiableDensity, params::Array)
  L = convert(Int64, params[2])
  hmc_sampler(x,dd,params[1],L)  # TODO Use a different data structure for params
end

function hmc_sampler(current_q::Array{Float64}, U::Function, grad_U::Function, epsilon::Float64, L::Int64, )

  q = current_q
  p = randn(length(q))'  # independent standard normal variates (row vector)
  current_p = p

  # Make a half step for momentum at the beginning
  p = p - epsilon * grad_U(q) / 2

  # Alternate full steps for position and momentum
  for i in 1:L
    # Make a full step for the position
    q = q + epsilon * p
    # Make a full step for the momentum, except at end of trajectory
    if i!=L
        p = p - epsilon * grad_U(q)
    end
  end

  # Make a half step for momentum at the end.
  p = p - epsilon * grad_U(q) / 2

  # Negate momentum at end of trajectory to make the proposal symmetric
  p = -p

  # Evaluate potential and kinetic energies at start and end of trajectory
  current_U = U(current_q)
  current_K = sum(current_p.^2) / 2
  proposed_U = U(q)
  proposed_K = sum(p.^2) / 2

  # Accept or reject the state at end of trajectory, returning either
  # the position at the end of the trajectory or the initial position
  if rand() < exp(current_U-proposed_U+current_K-proposed_K)
    return q, proposed_U  # accept
  else
    return current_q, current_U  # reject
  end
end

function bounded_hmc_sampler(x::Array{Float64}, dd::DifferentiableDensity, opts::Options)
  bounded_hmc_sampler(x,dd.f,dd.gradient,opts)
end

function bounded_hmc_sampler(current_q::Array{Float64}, U::Function, grad_U::Function, opts::Options)

  @defaults opts stepsize=.01 numsteps=50 bounded=[] lower_bounds=[] upper_bounds=[]

  epsilon = stepsize
  L = numsteps

  q = current_q
  p = randn(length(q))'  # independent standard normal variates (row vector)
  current_p = p

  # Make a half step for momentum at the beginning
  p = p - epsilon * grad_U(q) / 2

  # Alternate full steps for position and momentum
  for i in 1:L
    # Make a full step for the position
    q = q + epsilon * p
    # Make a full step for the momentum, except at end of trajectory
    if i!=L
        p = p - epsilon * grad_U(q)
    end

    # If out of bounds, reflect
    for j in bounded
        if q[j] < lower_bounds[j]
            q[j] = lower_bounds[j] + (lower_bounds[j] - q[j])
            p[j] = -p[j]
        end
        if q[j] > upper_bounds[j]
            q[j] = upper_bounds[j] - (q[j] - upper_bounds[j])
            p[j] = -p[j]
        end
    end
  end

  # Make a half step for momentum at the end.
  p = p - epsilon * grad_U(q) / 2

  # Negate momentum at end of trajectory to make the proposal symmetric
  p = -p

  # Evaluate potential and kinetic energies at start and end of trajectory
  current_U = U(current_q)
  current_K = sum(current_p.^2) / 2
  proposed_U = U(q)
  proposed_K = sum(p.^2) / 2

  @check_used opts
  # Accept or reject the state at end of trajectory, returning either
  # the position at the end of the trajectory or the initial position
  if rand() < exp(current_U-proposed_U+current_K-proposed_K)
    return q, proposed_U  # accept
  else
    return current_q, current_U  # reject
  end
end


mutable struct UMVE{T<:Real}
    A::Matrix{T}
    B::Matrix{T}
    C::OutputMatrix{T}
    G::Matrix{T}
    W::Matrix{T}
    V::Matrix{T}
    Ebar::Matrix{T}
    z::Vector{T}
    xi::Vector{T}
    Pi::Matrix{T} # state estimation error covariance
    Delta::Union{Matrix{T}, Missing} # input estimation error covariance

    function UMVE(A::S,
        B::S,
        C::OutputMatrix{T},
        E::S,
        W::AbstractMatrix{T},
        V::AbstractMatrix{T},
        decompose::Bool = true) where {T<:Real, S<:Matrix{T}}
        if isposdef(-W) #this does not exist
            throw(DomainError(W, "W must be positive definite"))
        end
        if !isposdef(V)
            throw(DomainError(V, "V must be positive definite"))
        end
        W = Matrix(W)
        V = Matrix(V)
        feasible, G, Ebar = checkUMVEFeasibility(C,E, decompose)
        nx = size(A,1)
        nxi = size(G,2)
        if feasible
            Pk = W
            Rk = R(A, C, Pk, W, V)
            Mk = M(G, C, Rk)
            new{T}(
                A,
                B,
                C,
                G,
                W,
                V,
                Ebar,
                zeros(T, nx),
                zeros(T, nxi),
                W,
                missing
            )
        else
            error("unfeasible")
        end
    end
end

function checkUMVEFeasibility(C::OutputMatrix, E::Matrix, decompose::Bool)
    rE = rank(E)
    feasible = rank(C*E) == rE
    if decompose
        SVD = svd(E)
        G = SVD.U[:,1:rE]
        Ebar = (Diagonal(SVD.S) * SVD.Vt)[1:rE,:]
    else
        nz = size(E,1)
        G = isempty(E) ? zeros(nz, 0) : Matrix{eltype(E)}(I, nz, nz)
        Ebar = E
    end
    feasible, G, Ebar
end

R(A::T, C::OutputMatrix{S}, Pk::T, W::T, V::T) where {S<:Real, T<:Matrix{S}} = C*(A*Pk*A' + W)*C' + V
M(G::T, C::OutputMatrix{S}, Rk::T) where {S<:Real, T<:Matrix{S}} = inv( G'C'*inv(Rk)*C*G ) * G'C'inv(Rk)
K(A::T, C::OutputMatrix{S}, Pk::T, W::T, Rk::T) where {S<:Real, T<:Matrix{S}} = C*(A*Pk*A' + W)*C'inv(Rk)
barA(Kk::T, C::OutputMatrix{S}, G::T, Mk::T) where {S<:Real, T<:Matrix{S}} = (I - Kk*C)*(I - G*Mk*C)
barL(Kk::T, C::OutputMatrix{S}, G::T, Mk::T) where {S<:Real, T<:Matrix{S}} = Kk + (I - Kk*C)*G*Mk

function updateState(o::UMVE{T}, u::Vector{T}, y::Vector{T}) where {T}
    Rk = R(o.A, o.C, o.Pi, o.W, o.V)
    # setfield!(o, :R, Rk)
    Mk = M(o.G, o.C, Rk)
    setfield!(o, :Delta, Mk*Rk*Mk')
    Kk = K(o.A, o.C, o.Pi, o.W, Rk)
    Abar = barA(Kk, o.C, o.G, Mk)
    Lbar = barL(Kk, o.C, o.G, Mk)
    biased_z = o.A*o.z + o.B*u
    setfield!(o, :z, Abar*biased_z + Lbar*y)
    setfield!(o, :xi, Mk*(y - o.C*biased_z))
    setfield!(o, :Pi, Abar*o.A*o.Pi*o.A'Abar' + Abar*o.W*Abar' + Lbar*o.V*Lbar')
end
updateState(o::UMVE{T}, u::T, y::Vector{T}) where {T} = updateState(o, [u], y)

outputMap(o::UMVE) = o.C*o.z

function Base.getproperty(o::UMVE, s::Symbol)
    if s === :nz
        return length(o.z)
    elseif s === :nxi
        return length(o.xi)
    else
        return getfield(o, s)
    end
end


using LatentDiffEq
using OrdinaryDiffEq
using BSON:@save
using Parameters: @with_kw
using Random
using Statistics
using Distributions
using ModelingToolkit
using Flux
using Luxor
using Images

@with_kw mutable struct Args_gen
      ## Latent Differential Equations
      # diffeq = Pendulum()

      tspan = (0.0f0, 4.95f0)         # time span
      dt = 0.05                       # timestep for ode solve
      u₀_range = (0, π/6)             # initial value range
      p₀_range = (1.0, 2.0)           # parameter value range
      n_traj = 450                    # Number of trajectories
      seed = 1                        # random seed

      ## High dimensional data arguments
      high_dim_args = (19, 1.75, 3.5)    # for the the simple pendulum
                                         # (pendulumlength, radius, rodthickness)
end

function generate_dataset(; diffeq = Pendulum(), kws...)
      args = Args_gen(; kws...)
      @unpack_Args_gen args
      Random.seed!(seed)

      ## Problem definition
      prob = remake(diffeq.prob, tspan = tspan)

      ## Ensemble functions definition
      prob_func(prob,i,repeat) = remake(prob, u0 = u0s[i], p = ps[i])

      # Sample initial condition and parameters from a uniform distribution
      ps = [rand_uniform(p₀_range, length(prob.p)) for i in 1:n_traj]
      # u0s = [rand_uniform(u₀_range, length(prob.u0)) for i in 1:n_traj]
      u0s = [vcat(rand_uniform_u0(u₀_range, length(prob.u0)), 0) for i in 1:n_traj]

      # Build and solve EnsembleProblem data
      ensemble_prob = EnsembleProblem(prob, prob_func = prob_func)
      sols = solve(ensemble_prob, diffeq.solver, saveat = args.dt, trajectories = n_traj)

      # Store the trajectories variables
      latent_data = [dropdims(Array(sol), dims = 2) for sol in sols]

      # Create animations with Luxor
      high_dim_data = [create_high_dim_data(sol, high_dim_args) for sol in sols]

      return latent_data, u0s, ps, high_dim_data
end

## util functions
rand_uniform(range::Tuple, size) = rand(Uniform(range...),(size,1))
rand_uniform(range::Array, size) = [rand(Uniform(r[1], r[2])) for r in eachrow(range)]
rand_uniform_u0(range::Tuple, size) = rand(Uniform(range...),(Int(size/2),1))

## Luxor util functions
# convert a timeseries to two trajectories
function maketrajectories(tseries, pendulumlength)
      innertrajectory = Point[]
      @inbounds for i in 1:length(tseries)
            angle1 = π/2 + tseries[i][1]
            pos1 = Point(pendulumlength * cos(angle1), pendulumlength * sin(angle1))
            push!(innertrajectory, pos1)
      end
      return innertrajectory
end

function drawrod(startpos, endpos, linewidth)
      rodlength = distance(startpos, endpos)
      setline(linewidth)
      line(startpos, endpos, :stroke)

      p = (startpos + endpos) / 2
      θ = atan(p[2] - startpos[2], p[1] - startpos[1])
      sethue("black")
      fontsize(8)
      # text("julia", p, valign=:middle, halign=:center, angle = θ)
      Luxor.text("|", p, valign=:middle, halign=:center, angle = θ)
end

function frame(pos, radius, rodthickness)
      background("black")
      sethue("white")
      offset = Point(0,-8.5)
      pos += offset
      sethue("white")
      circle(pos, radius, :fill)
      circle(offset, radius, :fill)
      drawrod(offset, pos, rodthickness)
      sethue("black")
      circle(offset, radius/2, :fill)
end

function create_frames(trajectories, radius, rodthickness, w=28, h=28)
      frames = Matrix{Float32}[]
      buffer = zeros(UInt32, w, h)
      for pos in trajectories
          a = @imagematrix! buffer frame(pos, radius, rodthickness) w h
          push!(frames, Float32.(Gray.(a)))
      end
      frames
end

function create_high_dim_data(sol, high_dim_args)
      pendulumlength, radius, rodthickness = high_dim_args
      trajectories = maketrajectories(sol, pendulumlength)
      frames = create_frames(trajectories, radius, rodthickness)
end

# latent_data, u0s, ps, high_dim_data = generate_dataset()

# data = (latent_data, u0s, ps, high_dim_data)
# root_dir = @__DIR__
# mkpath("$root_dir/data")
# @save "$root_dir/data/data.bson" data


# solve the ground state of Lieb-Liniger model
# also demonstrates the usage of preconditioners in Optim.jl

using TensorKit
using TensorOperations
using TensorKitAD
using KrylovKit
using Zygote
using KrylovKit
using LinearAlgebra
using Plots
using Optim
using ChainRules
using ChainRulesCore

using Revise
using statmech_tm_solver

#########################################################

# parameter for lieb-liniger
c = 1
μ = 1
L = 50

# bond dimensions
d, χ = 1, 8

# random initial guess
ψ_arr = rand(ComplexF64, (χ, d+1, χ))
convert_to_cmps(ψ_arr)

# cost function (energy) and its gradient
function f(ψ_arr::Array{ComplexF64, 3})
    ψ = convert_to_cmps(ψ_arr)
    return energy_lieb_liniger(ψ, c, L, μ)
end

function g!(ψg_arr::Array{ComplexF64, 3}, ψ_arr::Array{ComplexF64, 3})
    copyto!(ψg_arr, f'(ψ_arr))
end

# use LBFGS to optimize the energy for 250 steps
res = optimize(f, g!, ψ_arr, LBFGS(), Optim.Options(show_trace=true, iterations=250))
ψm = convert_to_cmps(Optim.minimizer(res));
ψm_arr = convert_to_array(ψm);

# calculate observables. total, kinetic and potential energies
E = energy_lieb_liniger(ψm, c, L, μ)
env = finite_env(K_mat(ψm, ψm), L)
env = permute(env, (2, 3), (4, 1))

op_ρ = particle_density(ψm)
op_k = kinetic(ψm)
tr(env * op_ρ)
tr(env * op_k)

#########################################################
# now we test the performance of the preconditioner

# construct the preconditioner
function _precondprep!(P::preconditioner, ψ_arr::Array{ComplexF64, 3})
    ψ = convert_to_cmps(ψ_arr)
    P.map = tangent_map(ψ, L)
    P.proj = gauge_fixing_proj(ψ, L)
end

# use preconditioned LBFGS, optimize the energy for additional 50 steps
P0 = preconditioner(ψm, L)
res1 = optimize(f, g!, ψm_arr, LBFGS(P = P0, precondprep=_precondprep!), Optim.Options(show_trace=true, iterations=50))
ψm1 = convert_to_cmps(Optim.minimizer(res1));
E1 = energy_lieb_liniger(ψm1, c, L, μ)
@show E1 - E

# use ordinary LBFGS, optimize the energy for additional 50 steps
res2 = optimize(f, g!, ψm_arr, LBFGS(), Optim.Options(show_trace=true, iterations=50))
ψm2 = convert_to_cmps(Optim.minimizer(res2));
E2 = energy_lieb_liniger(ψm2, c, L, μ)
@show E2 - E

# from original random init guess, 100 steps of ordinary LBFGS + 100 steps of preconditioned GradientDescent + 100 steps of preconditioned LBFGS
res3 = optimize(f, g!, ψ_arr, LBFGS(), Optim.Options(show_trace=true, iterations=100))
ψ3_arr = Optim.minimizer(res3)
ψm3 = convert_to_cmps(ψ3_arr)

P0 = preconditioner(ψm3, L)
res3 = optimize(f, g!, ψ3_arr, GradientDescent(P = P0, precondprep=_precondprep!), Optim.Options(show_trace=true, iterations=100))
ψ3_arr = Optim.minimizer(res3)
ψm3 = convert_to_cmps(ψ3_arr)

P0 = preconditioner(ψm3, L)
res3 = optimize(f, g!, ψ3_arr, LBFGS(P = P0, precondprep=_precondprep!), Optim.Options(show_trace=true, iterations=100))
ψm3 = convert_to_cmps(Optim.minimizer(res3));

E3 = energy_lieb_liniger(ψm3, c, L, μ)
@show E3 - E


import PaddedMatrices: logdensity, ∂logdensity!

using ReverseDiffExpressionsBase: adj, InitializedVarTracker, initialize!

struct BernoulliLogit{T} <: AbstractProbabilityDistribution{T} end

function logdensity(::BernoulliLogit{(false,true)}, y::BitVector, α::AbstractVector)
    t = target()
    @avx for i ∈ eachindex(α)
        αᵢ = α[i]
        invOmP = 1 + exp(αᵢ)
        nlogOmp = log(invOmP)
        nlogP = nlogOmp - αᵢ
        t -= y[i] ? nlogP : nlogOmP
    end
    extract_data(t)
end
logdensity(::BernoulliLogit{(false,false)}, y::BitVector, α::AbstractVector{T}) where {T} = zero(T)
constant(::BernoulliLogit, y::BitVector, α::AbstractVector{T}) where {T} = zero(T)

function ∂logdensity!((_,∂α)::Tuple{Nothing,<:AbstractVector}, ::BernoulliLogit{(false,true)}, y::BitVector, α::AbstractVector{T}) where {T}
    t = target()
    @avx for i ∈ eachindex(α)
        αᵢ = α[i]
        invOmP = 1 + exp(αᵢ)
        ∂logP = 1 / invOmP
        nlogOmp = log(invOmP)
        nlogP = nlogOmp - αᵢ
        t -= y[i] ? nlogP : nlogOmP
        ∂α[i] += y[i] ? ∂logP : ∂logP - one(T)
    end
    extract_data(t), (nothing, initialized(∂α))
end

struct BinomialLogit{T} <: AbstractProbabilityDistribution{T} end

function logdensity(::BinomialLogit{(false, true, false)}, y::AbstractVector, α::AbstractVector, N::AbstractVector)
    t = target()
    @avx for i ∈ eachindex(α)
        αᵢ = α[i]
        invOmP = 1 + exp( αᵢ )
        nlogOmP = log(invOmP)
        t -= N[i] * nlogOmP - y[i] * αᵢ
    end
    extract_data(t)
end

# faster, works for Floats that can't be convered to Int
# lbinom(n, p) = lgamma(n+1) - lgamma(n-p+1) - lgamma(p+1)
# more accurate, but doesn't allow n or p to have decimals.
lbinom(n, p) = first(logabsbinomial(trunc(Int,n), trunc(Int,p)))
function Binomial_logit_constant(
    ::Val{track}, s::AbstractVector{T}, α::AbstractVector{T}, N::AbstractVector{T}
) where {track,T}
    c = zero(T)
    @inbounds for i ∈ eachindex(s, N)
        c += lbinom(N[i],s[i])
    end
    c
end
function Binomial_logit_constant(
    ::Val{track}, s::AbstractVector{T}, α::AbstractVector{T}, N::T
) where {track,T}
    c = zero(T)
    @inbounds for i ∈ eachindex(s)
        c += lbinom(N,s[i])
    end
    c
end
function Binomial_logit_constant(
    s::AbstractVector{T}, α::AbstractVector{T}, N::Union{T,<:AbstractVector{T}}
) where {track,T}
    Binomial_logit_constant(Val{(false,true,false)}(), s, α, N)
end

function ∂logdensity!(::Nothing, ∂α::AbstractVector, ::Nothing, ::BernoulliLogit{(false,true,false)}, y::AbstractVector, α::AbstractVector{T}, N::AbstractVector) where {T}
    t = target()
    @avx for i ∈ eachindex(α)
        αᵢ = α[i]
        invOmP = 1 + exp(αᵢ)
        ∂logP = 1 / invOmP
        nlogOmp = log(invOmP)
        nlogP = nlogOmp - αᵢ
        t -= N[i] * nlogOmp - y[i] * αᵢ
        ∂α[i] += y[i] + (N[i]) * (∂logP - one(T))
    end
    extract_data(t)
end

function logdensity(
    ::BernoulliLogit, y::BitVector, X::AbstractMatrix{T}, β::AbstractVector{T}, α::AbstractFloat
)
    t = target()
    for i ∈ 1:size(X,1)
        a = α
        for p ∈ 1:size(X,2)
            a += X[i,p] * β[p]
        end
        invOmP = 1 + exp(a)
        nlogOmp = log(invOmP)
        nlogP = nlogOmp - a
        t -= y[i] ? nlogP : nlogOmP
    end
    extract_data(t)
end
Bernoulli_logit_constant(y, X, β, α, ::Any) = zero(eltype(β))

@generated function ∂logdensity!(
    ::Nothing, ::Nothing, ∂β::∂Β, ∂α::∂Α, ::BernoulliLogit,
    y::BitVector, X::AbstractMatrix{T}, β::AbstractVector{T}, α::AbstractFloat,
) where {T, ∂Β, ∂Α}
    β_is_param = ∂Β !== Nothing
    α_is_param = ∂Α !== Nothing
    if PaddedMatrices.is_sized(β)
        N_β = PaddedMatrices.number_elements(β)
        init_q = quote
            # T = promote_type(eltype(α),eltype(β),eltype(X))
#            target = zero($T)
            target = vbroadcast(Vec{$(VectorizationBase.pick_vector_width(T)),$T}, zero($T))
        end
        out_expr = quote end
        partial_exprs = Expr[]
        ∂P_undefined = true
        if β_is_param
            L_β = PaddedMatrices.pick_L(N_β, T)
            ∂P_undefined = false
            # push!(partial_exprs, :(∂P = OmP * P))
            push!(partial_exprs, :(∂logP = OmP ))
            push!(partial_exprs, :(∂logOmP = - P))
            out∂β = if isinitialized(∂Β)
                quote $([:(@inbounds ∂β[$n] += $(Symbol(:∂βP_, n))) for n ∈ 1:N_β]...) end
            else
                quote $([:(@inbounds ∂β[$n] = $(Symbol(:∂βP_, n))) for n ∈ 1:N_β]...) end
            end
            push!(out_expr.args, out∂β)
            for n ∈ 1:N_β
                push!(init_q.args, :($(Symbol(:∂βP_, n)) = zero($T)) )
            end
            partial_exprs_q = quote end
            for n ∈ 1:N_β
                push!(partial_exprs_q.args, :($(Symbol(:∂βP_, n)) += y[i] ? ∂logP * X[i,$n] : ∂logOmP * X[i,$n] ) )
            end
            push!(partial_exprs, partial_exprs_q)
        end
        # if X_is_param
        #     ∂P_undefined && push!(partial_exprs, :(∂P = OmP * P))
        #     ∂P_undefined = false
        #     push!(out_expr.args, Expr(:tuple, [Symbol(:∂XP_, n) for n ∈ 1:N_β]...))
        #     for n ∈ 1:N_β
        #         push!(init_q.args, :($(Symbol(:∂XP_, n)) = zero($T)) )
        #     end
        #     push!(partial_exprs, quote
        #         $([:($(Symbol(:∂XP_, n)) += y[i] ? β[$n] * ∂P : - β[$n] * ∂P) for n ∈ 1:N_β]...)
        #     end)
        # end
        if α_is_param
            # ∂P_undefined && push!(partial_exprs, :(∂P = OmP * P))
            ∂P_undefined && push!(partial_exprs, :(∂logP = OmP ))
            ∂P_undefined && push!(partial_exprs, :(∂logOmP = - P))
            ∂P_undefined = false
            push!(out_expr.args, isinitialized(∂Α) ? :(∂α[] += ∂αP) : :(∂α[] = ∂αP))
            push!(init_q.args, :(∂αP = zero($T)))
            push!(partial_exprs, :(∂αP += y[i] ? ∂logP : ∂logOmP))
        end
        W, Wshift = VectorizationBase.pick_vector_width_shift(T)
        # unroll_factor = max(8 >>> Wshift, 1)
        q = quote
            # $(Expr(:meta, :inline))
            $init_q
            @avx for i ∈ eachindex(y)
                # a = $(Expr(:call, :+, :α, [:(X[i,$n] * β[$n]) for n ∈ 1:N_β]...))
                a = vfmadd(X[i,1], β[1], α)
                $([:(a = vfmadd(X[i,$n], β[$n], a)) for n ∈ 2:N_β]...)
                OmP = one($T) / (one($T) + SLEEFPirates.exp( a ))
                P = one($T) - OmP
                logOmP = log(OmP)
                logP = a + logOmP
                target +=  y[i] ? logP : logOmP
                $(partial_exprs...)
            end
            $out_expr
            target
        end
    else
        throw("""
            Dynamically sized coefficient vector β is not yet supported.
            Feel free to file an issue, so that adding support will be prioritized.
            As a workaround, use a β with type parameterized by size.
            Eg, a padded vector from PaddedMatrices.jl or SVector from StaticArrays.jl.
        """)
    end
    q
end

struct LKJ{T} <: AbstractProbabilityDistribution{T} end

function logdensity(::LKJ{(true,false)}, L::AbstractCorrCholesky{N,T}, η::T) where {N,T}
    t = target()
    logd = logdiag(L)
    off = (N - 2) + 2η
    @avx for n ∈ eachindex(logd)
        t += (off - n) * logd[n]
    end
    extract_data(t)
end
function logdensity(::LKJ{(true,false)}, L::AbstractCorrCholesky{N,T}, η::T) where {N,T}
    t = target()
    logd = logdiag(L)
    off = (N - 2) + 2η
    @avx for n ∈ eachindex(logd)
        t += (off - n) * logd[n]
    end
    extract_data(t)
end

@generated function logdensity(::LKJ{track}, L::AbstractCorrCholesky{N,T}, η::T) where {N,T,track}
    W = VectorizationBase.pick_vector_width(N-1,T)
    # If the remainder is 1, we'll skip the first element (which equals 1, so the log is 0)
    # otherwise, we wont skip, because memory should be aligned with accessing the first element.
    rem1 = (N % W) == 1
    track_L, track_η = track
    if !track_L && !track_η
        return zero(T)
    end
    if track_L
        q = quote
            # $(Expr(:meta, :inline))
            target = vbroadcast(SVec{$W,$T}, zero($T))
            logd = logdiag(L)
            @avx for n ∈ 1:$(rem1 ? N-1 : N)
                target = vfmadd( ($(N - 2 - rem1) - n + 2η), logd[$(rem1 ? :(n+1) : :n)], target)
            end
            vtarget = extract_data(target)
        end
    end
    if track_η
        ηtarget = track_L ? :starget : :target
        ηoneadjust = if isodd(N)
            T(0.25*(N - 1)*log(π) - 0.25abs2(N-1) * log(2) - (N-1) * first(logabsgamma(0.5*(N+1))))
        else
            T(-0.25*(N)*log(π) + 0.25N * (3N - 4)*log(2) + N*first(logabsgamma(0.5N)) - (N-1)*first(logabsgamma(N)))
        end
        q_η = quote
            $ηtarget = zero($T)
            if η == one($T)
                for n ∈ 1:$((N-1) >>> 1)
                    $ηtarget -= first(logabsgamma($(T(2))*n))
                end
                # more accurate to sum this way?
                $ηtarget -= $ηoneadjust
            else
                for n ∈ 1:$(N-1)
                    $ηtarget -= first(logabsgamma(η + 0.5*($(N-1) - n)))
                end
                $ηtarget += $(N-1)*first(logabsgamma(η + $(0.5*(N-1))))
            end
        end
        push!(q.args, q_η)
        if track_L
            vstarget = :(Base.setindex(vbroadcast(Vec{$W,$T}, zero($T)), Core.VecElement($ηtarget), 1))
            push!(q.args, :(vadd(vtarget, $vstarget)))
        else
            push!(q.args, ηtarget)
        end
    end
    q
end
@generated function LKJ_constant(L::AbstractCorrCholesky{N,T}, η::T, ::Val{track}) where {N,T,track}
    :(LKJ(L,η,Val{$((!track[1],!track[2]))}()) - $(0.5*log(π)*StructuredMatrices.binomial2(N)))
end
@generated function ∂LKJ!(
    ∂L::Nothing, ∂η::Pη,
    L::AbstractCorrCholesky{N,T,M}, η::T
) where {N,T,PL,Pη,M}
    track_η = Pη !== Nothing
    if track_η
        η_uninit = !isinitialized(Pη)
        ∂ηop = η_uninit ? :(=) : :(+=)
    end
    W = VectorizationBase.pick_vector_width(N-1,T)
    if track_η
        ∂η_q = quote
            starget = zero($T)
            for n ∈ 1:$(N-1)
                c = η + 0.5*($(N-1) - n) # ∂c/∂η = 1
                starget -= first(logabsgamma(c))
                ∂ηs -= digamma(c) # ∂target/∂η = ∂target/∂c * ∂c/∂η = -digamma(c) * 1
            end
            $(Expr(∂ηop, :(∂η[]), ∂ηs))
            vadd(extract_data(target), Base.setindex(vbroadcast(Vec{$W,$T},zero($T)), startget, 1))
        end
    end
    if track_η
        q = quote
            # $(Expr(:meta, :inline))
            target = vbroadcast(SVec{$W,$T}, zero($T))
            logd = logdiag(L)
            ∂ηs = zero($T)
            @vvectorize $T for n ∈ 1:$(N-1)
                ∂ηₙ = logd[n+1]
                coef = ($(N - 3) - n + 2η)
                target = vfmadd( coef, ∂ηₙ, target )
                ∂ηs += $(T(2))*∂ηₙ
            end
            $∂η_q
        end
    else
        return zero(T)
    end
    simplify_expr(q)
end
@generated function ∂LKJ!(
    ∂L::PL, ∂η::Pη,
    L::AbstractCorrCholesky{N,T}, η::T
) where {N,T,Pη,M,PL <: StructuredMatrices.AbstractLowerTriangularMatrix{N,T,M}}
    track_η = Pη !== Nothing
    L_uninit = !isinitialized(PL)
    ∂Lop = L_uninit ? :(=) : :(+=)
    zeroq = quote
        @vvectorize $T for n ∈ 1:$(M-N)
            ∂L[n+$N] = zero($T)
        end
    end
    if track_η
        η_uninit = !isinitialized(Pη)
        ∂ηop = η_uninit ? :(=) : :(+=)
    end
    W = VectorizationBase.pick_vector_width(N-1,T)
    if track_η
        ∂η_q = quote
            starget = zero($T)
            for n ∈ 1:$(N-1)
                c = η + 0.5*($(N-1) - n) # ∂c/∂η = 1
                starget -= first(logabsgamma(c))
                ∂ηs -= digamma(c) # ∂target/∂η = ∂target/∂c * ∂c/∂η = -digamma(c) * 1
            end
            $(Expr(∂ηop, :(∂η[]), ∂ηs))
            vadd(extract_data(target), Base.setindex(vbroadcast(Vec{$W,$T},zero($T)), startget, 1))
        end
        q = quote
            target = vbroadcast(SVec{$W,$T}, zero($T))
            logd = logdiag(L)
            ∂ηs = zero($T)
            @vvectorize $T for n ∈ 1:$(N-1)
                ∂ηₙ = logd[n+1]
                coef = ($(N - 3) - n + 2η)
                target = vfmadd( coef, ∂ηₙ, target )
                $(Expr(∂Lop, :(∂L[n+1]), :(coef / L[n+1])))
                ∂ηs += $(T(2))*∂ηₙ
            end
            $∂η_q
        end
        if L_uninit
            pushfirst!(q.args, zeroq)
            pushfirst!(q.args, :(@inbounds ∂L[1] = zero($T)))
        end
        # pushfirst!(q.args, Expr(:meta,:inline))
    else
        q = quote
            target = vbroadcast(SVec{$W,$T}, zero($T))
            logd = logdiag(L)
            @vvectorize $T for n ∈ 1:$(N-1)
                ∂ηₙ = logd[n+1]
                coef = ($(N - 3) - n + 2η)
                target = vfmadd(coef, ∂ηₙ, target)
                $(Expr(∂Lop, :(∂L[n+1]), :(coef / L[n+1])))
            end
        end
        if L_uninit # need to zero skipped elements
            pushfirst!(q.args, :(@inbounds ∂L[1] = zero($T)))
            push!(q.args, zeroq)
        end
        # pushfirst!(q.args, Expr(:meta,:inline))
        push!(q.args, :(extract_data(target)))
    end
    simplify_expr(q)
end
@generated function ∂LKJ(sp::PaddedMatrices.StackPointer, ::Val{track}, L::AbstractCorrCholesky{N,T}, η::T) where {N,T,track}
    track_L, track_η = track
    q = quote end
    ret_expr = Expr(:tuple, :target)
    if track_L
        push!(q.args, :((sp, ∂L) = PtrVector{$N,$T}(sp)))
        push!(ret_expr.args, :(Diagonal(∂L)))
    else
        push!(q.args, :(∂L = nothing))
    end
    if track_η
        push!(q.args, :(∂η = Ref{$T}()))
        push!(ret_expr.args, :(∂η[]))
    else
        push!(q.args, :(∂η = nothing))
    end
    push!(q.args, :(∂LKJ!(uninitialized(∂L), uninitialized(∂η), L, η)))
    push!(q.args, :(sp, $ret_expr))
    q
end
push!(DISTRIBUTION_DIFF_RULES, :LKJ)

function gamma_quote(
    M::Int, ::Type{T}, (yisvec, αisvec, βisvec)::NTuple{3,Bool},
    (track_y, track_α, track_β)::NTuple{3,Bool}, partial::Bool,
    (inity, initα, initβ)::NTuple{3,Bool} = (true,true,true)
) where {T}
    q = quote end
    pre_quote = quote end
    return_expr = quote end
    loop = any((yisvec, αisvec, βisvec))
    # set initialized_target to loop; if we are looping, we'll start out at zero
    initialized_target = loop
    if yisvec
        yexpr = :(y[i])
        push!(pre_quote.args, :(logy = PaddedMatrices.LazyMap(log, y)))
        logyexpr = :(logy[i])
    else
        yexpr = :y
        push!(pre_quote.args, :(logy = Base.log(y)))
        logyexpr = :logy
    end
    if αisvec
        αexpr = :(α[i])
        lgammaαexpr = :(first(SpecialFunctions.logabsgamma(α[i])))
        αm1expr = :(α[i] - one(eltype(α)))
    else
        αexpr = :α
        lgammaαexpr = :(lgammaα)
        αm1expr = :(αm1)
        push!(pre_quote.args, :(lgammaα = first(SpecialFunctions.logabsgamma(α))))
        push!(pre_quote.args, :(αm1 = α - one(eltype(α))))
    end
    if βisvec
        βexpr = :(β[i])
        push!(pre_quote.args, :(logβ = PaddedMatrices.LazyMap(log, y)))
        logβexpr = :(logβ[i])
    else
        βexpr = :β
        logβexpr = :logβ
        push!(pre_quote.args, :(logβ = Base.log(β)))
    end
    if partial
        if track_y
            if yisvec
                ∂yassignment = :(=)
                ∂ystorage = :(∂yᵢ)
                # push!(return_expr.args, :(∂y'))
            else
                ∂ystorage = :(∂ys)
                if inity
                    push!(return_expr.args, :(∂y[] = ∂ys))
                else
                    push!(return_expr.args, :(∂y[] += ∂ys))
                end
                if loop
                    ∂yassignment = :(+=)
                    push!(pre_quote.args, :(∂ys = zero($T)))
                else
                    ∂yassignment = :(=)
                end
            end
        end
        if track_α
            if αisvec
                ∂αassignment = :(=)
                ∂αstorage = :(∂αᵢ)
            else
                ∂αstorage = :(∂αs)
                if initα
                    push!(return_expr.args, :(∂α[] = ∂αs))
                else
                    push!(return_expr.args, :(∂α[] += ∂αs))
                end
                if loop
                    ∂αassignment = :(+=)
                    push!(pre_quote.args, :(∂αs = zero($T)))
                else
                    ∂αassignment = :(=)
                end
            end
        end
        if track_β
            if βisvec
                ∂βassignment = :(=)
                ∂βstorage = :(∂βᵢ)
            else
                ∂βstorage = :(∂βs)
                if initβ
                    push!(return_expr.args, :(∂β[] = ∂βs))
                else
                    push!(return_expr.args, :(∂β[] += ∂βs))
                end
                if loop
                    ∂βassignment = :(+=)
                    push!(pre_quote.args, :(∂βs = zero($T)))
                else
                    ∂βassignment = :(=)
                end
            end
        end
    end
    if track_α || track_β
        push!(q.args, :( lβ = $logβexpr))
        if initialized_target
            push!(q.args, :( target = vmuladd($αexpr, lβ, target) ) )
        else
            push!(q.args, :( target = $αexpr * lβ) )
            initialized_target = true
        end
        if partial
            track_α && push!(q.args, Expr(∂αassignment, ∂αstorage, :lβ) )
            track_β && push!(q.args, Expr(∂βassignment, ∂βstorage, αexpr) )
        end
    end
    if track_α || track_y
        push!(q.args, :(ly = $logyexpr))
        if initialized_target
            push!(q.args, :( target = vmuladd($αm1expr, ly, target) ) )
        else
            push!(q.args, :( target = $αm1expr * ly ) )
            initialized_target = true
        end
        if partial
            if track_α
                if αisvec
                    push!(q.args, :(∂α₂ = $∂αstorage + ly))
                else
                    push!(q.args, Expr(:(+=), ∂αstorage, :ly) )
                end
            end
            track_y && push!(q.args, Expr(∂yassignment, ∂ystorage, :($αm1expr / $yexpr) ) )
        end
    end
    if track_β || track_y
        # initialized_target == true
        push!(q.args, :(target = vfnmadd($βexpr, $yexpr, target)))
        if partial
            if track_β
                if βisvec
                    push!(q.args, Expr(initβ ? :(=) : :(+=), :(∂β[i]), :($∂βstorage - $yexpr) ) )
                else
                    push!(q.args, Expr(:(-=), ∂βstorage, yexpr ) )
                end
            end
            if track_y
                if yisvec
                    push!(q.args, Expr(inity ? :(=) : :(+=), :(∂y[i]), :($∂ystorage - $βexpr) ) )
                else
                    push!(q.args, Expr(:(-=), ∂ystorage, βexpr ) )
                end
            end
        end
    end
    if track_α
        # initialized_target == true, because (if track_α || track_β) == true
        push!(q.args, :(target = vsub(target, $lgammaαexpr)))
        if partial
            if αisvec
                push!(q.args, Expr(initα ? :(=) : :(+=), :(∂α[i]), :(∂α₂ - SpecialFunctions.digamma($αexpr) ) ) )
            else
                push!(pre_quote.args, Expr(:(-=), ∂αstorage, :($M * SpecialFunctions.digamma(α)) ) )
            end
        end
    end
    q = if loop
        quote
            # $(Expr(:meta,:inline))
            @fastmath begin
                $pre_quote
            end
            target = vbroadcast(Vec{$(VectorizationBase.pick_vector_width(M,T)),$T}, zero($T))
#            out = zero($T)
            @vvectorize $T for i ∈ 1:$M
                $q
            end
            @fastmath begin
                $return_expr
            end
            target
        end
    else
        quote
            # $(Expr(:meta,:inline))
            @fastmath begin
                $pre_quote
                $q
                $return_expr
                target
            end
        end
    end
    simplify_expr(q)
end

function gamma_alloc_quote(
    M::Int, T, (yisvec, αisvec, βisvec)::NTuple{3,Bool},
    (track_y, track_α, track_β)::NTuple{3,Bool}, sp::Bool = true
)
    q = quote end
    return_expr = Expr(:tuple, :target)
    if track_y
        if yisvec
            if sp
                push!(q.args, :((sp,∂y) = PaddedMatrices.PtrVector{$M,$T}(sp)))
            else
                push!(q.args, :(∂y = PaddedMatrices.FixedSizeVector{$M,$T}(undef)))
            end
            push!(return_expr.args, :(∂y))
        else
            push!(q.args, :(∂y = Ref{$T}()))
            push!(return_expr.args, :(∂y[]))
        end
    else
        push!(q.args, :(∂y = nothing))
    end
    if track_α
        if αisvec
            if sp
                push!(q.args, :(∂α = PaddedMatrices.PtrVector{$M,$T}(sp)))
            else
                push!(q.args, :(∂α = PaddedMatrices.FixedSizeVector{$M,$T}(undef)))
            end
            push!(return_expr.args, :(∂α))
        else
            push!(q.args, :(∂α = Ref{$T}()))
            push!(return_expr.args, :(∂α[]))
        end
    else
        push!(q.args, :(∂α = nothing))
    end
    if track_β
        if βisvec
            if sp
                push!(q.args, :(∂β = PaddedMatrices.PtrVector{$M,$T}(sp)))
            else
                push!(q.args, :(∂β = PaddedMatrices.FixedSizeVector{$M,$T}(undef)))
            end
            push!(return_expr.args, :(∂β))
        else
            push!(q.args, :(∂β = Ref{$T}()))
            push!(return_expr.args, :(∂β[]))
        end
    else
        push!(q.args, :(∂β = nothing))
    end
    push!(q.args, :(target = ∂Gamma!(uninitialized(∂y), uninitialized(∂α), uninitialized(∂β), y, α, β)))
    push!(q.args, return_expression(return_expr, sp))
    q
end

# α * log(β) + (α-1) * log(y) - β*y - lgamma(α)
@generated function Gamma(
    ::Val{track},
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T, <: PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T, <: PaddedMatrices.AbstractFixedSizeVector{M,T}}
) where {track,M,T}
#) where {track,T,M}
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    gamma_quote(M, T, (true, αisvec, βisvec), track, false)
end

@generated function ∂Gamma!(
    ∂y::∂YN, ∂α::∂ΑN, ∂β::∂ΒN,
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T,<:PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T,<:PaddedMatrices.AbstractFixedSizeVector{M,T}}
) where {∂YN, ∂ΑN, ∂ΒN, M, T}
    yisvec = true
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    track_y = ∂YN !== Nothing
    track_α = ∂ΑN !== Nothing
    track_β = ∂ΒN !== Nothing
    inity = !isinitialized(∂YN)
    initα = !isinitialized(∂ΑN)
    initβ = !isinitialized(∂ΒN)
    gamma_quote(
        M, T, (yisvec, αisvec, βisvec),
        (track_y, track_α, track_β), true,
        (inity, initα, initβ)
    )
end

@generated function ∂Gamma(
    ::Val{track},
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T,<:PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T,<:PaddedMatrices.AbstractFixedSizeVector{M,T}}
#) where {track,M,T}
) where {track,T,M}
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    gamma_alloc_quote(M, T, (true, αisvec, βisvec), track, false)
end
@generated function ∂Gamma(
    sp::StackPointer,
    ::Val{track},
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T,<:PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T,<:PaddedMatrices.AbstractFixedSizeVector{M,T}}
) where {track,M,T}
            # ::Val{track}) where {track,T,M}
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    gamma_alloc_quote(M, T, (true, αisvec, βisvec), track, true)
end
@generated function Gamma(::Val{track}, y::Union{T,RealFloat{<:Any,T}}, α::T, β::T) where {track,T <: Real}
    gamma_quote(1, T, (false, false, false), track, false)
end
@generated function ∂Gamma(::Val{track}, y::Union{T,RealFloat{<:Any,T}}, α::T, β::T) where {track,T <: Real}
    gamma_alloc_quote(1, T, (false, false, false), track, false)
end
@generated function ∂Gamma(sp::StackPointer, ::Val{track}, y::Union{T,RealFloat{<:Any,T}}, α::T, β::T) where {track,T <: Real}
    gamma_alloc_quote(1, T, (false, false, false), track, true)
end
@generated function ∂Gamma!(
    ∂y::∂YN, ∂α::∂ΑN, ∂β::∂ΒN,
    y::Union{T,<:RealFloat{<:Any,T}},
    α::Union{T,<:RealFloat{<:Any,T}},
    β::Union{T,<:RealFloat{<:Any,T}}
) where {∂YN, ∂ΑN, ∂ΒN, T}
    track_y = ∂YN !== Nothing
    track_α = ∂ΑN !== Nothing
    track_β = ∂ΒN !== Nothing
    inity = !isinitialized(∂YN)
    initα = !isinitialized(∂ΑN)
    initβ = !isinitialized(∂ΒN)
    gamma_quote(
        1, T, (false,false,false),
        (track_y, track_α, track_β), true,
        (inity, initα, initβ)
    )
end

push!(DISTRIBUTION_DIFF_RULES, :Gamma)

function beta_quote(
    M::Int, ::Type{T}, (yisvec, αisvec, βisvec)::NTuple{3,Bool},
    (track_y, track_α, track_β)::NTuple{3,Bool}, (inity, initα, initβ)::NTuple{3,Bool}, partial::Bool) where {T}
    q = quote end
    pre_quote = quote end
    return_expr = quote end
    loop = any((yisvec, αisvec, βisvec))
    sp &= loop
    # set initialized_target to loop; if we are looping, we'll start out at zero
    initialized_target = loop
    if yisvec
        yexpr = :(y[i])
        push!(pre_quote.args, :(logity = PaddedMatrices.LazyMap(SLEEFPirates.logit, y)))
        logyexpr = :(SLEEFPirates.log(y[i]))
        logomyexpr = :(logy - logity[i])
    else
        yexpr = :y
        logyexpr = :logy
        logomyexpr = :logomy
        push!(pre_quote.args, :(logy = Base.log(y)))
        push!(pre_quote.args, :(logomy = SLEEFPirates.logit(y) - logy))
    end
    if αisvec
        αexpr = :(α[i])
        αm1expr = :(α[i] - one($T))
    else
        αexpr = :α
        αm1expr = :(αm1)
        push!(pre_quote.args, :(αm1 = α - one(α)))
    end
    if βisvec
        βexpr = :(β[i])
        βm1expr = :(β[i] - one($T))
    else
        βexpr = :β
        βm1expr = :(βm1)
        push!(pre_quote.args, :(βm1 = β - one($T)))
    end
    if αisvec || βisvec
        lbetaβexpr = :(first(SpecialFunctions.logabsbeta($αexpr, $βexpr)))
    else # neither are vectors
        lbetaβexpr = :lbetaαβ
        push!(pre_quote.args, :(lbetaαβ = first(SpecialFunctions.logabsbeta(α, β))))
    end
    if partial
        if track_y
            if yisvec
                yassignment = :(=)
                ∂ystorage = :∂yᵢ
            else
                ∂ystorage = :∂ys
                if inity
                    push!(return_expr.args, :(∂y[] = ∂ys))
                else
                    push!(return_expr.args, :(∂y[] += ∂ys))
                end
                if loop
                    push!(pre_quote.args, :(∂ys = zero($T)))
                    yassignment = :(+=)
                else
                    yassignment = :(=)
                end
            end
        end
        if track_α
            if αisvec
                αassignment = :(=)
                ∂αstorage = :∂αᵢ
                dgαexpr = :(SpecialFunctions.digamma(α[i]))
            else
                ∂αstorage = :∂αs
                dgαexpr = :dgα
                push!(pre_quote.args, :(dgα = SpecialFunctions.digamma(α)))
                if initα
                    push!(return_expr.args, :(∂α[] = ∂αs))
                else
                    push!(return_expr.args, :(∂α[] += ∂αs))
                end
                if loop
                    push!(pre_quote.args, :(∂αs = zero($T)))
                    αassignment = :(+=)
                else
                    αassignment = :(=)
                end
            end
        end
        if track_β
            if βisvec
                βassignment = :(=)
                ∂βstorage = :(∂βᵢ)
                dgβexpr = :(SpecialFunctions.digamma(β[i]))
            else
                ∂βstorage = :∂βs
                dgβexpr = :dgβ
                push!(pre_quote.args, :(dgβ = SpecialFunctions.digamma(β)))
                if initβ
                    push!(return_expr.args, :(∂β[] = ∂βs))
                else
                    push!(return_expr.args, :(∂β[] += ∂βs))
                end
                if loop
                    push!(pre_quote.args, :(∂βs = zero($T)))
                    βassignment = :(+=)
                else
                    βassignment = :(=)
                end
            end
        end
        if track_α || track_β
            if αisvec || βisvec
                push!(q.args, :(dgαβ = SpecialFunctions.digamma($αexpr + $βexpr) ))
            else # both are scalars
                push!(pre_quote.args, :(dgαβ = SpecialFunctions.digamma(α + β) ))
            end
        end
    end
    if track_α || track_y
        push!(q.args, :( am1 = $αm1expr))
        push!(q.args, :( logy = $logyexpr))
        if initialized_target
            push!(q.args, :( target = vmuladd(am1, logy, target) ) )
        else
            push!(q.args, :( target = am1 * logy ) )
            initialized_target = true
        end
        if partial
            track_α && push!(q.args, Expr(αassignment, ∂αstorage, :logy) )
            track_y && push!(q.args, Expr(yassignment, ∂ystorage, :(am1 / $yexpr) ) )
        end
    end
    if track_β || track_y
        push!(q.args, :( bm1 = $βm1expr))
        push!(q.args, :( logomy = $logomyexpr))
        if initialized_target
            push!(q.args, :( target = vmuladd(bm1, logomy, target) ) )
        else
            push!(q.args, :( target = bm1 * logomy ) )
            initialized_target = true
        end
        if partial
            track_β && push!(q.args, Expr(βassignment, ∂βstorage, :logomy) )
            if track_y
                if yisvec
                    push!(q.args, Expr(inity ? :(=) : :(+=), :(∂y[i]), :($∂ystorage - bm1 / (one($T) - $yexpr) ) ))
                else
                    push!(q.args, Expr(:(=), ∂ystorage, :($∂ystorage - bm1 / (one($T) - $yexpr) ) ))
                end
            end
        end
    end
    if track_α || track_β
        push!(q.args, :(ly = $logyexpr))
        if initialized_target
            push!(q.args, :( target = vsub(target, $lbetaβexpr) ) )
        else
            push!(q.args, :( target = -$lbetaβexpr ) )
            initialized_target = true
        end
        if partial
            if track_α
                if αisvec
                    push!(q.args, Expr(initα ? :(=) : :(+=), :(∂α[i]), :($∂αstorage + dgαβ - $dgαexpr) ) )
                else
                    push!(q.args, Expr(:(=), ∂αstorage, :($∂αstorage + dgαβ - $dgαexpr) ) )
                end
            end
            if track_β
                if βisvec
                    push!(q.args, Expr(initβ ? :(=) : :(+=), :(∂β[i]), :($∂βstorage + dgαβ - $dgβexpr) ) )
                else
                    push!(q.args, Expr(:(=), ∂βstorage, :($∂βstorage + dgαβ - $dgβexpr) ) )
                end
            end
        end
    end
    q = if loop
        quote
            # $(Expr(:meta,:inline))
            @fastmath begin
                $pre_quote
                target = vbroadcast(Vec{$(VectorizationBase.pick_vector_width(T)),$T}, zero($T))
#                out = zero($T)
            end
            @vectorize $T for i ∈ 1:$M
                $q
            end
            @fastmath begin
                $return_expr
            end
            target
        end
    else
        quote
            # $(Expr(:meta,:inline))
            @fastmath begin
                $pre_quote
                $q
                $return_expr
                target
            end
        end
    end
    simplify_expr(q)
end
function beta_alloc_quote(M, T, (yisvec, αisvec, βisvec), (track_y, tack_α, track_β), sp::Bool = true)
    q = quote end
    if track_y
        if yisvec
            if sp
                push!(q.args, :((sp,∂y) = PaddedMatrices.PtrVector{$M,$T}(sp)))
            else
                push!(q.args, :(∂y = PaddedMatrices.FixedSizeVector{$M,$T}(undef)))
            end
            push!(return_expr.args, :(∂y))
        else
            push!(q.args, :(∂y = Ref{$T}()))
            push!(return_expr.args, :∂y[])
        end
    else
        push!(q.args, :(∂y = nothing))
    end
    if track_α
        if αisvec
            if sp
                push!(q.args, :((sp,∂α) = PaddedMatrices.PtrVector{$M,$T}(undef)))
            else
                push!(q.args, :(∂α = PaddedMatrices.FixedSizeVector{$M,$T}(undef)))
            end
            push!(return_expr.args, :(∂α))
        else
            push!(q.args, :(∂α = Ref{$T}()))
            push!(return_expr.args, :(∂α[]))
        end
    else
        push!(q.args, :(∂α = nothing))
    end
    if track_β
        if βisvec
            if sp
                push!(pre_quote.args, :((sp,∂β) = PaddedMatrices.PtrVector{$M,$T}(undef)))
            else
                push!(pre_quote.args, :(∂β = PaddedMatrices.FixedSizeVector{$M,$T}(undef)))
            end
            push!(return_expr.args, :(∂β))
        else
            push!(q.args, :(∂β = Ref{$T}()))
            push!(return_expr.args, :(∂β[]))
        end
    else
        push!(q.args, :(∂β = nothing))
    end
    push!(q.args, :(target = ∂Beta!(uninitialized(∂y), uninitialized(∂α), uninitialized(∂β), y, α, β)))
    push!(q.args, return_expression(return_expr, sp))
    q
end

# α * log(β) + (α-1) * log(y) - β*y - lgamma(α)
@generated function Beta(
    ::Val{track},
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}}
            # ::Val{track}) where {track,T,M}
) where {track,M,T}
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    beta_quote(M, T, (true, αisvec, βisvec), track, (false,false,false), false)
end
@generated function ∂Beta!(
    ∂y::∂YN, ∂α::∂ΑN, ∂β::∂ΒN,
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}}
) where {∂YN, ∂ΑN, ∂ΒN, M, T}
    yisvec = true
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    track_y = ∂YN !== Nothing
    track_α = ∂YΑ !== Nothing
    track_β = ∂YΒ !== Nothing
    inity = !isinitialized(∂YN)
    initα = !isinitialized(∂YΑ)
    initβ = !isinitialized(∂YΒ)
    beta_quote(
        M, T, (yisvec, αisvec, βisvec), (track_y, track_α, track_β), (inity, initα, initβ), true
    )
end
@generated function ∂Beta(
    ::Val{track},
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}}
) where {track,M,T}
            # ::Val{track}) where {track,T,M}
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    beta_alloc_quote(M, T, (true, αisvec, βisvec), track, false)
end
@generated function ∂Beta(
    sp::StackPointer,
    ::Val{track},
    y::PaddedMatrices.AbstractFixedSizeVector{M,T},
    α::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}},
    β::Union{T,Int,<:PaddedMatrices.AbstractFixedSizeVector{M,T}}
) where {track,M,T}
# ) where {track,T,M}
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    βisvec = isa(β, PaddedMatrices.AbstractFixedSizeVector)
    beta_quote(M, T, (true, αisvec, βisvec), track, true)
end
@generated function Beta(::Val{track}, y::T, α::Union{T,Int}, β::Union{T,Int}) where {T <: Real,track}
    beta_quote(1, T, (false, false, false), track, false)
end
@generated function ∂Beta(::Val{track}, y::T, α::Union{T,Int}, β::Union{T,Int}) where {track,T <: Real}
    beta_quote(1, T, (false, false, false), track, true)
end
push!(DISTRIBUTION_DIFF_RULES, :Beta)

function lsgg_quote(
    M::Int, T,
    (yisvec, αisvec, ϕisvec, δisvec, σisvec)::NTuple{5,Bool},
    (track_y, track_α, track_ϕ, track_δ, track_σ)::NTuple{5,Bool},
    (inity, initα, initϕ, initδ, initσ)::NTuple{5,Bool},
    partial::Bool
)
    @assert !any((track_α, track_ϕ, track_δ, track_σ))
    q = quote end
    yexpr = yisvec ? :(y[m]) : :y
    αexpr = αisvec ? :(α[m]) : :α
    ϕexpr = ϕisvec ? :(ϕ[m]) : :ϕ
    δexpr = δisvec ? :(δ[m]) : :δ
    σexpr = σisvec ? :(σ[m]) : :σ
    push!(q.args, :(smy = $σexpr - $yexpr) )
    push!(q.args, :(asmy = $αexpr * smy))
    push!(q.args, :(smyd = smy - $ϕexpr) )
    push!(q.args, :(ptsmyd = smyd * $δexpr) )
    push!(q.args, :(eptsmyd = $(any((yisvec, ϕisvec, δisvec, σisvec)) ? :SLEEFPirates : :Base).exp(ptsmyd)) )
    W = VectorizationBase.pick_vector_width(M,T)
    if partial
        ∂yop = inity ? :(=) : :(+=)
        if yisvec
            return quote
                # Inlined because of Julia SIMD corruption bug (if sp)
                # inlined to avoid heap allocation of mvector (if !sp)
                # $(Expr(:meta,:inline))
                target = SIMDPirates.vbroadcast(Vec{$W,$T}, zero($T))
                LoopVectorization.@vvectorize for m ∈ 1:$M
                    $q
                    target = LoopVectorization.SIMDPirates.vadd(SIMDPirates.vsub(asmy, eptsmyd), target)
                    $(Expr(∂yop, :(∂y[m]), :(SIMDPirates.vfmsub($δexpr, eptsmyd, $αexpr))))
                end
                target
            end
        else
            @assert !any((αisvec, ϕisvec, δisvec, σisvec))
            return quote
                @fastmath begin
                    $q
                    target = asmy - eptsmyd
                    $(Expr(∂yop, :(∂y[]), :($δexpr * eptsmyd - $αexpr)))
                end
                target
            end
        end
    else
        if yisvec
            return quote
                # Inlined because of Julia SIMD corruption bug
                # $(Expr(:meta,:inline))
                target = SIMDPirates.vbroadcast(Vec{$W,$T}, zero($T))
                LoopVectorization.@vvectorize for m ∈ 1:$M
                    $q
                    target = LoopVectorization.SIMDPirates.vadd(SIMDPirates.vsub(asmy, eptsmyd), target)
                end
                target
            end
        else
            @assert !any((αisvec, ϕisvec, δisvec, σisvec))
            return quote
                @fastmath begin
                    $q
                    target = asmy - eptsmyd
                end
                target
            end
        end
    end
end
function lsgg_alloc_quote(
    M::Int, T,
    (yisvec, αisvec, ϕisvec, δisvec, σisvec)::NTuple{5,Bool},
    (track_y, track_α, track_ϕ, track_δ, track_σ)::NTuple{5,Bool},
    sp::Bool
)
    @assert !any((track_α, track_ϕ, track_δ, track_σ))
    if yisvec
        return quote
            # Inlined because of Julia SIMD corruption bug (if sp)
            # inlined to avoid heap allocation of mvector (if !sp)
            # $(Expr(:meta,:inline))
            $(sp ? :((sp, ∂y) = PtrVector{$M,$T}(sp)) : :( ∂y = FixedSizeVector{$M,$T}(undef)))
            target = ∂lsgg!(uninitialized(∂y),nothing,nothing,nothing,nothing,y,α,ϕ,δ,σ)
            $(sp ? :(sp, (target, ∂y)) : :(target, ∂y))
        end
    else
        @assert !any((αisvec, ϕisvec, δisvec, σisvec))
        return quote
            ∂y = Ref{$T}()
            target = ∂lsgg!(uninitialized(∂y),nothing,nothing,nothing,nothing,y,α,ϕ,δ,σ)
            $(sp ? :(sp, (target, ∂y[])) : :(target, ∂y[]))
        end
    end
end

@generated function lsgg(::Val{track}, y::AbstractFixedSizeVector{M,T}, α, ϕ, δ, σ) where {M,T,track}
    lsgg_quote(
        M, T, (true,
        α <: AbstractArray,
        ϕ <: AbstractArray,
        δ <: AbstractArray,
        σ <: AbstractArray),
        track, ntuple(_ -> false, Val(5)), false
    )
end
@generated function ∂lsgg!(
    ∂y::∂YN, ∂α::Nothing, ∂ϕ::Nothing, ∂δ::Nothing, ∂σ::Nothing,
    y::AbstractFixedSizeVector{M,T}, α, ϕ, δ, σ
) where {∂YN, M, T}
    yisvec = true
    αisvec = isa(α, PaddedMatrices.AbstractFixedSizeVector)
    ϕisvec = isa(ϕ, PaddedMatrices.AbstractFixedSizeVector)
    δisvec = isa(δ, PaddedMatrices.AbstractFixedSizeVector)
    σisvec = isa(σ, PaddedMatrices.AbstractFixedSizeVector)
    track = (true, false, false, false, false)
    init = (!isinitialized(∂YN), false, false, false, false)
    lsgg_quote(
       M, T, (yisvec, αisvec, ϕisvec, δisvec, σisvec), track, init, true
    )
end
@generated function ∂lsgg(::Val{track}, y::AbstractFixedSizeVector{M,T}, α, ϕ, δ, σ) where {M,T,track}
    lsgg_alloc_quote(
        M, T, (true,
        α <: AbstractArray,
        ϕ <: AbstractArray,
        δ <: AbstractArray,
        σ <: AbstractArray),
        track, false
    )
end
@generated function ∂lsgg(sp::StackPointer, ::Val{track}, y::AbstractFixedSizeVector{M,T}, α, ϕ, δ, σ) where {M,T,track}
    lsgg_alloc_quote(
        M, T, (true,
        α <: AbstractArray,
        ϕ <: AbstractArray,
        δ <: AbstractArray,
        σ <: AbstractArray),
        track, true
    )
end
@generated function lsgg(::Val{track}, y::T, α, ϕ, δ, σ) where {T<:Real,track}
    @assert !any((
        α <: AbstractArray,
        ϕ <: AbstractArray,
        δ <: AbstractArray,
        σ <: AbstractArray,
    ))
    lsgg_quote(
        1, T, (false,false,false,false,false),
        track, (false,false,false,false,false), false
    )
end
@generated function ∂lsgg!(
    ∂y::∂YN, ∂α::Nothing, ∂ϕ::Nothing, ∂δ::Nothing, ∂σ::Nothing,
    y::T, α, ϕ, δ, σ
) where {T<:Real,∂YN}
    @assert !any((
        α <: AbstractArray,
        ϕ <: AbstractArray,
        δ <: AbstractArray,
        σ <: AbstractArray,
    ))
    lsgg_alloc_quote(
        1, T, (false,false,false,false,false),
        track, (!isinitialized(T),false,false,false,false), true
    )
end
@generated function ∂lsgg(::Val{track}, y::T, α, ϕ, δ, σ) where {T<:Real,track}
    @assert !any((
        α <: AbstractArray,
        ϕ <: AbstractArray,
        δ <: AbstractArray,
        σ <: AbstractArray,
    ))
    lsgg_alloc_quote(
        1, T, (false,false,false,false,false), track, false
    )
end
@generated function ∂lsgg(sp::StackPointer, ::Val{track}, y::T, α, ϕ, δ, σ) where {T,track}
    @assert !any((
        α <: AbstractArray,
        ϕ <: AbstractArray,
        δ <: AbstractArray,
        σ <: AbstractArray,
    ))
    lsgg_alloc_quote(
        1, T, (false,false,false,false,false), track, true
    )
end
push!(DISTRIBUTION_DIFF_RULES, :lsgg)

emax(a, em, ed50) = em * a / (a + ed50)
emaxi(a, em, ed50) = @fastmath a * em * ed50 / (1 + ed50 * a)
emaxn(a, em, ed50) = @fastmath a*em*ed50 / (a + ed50)
normal_lpdf(y, μ, nhτ, nlogrootτ) = nhτ * abs2(y-μ) - nlogrootτ

# Val{(true,true,true,true,true,false,false)}()
@generated function EₘₐₓNMA(
    ::Val{track}, α::AbstractVector{T}, σᵤ::T, Eₘₐₓ::AbstractVector{T}, ED₅₀::AbstractVector{T}, σ::T,
    Treatments::StructuredMatrices.FixedSizeRaggedMatrix{M,N,P,<:Integer,<:Integer}, dose::AbstractVector{T}
) where {T, track, M, N, P}
    track_α, track_σᵤ, track_Eₘₐₓ, track_ED₅₀, track_σ, track_treat, track_dose = track
    @assert track_treat == false && track_dose == false
    padM = PaddedMatrices.calc_padding(M-1, T)
    W, Wshift = VectorizationBase.pick_vector_width_shift(M-1)

    quote
        target = zero($T)
        ταr = 1 / σᵤ
        τ = 1 / (σ*σ)
        j = 1
        col_lengths = Treatments.column_lengths
        vτ = SIMDPirates.vbroadcast(Vec{$W,$T}, τ)
        vnh = SIMDPirates.vbroadcast(Vec{$W,$T}, $(T(-0.5)))
        $([:($(Symbol(:vτma_,k)) = vmul(vτ, $(Expr(:tuple, [Core.VecElement{T}(2m/(m+1)) for m in k*W+1:(k+1)*W]...)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnhτ_,k)) = vmul(vnh, $(Symbol(:vτma_,k)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnlogrootτ_,k)) = vmul(vnh, SLEEFPirates.log($(Symbol(:vτma_,k))))) for k in 0:((padM>>>Wshift)-1)]...)
        @inbounds nhτ = $(Expr(:tuple, [:(($(Symbol(:vnhτ_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @inbounds nlogrootτ = $(Expr(:tuple, [:(($(Symbol(:vnlogrootτ_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @fastmath @inbounds for i in 1:$N
            t = Treatments[j]
            αᵢ = α[j]
            emi1 = αᵢ
            t == 0 || (emi1 -= emax(dose[j], Eₘₐₓ[t], ED₅₀[t]))
            αᵢ *= ταr
            target -= 0.5αᵢ*αᵢ
            s = zero($T)
            j += 1
            for k in 1:col_lengths[i]-1
                t = Treatments[j]
                emik = emax(dose[j], Eₘₐₓ[t], ED₅₀[t])
                δik = α[j] - emik - emi1 # - αi1
                # target += normal_lpdf(δik, s/(k-1), τ * (2*(k-1)) / k )
                # @show δik, s/k, nhτ[k], nlogrootτ[k]
                target += normal_lpdf(δik, s/k, nhτ[k], nlogrootτ[k])
                s += δik
                j += 1
            end
        end
        muladd($(T(N)), log(ταr), target)
    end
end
@inline function ∂emax(a, em, ed50)
    #    f(a,b,c) =  a*c   / (b + c)
    # ∂f∂a(a,b,c) =    c   / (b + c)
    # ∂f∂b(a,b,c) = -a*c   / (b + c)^2
    # ∂f∂c(a,b,c) =  a * b / (b + c)^2
    @fastmath begin
        invaped50 = 1 / ( a + ed50 )
        eminvaped50 = em * invaped50
        f = a * eminvaped50
        dfda = ed50 * eminvaped50 * invaped50
        dfdemax = a * invaped50
        dfded50 = - f * invaped50
    end
    f, dfda, dfdemax, dfded50
end
@inline function ∂emaxi(a, em, ed50)
    #    f(a,b,c) = a*b*c / (1 + a*c)
    # ∂f∂a(a,b,c) =   b*c / (1 + a*c)^2
    # ∂f∂b(a,b,c) = a * c / (1 + a*c)
    # ∂f∂c(a,b,c) = a*b   / (1 + a*c)^2
    @fastmath begin
        invdenom = 1 / muladd(ed50, a, one(ed50))
        doseed50 = a * ed50
        dfdemax = doseed50 * invdenom
        f = dfdemax * em
        emdivdenomdivdenom = em * invdenom * invdenom
        dfda = emdivdenomdivdenom * ed50
        dfded50 = emdivdenomdivdenom * a
    end
    f, dfda, dfdemax, dfded50
end
@inline function ∂emaxn(a, em, ed50)
    @fastmath begin
        invdenom = 1 / (ed50 + a)
        dfdem = ed50 * a * invdenom
        f = dfdem * em
        eminvdenominvdenom = em * invdenom * invdenom
        dfded50 = eminvdenominvdenom * a * a
        dfda = eminvdenominvdenom * ed50 * ed50
    end
    f, dfda, dfdem, dfded50
end
function ∂normal_lpdf(y, μ, nhτ, nlogrootτ, nhσ²)
    @fastmath begin
        z = y - μ
        nτ = nhτ + nhτ
        z² = z * z
        f = nhτ * z² - nlogrootτ
        ∂f∂y = nτ * z
        ∂f∂μ = -∂f∂y
        ∂f∂τ =  - 0.5 * z² - nhσ²
    end
    f, ∂f∂y, ∂f∂μ, ∂f∂τ
end

@generated function ∂EₘₐₓNMA!(
    dαv::∂ΑN, ∂σᵤ::∂ΣμN, demaxv::DEMAX, ded50v::DED50, ∂σ::∂ΣN,
    α::AbstractVector{T}, σᵤ::T, Eₘₐₓ::AbstractFixedSizeVector{C,T}, ED₅₀::AbstractFixedSizeVector{C,T}, σ::T,
    treatments::StructuredMatrices.FixedSizeRaggedMatrix{M,N,P,<:Integer,<:Integer}, dose::AbstractVector{T}
) where {T,C,M,N,P,∂ΑN, ∂ΣμN, DEMAX, DED50, ∂ΣN}
    padM = PaddedMatrices.calc_padding(M-1, T)
    W, Wshift = VectorizationBase.pick_vector_width_shift(M-1)
    inits_q = quote target = zero($T) end
    if !isinitialized(DEMAX) || !isinitialized(DED50)
        loop_body = quote end
        isinitialized(DEMAX) || push!(loop_body.args, :(demaxv[i] = zero($T)))
        isinitialized(DED50) || push!(loop_body.args, :(ded50v[i] = zero($T)))
        loop_q = quote
            @inbounds for i ∈ eachindex(demaxv)
                $loop_body
            end
        end
        push!(inits_q.args, loop_q)
    end
    quote
        dsdemax = PtrMatrix{2,$M,$T}(SIMDPirates.alloca(Val{$(2M)}(), $T))
        $inits_q
        invσ = 1/σ
        # σ² =  σ * σ
        τ = abs2(invσ)
        dτdσ = -2τ*invσ
        # @show dτdσ
        # tau, dtaudsigma = 1/abs2(sigma), -2/sigma^3
        j = 1
        dτ = zero(σ)
        col_lengths = treatments.column_lengths
        ταr = 1 / σᵤ
        ταr² = ταr * ταr
        dταrdσᵤ = -ταr * ταr
        dταr = zero($T)
        vτ = SIMDPirates.vbroadcast(Vec{$W,$T}, τ)
        vnh = SIMDPirates.vbroadcast(Vec{$W,$T}, $(T(-0.5)))
        $([:($(Symbol(:vτma_,k)) = vmul(vτ, $(Expr(:tuple, [Core.VecElement{T}(2m/(m+1)) for m in k*W+1:(k+1)*W]...)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnhτ_,k)) = vmul(vnh, $(Symbol(:vτma_,k)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnhσ²_,k)) = SIMDPirates.vfdiv(vnh,$(Symbol(:vτma_,k)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnlogrootτ_,k)) = vmul(vnh, SLEEFPirates.log($(Symbol(:vτma_,k))))) for k in 0:((padM>>>Wshift)-1)]...)
        τscales = $(Expr(:tuple, [T(2k/(k+1)) for k in 1:padM]...))
        @inbounds nhτ = $(Expr(:tuple, [:(($(Symbol(:vnhτ_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @inbounds nlogrootτ = $(Expr(:tuple, [:(($(Symbol(:vnlogrootτ_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @inbounds nhσ² = $(Expr(:tuple, [:(($(Symbol(:vnhσ²_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @fastmath @inbounds for i in 1:$N
            ti = treatments[j]
            ji = j
            αi1 = α[ji]
            if ti == 0
                emi1, demi1ddose, demi1demaxv, demi1ded50v = (zero($T),zero($T),zero($T),zero($T))
            else
                emi1, demi1ddose, demi1demaxv, demi1ded50v = ∂emax(dose[j], Eₘₐₓ[ti], ED₅₀[ti])
            end
            αi1²ταr = αi1 * αi1 * ταr
            target -= 0.5 * αi1²ταr * ταr
            dταr -= αi1²ταr
            dtargetdαin = αi1*ταr²
            emi1minusαi1 = emi1 - αi1
            dtargetdαi = zero($T)
            dtargetdemaxv = zero($T)
            dtargetded50v = zero($T)
            s = zero($T)
            j += 1
            # fill!(dsdemax, 0.0)
            for k in 1:col_lengths[i]-1
                t = treatments[j]
                emik, demikddose, demikdemaxv, demikded50v = ∂emax(dose[j], Eₘₐₓ[t], ED₅₀[t])
                δik = α[j] - emik + emi1minusαi1
                sscale = 1 / k
                # f, dfdy, dfdμ, dfdτ = dnormal_lpdf( δik, s*sscale, τ * τscale, logτs[k-1] )
                f, dfdy, dfdμ, dfdτ = ∂normal_lpdf( δik, s*sscale, nhτ[k], nlogrootτ[k], nhσ²[k] )
                target += f
                # @show dfdτ, τscales[k]
                dτ += dfdτ * τscales[k]
                dtargetdαi -= dfdy * sscale
                if ti > 0
                    dtargetdemaxv += dfdy*demi1demaxv*sscale
                    dtargetded50v += dfdy*demi1ded50v*sscale
                end
                for p in 1:k-1
                    tp = treatments[j + p - k]
                    dαv[j + p - k] += dfdμ*sscale
                    demaxv[tp] -= dsdemax[1,p]*dfdμ*sscale
                    ded50v[tp] -= dsdemax[2,p]*dfdμ*sscale
                end
                $(Expr(isinitialized(∂ΑN) ? :(+=) : :(=), :(dαv[j]), :dfdy))
                demaxv[t] -= dfdy*demikdemaxv
                ded50v[t] -= dfdy*demikded50v
                dsdemax[1,k] = demikdemaxv
                dsdemax[2,k] = demikded50v
                s += δik
                j += 1
            end
            $(Expr(isinitialized(∂ΑN) ? :(+=) : :(=), :(dαv[ji]), :(dtargetdαi - dtargetdαin)))
            if ti > 0
                demaxv[ti] += dtargetdemaxv
                ded50v[ti] += dtargetded50v
            end
        end
        target = muladd($(T(N)), Base.log(ταr), target)
        dταr = muladd($(T(N)), σᵤ, dταr)
        $(Expr(isinitialized(∂ΣμN) ? :(+=) : :(=), :(∂σᵤ[]), :(dταrdσᵤ*dταr)))
        $(Expr(isinitialized(∂ΣN) ? :(+=) : :(=), :(∂σ[]), :(dτ*dτdσ)))
        target
    end
end
# Val{(true,true,true,true,false,false)}()
@generated function ∂EₘₐₓNMA(
    sp::StackPointer, ::Val{track}, α::AbstractVector{T}, σᵤ::T, Eₘₐₓ::AbstractFixedSizeVector{C,T}, ED₅₀::AbstractFixedSizeVector{C,T}, σ::T,
    treatments::StructuredMatrices.FixedSizeRaggedMatrix{M,N,P,<:Integer,<:Integer}, dose::AbstractVector{T}
) where {T,C,M,N,P,track}
    padM = PaddedMatrices.calc_padding(M-1, T)
    W, Wshift = VectorizationBase.pick_vector_width_shift(M-1)
    quote
        sp, dαv = PtrVector{$P,$T}(sp)
        sp, demaxv = PtrVector{$C,$T}(sp)
        sp, ded50v = PtrVector{$C,$T}(sp)
        dsdemax = PtrMatrix{2,$M,$T}(pointer(sp,$T))
        target = zero($T)
        @inbounds for i in eachindex(demaxv)#, ded50v)
            ded50v[i] = zero($T)
            demaxv[i] = zero($T)
        end
        invσ = 1/σ
        # σ² =  σ * σ
        τ = abs2(invσ)
        dτdσ = -2τ*invσ
        # @show dτdσ
        # tau, dtaudsigma = 1/abs2(sigma), -2/sigma^3
        j = 1
        dτ = zero(σ)
        col_lengths = treatments.column_lengths
        ταr = 1 / σᵤ
        ταr² = ταr * ταr
        dταrdσᵤ = -ταr * ταr
        dταr = zero($T)
        vτ = SIMDPirates.vbroadcast(Vec{$W,$T}, τ)
        vnh = SIMDPirates.vbroadcast(Vec{$W,$T}, $(T(-0.5)))
        $([:($(Symbol(:vτma_,k)) = vmul(vτ, $(Expr(:tuple, [Core.VecElement{T}(2m/(m+1)) for m in k*W+1:(k+1)*W]...)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnhτ_,k)) = vmul(vnh, $(Symbol(:vτma_,k)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnhσ²_,k)) = SIMDPirates.vfdiv(vnh,$(Symbol(:vτma_,k)))) for k in 0:((padM>>>Wshift)-1)]...)
        $([:($(Symbol(:vnlogrootτ_,k)) = vmul(vnh, SLEEFPirates.log($(Symbol(:vτma_,k))))) for k in 0:((padM>>>Wshift)-1)]...)
        τscales = $(Expr(:tuple, [T(2k/(k+1)) for k in 1:padM]...))
        @inbounds nhτ = $(Expr(:tuple, [:(($(Symbol(:vnhτ_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @inbounds nlogrootτ = $(Expr(:tuple, [:(($(Symbol(:vnlogrootτ_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @inbounds nhσ² = $(Expr(:tuple, [:(($(Symbol(:vnhσ²_,m>>>Wshift))[$(1+m&(W-1))]).value) for m in 0:padM-1]...))
        @fastmath @inbounds for i in 1:$N
            ti = treatments[j]
            ji = j
            αi1 = α[ji]
            if ti == 0
                emi1, demi1ddose, demi1demaxv, demi1ded50v = (zero($T),zero($T),zero($T),zero($T))
            else
                emi1, demi1ddose, demi1demaxv, demi1ded50v = ∂emax(dose[j], Eₘₐₓ[ti], ED₅₀[ti])
            end
            αi1²ταr = αi1 * αi1 * ταr
            target -= 0.5 * αi1²ταr * ταr
            dταr -= αi1²ταr
            dtargetdαin = αi1*ταr²
            emi1minusαi1 = emi1 - αi1
            dtargetdαi = zero($T)
            dtargetdemaxv = zero($T)
            dtargetded50v = zero($T)
            s = zero($T)
            j += 1
            # fill!(dsdemax, 0.0)
            for k in 1:col_lengths[i]-1
                t = treatments[j]
                emik, demikddose, demikdemaxv, demikded50v = ∂emax(dose[j], Eₘₐₓ[t], ED₅₀[t])
                δik = α[j] - emik + emi1minusαi1
                sscale = 1 / k
                # f, dfdy, dfdμ, dfdτ = dnormal_lpdf( δik, s*sscale, τ * τscale, logτs[k-1] )
                f, dfdy, dfdμ, dfdτ = ∂normal_lpdf( δik, s*sscale, nhτ[k], nlogrootτ[k], nhσ²[k] )
                target += f
                # @show dfdτ, τscales[k]
                dτ += dfdτ * τscales[k]
                dtargetdαi -= dfdy * sscale
                if ti > 0
                    dtargetdemaxv += dfdy*demi1demaxv*sscale
                    dtargetded50v += dfdy*demi1ded50v*sscale
                end
                for p in 1:k-1
                    tp = treatments[j + p - k]
                    dαv[j + p - k] += dfdμ*sscale
                    demaxv[tp] -= dsdemax[1,p]*dfdμ*sscale
                    ded50v[tp] -= dsdemax[2,p]*dfdμ*sscale
                end
                dαv[j] = dfdy
                demaxv[t] -= dfdy*demikdemaxv
                ded50v[t] -= dfdy*demikded50v
                dsdemax[1,k] = demikdemaxv
                dsdemax[2,k] = demikded50v
                s += δik
                j += 1
            end
            dαv[ji] = dtargetdαi - dtargetdαin
            if ti > 0
                demaxv[ti] += dtargetdemaxv
                ded50v[ti] += dtargetded50v
            end
        end
        target = muladd($(T(N)), Base.log(ταr), target)
        dταr = muladd($(T(N)), σᵤ, dταr)
        sp, (target, dαv', dταrdσᵤ*dταr, demaxv', ded50v', dτ*dτdσ)
    end
end

push!(DISTRIBUTION_DIFF_RULES, :EₘₐₓNMA)

#=

function multinomial_quote(M, isvararg::Bool, T)
    q = quote end
    W, Wshift = VectorizationBase.pick_vector_width_shift(M, T)
    Wm1 = W - 1
    niter = M >>> Wshift
    nrem = (M + Wm1) & Wm1
    if isvararg
        push!(q.args, :(target = SIMDPirates.vmul(y, SLEEFPirates.log()) ))
        for i ∈ 1:niter-1
            push!(q.args, :( ))
        end

    else
        for i ∈ 0:niter-1

        end
    end

end

=#


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   ModiaMath.Frames (ModiaMath/Frames/_module.jl)
#
import ..StaticArrays: SVector, @SMatrix

"""
    const ModiaMath.Vector3D = SVector{3,Float64}

Type of a vector in 3D space (e.g. position vector of the origin of a frame)
"""
const Vector3D = SVector{3,Float64}

"""
    const ModiaMath.ZeroVector3D = Vector3D(0.0, 0.0, 0.0)

Constant of a Vector3D where all elements are zero
"""
const ZeroVector3D = Vector3D(0.0, 0.0, 0.0)

"""
    M = ModiaMath.skew(e::AbstractVector)

Return the skew symmetric matrix `M::SMatrix{3,3,Float64,9}` of vector `e` (`length(e) = 3`)
"""
skew(e::AbstractVector) = @SMatrix([  0.0   -e[3]    e[2];
                                      e[3]   0.0    -e[1];
                                     -e[2]   e[1]    0.0 ])


using LinearAlgebra

mutable struct Bird <: AbstractAgent
    id::Int
    pos::NTuple{2,Float64}
    vel::NTuple{2,Float64}
    speed::Float64
    cohere_factor::Float64
    separation::Float64
    separate_factor::Float64
    match_factor::Float64
    visual_distance::Float64
end

"""
```julia
flocking(;
    n_birds = 100,
    speed = 1.0,
    cohere_factor = 0.25,
    separation = 4.0,
    separate_factor = 0.25,
    match_factor = 0.01,
    visual_distance = 5.0,
    extent = (100, 100),
    spacing = visual_distance / 1.5
)
```
Same as in [Flock model](@ref).
"""
function flocking(;
    n_birds = 100,
    speed = 1.0,
    cohere_factor = 0.25,
    separation = 4.0,
    separate_factor = 0.25,
    match_factor = 0.01,
    visual_distance = 5.0,
    extent = (100, 100),
    spacing = visual_distance / 1.5,
)
    space2d = ContinuousSpace(extent, spacing)
    model = ABM(Bird, space2d, scheduler = random_activation)
    for _ in 1:n_birds
        vel = Tuple(rand(2) * 2 .- 1)
        add_agent!(
            model,
            vel,
            speed,
            cohere_factor,
            separation,
            separate_factor,
            match_factor,
            visual_distance,
        )
    end
    return model, flocking_agent_step!, dummystep
end

function flocking_agent_step!(bird, model)
    neighbor_ids = nearby_ids(bird, model, bird.visual_distance)
    N = 0
    match = separate = cohere = (0.0, 0.0)
    for id in neighbor_ids
        N += 1
        neighbor = model[id].pos
        heading = neighbor .- bird.pos

        cohere = cohere .+ heading
        if edistance(bird.pos, neighbor, model) < bird.separation
            separate = separate .- heading
        end
        match = match .+ model[id].vel
    end
    N = max(N, 1)
    cohere = cohere ./ N .* bird.cohere_factor
    separate = separate ./ N .* bird.separate_factor
    match = match ./ N .* bird.match_factor
    bird.vel = (bird.vel .+ cohere .+ separate .+ match) ./ 2
    bird.vel = bird.vel ./ norm(bird.vel)
    move_agent!(bird, model, bird.speed)
end


# svrg implementation

export SVRG

"""
# struct SVRG

The structure describing iterations for stochastic variance reduced gradient.

## Method

We have the following sequence:
   - a batch_gradient as the full gradient at current position
   - storing gradient and position before the mini batch sequence
   - then for `nb_mini_batch` :
       - uniform sampling of **one** term of the summation
       - computation of the gradient of the term at current position and the gradient
                        at position before mini batch
       - computation of direction of propagation as the batch gradient + gradient of term at current
                        position -  gradient of term at position before mini batch sequence
       - update of position with adequate step size

   The step size used in the algorithm is constant and according to the ref paper it should be of the order of
   L/4 where L is the lipschitz constant of the function to minimize

## FIELDS

- nb\\_minibatch : The number of minibatch to run after one full batch run
- stepsize: the step to use in
"""
mutable struct SVRG
    # number of minibatch of one term.
    nb_minibatch :: Int64
    # step_size
    stepsize :: Float64
end

"""
# function minimize(svrgpb::SVRG, evaluation::Evaluator{F,G}, max_iter::Int64, initialposition::Array{Float64,N})

Generic function used in minimisations for all algorithms.

The function has 3 type parameters F, G and N:

- F and G corresponds to the types of function used in instantiating the structure TermFunction{F} and TermmGradient{G}
(see in test and examples directories the cases TestLinearRegression or TestLogisticRegression).

- N is the dimension of array for searched parameters/gradients which depend upon the problem and the way we modelize it.
(See Logistic regression where we used N=2)

## Args
- svrgpb : the structure describing main parameters of the batch strategy
- Evaluator{F,G}
- max_iter : maximum number of iterations
- initialposition : initial position at iteration beginning

"""
function minimize(svrgpb::SVRG, evaluation::Evaluator{F,G}, max_iterations, initial_position::Array{Float64,N}) where {F,G,N}
    direction = zeros(Float64, size(initial_position))
    full_batch_gradient = zeros(Float64, size(initial_position))
    mini_batch_gradient_current = zeros(Float64, size(initial_position))
    mini_batch_gradient_origin = zeros(Float64, size(initial_position))
    nbterms = get_nbterms(evaluation)
    #
    mt = MersenneTwister(117);
    position = Array{Float64,N}(initial_position)
    iteration = 0
    initial_value = 0.
    more = true
    while more
        iteration += 1
        # get iteration parameters
        # compute gradient on large batch index set and store initial position
        compute_gradient!(evaluation, position ,full_batch_gradient)
        # get number of mini batch
        nb_mini_batch = svrgpb.nb_minibatch
        position_before_mini_batch = Array{Float64,N}(position)
        # loop on small batch iterations
        for i in 1:nb_mini_batch
            # sample mini batch terms
            term = 1 + floor(Int64, nbterms * rand(mt))
            compute_gradient!(evaluation, position , term, mini_batch_gradient_current)
            compute_gradient!(evaluation, position_before_mini_batch , term, mini_batch_gradient_origin)
            direction = mini_batch_gradient_current - mini_batch_gradient_origin + full_batch_gradient;
            position = position - svrgpb.stepsize * direction;
        end
        value = compute_value(evaluation, position)
        if iteration == 1
            initial_value = value
        end
        if iteration >= max_iterations
            @info "Reached maximal number of iterations required, initial_value final value:" initial_value value
            return position, value
        end
    end

end