text:
# Copyright (c) 2020-2021, Bank of Canada
# All rights reserved.

# -------------------------------------------------------------------------------
# TSeries struct
# -------------------------------------------------------------------------------

"""
    mutable struct TSeries{F, T, C} <: AbstractVector{T}
        firstdate::MIT{F}
        values::C
    end

Time series with frequency `F` with values of type `T` stored in a container of
type `C`. By default the type is `Float64` and the container is
`Vector{Float64}`.

Construction:
    ts = TSeries(args...)

    The standard construction is `TSeries(firstdate::MIT, values::AbstractVector)`

    If the first argument is an MIT-range (instead or an MIT), then the length
    of the `values` container must match the length of the given range.

    In the case of a range argument, the `values` can be omitted, in which case
    the container is initializes with `undef`. Or you can also pass a constant
    and then the `values` will be filled with that constant. To accomplish this,
    you can also use `fill`, e.g., `TSeries(20Q1:20Q4, 5)` is the same as
    `fill(5, 20Q1:20Q4)`.

    If only a `firstdate::MIT` is given, the `values` container is initialized
    to an empty `Vector`.

    If only an `n::Integer` is given, it is the same as passing the range
    `0U .+ (1:n)`. An initialization argument is not allowed in this case.

    A `TSeries` can also be constructed with `copy`, `similar`, and `fill`.

Indexing:

    Indexing with an `MIT` or a range of `MIT` works as you'd expect.

    Indexing with `Integer`s works the same as with `Vector`.

    Indexing with `Bool`-array works as you'd expect. For example,
    `s[s .< 0.0] .*= -1` multiplies in place the negative entries of `s` by -1,
    so effectively it's the same as `s .= abs.(s)`.

    There are important differences between indexing with MIT and not
    using MIT (i.e., using Integer or Bool-array).

    * with MIT-range we return a TSeries with the given range, otherwise we
      return a `Vector`

    * the range can be extended (the TSeries resized appropriately) by assigning
      outside the current range. This works only with MIT (you get a BoundsError
      if you try to assign outside the Integer range).

    * `begin` and `end` are MIT, so either use both or none of them. For example
      `s[2:end]` doesn't work because 2 is an `Int` and `end` is an `MIT`. You
      should use `s[begin+1:end]`.

"""
mutable struct TSeries{F<:Frequency,T<:Number,C<:AbstractVector{T}} <: AbstractVector{T}
    firstdate::MIT{F}
    values::C
end

Base.values(t::TSeries) = values(t.values)
@inline firstdate(t::TSeries) = t.firstdate
@inline lastdate(t::TSeries) = t.firstdate + length(t.values) - one(t.firstdate)

@inline frequencyof(::Type{<:TSeries{F}}) where {F<:Frequency} = F

"""
    rangeof(s)

Return the stored range of the given time series object.
"""
@inline rangeof(t::TSeries) = firstdate(t):lastdate(t)

"""
    firstdate(ts), lastdate(ts)

Return the first and last date of the allocated data for the given `TSeries`.
These are identical to `firstindex` and `lastindex`.
"""
firstdate, lastdate

# -------------------------------------------------------------------------------
# some methods that make the AbstractArray infrastructure of Julia work with TSeries

@inline Base.size(t::TSeries) = size(t.values)
@inline Base.axes(t::TSeries) = (firstdate(t):lastdate(t),)
@inline Base.axes1(t::TSeries) = firstdate(t):lastdate(t)

# the following are needed for copy() and copyto!() (and a bunch of Julia internals that use them)
Base.IndexStyle(::TSeries) = IndexLinear()
Base.dataids(t::TSeries) = Base.dataids(getfield(t, :values))

# normally only the first of the following is sufficient.
# we add few other versions of similar below
Base.similar(t::TSeries) = TSeries(t.firstdate, similar(t.values))

# -------------------------------------------------------------------------------
# Indexing with integers and booleans - same as vectors

# indexing with integers is plain and simple
Base.getindex(t::TSeries, i::Int) = getindex(t.values, i)
Base.setindex!(t::TSeries, v::Number, i::Int) = (setindex!(t.values, v, i); t)

# indexing with integer arrays, ranges of integers, and Bool arrays
Base.getindex(t::TSeries, i::AbstractRange{Int}) = getindex(t.values, i)
Base.getindex(t::TSeries, i::AbstractArray{Int}) = getindex(t.values, values(i))
Base.getindex(t::TSeries, i::AbstractArray{Bool}) = getindex(t.values, values(i))
Base.setindex!(t::TSeries, v, i::AbstractRange{Int}) = (setindex!(t.values, v, i); t)
Base.setindex!(t::TSeries, v, i::AbstractArray{Int}) = (setindex!(t.values, v, values(i)); t)
Base.setindex!(t::TSeries, v, i::AbstractArray{Bool}) = (setindex!(t.values, v, values(i)); t)

# -------------------------------------------------------------
# Some constructors
# -------------------------------------------------------------

# construct undefined from range
TSeries(T::Type{<:Number}, rng::UnitRange{<:MIT}) = TSeries(first(rng), Vector{T}(undef, length(rng)))
TSeries(rng::UnitRange{<:MIT}) = TSeries(Float64, rng)
TSeries(fd::MIT) = TSeries(fd .+ (0:-1))
TSeries(T::Type{<:Number}, fd::MIT) = TSeries(T, fd .+ (0:-1))
TSeries(n::Integer) = TSeries(1U:n*U)
TSeries(T::Type{<:Number}, n::Integer) = TSeries(T, 1U:n*U)
TSeries(rng::UnitRange{<:Integer}) = TSeries(0U .+ rng)
TSeries(T::Type{<:Number}, rng::UnitRange{<:Integer}) = TSeries(T, 0U .+ rng)
TSeries(rng::AbstractRange, ::UndefInitializer) = TSeries(Float64, rng)
TSeries(T::Type{<:Number}, rng::AbstractRange, ::UndefInitializer) = TSeries(T, rng)
TSeries(rng::UnitRange{<:MIT}, ini::Function) = TSeries(first(rng), ini(length(rng)))

Base.similar(::Type{<:AbstractArray}, T::Type{<:Number}, shape::Tuple{UnitRange{<:MIT}}) = TSeries(T, shape[1])
Base.similar(::Type{<:AbstractArray{T}}, shape::Tuple{UnitRange{<:MIT}}) where {T<:Number} = TSeries(T, shape[1])
Base.similar(::AbstractArray, T::Type{<:Number}, shape::Tuple{UnitRange{<:MIT}}) = TSeries(T, shape[1])
Base.similar(::AbstractArray{T}, shape::Tuple{UnitRange{<:MIT}}) where {T<:Number} = TSeries(T, shape[1])

# construct from range and fill with the given constant or array
Base.fill(v::Number, rng::UnitRange{<:MIT}) = TSeries(first(rng), fill(v, length(rng)))
TSeries(rng::UnitRange{<:MIT}, v::Number) = fill(v, rng)
TSeries(rng::UnitRange{<:MIT}, v::AbstractVector{<:Number}) =
    length(rng) == length(v) ? TSeries(first(rng), v) : throw(ArgumentError("Range and data lengths mismatch."))

# -------------------------------------------------------------
# Pretty printing
# -------------------------------------------------------------

function Base.summary(io::IO, t::TSeries)
    et = eltype(t) === Float64 ? "" : ",$(eltype(t))"
    ct = "" # ct = typeof(t.values) === Array{eltype(t),1} ? "" : ",$(typeof(t.values))"
    typestr = "TSeries{$(frequencyof(t))$(et)$(ct)}"
    if isempty(t)
        print(io, "Empty ", typestr, " starting ", t.firstdate)
    else
        print(IOContext(io, :compact => true), length(t.values), "-element ", typestr, " with range ", Base.axes1(t))
    end
end

Base.show(io::IO, ::MIME"text/plain", t::TSeries) = show(io, t)
function Base.show(io::IO, t::TSeries)
    summary(io, t)
    isempty(t) && return
    print(io, ":")
    limit = get(io, :limit, true)
    nval = length(t.values)
    from = t.firstdate
    nrow, ncol = displaysize(io)
    if limit && nval > nrow - 5
        top = div(nrow - 5, 2)
        bot = nval - nrow + 6 + top
        for i = 1:top
            print(io, "\n", lpad(from + (i - 1), 8), " : ", t.values[i])
        end
        print(io, "\n    ⋮")
        for i = bot:nval
            print(io, "\n", lpad(from + (i - 1), 8), " : ", t.values[i])
        end
    else
        for i = 1:nval
            print(io, "\n", lpad(from + (i - 1), 8), " : ", t.values[i])
        end
    end
end

# ------------------------------------------------------------------
# indexing with MIT
# ------------------------------------------------------------------

# this part is tricky!
# - When querying an index that falls outside the allocated range we throw a
#   BoundsError
# - When setting a value at index outside the allocated range we resize the
#   allocation to include the given index (setting new locations to NaN)
#

Base.getindex(t::TSeries, m::MIT) = mixed_freq_error(t, m)
@inline function Base.getindex(t::TSeries{F}, m::MIT{F}) where {F<:Frequency}
    @boundscheck checkbounds(t, m)
    fi = firstindex(t.values)
    getindex(t.values, fi + oftype(fi, m - firstdate(t)))
end

@inline _ind_range_check(x, rng::MIT) = _ind_range_check(x, rng:rng)
function _ind_range_check(x, rng::UnitRange{<:MIT})
    fi = firstindex(x.values, 1)
    fd = firstdate(x)
    stop = oftype(fi, fi + (last(rng)-fd))
    start = oftype(fi, fi + (first(rng)-fd))
    if start < fi || stop > lastindex(x.values, 1)
        Base.throw_boundserror(x, rng)
    end
    return (start, stop)
end

Base.getindex(t::TSeries, rng::AbstractRange{<:MIT}) = mixed_freq_error(t, rng)
function Base.getindex(t::TSeries{F}, rng::StepRange{MIT{F},Duration{F}}) where {F<:Frequency}
    start, stop = _ind_range_check(t, rng)
    step = oftype(stop-start, rng.step)
    return t.values[start:step:stop]
end
function Base.getindex(t::TSeries{F}, rng::UnitRange{MIT{F}}) where {F<:Frequency}
    start, stop = _ind_range_check(t, rng)
    return TSeries(first(rng), getindex(t.values, start:stop))
end

Base.setindex!(t::TSeries, ::Number, m::MIT) = mixed_freq_error(t, m)
function Base.setindex!(t::TSeries{F}, v::Number, m::MIT{F}) where {F<:Frequency}
    # @boundscheck checkbounds(t, m)
    if m ∉ rangeof(t)
        # !! resize!() doesn't work for TSeries out of the box. we implement it below.
        resize!(t, union(m:m, rangeof(t)))
    end
    fi = firstindex(t.values)
    setindex!(t.values, v, fi + oftype(fi, m - firstdate(t)))
end

Base.setindex!(t::TSeries, ::AbstractVector{<:Number}, rng::AbstractRange{<:MIT}) = mixed_freq_error(t, rng)
function Base.setindex!(t::TSeries{F}, vec::AbstractVector{<:Number}, rng::AbstractRange{MIT{F}}) where {F<:Frequency}
    if !issubset(rng, rangeof(t))
        # !! resize!() doesn't work for TSeries out of the box. we implement it below.
        resize!(t, union(rangeof(t), rng))
    end
    if rng isa AbstractUnitRange
        start, stop = _ind_range_check(t, rng)
        setindex!(t.values, vec, start:stop)
    elseif rng isa StepRange
        start, stop = _ind_range_check(t, rng)
        setindex!(t.values, vec, start:oftype(stop-start,rng.step):stop)
    else
        fd = firstdate(t)
        fi = firstindex(t.values, 1)
        inds = [oftype(fi, fi + (ind-fd)) for ind in rng]
        setindex!(t.values, vec, inds)
    end
end

Base.setindex!(t::TSeries{F1}, src::TSeries{F2}, rng::AbstractRange{MIT{F3}}) where {F1<:Frequency,F2<:Frequency,F3<:Frequency} = mixed_freq_error(t, src, rng)
@inline Base.setindex!(t::TSeries{F}, src::TSeries{F}, rng::AbstractRange{MIT{F}}) where {F<:Frequency} = copyto!(t, rng, src)

"""
    typenan(x), typenan(T)

Return a value that indicates Not-A-Number of the same type as the given `x` or
of the given type `T`.

For floating point types, this is the IEEE-defined NaN.
For integer types, we use typemax(). This is not ideal, but it'll do for now.
"""
function typenan end

typenan(x::T) where {T<:Real} = typenan(T)
typenan(T::Type{<:AbstractFloat}) = T(NaN)
typenan(T::Type{<:Integer}) = typemax(T)
typenan(T::Type{<:Union{MIT,Duration}}) = T(typemax(Int64))

istypenan(x) = false
istypenan(x::Integer) = x == typenan(x)
istypenan(x::AbstractFloat) = isnan(x)

# n::Integer - only the length changes. We keep the starting date
function Base.resize!(t::TSeries, n::Integer)
    lt = length(t)  # the old length
    if lt ≠ n
        resize!(t.values, Int64(n))
        # fill new locations with NaN
        t.values[lt+1:end] .= typenan(eltype(t))
    end
    return t
end

