text:
module Fields

import LinearAlgebra
using UnPack
using ..Grids

using Base.Broadcast: Broadcasted, BroadcastStyle, ArrayStyle

export AbstractField, CellCenter, CellCorner, CellFace, CellEdge
export CellFaces, CellEdges, CellCenters, CellCorners
export dual_fields
export average, midslice, midline
export primary_along, dual_along
export GridField

abstract type FieldLocations end
struct Center <: FieldLocations end
struct Face <: FieldLocations end
struct Edge <: FieldLocations end
struct Corner <: FieldLocations end

abstract type GridType end
struct Primary <: GridType end
struct Dual <: GridType end

"""
    get_dl

A 3-element tuple of `DataLocation`'s,
based on the grid type and direction.
"""
function get_dl end

get_dl(::Primary) = (Center1D(),Center1D(),Center1D())
get_dl(::Dual) = (Vertex1D(),Vertex1D(),Vertex1D())
get_dl(::Primary, dir::Val{1}) = (Vertex1D(),Center1D(),Center1D())
get_dl(::Primary, dir::Val{2}) = (Center1D(),Vertex1D(),Center1D())
get_dl(::Primary, dir::Val{3}) = (Center1D(),Center1D(),Vertex1D())
get_dl(::Dual, dir::Val{1}) = (Center1D(),Vertex1D(),Vertex1D())
get_dl(::Dual, dir::Val{2}) = (Vertex1D(),Center1D(),Vertex1D())
get_dl(::Dual, dir::Val{3}) = (Vertex1D(),Vertex1D(),Center1D())
get_dl(gt::GridType, dir::Int) = get_dl(gt, Val(dir))

abstract type AbstractField{FT,NDIMS} <: AbstractArray{FT, NDIMS} end

Base.size(f::AbstractField) = Base.size(f.data)
Base.size(f::AbstractField, dim) = Base.size(f.data, dim)
Base.length(f::AbstractField) = Base.length(f.data)
Base.eltype(f::AbstractField) = Base.eltype(f.data)
Base.ndims(::Type{T}) where {NDIMS,FT,T<:AbstractField{FT,NDIMS}} = NDIMS
Base.iterate(f::AbstractField, state=1) =
    state > length(f) ? nothing : (f[state], state+1)
Base.lastindex(f::AbstractField, dim) = size(f, dim)

Base.copyto!(dest::Array, f::AbstractField) = copyto!(dest, f.data)
Base.BroadcastStyle(::Type{<:AbstractField}) = ArrayStyle{AbstractField}()

find_field(bc::Broadcasted) = find_field(bc.args)
find_field(args::Tuple) = find_field(find_field(args[1]), Base.tail(args))
find_field(x) = x
find_field(a::AbstractField, rest) = a
find_field(::Any, rest) = find_field(rest)

Base.similar(bc::Broadcasted{ArrayStyle{AbstractField}}, ::Type{FT}) where {FT} =
    similar(find_field(bc), FT)

function Base.copyto!(dest::AbstractField, src::AbstractField)
    copyto!(dest.data, src.data)
    dest
end

transform_array(f::AbstractField) = f.data
transform_array(x) = x

@inline function Base.copyto!(dest::AbstractField, bc::Broadcasted{Nothing})
    # check for the case a .= b, where b is an array
    if bc.f === identity && bc.args isa Tuple{AbstractArray}
        copyto!(dest.data, bc.args[1])
    else
        copyto!(dest.data, Broadcasted(bc.f, transform_array.(bc.args), bc.axes))
    end
    dest
end

LinearAlgebra.norm(f::AbstractField) = LinearAlgebra.norm(f.data)
Base.@propagate_inbounds Base.getindex(f::AbstractField, key...) =
    Base.getindex(f.data, key...)
Base.@propagate_inbounds Base.setindex!(f::AbstractField, v, key...) =
    Base.setindex!(f.data, v, key...)

struct CellCenter{A,FT,NDIMS} <: AbstractField{FT,NDIMS}
    data::A
    function CellCenter(grid::Grid{FT}) where {FT}
        s = size(grid, get_dl(Primary()))
        data = zeros(FT, s...)
        new{typeof(data),FT,length(s)}(data)
    end
end
CellCenters(grid::Grid) = (CellCenter(grid),CellCenter(grid),CellCenter(grid))

struct CellCorner{A,FT,NDIMS} <: AbstractField{FT,NDIMS}
    data::A
    function CellCorner(grid::Grid{FT}) where {FT}
        s = size(grid, get_dl(Dual()))
        data = zeros(FT, s...)
        new{typeof(data),FT,length(s)}(data)
    end
end
CellCorners(grid::Grid) = (CellCorner(grid),CellCorner(grid),CellCorner(grid))

struct CellFace{dir,A,FT,NDIMS} <: AbstractField{FT,NDIMS}
    data::A
    function CellFace(grid::Grid{FT}, dir::Int) where {FT}
        @assert 1 <= dir <= 3
        s = size(grid, get_dl(Primary(), dir))
        data = zeros(FT, s...)
        new{dir,typeof(data),FT,length(s)}(data)
    end
end
CellFaces(grid::Grid) = (CellFace(grid, 1),CellFace(grid, 2),CellFace(grid, 3))

struct CellEdge{dir,A,FT,NDIMS} <: AbstractField{FT,NDIMS}
    data::A
    function CellEdge(grid::Grid{FT}, dir::Int) where {FT}
        @assert 1 <= dir <= 3
        s = size(grid, get_dl(Dual(), dir))
        data = zeros(FT, s...)
        new{dir,typeof(data),FT,length(s)}(data)
    end
end
CellEdges(grid::Grid) = (CellEdge(grid, 1),CellEdge(grid, 2),CellEdge(grid, 3))

direction(f::CellCenter) = nothing
direction(f::CellCorner) = nothing
direction(f::CellFace{dir}) where {dir} = dir
direction(f::CellEdge{dir}) where {dir} = dir

midslice(f, dir::Int) = midslice(f, Val(dir))
midslice(f, ::Val{1}) = f[round(Int,size(f, 1)/2),:,:]
midslice(f, ::Val{2}) = f[:,round(Int,size(f, 2)/2),:]
midslice(f, ::Val{3}) = f[:,:,round(Int,size(f, 3)/2)]

average(f, dir::Int) = average(f, Val(dir))
average(f, ::Val{1}) = dropdims(sum(f; dims=1), dims=1) ./ size(f, 1)
average(f, ::Val{2}) = dropdims(sum(f; dims=2), dims=2) ./ size(f, 2)
average(f, ::Val{3}) = dropdims(sum(f; dims=3), dims=3) ./ size(f, 3)

midline(f, dir::Int) = midline(f, Val(dir))
midline(f, ::Val{1}) = f[:,round(Int,size(f, 2)/2),round(Int,size(f, 3)/2)]
midline(f, ::Val{2}) = f[round(Int,size(f, 1)/2),:,round(Int,size(f, 3)/2)]
midline(f, ::Val{3}) = f[round(Int,size(f, 1)/2),round(Int,size(f, 2)/2),:]

struct GridField{G,F,CI}
    grid::G
    f::F
    ci::CI
    function GridField(grid::Grid, f::AbstractField)
        ci = CartesianIndices(size(f))
        new{typeof(grid), typeof(f),typeof(ci)}(grid, f, ci)
    end
end

get_dl(::CellCenter, _) = get_dl(Primary())
get_dl(::CellCorner, _) = get_dl(Dual())
get_dl(::CellFace, ::Val{N}) where {N} = get_dl(Primary(), Val(N))
get_dl(::CellEdge, ::Val{N}) where {N} = get_dl(Dual(), Val(N))
get_dl(f::AbstractField, dir::Int) = get_dl(f, Val(dir))

primary_along(f::AbstractField) = (primary_along(f, 1),primary_along(f, 2),primary_along(f, 3))
primary_along(f::AbstractField, dir::Int) = primary_along(f, Val(dir))
primary_along(f::CellCenter, ::Val{dim}) where {dim} = Center1D()
primary_along(f::CellCorner, ::Val{dim}) where {dim} = Vertex1D()
primary_along(f::CellFace{dim}, ::Val{dim}) where {dim} = Vertex1D()
primary_along(f::CellFace{fdim}, ::Val{dim}) where {dim,fdim} = Center1D()
primary_along(f::CellEdge{dim}, ::Val{dim}) where {dim} = Center1D()
primary_along(f::CellEdge{fdim}, ::Val{dim}) where {dim,fdim} = Vertex1D()

dual_along(f::AbstractField) = (dual_along(f, 1),dual_along(f, 2),dual_along(f, 3))
dual_along(f::AbstractField, dir::Int) = dual_along(f, Val(dir))
dual_along(f::CellCenter, ::Val{dim}) where {dim} = Vertex1D()
dual_along(f::CellCorner, ::Val{dim}) where {dim} = Center1D()
dual_along(f::CellFace{dim}, ::Val{dim}) where {dim} = Center1D()
dual_along(f::CellFace{fdim}, ::Val{dim}) where {dim,fdim} = Vertex1D()
dual_along(f::CellEdge{dim}, ::Val{dim}) where {dim} = Vertex1D()
dual_along(f::CellEdge{fdim}, ::Val{dim}) where {dim,fdim} = Center1D()

space(grid::Grid, ::Center1D) = grid.cent
space(grid::Grid, ::Vertex1D) = grid.corn

dual_fields(grid::Grid, f::CellCenter) = CellFaces(grid)
dual_fields(grid::Grid, f::CellCorner) = CellEdges(grid)

function Base.iterate(gf::GridField{G,F}, state=1) where {G,F<:AbstractField}
    if state > length(gf)
        return nothing
    else
        ijk = gf.ci[state]
        dl = get_dl(gf.f, direction(gf.f))
        x = space(gf.grid, dl[1]).x[Tuple(ijk)...]
        y = space(gf.grid, dl[2]).y[Tuple(ijk)...]
        z = space(gf.grid, dl[3]).z[Tuple(ijk)...]
        return ((;x,y,z,ijk), state+1)
    end
end
Base.length(gf::GridField) = length(gf.f)

function sweep(A, dim, iR = 1:size(A, dim))
    Rpre = CartesianIndices(size(A)[1:dim-1])
    Rpost = CartesianIndices(size(A)[dim+1:end])
    return Iterators.product(Rpre, iR, Rpost)
end

# TODO: improve implementation
export assign!
function assign!(f::AbstractField, val, ::Ghost)
    f[1,:,:] .= val; f[end,:,:] .= val
    f[:,1,:] .= val; f[:,end,:] .= val
    f[:,:,1] .= val; f[:,:,end] .= val
end

include("interpolations_base.jl")
include("interpolations_md.jl")
include("extrapolations.jl")
include("discrete_ops.jl")

end


