text:
# Wrapper function for log likelihood

"""
    ll(y, model, θ; initiate = "first")
Likelihood function for Count Data models.

* `y`: Time series
* `model`: Model Specification
* `θ`: Parameters (Vector or parameter)
* `initiate`: How is the time series initiated?

The time series can be initiated by "first", the first observed value, "intercept",
the intercept parameter (possibly exponentiated if log-link), or "marginal", the marginal
mean of the process.

# Example
```julia-repl
ll(y, model, pars)
```
"""
function ll(y::Array{T, 1} where T<:Integer,
    model::INGARCHModel,
    θ::parameter;
    initiate = "first")::Tuple{Float64, Vector{Float64}}

    T = length(y)
    p = length(model.pastObs)
    q = length(model.pastMean)
    if p == 0
        P = 0
    else
        P = maximum(model.pastObs)
    end
    if q == 0
        Q = 0
    else
        Q = maximum(model.pastMean)
    end

    nb = model.distr == "NegativeBinomial"
    zi = model.zi

    M = convert(Int64, maximum([P, Q]))

    if !parametercheck(θ, model)
        return -Inf, fill(-Inf, T)
    end

    ϕ = θ.ϕ
    ω = θ.ω

    indrel = (M+1):T

    λ = LinPred(y, model, θ, initiate)
    lls = fill(0.0, T)

    if any(λ .< 0)
        return (-Inf, fill(-Inf, T))
    end

    if !nb
        d = Poisson.(λ[indrel])
        if ω > 0
            lls[indrel] = log.(pdf.(d, y[indrel]).*(1 - ω) .+ (y[indrel] .== 0).*ω)
        else
            lls[indrel] = logpdf.(d, y[indrel])
        end
        LL = sum(lls)
    else
        ptemp = @. ϕ[1]/(ϕ[1] + λ[indrel])
        d = NegativeBinomial.(ϕ[1], ptemp)
        if ω > 0
            lls[indrel] = log.(pdf.(d, y[indrel]).*(1 - ω) .+ (y[indrel].== 0).*ω)
        else
            lls[indrel] = logpdf.(d, y[indrel])
        end

        LL = sum(lls)
    end

    return LL, lls
end

function ll(y::Array{T, 1} where T<:Integer,
    model::INARCHModel,
    θ::parameter;
    initiate = "first")::Tuple{Float64, Vector{Float64}}

    T = length(y)
    X = model.X
    r = length(model.external)

    p = length(model.pastObs)
    if p == 0
        P = 0
    else
        P = maximum(model.pastObs)
    end

    M = P

    logl = (model.link == "Log")
    lin = !logl

    nb = model.distr == "NegativeBinomial"
    zi = model.zi

    M = P
    nObs = T - M
    nPar = 1 + p + r + nb + zi

    if !parametercheck(θ, model)
        return -Inf, fill(-Inf, T)
    end

    β0 = θ.β0
    αRaw = θ.α
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    indrel = (M+1):T

    λ = LinPred(y, model, θ, initiate)
    if any(λ .< 0)
        return (-Inf, fill(-Inf, T))
    end

    lls = zeros(T)
    if !nb
        d = Poisson.(λ[indrel])
        if length(ω) == 1
            lls[indrel] = log.(pdf.(d, y[indrel]).*(1 - ω) .+ (y[indrel].== 0).*ω)
        elseif length(ω) == 0
            lls[indrel] = logpdf.(d, y[indrel])
        end
        LL = sum(lls)
    else
        ptemp = ϕ[1]./(ϕ[1] .+ λ[indrel])
        d = NegativeBinomial.(ϕ[1], ptemp)
        if length(ω) == 1
            lls[indrel] = log.(pdf.(d, y[indrel]).*(1 - ω) .+ (y[indrel].== 0).*ω)
        elseif length(ω) == 0
            lls[indrel] = logpdf.(d, y[indrel])
        end

        LL = sum(lls)
    end

    return LL, lls
end

function ll(y::Array{T, 1} where T<:Integer,
    model::IIDModel,
    θ::parameter;
    initiate = "first")::Tuple{Float64, Vector{Float64}}

    T = length(y)
    X = model.X

    r = length(model.external)

    logl = (model.link == "Log")
    lin = !logl

    nb = model.distr == "NegativeBinomial"
    zi = model.zi

    nObs = T
    nPar = 1 + r + nb + zi

    if !parametercheck(θ, model)
        return -Inf, fill(-Inf, T)
    end

    β0 = θ.β0
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    λ = LinPred(y, model, θ, initiate)
    if any(λ .< 0)
        return (-Inf, fill(-Inf, T))
    end

    if !nb
        d = Poisson.(λ)
        if length(ω) == 1
            lls = log.(pdf.(d, y).*(1 - ω) .+ (y.== 0).*ω)
        elseif length(ω) == 0
            lls = logpdf.(d, y)
        end
        LL = sum(lls)
    else
        ptemp = ϕ[1]./(ϕ[1] .+ λ)
        d = NegativeBinomial.(ϕ[1], ptemp)
        if length(ω) == 1
            lls = log.(pdf.(d, y).*(1 - ω) .+ (y.== 0).*ω)
        elseif length(ω) == 0
            lls = logpdf.(d, y)
        end

        LL = sum(lls)
    end

    return LL, lls
end

function ll(y::Array{T, 1} where T<:Integer,
    model::INMAModel,
    θ::parameter)::Tuple{Float64, Vector{Float64}}

    T = length(y)
    ymax = maximum(y)
    X = model.X
    r = length(model.external)

    logl1 = model.link[1] == "Log"
    lin1 = !logl1

    logl2 = model.link[2] == "Log"
    lin2 = !logl2

    nb1 = model.distr[1] == "NegativeBinomial"
    nb2 = model.distr[2] == "NegativeBinomial"
    if r == 0
        nb2 = false
    else
        if sum(model.external) == 0
            nb2 = false
        end
    end

    zi = model.zi

    q = length(model.pastMean)
    if q == 0
        Q = 0
    else
        Q = maximum(model.pastMean)
    end
    M = Q

    if !parametercheck(θ, model)
        return -Inf, fill(-Inf, T)
    end

    β0 = θ.β0
    βRaw = θ.β
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    if q > 0
        if q < Q
            β = zeros(Q)
            β[model.pastMean] = βRaw
        else
            β = βRaw
        end
    end

    rE = sum(model.external)
    rI = sum(.!model.external)

    if r > 0
        iE = findall(model.external)
        iI = findall(.!model.external)

        ηI = η[iI]
        ηE = η[iE]
    end

    μ = zeros(T)
    λ = fill(β0, T)

    if rI > 0
        for i = 1:rI
            λ .+= η[iI[i]].*X[iI[i], :]
        end
    end

    if logl1
        λ = exp.(λ)
    end

    if rE > 0
        for i = 1:rE
            μ .+= η[iE[i]].*X[iE[i], :]
        end
        if logl2
            μ = exp.(μ)
        end
    end

    PR = zeros(ymax + 1, T)
    PZ = zeros(ymax + 1, T)
    for t = 1:T
        if nb1
            PR[:, t] = pdf.(NegativeBinomial(ϕ[1], ϕ[1]/(ϕ[1] + λ[t])), 0:ymax)
        else
            PR[:, t] = pdf.(Poisson(λ[t]), 0:ymax)
        end

        if zi
            PR[2:ymax + 1, t] = PR[2:ymax + 1, t].*(1 .- ω./(1 .- PR[1, t]))
            PR[1, t] += ω
        end

        if rE > 0
            if nb2
                PZ[:, t] = pdf.(NegativeBinomial(ϕ[2], ϕ[2]/(ϕ[2] + μ[t])), 0:ymax)
            else
                PZ[:, t] = pdf.(Poisson(μ[t]), 0:ymax)
            end
        else
            PZ[:, t] = [1; fill(0, ymax)]
        end
    end

    lls = zeros(T)

    if q == 0
        for t = M+1:T
            lls[t] = log(dot(PR[1:y[t]+1], PZ[y[t]+1:-1:1]))
        end
        LL = sum(lls[M+1:end])

        return LL, lls
    end

    # Compute P(Σ α_i∘Y_{t-i} + Z_t = k) first --> pARZ
    PARZ = PZ

    if Q == 1
        QQ = zeros(ymax + 1, ymax + 1, T)
        a = zeros(ymax + 1, T)
        ϕmat = zeros(ymax + 1, T)
        u = zeros(ymax + 1, T)
        lls = zeros(T)
        B = zeros(ymax + 1, ymax + 1)

        for k = 0:ymax
            B[:, k + 1] = pdf.(Binomial(k, β[1]), 0:ymax)
        end

        a[1:y[M] + 1, M] = PR[1:y[M] + 1, M]./sum(PR[1:y[M] + 1, M])
        lls[M] = sum(a[:, M])
        ϕmat[:, M] = a[:, M]./lls[M]
        for t = M+1:T
            for k = 0:y[t]
                temp = PARZ[1:y[t] - k + 1, t]
                for l = 0:y[t-1]
                    QQ[k + 1, l+1, t] = dot(temp, B[y[t] - k + 1:-1:1, l+1])
                end
            end

            u[:, t] = diagm(PR[:, t])*QQ[:, :, t]*ϕmat[:, t-1]
            lls[t] = sum(u[:, t])
            ϕmat[:, t] = u[:, t]./lls[t]
        end

        lls[M:end] = log.(lls[M:end])
        LL = sum(lls[M+1:end])

        return LL, lls
    end

    if Q == 2
        b = zeros(T, ymax + 1, ymax + 1, ymax + 1)
        w = zeros(T)

        B1 = zeros(ymax + 1, ymax + 1)
        B2 = zeros(ymax + 1, ymax + 1)

        for i = 0:ymax
            B1[:, i + 1] = pdf.(Binomial(i, β[1]), 0:ymax)
            B2[:, i + 1] = pdf.(Binomial(i, β[2]), 0:ymax)
        end
        B = zeros(ymax + 1, ymax + 1, ymax + 1)
        for i = 0:ymax, j = 0:ymax
            B[i + 1, j + 1, :] = convolution(B1[i + 1, :], B2[j + 1, :])
        end

        for r1 = 0:y[M]
            temp1 = convolution(B1[:, r1 + 1], PARZ[:, M+1])
            trunc1 = PR[r1 + 1, M]/sum(PR[1:y[M]+1, M])
            for r2 = 0:y[M-1]
                temp = convolution(temp1, B2[:, r2 + 1])
                trunc2 = PR[r2 + 1, M - 1]/sum(PR[1:y[M-1]+1, M - 1])
                b[M+1, 1:y[M+1]+1, r1 + 1, r2 + 1] = PR[1:y[M+1]+1, M+1] .* trunc1 .* trunc2 .*temp[y[M+1]+1:-1:1]
            end
        end

        w[M+1] = sum(b[M+1, :, :, :])
        b[M+1, :, :, :] = b[M+1, :, :, :]./w[M+1]

        for t = M+2:T
            @inbounds for r1 = 0:ymax
                temp1 = convolution(B1[:, r1 + 1], PARZ[:, t])

                for r2 = 0:ymax
                    temp = convolution(temp1, B2[:, r2 + 1], y[t])
                    b[t, 1:y[t] + 1, r1 + 1, r2 + 1] = PR[1:y[t]+1, t].*sum(b[t-1, r1 + 1, r2 + 1, 1:y[t-3] + 1]).*temp[y[t]+1:-1:1]
                end
            end
            w[t] = sum(b[t, :, :, :])
            b[t, :, :, :] = b[t, :, :, :]./w[t]
        end

        lls = [fill(0, M); log.(w[M+1:end])]
        LL = sum(lls)
        bFinal = b[T, :, :, :]
    end

    if Q >= 3
        error("Exact likelihood evaluation for INARMA(p, q) with q ≥ 3 not supported.")
    end

    return LL, lls
end

function ll(y::Array{T, 1} where T<:Integer,
    model::T where T<:INGARCH,
    θ::Array{T, 1} where T<: AbstractFloat;
    initiate = "first")::Tuple{Float64, Vector{Float64}}

    θNew = θ2par(θ, model)
    return ll(y, model, θNew, initiate = initiate)
end

function ll(y::Array{T, 1} where T<:Integer,
    model::T where T<:INARMA,
    θ::Array{T, 1} where T<: AbstractFloat)::Tuple{Float64, Vector{Float64}}

    θNew = θ2par(θ, model)
    return ll(y, model, θNew)