# if range is given
Base.resize!(t::TSeries, rng::UnitRange{<:MIT}) = mixed_freq_error(t, eltype(rng))
function Base.resize!(t::TSeries{F}, rng::UnitRange{MIT{F}}) where {F<:Frequency}
    orng = rangeof(t)  # old range
    if first(rng) == first(orng)
        # if the beginning doesn't change we fallback on resize!(t, n)
        return resize!(t, length(rng))
    end
    tvals = copy(t.values) # old values - keep them safe for now
    inds_to_copy = intersect(rng, orng)
    # nrng = min(first(rng), first(orng)):max(last(rng), last(orng))
    _do = convert(Int, first(inds_to_copy) - first(rng)) + 1
    _so = convert(Int, first(inds_to_copy) - first(orng)) + 1
    _n = length(inds_to_copy)
    resize!(t.values, length(rng))
    t.firstdate = first(rng)
    # t[begin:first(inds_to_copy) - 1] .= typenan(eltype(t))
    # t[last(inds_to_copy) + 1:end] .= typenan(eltype(t))
    fill!(t.values, typenan(eltype(t)))
    copyto!(t.values, _do, tvals, _so, _n)
    return t
end

#
Base.copyto!(dest::TSeries, src::TSeries) = mixed_freq_error(dest, src)
@inline Base.copyto!(dest::TSeries{F}, src::TSeries{F}) where {F<:Frequency} = copyto!(dest, rangeof(src), src)

#
Base.copyto!(dest::TSeries, drng::AbstractRange{<:MIT}, src::TSeries) = mixed_freq_error(dest, drng, src)
function Base.copyto!(dest::TSeries{F}, drng::AbstractRange{MIT{F}}, src::TSeries{F}) where {F<:Frequency}
    fullindex = union(rangeof(dest), drng)
    resize!(dest, fullindex)
    copyto!(dest.values, Int(first(drng) - firstindex(dest) + 1), src[drng].values, 1, length(drng))
    return dest
end

# nothing

# view with MIT indexing
Base.view(t::TSeries, I::AbstractRange{<:MIT}) = mixed_freq_error(t, I)
function Base.view(t::TSeries{F}, I::AbstractRange{MIT{F}}) where {F<:Frequency}
    fi = firstindex(t.values)
    TSeries(first(I), view(t.values, oftype(fi, first(I) - firstindex(t) + fi):oftype(fi, last(I) - firstindex(t) + fi)))
end

# view with Int indexing
function Base.view(t::TSeries, I::AbstractRange{<:Integer})
    fi = firstindex(t.values)
    TSeries(firstindex(t) + first(I) - one(first(I)), view(t.values, oftype(fi, first(I)):oftype(fi, last(I))))
end

@inline Base.diff(x::TSeries) = x - lag(x)

# """
#     pct(x::TSeries, shift_value::Int, islog::Bool)

# Calculate percentage growth in `x` given a `shift_value`.

# __Note:__ The implementation is similar to IRIS.

# Examples
# ```julia-repl
# julia> x = TSeries(yy(2000), Vector(1:4));

# julia> pct(x, -1)
# TSeries{Yearly} of length 3
# 2001Y: 100.0
# 2002Y: 50.0
# 2003Y: 33.33333333333333
# ```
# See also: [`apct`](@ref)
# """
function pct(ts::TSeries, shift_value::Int; islog::Bool = false)
    if islog
        a = exp.(ts);
        b = shift(exp.(ts), shift_value);
    else
        a = ts;
        b = shift(ts, shift_value);
    end

    result = @. ( (a - b)/b ) * 100

    TSeries(result.firstdate, result.values)
end
export pct

# """
#     apct(x::TSeries, islog::Bool)

# Calculate annualised percent rate of change in `x`.

# __Note:__ The implementation is similar to IRIS.

# Examples
# ```julia-repl
# julia> x = TSeries(qq(2018, 1), Vector(1:8));

# julia> apct(x)
# TSeries{Quarterly} of length 7
# 2018Q2: 1500.0
# 2018Q3: 406.25
# 2018Q4: 216.04938271604937
# 2019Q1: 144.140625
# 2019Q2: 107.35999999999999
# 2019Q3: 85.26234567901243
# 2019Q4: 70.59558517284461
# ```

# See also: [`pct`](@ref)
# """
# function apct(ts::TSeries, islog::Bool = false)
#     if islog
#         a = exp(ts);
#         b = shift(exp(ts), - 1);
#     else
#         a = ts;
#         b = shift(ts, -1);
#     end

#     values_change = a/b
#     firstdate = values_change.firstdate

#     values = (values_change.^ppy(ts) .- 1) * 100

#     TSeries( (a/b).firstdate, values)
# end

# function Base.cumsum(s::TSeries)
#     TSeries(s.firstdate, cumsum(s.values))
# end

# function Base.cumsum!(s::TSeries)
#     s.values = cumsum(s.values)
#     return s
# end

# """
#     leftcropnan!(x::TSeries)

# Remove `NaN` values from starting at the beginning of `x`, in-place.

# __Note__: an internal function.
# """
# function leftcropnan!(s::TSeries)
#     while isequal(s[firstdate(s)], NaN)
#         popfirst!(s.values)
#         s.firstdate = s.firstdate + 1
#     end
#     return s
# end

# """
# rightcropnan!(x::TSeries)

# Remove `NaN` values from the end of `x`

# __Note__: an internal function.
# """
# function rightcropnan!(s::TSeries)
#     while isequal(s[lastdate(s)], NaN)
#         pop!(s.values)
#     end
#     return s
# end

# """
#     nanrm!(s::TSeries, type::Symbol)

# Remove `NaN` values that are either at the beginning of the `s` and/or end of `x`.

# Examples
# ```
# julia> s = TSeries(yy(2018), [NaN, NaN, 1, 2, NaN]);

# julia> nanrm!(s);

# julia> s
# TSeries{Yearly} of length 2
# 2020Y: 1.0
# 2021Y: 2.0
# ```
# """
# function nanrm!(s::TSeries, type::Symbol=:both)
#     if type == :left
#         leftcropnan!(s)
#     elseif type == :right
#         rightcropnan!(s)
#     elseif type == :both
#         leftcropnan!(s)
#         rightcropnan!(s)
#     else
#         error("Please select between :left, :right, or :both.")
#     end
#     return s
# end


text:
# For classification of an object x we want to know for which multiset the addition
# of x would increase the NCD1 (left-most part of NCDL) the least. So we need a way to calc
# this delta_diameter(A, x). First lets create a type for NCD multi-sets.
type NCDMultiset{T <: Any}
  clen::Function # Function to calc the compressed length of strings
  Xs::Vector{T}

  # We save the strings in one consecutive string together with their start indices:
  Xstr::ASCIIString
  Xstrstarts::Vector{Int64}

  # Intermediate values used in calculating NCD1:
  Gxstrs::Vector{Int64} # Length of compressed versions of objects
  minGxstr::Int64       # Min length of compressed versions of objects

  # Pre-calc'ed values:
  Gx::Float64
  ncd1::Float64

  NCDMultiset{T <: Any}(A::Array{T, 1}, clen = zliblength) = begin
    Astrs = map(string, A)
    Astrlengths = map(length, Astrs)
    Gxstrs = map(clen, Astrs)

    nm = new(clen, A,
      join(Astrs, ""), cumsum(Astrlengths),
      Gxstrs, minimum(Gxstrs), -1.0)
    nm
  end
end

# Cached version of methods
ncd1(nm::NCDMultiset, x) = (nm.ncd1 >= 0.0) ? nm.ncd1 : (nm.ncd1 = calc_ncd1(nm))

# NCD1 = (G(X) - minimum(G(X))) / maxGxminus1
function calc_ncd1(nm::NCDMultiset)
  ncd1value, nm.Gx, gxminu1s = ncd1(nm.Xstr, nm.Xstrstarts, nm.minGxstr, nm.clen)
  nm.maxGxminus1 = maximum(gxminus1s)
  ncd1value
end

function ncd1(Xstr, Xstrstarts, minGxstr, clen = zliblength)
  Gx = clen(Xstr)
  gxminus1s = Gxminus1s(Xstr, Xstrstarts, clen)
  maxGxminus1 = maximum(gxminus1s)
  ncd1value = (Gx - minGxstr) / maxGxminus1
  return (ncd1value, Gx, gxminus1s)
end

# Calculate G(X \ {x}) for each x in X.
function Gxminus1s(Xstr, Xstrstarts, clen = zliblength)
  n = length(Xstrstarts)
  result = zeros(n)
  result[1] = clen(Xstr[Xstrstarts[2]:end])
  result[n] = clen(Xstr[1:(Xstrstarts[n]-1)])
  map(2:(n-1)) do i
    Xminus1str = Xstr[1:(Xstrstarts[i]-1)] * Xstr[Xstrstarts[i+1]:end]
    result[i] = clen(Xminus1str)
  end
  result
end

function delta_diameter{T <: Any}(A::Array{T, 1}, x::T, clen = zliblength)
  ncdmultiset = NCDMultiset(A, clen)
  delta_diameter(ncdmultiset, x)
end

# The delta_diameter is NCD1(A + [x]) - NCD1(A)
function delta_diameter(nm::NCDMultiset, x)
  ncd1_a = ncd1(nm) # To calc intermediate arrays that we can then reuse.
  xstr = string(x)
  # Extend string with xstr and recalc ncd1
  Axstr = nm.Xstr * xstr
  starts = copy(nm.Xstrstarts)
  push!(starts, length(nm.Xstr))
  AxGxminus1 = Gxminus1s(Axstr, starts, nm.clen)
  ncd1_ax = (nm.clen(Axstr) - min(nm.minGXstr, length(nm.clen(xstr)))) / maximum(AxGxminus1)
  ncd1_ax - ncd1_a
end

ext: jl


text:
default_compile_dir() = joinpath(homedir(),
                                 ".julia",
                                 "sysimages")

default_compile_filename() = "sys_itensors.so"

default_compile_path() = joinpath(default_compile_dir(),
                                  default_compile_filename())

function compile_note(; dir = default_compile_dir(),
                        filename = default_compile_filename())
  path = joinpath(dir, filename)
  return """
  You will be able to start Julia with a compiled version of ITensors using:
  ```
  ~ julia --sysimage $path
  ```
  and you should see that the startup times and JIT compilation times are substantially improved when you are using ITensors.

  In unix, you can create an alias with the Bash command:
  ```
  ~ alias julia_itensors="julia --sysimage $path -e 'using ITensors' -i"
  ```
  which you can put in your `~/.bashrc`, `~/.zshrc`, etc. This also executes `using ITensors` so that ITensors is loaded and ready to use, you can leave off ` -e 'using ITensors' -i` if you don't want that. Then you can start Julia with a version of ITensors installed with the command:
  ```
  ~ julia_itensors
  ```

  Note that if you update ITensors to a new version, for example with `using Pkg; Pkg.update("ITensors")`, you will need to run the `ITensors.compile()` command again to recompile the new version of ITensors.
  """
end

function compile(; dir::AbstractString = default_compile_dir(),
                 filename::AbstractString = default_compile_filename())
  if !isdir(dir)
    println("""The directory "$dir" doesn't exist yet, creating it now.""")
    println()
    mkdir(dir)
  end
  path = joinpath(dir, filename)
  println("""Creating the system image "$path" containing the compiled version of ITensors. This may take a few minutes.""")
  create_sysimage(:ITensors;
                  sysimage_path = path,
                  precompile_execution_file = joinpath(@__DIR__,
                                                       "precompile_itensors.jl"))
  println(compile_note(; dir = dir, filename = filename))
  return path
end

@doc """
    ITensors.compile(; dir = "$(default_compile_dir())",
                       filename = "$(default_compile_filename())")

Compile ITensors.jl with [PackageCompiler](https://julialang.github.io/PackageCompiler.jl/dev/). This will take some time, perhaps a few minutes.

This will create a system image containing the compiled version of ITensors located at `dir/filename`, by default `$(default_compile_path())`.

$(compile_note())
""" compile


text:
mylevel(v::BilevelVariableRef) = v.level
in_level(v::BilevelVariableRef, level::Level) = (
    v.level === LOWER_BOTH ||
    v.level === UPPER_BOTH ||
    v.level === level ||
    (v.level === DUAL_OF_LOWER && level === UPPER_ONLY))

in_level(v::BilevelVariableRef, ::UpperModel) = in_upper(v)
in_level(v::BilevelVariableRef, ::LowerModel) = in_lower(v)
in_upper(v::BilevelVariableRef) = in_upper(mylevel(v))
in_lower(v::BilevelVariableRef) = in_lower(mylevel(v))
upper_ref(v::BilevelVariableRef) = v.model.var_upper[v.idx]
lower_ref(v::BilevelVariableRef) = v.model.var_lower[v.idx]

const BilevelAffExpr = GenericAffExpr{Float64, BilevelVariableRef}
const BilevelQuadExpr = GenericQuadExpr{Float64, BilevelVariableRef}

function jump_var_ref(v::BilevelVariableRef)
    level = mylevel(v)
    if level == LOWER_ONLY || level == LOWER_BOTH
        return lower_ref(v)
    else
        return upper_ref(v)
    end
end

function solver_ref(v::BilevelVariableRef)
    m = v.model
    if mylevel(v) == LOWER_ONLY
        return m.sblm_to_solver[
            m.lower_to_sblm[JuMP.index(lower_ref(v))]]
    else
        return m.sblm_to_solver[
            m.upper_to_sblm[JuMP.index(upper_ref(v))]]
    end
end

Base.broadcastable(v::BilevelVariableRef) = Ref(v)
Base.copy(v::BilevelVariableRef) = v
Base.:(==)(v::BilevelVariableRef, w::BilevelVariableRef) =
    v.model === w.model && v.idx == w.idx && v.level == w.level
JuMP.owner_model(v::BilevelVariableRef) = v.model
JuMP.isequal_canonical(v::BilevelVariableRef, w::BilevelVariableRef) = v == w
JuMP.variable_type(::AbstractBilevelModel) = BilevelVariableRef
# add in both levels
function JuMP.add_variable(inner::InnerBilevelModel, v::JuMP.AbstractVariable, name::String="")
    m = bilevel_model(inner)
    m.last_variable_index += 1
    vref = BilevelVariableRef(m, m.last_variable_index, level_both(inner))
    # break info so that bounds go to correct level
    var_upper, var_lower = split_variable(inner, v)
    v_upper = JuMP.add_variable(m.upper, var_upper, name)
    m.var_upper[vref.idx] = v_upper
    v_lower = JuMP.add_variable(m.lower, var_lower, name)
    m.var_lower[vref.idx] = v_lower
    set_link!(inner, v_upper, v_lower)
    m.var_info[vref.idx] = empty_info(level_both(inner))
    JuMP.set_name(vref, name)
    m.var_upper_rev = nothing
    m.var_lower_rev = nothing
    vref
