text:
# --------------------------------------------------------------------------
# ACE1.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

module OrthPolys

using SparseArrays
using LinearAlgebra: dot

import JuLIP: evaluate!, evaluate_d!, JVec, cutoff, fltype
import JuLIP.FIO: read_dict, write_dict
import JuLIP.MLIPs: alloc_B, alloc_dB, IPBasis

import ACE1
using ACE1.Transforms: DistanceTransform, transform, transform_d,
                        inv_transform, MultiTransform

import Base: ==

export transformed_jacobi

# this is a hack to prevent a weird compiler error that I don't understand
# at all yet
___f___(D::Dict) = (@show D)

function _fcut_(pl, tl, pr, tr, t)
   if (pl > 0 && t < tl) || (pr > 0 && t > tr)
      return zero(t)
   end
   return (t - tl)^pl * (t - tr)^pr
end

function _fcut_d_(pl, tl, pr, tr, t)
   if (pl > 0 && t < tl) || (pr > 0 && t > tr)
      return zero(t)
   end
   df = 0.0
   if pl > 0; df += pl * (t - tl)^(pl-1) * (t-tr )^pr ; end
   if pr  > 0; df += pr  * (t -  tr)^(pr -1) * (t-tl)^pl; end
   return df
end

struct OrthPolyBasis{T} <: ACE1.ScalarBasis{T}
   # ----------------- the parameters for the cutoff function
   pl::Int        # cutoff power left
   tl::T          # cutoff left (transformed variable)
   pr::Int        # cutoff power right
   tr::T          # cutoff right (transformed variable)
   # ----------------- the recursion coefficients
   A::Vector{T}
   B::Vector{T}
   C::Vector{T}
   # ----------------- used only for construction ...
   #                   but useful to have since it defines the notion or orth.
   tdf::Vector{T}
   ww::Vector{T}
end

fltype(P::OrthPolyBasis{T}) where {T} = T
Base.length(P::OrthPolyBasis) = length(P.A)

==(J1::OrthPolyBasis, J2::OrthPolyBasis) =
      all( getfield(J1, sym) == getfield(J2, sym)
           for sym in (:pr, :tr, :pl, :tl, :A, :B, :C) )

write_dict(J::OrthPolyBasis{T}) where {T} = Dict(
      "__id__" => "ACE1_OrthPolyBasis",
      "T" => write_dict(T),
      "pr" => J.pr,
      "tr" => J.tr,
      "pl" => J.pl,
      "tl" => J.tl,
      "A" => J.A,
      "B" => J.B,
      "C" => J.C
   )

OrthPolyBasis(D::Dict, T=read_dict(D["T"])) =
   OrthPolyBasis(
      D["pl"], D["tl"], D["pr"], D["tr"],
      Vector{T}(D["A"]), Vector{T}(D["B"]), Vector{T}(D["C"]),
      T[], T[]
   )

read_dict(::Val{:ACE1_OrthPolyBasis}, D::Dict) = OrthPolyBasis(D)

# rand applied to a J will return a random transformed distance drawn from
# the measure w.r.t. which the polynomials were constructed.
# TODO: allow non-constant weights!
function ACE1.rand_radial(J::OrthPolyBasis)
   @assert maximum(abs, diff(J.ww)) == 0
   return rand(J.tdf)
end

function OrthPolyBasis(N::Integer,
                       pcut::Integer,
                       tcut::T,
                       pin::Integer,
                       tin::T,
                       tdf::AbstractVector{T},
                       ww::AbstractVector{T} = ones(T, length(tdf))
                       ) where {T <: AbstractFloat}
   @assert pcut >= 0  && pin >= 0
   @assert N > 2

   if tcut < tin
      tl, tr = tcut, tin
      pl, pr = pcut, pin
   else
      tl, tr = tin, tcut
      pl, pr = pin, pcut
   end

   if minimum(tdf) < tl || maximum(tdf) > tr
      @warn("OrthoPolyBasis: t range outside [tl, tr]")
   end

   A = zeros(T, N)
   B = zeros(T, N)
   C = zeros(T, N)

   # normalise the weights s.t. <1, 1> = 1
   ww = ww ./ sum(ww)
   # define inner products
   dotw = (f1, f2) -> dot(f1, ww .* f2)

   # start the iteration
   _J1 = _fcut_.(pl, tl, pr, tr, tdf)
   a = sqrt( dotw(_J1, _J1) )
   A[1] = 1/a
   J1 = A[1] * _J1

   # a J2 = (t - b) J1
   b = dotw(tdf .* J1, J1)
   _J2 = (tdf .- b) .* J1
   a = sqrt( dotw(_J2, _J2) )
   A[2] = 1/a
   B[2] = -b / a
   J2 = (A[2] * tdf .+ B[2]) .* J1

   # keep the last two for the 3-term recursion
   Jprev = J2
   Jpprev = J1

   for n = 3:N
      # a Jn = (t - b) J_{n-1} - c J_{n-2}
      b = dotw(tdf .* Jprev, Jprev)
      c = dotw(tdf .* Jprev, Jpprev)
      _J = (tdf .- b) .* Jprev -c * Jpprev
      a = sqrt( dotw(_J, _J) )
      A[n] = 1/a
      B[n] = - b / a
      C[n] = - c / a
      Jprev, Jpprev = _J / a, Jprev
   end

   return OrthPolyBasis(pl, tl, pr, tr, A, B, C, collect(tdf), collect(ww))
end

alloc_B( J::OrthPolyBasis{T}) where {T} = zeros(T, length(J))
alloc_dB(J::OrthPolyBasis{T}) where {T} = zeros(T, length(J))
alloc_B( J::OrthPolyBasis{T}, ::Integer) where {T} = zeros(T, length(J))
alloc_dB(J::OrthPolyBasis{T}, ::Integer) where {T} = zeros(T, length(J))

# TODO: revisit this to allow type genericity!!!
alloc_B( J::OrthPolyBasis, ::TX) where {TX <: Number} = zeros(TX, length(J))
alloc_dB(J::OrthPolyBasis, ::TX) where {TX <: Number} = zeros(TX, length(J))
alloc_B( J::OrthPolyBasis,
         ::Union{JVec{TX}, AbstractVector{JVec{TX}}}) where {TX} =
   zeros(TX, length(J))
alloc_dB( J::OrthPolyBasis,
         ::Union{JVec{TX}, AbstractVector{JVec{TX}}}) where {TX} =
   zeros(TX, length(J))

function evaluate!(P, tmp, J::OrthPolyBasis, t; maxn=length(J))
   @assert length(P) >= maxn
   P[1] = J.A[1] * _fcut_(J.pl, J.tl, J.pr, J.tr, t)
   if maxn == 1; return P; end
   P[2] = (J.A[2] * t + J.B[2]) * P[1]
   if maxn == 2; return P; end
   @inbounds for n = 3:maxn
      P[n] = (J.A[n] * t + J.B[n]) * P[n-1] + J.C[n] * P[n-2]
   end
   return P
end

function evaluate_d!(P, dP, tmp, J::OrthPolyBasis, t; maxn=length(J))
   @assert maxn <= min(length(P), length(dP))

   P[1] = J.A[1] * _fcut_(J.pl, J.tl, J.pr, J.tr, t)
   dP[1] = J.A[1] * _fcut_d_(J.pl, J.tl, J.pr, J.tr, t)
   if maxn == 1; return dP; end

   α = J.A[2] * t + J.B[2]
   P[2] = α * P[1]
   dP[2] = α * dP[1] + J.A[2] * P[1]
   if maxn == 2; return dP; end

   @inbounds for n = 3:maxn
      α = J.A[n] * t + J.B[n]
      P[n] = α * P[n-1] + J.C[n] * P[n-2]
      dP[n] = α * dP[n-1] + J.C[n] * dP[n-2] + J.A[n] * P[n-1]
   end
   return dP
end

"""
`discrete_jacobi(N; pcut=0, tcut=1.0, pin=0, tin=-1.0, Nquad = 1000)`

A utility function to generate a jacobi-type basis
"""
function discrete_jacobi(N; pcut=0, tcut=1.0, pin=0, tin=-1.0, Nquad = 1000)
   tl, tr = minmax(tin, tcut)
   dt = (tr - tl) / Nquad
   tdf = range(tl + dt/2, tr - dt/2, length=Nquad)
   return OrthPolyBasis(N, pcut, tcut, pin, tin, tdf)
end

# ----------------------------------------------------------------
#   Transformed Polynomials Basis
# ----------------------------------------------------------------

struct TransformedPolys{T, TT, TJ} <: ACE1.ScalarBasis{T}
   J::TJ          # the actual basis
   trans::TT      # coordinate transform
   rl::T          # lower bound r
   ru::T          # upper bound r = rcut
end

==(J1::TransformedPolys, J2::TransformedPolys) = (
   (J1.J == J2.J) &&
   (J1.trans == J2.trans) &&
   (J1.rl == J2.rl) &&
   (J1.ru == J2.ru) )

function TransformedPolys(J, trans, rl, ru)
   # get the combined type if rl, ru are different
   T = promote_type(typeof(rl), typeof(ru))
   # integer is not allowed and we then default to float64
   if !(T <: AbstractFloat)
      T = Float64
   end
   rl_ = convert(T, rl)
   ru_ = convert(T, ru)
   return TransformedPolys(J, trans, rl_, ru_)
end

write_dict(J::TransformedPolys) = Dict(
      "__id__" => "ACE1_TransformedPolys",
      "J" => write_dict(J.J),
      "rl" => J.rl,
      "ru" => J.ru,
      "trans" => write_dict(J.trans)
   )

TransformedPolys(D::Dict) =
   TransformedPolys(
      read_dict(D["J"]),
      read_dict(D["trans"]),
      D["rl"],
      D["ru"]
   )

read_dict(::Val{:ACE1_TransformedPolys}, D::Dict) = TransformedPolys(D)

Base.length(J::TransformedPolys) = length(J.J)
fltype(P::TransformedPolys{T}) where {T} = T

function ACE1.rand_radial(J::TransformedPolys)
   t = ACE1.rand_radial(J.J)
   return inv_transform(J.trans, t)
end

cutoff(J::TransformedPolys) = J.ru

alloc_B( J::TransformedPolys, args...) = alloc_B(J.J, args...)
alloc_dB(J::TransformedPolys) = alloc_dB(J.J)
alloc_dB(J::TransformedPolys, N::Integer) = alloc_dB(J.J)

# in evaluate! and evaluate_d!: args... can be nothing or z, z0

function evaluate!(P, tmp, J::TransformedPolys, r, args...; maxn=length(J))
   # transform coordinates
   t = transform(J.trans, r, args...)
   # evaluate the actual polynomials
   evaluate!(P, nothing, J.J, t; maxn=maxn)
   return P
end

function evaluate_d!(P, dP, tmp, J::TransformedPolys, r, args...; maxn=length(J))
   # transform coordinates
   t = transform(J.trans, r, args...)
   dt = transform_d(J.trans, r, args...)
   # evaluate the actual Jacobi polynomials + derivatives w.r.t. x
   evaluate_d!(P, dP, nothing, J.J, t, maxn=maxn)
   @. dP *= dt
   return dP
end

"""
`transformed_jacobi(maxdeg, trans, rcut, rin = 0.0; kwargs...)` : construct
a `TransformPolys` basis with an inner polynomial basis of `OrthPolys` type.

* `maxdeg` : maximum degree
* `trans` : distance transform; normally `PolyTransform(...)`
* `rin, rcut` : inner and outer cutoff

**Keyword arguments:**

* `pcut = 2` : cutoff parameter
* `pin = 0` : inner cutoff parameter
* `Nquad = 1000` : number of quadrature points
"""
function transformed_jacobi(maxdeg::Integer,
                            trans::DistanceTransform,
                            rcut::Real, rin::Real = 0.0;
                            kwargs...)
   J =  discrete_jacobi(maxdeg; tcut = transform(trans, rcut),
                                tin = transform(trans, rin),
                                pcut = 2,
                                kwargs...)
   return TransformedPolys(J, trans, rin, rcut)
end

function transformed_jacobi(maxdeg::Integer,
                           trans::MultiTransform;
                           pcut = 2,
                           kwargs...)
   # this construction can only work if the transforms are then
   # sent to a single unified domain.
   @assert eltype(trans.transforms) <: ACE1.Transforms.AffineT
   @assert all( (t.y1 == -1) && (t.y2 == 1)
                 for t in trans.transforms )
   # obtain the maximum outer cutoff and minimum inner cutoff
   rin, rcut = ACE1.Transforms.cutoff_extrema(trans)
   # now construct the orthogonal polynomials with the [-1,1] domain.
   J =  discrete_jacobi(maxdeg; tcut = 1.0, tin = -1.0,
                                pcut = pcut, kwargs...)
   return TransformedPolys(J, trans, rin, rcut)
end

end


text:
type VertexArray{Vertex, Face, IT}
    id::GLuint
    length::Int
    buffer::Vector
    indices::IT
    context::GLAbstraction.GLContext
    function (::Type{VertexArray{Vertex, Face}}){Vertex, Face, IT}(id, bufferlength, buffers, indices::IT)
        new{Vertex, Face, IT}(id, bufferlength, buffers, indices, GLAbstraction.current_context())
    end
end
# TODO
Base.convert(::Type{VertexArray}, x) = VertexArray(x)
Base.convert(::Type{VertexArray}, x::VertexArray) = x

gl_face_enum{V, IT, T <: Integer}(::VertexArray{V, T, IT}) = GL_POINTS
gl_face_enum{V, IT, I}(::VertexArray{V, Face{1, I}, IT}) = GL_POINTS
gl_face_enum{V, IT, I}(::VertexArray{V, Face{2, I}, IT}) = GL_LINES
gl_face_enum{V, IT, I}(::VertexArray{V, Face{3, I}, IT}) = GL_TRIANGLES

gl_face_type(::Type{<: NTuple{2, <: AbstractVertex}}) = Face{2, Int}
gl_face_type(::Type) = Face{1, Int} # Default to Point
gl_face_type(::Type{T}) where T <: Face = T

# get_facetype(x::SubArray) = get_facetype(x.indices[1])

function draw_vbo{V, T, IT <: GLBuffer}(vbo::VertexArray{V, T, IT})
    glDrawElements(
        gl_face_enum(vbo),
        length(vbo.indices) * GLAbstraction.cardinality(vbo.indices),
        GLAbstraction.julia2glenum(eltype(IT)), C_NULL
    )
end
function draw_vbo{V, T}(vbo::VertexArray{V, T, DataType})
    glDrawArrays(gl_face_enum(vbo), 0, length(vbo))
end

Base.eltype{T, IT, N}(::VertexArray{T, IT, N}) = T
Base.length(x::VertexArray) = x.length

function VertexArray{T}(buffer::AbstractArray{T}, attrib_location = 0; face_type = gl_face_type(T))
    VertexArray(GLBuffer(buffer), face_type, attrib_location)
end
function VertexArray{T, AT <: AbstractArray, IT <: AbstractArray}(
        view::SubArray{T, 1, AT, Tuple{IT}, false}, attrib_location = 0; face_type = nothing # TODO figure out better ways then ignoring face type
    )
    indexes = view.indexes[1]
    buffer = view.parent
    VertexArray(GLBuffer(buffer), GLAbstraction.indexbuffer(indexes), attrib_location)