end

####################################
### Ab hier wird alles schneller ###
####################################

function Pmat(λ::Vector{Float64},
    μ::Vector{Float64},
    ymax::Int64,
    T::Int64,
    nb1::Bool,
    nb2::Bool,
    ϕ::Vector{Float64},
    zi::Bool,
    ω::Float64,
    rE::Int64)::Tuple{Matrix{Float64}, Matrix{Float64}}

    PR = zeros(ymax + 1, T)
    PZ = zeros(ymax + 1, T)

    if nb1 & !zi
        for t = 1:T
            PR[:, t] = pdf.(NegativeBinomial(ϕ[1], ϕ[1]/(ϕ[1] + λ[t])), 0:ymax)
        end
    end
    if !nb1 & !zi
        for t = 1:T
            PR[:, t] = pdf.(Poisson(λ[t]), 0:ymax)
        end
    end

    if nb1 & zi
        PRtemp = zeros(T)
        for t = 1:T
            PRtemp = pdf.(NegativeBinomial(ϕ[1], ϕ[1]/(ϕ[1] + λ[t])), 0:ymax)
            if ω + PRtemp[1] < 1
                PR[2:ymax + 1, t] = PRtemp[2:ymax + 1].*(1 .- ω./(1 .- PRtemp[1]))
                PR[1, t] = PRtemp[1] + ω
            else
                PR[2:ymax + 1, t] .= 0.0
                PR[1, t] = 1.0
            end

        end
    end

    if !nb1 & zi
        PRtemp = zeros(T)
        for t = 1:T
            PRtemp = pdf.(Poisson(λ[t]), 0:ymax)
            if ω + PRtemp[1] < 1
                PR[2:ymax + 1, t] = PRtemp[2:ymax + 1].*(1 .- ω./(1 .- PRtemp[1]))
                PR[1, t] = PRtemp[1] + ω
            else
                PR[2:ymax + 1, t] .= 0.0
                PR[1, t] = 1.0
            end
        end
    end

    if nb2 & (rE > 0)
        for t = 1:T
            PZ[:, t] = pdf.(NegativeBinomial(ϕ[2], ϕ[2]/(ϕ[2] + μ[t])), 0:ymax)
        end
    end

    if !nb2 & (rE > 0)
        for t = 1:T
            PZ[:, t] = pdf.(Poisson(μ[t]), 0:ymax)
        end
    end

    if rE == 0
        PZ[1, :] .= 1.0
    end

    return PR, PZ
end

function Q1funInner(y::Vector{Int64},
                    t::Int64,
                    ymax::Int64,
                    PARZt::Vector{Float64},
                    PRt::Vector{Float64},
                    B::Matrix{Float64},
                    ϕmatOld::Vector{Float64})::Vector{Float64}
    #
    QQ = zeros(ymax + 1, ymax + 1)
    @inbounds for k = 0:y[t]
        temp = PARZt[y[t] - k + 1:-1:1]
        QQ[k+1, 1:y[t-1]+1] = temp'B[1:y[t] - k + 1, 1:y[t-1]+1]
    end

    # return diagm(PRt)*QQ*ϕmatOld # u
    return PRt.*(QQ*ϕmatOld) # u
end

function Q1fun(y::Vector{Int64},
               ymax::Int64,
               T::Int64,
               β::Vector{Float64},
               PR::Matrix{Float64},
               M::Int64,
               PARZ::Matrix{Float64})::Tuple{Float64, Vector{Float64}}
    a = zeros(ymax + 1, T)
    lls = zeros(T)
    B = zeros(ymax + 1, ymax + 1)

    for k = 0:ymax
        B[:, k + 1] = pdf.(Binomial(k, β[1]), 0:ymax)
    end

    a[1:y[M] + 1, M] = PR[1:y[M] + 1, M]./sum(PR[1:y[M] + 1, M])
    lls[M] = sum(a[:, M])
    ϕmatOld = a[:, M]./lls[M]

    @inbounds for t = M+1:T
        u = Q1funInner(y, t, ymax, PARZ[:, t], PR[:, t], B, ϕmatOld)
        lls[t] = sum(u)
        ϕmatOld = u./lls[t]
    end

    lls[M:end] = log.(lls[M:end])
    LL = sum(lls[M+1:end])

    return LL, lls
end

function createB(ymax::Int64, β::Vector{Float64})::Tuple{Matrix{Float64}, Matrix{Float64}}
    B1 = zeros(ymax + 1, ymax + 1)
    B2 = zeros(ymax + 1, ymax + 1)

    for i = 0:ymax
        B1[:, i + 1] = pdf.(Binomial(i, β[1]), 0:ymax)
        B2[:, i + 1] = pdf.(Binomial(i, β[2]), 0:ymax)
    end
    B1, B2
end

function Q2funDeeper(t::Int64,
    y::Vector{Int64},
    ymax::Int64,
    PARZt::Vector{Float64},
    PRt::Vector{Float64},
    B1::Matrix{Float64},
    B2::Matrix{Float64},
    bOldSum::Matrix{Float64})::Tuple{Float64, Array{Float64, 2}}

    bNew = zeros(y[t-2] + 1, ymax + 1, y[t] + 1)

    bOldSumRel = bOldSum[1:y[t-2] + 1, :]
    B2rel = B2[:, 1:y[t-2]+1]
    @simd for r1 = 0:y[t-1]
        # Convolution of β_1 ∘ R_{t-1}, AR part and regressors Z_t
        temp1 = convolution(B1[:, r1 + 1], PARZt)
        tempMat = convolutionRevScale(temp1, B2rel, y[t], PRt)
        bNew[:, r1 + 1, :] = bOldSumRel[:, r1 + 1].*tempMat
    end

    wt = sum(bNew)

    return wt, sum(bNew ./ wt, dims = 1)[1, :, :]
end

function Q2funInner(y::Vector{Int64},
                    ymax::Int64,
                    M::Int64,
                    T::Int64,
                    B1::Matrix{Float64},
                    B2::Matrix{Float64},
                    PARZ::Matrix{Float64},
                    PR::Matrix{Float64})::Vector{Float64}
    bOld = zeros(ymax + 1, ymax + 1, ymax + 1)
    w = zeros(T)

    for r1 = 0:y[M]
        temp1 = convolution(B1[:, r1 + 1], PARZ[:, M+1])
        trunc1 = PR[r1 + 1, M]/sum(PR[1:y[M]+1, M])
        for r2 = 0:y[M-1]
            temp = convolution(temp1, B2[:, r2 + 1])
            trunc2 = PR[r2 + 1, M - 1]/sum(PR[1:y[M-1]+1, M - 1])
            bOld[r2 + 1, r1 + 1, 1:y[M+1]+1] = PR[1:y[M+1]+1, M+1] .* trunc1 .* trunc2 .*temp[y[M+1]+1:-1:1]
        end
    end

    w[M+1] = sum(bOld)
    bOld = bOld./w[M+1]
    bOldSum = sum(bOld[1:y[M-1] + 1, :, :], dims = 1)[1, :, :]

    @inbounds for t = M+2:T
        PARZt = PARZ[:, t]
        PRt = PR[:, t]
        w[t], bOldSum = Q2funDeeper(t, y, ymax, PARZt, PRt[1:y[t] + 1], B1, B2, bOldSum)
    end

    return w
end

function Q2fun(y::Vector{Int64},
               ymax::Int64,
               T::Int64,
               β::Vector{Float64},
               PARZ::Matrix{Float64},
               PR::Matrix{Float64},
               M::Int64)::Tuple{Float64, Vector{Float64}}

    B1, B2 = createB(ymax, β)

    w = Q2funInner(y, ymax, M, T, B1, B2, PARZ, PR)

    lls = [fill(0, M); log.(w[M+1:end])]
    LL = sum(lls)
    return LL, lls
end

function ll(y::Array{T, 1} where T<:Integer,
    model::INARMAModel,
    θ::parameter)::Tuple{Float64, Vector{Float64}}

    T = length(y)
    ymax = maximum(y)
    X = model.X
    r = Int64(length(model.external))

    logl1 = model.link[1] == "Log"
    lin1 = !logl1

    logl2 = model.link[2] == "Log"
    lin2 = !logl2

    nb1 = model.distr[1] == "NegativeBinomial"
    nb2 = model.distr[2] == "NegativeBinomial"
    if r == 0
        nb2 = false
    else
        if sum(model.external) == 0
            nb2 = false
        end
    end

    zi = model.zi

    p = length(model.pastObs)
    if p == 0
        P = 0
    else
        P = Int64(maximum(model.pastObs))
    end

    q = Int64(length(model.pastMean))
    if q == 0
        Q = 0
    else
        Q = Int64(maximum(model.pastMean))
    end
    M = maximum([P, Q])

    if !parametercheck(θ, model)
        return -Inf, fill(-Inf, T)
    end

    β0 = Float64(θ.β0)
    αRaw = Vector{Float64}(θ.α)
    βRaw = Vector{Float64}(θ.β)
    η = Vector{Float64}(θ.η)
    ϕ = Vector{Float64}(θ.ϕ)
    ω = Float64(θ.ω)

    if p > 0
        if p < P
            α = zeros(P)
            α[model.pastObs] = αRaw
        else
            α = αRaw
        end
    end

    if q > 0
        if q < Q
            β = zeros(Q)
            β[model.pastMean] = βRaw
        else
            β = βRaw
        end
    end

    rE = sum(model.external)
    rI = r - rE

    if r > 0
        iE = findall(model.external)
        iI = findall(.!model.external)

        ηI = η[iI]
        ηE = η[iE]

        λ, μ = createλμ(T, β0, rI, iI, η, X, rE, iE, logl1, logl2)
    else
        λ, μ = createλμ(T, β0, logl1)
    end

    PR, PZ = Pmat(λ, μ, ymax, T, nb1, nb2, ϕ, zi, ω, rE)

    lls = zeros(T)

    # Below only for INAR Process, rather irrelevant
    if q == 0
        PY = zeros(p, ymax + 1, ymax + 1)
        for i = 1:p
            for j = sort(unique(y))
                PY[i, j + 1, 1:j + 1] = pdf.(Binomial.(j, α[i]), 0:j)
            end
        end

        PAR = zeros(ymax + 1, P + 1)

        for t = M+1:T
            for i = 1:P
                PAR[:, i] = PY[i, y[t-i] + 1, :]
            end
            PAR[:, P + 1] = PZ[:, t]

            ptemp = GetProbFromP(PAR)
            lls[t] = log(dot(PR[1:y[t]+1], ptemp[y[t]+1:-1:1]))
        end
        LL = sum(lls[M+1:end])

        return LL, lls
    end

    # Compute P(Σ α_i∘Y_{t-i} + Z_t = k) first --> pARZ
    PAR = zeros(ymax + 1, P + 1)
    PARZ = zeros(ymax + 1, T)
    for t = M+1:T
        for i = 1:P
            PAR[:, i] = pdf.(Binomial(y[t - i], α[i]), 0:ymax)
        end
        PAR[:, end] = PZ[:, t]
        PARZ[:, t] = GetProbFromP(PAR)
    end

    if Q == 1
        return Q1fun(y, ymax, T, β, PR, M, PARZ)
    end

    if Q == 2
        return Q2fun(y, ymax, T, β, PARZ, PR, M)
    end

    if Q >= 3
        error("Exact likelihood evaluation for INARMA(p, q) with q ≥ 3 not supported.")
    end

    return LL, lls
end

function CoreINAR(y::Vector{Int64},
                  ymax::Int64,
                  P::Int64,
                  M::Int64,
                  T::Int64,
                  PY::Array{Float64, 3},
                  PZ::Matrix{Float64},
                  PR::Matrix{Float64})::Vector{Float64}
    #
    lls = zeros(T)
    PAR = zeros(ymax + 1, P + 1)
    @inbounds for t = M+1:T
        out = 0.0
        for i = 1:P
            PAR[:, i] = PY[i, y[t-i] + 1, :]
        end
        PAR[:, P + 1] = PZ[:, t]

        ptemp = GetProbFromP(PAR)
        for i = 0:y[t]
            out += PR[i+1]*ptemp[y[t] - i + 1]
        end
        lls[t] = log(out)
    end
    lls