end
function JuMP.add_variable(single::SingleBilevelModel, v::JuMP.AbstractVariable, name::String="")
    m = bilevel_model(single)
    m.last_variable_index += 1
    vref = BilevelVariableRef(m, m.last_variable_index, level(single))
    v_level = JuMP.add_variable(mylevel_model(single), v, name)
    mylevel_var_list(single)[vref.idx] = v_level
    push_single_level_variable!(single, v_level)
    m.var_info[vref.idx] = empty_info(level(single))
    JuMP.set_name(vref, name)
    m.var_upper_rev = nothing
    m.var_lower_rev = nothing
    vref
end
function JuMP.delete(::BilevelModel, vref::BilevelVariableRef)
    model = vref.model
    idx = vref.idx
    delete!(model.var_info, idx)
    if haskey(model.var_upper, idx)
        v_up = model.var_upper[idx]
        delete!(model.var_upper, idx)
        delete!(model.upper_only, v_up)
        delete!(model.upper_to_lower_link, v_up)
        delete!(model.upper_var_to_lower_ctr_link, v_up)
        delete!(model.link, v_up)
        JuMP.delete(model.upper, v_up)
    end
    if haskey(model.var_lower, idx)
        v_lo = model.var_lower[idx]
        delete!(model.var_lower, idx)
        delete!(model.lower_only, v_lo)
        delete!(model.lower_to_upper_link, v_lo)
        JuMP.delete(model.lower, v_lo)
    end
    model.var_upper_rev = nothing
    model.var_lower_rev = nothing
    return nothing
end
JuMP.is_valid(m::BilevelModel, vref::BilevelVariableRef) = vref.idx in keys(m.var_info)
JuMP.is_valid(m::InnerBilevelModel, vref::BilevelVariableRef) =
    JuMP.is_valid(bilevel_model(m), vref) && in_level(vref, level(m))
JuMP.num_variables(m::BilevelModel) = length(m.var_info)
JuMP.num_variables(m::UpperModel) = length(m.m.var_upper)
JuMP.num_variables(m::LowerModel) = length(m.m.var_lower)
function empty_info(level::Level)
    return BilevelVariableInfo(level)
end

function JuMP.all_variables(m::InnerBilevelModel)
    list = keys(mylevel_var_list(m))
    model = bilevel_model(m)
    ret = BilevelVariableRef[]
    for idx in list
        var = BilevelVariableRef(model, idx)
        if in_level(var, m)
            push!(ret, var)
        end
    end
    return ret
end
function JuMP.all_variables(model::BilevelModel)
    vec = JuMP.all_variables(Upper(model))
    append!(vec, JuMP.all_variables(Lower(model)))
    return unique(vec)
end

"""
Split variable because actual owner of the variable should be the one holding the bounds.
"""
function split_variable(::UpperModel, v::JuMP.AbstractVariable)
    var_upper = v
    var_lower = JuMP.ScalarVariable(JuMP.VariableInfo(
        false, NaN,
        false, NaN,
        false, NaN,
        v.info.has_start, v.info.start,
        false, false))
    return var_upper, var_lower
end
function split_variable(::LowerModel, v::JuMP.AbstractVariable)
    var_lower = v
    var_upper = JuMP.ScalarVariable(JuMP.VariableInfo(
        false, NaN,
        false, NaN,
        false, NaN,
        v.info.has_start, v.info.start,
        false, false))
    return var_upper, var_lower
end

function JuMP.has_lower_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return !isnan(get_dual_lower_bound(get_constrain_ref(vref)))
    end
    JuMP.has_lower_bound(jump_var_ref(vref))
end
function JuMP.lower_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return get_dual_lower_bound(get_constrain_ref(vref))
    end
    # @assert !JuMP.is_fixed(vref)
    JuMP.lower_bound(jump_var_ref(vref))
end
function JuMP.set_lower_bound(vref::BilevelVariableRef, lower::Number)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_lower_bound(get_constrain_ref(vref), lower)
        return
    end
    JuMP.set_lower_bound(jump_var_ref(vref), lower)
    return
end
function JuMP.delete_lower_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_lower_bound(get_constrain_ref(vref), NaN)
        return
    end
    JuMP.delete_lower_bound(jump_var_ref(vref))
    return
end
function JuMP.has_upper_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return !isnan(get_dual_upper_bound(get_constrain_ref(vref)))
    end
    JuMP.has_upper_bound(jump_var_ref(vref))
end
function JuMP.upper_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return get_dual_upper_bound(get_constrain_ref(vref))
    end
    # @assert !JuMP.is_fixed(vref)
    JuMP.upper_bound(jump_var_ref(vref))
end
function JuMP.set_upper_bound(vref::BilevelVariableRef, upper)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_upper_bound(get_constrain_ref(vref), upper)
        return
    end
    JuMP.set_upper_bound(jump_var_ref(vref), upper)
    return
end
function JuMP.delete_upper_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_upper_bound(get_constrain_ref(vref), NaN)
        return
    end
    JuMP.delete_upper_bound(jump_var_ref(vref))
end

JuMP.is_fixed(vref::BilevelVariableRef) = JuMP.is_fixed(jump_var_ref(vref))
JuMP.fix_value(vref::BilevelVariableRef) = JuMP.fix_value(jump_var_ref(vref))
function JuMP.fix(vref::BilevelVariableRef, value; force::Bool=false)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be fixed.")
    end
    JuMP.fix(jump_var_ref(vref), value; force=force)