text:
"""
    decouple_covmat(Y, B1, B2; lmin1=2, lmin2=2)

Decouples a covariance matrix Y, performing B₁⁻¹ × Y × (B₂⁻¹)^†
by mutating Y.
"""
function decouple_covmat(Y::SA, B1::SA, B2::SA) where {T, SA <: SpectralArray{T,2}}
    M = deepcopy(Y)
    C = parent(M)
    rdiv!(C', lu(parent(B1)'))
    rdiv!(C, lu(parent(B2)'))
    return M
end

"""
    coupledcov(ch1, ch2, workspace, spectra;
               noiseratios=Dict(), lmax=0) where T

# Arguments:
- `ch1::Symbol`: spectrum type of first spectrum (i.e. :TT, :TE, :EE)
- `ch2::Symbol`: spectrum type of second spectrum (i.e. :TT, :TE, :EE)
- `workspace`: cache for working with covariances
- `spectra`: signal spectra

# Keywords
- `noiseratios::AbstractDict`: ratio of noise spectra to white noise
- `lmax=0`: maximum multipole moment for covariance matrix

# Returns:
- `SpectralArray{T,2}`: covariance matrix (0-indexed)
"""
function coupledcov(ch1::Symbol, ch2::Symbol, workspace::CovarianceWorkspace{T},
                    spectra::AbstractDict, noiseratios::AbstractDict=Dict();
                    lmin=0, lmax=nothing) where T

    lmax = isnothing(lmax) ? workspace.lmax : lmax
    𝐂 = spectralzeros(lmin:lmax, lmin:lmax)

    if length(noiseratios) == 0  # by default, do not rescale for noise
        identity_spectrum = spectralones(0:lmax)
        noiseratios = ConstantDict{SpectrumName,typeof(identity_spectrum)}(
            identity_spectrum)
    end

    if (ch1==:TT) && (ch2==:TT)
        return coupledcovTTTT!(𝐂, workspace, spectra, noiseratios)
    elseif (ch1==:EE) && (ch2==:EE)
        return coupledcovEEEE!(𝐂, workspace, spectra, noiseratios)
    elseif (ch1==:TE) && (ch2==:TE)
        return coupledcovTETE!(𝐂, workspace, spectra, noiseratios)
    elseif (ch1==:TT) && ( ch2==:TE)
        return coupledcovTTTE!(𝐂, workspace, spectra, noiseratios)
    elseif (ch1==:TT) && ( ch2==:EE)
        return coupledcovTTEE!(𝐂, workspace, spectra, noiseratios)
    elseif (ch1==:TE) && (ch2==:EE)
        return coupledcovTEEE!(𝐂, workspace, spectra, noiseratios)
    end
    print("$(ch1),$(ch2) not implemented")
end

function coupledcovTTTT!(𝐂::SpectralArray, workspace::CovarianceWorkspace,
                         spectra, noiseratios)

    @assert axes(𝐂, 1) == axes(𝐂, 2)
    i, j, p, q = workspace.field_names

    r_ℓ_ip = noiseratios[:TT, i, p]
    r_ℓ_jq = noiseratios[:TT, j, q]
    r_ℓ_iq = noiseratios[:TT, i, q]
    r_ℓ_jp = noiseratios[:TT, j, p]

    loop_covTTTT!(𝐂,
        spectra[:TT,i,p], spectra[:TT,j,q], spectra[:TT,i,q], spectra[:TT,j,p],
        r_ℓ_ip, r_ℓ_jq, r_ℓ_iq, r_ℓ_jp,
        window_function_W!(workspace, :∅∅, :∅∅, i, p, :TT, j, q, :TT),
        window_function_W!(workspace, :∅∅, :∅∅, i, q, :TT, j, p, :TT),
        window_function_W!(workspace, :∅∅, :TT, i, p, :TT, j, q, :TT),
        window_function_W!(workspace, :∅∅, :TT, j, q, :TT, i, p, :TT),
        window_function_W!(workspace, :∅∅, :TT, i, q, :TT, j, p, :TT),
        window_function_W!(workspace, :∅∅, :TT, j, p, :TT, i, q, :TT),
        window_function_W!(workspace, :TT, :TT, i, p, :TT, j, q, :TT),
        window_function_W!(workspace, :TT, :TT, i, q, :TT, j, p, :TT))

    return 𝐂
end

# inner loop
function loop_covTTTT!(𝐂::SpectralArray{T,2},
                       TTip::SpectralVector{T}, TTjq::SpectralVector{T},
                       TTiq::SpectralVector{T}, TTjp::SpectralVector{T},
                       r_ℓ_ip::SpectralVector{T}, r_ℓ_jq::SpectralVector{T},
                       r_ℓ_iq::SpectralVector{T}, r_ℓ_jp::SpectralVector{T},
                       W1, W2, W3, W4, W5, W6, W7, W8) where {T}

    lmin, lmax = first(axes(𝐂, 1)), last(axes(𝐂, 1))
    thread_buffers = get_thread_buffers(T, 2 * lmax + 1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w = WignerF(T, ℓ₁, ℓ₂, 0, 0)  # set up the wigner recurrence
            buffer_view = uview(buffer, 1:length(w.nₘᵢₙ:w.nₘₐₓ))  # preallocated buffer
            w3j² = WignerSymbolVector(buffer_view, w.nₘᵢₙ:w.nₘₐₓ)
            wigner3j_f!(w, w3j²)  # deposit symbols into buffer
            w3j².symbols .= w3j².symbols .^ 2  # square the symbols
            𝐂[ℓ₁, ℓ₂] = (
                sqrt(TTip[ℓ₁] * TTip[ℓ₂] * TTjq[ℓ₁] * TTjq[ℓ₂]) * Ξ_TT(W1, w3j², ℓ₁, ℓ₂) +
                sqrt(TTiq[ℓ₁] * TTiq[ℓ₂] * TTjp[ℓ₁] * TTjp[ℓ₂]) * Ξ_TT(W2, w3j², ℓ₁, ℓ₂) +
                sqrt(TTip[ℓ₁] * TTip[ℓ₂]) * Ξ_TT(W3, w3j², ℓ₁, ℓ₂) * r_ℓ_jq[ℓ₁] * r_ℓ_jq[ℓ₂] +
                sqrt(TTjq[ℓ₁] * TTjq[ℓ₂]) * Ξ_TT(W4, w3j², ℓ₁, ℓ₂) * r_ℓ_ip[ℓ₁] * r_ℓ_ip[ℓ₂] +
                sqrt(TTiq[ℓ₁] * TTiq[ℓ₂]) * Ξ_TT(W5, w3j², ℓ₁, ℓ₂) * r_ℓ_jp[ℓ₁] * r_ℓ_jp[ℓ₂]  +
                sqrt(TTjp[ℓ₁] * TTjp[ℓ₂]) * Ξ_TT(W6, w3j², ℓ₁, ℓ₂) * r_ℓ_iq[ℓ₁] * r_ℓ_iq[ℓ₂]  +
                Ξ_TT(W7, w3j², ℓ₁, ℓ₂) * r_ℓ_ip[ℓ₁] * r_ℓ_jq[ℓ₁] * r_ℓ_ip[ℓ₂] * r_ℓ_jq[ℓ₂] +
                Ξ_TT(W8, w3j², ℓ₁, ℓ₂) * r_ℓ_iq[ℓ₁] * r_ℓ_jp[ℓ₁] * r_ℓ_iq[ℓ₂] * r_ℓ_jp[ℓ₂])
            𝐂[ℓ₂, ℓ₁] = 𝐂[ℓ₁, ℓ₂]
        end
    end
end

function coupledcovEEEE!(𝐂::SpectralArray, workspace::CovarianceWorkspace, spectra,
                         noiseratios)

    @assert axes(𝐂, 1) == axes(𝐂, 2)
    i, j, p, q = workspace.field_names

    r_ℓ_ip = noiseratios[:EE, i, p]
    r_ℓ_jq = noiseratios[:EE, j, q]
    r_ℓ_iq = noiseratios[:EE, i, q]
    r_ℓ_jp = noiseratios[:EE, j, p]

    loop_covEEEE!(𝐂,
        spectra[:EE,i,p], spectra[:EE,j,q], spectra[:EE,i,q], spectra[:EE,j,p],
        r_ℓ_ip, r_ℓ_jq, r_ℓ_iq, r_ℓ_jp,
        window_function_W!(workspace, :∅∅, :∅∅, i, p, :PP, j, q, :PP),
        window_function_W!(workspace, :∅∅, :∅∅, i, q, :PP, j, p, :PP),
        window_function_W!(workspace, :∅∅, :PP, i, p, :PP, j, q, :PP),
        window_function_W!(workspace, :∅∅, :PP, j, q, :PP, i, p, :PP),
        window_function_W!(workspace, :∅∅, :PP, i, q, :PP, j, p, :PP),
        window_function_W!(workspace, :∅∅, :PP, j, p, :PP, i, q, :PP),
        window_function_W!(workspace, :PP, :PP, i, p, :PP, j, q, :PP),
        window_function_W!(workspace, :PP, :PP, i, q, :PP, j, p, :PP))

    return 𝐂
end

# inner loop
function loop_covEEEE!(𝐂::SpectralArray{T,2},
                       EEip::SpectralVector{T}, EEjq::SpectralVector{T},
                       EEiq::SpectralVector{T}, EEjp::SpectralVector{T},
                       r_ℓ_ip::SpectralVector{T}, r_ℓ_jq::SpectralVector{T},
                       r_ℓ_iq::SpectralVector{T}, r_ℓ_jp::SpectralVector{T},
                       W1, W2, W3, W4, W5, W6, W7, W8) where {T}

    lmin, lmax = first(axes(𝐂, 1)), last(axes(𝐂, 1))
    thread_buffers = get_thread_buffers(T, 2 * lmax + 1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w = WignerF(T, ℓ₁, ℓ₂, -2, 2)  # set up the wigner recurrence
            buffer_view = uview(buffer, 1:length(w.nₘᵢₙ:w.nₘₐₓ))  # preallocated buffer
            w3j² = WignerSymbolVector(buffer_view, w.nₘᵢₙ:w.nₘₐₓ)
            wigner3j_f!(w, w3j²)  # deposit symbols into buffer
            w3j².symbols .= w3j².symbols .^ 2  # square the symbols
            𝐂[ℓ₁, ℓ₂] = (
                sqrt(EEip[ℓ₁] * EEip[ℓ₂] * EEjq[ℓ₁] * EEjq[ℓ₂]) * Ξ_EE(W1, w3j², ℓ₁, ℓ₂) +
                sqrt(EEiq[ℓ₁] * EEiq[ℓ₂] * EEjp[ℓ₁] * EEjp[ℓ₂]) * Ξ_EE(W2, w3j², ℓ₁, ℓ₂) +
                sqrt(EEip[ℓ₁] * EEip[ℓ₂]) * Ξ_EE(W3, w3j², ℓ₁, ℓ₂) * r_ℓ_jq[ℓ₁] * r_ℓ_jq[ℓ₂] +
                sqrt(EEjq[ℓ₁] * EEjq[ℓ₂]) * Ξ_EE(W4, w3j², ℓ₁, ℓ₂) * r_ℓ_ip[ℓ₁] * r_ℓ_ip[ℓ₂] +
                sqrt(EEiq[ℓ₁] * EEiq[ℓ₂]) * Ξ_EE(W5, w3j², ℓ₁, ℓ₂) * r_ℓ_jp[ℓ₁] * r_ℓ_jp[ℓ₂] +
                sqrt(EEjp[ℓ₁] * EEjp[ℓ₂]) * Ξ_EE(W6, w3j², ℓ₁, ℓ₂) * r_ℓ_iq[ℓ₁] * r_ℓ_iq[ℓ₂] +
                Ξ_EE(W7, w3j², ℓ₁, ℓ₂) * r_ℓ_ip[ℓ₁] * r_ℓ_jq[ℓ₁] * r_ℓ_ip[ℓ₂] * r_ℓ_jq[ℓ₂] +
                Ξ_EE(W8, w3j², ℓ₁, ℓ₂) * r_ℓ_iq[ℓ₁] * r_ℓ_jp[ℓ₁] * r_ℓ_iq[ℓ₂] * r_ℓ_jp[ℓ₂])
            𝐂[ℓ₂, ℓ₁] = 𝐂[ℓ₁, ℓ₂]
        end
    end
end

function coupledcovTTTE!(𝐂::SpectralArray, workspace::CovarianceWorkspace, spectra,
                                     noiseratios)

    @assert axes(𝐂, 1) == axes(𝐂, 2)
    i, j, p, q = workspace.field_names

    r_ℓ_ip = noiseratios[:TT, i, p]
    r_ℓ_jp = noiseratios[:TT, j, p]

    loop_covTTTE!(𝐂,
        spectra[:TT,i,p], spectra[:TT,j,p], spectra[:TE,i,q], spectra[:TE,j,q],
        r_ℓ_ip, r_ℓ_jp,
        window_function_W!(workspace, :∅∅, :∅∅, i, p, :TT, j, q, :TP),
        window_function_W!(workspace, :∅∅, :∅∅, i, q, :TP, j, p, :TT),
        window_function_W!(workspace, :∅∅, :TT, j, q, :TP, i, p, :TT),
        window_function_W!(workspace, :∅∅, :TT, i, q, :TP, j, p, :TT))

    return 𝐂
end

# inner loop
function loop_covTTTE!(𝐂::SpectralArray{T,2},
                       TTip::SpectralVector{T}, TTjp::SpectralVector{T},
                       TEiq::SpectralVector{T}, TEjq::SpectralVector{T},
                       r_ℓ_ip::SpectralVector{T}, r_ℓ_jp::SpectralVector{T},
                       W1, W2, W3, W4) where {T}

    lmin, lmax = first(axes(𝐂, 1)), last(axes(𝐂, 1))
    thread_buffers = get_thread_buffers(T, 2lmax + 1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w = WignerF(T, ℓ₁, ℓ₂, 0, 0)  # set up the wigner recurrence
            buffer_view = uview(buffer, 1:length(w.nₘᵢₙ:w.nₘₐₓ))  # preallocated buffer
            w3j² = WignerSymbolVector(buffer_view, w.nₘᵢₙ:w.nₘₐₓ)
            wigner3j_f!(w, w3j²)  # deposit symbols into buffer
            w3j².symbols .= w3j².symbols .^ 2  # square the symbols
            𝐂[ℓ₁, ℓ₂] = (
                sqrt(TTip[ℓ₁] * TTip[ℓ₂]) * (TEjq[ℓ₁] + TEjq[ℓ₂]) * Ξ_TT(W1, w3j², ℓ₁, ℓ₂) +
                sqrt(TTjp[ℓ₁] * TTjp[ℓ₂]) * (TEiq[ℓ₁] + TEiq[ℓ₂]) * Ξ_TT(W2, w3j², ℓ₁, ℓ₂) +
                (TEjq[ℓ₁] + TEjq[ℓ₂]) * Ξ_TT(W3, w3j², ℓ₁, ℓ₂)  * r_ℓ_ip[ℓ₁] * r_ℓ_ip[ℓ₂]  +
                (TEiq[ℓ₁] + TEiq[ℓ₂]) * Ξ_TT(W4, w3j², ℓ₁, ℓ₂)  * r_ℓ_jp[ℓ₁] * r_ℓ_jp[ℓ₂]
            ) / 2

            𝐂[ℓ₂, ℓ₁] = 𝐂[ℓ₁, ℓ₂]
        end
    end
end

function coupledcovTETE!(𝐂::SpectralArray, workspace::CovarianceWorkspace, spectra,
                                     noiseratios)

    @assert axes(𝐂, 1) == axes(𝐂, 2)
    i, j, p, q = workspace.field_names

    r_TT_ip = noiseratios[:TT, i, p]
    r_PP_jq = noiseratios[:EE, j, q]

    loop_covTETE!(𝐂,
        spectra[:TT,i,p], spectra[:EE,j,q], spectra[:TE,i,q], spectra[:TE,j,p],
        r_TT_ip, r_PP_jq,
        window_function_W!(workspace, :∅∅, :∅∅, i, p, :TT, j, q, :PP),
        window_function_W!(workspace, :∅∅, :∅∅, i, q, :TP, j, p, :PT),
        window_function_W!(workspace, :∅∅, :PP, i, p, :TT, j, q, :PP),
        window_function_W!(workspace, :∅∅, :TT, j, q, :PP, i, p, :TT),
        window_function_W!(workspace, :TT, :PP, i, p, :TT, j, q, :PP))

    return 𝐂
end

# inner loop
function loop_covTETE!(𝐂::SpectralArray{T,2},
                       TTip::SpectralVector{T}, EEjq::SpectralVector{T},
                       TEiq::SpectralVector{T}, TEjp::SpectralVector{T},
                       r_TT_ip::SpectralVector{T}, r_PP_jq::SpectralVector{T},
                       W1, W2, W3, W4, W5) where {T}

    lmin, lmax = first(axes(𝐂, 1)), last(axes(𝐂, 1))
    thread_buffers_0 = get_thread_buffers(T, 2*lmax+1)
    thread_buffers_2 = get_thread_buffers(T, 2*lmax+1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer0 = thread_buffers_0[Threads.threadid()]
        buffer2 = thread_buffers_2[Threads.threadid()]

        for ℓ₂ in ℓ₁:lmax
            w00 = WignerF(T, ℓ₁, ℓ₂, 0, 0)  # set up the wigner recurrence
            w22 = WignerF(T, ℓ₁, ℓ₂, -2, 2)  # set up the wigner recurrence
            buffer_view_0 = uview(buffer0, 1:(w00.nₘₐₓ - w00.nₘᵢₙ + 1))  # preallocated buffer
            buffer_view_2 = uview(buffer2, 1:(w22.nₘₐₓ - w22.nₘᵢₙ + 1))  # preallocated buffer
            w3j_00 = WignerSymbolVector(buffer_view_0, w00.nₘᵢₙ:w00.nₘₐₓ)
            w3j_22 = WignerSymbolVector(buffer_view_2, w22.nₘᵢₙ:w22.nₘₐₓ)
            wigner3j_f!(w00, w3j_00)  # deposit symbols into buffer
            wigner3j_f!(w22, w3j_22)  # deposit symbols into buffer

            # varied over ℓ₃
            w3j_00_22 = w3j_22  # buffer 2
            w3j_00_22.symbols .*= w3j_00.symbols   # buffer2 = (buffer 2) * (buffer 1)
            w3j_00_00 = w3j_00
            w3j_00_00.symbols .*= w3j_00.symbols # (buffer 1) = (buffer 1) * (buffer 1)

            𝐂[ℓ₁, ℓ₂] = (
                sqrt(TTip[ℓ₁] * TTip[ℓ₂] * EEjq[ℓ₁] * EEjq[ℓ₂]) * Ξ_TE(W1, w3j_00_22, ℓ₁, ℓ₂) +
                0.5 * (TEiq[ℓ₁] * TEjp[ℓ₂] + TEjp[ℓ₁] * TEiq[ℓ₂]) * Ξ_TT(W2, w3j_00_00, ℓ₁, ℓ₂) +
                sqrt(TTip[ℓ₁] * TTip[ℓ₂]) * Ξ_TE(W3, w3j_00_22, ℓ₁, ℓ₂) * r_PP_jq[ℓ₁] * r_PP_jq[ℓ₂] +
                sqrt(EEjq[ℓ₁] * EEjq[ℓ₂]) * Ξ_TE(W4, w3j_00_22, ℓ₁, ℓ₂) * r_TT_ip[ℓ₁] * r_TT_ip[ℓ₂] +
                Ξ_TE(W5, w3j_00_22, ℓ₁, ℓ₂) * r_TT_ip[ℓ₁] * r_TT_ip[ℓ₂] * r_PP_jq[ℓ₁] * r_PP_jq[ℓ₂])

            𝐂[ℓ₂, ℓ₁] = 𝐂[ℓ₁, ℓ₂]
        end
    end
end

function coupledcovTEEE!(𝐂::SpectralArray, workspace::CovarianceWorkspace, spectra,
                                     noiseratios; planck=true)

    @assert axes(𝐂, 1) == axes(𝐂, 2)
    i, j, p, q = workspace.field_names

    r_EE_jq = noiseratios[:EE, j, q]
    r_EE_jp = noiseratios[:EE, j, p]

    if planck
        loop_covTEEE_planck!(𝐂,
            spectra[:EE,j,q], spectra[:EE,j,p], spectra[:TE,i,p], spectra[:TE,i,q],
            r_EE_jq, r_EE_jp,
            window_function_W!(workspace, :∅∅, :∅∅, i, p, :TP, j, q, :PP),
            window_function_W!(workspace, :∅∅, :∅∅, i, q, :TP, j, p, :PP),
            window_function_W!(workspace, :∅∅, :PP, i, p, :TP, j, q, :PP),
            window_function_W!(workspace, :∅∅, :PP, i, q, :TP, j, p, :PP))
    else
        loop_covTEEE!(𝐂,
            spectra[:EE,j,q], spectra[:EE,j,p], spectra[:TE,i,p], spectra[:TE,i,q],
            r_EE_jq, r_EE_jp,
            window_function_W!(workspace, :∅∅, :∅∅, i, p, :TP, j, q, :PP),
            window_function_W!(workspace, :∅∅, :∅∅, i, q, :TP, j, p, :PP),
            window_function_W!(workspace, :∅∅, :PP, i, p, :TP, j, q, :PP),
            window_function_W!(workspace, :∅∅, :PP, i, q, :TP, j, p, :PP))
    end

    return 𝐂
end

# inner loop
function loop_covTEEE!(𝐂::SpectralArray{T,2},
                       EEjq::SpectralVector{T}, EEjp::SpectralVector{T},
                       TEip::SpectralVector{T}, TEiq::SpectralVector{T},
                       r_EE_jq::SpectralVector{T}, r_EE_jp::SpectralVector{T},
                       W1, W2, W3, W4) where {T}

    lmin, lmax = first(axes(𝐂, 1)), last(axes(𝐂, 1))
    thread_buffers_0 = get_thread_buffers(T, 2*lmax+1)
    thread_buffers_2 = get_thread_buffers(T, 2*lmax+1)
    @qthreads for ℓ₁ in lmin:lmax
        buffer0 = thread_buffers_0[Threads.threadid()]
        buffer2 = thread_buffers_2[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w00 = WignerF(T, ℓ₁, ℓ₂, 0, 0)  # set up the wigner recurrence
            w22 = WignerF(T, ℓ₁, ℓ₂, -2, 2)  # set up the wigner recurrence
            buffer_view_0 = uview(buffer0, 1:(w00.nₘₐₓ - w00.nₘᵢₙ + 1))  # preallocated buffer
            buffer_view_2 = uview(buffer2, 1:(w22.nₘₐₓ - w22.nₘᵢₙ + 1))  # preallocated buffer
            w3j_00 = WignerSymbolVector(buffer_view_0, w00.nₘᵢₙ:w00.nₘₐₓ)
            w3j_22 = WignerSymbolVector(buffer_view_2, w22.nₘᵢₙ:w22.nₘₐₓ)
            wigner3j_f!(w00, w3j_00)  # deposit symbols into buffer
            wigner3j_f!(w22, w3j_22)  # deposit symbols into buffer
            # varied over ℓ₃
            w3j² = w3j_22  # buffer 2
            w3j².symbols .*= w3j_00.symbols   # buffer2 = (buffer 2) * (buffer 1)

            𝐂[ℓ₁, ℓ₂] = (
                sqrt(EEjq[ℓ₁] * EEjq[ℓ₂]) * (TEip[ℓ₁] + TEip[ℓ₂]) * Ξ_EE(W1, w3j², ℓ₁, ℓ₂) +
                sqrt(EEjp[ℓ₁] * EEjp[ℓ₂]) * (TEiq[ℓ₁] + TEiq[ℓ₂]) * Ξ_EE(W2, w3j², ℓ₁, ℓ₂) +
                (TEip[ℓ₁] + TEip[ℓ₂]) * Ξ_EE(W3, w3j², ℓ₁, ℓ₂) * r_EE_jq[ℓ₁] * r_EE_jq[ℓ₂] +
                (TEiq[ℓ₁] + TEiq[ℓ₂]) * Ξ_EE(W4, w3j², ℓ₁, ℓ₂) * r_EE_jp[ℓ₁] * r_EE_jp[ℓ₂]
            ) / 2

            𝐂[ℓ₂, ℓ₁] = 𝐂[ℓ₁, ℓ₂]
        end
    end
end

# inner loop
function loop_covTEEE_planck!(𝐂::SpectralArray{T,2},
                              EEjq::SpectralVector{T}, EEjp::SpectralVector{T},
                              TEip::SpectralVector{T}, TEiq::SpectralVector{T},
                              r_EE_jq::SpectralVector{T}, r_EE_jp::SpectralVector{T},
                              W1, W2, W3, W4) where {T}

    lmin, lmax = first(axes(𝐂, 1)), last(axes(𝐂, 1))
    thread_buffers = get_thread_buffers(T, 2lmax + 1)
    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w = WignerF(T, ℓ₁, ℓ₂, -2, 2)  # set up the wigner recurrence
            buffer_view = uview(buffer, 1:length(w.nₘᵢₙ:w.nₘₐₓ))  # preallocated buffer
            w3j² = WignerSymbolVector(buffer_view, w.nₘᵢₙ:w.nₘₐₓ)
            wigner3j_f!(w, w3j²)  # deposit symbols into buffer
            w3j².symbols .= w3j².symbols .^ 2  # square the symbols
            𝐂[ℓ₁, ℓ₂] = (
                sqrt(EEjq[ℓ₁] * EEjq[ℓ₂]) * (TEip[ℓ₁] + TEip[ℓ₂]) * Ξ_EE(W1, w3j², ℓ₁, ℓ₂) +
                sqrt(EEjp[ℓ₁] * EEjp[ℓ₂]) * (TEiq[ℓ₁] + TEiq[ℓ₂]) * Ξ_EE(W2, w3j², ℓ₁, ℓ₂) +
                (TEip[ℓ₁] + TEip[ℓ₂]) * Ξ_EE(W3, w3j², ℓ₁, ℓ₂) * r_EE_jq[ℓ₁] * r_EE_jq[ℓ₂] +
                (TEiq[ℓ₁] + TEiq[ℓ₂]) * Ξ_EE(W4, w3j², ℓ₁, ℓ₂) * r_EE_jp[ℓ₁] * r_EE_jp[ℓ₂]
            ) / 2

            𝐂[ℓ₂, ℓ₁] = 𝐂[ℓ₁, ℓ₂]
        end
    end
end

function coupledcovTTEE!(𝐂::SpectralArray{T}, workspace::CovarianceWorkspace, spectra,
                         noiseratios) where T

    @assert axes(𝐂, 1) == axes(𝐂, 2)
    i, j, p, q = workspace.field_names

    loop_covTTEE!(𝐂,
        spectra[:TE,i,p], spectra[:TE,i,q], spectra[:TE,j,q], spectra[:TE,j,p],
        window_function_W!(workspace, :∅∅, :∅∅, i, p, :TP, j, q, :TP),
        window_function_W!(workspace, :∅∅, :∅∅, i, q, :TP, j, p, :TP),
    )

    return 𝐂
end

# inner loop
function loop_covTTEE!(𝐂::SpectralArray{T,2},
                       TEip::SpectralVector{T}, TEiq::SpectralVector{T},
                       TEjq::SpectralVector{T}, TEjp::SpectralVector{T},
                       W1, W2) where {T}

    lmin, lmax = first(axes(𝐂, 1)), last(axes(𝐂, 1))
    thread_buffers = get_thread_buffers(T, 2 * lmax + 1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w = WignerF(T, ℓ₁, ℓ₂, 0, 0)  # set up the wigner recurrence
            buffer_view = uview(buffer, 1:length(w.nₘᵢₙ:w.nₘₐₓ))  # preallocated buffer
            w3j² = WignerSymbolVector(buffer_view, w.nₘᵢₙ:w.nₘₐₓ)
            wigner3j_f!(w, w3j²)  # deposit symbols into buffer
            w3j².symbols .= w3j².symbols .^ 2  # square the symbols
            𝐂[ℓ₁, ℓ₂] = (
                (TEip[ℓ₁] * TEjq[ℓ₂] + TEjq[ℓ₁] * TEip[ℓ₂]) * Ξ_TT(W1, w3j², ℓ₁, ℓ₂) +
                (TEiq[ℓ₁] * TEjp[ℓ₂] + TEjp[ℓ₁] * TEiq[ℓ₂]) * Ξ_TT(W2, w3j², ℓ₁, ℓ₂)
            ) / 2

            𝐂[ℓ₂, ℓ₁] = 𝐂[ℓ₁, ℓ₂]
        end
    end
end

ext: jl


text:
function F(E, p, q, λᵢ, λₛ, Eₑ, Eₛ, Eᵢ, energy, add, j, i)
    Eₘ = maximum(E)
    ϵ = 0.0
    N, M = size(E)
    if add == true
        push!(Eₑ, 0.0)
        push!(Eₛ, 0.0)
        push!(Eᵢ, 0.0)
    end

    for k = (i-1):(i+1)
        for r = (j-1):(j+1)
            if k != N && r != M && k != 1 && r != 1
                # @show r,k
                # Optimize this!!!
                s = CartesianIndex(r,k)
                s1 = CartesianIndex(r+1,k)
                s2 = CartesianIndex(r,k+1)
                s3 = CartesianIndex(r-1,k)
                s4 = CartesianIndex(r,k-1)
                s5 = CartesianIndex(r+1,k-1)
                s6 = CartesianIndex(r-1,k+1)
                ϵ₀ = (Eₘ/sqrt(1+p[s]^2+q[s]^2)-E[s])^2
                ϵₛ = λₛ*((p[s1] - p[s])^2 + (q[s1] - q[s])^2 + (p[s2] - p[s])^2 + (q[s2] - q[s])^2 + (p[s3] - p[s])^2 + (q[s3] - q[s])^2 + (p[s4] - p[s])^2 + (q[s4] - q[s])^2)
                ϵᵢ = λᵢ*(((p[s1] - p[s]) - (q[s2] - q[s]))^2 + ((p[s] - p[s3]) - (q[s6] - q[s3]))^2 + ((p[s5] - p[s4]) - (q[s] - q[s4]))^2)
                ϵ += ϵ₀ + ϵᵢ + ϵₛ
                # @show ϵ₀, ϵᵢ, ϵₛ
                energy[i,j] = ϵ₀ + ϵᵢ + ϵₛ
                # energy[i,j] = ϵ₀
                if add == true
                    Eₑ[end] += ϵ₀
                    Eₛ[end] += ϵₛ
                    Eᵢ[end] += ϵᵢ
                end
                # energy[i,j] = ϵ₀
                # energy[i,j] = ϵᵢ
                # energy[i,j] = ϵₛ
            end
        end
    end
    return ϵ
end

function P(ϵ, ϵₙ, ρ, T)
    R = 1.0
    if ϵ < ϵₙ
        R = exp(-(ϵₙ-ϵ)/T)
        # @show R
        R = R*(ρ^(1/T-1))
        # R = 0.0
    end
    return R
end

function (algorithm::SimulatedAnnealing)(img::AbstractArray)

    pin = algorithm.pin
    qin = algorithm.qin
    T₀ = algorithm.T₀
    α = algorithm.α
    integration_factor = algorithm.integration_factor
    smoothness = algorithm.smoothness
    max_iter = algorithm.max_iter
    debug = algorithm.debug

    ϵₑ = zeros(Float64, 0)
    ϵᵢ = zeros(Float64, 0)
    ϵₛ = zeros(Float64, 0)

    # Setup initial variables
    E = Float64.(copy(img)).*255
    Eₘ = maximum(E)
    p = copy(pin)
    q = copy(qin)
    pmax = maximum(abs.(p))
    qmax = maximum(abs.(q))
    ρₘ = sqrt(maximum(p.^2) + maximum(q.^2))
    pₙ = copy(pin)
    qₙ = copy(qin)
    p₀ = copy(pin)
    q₀ = copy(qin)
    λᵢ = integration_factor
    λₛ = smoothness
    k = 1
    Tₖ = T₀
    energy = zeros(size(E))
    ϵ = F(E, p, q, λᵢ, λₛ, ϵₑ, ϵₛ, ϵᵢ, energy, true, 20, 20)
    ϵ₀ = F(E, p, q, λᵢ, λₛ, ϵₑ, ϵₛ, ϵᵢ, energy, false, 20, 20)
    N, M = size(E)
    r2 = 1.0:1:32
    if debug == true
        Z = convert_gradient(Horn(),p,q)
        img = generate_surface(p, q, 1, [0,0,1], radius = 50, img_size=32, noise = 0, background = false)
        display(AbstractPlotting.PlotDisplay(), Makie.vbox(surface(r2, r2, Z), image(img'), heatmap(energy), arrows(r2, r2, pin[end:-1:1,:]', qin[end:-1:1,:]')))
    end

    while k < max_iter
        for i in CartesianIndices(E)
            if i[1] != N && i[2] != M && i[1] != 1 && i[2] != 1
                ϵ = F(E, p, q, λᵢ, λₛ, ϵₑ, ϵᵢ, ϵₛ, energy, false, i[1], i[2])
                θ = rand(0:10^-6:π/2)
                Φ = rand(0:10^-6:2*π)
                # ρₙ = tan(θ)
                ρₙ = sqrt(((Eₘ^2)/(E[i]^2)) - 1)
                if ρₙ > 2*ρₘ
                    ρₙ = 2*ρₘ
                end
                pₙ[i] = ρₙ*cos(Φ)
                qₙ[i] = ρₙ*sin(Φ)
                ϵₙ = F(E, pₙ, qₙ, λᵢ, λₛ, ϵₑ, ϵₛ, ϵᵢ, energy, false, i[1], i[2])
                ρᵢ = sqrt(p[i]^2 + q[i]^2)
                # ρₙ = sqrt(pₙ[i]^2 + qₙ[i]^2)
                ρ = (ρₙ * (1 + ρᵢ^2)^(3/2)) / (ρᵢ * (1 + ρₙ^2)^(3/2))
                R = P(ϵ, ϵₙ, ρ, Tₖ)
                prob = rand(0:0.00001:1)
                if prob < R
                    copyto!(p,pₙ)
                    copyto!(q,qₙ)
                else
                    copyto!(pₙ,p)
                    copyto!(qₙ,q)
                end
            end
        end
        ϵ = F(E, pₙ, qₙ, λᵢ, λₛ, ϵₑ, ϵₛ, ϵᵢ, energy, true, 20, 20)
        if k % 1000 == 0
            @show k
        end
        if k % 200 == 0 && debug == true
            r2 = 1.0:1:32
            Z = convert_gradient(Horn(),p,q)
            img = generate_surface(p, q, 1, [0,0,1], radius = 50, img_size=32, noise = 0, background = false)
            display(Makie.vbox(surface(r2, r2, Z), image(img'), heatmap((energy./maximum(abs.(energy)))), arrows(r2, r2, -p', q')))
            display(AbstractPlotting.PlotDisplay(), Makie.vbox(surface(r2, r2, Z), image(img'), heatmap((energy./maximum(abs.(energy)))), arrows(r2, r2, p[end:-1:1,:]', q[end:-1:1,:]')))
            # return energy, 0
        end
        if k % 100 == 0 && debug == true
            @show ϵ, Tₖ, k#, sum(energy)
        end
        k += 1
        Tₖ = (α^k)*T₀
    end
    if debug == true
        @show ϵ₀
    end
    return p,q,ϵₑ,ϵₛ,ϵᵢ
end


text:
using PyPlot, Statistics, NbodyGradient, LinearAlgebra

export State

GNEWT = NbodyGradient.GNEWT

function compute_energy(m::Array{T,1},x::Array{T,2},v::Array{T,2},n::Int64) where {T <: Real}
  KE = 0.0
  for j=1:n
    KE += 0.5*m[j]*(v[1,j]^2+v[2,j]^2+v[3,j]^2)
  end
  PE = 0.0
  for j=1:n-1
    for k=j+1:n
       PE += -GNEWT*m[j]*m[k]/norm(x[:,j] .- x[:,k])
    end
  end
  ang_mom = zeros(3)
  for j=1:n
    ang_mom .+= m[j]*cross(x[:,j],v[:,j])
  end
  return KE,PE,ang_mom
end

n=5
xout = zeros(3,n)
# Positions at time September 5, 1994 at 0h00 in days (from Hairer, Lubich & Wanner
# 2006, Geometric Numerical Integration, 2nd Edition, Springer, pp. 13-14):

xout .= transpose([0.0 0.0 0.0;
  -3.5023653  -3.8169847  -1.5507963;
  9.0755314  -3.0458353  -1.6483708;
  8.3101420  -16.2901086  -7.2521278;
  11.4707666  -25.7294829  -10.8169456])

vout = zeros(3,n)
vout .= transpose([0.0 0.0 0.0;
  0.00565429  -0.00412490  -0.00190589;
  .00168318  .00483525  .00192462;
  .00354178  .00137102  .00055029;
  .00288930  .00114527  .00039677])

# Specify masses, including terrestrial planets in the Sun:
m = [1.00000597682,0.000954786104043,0.000285583733151,
         0.0000437273164546,0.0000517759138449] #,6.58086572e-9];

# Compute the center-of-mass:
vcm = zeros(3);
xcm = zeros(3);
for j=1:n
    vcm .+= m[j]*vout[:,j];
    xcm .+= m[j]*xout[:,j];
end
vcm ./= sum(m);
xcm ./= sum(m);
# Adjust so CoM is stationary
for j=1:n
    vout[:,j] .-= vcm[:];
    xout[:,j] .-= xcm[:];
end

#=
struct CartesianElements{T} <: NbodyGradient.InitialConditions{T}
    x::Matrix{T}
    v::Matrix{T}
    m::Vector{T}
    nbody::Int64
end

function NbodyGradient.State(ic::NbodyGradient.InitialConditions{T}) where T<:AbstractFloat
    n = ic.nbody
    x = copy(ic.x)
    v = copy(ic.v)
    jac_init = zeros(7*n,7*n)
    xerror = zeros(T,size(x))
    verror = zeros(T,size(v))
    jac_step = Matrix{T}(I,7*n,7*n)
    dqdt = zeros(T,7*n)
    dqdt_error = zeros(T,size(dqdt))
    jac_error = zeros(T,size(jac_step))
    pair = zeros(Bool,n,n)

    rij = zeros(T,3)
    a = zeros(T,3,n)
    aij = zeros(T,3)
    x0 = zeros(T,3)
    v0 = zeros(T,3)
    input = zeros(T,8)
    delxv = zeros(T,6)
    rtmp = zeros(T,3)
    return State(x,v,[0.0],copy(ic.m),jac_step,dqdt,jac_init,xerror,verror,dqdt_error,jac_error,n,
    pair,rij,a,aij,x0,v0,input,delxv,rtmp)
end
=#

ic = CartesianIC(xout,vout,m,n,0.0);

# Now, integrate this forward in time:
xerror = zeros(3,n); verror = zeros(3,n);
hgrid = [1.5625,3.125,6.25,12.5,25.0,50.0,100.0,200.0]

power = 20
nstep = 2^power
grad = false
ngrid = length(hgrid)

# 4e6 time steps ~ 2 Myr (takes about 9 minutes to run)
# I've reducd that to ~400,000 to take ~1 minute.

# Set up array to save the state as a function of time:
# Save the potential & kinetic energy, as well as angular momentum:
egrid = zeros(nstep); ang_mom = zeros(3,nstep)
sig_energy = zeros(ngrid); sig_ang_mom = zeros(3,ngrid)
mean_energy = zeros(ngrid); mean_ang_mom = zeros(3,ngrid)
# Time the integration:
for j=1:length(hgrid)
  h = hgrid[j]
  s = State(ic)
  # Specify Kepler step only between planets & star:
  s.pair .= true
  for i=1:s.n
    s.pair[1,i] = false
    s.pair[i,1] = false
  end
  if grad; d = Derivatives(T,s.n); end

  tstart = time()
  # Carry out the integration:
  for i=1:nstep
       if grad
      ahl21!(s,d,h)
    else
      ahl21!(s,h)
    end
    KE_step,PE_step,ang_mom_step=compute_energy(s.m,s.x,s.v,n)
    egrid[i] = KE_step+PE_step
    ang_mom[:,i] = ang_mom_step
  end
  mean_energy[j] = mean(egrid)
  sig_energy[j] = std(egrid)
  mean_ang_mom[1,j] = mean(ang_mom[1,:])
  mean_ang_mom[2,j] = mean(ang_mom[2,:])
  mean_ang_mom[3,j] = mean(ang_mom[3,:])
  sig_ang_mom[1,j] = std(ang_mom[1,:])
  sig_ang_mom[2,j] = std(ang_mom[2,:])
  sig_ang_mom[3,j] = std(ang_mom[3,:])
  telapse = time()- tstart
  println(j," ",telapse," ",sig_energy[j]," ",sig_ang_mom[:,j])
end

# These integrations have already been carried out:  so,
# just plotting the results:

clf()
tstep = hgrid
loglog(tstep,sig_energy,"o",label="Measured scatter in energy")
loglog(tstep,1.403145134013028e-16 .* (tstep ./ tstep[8]).^4,label="Quartic scaling with time step")
xlabel("Time step [days]")
ylabel(L"Energy error [$M_\odot$ AU$^2$ day$^{-2}$]")
loglog([1.5625,200],abs(mean_energy[end])*2.2e-16*[1,1],linestyle=":",label="Double precision floor")
legend()

tight_layout()
savefig("Energy_error_vs_timestep_mixed.pdf",bbox_inches="tight")

clf()

loglog(tstep,sig_ang_mom[1,:],label=L"$L_x$")
loglog(tstep,sig_ang_mom[2,:],label=L"$L_y$")
loglog(tstep,sig_ang_mom[3,:],label=L"$L_z$")
loglog(tstep,mean_ang_mom[1,end]*2.2e-16*ones(8),linestyle=":",color="blue")
loglog(tstep,mean_ang_mom[2,end]*2.2e-16*ones(8),linestyle=":",color="orange")
loglog(tstep,mean_ang_mom[3,end]*2.2e-16*ones(8),linestyle=":",color="green")
loglog(tstep,1e-21 .* tstep, color="k",linestyle="--")
xlabel("Time step [days]")
ylabel(L"Angular momentum error [$M_\odot$ AU$^2$ day$^{-1}$]")
legend()

tight_layout()
savefig("Angular_momentum_error_vs_timestep_mixed.pdf",bbox_inches="tight")

max_stars_repo_name: langfzac/nbg-papers

max_issues_repo_name: langfzac/nbg-papers

max_forks_repo_name: langfzac/nbg-papers


text:
module LoggerFormatTest

using MiniLoggers
using ReTest

conts(s, sub) = match(Regex(sub), s) !== nothing

@testset "basic format" begin
    io = IOBuffer()
    logger = MiniLogger(io = io, minlevel = MiniLoggers.Debug,
                        format = "{datetime}:{level}:{module}:{basename}:{filepath}:{line}:{group}:{module}:{id}:{message}")
    with_logger(logger) do
        @info "foo"
        s = String(take!(io))
        @test !conts(s, "datetime")
        @test !conts(s, "level")
        @test !conts(s, "module")
        @test !conts(s, "basename")
        @test !conts(s, "filepath")
        @test !conts(s, "line")
        @test !conts(s, "group")
        @test !conts(s, "module")
        @test !conts(s, "id")
        @test conts(s, "foo")
        @test conts(s, "Info")
    end
end

end # module


text:
#warmup feasibility
function find_best_itemᵥ(
    NH::Neighbourhood,
    v::Int64,
    iu::ItemUse,
    fs::FriendSets,
    constraints::Constraint,
    x_forced0ₜ::Vector{Bool},
    n_items::Int64,
    n_fs::Int64,
    ol_maxₜ::Vector{Float64},
    to_apply::Vector{Bool},
)
    f₀, infeas₀ = Inf, Inf
    T = size(NH.x, 2)
    x₀ = copy(NH.x)
    xₜ = x₀[:, v]
    idxₜ₂ = findall(iszero.(xₜ)) #i₂ = 1, ..., I
    idxₜ₂ = Random.shuffle!(idxₜ₂) #removed
    i⁺ = copy(idxₜ₂[1])
    if to_apply[1]
        iu_max₀ = sum(NH.x, dims = 2)[:, 1] - iu.max
    else
        iu_max₀ = 0
    end
    if to_apply[2]
        iu_min₀ = -sum(NH.x, dims = 2)[:, 1] + iu.min
    else
        iu_min₀ = 0
    end
    iu_max⁺ = 0
    iu_min⁺ = 0
    ol₀ₜ = 0
    for i₂ in idxₜ₂
        if x_forced0ₜ[i₂]
            x₁, infeas₁, iu_max, iu_min =
                copy(x₀), copy(infeas₀), copy(iu_max₀), copy(iu_min₀)
            x₁[i₂, v] = one(Float64)
            if to_apply[1]
                iu_max[i₂] = iu_max[i₂] + 1
                iu_max = sum_pos(iu_max)
            end
            if to_apply[2]
                iu_min[i₂] = iu_min[i₂] - 1
                iu_min = sum_pos(iu_min)
            end
            xₜ = copy(x₁[:, v])
            infeas₁, x_Iₜ = check_feas(fs, constraints, xₜ, n_fs, n_items)
            if to_apply[3]
                ol = eval_overlapᵥ(xₜ, x₁, fs.counts, ol_maxₜ, v)
            else
                ol = 0
            end
            f₁ = infeas₁ + iu_min + iu_max + ol
            if (f₁ < f₀)
                f₀, infeas₀ = copy(f₁), copy(infeas₁)
                iu_max⁺ = copy(iu_max)
                iu_min⁺ = copy(iu_min)
                ol₀ₜ = copy(ol)
                i⁺ = copy(i₂)
            end
        end
    end
    NH.ol[v] = copy(ol₀ₜ)
    NH.infeas[v] = copy(infeas₀)
    NH.iu = iu_min⁺ + iu_max⁺
    NH.x[i⁺, v] = one(Float64)
    return NH::Neighbourhood
end


text:
using MiscUtils
using Base.Test

@def makesomevars begin
  x = 1
end
@makesomevars
@test x == 1


text:
using COSMO, LinearAlgebra, JuMP, Random, SparseArrays

# Define problem matrices
n = 9
m = 2

A1 = [-4.0 0.0 -2.0 0.0 0.0 -1.0 0.0 0.0 0.0; 0.0 -3.0 -1.0 0.0 0.0 0.0 0.0 0.0 0.0; -2.0 -1.0 -2.0 0.0 0.0 5.0 4.0 -4.0 0.0; 0.0 0.0 0.0 -4.0 -5.0 0.0 0.0 3.0 0.0; 0.0 0.0 0.0 -5.0 4.0 0.0 0.0 2.0 0.0; -1.0 0.0 5.0 0.0 0.0 5.0 -4.0 -4.0 -5.0; 0.0 0.0 4.0 0.0 0.0 -4.0 -1.0 -1.0 -3.0; 0.0 0.0 -4.0 3.0 2.0 -4.0 -1.0 2.0 -2.0; 0.0 0.0 0.0 0.0 0.0 -5.0 -3.0 -2.0 -3.0];
A2 = [-5.0 0.0 3.0 0.0 0.0 -2.0 0.0 0.0 0.0; 0.0 -3.0 -5.0 0.0 0.0 0.0 0.0 0.0 0.0; 3.0 -5.0 3.0 0.0 0.0 5.0 -4.0 -5.0 0.0; 0.0 0.0 0.0 3.0 2.0 0.0 0.0 -2.0 0.0; 0.0 0.0 0.0 2.0 4.0 0.0 0.0 -3.0 0.0; -2.0 0.0 5.0 0.0 0.0 1.0 -5.0 -2.0 -4.0; 0.0 0.0 -4.0 0.0 0.0 -5.0 -2.0 -3.0 3.0; 0.0 0.0 -5.0 -2.0 -3.0 -2.0 -3.0 5.0 3.0; 0.0 0.0 0.0 0.0 0.0 -4.0 3.0 3.0 -4.0];
B1 = [-0.11477375644968069 0.0 6.739182490600791 0.0 0.0 -1.2185593245043502 0.0 0.0 0.0; 0.0 1.2827680528587497 -5.136452036888789 0.0 0.0 0.0 0.0 0.0 0.0; 6.739182490600791 -5.136452036888789 7.344770673489607 0.0 0.0 -0.2224400187044442 -10.505300166831221 -1.2627361794562273 0.0; 0.0 0.0 0.0 10.327710040060499 8.91534585379813 0.0 0.0 -6.525873789637007 0.0; 0.0 0.0 0.0 8.91534585379813 0.8370459338528677 0.0 0.0 -6.210900615408826 0.0; -1.2185593245043502 0.0 -0.2224400187044442 0.0 0.0 -3.8185953011245024 -0.994033914192722 2.8156077981712997 1.4524716674219218; 0.0 0.0 -10.505300166831221 0.0 0.0 -0.994033914192722 0.029162208619863517 -2.8123790276830745 7.663416446183705; 0.0 0.0 -1.2627361794562273 -6.525873789637007 -6.210900615408826 2.8156077981712997 -2.8123790276830745 4.71893305728242 6.322431630550857; 0.0 0.0 0.0 0.0 0.0 1.4524716674219218 7.663416446183705 6.322431630550857 0.5026094532322212];
c = [-0.21052661285686525, -1.263324575834677];

# Solve with JuMP + COSMO

# Solve without clique merging
model = JuMP.Model(optimizer_with_attributes(COSMO.Optimizer, "merge_strategy" => COSMO.NoMerge, "complete_dual" => true));
@variable(model, x[1:m]);
@objective(model, Min, c' * x )
@constraint(model, Symmetric(B - A1  .* x[1] - A2 .* x[2] )  in JuMP.PSDCone());
JuMP.optimize!(model)

# Solve with a clique graph based strategy
model = JuMP.Model(optimizer_with_attributes(COSMO.Optimizer, "merge_strategy" => COSMO.CliqueGraphMerge, "complete_dual" => true));
@variable(model, x[1:m]);
@objective(model, Min, c' * x )
@constraint(model, Symmetric(B - A1  .* x[1] - A2 .* x[2] )  in JuMP.PSDCone());
JuMP.optimize!(model)

# Solve with a parent-child merge strategy
pc_strategy = with_options(COSMO.ParentChildMerge, t_size = 3, t_fill = 3)
model = JuMP.Model(optimizer_with_attributes(COSMO.Optimizer, "merge_strategy" => pc_strategy, "complete_dual" => true));
@variable(model, x[1:m]);
@objective(model, Min, c' * x )
@constraint(model, Symmetric(B - A1  .* x[1] - A2 .* x[2] )  in JuMP.PSDCone());
JuMP.optimize!(model)


text:
function trailingmeasure!(F::TF, crsp::DataFrame, calendarinterval::DatePeriod,
  minfraction::Float64;
  Ftarget::Symbol = error("Ftarget is required"),
  Fmeasure::Symbol = error("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing,)::Nothing where TF

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{String} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (error("$Fnew already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || error("dataframe not sorted by [:permno, :date]")

  #find the minimum date in the entire dataset
  alldates = unique(crsp.date) |> sort!
  initialminpoints = round(
    searchsortedfirst(alldates, alldates[1] + calendarinterval)*minfraction,RoundUp) |> Int

  #now compute the min number of points for each date
  function getminpoints(i)
    (alldates[i] ≤ alldates[1] + calendarinterval) && return initialminpoints
    lowerdateindex = searchsortedfirst(alldates, alldates[i] + Day(1) - calendarinterval)
    return round((i - lowerdateindex + 1) * minfraction,RoundUp) |> Int
  end
  minpointindex = Dict{Date, Int}(
    alldates[i] => getminpoints(i) for i ∈ initialminpoints:length(alldates))
  minstartdate = alldates[initialminpoints]
  @assert (alldates[alldates .≥ minstartdate] .|> (d)->haskey(minpointindex, d)) |> all

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #when we are looking at eom values only
  @inline getj₀(d::Date, dates::AbstractVector{Date}) = searchsortedfirst(
    dates, d + Day(1) - calendarinterval)

  #the below is a low-effort performance booster- don't check arrays below this
  absminpoints::Int = minimum(values(minpointindex) |> collect)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = size(scrsp, 1)
    (absminpoints < nrows) || continue #if there are enough  entries to compute the trailing returns

    #a hack to only compute calcualtions at end of month
    for j ∈ absminpoints:nrows
      applyconditional && (!scrsp[j,Fconditional]) && continue #if we are only computing end of month values
      currentdate = scrsp[j, :date]
      (currentdate < minstartdate) && continue
      j₀::Int = getj₀(currentdate, scrsp.date)
      N::Int = j-j₀ +1
      #if the returns are not stale and complete enough, record them
      (N < minpointindex[currentdate]) && continue
      scrsp[j, Fmeasure] = F(scrsp[j₀:j , Ftarget])
    end
  end

  return nothing
end

#applies a function to a vector of trailing variables
#this version uses a rolling calendar interval with a set number
#of minimum points
function trailingmeasure!(F::Function, crsp::DataFrame, calendarinterval::DatePeriod,
  minpoints::Int;
  Ftarget::Symbol = error("Ftarget is required"),
  Fmeasure::Symbol = error("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing)::Nothing

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{String} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (error("$Fmeasure already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || error("dataframe not sorted by [:permno, :date]")

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #when we are looking at eom values only
  @inline getj₀(d::Date, dates::AbstractVector{Date}) = searchsortedfirst(
    dates, d + Day(1) - calendarinterval)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = size(scrsp, 1)
    (minpoints ≤ nrows) || continue #if there are enough  entries to compute the trailing returns

    #a hack to only compute calcualtions at end of month
    for j ∈ minpoints:nrows
      applyconditional && (!scrsp[j,Fconditional]) && continue #if we are only computing end of month values
      j₀::Int = getj₀(scrsp[j, :date], scrsp.date)
      N::Int = j-j₀ +1
      #if the returns are not stale and complete enough, record them
      (N < minpoints) && continue
      scrsp[j, Fmeasure] = F(scrsp[j₀:j , Ftarget])
    end
  end

  return nothing
end

#trailing index measure
function stricttrailingmeasure!(F::TF, crsp::DataFrame, calendarinterval::DatePeriod;
  Ftarget::Symbol = error("Ftarget is required"),
  Fmeasure::Symbol = error("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing,)::Nothing where TF

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{String} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (error("$Fnew already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || error("dataframe not sorted by [:permno, :date]")

  #find the lagged dates in the entire dataset
  alldates = unique(crsp.date) |> sort!
  alldatesidx = Dict{Date, Int}(alldates[i] => i for i ∈ 1:length(alldates))
  initialminpoints = searchsortedfirst(alldates, alldates[1] + calendarinterval) |> Int

  #now compute the lagged dates
  getlaggedpoint(i) = alldates[searchsortedfirst(alldates, alldates[i] - calendarinterval)]

  laggeddateidx = Dict{Date, Date}(
    alldates[i] => getlaggedpoint(i) for i ∈ initialminpoints:length(alldates))
  minstartdate = alldates[initialminpoints]
  @assert (alldates[alldates .≥ minstartdate] .|> (d)->haskey(laggeddateidx, d)) |> all

  #low effort check to only look at rows after the minimum days across all intervals
  minintervaldays = minimum(((k,v)->alldatesidx[k]-alldatesidx[v]
    ).(keys(laggeddateidx), values(laggeddateidx)))

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = length(scrsp.date)

    #if there aren't enough dates within a permno to create lagged dates, skip the permno
    (firstdate, lastdate) = (scrsp[1,:date], scrsp[end,:date])
    if ((nrows < minintervaldays) ||
        (lastdate < minstartdate) ||
        (laggeddateidx[lastdate] < firstdate))

      continue
    end

    #a hack to only compute calcualtions at end of month
    for j ∈ minintervaldays:length(scrsp.date)
      #if we are only computing end of month values, make sure we are at a month end
      applyconditional && (!scrsp[j,Fconditional]) && continue
      currentdate = scrsp[j, :date]

       #skip if currentdate is before the first entry in the lookup table
      currentdate < minstartdate && continue

      #now check if the lagged date is a valid index entry for this permno
      targetlaggeddate = laggeddateidx[currentdate]
      j₀ = searchsortedfirst(scrsp.date, targetlaggeddate)
      actuallaggeddate = scrsp[j₀, :date]
      (actuallaggeddate ≠ targetlaggeddate) && continue
      scrsp[j, Fmeasure] = F(scrsp[j₀:j , Ftarget])
    end
  end

  return nothing
end

#applies a function to a vector of trailing variables
#this version uses a rolling window computed by the number of entries as opposed to
#one computed by the dates
function trailingmeasure!(F::TF, crsp::DataFrame, windowsize::Int;
  Ftarget::Symbol = throw("Ftarget is required"),
  Fmeasure::Symbol = throw("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing,
  maxcalendarinterval::DatePeriod = throw("maxcalendarinterval is required")
  )::Nothing where TF

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{Symbol} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (throw("$Fnew already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || throw("dataframe not sorted by [:permno, :date]")

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = size(scrsp, 1)
    (nrows<windowsize) || continue #if there are enough  entries to compute the trailing returns

    for j ∈ windowsize:nrows
      applyconditional && (!scrsp[j,Fconditional]) && continue #if we are only computing end of month values
      staledate::Date = scrsp[j, :date] - maxcalendarinterval
      j₀::Int = j-windowsize+1

      (scrsp[j₀, :date] ≤ staledate)  && continue #elliminates stale windows
      scrsp[j, Fmeasurename] = F(scrsp[j₀:j, Ftarget])
    end
  end

  return nothing
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
include("Domain.jl")
include("Field.jl")

################################################################################
################################################################################
#Fourier and Hartley transformation
function fft(f::Field{T, N, Dom}) where {T, N, Dom <: RGDomain}
    dom = getdomain(f)
    if !isharmonic(dom)
        error("Fourier Transform requires harmonic domain")
    end
    new_dom = getcodomain(dom)
    val = fft(f.val)
    return field(new_dom, val)
end

function ifft(f::Field{T, N, Dom}) where {T, N, Dom <: RGDomain}
    dom = getdomain(f)
    if isharmonic(dom)
        error("Inverse Fourier Transform requires real domain")
    end
    new_dom = getcodomain(dom)
    val = ifft(f.val)
    return field(new_dom, val)
end

function hartley(f::Array)
    return FFTW.r2r(f, FFTW.DHT)
end

function ihartley(f::Array)
    return hartley(f) ./ size(f)
end


text:
using CmdStan, Documenter

DOC_ROOT = rel_path_cmdstan("..", "docs")
DocDir =  joinpath(DOC_ROOT, "src")

page_list = Array{Pair{String, Any}, 1}();
append!(page_list, [Pair("Introduction", "INTRO.md")]);
append!(page_list, [Pair("Installation", "INSTALLATION.md")]);
append!(page_list, [Pair("Walkthrough", "WALKTHROUGH.md")]);
append!(page_list, [Pair("Versions", "VERSIONS.md")]);
append!(page_list, [Pair("Functions", "index.md")]);

makedocs(
  format = Documenter.HTML(prettyurls = haskey(ENV, "GITHUB_ACTIONS")),
  root = DOC_ROOT,
  modules = Module[],
  sitename = "CmdStan.jl",
  authors = "Rob J Goedman",
  pages = page_list,
)

deploydocs(
    root = DOC_ROOT,
    repo = "github.com/StanJulia/CmdStan.jl.git",
    devbranch = "master",
    push_preview = true,
 )


text:
using JLD
include("main.jl")

# read in solution. Defines m, vf, c, kp, coefs
@load "BFZ_additive.jld"

function get_J_coefs_Jk(m::X.BFZ, vf)
    ρ, α, β = X._unpack(m.agent)
    J = vf.^(ρ) ./ ρ
    J_coefs = X.get_coefs(m.basis, m.basis_struc, J)
    J_k = X.funeval(J_coefs, m.basis, m.grid, [1 0 0])
    J, J_coefs, J_k
end

max_stars_repo_name: NYUEcon/GrowthModels

max_issues_repo_name: NYUEcon/GrowthModels

max_forks_repo_name: NYUEcon/GrowthModels


text:
export current_context, with_context

using UUIDs: uuid4

const CONTEXT_KEY = :OPEN_TELEMETRY_CONTEXT

# ??? performace issue with NamedTuple
struct Context{T<:NamedTuple}
    kv::T
end

Context(;kw...) = Context(kw.data)

Base.merge(c1::Context, c2::Context) = Context(merge(c1.kv, c2.kv))

create_key(s) = Symbol(s, '-', uuid4())

for f in (:getindex, :haskey, :get)
    @eval Base.$f(ctx::Context, args...) = $f(ctx.kv, args...)
end

with_context(f, kv::Pair{Symbol}...) = with_context(f, current_context(), kv...)

"""
    with_context(f, [context], [k=>v]...)

Run function `f` in the `context`. If extra `kv` pairs are provided, they will be merged with the `context` to form a new context. When `context` is not provided, the [`current_context`](@ref) will be used.
"""
function with_context(f, ctx::Context, kv::Pair{Symbol}...)
    task_local_storage(CONTEXT_KEY, merge(ctx,Context(;kv...))) do
        f()
    end
end

"""
Return the `Context` associated with the caller's current execution unit.
"""
current_context() = get!(task_local_storage(), CONTEXT_KEY, Context())

# !!! intentional type piracy
function Base.schedule(t::Task)
    if isnothing(t.storage)
        t.storage = IdDict()
    end
    t.storage[CONTEXT_KEY] = current_context()
    Base.enq_work(t)
end


text:
# https://discourse.julialang.org/t/write-a-rest-interface-like-flask/18538/16

# Bukdu v0.4.4
using Bukdu

struct RESTController <: ApplicationController
    conn::Conn
end

function create(c::RESTController)
    @info :payload (c.params.message, c.params.x, c.params.y)
    render(JSON, "OK")
end

routes() do
    post("/messages", RESTController, create)
    plug(Plug.Parsers, [:json])
end

Bukdu.start(8080)

#=
curl -H "Content-Type: application/json" http://127.0.0.1:8080/messages -d '{"message": "Hello Data"}'
curl -H "Content-Type: application/json" http://127.0.0.1:8080/messages -d '{"x": 1.2, "y": 2.3}'

On Windows:
- `"` needs to be escaped either by `\"` or `"""`
- initial `'` needs to be replaces with `"`
curl -H "Content-Type: application/json" http://127.0.0.1:8080/messages -d "{\"message\": \"Hello Data\"}"
curl -H "Content-Type: application/json" http://127.0.0.1:8080/messages -d "{"""x""": 1.2, """y""": 2.3}"
For more information, see: https://stackoverflow.com/questions/7172784/how-do-i-post-json-data-with-curl
=#


text:
import Base: isempty,copy,getindex,setindex!,firstindex,lastindex,copyto!,length,*,+,IndexStyle,iterate
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# FILE CONTENTS:
#     SUMMARY: Results types for storing arrays used during computation
#
#     TYPES                                        Tree
#        SolverResults                           ---------
#        SolverIterResults                      SolverResults
#        ConstrainedResults                       ↙     ↘
#        UnconstrainedResults          ResultsCache   SolverIterResults
#                                                          ↙     ↘
#                                      UnconstrainedResults    ConstrainedResults
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# Trajectory Types
Trajectory = Vector{T} where T <: AbstractArray
TrajectoryVectors = Vector{Vector{T}} where T <: Real
TrajectoryMatrices = Vector{Matrix{T}} where T <: Real
TrajectoryDiagonals = Vector{Diagonal{Vector{T}}} where T <: Real

"""
$(TYPEDEF)
    Abstract type for the backward pass expansion terms
"""
abstract type AbstractBackwardPass end

struct BackwardPass <: AbstractBackwardPass
    Qx::Vector{Vector{Float64}}
    Qu::Vector{Vector{Float64}}
    Qxx::Vector{Matrix{Float64}}
    Qux::Vector{Matrix{Float64}}
    Quu::Vector{Matrix{Float64}}

    Qux_reg::Vector{Matrix{Float64}}
    Quu_reg::Vector{Matrix{Float64}}
end

function BackwardPass(n::Int,m::Int,N::Int)
    Qx = [zeros(n) for i = 1:N]
    Qu = [zeros(m) for i = 1:N-1]
    Qxx = [zeros(n,n) for i = 1:N]
    Qux = [zeros(m,n) for i = 1:N-1]
    Quu = [zeros(m,m) for i = 1:N-1]

    Qux_reg = [zeros(m,n) for i = 1:N-1]
    Quu_reg = [zeros(m,m) for i = 1:N-1]

    BackwardPass(Qx,Qu,Qxx,Qux,Quu,Qux_reg,Quu_reg)
end

function copy(bp::BackwardPass)
    BackwardPass(deepcopy(bp.Qx),deepcopy(bp.Qu),deepcopy(bp.Qxx),deepcopy(bp.Qux),deepcopy(bp.Quu),deepcopy(bp.Qux_reg),deepcopy(bp.Quu_reg))
end

"""
$(TYPEDEF)
Abstract type for the output of solving a trajectory optimization problem
"""
abstract type SolverResults end

"""
$(TYPEDEF)
Abstract type for the output of a single iteration step
"""
abstract type SolverIterResults <: SolverResults end

# abstract type SolverIterResultsStatic <: SolverIterResults end

abstract type SolverVectorResults <: SolverIterResults end
abstract type UnconstrainedIterResults <: SolverVectorResults end
abstract type ConstrainedIterResults <: SolverVectorResults end

################################################################################
#                                                                              #
#                   UNCONSTRAINED RESULTS STRUCTURE                            #
#                                                                              #
################################################################################

struct UnconstrainedVectorResults <: UnconstrainedIterResults
    X::Trajectory  # States (n,N)
    U::Trajectory  # Controls (m,N)

    K::Trajectory # Feedback (state) gain (m,n,N)
    d::Trajectory  # Feedforward gain (m,N)

    X_::Trajectory # Predicted states (n,N)
    U_::Trajectory # Predicted controls (m,N)

    S::Trajectory  # Cost-to-go hessian (n,n)
    s::Trajectory  # Cost-to-go gradient (n,1)

    fdx::Trajectory # Discrete dynamics state jacobian (n,n,N)
    fdu::Trajectory # Discrete dynamics control jacobian (n,m,N-1)

    ρ::Vector{Float64}
    dρ::Vector{Float64}

    bp::BackwardPass

    function UnconstrainedVectorResults(X,U,K,d,X_,U_,S,s,fdx,fdu,ρ,dρ,bp)
        new(X,U,K,d,X_,U_,S,s,fdx,fdu,ρ,dρ,bp)
    end
end

"""
$(SIGNATURES)
Construct results from sizes
# Arguments
* n: number of states
* m: number of controls
* N: number of time steps
"""
function UnconstrainedVectorResults(n::Int,m::Int,N::Int)
    X  = [zeros(n)   for i = 1:N]
    U  = [zeros(m)   for i = 1:N-1]

    K  = [zeros(m,n) for i = 1:N-1]
    d  = [zeros(m)   for i = 1:N-1]

    X_ = [zeros(n)   for i = 1:N]
    U_ = [zeros(m)   for i = 1:N-1]

    S  = [zeros(n,n) for i = 1:N]
    s  = [zeros(n)   for i = 1:N]

    fdx = [zeros(n,n) for i = 1:N-1]
    fdu = [zeros(n,m) for i = 1:N-1]

    ρ = zeros(1)
    dρ = zeros(1)

    bp = BackwardPass(n,m,N)

    UnconstrainedVectorResults(X,U,K,d,X_,U_,S,s,fdx,fdu,ρ,dρ,bp)
end

function copy(r::UnconstrainedVectorResults)
    UnconstrainedVectorResults(copy(r.X),copy(r.U),copy(r.K),copy(r.d),copy(r.X_),copy(r.U_),copy(r.S),copy(r.s),copy(r.fdx),copy(r.fdu),copy(r.ρ),copy(r.dρ),copy(r.bp))
end

################################################################################
#                                                                              #
#                     CONSTRAINED RESULTS STRUCTURE                            #
#                                                                              #
################################################################################

struct ConstrainedVectorResults <: ConstrainedIterResults
    X::Trajectory  # States (n,N)
    U::Trajectory  # Controls (m,N)

    K::Trajectory # Feedback (state) gain (m,n,N)
    d::Trajectory  # Feedforward gain (m,N)

    X_::Trajectory # Predicted states (n,N)
    U_::Trajectory # Predicted controls (m,N)

    S::Trajectory  # Cost-to-go hessian (n,n)
    s::Trajectory  # Cost-to-go gradient (n,1)

    fdx::Trajectory # State jacobian (n,n,N)
    fdu::Trajectory # Control (k) jacobian (n,m,N-1)

    C::Trajectory      # Constraint values (p,N)
    C_prev::Trajectory # Previous constraint values (p,N)
    Iμ::Trajectory        # fcxtive constraint penalty matrix (p,p,N)
    λ::Trajectory # Lagrange multipliers (p,N)
    μ::Trajectory     # Penalty terms (p,N)

    Cx::Trajectory # State jacobian (n,n,N)
    Cu::Trajectory # Control (k) jacobian (n,m,N-1)

    t_prev::Trajectory
    λ_prev::Trajectory

    nesterov::Vector{Float64}

    active_set::Trajectory

    ρ::Array{Float64,1}
    dρ::Array{Float64,1}

    bp::BackwardPass

    function ConstrainedVectorResults(
            X,U,
            K,d,
            X_,U_,
            S,s,
            fdx,fdu,
            C,C_prev,
            Iμ,λ,μ,
            Cx,Cu,
            t_prev,λ_prev,nesterov,
            active_set,
            ρ,dρ,bp)
        new(X,U,K,d,X_,U_,S,s,fdx,fdu,C,C_prev,Iμ,λ,μ,Cx,Cu,t_prev,λ_prev,nesterov,active_set,ρ,dρ,bp)
    end
end

isempty(res::SolverIterResults) = isempty(res.X) && isempty(res.U)

ConstrainedVectorResults() = ConstrainedVectorResults(0,0,0,0,0)

"""
$(SIGNATURES)
Construct results from sizes
# Arguments
* n: number of states
* m: number of controls
* p: number of constraints
* N: number of time steps
* p_N (default=n): number of terminal constraints
"""
function ConstrainedVectorResults(n::Int,m::Int,p::Int,N::Int,p_N::Int)

    X  = [zeros(n)   for i = 1:N]
    U  = [zeros(m)   for i = 1:N-1]

    K  = [zeros(m,n) for i = 1:N-1]
    d  = [zeros(m)   for i = 1:N-1]

    X_ = [zeros(n)   for i = 1:N]
    U_ = [zeros(m)   for i = 1:N-1]

    S  = [zeros(n,n) for i = 1:N]
    s  = [zeros(n)   for i = 1:N]

    fdx = [zeros(n,n) for i = 1:N-1]
    fdu = [zeros(n,m) for i = 1:N-1]

    # Stage Constraints
    C      = [i != N ? zeros(p) : zeros(p_N)  for i = 1:N]
    C_prev = [i != N ? zeros(p) : zeros(p_N)  for i = 1:N]
    Iμ     = [i != N ? Diagonal(ones(p)) : Diagonal(ones(p_N)) for i = 1:N]
    λ      = [i != N ? zeros(p) : zeros(p_N)  for i = 1:N]
    μ      = [i != N ? ones(p) : ones(p_N)  for i = 1:N]

    Cx  = [i != N ? zeros(p,n) : zeros(p_N,n)  for i = 1:N]
    Cu  = [i != N ? zeros(p,m) : zeros(p_N,0)  for i = 1:N]

    t_prev      = [i != N ? ones(p) : ones(p_N)  for i = 1:N]
    λ_prev      = [i != N ? zeros(p) : zeros(p_N)  for i = 1:N]

    nesterov = [0.;1.]

    active_set = [i != N ? zeros(Bool,p) : zeros(Bool,p_N)  for i = 1:N]

    ρ = zeros(1)
    dρ = zeros(1)

    bp = BackwardPass(n,m,N)

    ConstrainedVectorResults(X,U,K,d,X_,U_,S,s,fdx,fdu,
        C,C_prev,Iμ,λ,μ,Cx,Cu,t_prev,λ_prev,nesterov,active_set,ρ,dρ,bp)
end

function ConstrainedVectorResults(solver::Solver,X::Trajectory,U::Trajectory)
    p,pI,pE = get_num_constraints(solver)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver)
    n,m,N = get_sizes(solver)
    results = ConstrainedVectorResults(n,m,p,N,p_N)
    if solver.opts.square_root
        for k = 1:N
            results.S[k] = zeros(n+m,n)
        end
    end
    copyto!(results.X,X)
    copyto!(results.U,U)
    update_constraints!(results,solver)
    update_jacobians!(results,solver)
    return results
end

function ConstrainedVectorResults(solver::Solver,X::AbstractMatrix,U::AbstractMatrix)
    if size(U,2) == solver.N
        U = U[:,1:N-1]
    end
    ConstrainedVectorResults(solver,to_dvecs(X),to_dvecs(U))
end

function copy(r::ConstrainedVectorResults)
    ConstrainedVectorResults(deepcopy(r.X),deepcopy(r.U),deepcopy(r.K),deepcopy(r.d),deepcopy(r.X_),deepcopy(r.U_),deepcopy(r.S),deepcopy(r.s),deepcopy(r.fdx),deepcopy(r.fdu),
        deepcopy(r.C),deepcopy(r.C_prev),deepcopy(r.Iμ),deepcopy(r.λ),deepcopy(r.μ),
        deepcopy(r.Cx),deepcopy(r.Cu),deepcopy(r.t_prev),deepcopy(r.λ_prev),copy(r.nesterov),deepcopy(r.active_set),copy(r.ρ),copy(r.dρ),copy(r.bp))
end

#############
# Utilities #
#############
function remove_infeasible_controls!(results::SolverIterResults,solver::Solver)
    # turn off infeasible functionality
    p_inf,pI_inf,pE_inf = get_num_constraints(solver)

    solver.state.infeasible = false

    # get sizes
    n,m,N = get_sizes(solver)
    m̄,mm = get_num_controls(solver)
    n̄,nn = get_num_states(solver)

    if solver.state.minimum_time
        idx = 1:p_inf-n-1
        idx = [idx;p_inf]
    else
        idx = 1:p_inf-n
    end

    for k = 1:N-1
        results.U[k] = results.U[k][1:m̄]
        results.U_[k] = results.U_[k][1:m̄]
        results.K[k] = results.K[k][1:m̄,1:nn]
        results.d[k] = results.d[k][1:m̄]
        results.fdu[k] = results.fdu[k][1:nn,1:m̄]

        results.C[k] = results.C[k][idx]
        results.C_prev[k] = results.C_prev[k][idx]
        results.Cx[k] = results.Cx[k][idx,1:nn]
        results.Cu[k] = results.Cu[k][idx,1:m̄]
        results.λ[k] = results.λ[k][idx]
        results.λ_prev[k] = results.λ_prev[k][idx]
        results.μ[k] = results.μ[k][idx]
        results.Iμ[k] = Diagonal(Array(results.Iμ[k])[idx,idx]) # TODO there should be a more efficient way to do this
        results.active_set[k] = results.active_set[k][idx]

        results.bp.Qu[k] = results.bp.Qu[k][1:m̄]
        results.bp.Quu[k] = results.bp.Quu[k][1:m̄,1:m̄]
        results.bp.Qux[k] = results.bp.Qux[k][1:m̄,1:nn]
        results.bp.Quu_reg[k] = results.bp.Quu_reg[k][1:m̄,1:m̄]
        results.bp.Qux_reg[k] = results.bp.Qux_reg[k][1:m̄,1:nn]
    end
    # Don't need to modify terminal results C,Cx,Cu,λ,μ,Iμ since they are uneffected by u_infeasible
    return nothing
end

function init_results(solver::Solver,X::AbstractArray,U::AbstractArray; λ=Array{Float64,2}(undef,0,0), μ=Array{Float64,2}(undef,0,0))
    n,m,N = get_sizes(solver)

    if !isempty(X)
        solver.state.infeasible = true
    end

    # Chop off last control if N controls are passed in
    if size(U,2) == N
        U = U[:,1:N-1]
    end

    # Generate initial trajectoy (tacking on infeasible and minimum time controls)
    X_init, U_init = get_initial_trajectory(solver, X, U)

    # Get sizes
    m̄,mm = get_num_controls(solver)
    n̄,nn = get_num_states(solver)

    if solver.state.constrained
        p,pI,pE = get_num_constraints(solver)
        p_N,pI_N,pE_N = get_num_terminal_constraints(solver)

        results = ConstrainedVectorResults(nn,mm,p,N,p_N)

        # Set initial penalty term values
        copyto!(results.μ, results.μ*solver.opts.penalty_initial) # TODO change to assign, not multiply: μ_initial needs to be initialized as an array instead of float

        # Special penalty initializations #TODO indexing needs to be fixed here to account for less control constraints than controls
        # if solver.state.minimum_time
        #     for k = 1:N-1
        #         k != 1 ? results.μ[k][p] = solver.opts.penalty_initial_minimum_time_equality : nothing
        #         results.μ[k][m̄] = solver.opts.penalty_initial_minimum_time_inequality
        #         results.μ[k][m̄+m̄] = solver.opts.penalty_initial_minimum_time_inequality
        #     end
        # end
        if solver.state.infeasible
            nothing #TODO
        end

        # Initial Lagrange multipliers (warm start)
        if ~isempty(λ)
            copy_λ!(solver, results, λ)
        end
        if ~isempty(μ)
            # error("penalty warm start not implemented")
            copyto!(results.μ,μ)
        end

        # Set initial regularization
        results.ρ[1] = solver.opts.bp_reg_initial

    else
        results = UnconstrainedVectorResults(n,m,N)
    end

    if solver.opts.square_root
        for k = 1:N
            results.S[k] = zeros(nn+mm,nn)
        end
    end

    copyto!(results.X, X_init)
    copyto!(results.U, U_init)
    return results
end

function copy_λ!(solver, results, λ)
    N = solver.N
    p_new = length(λ[1])
    p_N_new = length(λ[end])

    p, = get_num_constraints(solver)
    p_N, = get_num_terminal_constraints(solver)

    if p_new == p  # all constraint λs passed in
        cid = trues(p)
    elseif p_new == solver.obj.p  # only "original" constraint λs passed
        cid = original_constraint_inds(solver)
    else
        err = ArgumentError("λ is not the correct dimension ($p_new). It must be either size $p or $(solver.obj.p)")
        throw(err)
    end
    for k = 1:N-1
        results.λ[k][cid] = λ[k]
    end
    results.λ[N] = λ[N]
end


text:
import Dolo
path = Dolo.pkg_path

fn = joinpath(path,"examples","models","rbc_mc.yaml")
model_mc = Dolo.yaml_import(fn)

drc = Dolo.ConstantDecisionRule(model_mc.calibration[:controls])
@time dr0, drv0 = Dolo.solve_policy(model_mc, drc; n_eval=1000) #, verbose=true, maxit=10000 )

# two loops:
# - VFI:
    # stops when err_v M tol_v

# optimize then inner_loop
# - inner: update until:
        - it_inner > maxit_inner
        - err_h < tol_h


text:
module Generator

using Pkg
using Genie, Genie.Generator

function newapp(app_name::String; ui::Bool = true, charts::Bool = true, autostart::Bool = true, kwargs...)
  Genie.Generator.newapp(app_name; autostart = false, kwargs...)

  ui && Pkg.add("StippleUI")
  charts && Pkg.add("StippleCharts")

  autostart && Genie.up()
end

end


text:
function Base._cpow(z::Complex{DoubleFloat{T}}, p::Complex{DoubleFloat{T}}) where {T<:IEEEFloat}
    pᵣ, pᵢ = reim(p)
    r = abs(z)
    θ = angle(z)
    rᵖ = r^pᵣ * exp(-pᵢ*θ)
    ϕ = pᵣ*θ + pᵢ*log(r)
    if isfinite(ϕ)
        return rᵖ * cis(ϕ)
    else
        iszero(rᵖ) && return zero(Complex{DoubleFloat{T}}) # no way to get correct signs of 0.0
        return Complex{DoubleFloat{T}}(NaN) # non-finite phase angle or NaN input
    end
end


text:
using ScrapeSEC
using Test

tests = ["download_metadata",
         "download_filings",
         "main_index"]

println("Running tests:")
for t in tests
    fp = "$(t).jl"
    println("* $fp ...")
    include(fp)
end


text:
#using Queryverse
using DataFrames
using DataFramesMeta
using Dates
using DelimitedFiles
using JuMP
using Logging
using CSV
using Query
using Statistics

# Parameters
const N_Gens =  64 # number of conventional generators
const N_Peakers =  80 # number of conventional generators
const N_StorgUs = 8 # number of storage units
const N_Blocks =7

const FILE_ACCESS_OVER = "w+"
const FILE_ACCESS_APPEND = "a+"
const POUNDS_TO_KG = 0.45359237

include("import_post_data.jl")

""""
Generators cost of unit i (variable cost)
Cost Unitᵢ[USD]= Avg. heat rate[mmbtu/MWh] * Avg. Fuel Price [USD/mmbtu] * Power Out[MWh]

    Power Out is obtained from:
    - UnitSchedule_(DEC-DEP 2019).csv

    Average Fuel Price is estimated from the data in the files:
    - data_fuel_price.csv & df_fuel_peakers.csv
    (both files are used as inputs of the UC model and the data comes from the
    files provided by Duke Energy, files UNIT_FUEL_PRICE(DEC 2019).csv and
    UNIT_FUEL_PRICE(DEP 2019).csv)

    Average heat rate is estimated from the data in the files:
    - df_gens.csv and df_peakers2.csv
    (both files are used as inputs of the UC model and the data comes from the
    files provided by Duke Energy, HR_Breakpoints_Summer_DEC-DEP_(2019) V2.csv &
    HR_Breakpoints_Winter_DEC-DEP_(2019) V2.csv)`

Parameters:
    df_gens:
    df_peakers:
    df_fuel_gens:
    df_fuel_peakers:
    df_gen_sch:
    ini_date: String with dateformat format 'MM/DD/YYYY HH:MM:SS', inclusive
    end_date: String with dateformat format 'MM/DD/YYYY HH:MM:SS', inclusive
"""
#TODO: Add summer data
function calc_model_costs(df_gens::DataFrame, df_peakers::DataFrame, df_fuel_gens::DataFrame,
    df_fuel_peakers::DataFrame, df_gen_sch::DataFrame, ini_date::String, end_date::String)

    #Combines the slow generators and peakers specs into one dataframe
    df_gens_all = vcat(df_gens, df_peakers);
    # Add the average heat rate data to the all gens' specs dataframe
    df_gens_all.Avg_HR =(df_gens_all.IHRC_B1_HR+df_gens_all.IHRC_B2_HR +
        df_gens_all.IHRC_B3_HR+df_gens_all.IHRC_B4_HR+df_gens_all.IHRC_B5_HR +
        df_gens_all.IHRC_B6_HR+df_gens_all.IHRC_B7_HR) / 7.0;

    #Combines the slow generators and peakers fuel data into one dataframe
    df_fuel_all = vcat(df_fuel_gens, df_fuel_peakers);
    # Adds the average fuel price to the fuel price dataframe as a new column
    insertcols!(df_fuel_all, 4, :avg_fuelprice => zeros(Float64), makeunique=true);

    #TODO: Check if instead of getting an average fuel price for the whole year
    # we could use  daily prices

    # Create a matrix of fuel prices to estimate the avg price and use it for cost calculatrion
    mat_fuelprice = Matrix(df_fuel_all[:,4:368]);
    avg_fuelprice = mean(mat_fuelprice, dims=2);
    df_avg_fuelprice = DataFrame(avg_fuelprice);

    rename!(df_avg_fuelprice,:x1 => :Avg_FP)
    df_fp_all = hcat(df_fuel_all, df_avg_fuelprice)

    # Adds the UNIT_ID that includes the cc key to the fuel price dataframe as a new column
    insertcols!(df_fp_all, 2, :UNIT_NAME => "", makeunique=true);

    for j in eachrow(df_fp_all)
        Idk = findall(x -> x == j.Column2, vec(df_gens_all.UNIT_ID));
        if isempty(Idk) == false
            j.UNIT_NAME = df_gens_all.UNIT_NAME[Idk[1]];
        end
    end

    # Filtering the Actual generation outcomes from UNIT_Schedules.csv
    # Reformats the date-time data given in column DATE
    #df_gen_sch.DATE = DateTime.(df_gen_sch.DATE, "mm/dd/yyyy HH:MM:SS p");

    # Call function
    #Selects a subset of schedules that fall into the months that users specify.
    # @where could be extended to include other month using the OR operation (||)
    ini_date = DateTime(ini_date, "mm/dd/yyyy HH:MM:SS");
    end_date = DateTime(end_date, "mm/dd/yyyy HH:MM:SS");

    df_convgen_sch = @from i in df_gen_sch begin
                @where DateTime(i.DATE) >= ini_date && DateTime(i.DATE) <= end_date
                @select {i.DATE, i.UNIT_ID, i.Output, i.On_off, i.ShutDown, i.Startup}
                @collect DataFrame
           end

    insertcols!(df_convgen_sch, size(df_convgen_sch,2), :Avg_HR => zeros(Float64));
    insertcols!(df_convgen_sch, size(df_convgen_sch,2), :Avg_FP => zeros(Float64));
    insertcols!(df_convgen_sch, 3, :UNIT_NAME => "");
    insertcols!(df_convgen_sch, size(df_convgen_sch,2), :TYPE_MAIN_FUEL => "");

    for row in eachrow(df_convgen_sch)
          Idx = findall(x -> x == row.UNIT_ID, vec(df_gens_all.UNIT_ID));
          if isempty(Idx) == false
                row.Avg_HR = df_gens_all.Avg_HR[Idx[1]];
                row.UNIT_NAME = df_gens_all.UNIT_NAME[Idx[1]];
          end
    end

    for row in eachrow(df_convgen_sch)
          Idx = findall(x -> x == row.UNIT_NAME, vec(df_fp_all.UNIT_NAME));
          if isempty(Idx) == false
                row.Avg_FP = df_fp_all.Avg_FP[Idx[1]];
                row.TYPE_MAIN_FUEL = df_fp_all.Column3[Idx[1]]; # Column3 = type of main fuel
          end
    end

    insertcols!(df_convgen_sch, size(df_convgen_sch,2), :GenCost => df_convgen_sch[!,:Output] .*df_convgen_sch[!,:Avg_HR] .* df_convgen_sch[!,:Avg_FP]);

    return df_convgen_sch;
end; # end function

"""
CO₂ emissions [metric ton] = avg. heat rate[mmbtu/MWh] * CO₂ emissions rate[kg/mmbu] * Power Out[MWh] * 10⁻³
 The CO₂ factor was obtained from the EIA site:
 [https://www.eia.gov/electricity/annual/html/epa_a_03.html]
"""
function calc_model_CO₂_emissions(df_co2::DataFrame, df_gen::DataFrame)
    # Calculating CO2 emissions for Peakers and CC units
    tonCO2_mmbtu = (@where(df_co2, occursin("NG").(:Fuel_code)).KgCO2_MMBTU ) / 1000.0;

    ng_tonCO2_mmbtu = (@where(df_co2, occursin("NG").(:Fuel_code)).KgCO2_MMBTU) / 1000.0;
    coal_tonCO2_mmbtu = (@where(df_co2, occursin("BIT").(:Fuel_code)).KgCO2_MMBTU)/ 1000.0;
    oil_tonCO2_mmbtu = (@where(df_co2, occursin("DFO").(:Fuel_code)).KgCO2_MMBTU)/ 1000.0;

    insertcols!(df_gen, size(df_gen,2), :EmissCO2_ton => zeros(Float64));

    #Options: COAL, CoGen, NGAS, Nuclear, and LOIL
    for unit in eachrow(df_gen)
        if (unit.TYPE_MAIN_FUEL == "COAL")
            unit.EmissCO2_ton = unit.Avg_HR * coal_tonCO2_mmbtu[1] * unit.Output;
        elseif (unit.TYPE_MAIN_FUEL == "NGAS")
            unit.EmissCO2_ton = unit.Avg_HR  * ng_tonCO2_mmbtu[1] * unit.Output;
        elseif (unit.TYPE_MAIN_FUEL == "LOIL")
            unit.EmissCO2_ton = unit.Avg_HR  * oil_tonCO2_mmbtu[1] * unit.Output;
      end;
    end

    return df_gen;
end

"""
NOₓ emissions [kg] = avg. heat rate [mmbtu/MWh] * avg NOₓ [lb/mmbu] * power out [MWh] * conv factor[Kg/lb]

The NOₓ factor is estimated form the files:
    - NOX_Breakpoints_Winter_Results.csv
    (the data  in this file comes from the files provided by Duke Energy,
    HR_Breakpoints_Summer_DEC-DEP_(2019) V2.csv &
    HR_Breakpoints_Winter_DEC-DEP_(2019) V2.csv)
"""
#TODO: Add summer data
function calc_model_NOx_emissions(df_nox::DataFrame, df_gen::DataFrame)
    df_nox.Avg_NOX = (df_nox.IHRC_B1_NOX + df_nox.IHRC_B2_NOX +
            df_nox.IHRC_B3_NOX + df_nox.IHRC_B4_NOX +
            df_nox.IHRC_B5_NOX + df_nox.IHRC_B6_NOX +
            df_nox.IHRC_B7_NOX) / 7.0;

    insertcols!(df_gen, size(df_gen,2), :Avg_NOX => zeros(Float64));

    for i in eachrow(df_gen)
        Idx = findall(x -> x == i.UNIT_NAME, vec(df_nox.UNIT_NAME));
        if isempty(Idx) == false
            i.Avg_NOX = df_nox.Avg_NOX[Idx[1]] * POUNDS_TO_KG;
        end
    end

    #  NOx Emissions
    insertcols!(df_gen, size(df_gen,2), :EmissNOx_kg => df_gen[!,:Avg_HR] .* df_gen[!,:Avg_NOX] .* df_gen[!,:Output]);

    return df_gen;
end

"""
SO₂ emissions [kg] = Avg. heat rate[mmbtu/MWh] * SO₂ factor[kg/mmbu] * Power Out[MWh]
The SO₂ factor was obtained from the data provided by Duke Energy in the files:
    - SO2_Breakpoints_Winter_Results.csv
    (the data  in this file from the files provided by Duke Energy,
    HR_Breakpoints_Summer_DEC-DEP_(2019) V2.csv &
    HR_Breakpoints_Winter_DEC-DEP_(2019) V2.csv)
"""
#TODO: Add summer data
function calc_model_SO₂_emissions(df_so2::DataFrame, df_gen::DataFrame)
    df_so2.Avg_SO2 = (df_so2.IHRC_B1_SO2 + df_so2.IHRC_B2_SO2 +
                df_so2.IHRC_B3_SO2 + df_so2.IHRC_B4_SO2 +
                df_so2.IHRC_B5_SO2 + df_so2.IHRC_B6_SO2 +
                df_so2.IHRC_B7_SO2) / 7.0;

    insertcols!(df_gen, size(df_gen,2), :Avg_SO2 => zeros(Float64));
    for i in eachrow(df_gen)
        Idx = findall(x -> x == i.UNIT_NAME, vec(df_so2.UNIT_NAME));
        if isempty(Idx) == false
            i.Avg_SO2 = df_so2.Avg_SO2[Idx[1]] * POUNDS_TO_KG;
        end
    end

    insertcols!(df_gen, size(df_gen,2), :EmissSO2_kg => df_gen[!,:Avg_HR] .* df_gen[!,:Avg_SO2] .* df_gen[!,:Output]);

    return df_gen;
end

df_bucr_gens_out = CSV.read(".//outputs//BUCR_GenOutputs.csv", header=true, DataFrame);
df_bucr_peakers_out = CSV.read(".//outputs//BUCR_PeakerOutputs.csv", header=true, DataFrame);

df_bucr_all_out = vcat(df_bucr_gens_out, df_bucr_peakers_out);

df_bucr_all_out = vcat(df_bucr_gens_out, df_bucr_peakers_out);

sim_year = 2019;

base_date = DateTime(sim_year, 1, 1, 00, 00, 00);
insertcols!(df_bucr_all_out, 1, :DATE => base_date);

for row in eachrow(df_bucr_all_out)
      if row.Hour >= 24
                row.Day = row.Day + 1;
                row.Hour = row.Hour - 24;
      end
      row.DATE = (row.DATE + Dates.Hour(row.Hour)) + Dates.Day(row.Day-1);
end

#insertcols!(df_bucr_all_out, 1, :DATE => (base_date + Dates.Hour.(df_bucr_all_out[!,:Hour])) + Dates.Day.(df_bucr_all_out[!,:Day]-1));
insertcols!(df_bucr_all_out, 2, :DATE_FORMAT => Dates.format.(df_bucr_all_out.DATE, "mm-dd-yyyy HH:MM:SS"));

rename!(df_bucr_all_out, Dict(Symbol("On/off") => "On_off"));
rename!(df_bucr_all_out, Dict(Symbol("UNIT_NAME") => "UNIT_ID"));

df_bucr_all_out = df_bucr_all_out[!, [:DATE, :DATE_FORMAT, :GeneratorID, :UNIT_ID,
    :MinPowerOut, :MaxPowerOut, :Output, :On_off, :ShutDown, :Startup, :TotalGenPosV,
    :TotalGenNegV, :MaxGenVio, :MinGenVio]];

CSV.write(".//outputs//postprocess//postproc_BUCR_Outputs.csv", df_bucr_all_out);

# Date format MM/DD/YYYY HH:MM:SS
Start_date = "01/01/2019 07:00:00"
#Final_date = "01/31/2019 23:00:00"
Final_date = "12/25/2019 06:00:00"

DF_Model_Results = DataFrame();

DF_Model_Results = calc_model_costs(df_gens, df_peakers, df_fuel_gens, df_fuel_peakers,
    df_bucr_all_out, Start_date, Final_date)

DF_Model_Results = calc_model_CO₂_emissions(df_eia_co2, DF_Model_Results);
#TODO: Add calculation for summer
DF_Model_Results = calc_model_NOx_emissions(df_nox_winter, DF_Model_Results);
DF_Model_Results = calc_model_SO₂_emissions(df_so2_winter, DF_Model_Results);

#Format strings to save csv file
Start_date = replace(Start_date, "/" => "-")
Final_date = replace(Final_date, "/" => "-")

Start_date = Start_date[1:findfirst(isequal(' '), Start_date)-1]
Final_date = Final_date[1:findfirst(isequal(' '), Final_date)-1]

CSV.write(".//outputs//postprocess//postproc_model_costs_$(Start_date)_$(Final_date).csv", DF_Model_Results)

#=
df_gens_cost_emiss = CSV.read(".//outputs//postprocess//postproc_gens_cost_emiss_01-01-2019_01-31-2019.csv", header=true, DataFrame);

#df_merge = innerjoin(df_bucr_all_out, df_gens_cost_emiss, on = [:DATE => :BEGIN_DATE, :UNIT_NAME => :UNIT_ID], makeunique=true);
df_merge = outerjoin(df_bucr_all_out, df_gens_cost_emiss,
            on = [:DATE => :BEGIN_DATE, :UNIT_ID => :UNIT_ID],
            validate=(true, true), makeunique=true);

CSV.write(".//outputs//postprocess//postproc_merge_out.csv", df_merge);
=#

#Reference: https://julia-data-query.readthedocs.io/en/latest/dplyr.html
DF_Model_Results_Byname = groupby(DF_Model_Results, [:UNIT_NAME, :TYPE_MAIN_FUEL]);
DF_Model_Results_Byfuel = groupby(DF_Model_Results, [:TYPE_MAIN_FUEL]);

#=DF_Model_Results_Byname = @combine(DF_Model_Results_Byname,
    Sum_GenCost = sum(:GenCost),
    Sum_Output = sum(:Output));
=#
DF_Model_Results_Byname = @combine(DF_Model_Results_Byname,
    Sum_GenCost = sum(:GenCost),
    Sum_MW = sum(:Output),
    Sum_EmissCO2_ton = sum(:EmissCO2_ton),
    Sum_EmissSO2_kg = sum(:EmissSO2_kg),
    Sum_EmissNOx_kg = sum(:EmissNOx_kg));

CSV.write(".//outputs//postprocess//postproc_model_cost_byUnit_$(Start_date)_$(Final_date).csv", DF_Model_Results_Byname)

DF_Model_Results_Byfuel = @combine(DF_Model_Results_Byfuel,
    Sum_GenCost = sum(:GenCost),
    Sum_MW = sum(:Output),
    Sum_EmissCO2_ton = sum(:EmissCO2_ton),
    Sum_EmissSO2_kg = sum(:EmissSO2_kg),
    Sum_EmissNOx_kg = sum(:EmissNOx_kg));

CSV.write(".//outputs//postprocess//postproc_model_cost_emiss_byFuel_$(Start_date)_$(Final_date).csv", DF_Model_Results_Byfuel)

#TODO:
# Check why this is happening:
# ASHV_CT03_0 and ASHV_CT04_0 are in data_fuel_price.csv and data_fuel_price_peakers.csv input files
# ASHV_CC01, ASHV_CC02 are not in data_fuel_price.csv or data_fuel_price_peakers.csv
# input files. In the fuel price files provided by Duke there is only info of them from August 2019.

# Getting time-series statistics
# Reference: https://www.tutorialspoint.com/time_series/time_series_error_metrics.htm

# Root Mean Square Error:
# RMSE = ⎷ (1/n ∑(yₜ' - yₜ)²)
# Mean Absolute Error

#include(".//postprocessing//calc_model_costs.jl")

max_stars_repo_name: power-systems-op/grace-arpae

max_issues_repo_name: power-systems-op/grace-arpae

max_forks_repo_name: power-systems-op/grace-arpae


text:
import JukaiNLP.IdDict

dict = IdDict()
id1 = push!(dict, "abc")
id2 = push!(dict, "def")
id3 = push!(dict, "abc")
@test id1 == 1
@test id2 == 2
@test id1 == id3
@test dict["abc"] == 1

@test length(dict) == 2

@test count(dict, id1) == 2
@test count(dict, id2) == 1

@test getkey(dict, id1) == "abc"
@test getkey(dict, id2) == "def"

@test get(dict, "abc") == 1
@test get(dict, "abd") == 0

max_stars_repo_name: hshindo/JukaiNLP

max_issues_repo_name: hshindo/JukaiNLP

max_forks_repo_name: hshindo/JukaiNLP


text:
using Documenter, UnitRoot

makedocs(modules=UnitRoot,
        doctest=true)

deploydocs(deps = Deps.pip("mkdocs", "python-markdown-math"),
    repo = "https://github.com/p-chaim/UnitRoot.jl",
    julia  = "0.6.0",
    osname = "linux")


text:
function export_EIS_to_Z_file(path, f_range, Z_range)
  if length(f_range) != length(Z_range)
    println("ERROR: length(f_range) != length(Z_range) >> $(length(f_range)) != $(length(Z_range))")
    return throw(Exception)
  end
  src_dir = @__DIR__
  open(path, "w") do f
    open(src_dir*"/Z_view_template.txt", "r") do template
      for line in eachline(template)
        write(f, line*"\n")
      end
    end
    for i in 1:length(f_range)
      write(f, @sprintf("%E\t%E\t%E\t%E\t%E\t%E\t%E\t%i\t%i\t",
                f_range[i], 0, 0, 0, real(Z_range[i]), imag(Z_range[i]), 0, 0, 0
                )*"\n"
      )
    end
  end
end

max_stars_repo_name: Masicko/ImageToEIS

max_issues_repo_name: Masicko/ImageToEIS

max_forks_repo_name: Masicko/ImageToEIS


text:
using SimpleDirectMediaLayer

# import SimpleDirectMediaLayer:
#     Window, Renderer

export World
export gen_id

const DISPLAY_RATIO = 16.0 / 9.0
const WIDTH = 1024 + 512
# Larget number causes the view to encompass more of the world
# which makes objects appear smaller.
const VIEW_SCALE = 1.5

const WINDOW_POSITION = (1000, 100)

mutable struct World
    title::String

    # Device properties
    window_position_x::UInt32
    window_position_y::UInt32
    window_width::UInt32
    window_height::UInt32

    window_centered::Bool

    view_width::Float64
    view_height::Float64
    view_centered::Bool

    # view space to device-space projection
    view_space::Math.AffineTransform
    inv_view_space::Math.AffineTransform

    ids::UInt32

    window #::Ptr{Window}
    renderer #::Ptr{Renderer}

    function World(title::String)
        o = new()
        o.title = title
        o.window_position_x = WINDOW_POSITION[1]
        o.window_position_y = WINDOW_POSITION[2]
        o.window_width = WIDTH
        o.window_height = UInt32(Float64(WIDTH) / DISPLAY_RATIO)
        o.window_centered = true
        o.view_centered = true
        o.view_width = Float64(o.window_width) * VIEW_SCALE
        o.view_height = Float64(o.window_height) * VIEW_SCALE
        o.ids = 0
        o.window = o    # Not assigned yet
        o.renderer = o  # Not assigned yet

        o.view_space = Math.AffineTransform{Float64}()
        o.inv_view_space = Math.AffineTransform{Float64}()

        set_view_space!(o)

        println("Display dimensions: [", o.window_width, " x ", o.window_height, "]")
        println("View dimensions: [", o.view_width, " x ", o.view_height, "]")

        o
    end
end

function gen_id(world::World)
    world.ids += 1
    world.ids
end

function set_view_space!(world::Ranger.World)
    center = Math.AffineTransform{Float64}()

    # What separates world from view is the ratio between the device (aka window)
    # and an optional centering translation.
    width_ratio = Float64(world.window_width) / world.view_width
    height_ratio = Float64(world.window_height) / world.view_height

    if world.view_centered
        Math.make_translate!(center, Float64(world.window_width) / 2.0, Float64(world.window_height) / 2.0)
    end

    Math.scale!(center, width_ratio, height_ratio)
    Math.set!(world.view_space, center)

    Math.invert!(world.view_space, world.inv_view_space)
end

max_stars_repo_name: wdevore/Ranger-Julia-SDL

max_issues_repo_name: wdevore/Ranger-Julia-SDL

max_forks_repo_name: wdevore/Ranger-Julia-SDL


text:
module Sets
using LinearAlgebra
using RecipesBase
import MultivariatePolynomials
const MP = MultivariatePolynomials
import DynamicPolynomials
import MultivariateBases
const MB = MultivariateBases
const MonoBasis = MB.MonomialBasis{DynamicPolynomials.Monomial{true}, DynamicPolynomials.MonomialVector{true}}
using Polyhedra

const SpaceVariable = DynamicPolynomials.PolyVar{true}

abstract type AbstractSet{T} end

"""
    struct Polar{T, S<:AbstractSet{T}} <: AbstractSet{T}
        set::S
    end

The polar of the set `set`.
"""
struct Polar{T, S<:AbstractSet{T}} <: AbstractSet{T}
    set::S
end

const PolarOf{S} = Polar{<:Any, S}

const PolarOrNot{S} = Union{S, PolarOf{S}}

"""
    polar(set::AbstractSet)

Return the polar of `set`.
"""
Polyhedra.polar(set::AbstractSet) = Polar(set)
Polyhedra.polar(set::Polar) = set.set

function print_support_function(set::Polar; kws...)
    print("h(S, x) =")
    _print_gauge_function(polar(set); kws...)
end

"""
    polar_representation(set::AbstractSet)

Return a representation of the same set but in the polar space of the current
representation of `set`.
"""
function polar_representation end

"""
    gauge1(set::AbstractSet)

Function `f(x)` such that `set` is the 1-sublevel set of `f`, i.e.
`{ x | f(x) ≤ 1 }`.
"""
function gauge1 end

"""
    struct PerspectiveDual{T, S <: AbstractSet{T}} <: AbstractSet{T}
        set::S
    end

Set determined by the dual of the perspective cone of `set`.
"""
struct PerspectiveDual{T, S <: AbstractSet{T}} <: AbstractSet{T}
    set::S
end

const PerspectiveDualOf{S} = PerspectiveDual{<:Any, S}

const PerspectiveDualOrPolarOrNot{S} = Union{PerspectiveDualOf{S}, PolarOrNot{S}}

"""
    perspective_dual(set::AbstractSet)

Return the set determined by the dual of the perspective cone of `set`.
"""
perspective_dual(set::AbstractSet) = PerspectiveDual(set)
perspective_dual(set::PerspectiveDual) = set.set

function scaling_function(set::PerspectiveDual)
    @assert length(space_variables(set)) == 2
    vars = [perspective_variable(set); space_variables(set)]
    # z is a halfspace of the primal so a ray of the dual
    z = [1.0, 0.0, 0.0]
    in_set(Δ::Vector) = perspective_gauge0(set.set)(vars => z + Δ) < 0
    @assert in_set(zeros(3))
    return (Δz, Δx, Δy) -> begin
        Δ = [Δz, Δx, Δy]
        _in_set(λ::Real) = in_set(Δ * λ)
        λ = 1.0
        while _in_set(λ)
            if λ > 1e10
                error("Error in plotting : the `InteriorPoint` seems to be on the boundary")
            end
            λ *= 2
        end
        λmin = 0.0
        λmax = λ
        # Binary search. Invariant: in_set(λmin) and !in_set(λmax)
        while abs(λmin - λmax) > 1e-8
            λ = (λmin + λmax) / 2
            if _in_set(λ)
                λmin = λ
            else
                λmax = λ
            end
        end
        λ = (λmin + λmax) / 2
        return 1 / λ
    end
end

# TODO rename space_dimension to avoid confusion with Polyhedra.dimension
"""
    dimension(set::AbstractSet)

Return the dimension of the space where the set `set` is defined.
"""
dimension(set::AbstractSet) = length(space_variables(set))

function dimension(set::Union{Polar, PerspectiveDual})
    return dimension(set.set)
end

"""
    space_variables(set::AbstractSet)

Return the variables of the space where the set `set` is defined or `nothing`
if none are used.
"""
function space_variables end

function space_variables(set::Union{Polar, PerspectiveDual})
    return space_variables(set.set)
end

function perspective_variable(set::Union{Polar, PerspectiveDual})
    return perspective_variable(set.set)
end
function perspective_gauge0 end
function perspective_gauge1 end

convexity_proof(set::Union{Polar, PerspectiveDual}) = convexity_proof(set.set)

struct UnknownSet{T} <: AbstractSet{T} end
include("transformations.jl")

struct Piecewise{T, S<:AbstractSet{T}, U, Po <: Polyhedra.Polyhedron{U}, Pi} <: AbstractSet{T}
    sets::Vector{S}
    polytope::Po
    pieces::Vector{Pi}
    # TODO add adjacency graph in Polyhedra
    graph::Vector{Vector{Tuple{Int, Vector{U}}}}
end
function Piecewise(sets::Vector{<:AbstractSet}, polytope::Polyhedra.Polyhedron{U}) where U
    @assert length(sets) == nhalfspaces(polytope)
    points = [Set(incidentpointindices(polytope, hidx)) for hidx in eachindex(halfspaces(polytope))]
    graph = [Tuple{Int, Vector{U}}[] for i in eachindex(halfspaces(polytope))]
    if !(Polyhedra.origin(Polyhedra.pointtype(polytope), Polyhedra.fulldim(polytope)) in polytope)
        error("The origin is not in the polytope")
    end
    for (i, hi) in enumerate(halfspaces(polytope))
        hi.β > 0 || error("The origin is not in the polytope")
        vi = hi.a / hi.β
        for (j, hj) in enumerate(halfspaces(polytope))
            hi.β > 0 || error("The origin is not in the polytope")
            vj = hj.a / hj.β
            i == j && break
            if length(points[i] ∩ points[j]) ≥ fulldim(polytope) - 1
                push!(graph[i], (j, vj - vi))
                push!(graph[j], (i, vi - vj))
            end
        end
    end
    function piece(i, h)
        # Need to be a polyhedron as `detecthlinearity` needs a solver in `add_constraint_inclusion_domain`
        return Polyhedra.polyhedron(hrep([HalfSpace(edge[2], zero(U)) for edge in graph[i]]),
                                    Polyhedra.DefaultLibrary{U}(Polyhedra.default_solver(polytope)))
    end
    pieces = [piece(i, h) for (i, h) in enumerate(halfspaces(polytope))]
    return Piecewise(sets, polytope, pieces, graph)
end
dimension(set::Piecewise) = Polyhedra.fulldim(set.polytope)
space_variables(set::Piecewise) = space_variables(set.sets[1])
function scaling_function(set::Piecewise)
    g = scaling_function.(set.sets)
    return (x, y) -> begin
        v = [x, y]
        i = findfirst(piece -> v in piece, set.pieces)
        return g[i](x, y)
    end
end
function zero_eliminate(set::Piecewise, I)
    _elim(p) = Polyhedra.fixandeliminate(p, I, zeros(Polyhedra.coefficient_type(p), length(I)))
    J = setdiff(1:dimension(set), I)
    return Piecewise(
        [zero_eliminate(s, I) for s in set.sets],
        _elim(set.polytope),
        _elim.(set.pieces),
        map(set.graph) do adj
            map(adj) do iv
                iv[1], iv[2][J]
            end
        end
    )
end

function _print_gauge_function(set::Piecewise; digits=6)
    DynamicPolynomials.@polyvar x[1:2]
    println()
    for (set, piece) in zip(set.sets, set.pieces)
        print("         ")
        _print_gauge_function(set, digits=digits)
        print("              if ")
        for (i, h) in enumerate(halfspaces(piece))
            if i > 1
                print(", ")
            end
            a = -h.a
            if count(!iszero, a) == 1
                a /= abs(sum(a)) # Simplify printing
            end
            if digits !== nothing
                a = round.(a, digits=digits)
            end
            print(a'x)
            print(" ≥ 0")
        end
        println()
    end
end

include("polytope.jl")
include("ellipsoids.jl")
include("polynomials.jl")
include("recipe.jl")

end


text:
"""
    ParametricEntity <: AbstractEntity

A geometric entity given by an explicit `parametrization` field. The valid
arguments should lie in its `domain`. This differs from an `ElementaryEntity` in
that the underlying representation of the geometric entity is available.
"""
struct ParametricEntity <: AbstractEntity
    dim::UInt8
    tag::Int
    parametrization
    domain
    boundary::Vector{<:AbstractEntity}
    function ParametricEntity(dim, tag, f, d)
        ent = new(dim, tag, f, d,AbstractEntity[])
        # every entity gets added to a global variable ENTITIES so that we can
        # ensure the (d,t) pair is a UUID for an entity, and to easily retrieve
        # different entities.
        global_add_entity!(ent)
        return ent
    end
end

domain(p::ParametricEntity)              = p.domain
parametrization(p::ParametricEntity)     = p.parametrization
geometric_dimension(p::ParametricEntity) = p.dim
tag(p::ParametricEntity) = p.tag
boundary(p::ParametricEntity) = p.boundary

function flip_normal(ent::ParametricEntity)
    @assert ambient_dimension(ent) == geometric_dimension(ent) + 1
    ParametricEntity(geometric_dimension(ent),-tag(ent),parametrization(ent),domain(ent))
end

function ParametricEntity(f,dom)
    d = geometric_dimension(dom)
    t = new_tag(d) # automatically generate a new (valid) tag
    ParametricEntity(d, t, f, dom)
end

function return_type(p::ParametricEntity)
    # NOTE: this relies on the brittle promote_op
    d = domain(p)
    x = center(d)
    f = parametrization(p)
    T = Base.promote_op(f, typeof(x))
    return T
end

ambient_dimension(p::ParametricEntity) = length(return_type(p))

function (par::ParametricEntity)(x)
    @assert x in domain(par)
    par.parametrization(x)
end

jacobian(psurf::ParametricEntity,s::SVector)  = ForwardDiff.jacobian(psurf.parametrization, s::SVector)
jacobian(psurf::ParametricEntity,s)           = jacobian(psurf, SVector(s))

"""
    line(a,b)

Create a straight line connecting points `a` and `b`. Returns an instance
of [`ParametricEntity`](@ref).
"""
function line(a::SVector,b::SVector)
    f = (u) -> a + u[1]*(b-a)
    d = HyperRectangle(0.,1.)
    ParametricEntity(f,d)
end
line(a,b) = line(SVector(a),SVector(b))

max_stars_repo_name: WaveProp/ParametricSurfaces

max_issues_repo_name: WaveProp/ParametricSurfaces

max_forks_repo_name: WaveProp/ParametricSurfaces