text:
export isphysical, reset_billiard!
export RaySplitter, raysplit_indices
export law_of_refraction

#####################################################################################
# RaySplitter structures
#####################################################################################
"""
    RaySplitter(idxs, transmission, refraction [, newangular]; affect)
Return a `RaySplitter` instance, used to perform raysplitting.
`idxs` is a `Vector{Int}` with the indices of the obstacles
that this `RaySplitter` corresponds to.

`transmission`, `refraction` and `newangular` are **functions**. Let
`φ` be the angle of incidence and `ω` be the angular velocity
and `pflag` the propagation flag (before transmission).
The functions have the following signatures:

1. `transmission(φ, pflag, ω) -> T`, transmission probability.
2. `refraction(φ, pflag, ω) -> θ`, refraction angle.
   This angle is *relative* to the normal vector.
3. `newangular(ω, pflag) -> newω`, new angular velocity after transmission.

The above three functions use the **same convention**: the argument `pflag` is the
one the obstacle has **before transmission**. For example, if a particle is
outside an [`Antidot`](@ref) (with `pflag = true` here) and is transmitted inside
the `Antidot` (`pflag` becomes `false` here), then all three functions will be
given their second argument (the Boolean one) as `true`!

`affect` is a function, and denotes which obstacles
of the billiard are affected when transmission occurs at obstacle
`i` (for which obstacles should the field `pflag` be reversed).
Defaults to `idxs = (i) -> i`, i.e. only the colliding obstacle is affected.
If you want many obstacles to be affected you could write
`idxs = (i) -> SVector(2,3,5)`, etc.
Keep in mind that the only values of `i` that can be passed into this function
are the ones that are given in the argument `idxs`!
"""
struct RaySplitter{T, Φ, Ω, A}
    oidx::Vector{Int}
    transmission::T
    refraction::Φ
    newω::Ω
    affect::A
end

@inline default_affect(i) = i
@inline default_angular(ω, pflag) = ω

function RaySplitter(idxs, tr, ref, newangular = default_angular;
                                    affect = default_affect)
    for i ∈ idxs
        i ∈ affect(i) || throw(ArgumentError(
        "All indices that correspond to this RaySplitter must also be affected!"))
    end
    typeof(idxs) == Int && (idxs = [idxs])
    return RaySplitter(sort(idxs), tr, ref, newangular, affect)
end

#pretty print:
function Base.show(io::IO, ::MIME"text/plain", ray::RaySplitter)
    ps = 15
    angtext = ray.newω == default_angular ? "default" : "$(ray.newω)"
    afftext = ray.affect == default_affect ? "default" : "$(ray.affect)"
    print(io, "RaySplitter for indices $(ray.oidx)"*"\n",
    rpad(" transmission: ", ps)*"$(ray.transmission)\n",
    rpad(" refraction: ", ps)*"$(ray.refraction)\n",
    rpad(" new angular: ", ps)*angtext*"\n",
    rpad(" affect: ", ps)*afftext*"\n"
    )
end

function Base.show(io::IO, ray::RaySplitter)
    print(io, "RaySplitter for indices $(ray.oidx)")
end

"""
    raysplit_indices(bd::Billiard, raysplitters::Tuple)
Create a vector of integers. The `i`th entry tells you which entry of the
`raysplitters` tuple is associated with the `i`th obstacle of the billiard.

If the `i`th entry is `0`, this means that the obstacle does not do raysplitting.
"""
function raysplit_indices(bd::Billiard, raysplitters::Tuple)
    O = zeros(Int, length(bd.obstacles))
    for (k, rayspl) ∈ enumerate(raysplitters)
        O[rayspl.oidx] .= k
    end
    return O
end

allaffected(ray::RaySplitter) = union(ray.affect(i) for i in ray.oidx)

"""
    acceptable_raysplitter(raysplitters, bd::Billiard)
Return `true` if the given `raysplitters`
can be used in conjuction with given billiard `bd`.
"""
function acceptable_raysplitter(raysplitters::Tuple, bd::Billiard)

    for ray in raysplitters
        for i in (ray.oidx ∪ ray.affect(ray.oidx))
            if !supports_raysplitting(bd[i])
                error("Obstacle at index $i of given billiard table"*
                " does not have a field `pflag`"*
                " and therefore does not support ray-splitting."*
                " However a `RaySplitter` uses it!")
            end
        end
    end

    idx = raysplit_indices(bd, raysplitters)

    # Make sure that no indices are shared in the corresponding obstacles
    vecs = [ray.oidx for ray in raysplitters]
    if length(unique(vcat(vecs...))) != sum(length.(vecs))
        error("Different `RaySplitter`s cannot correspond to the same "*
        "obstacles!")
    end

    true
end

#####################################################################################
# Resolve collisions
#####################################################################################
timeprec_rayspl(::Particle{T}) where {T} = timeprec(T)
timeprec_rayspl(::MagneticParticle{T}) where {T} = timeprec_forward(T)

const CLAMPING_ANGLE = 0.1
angleclamp(φ::T) where {T} = clamp(φ, -π/2 + T(CLAMPING_ANGLE), π/2 - T(CLAMPING_ANGLE))

function incidence_angle(p::AbstractParticle{T}, a::Obstacle{T})::T where {T}
    # Raysplit Algorithm step 1: Determine incidence angle (0 < φ < π/4)
    n = normalvec(a, p.pos)
    inverse_dot = clamp(dot(p.vel, -n), -1.0, 1.0)
    φ = acos(inverse_dot)
    # Raysplit Algorithm step 2: get correct sign
    if cross2D(p.vel, n) < 0
        φ *= -1
    end
    return φ
end

function istransmitted(p::AbstractParticle{T}, a::Obstacle{T}, Tr::F) where {T, F}
    ω = typeof(p) <: MagneticParticle ? p.omega : zero(T)
    φ = incidence_angle(p, a)
    # Raysplit Algorithm step 3: check transmission probability
    trans = Tr(φ, a.pflag, ω) > rand()
end

# Raysplit Algorithm step 4: relocate the particle _inside_ the obstacle
# if ray-splitting happens (see `ineq` variable)
function relocate_rayspl!(p::AbstractParticle{T}, o::Obstacle{T},
    trans::Bool = false) where {T}

    ineq = 2trans - 1
    d = distance(p.pos, o)
    notokay = ineq*d > 0.0
    if notokay
        n = normalvec(o, p.pos)
        p.pos -= d * n
    end
    return notokay
end

function resolvecollision!(p::AbstractParticle{T}, bd::Billiard{T}, colidx::Int,
    trans::Bool, rayspl::RaySplitter) where {T<:AbstractFloat}

    a = bd[colidx]
    ismagnetic = typeof(p) <: MagneticParticle
    ω = ismagnetic ? p.omega : T(0)

    # Raysplit Algorithm step 5: recompute angle of incidence
    φ = incidence_angle(p, a)

    if trans #perform raysplitting
        # Raysplit Algorithm step 6: find transmission angle in relative angles
        theta = angleclamp(rayspl.refraction(φ, a.pflag, ω))

        # Raysplit Algorithm step 7: reverse the Obstacle propagation flag
        # for all obstacles dictated by the RaySplitter
        # (this also reverses `normalvec` !)
        for oi ∈ rayspl.affect(colidx)
            bd[oi].pflag = ! bd[oi].pflag
        end

        # Raysplit Algorithm step 8: find transmission angle in real-space angles
        n = normalvec(a, p.pos) #notice that this is reversed! It's the new normalvec!
        Θ = theta + atan(n[2], n[1])

        # Raysplit Algorithm step 9: Perform refraction
        p.vel = SVector{2,T}(cos(Θ), sin(Θ))

        # Raysplit Algorithm step 10: Set new angular velocity
        if ismagnetic
            ω = rayspl.newω(p.omega, !a.pflag)  # notice the exclamation mark
            p.omega = ω
            p.r = abs(1/ω)
        end
    else # No ray-splitting:
        #perform specular
        specular!(p, a)
        end
    return
end

# Ray-splitting version of bounce!
# raysplitters is a tuple of RaySplitter. raysidx is a vector that given the obstacle
# index it tells you which raysplitter to choose from the tuple OR to not
# do raysplitting at all (for returned index 0)
function bounce!(p::AbstractParticle{T}, bd::Billiard{T},
    raysidx::Vector{Int}, raysplitters::Tuple) where {T}

    i::Int, tmin::T, cp::SV{T} = next_collision(p, bd)

    if tmin == Inf
        return i, tmin, p.pos, p.vel

    elseif raysidx[i] != 0

        propagate!(p, cp, tmin)

        trans::Bool = istransmitted(p, bd[i], raysplitters[raysidx[i]].transmission)

        relocate_rayspl!(p, bd[i], trans)
        resolvecollision!(p, bd, i, trans, raysplitters[raysidx[i]])
    else
        relocate!(p, bd[i], tmin, cp)
        resolvecollision!(p, bd[i])
    end
    typeof(p) <: MagneticParticle && (p.center = find_cyclotron(p))
    return i, tmin, p.pos, p.vel
end

#####################################################################################
# Evolve raysplitting
#####################################################################################
@inline raysplit_indices(bd, ::Nothing) = nothing
@inline bounce!(p, bd, ::Nothing, ::Nothing) = bounce!(p, bd)
@inline acceptable_raysplitter(::Nothing, bd::Billiard) = true

evolve!(p, bd, t, ray::RaySplitter; warning = false) =
    evolve!(p, bd, t, (ray, ); warning = warning)
function evolve!(p::Matrix{T}, bd::Billiard{T}, t, raysplitters::Tuple;
    warning = false) where {T}

    if t <= 0
        throw(ArgumentError("`evolve!()` cannot evolve backwards in time."))
    end
    # Check if raysplitters are acceptable
    acceptable_raysplitter(raysplitters, bd)

    ismagnetic = typeof(p) <: MagneticParticle

    raysidx = raysplit_indices(bd, raysplitters)

    rt = T[]; push!(rt, 0)
    rpos = SVector{2,T}[]; push!(rpos, p.pos)
    rvel = SVector{2,T}[]; push!(rvel, p.vel)
    ismagnetic && (omegas = T[]; push!(omegas, p.omega))

    count = zero(t)
    t_to_write = zero(T)

    while count < t

        i, tmin, pos, vel = bounce!(p, bd, raysidx, raysplitters)
        t_to_write += tmin

        if ismagnetic && tmin == Inf
            warning && warn("Pinned particle in evolve! (Inf. col. t)")
            push!(rpos, pos); push!(rvel, vel)
            push!(rt, tmin); push!(omegas, p.omega)
            return (rt, rpos, rvel, omegas)
        end

        if isperiodic(bd) && i ∈ bd.peridx
            # Pinned particle:
            if ismagnetic && t_to_write ≥ 2π/absω
                warning && warn("Pinned particle in evolve! (completed circle)")
                push!(rpos, pos); push!(rvel, vel)
                push!(rt, tmin); push!(omegas, p.omega)
                return (rt, rpos, rvel, omegas)
            end
            #If not pinned, continue (do not write for PeriodicWall)
            continue
        else
            push!(rpos, p.pos + p.current_cell)
            push!(rvel, p.vel); push!(rt, t_to_write);
            ismagnetic && push!(omegas, p.omega)
            count += increment_counter(t, t_to_write)
            t_to_write = zero(T)
        end

    end#time loop
    if ismagnetic
        return (rt, rpos, rvel, omegas)
    else
        return (rt, rpos, rvel)
    end
end

########################
# is physical, etc.
########################

function supports_raysplitting(obst::Obstacle)
  n = fieldnames(typeof(obst))
  in(:pflag, n)
end

"""
    reset_billiard!(bd)
Sets the `pflag` field of all ray-splitting obstacles of a billiard table
to `true`.
"""
function reset_billiard!(bd::Billiard)
    for obst in bd
        supports_raysplitting(obst) && (obst.pflag = true)
    end
end

"""
    isphysical(raysplitter(s))
Return `true` if the given `raysplitters` have physically
plausible properties.

Specifically, check if (φ is the incidence angle, θ the refraction angle):

* Critical angle means total reflection: If θ(φ) ≥ π/2 then Tr(φ) = 0
* Transmission probability is even function: Tr(φ) ≈ Tr(-φ) at ω = 0
* Refraction angle is odd function: θ(φ) ≈ -θ(-φ) at ω = 0
* Ray reversal is true: θ(θ(φ, pflag, ω), !pflag, ω) ≈ φ
* Magnetic conservation is true: (ω_new(ω_new(ω, pflag), !pflag) ≈ ω
"""
isphysical(ray::RaySplitter) = isphysical((ray,))
function isphysical(rays::Tuple)
  for (i, ray) in enumerate(rays)
    scatter = ray.refraction
    tr = ray.transmission
    om = ray.newω
    range = -1.5:0.01:1.5
    orange = -1.0:0.1:1.0
    for pflag in (true, false)
      for ω in orange
        for φ in range
          # Calculate refraction angle:
          θ = scatter(φ, pflag, ω)
          # Calculate transmission probability:
          T = tr(φ, pflag, ω)
          # Check critical angle:
          if θ >= π/2 && T > 0
            es = "Refraction angle ≥ π/2 and T > 0 !\n"
            es*= "For index = $i, tested with φ = $φ, pflag = $pflag, ω = $ω"
            println(es)
            return false
          end
            # Check symmetry:
            if ω==0
              if !isapprox(θ, -scatter(-φ, pflag, ω))
                es = "Scattering angle function is not odd!\n"
                es *="For index = $i, tested with φ = $φ, pflag = $pflag, ω = $ω"
                println(es)
                return false
              end
              if !isapprox(T, tr(-φ, pflag, ω))
                es = "Transmission probability function is not even!\n"
                es *="For index = $i, tested with φ = $φ, pflag = $pflag, ω = $ω"
                println(es)
                return false
              end
            end
            # Check ray-reversal:
            if !isapprox(scatter(θ, !pflag, ω), φ)
              es = "Ray-reversal does not hold!\n"
              es *="For index = $i, tested with φ = $φ, pflag = $pflag, ω = $ω"
              println(es)
              return false
            end
            if !isapprox(om(om(ω, pflag), !pflag), ω)
              es = "Magnetic reversal does not hold!\n"
              es *="For index = $i, tested with φ = $φ, pflag = $pflag, ω = $ω"
              println(es)
              return false
            end
        end#φ range
      end#ω range
    end#pflag range
  end#obstacle range
  return true