end

is_struct{T}(::Type{T}) = !(sizeof(T) != 0 && nfields(T) == 0)
is_glsl_primitive{T <: StaticVector}(::Type{T}) = true
is_glsl_primitive{T <: Union{Float32, Int32}}(::Type{T}) = true
is_glsl_primitive(T) = false

_typeof{T}(::Type{T}) = Type{T}
_typeof{T}(::T) = T
function VertexArray{T}(buffer::GLBuffer{T}, indices, attrib_location)
    id = glGenVertexArrays()
    glBindVertexArray(id)
    face_type = if isa(indices, GLBuffer)
        GLAbstraction.bind(indices)
        eltype(indices)
    elseif isa(indices, DataType) && indices <: Face
        indices
    elseif isa(indices, Integer)
        Face{1, OffsetInteger{1, GLint}}
    else
        error("indices must be Int, GLBuffer or Face type")
    end
    GLAbstraction.bind(buffer)
    if !is_glsl_primitive(T)
        for i = 1:nfields(T)
            FT = fieldtype(T, i); ET = eltype(FT)
            glVertexAttribPointer(
                attrib_location,
                GLAbstraction.cardinality(FT), GLAbstraction.julia2glenum(ET),
                GL_FALSE, sizeof(T), Ptr{Void}(fieldoffset(T, i))
            )
            glEnableVertexAttribArray(attrib_location)
            attrib_location += 1
        end
    else
        FT = T; ET = eltype(FT)
        glVertexAttribPointer(
            attrib_location,
            GLAbstraction.cardinality(FT), GLAbstraction.julia2glenum(ET),
            GL_FALSE, 0, C_NULL
        )
        glEnableVertexAttribArray(attrib_location)
    end
    glBindVertexArray(0)
    obj = VertexArray{T, face_type}(id, length(buffer), [buffer], indices)
    #finalizer(obj, GLAbstraction.free)
    obj
end

function compile_program(shaders...)
    program = GLAbstraction.createprogram()
    #attach new ones
    foreach(shaders) do shader
        glAttachShader(program, shader.id)
    end
    #link program
    glLinkProgram(program)
    if !GLAbstraction.islinked(program)
        for shader in shaders
            write(STDOUT, shader.source)
            println("---------------------------")
        end
        error(
            "program $program not linked. Error in: \n",
            join(map(x-> string(x.name), shaders), " or "), "\n", GLAbstraction.getinfolog(program)
        )
    end
    program
end

const max_batch_size = 1024

"""
Statically sized uniform buffer.
Supports push!, but with fixed memory, so it will error after reaching
it's preallocated length.
"""
type UniformBuffer{T, N}
    buffer::GLBuffer{T}
    offsets::NTuple{N, Int}
    elementsize::Int
    length::Int
end
const GLSLScalarTypes = Union{Float32, Int32, UInt32}
Base.eltype{T, N}(::UniformBuffer{T, N}) = T

function glsl_alignement_size(T)
    T <: Bool && return sizeof(Int32), sizeof(Int32)
    N = sizeof(T)
    T <: GLSLScalarTypes && return N, N
    T <: Function && return sizeof(Vec4f0), sizeof(Vec4f0) # sizeof(EmptyStruct) padded to Vec4f0
    ET = eltype(T)
    if T <: Mat4f0
        a, s = glsl_alignement_size(Vec4f0)
        return a, 4s
    end
    N = sizeof(ET)
    if T <: Vec2f0
        return 2N, 2N
    end
    if T <: Vec4f0
        return 4N, 4N
    end
    if T <: Vec3f0
        return 4N, 3N
    end
    error("Struct $T not supported yet. Please help by implementing all rules from https://khronos.org/registry/OpenGL/specs/gl/glspec45.core.pdf#page=159")
end

function std140_offsets{T}(::Type{T})
    elementsize = 0
    offsets = if T <: GLSLScalarTypes
        elementsize = sizeof(T)
        (0,)
    else
        offset = 0
        offsets = ntuple(nfields(T)) do i
            ft = fieldtype(T, i)
            alignement, sz = glsl_alignement_size(ft)
            if offset % alignement != 0
                offset = (div(offset, alignement) + 1) * alignement
            end
            of = offset
            offset += sz
            of
        end
        elementsize = offset
        offsets
    end
    offsets, elementsize
end

"""
    Pre allocates an empty buffer with `max_batch_size` size
    which can be used to store multiple uniform blocks of type T
"""
function UniformBuffer{T}(::Type{T}, max_batch_size = 1024, mode = GL_STATIC_DRAW)
    offsets, elementsize = std140_offsets(T)
    buffer = GLBuffer{T}(
        max_batch_size,
        elementsize * max_batch_size,
        GL_UNIFORM_BUFFER, mode
    )
    UniformBuffer(buffer, offsets, elementsize, 0)
end
Base.convert(::Type{UniformBuffer}, x) = UniformBuffer(x)
Base.convert(::Type{UniformBuffer}, x::UniformBuffer) = x
"""
    Creates an Uniform buffer with the contents of `data`
"""
function UniformBuffer{T}(data::T, mode = GL_STATIC_DRAW)
    buffer = UniformBuffer(T, 1, mode)
    push!(buffer, data)
    buffer
end

function assert_blocksize(buffer::UniformBuffer, program, blockname::String)
    block_index = glGetUniformBlockIndex(program, blockname)
    blocksize_ref = Ref{GLint}(0)
    glGetActiveUniformBlockiv(
        program, block_index,
        GL_UNIFORM_BLOCK_DATA_SIZE, blocksize_ref
    )
    blocksize = blocksize_ref[]
    @assert buffer.elementsize * length(buffer.buffer) == blocksize
end

_getfield(x::GLSLScalarTypes, i) = x
_getfield(x, i) = getfield(x, i)

function iterate_fields{T, N}(buffer::UniformBuffer{T, N}, x, index)
    offset = buffer.elementsize * (index - 1)
    x_ref = isimmutable(x) ? Ref(x) : x
    base_ptr = Ptr{UInt8}(pointer_from_objref(x_ref))
    ntuple(Val{N}) do i
        offset + buffer.offsets[i], base_ptr + fieldoffset(T, i), sizeof(fieldtype(T, i))
    end
end

function Base.setindex!{T, N}(buffer::UniformBuffer{T, N}, element::T, idx::Integer)
    if idx > length(buffer.buffer)
        throw(BoundsError(buffer, idx))
    end
    buff = buffer.buffer
    glBindBuffer(buff.buffertype, buff.id)
    dptr = Ptr{UInt8}(glMapBuffer(buff.buffertype, GL_WRITE_ONLY))
    for (offset, ptr, size) in iterate_fields(buffer, element, idx)
        unsafe_copy!(dptr + offset, ptr, size)
    end
    glUnmapBuffer(buff.buffertype)
    GLAbstraction.bind(buff, 0)
    element
end
extract_val(::Val{X}) where X = X
function Base.setindex!{T <: Composable, N, TF}(x::UniformBuffer{T, N}, val::TF, field::Type{<: Field})
    index = extract_val(FieldTraits.fieldindex(T, field))
    if index === 0
        throw(BoundsError(x, field))
    end
    val_conv = convert(fieldtype(T, index), val)
    val_ref = if isbits(val)
        Base.RefValue(val)
    elseif isimmutable(val)
        error("Struct $TF contains pointers and can't be transferred to GPU")
    else
        pointer_from_objref(val)
    end
    buff = x.buffer
    GLAbstraction.bind(buff) do
        glBufferSubData(buff.buffertype, x.offsets[index], sizeof(val_conv), val_ref)
    end
    x
end

function Base.getindex{T <: Composable, N}(x::UniformBuffer{T, N}, field::Type{<: Field})
    index = extract_val(FieldTraits.fieldindex(T, field)[1])
    if index == 0
        throw(BoundsError(x, field))
    end
    ET = fieldtype(T, index)
    val_ref = Ref{ET}()
    GLAbstraction.bind(x.buffer) do
        glGetBufferSubData(x.buffer.buffertype, x.offsets[index], sizeof(ET), val_ref)
    end
    val_ref[]
end

function Base.push!{T, N}(buffer::UniformBuffer{T, N}, element::T)
    buffer.length += 1
    buffer[buffer.length] = element
    buffer
end

immutable Command
    vertexCount::GLuint
    instanceCount::GLuint
    firstIndex::GLuint
    baseVertex::GLuint
    baseInstance::GLuint
end

export Command

ext: jl


text:
"""
    Tcmb(z::Real)

Computes the temperature of the CMB at redshift `z`.
"""
Tcmb(z::T) where T = ( 2.728 * ( 1.0 + z ) )

"""
    kSzPrefac(ν::Real, z::Real, DI_over_I::Bool)

Prefactor for the kinetic Sunyaev-Zel'dovich effect.
"""
function kSzPrefac(ν::Real, z::Real, DI_over_I::Bool)

    kSzPrefac  = -1.0 * σ_T / c_light

    if DI_over_I
        x   = h_planck * ν / ( k_B * Tcmb(z) )

        kSzPrefac *= exp(x) - 1 / (x * exp(x))
    end

    return kSzPrefac
end

"""
    kinetic_SZ(n_cm3::Real, vel_y_cgs::Real,
                    ν::Real=1.e9, z::Real=0.0;
                    DI_over_I::Bool=false)

Computes the kinetic Sunyaev-Zel'dovich effect from electron density `n_cm3` and velocity in y-direction to the projection plane in cgs units `vel_y_cgs`.
If `DI_over_I` is set to `true` you also need to provide an observation frequency `ν` and redshift `z`.
"""
function kinetic_SZ(n_cm3::Real, vel_y_cgs::Real,
                    ν::Real=1.e9, z::Real=0.0;
                    DI_over_I::Bool=false)
    return kSzPrefac(ν, z, DI_over_I) * n_cm3 * vel_y_cgs
end

"""
    comptonY(n_cm3::Real, T_K::Real, z::Real)

Computes the Compton-Y parameter from electron density `n_cm3` and temperature `T` in Kelvin at redshift `z`.
"""
function comptonY(n_cm3::Real, T_K::Real, z::Real)
    return yPrefac * n_cm3 * ( T_K - Tcmb(z) )
end

"""
    tSzPrefac(ν::Real, z::Real)

Computes the prefactor for the thermal Sunyaev-Zel'dovich effect.
"""
function tSzPrefac(ν::Real, z::Real, DI_over_I::Bool)

    x   = h_planck * ν / ( k_B * Tcmb(z) )
    tSzPrefac = (x * (exp(x) + 1.0) / (exp(x) - 1.0) - 4.0)

    if DI_over_I
        tSzPrefac *= exp(x) - 1 / (x * exp(x))
    end

    return tSzPrefac
end

"""
    thermal_SZ(n_cm3::Real, T_K::Real, z::Real=0.0, ν::Real=1.44e9)

Computes the thermal Sunyaev-Zel'dovich effect for electron density `n_cm3` and temperature `T_K` in Kelvin at redshift `z` and observer frequency `ν`.
"""
function thermal_SZ(n_cm3::Real, T_K::Real,
                    z::Real=0.0, ν::Real=1.44e9;
                    DI_over_I::Bool=false)

    return tSzPrefac(ν, z, DI_over_I) * comptonY(n_cm3, T_K, z)
end


text:
# ------------ Address map ------------ #

abstract type AddressMap{K} end

@inline get_address_schema(am) = ()

# Leaves.
abstract type Leaf{K} <: AddressMap{K} end

struct Empty <: Leaf{Empty} end

abstract type Select <: Leaf{Select} end
struct SelectAll <: Select end
struct Special <: Leaf{Special} end

struct Value{K} <: Leaf{Value}
    val::K
end

struct Choice{S, K} <: Leaf{Choice}
    score::S
    val::K
end

@inline convert(::Type{Value}, c::Choice) = Value(get_value(c))
@inline convert(::Type{Choice}, c::Value) = Choice(missing, get_value(c))

@inline length(e::Empty) = 0
@inline length(v::Value) = 1
@inline length(c::Choice) = 1

@inline value_length(e::Empty) = 0
@inline value_length(v::Value{K}) where K = length(get_value(v))
@inline value_length(c::Choice{K}) where K = length(get_value(v))

@inline ndims(e::Empty) = 0
@inline ndims(v::Value{K}) where K = ndims(get_value(v))
@inline ndims(v::Choice{K}) where K = ndims(get_value(v))

@inline projection(c::Choice, tg::Empty) = 0.0, Empty()
@inline projection(c::Choice, tg::SelectAll) = c.score, c

@inline filter(fn, l::Leaf) = Empty()
@inline filter(fn, addr, l::Leaf) = Empty()

@inline select(c::Select) = c
@inline select(c::Value) = SelectAll()
@inline select(c::Choice) = SelectAll()

@inline get_score(c::Choice) = c.score

# ------------- Interfaces ------------ #

@inline has_sub(::Leaf, _) = false
@inline has_sub(::Leaf, ::Tuple{}) = false
@inline has_sub(::AddressMap, ::Tuple{}) = false

@inline haskey(::Leaf, _) = false
@inline haskey(::Leaf, ::Tuple{}) = false
@inline haskey(::AddressMap, ::Tuple{}) = false
@inline haskey(am::AddressMap, addr::Tuple{A}) where A <: Address = haskey(am, addr[1])
@inline function haskey(am::AddressMap, addr::Tuple)
    hd, tl = addr[1], addr[2 : end]
    has_sub(am, hd) && haskey(get_sub(am, hd), tl)
end

@inline set_sub!(::Leaf, args...) = error("(set_sub!): trying to set submap of an instance of Leaf type.\nThis normally happens because you've already assigned to this address, or part of the prefix of this address.")
@inline function set_sub!(am::AddressMap{K}, addr::Tuple{}, v::AddressMap{<: K}) where {A <: Address, K} end
@inline function set_sub!(am::AddressMap{K}, addr, ::Empty) where {A <: Address, K} end

@inline setindex!(am::AddressMap, v, a) = set_sub!(am, a, v)

@inline get_sub(::Leaf, _) = Empty()
function get_sub(am::AddressMap{K}, addr::Tuple{T})::Union{Empty, AddressMap{K}} where {K, T}
    get_sub(am, addr[1])
end
function get_sub(am::AddressMap{K}, addr::Tuple)::Union{Empty, AddressMap{K}} where K
    get_sub(get_sub(am, addr[1]), addr[2 : end])
end
#Zygote.@adjoint get_sub(a, addr) = get_sub(a, addr), ret_grad -> (ret_grad, nothing)

function shallow_iterator(::AddressMap{K}) where K end
@inline shallow_iterator(::Leaf) = ()

@inline isempty(::Empty) = true
@inline isempty(::Value) = false
@inline isempty(::Choice) = false
@inline isempty(::SelectAll) = false
@inline isempty(am::AddressMap) = all(shallow_iterator(am)) do (_, v)
    isempty(v)
end

@inline has_value(am::AddressMap, addr) = has_value(get_sub(am, addr))
@inline has_value(am::AddressMap) = false