end
function JuMP.unfix(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be fixed.")
    end
    JuMP.unfix(jump_var_ref(vref))
end

function JuMP.start_value(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return JuMP.dual_start_value(get_constrain_ref(vref))
    end
    JuMP.start_value(jump_var_ref(vref))
end
function JuMP.set_start_value(vref::BilevelVariableRef, start)
    if mylevel(vref) == DUAL_OF_LOWER
        JuMP.set_dual_start_value(get_constrain_ref(vref), start)
        return
    end
    in_upper(vref) && JuMP.set_start_value(upper_ref(vref), start)
    in_lower(vref) && JuMP.set_start_value(lower_ref(vref), start)
    return
end

JuMP.is_binary(vref::BilevelVariableRef) = JuMP.is_binary(jump_var_ref(vref))
function JuMP.set_binary(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be binary.")
    end
    JuMP.set_binary(jump_var_ref(vref))
end
function JuMP.unset_binary(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be binary.")
    end
    JuMP.unset_binary(jump_var_ref(vref))
end
JuMP.is_integer(vref::BilevelVariableRef) = JuMP.is_integer(jump_var_ref(vref))
function JuMP.set_integer(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be integer.")
    end
    JuMP.set_integer(jump_var_ref(vref))
end
function JuMP.unset_integer(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be integer.")
    end
    JuMP.unset_integer(jump_var_ref(vref))
end

function JuMP.value(v::BilevelVariableRef; result::Int = 1)::Float64
    m = owner_model(v)
    solver = m.solver
    ref = solver_ref(v)
    return MOI.get(solver, MOI.VariablePrimal(result), ref)
end


text:
#=
    this splitting criterion considers the  ratio = sum(numerator)/sum(denominator) (for each split)
    we minimize the (ratio.-meanLeft)^2 of both children
=#

function calculateSplitValue(a::sseSplit, fname::Symbol, number_of_char_features::Int, labellist::Vector{T}, sumnumerator::Array{Float64,1}, sumdenominator::Array{Float64,1}, sumweight::Array{Float64,1}, countlistfloat::Array{Float64,1}, minweight::Float64, subs::DTSubsets, moments_per_pdaclass) where T <: Unsigned
    # here randomweight==0
    # for subsets, exhaustive search with flipping members (gray code) or "increasing" subset search ({1}, {1,2}, {1,2,3}, .... {1,2,3, ....., n-1,2})
    # all input lists (labellist,sumnumerator,sumdenominator,sumweight,countlistfloat) need to be sorted in the same manner
    # convention: in the beginning everything is on the right side
    elementsInLeftChildBV = BitVector(undef, length(labellist));fill!(elementsInLeftChildBV, false) # indicates which classes belong to right child
    chosen_subset_bitarray = BitVector(undef, 0)

    totalMeanAndVariance = copy(emptyCustomVariance)
    for idx = 1:length(moments_per_pdaclass)
        @inbounds merge!(totalMeanAndVariance, moments_per_pdaclass[idx])
    end

    weighttot = sum(sumweight)
    weighttot_minw = weighttot - minweight

    momentsl = copy(emptyCustomVariance)
    # momentsr needs to be the 'total' at the start of this.
    momentsr = totalMeanAndVariance

    sumwl = 0.0
    varl = varr = val = -Inf
    chosen_sumwl = NaN
    for i in subs
      # @show "vv=$(i)" #should be the element which switches
      # i is an index, it indicates which element of labellist will flip sides for the next calculation
        @inbounds switching_class = moments_per_pdaclass[i]
        @inbounds if elementsInLeftChildBV[i]
        # update elementsInLeftChildBV, i.e. toggle the value of the ith component, $=XOR
            @inbounds  elementsInLeftChildBV[i] = xor(elementsInLeftChildBV[i], true) # updating needs to occur here before we copy the array (in case it is a better split than what we have seen so far)
        # move class from left to right side
            unmerge!(momentsl, switching_class) # switching_class is 'removed' from the left variance/mean statistcs momentsl
            merge!(momentsr, switching_class) # switching_class is 'added' to the right variance/mean statistcs momentsl
        # sumnl-=sumnumerator[i]
        # sumdl-=sumdenominator[i]
            @inbounds  sumwl -= sumweight[i]
        else
          # update elementsInLeftChildBV, i.e. toggle the value of the ith component, $=XOR
            @inbounds  elementsInLeftChildBV[i] = xor(elementsInLeftChildBV[i], true)
          # move class from right to left side
            merge!(momentsl, switching_class)
            unmerge!(momentsr, switching_class)
          # sumnl+=sumnumerator[i]
            # sumdl+=sumdenominator[i]
            @inbounds  sumwl += sumweight[i]
        end
        if (sumwl > minweight) && (weighttot_minw > sumwl) # do we have enough exposure? is the split valid?
          # vold=valnew
            varl = momentsl.m2
            varr = momentsr.m2
          # @show fname,sumwl,weighttot-sumwl
            valnew = -(varl + varr)
          # if vold==valnew;@show i,valnew,sumnl,sumdl;end;
            if valnew > val
                val = valnew
                chosen_subset_bitarray = copy(elementsInLeftChildBV)
                chosen_sumwl = sumwl
            end
        end

    end
    if isfinite(val)
      # warning: if the labellist is not sorted 1:n we need to be careful here!
        chosen_subset = labellist[chosen_subset_bitarray]
    else
        chosen_subset = Array{UInt}(undef, 0)
    end
    return val, chosen_subset, chosen_sumwl, weighttot - chosen_sumwl
end

function calculateSplitValue(a::mseSplit, fname::Symbol, number_of_char_features::Int, labellist::Vector{T}, sumnumerator::Array{Float64,1}, sumdenominator::Array{Float64,1}, sumweight::Array{Float64,1}, countlistfloat::Array{Float64,1}, minweight::Float64, subs::DTSubsets, moments_per_pdaclass) where T <: Unsigned
# here randomweight==0
# for subsets, exhaustive search with flipping members (gray code) or "increasing" subset search ({1}, {1,2}, {1,2,3}, .... {1,2,3, ....., n-1,2})
# all input lists (labellist,sumnumerator,sumdenominator,sumweight,countlistfloat) need to be sorted in the same manner
# convention: in the beginning everything is on the right side
    elementsInLeftChildBV = BitVector(undef, length(labellist));fill!(elementsInLeftChildBV, false) # indicates which classes belong to right child
    chosen_subset_bitarray = BitVector(undef, 0)

    totalMeanAndVariance = copy(emptyCustomVariance)
    for idx = 1:length(moments_per_pdaclass)
        @inbounds merge!(totalMeanAndVariance, moments_per_pdaclass[idx])
    end

    weighttot = sum(sumweight)
    weighttot_minw = weighttot - minweight

    momentsl = copy(emptyCustomVariance)
# momentsr needs to be the 'total' at the start of this.
    momentsr = totalMeanAndVariance

    sumwl = 0.0
    varl = varr = val = -Inf
    chosen_sumwl = NaN
    for i in subs
  # @show "vv=$(i)" #should be the element which switches
  # i is an index, it indicates which element of labellist will flip sides for the next calculation
        @inbounds switching_class = moments_per_pdaclass[i]
        @inbounds if elementsInLeftChildBV[i]
    # update elementsInLeftChildBV, i.e. toggle the value of the ith component, $=XOR
            @inbounds  elementsInLeftChildBV[i] = xor(elementsInLeftChildBV[i], true) # updating needs to occur here before we copy the array (in case it is a better split than what we have seen so far)
    # move class from left to right side
            unmerge!(momentsl, switching_class) # switching_class is 'removed' from the left variance/mean statistcs momentsl
            merge!(momentsr, switching_class) # switching_class is 'added' to the right variance/mean statistcs momentsl
    # sumnl-=sumnumerator[i]
    # sumdl-=sumdenominator[i]
            @inbounds  sumwl -= sumweight[i]
        else
      # update elementsInLeftChildBV, i.e. toggle the value of the ith component, $=XOR
            @inbounds  elementsInLeftChildBV[i] = xor(elementsInLeftChildBV[i], true)
      # move class from right to left side
            merge!(momentsl, switching_class)
            unmerge!(momentsr, switching_class)
      # sumnl+=sumnumerator[i]
    	# sumdl+=sumdenominator[i]
    	@inbounds  sumwl += sumweight[i]
        end
        if (sumwl > minweight) && (weighttot_minw > sumwl) # do we have enough exposure? is the split valid?
	  # vold=valnew
            varl = momentsl.var
            varr = momentsr.var
	  # @show fname,sumwl,weighttot-sumwl
            valnew = -(varl + varr)
      # if vold==valnew;@show i,valnew,sumnl,sumdl;end;
	  if valnew > val
                val = valnew
                chosen_subset_bitarray = copy(elementsInLeftChildBV)
                chosen_sumwl = sumwl
            end
        end

    end
    if isfinite(val)
  # warning: if the labellist is not sorted 1:n we need to be careful here!
        chosen_subset = labellist[chosen_subset_bitarray]
    else
        chosen_subset = Array{UInt}(undef, 0)
    end
    return val, chosen_subset, chosen_sumwl, weighttot - chosen_sumwl
end

function calculateSplitValue(a::mseSplit, number_of_char_features::Int, labellist::Array{UInt8,1}, sumnumerator::Array{Float64,1}, sumdenominator::Array{Float64,1}, sumweight::Array{Float64,1}, countlistfloat::Array{Float64,1}, minweight::Float64, subs::DTSubsets, feature_column_id::Int, moments_per_pdaclass)
    error("221x this will not work yet.....")
# here randomweight>0
# for subsets, exhaustive search with flipping members (gray code) or "increasing" subset search ({1}, {1,2}, {1,2,3}, .... {1,2,3, ....., n-1,2})
# all input lists (labellist,sumnumerator,sumdenominator,sumweight,countlistfloat) need to be sorted in the same manner
# convention, in the beginning everything is on the right side
end

function aggregate_data_mse(f, numerator::Array{Float64,1}, denominator::Array{Float64,1}, weight::Array{Float64,1})
    # this is the core function of the modelling process
    # besides copying of the data, the vast majority of time is spent in here!
    # most of the time is spent here, if we can improve the for loop below, that would improve performance greatly!
    # one possibility would be to introduce parallelization here (which is not straightforward, I think....)
    lo = one(eltype(f.parent.refs))
    hi = convert(eltype(f.parent.refs), length(f.parent.pool))
    ooo = one(lo) - lo
    vecsize = hi + ooo
    cnt = zeros(Int, vecsize)
    sumnumerator = zeros(Float64, vecsize)
    sumdenominator = zeros(Float64, vecsize)
    sumweight = zeros(Float64, vecsize)
    moments_per_pdaclass = Array{CustomVariance}(undef, vecsize)
    for i = 1:vecsize
        @inbounds moments_per_pdaclass[i] = copy(emptyCustomVariance)
    end
    @inbounds for count in f.indices[1]
        @inbounds idx = f.parent.refs[count] + ooo
        @inbounds cnt[idx] += 1
        @inbounds this_num = numerator[count]
        @inbounds sumnumerator[idx] += this_num
        @inbounds this_denom = denominator[count]
        @inbounds sumdenominator[idx] += this_denom
        @inbounds sumweight[idx] += weight[count]
        @inbounds addSinglePoint!(moments_per_pdaclass[idx], this_num, this_denom)
    end
    return cnt, sumnumerator, sumdenominator, sumweight, moments_per_pdaclass
end

mutable struct CustomVariance
    # considers the variance of d_i*(n_i/d_i-fittedMean)^2 where n_i/d_i is calculated pointwise and fitted mean is (sum(n_i)/sum(d_i))
    meanratio::Float64
    sn::Float64 # sum nominator
    sd::Float64 # sum denominator
    n::Int # count
    m2::Float64 # var*(n-1)
    var::Float64 # weigthedVariance
    function CustomVariance(num, denom)
        n = length(num)
        @assert n == length(denom) # length(w)
        sn = sum(num)
        sd = sum(denom)
        meanratio = sn / sd
        m2 = zero(eltype(num))
        for i = 1:n
            @inbounds m2 += denom[i] * (num[i] / denom[i] - meanratio)^2
        end
        if n > 1
            var = m2 / (n - 1)
        else
            var = zero(eltype(num))
        end
        return new(meanratio, sn, sd, n, m2, var)
    end
    function CustomVariance()
        return new(0.0, 0.0, 0.0, 0, 0.0, 0.0)
    end
end

function removeSinglePoint!(a::CustomVariance, num, denom)
    sn = a.sn - num
    sd = a.sd - denom
    meanratio = sn / sd
    # merge: m2 = a.m2 + b.m2 + 2*(a.sn*(a.meanratio-meanratio)+§*(num/denom-meanratio)) + a.sd*(meanratio^2-a.meanratio^2) + denom*(meanratio^2-num/denom^2)
    # unmerge:
    a.m2 = a.m2 + 2 * (num * (a.meanratio - num / denom) + sn * (a.meanratio - meanratio)) + denom * ((num / denom)^2 - a.meanratio^2) + sd * (meanratio^2 - a.meanratio^2)
    a.sn = sn
    a.sd = sd
    a.n -= 1
    a.meanratio = meanratio
    a.var = a.m2 / (a.n - 1)
    return nothing
end

function addSinglePoint!(a::CustomVariance, num, denom)
    # considers the variance of d_i*(n_i/d_i-fittedMean)^2 where n_i/d_i is calculated pointwise and fitted mean is (sum(n_i)/sum(d_i))
    sn = a.sn + num
    sd = a.sd + denom
    a.n = a.n + 1
    meanratio = sn / sd
    a.m2 = a.m2 + 2 * (a.sn * (a.meanratio - meanratio) + num * (num / denom - meanratio)) + a.sd * (meanratio^2 - a.meanratio^2) + denom * (meanratio^2 - (num / denom)^2)
    a.sn = sn
    a.sd = sd
    a.meanratio = meanratio
    a.var = a.m2 / (a.n - 1)
    return nothing
end

function Base.merge!(a::CustomVariance, b::CustomVariance)
    # considers the variance of d_i*(n_i/d_i-fittedMean)^2 where n_i/d_i is calculated pointwise and fitted mean is (sum(n_i)/sum(d_i))
    # merges b into a
    # i.e. the points form b are 'added' to a
    sn = a.sn + b.sn
    sd = a.sd + b.sd
    n = a.n + b.n
    meanratio = sn / sd
    m2 = a.m2 + b.m2 + 2 * (a.sn * (a.meanratio - meanratio) + b.sn * (b.meanratio - meanratio)) + a.sd * (meanratio^2 - a.meanratio^2) + b.sd * (meanratio^2 - b.meanratio^2)
    var = m2 / (n - 1)
    a.sn = sn
    a.sd = sd
    a.n = n
    a.meanratio = meanratio
    a.m2 = m2
    a.var = var
    return nothing
end

function Base.merge(a::CustomVariance, b::CustomVariance)
    # considers the variance of d_i*(n_i/d_i-fittedMean)^2 where n_i/d_i is calculated pointwise and fitted mean is (sum(n_i)/sum(d_i))
    c = copy(a)
    merge!(c, b)
    return c
end

"""
function unmerge!(a::CustomVariance,b::CustomVariance)
this function updates the mean and variance of a
The assumption is that the underlying data of b is a subset of a.
If this is not the case, the result of this function is most likely meaningless
That is to say "subtract" the series b from the series a and update the mean and variance of a accordingly
"""
function unmerge!(a::CustomVariance, b::CustomVariance)
    # considers the variance of d_i*(n_i/d_i-fittedMean)^2 where n_i/d_i is calculated pointwise and fitted mean is (sum(n_i)/sum(d_i))
    sn = a.sn - b.sn
    sd = a.sd - b.sd
    n = a.n - b.n
    meanratio = sn / sd
    # merge: m2 = a.m2 + b.m2 + 2*(a.sn*(a.meanratio-meanratio)+b.sn*(b.meanratio-meanratio)) + a.sd*(meanratio^2-a.meanratio^2) + b.sd*(meanratio^2-b.meanratio^2)
    # unmerge:
    m2 = a.m2 - b.m2 + 2 * (b.sn * (a.meanratio - b.meanratio) + sn * (a.meanratio - meanratio)) + b.sd * (b.meanratio^2 - a.meanratio^2) + sd * (meanratio^2 - a.meanratio^2)
    var = m2 / (n - 1)
    a.sn = sn
    a.sd = sd
    a.n = n
    a.meanratio = meanratio
    a.m2 = m2
    a.var = var
    return nothing
end

function Base.copy(a::CustomVariance)
    b = CustomVariance()
    b.sn = a.sn
    b.sd = a.sd
    b.n = a.n
    b.meanratio = a.meanratio
    b.m2 = a.m2
    b.var = a.var
    return b
end


text:
# Exponential
ϕ = @SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(ExponentialMap(x)),ϕ) ≈ jacobian(ExponentialMap,ϕ)
@test retraction_map(ExponentialMap(ϕ)) == ExponentialMap

ϕ = 1e-6*@SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(ExponentialMap(x)),ϕ) ≈ jacobian(ExponentialMap,ϕ)
@test retraction_map(ExponentialMap(ϕ)) == ExponentialMap

# Vector Part
v = 0.1*@SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(VectorPart(x)),v) ≈
    jacobian(VectorPart, v)
@test retraction_map(VectorPart(v)) == VectorPart

# Gibbs Vectors
g = @SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(CayleyMap(x)),g) ≈ jacobian(CayleyMap, g)
@test retraction_map(CayleyMap(g)) == CayleyMap

# MRPs
p = SVector(rand(MRP))
@test ForwardDiff.jacobian(x->SVector(MRPMap(x)),p) ≈
    jacobian(MRPMap, p)
@test retraction_map(MRPMap(p)) == MRPMap

μ0 = 1/TO.scaling(VectorPart)
jac_eye = [@SMatrix zeros(1,3); μ0*Diagonal(@SVector ones(3))];
@test jacobian(ExponentialMap, p*1e-10) ≈ jac_eye
@test jacobian(MRPMap, p*1e-10) ≈ jac_eye
@test jacobian(CayleyMap, p*1e-10) ≈ jac_eye
@test jacobian(VectorPart, p*1e-10) ≈ jac_eye

############################################################################################
#                                 INVERSE RETRACTION MAPS
############################################################################################

# Exponential Map
Random.seed!(1);
q = rand(UnitQuaternion)
q = UnitQuaternion{ExponentialMap}(q)
qval = SVector(q)
@test ExponentialMap(q) == TO.scaling(ExponentialMap)*logm(q)
@test ExponentialMap(ExponentialMap(q)) ≈ q
@test ExponentialMap(ExponentialMap(ϕ)) ≈ ϕ

function invmap(q)
    μ = TO.scaling(ExponentialMap)
    v = @SVector [q[2], q[3], q[4]]
    s = q[1]
    θ = norm(v)
    M = μ*2atan(θ, s)/θ
    return M*v
end
@test invmap(qval) ≈ TO.scaling(ExponentialMap)*logm(q)

qI = VectorPart(v*1e-5)
@test ForwardDiff.jacobian(invmap, qval) ≈ jacobian(ExponentialMap, q)
@test ForwardDiff.jacobian(invmap, SVector(qI)) ≈ jacobian(ExponentialMap, qI)

# Vector Part
@test VectorPart(q) == TO.scaling(VectorPart)*qval[2:4]
@test VectorPart(VectorPart(q)) ≈ q
@test VectorPart(VectorPart(v)) ≈ v

# Cayley
invmap(q) = 1/q[1] * @SVector [q[2], q[3], q[4]]
@test CayleyMap(q) ≈ TO.scaling(CayleyMap)*invmap(qval)
@test ForwardDiff.jacobian(invmap, qval) ≈ jacobian(CayleyMap, q)
@test CayleyMap(CayleyMap(q)) ≈ q
@test CayleyMap(CayleyMap(g)) ≈ g

# MRP
invmap(q) = TO.scaling(MRPMap)/(1+q[1]) * @SVector [q[2], q[3], q[4]]
@test MRPMap(q) ≈ invmap(qval)
@test ForwardDiff.jacobian(invmap, qval) ≈ jacobian(MRPMap, q)
@test MRPMap(MRPMap(q)) ≈ q
@test MRPMap(MRPMap(p)) ≈ p

# Test near origin
μ0 = TO.scaling(VectorPart)
jacT_eye = [@SMatrix zeros(1,3); μ0*Diagonal(@SVector ones(3))]';
@test isapprox(jacobian(ExponentialMap,qI), jacT_eye, atol=1e-5)
@test isapprox(jacobian(VectorPart,qI), jacT_eye, atol=1e-5)
@test isapprox(jacobian(CayleyMap,qI), jacT_eye, atol=1e-5)
@test isapprox(jacobian(MRPMap,qI), jacT_eye, atol=1e-5)


text:
function foo()
    lines = readlines(open("data/6.txt", "r"))
    points = zeros(Int, length(lines), 2)
    maxx = 0
    for (i, line) in enumerate(lines)
        point = parse.(Int, split(line, ','))
        if any(point .> maxx)
            maxx = maximum(point)
        end
        points[i, :] = point
    end
    areas = Dict{Int, Int}()

    bad_idx = Set{Int}()
    stop = maxx
    start = 0
    for i in start:stop
        for j in start:stop
            dist = abs.(points[:, 1] .- i) .+ abs.(points[:, 2] .- j)
            min_dist = minimum(dist)
            if sum(dist .== min_dist) == 1
                idx = argmin(dist)
                if !(idx in keys(areas))
                    areas[idx] = 0
                end
                areas[idx] += 1
                if i == start || i == stop || j == start || j == stop
                    push!(bad_idx, idx)
                end
            end
        end
    end
    for idx in bad_idx
        delete!(areas, idx)
    end
    return maximum(values(areas))
end

println(foo())


text:
## ASCII display

function irange(u, W, H)
    x, y = u.x.val, u.y.val
    rx = (1 + trunc(Integer, x.lo)):(1 + trunc(Integer, x.hi) - 1)
    ry = (1 + H - trunc(Integer, y.hi)):(1 + H - trunc(Integer, y.lo) - 1)
    rx, ry
end

"""

Function to graph, using ascii characters, a predicate related to a
two-dimensional equation, such as `f == 0` or `f < g`, where `f` and
`g` are functions, e.g., `f(x,y) = y - sqrt(x)`. The positional
arguments are `L`, `R`, `B`, `T` which indicate the x-y range of the
graph of the equation. The keywork arguments `W` and `H` indicate the number
of pixels to use.

Set `offset=0` to see algorithm

A pixel is important, as the graph will color a pixel

- "white" (e.g., " ") if no solution exists in the region indicated by the pixel
- "black" (e.g., "x") if a solution is known to exist in the region
- "red" (e.g. ".") if a solution _might_ exist

"""
function asciigraph(r, L=-5, R=5, B=-5, T=5; W=2^4, H=2^4)

    cols = Dict()
    cols[:red]="."; cols[:white]=" ";cols[:black]="x"

    red, black, white = GRAPH(r, L, R, B, T, W, H)

    graph = repmat([:red], H, W)
    for u in white
        r1, r2 = irange(u, W, H)
        graph[r2, r1] = :white
    end
    for u in black
        r1, r2 = irange(u, W, H)
        graph[r2, r1] = :black
    end

    mat = graph
    m,n = size(mat)
    for i in 1:m
        for j in 1:n
            print(cols[mat[i,j]])
        end
        print("\n")
    end
end


text:
import RecipesBase

"""
    plot(s::SeisSplit.Result; max_samples=1000, antialias=true) -> ::Plots.Plot

Plot the results of shear wave splitting analysis from the `SeisSplit.splitting` function.

# Keyword arguments
- `antialias = true`: Control whether decimation of traces for plotting purposes
  uses antialising to avoid spurious signals.  `true` will avoid spurious aliasing
  signals, whilst `false` may be quicker.
- `max_samples = 1000`: Maximum number of samples to plot in the trace and particle
  motion subplots.  Windows or traces with more than this number are decimated to
  below this number to speed up plotting.
"""
plot

RecipesBase.@recipe function f(s::Result; max_samples=1000, antialias=true)
    grid --> false
    framestyle --> :box
    colorbar := false
    layout := 8
    title := hcat("Input", "Source pol. pre-corr", "Source pol. post-corr",
                  "Fast-slow pre-corr", "Fast-slow corr", "Small eigenvalue",
                  "PM pre-corr", "PM post-corr", "")
    titlefontsize --> 11
    legendfontsize --> 8

    # Input traces in N-E orientation
    n_orig, e_orig = Seis.rotate_through(s.trace1, s.trace2, -s.trace1.sta.azi)

    # Tuple of window limits
    window = s.window_start, s.window_end

    ## Input traces in original orientation with analysis window highlighted
    # Analysis window in colour behind traces
    # RecipesBase.@series begin
    #     subplot := 1
    #     w1, w2 =s.window_start, s.window_end
    #     []
    # end
    # Traces
    RecipesBase.@series begin
        subplot := 1
        label := [s.trace1.sta.cha s.trace2.sta.cha]
        trace1, trace2 = _decimate_to_max_samples.((s.trace1, s.trace2), max_samples,
            antialias)
        Seis.times(trace1), [Seis.trace(trace1) Seis.trace(trace2)]
    end
    # Window limits
    RecipesBase.@series begin
        subplot := 1
        seriestype := :vline
        label := ""
        linecolor := :red
        linewidth := 1
        [window...]
    end

    ## Rotated to spol
    RecipesBase.@series begin
        subplot := 2
        spol, spol_90 = _decimate_to_max_samples.(
            Seis.rotate_through(n_orig, e_orig, 180-s.spol),
            max_samples, antialias, Ref(window))
        label := ""
        linecolor := [:blue :red]
        Seis.times(spol), [Seis.trace(spol), Seis.trace(spol_90)]
    end

    ## Corrected and rotated to spol
    RecipesBase.@series begin
        subplot := 3
        label := ""
        linecolor := [:blue :red]
        n, e = _decimate_to_max_samples.((n_orig, e_orig), max_samples, antialias, Ref(window))
        SeisSplit.apply_split!(n, e, s.phi_best, -s.dt_best)
        fast, slow = Seis.rotate_through(n, e, 180-s.spol)
        Seis.times(fast), [Seis.trace(fast), Seis.trace(slow)]
    end

    ## Fast and slow before dt shift
    sfast, sslow = Seis.rotate_through!.(_decimate_to_max_samples.((n_orig, e_orig), max_samples,
        antialias, Ref(window))..., s.phi_best)
    # Flip slow wave if opposite polarity to make it easier to see if traces line up
    slow_polarity = let p = _apparent_polarity(sfast, sslow)
        # Account for zero polarity if traces not correlated at all
        p == 0 ? 1 : p
    end
    RecipesBase.@series begin
        subplot := 4
        label := ""
        legend := false
        linecolor := [:blue :red]
        xlims := window
        Seis.times(sfast), [Seis.trace(sfast), slow_polarity.*Seis.trace(sslow)]
    end

    ## Fast and slow after dt shift
    RecipesBase.@series begin
        subplot := 5
        label := ""
        legend := false
        linecolor := [:blue :red]
        xlims := window
        n, e = _decimate_to_max_samples.((n_orig, e_orig), max_samples, antialias, Ref(window))
        # Note we are taking the splitting off, hence -δt
        SeisSplit.apply_split!(n, e, s.phi_best, -s.dt_best)
        fast, slow = Seis.rotate_through(n, e, s.phi_best)
        [Seis.times(fast), Seis.times(slow)], [Seis.trace(fast), slow_polarity.*Seis.trace(slow)]
    end

    ## λ₂ surface
    # 10σ - 100σ contours
    RecipesBase.@series begin
        subplot := 6
        seriestype := :contour
        levels := 10:10:100
        linewidth := 1
        linecolor := :gray
        linealpha := 0.6
        s.dt, s.phi, s.lam2
    end

    # 1σ – 10σ contours
    RecipesBase.@series begin
        subplot := 6
        seriestype := :contour
        levels := 1:10
        linewidth := 1
        linecolor := :black
        s.dt, s.phi, s.lam2
    end
    # 1σ contour
    RecipesBase.@series begin
        subplot := 6
        seriestype := :contour
        levels := [1]
        linewidth := 3
        linecolor := :black
        s.dt, s.phi, s.lam2
    end
    # Best value
    RecipesBase.@series begin
        subplot := 6
        seriestype := :scatter
        xguide --> "dt / s"
        yguide --> "phi / deg"
        xlims := extrema(s.dt)
        ylims := extrema(s.phi)
        label := ""
        [s.dt_best], [s.phi_best]
    end

    ## Particle motion before correction
    amax = 1.2*maximum(x -> sqrt(x[1]^2 + x[2]^2), zip(Seis.trace(e_orig), Seis.trace(n_orig)))
    RecipesBase.@series begin
        subplot := 7
        aspect_ratio := :equal
        label := ""
        linecolor --> :black
        xguide := "East"
        yguide := "North"
        xticks --> nothing
        yticks --> nothing
        xlims := (-amax, amax)
        ylims := (-amax, amax)
        x, y = _decimate_to_max_samples.((e_orig, n_orig), max_samples, antialias,
            Ref(window))
        Seis.trace(x), Seis.trace(y)
    end

    ## Particle motion after correction
    RecipesBase.@series begin
        subplot := 8
        aspect_ratio := :equal
        label := ""
        linecolor --> :black
        xguide := "East"
        yguide := "North"
        xticks --> nothing
        yticks --> nothing
        xlims := (-amax, amax)
        ylims := (-amax, amax)
        n, e = deepcopy.((n_orig, e_orig))
        SeisSplit.apply_split!(n, e, s.phi_best, -s.dt_best)
        y, x = _decimate_to_max_samples.((n, e), max_samples, antialias, Ref(window))
        Seis.trace(x), Seis.trace(y)
    end

    ## Empty plot at bottom right
    # RecipesBase.@series begin
    #     subplot := 9
    #     framestyle := :empty
    #     nothing
    # end
end

"""
    _decimate_to_max_samples(t, max_samples, antialias[, cut]) -> t′

Return a new trace `t′` cut between `cut[1]` and `cut[2]`, and decimated such that
its number of samples is less than or equal to `max_samples`.  If `antialias`
is `true`, then an antialiasing filter is applied; otherwise it is not.
"""
function _decimate_to_max_samples(t, max_samples, antialias, cut=(Seis.starttime(t), Seis.endtime(t)))
    max_samples > 10 || throw(ArgumentError("`max_samples` must be 10 or more"))
    t_cut = Seis.cut(t, cut...)
    n = Seis.nsamples(t_cut)
    decimation = 1
    while n/decimation > max_samples
        decimation += 1
    end
    Seis.decimate!(t_cut, decimation, antialias=antialias)
end

"""
    _apparent_polarity(fast, slow) -> -1/0/1

Return `1` if the `fast` and `slow` traces are positively correlated,
`-1` if they are negatively correlated, and `0` if their maximum
absolute cross-correlation is exactly 0.
"""
function _apparent_polarity(fast, slow)
    minval, maxval = extrema(DSP.xcorr(Seis.trace(fast), Seis.trace(slow)))
    max_xc = abs(minval) > abs(maxval) ? minval : maxval
    sign(max_xc)
end

max_stars_repo_name: anowacki/SeisSplit

max_issues_repo_name: anowacki/SeisSplit

max_forks_repo_name: anowacki/SeisSplit


text:
using LinearAlgebra
using Flux
using Flux: @epochs
using SPECTrecon
using MIRTjim
using CUDA
using BSON: @save, @load
using Plots: savefig

include("model.jl")
include("reg_elem.jl")
include("load_data.jl")
include("DVK.jl")
include("optim.jl")

dy = T(4.7952)
G = linearmap(mumap, psfs, dy; interpmeth = :two, mode = :fast)
model = residual_net()
niter = 4
β = 2

optim_cnn(xrecon16,
          density,
          G,
          model,
          yi,
          ri,
          niter,
          β,
          xtrue,
          doseGT,
          fullkernel,
          scale_factor_den,
          scale_factor_dose;
          nepoch = 200)

# Test code
@load "mymodel_recon_16th.bson" model
# run ML-EM recon
xrecon_cnn = reg_eml_em(xrecon16, density, G, model, yi, ri, niter; β)
# generate DVK dose-rate map
doseDVK_xrecon = gen_dvk(xrecon16, density, fullkernel, scale_factor_den, scale_factor_dose)
doseDVK_xtrue = gen_dvk(xtrue, density, fullkernel, scale_factor_den, scale_factor_dose)
# input to the CNN
data_input = Flux.unsqueeze(Flux.unsqueeze(xrecon_cnn, 4), 5)
# output of the CNN
pred = model(data_input)
pred_act = max.(0, pred[:,:,:,1,1])
# pred_dose = max.(0, pred[:,:,:,2,1] .+ doseDVK_xrecon)

idx = 20
jim(jim(pred_act[:,:,idx], "pred act"),
    jim(xtrue[:,:,idx], "xtrue"),
    jim(pred_act[:,:,idx] .- xtrue[:,:,idx], "diff"))
# residual diff of dose-rate map
jim(pred[:,:,idx,2,1], "residual dose-rate")
# savefig("activity_pred_100_xtrue_idx_40.pdf")
jim(jim(pred_dose[:,:,idx], "pred dose"),
    jim(doseGT[:,:,idx], "dose GT"),
    jim(pred_dose[:,:,idx] .- doseGT[:,:,idx], "diff"))
savefig("dose_pred_100_xtrue_idx_40.pdf")
jim(jim(doseDVK_xrecon[:,:,idx], "dose DVK xrecon"),
    jim(doseDVK_xtrue[:,:,idx], "dose DVK xtrue"),
    jim(doseDVK_xrecon[:,:,idx] .- doseDVK_xtrue[:,:,idx], "diff"))
savefig("doseDVK_xrecon_100_xtrue_idx_40.pdf")
jim(jim(xrecon_cnn[:,:,idx], "xrecon"),
    jim(pred_act[:,:,idx], "pred act"),
    jim(xrecon_cnn[:,:,idx] .- pred_act[:,:,idx], "diff"))
savefig("activity_pred_100_xrecon_idx_40.pdf")


text:
module RecursiveNodes

export recursive_nodes
export MultiIndexSet

include("./recursive_nodes.jl")
include("./multi_indices.jl")

end


text:
using Random

#Generate a matrix of given rank and ratio kappa between
#largest and smallest nonzero singular values
function generate_matrix(
    m :: Int,       #Number of rows
    n :: Int,       #Number of columns
    r :: Int,       #Rank
    κ :: Float64    #Condition number
)
    #Compute Rotations
    u, _ = qr(randn(m, m))
    vt, _ = qr(randn(n, n))

    #Compute singular values
    sig = sort(rand(r), rev=true) |> λ -> λ./maximum(λ)
    a = (1 - 1/κ)/(1 - minimum(sig))
    b = 1 - a
    σ =  a.*sig .+ b
    (min(m, n) > r) && append!(σ, zeros(min(m, n) - r))

    #Return Matrix
    return u*diagm(m, n, 0 => σ)*vt
end

# Generate a set of isotropic vector
function isotropic_vector(n::Int)
    return randn(n)
end

function isotropic_vector!(x::AbstractVector)
    randn!(x)
    return nothing
end

function randomized_trace(A::AbstractMatrix, nsamples::Int)
    n = size(A, 1)
    w = zeros(n)
    x = zeros(n)
    estimator = 0.0
    # TODO: use compensated summation technique for large nsamples
    factor = 1.0/nsamples

    for i=1:nsamples
        isotropic_vector!(w)
        x .= A*w
        estimator += factor*(w'*x)
    end

    return estimator
end


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Functions to compute the precession-nutation rotation according to IAU-2006
#   theory.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
#   [2] Vallado, D. A (06-Feb-2018). Consolidated Errata of Fundamentals of
#       Astrodynamics and Applications 4th Ed.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export precession_nutation_iau2006

"""
    function precession_nutation_iau2006(JD_TT::Number)

Compute the coordinates `X`, `Y`, and `s` related to the Celestial Intermediate
Pole (CIP) with respect to the Geocentric Celestial Reference Frame (GCRF). This
accounts for the effects of both precession and nutation of the CIP.

# Returns

* The coordinate `X` of the CIP w.r.t. the GCRF;
* The coordinate `Y` of the CIP w.r.t. the GCRF;
* The CIO locator `s` that provides the position of the CIO on the Equator of
  the CIP corresponding to the kinematical definition of the non-rotation origin
  in the GCRS when the CIP is moving with respect to the GCRS between the
  reference epoch and the epoch due to precession and nutation [1, p. 214].

"""
function precession_nutation_iau2006(JD_TT::Number)
    # Compute the Julian Centuries from `JD_TT`.
    T_TT = (JD_TT - JD_J2000)/36525

    # Auxiliary variables
    # ===================

    a2d = 1/3600
    d2r = π/180
    a2r = a2d*d2r

    # Delaunay arguments of the Sun and Moon
    # ======================================
    #
    # Evaluate the Delaunay arguments associated with the Moon and the Sun in
    # [arcsec] [1, p. 210].

    M_m  = @evalpoly(T_TT, +485868.249036, +1717915923.2178, +31.8792, +0.051635, -0.00024470)
    M_s  = @evalpoly(T_TT, +1287104.79305, +129596581.0481,   -0.5532, +0.000136, -0.00001149)
    u_Mm = @evalpoly(T_TT, +335779.526232, +1739527262.8478, -12.7512, -0.001037, +0.00000417)
    D_s  = @evalpoly(T_TT, +1072260.70369, +1602961601.2090,  -6.3706, +0.006593, -0.00003169)
    Ω_m  = @evalpoly(T_TT, +450160.398036,    -6962890.5431,  +7.4722, +0.007702, -0.00005939)

    # Convert to the interval [0,2π].
    M_s  = mod(M_s*a2d,  360)*d2r
    M_m  = mod(M_m*a2d,  360)*d2r
    u_Mm = mod(u_Mm*a2d, 360)*d2r
    D_s  = mod(D_s*a2d,  360)*d2r
    Ω_m  = mod(Ω_m*a2d,  360)*d2r

    # Planetary effects of the nutation and obliquity of the ecliptic
    # ===============================================================

    # Mean Heliocentric longitudes of the planets.
    #
    # TODO: In the example in [1, p. 221], the value related to Uranus is slight
    # different. The value used here is shown in [1, p. 211].

    λ_M☿ = @evalpoly(T_TT, 4.402_608_842, 2608.790_314_157_4)
    λ_M♀ = @evalpoly(T_TT, 3.176_146_697, 1021.328_554_621_1)
    λ_Me = @evalpoly(T_TT, 1.753_470_314,  628.307_584_999_1)
    λ_M♂ = @evalpoly(T_TT, 6.203_480_913,  334.061_242_670_0)
    λ_M♃ = @evalpoly(T_TT, 0.599_546_497,   52.969_096_264_1)
    λ_M♄ = @evalpoly(T_TT, 0.874_016_757,   21.329_910_496_0)
    λ_M⛢ = @evalpoly(T_TT, 5.481_293_872,    7.478_159_856_7)
    λ_M♆ = @evalpoly(T_TT, 5.311_886_287,    3.813_303_563_8)

    # General precession in longitude.
    p_λ = @evalpoly(T_TT, 0, 0.024_381_75, 0.000_005_386_91)

    # Convert to the interval [0,2π].
    λ_M☿ = mod(λ_M☿, 2π)
    λ_M♀ = mod(λ_M♀, 2π)
    λ_Me = mod(λ_Me, 2π)
    λ_M♂ = mod(λ_M♂, 2π)
    λ_M♃ = mod(λ_M♃, 2π)
    λ_M♄ = mod(λ_M♄, 2π)
    λ_M⛢ = mod(λ_M⛢, 2π)
    λ_M♆ = mod(λ_M♆, 2π)
    p_λ  = mod(p_λ,  2π)

    # X position of the CIP
    # ==========================================================================

    X = @evalpoly(T_TT,   -0.016_617,
                       +2004.191_898,
                          -0.429_782_9,
                          -0.198_618_34,
                          +0.000_007_578,
                          +0.000_005_928_5)

    Xc = zeros(typeof(X), 5)
    coefs_v = (nut_coefs_iau2006_X0, nut_coefs_iau2006_X1, nut_coefs_iau2006_X2,
               nut_coefs_iau2006_X3, nut_coefs_iau2006_X4)

    @inbounds for i = 1:length(coefs_v)

        coefs = coefs_v[i]
        Xp    = zero(typeof(X))

        for j = 1:size(coefs,1)
            As     = coefs[j, 2]
            Ac     = coefs[j, 3]
            ap     = coefs[j, 4]*M_m  + coefs[j, 5]*M_s  + coefs[j, 6]*u_Mm +
                     coefs[j, 7]*D_s  + coefs[j, 8]*Ω_m  + coefs[j, 9]*λ_M☿ +
                     coefs[j,10]*λ_M♀ + coefs[j,11]*λ_Me + coefs[j,12]*λ_M♂ +
                     coefs[j,13]*λ_M♃ + coefs[j,14]*λ_M♄ + coefs[j,15]*λ_M⛢ +
                     coefs[j,16]*λ_M♆ + coefs[j,17]*p_λ
            sj,cj  = sincos(ap)
            Xp    += (As*sj + Ac*cj)/1e6
        end

        Xc[i] = Xp
    end

    X += @evalpoly(T_TT, Xc[1], Xc[2], Xc[3], Xc[4], Xc[5])

    # Convert to [rad].
    X *= a2r

    # Y position of the CIP
    # ==========================================================================

    Y = @evalpoly(T_TT, -0.006_951,
                        -0.025_896,
                       -22.407_274_7,
                        +0.001_900_59,
                        +0.001_112_526,
                        +0.000_000_135_8)

    Yc = zeros(typeof(Y), 5)
    coefs_v = (nut_coefs_iau2006_Y0, nut_coefs_iau2006_Y1, nut_coefs_iau2006_Y2,
               nut_coefs_iau2006_Y3, nut_coefs_iau2006_Y4)

    @inbounds for i = 1:length(coefs_v)

        coefs = coefs_v[i]
        Yp    = zero(typeof(Y))

        for j = 1:size(coefs,1)
            As     = coefs[j, 2]
            Ac     = coefs[j, 3]
            ap     = coefs[j, 4]*M_m  + coefs[j, 5]*M_s  + coefs[j, 6]*u_Mm +
                     coefs[j, 7]*D_s  + coefs[j, 8]*Ω_m  + coefs[j, 9]*λ_M☿ +
                     coefs[j,10]*λ_M♀ + coefs[j,11]*λ_Me + coefs[j,12]*λ_M♂ +
                     coefs[j,13]*λ_M♃ + coefs[j,14]*λ_M♄ + coefs[j,15]*λ_M⛢ +
                     coefs[j,16]*λ_M♆ + coefs[j,17]*p_λ
            sj,cj  = sincos(ap)
            Yp    += (As*sj + Ac*cj)/1e6
        end

        Yc[i] = Yp
    end

    Y += @evalpoly(T_TT, Yc[1], Yc[2], Yc[3], Yc[4], Yc[5])

    # Convert to [rad].
    Y *= a2r

    # Parameter `s` (CIO locator)
    # ==========================================================================
    #
    # The value `s` provides the position of the CIO on the Equator of the CIP
    # corresponding to the kinematical definition of the non-rotation origin in
    # the GCRS when the CIP is moving with respect to the GCRS between the
    # reference epoch and the epoch due to precession and nutation [1, p. 214].

    # We must convert this term to [arcsec] to match the units.
    #                      ||
    #                   |------|
    s = @evalpoly(T_TT, (-X*Y/2)/a2r + 0.000_094,
                                     +0.003_808_65,
                                     -0.000_122_68,
                                     -0.072_574_11,
                                     +0.000_027_98,
                                     +0.000_015_65)

    sc = zeros(typeof(s), 5)
    coefs_v = (nut_coefs_iau2006_s0, nut_coefs_iau2006_s1, nut_coefs_iau2006_s2,
               nut_coefs_iau2006_s3, nut_coefs_iau2006_s4)

    @inbounds for i = 1:length(coefs_v)

        coefs = coefs_v[i]
        sp    = zero(typeof(s))

        for j = 1:size(coefs,1)
            As     = coefs[j, 2]
            Ac     = coefs[j, 3]
            ap     = coefs[j, 4]*M_m  + coefs[j, 5]*M_s  + coefs[j, 6]*u_Mm +
                     coefs[j, 7]*D_s  + coefs[j, 8]*Ω_m  + coefs[j, 9]*λ_M☿ +
                     coefs[j,10]*λ_M♀ + coefs[j,11]*λ_Me + coefs[j,12]*λ_M♂ +
                     coefs[j,13]*λ_M♃ + coefs[j,14]*λ_M♄ + coefs[j,15]*λ_M⛢ +
                     coefs[j,16]*λ_M♆ + coefs[j,17]*p_λ
            sj,cj  = sincos(ap)
            sp    += (As*sj + Ac*cj)/1e6
        end

        sc[i] = sp
    end

    s += @evalpoly(T_TT, sc[1], sc[2], sc[3], sc[4], sc[5])

    # Convert to [rad].
    s *= a2r

    return X, Y, s
end


text:
@testset "Visualization Forecast" begin
    air_passengers = CSV.File(StateSpaceModels.AIR_PASSENGERS) |> DataFrame
    log_air_passengers = log.(air_passengers.passengers)
    model = BasicStructural(log_air_passengers, 12)
    # forecasting
    b = cross_validation(model, 24, 110)
    rec = RecipesBase.apply_recipe(Dict{Symbol, Any}(), b, "str")
    @test length(rec) == 2
end


text:
function execute(domain::GroundDomain, state::State, term::Term, options...)
    if term.name == get_name(no_op) return state end
    return execute!(domain, copy(state), term; options...)
end

function execute(domain::GroundDomain, state::State,
                 action::GroundActionGroup, args; options...)
    return execute!(domain, copy(state), action, args; options...)
end

function execute(domain::GroundDomain, state::State,
                 action::GroundAction; options...)
    return execute!(domain, copy(state), action; options...)
end

function execute!(domain::GroundDomain, state::State, term::Term; options...)
   if term.name == get_name(no_op) return state end
   if (term isa Const) term = Compound(term.name, []) end
   action = domain.actions[term.name].actions[term]
   return execute!(domain, state, action; options...)
end

function execute!(domain::GroundDomain, state::State,
                  group::GroundActionGroup, args; options...)
   term = Compound(group.name, args)
   return execute!(domain, state, group.actions[term]; options...)
end

function execute!(domain::GroundDomain, state::State,
                  action::GroundAction; check::Bool=true)
    # Check whether preconditions hold
    if check && !available(domain, state, action)
        error("Precondition $(get_precond(action)) does not hold.")
    end
    # Update state with diff
    return update!(domain, state, action.effect)
end


text:
using HDF5
using ArgParse

require("slbfgs.jl")
require("sgd.jl")

s = ArgParseSettings()

@add_arg_table s begin
  "--hdf"
  help = "Name of the HDF5 where results should be stored"
  arg_type = String
  "--seed"
  help = "Random seed"
  arg_type = Int
  default = -1
  "--exp"
  help = "Filename of a Julia file describing the experiment"
  # A Julia file describing an experiment must export the following symbols:
  # data for the dataset, loss for the loss function, grad! for the gradient
  # and hvp! for Hessian vector products  and theta for the initial
  # parameter vector (see blog.jl)
  arg_type = String
  default = "blog.jl"
  "--method"
  help = "The optimization method to be used, can be 'lbfgs' or 'sgd'"
  arg_type = String
  default = "lbfgs"
  "--epochs"
  help = "Number of passes through the data"
  arg_type = Int
  default = 20
  "--stepsize"
  help = "The stepsize to be used"
  arg_type = Float64
  default = 0.0001
  "--lambda"
  help = "Regularization parameter"
  arg_type = Float64
  default = 0.001
  "--memory_size"
  help = "The memory size of the LBFGS matrix"
  arg_type = Int
  default = 10
  "--delta"
  help = "Regularization parameter for LBFGS matrix, think B_k^{-1} is replaced by (B_k + delta I)^{-1}"
  arg_type = Float64
  default = 0.0
  "--hess_samples"
  help = "Number of samples used to estimate the Hessian matrix"
  arg_type = Int
  default = 20
  "--hess_period"
  help = "Number of stochastic steps before a new vector pair is added to the LBFGS matrix"
  arg_type = Int
  default = 10
  "--reduce_grad_var"
  help = "Whether variance reduced gradients are used"
  arg_type = Bool
  default = true
  "--use_gradient_step"
  help = "When updating theta, take B_k as the identity matrix"
  arg_type = Bool
  default = false
  "--use_hvp"
  help = "If true, use hessian vector products to approximate the LBFGS matrix"
  arg_type = Bool
  default = true
end

parsed_args = parse_args(ARGS, s)

println(parsed_args)

function save_dict(dict::Dict, group::HDF5Group)
  for (key, value) in dict
    group[key] = string(value)
  end
end

function do_experiment(config::Dict)
  if config["seed"] == -1
    config["seed"] = time_ns()
  end
  srand(config["seed"])

  require(config["exp"])

  h5open(config["hdf"], "w") do hdf
      params = g_create(hdf, "params")
      save_dict(config, params)
      diagnostics = g_create(hdf, "diagnostics")
      dataset = g_create(hdf, "dataset")
      save_dict(dataset_info, dataset)
      if config["method"] == "sgd"
        sgd(loss, grad!, data, theta, config, diagnostics)
      end
      if config["method"] == "lbfgs"
        res = slbfgs(loss, grad!, hvp!, data, theta, config, diagnostics)
      end
  end
end

do_experiment(parsed_args)

max_stars_repo_name: pcmoritz/slbfgs

max_issues_repo_name: pcmoritz/slbfgs

max_forks_repo_name: pcmoritz/slbfgs


text:
import Revise
import GLRenderer as GL
import ThreeDP3 as T
import Images as I
import MiniGSG as S
import PoseComposition: Pose
import Serialization
import NearestNeighbors as NN
import Statistics: mean
using Gen
import MeshCatViz

import MeshCatViz
MeshCatViz.setup_visualizer()

function compute_ADDS_error(xyz_models, gt_ids, gt_poses, pred_ids, pred_poses)
    add_s = []
    for (id,gt_pose) in zip(gt_ids, gt_poses)
        idx = findfirst(pred_ids .== id)
        if isnothing(idx)
            push!(add_s, 1.0)
            continue
        end
        pred_pose = pred_poses[idx]

        model = xyz_models[id]

        tree = NN.KDTree(T.move_points_to_frame_b(model, pred_pose))
        _, dists = NN.nn(tree, T.move_points_to_frame_b(model, gt_pose))

        push!(add_s, min(mean(dists)))
    end
    add_s
end

YCB_DIR = "/home/nishadg/mcs/ThreeDVision.jl/data/ycbv2"
world_scaling_factor = 1.0
xyz_models = T.load_ycbv_point_xyz(YCB_DIR);
all_ids = collect(1:21);

RESULTS_DIR = "/home/nishadg/mcs/ThreeDVision.jl/pre_neurips_scripts/pybullet_inference/full_massive_run_1/"
TS = collect(0:499);

# +
errors_dense_fusion = Dict([id => [] for id in all_ids])
errors_robust_6d = Dict([id => [] for id in all_ids])
errors_3dp3_star = Dict([id => [] for id in all_ids])
errors_3dp3 = Dict([id => [] for id in all_ids])

for t in TS
    metadata = Serialization.deserialize(joinpath(RESULTS_DIR,"data",lpad(t,6,"0")*"-meta.data"))
    cam_pose = metadata.camera_pose
    ids = metadata.ids
    gt_poses = metadata.object_poses

    dense_poses, dense_ids = T.load_synthetic_dense_fusion_predictions(RESULTS_DIR, t);
    dense_poses = [T.get_c_relative_to_a(cam_pose,p) for p in dense_poses]

    objnet_poses, objnet_ids = T.load_synthetic_objposenet_fusion_predictions(RESULTS_DIR, t)
    objnet_poses = [T.get_c_relative_to_a(cam_pose,p) for p in objnet_poses]

    world_scaling_factor = 100.0
    inf_poses_no_sg, inf_poses = Serialization.deserialize(joinpath(RESULTS_DIR,"sg_results/$(lpad(t,4,"0")).poses"));
    inf_poses = [Pose(p.pos ./ world_scaling_factor, p.orientation) for p in inf_poses]
    inf_poses_no_sg = [Pose(p.pos ./ world_scaling_factor, p.orientation) for p in inf_poses_no_sg]

    df = compute_ADDS_error(xyz_models, ids, gt_poses, dense_ids, dense_poses)
    r6d = compute_ADDS_error(xyz_models, ids, gt_poses, objnet_ids, objnet_poses)
    ablation = compute_ADDS_error(xyz_models, ids, gt_poses, ids, inf_poses_no_sg)
    full = compute_ADDS_error(xyz_models, ids, gt_poses, ids, inf_poses)

    for (i,id) in enumerate(ids)
        push!(errors_dense_fusion[id], df[i])
        push!(errors_robust_6d[id], r6d[i])
        push!(errors_3dp3_star[id], ablation[i])
        push!(errors_3dp3[id], full[i])
    end
end

# +
ids = [1,2,3,4,9]

id = 1
df= errors_dense_fusion[id]
r6d= errors_robust_6d[id]
ablation= errors_3dp3_star[id]
full= errors_3dp3[id]

max_plot_val = 0.03
valid_idxs = (df .< max_plot_val) .&  (ablation .< max_plot_val)  .&  (full .< max_plot_val) .&  (full .< max_plot_val);

Plots.scatter(df[valid_idxs], full[valid_idxs],alpha=0.4)
Plots.scatter!(r6d[valid_idxs], full[valid_idxs],alpha=0.4)
Plots.plot!([0.0, max_plot_val],[0.0, max_plot_val],legend=false,
    xlabel="Dense Fusion ADD-S Error",
    ylabel="3DP3 ADD-S Error")
# -

Plots.scatter(r6d[valid_idxs], full[valid_idxs])
Plots.plot!([0.0, max_plot_val],[0.0, max_plot_val],legend=false,
    xlabel="Dense Fusion ADD-S Error",
    ylabel="3DP3 ADD-S Error")

Plots.scatter(df[valid_idxs], full[valid_idxs])
Plots.plot!([0.0, max_plot_val],[0.0, max_plot_val],legend=false,
    xlabel="Dense Fusion ADD-S Error",
    ylabel="3DP3 ADD-S Error")

Plots.scatter(df[valid_idxs], ablation[valid_idxs])
Plots.plot!([0.0, max_plot_val],[0.0, max_plot_val],legend=false,
    xlabel="Dense Fusion ADD-S Error",
    ylabel="3DP3 Ablation ADD-S Error")

Plots.scatter(ablation[valid_idxs], full[valid_idxs])
Plots.plot!([0.0, max_plot_val],[0.0, max_plot_val],legend=false,
    xlabel="3DP3 Ablation ADD-S Error",
    ylabel="3DP3 ADD-S Error")

# +
import Serialization: deserialize

function compute_ADDS_error(xyz_models, gt_ids, gt_poses, pred_ids, pred_poses)
    add_s = []
    for (id,gt_pose) in zip(gt_ids, gt_poses)
        idx = findfirst(pred_ids .== id)
        if isnothing(idx)
            push!(add_s, 1.0)
            continue
        end
        pred_pose = pred_poses[idx]

        model = xyz_models[id]

        tree = KDTree(move_points_to_frame_b(model, pred_pose))
        _, dists = nn(tree, move_points_to_frame_b(model, gt_pose))

        push!(add_s, min(mean(dists)))
    end
    add_s
end

function load_synthetic_results(RESULTS_DIR, T; world_scaling_factor=100.0)
    metadata = Serialization.deserialize(joinpath(RESULTS_DIR,"data",lpad(T,6,"0")*"-meta.data"))

    cam_pose = metadata.camera_pose
    ids = metadata.ids
    gt_poses = metadata.object_poses

    dense_poses, dense_ids = T.load_synthetic_dense_fusion_predictions(RESULTS_DIR, T)
    dense_poses = [Pose(p.pos, p.orientation) for p in dense_poses]
    dense_poses = [T.get_c_relative_to_a(cam_pose,p) for p in dense_poses]

    inf_poses_no_sg, inf_poses = Serialization.deserialize(joinpath(RESULTS_DIR,"sg_results/$(lpad(T,4,"0")).poses"));
    inf_poses = [Pose(p.pos ./ world_scaling_factor, p.orientation) for p in inf_poses]
    inf_poses_no_sg = [Pose(p.pos ./ world_scaling_factor, p.orientation) for p in inf_poses_no_sg]

    ids, gt_poses, dense_ids, dense_poses, inf_poses, inf_poses_no_sg
end

function evaluate_scene(RESULTS_DIR, xyz_models, T; world_scaling_factor=100.0)
    metadata = Serialization.deserialize(joinpath(RESULTS_DIR,"data",lpad(T,6,"0")*"-meta.data"))

    cam_pose = metadata.camera_pose
    ids = metadata.ids
    gt_poses = metadata.object_poses

    dense_poses, dense_ids = load_synthetic_dense_fusion_predictions(RESULTS_DIR, T)
    dense_poses = [Pose(p.pos, p.orientation) for p in dense_poses]
    dense_poses = [T.get_c_relative_to_a(cam_pose,p) for p in dense_poses]

    objnet_poses, objnet_ids = load_synthetic_objposenet_fusion_predictions(RESULTS_DIR, T)
    objnet_poses = [Pose(p.pos, p.orientation) for p in objnet_poses]
    objnet_poses = [T.get_c_relative_to_a(cam_pose,p) for p in objnet_poses]

    inf_poses_no_sg, inf_poses = deserialize(joinpath(RESULTS_DIR,"sg_results/$(lpad(T,4,"0")).poses"));
    inf_poses = [Pose(p.pos ./ world_scaling_factor, p.orientation) for p in inf_poses]
    inf_poses_no_sg = [Pose(p.pos ./ world_scaling_factor, p.orientation) for p in inf_poses_no_sg]

    dense_add = compute_ADDS_error(xyz_models, ids, gt_poses, dense_ids, dense_poses);
    objnet_add = compute_ADDS_error(xyz_models, ids, gt_poses, objnet_ids, objnet_poses);
    my_wo_refinement_add = compute_ADDS_error(xyz_models, ids, gt_poses, ids, inf_poses_no_sg);
    my_add = compute_ADDS_error(xyz_models, ids, gt_poses, ids, inf_poses);

    ids, objnet_add, dense_add, my_wo_refinement_add, my_add
end

# +
YCB_DIR = "/home/nishadg/mcs/ThreeDVision.jl/data/ycbv2"
original_id_to_cloud = T.load_ycbv_models(YCB_DIR);
xyz_models = T.load_ycbv_point_xyz(YCB_DIR);
model_list = readlines(joinpath(YCB_DIR, "model_list.txt"))

# # +
RESULTS_DIRS = [
    "/home/nishadg/mcs/ThreeDVision.jl/scripts/pybullet_inference/full_massive_run_1/",
    "/home/nishadg/mcs/ThreeDVision.jl/scripts/pybullet_inference/full_massive_run_2/",
]
TS_LIST = [collect(0:499),collect(500:999)]
# -


text:
# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

function Base.copy(::Evaluator)
    return error("Copying nonlinear problems not yet implemented")
end

"""
    ordinal_index(evaluator::Evaluator, c::ConstraintIndex)::Int

Return the 1-indexed value of the constraint index `c` in `evaluator`.

## Examples

```julia
model = Model()
x = MOI.VariableIndex(1)
c1 = add_constraint(model, :(\$x^2), MOI.LessThan(1.0))
c2 = add_constraint(model, :(\$x^2), MOI.LessThan(1.0))
evaluator = Evaluator(model)
MOI.initialize(evaluator, Symbol[])
ordinal_index(evaluator, c2)  # Returns 2
delete(model, c1)
evaluator = Evaluator(model)
MOI.initialize(evaluator, Symbol[])
ordinal_index(model, c2)  # Returns 1
```
"""
function ordinal_index(evaluator::Evaluator, c::ConstraintIndex)
    # TODO(odow): replace with a cache that maps indices to their 1-indexed
    # row in the constraint matrix. But failing that, since we know that
    # constraints are added in increasing order and that they can be deleted, we
    # know that index `i` must appear as constraint `1` to `i`. So we start at
    # `i` and backtrack (to account for deleted constraints) until we find it.
    # In the typical case with no deletion, there should be no overhead.
    for i in min(c.value, length(evaluator.ordered_constraints)):-1:1
        if evaluator.ordered_constraints[i] == c
            return i
        end
    end
    return error("Invalid constraint index $(c)")
end

function MOI.features_available(evaluator::Evaluator)
    features = Symbol[]
    if evaluator.backend !== nothing
        append!(features, MOI.features_available(evaluator.backend))
    end
    if !(:ExprGraph in features)
        push!(features, :ExprGraph)
    end
    return features
end

function Base.show(io::IO, evaluator::Evaluator)
    Base.print(io, "Nonlinear.Evaluator with available features:")
    for feature in MOI.features_available(evaluator)
        print(io, "\n  * :", feature)
    end
    return
end

function MOI.initialize(evaluator::Evaluator, features::Vector{Symbol})
    empty!(evaluator.ordered_constraints)
    empty!(evaluator.julia_expressions)
    evaluator.eval_objective_timer = 0.0
    evaluator.eval_objective_gradient_timer = 0.0
    evaluator.eval_constraint_timer = 0.0
    evaluator.eval_constraint_jacobian_timer = 0.0
    evaluator.eval_hessian_lagrangian_timer = 0.0
    append!(evaluator.ordered_constraints, keys(evaluator.model.constraints))
    if :ExprGraph in features
        for i in 1:length(evaluator.model.expressions)
            push!(
                evaluator.julia_expressions,
                convert_to_expr(
                    evaluator,
                    evaluator.model.expressions[i];
                    moi_output_format = true,
                ),
            )
        end
        filter!(f -> f != :ExprGraph, features)
    end
    if evaluator.backend !== nothing
        MOI.initialize(evaluator.backend, features)
    end
    return
end

function MOI.objective_expr(evaluator::Evaluator)
    if evaluator.model.objective === nothing
        error(
            "Unable to query objective_expr because no nonlinear objective " *
            "was set",
        )
    end
    return convert_to_expr(
        evaluator,
        evaluator.model.objective;
        moi_output_format = true,
    )
end

function MOI.constraint_expr(evaluator::Evaluator, i::Int)
    constraint = evaluator.model[evaluator.ordered_constraints[i]]
    f = convert_to_expr(
        evaluator,
        constraint.expression;
        moi_output_format = true,
    )
    if constraint.set isa MOI.LessThan
        return :($f <= $(constraint.set.upper))
    elseif constraint.set isa MOI.GreaterThan
        return :($f >= $(constraint.set.lower))
    elseif constraint.set isa MOI.EqualTo
        return :($f == $(constraint.set.value))
    else
        @assert constraint.set isa MOI.Interval
        return :($(constraint.set.lower) <= $f <= $(constraint.set.upper))
    end
end

function MOI.eval_objective(evaluator::Evaluator, x)
    start = time()
    obj = MOI.eval_objective(evaluator.backend, x)
    evaluator.eval_objective_timer += time() - start
    return obj
end

function MOI.eval_objective_gradient(evaluator::Evaluator, g, x)
    start = time()
    MOI.eval_objective_gradient(evaluator.backend, g, x)
    evaluator.eval_objective_gradient_timer += time() - start
    return
end

function MOI.eval_constraint(evaluator::Evaluator, g, x)
    start = time()
    MOI.eval_constraint(evaluator.backend, g, x)
    evaluator.eval_constraint_timer += time() - start
    return
end

function MOI.jacobian_structure(evaluator::Evaluator)
    return MOI.jacobian_structure(evaluator.backend)
end

function MOI.eval_constraint_jacobian(evaluator::Evaluator, J, x)
    start = time()
    MOI.eval_constraint_jacobian(evaluator.backend, J, x)
    evaluator.eval_constraint_jacobian_timer += time() - start
    return
end

function MOI.hessian_lagrangian_structure(evaluator::Evaluator)
    return MOI.hessian_lagrangian_structure(evaluator.backend)
end

function MOI.eval_hessian_lagrangian(evaluator::Evaluator, H, x, σ, μ)
    start = time()
    MOI.eval_hessian_lagrangian(evaluator.backend, H, x, σ, μ)
    evaluator.eval_hessian_lagrangian_timer += time() - start
    return
end

function MOI.eval_constraint_jacobian_product(evaluator::Evaluator, y, x, w)
    start = time()
    MOI.eval_constraint_jacobian_product(evaluator.backend, y, x, w)
    evaluator.eval_constraint_jacobian_timer += time() - start
    return
end

function MOI.eval_constraint_jacobian_transpose_product(
    evaluator::Evaluator,
    y,
    x,
    w,
)
    start = time()
    MOI.eval_constraint_jacobian_transpose_product(evaluator.backend, y, x, w)
    evaluator.eval_constraint_jacobian_timer += time() - start
    return
end

function MOI.eval_hessian_lagrangian_product(
    evaluator::Evaluator,
    H,
    x,
    v,
    σ,
    μ,
)
    start = time()
    MOI.eval_hessian_lagrangian_product(evaluator.backend, H, x, v, σ, μ)
    evaluator.eval_hessian_lagrangian_timer += time() - start
    return
end

"""
    adjacency_matrix(nodes::Vector{Node})

Compute the sparse adjacency matrix describing the parent-child relationships in
`nodes`.

The element `(i, j)` is `true` if there is an edge *from* `node[j]` to
`node[i]`. Since we get a column-oriented matrix, this gives us a fast way to
look up the edges leaving any node (i.e., the children).
"""
function adjacency_matrix(nodes::Vector{Node})
    N = length(nodes)
    I, J = Vector{Int}(undef, N), Vector{Int}(undef, N)
    numnz = 0
    for (i, node) in enumerate(nodes)
        if node.parent < 0
            continue
        end
        numnz += 1
        I[numnz] = i
        J[numnz] = node.parent
    end
    resize!(I, numnz)
    resize!(J, numnz)
    return SparseArrays.sparse(I, J, ones(Bool, numnz), N, N)
end

"""
    convert_to_expr(
        evaluator::Evaluator,
        expr::Expression;
        moi_output_format::Bool,
    )

Convert the [`Expression`](@ref) `expr` into a Julia `Expr`.

If `moi_output_format = true`:
 * subexpressions will be converted to Julia `Expr` and substituted into the
   output expression.
 * the current value of each parameter will be interpolated into the expression
 * variables will be represented in the form `x[MOI.VariableIndex(i)]`

If `moi_output_format = false`:
 * subexpressions will be represented by a [`ExpressionIndex`](@ref) object.
 * parameters will be represented by a [`ParameterIndex`](@ref) object.
 * variables will be represennted by an [`MOI.VariableIndex`](@ref) object.

!!! warning
    To use `moi_output_format = true`, you must have first called
    [`MOI.initialize`](@ref) with `:ExprGraph` as a requested feature.
"""
function convert_to_expr(
    evaluator::Evaluator,
    expression::Expression;
    moi_output_format::Bool,
)
    expr = convert_to_expr(evaluator.model, expression)
    if moi_output_format
        return _convert_to_moi_format(evaluator, expr)
    end
    return expr
end

_convert_to_moi_format(::Evaluator, x::MOI.VariableIndex) = :(x[$x])

function _convert_to_moi_format(evaluator::Evaluator, p::ParameterIndex)
    return evaluator.model[p]
end

function _convert_to_moi_format(evaluator::Evaluator, x::ExpressionIndex)
    return evaluator.julia_expressions[x.value]
end

_convert_to_moi_format(::Evaluator, x) = x

function _convert_to_moi_format(evaluator::Evaluator, x::Expr)
    for i in 1:length(x.args)
        x.args[i] = _convert_to_moi_format(evaluator, x.args[i])
    end
    return x
end

max_stars_repo_name: JuliaOpt/MathOptInterface

max_issues_repo_name: JuliaOpt/MathOptInterface

max_forks_repo_name: JuliaOpt/MathOptInterface


text:
using ASCIIPlots
using Base.Test

x = randn(20)
y = 21.0 * x + randn(20)
lineplot(x, y)


text:
using QuantumWalk

import QuantumWalk: check_qwdynamics, QWDynamics, Szegedy, AbstractSzegedy, QWModelDiscr

##

struct QWPeriod{T} <: QWDynamics{T}
  model::T
  parameters::Dict{Symbol,Any}

  function QWPeriod(model::T,
                    parameters::Dict{Symbol}) where T<:QWModelDiscr
    check_qwdynamics(QWPeriod, model, parameters)
    new{T}(model, parameters)
  end
end

function determine_period(qwp::QWPeriod,
                          init_state,
                          state_diff_val::Real,
                          tmax::Int=nv(graph(qwp)))
  state = evolve(qwp, init_state)
  for t=1:tmax
    if state_diff(qwp, state, init_state) < state_diff_val
      return t
    end
    state = evolve(qwp, state)
  end
  return -1
end

function QWPeriod(szegedy::AbstractSzegedy)
  operators = QuantumWalk.szegedy_walk_operators(szegedy)
  parameters = Dict{Symbol,Any}(:operators => operators)
  QWPeriod(szegedy, parameters)
end

function check_qwdynamics(::Type{QWPeriod}, szegedy::AbstractSzegedy, parameters::Dict{Symbol})
  QuantumWalk.check_qwmodel(szegedy, parameters)
end

function state_diff(qwp::QWPeriod{<:AbstractSzegedy},
                    state1::SparseVector{T},
                    state2::SparseVector{T}) where T<:Number
  1-abs(sum(state1 .* conj.(state2)))
end
##
n = 20
qwp = QWPeriod(Szegedy(barabasi_albert(n, 3)));
determine_period(qwp, sparse(fill(1/n, n^2)), 0.01)
state = sparse(randn(n^2)); state /= norm(state);
determine_period(qwp, state, 0.05, 8000)


text:
module Logging

import Base: show, info, warn

export debug, info, warn, err, critical, log,
       @debug, @info, @warn, @err, @error, @critical, @log,
       Logger,
       LogLevel, DEBUG, INFO, WARNING, ERROR, CRITICAL, OFF,
       LogFacility,
       SysLog

@enum LogLevel OFF=-1 CRITICAL=2 ERROR WARNING INFO=6 DEBUG

function Base.convert(::Type{LogLevel}, x::AbstractString)
    Dict("OFF"=>OFF,
         "CRITICAL"=>CRITICAL,
         "ERROR"=>ERROR,
         "WARNING"=>WARNING,
         "INFO"=>INFO,
         "DEBUG"=>DEBUG)[x]
end

@enum LogFacility LOG_KERN LOG_USER LOG_MAIL LOG_DAEMON LOG_AUTH LOG_SYSLOG LOG_LPR LOG_NEWS LOG_UUCP LOG_CRON LOG_AUTHPRIV LOG_LOCAL0=16 LOG_LOCAL1 LOG_LOCAL2 LOG_LOCAL3 LOG_LOCAL4 LOG_LOCAL5 LOG_LOCAL6 LOG_LOCAL7

type SysLog
    socket::UDPSocket
    ip::IPv4
    port::UInt16
    facility::LogFacility
    machine::AbstractString
    user::AbstractString
    maxlength::UInt16

    SysLog(host::AbstractString,
           port::Int,
           facility::LogFacility=LOG_USER,
           machine::AbstractString=gethostname(),
           user::AbstractString=Base.source_path()==nothing ? "" : basename(Base.source_path()),
           maxlength::Int=1024) = new(
        UDPSocket(),
        getaddrinfo(host),
        UInt16(port),
        facility,
        machine,
        user,
        UInt16(maxlength)
    )
end

LogOutput = Union{IO,SysLog}

type Logger
    name::AbstractString
    level::LogLevel
    output::Array{LogOutput,1}
    parent::Logger

    Logger(name::AbstractString, level::LogLevel, output::IO, parent::Logger) = new(name, level, [output], parent)
    Logger(name::AbstractString, level::LogLevel, output::IO) = (x = new(); x.name = name; x.level=level; x.output=[output]; x.parent=x)
    Logger{T<:LogOutput}(name::AbstractString, level::LogLevel, output::Array{T,1}, parent::Logger) = new(name, level, output, parent)
    Logger{T<:LogOutput}(name::AbstractString, level::LogLevel, output::Array{T,1}) = (x = new(); x.name = name; x.level=level; x.output=output; x.parent=x)
end

show(io::IO, logger::Logger) = print(io, "Logger(", join(Any[logger.name,
                                                             logger.level,
                                                             logger.output,
                                                             logger.parent.name], ","), ")")

const _root = Logger("root", WARNING, STDERR)
Logger(name::AbstractString;args...) = configure(Logger(name, WARNING, STDERR, _root); args...)
Logger() = Logger("logger")

write_log(syslog::SysLog, color::Symbol, msg::AbstractString) = send(syslog.socket, syslog.ip, syslog.port, length(msg) > syslog.maxlength ? msg[1:syslog.maxlength] : msg)
write_log{T<:IO}(output::T, color::Symbol, msg::AbstractString) = (print(output, msg); flush(output))
write_log(output::Base.TTY, color::Symbol, msg::AbstractString) = Base.print_with_color(color, output, msg)

function log(syslog::SysLog, level::LogLevel, color::Symbol, logger_name::AbstractString, msg...)
    # syslog needs a timestamp in the form: YYYY-MM-DDTHH:MM:SS-TZ:TZ
    t = time()
    timestamp = string(Libc.strftime("%Y-%m-%dT%H:%M:%S",t), Libc.strftime("%z",t)[1:end-2], ":", Libc.strftime("%z",t)[end-1:end])
    logstring = string("<", (UInt16(syslog.facility) << 3) + UInt16(level), ">1 ", timestamp, " ", syslog.machine, " ", syslog.user, " - - - ", level, ":", logger_name,":", msg...)
    write_log(syslog, color, logstring)
end

function log{T<:IO}(output::T, level::LogLevel, color::Symbol, logger_name::AbstractString, msg...)
    logstring = string(Libc.strftime("%d-%b %H:%M:%S",time()),":",level, ":",logger_name,":", msg...,"\n")
    write_log(output, color, logstring)
end

for (fn,lvl,clr) in ((:debug,    DEBUG,    :cyan),
                     (:info,     INFO,     :blue),
                     (:warn,     WARNING,  :magenta),
                     (:err,      ERROR,    :red),
                     (:critical, CRITICAL, :red))

    @eval function $fn(logger::Logger, msg...)
        if $lvl <= logger.level
            for output in logger.output
                log(output, $lvl, $(Expr(:quote, clr)), logger.name, msg...)
            end
        end
    end

    @eval $fn(msg...) = $fn(_root, msg...)

end

function configure(logger=_root; args...)
    for (tag, val) in args
        if tag == :parent
            logger.parent = parent = val::Logger
            logger.level = parent.level
            logger.output = parent.output
        end
    end

    for (tag, val) in args
        tag == :io            ? typeof(val) <: AbstractArray ? (logger.output = val) :
                                                               (logger.output = [val::LogOutput]) :
        tag == :output        ? typeof(val) <: AbstractArray ? (logger.output = val) :
                                                               (logger.output = [val::LogOutput]) :
        tag == :filename      ? (logger.output = [open(val, "a")]) :
        tag == :level         ? (logger.level  = val::LogLevel) :
        tag == :override_info ? nothing :  # handled below
        tag == :parent        ? nothing :  # handled above
                                (Base.error("Logging: unknown configure argument \"$tag\""))
    end

    logger
end

override_info(;args...) = (:override_info, true) in args

macro configure(args...)
    _args = gensym()
    quote
        logger = Logging.configure($(args...))
        if Logging.override_info($(args...))
            function Base.info(msg::AbstractString...)
                Logging.info(Logging._root, msg...)
            end
        end
        include(joinpath(Pkg.dir("Logging"), "src", "logging_macros.jl"))
        logger
    end
end

end # module