end

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

"""
    law_of_refraction(n1, n2 = 1.0) -> t, r

Create transmission and refraction functions `t, r` that follow Snell's
law, i.e. the transmission probability is set to 1.0 except
for the case of total internal reflection.

`n1` is the index of refraction for the `pflag = false` side of an obstacle,
while `n2` is the index of refraction for `pflag = true`.
"""
function law_of_refraction(n1, n2 = 1.0)
    # n1 is "inside" the obstacle, n2 is "outside"

    if n1 < 1.0 || n2 < 1.0
        error("You have just given a physicist a headache.")
    end

    # Snell's law
    refraction(φ, pflag, ω) =
        asin(clamp((pflag ? (n2/n1) : (n1/n2) )* sin(φ), -1.0, 1.0))

    # total internal reflection
    function transmission(φ, pflag, ω)
        nratio = pflag ? (n1/n2) : (n2/n1)

        if abs(nratio) < 1 &&  abs(φ) >= asin(nratio)
            return 0.0
        else
            return 1.0
        end
    end

    return transmission, refraction
end


text:
# abbreviations of types with long names
const AbstractScalarOrVec{T} = Union{T, AbstractVector{T}} where T<:Real
const AV = AbstractVector
const AM = AbstractMatrix
const AVoM = AbstractVecOrMat
const ASoV = AbstractScalarOrVec
const AF = AbstractFloat

ext: jl


text:
struct Lambertian <: Material
    albedo::Texture
end

function scatter(lambertian::Lambertian, ray::Ray, hp::HitPoint)::Tuple{Bool, Union{Nothing, AbstractVec}, Union{Nothing, Ray}}
    attenuation::AbstractVec = value(lambertian.albedo, hp.u, hp.v, hp.point)
    scattered::Ray = Ray(hp.point, hp.normal + random_in_unit_sphere())
    true, attenuation, scattered
end

max_stars_repo_name: goooxu/JuliaRayTracing

max_issues_repo_name: goooxu/JuliaRayTracing

max_forks_repo_name: goooxu/JuliaRayTracing


text:
# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).

"""
    VectorOfArrays{T,N,M} <: AbstractVector{<:AbstractArray{T,N}}

An `VectorOfArrays` represents a vector of `N`-dimensional arrays (that may
differ in size). Internally, `VectorOfArrays` stores all elements of all
arrays in a single flat vector. `M` must equal `N - 1`

The `VectorOfArrays` itself supports `push!`, `unshift!`, etc., but the size
of each individual array in the vector is fixed. `resize!` can be used to
shrink, but not to grow, as the size of the additional element arrays in the
vector would be unknown. However, memory space for up to `n` arrays with a
maximum size `s` can be reserved via
`sizehint!(A::VectorOfArrays, n, s::Dims{N})`.

Constructors:

```julia
VectorOfArrays{T,N}()

VectorOfArrays(A::AbstractVector{<:AbstractArray})
VectorOfArrays{T}(A::AbstractVector{<:AbstractArray})
VectorOfArrays{T,N}(A::AbstractVector{<:AbstractArray})

VectorOfArrays(
    data::AbstractVector,
    elem_ptr::AbstractVector{Int},
    kernel_size::AbstractVector{<:Dims}
    checks::Function = ArraysOfArrays.full_consistency_checks
)
```

Other suitable values for `checks` are `ArraysOfArrays.simple_consistency_checks`
and `ArraysOfArrays.no_consistency_checks`.

`VectorOfVectors` is defined as an type alias:

```julia
`VectorOfVectors{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}`
```
"""
struct VectorOfArrays{
    T, N, M,
    VT<:AbstractVector{T},
    VI<:AbstractVector{Int},
    VD<:AbstractVector{Dims{M}}
} <: AbstractVector{Array{T,N}}
    data::VT
    elem_ptr::VI
    kernel_size::VD

    function VectorOfArrays{T,N}() where {T,N}
        M = length(Base.front(ntuple(_ -> 0, Val{N}())))
        data = Vector{T}()
        elem_ptr = [firstindex(data)]
        kernel_size = Vector{Dims{M}}()

        new{
            T, N, M,
            typeof(data),
            typeof(elem_ptr),
            typeof(kernel_size)
        }(data, elem_ptr, kernel_size)
    end

    function VectorOfArrays(
        data::VT,
        elem_ptr::VI,
        kernel_size::VD,
        checks::Function = full_consistency_checks
    ) where {
        T, M,
        VT<:AbstractVector{T},
        VI<:AbstractVector{Int},
        VD<:AbstractVector{Dims{M}}
    }
        N = length((ntuple(_ -> 0, Val{M}())..., 0))
        A = new{T,N,M,VT,VI,VD}(data, elem_ptr, kernel_size)
        checks(A)
        A
    end
end

export VectorOfArrays

function VectorOfArrays{T,N}(A::AbstractVector{<:AbstractArray{U,N}}) where {T,N,U}
    R = VectorOfArrays{T,N}()
    append!(R, A)
end

VectorOfArrays{T}(A::AbstractVector{<:AbstractArray{U,N}}) where {T,U,N} = VectorOfArrays{T,N}(A)

VectorOfArrays(A::AbstractVector{<:AbstractArray{T,N}}) where {T,N} = VectorOfArrays{T,N}(A)

Base.convert(VA::Type{VectorOfArrays{T,N}}, A::AbstractVector{AbstractArray{U,N}}) where {T,N,U} = VA(A)
Base.convert(VA::Type{VectorOfArrays}, A::AbstractVector{AbstractArray{T,N}}) where {T,N} = VA(A)

"""
    internal_element_ptr(A::VectorOfArrays)

Returns the internal element pointer vector of `A`.

Do *not* change modify the returned vector in any way, as this would break the
inner consistency of `A`.

Use with care, see [`element_ptr`](@ref) for a safe version of this function.
"""
internal_element_ptr(A::VectorOfArrays) = A.elem_ptr

"""
    element_ptr(A::VectorOfArrays)

Returns a copy of the internal element pointer vector of `A`.
"""
element_ptr(A::VectorOfArrays) = deepcopy(internal_element_ptr(A))

function full_consistency_checks(A::VectorOfArrays)
    simple_consistency_checks(A)
    all(eachindex(A.kernel_size)) do i
        len = A.elem_ptr[i+1] - A.elem_ptr[i]
        klen = prod(A.kernel_size[i])
        len >= 0 && (klen == 1 || mod(len, klen) == 0)
    end || throw(ArgumentError("VectorOfArrays inconsistent: Content of elem_ptr and kernel_size is inconsistent"))
    nothing
end

function simple_consistency_checks(A::VectorOfArrays{T,N,M}) where {T,N,M}
    M == N - 1 || throw(ArgumentError("VectorOfArrays{T,N,M} inconsistent: M must equal N - 1"))
    firstindex(A.elem_ptr) == firstindex(A.kernel_size) || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible indexing"))
    size(A.elem_ptr, 1) == size(A.kernel_size, 1) + 1 || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible size"))
    first(A.elem_ptr) >= firstindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: First elem_ptr inconsistent with data indices"))
    last(A.elem_ptr) - 1 <= lastindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: Last elem_ptr inconsistent with data indices"))
    nothing
end

function no_consistency_checks(A::VectorOfArrays)
    nothing
end

Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer)
    elem_ptr = A.elem_ptr

    from = elem_ptr[i]
    until = elem_ptr[i+1]
    to = until - 1
    len = until - from

    ksize = A.kernel_size[i]
    klen = prod(ksize)
    len_p, klen_p = promote(len, klen)
    sz_lastdim = len == 0 ? len_p : div(len_p, klen_p)
    sz = (ksize..., sz_lastdim)

    (from:to, sz)
end

import Base.==
(==)(A::VectorOfArrays, B::VectorOfArrays) =
    A.data == B.data && A.elem_ptr == B.elem_ptr && A.kernel_size == B.kernel_size

"""
    flatview(A::VectorOfArrays{T})::Vector{T}

Returns the internal serialized representation of all element arrays of `A` as
a single vector. Do *not* change the length of the returned vector, as it
would break the inner consistency of `A`.
"""
flatview(A::VectorOfArrays) = A.data

Base.size(A::VectorOfArrays) = size(A.kernel_size)

Base.IndexStyle(A::VectorOfArrays) = IndexLinear()

Base.@propagate_inbounds _reshape_dataview(dataview::AbstractArray, s::NTuple{1,Integer}) = dataview

Base.@propagate_inbounds _reshape_dataview(dataview::AbstractArray, s::NTuple{N,Integer}) where {N} =
    Base.__reshape((dataview, IndexStyle(dataview)), s)

Base.@propagate_inbounds function Base.getindex(A::VectorOfArrays, i::Integer)
    @boundscheck checkbounds(A, i)
    r, s = _elem_range_size(A, i)
    dataview = view(A.data, r)
    _reshape_dataview(dataview, s)
end

Base.@propagate_inbounds function Base._getindex(l::IndexStyle, A::VectorOfArrays, idxs::AbstractUnitRange{<:Integer})
    from = first(idxs)
    to = last(idxs)
    elem_ptr = A.elem_ptr[from:(to+1)]
    kernel_size = A.kernel_size[from:to]
    data = A.data[first(elem_ptr):(last(elem_ptr) - 1)]
    broadcast!(+, elem_ptr, elem_ptr, firstindex(data) - first(elem_ptr))
    VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks)
end

Base.@propagate_inbounds function Base._getindex(l::IndexStyle, A::VectorOfArrays, idxs::AbstractVector{<:Integer})
    @boundscheck checkbounds(A, idxs)

    A_ep = A.elem_ptr
    A_data = A.data

    elem_ptr = similar(A_ep, length(eachindex(idxs)) + 1)
    delta_i = firstindex(elem_ptr) - firstindex(idxs)

    elem_ptr[firstindex(elem_ptr)] = firstindex(A_data)
    for i in eachindex(idxs)
        idx = idxs[i]
        l = A_ep[idx + 1] - A_ep[idx]
        elem_ptr[i + 1 + delta_i] = elem_ptr[i + delta_i] + l
    end

    data = similar(A_data, last(elem_ptr) - first(elem_ptr))
    if firstindex(data) != firstindex(A_data)
        @assert firstindex(data) != first(elem_ptr)
        broadcast!(+, elem_ptr, elem_ptr, firstindex(data) - first(elem_ptr))
    end

    for i in eachindex(idxs)
        idx = idxs[i]
        l = A_ep[idx + 1] - A_ep[idx]

        # Sanity check:
        @assert l == elem_ptr[i + 1 + delta_i] - elem_ptr[i + delta_i]

        copyto!(data, elem_ptr[i + delta_i], A_data, A_ep[idx], l)
    end

    kernel_size = A.kernel_size[idxs]

    VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks)
end

Base.@propagate_inbounds function Base.setindex!(A::VectorOfArrays{T,N}, x::AbstractArray{U,N}, i::Integer) where {T,N,U}
    r, s = _view_reshape_spec(A, i)
    @boundscheck s == size(x) || throw(DimensionMismatch("Can't assign array to element $i of VectorOfArrays, array size is incompatible"))
    A.data[rng] = x
    A
end

Base.length(A::VectorOfArrays) = length(A.kernel_size)

@inline function Base.resize!(A::VectorOfArrays{T,N,M}, n::Integer) where {T,M,N}
    old_n = length(A)
    if n > old_n
        throw(ArgumentError("Cannot resize VectorOfArrays from length $old_n to $n, can only shrink, not grow"))
    elseif n < old_n
        resize!(A.data, A.elem_ptr[n+1] - 1)
        resize!(A.elem_ptr, n + 1)
        resize!(A.kernel_size, n)
    end
    A
end

function append_elemptr!(A::AbstractVector{<:Integer}, B::AbstractVector{<:Integer})
    idxs_A = LinearIndices(A)
    idxs_B = LinearIndices(B)
    length_B = length(idxs_B)

    A_from = last(idxs_A) + 1
    A_to = A_from - 1 + length_B - 1
    resize!(A, length(first(idxs_A):A_to))

    B_from = first(idxs_B) + 1
    B_to = last(idxs_B)

    A_idxs_offs = A_from - B_from

    checkindex(Bool, eachindex(B), B_from:B_to)
    checkindex(Bool, eachindex(A), (B_from:B_to) .+ A_idxs_offs)
    @inbounds begin
        value_offset = A[A_from - 1] - B[B_from - 1]
        @simd for i in B_from:B_to
            A[i + A_idxs_offs] = B[i] + value_offset
        end
    end

    A
end

function Base.append!(A::VectorOfArrays{T,N}, B::VectorOfArrays{U,N}) where {T,N,U}
    if !isempty(B)
        append!(A.data, B.data)
        append_elemptr!(A.elem_ptr, B.elem_ptr)
        append!(A.kernel_size, B.kernel_size)
    end
    A