@inline get_value(v::Empty) = missing
@inline get_value(v::Value) = v.val
@inline get_value(v::Choice) = v.val
@inline has_value(v::Value) = true
@inline has_value(v::Choice) = true

@inline Base.:(==)(a::Value, b::Value) = a.val == b.val

# TODO: inefficient.
function Base.:(==)(a::AddressMap, b::AddressMap)
    for (k, sub) in shallow_iterator(a)
        get_sub(b, k) != sub && return false
    end
    for (k, sub) in shallow_iterator(b)
        get_sub(a, k) != sub && return false
    end
    return true
end

function Base.length(a::AddressMap)
    l = 0
    for (_, sub) in shallow_iterator(a)
        l += length(sub)
    end
    l
end

function value_length(a::AddressMap)
    l = 0
    for (_, sub) in shallow_iterator(a)
        l += value_length(sub)
    end
    l
end

function Base.ndims(a::AddressMap)
    l = 0
    for (_, sub) in shallow_iterator(a)
        l += ndims(sub)
    end
    l
end

@inline Base.merge(am::AddressMap, ::Empty) = deepcopy(am), false
@inline Base.merge(::Empty, am::AddressMap) = deepcopy(am), false
@inline Base.merge(l::Leaf, ::Empty) = deepcopy(l), false
@inline Base.merge(::Empty, l::Leaf) = deepcopy(l), false
@inline Base.merge(c::Choice, v::Value) = deepcopy(c), true
@inline Base.merge(v::Value, c::Choice) = Value(get_value(c)), true
@inline Base.merge(v1::Value, v2::Value) = Value(get_value(v2)), true
@inline Base.merge!(am::AddressMap, ::Empty) = false
@inline Base.merge!(l::Leaf, ::Empty) = false
@inline function Base.merge!(c::Choice, v::Value)
    c.value = get_value(v)
    c.score = missing
    true
end
@inline function Base.merge!(v::Value, c::Choice)
    v.val = get_value(c)
    true
end

function collect!(par::T, addrs::Vector, chd::Dict, v::V, meta) where {T <: Tuple, V <: Leaf{Choice}}
    push!(addrs, par)
    chd[par] = get_value(v)
end

function collect!(par::T, addrs::Vector, chd::Dict, v::V, meta) where {T <: Tuple, V <: Leaf{Value}}
    push!(addrs, par)
    chd[par] = get_value(v)
end

function collect!(par::T, addrs::Vector, chd::Dict, v::SelectAll, meta) where T <: Tuple
    push!(addrs, par)
end

function collect!(par::T, addrs::Vector, chd::Dict, tr, meta) where T <: Tuple
    tr = convert(DynamicMap{Value}, tr)
    collect!(par, addrs, chd, tr, meta)
end

function collect(am::M) where M <: AddressMap
    addrs = Any[]
    chd = Dict()
    meta = Dict()
    collect!(addrs, chd, am, meta)
    return addrs, chd, meta
end

function iterate(fn, am::M) where M <: AddressMap
    for (k, v) in shallow_iterator(am)
        fn((k, v))
    end
end

function flatten(am::M) where M <: AddressMap
    addrs, chd, _ = collect(am)
    arr = array(am, Float64)
    addrs, arr, chd
end

function Base.display(am::M;
                      show_values = true,
                      show_types = false) where M <: AddressMap
    println(" ___________________________________\n")
    println("             Address Map\n")
    addrs, chd, meta = collect(am)
    if show_values
        for a in addrs
            if haskey(meta, a) && haskey(chd, a)
                println(" $(meta[a]) $(a) = $(chd[a])")
            elseif haskey(chd, a)
                println(" $(a) = $(chd[a])")
            else
                println(" $(a)")
            end
        end
    elseif show_types
        for a in addrs
            if haskey(meta, a) && haskey(chd, a)
                println(" $(meta[a]) $(a) = $(typeof(chd[a]))")
            elseif haskey(chd, a)
                println(" $(a) = $(typeof(chd[a]))")
            else
                println(" $(a)")
            end
        end
    elseif show_types && show_values
        for a in addrs
            if haskey(meta, a) && haskey(chd, a)
                println(" $(meta[a]) $(a) = $(chd[a]) : $(typeof(chd[a]))")
            elseif haskey(chd, a)
                println(" $(a) = $(chd[a]) : $(typeof(chd[a]))")
            else
                println(" $(a)")
            end
        end
    else
        for a in addrs
            println(" $(a)")
        end
    end
    println(" ___________________________________\n")
end

# ------------ includes ------------ #

include("address_maps/dynamic.jl")
include("address_maps/static.jl")
include("address_maps/solo.jl")
include("array_compat.jl")


text:
#=
Copyright (c) 2016 Saurav Sachidanand

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
=#

cbor_tag(::UInt8) = ADDNTL_INFO_UINT8
cbor_tag(::UInt16) = ADDNTL_INFO_UINT16
cbor_tag(::UInt32) = ADDNTL_INFO_UINT32
cbor_tag(::UInt64) = ADDNTL_INFO_UINT64

cbor_tag(::Float64) = ADDNTL_INFO_FLOAT64
cbor_tag(::Float32) = ADDNTL_INFO_FLOAT32
cbor_tag(::Float16) = ADDNTL_INFO_FLOAT16

function encode_unsigned_with_type(
        io::IO, typ::UInt8, num::Unsigned
    )
    write(io, typ | cbor_tag(num))
    write(io, bswap(num))
end

function encode_length(io::IO, typ::UInt8, x)
    encode_smallest_int(io, typ, x isa String ? sizeof(x) : length(x))
end

"""
Array lengths and other integers (e.g. tags) in CBOR are encoded with smallest integer type,
which we do with this method!
"""
function encode_smallest_int(io::IO, typ::UInt8, num::Integer)
    @assert num >= 0 "array lengths must be greater 0. Found: $num"
    if num < SINGLE_BYTE_UINT_PLUS_ONE
        write(io, typ | UInt8(num)) # smaller 24 gets directly stored in type tag
    elseif num < UINT8_MAX_PLUS_ONE
        encode_unsigned_with_type(io, typ, UInt8(num))
    elseif num < UINT16_MAX_PLUS_ONE
        encode_unsigned_with_type(io, typ, UInt16(num))
    elseif num < UINT32_MAX_PLUS_ONE
        encode_unsigned_with_type(io, typ, UInt32(num))
    elseif num < UINT64_MAX_PLUS_ONE
        encode_unsigned_with_type(io, typ, UInt64(num))
    else
        error("128-bits ints can't be encoded in the CBOR format.")
    end
end

function encode(io::IO, float::Union{Float64, Float32, Float16})
    write(io, TYPE_7 | cbor_tag(float))
    # hton only works for 32 + 64, while bswap works for all
    write(io, Base.bswap_int(float))
end

# ------- straightforward encoding for a few Julia types
function encode(io::IO, bool::Bool)
    write(io, CBOR_FALSE_BYTE + bool)
end

function encode(io::IO, num::Unsigned)
    encode_unsigned_with_type(io, TYPE_0, num)
end

function encode(io::IO, num::T) where T <: Signed
    encode_unsigned_with_type(io, TYPE_1, unsigned(-num - one(T)))
end

function encode(io::IO, byte_string::Vector{UInt8})
    encode_length(io, TYPE_2, byte_string)
    write(io, byte_string)
end

function encode(io::IO, string::String)
    encode_length(io, TYPE_3, string)
    write(io, string)
end

function encode(io::IO, list::Vector)
    encode_length(io, TYPE_4, list)
    for e in list
        encode(io, e)
    end
end

function encode(io::IO, map::Dict)
    encode_length(io, TYPE_5, map)
    for (key, value) in map
        encode(io, key)
        encode(io, value)
    end
end

function encode(io::IO, big_int::BigInt)
    tag = if big_int < 0
        big_int = -big_int - 1
        NEG_BIG_INT_TAG
    else
        POS_BIG_INT_TAG
    end
    hex_str = hex(big_int)
    if isodd(length(hex_str))
        hex_str = "0" * hex_str
    end
    encode(io, Tag(tag, hex2bytes(hex_str)))
end

function encode(io::IO, tag::Tag)
    tag.id >= 0 || error("Tag needs to be a positive integer")
    encode_with_tag(io, Unsigned(tag.id), tag.data)
end

"""
Wrapper for collections with undefined length, that will then get encoded
in the cbor format. Underlying is just
"""
struct UndefLength{ET, A}
    iter::A
end

function UndefLength(iter::T) where T
    UndefLength{eltype(iter), T}(iter)
end

function UndefLength{T}(iter::A) where {T, A}
    UndefLength{T, A}(iter)
end

Base.iterate(x::UndefLength) = iterate(x.iter)
Base.iterate(x::UndefLength, state) = iterate(x.iter, state)

# ------- encoding for indefinite length collections
function encode(
        io::IO, iter::UndefLength{ET}
    ) where ET
    if ET in (Vector{UInt8}, String)
        typ = ET == Vector{UInt8} ? TYPE_2 : TYPE_3
        write(io, typ | ADDNTL_INFO_INDEF)
        foreach(x-> encode(io, x), iter)
    else
        typ = ET <: Pair ? TYPE_5 : TYPE_4 # Dict or any array
        write(io, typ | ADDNTL_INFO_INDEF)
        for e in iter
            if e isa Pair
                encode(io, e[1])
                encode(io, e[2])
            else
                encode(io, e)
            end
        end
    end
    write(io, BREAK_INDEF)
end

# ------- encoding with tags

function encode_with_tag(io::IO, tag::Unsigned, data)
    encode_smallest_int(io, TYPE_6, tag)
    encode(io, data)
end

struct Undefined
end

function encode(io::IO, null::Nothing)
    write(io, CBOR_NULL_BYTE)
end

function encode(io::IO, undef::Undefined)
    write(io, CBOR_UNDEF_BYTE)
end

struct SmallInteger{T}
    num::T
end
Base.convert(::Type{<: SmallInteger}, x) = SmallInteger(x)
Base.convert(::Type{<: SmallInteger}, x::SmallInteger) = x
Base.:(==)(a::SmallInteger, b::SmallInteger) = a.num == b.num
Base.:(==)(a::Number, b::SmallInteger) = a == b.num
Base.:(==)(a::SmallInteger, b::Number) = a.num == b

function encode(io::IO, small::SmallInteger{<: Unsigned})
    encode_smallest_int(io, TYPE_0, small.num)
end
function encode(io::IO, small::SmallInteger{<: Signed})
    if small.num >= 0
        encode_smallest_int(io, TYPE_0, unsigned(small.num))
    else
        encode_smallest_int(io, TYPE_1, unsigned(-small.num - one(small.num)))
    end
end

function fields2array(typ::T) where T
    fnames = fieldnames(T)
    getfield.((typ,), [fnames...])
end

"""
Any Julia type get's serialized as Tag 27
Tag             27
Data Item       array [typename, constructargs...]
Semantics       Serialised language-independent object with type name and constructor arguments
Reference       http://cbor.schmorp.de/generic-object
Contact         Marc A. Lehmann <cbor@schmorp.de>
"""
function encode(io::IO, struct_type::T) where T
    # TODO don't use Serialization for the whole struct!
    # It almost works to deserialize from just the fields and type,
    # but that ends up being problematic for
    # anonymous functions (the type changes between serialization & deserialization)
    tio = IOBuffer(); serialize(tio, struct_type)
    encode(
        io,
        Tag(
            CUSTOM_LANGUAGE_TYPE,
            [string("Julia/", T), take!(tio), fields2array(struct_type)]
        )
    )
end


text:
using Random, Turing, Test
import AbstractMCMC
import MCMCChains
import Turing.Inference
using Turing.RandomMeasures

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

@testset "gibbs.jl" begin
    @turing_testset "gibbs constructor" begin
        N = 500
        s1 = Gibbs(HMC(0.1, 5, :s, :m))
        s2 = Gibbs(PG(10, :s, :m))
        s3 = Gibbs(PG(3, :s), HMC( 0.4, 8, :m))
        s4 = Gibbs(PG(3, :s), HMC(0.4, 8, :m))
        s5 = Gibbs(CSMC(3, :s), HMC(0.4, 8, :m))
        s6 = Gibbs(HMC(0.1, 5, :s), ESS(:m))

        c1 = sample(gdemo_default, s1, N)
        c2 = sample(gdemo_default, s2, N)
        c3 = sample(gdemo_default, s3, N)
        c4 = sample(gdemo_default, s4, N)
        c5 = sample(gdemo_default, s5, N)
        c6 = sample(gdemo_default, s6, N)

        # Test gid of each samplers
        g = Turing.Sampler(s3, gdemo_default)

        @test g.state.samplers[1].selector != g.selector
        @test g.state.samplers[2].selector != g.selector
        @test g.state.samplers[1].selector != g.state.samplers[2].selector

        # run sampler: progress logging should be disabled and
        # it should return a Chains object
        @test sample(gdemo_default, g, N) isa MCMCChains.Chains
    end
    @numerical_testset "gibbs inference" begin
        Random.seed!(100)
        alg = Gibbs(
            CSMC(10, :s),
            HMC(0.2, 4, :m))
        chain = sample(gdemo(1.5, 2.0), alg, 3000)
        check_numerical(chain, [:s, :m], [49/24, 7/6], atol=0.1)

        Random.seed!(100)

        alg = Gibbs(
            MH(:s),
            HMC(0.2, 4, :m))
        chain = sample(gdemo(1.5, 2.0), alg, 5000)
        check_numerical(chain, [:s, :m], [49/24, 7/6], atol=0.1)

        alg = Gibbs(
            CSMC(15, :s),
            ESS(:m))
        chain = sample(gdemo(1.5, 2.0), alg, 10_000)
        check_numerical(chain, [:s, :m], [49/24, 7/6], atol=0.1)

        alg = CSMC(10)
        chain = sample(gdemo(1.5, 2.0), alg, 5000)
        check_numerical(chain, [:s, :m], [49/24, 7/6], atol=0.1)

        setadsafe(true)

        Random.seed!(200)
        gibbs = Gibbs(
            PG(10, :z1, :z2, :z3, :z4),
            HMC(0.15, 3, :mu1, :mu2))
        chain = sample(MoGtest_default, gibbs, 1500)
        check_MoGtest_default(chain, atol = 0.15)

        setadsafe(false)

        Random.seed!(200)
        gibbs = Gibbs(
            PG(10, :z1, :z2, :z3, :z4),
            ESS(:mu1), ESS(:mu2))
        chain = sample(MoGtest_default, gibbs, 1500)
        check_MoGtest_default(chain, atol = 0.15)
    end

    @turing_testset "transitions" begin
        @model gdemo_copy() = begin
            s ~ InverseGamma(2, 3)
            m ~ Normal(0, sqrt(s))
            1.5 ~ Normal(m, sqrt(s))
            2.0 ~ Normal(m, sqrt(s))
            return s, m
        end
        model = gdemo_copy()

        function AbstractMCMC.sample_end!(
            ::AbstractRNG,
            ::typeof(model),
            ::Turing.Sampler{<:Gibbs},
            ::Integer,
            transitions::Vector;
            kwargs...
        )
            transitions isa Vector{<:Inference.Transition} ||
                error("incorrect transitions")
            return
        end

        function callback(rng, model, sampler, transition, i; kwargs...)
            transition isa Inference.GibbsTransition || error("incorrect transition")
            return
        end

        alg = Gibbs(MH(:s), HMC(0.2, 4, :m))
        sample(model, alg, 100; callback = callback)
    end

    @turing_testset "dynamic model" begin
        @model imm(y, alpha, ::Type{M}=Vector{Float64}) where {M} = begin
            N = length(y)
            rpm = DirichletProcess(alpha)

            z = tzeros(Int, N)
            cluster_counts = tzeros(Int, N)
            fill!(cluster_counts, 0)

            for i in 1:N
                z[i] ~ ChineseRestaurantProcess(rpm, cluster_counts)
                cluster_counts[z[i]] += 1
            end

            Kmax = findlast(!iszero, cluster_counts)
            m = M(undef, Kmax)
            for k = 1:Kmax
                m[k] ~ Normal(1.0, 1.0)
            end
        end
        model = imm(randn(100), 1.0);
        sample(model, Gibbs(MH(10, :z), HMC(0.01, 4, :m)), 100);
        sample(model, Gibbs(PG(10, :z), HMC(0.01, 4, :m)), 100);
    end