end

function CreatePY(y::Vector{Int64},
    ymax::Int64,
    p::Int64,
    α::Vector{Float64})::Array{Float64, 3}

    PY = zeros(p, ymax + 1, ymax + 1)
    allVals = sort(unique(y))
    @inbounds for i = 1:p
        for j = allVals
            PY[i, j + 1, 1:j + 1] = pdf.(Binomial.(j, α[i]), 0:j)
        end
    end
    PY
end

function createλμ(T::Int64,
    β0::Float64,
    rI::Int64,
    iI::Vector{Int64},
    η::Vector{Float64},
    X::Matrix{Float64},
    rE::Int64,
    iE::Vector{Int64},
    logl1::Bool,
    logl2::Bool)::Tuple{Vector{Float64}, Vector{Float64}}

    μ = zeros(T)
    λ = fill(β0, T)

    if rI > 0
        λ .+= (η[iI]'X[iI, :])[1, :]
    end

    if logl1
        λ = exp.(λ)
    end

    if rE > 0
        μ .+= (η[iE]'X[iE, :])[1, :]

        if logl2
            μ = exp.(μ)
        end
    end
    λ, μ
end

function createλμ(T::Int64,
    β0::Float64,
    logl1::Bool)::Tuple{Vector{Float64}, Vector{Float64}}

    if logl1
        λ = fill(exp(β0), T)
    else
        λ = fill(β0, T)
    end

    λ, zeros(T)
end

function ll(y::Array{T, 1} where T<:Integer,
    model::INARModel,
    θ::parameter)::Tuple{Float64, Vector{Float64}}

    T = length(y)
    ymax = maximum(y)
    X = model.X

    r = length(model.external)

    logl1 = model.link[1] == "Log"
    lin1 = !logl1

    logl2 = model.link[2] == "Log"
    lin2 = !logl2

    nb1 = model.distr[1] == "NegativeBinomial"
    nb2 = model.distr[2] == "NegativeBinomial"
    if r == 0
        nb2 = false
    else
        if sum(model.external) == 0
            nb2 = false
        end
    end

    zi = model.zi

    p = length(model.pastObs)
    if p == 0
        P = 0
    else
        P = maximum(model.pastObs)
    end

    M = P

    if !parametercheck(θ, model)
        return -Inf, fill(-Inf, T)
    end

    β0 = θ.β0
    αRaw = θ.α
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    if p > 0
        if p < P
            α = zeros(P)
            α[model.pastObs] = αRaw
        else
            α = αRaw
        end
    end

    rE = sum(model.external)
    rI = sum(.!model.external)

    if r > 0
        iE = findall(model.external)
        iI = findall(.!model.external)

        ηI = η[iI]
        ηE = η[iE]

        λ, μ = createλμ(T, β0, rI, iI, η, X, rE, iE, logl1, logl2)
    else
        λ, μ = createλμ(T, β0, logl1)
    end
    PR, PZ = Pmat(λ, μ, ymax, T, nb1, nb2, ϕ, zi, ω, rE)
    PY = CreatePY(y, ymax, p, α)
    lls = CoreINAR(y, ymax, P, M, T, PY, PZ, PR)

    LL = sum(lls[M+1:end])

    return LL, lls
end


text:
using StagedFilters
using Test

using PyCall, BenchmarkTools;

N = 1_000_000

data = convert.(Float64,collect(range(1,N,length=N)));
smoothed = zeros(eltype(data),length(data)); # <--- wholesome, type stable code.
savgol = pyimport("scipy.signal")."savgol_filter";
x = PyObject(data);

@info "Julia f32x$N"
@btime smooth!(SavitzkyGolayFilter{2,2}, $data, $smoothed);
@info "SciPy f32x$N"
@btime pycall($savgol, PyObject, $x,5,2,mode="wrap");

data = convert.(Float32,collect(range(1,N,length=N)));
smoothed = zeros(eltype(data),length(data)); # <--- wholesome, type stable$
savgol = pyimport("scipy.signal")."savgol_filter";
x = PyObject(data);

@info "Julia f64x$N"
@btime smooth!(SavitzkyGolayFilter{2,2}, $data, $smoothed);
@info "SciPy f64x$N"
@btime pycall($savgol, PyObject, $x,5,2,mode="wrap");

"""
In [17]: import scipy.signal

In [18]: data = [i / 1.0 for i in range(1000000)]

In [19]: f = scipy.signal.savgol_filter

In [20]: %%timeit
    ...: f(data, 5, 2, mode="wrap")
        ...:
            ...:
            58 ms ± 257 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
"""

ext: jl


text:
"""
    Trues = Ones{Bool, N, Axes} where {N, Axes}

Lazy version of `trues` with axes.
Typically created using `Trues(dims)` or `Trues(dims...)`

# Example
```jldoctest
julia> Trues(1,3)
1×3 Ones{Bool,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}} = true

julia> Trues((2,3))
2×3 Ones{Bool,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}} = true
```
"""
const Trues = Ones{Bool, N, Axes} where {N, Axes}

""" `Falses = Zeros{Bool, N, Axes}` (see `Trues`) """
const Falses = Zeros{Bool, N, Axes} where {N, Axes}

# y[mask] = x when mask isa Trues (cf y[:] = x)
# Supported here only for arrays with standard OneTo axes.
function Base.setindex!(y::AbstractArray{T,N}, x, mask::Trues{N, NTuple{N,Base.OneTo{Int}}}) where {T,N}
    if axes(x) isa NTuple{N,Base.OneTo{Int}} && axes(y) isa NTuple{N,Base.OneTo{Int}}
        @boundscheck size(y) == size(mask) || throw(BoundsError(y, mask))
        @boundscheck size(x) == size(mask) || throw(DimensionMismatch(
            "tried to assign $(length(x)) elements to $(length(y)) destinations"))
        @boundscheck checkbounds(y, mask)
        return copyto!(y, x)
    end
    return setindex!(y, x, trues(size(mask))) # fall back on usual setindex!
end

# x[mask] when mask isa Trues (cf x[trues(size(x))] or x[:])
# Supported here only for arrays with standard OneTo axes.
function Base.getindex(x::AbstractArray{T,N}, mask::Trues{N, NTuple{N,Base.OneTo{Int}}}) where {T,N}
    if axes(x) isa NTuple{N,Base.OneTo{Int}}
       @boundscheck size(x) == size(mask) || throw(BoundsError(x, mask))
       return vec(x)
    end
    return x[trues(size(x))] # else revert to usual getindex method
end


text:
module SomePackage

greet() = print("Hello World!")

end # module


text:
# ------------------------------------------------------------------
# Code to simulate logit data and estimate nonparametrically
# Plots estimates of own and cross-price elasticities
# Written by James Brand
# ------------------------------------------------------------------
using Statistics, NPDemand
using RCall, DataFrames
@rlibrary ggplot2

J=2; # number of products
T =5000;
beta = -0.4; # price coefficient in utility function
sdxi = 0.15; # standard deviation of xi

S = 50; # number of simulations
G = 10; # size of grid on which to evaluate price elasticities
elast_own = zeros(S,G);
elast_cross = zeros(S,G);
elast_own_dist = zeros(S,T);
elastTrue = zeros(S,G);

s, pt, zt, xi = simulate_logit(J,T, beta, sdxi);

p_points = range(quantile(pt[:,1],.25),stop = quantile(pt[:,1],.75),length = G);
p_points = convert.(Float64, p_points)

# ------------------------------------------------------
# Set options for estimation and elasticity calculation
# ------------------------------------------------------
bernO = 2*ones(2J,1);        # Order of Bernstein Polynomial
iv=0;                       # Order of IV Polynomial = (bernO + iv)
trueS = false;                    # Evaluate at true market shares or not
own = [1,1];                # [derivative of j, with respect to price of k]

nfolds = 5; # number of cross-validation folds
nlam = 10; # number of regularization parameters to try. Actual values chosen automatically by hierNet
strong = true; # boolean for whether or not to impose strong hierarchy constraint
# Note: "strong = true" takes much longer than "strong = false."
nboot = 1; # number of bootstrapped samples to run hierarchical lasso on

# ------------------------------------------------------
# Simulation
# ------------------------------------------------------
included_symmetric_pct = zeros(2J,2J)
included_pct = zeros(2J,2J)
for si = 1:1:S
    # Simulate demand in two groups -- each product only substitutes to J-1 others
    s, pt, zt = simulate_logit(J,T, beta, sdxi);
    s2, pt2, zt2  = simulate_logit(J,T, beta, sdxi);

    s = [s s2];
    s = s ./ 2;
    pt = [pt pt2];
    zt = [zt zt2];
    df = toDataFrame(s,pt,zt);

    # hierNet() Returns two matrices: one, the "raw" selected model, and another
    #   which imposes symmetry. I.e. if j is a substute for k, then k is
    #   a substitute for j as well (this can drastically increase the # parameters
    #    to estimate when s has many columns)
    included, included_symmetric = hierNet_boot(df; nfolds = nfolds, nlam = 10, strong = false, nboot = 1);

    # Estimate demand nonparametrically
        # If you want to include an additional covariate in all demand
        # functions, add an additional argument "marketvars" after included. If it is an
        # additional product characteristic, marketvars should be T x J
    inv_sigma, designs = inverse_demand(df; included = included_symmetric);

    # Calculate price elasticities
    prices = median(pt).*ones(G,2J);
    prices[:,1] = p_points;
    elast, Jacobians, share_vec = price_elasticity(inv_sigma, df, prices; whichProducts = own,
        included = included_symmetric, trueS = trueS);

    trueelast = beta.*p_points.*(1 .- 2 .* share_vec[:,1])

    elast_own[si,:] = elast;
    elastTrue[si,:] = trueelast;
    included_pct[:,:] += included./S; # summarizing selection patterns. See pure_model_selection.jl
    included_symmetric_pct[:,:] += included_symmetric./S;
end

elast025 = zeros(G,1)
elast50 = zeros(G,1)
elast975 = zeros(G,1)
for i = 1:G
    elast025[i] = quantile(elast_own[:,i], 0.1)
    elast50[i] = quantile(elast_own[:,i], 0.5)
    elast975[i] = quantile(elast_own[:,i], 0.9)
end
elast025 = dropdims(elast025,dims = 2)
elast50 = dropdims(elast50,dims = 2)
elast975 = dropdims(elast975,dims = 2)

df = DataFrame(p = p_points, e025 = elast025, e50 = elast50, e975 = elast975)
ggplot(df, aes(x=:p, y=:e50)) + geom_line() + geom_line(aes(y=:e975), color = "gray", linetype = "dashed") +
    geom_line(aes(y=:e025), color = "gray", linetype = "dashed") +
    xlab("Price") + ylab("Own-Elasticity") + theme_light()


text:
include("gatherApi.jl")
include("buildApi.jl")


text:
"""
    PaduaTransforms

an implementation of the Padua transform and its inverse via the fast Fourier transform.
"""
module PaduaTransforms

using FFTW
using LinearAlgebra: rmul!

export getpaduanum, getdegree, nextdegree, nextpaduanum
export getpaduapoints
export PaduaTransformPlan, paduatransform!
export InvPaduaTransformPlan, invpaduatransform!

## Number of Padua Points and Degree ##

"""
    getpaduanum(n)

calculates number of Padua points needed to approximate a function using Chebyshev
polynomials up to total degree `n`. This number is equal to the number of coefficients.
The formula is

```math
N = (n + 1) ⋅ (n + 2) ÷ 2
```

# Examples
```jldoctest
julia> getpaduanum(13)
105
```
"""
getpaduanum(degree) = (degree + 1) * (degree + 2) ÷ 2

"""
    getdegree(N)

calculates total degree, given the number of coefficients or Padua points `N`.
Throws an error if `N` is not a possible number of Padua points. The formula is

```math
n = \\frac{\\sqrt{1 + 8N} - 3}{2}
```

# Examples
```jldoctest
julia> getdegree(105)
13

julia> getdegree(104)
ERROR: ArgumentError: number of Padua points or coeffs must be (n + 1) * (n + 2) ÷ 2
[...]
```
"""
function getdegree(N)
    d = (sqrt(1 + 8N) - 3) / 2
    isinteger(d) ? Int(d) : throw(ArgumentError(
        "number of Padua points or coeffs must be (n + 1) * (n + 2) ÷ 2"))
end

"""
    nextdegree(N)

get degree of `nextpaduanum(N)`.

# Examples

```jldoctest
julia> nextdegree(104)
13
```
"""
nextdegree(N) = Int(cld(sqrt(1 + 8N) - 3, 2))

"""
    nextpaduanum(N)

get next valid number of Padua points ≥ `N`.

# Examples

```jldoctest
julia> nextpaduanum(104)
105
```
"""
nextpaduanum(N) = getpaduanum(nextdegree(N))

## Padua Points ##

"""
    paduapoint(T::Type, j::Integer, i::Integer, n::Integer)

returns the Padua point ``z_{ij}``, where

```math
z_{ij} = (\\cos{\\frac{jπ}{n}}, \\cos{\\frac{iπ}{n+1}})
```

Note, that only points with ``i-j`` even are actually Padua points.
Check with [`ispadua`](@ref).

# Examples
```jldoctest
julia> [PaduaTransforms.paduapoint(Float32, x, y, 1) for y in 0:1+1, x in 0:1]
3×2 Matrix{Tuple{Float32, Float32}}:
 (1.0, 1.0)   (-1.0, 1.0)
 (1.0, 0.0)   (-1.0, 0.0)
 (1.0, -1.0)  (-1.0, -1.0)
```

"""
function paduapoint(::Type{T}, j, i, n) where T
    x = cospi(T(j) / T(n))
    y = cospi(T(i) / T(n + 1))
    return x, y
end

"""
    ispadua(i, j)

returns if [`paduapoint`](@ref) at position `(i, j)` is a Padua point.

# Examples
```jldoctest
julia> pointornothing(i, j, n) = PaduaTransforms.ispadua(i, j) ? PaduaTransforms.paduapoint(Float64, j, i, n) : nothing
pointornothing (generic function with 1 method)

julia> [pointornothing(y, x, 2) for y in 0:3, x in 0:2]
4×3 Matrix{Union{Nothing, Tuple{Float64, Float64}}}:
 (1.0, 1.0)   nothing      (-1.0, 1.0)
 nothing      (0.0, 0.5)   nothing
 (1.0, -0.5)  nothing      (-1.0, -0.5)
 nothing      (0.0, -1.0)  nothing
```

"""
ispadua(i, j) = iseven(i - j)

"""
    getpaduapoints([f::Function,], [T=Float64,] n)

evaluates the function `f` on the Padua points (of type `T`) for degree `n`.

If no function `f` is provided, `getpaduapoints` returns a matrix containing
the x and y components of the paduapoints as columns. If `f` returns a single value,
`getpaduapoints` returns a `Vector{T}`. If `f` returns a tuple or other iterable,
`getpaduapoints` returns a `Matrix{T}` where the i-th column represents the i-th
entry in `f` applied to all the Padua points.

# Examples
```jldoctest
julia> getpaduapoints(2)
6×2 Matrix{Float64}:
  1.0   1.0
  1.0  -0.5
  0.0   0.5
  0.0  -1.0
 -1.0   1.0
 -1.0  -0.5

julia> getpaduapoints(Float32, 1)
3×2 Matrix{Float32}:
  1.0   1.0
  1.0  -1.0
 -1.0   0.0

julia> getpaduapoints(Float32, 2) do x, y; x*y; end
6-element Vector{Float32}:
  1.0
 -0.50000006
  0.0
 -0.0
 -1.0
  0.50000006

julia> getpaduapoints(2) do x, y; x*y, 5*x*y; end
6×2 Matrix{Float64}:
  1.0   5.0
 -0.5  -2.5
  0.0   0.0
 -0.0  -0.0
 -1.0  -5.0
  0.5   2.5

```
"""
function getpaduapoints(f::Function, ::Type{T}, n) where T
    D = length(f(zero(T), zero(T)))
    out = Matrix{T}(undef, getpaduanum(n), D)

    i = 1
    for x in 0:n
        for y in 0:n+1
            if ispadua(x, y)
                px, py = paduapoint(T, x, y, n)
                out[i, :] .= f(px, py)
                i += 1
            end
        end
    end

    # Should return vector instead of matrix in 1D case
    return D == 1 ? vec(out) : out
end

getpaduapoints(f::Function, n) = getpaduapoints(f, Float64, n)
getpaduapoints(T, n) = getpaduapoints((x, y) -> (x, y), T, n)
getpaduapoints(n) = getpaduapoints(Float64, n)

## Padua Transform ##

struct PaduaTransformPlan{T, P}
    degree::Int
    vals::Matrix{T}
    dctplan::P
end

"""
    PaduaTransformPlan{T}(n::Integer)

create plan to compute coefficients of Chebyshev polynomials in 2D up to total degree `n`
using the Padua transform.

"""
function PaduaTransformPlan{T}(degree::Integer) where T
    vals = Matrix{T}(undef, degree + 2, degree + 1)
    plan = FFTW.plan_r2r!(vals, FFTW.REDFT00)

    PaduaTransformPlan{T, typeof(plan)}(degree, vals, plan)
end

"""
    weight!(mat::AbstractMatrix, degree::Integer)

weight fourier coefficients to obtain Chebyshev coefficients as part of a [`paduatransform!`](@ref).
The weighting factor applied to the coefficients is

```math
w = \\frac{1}{n(n+1)} ⋅ \\begin{cases}
    \\frac{1}{2} & \\textrm{if on vertex}   \\\\
    1           & \\textrm{if on edge}    \\\\
    2           & \\textrm{if in interior} \\\\
\\end{cases}
```

# Examples
```julia-repl
julia> weight!(ones(4+2, 4+1), 4)
6×5 Matrix{Float64}:
 0.025  0.05  0.05  0.05  0.025
 0.05   0.1   0.1   0.1   0.05
 0.05   0.1   0.1   0.1   0.05
 0.05   0.1   0.1   0.1   0.05
 0.05   0.1   0.1   0.1   0.05
 0.025  0.05  0.05  0.05  0.025
```
"""
function weight!(mat::AbstractMatrix{T}, degree::Integer) where T
    rmul!(mat, T(2) / T(degree * (degree + 1)))
    rmul!(@view(mat[1, :]), T(0.5))
    rmul!(@view(mat[end, :]), T(0.5))
    rmul!(@view(mat[:, 1]), T(0.5))
    rmul!(@view(mat[:, end]), T(0.5))

    mat
end

"""
    tovalsmat!(mat::Matrix, from::AbstractVector, degree::Integer)

write values of function evaluated at Padua points from `from` to matrix `mat`.

# Examples
```jldoctest
julia> PaduaTransforms.tovalsmat!(ones(3 + 2, 3 + 1), 1:getpaduanum(3), 3)
5×4 Matrix{Float64}:
 1.0  0.0  6.0   0.0
 0.0  4.0  0.0   9.0
 2.0  0.0  7.0   0.0
 0.0  5.0  0.0  10.0
 3.0  0.0  8.0   0.0

julia> PaduaTransforms.tovalsmat!(ones(2 + 2, 2 + 1), 1:getpaduanum(2), 2)
4×3 Matrix{Float64}:
 1.0  0.0  5.0
 0.0  3.0  0.0
 2.0  0.0  6.0
 0.0  4.0  0.0
```
"""
function tovalsmat!(mat::Matrix{T}, from::AbstractVector, degree::Integer) where T
    axes(from, 1) == 1:getpaduanum(degree) || throw(ArgumentError(
        "axes of from must be (1:getpaduanum(degree),)"))
    size(mat) == (degree + 2, degree + 1) || throw(ArgumentError(
        "axes of vals mat must be (degree + 2, degree + 1)"))

    if isodd(degree)
        # x 0
        # 0 x
        # x 0

        @inbounds for i in 1:length(from)
            mat[2i - 1] = from[i]
            mat[2i] = zero(T)
        end
    else
        @assert iseven(degree)
        # x 0 x
        # 0 x 0
        # x 0 x
        # 0 x 0

        valspercol = (degree + 2) ÷ 2

        # odd columns (j is column index)
        for j in 1:2:degree + 1, i in 1:valspercol
            k = (j - 1) * valspercol + i
            @inbounds mat[2i - 1, j] = from[k]
            @inbounds mat[2i, j] = zero(T)
        end

        # even columns
        for j in 2:2:degree + 1, i in 1:valspercol
            k = (j - 1) * valspercol + i
            @inbounds mat[2i - 1, j] = zero(T)
            @inbounds mat[2i, j] = from[k]
        end
    end

    mat
end

"""
    fromcoeffsmat!(to::AbstractVector, mat::Matrix, degree::Integer, ::Val{lex})

write Chebyshev coefficients from `mat` into vector `to`. `lex::Bool` determines whether
coefficients should be written in lexigraphical order or not. The lower right triangle does
not get written into `to`. These would represent greater polynomial degrees than `degree`.

If `lex` is `Val(true)` the coefficients correspond to the following basis polynomials

```math
T_0(x) T_0(y), T_1(x) T_0(y), T_0(x) T_1(y), T_2(x) T_0(y), T_1(x) T_1(y), T_0(x) T_2(y), ...
```

else if `lex` is `Val(false)` they correspond to

```math
T_0(x) T_0(y), T_0(x) T_1(y), T_1(x) T_0(y), T_0(x) T_2(y), T_1(x) T_1(y), T_2(x) T_0(y), ...
```

# Examples
```julia-repl
julia> mat = [(x, y) for y in 0:2+1, x in 0:2]
4×3 Matrix{Tuple{Int64, Int64}}:
 (0, 0)  (1, 0)  (2, 0)
 (0, 1)  (1, 1)  (2, 1)
 (0, 2)  (1, 2)  (2, 2)
 (0, 3)  (1, 3)  (2, 3)

julia> to1 = similar(mat, getpaduanum(2)); to2 = similar(mat, getpaduanum(2));

julia> fromcoeffsmat!(to1, mat, 2, Val(true))
6-element Vector{Tuple{Int64, Int64}}:
 (0, 0)
 (1, 0)
 (0, 1)
 (2, 0)
 (1, 1)
 (0, 2)

julia> fromcoeffsmat!(to2, mat, 2, Val(false))
6-element Vector{Tuple{Int64, Int64}}:
 (0, 0)
 (0, 1)
 (1, 0)
 (0, 2)
 (1, 1)
 (2, 0)
```
"""
function fromcoeffsmat!(to::AbstractVector, mat::Matrix, degree::Integer, ::Val{false})
    length(to) == getpaduanum(degree) || throw(ArgumentError(
        "length of output vector must equal paduanum(degree)"))
    axes(mat) == (1:(degree + 2), 1:(degree + 1)) || throw(ArgumentError(
        "axes of coeffs mat must be (1:(degree + 2), 1:(degree + 1))"))

    n = firstindex(to)
    for d in 1:degree + 1
        for ix in 1:d
            iy = d - ix + 1
            @assert ix + iy == d + 1 "ix and iy must lie on d-th diagonal"

            to[n] = mat[iy, ix]
            n += 1
        end
    end

    to
end

function fromcoeffsmat!(to::AbstractVector, mat::Matrix, degree::Integer, ::Val{true})
    length(to) == getpaduanum(degree) || throw(ArgumentError(
        "length of output must equal paduanum(degree)"))
    size(mat) == (degree + 2, degree + 1) || throw(ArgumentError(
        "axes of coeffs mat must be (1:(degree + 2), 1:(degree + 1))"))

    n = firstindex(to)
    for d in 1:degree + 1
        for iy in 1:d
            ix = d - iy + 1
            @assert ix + iy == d + 1 "ix and iy must lie on d-th diagonal"

            to[n] = mat[iy, ix]
            n += 1
        end
    end

    to
end

"""
    fromcoeffsmat!(to::AbstractMatrix, mat::Matrix, degree::Integer)

copy Chebyshev coefficients from `mat` to `to` without copying coefficients coresponding to
total degree greater than `degree`.

# Examples
```jldoctest
julia> PaduaTransforms.fromcoeffsmat!(zeros(4, 4), reshape(1:20, 5, 4), 3)
4×4 Matrix{Float64}:
 1.0  6.0  11.0  16.0
 2.0  7.0  12.0   0.0
 3.0  8.0   0.0   0.0
 4.0  0.0   0.0   0.0
```
"""
function fromcoeffsmat!(to::AbstractMatrix, mat::AbstractMatrix, degree::Integer)
    axes(to) == (1:degree + 1, 1:degree + 1) || throw(ArgumentError(
        "axes of output must be (1:(degree + 2), 1:(degree + 1))"))
    axes(mat) == (1:degree + 2, 1:degree + 1) || throw(ArgumentError(
        "axes of coeffs mat must be (1:(degree + 2), 1:(degree + 1))"))

    for j in 1:degree + 1
        for i in 1:degree + 2 - j
            @inbounds to[i, j] = mat[i, j]
        end
    end

    to
end

"""
    paduatransform!(coeffs, P::PaduaTransformPlan, vals[, lex])

obtain coefficients of Chebyshev polynomials on 2D via the Padua transform, given values
`vals` evaluated at the Padua points. Coefficients will be written into `coeffs`, which
should either be a matrix, a vector or an iterable of either.

If a matrix of `vals` and a corresponding iterable of `coeffs` matrixes or
vectors is given, each `vals` column will be transformed seperately in a multidimensional
Padua transform.

`lex` determines the order in which coefficients are written into `out` if `out` is a
vector. See [`fromcoeffsmat!`](@ref) for details.

!!! warning
    if `coeffs` is a matrix, make sure that all entries in the lower right diagonal are zero
    as these will not get overwritten.

# Examples
```jldoctest
julia> plan = PaduaTransformPlan{Float64}(2);

julia> f(x, y) = 3 + 4x + 5 * x*y, 6 + 7y
f (generic function with 1 method)

julia> vals = getpaduapoints(f, 2)
6×2 Matrix{Float64}:
 12.0  13.0
  4.5   2.5
  3.0   9.5
  3.0  -1.0
 -6.0  13.0
  1.5   2.5

julia> paduatransform!(zeros(3, 3), plan, vals[:, 1])
3×3 Matrix{Float64}:
 3.0  4.0  0.0
 0.0  5.0  0.0
 0.0  0.0  0.0

julia> paduatransform!(zeros(6), plan, vals[:, 2], Val(true))
6-element Vector{Float64}:
 6.0
 0.0
 7.0
 0.0
 0.0
 0.0

julia> paduatransform!((zeros(3, 3), zeros(3, 3)), plan, vals)
([3.0 4.0 0.0; 0.0 5.0 0.0; 0.0 0.0 0.0], [6.0 0.0 0.0; 7.0 0.0 0.0; 0.0 0.0 0.0])
```
"""
function paduatransform!(P::PaduaTransformPlan)
    coeffs = P.dctplan * P.vals
    weight!(coeffs, P.degree)

    coeffs
end

function paduatransform!(out, P::PaduaTransformPlan, vals::AbstractVector{<: Number}, args...)
    tovalsmat!(P.vals, vals, P.degree)
    coeffs = paduatransform!(P)
    fromcoeffsmat!(out, coeffs, P.degree, args...)
end

function paduatransform!(coeffs, P::PaduaTransformPlan, vals, args...)
    length(coeffs) == length(vals)|| throw(ArgumentError(
        "the length of coeffs and vals iterables must match"))

    # can't use eltype to check eltype beforehand because eachcol/eachrow are type unstable
    # and I want eachcol/eachrow to be usable as an iterable of vectors
    for (c, v::AbstractVector{<:Number}) in zip(coeffs, vals)
        paduatransform!(c, P, v, args...)
    end

    coeffs
end

function paduatransform!(coeffs, P::PaduaTransformPlan, vals::AbstractMatrix{<:Number}, args...)
    paduatransform!(coeffs, P, eachcol(vals), args...)
end

## Inverse Padua Transform ##

struct InvPaduaTransformPlan{T, P}
    degree::Int
    coeffs::Matrix{T}
    dctplan::P
end

"""
    InvPaduaTransformPlan{T}(n::Integer)

create plan to compute values on Padua points, given coefficients of Chebyshev polynomials
up to total degree `n`.
"""
function InvPaduaTransformPlan{T}(degree::Integer) where T
    coeffs = Matrix{T}(undef, degree + 2, degree + 1)
    iplan = FFTW.plan_r2r!(coeffs, FFTW.REDFT00)

    InvPaduaTransformPlan{T, typeof(iplan)}(degree, coeffs, iplan)
end

"""
    tocoeffsmat!(mat::AbstractMatrix, coeffs::AbstractMatrix)

writes coefficients in `coeffs` into matrix `mat` for the [`invpaduatransform!`](@ref).

# Examples
```jldoctest
julia> PaduaTransforms.tocoeffsmat!(zeros(5, 4), reshape(1:16, 4, 4))
5×4 Matrix{Float64}:
 1.0  5.0   9.0  13.0
 2.0  6.0  10.0  14.0
 3.0  7.0  11.0  15.0
 4.0  8.0  12.0  16.0
 0.0  0.0   0.0   0.0
```
"""
function tocoeffsmat!(mat::AbstractMatrix{T}, coeffs::AbstractMatrix) where T
    mat[1:end-1, :]  = coeffs
    mat[    end, :] .= zero(T)

    mat
end

"""
    invweight!(coeffs::AbstractMatrix)

weight Chebyshev coefficients before the Fourier transform for the [`invpaduatransform!`](@ref).
using the weighting

```math
w = \\begin{cases}
    1            & \\textrm{if on vertex}   \\\\
    \\frac{1}{2} & \\textrm{if on edge}     \\\\
    \\frac{1}{4} & \\textrm{if in interior} \\\\
\\end{cases}
```

# Examples
```jldoctest
julia> PaduaTransforms.invweight!(ones(5, 5))
5×5 Matrix{Float64}:
 1.0  0.5   0.5   0.5   1.0
 0.5  0.25  0.25  0.25  0.5
 0.5  0.25  0.25  0.25  0.5
 0.5  0.25  0.25  0.25  0.5
 1.0  0.5   0.5   0.5   1.0
```
"""
function invweight!(coeffs::AbstractMatrix{T}) where T
    rmul!(@view(coeffs[:,2:end-1]), T(0.5))
    rmul!(@view(coeffs[2:end-1, :]), T(0.5))

    coeffs
end

"""
    fromvalsmat!(to::AbstractVector, mat::AbstractMatrix, n::Integer)

write values from `mat` into the vector `to` after an [`invpaduatransform!`](@ref) of total
degree `n`.
"""
function fromvalsmat!(to::AbstractVector, mat::AbstractMatrix, degree::Integer)
    axes(to, 1) == 1:getpaduanum(degree) || throw(Argumenterror(
        "axes of output must be (1:getpaduanum(degree),)"))
    axes(mat) == (1:degree + 2, 1:degree + 1) || throw(Argumenterror(
        "axes of vals mat must be (1:degree + 2, 1:degree + 1)"))

    if isodd(degree)
        # x 0
        # 0 x
        # x 0

        @inbounds for i in 1:length(to)
            to[i] = mat[2i - 1]
        end
    else
        @assert iseven(degree)
        # x 0 x
        # 0 x 0
        # x 0 x
        # 0 x 0

        valspercol = (degree + 2) ÷ 2

        # odd columns (j is column index)
        for j in 1:2:degree + 1, i in 1:valspercol
            k = (j - 1) * valspercol + i
            @inbounds to[k] = mat[2i - 1, j]
        end

        # even columns
        for j in 2:2:degree + 1, i in 1:valspercol
            k = (j - 1) * valspercol + i
            @inbounds to[k] = mat[2i, j]
        end
    end

    to
end

function invpaduatransform!(IP::InvPaduaTransformPlan)
    invweight!(IP.coeffs)
    IP.dctplan * IP.coeffs

    IP.coeffs
end

"""
    invpaduatransform!(vals::AbstractVector, IP::InvPaduaTransformPlan, coeffs::AbstractMatrix)

evaluates the polynomial defined by the coefficients of Chebyshev polynomials `coeffs` on
the Padua points using the inverse transform plan `IP` and writes the resulting values into
`vals`.

If a `vals` matrix and a corresponding iterable of `coeffs` matrixes is given,
each `coeffs` matrix will be transformed seperately into a column in `vals`.

# Examples
```jldoctest
julia> iplan = InvPaduaTransformPlan{Float64}(2);

julia> coeffs = [3 4 0; 0 5 0; 0 0 0]
3×3 Matrix{Int64}:
 3  4  0
 0  5  0
 0  0  0

julia> invpaduatransform!(zeros(6), iplan, coeffs)
6-element Vector{Float64}:
 12.0
  4.5
  3.0
  3.0
 -6.0
  1.5

julia> getpaduapoints(2) do x, y; 3 + 4x + 5 * x*y; end
6-element Vector{Float64}:
 12.0
  4.5
  3.0
  3.0
 -6.0
  1.4999999999999996
```
"""
function invpaduatransform!(vals::AbstractVector{<:Number}, IP::InvPaduaTransformPlan, coeffs)
    tocoeffsmat!(IP.coeffs, coeffs)
    invpaduatransform!(IP)
    fromvalsmat!(vals, IP.coeffs, IP.degree)
end

function invpaduatransform!(vals, IP::InvPaduaTransformPlan, coeffs)
    length(coeffs) == length(vals)|| throw(ArgumentError(
        "the length of coeffs and vals iterables must match"))

    for (v::AbstractVector{<:Number}, c) in zip(vals, coeffs)
        invpaduatransform!(v, IP, c)
    end

    vals
end

function invpaduatransform!(vals::AbstractMatrix{<:Number}, IP::InvPaduaTransformPlan, coeffs, args...)
    invpaduatransform!(eachcol(vals), IP, coeffs, args...)
end

end  # PaduaTransforms


text:
using Pkg
Pkg.activate("/home/Morgan/FrankWolfe.jl")
# Pkg.activate("/home/Morgan/FrankWolfe.jl/fw-rde")

using FrankWolfe
# using FrankWolfe


text:
using AMDGPU

@show AMDGPU.agents()

# easiest way to program GPUs is using array operations

a = ROCArray([1 2 3 4])


text:
using CSVFiles
using DataFrames
using Distributions
using Mimi
using Statistics

##################################################################################################################################
# RUN VERSION OF SNEASY+CH4 THAT NEGLECTS POSTERIOR PARAMETER CORRELATIONS.
##################################################################################################################################
# Description: This creates a function that runs two versions of SNEASY+CH4 (a standard run, and a run with an extra pulse
#              of CH₄ emissions in a user-specified year) and saves the key model projection output. It samples the posterior
#              parameters independently to remove any posterior correlations that emerged during model calibration.
#
# Function Arguments:
#
#       climate_model = A symbol identifying the specific version of SNEASY+CH4 (options are :sneasy_fair, :sneasy_fund,
#                       :sneasy_hector, and :sneasy_magicc).
#       rcp           = A string identifying the RCP emissions and forcing scenario to use (options are "RCP26" and "RCP85").
#       pulse_year    = The year to add a pulse of methane emissions.
#       pulse_size    = The size of the methane emissions pulse in MtCH₄.
#       end_year      = The final year to run the model for.
#----------------------------------------------------------------------------------------------------------------------------------

function construct_sneasych4_remove_correlations(climate_model::Symbol, rcp::String,  pulse_year::Int, pulse_size::Float64, end_year::Int)

    #-------------------------------------------
    # Load and clean up forcing scenario data.
    #-------------------------------------------

    # Load RCP scenario emissions data.
    rcp_emissions = DataFrame(load(joinpath(@__DIR__, "..", "..", "data", "model_data", rcp*"_emissions.csv"), skiplines_begin=36))

    # Crop emissions to proper time periods (1765-end_year).
    rcp_indices = findall((in)(collect(1765:end_year)), rcp_emissions.YEARS)
    rcp_emissions = rcp_emissions[rcp_indices, :]

    # Get years the model is run for (initializing it in 1765).
    model_years  = collect(1765:end_year)
    number_years = length(model_years)

    # Get indices for 1861-1880 to normalize temperature projections.
    indices_1861_1880 = findall((in)(collect(1861:1880)), model_years)

    # Load calibration data from 1765-2017 (measurement errors used in simulated noise).
    calibration_data = load_calibration_data(2017)

    # Pre-allocate vectors to hold simulated CAR(1) & AR(1) with measurement error noise.
    norm_oceanco2   = zeros(number_years)
    ar1_temperature = zeros(number_years)
    ar1_oceanheat   = zeros(number_years)
    car1_co2        = zeros(number_years)
    car1_ch4        = zeros(number_years)

    # Replicate errors for years without observations over model time horizon (note, CO₂ and CH₄ ice-core have constant error estimates).
    obs_error_temperature = replicate_errors(1765, end_year, calibration_data.hadcrut_temperature_sigma)
    obs_error_oceanheat   = replicate_errors(1765, end_year, calibration_data.ocean_heat_sigma)
    # Use constant Law Dome CH₄ observation errors for start period to 1983 (after which time-varying NOAA flask measurements start).
    obs_error_ch4 = replicate_errors(1765, end_year, calibration_data.noaa_ch4_sigma)
    obs_error_ch4[1:findfirst(isequal(1983), model_years)] .= unique(skipmissing(calibration_data.lawdome_ch4_sigma))[1]
    # Set constant CO₂ observations errors for start year-1958 (Law Dome), and 1959-end year (Mauna Loa).
    obs_error_co2 = ones(number_years) .* unique(skipmissing(calibration_data.maunaloa_co2_sigma))[1]
    obs_error_co2[1:findfirst(isequal(1958), model_years)] .= unique(skipmissing(calibration_data.lawdome_co2_sigma))[1]

    # Set up marginal CH₄ emissions time series to have an extra emission pulse in a user-specified year.
    ch4_emissions_pulse = rcp_emissions.CH4
    pulse_year_index = findall(x -> x == pulse_year, rcp_emissions.YEARS)[1]
    ch4_emissions_pulse[pulse_year_index] = ch4_emissions_pulse[pulse_year_index] + pulse_size

    #-----------------------------------------------------------------------------------
    # Create two versions of SNEASY+CH4 based on the specific methane cycle selected.
    #-----------------------------------------------------------------------------------
    if climate_model == :sneasy_fair
        sneasych4_base  = create_sneasy_fairch4(rcp_scenario=rcp)
        sneasych4_pulse = create_sneasy_fairch4(rcp_scenario=rcp)
        update_param!(sneasych4_pulse, :fossil_emiss_CH₄, ch4_emissions_pulse)

    elseif climate_model == :sneasy_fund
        sneasych4_base  = create_sneasy_fundch4(rcp_scenario=rcp)
        sneasych4_pulse = create_sneasy_fundch4(rcp_scenario=rcp)
        update_param!(sneasych4_pulse, :globch4, ch4_emissions_pulse)

    elseif climate_model == :sneasy_hector
        sneasych4_base  = create_sneasy_hectorch4(rcp_scenario=rcp)
        sneasych4_pulse = create_sneasy_hectorch4(rcp_scenario=rcp)
        update_param!(sneasych4_pulse, :CH4_emissions, ch4_emissions_pulse)

    elseif climate_model == :sneasy_magicc
        sneasych4_base  = create_sneasy_magiccch4(rcp_scenario=rcp)
        sneasych4_pulse = create_sneasy_magiccch4(rcp_scenario=rcp)
        update_param!(sneasych4_pulse, :CH4_emissions, ch4_emissions_pulse)

    else
        error("Incorrect climate model selected. Options include :sneasy_fair, :sneasy_fund, :sneasy_hector, or :sneasy_magicc.")
    end

    # F2x_CO₂ has a default value in the component definition and therefore no
    # parameter in the model is created. The following line creates a parameter
    # in the model that can then be updated with update_param! in the inner loop.
    set_param!(sneasych4_base, :doeclim, :F2x_CO₂, 3.7)
    set_param!(sneasych4_pulse, :doeclim, :F2x_CO₂, 3.7)

    #---------------------------------------------------------------------------------------------------------------------
    # Create functions to set parameters and return results that differ across the four methane cycle models.
    #---------------------------------------------------------------------------------------------------------------------

    # Create a function to set the uncertain methane cycle parameters.
    # Note: This function assumes parameters follow the same ordering as in the `create_log_posterior.jl` file.
    update_ch4_params! = create_update_ch4_function(climate_model)

    # Create a function to return projected CH₄ concentrations (component and variable names differ across models).
    get_ch4_results! = create_get_ch4_results_function(climate_model)

    # Create model instances
    sneasych4_base_mi = Mimi.build(sneasych4_base)
    sneasych4_pulse_mi = Mimi.build(sneasych4_pulse)

    #-------------------------------------------------------------------------------------------------------------------------
    # Given user-specified settings, create a function to run SNEASY+CH4 over the independently sampled posterior parameters.
    #-------------------------------------------------------------------------------------------------------------------------

    function sneasych4_remove_correlations(calibrated_parameters::Array{Float64,2}, ci_interval_1::Float64, ci_interval_2::Float64)

        # Calculate number of calibrated parameter samples (row = sample from joint posterior distribution, column = specific parameter).
        number_samples, number_parameters = size(calibrated_parameters)

        # Calculate random indices to sample each posterior parameter independently (without repeating values).
        random_indices    = zeros(Int, number_samples, number_parameters)
        random_parameters = zeros(number_samples, number_parameters)

        for i = 1:number_parameters
            # For each uncertain parameter, create a vector of random indices (from 1:n_samples).
            random_indices[:,i] = sample(1:number_samples, number_samples, replace=false)
            # Use these indices to shuffle up the calibrated parameters to remove posterior correlations.
            random_parameters[:,i] = calibrated_parameters[random_indices[:,i], i]
        end

        # Pre-allocate arrays to store SNEASY+CH4 results.
        base_temperature  =   zeros(Union{Missing, Float64}, number_samples, number_years)
        base_co2          =   zeros(Union{Missing, Float64}, number_samples, number_years)
        base_ch4          =   zeros(Union{Missing, Float64}, number_samples, number_years)
        base_ocean_heat   =   zeros(Union{Missing, Float64}, number_samples, number_years)
        base_oceanco2     =   zeros(Union{Missing, Float64}, number_samples, number_years)

        pulse_temperature =   zeros(Union{Missing, Float64}, number_samples, number_years)
        pulse_co2         =   zeros(Union{Missing, Float64}, number_samples, number_years)

        # Use independently sampled marginal posterior parameters to run the base and pulse versions of SNEASY+CH4.
        for i in 1:number_samples

            σ_temperature      = random_parameters[i,1]
            σ_ocean_heat       = random_parameters[i,2]
            σ²_white_noise_CO₂ = random_parameters[i,3]
            σ²_white_noise_CH₄ = random_parameters[i,4]
            ρ_temperature      = random_parameters[i,5]
            ρ_ocean_heat       = random_parameters[i,6]
            α₀_CO₂             = random_parameters[i,7]
            α₀_CH₄             = random_parameters[i,8]
            temperature_0      = random_parameters[i,9]
            ocean_heat_0       = random_parameters[i,10]
            CO₂_0              = random_parameters[i,11]
            CH₄_0              = random_parameters[i,12]
            N₂O_0              = random_parameters[i,13]
            ECS                = random_parameters[i,14]
            heat_diffusivity   = random_parameters[i,15]
            rf_scale_aerosol   = random_parameters[i,16]
            rf_scale_CH₄       = random_parameters[i,17]
            F2x_CO₂            = random_parameters[i,18]
            Q10                = random_parameters[i,19]
            CO₂_fertilization  = random_parameters[i,20]
            CO₂_diffusivity    = random_parameters[i,21]

            # Set parameters for base version of SNEASY+CH4.
            update_param!(sneasych4_base_mi, :t2co, ECS)
            update_param!(sneasych4_base_mi, :kappa, heat_diffusivity)
            update_param!(sneasych4_base_mi, :F2x_CO₂, F2x_CO₂)
            update_param!(sneasych4_base_mi, :Q10, Q10)
            update_param!(sneasych4_base_mi, :Beta, CO₂_fertilization)
            update_param!(sneasych4_base_mi, :Eta, CO₂_diffusivity)
            update_param!(sneasych4_base_mi, :atmco20, CO₂_0)
            update_param!(sneasych4_base_mi, :CO₂_0, CO₂_0)
            update_param!(sneasych4_base_mi, :N₂O_0, N₂O_0)
            update_param!(sneasych4_base_mi, :rf_scale_CO₂, co2_rf_scale(F2x_CO₂, CO₂_0, N₂O_0))
            update_param!(sneasych4_base_mi, :α, rf_scale_aerosol)
            update_ch4_params!(sneasych4_base_mi, Vector(random_parameters[i,:]))

            # Set parameters for version of SNEASY+CH4 with CH₄ emissions pulse.
            update_param!(sneasych4_pulse_mi, :t2co, ECS)
            update_param!(sneasych4_pulse_mi, :kappa, heat_diffusivity)
            update_param!(sneasych4_pulse_mi, :F2x_CO₂, F2x_CO₂)
            update_param!(sneasych4_pulse_mi, :Q10, Q10)
            update_param!(sneasych4_pulse_mi, :Beta, CO₂_fertilization)
            update_param!(sneasych4_pulse_mi, :Eta, CO₂_diffusivity)
            update_param!(sneasych4_pulse_mi, :atmco20, CO₂_0)
            update_param!(sneasych4_pulse_mi, :CO₂_0, CO₂_0)
            update_param!(sneasych4_pulse_mi, :N₂O_0, N₂O_0)
            update_param!(sneasych4_pulse_mi, :rf_scale_CO₂, co2_rf_scale(F2x_CO₂, CO₂_0, N₂O_0))
            update_param!(sneasych4_pulse_mi, :α, rf_scale_aerosol)
            update_ch4_params!(sneasych4_pulse_mi, Vector(random_parameters[i,:]))

            # Wrap code in a try/catch statement in case non-physical parameter combinations produce a model error.
            try

                # Run both models.
                run(sneasych4_base_mi)
                run(sneasych4_pulse_mi)

                # Create noise to superimpose on results using calibrated statistical parameters and measurement noise (note: Both models use same estimated noise).
                ar1_temperature[:] = simulate_ar1_noise(number_years, σ_temperature, ρ_temperature, obs_error_temperature)
                ar1_oceanheat[:]   = simulate_ar1_noise(number_years, σ_ocean_heat,  ρ_ocean_heat,  obs_error_oceanheat)
                norm_oceanco2[:]   = rand(Normal(0,0.4*sqrt(10)), number_years)

                # CO₂ and CH₄ use CAR(1) statistical process parameters.
                car1_co2[:] = simulate_car1_noise(number_years, α₀_CO₂, σ²_white_noise_CO₂, obs_error_co2)
                car1_ch4[:] = simulate_car1_noise(number_years, α₀_CH₄, σ²_white_noise_CH₄, obs_error_ch4)

                # Store model projections resulting from parameter sample `i` for base model.
                base_temperature[i,:]  = sneasych4_base_mi[:doeclim, :temp] .- mean(sneasych4_base_mi[:doeclim, :temp][indices_1861_1880]) .+ ar1_temperature .+ temperature_0
                base_co2[i,:]          = sneasych4_base_mi[:ccm, :atmco2] .+ car1_co2
                base_ocean_heat[i,:]   = sneasych4_base_mi[:doeclim, :heat_interior] .+ ar1_oceanheat .+ ocean_heat_0
                base_oceanco2[i,:]     = sneasych4_base_mi[:ccm, :atm_oc_flux] .+ norm_oceanco2
                base_ch4[i,:]          = get_ch4_results!(sneasych4_base_mi) .+ car1_ch4

                # Store tempeature and CO₂ projections resulting from parameter sample `i` for pulse model (used for estimating the SC-CH₄).
                pulse_temperature[i,:] = sneasych4_pulse_mi[:doeclim, :temp] .- mean(sneasych4_pulse_mi[:doeclim, :temp][indices_1861_1880]) .+ ar1_temperature .+ temperature_0
                pulse_co2[i,:]         = sneasych4_pulse_mi[:ccm, :atmco2] .+ car1_co2

            catch

                # Set values to -99999.99 if non-physical parameter combinations produce a model error.
                base_temperature[i,:]  .= -99999.99
                base_co2[i,:]          .= -99999.99
                base_ocean_heat[i,:]   .= -99999.99
                base_oceanco2[i,:]     .= -99999.99
                base_ch4[i,:]          .= -99999.99
                pulse_temperature[i,:] .= -99999.99
                pulse_co2[i,:]         .= -99999.99
            end
        end

        # Identify model indices that caused a model error or yield non-physical outcomes (i.e. strongly negative temperatures in 2300 under RCP 8.5).
        error_indices = findall(x-> x < -10.0, base_temperature[:,end])
        good_indices  = findall(!in(error_indices), collect(1:number_samples))

        # Calculate credible intervals for base model projections using model indices that did not cause errors.
        ci_temperature = get_confidence_interval(collect(1765:end_year), base_temperature[good_indices,:], ci_interval_1, ci_interval_2)
        ci_co2         = get_confidence_interval(collect(1765:end_year), base_co2[good_indices,:],         ci_interval_1, ci_interval_2)
        ci_ocean_heat  = get_confidence_interval(collect(1765:end_year), base_ocean_heat[good_indices,:],  ci_interval_1, ci_interval_2)
        ci_oceanco2    = get_confidence_interval(collect(1765:end_year), base_oceanco2[good_indices,:],    ci_interval_1, ci_interval_2)
        ci_ch4         = get_confidence_interval(collect(1765:end_year), base_ch4[good_indices,:],         ci_interval_1, ci_interval_2)

        return base_temperature[good_indices,:], base_co2[good_indices,:], base_ch4[good_indices,:], base_ocean_heat[good_indices,:], base_oceanco2[good_indices,:], pulse_temperature[good_indices,:], pulse_co2[good_indices,:],
               ci_temperature, ci_co2, ci_ocean_heat, ci_oceanco2, ci_ch4,
               error_indices, good_indices, random_indices

    end

    # Return function with user model specifications.
    return sneasych4_remove_correlations
end


text:
# ## Data

import CSV

using TuringModels: project_root
using DataFrames

path = joinpath(project_root, "data", "chimpanzees.csv")
df = CSV.read(path, DataFrame; delim=';');

# ## Model

using Turing

@model m12_4(pulled_left, actor, condition, prosoc_left) = begin
    ## Total num of y
    N = length(pulled_left)

    ## Separate σ priors for each actor
    σ_actor ~ truncated(Cauchy(0, 1), 0, Inf)

    ## Number of unique actors in the data set
    N_actor = length(unique(actor)) #7

    ## Vector of actors (1,..,7) which we'll set priors on
    α_actor ~ filldist(Normal(0, σ_actor), N_actor)

    ## Prior for intercept, prosoc_left, and the interaction
    α ~ Normal(0, 10)
    βp ~ Normal(0, 10)
    βpC ~ Normal(0, 10)

    logitp = α .+ α_actor[actor] .+
            (βp .+ βpC * condition) .* prosoc_left

    pulled_left .~ BinomialLogit.(1, logitp)
end;

# ## Output

chains = sample(
    m12_4(df.pulled_left, df.actor, df.condition, df.prosoc_left),
    NUTS(0.65),
    1000
)

# \defaultoutput{}

# ## Original output

m124rethinking = "
             Mean StdDev lower 0.89 upper 0.89 n_eff Rhat
 a_actor[1]  -1.13   0.95      -2.62       0.27  2739    1
 a_actor[2]   4.17   1.66       1.80       6.39  3958    1
 a_actor[3]  -1.44   0.95      -2.90      -0.02  2720    1
 a_actor[4]  -1.44   0.94      -2.92      -0.04  2690    1
 a_actor[5]  -1.13   0.94      -2.58       0.31  2727    1
 a_actor[6]  -0.19   0.94      -1.64       1.25  2738    1
 a_actor[7]   1.34   0.97      -0.09       2.87  2889    1
 a            0.42   0.93      -1.00       1.81  2622    1
 bp           0.83   0.26       0.41       1.25  8594    1
 bpC         -0.13   0.30      -0.62       0.34  8403    1
 sigma_actor  2.26   0.94       1.07       3.46  4155    1
";


text:
using Py2JlRuntime
# .\runtests\example.jl, line 1
function var".f_1"(_x′, _y′)
    x = _x′
    y = _y′
    local x
    local y
    # .\runtests\example.jl, line 2
    return jpy_add(x, y)
    jpy_none;
end
const f = var".f_1"
# .\runtests\example.jl, line 4
function var".g_2"(_x′ ; _y′ = jpy_literal(2))
    x = _x′
    y = _y′
    local x
    local y
    # .\runtests\example.jl, line 5
    return jpy_call(f, (x, y, ), (;))
    jpy_none;
end
const g = var".g_2"
# .\runtests\example.jl, line 7
function var".range_3"(_n′)
    Channel() do PyContinuation
      n = _n′
      local n
      local i
      # .\runtests\example.jl, line 8
      i = jpy_literal(0)
      # .\runtests\example.jl, line 9
      @noscope while @jpy_all(jpy_lt(i, n))
          # .\runtests\example.jl, line 10
          @jpy_yield(i);
          # .\runtests\example.jl, line 11
          i = jpy_iadd(i, jpy_literal(1))
          jpy_none;
      end
      jpy_none;
    end
end
const range = var".range_3"
# .\runtests\example.jl, line 12
s = jpy_literal(0)
# .\runtests\example.jl, line 13
var".iter_4" = jpy_getiter(jpy_list(jpy_literal(1), jpy_literal(2), jpy_literal(3)))
@noscope while jpy_movenext(var".iter_4")
    e = jpy_getcurrent(var".iter_4")
    # .\runtests\example.jl, line 14
    s = jpy_add(s, jpy_call(g, (e, ), (;)))
    jpy_none;
end
# .\runtests\example.jl, line 16
d = jpy_dict(jpy_literal(1)=>jpy_literal(2), jpy_literal(3)=>jpy_literal(4))
# .\runtests\example.jl, line 18
jpy_call(println, (s, ), (;));
# .\runtests\example.jl, line 19
jpy_call(println, (d[jpy_literal(3)], ), (;));
jpy_none;


text:
module MyTestRepo

using Plots

function lsq_fit(x, y)
    @assert length(x) == length(y)
    N = length(x)
    X = zeros(N, 2)
    for j=1:N
        X[j,1] = 1.0
        X[j,2] = x[j]
    end
    C = X'*X
    alpha = inv(C)*X'*y
    return alpha
end

function plot_fit(x, y, fn)
    # x = [0.0, 2.0, 4.0]
    # y = [0.0, 3.0, 5.0]
    α0, α1 = lsq_fit(x, y)
    scatter(x, y, color=:black, label="Data")
    plot!(x, α0 .+ α1*x, color=:black, lw=3)
    savefig(fn)
end

export lsq_fit, plot_fit

#=
comment
long comment
=#

end


text:
# To view in browser start a server in the build dir:
# python -m http.server --bind localhost

if !("./" in LOAD_PATH)
    push!(LOAD_PATH, "./")
end

using Pkg
Pkg.add("Documenter")

using Documenter
using QuasinormalModes

makedocs(sitename="QuasinormalModes.jl",
         modules=[QuasinormalModes],
         pages=[
             "index.md",
             "intro.md",
             "org.md",
             "schw.md",
             "sho.md",
             "api_ref.md"
             ]
         )

deploydocs(
    repo="github.com/lucass-carneiro/QuasinormalModes.jl.git",
    target="build",
    devurl="./",
    push_preview=false,
)


text:
# As your probe drifted down through this area, it released an assortment of
# beacons and scanners into the water. It's difficult to navigate in the pitch
# black open waters of the ocean trench, but if you can build a map of the
# trench using data from the scanners, you should be able to safely reach the
# bottom.
#
# The beacons and scanners float motionless in the water; they're designed to
# maintain the same position for long periods of time. Each scanner is capable
# of detecting all beacons in a large cube centered on the scanner; beacons
# that are at most 1000 units away from the scanner in each of the three axes
# (x, y, and z) have their precise position determined relative to the scanner.
# However, scanners cannot detect other scanners. The submarine has
# automatically summarized the relative positions of beacons detected by each
# scanner (your puzzle input).
#
# For example, if a scanner is at x,y,z coordinates 500,0,-500 and there are
# beacons at -500,1000,-1500 and 1501,0,-500, the scanner could report that the
# first beacon is at -1000,1000,-1000 (relative to the scanner) but would not
# detect the second beacon at all.
#
# Unfortunately, while each scanner can report the positions of all detected
# beacons relative to itself, the scanners do not know their own position.
# You'll need to determine the positions of the beacons and scanners yourself.
#
# The scanners and beacons map a single contiguous 3d region. This region can
# be reconstructed by finding pairs of scanners that have overlapping detection
# regions such that there are at least 12 beacons that both scanners detect
# within the overlap. By establishing 12 common beacons, you can precisely
# determine where the scanners are relative to each other, allowing you to
# reconstruct the beacon map one scanner at a time.
#
# For a moment, consider only two dimensions. Suppose you have the following
# scanner reports:
#
# --- scanner 0 ---
# 0,2
# 4,1
# 3,3
#
# --- scanner 1 ---
# -1,-1
# -5,0
# -2,1
#
# Drawing x increasing rightward, y increasing upward, scanners as S, and
# beacons as B, scanner 0 detects this:
#
# ...B.
# B....
# ....B
# S....
#
# Scanner 1 detects this:
#
# ...B..
# B....S
# ....B.
#
# For this example, assume scanners only need 3 overlapping beacons. Then, the
# beacons visible to both scanners overlap to produce the following complete
# map:
#
# ...B..
# B....S
# ....B.
# S.....
#
# Unfortunately, there's a second problem: the scanners also don't know their
# rotation or facing direction. Due to magnetic alignment, each scanner is
# rotated some integer number of 90-degree turns around all of the x, y, and z
# axes. That is, one scanner might call a direction positive x, while another
# scanner might call that direction negative y. Or, two scanners might agree on
# which direction is positive x, but one scanner might be upside-down from the
# perspective of the other scanner. In total, each scanner could be in any of
# 24 different orientations: facing positive or negative x, y, or z, and
# considering any of four directions "up" from that facing.
#
# For example, here is an arrangement of beacons as seen from a scanner in the
# same position but in different orientations:
#
# --- scanner 0 ---
# -1,-1,1
# -2,-2,2
# -3,-3,3
# -2,-3,1
# 5,6,-4
# 8,0,7
#
# --- scanner 0 ---
# 1,-1,1
# 2,-2,2
# 3,-3,3
# 2,-1,3
# -5,4,-6
# -8,-7,0
#
# --- scanner 0 ---
# -1,-1,-1
# -2,-2,-2
# -3,-3,-3
# -1,-3,-2
# 4,6,5
# -7,0,8
#
# --- scanner 0 ---
# 1,1,-1
# 2,2,-2
# 3,3,-3
# 1,3,-2
# -4,-6,5
# 7,0,8
#
# --- scanner 0 ---
# 1,1,1
# 2,2,2
# 3,3,3
# 3,1,2
# -6,-4,-5
# 0,7,-8
#
# By finding pairs of scanners that both see at least 12 of the same beacons,
# you can assemble the entire map. For example, consider the following report:
#
# --- scanner 0 ---
# 404,-588,-901
# 528,-643,409
# -838,591,734
# 390,-675,-793
# -537,-823,-458
# -485,-357,347
# -345,-311,381
# -661,-816,-575
# -876,649,763
# -618,-824,-621
# 553,345,-567
# 474,580,667
# -447,-329,318
# -584,868,-557
# 544,-627,-890
# 564,392,-477
# 455,729,728
# -892,524,684
# -689,845,-530
# 423,-701,434
# 7,-33,-71
# 630,319,-379
# 443,580,662
# -789,900,-551
# 459,-707,401
#
# --- scanner 1 ---
# 686,422,578
# 605,423,415
# 515,917,-361
# -336,658,858
# 95,138,22
# -476,619,847
# -340,-569,-846
# 567,-361,727
# -460,603,-452
# 669,-402,600
# 729,430,532
# -500,-761,534
# -322,571,750
# -466,-666,-811
# -429,-592,574
# -355,545,-477
# 703,-491,-529
# -328,-685,520
# 413,935,-424
# -391,539,-444
# 586,-435,557
# -364,-763,-893
# 807,-499,-711
# 755,-354,-619
# 553,889,-390
#
# --- scanner 2 ---
# 649,640,665
# 682,-795,504
# -784,533,-524
# -644,584,-595
# -588,-843,648
# -30,6,44
# -674,560,763
# 500,723,-460
# 609,671,-379
# -555,-800,653
# -675,-892,-343
# 697,-426,-610
# 578,704,681
# 493,664,-388
# -671,-858,530
# -667,343,800
# 571,-461,-707
# -138,-166,112
# -889,563,-600
# 646,-828,498
# 640,759,510
# -630,509,768
# -681,-892,-333
# 673,-379,-804
# -742,-814,-386
# 577,-820,562
#
# --- scanner 3 ---
# -589,542,597
# 605,-692,669
# -500,565,-823
# -660,373,557
# -458,-679,-417
# -488,449,543
# -626,468,-788
# 338,-750,-386
# 528,-832,-391
# 562,-778,733
# -938,-730,414
# 543,643,-506
# -524,371,-870
# 407,773,750
# -104,29,83
# 378,-903,-323
# -778,-728,485
# 426,699,580
# -438,-605,-362
# -469,-447,-387
# 509,732,623
# 647,635,-688
# -868,-804,481
# 614,-800,639
# 595,780,-596
#
# --- scanner 4 ---
# 727,592,562
# -293,-554,779
# 441,611,-461
# -714,465,-776
# -743,427,-804
# -660,-479,-426
# 832,-632,460
# 927,-485,-438
# 408,393,-506
# 466,436,-512
# 110,16,151
# -258,-428,682
# -393,719,612
# -211,-452,876
# 808,-476,-593
# -575,615,604
# -485,667,467
# -680,325,-822
# -627,-443,-432
# 872,-547,-609
# 833,512,582
# 807,604,487
# 839,-516,451
# 891,-625,532
# -652,-548,-490
# 30,-46,-14
#
# Because all coordinates are relative, in this example, all "absolute"
# positions will be expressed relative to scanner 0 (using the orientation of
# scanner 0 and as if scanner 0 is at coordinates 0,0,0).
#
# Scanners 0 and 1 have overlapping detection cubes; the 12 beacons they both
# detect (relative to scanner 0) are at the following coordinates:
#
# -618,-824,-621
# -537,-823,-458
# -447,-329,318
# 404,-588,-901
# 544,-627,-890
# 528,-643,409
# -661,-816,-575
# 390,-675,-793
# 423,-701,434
# -345,-311,381
# 459,-707,401
# -485,-357,347
#
# These same 12 beacons (in the same order) but from the perspective of scanner
# 1 are:
#
# 686,422,578
# 605,423,415
# 515,917,-361
# -336,658,858
# -476,619,847
# -460,603,-452
# 729,430,532
# -322,571,750
# -355,545,-477
# 413,935,-424
# -391,539,-444
# 553,889,-390
#
# Because of this, scanner 1 must be at 68,-1246,-43 (relative to scanner 0).
#
# Scanner 4 overlaps with scanner 1; the 12 beacons they both detect (relative
# to scanner 0) are:
#
# 459,-707,401
# -739,-1745,668
# -485,-357,347
# 432,-2009,850
# 528,-643,409
# 423,-701,434
# -345,-311,381
# 408,-1815,803
# 534,-1912,768
# -687,-1600,576
# -447,-329,318
# -635,-1737,486
#
# So, scanner 4 is at -20,-1133,1061 (relative to scanner 0).
#
# Following this process, scanner 2 must be at 1105,-1205,1229 (relative to
# scanner 0) and scanner 3 must be at -92,-2380,-20 (relative to scanner 0).
#
# The full list of beacons (relative to scanner 0) is:
#
# -892,524,684
# -876,649,763
# -838,591,734
# -789,900,-551
# -739,-1745,668
# -706,-3180,-659
# -697,-3072,-689
# -689,845,-530
# -687,-1600,576
# -661,-816,-575
# -654,-3158,-753
# -635,-1737,486
# -631,-672,1502
# -624,-1620,1868
# -620,-3212,371
# -618,-824,-621
# -612,-1695,1788
# -601,-1648,-643
# -584,868,-557
# -537,-823,-458
# -532,-1715,1894
# -518,-1681,-600
# -499,-1607,-770
# -485,-357,347
# -470,-3283,303
# -456,-621,1527
# -447,-329,318
# -430,-3130,366
# -413,-627,1469
# -345,-311,381
# -36,-1284,1171
# -27,-1108,-65
# 7,-33,-71
# 12,-2351,-103
# 26,-1119,1091
# 346,-2985,342
# 366,-3059,397
# 377,-2827,367
# 390,-675,-793
# 396,-1931,-563
# 404,-588,-901
# 408,-1815,803
# 423,-701,434
# 432,-2009,850
# 443,580,662
# 455,729,728
# 456,-540,1869
# 459,-707,401
# 465,-695,1988
# 474,580,667
# 496,-1584,1900
# 497,-1838,-617
# 527,-524,1933
# 528,-643,409
# 534,-1912,768
# 544,-627,-890
# 553,345,-567
# 564,392,-477
# 568,-2007,-577
# 605,-1665,1952
# 612,-1593,1893
# 630,319,-379
# 686,-3108,-505
# 776,-3184,-501
# 846,-3110,-434
# 1135,-1161,1235
# 1243,-1093,1063
# 1660,-552,429
# 1693,-557,386
# 1735,-437,1738
# 1749,-1800,1813
# 1772,-405,1572
# 1776,-675,371
# 1779,-442,1789
# 1780,-1548,337
# 1786,-1538,337
# 1847,-1591,415
# 1889,-1729,1762
# 1994,-1805,1792
#
# In total, there are 79 beacons.
#
# Assemble the full map of beacons. How many beacons are there?

# read data
iter = eachline("input.txt")
scanners = Vector{Tuple{Int32, Int32, Int32}}[]
while !isnothing(iterate(iter))
  # while condition will consume the "--- scanner X ---" lines
  beacons = map(l -> tuple(parse.(Int32, split(l, ','))...), Iterators.takewhile(l -> l != "", iter))
  push!(scanners, beacons)
end

# Utility iterator to compute all combinations of 2 indices
struct Permutations
  len::Int
end

function Base.iterate(P::Permutations, (i,j)=(1,1))
  j += 1
  if j > P.len
    i += 1
    j = i + 1
    j > P.len && return nothing
  end
  idxs = (i,j)
  return (idxs, idxs)
end

Base.length(P::Permutations) = binomial(P.len, 2)

Base.eltype(P::Permutations) = Tuple{Int, Int}

# Distance is the sqrt of the sum of squares of differences. I'm going to leave
# off the sqrt bit so I'm not beholden to floating point errors.
function distance(scanneridx, i, j)
  (x1,y1,z1) = scanners[scanneridx][i]
  (x2,y2,z2) = scanners[scanneridx][j]
  (widen(x2)-x1)^2 + (widen(y2)-y1)^2 + (widen(z2)-z1)^2
end

distance(scanneridx, (i,j)) = distance(scanneridx, i, j)

distance(scanneridx) = (p) -> distance(scanneridx, p)

# Utility thot finds the distance between all beacons in the a scanner. This is
# used to find points that potentially match between two scanners.
struct RefDistances
  idxs::Vector{Tuple{Int, Int}}
  dists::Vector{Int}
  function RefDistances(baseidx)
    base = scanners[baseidx]
    idxs = collect(Permutations(length(base)))
    dists = map(distance(1), idxs)
    sortedidx = sortperm(dists)
    permute!(idxs, sortedidx)
    permute!(dists, sortedidx)
    new(idxs, dists)
  end
end

Base.getindex(R::RefDistances, i::Int) = R.idxs[i]

function Base.findfirst(R::RefDistances, dist)
  distidx = searchsortedfirst(R.dists, dist)
  distidx <= length(R.dists) && R.dists[distidx] == dist ? distidx : nothing
end

# Functions to re-orient points
const reorientfuncs = [
  identity, ((x,y,z),) -> (-x,-y,z), ((x,y,z),) -> (-x,-z,-y),
  ((x,y,z),) -> (-x,y,-z), ((x,y,z),) -> (-x,z,y), ((x,y,z),) -> (-y,-x,-z),
  ((x,y,z),) -> (-y,-z,x), ((x,y,z),) -> (-y,x,z), ((x,y,z),) -> (-y,z,-x),
  ((x,y,z),) -> (-z,-x,y), ((x,y,z),) -> (-z,-y,-x), ((x,y,z),) -> (-z,x,-y),
  ((x,y,z),) -> (-z,y,x), ((x,y,z),) -> (x,-y,-z), ((x,y,z),) -> (x,-z,y),
  ((x,y,z),) -> (x,z,-y), ((x,y,z),) -> (y,-x,z), ((x,y,z),) -> (y,-z,-x),
  ((x,y,z),) -> (y,x,-z), ((x,y,z),) -> (y,z,x), ((x,y,z),) -> (z,-x,-y),
  ((x,y,z),) -> (z,-y,x), ((x,y,z),) -> (z,x,y), ((x,y,z),) -> (z,y,-x),
]

# Utility that re-orients points and moves them so they are relative to the
# base scanner
struct Positioner
  reorientfunc
  adj::Tuple{Int, Int, Int}
  matched::BitVector
end

(P::Positioner)(t) = P.reorientfunc(t) .+ P.adj

function findpositioner(base, scanner, matches)
  len = length(scanner)
  best = nothing
  bestcnt = 0

  # if there are more than 12 total matches, we can't assume the first match is
  # actually valid
  for tryidx in 1:(length(matches)-11)
    for reorientfunc in reorientfuncs
      # We'll try reorienting the first matched point, then calculate an
      # adjustment. We'll then try the same reorientation and adjustment on the
      # rest of the matches.
      (b2,b1) = matches[tryidx]
      adj = base[b1] .- reorientfunc(scanner[b2])
      matched = falses(len)
      matched[b2] = true
      for k in (tryidx+1):length(matches)
        (b2,b1) = matches[k]
        if isequal(base[b1], reorientfunc(scanner[b2]) .+ adj)
          matched[b2] = true
        end
      end

      cnt = count(matched)
      if cnt > bestcnt && cnt >= 12
        best = Positioner(reorientfunc, adj, matched)
        bestcnt = cnt
      end
    end
  end
  return best
end

# utility to sort a list of numbers by the frequency - also reduces the set to
# unique values
function findmostcommon(v)
  cnts = Dict{eltype(v), Int}()
  common = nothing
  commoncnt = 0
  for e in v
    cnts[e] = get(cnts, e, 0) + 1
    if cnts[e] > commoncnt
      common = e
      commoncnt = cnts[e]
    end
  end
  return common
end

# loop over scanners until we've figured them all out - if the loop ever makes
# no progress, it tries to change the "base" and find other matches that way
finished = falses(length(scanners))
finished[1] = true
baseidx = 1
loopfinishedcnt = 1
refdist = RefDistances(baseidx)
while !all(finished)
  base = scanners[baseidx]
  progress = false
  for i in (baseidx+1):length(scanners)
    finished[i] && continue

    # find matching distances
    scanner = scanners[i]
    matches = Dict{Int, Vector{Int}}()
    for (j,k) in Permutations(length(scanner))
      dist = distance(i, j, k)
      distidx = findfirst(refdist, dist)
      isnothing(distidx) && continue
      (pot1, pot2) = refdist[distidx]
      if haskey(matches, j)
        push!(matches[j], pot1, pot2)
      else
        matches[j] = [pot1, pot2]
      end
      if haskey(matches, k)
        push!(matches[k], pot1, pot2)
      else
        matches[k] = [pot1, pot2]
      end
    end

    # do we have at least 12 matches?
    length(matches) < 12 && continue

    # figure out how this scanner is oriented relative to the first
    positioner = findpositioner(base, scanner, [(k,findmostcommon(v)) for (k,v) in matches])
    if isnothing(positioner)
      println("Had 12+ matches, but couldn't find a positioner for scanner ", i)
      continue
    end

    # add new points
    append!(base, map(positioner, scanner[.!positioner.matched]))
    println("Number of beacons after ", count(finished), " scanners: ", length(base))

    # rebuild distances and mark this scanner finished
    global refdist = RefDistances(baseidx)
    finished[i] = true
    progress = true
  end

  if !progress
    global baseidx = findnext(i -> !i, finished, baseidx+1)
    if isnothing(baseidx)
      global baseidx = 1
      cnt = count(finished)
      loopfinishedcnt == cnt && error("stuck")
      global loopfinishedcnt = cnt
    end
    println("Switching base to ", baseidx)
    global refdist = RefDistances(baseidx)
  end
end

println("Number of beacons: ", length(scanners[1]))