end

function Base.append!(A::VectorOfArrays{T,N}, B::AbstractVector{<:AbstractArray{U,N}}) where {T,N,U}
    if !isempty(B)
        n_A = length(eachindex(A))
        n_B = length(eachindex(B))
        datalen_A = length(eachindex(A.data))
        datalen_B = zero(Int)
        for i in eachindex(B)
            datalen_B += Int(length(eachindex(B[i])))
        end

        sizehint!(A.data, datalen_A + datalen_B)
        sizehint!(A.elem_ptr, n_A + n_B + 1)
        sizehint!(A.kernel_size, n_A + n_B)

        for i in eachindex(B)
            push!(A, B[i])
        end
    end
    A
end

Base.vcat(V::VectorOfArrays) = V

function Base.vcat(Vs::(VectorOfArrays{U,N} where U)...) where {N}
    data = vcat(map(x -> x.data, Vs)...)

    elem_ptr = similar(Vs[1].elem_ptr, 1)
    elem_ptr[1] = firstindex(data)
    @inbounds for i in eachindex(Vs)
        append_elemptr!(elem_ptr, Vs[i].elem_ptr)
    end

    kernel_size = vcat(map(x -> x.kernel_size, Vs)...)

    VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks)
end

function Base.copy(V::VectorOfArrays)
    VectorOfArrays(copy(V.data), copy(V.elem_ptr), copy(V.kernel_size), no_consistency_checks)
end

Base.@propagate_inbounds function Base.unsafe_view(A::VectorOfArrays, idxs::AbstractUnitRange{<:Integer})
    from = first(idxs)
    to = last(idxs)
    VectorOfArrays(
        A.data,
        view(A.elem_ptr, from:(to+1)),
        view(A.kernel_size, from:to),
        no_consistency_checks
    )
end

function Base.sizehint!(A::VectorOfArrays{T,N}, n, s::Dims{N}) where {T,N}
    sizehint!(A.data, n * mul(s))
    sizehint!(A.elem_ptr, n + 1)
    sizehint!(A.kernel_size, n)
    A
end

function Base.push!(A::VectorOfArrays{T,N}, x::AbstractArray{U,N}) where {T,N,U}
    @assert last(A.elem_ptr) == lastindex(A.data) + 1
    append!(A.data, x)
    push!(A.elem_ptr, lastindex(A.data) + 1)
    push!(A.kernel_size, Base.front(size(x)))
    A
end

function UnsafeArrays.uview(A::VectorOfArrays)
    VectorOfArrays(
        uview(A.data),
        uview(A.elem_ptr),
        uview(A.kernel_size),
        no_consistency_checks
    )
end

function innermap(f::Base.Callable, A::VectorOfArrays)
    new_data = map(f, A.data)
    VectorOfArrays(new_data, A.elem_ptr, A.kernel_size, simple_consistency_checks)
end

function deepmap(f::Base.Callable, A::VectorOfArrays)
    new_data = deepmap(f, A.data)
    VectorOfArrays(new_data, A.elem_ptr, A.kernel_size, simple_consistency_checks)
end

"""
    VectorOfVectors{T,...} = VectorOfArrays{T,1,...}

Constructors:

```julia
VectorOfVectors(A::AbstractVector{<:AbstractVector})
VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T}

VectorOfVectors(
    data::AbstractVector, elem_ptr::AbstractVector{Int},
    checks::Function = full_consistency_checks
)

See also [VectorOfArrays](@ref).
```
"""
const VectorOfVectors{
    T,
    VT<:AbstractVector{T},
    VI<:AbstractVector{Int},
    VD<:AbstractVector{Dims{0}}
} = VectorOfArrays{T,1,0,VT,VI,VD}

export VectorOfVectors

VectorOfVectors{T}() where {T} = VectorOfArrays{T,1}()

VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T} = VectorOfArrays{T,1}(A)
VectorOfVectors(A::AbstractVector{<:AbstractVector}) = VectorOfArrays(A)

VectorOfVectors(
    data::AbstractVector,
    elem_ptr::AbstractVector{Int},
    checks::Function = full_consistency_checks
) = VectorOfArrays(
    data,
    elem_ptr,
    similar(elem_ptr, Dims{0}, size(elem_ptr, 1) - 1),
    checks
)

"""
    consgrouped_ptrs(A::AbstractVector)

Compute an element pointer vector, suitable for creation of a
`VectorOfVectors` that implies grouping equal consecutive entries of
`A`.

Example:

```julia
    A = [1, 1, 2, 3, 3, 2, 2, 2]
    elem_ptr = consgrouped_ptrs(A)
    first.(VectorOfVectors(A, elem_ptr)) == [1, 2, 3, 2]
```
consgrouped_ptrs
Typically, `elem_ptr` will be used to apply the computed grouping to other
data:

```julia
    B = [1, 2, 3, 4, 5, 6, 7, 8]
    VectorOfVectors(B, elem_ptr) == [[1, 2], [3], [4, 5], [6, 7, 8]]
```
"""
function consgrouped_ptrs end
export consgrouped_ptrs

function consgrouped_ptrs(A::AbstractVector)
    elem_ptr = Vector{Int}()
    idxs = eachindex(A)
    push!(elem_ptr, first(idxs))
    if !isempty(A)
        prev_0 = A[first(idxs)]
        prev::typeof(prev_0) = prev_0
        @inbounds for i in (first(idxs) + 1):last(idxs)
            curr = A[i]
            if (curr != prev)
                push!(elem_ptr, i)
                prev = curr
            end
        end
        push!(elem_ptr, last(idxs) + 1)
    end
    elem_ptr
end

"""
    consgroupedview(source::AbstractVector, target)

Compute a grouping of equal consecutive elements on `source` via
[`consgrouped_ptrs`](@ref) and apply the grouping to target, resp. each
element of `target`. `target` may be an vector or a named or unnamed tuple of
vectors. The result is a `VectorOfVectors`, resp. a tuple of such.

Example:

    A = [1, 1, 2, 3, 3, 2, 2, 2]
    B = [1, 2, 3, 4, 5, 6, 7, 8]
    consgroupedview(A, B) == [[1, 2], [3], [4, 5], [6, 7, 8]]

`consgroupedview` plays well with columnar tables, too:

```julia
    using Tables, TypedTables
    data = Table(
        a = [1, 1, 2, 3, 3, 2, 2, 2],
        b = [1, 2, 3, 4, 5, 6, 7, 8],
        c = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]
    )

    result = Table(consgroupedview(data.a, Tables.columns(data)))
```

will return

```
     a          b          c
   ┌──────────────────────────────────────
 1 │ [1, 1]     [1, 2]     [1.1, 2.2]
 2 │ [2]        [3]        [3.3]
 3 │ [3, 3]     [4, 5]     [4.4, 5.5]
 4 │ [2, 2, 2]  [6, 7, 8]  [6.6, 7.7, 8.8]
```

without copying any data:

```
    flatview(result.a) === data.a
    flatview(result.b) === data.b
    flatview(result.c) === data.c
```
"""
function consgroupedview end
export consgroupedview

function consgroupedview(source::AbstractVector, target::AbstractVector)
    elem_ptr = consgrouped_ptrs(source)
    VectorOfVectors(target, elem_ptr)
end

function consgroupedview(source::AbstractVector, target::NTuple{N,AbstractVector}) where {N}
    elem_ptr = consgrouped_ptrs(source)
    map(X -> VectorOfVectors(X, elem_ptr), target)
end

function consgroupedview(source::AbstractVector, target::NamedTuple{syms,<:NTuple{N,AbstractVector}}) where {syms,N}
    elem_ptr = consgrouped_ptrs(source)
    map(X -> VectorOfVectors(X, elem_ptr), target)
end


text:
using LinearAlgebra

# functions
function pprint(mat::Matrix)
    show(stdout, "text/plain", mat)
    println("\n")
end

function pprint(vec::Vector)
    println(vec)
    println("\n")
end

# http://cgkit.sourceforge.net/doc2/mat3.html
function mat3(x, y, z)
    hcat(x, y, z)
end

"""
Calculate the rotation of reference frame (ax, ay, az) so that it matches the reference frame
(bx, by, bz).
All parameters must be 3-element vectors. Both refrence frames must be orthogonal,
all vectors must already be normalized.
Source: http://en.wikipedia.org/wiki/User:Snietfeld/TRIAD_Algorithm
"""
function rot3d(ax, ay, az, bx, by, bz)
    R_ai = mat3(ax, az, ay)
    R_bi = mat3(bx, bz, by)
    return R_bi * R_ai'
end

# main
ax, ay, az = [1.,0,0], [ 0,1.,0], [0,0,1.]
bx, by, bz = [0,1.,0], [-1.,0,0], [0,0,1.]
vec = [1.,0,0]

pprint(vec)

pprint(rot3d(ax, ay, az, bx, by, bz))
pprint(rot3d(ax, ay, az, bx, by, bz) * vec)
pprint(mat3(bx, by, bz))


text:
include("common.jl")

if DFTK.mpi_nprocs() > 1
    disable_threading()
else
    setup_threading()
end

function plots_noniteractive()
    # Setup environment for making automated plots
    ENV["GKS_ENCODING"] = "utf8"
    ENV["GKSwstype"]    = "100"
    ENV["PLOTS_TEST"]   = "true"
end

config = Dict{String, Any}()
for d in readdir()
    if isdir(d) && isfile(joinpath(d, "config.jl"))
        include(joinpath(d, "config.jl"))
    end
end

if !isinteractive()
    plots_noniteractive()
    run_step_on_cases(dump_guess, config)
    run_step_on_cases(scf, config)
end

max_stars_repo_name: mfherbst/supporting-adaptive-damping

max_issues_repo_name: mfherbst/supporting-adaptive-damping

max_forks_repo_name: mfherbst/supporting-adaptive-damping


text:
function process_resource(request::HTTP.Request, state::HandlerState; namespace::AbstractString, path::AbstractString)
    (relative_path, is_fp) = parse_fingerprint_path(path)
    registered_files = state.cache.resources.files
    if !haskey(registered_files, namespace)
        #TODO Exception like in python
        return HTTP.Response(404)
    end
    namespace_files = registered_files[namespace]
    if !in(relative_path, namespace_files.files)
        #TODO Exception like in python
        return HTTP.Response(404)
    end

    filepath = joinpath(namespace_files.base_path, relative_path)

    try
        headers = Pair{String,String}[]
        file_contents = read(joinpath(namespace_files.base_path, relative_path))
        mimetype = mime_by_path(relative_path)
        !isnothing(mimetype) && push!(headers, "Content-Type" => mimetype)
        if is_fp
            push!(headers,
                "Cache-Control" => "public, max-age=31536000" # 1 year
            )
        else
            etag = bytes2hex(MD5.md5(file_contents))
            push!(headers, "ETag" => etag)
            request_etag = HTTP.header(request, "If-None-Match", "")
            request_etag == etag && return HTTP.Response(304)
        end
        return HTTP.Response(200, headers; body = file_contents)

    catch e
        !(e isa SystemError) && rethrow(e)
        #TODO print to log
        return HTTP.Response(404)
    end
end


text:
mutable struct SGD <: Optimizer
    xparams::Vector{XVariable}
    lr::AbstractFloat
    name::String
    function SGD(xparams::Vector{XVariable}; lr=1e-4)
        new(xparams, lr, "SGD")
    end
end

# pretty printing
function Base.show(io::IO, O::SGD)
    print("SGD(lr=$(O.lr))")
end

function update!(O::SGD; clipfn::Function=LPInfNormClip, clipvalue=10.0)
    μ = - O.lr
    for i = 1:length(O.xparams)
        c , θ = O.xparams[i]
        ∇ = clipfn(setNanInfZero(θ.delta), clipvalue)
        @. θ.value += μ * ∇
    end
end

max_stars_repo_name: sonosole/Delta

max_issues_repo_name: sonosole/Delta

max_forks_repo_name: sonosole/Delta


text:
@with_kw struct IDM{R<:Real} <: FieldVector{6, R}
   	a::R       # max comfy acceleration
	b::R       # max comfy brake speed
	T::R       # desired safety time headway
	v0::R      # desired speed
	s0::R      # minimum headway (e.g. if x is less than this then you crashed)
	delta::R     # accel exponent
end

const NORMAL_IDM =     IDM(1.4, 2.0, 1.5, 33.3, 2.0, 4.0)
const AGGRESSIVE_IDM = IDM(2.0, 3.0, 1.0, 38.9, 0.0, 4.0)
const TIMID_IDM =      IDM(0.8, 1.0, 2.0, 27.8, 4.0, 4.0)

function action(idm::IDM, c, idx, ::Type{LongAccel})
    fidx = fore_index(c, idx)
    xdot = vel(state(c, idx))
    if fidx == 0
        xddot = idm.a*(1.0 - (xdot/idm.v0))
    else
        delta_x = headway(c, idx) # positive distance to car in front
        delta_xdot = xdot - vel(state(c, fidx))
        ss = sstar(idm, xdot, delta_xdot)
        xddot = idm.a*(1.0 - (xdot/idm.v0)^idm.delta - (ss/delta_x)^2)
    end
    return LongAccel(xddot)
end

sstar(idm::IDM, xdot, delta_xdot) = idm.s0 + idm.T*xdot + xdot*delta_xdot/(2*sqrt(idm.a*idm.b))
max_acc(idm::IDM) = idm.a