end


text:
# CSCI5573Final/main.jl
# Licensed under the MIT License. See LICENSE.md file in the project root for
# full license information.

using SimLynx

include("config.jl")
include("task.jl")
include("daggen.jl")
include("utilities.jl")
include("components.jl")
include("schedulers.jl")

debug          = false
verbose        = false
output_results = false
stochasticish  = false

COMMS        = nothing
TASKS        = nothing
PROCESSORS   = nothing
READY_QUEUE  = nothing
TERMINATED   = nothing

COMPLETE     = false
FINISH_TIME  = ""

function run_sim(dag_file="")
    RUN_NAME = ""
    tasklist, num_tasks = nothing, nothing

    if stochasticish
        RUN_NAME = "RandomDag$(rand(0:10000))"
        tasklist, num_tasks = daggen(num_tasks = MAX_TASKS)
    else
        RUN_NAME = replace(dag_file, "."=>"")
        tasklist, num_tasks = read_daggen("SampleDAGS/"*dag_file)
    end

    RUN_PATH = SIM_RUN * "/" * RUN_NAME

    if !isdir(RUN_PATH)
        mkdir(RUN_PATH)
    end

    @simulation begin
        if debug
            current_trace!(true)
        end

        global TASKS, COMMS = ListToDictDAG(tasklist, "$RUN_PATH/dagGraph.dot")

        global READY_QUEUE = FifoQueue{Int64}()

        global TERMINATED = []

        global PROCESSORS = []

        for i = 1:length(PROCESSOR_POWERS)
            r = Resource(1, "PROCESSOR $i")
            m = PROCESSOR_POWERS[i]
            push!(PROCESSORS, Processor(m, r))
        end

        @schedule now Enqueuer()

        @schedule at 0 FCFSScheduler()
        # @schedule at 0 HEFTScheduler()
        # @schedule at 0 PEFTScheduler()
        # @schedule at 0 ProposedScheduler()

        start_simulation()

        if output_results
            print_results(RUN_PATH)
        end
    end
    return RUN_PATH
end

function main()
    if !isdir(SIM_RUN)
        mkdir(SIM_RUN)
    end

    println("Starting simulation")
    for dagfile in DAGS
        for i in 1:TRIALS
            if i === TRIALS
                global output_results = true
            end
            run_path = run_sim(dagfile)
            open("$(run_path)/runtimes.txt", "a") do file
                write(file, "$(FINISH_TIME)\n")
            end
        end
        global output_results = false
    end

    return
end

main()


text:
using Test
using CartesianMesh

CM = CartesianMesh

function allequal(v1,v2)
    return all(v1 .≈ v2)
end

x0 = [5.,8.]
widths = [2.,3.]
nelements = [5,3]
mesh = CM.UniformMesh(x0,widths,nelements)

@test_throws BoundsError CM.linear_to_cartesian_index(mesh,0)
@test_throws BoundsError CM.linear_to_cartesian_index(mesh,16)

elx,ely = CM.linear_to_cartesian_index(mesh,8)
@test (elx,ely) == (3,2)
@test CM.linear_to_cartesian_index(mesh,6) == (2,3)
@test CM.linear_to_cartesian_index(mesh,1) == (1,1)
@test CM.linear_to_cartesian_index(mesh,15) == (5,3)

mesh = CM.UniformMesh(x0,widths,[1,1])
@test CM.linear_to_cartesian_index(mesh,1) == (1,1)

xL,xR = CM.element(mesh,1)
@test allequal(xL,[5.,8.])
@test allequal(xR,[7.,11.])

mesh = CM.UniformMesh(x0,widths,[5,3])
xL,xR = CM.element(mesh,9)
@test allequal(xL,[5.8,10.])
@test allequal(xR,[6.2,11.])

nbrs = CM.neighbors(mesh,1)
@test allequal(nbrs,[mesh.bdryflag,4,2,mesh.bdryflag])
nbrs = CM.neighbors(mesh,5)
@test allequal(nbrs,[4,8,6,2])
nbrs = CM.neighbors(mesh,3)
@test allequal(nbrs,[2,6,mesh.bdryflag,mesh.bdryflag])

@test CM.faces_per_cell(mesh) == 4


text:
using Pkg
Pkg.activate(joinpath(@__DIR__,".."))
using HDF5

# the file layout coming from plmc is
# nsites            int32   (1)
# ncodes            int32   (1)
# nseqs             int32   (1)
# nskippedseqs      int32   (1)
# niter             int32   (1)
# theta             cfloat  (1)
# lh                cfloat  (1)
# le                cfloat  (1)
# lg                cfloat  (1)
# neff              cfloat  (1)
# alphabet          int8    (ncodes)
# ??                cfloat  (nseqs + nskippedseqs)
# focus seq         int8    (nsites)
# offset map        int32   (nsites)
# 1pmarginals       cfloat  (nsites*ncodes)
# fields            cfloat  (nsites*ncodes)
# 2pmarginals       cfloat  (choose(sites,2) * ncodes * ncodes)
# couplings         cfloat  (choose(sites,2) * ncodes * ncodes

function parse_plmc(plmc_file, out_file)

    # open file
    fid = open(plmc_file)

    # get nsites, ncodes, nseqs and nskippedseqs
    buffer = zeros(UInt8, 4*sizeof(Cint))
    readbytes!(fid, buffer)
    nsites, ncodes, nseqs, nskippedseqs = Int.(reinterpret(Cint, buffer))

    @show nsites, ncodes, nseqs, nskippedseqs

    # skip the niter, theta, lh, le, lg, neff, ??? and focus_seq
    to_skip = sizeof(Cint)*(1+nsites) + sizeof(UInt8)*(ncodes+nsites) +  sizeof(Cfloat)*(5 + nseqs + nskippedseqs)
    skip(fid, to_skip)

    # get 1pmarginals (we only use them as consistency check)
    nfields = nsites * ncodes
    buffer = zeros(UInt8, sizeof(Cfloat)*nfields)
    readbytes!(fid, buffer)

    # the sum of the 1pmarginals should be equal to nsites
    @assert abs(sum(reinterpret(Cfloat, buffer)) - nsites) < 0.01

    # read fields
    readbytes!(fid, buffer)
    fields = Float64.(reinterpret(Cfloat, buffer))

    # get 2pmarginals
    npairs = binomial(nsites, 2)
    ncouplings = npairs * ncodes * ncodes
    buffer = zeros(UInt8, sizeof(Cfloat)*ncouplings)
    readbytes!(fid, buffer)

    # the sum of the 2pmarginals should be equal to binomial(nsites, 2)
    @assert (sum(reinterpret(Cfloat, buffer)) - npairs) < 0.01

    # read couplings
    readbytes!(fid, buffer)
    couplings_flat = Float64.(reinterpret(Cfloat, buffer))

    # close file
    close(fid)

    # format correctly
    fields = permutedims(reshape(fields, ncodes, nsites))

    couplings = zeros(npairs, ncodes, ncodes)
    k = 1
    for p in 1:npairs
        for a in 1:ncodes
            for b in 1:ncodes
                couplings[p, a, b] = couplings_flat[k]
                k += 1
            end
        end
    end

    if isfile(out_file)
        rm(out_file)
    end
    h5write(out_file, "couplings", couplings)
    h5write(out_file, "fields", fields)
    h5write(out_file, "constant", 0.0)

end

plmc_file = ARGS[1]
out_file = ARGS[2]

if !isfile(plmc_file)
    println("File not found: $plmc_file")
    exit()
end

parse_plmc(plmc_file, out_file)

max_stars_repo_name: christophfeinauer/PairwiseDistillations

max_issues_repo_name: christophfeinauer/PairwiseDistillations

max_forks_repo_name: christophfeinauer/PairwiseDistillations


text:
import Base: twiceprecision, rat, narrow, isbetween

function splitprec(::Type{F}, i::SInteger) where {F<:AbstractFloat}
    hi = truncbits(F(i), cld(precision(F), F(2)))
    ihi = oftype(i, hi)
    hi, i - ihi
end

truncmask(x::SFloat16, mask) = SFloat16{reinterpret(Float16, values(mask) & reinterpret(UInt16, values(x)))}()
truncmask(x::SFloat32, mask) = SFloat32{reinterpret(Float32, values(mask) & reinterpret(UInt32, values(x)))}()
truncmask(x::SFloat64, mask) = SFloat64{reinterpret(Float64, values(mask) & reinterpret(UInt64, values(x)))}()

truncbits(x::F, nb) where {F<:Union{SFloat16, SFloat32, SFloat64}} =
    truncmask(x, Base.typemax(Base.uinttype(F)) << nb)

const AbstractTPI64 = Union{TPVal{<:SInt64,<:SInt64},TwicePrecision{Int64}}
const AbstractTPF64 = Union{TPVal{<:SFloat64,<:SFloat64},TwicePrecision{Float64}}
const AbstractTPF32 = Union{TPVal{<:SFloat32,<:SFloat32},TwicePrecision{Float32}}
const AbstractTPF16 = Union{TPVal{<:SFloat16,<:SFloat16},TwicePrecision{Float16}}

const AbstractTwicePrecision = Union{AbstractTPI64,AbstractTPF64,AbstractTPF32,AbstractTPF16}

"gethi - retrieve the the hi bits (most significant bits) from a twice precision number"
function gethi end

"getlo - retrieve the the low bits (least significant bits) from a twice precision number"
function getlo end

gethi(x::TwicePrecision) = x.hi
getlo(x::TwicePrecision) = x.lo

Base.@pure gethi(::TPVal{H,L}) where {H,L} = H()::H
Base.@pure gethi(::Type{<:TPVal{H,L}}) where {H,L} = H()::H

Base.@pure getlo(::TPVal{H,L}) where {H,L} = L()::L
Base.@pure getlo(::Type{<:TPVal{H,L}}) where {H,L} = L()::L

eltype(::TPVal{H,L}) where {H,L} = eltype(H)

Base.@pure values(::TPVal{H,L}) where {H,L} = H()::H + L()::L
Base.@pure values(::Type{<:TPVal{H,L}}) where {H,L} = H()::H + L()::L

TPVal(hi::T) where T = TPVal(hi, zero(eltype(hi)))

TPVal(hi::SReal, lo::SReal) = TPVal(promote(hi, lo)...)

TPVal(hi::SFloat64{H}, lo::SFloat64{L}) where {H,L} = TPVal{SFloat64{H},SFloat64{L}}()
TPVal(hi::SFloat32{H}, lo::SFloat32{L}) where {H,L} = TPVal{SFloat32{H},SFloat32{L}}()
TPVal(hi::SFloat16{H}, lo::SFloat16{L}) where {H,L} = TPVal{SFloat16{H},SFloat16{L}}()
TPVal(hi::SInt64{H}, lo::SInt64{L}) where {H,L} = TPVal{SInt64{H},SInt64{L}}()

TPVal(x::TwicePrecision{T}) where T = TPVal(x.hi, x.lo)

TPVal(::Type{T}, ::Tuple{N,D}) where {N<:SInteger,D<:SInteger,T<:Union{Float16,Float32}} =
    TPVal(ofeltype(T, N()/D()))

TPVal(::Type{T}, ::Tuple{N,D}) where {T,N<:SReal,D<:SReal} = TPVal(ofeltype(T, N())) / D()
TPVal(::Type{Float64}, nd::Tuple{N,D}) where {T,N<:SReal,D<:SReal} =
    TPVal(ofeltype(Float64, nd[1]::N)) / nd[2]::D

TPVal(::Type{T}, nd::Tuple, nb::Integer) where {T} = Base.twiceprecision(TPVal(T, nd), nb)

TPVal(hi::T, lo::T) where T<:BaseNumber = TPVal(SVal(hi), SVal(lo))

function TPVal(::Type{T}, x::SReal{X}) where {X,T}
    xT = ofeltype(T, x)
    TPVal(xT, ofeltype(T, x - xT))
end
#TPVal{T}(x::SReal) where T = TPVal{T}(x, zero(x))

TPVal(x::SReal) = TPVal(eltype(x), x)

TPVal(::Type{T}, i::SInteger{X}) where {X,T<:AbstractFloat} =
    TPVal(T, canonicalize2(splitprec(T, i)...)...)

function TPVal(::Type{T}, hi::H, lo::L) where {T,H<:SReal,L<:SReal}
    if T === eltype(hi) === eltype(lo)
        return TPVal{H,L}()
    else
        return TPVal(T, ofeltype(T, hi), ofeltype(T, lo))
    end
end

SReal(x::TPVal) = gethi(x) + getlo(x)

#---
# Handle ambiguities created by mixed static base types

"""
    tpval

constructs a twice precision value
"""
tpval(::Type{T}, hi::Real) where T = tpval(T, hi, zero(hi))

# ↓ this accounts for mixed static/base types
tpval(::Type{T}, hi::Real, lo::Real) where T = tpval(T, promote(hi, lo)...)

tpval(::Type{T}, hi::T, lo::T) where T<:BaseReal = TwicePrecision{T}(hi, lo)
tpval(::Type{T}, frac::Tuple{<:SReal,<:SReal}, n::SInteger) where T = TPVal(T, frac, n)
tpval(::Type{T}, frac::Tuple{<:BaseNumber,<:BaseNumber}, n::Number) where T = TwicePrecision{T}(frac, values(n))

tpval(hi::T, lo::T) where T= TwicePrecision(hi, lo)
# Numerator/Denominator constructors
function tpval(::Type{T}, nd::Tuple{BaseInteger,BaseInteger}) where  {T<:Union{Float16,Float32}}
    n, d = nd
    TwicePrecision{T}(n/d)
end

function tpval(::Type{T}, nd::Tuple{<:SInteger,BaseInteger}) where  {T<:Union{Float16,Float32}}
    n, d = nd
    TwicePrecision{T}(values(n)/d)
end

function tpval(::Type{T}, nd::Tuple{BaseInteger,<:SInteger}) where  {T<:Union{Float16,Float32}}
    n, d = nd
    TwicePrecision{T}(n/values(d))
end

function tpval(::Type{T}, nd::Tuple{<:SInteger,<:SInteger}) where  {T<:Union{Float16,Float32}}
    n, d = nd
    tpval(T, n/d)