struct IDMGaussianCopula
    gc::GaussianCopula
    timid::IDM
    aggressive::IDM
end

IDMGaussianCopula(rho::Float64) = IDMGaussianCopula(GaussianCopula(length(IDM), rho), TIMID_IDM, AGGRESSIVE_IDM)

function rand(rng::AbstractRNG, d::IDMGaussianCopula)
    return d.timid + (d.aggressive-d.timid).*rand(rng(d.gc))
end


text:
using HDF5 # To have access to .hf5 files
using Plots # To be able to plot data

########################################
# Function that read the .hf5 file
########################################

namefile = "../../data/Dump_Diffusion_Coefficients_Cut.hf5"

"""
    openCutData(namefile)

Recovers the data from the .hf5 file `namefile` and returns it as a 4-uple.
In the order nbj, aMeasure, tabj, tabDRRjj.
"""
function openCutData(namefile)
    file = h5open(namefile, "r")
    nbj = read(file,"nbj")
    aMeasure = read(file,"aMeasure")
    tabj = read(file,"tabj")
    tabDRRjj = read(file,"tabDRRjj")
    jlc = read(file,"jlc")
    close(file)
    return jlc, nbj, aMeasure, tabj, tabDRRjj
end

########################################
# Getting the data
########################################

println("Recovering plot data...")
jlc, nbj, aMeasure, tabj, tabDRRjj = openCutData(namefile)

########################################
# Removing diffusion coefficients within the loss-cone
########################################

ij = 1

while (tabj[ij] < jlc)
    global ij += 1
end

tabj = tabj[ij:nbj]
tabDRRjj = tabDRRjj[ij:nbj]

########################################
# Converting the data into log-scaling
########################################

println("Converting data into log-log scaling...")

for ij=1:length(tabj)
    tabj[ij] = log10(tabj[ij])
end

for ij=1:length(tabDRRjj)
    tabDRRjj[ij] = log10(tabDRRjj[ij])
end

########################################
# Plotting the data
########################################

println("Plotting the data...")

p = plot(tabj, tabDRRjj, legend=false)#, scale=:log10)
savefig(p,"../../graphs/Julia/DjjCut.png") # Saves the figure
display(p) # Display plot

readline() # Plot window stays open until we press "Enter"


text:
function *(loc::LocationalCode, bits::String)
    xloc = loc.x * bits[1]
    yloc = loc.y * bits[2]
    zloc = loc.z * bits[3]
    return LocationalCode(xloc, yloc, zloc)
end

*(loc::LocationalCode, bits::Vector{Int}) = *(loc, join(string.(bits)))
*(loc::LocationalCode, bits::Tuple{Int,Int,Int}) = *(loc, join(string.(bits)))

function LocationalCode(octree::Octree, p)
    binsize = 2^octree.maxdepth
    octree_size = octree.root.boundary.widths[1]

    # Convert to code
    xloc = floor(Int, (p[1] - octree.root.boundary.origin[1]) / octree_size * binsize)
    yloc = floor(Int, (p[2] - octree.root.boundary.origin[2]) / octree_size * binsize)
    zloc = floor(Int, (p[3] - octree.root.boundary.origin[3]) / octree_size * binsize)

    # Convert to strings
    xloc = bitstring(xloc)[end-octree.maxdepth+1:end]
    yloc = bitstring(yloc)[end-octree.maxdepth+1:end]
    zloc = bitstring(zloc)[end-octree.maxdepth+1:end]
    LocationalCode(xloc, yloc, zloc)
end

function Base.getindex(code::LocationalCode, inds...)
    return LocationalCode(code.x[inds...], code.y[inds...], code.z[inds...])
end

Base.lastindex(code::LocationalCode) = Base.lastindex(code.x)

Base.iterate(code::LocationalCode) =
    ((parse(Int, code.x[1]), parse(Int, code.y[1]), parse(Int, code.z[1])), 2)

function Base.iterate(code::LocationalCode, state)
    if state <= length(code.x)
        return (
            (
             parse(Int, code.x[state]),
             parse(Int, code.y[state]),
             parse(Int, code.z[state]),
            ),
            state + 1,
        )
    else
        return nothing
    end
end

"""Check if code is blank (corresponding to root cell)."""
isroot(code::LocationalCode) = (code == LocationalCode("", "", ""))

function code_to_cell(octree::Octree, code::LocationalCode)
    cell = octree.root

    # If cell is root, then return
    isroot(code) && return cell

    for (i, (xloc, yloc, zloc)) in enumerate(code)
        i > octree.maxdepth && break
        cell = cell[xloc+1, yloc+1, zloc+1]
        isnothing(cell.children) && break
    end
    return cell
end

toint(code::LocationalCode) = parse(Int, code.x, base = 2)


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""NonResourceAttributes includes the authorization attributes available for non-resource requests to the Authorizer interface

    IoK8sApiAuthorizationV1NonResourceAttributes(;
        path=nothing,
        verb=nothing,
    )

    - path::String : Path is the URL path of the request
    - verb::String : Verb is the standard HTTP verb
"""
mutable struct IoK8sApiAuthorizationV1NonResourceAttributes <: SwaggerModel
    path::Any # spec type: Union{ Nothing, String } # spec name: path
    verb::Any # spec type: Union{ Nothing, String } # spec name: verb

    function IoK8sApiAuthorizationV1NonResourceAttributes(;path=nothing, verb=nothing)
        o = new()
        validate_property(IoK8sApiAuthorizationV1NonResourceAttributes, Symbol("path"), path)
        setfield!(o, Symbol("path"), path)
        validate_property(IoK8sApiAuthorizationV1NonResourceAttributes, Symbol("verb"), verb)
        setfield!(o, Symbol("verb"), verb)
        o
    end
end # type IoK8sApiAuthorizationV1NonResourceAttributes

const _property_map_IoK8sApiAuthorizationV1NonResourceAttributes = Dict{Symbol,Symbol}(Symbol("path")=>Symbol("path"), Symbol("verb")=>Symbol("verb"))
const _property_types_IoK8sApiAuthorizationV1NonResourceAttributes = Dict{Symbol,String}(Symbol("path")=>"String", Symbol("verb")=>"String")
Base.propertynames(::Type{ IoK8sApiAuthorizationV1NonResourceAttributes }) = collect(keys(_property_map_IoK8sApiAuthorizationV1NonResourceAttributes))
Swagger.property_type(::Type{ IoK8sApiAuthorizationV1NonResourceAttributes }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiAuthorizationV1NonResourceAttributes[name]))}
Swagger.field_name(::Type{ IoK8sApiAuthorizationV1NonResourceAttributes }, property_name::Symbol) =  _property_map_IoK8sApiAuthorizationV1NonResourceAttributes[property_name]

function check_required(o::IoK8sApiAuthorizationV1NonResourceAttributes)
    true
end

function validate_property(::Type{ IoK8sApiAuthorizationV1NonResourceAttributes }, name::Symbol, val)
end


text:
# See https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array-1
using BinaryTraits: @trait, @implement, @check
using BinaryTraits.Prefix: Has, No

import Base: size, getindex, setindex!

# -----------------------------------------------------------------------------
# AbstractArray interface
# -----------------------------------------------------------------------------

@trait Dimension
@implement Has{Dimension} by size(_)::Tuple

# Bottom is used due to duck typing of the value `v` below
const Bottom = Base.Bottom

@trait LinearIndexing
@implement Has{LinearIndexing} by getindex(_, i::Int)
@implement Has{LinearIndexing} by setindex!(_, v::Bottom, i::Int)

const IntVarArg = Vararg{Int, N} where N

@trait CartesianIndexing
@implement Has{CartesianIndexing} by getindex(_, I::IntVarArg)
@implement Has{CartesianIndexing} by setindex!(_, v::Bottom, I::IntVarArg)

# -----------------------------------------------------------------------------
# Example: 1-D Int array
# -----------------------------------------------------------------------------
const Array1DInt = Array{Int,1}

@assign Array1DInt with Has{Dimension}
@check(Array1DInt)
#=
julia> @check(Array1DInt)
✅ Array{Int64,1} has implemented:
1. BinaryTrait{Dimension}: Positive{Dimension} ⇢ size(🔹)::Tuple
=#

@assign Array1DInt with Has{LinearIndexing}
@check(Array1DInt)
#=
julia> @check(Array1DInt)
✅ Array{Int64,1} has implemented:
1. BinaryTrait{Dimension}: Positive{Dimension} ⇢ size(🔹)::Tuple
2. BinaryTrait{LinearIndexing}: Positive{LinearIndexing} ⇢ getindex(🔹, ::Int64)::Any
3. BinaryTrait{LinearIndexing}: Positive{LinearIndexing} ⇢ setindex!(🔹, ::Union{}, ::Int64)::Any
=#

# 1D array is a specialized version of CartesianIndexing.
# Let's verify.
@assign Array1DInt with Has{CartesianIndexing}
@check(Array1DInt)
#=
julia> @check(Array1DInt)
✅ Array{Int64,1} has implemented:
1. BinaryTrait{Dimension}: Positive{Dimension} ⇢ size(🔹)::Tuple
2. BinaryTrait{LinearIndexing}: Positive{LinearIndexing} ⇢ getindex(🔹, ::Int64)::Any
3. BinaryTrait{CartesianIndexing}: Positive{CartesianIndexing} ⇢ getindex(🔹, ::Vararg{Int64,N} where N)::Any
4. BinaryTrait{CartesianIndexing}: Positive{CartesianIndexing} ⇢ setindex!(🔹, ::Union{}, ::Vararg{Int64,N} where N)::Any
5. BinaryTrait{LinearIndexing}: Positive{LinearIndexing} ⇢ setindex!(🔹, ::Union{}, ::Int64)::Any
=#

# -----------------------------------------------------------------------------
# Example: SquaresVector
# -----------------------------------------------------------------------------

struct SquaresVector <: AbstractArray{Int, 1}
    count::Int
end

Base.size(S::SquaresVector) = (S.count,)
Base.getindex(S::SquaresVector, i::Int) = i*i

@assign SquaresVector with HasDimension,HasLinearIndexing
@check(SquaresVector)
#=
julia> @check(SquaresVector)
✅ SquaresVector has implemented:
1. DimensionTrait: HasDimension ⇢ size(🔹)::Tuple
2. LinearIndexingTrait: HasLinearIndexing ⇢ getindex(🔹, ::Int64)::Any
3. LinearIndexingTrait: HasLinearIndexing ⇢ setindex!(🔹, ::Union{}, ::Int64)::Any
=#


text:
import JuliaInterpreter.eval_code

# Simple evaling of function argument
function evalfoo1(x,y)
    x+y
end
frame = JuliaInterpreter.enter_call(evalfoo1, 1, 2)
@test eval_code(frame, "x") == 1
@test eval_code(frame, "y") == 2

# Evaling with sparams
evalsparams(x::T) where T = x
frame = JuliaInterpreter.enter_call(evalsparams, 1)
@test eval_code(frame, "x") == 1
eval_code(frame, "x = 3")
@test eval_code(frame, "x") == 3
@test eval_code(frame, "T") == Int
eval_code(frame, "T = Float32")
@test eval_code(frame, "T") == Float32

# Evaling with keywords
evalkw(x; bar=true) = x
frame = JuliaInterpreter.enter_call(evalkw, 2)
frame = JuliaInterpreter.maybe_step_through_wrapper!(frame)
@test eval_code(frame, "x") == 2
@test eval_code(frame, "bar") == true
eval_code(frame, "bar = false")
@test eval_code(frame, "bar") == false

# Evaling with symbols
evalsym() = (x = :foo)
frame = JuliaInterpreter.enter_call(evalsym)
# Step until the local actually end up getting defined
JuliaInterpreter.step_expr!(frame)
JuliaInterpreter.step_expr!(frame)
@test eval_code(frame, "x") == :foo

# Evaling multiple statements (https://github.com/JuliaDebug/Debugger.jl/issues/188)
frame = JuliaInterpreter.enter_call(evalfoo1, 1, 2)
@test eval_code(frame, "x = 1; y = 2") == 2
@test eval_code(frame, "x") == 1
@test eval_code(frame, "y") == 2

# https://github.com/JuliaDebug/Debugger.jl/issues/177
function f()
    x = 1
    f = ()->(x = 2)
    f()
    x
end
frame = JuliaInterpreter.enter_call(f)
JuliaInterpreter.step_expr!(frame)
JuliaInterpreter.step_expr!(frame)
@test eval_code(frame, "x") == 1
eval_code(frame, "x = 3")
@test eval_code(frame, "x") == 3
JuliaInterpreter.finish!(frame)
@test JuliaInterpreter.get_return(frame) == 2

function debugfun(non_accessible_variable)
    garbage = ones(10)
    map(1:10) do i
        1+1
        a = 5
        @bp
        garbage[i] + non_accessible_variable[i]
        non_accessible_variable = 2
    end
end
fr = JuliaInterpreter.enter_call(debugfun, [1,2])
fr, bp = debug_command(fr, :c)
@test eval_code(fr, "non_accessible_variable") == [1,2]
@test eval_code(fr, "garbage") == ones(10)
eval_code(fr, "non_accessible_variable = 5.0")
@test eval_code(fr, "non_accessible_variable") == 5.0

if VERSION >= v"1.4" # for var"" syntax
    # Evaluating SSAValues
    f(x) = x^2
    frame = JuliaInterpreter.enter_call(f, 5)
    JuliaInterpreter.step_expr!(frame)
    JuliaInterpreter.step_expr!(frame)
    # This could change with changes to Julia lowering
    @test eval_code(frame, "var\"%2\"") == Val(2)
    @test eval_code(frame, "var\"@_1\"") == f
end

function fun(;output=:sym)
   x = 5
   y = 3
end
fr = JuliaInterpreter.enter_call(fun)
fr = JuliaInterpreter.maybe_step_through_wrapper!(fr)
JuliaInterpreter.step_expr!(fr)
@test eval_code(fr, "x") == 5
@test eval_code(fr, "output") == :sym
eval_code(fr, "output = :foo")
@test eval_code(fr, "output") == :foo


text:
function metropolis(T::Real, initial::Real, candidate::Real)
    if candidate < initial
        true
    else
        rand() < exp((initial - candidate) / T)
    end
end

function probabilistic_improvement(tuning_run::Run;
                                   threshold::AbstractFloat = 2.,
                                   acceptance_criterion::Function = metropolis)
    initial_cost  = tuning_run.starting_cost
    x             = deepcopy(tuning_run.starting_point)
    x_proposal    = deepcopy(tuning_run.starting_point)
    name          = "Probabilistic Improvement"
    cost_calls    = 0
    iteration     = 1
    neighbor!(x_proposal)
    proposal      = @fetch (tuning_run.measurement_method(tuning_run, x_proposal))
    cost_calls   += tuning_run.cost_evaluations
    if acceptance_criterion(threshold, initial_cost, proposal)
        update!(x, x_proposal.parameters)
        initial_cost = proposal
    end
    Result(name, tuning_run.starting_point, x, initial_cost, iteration,
           iteration, cost_calls, false)
end


text:
using Poisson2D
using LinearAlgebra
using BenchmarkTools

@show n_threads = Threads.nthreads();
# without the next line, it will use the max number of threads available
BLAS.set_num_threads(n_threads);
# set problem parameters
a = 0;
b = 1;
c = 0;
d = 1;
n = 999;

kx = 5;
ky = 4;
f = (x,y)->((kx*π/(b-a))^2 + (ky * π/(d-c))^2) * sin(kx*π*(x-a)/(b-a)) * sin(ky*π*(y-c)/(d-c));

# set up and solve problem
problem = FDPoisson2DProblem(a,b,c,d,n,f);
assemble_system!(problem);
@btime U = direct_solve_poisson2d(problem);


text:
function add_variables!(m::JuMP.AbstractModel, ps::Array{REoptInputs})

	dvs_idx_on_techs = String[
		"dvSize",
		"dvPurchaseSize",
	]
	dvs_idx_on_techs_timesteps = String[
        "dvCurtail",
		"dvRatedProduction",
	]
	dvs_idx_on_storagetypes = String[
		"dvStoragePower",
		"dvStorageEnergy",
	]
	dvs_idx_on_storagetypes_timesteps = String[
		"dvDischargeFromStorage",
		"dvGridToStorage",
	]
	for p in ps
		_n = string("_", p.s.site.node)
		for dv in dvs_idx_on_techs
			x = dv*_n
			m[Symbol(x)] = @variable(m, [p.techs.all], base_name=x, lower_bound=0)
		end

		for dv in dvs_idx_on_techs_timesteps
			x = dv*_n
			m[Symbol(x)] = @variable(m, [p.techs.all, p.time_steps], base_name=x, lower_bound=0)
		end

		for dv in dvs_idx_on_storagetypes
			x = dv*_n
			m[Symbol(x)] = @variable(m, [p.s.storage.types], base_name=x, lower_bound=0)
		end

		for dv in dvs_idx_on_storagetypes_timesteps
			x = dv*_n
			m[Symbol(x)] = @variable(m, [p.s.storage.types, p.time_steps], base_name=x, lower_bound=0)
		end

		dv = "dvGridPurchase"*_n
		m[Symbol(dv)] = @variable(m, [p.time_steps], base_name=dv, lower_bound=0)

		dv = "dvPeakDemandTOU"*_n
		m[Symbol(dv)] = @variable(m, [p.ratchets, 1], base_name=dv, lower_bound=0)

		dv = "dvPeakDemandMonth"*_n
		m[Symbol(dv)] = @variable(m, [p.months, 1], base_name=dv, lower_bound=0)

		dv = "dvProductionToStorage"*_n
		m[Symbol(dv)] = @variable(m, [p.s.storage.types, p.techs.all, p.time_steps], base_name=dv, lower_bound=0)

		dv = "dvStoredEnergy"*_n
		m[Symbol(dv)] = @variable(m, [p.s.storage.types, 0:p.time_steps[end]], base_name=dv, lower_bound=0)

		dv = "MinChargeAdder"*_n
		m[Symbol(dv)] = @variable(m, base_name=dv, lower_bound=0)

        if !isempty(p.s.electric_tariff.export_bins)
            dv = "dvProductionToGrid"*_n
            m[Symbol(dv)] = @variable(m, [p.techs.elec, p.s.electric_tariff.export_bins, p.time_steps], base_name=dv, lower_bound=0)
        end

		ex_name = "TotalTechCapCosts"*_n
		m[Symbol(ex_name)] = @expression(m, p.third_party_factor *
			sum( p.cap_cost_slope[t] * m[Symbol("dvPurchaseSize"*_n)][t] for t in p.techs.all )
		)

		ex_name = "TotalStorageCapCosts"*_n
		m[Symbol(ex_name)] = @expression(m, p.third_party_factor *
			sum(  p.s.storage.installed_cost_per_kw[b] * m[Symbol("dvStoragePower"*_n)][b]
				+ p.s.storage.installed_cost_per_kwh[b] * m[Symbol("dvStorageEnergy"*_n)][b] for b in p.s.storage.types )
		)

		ex_name = "TotalPerUnitSizeOMCosts"*_n
		m[Symbol(ex_name)] = @expression(m, p.third_party_factor * p.pwf_om *
			sum( p.om_cost_per_kw[t] * m[Symbol("dvSize"*_n)][t] for t in p.techs.all )
		)

        ex_name = "TotalPerUnitProdOMCosts"*_n
		m[Symbol(ex_name)] = 0

		add_elec_utility_expressions(m, p; _n=_n)

		#################################  Objective Function   ########################################
		m[Symbol("Costs"*_n)] = @expression(m,
			# Capital Costs
			m[Symbol("TotalTechCapCosts"*_n)] + m[Symbol("TotalStorageCapCosts"*_n)] +

			## Fixed O&M, tax deductible for owner
			m[Symbol("TotalPerUnitSizeOMCosts"*_n)] * (1 - p.s.financial.owner_tax_pct) +

			# Utility Bill, tax deductible for offtaker, including export benefit
			m[Symbol("TotalElecBill"*_n)] * (1 - p.s.financial.offtaker_tax_pct)
		);
    end
    add_bounds(m, ps)
end

"""
add non-negative bounds to decision variables
"""
function add_bounds(m::JuMP.AbstractModel, ps::Array{REoptInputs})

	dvs_idx_on_techs = String[
		"dvSize",
		"dvPurchaseSize",
	]
	dvs_idx_on_techs_timesteps = String[
        "dvCurtail",
		"dvRatedProduction",
	]
	dvs_idx_on_storagetypes = String[
		"dvStoragePower",
		"dvStorageEnergy",
	]
	dvs_idx_on_storagetypes_timesteps = String[
		"dvDischargeFromStorage",
		"dvGridToStorage",
	]
	for p in ps
        _n = string("_", p.s.site.node)

		for dv in dvs_idx_on_techs
			x = dv*_n
			@constraint(m, [tech in p.techs.all], -m[Symbol(x)][tech] ≤ 0 )
		end

		for dv in dvs_idx_on_techs_timesteps
			x = dv*_n
            @constraint(m, [tech in p.techs.all, ts in p.time_steps],
                -m[Symbol(x)][tech, ts] ≤ 0
            )
		end

		for dv in dvs_idx_on_storagetypes
			x = dv*_n
            @constraint(m, [b in p.s.storage.types],
                -m[Symbol(x)][b] ≤ 0
            )
		end

		for dv in dvs_idx_on_storagetypes_timesteps
			x = dv*_n
            @constraint(m, [b in p.s.storage.types, ts in p.time_steps],
                -m[Symbol(x)][b, ts] ≤ 0
            )
		end

		dv = "dvGridPurchase"*_n
		@constraint(m, [ts in p.time_steps], -m[Symbol(dv)][ts] ≤ 0)

		dv = "dvPeakDemandTOU"*_n
		@constraint(m, [r in p.ratchets], -m[Symbol(dv)][r, 1] ≤ 0)

		dv = "dvPeakDemandMonth"*_n
		@constraint(m, [mth in p.months], -m[Symbol(dv)][mth, 1] ≤ 0)

		dv = "dvProductionToStorage"*_n
        @constraint(m, [b in p.s.storage.types, tech in p.techs.all, ts in p.time_steps],
            -m[Symbol(dv)][b, tech, ts] ≤ 0
        )

		dv = "dvStoredEnergy"*_n
        @constraint(m, [b in p.s.storage.types, ts in 0:p.time_steps[end]],
            -m[Symbol(dv)][b, ts] ≤ 0
        )

		dv = "MinChargeAdder"*_n
		@constraint(m, -m[Symbol(dv)] ≤ 0)
    end
end

function build_reopt!(m::JuMP.AbstractModel, ps::Array{REoptInputs})
    add_variables!(m, ps)
    @warn "Outages are not currently modeled in multinode mode."
    @warn "Diesel generators are not currently modeled in multinode mode."
    for p in ps
        _n = string("_", p.s.site.node)

        for b in p.s.storage.types
            if p.s.storage.max_kw[b] == 0 || p.s.storage.max_kwh[b] == 0
                @constraint(m, [ts in p.time_steps], m[Symbol("dvStoredEnergy"*_n)][b, ts] == 0)
                @constraint(m, m[Symbol("dvStorageEnergy"*_n)][b] == 0)
                @constraint(m, m[Symbol("dvStoragePower"*_n)][b] == 0)
                @constraint(m, [t in p.techs.elec, ts in p.time_steps_with_grid],
                            m[Symbol("dvProductionToStorage"*_n)][b, t, ts] == 0)
                @constraint(m, [ts in p.time_steps], m[Symbol("dvDischargeFromStorage"*_n)][b, ts] == 0)
                @constraint(m, [ts in p.time_steps], m[Symbol("dvGridToStorage"*_n)][b, ts] == 0)
            else
                add_storage_size_constraints(m, p, b; _n=_n)
                add_storage_dispatch_constraints(m, p, b; _n=_n)
            end
        end

        if any(max_kw->max_kw > 0, (p.s.storage.max_kw[b] for b in p.s.storage.types))
            add_storage_sum_constraints(m, p; _n=_n)
        end

        add_production_constraints(m, p; _n=_n)

        if !isempty(p.techs.all)
            add_tech_size_constraints(m, p; _n=_n)
            if !isempty(p.techs.no_curtail)
                add_no_curtail_constraints(m, p; _n=_n)
            end
        end

        add_elec_load_balance_constraints(m, p; _n=_n)

        if !isempty(p.s.electric_tariff.export_bins)
            add_export_constraints(m, p; _n=_n)
        end

        if !isempty(p.s.electric_tariff.monthly_demand_rates)
            add_monthly_peak_constraint(m, p; _n=_n)
        end

        if !isempty(p.s.electric_tariff.tou_demand_ratchet_timesteps)
            add_tou_peak_constraint(m, p; _n=_n)
        end

		if !(p.s.electric_utility.allow_simultaneous_export_import) & !isempty(p.s.electric_tariff.export_bins)
			add_simultaneous_export_import_constraint(m, p; _n=_n)
		end

        if p.s.electric_tariff.demand_lookback_percent > 0
            add_demand_lookback_constraints(m, p; _n=_n)
        end

        if !isempty(p.s.electric_tariff.coincpeak_periods)
            add_coincident_peak_charge_constraints(m, p; _n=_n)
        end

    end
end

function add_objective!(m::JuMP.AbstractModel, ps::Array{REoptInputs})
	if !(any(p.s.settings.add_soc_incentive for p in ps))
		@objective(m, Min, sum(m[Symbol(string("Costs_", p.s.site.node))] for p in ps))
	else # Keep SOC high
		@objective(m, Min, sum(m[Symbol(string("Costs_", p.s.site.node))] for p in ps)
        - sum(sum(m[Symbol(string("dvStoredEnergy_", p.s.site.node))][:elec, ts]
            for ts in p.time_steps) for p in ps) / (8760. / ps[1].hours_per_timestep))
	end  # TODO need to handle different hours_per_timestep?
	nothing
end

function run_reopt(m::JuMP.AbstractModel, ps::Array{REoptInputs})

	build_reopt!(m, ps)

	add_objective!(m, ps)

	@info "Model built. Optimizing..."
	tstart = time()
	optimize!(m)
	opt_time = round(time() - tstart, digits=3)
	if termination_status(m) == MOI.TIME_LIMIT
		status = "timed-out"
    elseif termination_status(m) == MOI.OPTIMAL
        status = "optimal"
    else
        status = "not optimal"
        @warn "REopt solved with " termination_status(m), ", returning the model."
        return m
	end
	@info "REopt solved with " termination_status(m)
	@info "Solving took $(opt_time) seconds."

	tstart = time()
	results = reopt_results(m, ps)
	time_elapsed = time() - tstart
	@info "Total results processing took $(round(time_elapsed, digits=3)) seconds."
	results["status"] = status
	results["solver_seconds"] = opt_time
	return results
end

function reopt_results(m::JuMP.AbstractModel, ps::Array{REoptInputs})
	# TODO address Warning: The addition operator has been used on JuMP expressions a large number of times.
	results = Dict{Union{Int, String}, Any}()
	for p in ps
		results[p.s.site.node] = reopt_results(m, p; _n=string("_", p.s.site.node))
	end
	return results
end


text:
using Gadfly

n = 10
plot(x=1+rand(n), y=1+rand(n),
     Scale.x_log10(minvalue=1.0, maxvalue=10),
     Scale.y_log10(minvalue=1.0, maxvalue=10))


text:
function check_error(m, x0)
    if length(x0) != dim(m)
        throw(ArgumentError("initial variable must have same length as number of design variables"))
    end

    for j in 1:length(x0)
        # x is not in box
        if !(min(m, j) <= x0[j] <= max(m,j))
            throw(ArgumentError("initial variable at index $j outside box constraint"))
        end
    end
end

# From Optim.jl
function assess_convergence(x::Array,
                            x_previous::Array,
                            f_x::Real,
                            f_x_previous::Real,
                            gr::Array,
                            xtol::Real,
                            ftol::Real,
                            grtol::Real)
    x_converged, f_converged, gr_converged, f_increased = false, false, false, false

    x_residual = maxdiff(x, x_previous)
    f_residual = abs(f_x - f_x_previous)
    gr_residual = norm(vec(gr), Inf)

    if x_residual < xtol
        x_converged = true
    end

    # Absolute Tolerance
    # if abs(f_x - f_x_previous) < ftol
    # Relative Tolerance
    if f_residual / (abs(f_x) + ftol) < ftol
        f_converged = true
    end

    if f_x > f_x_previous
        f_increased = true
    end

    if gr_residual < grtol
        gr_converged = true
    end

    converged = x_converged || f_converged || gr_converged

    return x_converged, f_converged, gr_converged, x_residual, f_residual, gr_residual, f_increased, converged
end


text:
function arraydiff(a, b)
  filter(it -> it ∉ b, a)
end

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


max_stars_repo_path: test/7 Inverse hyperbolic functions/7.1 Inverse hyperbolic sine/7.1.5 Inverse hyperbolic sine functions.jl


text:
module NaiveBayes

using Distributions
using HDF5
using KernelDensity
using Grid
using StatsBase

export NBModel,
        MultinomialNB,
        GaussianNB,
        KernelNB,
        HybridNB,
        fit,
        predict,
        predict_proba,
        predict_logprobs,
        restructure_matrix,
        to_matrix,
        write_model,
        load_model,
        get_feature_names,
        train

include("nbtypes.jl")
include("common.jl")
include("hybrid.jl")
include("gaussian.jl")
include("multinomial.jl")

end


text:
function empty_estar_table()
    (
    E             = empty_col(uenergy ),
    Collision     = empty_col(ustoppow),
    Radiative     = empty_col(ustoppow),
    Total         = empty_col(ustoppow),
    CSDA          = empty_col(umassrange),
    RadiationYield= empty_col(unone   ),
    DensityEffect = empty_col(unone   ),
   )
end

struct ESTARData  <: DataSource
    element_tables::Vector{typeof(empty_estar_table())}
end
emptytable(::Type{ESTARData}) = empty_estar_table()

const ESTAR  = load(ESTARData, Zs=1:100, dir=datapath("ESTAR"))

for P in columnnames(ESTARData)
    P == :E && continue
    @eval default_data_source(mat, pt::Electron, ::$P) = ESTAR
end


text:
println("housing")

using Revise

using Pkg
Pkg.activate(".");
Pkg.resolve();
Pkg.instantiate();

using ArgParse
using Bijectors
using Clustering
using CSV
using BenchmarkTools
using DrWatson
using FileIO
using Formatting
using GADM
using GeoDatasets
using Images
using LinearAlgebra
using LineSearches
using Optim
using ParameterHandling
using Plots
using Random
using Stheno
using TemporalGPs
using TexTables
using Zygote

using ParameterHandling: positive, deferred, flatten
using TemporalGPs: approx_posterior_marginals, RegularInTime, Separable
using Zygote: dropgrad

const exp_dir_name = joinpath("apartment_prices");
const data_dir = joinpath(datadir(), exp_dir_name);
const metadata_dir = joinpath(data_dir, "meta");
const training_output_dir = joinpath(data_dir, "training");
const prediction_output_dir = joinpath(data_dir, "predictions");
const results_dir = joinpath(projectdir(), "results", exp_dir_name);

include(joinpath(srcdir(), "housing_data.jl"))

# Pull out the period of time specified from the full data set, sort it, and return it.
function generate_raw_data_set(all_data, first_date::Date, last_date::Date)
    data = all_data[first_date .< all_data.date .< last_date, :]
    return sort(data, :date)
end

# Log and normalise the data set, and return the mean / std used to transform it.
function log_normalised_price(data, train_indices)
    log_price = log.(data.price)
    return (log_price .- mean(log_price[train_indices])) ./ std(log_price[train_indices])
end

# Take a vector of non-sequential `Int`s and construct a vector of sequential `Int`s with
# the same start and end numbers. Return two addition vectors, one containing the indices of
# the elements of the returned vector that are newly added, and the other the indices of
# the existing `Int`s.
function fill_in_empties(x::AbstractVector{Int}, y::AbstractVector, fill_val)
    return vcat(
        map((y, d) -> vcat([y], fill(fill_val, d - 1)), y, diff(x))...,
        [y[end]],
    )
end

smse(predicted, observed) = mean(abs2, filter(!ismissing, mean.(predicted) .- observed))

rsmse(predicted, observed) = sqrt(smse(predicted, observed))

function pplp(predicted, observed)
    present_idx = findall(!ismissing, observed)
    return mean(logpdf.(predicted[present_idx], observed[present_idx]))
end

# Encodes whether or not a property is a new-build as a binary variable.
encode_new_build(new_build_flags::Vector) = map(x -> x == "N" ? 0.0 : 1.0, new_build_flags)

# Converts property types to a 1-hot encoding.
function encode_property_type(property_types::Vector{String})
    ks = unique(property_types)
    vals = eachindex(ks)
    mapping = Dict([k => val for (k, val) in zip(ks, vals)])

    feature_mat = fill(0.0, length(ks), length(property_types))
    for n in eachindex(property_types)
        feature_mat[mapping[property_types[n]], n] = 1.0
    end
    return feature_mat
end

function train_indices_name(start_date::Date, finish_date::Date, N_train::Int)
    return joinpath(data_dir, "$start_date-$finish_date-$N_train.jld2")
end

function maybe_create_train_indices(
    rng::AbstractRNG, start_date::Date, finish_date::Date, N_train::Int, N_total::Int,
)
    fname = train_indices_name(start_date, finish_date, N_train)

    if !isfile(fname)
        @info "CREATING NEW TEST INDICES FOR DATES $start_date-$finish_date-$N_train."
        train_indices = randperm(rng, N_total)[1:N_train]
        FileIO.save(fname, Dict("train_indices" => train_indices))
    else
        @info "NOT CREATING NEW TEST INDICES AS THEY ALREADY EXIST."
    end
end

function get_train_indices(start_date::Date, finish_date::Date, N_train::Int)
    return FileIO.load(
        train_indices_name(start_date, finish_date, N_train),
    )["train_indices"]
end

function generate_data_set(
    data, first_date::Date, last_date::Date, train_indices::Vector{Int},
)
    # Mark train and test data.
    data.is_train = fill(false, size(data, 1))
    data.is_train[train_indices] .= true

    # Standardise the data using only statistics from the training data.
    data.normalised_price = log_normalised_price(data, train_indices)

    # Get new build encoding.
    new_build = encode_new_build(collect(data.new_build))

    # Get property type encoding.
    property_type = encode_property_type(collect(data.property_type))

    data.new_build_enc = new_build
    data.property_type_enc = ColVecs(property_type)

    day_1 = minimum(data.date)
    N_days = (maximum(data.date) - day_1).value + 1
    ts = RegularSpacing(1.0, 1.0, N_days)
    data.integer_day = map(d -> (d - day_1).value + 1, data.date)

    train_data = build_processed_dataset(filter(row -> row.is_train, data), ts)
    test_data = build_processed_dataset(filter(row -> !row.is_train, data), ts)
    return train_data, test_data
end

function build_processed_dataset(data::DataFrame, ts::RegularSpacing)

    # Group the training data by day / integer date.
    daily_data = map(collect(groupby(data, :date))) do d
        X = collect(hcat(d.lat, d.lon)')
        return (
            t = only(unique(d.integer_day)),
            x = ColVecs(X),
            y = collect(d.normalised_price),
        )
    end

    # Construct dictionary from the vector.
    daily_data_dict = Dict([d.t => d for d in daily_data])

    # Define fallbacks for what should be filled in if no data is provided.
    default = (
        x = ColVecs(collect(reshape(first(first(daily_data).x), :, 1))),
        y = [missing],
    )

    # Construct inputs.
    x_space = map(t -> get(daily_data_dict, t, default).x, ts)
    inputs = RegularInTime(ts, x_space)

    # Construct outputs.
    outputs = vcat(map(t -> get(daily_data_dict, t, default).y, ts)...)

    return (inputs=inputs, outputs=outputs)
end

# Utility functionality for working with `ParameterHandling` a bit more conveniently.
function wrap_objective(f, θ)
    θ_flat, unflatten = flatten(θ)
    f_flat(θ_flat) = f(ParameterHandling.value(unflatten(θ_flat)))
    return f_flat, θ_flat, unflatten
end

function ParameterHandling.flatten(x::ColVecs)
    X_flat, unflatten_X = flatten(x.X)
    unflatten_ColVecs(X_flat) = ColVecs(unflatten_X(X_flat))
    return X_flat, unflatten_ColVecs
end

# Make it possible to parse strings from the command line into Vector{Int}s.
function ArgParse.parse_item(::Type{Vector{Int}}, x::String)
    @assert x[1] == '[' && x[end] == ']'
    return CSV.read(IOBuffer(x[2:prevind(x, end)]); header=false, transpose=true).Column1
end

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--run-experiments"
            help = "true / false - run the experiments or not"
            arg_type = Bool
            default = true
        "--M"
            help = "number of pseudo-points"
            arg_type = Int
            default = 50
    end
    return parse_args(s)
end

const args = parse_commandline()

# Just load the big data set once and re-use it lots. It's quite time consuming to load it
# back up lots of times, particularly when iterating on this script.
const all_data = @time generate_housing_data(data_dir);

# These aren't actually London's coordinates, they're for the whole of England + Wales.
london_lons = (-6, 2);
london_lats = (50, 56);

struct GridBox{Tlats<:Tuple{<:Real, <:Real}, Tlons<:Tuple{<:Real, <:Real}}
    lats::Tlats
    lons::Tlons
end

function in_box(box::GridBox, lat::Real, lon::Real)
    return box.lats[1] <= lat && box.lats[2] > lat &&
        box.lons[1] <= lon && box.lons[2] > lon
end

# Filter observational data and restrict to a box around Greater London.
london_data = filter(
    row -> in_box(GridBox(london_lats, london_lons), row.lat, row.lon) &&
        row.property_type == "F",
    all_data,
);

# Define locations at which to compute the posterior marginals.
N_base = 50; # related to how densely we plot the points.
N_lons = Int(floor(N_base * (london_lons[2] - london_lons[1])));
N_lats = Int(floor(N_base * (london_lats[2] - london_lats[1])));
const london_lats_range = range(london_lats[1], london_lats[2]; length=N_lats);
const london_lons_range = range(london_lons[1], london_lons[2]; length=N_lons);
const lat_lon_pairs = tuple.(london_lats_range, london_lons_range');

BLAS.set_num_threads(8)

# Set proportion of data to be used as training.
const train_prop = 0.5

#
# Specify experiments to run.
#

mkpath(metadata_dir)
const settings_fname = joinpath(metadata_dir, "settings.bson")

let
    wsave(
        settings_fname,
        :settings => dict_list(
            Dict(
                # The total number of pseudo-points to use per time step.
                :M => [75],

                # Kernels to try, specified using the `ParameterHandling.jl` interface.
                :kernel_info => [
                    (
                        name = "separable",
                        value = deferred(
                            (λ_time, λ_space, s) -> Separable(
                                s * stretch(EQ(), λ_space),
                                stretch(Matern32(), λ_time),
                            ),
                            positive(1e-2),
                            fill(positive(1.0), 2),
                            positive(1.0),
                        ),
                    ),
                    (
                        name = "sum-separable",
                        value = deferred(
                            (λs_time, λs_space, ss) -> begin
                                k1 = Separable(
                                    ss[1] * stretch(EQ(), λs_space[1]),
                                    stretch(Matern32(), λs_time[1]),
                                )
                                k2 = Separable(
                                    ss[2] * stretch(EQ(), λs_space[2]),
                                    stretch(Matern32(), λs_time[2]),
                                )
                                return k1 + k2
                            end,
                            (positive(1e-3), positive(1e-1)),
                            (
                                fill(positive(1.0), 2),
                                fill(positive(5.0), 2),
                            ),
                            (positive(0.7), positive(0.3)),
                        ),
                    ),
                ],

                # The dates between which the training data is selected.
                :date_range => [
                    (start=Date(2010), stop=Date(2020)),
                ]
            ),
        ),
    );
end

#
# Run experiments. This is split into two sections:
# 1. training: this takes by far the longest amount of time. The entire model is saved as an
#     output from this step.
# 2. prediction: loads the results from step 1 and uses them to make predictions. Not run in
#     the training block because you don't want to have to re-run the training procedure to
#     make slightly different kinds of predictions. Not in the results block because it can
#     take a while, which is less than helpful when tweaking plots.
#

if args["run-experiments"]

println("Running experiments")

mkpath(training_output_dir)

let
    println("Loading housing data")

    settings = load(settings_fname)[:settings];

    for (n, setting) in enumerate(settings)

        println("Experiment $n / $(length(settings))")
        display(setting)
        println()

        # Generate the data set.
        first_date = setting[:date_range].start
        last_date = setting[:date_range].stop

        london_data_local = generate_raw_data_set(london_data, first_date, last_date);
        N_total = size(london_data_local, 1);
        N_train = Int(floor(train_prop * N_total));
        maybe_create_train_indices(
            MersenneTwister(123456), first_date, last_date, N_train, N_total,
        )

        train_indices = get_train_indices(first_date, last_date, N_train);
        train_data, test_data = generate_data_set(
            london_data_local, first_date, last_date, train_indices,
        );

        inputs = train_data.inputs
        outputs = train_data.outputs
        println("Total number of training data is $(length(outputs)).")
        println("Total number of test data is $(length(test_data.outputs))")

        # Compute pseudo-point initialisation.
        x_r_full = ColVecs(hcat(getfield.(inputs.vs, :X)...))
        Random.seed!(123456)
        clustering_result = kmeans(x_r_full.X, setting[:M])
        Z_r = clustering_result.centers

        # Specify the initial parameters in structured form.
        initial_parameters = (
            k = setting[:kernel_info].value,
            σ² = bounded(0.5, 1e-2, 2.0),
            Z_r = fixed(Z_r),
        )

        # Create the objective function, wrapping it in `ParameterHandling.jl` stuff to
        # make it play nicely with Optim.
        objective, flat_initial_parameters, unflatten = wrap_objective(
            θ -> begin
                f = to_sde(GP(θ.k, GPC()))
                return -elbo(f(dropgrad(inputs), θ.σ²), dropgrad(outputs), ColVecs(θ.Z_r))
            end,
            initial_parameters,
        )

        # Construct the optimiser.
        training_results = Optim.optimize(
            objective,
            θ -> Zygote.gradient(objective, θ)[1],
            flat_initial_parameters,
            LBFGS(
                alphaguess = LineSearches.InitialStatic(scaled=true),
                linesearch = LineSearches.BackTracking(),
                m=50,
            ),
            # The numbers involved in my objective function are generally quite large,
            # so looser tolerances are appropriate.
            Optim.Options(
                x_tol = 1e-4,
                f_tol = 1e-8,
                g_tol = 1e-1,
                f_calls_limit = 1_000,
                g_calls_limit = 100,
                iterations = 1_000,
                show_trace = true,
                extended_trace=true,
                show_every=10,
            );
            inplace=false,
        )

        wsave(
            joinpath(
                training_output_dir,
                savename(
                    Dict(
                        :M => setting[:M],
                        :kernel => setting[:kernel_info].name,
                        :start_date => first_date,
                    ),
                    "bson",
                ),
            ),
            Dict(
                :setting => setting,
                :training_results => training_results,
                :unflatten => unflatten,
            ),
        )
    end
end

end

#
# Make predictions at locations that we wish to view them, and compute RMSE on data etc.
#

println("Generating predictions.")

mkpath(prediction_output_dir)

let
    trained_models = collect_results(training_output_dir);

    for (n, trained_model_) in enumerate(eachrow(trained_models))

        println("n, $(trained_model_.path)")

        # Defensive copy. Something in here is mutating for some unknown reason...
        trained_model = deepcopy(trained_model_);

        # Generate the data set.
        first_date = trained_model.setting[:date_range].start
        last_date = trained_model.setting[:date_range].stop

        # Generate the dataset.
        london_data_local = generate_raw_data_set(london_data, first_date, last_date);
        N_total = size(london_data_local, 1);
        N_train = Int(floor(train_prop * N_total));
        train_indices = get_train_indices(first_date, last_date, N_train);
        train_data, test_data = generate_data_set(
            london_data_local, first_date, last_date, train_indices,
        );
        println("Total number of training data is $(length(train_data.outputs)).")
        println("Total number of test data is $(length(test_data.outputs))")

        # Build the model.
        unflatten = trained_model.unflatten
        results = trained_model.training_results
        learned_parameters = ParameterHandling.value(unflatten(results.minimizer))

        f = to_sde(GP(learned_parameters.k, GPC()))
        fx = f(train_data.inputs, learned_parameters.σ²);

        # Make predictions at entire grid of points at which we might consider predicting.
        get_vector(x::Tuple) = [x[1], x[2]]
        x_r_pr = ColVecs(
            vcat(
                hcat(get_vector.(lat_lon_pairs)...),
                # repeat([0.0, 0.0, 1.0, 0.0, 0.0, 0.0], 1, length(lat_lon_pairs)),
            ),
        );

        # Compute posterior marginals for the sake of plotting.
        z_r = ColVecs(learned_parameters.Z_r)
        pos = length(TemporalGPs.get_time(train_data.inputs)) - 100
        @show pos, length(TemporalGPs.get_time(train_data.inputs))
        preds = @time approx_posterior_marginals(
            dtc, fx, train_data.outputs, z_r, x_r_pr, pos,
        )

        # Compute posterior marginals at the training data and the test data.
        preds_train = @time approx_posterior_marginals(
            dtc, fx, train_data.outputs, z_r, train_data.inputs,
        )
        preds_test = @time approx_posterior_marginals(
            dtc, fx, train_data.outputs, z_r, test_data.inputs,
        )

        # Write the results.
        wsave(
            joinpath(
                prediction_output_dir,
                savename(
                    Dict(
                        :M => trained_model.setting[:M],
                        :kernel => trained_model.setting[:kernel_info].name,
                        :start_date => first_date,
                    ),
                    "bson",
                ),
            ),
            Dict(
                "trained_model" => trained_model,
                "preds_train" => preds_train,
                "preds_test" => preds_test,
                "preds_extrapolate" => preds,
                "x_r_pr" => x_r_pr,
                "lat_lon_pairs" => lat_lon_pairs,
                "smse_train" => rsmse(preds_train, train_data.outputs),
                "smse_test" => rsmse(preds_test, test_data.outputs),
                "pplp_train" => pplp(preds_train, train_data.outputs),
                "pplp_test" => pplp(preds_test, test_data.outputs),
            ),
        )
    end
end

#
# Visualise results.
#

# Load shoreline data.
# const shore_data = GeoDatasets.gshhg('f');

my_format(x) = format(round(x; sigdigits=3); stripzeros=true)

let
    gr()

    results = collect_results(prediction_output_dir)

    name_map = Dict(
        "separable" => "Separable",
        "sum-separable" => "Sum-Separable",
    );

    # Compute some additional fields necessary for table creation.
    results.M = map(x -> x.setting[:M], results.trained_model);
    results.start_date = map(x -> x.setting[:date_range].start, results.trained_model);
    results.rsmse_test_string = map(my_format, results.smse_test);
    results.npplp_test_string = map(my_format, -results.pplp_test);
    results.row_keys = map(
        x -> "($(x.setting[:M])) - " * name_map[x.setting[:kernel_info].name],
        results.trained_model,
    );

    # Construct tables for each training window.
    for group in groupby(results, :start_date)

        headers = [
            "RSMSE",
            "NPPLP",
        ]
        col_data = [
            group.rsmse_test_string,
            group.npplp_test_string,
        ]
        cols = [
            TableCol(header, collect(group.row_keys), collect(data)) for
            (header, data) in zip(headers, col_data)
        ]
        results_tex_table = hcat(cols...)

        # Generate paper-ready tables of results, comparing each model.
        start_date = only(unique(group.start_date))
        open(joinpath(results_dir, "performance_results_date=$start_date.txt"), "w") do io
            write(io, to_ascii(results_tex_table))
        end
        open(joinpath(results_dir, "performance_results_date=$start_date.tex"), "w") do io
            write(io, to_tex(results_tex_table))
        end
    end

    box = GridBox(london_lats, london_lons);

    # London boroughs info, including geometries.
    # boroughs = GADM.get("GBR", "England", "Greater London"; children=true)[2];
    boroughs = vcat(
        GADM.get("GBR", "England"; children=true)[2].geom,
        GADM.get("GBR", "Wales"; children=true)[2].geom,
    );

    for result in eachrow(results)

        trained_model = result.trained_model;
        preds_mat = reshape(result.preds_extrapolate, N_lats, :);

        # Generate the dataset.
        first_date = trained_model.setting[:date_range].start
        last_date = trained_model.setting[:date_range].stop
        london_data_local = generate_raw_data_set(london_data, first_date, last_date);
        N_total = size(london_data_local, 1);
        N_train = Int(floor(train_prop * N_total));
        train_indices = get_train_indices(first_date, last_date, N_train);
        train_data, test_data = generate_data_set(
            london_data_local, first_date, last_date, train_indices,
        );

        x_r = train_data.inputs.vs[end-20];

        unflatten = trained_model.unflatten
        results = trained_model.training_results
        learned_parameters = ParameterHandling.value(unflatten(results.minimizer))
        z_r = ColVecs(learned_parameters.Z_r)

        colour_scheme = :viridis;

        let
            plt = plot(
                xlims=london_lons,
                ylims=london_lats,
                xticks=nothing,
                yticks=nothing,
            )

            heatmap!(
                plt, london_lons_range, london_lats_range, mean.(preds_mat);
                label="", c=cgrad(colour_scheme), right_margin=6Plots.mm,
            );

            scatter!(
                plt, getindex.(z_r, 2), first.(z_r);
                label="z",
                markersize=4,
                color=colorant"#F5793A", # Orange
            )

            # Plot the Greater London boroughs.
            for b in boroughs
                plot!(plt, b; fillcolor=nothing)
            end

            mkpath(results_dir)
            savefig(
                plt,
                joinpath(
                    results_dir,
                    savename(
                        Dict(
                            :M => trained_model.setting[:M],
                            :kernel => trained_model.setting[:kernel_info].name,
                            :start_date => first_date,
                        ),
                        "posterior_mean.png",
                    ),
                )
            )
        end

        let
            plt = plot(
                xlims=london_lons,
                ylims=london_lats,
                xticks=nothing,
                yticks=nothing,
            )

            heatmap!(
                london_lons_range, london_lats_range, std.(preds_mat);
                label="", c=colour_scheme, right_margin=3Plots.mm,
            )

            scatter!(
                plt, getindex.(z_r, 2), first.(z_r);
                label="z",
                markersize=4,
                color=colorant"#F5793A", # Orange
            )

            # Plot the Greater London boroughs.
            for b in boroughs
                plot!(plt, b; fillcolor=nothing)
            end

            mkpath(results_dir)
            savefig(
                plt,
                joinpath(
                    results_dir,
                    savename(
                        Dict(
                            :M => trained_model.setting[:M],
                            :kernel => trained_model.setting[:kernel_info].name,
                            :start_date => first_date,
                        ),
                        "posterior_std.png",
                    ),
                )
            )
        end
    end
end


text:
module ZfpCompression

    using zfp_jll

    export zfp_compress, zfp_decompress, zfp_decompress!

    include("zfp.jl")

end


text:
module Proj4

using Libdl

export Projection, # proj_types.jl
       transform, transform!,  # proj_functions.jl
       is_latlong, is_geocent, compare_datums, spheroid_params,
       xy2lonlat, xy2lonlat!, lonlat2xy, lonlat2xy!

# Load in `deps.jl`, complaining if it does not exist
const depsjl_path = joinpath(@__DIR__, "..", "deps", "deps.jl")
if !isfile(depsjl_path)
    error("Proj4 not installed properly, run Pkg.build(\"Proj4\"), restart Julia and try again")
end
include(depsjl_path)

# Module initialization function
function __init__()
    # Always check your dependencies from `deps.jl`
    check_deps()
end

include("projection_codes.jl") # ESRI and EPSG projection strings
include("proj_capi.jl") # low-level C-facing functions (corresponding to src/proj_api.h)

function _version()
    m = match(r"(\d+).(\d+).(\d+),.+", _get_release())
    VersionNumber(parse(Int, m[1]), parse(Int, m[2]), parse(Int, m[3]))
end

"Parsed version number for the underlying version of libproj"
const version = _version()

# Detect underlying libproj support for geodesic calculations
const has_geodesic_support = version >= v"4.9.0"

if has_geodesic_support
    export geod_direct, geod_inverse, geod_destination, geod_distance
    include("proj_geodesic.jl") # low-level C-facing functions (corresponding to src/geodesic.h)
end

include("proj_types.jl") # type definitions for proj objects
include("proj_functions.jl") # user-facing proj functions

"Get a global error string in human readable form"
error_message() = _strerrno()

end # module


text:
# TODO: disabling three-body terms
"""
    Transcorrelated1D(address; t=1.0, v=1.0, v_ho=0.0, cutoff=1, three_body_term=true)