end

function tpval(::Type{T}, nd::Tuple{<:SReal,<:SReal}) where {T}
    n, d = nd
    TPVal(ofeltype(T, n)) / d
end

function tpval(::Type{T}, nd::Tuple{<:SReal,BaseReal}) where {T}
    n, d = nd
    TwicePrecision{T}(values(n)) / d
end

function tpval(::Type{T}, nd::Tuple{BaseReal,<:SReal}) where {T}
    n, d = nd
    TwicePrecision{T}(n) / values(d)
end

function tpval(::Type{T}, nd::Tuple{BaseReal,BaseReal}) where {T}
    n, d = nd
    TwicePrecision{T}(n) / d
end

#=
for N in (SReal,BaseReal)
for D in (SReal,BaseReal)
    if N == BaseReal || D == BaseReal
        @eval begin
            tpval(::Type{T}, x::Tuple{$N,$D}) where T = TwicePrecision{T}(values(x[1])) / values(x[2])
        end
    else
        @eval begin
            tpval(::Type{T}, x::Tuple{$N,$D}) where T = TPVal(ofeltype(T, x[1])) / x[2]
        end
    end
for I in (SInteger,BaseInteger)
    if N == BaseReal || D == BaseReal || I == BaseInteger
        @eval begin
            tpval(::Type{T}, x::Tuple{$N,$D}, nb::$I) where T =
                Base.twiceprecision(TwicePrecision{T}(values(x[1]),values(x[2])), values(nb))
        end
    else
        @eval begin
            tpval(::Type{T}, x::Tuple{$N,$D}, nb::$I) where T = Base.twiceprecision(TPVal(T, x[1], x[2]), nb)
        end
    end
end
end
end
=#

tpval(hi::SFloat64{H}, lo::SFloat64{L}) where {H,L} = TPVal{SFloat64{H},SFloat64{L}}()
tpval(hi::SFloat32{H}, lo::SFloat32{L}) where {H,L} = TPVal{SFloat32{H},SFloat32{L}}()
tpval(hi::SFloat16{H}, lo::SFloat16{L}) where {H,L} = TPVal{SFloat16{H},SFloat16{L}}()
tpval(x::SReal) = TPVal(x)
tpval(x::BaseNumber) = TwicePrecision(x)

tpval(nd::Tuple{N,D}, nb::Integer) where {T,N,D} = Base.twiceprecision(tpval(nd), nb)

#---
# Numerator/Denominator constructors
TPVal(nd::Tuple{<:SReal,<:SReal}, nb::SInteger) where T =
    Base.twiceprecision(TPVal(nd), nb)

# Truncating constructors. Useful for generating values that can be
# exactly multiplied by small integers.
function Base.twiceprecision(v::SReal, nb::SInteger)
    hi = Base.truncbits(v, nb)
    TPVal(hi, v-hi)
end

for T in (SFloat16,SFloat32,SFloat64)
    @eval begin
        function Base.twiceprecision(val::TPVal{<:$T,<:$T}, b::SInteger)
            hi = Base.truncbits(gethi(val), b)
            TPVal(gethi(val), (gethi(val) - hi) + getlo(val))
        end
    end
end

#---conversion
TwicePrecision(::TPVal{H,L}) where {H,L,T} = TwicePrecision(values(H),values(L))

(::Type{T})(x::TPVal{H,L}) where {T<:Number,H,L} = T(values(x))::T

for (ST1,BT1) in S2B
    for (ST2,BT2) in S2B
        if ST1 == ST2
            @eval begin
                (::Type{<:TPVal{<:$ST1,<:$ST1}})(val::TPVal{$ST2{H},$ST2{L}}) where {H,L} = val
            end
        else
            @eval begin
                (::Type{<:TPVal{<:$ST1,<:$ST1}})(val::TPVal{$ST2{H},$ST2{L}}) where {H,L} =
                    TPVal{$ST1{$BT1(H::$BT2)},$ST1{$BT1(L::$BT2)}}()
            end
        end
    end
end

Base.float(x::TPVal{Tuple{H,L},T}) where {H,L,T<:AbstractFloat} = x
Base.float(x::TPVal{Tuple{H,L},T}) where {H,L,T} = TPVal{Tuple{float(H::T),float(L::T)},float(T)}()

#Base.big(::TPVal{T,H,L}) where {T,H,L} = big(H) + big(L)

-(x::TPVal) = TPVal(-gethi(H), -getlo(L))

Base.zero(::Type{<:TPVal{T}}) where {T} = TPVal(zero(T), zero(T))
Base.one(::Type{<:TPVal{T}}) where {T} = TPVal(one(T), one(T))

# Arithmetic

@inline function +(x::TPVal{H,L}, y::SReal) where {H,L,T}
    s_hi, s_lo = Base.add12(gethi(x)::H, y)
    hnew, hlow = Base.canonicalize2(s_hi, s_lo+getlo()::L)
    TPVal(hnew, hlow)
end

+(x::SReal, y::TPVal) = y+x

@inline function +(x::TPVal{Hx,Lx}, y::TPVal{Hy,Ly}) where {Hx,Lx,Hy,Ly}
    hx = gethi(x)::Hx
    lx = getlo(x)::Lx
    hy = gethi(y)::Hy
    ly = getlo(y)::Ly
    r = hx + hy
    s = abs(hx) > abs(hy) ? (((hx - r) + hy) + ly) + lx : (((hy - r) + hx) + lx) + ly
    hnew, lnew = canonicalize2(r, s)
    TPVal(hnew,lnew)
end

@inline function +(x::TPVal{H,L}, y::TwicePrecision) where {H,L}
    hx = gethi(x)::H
    lx = getlo(x)::L
    r = hx + y.hi
    s = abs(hx) > abs(y.hi) ? (((hx - r) + y.hi) + y.lo) + lx : (((y.hi - r) + hx) + lx) + y.lo
    hnew, lnew = canonicalize2(r, s)
    TwicePrecision(hnew, lnew)
end

@inline +(x::TwicePrecision, y::TPVal) = y + x

-(x::TPVal, y::TPVal) = x + (-y)
-(x::BaseReal, y::TPVal) = x + (-y)
-(x::TPVal, y::BaseReal) = x + (-y)

#
*(x::TPVal, v::BaseNumber) = TwicePrecision(x) * v
*(x::TPVal, v::TwicePrecision) = TwicePrecision(x) * v
*(v::TwicePrecision, x::TPVal) = x * v
*(v::Number, x::TPVal) = x * v

function *(x::TPVal{H,L}, v::SReal) where {H,L}
    v == 0 && return TPVal(H() * v, L() * v)
    x * TPVal(oftype(H() * v, v))
end

for T in (SFloat16, SFloat32, SFloat64)
    @eval begin
        function *(x::TPVal{H,L}, v::SInteger) where {H<:$T,L<:$T}
            v == 0 && return TPVal(H() * v, L() * v)
            nb = ceil(Int, log2(abs(v)))
            u = Base.truncbits(H(), nb)
            tpval(canonicalize2(u * v, H() - u + L() * v)...)
        end
    end
end

@inline function *(x::TPVal{Hx,Lx}, y::TPVal{Hy,Ly}) where {Hx,Lx,Hy,Ly}
    zh, zl = Base.mul12(gethi(x), gethi(y))
    ifelse(iszero(zh) | !isfinite(zh), TPVal(zh, zh), TPVal(canonicalize2(zh, gethi(x) * getlo(y) + getlo(x) * gethi(y))...))
end

#= TODO: shouldn't need this anymore
function *(x::TPVal{H,L}, s::SInteger) where {H,L,T<:Union{Float16, Float32, Float64}}
    v == 0 && return TPVal(H::T1 * s, L::T1 * s)
    nb = ceil(Int, log2(abs(s)))
    u = Base.truncbits(H::T1, get(nb))
    TPVal(canonicalize2(SReal(u* s), SReal((H::T1-u) + L::T1*s))...)
end
=#

# TODO: check this

#*(x::TPVal{V1,T1}, y::TPVal{V2,T2}) where {V1,V2,T1,T2} = *(promote(x, y)...)
/(x::TPVal{H,L}, v::Real) where {H,L} = x / tpval(ofeltype(gethi(x)::H/v, v))

/(x::TPVal, v::BaseNumber) = x / TwicePrecision(oftype(gethi(x)/v, v))
/(x::TPVal, v::SNumber) = x / TPVal(ofeltype(gethi(x)/v, v))

/(v::BaseNumber, x::TPVal) = TwicePrecision(oftype(gethi(x)/v, v)) / x
/(v::SNumber, x::TPVal) = TPVal(ofeltype(gethi(x)/v, v)) / x

/(v::TwicePrecision, x::TPVal) = v / TwicePrecision(x)

/(v::TPVal, x::TwicePrecision) = TwicePrecision(v) / x

function /(x::TPVal{Hx,Lx}, y::TPVal{Hy,Ly}) where {Hx,Lx,Hy,Ly}
    hi = gethi(x) / gethi(y)
    uh, ul = Base.mul12(hi, gethi(y))
    lo = ((((gethi(x) - uh) - ul) + getlo(x)) - hi*getlo(y))/gethi(y)
    ret = TPVal(canonicalize2(hi, lo)...)
    ifelse(iszero(hi) | !isfinite(hi), TPVal(hi, hi), ret)
end

const nbits16 = SInt(Base.cld(precision(Float16), 2))
const nbits32 = SInt(Base.cld(precision(Float32), 2))
const nbits64 = SInt(Base.cld(precision(Float64), 2))

# lack of specificity in base requires that these be more verbosely written out
Base.nbitslen(::Type{Float16}, l::SReal{V1}, f::SReal{V2}) where {V1,V2} =
    min(nbits16, Base.nbitslen(l, f))
Base.nbitslen(::Type{Float32}, l::SReal{V1}, f::SReal{V2}) where {V1,V2} =
    min(nbits32, Base.nbitslen(l, f))
Base.nbitslen(::Type{Float64}, l::SReal{V1}, f::SReal{V2}) where {V1,V2} =
    min(nbits64, Base.nbitslen(l, f))

Base.nbitslen(::Type{<:SFloat16}, l::SReal{V1}, f::SReal{V2}) where {V1,V2} =
    min(nbits16, Base.nbitslen(l, f))
Base.nbitslen(::Type{<:SFloat32}, l::SReal{V1}, f::SReal{V2}) where {V1,V2} =
    min(nbits32, Base.nbitslen(l, f))
Base.nbitslen(::Type{<:SFloat64}, l::SReal{V1}, f::SReal{V2}) where {V1,V2} =
    min(nbits64, Base.nbitslen(l, f))

Base.nbitslen(len::SReal, offset::SReal) = len < 2 ? SZero : ceil(Int, log2(max(offset-SOne, len-offset))) + SOne

#nbitslen(r::StaticRange{T,B,E,S,F,L}) where {T,B,E,S,F,L} = nbitslen(T, L, F)
#@inline ($f)(x::Float64) = nan_dom_err(ccall(($(string(f)), libm), Float64, (Float64,), x), x)

#

#nbithelper(::Type{T}) where T = cld(SReal(precision(T)), SReal(2))