Implements a transcorrelated Hamiltonian for contact interactions in one dimensional
momentum space from [Jeszenski *et al.* (2018)](http://arxiv.org/abs/1806.11268).
Currently limited to two component fermionic addresses.

```math
\\begin{aligned}

\\tilde{H} &= t \\sum_{kσ}k^2 n_{k,σ} \\\\
    &\\quad + \\sum_{pqkσσ'} T_{pqk} a^†_{p-k,σ} a^†_{q+k,σ'} a_{q,σ'} a_{p,σ} \\\\
    &\\quad + \\sum_{pqskk'σσ'} Q_{kk'}a^†_{p-k,σ} a^†_{q+k,σ} a^†_{s+k-k',σ'}
                                       a_{s,σ'} a_{q,σ} a_{p,σ} \\\\
    &\\quad + V̂_\\mathrm{ho}
\\end{aligned}
```

where

```math
\\begin{aligned}
\\tilde{u}(k) &= \\begin{cases} -\\frac{2}{k^2} &\\mathrm{if\\ } |k| ≥ k_c\\\\
0 & \\mathrm{otherwise}
\\end{cases}
\\\\

T_{pqk} &= \\frac{v}{M} + \\frac{2v}{M}\\left[k^2\\tilde{u}(k)
          - (p - q)k\\tilde{u}(k)\\right] + \\frac{2v^2}{t}W(k)\\\\
W(k) &= \\frac{1}{M^2}\\sum_{q} (k - q)q\\, \\tilde{u}(q)\\,\\tilde{u}(k - q) \\\\
Q_{kl} &= -\\frac{v^2}{t M^2}k \\tilde{u}(k)\\,l\\tilde{u}(l),
\\end{aligned}
```

# Arguments

* `address`: The starting address, defines number of particles and sites.
* `v`: The interaction parameter.
* `t`: The kinetic energy prefactor.
* `v_ho`: Strength of the external harmonic oscillator potential ``V̂_\\mathrm{ho}``.
  See [`HubbardMom1DEP`](@ref).
* `cutoff` controls ``k_c`` in equations above. Note: skipping generating
  off-diagonal elements below the cutoff is not implemented - zero-valued elements
  are returned instead.
* `three_body_term`: If set to false, generating three body excitations is skipped.
  Note: when disabling three body terms, cutoff should be set to a higher value for
  best results.

# See also

* [`HubbardMom1D`](@ref)
* [`HubbardMom1DEP`](@ref)
"""
struct Transcorrelated1D{
    M,F<:CompositeFS{2,<:Any,M,<:Tuple{FermiFS,FermiFS}}, # TODO: relax this to allow bosons
    P<:Union{Nothing,SVector{M,Float64}}
} <: AbstractHamiltonian{Float64}
    address::F
    cutoff::Int
    v::Float64
    t::Float64
    v_ho::Float64
    ks::SVector{M,Float64} # wave numbers
    kes::SVector{M,Float64} # single-particle dispersion
    ws::SVector{M,Float64} # pre-computed W(k)
    us::SVector{M,Float64} # correlation factor
    potential::P # external potential
    three_body_term::Bool
end

function Transcorrelated1D(address; t=1.0, v=1.0, v_ho=0.0, cutoff=1, three_body_term=true)
    M = num_modes(address)
    cutoff < 1 && throw(ArgumentError("`cutoff` must be a positive integer"))
    ks = SVector{M}(i_to_k.(1:M, M))
    kes = t .* ks.^2
    ws = SVector{M}(w_function.(0:M-1, cutoff))
    us = SVector{M}(correlation_factor.(1:M, cutoff, M))
    if iszero(v_ho)
        potential = nothing
    else
        potential = momentum_space_harmonic_potential(M, v_ho)
    end

    return Transcorrelated1D{M,typeof(address),typeof(potential)}(
        address, cutoff, float(v), float(t), float(v_ho), ks, kes, ws, us, potential,
        three_body_term
    )
end

function Base.show(io::IO, h::Transcorrelated1D)
    print(io, "Transcorrelated1D(", starting_address(h), ", t=$(h.t), v=$(h.v)")
    if !iszero(h.v_ho)
        print(io, ", v_ho=$(h.v_ho))")
    else
        print(io, ")")
    end
end

LOStructure(::Type{<:Transcorrelated1D}) = AdjointUnknown()

function get_offdiagonal(h::Transcorrelated1D{<:Any,F}, add::F, i) where {F}
    return offdiagonals(h, add)[i]
end
function num_offdiagonals(h::Transcorrelated1D{<:Any,F}, add::F) where {F}
    return length(offdiagonals(h, add))
end

# Note: n starts at 0 in the middle of the address, and can be positive or negative
@inline i_to_k(i, M) = i_to_n(i, M) * 2π/M
@inline i_to_n(i, M) = i - M ÷ 2 - isodd(M)
@inline n_to_i(n, M) = n + M ÷ 2 + isodd(M)
@inline n_to_k(n, M) = i_to_k(n_to_i(n, M), M)

"""
    correlation_factor(n, cutoff, M)

Compute the (dimensionless) correlation factor multiplied by ``k``:

```math
\\begin{aligned}
k \\tilde{u}(k) &= \\begin{cases} -\\frac{2}{k} &\\mathrm{if\\ } |k| ≥ k_c \\\\
0 & \\mathrm{otherwise}
\\end{cases}
\\end{aligned}
```

where ``k = π + 2πn/M`` and `k_c = π + 2π/M * cutoff`.

See also [`Transcorrelated1D`](@ref).
"""
function correlation_factor(n, cutoff, M)
    if abs(n) ≥ cutoff
        -1/2n_to_k(n, M)
    else
        0.0
    end
end
function correlation_factor(h::Transcorrelated1D{M}, n) where {M}
    absn = abs(n)
    if absn > 0
        return sign(n) * h.us[absn]
    else
        return 0.0
    end
end

"""
    w_function(n::Integer, nc::Integer)
    w_function(h::Transcorrelated1D, n::Integer)

Compute the (dimensionless) function
```math
W(k) = \\frac{1}{M^2}\\sum_{q} (k - q)q\\, \\tilde{u}(q)\\,\\tilde{u}(k - q) .
```
where ``k = π + 2π\\mathtt{n}/M``,  ``k_c = π + 2π\\mathtt{nc}/M``,
and ``k\\tilde{u}(k)`` is the [`correlation_factor`](@ref).

See also [`Transcorrelated1D`](@ref).
"""
function w_function(n::Integer, nc::Integer)
    prefactor = -1 / (8π^2)
    n = abs(n)

    if n == 0
        x = π^2/6 - sum(1/(np^2) for np in 1:nc-1; init=0.0)
    elseif 2nc > n > 0
        x = 1/n * sum(1/np for np in nc:n+nc-1)
    else
        x = 1/n * sum(1/np for np in nc:n+nc-1) -
            1/2 * sum(1/(np * (n - np)) for np in nc:n-nc)
    end

    return prefactor * x
end
w_function(h::Transcorrelated1D, n::Integer) = h.ws[abs(n) + 1]

"""
    t_function(h::Transcorrelated1D, p, q, k)

Compute

```math
T_{pqk} = \\frac{v}{M} + \\frac{2v}{M}\\left[k^2\\tilde{u}(k)
          - (p - q)k\\tilde{u}(k)\\right] + \\frac{2v^2}{t}W(k)
```

where ``k\\tilde{u}(k)`` is the [`correlation_factor`](@ref) and ``W(k)`` is the
[`w_function`](@ref).

See also [`Transcorrelated1D`](@ref).
"""
function t_function(h::Transcorrelated1D{M}, p, q, k) where {M}
    @unpack t, v = h
    k_pi = n_to_k(k, M)
    pmq_pi = n_to_k(p - q, M)
    cor_k = correlation_factor(h, k)
    return v/M + 2v/M * (cor_k * k_pi - cor_k * pmq_pi) +
        2v^2/t * w_function(h, k)
end

"""
    q_function(h::Transcorrelated1D, k, l)

Compute

```math
Q_{kl} = -\\frac{v^2}{t M^2}k \\tilde{u}(k)\\,l\\tilde{u}(l),
```

where ``k\\tilde{u}(k)`` is the [`correlation_factor`](@ref).

See also [`Transcorrelated1D`](@ref).
"""
function q_function(h::Transcorrelated1D{M}, k, l) where {M}
    @unpack t, v = h
    cor_k = correlation_factor(h, k)
    cor_l = correlation_factor(h, l)

    return -v^2/(t * M^2) * cor_k * cor_l
end

starting_address(ham::Transcorrelated1D) = ham.address

function momentum_transfer_diagonal(h::Transcorrelated1D{M}, map1, map2) where {M}
    @unpack v, t = h
    return momentum_transfer_diagonal(map1, map2) * (v/M + 2v^2/t * w_function(h, 0)) / 2
end

"""
    transcorrelated_diagonal(h::Transcorrelated1D, map1, map2)

The diagonal part of [`transcorrelated_three_body_excitation`](@ref).
"""
function transcorrelated_diagonal(h::Transcorrelated1D{M}, map1, map2) where {M}
    value = 0.0
    for p in 1:length(map1)
        for q in 1:p-1
            k = map1[p].mode - map1[q].mode
            qkk = q_function(h, -k, k)
            value += 2 * qkk * length(map2)
        end
    end
    return value
end

function diagonal_element(h::Transcorrelated1D{<:Any,F}, add::F) where {F}
    c1, c2 = add.components
    map1 = OccupiedModeMap(c1)
    map2 = OccupiedModeMap(c2)

    value = dot(h.kes, map1) + dot(h.kes, map2) +
        momentum_transfer_diagonal(h, map1, map2)
    if h.three_body_term
        value += transcorrelated_diagonal(h, map1, map2) +
            transcorrelated_diagonal(h, map2, map1)
    end
    if !isnothing(h.potential)
        value += momentum_external_potential_diagonal(h.potential, c1, map1) +
            momentum_external_potential_diagonal(h.potential, c2, map2)
    end

    return value
end

struct Transcorrelated1DOffdiagonals{H,A,O1,O2}<:AbstractOffdiagonals{A,Float64}
    hamiltonian::H
    address::A
    map1::O1
    map2::O2
    length::Int
end

function offdiagonals(h::Transcorrelated1D{M,F}, add::F) where {M,F}
    offdiags = Tuple{F,Float64}[]
    c1, c2 = add.components
    map1 = OccupiedModeMap(c1)
    map2 = OccupiedModeMap(c2)
    N1 = length(map1)
    N2 = length(map2)
    n_mom = N1 * N2 * (M - 1)

    three_body_term = h.three_body_term
    n_trans1 = three_body_term ? N1 * (N1 - 1) * N2 * M * M : 0
    n_trans2 = three_body_term ? N2 * (N2 - 1) * N1 * M * M : 0

    if !isnothing(h.potential)
        n_pot1 = N1 * (M - 1)
        n_pot2 = N2 * (M - 1)
    else
        n_pot1 = 0
        n_pot2 = 0
    end
    len = n_mom + n_trans1 + n_trans2 + n_pot1 + n_pot2
    return Transcorrelated1DOffdiagonals(h, add, map1, map2, len)
end

Base.size(od::Transcorrelated1DOffdiagonals) = (od.length,)

function Base.getindex(od::Transcorrelated1DOffdiagonals, i)
    @unpack address, map1, map2 = od
    c1, c2 = address.components
    h = od.hamiltonian
    N1 = length(map1)
    N2 = length(map2)
    M = num_modes(c1)

    n_mom = N1 * N2 * (M - 1)
    three_body_term = od.hamiltonian.three_body_term
    n_trans1 = three_body_term ? N1 * (N1 - 1) * N2 * M * M : 0
    n_trans2 = three_body_term ? N2 * (N2 - 1) * N1 * M * M : 0

    # This should be efficient as it depends on the type of the potential
    if !isnothing(od.hamiltonian.potential)
        n_pot1 = N1 * (M - 1)
        n_pot2 = N2 * (M - 1)
    else
        n_pot1 = 0
        n_pot2 = 0
    end

    # Fallback on zero values
    new_c = CompositeFS(c1, c2)

    if i ≤ n_mom
        # Momentum transfer
        new_c1, new_c2, value, p, q, k = momentum_transfer_excitation(
            c1, c2, i, map1, map2; fold=false
        )
        if !iszero(value)
            @assert new_c1 ≠ c1 || new_c2 ≠ c2
            value *= t_function(h, p, q, k)
            new_c = CompositeFS(new_c1, new_c2)
        end
    elseif i ≤ n_mom + n_trans1
        # Transcorrelated excitation from first to second component
        i -= n_mom

        new_c1, new_c2, value, k, l = transcorrelated_three_body_excitation(
            c1, c2, i, map1, map2
        )
        value *= q_function(h, k, l)
        if !iszero(value)
            @assert new_c1 ≠ c1 || new_c2 ≠ c2
            new_c = CompositeFS(new_c1, new_c2)
        end
    elseif i ≤ n_mom + n_trans1 + n_trans2
        # Transcorrelated excitation from second to first component
        i -= n_mom + n_trans1

        new_c2, new_c1, value, k, l = transcorrelated_three_body_excitation(
            c2, c1, i, map2, map1
        )
        value *= q_function(h, k, l)
        if !iszero(value)
            @assert new_c1 ≠ c1 || new_c2 ≠ c2
            new_c = CompositeFS(new_c1, new_c2)
        end
    elseif i ≤ n_mom + n_trans1 + n_trans2 + n_pot1
        # Potential acting on first component
        i -= n_mom + n_trans1 + n_trans2

        new_c1, value = momentum_external_potential_excitation(
            od.hamiltonian.potential, c1, i, map1
        )
        if !iszero(value)
            new_c = CompositeFS(new_c1, c2)
        end
    elseif i ≤ n_mom + n_trans1 + n_trans2 + n_pot1 + n_pot2
        # Potential acting on second component
        i -= n_mom + n_trans1 + n_trans2 + n_pot1

        new_c2, value = momentum_external_potential_excitation(
            od.hamiltonian.potential, c2, i, map2
        )
        if !iszero(value)
            new_c = CompositeFS(c1, new_c2)
        end
    else
        throw(BoundsError(od, i))
    end
    return new_c, value
end