#=
function _rat(x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    f = trunc(Int, y)
    ynew = y
    ynew -= f
    anew = f*a + c
    cnew = a
    bnew = f*b + d
    dnew = b
    if max(abs(anew), abs(bnew)) <= Base.convert(Int,m)
        return cnew, dnew
    elseif oftype(x, anew)/oftype(x, bnew) == x
        return anew, bnew
    elseif ynew <= m
        ynew = inv(ynew)
        _rat(x, ynew, m, anew, bnew, cnew, dnew)
    else
        return anew, bnew
    end
end
=#

narrow(::Type{T}) where {T<:SFloat} = SFloat64
narrow(::Type{<:SFloat64}) = SFloat32
narrow(::Type{<:SFloat32}) = SFloat16
narrow(::Type{<:SFloat16}) = SFloat16

#=
function _rat(x::SVal{X}, y::SVal{Y}, m::SVal{M}, a::SVal{A}, b::SVal{B}, c::SVal{C}, d::SVal{D}) where {X,Y,M,A,B,C,D}
    f = trunc(Int, y)
    ynew = y
    ynew -= f
    anew = f*a + c
    cnew = a
    bnew = f*b + d
    dnew = b
    if max(abs(anew), abs(bnew)) <= ofeltype(Int, m)
        return cnew, dnew
    elseif ofeltype(x, anew)/oftype(x, bnew) == x
        return anew, bnew
    elseif abs(ynew) <= m
        ynew = inv(ynew)
        _rat(x, ynew, m, anew, bnew, cnew, dnew)
    else
        return anew, bnew
    end
end
=#
#=
function rat(v::SReal)
    _rat(v, v, maxintfloat(narrow(typeof(v)), Int), SOne, SZero, SZero, SOne)
end

function _rat(x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
   while abs(y) <= m
        f = trunc(Int, y)
        y -= f
        a, c = f*a + c, a
        b, d = f*b + d, b
        max(abs(a), abs(b)) <= convert(Int,m) || return c, d
        ofeltype(x,a) / ofeltype(x,b) == x && break
        y = inv(y)
    end
    return a, b
end
=#
@generated function rat(x::SReal{X}) where X
    n, d = rat(X)  #_getf(x, x, maxintfloat(narrow(typeof(x)), Int), SOne, SZero, SZero, SOne)
    quote
        $(SInt(n), SInt(d))
    end
end

function _getf(x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    _sub_y(trunc(Int, y), x, y, m, a, b, c, d)
end

function _sub_y(f::SReal, x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    if abs(y) > m
        a, b
    else
        _update_abcd(f, x, y - f, m, a, b, c, d)
    end
end

function _update_abcd(f::SReal, x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    _inv_y(x, y, m, f*a+c, f*b+d, a, b)
end

function _inv_y(x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    if max(abs(a), abs(b)) > ofeltype(Int, m)
        c, d
    elseif ofeltype(x,a) / ofeltype(x,b) == x
        a, b
    else
        _getf(x, inv(y), m, a, b, c, d)
    end
end
#=
function _rat(x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    if abs(y) <= m
        a, b
    else
        _ratsub(trunc(Int, y), x, y, m, a, b, c, d)
    end
end

function _ratsub(f::SReal, x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    _ratsub2(x, y, m, f*a+c, a, f*b+d, b)
end

function _ratsub2(x::SReal, y::SReal, m::SReal, a::SReal, b::SReal, c::SReal, d::SReal)
    if max(abs(a), abs(b)) <= ofeltype(Int,m)
        c, d
    elseif oftype(x,a)/oftype(x,b) == x
        a, b
    else
        _rat(x, y, m, a, b, c, d)
    end
end
=#

# this should be taken care of using promotion
#<(x::TPVal, y::SReal) = <(x, y)
#<(x::SReal, y::TPVal) = <(TPVal(x), y)

<(x::TPVal{Hx,Lx}, y::TPVal{Hy,Ly}) where {Hx,Lx,Hy,Ly} =
    gethi(x)::Hx < gethi(y)::Hy || ((gethi(x)::Hx == gethi(y)::Hy) & (getlo(x)::Lx < getlo(y)::Ly))

Base.show(io::IO, r::TPVal) = showsval(io, r)
Base.show(io::IO, ::MIME"text/plain", r::TPVal) = showsval(io, r)

showsval(io::IO, r::TPVal{H,L}) where {H,L} = print(io, "TPVal($H, $L)")


text:
module TwoCollidingBoxes

using ModiaLang
import Modia3D
using  Modia3D.ModiaInterface

vmat1 = Modia3D.VisualMaterial(color="Green", transparency=0.6)
vmat2 = Modia3D.VisualMaterial(color="Red"  , transparency=0.6)
cmat = "Steel"

collidingBoxes = Model(
    world = Object3D(feature=Scene(gravityField=UniformGravityField(g=9.81, n=[-1, 0, 0]),
                                          enableContactDetection=true)),

    fixedBox = Object3D(parent=:world,
                        translation=[-3.0, 0.0, -2.5],
                        feature=Solid(shape=Box(lengthX=3.0, lengthY=2.0, lengthZ=5.0),
                                      solidMaterial="Aluminium",
                                      visualMaterial=vmat1,
                                      contactMaterial=cmat,
                                      collision=true,
                                      collisionSmoothingRadius=0.001)),

    movingBox = Object3D(feature=Solid(shape=Box(lengthX=0.5, lengthY=0.5, lengthZ=0.5),
                                       solidMaterial="Aluminium",
                                       visualMaterial=vmat2,
                                       contactMaterial=cmat,
                                       collision=true,
                                       collisionSmoothingRadius=0.001)),
    movingBoxJoint = FreeMotion(obj1=:world, obj2=:movingBox,
                                r=Var(init=[0.0, 0.0, 0.0]),
                                rot=Var(init=[pi/2, 0.0, 0.0]))
)

twoCollidingBoxes = @instantiateModel(buildModia3D(collidingBoxes), unitless=true, log=false, logStateSelection=false, logCode=false)

stopTime = 2.0
tolerance = 1e-6
interval = 0.001
requiredFinalStates = [-3.3213935067112574, 3.837601056986228e-6, 1.6309713914861226, -6.472201388283604, 4.337577865425163e-6, 1.9041135911575697, 1.5707911214094912, -3.1891125895973017e-6, 4.024722484494325, -1.8500097559885116e-6, 1.6544557228815132e-5, 4.481758166130571]
simulate!(twoCollidingBoxes, stopTime=stopTime, tolerance=tolerance, interval=interval, log=true, logStates=true, logEvents=true, requiredFinalStates=requiredFinalStates)

@usingModiaPlot
plot(twoCollidingBoxes, ["movingBoxJoint.r" "movingBoxJoint.rot"; "movingBoxJoint.v" "movingBoxJoint.w"], figure=1)

end


text:
"""
    AbstractRuleset <: ModelParameters.AbstractModel

Abstract supertype for [`Ruleset`](@ref) objects and variants.
"""
abstract type AbstractRuleset <: AbstractModel end

# Getters
ruleset(rs::AbstractRuleset) = rs
rules(rs::AbstractRuleset) = rs.rules
boundary(rs::AbstractRuleset) = rs.boundary
proc(rs::AbstractRuleset) = rs.proc
opt(rs::AbstractRuleset) = rs.opt
cellsize(rs::AbstractRuleset) = rs.cellsize
timestep(rs::AbstractRuleset) = rs.timestep
radius(set::AbstractRuleset) = radius(rules(set))

Base.step(rs::AbstractRuleset) = timestep(rs)

# ModelParameters interface
Base.parent(rs::AbstractRuleset) = rules(rs)
ModelParameters.setparent!(rs::AbstractRuleset, rules) = rs.rules = rules
ModelParameters.setparent(rs::AbstractRuleset, rules) = @set rs.rules = rules

"""
    Rulseset <: AbstractRuleset

    Ruleset(rules...; kw...)

A container for holding a sequence of `Rule`s and simulation
details like boundary handing and optimisation.
Rules will be run in the order they are passed, ie. `Ruleset(rule1, rule2, rule3)`.

# Keywords

- `proc`: a [`Processor`](@ref) to specificy the hardware to run simulations on,
    like [`SingleCPU`](@ref), [`ThreadedCPU`](@ref) or [`CuGPU`](@ref) when
    KernelAbstractions.jl and a CUDA gpu is available.
- `opt`: a [`PerformanceOpt`](@ref) to specificy optimisations like
    [`SparseOpt`](@ref). Defaults to [`NoOpt`](@ref).
- `boundary`: what to do with boundary of grid edges.
    Options are `Remove()` or `Wrap()`, defaulting to [`Remove`](@ref).
- `cellsize`: size of cells.
- `timestep`: fixed timestep where this is required for some rules.
    eg. `Month(1)` or `1u"s"`.
"""
Base.@kwdef mutable struct Ruleset{B<:BoundaryCondition,P<:Processor,Op<:PerformanceOpt,C,T} <: AbstractRuleset
    # Rules in Ruleset are intentionally not type-stable.
    # But they are when rebuilt in a StaticRuleset later
    rules::Tuple{Vararg{<:Rule}} = ()
    boundary::B                  = Remove()
    proc::P                      = SingleCPU()
    opt::Op                      = NoOpt()
    cellsize::C                  = 1
    timestep::T                  = nothing
end
Ruleset(rules::Rule...; kw...) = Ruleset(; rules=rules, kw...)
Ruleset(rules::Tuple; kw...) = Ruleset(; rules=rules, kw...)
function Ruleset(rs::AbstractRuleset)
    Ruleset(
        rules(rs), boundary(rs), proc(rs), opt(rs), cellsize(rs), timestep(rs),
    )
end

struct StaticRuleset{R<:Tuple,B<:BoundaryCondition,P<:Processor,Op<:PerformanceOpt,C,T} <: AbstractRuleset
    rules::R
    boundary::B
    proc::P
    opt::Op
    cellsize::C
    timestep::T
end
function StaticRuleset(rs::AbstractRuleset)
    StaticRuleset(
        rules(rs), boundary(rs), proc(rs), opt(rs), cellsize(rs), timestep(rs),
    )
end


text:
# implementation of AbstractBinning for uniformly spaced bins

"""
  Implement uniformly spaced bins

  **Inputs**

  * vlow: the lowest value (the left endpoint of the first bin)
  * vhigh: the highest value (the right endpoint of the last bin)
  * N: the number of bins
  * comm: MPI communicator, defaults to MPI.COMM_WORLD
"""
mutable struct UniformBinning{K, V} <: AbstractBinning{K, V}
  vlow::V
  vhigh::V
  delta::V  # bin size
  N::Int    # number of bins
  bin_keys::Vector{Vector{K}}
  bin_vals::Vector{Vector{V}}

  # used for summation
  bin_keycount::Vector{Int}
  val_sum::Vector{V}

  comm::MPI.Comm

  function UniformBinning{K, V}(vlow, vhigh, N::Integer, comm::MPI.Comm=MPI.COMM_WORLD) where {K, V}

    @assert vhigh > vlow

    delta = (vhigh - vlow)/N
    bin_keys = Vector{Vector{K}}(N)
    bin_vals = Vector{Vector{V}}(N)
    for i=1:N
      bin_keys[i] = Vector{K}(0)
      bin_vals[i] = Vector{V}(0)
    end

    bin_keycount = zeros(Int, N)
    val_sum = zeros(Int, N)

    return new(vlow, vhigh, delta, N, bin_keys, bin_vals, bin_keycount, val_sum,
               comm)
  end
end

function calcBin(obj::UniformBinning, val)

  bin = Int(div(val - obj.vlow + obj.delta, obj.delta))

  # each bin is v_low <= v < v_high, so the value vhigh will be put in bin
  # N + 1.  Force it into bin N
  bin = min(bin, obj.N)

  return bin
end

function push!(obj::UniformBinning, key, val)

  bin = calcBin(obj, val)
  push!(obj.bin_keys[bin], key)
  push!(obj.bin_vals[bin], val)

  return nothing
end

function sumBins(obj::UniformBinning{K, V}) where {K, V}

  # the version of Allreduce with the same send and receive buffers is not
  # wrapped in MPI.jl, so use temporary arrays
  bin_keycount_local = zeros(Int, obj.N)
  val_sum_local = zeros(V, obj.N)

  for i=1:obj.N
    bin_keycount_local[i] = length(obj.bin_keys[i])
    for j=1:length(obj.bin_vals[i])
      val_sum_local[i] += obj.bin_vals[i][j]
    end
  end

  fill!(obj.bin_keycount, 0)
  fill!(obj.val_sum, zero(V))
  MPI.Allreduce!(bin_keycount_local, obj.bin_keycount, MPI.SUM, obj.comm)
  MPI.Allreduce!(val_sum_local, obj.val_sum, MPI.SUM, obj.comm)

  return ROView(obj.bin_keycount), ROView(obj.val_sum)
end


text:
# Prototype - copied from pluto_RayTracingWeekend.jl
# Adapted from [Ray Tracing In One Weekend by Peter Shirley](https://raytracing.github.io/books/RayTracingInOneWeekend.html) and
# [cshenton's Julia implementation](https://github.com/cshenton/RayTracing.jl)"
using Pkg
Pkg.activate(@__DIR__)

"""Next ideas:

- Incorporate ideas from https://discourse.julialang.org/t/why-is-this-raytracer-slow/59176/20
  - @inline to many functions
- save image, e.g. PNG
- The only remaining allocations that appear expensive are of `HitRecord`s (on the stack)
- pre-allocate X paths (ray bundles? not sure what the terminology is) then
  implement versions of hit, scatter, etc. that operate on an entire tensor at once.
  (i.e. efficiently parallelizable with multithreading, on SIMD or GPU)
  - use FieldVector?
- figure out the incorrect look in refraction of negatively scaled sphere
- use parameterizable structs for Vec but keep the same efficiency
- share with community, ask for feedback
- continue watching MIT course
"""

using BenchmarkTools, Images, InteractiveUtils, LinearAlgebra, StaticArrays
Option{T} = Union{Missing, T}
t_col = @SVector[0.4f0, 0.5f0, 0.1f0] # test color

squared_length(v::SVector{3,Float32}) = v ⋅ v

# Before optimization:
#   677-699 ns (41 allocations: 2.77 KiB) # squared_length(v::SVector) = v ⋅ v; @btime squared_length(t_col)
# inside a function:
#   894-906 ns (43 allocations: 2.83 KiB)
#
# don't define `Base.getproperty(vec::SVector{3}, sym::Symbol)`
#   208.527 ns (3 allocations: 80 bytes)
#
# just test `@btime squared_length(t_col)`:
#    12.115 ns (1 allocation: 16 bytes)
# test again after redefining: `Vec3 = SVector{3, Float64}` instead of `Vec3 = SVector{3}`
#    11.966 ns (1 allocation: 16 bytes)
# test again after redefining: `Vec3 = SVector{3, Float32}` instead of `Vec3 = SVector{3, Float64}`
#    13.652 ns (1 allocation: 16 bytes)
#    (SLOWER! I wonder if this would hold with larger vectors or if we become mem-bound)
# replace: squared_length(v::SVector{3,MyFloat})
#     3.256 ns (0 allocations: 0 bytes)
# replace: @btime squared_length($t_col) # notice the $
#     1.152 ns (0 allocations: 0 bytes)
#@btime squared_length($t_col)

@inline near_zero(v::SVector{3,Float32}) = squared_length(v) < 1e-5
#@btime near_zero($t_col) # 1.382 ns (0 allocations: 0 bytes)

# Test images

img = rand(4, 3)

img_rgb = rand(RGB{Float32}, 4, 4)

ex1 = [RGB{Float32}(1,0,0) RGB{Float32}(0,1,0) RGB{Float32}(0,0,1);
       RGB{Float32}(1,1,0) RGB{Float32}(1,1,1) RGB{Float32}(0,0,0)]

ex2 = zeros(RGB{Float32}, 2, 3)
ex2[1,1] = RGB{Float32}(1,0,0)
ex2

# The "Hello World" of graphics
function gradient(nx::Int, ny::Int)
	img = zeros(RGB{Float32}, ny, nx)
	for i in 1:ny, j in 1:nx # Julia is column-major, i.e. iterate 1 column at a time
		x = j; y = (ny-i);
		r = x/nx
		g = y/ny
		b = 0.2
		img[i,j] = RGB{Float32}(r,g,b)
	end
	img
end

gradient(200,100)

rgb(v::SVector{3,Float32}) = RGB{Float32}(v[1], v[2], v[3])
#rgb(v::Vec3) = RGB{MyFloat}(v...)
#@btime rgb($t_col) # 289.914 ns (4 allocations: 80 bytes)

@inline rgb_gamma2(v::SVector{3,Float32}) = RGB{Float32}(sqrt.(v)...)

#@btime rgb_gamma2($t_col) # 454.766 ns (11 allocations: 256 bytes)

struct Ray
	origin::SVector{3, Float32} # Point
	dir::SVector{3, Float32} # Vec3 # direction (unit vector)
end

# interpolates between blue and white
_origin = @SVector[0.0f0,0.0f0,0.0f0]
_v3_minusY = @SVector[0.0f0,-1.0f0,0.0f0]
_t_ray1 = Ray(_origin, _v3_minusY)
#@btime Ray($_origin, $_v3_minusY) # 6.341 ns (0 allocations: 0 bytes)

# equivalent to C++'s ray.at()
@inline point(r::Ray, t::Float32) = r.origin .+ t .* r.dir
#@btime point($_t_ray1, 0.5f0) # 1.392 ns (0 allocations: 0 bytes)

#md"# Chapter 4: Rays, simple camera, and background"

@inline function skycolor(ray::Ray)
	white = @SVector[1.0f0,1.0f0,1.0f0]
	skyblue = @SVector[0.5f0,0.7f0,1.0f0]
	t = 0.5f0(ray.dir[2] + 1.0f0)
    (1.0f0-t)*white + t*skyblue
end
#@btime skycolor($_t_ray1) # 1.402 ns (0 allocations: 0 bytes)

# before optim:
#   1.474 μs (23 allocations: 608 bytes)
# pre-allocate args:
#   1.150 μs (19 allocations: 480 bytes)
# `rgb(skycolor(t_ray1))`
#   1.143 μs (19 allocations: 480 bytes)
# `skycolor(t_ray1)`:
#       636.259 ns (9 allocations: 256 bytes)
#    extract t_white and t_skyblue from skycolor():
#       279.779 ns (5 allocations: 128 bytes)
#    pre-allocate the ray given in argument:
#        51.130 ns (5 allocations: 128 bytes)
#    skycolor!($t_col2, $t_ray1, $t_white, $t_skyblue) using @set to return the value:
#        14.727 ns (1 allocation: 32 bytes)
#    replace all the Color() by @SVector[...], don't use global colors (use local)
#         1.402 ns (0 allocations: 0 bytes)
# restore the rgb(), i.e. `@btime rgb(skycolor($t_ray1))`
#   291.492 ns (4 allocations: 80 bytes)
#@btime rgb(skycolor($t_ray1)) # 291.492 ns (4 allocations: 80 bytes)

# md"""# Random vectors
# C++'s section 8.1"""

@inline random_between(min=0.0f0, max=1.0f0) = rand(Float32)*(max-min) + min # equiv to random_double()
#@btime random_between(50f0, 100f0) # 4.519 ns (0 allocations: 0 bytes)

@inline random_vec3(min::Float32, max::Float32) = @SVector[random_between(min, max) for i ∈ 1:3]

# Before optimization:
#   352.322 ns (6 allocations: 224 bytes)
# Don't use list comprehension:
#   179.684 ns (5 allocations: 112 bytes)
# Return @SVector[] instead of Vec3():
#    12.557 ns (0 allocations: 0 bytes)
#@btime random_vec3(-1.0f0,1.0f0)

@inline random_vec2(min::Float32, max::Float32) = @SVector[random_between(min, max) for i ∈ 1:2]
#@btime random_vec2(-1.0f0,1.0f0) # 8.536 ns (0 allocations: 0 bytes)

@inline function random_vec3_in_sphere() # equiv to random_in_unit_sphere()
	while true
		p = random_vec3(-1f0, 1f0)
		if p⋅p <= 1
			return p
		end
	end
end
#@btime random_vec3_in_sphere() # 46.587 ns (0 allocations: 0 bytes)

squared_length(random_vec3_in_sphere())

"Random unit vector. Equivalent to C++'s `unit_vector(random_in_unit_sphere())`"
@inline random_vec3_on_sphere() = normalize(random_vec3_in_sphere())
#@btime random_vec3_on_sphere()

# Before optim:
#   517.538 ns (12 allocations: 418 bytes)... but random
# After reusing _rand_vec3:
#    92.865 ns (2 allocations: 60 bytes)
# Various speed-ups (don't use Vec3, etc.):
#    52.427 ns (0 allocations: 0 bytes)
random_vec3_on_sphere()
norm(random_vec3_on_sphere())

@inline function random_vec2_in_disk() # equiv to random_in_unit_disk()
	while true
		p = random_vec2(-1f0, 1f0)
		if p⋅p <= 1
			return p
		end
	end
end
#@btime random_vec2_in_disk() # 16.725 ns (0 allocations: 0 bytes)

""" Temporary function to shoot rays through each pixel. Later replaced by `render`

	Args:
		scene: a function that takes a ray, returns the color of any object it hit
"""
function main(nx::Int, ny::Int, scene)
	lower_left_corner = @SVector[-2,-1,-1]
	horizontal = @SVector[4,0,0]
	vertical = @SVector[0f0,2f0,0f0]
	origin = @SVector[0,0,0]

	img = zeros(RGB{Float32}, ny, nx)
	for i in 1:ny, j in 1:nx # Julia is column-major, i.e. iterate 1 column at a time
		u = j/nx
		v = (ny-i)/ny # Y-up!
		ray = Ray(origin, normalize(lower_left_corner + u*horizontal + v*vertical))
		#r = x/nx
		#g = y/ny
		#b = 0.2
		img[i,j] = rgb(scene(ray))
	end
	img
end

# Before optimizations:
#   28.630 ms (520010 allocations: 13.05 MiB)
# After optimizations (skycolor, rand, etc.)
#   10.358 ms (220010 allocations: 5.42 MiB) # at this stage, rgb() is most likely the culprit..
# After replacing Color,Vec3 by @SVector:
#    6.539 ms (80002 allocations: 1.75 MiB)
# Lots more optimizations, including @inline of low-level functions:
# 161.546 μs (     2 allocations: 234.45 KiB)
main(200,100, skycolor)

#md"# Chapter 5: Add a sphere"

@inline function hit_sphere1(center::SVector{3,Float32}, radius::Float32, r::Ray)
	oc = r.origin - center
	a = r.dir ⋅ r.dir
	b = 2oc ⋅ r.dir
	c = (oc ⋅ oc) - radius*radius
	discriminant = b*b - 4a*c
	discriminant > 0
end

@inline function sphere_scene1(r::Ray)
	if hit_sphere1(@SVector[0f0,0f0,-1f0], 0.5f0, r) # sphere of radius 0.5 centered at z=-1
		return @SVector[1f0,0f0,0f0] # red
	else
		skycolor(r)
	end
end

main(200,100,sphere_scene1) # 227.851 μs (2 allocations: 234.45 KiB)

#md"# Chapter 6: Surface normals and multiple objects"

@inline function hit_sphere2(center::SVector{3,Float32}, radius::Float32, r::Ray)
	oc = r.origin - center
	a = r.dir ⋅ r.dir
	b = 2oc ⋅ r.dir
	c = (oc ⋅ oc) - radius*radius
	discriminant = b*b - 4a*c
	if discriminant < 0
		return -1
	else
		return (-b - sqrt(discriminant)) / 2a
	end
end

@inline function sphere_scene2(r::Ray)
	sphere_center = SVector{3,Float32}(0f0,0f0,-1f0)
	t = hit_sphere2(sphere_center, 0.5f0, r) # sphere of radius 0.5 centered at z=-1
	if t > 0f0
		n⃗ = normalize(point(r, t) - sphere_center) # normal vector. typed n\vec
		return 0.5f0n⃗ + SVector{3,Float32}(0.5f0,0.5f0,0.5f0) # remap normal to rgb
	else
		skycolor(r)
	end
end

@btime main(200,100,sphere_scene2)  # 262.286 μs (2 allocations: 234.45 KiB)

"An object that can be hit by Ray"
abstract type Hittable end

"""Materials tell us how rays interact with a surface"""
abstract type Material end

struct NoMaterial <: Material end

const _no_material = NoMaterial()
const _y_up = @SVector[0f0,1f0,0f0]

"Record a hit between a ray and an object's surface"
mutable struct HitRecord
	t::Float32 # vector from the ray's origin to the intersection with a surface.

	# If t==Inf32, there was no hit, and all following values are undefined!
	#
	p::SVector{3,Float32} # point of the intersection between an object's surface and a ray
	n⃗::SVector{3,Float32} # surface's outward normal vector, points towards outside of object?

	# If true, our ray hit from outside to the front of the surface.
	# If false, the ray hit from within.
	front_face::Bool
	mat::Material

	@inline HitRecord() = new(Inf32) # no hit!
	@inline HitRecord(t,p,n⃗,front_face,mat) = new(t,p,n⃗,front_face,mat)
end

struct Sphere <: Hittable
	center::SVector{3,Float32}
	radius::Float32
	mat::Material
end

# md"""The geometry defines an `outside normal`. A HitRecord stores the `local normal`.
# ![Surface normal](https://raytracing.github.io/images/fig-1.06-normal-sides.jpg)
# """

"""Equivalent to `hit_record.set_face_normal()`"""
@inline function ray_to_HitRecord(t, p, outward_n⃗, r_dir::SVector{3,Float32}, mat::Material)
	front_face = r_dir ⋅ outward_n⃗ < 0
	n⃗ = front_face ? outward_n⃗ : -outward_n⃗
	rec = HitRecord(t,p,n⃗,front_face,mat)
end

struct Scatter
	r::Ray
	attenuation::SVector{3,Float32}

	# claforte: TODO: rename to "absorbed?", i.e. not reflected/refracted?
	reflected::Bool # whether the scattered ray was reflected, or fully absorbed
	@inline Scatter(r,a,reflected=true) = new(r,a,reflected)
end

#"Diffuse material"
mutable struct Lambertian<:Material
	albedo::SVector{3,Float32}
end

"""Compute reflection vector for v (pointing to surface) and normal n⃗.

	See [diagram](https://raytracing.github.io/books/RayTracingInOneWeekend.html#metal/mirroredlightreflection)"""
#reflect(v::SVector{3,Float32}, n⃗::SVector{3,Float32}) = normalize(v - (2v⋅n⃗)*n⃗) # claforte: normalize not needed?
@inline reflect(v::SVector{3,Float32}, n⃗::SVector{3,Float32}) = v - (2v⋅n⃗)*n⃗

@assert reflect(@SVector[0.6f0,-0.8f0,0f0], @SVector[0f0,1f0,0f0]) == @SVector[0.6f0,0.8f0,0f00] # diagram's example

"""Create a scattered ray emitted by `mat` from incident Ray `r`.

	Args:
		rec: the HitRecord of the surface from which to scatter the ray.

	Return missing if it's fully absorbed. """
@inline function scatter(mat::Lambertian, r::Ray, rec::HitRecord)::Scatter
	scatter_dir = rec.n⃗ + random_vec3_on_sphere()
	if near_zero(scatter_dir) # Catch degenerate scatter direction
		scatter_dir = rec.n⃗
	else
		scatter_dir = normalize(scatter_dir)
	end
	scattered_r = Ray(rec.p, scatter_dir)
	attenuation = mat.albedo
	return Scatter(scattered_r, attenuation)
end

const _no_hit = HitRecord()

@inline function hit(s::Sphere, r::Ray, tmin::Float32, tmax::Float32)#::Option{HitRecord} # the return type seems to cause excessive allocations
    oc = r.origin - s.center
    a = 1 #r.dir ⋅ r.dir # normalized vector - always 1
    half_b = oc ⋅ r.dir
    c = oc⋅oc - s.radius^2
    discriminant = half_b^2 - a*c
	if discriminant < 0 return _no_hit end
	sqrtd = √discriminant

	# Find the nearest root that lies in the acceptable range
	root = (-half_b - sqrtd) / a
	if root < tmin || tmax < root
		root = (-half_b + sqrtd) / a
		if root < tmin || tmax < root
			return _no_hit
		end
	end

	t = root
	p = point(r, t)
	n⃗ = (p - s.center) / s.radius
	return ray_to_HitRecord(t, p, n⃗, r.dir, s.mat)
end

struct HittableList <: Hittable
    list::Vector{Hittable}
end

#"""Find closest hit between `Ray r` and a list of Hittable objects `h`, within distance `tmin` < `tmax`"""
@inline function hit(hittables::HittableList, r::Ray, tmin::Float32, tmax::Float32)# ::HitRecord
    closest = tmax # closest t so far
    rec = _no_hit
    for h in hittables.list
        temprec = hit(h, r, tmin, closest)
        if temprec !== _no_hit
            rec = temprec
            closest = rec.t # i.e. ignore any further hit > this one's.
        end
    end
    rec
end

@inline color_vec3_in_rgb(v::SVector{3,Float32}) = 0.5normalize(v) + @SVector[0.5f,0.5f,0.5f]

#md"# Metal material"

mutable struct Metal<:Material
	albedo::SVector{3,Float32}
	fuzz::Float32 # how big the sphere used to generate fuzzy reflection rays. 0=none
	@inline Metal(a,f=0.0) = new(a,f)
end

@inline function scatter(mat::Metal, r_in::Ray, rec::HitRecord)::Scatter
	reflected = normalize(reflect(r_in.dir, rec.n⃗) + mat.fuzz*random_vec3_on_sphere())
	Scatter(Ray(rec.p, reflected), mat.albedo)
end

#md"# Scenes"

#"Scene with 2 Lambertian spheres"
function scene_2_spheres()::HittableList
	spheres = Sphere[]

	# small center sphere
	push!(spheres, Sphere(@SVector[0f0,0f0,-1f0], 0.5, Lambertian(@SVector[0.7f0,0.3f0,0.3f0])))

	# ground sphere (planet?)
	push!(spheres, Sphere(@SVector[0f0,-100.5f0,-1f0], 100, Lambertian(@SVector[0.8f0,0.8f0,0.0f0])))
	HittableList(spheres)
end

#"""Scene with 2 Lambertian, 2 Metal spheres.
#
#	See https://raytracing.github.io/images/img-1.11-metal-shiny.png"""
function scene_4_spheres()::HittableList
	scene = scene_2_spheres()

	# left and right Metal spheres
	push!(scene.list, Sphere(@SVector[-1f0,0f0,-1f0], 0.5f0, Metal(@SVector[0.8f0,0.8f0,0.8f0], 0.3f0)))
	push!(scene.list, Sphere(@SVector[1f0,0f0,-1f0], 0.5f0, Metal(@SVector[0.8f0,0.6f0,0.2f0], 0.8f0)))
	return scene
end

#md"""# Camera

# Adapted from C++'s sections 7.2, 11.1 """
mutable struct Camera
	origin::SVector{3,Float32}
	lower_left_corner::SVector{3,Float32}
	horizontal::SVector{3,Float32}
	vertical::SVector{3,Float32}
	u::SVector{3,Float32}
	v::SVector{3,Float32}
	w::SVector{3,Float32}
	lens_radius::Float32
end

"""
	Args:
		vfov: vertical field-of-view in degrees
		aspect_ratio: horizontal/vertical ratio of pixels
      aperture: if 0 - no depth-of-field
"""
function default_camera(lookfrom::SVector{3,Float32}=@SVector[0f0,0f0,0f0], lookat::SVector{3,Float32}=@SVector[0f0,0f0,-1f0],
						vup::SVector{3,Float32}=@SVector[0f0,1f0,0f0], vfov=90.0f0, aspect_ratio=16.0f0/9.0f0,
						aperture=0.0f0, focus_dist=1.0f0)
	viewport_height = 2.0f0 * tand(vfov/2f0)
	viewport_width = aspect_ratio * viewport_height

	w = normalize(lookfrom - lookat)
	u = normalize(vup × w)
	v = w × u

	origin = lookfrom
	horizontal = focus_dist * viewport_width * u
	vertical = focus_dist * viewport_height * v
	lower_left_corner = origin - horizontal/2f0 - vertical/2f0 - focus_dist*w
	lens_radius = aperture/2f0
	Camera(origin, lower_left_corner, horizontal, vertical, u, v, w, lens_radius)
end

default_camera()

#md"# Render

@inline function get_ray(c::Camera, s::Float32, t::Float32)
	rd = SVector{2,Float32}(c.lens_radius * random_vec2_in_disk())
	offset = c.u * rd[1] + c.v * rd[2] #offset = c.u * rd.x + c.v * rd.y
    Ray(c.origin + offset, normalize(c.lower_left_corner + s*c.horizontal +
									 t*c.vertical - c.origin - offset))
end

#@btime get_ray(default_camera(), 0.0f0, 0.0f0)

"""Compute color for a ray, recursively

	Args:
		depth: how many more levels of recursive ray bounces can we still compute?"""
@inline function ray_color(r::Ray, world::HittableList, depth=4)
    if depth <= 0
		return @SVector[0f0,0f0,0f0]
	end

	rec = hit(world, r, 1f-4, Inf32)
    if rec !== _no_hit
		# For debugging, represent vectors as RGB:
		# return color_vec3_in_rgb(rec.p) # show the normalized hit point
		# return color_vec3_in_rgb(rec.n⃗) # show the normal in RGB
		# return color_vec3_in_rgb(rec.p + rec.n⃗)
		# return color_vec3_in_rgb(random_vec3_in_sphere())
		#return color_vec3_in_rgb(rec.n⃗ + random_vec3_in_sphere())

        s = scatter(rec.mat, r, rec)
		if s.reflected
			return s.attenuation .* ray_color(s.r, world, depth-1)
		else
			return @SVector[0f0,0f0,0f0]
		end
    else
        skycolor(r)
    end
end

"""Render an image of `scene` using the specified camera, number of samples.

	Args:
		scene: a HittableList, e.g. a list of spheres
		n_samples: number of samples per pixel, eq. to C++ samples_per_pixel

	Equivalent to C++'s `main` function."""
function render(scene::HittableList, cam::Camera, image_width=400,
				n_samples=1)
	# Image
	aspect_ratio = 16.0f0/9.0f0 # TODO: use cam.aspect_ratio for consistency
	image_height = convert(Int64, floor(image_width / aspect_ratio))

	# Render
	img = zeros(RGB{Float32}, image_height, image_width)
	# Compared to C++, Julia is:
	# 1. column-major, i.e. iterate 1 column at a time, so invert i,j compared to C++
	# 2. 1-based, so no need to subtract 1 from image_width, etc.
	# 3. The array is Y-down, but `v` is Y-up
	for i in 1:image_height, j in 1:image_width
		accum_color = @SVector[0f0,0f0,0f0]
		for s in 1:n_samples
			u = convert(Float32, j/image_width)
			v = convert(Float32, (image_height-i)/image_height) # i is Y-down, v is Y-up!
			if s != 1 # 1st sample is always centered, for 1-sample/pixel
				# claforte: I think the C++ version had a bug, the rand offset was
				# between [0,1] instead of centered at 0, e.g. [-0.5, 0.5].
				u += convert(Float32, (rand()-0.5f0) / image_width)
				v += convert(Float32 ,(rand()-0.5f0) / image_height)
			end
			ray = get_ray(cam, u, v)
			accum_color += ray_color(ray, scene)
		end
		img[i,j] = rgb_gamma2(accum_color / n_samples)
	end
	img
end

# After some optimization:
#  46.506 ms (917106 allocations: 16.33 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#  14.137 ms (118583 allocations: 4.14 MiB)
# Don't specify return value of Option{HitRecord} in hit()
#  24.842 ms (527738 allocations: 16.63 MiB)
# Don't initialize unnecessary elements in HitRecord():
#  14.862 ms (118745 allocations: 4.15 MiB)  (but computer was busy...)
# Replace MyFloat by Float32:
#  11.792 ms ( 61551 allocations: 2.88 MiB)
# Remove ::HitRecord return value in remaining hit() method:
#  11.545 ms ( 61654 allocations: 2.88 MiB)
# with mutable HitRecord
#  11.183 ms ( 61678 allocations: 2.88 MiB) (insignificant?)
# @inline tons of stuff:
#   8.129 ms ( 61660 allocations: 2.88 MiB)
render(scene_2_spheres(), default_camera(), 96, 16)

render(scene_4_spheres(), default_camera(), 96, 16)

#md"""# Dielectrics

# from Section 10.2 Snell's Law:
# ![Ray refraction](https://raytracing.github.io/images/fig-1.13-refraction.jpg)

# Refracted angle `sinθ′ = (η/η′)⋅sinθ`, where η (\eta) are the refractive indices.

# Split the parts of the ray into `R′=R′⊥+R′∥` (perpendicular and parallel to n⃗′)."""

# """
# 	Args:
# 		refraction_ratio: incident refraction index divided by refraction index of
# 			hit surface. i.e. η/η′ in the figure above"""
@inline function refract(dir::SVector{3,Float32}, n⃗::SVector{3,Float32}, refraction_ratio::Float32)
	cosθ = min(-dir ⋅ n⃗, 1)
	r_out_perp = refraction_ratio * (dir + cosθ*n⃗)
	r_out_parallel = -√(abs(1-squared_length(r_out_perp))) * n⃗
	normalize(r_out_perp + r_out_parallel)
end

@assert refract(SVector{3,Float32}(0.6,-0.8,0), SVector{3,Float32}(0,1,0), 1.0f0) == SVector{3,Float32}(0.6,-0.8,0) # unchanged

t_refract_widerθ = refract(SVector{3,Float32}(0.6,-0.8,0), SVector{3,Float32}(0,1,0), 2.0f0) # wider angle
@assert isapprox(t_refract_widerθ, SVector{3,Float32}(0.87519, -0.483779, 0.0); atol=1e-3)

t_refract_narrowerθ = refract(SVector{3,Float32}(0.6,-0.8,0), SVector{3,Float32}(0,1,0), 0.5f0) # narrower angle
@assert isapprox(t_refract_narrowerθ, SVector{3,Float32}(0.3, -0.953939, 0.0); atol=1e-3)

mutable struct Dielectric <: Material
	ir::Float32 # index of refraction, i.e. η.
end

@inline function reflectance(cosθ::Float32, refraction_ratio::Float32)
	# Use Schlick's approximation for reflectance.
	# claforte: may be buggy? I'm getting black pixels in the Hollow Glass Sphere...
	r0 = (1f0-refraction_ratio) / (1f0+refraction_ratio)
	r0 = r0^2
	r0 + (1f0-r0)*((1f0-cosθ)^5)
end

@inline function scatter(mat::Dielectric, r_in::Ray, rec::HitRecord)
	attenuation = @SVector[1f0,1f0,1f0]
	refraction_ratio = rec.front_face ? (1.0f0/mat.ir) : mat.ir # i.e. ηᵢ/ηₜ
	cosθ = min(-r_in.dir⋅rec.n⃗, 1.0f0)
	sinθ = √(1.0f0 - cosθ^2)
	cannot_refract = refraction_ratio * sinθ > 1.0
	if cannot_refract || reflectance(cosθ, refraction_ratio) > rand()
		dir = reflect(r_in.dir, rec.n⃗)
	else
		dir = refract(r_in.dir, rec.n⃗, refraction_ratio)
	end
	Scatter(Ray(rec.p, dir), attenuation) # TODO: rename reflected -> !absorbed?
end

#"From C++: Image 15: Glass sphere that sometimes refracts"
@inline function scene_diel_spheres(left_radius=0.5f0)::HittableList # dielectric spheres
	spheres = Sphere[]

	# small center sphere
	push!(spheres, Sphere(SVector{3,Float32}(0f0,0f0,-1f0), 0.5f0, Lambertian(SVector{3,Float32}(0.1f0,0.2f0,0.5f0))))

	# ground sphere (planet?)
	push!(spheres, Sphere(SVector{3,Float32}(0f0,-100.5f0,-1f0), 100f0, Lambertian(SVector{3,Float32}(0.8f0,0.8f0,0.0f0))))

	# left and right spheres.
	# Use a negative radius on the left sphere to create a "thin bubble"
	push!(spheres, Sphere(SVector{3,Float32}(-1f0,0f0,-1f0), left_radius, Dielectric(1.5f0)))
	push!(spheres, Sphere(SVector{3,Float32}( 1f0,0f0,-1f0), 0.5f0, Metal(SVector{3,Float32}(0.8f0,0.6f0,0.2f0), 0.0f0)))
	HittableList(spheres)
end

render(scene_diel_spheres(), default_camera(), 96, 16)
#render(scene_diel_spheres(), default_camera(), 320, 32)

#md"# Positioning camera"

function scene_blue_red_spheres()::HittableList # dielectric spheres
	spheres = Sphere[]
	R = cos(pi/4)
	push!(spheres, Sphere(@SVector[-R,0f0,-1f0], R, Lambertian(@SVector[0f0,0f0,1f0])))
	push!(spheres, Sphere(@SVector[R,0f0,-1f0], R, Lambertian(@SVector[1f0,0f0,0f0])))
	HittableList(spheres)
end

#md"# Spheres with depth-of-field"

#md"# Random spheres"

function scene_random_spheres()::HittableList # dielectric spheres
	spheres = Sphere[]

	# ground
	push!(spheres, Sphere(@SVector[0f0,-1000f0,-1f0], 1000f0,
						  Lambertian(@SVector[0.5f0,0.5f0,0.5f0])))

	for a in -11:10, b in -11:10
		choose_mat = rand()
		center = @SVector[a + 0.9f0*rand(), 0.2f0, b + 0.90f0*rand()]

		# skip spheres too close?
		if norm(center - @SVector[4f0,0.2f0,0f0]) < 0.9f0 continue end

		if choose_mat < 0.8f0
			# diffuse
			albedo = @SVector[rand() for i ∈ 1:3] .* @SVector[rand() for i ∈ 1:3]
			push!(spheres, Sphere(center, 0.2f0, Lambertian(albedo)))
		elseif choose_mat < 0.95f0
			# metal
			albedo = @SVector[random_between(0.5f0,1.0f0) for i ∈ 1:3]
			fuzz = random_between(0.0f0, 5.0f0)
			push!(spheres, Sphere(center, 0.2f0, Metal(albedo, fuzz)))
		else
			# glass
			push!(spheres, Sphere(center, 0.2f0, Dielectric(1.5f0)))
		end
	end

	push!(spheres, Sphere(@SVector[0f0,1f0,0f0], 1.0f0, Dielectric(1.5f0)))
	push!(spheres, Sphere(@SVector[-4f0,1f0,0f0], 1.0f0,
						  Lambertian(@SVector[0.4f0,0.2f0,0.1f0])))
	push!(spheres, Sphere(@SVector[4f0,1f0,0f0], 1.0f0,
						  Metal(@SVector[0.7f0,0.6f0,0.5f0], 0.0f0)))
	HittableList(spheres)
end

scene_random_spheres()

# Hollow Glass sphere using a negative radius
# claforte: getting a weird black halo in the glass sphere... might be due to my
# "fix" for previous black spots, by moving the RecordHit point a bit away from
# the hit surface...
render(scene_diel_spheres(-0.5f0), default_camera(), 96, 16)

render(scene_blue_red_spheres(), default_camera(), 96, 16)

render(scene_diel_spheres(), default_camera(@SVector[-2f0,2f0,1f0], @SVector[0f0,0f0,-1f0],
							 				@SVector[0f0,1f0,0f0], 20.0f0), 96, 16)

t_lookfrom1 = @SVector[13.0f0,2.0f0,3.0f0]
t_lookat1 = @SVector[0.0f0,0.0f0,0.0f0]
t_cam1 = default_camera(t_lookfrom1, t_lookat1, @SVector[0.0f0,1.0f0,0.0f0], 20.0f0, 16.0f0/9.0f0,
                        0.1f0, 10.0f0)

# took ~20s (really rough timing) in REPL, before optimization
# after optimization:
#   880.997 ms (48801164 allocations: 847.10 MiB)
# after switching to Float32 + reducing allocations using rand_vec3!():
#    79.574 ms (  542467 allocations:  14.93 MiB)
# after optimizing skycolor, rand*, probably more stuff I forgot...
#    38.242 ms (  235752 allocations:   6.37 MiB)
# after removing all remaining Color, Vec3, replacing them with @SVector[]...
#    26.790 ms (   91486 allocations: 2.28 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#    25.856 ms (   70650 allocations: 1.96 MiB)
# Don't specify return value of Option{HitRecord} in hit()
# Don't initialize unnecessary elements in HitRecord():
#    38.961 ms (   70681 allocations: 1.96 MiB) # WORSE, probably because we're writing a lot more to stack...?
# Replace MyFloat by Float32:
#    36.726 ms (13889 allocations: 724.09 KiB)
# @inline lots of stuff:
#    14.690 ms (13652 allocations: 712.98 KiB)
render(scene_random_spheres(), t_cam1, 96, 1)

# took 5020s in Pluto.jl, before optimizations!
# after lots of optimizations, up to switching to Float32 + reducing allocations using rand_vec3!():
#   10.862032 seconds (70.09 M allocations: 1.858 GiB, 1.53% gc time)
# after optimizing skycolor, rand*, probably more stuff I forgot...
#    4.926 s (25,694,163 allocations: 660.06 MiB)
# after removing all remaining Color, Vec3, replacing them with @SVector[]...
#    3.541 s (9074843 allocations: 195.22 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#    3.222 s (2055106 allocations: 88.53 MiB)
# Don't specify return value of Option{HitRecord} in hit()
# Don't initialize unnecessary elements in HitRecord():
#    5.016 s (2056817 allocations: 88.61 MiB) #  WORSE, probably because we're writing a lot more to stack...?
# Replace MyFloat by Float32:
#    5.185 s (1819234 allocations: 83.55 MiB) # Increase of 1% is probably noise
# Remove normalize() in reflect() (that function assumes the inputs are normalized)
#    5.040 s (1832052 allocations: 84.13 MiB)
# @inline lots of stuff:
#    2.110 s (1823044 allocations: 83.72 MiB)
render(scene_random_spheres(), t_cam1, 200, 32)

# After some optimization, took ~5.6 hours:
#   20171.646846 seconds (94.73 G allocations: 2.496 TiB, 1.06% gc time)
# ... however the image looked weird... too blurry
# After removing all remaining Color, Vec3, replacing them with @SVector[]... took ~3.7 hours:
#   13326.770907 seconds (29.82 G allocations: 714.941 GiB, 0.36% gc time)
# Using convert(Float32, ...) instead of MyFloat(...):
# Don't specify return value of Option{HitRecord} in hit()
# Don't initialize unnecessary elements in HitRecord():
# Took ~4.1 hours:
#   14723.339976 seconds (5.45 G allocations: 243.044 GiB, 0.11% gc time) # WORSE, probably because we're writing a lot more to stack...?
# Replace MyFloat by Float32:
# Lots of other optimizations including @inline lots of stuff:
#    6018.101653 seconds (5.39 G allocations: 241.139 GiB, 0.21% gc time)
#render(scene_random_spheres(), t_cam1, 1920, 1000)

t_lookfrom2 = @SVector[3.0f0,3.0f0,2.0f0]
t_lookat2 = @SVector[0.0f0,0.0f0,-1.0f0]
dist_to_focus2 = norm(t_lookfrom2-t_lookat2)
t_cam2 = default_camera(t_lookfrom2, t_lookat2, @SVector[0.0f0,1.0f0,0.0f0], 20.0f0, 16.0f0/9.0f0,
                        2.0f0, dist_to_focus2)

# Before optimization:
#  5.993 s  (193097930 allocations: 11.92 GiB)
# after disabling: `Base.getproperty(vec::SVector{3}, sym::Symbol)`
#  1.001 s  ( 17406437 allocations: 425.87 MiB)
# after forcing Ray and point() to use Float64 instead of AbstractFloat:
#  397.905 ms (6269207 allocations: 201.30 MiB)
# after forcing use of Float32 instead of Float64:
#  487.680 ms (7128113 allocations: 196.89 MiB) # More allocations... something is causing them...
# after optimizing rand_vec3!/rand_vec2! to minimize allocations:
#  423.468 ms (6075725 allocations: 158.92 MiB)
# after optimizing skycolor, rand*, probably more stuff I forgot...
#  217.853 ms (2942272 allocations: 74.82 MiB)
# after removing all remaining Color, Vec3, replacing them with @SVector[]...
#   56.778 ms (1009344 allocations: 20.56 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#   23.870 ms (210890 allocations: 8.37 MiB)
# Replace MyFloat by Float32:
#   22.390 ms (153918 allocations: 7.11 MiB)
# Various other changes, e.g. remove unnecessary normalize
#   20.241 ms (153792 allocations: 7.10 MiB)
# @inline lots of stuff:
#   18.065 ms (153849 allocations: 7.10 MiB)
#render(scene_diel_spheres(), t_cam2, 96, 16)

# using Profile
# Profile.clear_malloc_data()
# render(scene_2_spheres(), default_camera(), 96, 16)