text:
"""
A package for handling lazily initialized fields.

### Exports:
* macros: `@lazy`, `@init!`, `@uninit!`, `@isinit`.
* functions: `init!` `uninit!`, `isinit`.
* objects: `uninit`.
* exceptions: `NonLazyFieldException`, `UninitializedFieldException`

### Example usage:

```julia-repl
julia> @lazy struct Foo{T}
           a::T
           @lazy b::Int
           @lazy c::Union{Float64, Nothing}
           @lazy d::Union{Int, Nothing}
           e::Float64
       end

julia> f = Foo(2, uninit, 2.0, nothing, 3.0)
Foo{Int64}(2, uninit, 2.0, nothing, 3.0)

julia> @isinit f.b
false

julia> @isinit f.c
true

julia> f.b
ERROR: field `b` in struct of type `Foo{Int64}` is not initialized
[...]

julia> @init! f.b = 4
4

julia> f.b
4

julia> @init! f.a=2
ERROR: field `a` in struct of type `Foo{Int64}` is not lazy, lazy fields are `b` ,`c` ,`d`.
[...]

julia> @uninit! f.b
uninit

julia> @isinit f.b
false
```

"""
module LazilyInitializedFields

export @lazy, uninit,
       @init!, @isinit, @uninit!,
        init!,  isinit,  uninit!,
        NonLazyFieldException, UninitializedFieldException

"""
    Uninitialized

A type with no fields that is the type of [`uninit`](@ref).
"""
struct Uninitialized end

"""
    uninit

The singleton instance of the type [`Uninitialized`](@ref), used
for fields that are currently uninitialized.
"""
const uninit = Uninitialized()
Base.show(io::IO, u::Uninitialized) = print(io, "uninit")

# The @lazy macro will extended this function
# for the struct getting defined so that we can use
# it to check if a field is lazy, for example:
# islazyfield(::Type{Foo}, s::Symbol) = s === :a || s === :b
function islazyfield end

struct NonLazyFieldException <: Exception
    T::DataType
    s::Symbol
end

Base.showerror(io::IO, err::NonLazyFieldException) =
    print(io, "field `$(err.s)` in struct of type `$(err.T)` is not lazy, lazy fields are ",
          join("`" .* string.(filter(x->islazyfield(err.T, x), fieldnames(err.T))) .* "`", ", "), ".")

struct UninitializedFieldException <: Exception
    T::DataType
    s::Symbol
end
Base.showerror(io::IO, err::UninitializedFieldException) =
    print(io, "field `", err.s, "` in struct of type `$(err.T)` is not initialized")

"""
    init!(a, s::Symbol)

Function version of [@init!](@ref).
"""
@inline function init!(x::T, s::Symbol, v) where {T}
    islazyfield(T, s) || throw(NonLazyFieldException(T, s))
    return setfield!(x, s, v)
end

_check_setproperty_expr(expr, s) =
    (expr isa Expr && expr.head === :(=) && expr.args[1].head === :.) || error("invalid usage of $s")
"""
    @init! a.b = v

Initialize the lazy field `b` in object `a` to `v`.
Throw a `NonLazyFieldException` if `b` is not a lazy field
of `a`.
Macro version of `init!(a, :b, v)`

```jldoctest
julia> using LazilyInitializedFields

julia> @lazy struct Foo
           @lazy b::Int
       end

julia> f = Foo(uninit)
Foo(uninit)

julia> f.b
ERROR: field `b` in struct of type `Foo` is not initialized
[...]

julia> @init! f.b = 3
3

julia> f.b
3
```
"""
macro init!(expr)
    _check_setproperty_expr(expr, "@init!")
    v = expr.args[2]
    body, sym = expr.args[1].args
    return :(init!($(esc(body)), $(esc(sym)), $(esc(v))))
end

_check_getproperty_expr(expr, s) =
    (expr isa Expr && expr.head === :.) || error("invalid usage of $s")

"""
    isinit(a, s::Symbol)

Function version of [@isinit](@ref).
"""
@inline function isinit(x::T, s) where {T}
    islazyfield(T, s) || throw(NonLazyFieldException(T, s))
    !(getfield(x, s) isa Uninitialized)
end
"""
    @isinit a.b

Check if the lazy field `b` in the object `a` is initialized.
Throw a `NonLazyFieldException` if `b` is not a lazy field
of `a`.
Macro version of [`isinit(a, :b)`](@ref)

```jldoctest
julia> using LazilyInitializedFields

julia> @lazy struct Foo
           @lazy b::Int
       end

julia> f = Foo(uninit)
Foo(uninit)

julia> @isinit f.b
false

julia> @init! f.b = 5
5

julia> @isinit f.b
true
```
"""
macro isinit(expr)
    _check_getproperty_expr(expr, "@isinit")
    return :(isinit($(esc.(expr.args)...)))
end

"""
    uninit!(a, s::Symbol)

Function version of [`@uninit!`](@ref).
"""

@inline function uninit!(x::T, s::Symbol) where {T}
    islazyfield(T, s) || throw(NonLazyFieldException(T, s))
    return setfield!(x, s, uninit)
end

"""
    @uninit! f.b

Uninitializes the field `b` in the object `f`
Throw a `NonLazyFieldException` if `b` is not a lazy field
of `a`.
Macro version of [`uninit`](@ref)

```jldoctest
julia> using LazilyInitializedFields

julia> @lazy struct Foo
           @lazy b::Int
       end

julia> f = Foo(uninit)
Foo(uninit)

julia> @isinit f.b
false

julia> @init! f.b = 5
5

julia> @isinit f.b
true
```
"""
macro uninit!(expr)
    _check_getproperty_expr(expr, "@uninit!")
    return :(uninit!($(esc.(expr.args)...)))
end

global in_lazy_struct
"""
    @lazy struct Foo
        a::Int
        @lazy b::Int
        @lazy c::Float64
    end

Make struct `Foo` have the lazy fields `b` and `c`.
"""
macro lazy(expr)
    if expr isa Expr && expr.head === :struct
        try
            global in_lazy_struct = true
            return lazy_struct(expr)
        finally
            global in_lazy_struct = false
        end
    elseif expr isa Expr && expr.head === :(::) && length(expr.args) == 2
        return lazy_field(expr)
    else
        _throw_invalid_usage()
    end
end

function lazy_field(expr)
    # expr is checked for correct form in @lazy
    in_lazy_struct || error("@lazy macro use outside of @lazy struct")
    name, T = expr.args
    :($(esc(name))::Union{$Uninitialized, $(esc(T))})
end

function lazy_struct(expr)
    mutable, structdef, body = expr.args
    structname = if structdef isa Symbol
        structdef
    elseif structdef isa Expr && structdef.head === :curly
        structdef.args[1]
    else
        error("internal error: unhandled expression $expr")
    end

    expr.args[1] = true # make mutable
    lazyfield = QuoteNode[]
    for (i, arg) in enumerate(body.args)
        if arg isa Expr && arg.head === :macrocall && arg.args[1] === Symbol("@lazy")
            body.args[i] = macroexpand(@__MODULE__, arg)
            name = body.args[i].args[1]
            @assert name isa Symbol
            push!(lazyfield, QuoteNode(name))
        end
    end

    if length(lazyfield) == 0
        error("expected a @lazy field inside the struct")
    end

    checks = foldr((a,b)->:(s === $a || $b), lazyfield[1:end-1]; init=:(s === $(lazyfield[end])))
    ret = Expr(:block)
    push!(ret.args, quote
        $(esc(expr))
        # is @pure overkill?
        Base.@pure $(LazilyInitializedFields).islazyfield(::Type{<:$(esc(structname))}, s::Symbol) = $checks
        function Base.getproperty(x::$(esc(structname)), s::Symbol)
            if $(LazilyInitializedFields).islazyfield($(esc(structname)), s)
                r = Base.getfield(x, s)
                r isa $Uninitialized && throw(UninitializedFieldException(typeof(x), s))
                return r
            end
            return Base.getfield(x, s)
        end
    end)
    if !mutable
        push!(ret.args, quote
            function Base.setproperty!(x::$(esc(structname)), s::Symbol, v)
                error("setproperty! for struct of type `", $(esc(structname)), "` has been disabled")
            end
        end)
    end
    return ret
end

end


text:
# NIfTI.jl
# Methods for reading NIfTI MRI files in Julia

module NIfTI

using CodecZlib, Mmap, MappedArrays, TranscodingStreams

import Base.getindex, Base.size, Base.ndims, Base.length, Base.write, Base64
export NIVolume, niread, niwrite, voxel_size, time_step, vox, getaffine, setaffine

include("parsers.jl")
include("extensions.jl")
include("volume.jl")

function define_packed(ty::DataType)
    packed_offsets = cumsum([sizeof(x) for x in ty.types])
    sz = pop!(packed_offsets)
    pushfirst!(packed_offsets, 0)

    @eval begin
        function Base.read(io::IO, ::Type{$ty})
            bytes = read!(io, Array{UInt8}(undef, $sz...))
            hdr = $(Expr(:new, ty, [:(unsafe_load(convert(Ptr{$(ty.types[i])}, pointer(bytes)+$(packed_offsets[i])))) for i = 1:length(packed_offsets)]...,))
            if hdr.sizeof_hdr == ntoh(Int32(348))
                return byteswap(hdr), true
            end
            hdr, false
        end
        function Base.write(io::IO, x::$ty)
            bytes = UInt8[]
            for name in fieldnames($ty)
                append!(bytes, reinterpret(UInt8, [getfield(x,name)]))
            end
            write(io, bytes)
            $sz
        end
    end
    nothing
end

mutable struct NIfTI1Header
    sizeof_hdr::Int32

    data_type::NTuple{10,UInt8}
    db_name::NTuple{18,UInt8}
    extents::Int32
    session_error::Int16
    regular::Int8

    dim_info::Int8
    dim::NTuple{8,Int16}
    intent_p1::Float32
    intent_p2::Float32
    intent_p3::Float32
    intent_code::Int16
    datatype::Int16
    bitpix::Int16
    slice_start::Int16
    pixdim::NTuple{8,Float32}
    vox_offset::Float32
    scl_slope::Float32
    scl_inter::Float32
    slice_end::Int16
    slice_code::Int8
    xyzt_units::Int8
    cal_max::Float32
    cal_min::Float32
    slice_duration::Float32
    toffset::Float32

    glmax::Int32
    glmin::Int32

    descrip::NTuple{80,UInt8}
    aux_file::NTuple{24,UInt8}

    qform_code::Int16
    sform_code::Int16
    quatern_b::Float32
    quatern_c::Float32
    quatern_d::Float32
    qoffset_x::Float32
    qoffset_y::Float32
    qoffset_z::Float32

    srow_x::NTuple{4,Float32}
    srow_y::NTuple{4,Float32}
    srow_z::NTuple{4,Float32}

    intent_name::NTuple{16,UInt8}

    magic::NTuple{4,UInt8}
end
define_packed(NIfTI1Header)

# byteswapping

function byteswap(hdr::NIfTI1Header)
    for fn in fieldnames(typeof(hdr))
        val = getfield(hdr, fn)
        if isa(val, Number) && sizeof(val) > 1
            setfield!(hdr, fn, ntoh(val))
        elseif isa(val, NTuple) && sizeof(eltype(val)) > 1
            setfield!(hdr, fn, map(ntoh, val))
        end
    end
    hdr
end

struct NIVolume{T<:Number,N,R} <: AbstractArray{T,N}
    header::NIfTI1Header
    extensions::Vector{NIfTIExtension}
    raw::R
end
NIVolume(header::NIfTI1Header, extensions::Vector{NIfTIExtension}, raw::R) where {R}=
    niupdate(new(header, extensions, raw))

NIVolume(header::NIfTI1Header, extensions::Vector{NIfTIExtension}, raw::AbstractArray{T,N}) where {T<:Number,N} =
    NIVolume{T,N,typeof(raw)}(header, extensions, raw)
NIVolume(header::NIfTI1Header, raw::AbstractArray{T,N}) where {T<:Number,N} =
    NIVolume{typeof(one(T)*1f0+1f0),N,typeof(raw)}(header, NIfTIExtension[], raw)
NIVolume(header::NIfTI1Header, extensions::Vector{NIfTIExtension}, raw::AbstractArray{Bool,N}) where {N} =
    NIVolume{Bool,N,typeof(raw)}(header, extensions, raw)
NIVolume(header::NIfTI1Header, raw::AbstractArray{Bool,N}) where {N} =
    NIVolume{Bool,N,typeof(raw)}(header, NIfTIExtension[], raw)

include("coordinates.jl")

# Always in mm
voxel_size(header::NIfTI1Header) =
    [header.pixdim[i] * SPATIAL_UNIT_MULTIPLIERS[header.xyzt_units & Int8(3)] for i = 2:min(header.dim[1], 3)+1]

# Always in ms
time_step(header::NIfTI1Header) =
    header.pixdim[5] * TIME_UNIT_MULTIPLIERS[header.xyzt_units >> 3]

function to_dim_info(dim_info::Tuple{Integer,Integer,Integer})
    if dim_info[1] > 3 || dim_info[1] < 0
        error("Invalid frequency dimension $(dim_info[1])")
    elseif dim_info[2] > 3 || dim_info[2] < 0
        error("Invalid phase dimension $(dim_info[2])")
    elseif dim_info[3] > 3 || dim_info[3] < 0
        error("Invalid slice dimension $(dim_info[3])")
    end

    return Int8(dim_info[1] | (dim_info[2] << 2) | (dim_info[3] << 4))
end

# Returns or sets dim_info as a tuple whose values are the frequency, phase, and slice dimensions
function dim_info(header::NIfTI1Header)
    return (
        header.dim_info & int8(3),
        (header.dim_info >> 2) & int8(3),
        (header.dim_info >> 4) & int8(3)
    )
end
function dim_info(header::NIfTI1Header, dim_info::Tuple{T, T, T}) where {T<:Integer}
    header.dim_info = to_dim_info(dim_info)
end

# Gets the size of a type in bits
nibitpix(t::Type) = Int16(sizeof(t)*8)
nibitpix(::Type{Bool}) = Int16(1)

# Constructor
function NIVolume(
    # Optional MRI volume; if not given, an empty volume is used
    raw::AbstractArray{T}=Int16[],
    extensions::Union{Vector{NIfTIExtension},Nothing}=nothing;

    # Fields specified as UNUSED in NIfTI1 spec
    data_type::AbstractString="", db_name::AbstractString="", extents::Integer=Int32(0),
    session_error::Integer=Int16(0), regular::Integer=Int8(0), glmax::Integer=Int32(0),
    glmin::Integer=Int16(0),

    # The frequency encoding, phase encoding and slice dimensions.
    dim_info::NTuple{3, Integer}=(0, 0, 0),
    # Describes data contained in the file; for valid values, see
    # http://nifti.nimh.nih.gov/nifti-1/documentation/nifti1fields/nifti1fields_pages/group__NIfTI1__INTENT__CODES.html
    intent_p1::Real=0f0, intent_p2::Real=0f0, intent_p3::Real=0f0,
    intent_code::Integer=Int16(0), intent_name::AbstractString="",
    # Information about which slices were acquired, in case the volume has been padded
    slice_start::Integer=Int16(0), slice_end::Integer=Int16(0), slice_code::Int8=Int8(0),
    # The size of each voxel and the time step. These are formulated in mm unless xyzt_units is
    # explicitly specified
    voxel_size::NTuple{3, Real}=(1f0, 1f0, 1f0), time_step::Real=0f0, xyzt_units::Int8=Int8(18),
    # Slope and intercept by which volume shoudl be scaled
    scl_slope::Real=1f0, scl_inter::Real=0f0,
    # These describe how data should be scaled when displayed on the screen. They are probably
    # rarely used
    cal_max::Real=0f0, cal_min::Real=0f0,
    # The amount of time taken to acquire a slice
    slice_duration::Real=0f0,
    # Indicates a non-zero start point for time axis
    toffset::Real=0f0,

    # "Any text you like"
    descrip::AbstractString="",
    # Name of auxiliary file
    aux_file::AbstractString="",

    # Parameters for Method 2. See the NIfTI spec
    qfac::Float32=0f0, quatern_b::Real=0f0, quatern_c::Real=0f0, quatern_d::Real=0f0,
    qoffset_x::Real=0f0, qoffset_y::Real=0f0, qoffset_z::Real=0f0,
    # Orientation matrix for Method 3
    orientation::Union{Matrix{Float32},Nothing}=nothing) where {T <: Number}
    local t
    if isempty(raw)
        raw = Int16[]
        t = Int16
    else
        t = T
    end

    if extensions == nothing
        extensions = NIfTIExtension[]
    end

    method2 = qfac != 0 || quatern_b != 0 || quatern_c != 0 || quatern_d != 0 || qoffset_x != 0 ||
        qoffset_y != 0 || qoffset_z != 0
    method3 = orientation != nothing

    if method2 && method3
        error("Orientation parameters for Method 2 and Method 3 are mutually exclusive")
    end

    if method3
        if size(orientation) != (3, 4)
            error("Orientation matrix must be of dimensions (3, 4)")
        end
    else
        orientation = zeros(Float32, 3, 4)
    end

    if slice_start == 0 && slice_end == 0 && dim_info[3] != 0
        slice_start = 0
        slice_end = size(raw, dim_info[3]) - 1
    end

    NIVolume(NIfTI1Header(SIZEOF_HDR1, string_tuple(data_type, 10), string_tuple(db_name, 18), extents, session_error,
        regular, to_dim_info(dim_info), to_dim_i16(size(raw)), intent_p1, intent_p2,
        intent_p3, intent_code, eltype_to_int16(t), nibitpix(t),
        slice_start, (qfac, voxel_size..., time_step, 0, 0, 0), 352,
        scl_slope, scl_inter, slice_end, slice_code,
        xyzt_units, cal_max, cal_min, slice_duration,
        toffset, glmax, glmin, string_tuple(descrip, 80), string_tuple(aux_file, 24), (method2 || method3),
        method3, quatern_b, quatern_c, quatern_d,
        qoffset_x, qoffset_y, qoffset_z, (orientation[1, :]...,),
        (orientation[2, :]...,), (orientation[3, :]...,), string_tuple(intent_name, 16), NP1_MAGIC), extensions, raw)
end

# Validates the header of a volume and updates it to match the volume's contents
function niupdate(vol::NIVolume{T}) where {T}
    vol.header.sizeof_hdr = SIZEOF_HDR1
    vol.header.dim = to_dim_i16(size(vol.raw))
    vol.header.datatype = eltype_to_int16(T)
    vol.header.bitpix = nibitpix(T)
    vol.header.vox_offset = isempty(vol.extensions) ? Int32(352) :
        Int32(mapreduce(esize, +, vol.extensions) + SIZEOF_HDR1)
    vol
end

# Avoid method ambiguity
write(io::Base64.Base64EncodePipe, vol::NIVolume{UInt8,1}) = invoke(write, (IO, NIVolume{UInt8,1}), io, vol)

# Write a NIfTI file
function write(io::IO, vol::NIVolume)
    write(io, niupdate(vol).header)
    if isempty(vol.extensions)
        write(io, Int32(0))
    else
        for ex in vol.extensions
            sz = esize(ex)
            write(io, Int32(sz))
            write(io, Int32(ex.ecode))
            write(io, ex.edata)
            write(io, zeros(UInt8, sz - length(ex.edata)))
        end
    end
    if eltype(vol.raw) == Bool
        write(io, BitArray(vol.raw))
    else
        write(io, vol.raw)
    end
end

# Convenience function to write a NIfTI file given a path
function niwrite(path::AbstractString, vol::NIVolume)
    if split(path,".")[end] == "gz"
        io = open(path, "w")
        stream = GzipCompressorStream(io)
        write(stream, vol)
        close(stream)
        close(io)
    else
        io = open(path, "w")
        write(io, vol)
        close(io)
    end
end

# Read header from a NIfTI file
function read_header(io::IO)
    header, swapped = read(io, NIfTI1Header)
    if header.sizeof_hdr != SIZEOF_HDR1
        error("This is not a NIfTI-1 file")
    end
    header, swapped
end

# Look for a gzip header in an IOStream
function isgz(io::IO)
    try
        ret = read(io, UInt8) == 0x1F && read(io, UInt8) == 0x8B
        seekstart(io)
        ret
    catch err
        if isa(err, EOFError)
            @debug "reading the file resulted in an EOF error and \nthe end of the file was read.\nNo more data was available to read from the filestream.\nIt is likely that the file was corrupted or is empty (0 bytes)."
            rethrow(err)
        end
    end
end

function niread(file::AbstractString; mmap::Bool=false, mode::AbstractString="r")
    io = niopen(file, mode)
    hdr, swapped = read_header(io)
    ex = read_extensions(io, hdr.vox_offset - 352)

    if hdr.magic === NP1_MAGIC
        vol = read_volume(io, to_eltype(hdr.datatype), to_dimensions(hdr.dim), mmap)
    else
        vol = read_volume(niopen(hdr_to_img(file), mode), to_eltype(hdr.datatype), to_dimensions(hdr.dim), mmap)
    end

    if swapped && sizeof(eltype(vol)) > 1
        return NIVolume(hdr, ex, mappedarray(ntoh, hton, vol))
    else
        return NIVolume(hdr, ex, vol)
    end
end

# Allow file to be indexed like an array, but with indices yielding scaled data
@inline getindex(f::NIVolume{T}, idx::Vararg{Int}) where {T} =
    getindex(f.raw, idx...,) * f.header.scl_slope + f.header.scl_inter

add1(x::Union{AbstractArray{T},T}) where {T<:Integer} = x + 1
add1(::Colon) = Colon()
@inline vox(f::NIVolume, args...,) = getindex(f, map(add1, args)...,)
size(f::NIVolume) = size(f.raw)
size(f::NIVolume, d) = size(f.raw, d)
ndims(f::NIVolume) = ndims(f.raw)
length(f::NIVolume) = length(f.raw)
lastindex(f::NIVolume) = lastindex(f.raw)

end

ext: jl


text:
#=
This script spikes a variant into a vcf file for testing
with VVP and VAAST3.

Use a bedfile to input the variant or multiple variants.

Usage: SpikeVCF.jl myVCFfile.vcf myBedfile.bed id_index

Written: 20180510 WSW GPL
=#

if length(ARGS) < 3
    print("\n\nUsage: spikeVCF.jl myVCFfile.vcf myBedfile.bed <id_idx>

    Where the bed file contains one variant in a tab-delimited line:

    7     150649881   G    A    R    C    KCNH2   ENST00000430723

    and id_indices are the 1-based vcf sample indices to have a
    0/1 genotype (e.g. 2,5,9).  Output is to STDOUT.\n\n")
    exit()
end

vcf = ARGS[1]
bed = ARGS[2]

s=read(`grep '^#C' "$vcf"`, String)
ss=split(s, '\t')
num=length(ss) - 9    #number of samples

#print("Processing VCF file: $vcf\n","Found $num samples in vcf\n", "Spiking variant in $bed\n")

genos=fill("0/0", num)
mutidxs=split(ARGS[3], ',')
ac = length(mutidxs)

for i in 1:length(mutidxs)
  smutidx=popfirst!(mutidxs)
  mutidx=parse(Int,smutidx)
  genos[mutidx] = "0/1"
end

sgenos = join(genos, '\t')

b = []
open(ARGS[2]) do bed
  for i in eachline(bed)
    if occursin(r"^[X|Y|M]", i)
      println("Sorry, only autosomes are allowed in this script.")
      exit()
    end
    push!(b,i)
  end
end

bb = split(b[1], '\t')
cp = bb[1] * ":" *  bb[2]
alt=bb[4]
gn=bb[7]
tid=bb[8]
aa = bb[5] * "/" * bb[6]
tidc=chomp(tid)
CSQ="AC=$ac;CSQ=$alt|||$gn|||$tidc|||||||||$aa||||||"

spike = join([bb[1], bb[2], cp, bb[3], alt, "9999.99", "PASS", "$CSQ", "GT", sgenos], '\t')

chrbed = parse(Int,bb[1])
posbed = parse(Int,bb[2])

open(ARGS[1]) do vcf    #read bedfile and make lookup and set

  k=0

  for i in eachline(vcf)

    if occursin(r"^\#", i)
      println(i)
      continue
    end

    if occursin(r"^X", i)
      println(i)
      continue
    end

    if occursin(r"^Y", i)
      println(i)
      continue
    end

    a = split(i, '\t')
    chr = parse(Int,a[1])

    if chrbed != chr
      println(i)
      continue
    end

    if chrbed == chr
      pos = parse(Int,a[2])

      if pos < posbed
         println(i)
      end

      if posbed < pos
         k = k + 1;
         if k == 1
            println(spike)
         else
            println(i)
         end
      end
    end
  end
end

max_stars_repo_name: ScottWatkins/Misc

max_issues_repo_name: ScottWatkins/Misc

max_forks_repo_name: ScottWatkins/Misc


text:
# SPDX-License-Identifier: X11
# 2020-08-29
# Five Variables (100pt)

println(findfirst(x -> x == 0, parse.(Int, split(readline(), " "))))

max_stars_repo_name: jthvai/competitive-programming

max_issues_repo_name: jthvai/competitive-programming

max_forks_repo_name: jthvai/competitive-programming


text:
function move_downloads()

    dl_dir = "/Users/mzgubic/Downloads/"
    for fname in readdir(dl_dir) # TODO: change to ~

        # order matters
        if occursin("Income", fname)
            println(fname)
            mv(dl_dir*fname, "data/income/"*fname, force=true)

        elseif occursin("Net", fname)
            println(fname)
            mv(dl_dir*fname, "data/networth/"*fname, force=true)

        elseif occursin("Expenses", fname)
            println(fname)
            mv(dl_dir*fname, "data/expense/"*fname, force=true)

        elseif occursin("exchange", fname)
            println(fname)
            mv(dl_dir*fname, "data/fx/"*fname, force=true)
        end

    end

end

function is_individual(fpath)::Bool
    return !occursin("Repeated", fpath)
end

function step(row)
    if row.frequency == "DAILY"
        return Dates.Day(1)
    elseif row.frequency == "MONTHLY"
        return Dates.Month(1)
    else
        throw("Error")
    end
end

function load_repeated(fpath::String)::DataFrame

    # repeated df
    rep_df = DataFrame(CSV.File(fpath))

    # individual df
    df = DataFrame(date=Date[],
                   amount=Float64[],
                   currency=String[],
                   category=String[],
                   description=String[])

    # loop over repeated and extract individuals
    for row in eachrow(rep_df)
        datestep = step(row)
        daterange = range(row.start, row.end, step=datestep)
        amount = row.amount / length(daterange)
        for date in daterange
            push!(df, [date, amount, row.currency, row.category, row.description])
        end
    end

    return df

end

function load_expenses(fpath::String)::DataFrame

    # treat individual and repeated expenses separately
    println(fpath)
    if is_individual(fpath)
        return DataFrame(CSV.File(fpath))
    else
        return load_repeated(fpath)
    end

end

function load_expenses()::DataFrame
    expense_dir = "data/expense/"
    dfs = [load_expenses(expense_dir*fname) for fname in readdir(expense_dir)]
    reverse!(dfs)
    return vcat(dfs...)
end

function load_incomes()::DataFrame
    income_dir = "data/income/"
    return DataFrame(CSV.File(income_dir*"Expenses - Income.csv"))
end

function load_net()::DataFrame
    net_dir = "data/networth/"
    return DataFrame(CSV.File(net_dir*"Expenses - NetWorth.csv", header=2))
end

function parse_from_to(fpath::String)::Tuple{String, String}

    fname = split(fpath, "/")[end]
    currencies = split(fname, "-exchange-rate")[1]

    first = split(currencies, "-")[1]
    second = split(currencies, "-")[2]

    d = Dict("euro"=>"EUR",
             "british"=>"GBP",
             "dollar"=>"USD")

    return d[first], d[second]
end

function load_fxs(fpath::String)::DataFrame

    # load the raw file
    header = 16 # ignore the first 15 lines in the file
    df = DataFrame(CSV.File(fpath, header=header))
    rename!(df, Dict(Symbol(" value")=>:value))

    # ignore old exchange rates
    old = Date(2015, 01, 01) # make this optional input (supply from expenses/income data)
    df = df[df[:, "date"] .> old, :]

    # add from to info
    from, to = parse_from_to(fpath)
    insertcols!(df, 3, :to_curr=>repeat([to], size(df)[1]))
    insertcols!(df, 3, :from_curr=>repeat([from], size(df)[1]))

    # add a df with reverse exchange rate as well
    df_reverse = copy(df)
    df_reverse[:, "value"] = 1 ./ df[:, "value"]
    rename!(df_reverse, Dict(:to_curr=>:from_curr, :from_curr=>:to_curr))

    return vcat(df, df_reverse)

end

function load_fxs()::DataFrame
    fx_dir = "data/fx/"
    dfs = [load_fxs(fx_dir*fname) for fname in readdir(fx_dir)]
    return vcat(dfs...)
end

function xchange_rate(from_curr::String, to_curr::String, date::Date, fx::DataFrame)::Float64

    # get the subset of rates from the particular currency combination
    xrates = @from i in fx begin
        @where (i.from_curr == from_curr) && (i.to_curr == to_curr)
        @select i
        @collect DataFrame
    end

    # if currency pairs do not exist
    if size(xrates)[1] == 0
        println("WARNING: No exchange rate between $(from_curr) and $(to_curr), using 1.")
        return 1.0
    end

    # try getting an exact rate
    exact_rate = xrates[xrates.date .== date, "value"]
    if length(exact_rate) == 1
        return exact_rate[1]

    # otherwise get an estimate from the closest available day
    else
        days = [abs(Dates.value(d)) for d in xrates.date-date]
        return xrates.value[argmin(days)]
    end

end

function exchange!(exp_df::DataFrame, to_curr::String, fx::DataFrame)::DataFrame

    # loop over all expenses
    for exp in eachrow(exp_df)

        # and convert the spending if needed
        if exp.currency != to_curr

            # determine the xchange rate
            xrate = xchange_rate(exp.currency, to_curr, exp.date, fx)

            # change the amount and the currency
            exp.amount = exp.amount * xrate
            exp.currency = to_curr

        end

    end

    return exp_df
end

max_stars_repo_name: mzgubic/exchequer

max_issues_repo_name: mzgubic/exchequer

max_forks_repo_name: mzgubic/exchequer


text:
function generictest_dixon_Fp(G, p = SymbolicWedderburn.dixon_prime(G))
    F = SymbolicWedderburn.FiniteFields.GF{p}
    ccG = conjugacy_classes(G)
    Ns = [SymbolicWedderburn.CMMatrix(ccG, i) for i = 1:length(ccG)]
    @test isdiag(SymbolicWedderburn.common_esd(Ns, F))
    esd = SymbolicWedderburn.common_esd(Ns, F)

    let W = esd.basis
        @test inv(W) isa Matrix{F}
        @test isone(inv(W) * W)
        iW = inv(W)
        for N in Ns
            # basis actually diagonalizes all Ns
            m = W * N * iW
            @test isdiag(m)
        end
    end

    @test SymbolicWedderburn.characters_dixon(ccG, F) isa
          Vector{<:SymbolicWedderburn.Character}

    chars = SymbolicWedderburn.characters_dixon(ccG, F)

    # checking the degrees
    @test sum(degree.(chars) .^ 2) == order(G)

    # orthogonality of characters over F
    @test [dot(χ, ψ) for χ in chars, ψ in chars] == I
end

function generictest_dixon_C(G, p = SymbolicWedderburn.dixon_prime(G))
    F = SymbolicWedderburn.FiniteFields.GF{p}
    ccG = conjugacy_classes(G)
    chars_Fp = SymbolicWedderburn.characters_dixon(ccG, F)

    degrees = degree.(chars_Fp)

    m = SymbolicWedderburn._multiplicities(chars_Fp)
    for i = 1:size(m, 1)
        @test all(Int.(m[i, :, :]) .<= degrees[i])
    end

    @test SymbolicWedderburn.complex_characters(chars_Fp) isa
          Vector{<:SymbolicWedderburn.Character}
    @test length(SymbolicWedderburn.complex_characters(chars_Fp[1:1])) == 1

    chars_CC = SymbolicWedderburn.complex_characters(chars_Fp)
    id = one(first(first(ccG)))
    @test [ψ(id) for ψ in chars_CC] == degrees

    # orthogonality of characters over ℂ
    @test [dot(χ, ψ) for χ in chars_CC, ψ in chars_CC] == I
end

@testset "Dixon Algorithm" begin
    @testset "DixonPrimes" begin
        @test SymbolicWedderburn.dixon_prime(20, 20) == 41

        @testset "random" begin
            for n in rand(2:1000, 10)
                F = SymbolicWedderburn.Primes.factor(n)
                e = lcm(collect(keys(F)))
                p = SymbolicWedderburn.dixon_prime(n, e)
                @test mod(p, e) == 1
                @test p > 2 * floor(sqrt(n))
            end
        end

        @testset "DixonPrimeGroups" begin
            G = PermutationGroups.SymmetricGroup(4)
            ccG = conjugacy_classes(G)
            @test exponent(G) == 12
            @test exponent(ccG) == 12
            @test SymbolicWedderburn.dixon_prime(G) ==
                  SymbolicWedderburn.dixon_prime(ccG)
        end
    end

    @testset "Dixon over GF{p}" begin
        @testset "Example: Alt(4)" begin
            G = PermGroup([perm"(1,2,3)(4)", perm"(2,3,4)"])
            ccG = let S = gens(G)
                [
                    Orbit(S, one(G)),
                    Orbit(S, perm"(2,3,4)"),
                    Orbit(S, perm"(2,4,3)"),
                    Orbit(S, perm"(1,2)(3,4)"),
                ] # the order is taken from GAP
            end

            let ccG = ccG, p = 29
                F = SymbolicWedderburn.FiniteFields.GF{p}
                Ns = [SymbolicWedderburn.CMMatrix(ccG, i) for i = 1:length(ccG)]
                @test_throws AssertionError SymbolicWedderburn.common_esd(Ns, F)
            end

            let ccG = ccG, p = 31
                F = SymbolicWedderburn.FiniteFields.GF{p}
                Ns = [SymbolicWedderburn.CMMatrix(ccG, i) for i = 1:length(ccG)]
                esd = SymbolicWedderburn.EigenSpaceDecomposition(F.(Ns[1]))
                esd = SymbolicWedderburn.refine(esd, F.(Ns[2]))
                esd = SymbolicWedderburn.refine(esd, F.(Ns[3]))
                @test isdiag(esd)

                @test isdiag(SymbolicWedderburn.common_esd(Ns, F))
            end

            # generic tests
            generictest_dixon_Fp(G)

            # tests specific to G
            p = SymbolicWedderburn.dixon_prime(ccG)
            F = SymbolicWedderburn.FiniteFields.GF{p}

            chars = SymbolicWedderburn.characters_dixon(ccG, F)

            @test sort(degree.(chars)) == [1, 1, 1, 3]

            @test Set(Int.(χ.vals) for χ in chars) ==
                  Set([[3, 0, 0, 6], [1, 4, 2, 1], [1, 2, 4, 1], [1, 1, 1, 1]])

        end
    end

    @testset "Lifting characters to ℂ" begin
        @testset "Example: Alt(4)" begin
            G = PermGroup([perm"(1,2,3)(4)", perm"(2,3,4)"])
            ccG = let S = gens(G)
                [
                    Orbit(S, one(G)),
                    Orbit(S, perm"(2,3,4)"),
                    Orbit(S, perm"(2,4,3)"),
                    Orbit(S, perm"(1,2)(3,4)"),
                ] # the order of cclasses is taken from GAP
            end

            @testset "PowerMap" begin
                pmG = SymbolicWedderburn.PowerMap(ccG)
                @test size(pmG) == (4, 6)
                @test pmG[:, 0] == ones(Int, 4)
                @test pmG[:, 1] == 1:4
                @test pmG[:, 2] == [1, 3, 2, 1]
                @test pmG[:, 3] == [1, 1, 1, 4]
                @test pmG[:, 4] == [1, 2, 3, 1]
                @test pmG[:, 5] == [1, 3, 2, 4]
                @test all(isone.(first.(ccG) .^ 6))
            end

            # generic tests
            generictest_dixon_C(G)

            chars_C = SymbolicWedderburn.characters_dixon(G)
            E = SymbolicWedderburn.Cyclotomics.E

            @test [χ.vals for χ in chars_C] == [
                E(3, 0) .* [3, 0, 0, -1],
                E(3, 0) .* [1, E(3, 2), E(3, 1), 1],
                E(3, 0) .* [1, E(3, 1), E(3, 2), 1],
                E(3, 0) .* [1, 1, 1, 1],
            ]
        end

        @testset "Example: Sym(4)" begin
            G = PermGroup([perm"(1,2)", perm"(1,2,3,4)"])
            ccG = let S = gens(G)
                ccG = [
                    Orbit(S, one(G)),
                    Orbit(S, perm"(1,2)(4)"),
                    Orbit(S, perm"(1,2)(3,4)"),
                    Orbit(S, perm"(1,2,3)(4)"),
                    Orbit(S, perm"(1,2,3,4)"),
                ] # the order of cclasses is taken from GAP
            end
            generictest_dixon_Fp(G)
            generictest_dixon_C(G)

            chars = SymbolicWedderburn.characters_dixon(ccG)

            @test sort(degree.(chars)) == [1, 1, 2, 3, 3]
            @test [χ.vals for χ in chars] == [
                [2, 0, 2, -1, 0],
                [3, 1, -1, 0, -1],
                [1, 1, 1, 1, 1],
                [1, -1, 1, 1, -1],
                [3, -1, -1, 0, 1],
            ]
        end

        @testset "example: C₅ ⋊ C₄" begin
            S = [perm"(1,2,4,5,3)", perm"(2,5,3,4)"]
            G = PermGroup(S)

            ccG = [
                Orbit(S, one(G)),
                Orbit(S, perm"(2,3)(4,5)"),
                Orbit(S, perm"(2,4,3,5)"),
                Orbit(S, perm"(2,5,3,4)"),
                Orbit(S, perm"(1,2,4,5,3)"),
            ]
            # the order of cclasses is taken from GAP
            generictest_dixon_Fp(G)
            generictest_dixon_C(G)

            chars = SymbolicWedderburn.characters_dixon(ccG)

            @test sort(degree.(chars)) == [1, 1, 1, 1, 4]
            @test [χ.vals for χ in chars] == [
                E(4, 0) .* [4, 0, 0, 0, -1],
                E(4, 0) .* [1, 1, 1, 1, 1],
                E(4, 0) .* [1, 1, -1, -1, 1],
                E(4, 0) .* [1, -1, E(4), -E(4), 1],
                E(4, 0) .* [1, -1, -E(4), E(4), 1],
            ]
        end
    end

    @testset "SmallPermGroups" begin
        include("smallgroups.jl")
        for (ord, groups) in SmallPermGroups
            @testset "SmallGroup($ord, $n)" for (n, G) in enumerate(groups)
                @test SymbolicWedderburn.characters_dixon(G) isa
                      Vector{<:SymbolicWedderburn.Character}
                generictest_dixon_Fp(G)
                generictest_dixon_C(G)
            end
        end
    end
end

@testset "Characters io" begin
    G = PermGroup([perm"(2,3)(4,5)"])
    chars = SymbolicWedderburn.characters_dixon(G)

    @test sprint(show, chars[1]) == "SymbolicWedderburn.Character: [1, 1]"
    @test sprint(show, chars[2]) == "SymbolicWedderburn.Character: [1, -1]"

    @test sprint(show, MIME"text/plain"(), chars[1]) ==
        "SymbolicWedderburn.Character over Cyclotomic{Int64,SparseArrays.SparseVector{Int64,Int64}}\n()^G         → \t 1*E(1)^0\n(2,3)(4,5)^G → \t 1*E(1)^0"
    @test sprint(show, MIME"text/plain"(), chars[2]) ==
        "SymbolicWedderburn.Character over Cyclotomic{Int64,SparseArrays.SparseVector{Int64,Int64}}\n()^G         → \t 1*E(1)^0\n(2,3)(4,5)^G → \t-1*E(1)^0"
end


text:
using Compat
using PDMats
using Distributions

if VERSION >= v"0.7"
  using Statistics
  using Distributed
  import Statistics: mean, median, maximum, minimum, quantile, std, var, cov, cor
else
  using Compat.Statistics
  using Compat.Distributed
  import Base: mean, median, maximum, minimum, quantile, std, var, cov, cor
end

#@compat abstract type GaussianMixtureModelCommonCovarAbstract <: Distribution end
@compat abstract type GaussianMixtureModelCommonCovarAbstract  <: Distribution{Multivariate,Continuous} end

struct GaussianMixtureModelCommonCovar <: GaussianMixtureModelCommonCovarAbstract
	mu::Array{Float64,2}
	probs::Vector{Float64}
        covar::AbstractPDMat
        aliastable::Distributions.AliasTable

    function GaussianMixtureModelCommonCovar(m::Array{Float64,2}, p::Vector{Float64}, ic::AbstractPDMat)
        if size(m,2) != length(p)
            error("means and probs must have the same number of elements")
        end
		if( size(ic,1) != size(ic,2) )
		    error("covariance matrix must be square")
		end
		if( size(m,1) != size(ic,1) )
		    error("means and covar matrix not compatible sizes: ",size(m)," vs ",size(ic) )
		end
        sump = 0.0
        for i in 1:length(p)
            if p[i] < 0.0
                error("MixtureModel: probabilities must be non-negative")
            end
            sump += p[i]
        end
        table = Distributions.AliasTable(p)
        new(m, p ./ sump, ic, table)
    end

end

function GaussianMixtureModelCommonCovar(m::Array{Float64,2}, p::Vector{Float64}, ic::AbstractMatrix)
    GaussianMixtureModelCommonCovar(m,p,make_matrix_pd(ic))
end

function mean(d::GaussianMixtureModelCommonCovarAbstract)
    np = size(d.mu,2)
    m = zeros(np)
    for i in 1:length(d.probs)
        m += vec(d.mu[:,i]) .* d.probs[i]
    end
    return m
end

function pdf(d::GaussianMixtureModelCommonCovar, x::Array{Float64,1} )
    p = 0.0
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        p += Distributions.pdf(Distributions.MvNormal(d.covar), x .- vec(d.mu[:,i]) ) * d.probs[i]
    end
    return p
end

function pdf(d::GaussianMixtureModelCommonCovar, x::Array{Float64,2} )
    np = size(x,2)
    p = zeros(np)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        p += Distributions.pdf(Distributions.MvNormal(d.covar), x .- vec(d.mu[:,i]) ) * d.probs[i]
    end
    return p
end

function logpdf(d::GaussianMixtureModelCommonCovar, x::Array{Float64,1})
    logp = -Inf
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        logp_i = Distributions.logpdf(Distributions.MvNormal(d.covar), x .- vec(d.mu[:,i]) ) + log(d.probs[i])
        logp = logp > logp_i ? logp + log1p(exp(logp_i-logp)) : logp_i + log1p(exp(logp-logp_i))
    end
    return logp
end

function logpdf(d::GaussianMixtureModelCommonCovar, x::Array{Float64,2})
    np = size(x,2)
    logp = fill(-Inf,np)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        logp_i = Distributions.logpdf(Distributions.MvNormal(d.covar), x .- vec(d.mu[:,i]) ) + log(d.probs[i])
        for j in 1:length(logp)
          logp[j] = logp[j] > logp_i[j] ? logp[j] + log1p(exp(logp_i[j]-logp[j])) : logp_i[j] + log1p(exp(logp[j]-logp_i[j]))
        end
    end
    return logp
end

function rand(d::GaussianMixtureModelCommonCovar)
    i = Distributions.rand(d.aliastable)
    return  Distributions.rand(Distributions.MvNormal(vec(d.mu[:,i]),d.covar))
end

struct GaussianMixtureModelCommonCovarTruncated <: GaussianMixtureModelCommonCovarAbstract
	mu::Array{Float64,2}
	probs::Vector{Float64}
        covar::AbstractPDMat
        aliastable::Distributions.AliasTable
        max_mahalanobis::Float64

    function GaussianMixtureModelCommonCovarTruncated(m::Array{Float64,2}, p::Vector{Float64}, ic::AbstractPDMat, mm::Float64)
        if size(m,2) != length(p)
            error("means and probs must have the same number of elements")
        end
		if( size(ic,1) != size(ic,2) )
		    error("covariance matrix must be square")
		end
		if( size(m,1) != size(ic,1) )
		    error("means and covar matrix not compatible sizes: ",size(m)," vs ",size(ic) )
		end
        sump = 0.0
        for i in 1:length(p)
            if p[i] < 0.0
                error("MixtureModel: probabilities must be non-negative")
            end
            sump += p[i]
        end
        table = Distributions.AliasTable(p)
        new(m, p ./ sump, ic, table, mm)
    end
end

function GaussianMixtureModelCommonCovarTruncated(m::Array{Float64,2}, p::Vector{Float64}, ic::AbstractMatrix, mm::Float64)
    GaussianMixtureModelCommonCovarTruncated(m, p, make_matrix_pd(ic), mm)
end

function pdf(d::GaussianMixtureModelCommonCovarTruncated, x::Array{Float64,1} )
    p = 0.0
    normalization = cdf(Distributions.Chisq(length(x)),d.max_mahalanobis)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        di = Distributions.MvNormal(vec(d.mu[:,i]),d.covar)
        if sqmahal(di, x) > d.max_mahalanobis continue end
        p += Distributions.pdf(di, x ) * d.probs[i] / normalization
    end
    return p
end

function pdf(d::GaussianMixtureModelCommonCovarTruncated, x::Array{Float64,2} )
    np = size(x,2)
    p = zeros(np)
    normalization = cdf(Distributions.Chisq(length(x)),d.max_mahalanobis)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        di = Distributions.MvNormal(vec(d.mu[:,i]),d.covar)
        p += sqmahal(di, x) .> d.max_mahalanobis ? Distributions.pdf(di, x ) * d.probs[i] / normalization : 0.0
    end
    return p
end

function logpdf(d::GaussianMixtureModelCommonCovarTruncated, x::Array{Float64,1})
    logp = -Inf
    log_normalization = logcdf(Distributions.Chisq(length(x)),d.max_mahalanobis)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        di = Distributions.MvNormal(vec(d.mu[:,i]),d.covar)
        if sqmahal(di, x) > d.max_mahalanobis continue end
        logp_i = Distributions.logpdf(di, x) + log(d.probs[i]) - log_normalization
        logp = logp > logp_i ? logp + log1p(exp(logp_i-logp)) : logp_i + log1p(exp(logp-logp_i))
    end
    return logp
end

function logpdf(d::GaussianMixtureModelCommonCovarTruncated, x::Array{Float64,2})
    np = size(x,2)
    logp = fill(-Inf,np)
    log_normalization = logcdf(Distributions.Chisq(length(x)),d.max_mahalanobis)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        di = Distributions.MvNormal(vec(d.mu[:,i]),d.covar)
        logp_i = Distributions.logpdf(di, x) + log(d.probs[i]) - log_normalization
        for j in 1:length(logp)
          if sqmahal(di, x[:,j]) > d.max_mahalanobis continue end
          logp[j] = logp[j] > logp_i[j] ? logp[j] + log1p(exp(logp_i[j]-logp[j])) : logp_i[j] + log1p(exp(logp[j]-logp_i[j]))
        end
    end
    return logp
end

function rand(d::GaussianMixtureModelCommonCovarTruncated)
    @assert d.max_mahalanobis > 0.1
    i = Distributions.rand(d.aliastable)
    di = Distributions.MvNormal(vec(d.mu[:,i]),d.covar)
    local rv
    dist = Inf
    while dist > d.max_mahalanobis
      rv = Distributions.rand(di)
      dist = sqmahal(di, rv)
    end
    return rv
end

struct GaussianMixtureModelCommonCovarSubset <: GaussianMixtureModelCommonCovarAbstract
	mu::Array{Float64,2}
	probs::Vector{Float64}
        covar::AbstractPDMat
        aliastable::Distributions.AliasTable
        param_active::Vector{Int64}

     function GaussianMixtureModelCommonCovarSubset(m::Array{Float64,2}, p::Vector{Float64}, ic::AbstractArray{Float64,2}, pact::Vector{Int64} )
        if size(m,2) != length(p)
            error("means and probs must have the same number of elements")
        end
		if( size(ic,1) != size(ic,2) )
		    error("covariance matrix must be square")
		end
        @assert(1<=length(pact)<=length(p))
        for idx in pact
            if ! (1<=idx<=length(p))
               error("active parameter ",idx," not in range 1:",length(p))
            end
        end
        sump = 0.0
        for i in 1:length(p)
            if p[i] < 0.0
                error("MixtureModel: probabilities must be non-negative")
            end
            sump += p[i]
        end
        table = Distributions.AliasTable(p)
        covar_subset = PDMat(make_matrix_pd(ic[pact,pact]))
        new( m, p ./sump, covar_subset, table, copy(pact) )
     end
     function GaussianMixtureModelCommonCovarSubset(m::Array{Float64,2}, p::Vector{Float64}, ic::AbstractArray{Float64,1}, pact::Vector{Int64} )
        if size(m,2) != length(p)
            error("means and probs must have the same number of elements")
        end
        @assert(1<=length(pact)<=length(p))
        for idx in pact
            if ! (1<=idx<=length(p))
               error("active parameters not in range 1:",length(p))
            end
        end
        sump = 0.0
        for i in 1:length(p)
            if p[i] < 0.0
                error("MixtureModel: probabilities must be non-negative")
            end
            sump += p[i]
        end
        table = Distributions.AliasTable(p)
        covar_subset = PDiagMat(ic[pact])
        new( m, p, covar_subset, table, copy(pact) )
     end
end

function mean(d::GaussianMixtureModelCommonCovarSubset)
    m = copy(x)
    for i in 1:length(d.probs)
        m += vec(d.mu[:,i]) .* d.probs[i]
    end
    return m
end

function pdf(d::GaussianMixtureModelCommonCovarSubset, x::Vector{Float64} )
    p = 0.0
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        p += Distributions.pdf(Distributions.MvNormal(d.covar), x[d.param_active] .- vec(d.mu[d.param_active,i]) ) * d.probs[i]
    end
    return p
end

function pdf(d::GaussianMixtureModelCommonCovarSubset, x::Array{Float64,2} )
    np = size(x,2)
    p = zeros(np)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        p += Distributions.pdf(Distributions.MvNormal(d.covar), x[d.param_active,:] .- vec(d.mu[d.param_active,i]) ) * d.probs[i]
    end
    return p
end

function logpdf(d::GaussianMixtureModelCommonCovarSubset, x::Vector{Float64})
    logp = -Inf
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        logp_i = Distributions.logpdf(Distributions.MvNormal(d.covar), x[d.param_active] .- vec(d.mu[d.param_active,i]) ) + log(d.probs[i])
        logp = logp > logp_i ? logp + log1p(exp(logp_i-logp)) : logp_i + log1p(exp(logp-logp_i))
    end
    return logp
end

function logpdf(d::GaussianMixtureModelCommonCovarSubset, x::Array{Float64,2})
    np = size(x,2)
    logp = fill(-Inf,np)
    for i in 1:length(d.probs)
        if d.probs[i]<=0.0 continue end
        logp_i = Distributions.logpdf(Distributions.MvNormal(d.covar), x[d.param_active,:] .- vec(d.mu[d.param_active,i]) ) + log(d.probs[i])
        for j in 1:length(logp)
          logp[j] = logp[j] > logp_i[j] ? logp[j] + log1p(exp(logp_i[j]-logp[j])) : logp_i[j] + log1p(exp(logp[j]-logp_i[j]))
        end
    end
    return logp
end

# Need to know dimension
function rand(d::GaussianMixtureModelCommonCovarSubset)
    i = Distributions.rand(d.aliastable)
    param = vec(d.mu[:,i])
    #println("# rand:  mu = ", d.mu[:,i], " Sigma = ", d.covar)
    param[d.param_active] = Distributions.rand(Distributions.MvNormal(vec(d.mu[d.param_active,i]),d.covar))
    return param
end


text:
#-----------------------------------------------------------
function TestOrthogonality(α;tol=1.)
	Δα = diff.(α)
	orthogonal = [sum(abs.(δα.-round.(δα)).>tol) for δα in Δα].==0
end

#-----------------------------------------------------------
function Orientation(project,polygons)
	L = nrow(project.dat["master"])

	orientation(val) = 180 ./π .* atan.(val)
	tmp = orientation.(polygons.grad)
	# Flag dwellings as orthogonal (true) or otherwise (false). Include in master
	tol_deg=1.
	project.dat["master"][!,"orthogonal"] = TestOrthogonality(tmp;tol=tol_deg)
	# Determine the first quadrant angle (subtract from 90deg for North bearing)
	polygons[!,:orient_quad1] = [mean(ang)+45 for ang in tmp]
	polygons[!,:orient_set] = [abs.(tmp[i].-polygons.orient_quad1[i]).<tol_deg for i in 1:L]
	return polygons
end

#-----------------------------------------------------------
function FindAdiabatics(project,polygons)
	INonSelfRef(vec,i) = findall(vec.!=vec[i])
	IAdiab(vec₁,vec₂,not_i,i) = intersect(findall(vec₁[not_i[i]].==vec₁[i]),findall(vec₂[not_i[i]].==vec₂[i]))
	Ranges(l) = collect(range(1,stop=l))
	IPoly(l,i) = Int.(ones(l)).*i
	IMap(iPolyⱼ,i) = findall(iPolyⱼ.==i)
	remap(iAdiab,iMapᵢ) = iAdiab[iMapᵢ]
	iRedundantVertices(angles,tol_deg) = findall(abs.(diff(angles)).<(tol_deg*π/180)).+1
	iSelectiveVertexRemoval(iAdiabⱼ,iRmⱼ) = iRmⱼ[[iAdiabⱼ[i].==iAdiabⱼ[i-1] for i in iRmⱼ]]
	RemoveRedundantVertices(nⱼ,iRmⱼ,offset) = nⱼ[Not(iRmⱼ.+offset)]

	# Generate list of all vertices (all polygons, one list)
	edges = DataFrame()
	edges[!,:cen_e] = vcat(polygons.cen_e...)
	edges[!,:cen_n] = vcat(polygons.cen_n...)

	# List of osgb refs
	L = length.(polygons.cen_e)
	edges[!,:osgbs] = Int.(vcat(ones.(L).*project.dat["master"].osgb...))
	edges[!,:iOtherOsgb] = INonSelfRef.((edges.osgbs,),1:length(edges.cen_e))

	# List of edge sequence order
	edges[!,:iEdge] = vcat(Ranges.(L)...)

	# List of parent polygon indicces
	edges[!,:iPoly] = vcat(IPoly.(L,1:length(L))...)

	# Find all adjacencies, based on matching edge centroids
	L = length(edges.cen_e)
	edges[!,:iAdiab] = IAdiab.((edges.cen_e,),(edges.cen_n,),(edges.iOtherOsgb,),1:L)
	edges[!,:adiabBool] = length.(edges.iAdiab).>0

	# Remap adiabatic edges and booleans back to parent polygons
	L = nrow(polygons)
	iMap = IMap.((edges.iPoly,),1:L)
	iAdiab′ = remap.((edges.iAdiab,),iMap)
	adiabBool′ = remap.((edges.adiabBool,),iMap)

	# Remove redundant vertices - first pass (matching centroids)
	iRm₁ = iRedundantVertices.(polygons.orient,8.)

	# Check that for the vertex targeted for removal, that the connected edges
	# ... are either both external, or both adiabatic (mixed should be retained)
	iRm₂ = iSelectiveVertexRemoval.(adiabBool′,iRm₁)

	# Remove redundant vertices
	n″ = RemoveRedundantVertices.(polygons.n,iRm₂,0)
	e″ = RemoveRedundantVertices.(polygons.e,iRm₂,0)
	iAdiab″ = RemoveRedundantVertices.(iAdiab′,iRm₂,-1)
	adiabBool″ = RemoveRedundantVertices.(adiabBool′,iRm₂,-1)

	return n″, e″, iAdiab″, adiabBool″
end

function CompilePolygons(n,e)
	Orient(nᵢ,eᵢ) = atan.(diff(nᵢ)./diff(eᵢ))
	L = nrow(project.dat["master"])
	# Evaluate edge orientations
	orient = Orient.(n,e)
	# Eval diffs and centroids
	Δn = diff.(n)
	Δe = diff.(e)
	cen_n = [n[i][1:end-1] .+ Δn[i]./2 for i in 1:L]
	cen_e = [e[i][1:end-1] .+ Δe[i]./2 for i in 1:L]
	# Combine in DataFrame
	polygons = DataFrame(
		:n=>n,
		:e=>e,
		:Δn=>Δn,
		:Δe=>Δe,
		:cen_n=>cen_n,
		:cen_e=>cen_e,
		:orient=>orient
	)
	return polygons
end
#-----------------------------------------------------------
function ProcessGeom(project)
	# Gather vertices, diffs, and centroids
	n = deepcopy(project.dat["master"].northings_GML)
	e = deepcopy(project.dat["master"].eastings_GML)

	# First pass (pre-clean)
	polygons = CompilePolygons(n,e)
	# Find adiabatics and remove redundant vertices
	n, e, iAdiab, adiabBool = FindAdiabatics(project,polygons)
	# Second pass (post-clean)
	polygons = CompilePolygons(n,e)
	polygons[!,:iAdiab] = iAdiab
	polygons[!,:adiabBool] = adiabBool

	project.dat["master"].northings_GML = polygons.n
	project.dat["master"].eastings_GML = polygons.e
	project.dat["master"][!,"iAdiab"] = iAdiab
	project.dat["master"][!,"adiabBool"] = adiabBool
	pyth(Δx,Δy) = sqrt.(Δx.^2 .+ Δy.^2)
	project.dat["master"][!,"dims"] = pyth.(polygons.Δe,polygons.Δn)

	project.dat["master"][!,"adiabBool"] = adiabBool

	# polygons = Orientation(project,polygons)

	return project
end


text:
grayencode(n::Integer) = n ⊻ (n >> 1)
function graydecode(n::Integer)
    r = n
    while (n >>= 1) != 0
        r ⊻= n
    end
    return r
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
module LagrangeFunction

export GvectorsT
include("LagrangeFunction/Gvectors.jl")

export LF1dGrid
export init_LF1d_p
export init_LF1d_c
export init_LF1d_sinc
include("LagrangeFunction/LF1d.jl")

export LF3dGrid
export init_LF3d_p
export init_LF3d_c
export init_LF3d_sinc
include("LagrangeFunction/LF3d.jl")

export ortho_gram_schmidt
include("LagrangeFunction/ortho_gram_schmidt.jl")

export orthonormalize
include("LagrangeFunction/orthonormalize.jl")

export print_Energies
include("LagrangeFunction/EnergiesT.jl")
include("LagrangeFunction/PotentialsT.jl")

include("LagrangeFunction/op_nabla2.jl")
include("LagrangeFunction/op_H.jl")

export calc_evals, calc_rho
include("LagrangeFunction/calc_Energies.jl")
include("LagrangeFunction/calc_evals.jl")
include("LagrangeFunction/calc_rho.jl")
include("LagrangeFunction/calc_grad.jl")

export KS_solve_scf, KS_solve_Emin_cg, KS_solve_Emin_pcg
include("LagrangeFunction/KS_solve_Emin_cg.jl")
include("LagrangeFunction/KS_solve_Emin_pcg.jl")
include("LagrangeFunction/KS_solve_scf.jl")

export prec_mkl_ilu0, get_Laplacian3d_kron
include("LagrangeFunction/sparse_LF3d.jl")
include("LagrangeFunction/prec_mkl_ilu0.jl")
include("LagrangeFunction/apply_prec_ilu0.jl")

include("LagrangeFunction/solve_poisson_FFT.jl")

include("LagrangeFunction/LDA_VWN.jl")
include("LagrangeFunction/diag_lobpcg.jl")

export build_nabla2
include("LagrangeFunction/build_nabla2.jl")

end


text:
module StaticBitSets

    include("sbitset.jl")

    export SBitSet, push, pop, ⟂, singleton_intersection

end


text:
# ---
# title: 61. Rotate List
# id: problem61
# author: Qling
# date: 2021-03-27
# difficulty: Medium
# categories: Linked List, Two Pointers
# link: <https://leetcode.com/problems/rotate-list/description/>
# hidden: true
# ---
#
# Given the `head` of a linked list, rotate the list to the right by `k` places.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/11/13/rotate1.jpg)
#
#
#
#     Input: head = [1,2,3,4,5], k = 2
#     Output: [4,5,1,2,3]
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/11/13/roate2.jpg)
#
#
#
#     Input: head = [0,1,2], k = 4
#     Output: [2,0,1]
#
#
#
#
# **Constraints:**
#
#   * The number of nodes in the list is in the range `[0, 500]`.
#   * `-100 <= Node.val <= 100`
#   * `0 <= k <= 2 * 109`
#
#
## @lc code=start
using LeetCode

function rotate_right(head::ListNode, k::Int)::ListNode
    isnothing(head) && return head

    last_element, length = head, 1
    while !isnothing(next(last_element))
        last_element = next(last_element)
        length += 1
    end

    k = k % length
    next!(last_element, head)

    temp_element = head
    for _ = 0:(length-k-2)
        temp_element = next(temp_element)
    end

    answer = next(temp_element)
    next!(temp_element, nothing)

    return answer
end
## @lc code=end


text:
MKLFloats = Union{Float32,Float64,ComplexF32,ComplexF64}

"""
    MKLcsr{T<:MKLFloats}

The opaque struct `sparse_matrix` from the MKL library in compressed sparse row (CSR) format
"""
mutable struct MKLcsr{T<:MKLFloats}
end

Base.eltype(m::MKLcsr{T}) where {T} = T

for (T, cr, ex, syrk) in (
    (Float32,    :mkl_sparse_s_create_csr, :mkl_sparse_s_export_csr, :mkl_sparse_s_syrkd,),
    (Float64,    :mkl_sparse_d_create_csr, :mkl_sparse_d_export_csr, :mkl_sparse_d_syrkd,),
    (ComplexF32, :mkl_sparse_c_create_csr, :mkl_sparse_c_export_csr, :mkl_sparse_c_syrkd,),
    (ComplexF64, :mkl_sparse_z_create_csr, :mkl_sparse_z_export_csr, :mkl_sparse_z_syrkd,),
    )
    @eval begin
        function csrptr(adjm::Adjoint{$T,SparseMatrixCSC{$T,BlasInt}})
            m = adjm.parent
            r = Ref(Ptr{MKLcsr{$T}}(0))
            ret = ccall(
                ($(string(cr)), libmkl_rt),
                sparse_status_t,
                (
                    Ref{Ptr{MKLcsr{$T}}}, sparse_index_base_t, BlasInt, BlasInt, Ptr{BlasInt},
                    Ptr{BlasInt}, Ptr{BlasInt}, Ptr{$T}
                ),
                r, SPARSE_INDEX_BASE_ONE, m.n, m.m, m.colptr, pointer(m.colptr, 2),
                m.rowval, m.nzval
            )
            ret == SPARSE_STATUS_SUCCESS || throw(ArgumentError(string(ret)))
            p = r[]
            value = unsafe_load(p)
            finalizer(value) do x
                ccall(
                    (:mkl_sparse_destroy, libmkl_rt),
                    sparse_status_t,
                    (Ref{MKLcsr{$T}},),
                    Ref(value),
                )
            end
            return p
        end

        function SparseArrays.SparseMatrixCSC(csrpt::Ptr{MKLcsr{$T}})
            indexing = Ref(Cint(0))
            rows = Ref(BlasInt(0))
            cols = Ref(BlasInt(0))
            rows_start = Ref(Ptr{BlasInt}(0))
            rows_end = Ref(Ptr{BlasInt}(0))
            col_indx = Ref(Ptr{BlasInt}(0))
            values = Ref(Ptr{$T}(0))
            ret = ccall(
                ($(string(ex)), libmkl_rt),
                sparse_status_t,
                (Ptr{MKLcsr{$T}}, Ref{Cint}, Ref{BlasInt}, Ref{BlasInt},
                Ref{Ptr{BlasInt}}, Ref{Ptr{BlasInt}}, Ref{Ptr{BlasInt}}, Ref{Ptr{$T}}),
                csrpt, indexing, rows, cols, rows_start, rows_end, col_indx, values,
            )
            ret == SPARSE_STATUS_SUCCESS || throw(ArgumentError(ret))
            rowptr = Vector{BlasInt}(undef, rows[] + 1)
            unsafe_copyto!(pointer(rowptr), rows_start[], rows[])
            rowend = Vector{BlasInt}(undef, rows[])
            unsafe_copyto!(pointer(rowend), rows_end[], rows[])
            nonzeros = last(rowend) - indexing[]
            first(rowptr) == indexing[] || throw(ArgumentError("indexing = $indexing ≠ $(first(rowptr)) = first(rowptr)"))
            view(rowptr, 2:rows[]) == view(rowend, 1:(rows[]-1)) || throw(ArgumentError("indices are not dense"))
            rowptr[rows[] + 1] = last(rowend)
            colvals = Vector{BlasInt}(undef, nonzeros)
            unsafe_copyto!(pointer(colvals), col_indx[], nonzeros)
            nzvals = Vector{$T}(undef, nonzeros)
            unsafe_copyto!(pointer(nzvals), values[], nonzeros)
            return SparseMatrixCSC{$T,BlasInt}(cols[], rows[], rowptr, colvals, nzvals)
        end

        function syrkd!(C::StridedMatrix{$T}, Apt::Ptr{MKLcsr{$T}}, op::Char, α::Number, β::Number)
            ret = ccall(
                ($(string(syrk)), libmkl_rt),
                sparse_status_t,
                (sparse_operation_t, Ptr{MKLcsr{$T}}, $T, $T, Ptr{$T}, sparse_layout_t, BlasInt),
                uppercase(op) == 'T' ? SPARSE_OPERATION_TRANSPOSE : SPARSE_OPERATION_NON_TRANSPOSE,
                Apt, α, β, C, SPARSE_LAYOUT_COLUMN_MAJOR, stride(C, 2),
            )
            ret == SPARSE_STATUS_SUCCESS || throw(ArgumentError(ret))
            return C
        end
    end # eval
end #loop on types


text:
using AlphaZero

gspec = Examples.games["bonbon-rectangle"]
mcts = MCTS.Env(gspec, MCTS.RolloutOracle(gspec))
computer = MctsPlayer(mcts, niters=10, timeout=15.0, τ=ConstSchedule(0.5))

interactive!(gspec, computer, Human())
#explore(computer, gspec)


text:
struct VanillaOption
    isCall::Bool
    strike::Float64
    maturity::Float64
end

#advancePath(gen, pathValues, t0, t1)
#advancePayoff(time, pathValues)
function evaluatePayoff(payoff::VanillaOption, x, df)
    if payoff.isCall
        return df * max(x - payoff.strike, 0)
    else
        return df * max(payoff.strike - x, 0)
    end
end

function specificTimes(payoff::VanillaOption)
    return [payoff.maturity]
end


text:
# Autogenerated wrapper script for VegaFusion_jll for aarch64-apple-darwin
export vegafusionserver

JLLWrappers.@generate_wrapper_header("VegaFusion")
JLLWrappers.@declare_executable_product(vegafusionserver)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_executable_product(
        vegafusionserver,
        "bin/vegafusion-server",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
module NGPtools

import Distributions.MvNormal
import SStools.simulate

export assemble_matrices, generate, sample

"""
Observation matrix for state space model. In this case, we assume that the
Np components of the observation vector each track the state of the same
index (more specifically, its U component). All other terms are ignored.
Np is the dimension of the observation vector.
Ns is the number of underlying states.
"""
function Z(Np, Ns)
    maxind = min(3Ns, Np)
    Zmat = sparse(1:maxind, 1:maxind, 1, Np, 3Ns)
    return Zmat
end

"""
State transition matrix. Defined in (5) of Zhu and Dunson. Order of entries
is all state variables, followed by all derivatives, then all latents.
Ns is the dimension of the underlying state space.
"""
function G(Ns, δ; approx=false)
    gmat = eye(3 * Ns)
    gmat += diagm(δ * ones(2Ns), Ns)
    if approx
        gmat += diagm(δ^2/2 * ones(Ns), 2Ns)
    end

    return sparse(gmat)
end

"""
Covariance of state disturbances. Defined in (5) of Zhu and Dunson. Order of
entries is all state variables, followed by all derivatives, then all latents.
"""
function W(Ns, δ, varU::Array{Float64}, varA::Array{Float64}; approx=false)
    if approx
        wmat = δ * spdiagm([varU ; varA])
    else
        wmat = zeros(Float64, 3Ns, 3Ns)
        for idx in 1:Ns
            wmat[idx, idx] = (δ^3 / 3)varU[idx] + (δ^5 / 20)varA[idx]
            wmat[idx, Ns + idx] = (δ^2 / 2)varU[idx] + (δ^4 / 8)varA[idx]
            wmat[idx, 2Ns + idx] = (δ^3 / 6)varA[idx]
            wmat[Ns + idx, idx] = wmat[idx, Ns + idx]
            wmat[Ns + idx, Ns + idx] = δ * varU[idx] + (δ^3 / 3)varA[idx]
            wmat[Ns + idx, 2Ns + idx] = (δ^2 / 2)varA[idx]
            wmat[2Ns + idx, idx] = wmat[idx, 2Ns + idx]
            wmat[2Ns + idx, Ns + idx] = wmat[Ns + idx, 2Ns + idx]
            wmat[2Ns + idx, 2Ns + idx] = δ * varA[idx]
        end
    end

    return wmat
end

function W(Ns, δ, varU::Array{Float64}, varA::Float64; approx=false)
    return W(Ns, δ, varU, varA * ones(Ns); approx=approx)
end

function W(Ns, δ, varU::Float64, varA::Array{Float64}; approx=false)
    return W(Ns, δ, varU * ones(Ns), varA; approx=approx)
end

function W(Ns, δ, varU::Float64, varA::Float64; approx=false)
    return W(Ns, δ, varU * ones(Ns), varA * ones(Ns); approx=approx)
end

"""
Take nested GP parameters and return matrices suitable for feeding into
state space model.

Np is the dimension of the observation vector
Ns is the number of underlying states
δ is a vector of differences between observation times. If observations are
uniformly spaced, δ can be a scalar or a length-1 vector.
"""
function assemble_matrices(Np, Ns, δ, σϵ, σU, σA, σμ, σα; approx=false)
    Nt = length(δ)  # number of time points

    # calculate dimensions
    # Note: for an observation vector of size Np, we introduce 3 latent
    # variables per observation and 2 or 3 noise variables (depending on
    # approximation)
    Nm = 3 * Ns
    if approx
        Nr = 2 * Ns
    else
        Nr = 3 * Ns
    end

    # preallocate arrays:
    H = σϵ^2 * eye(Ns)
    T = Array(Float64, Nm, Nm, Nt)
    Q = Array(Float64, Nr, Nr, Nt)

    if approx
        R = sparse([2:3:Nm ; 3:3:Nm], [1:2:Nr ; 2:2:Nr], 1., Nm, Nr)
    else
        R = speye(Float64, Nm)
    end

    for t in 1:Nt
        T[:, :, t] = G(Ns, δ[t], approx=approx)
        Q[:, :, t] = W(Ns, δ[t], σU.^2, σA.^2, approx=approx)
    end

    if Nt == 1
        T = sparse(squeeze(T, 3))
        Q = squeeze(Q, 3)
        if approx
            Q = sparse(Q)
        end
    end

    a0 = zeros(Nm)
    P0 = diagm(repeat([σμ, σμ, σα], inner=[Ns]))

    return Z(Np, Ns), H, T, R, Q, a0, P0
end

"""
Generate nGP data according to the state space model. Uses matrices
constructed in assemble_matrices.

Np is the dimension of the observation vector
Ns is the number of underlying states
δ is a spacing between observations
"""
function generate(Np, Ns, Nt, δ, σϵ, σU, σA, σμ, σα; approx=false)

    Z, H, T, R, Q, a0, P0 = assemble_matrices(Np, Ns, δ, σϵ, σU, σA, σμ, σα, approx=approx)

    Nm = size(Z, 2)

    α = Array(Float64, Nm, Nt)
    y = Array(Float64, Np, Nt)

    # initialize
    α[:, 1] = rand(MvNormal(full(a0), full(P0)))

    # simulate
    for t in 1:Nt
        ϵ = rand(MvNormal(H))
        local Q_t = ndims(Q) < 3 ? Q : Q[:, :, t]
        η = rand(MvNormal(Q_t))
        y[:, t] = Z * α[:, t] + ϵ

        if t < Nt
            local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)
            α[:, t + 1] = T_t * α[:, t] + R * η
        end
    end

    return y, α
end

"""
Sample Ns times from the nGP posterior, given observations y.
"""
function _sample(y, Ns, a0, P0, Z, H, T, R, Q; interleaved=true)
    Nt = size(y, 2)
    Nm = size(Z, 2)

    α_samples = Array(Float64, Nm, Nt, Ns)

    for idx in 1:Ns
        α_samples[:, :, idx] = simulate(y, a0, P0, Z, H, T, R, Q, interleaved=interleaved)
    end

    return α_samples

end

"""
Version where spacing is uniform, so number of time points must be specified
explicitly.
"""
function sample(y, Nsamples, Nstates, δ, σϵ, σU, σA, σμ, σα; approx=false)
    Np = size(y, 1)
    Z, H, T, R, Q, a0, P0 = assemble_matrices(Np, Nstates, δ, σϵ, σU, σA, σμ, σα, approx=approx)
    _sample(y, Nsamples, a0, P0, Z, H, T, R, Q)
end

"""
Version where observation model (Z and H) is specified externally.
"""
function sample(y, Nsamples, Z::Array{Float64}, H::Array{Float64}, δ::Float64,
    σU, σA, σμ, σα; approx=false, interleaved=true)
    Nstates = convert(Int64, size(Z, 2) / 3)
    Np = size(y, 1)
    _, _, T, R, Q, a0, P0 = assemble_matrices(Np, Nstates, δ, 0, σU, σA, σμ, σα, approx=approx)
    _sample(y, Nsamples, a0, P0, Z, H, T, R, Q; interleaved=interleaved)
end

end  # module

max_stars_repo_name: jmxpearson/labcr

max_issues_repo_name: jmxpearson/labcr

max_forks_repo_name: jmxpearson/labcr


text:
using Test
using OMEinsum
using OMEinsum: get_size_dict, Sum, Tr, PairWise, DefaultRule, IndexSize, Permutedims
using SymEngine
using LinearAlgebra: I

SymEngine.free_symbols(syms::Union{Real, Complex}) = Basic[]
SymEngine.free_symbols(syms::AbstractArray{T}) where {T<:Union{Real, Complex}} = Basic[]
function rand_assign(syms...)
    fs = union(free_symbols.(syms)...)
    Dict(zip(fs, randn(length(fs))))
end

function _basic_approx(x, y; atol=1e-8)
    diff = x-y
    assign = rand_assign(x, y)
    length(assign) > 0 && (diff = subs.(diff, Ref.((assign...,))...))
    nres = ComplexF64.(diff)
    all(isapprox.(nres, 0; atol=atol))
end

Base.:≈(x::AbstractArray{<:Basic}, y::AbstractArray; atol=1e-8) = _basic_approx(x, y, atol=atol)
Base.:≈(x::AbstractArray, y::AbstractArray{<:Basic}; atol=1e-8) = _basic_approx(x, y, atol=atol)
Base.:≈(x::AbstractArray{<:Basic}, y::AbstractArray{<:Basic}; atol=1e-8) = _basic_approx(x, y, atol=atol)
Base.Complex{T}(a::Basic) where T = T(real(a)) + im*T(imag(a))

@testset "tensor order check" begin
    ixs = ((1,2), (2,3))
    a = randn(3,3)
    @test_throws ArgumentError get_size_dict(ixs, (a,a,a))
    a = randn(3,3,3)
    @test_throws ArgumentError get_size_dict(ixs, (a,a))
end

@testset "get size dict" begin
    a = randn(3,2)
    sizedict = get_size_dict(((1,2), (2,3)), (a, a'))
    @test (sizedict[1], sizedict[2], sizedict[3]) == (3,2,3)
    sizedict = get_size_dict((('i','j'), ('k','j')), (a, a))
    @test (sizedict['i'], sizedict['j'], sizedict['k']) == (3,2,3)
    @test_throws DimensionMismatch get_size_dict((('i','j'), ('j','k')), (a, a))
end

@testset "einsum" begin
    # matrix and vector multiplication
    a,b,c = randn(2,2), rand(2,2), rand(2,2)
    v = rand(2)
    t = randn(2,2,2,2)
    @test einsum(ein"ijkl -> ijkl", (t,)) ≈ t
    @test einsum(ein"αβγδ -> αβγδ", (t,)) ≈ t
    @test einsum(EinCode(((1,2),(2,3),(3,4)),(1,4)), (a,b,c)) ≈ a * b * c
    @test einsum(EinCode(((1,20),(20,3),(3,4)), (1,4)), (a,b,c)) ≈ a * b * c
    @test einsum(EinCode(((1,2),(2,3),(3,4)),(4,1)), (a,b,c)) ≈ permutedims(a*b*c, (2,1))
    @test einsum(EinCode(((1,2),(2,)), (1,)), (a,v)) ≈ a * v

    # more matmul
    @test ein"ij,jk -> ik"(a,a) ≈ a * a
    @test ein"ij,jk -> ki"(a,a) ≈ transpose(a * a)
    @test ein"ij,kj -> ik"(a,a) ≈ a * transpose(a)
    @test ein"ij,kj -> ki"(a,a) ≈ transpose(a * transpose(a))
    @test ein"ji,jk -> ik"(a,a) ≈ transpose(a) * a
    @test ein"ji,jk -> ki"(a,a) ≈ transpose(transpose(a) * a)
    @test ein"ji,kj -> ik"(a,a) ≈ transpose(a) * transpose(a)
    @test ein"ji,kj -> ki"(a,a) ≈ transpose(transpose(a) * transpose(a))

    # contract to 0-dim array
    @test einsum(EinCode(((1,2),(1,2)), ()), (a,a))[] ≈ sum(a .* a)

    # trace
    @test einsum(EinCode(((1,1),),()), (a,))[] ≈ sum(a[i,i] for i in 1:2)
    @test einsum(EinCode(((1,1),),()), (a,)) isa Array
    aa = rand(2,4,4,2)
    @test einsum(EinCode(((1,2,2,1),), ()), (aa,))[] ≈ sum(aa[i,j,j,i] for i in 1:2, j in 1:4)

    # partial trace
    @test einsum(EinCode(((1,2,2,3),), (1,3)), (aa,)) ≈ sum(aa[:,i,i,:] for i in 1:4)
    # with permutation
    @test einsum(EinCode(((1,2,2,3),), (3,1)), (aa,)) ≈ transpose(sum(aa[:,i,i,:] for i in 1:4))

    # diag
    @test einsum(EinCode(((1,2,2,3),), (1,2,3)), (aa,)) ≈ aa[:,[CartesianIndex(i,i) for i in 1:4],:]

    # permutation
    @test einsum(EinCode(((1,2),), (2,1)), (a,)) ≈ permutedims(a,(2,1))
    @test einsum(EinCode(((1,2,3,4),),(2,3,1,4)), (t,)) ≈ permutedims(t,(2,3,1,4))

    # tensor contraction
    ta = zeros(size(t)[[1,2]]...)
    for (i,j,k,l) in Iterators.product(1:2,1:2,1:2,1:2)
        ta[i,l] += t[i,j,k,l] * a[j,k]
    end
    @test einsum(EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a)) ≈  ta

    ta = zeros(size(t)[[1,2]]...)
    for (i,j,k,l) in Iterators.product(1:2,1:2,1:2,1:2)
        ta[i,l] += t[l,k,j,i] * a[j,k]
    end
    @test einsum(EinCode(((4,3,2,1), (2,3)),(1,4)), (t,a)) ≈  ta

    # star-contraction
    aaa = zeros(2,2,2);
    for (i,j,k,l) in Iterators.product(1:2,1:2,1:2,1:2)
        aaa[j,k,l] += a[i,j] * a[i,k] * a[i,l]
    end
    @test aaa ≈ einsum(EinCode(((1,2),(1,3),(1,4)),(2,3,4)), (a,a,a))

    # star and contract
    aaa = zeros(2);
    for (i,j,l) in Iterators.product(1:2,1:2,1:2)
        aaa[l] += a[i,j] * a[i,j] * a[i,l]
    end
    @test einsum(EinCode(((1,2),(1,2),(1,3)), (3,)), (a,a,a)) ≈ aaa

    # index-sum
    a = rand(2,2,5)
    @test einsum(EinCode(((1,2,3),),(1,2)),(a,)) ≈ sum(a, dims=3)
    # with permutation
    @test ein"ijk -> ki"(a) ≈ transpose(dropdims(sum(a,dims=2),dims=2))
    @test ein"ijk -> "(a)[] ≈ sum(a)

    # Hadamard product
    a = rand(2,3)
    b = rand(2,3)
    @test einsum(EinCode(((1,2),(1,2)), (1,2)), (a,b)) ≈ a .* b

    # Outer
    a = rand(2,3)
    b = rand(2,3)
    @test einsum(EinCode(((1,2),(3,4)),(1,2,3,4)),(a,b)) ≈ reshape(a,2,3,1,1) .* reshape(b,1,1,2,3)

    # Broadcasting
    @test ein"->ii"(OMEinsum.asarray(1); size_info=IndexSize('i'=>5)) == Matrix(I, 5, 5)

    # trivil
    @test ein"->ii"(asarray(1); size_info=IndexSize('i'=>5)) == Matrix(I, 5, 5)
    res = ein"->"(asarray(3))
    @test res isa Array
    @test res[] === 3
    res = ein",->"(asarray(3), asarray(4))
    @test res isa Array
    @test res[] === 12

    # Projecting to diag
    a = rand(2,2)
    a2 = [a[1] 0; 0 a[4]]
    @test einsum(EinCode(((1,1),), (1,1)), (a,)) ≈ a2

    ## operations that can be combined
    a = rand(2,2,2,2)
    @test einsum(EinCode(((1,1,2,2),), ()), (a,))[] ≈ sum(a[[CartesianIndex(i,i) for i in 1:2], [CartesianIndex(i,i) for i in 1:2]])

    @test einsum(EinCode(((1,2,3,4), (3,4,5,6)), (1,2,5,6)), (a,a)) ≈ reshape(reshape(a,4,4) * reshape(a,4,4),2,2,2,2)
end

@testset "fallback" begin
    # while we expect some scaling in the allocations for multiple inputs, it
    # shouldn't increase too much
    a = rand(100,100)
    b = rand(100,100)
    einsum(EinCode(((1,2),(2,3)),(1,3)), (a,b))
    allocs1 = @allocated einsum(EinCode(((1,2),(2,3)),(1,3)), (a,b))
    @test allocs1 < 10^5
    einsum(EinCode(((1,2),(2,3),(3,4)),(1,4)), (a,b,b))
    allocs2 = @allocated einsum(EinCode(((1,2),(2,3),(3,4)),(1,4)), (a,b,b))
    # doing twice the work (two multiplications instead of one) shouldn't
    # incure much more than twice the allocations.
    @test_broken allocs2 < 2.1 * allocs1

    @test_throws MethodError einsum(((1,2),(2,3)), (a,a))
end

@testset "error handling" begin
    a = randn(3,3)
    b = randn(4,4)
    @test_throws DimensionMismatch einsum(EinCode(((1,2), (2,3)), (1,3)), (a, b))
end

@testset "string input" begin
    a = randn(3,3)
    @test einsum(ein"ij,jk -> ik", (a,a)) ≈ einsum(EinCode(((1,2),(2,3)), (1,3)), (a,a))
    @test ein"ij,jk -> ik"(a,a) ≈ einsum(EinCode(((1,2),(2,3)), (1,3)), (a,a))
    @test ein"αβ,βγ -> αγ"(a,a) ≈ a * a
    # Note: the following statement is nolonger testable, since will cause load error now!
    #@test_throws ArgumentError einsum(ein"ij,123 -> k", (a,a))
end

@testset "macro input" begin
    a = randn(2,2)
    @test a * a ≈ @ein [i,k] := a[i,j] * a[j,k]
    @test sum(a[i,i] for i in 1:2) ≈ (@ein [] := a[i,i])[]
    @test [a[1,1] 0; 0 a[2,2]] ≈ @ein [i,i] := a[i,i]
    @test permutedims(a) ≈ @ein [α,a] := a[a,α]
    @test permutedims(a) ≈ @ein [α,1] := a[1,α]
end

@testset "argument checks" begin
    @test_throws ArgumentError einsum(ein"ij,jk -> ik", (rand(2,2), rand(2,2), rand(2,2)))
    @test_throws ArgumentError einsum(ein"ij,jk,k -> ik", (rand(2,2), rand(2,2)))
    @test_throws ArgumentError einsum(ein"ij,ijk -> ik", (rand(2,2), rand(2,2)))
    @test_throws DimensionMismatch einsum(ein"ij,jk -> ik", (rand(2,3), rand(2,2)))
end

@testset "dispatched" begin
    # index-sum
    a = rand(2,2,5)
    ixs, xs = ((1,2,3),), (a,)
    @test einsum(Sum(), EinCode(ixs,(1,2)),xs, get_size_dict(ixs, xs)) ≈ sum(a, dims=3)
    a = rand(5,5)
    @test einsum(Tr(), EinCode(((1,1),),()), (a,), get_size_dict(((1,1),), (a,)))[] ≈ sum(a[i,i] for i in 1:5)
    t = rand(5,5,5,5)
    a = rand(5,5)
    size_dict = IndexSize((1,2,3,4,2,3), ((size(t)..., size(a)...)))
    ta = loop_einsum(EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a), size_dict)
    @test einsum(PairWise(), EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a), size_dict) ≈  ta
    @test einsum(DefaultRule(), EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a), size_dict) ≈  ta

    # index-sum
    a = Basic.(rand(2,2,5))
    ixs, xs = ((1,2,3),), (a,)
    @test einsum(Sum(), EinCode(ixs,(1,2)),xs, get_size_dict(ixs, xs)) ≈ sum(a, dims=3)
    a = Basic.(rand(5,5))
    @test isapprox(einsum(Tr(), EinCode(((1,1),),()), (a,), get_size_dict(((1,1),), (a,)))[], sum(a[i,i] for i in 1:5), rtol=1e-8)
    t = Basic.(rand(5,5,5,5))
    a = Basic.(rand(5,5))
    size_dict = IndexSize((1,2,3,4,2,3), ((size(t)..., size(a)...)))
    ta = loop_einsum(EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a), size_dict)
    @test einsum(PairWise(), EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a), size_dict) ≈  ta
    @test einsum(DefaultRule(), EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a), size_dict) ≈  ta
end

@testset "isbatchmul" begin
    for (ixs, iy) in [(((1,2), (2,3)), (1,3)), (((1,2,3), (2,3)), (1,3)),
                        (((7,1,2,3), (2,4,3,7)), (1,4,3)),
                        (((3,), (3,)), (3,)), (((3,1), (3,)), (3,1))
                        ]
        xs = ([randn(ComplexF64, fill(4,length(ix))...) for ix in ixs]...,)
        @test EinCode(ixs, iy)(xs...) ≈ loop_einsum(EinCode(ixs, iy), xs, OMEinsum.get_size_dict(ixs, xs))
        @test OMEinsum.batched_contract(Val(ixs[1]), xs[1], Val(ixs[2]), xs[2], Val(iy)) ≈ loop_einsum(EinCode(ixs, iy), xs, OMEinsum.get_size_dict(ixs, xs))
    end
end


text:
include("LatticeVars.jl")
include("AtomicVars.jl")

include("r3frac.jl")
include("r3mv.jl")

function create_lattice_vars()
    LatVecs = zeros(3,3)
    LatVecs[1,:] = [5.0, 5.0, 0.0]
    LatVecs[2,:] = [5.0, 0.0, 5.0]
    LatVecs[3,:] = [0.0, 5.0, 5.0]
    lattice_vars = LatticeVars( LatVecs )
    return lattice_vars
end

function test_AtomicVars()

    lattice_vars = create_lattice_vars()

    maxatoms = 200
    maxspecies = 8
    atposl = zeros(3,maxatoms,maxspecies)

    nspecies = 2
    natoms = [1,1]

    # species 1, atom 1
    atposl[:,1,1] = [0.1, 0.1, 0.1]
    # species 2, atom 1
    atposl[:,1,2] = [0.25, 0.25, 0.25]

    atomic_vars = AtomicVars(nspecies, natoms, atposl, lattice_vars)
end

test_AtomicVars()


text:
"""
    DerivedFunctionException(exc, salsa_trace)

Exceptions thrown from within any Salsa @derived function are wrapped in a
DerivedFunctionException(), which also captures the "Salsa Trace", which is the stack trace
of the Salsa functions (derived functions and input functions) that were being executed
when the exception was thrown.

The Salsa Trace is richer than a normal stacktrace: it includes the _values_ of all function
arguments (salsa keys). It is displayed in reverse order, so that the stack trace can be
read from the inside-out, together with the normal julia stacktrace.
"""
struct DerivedFunctionException{T} <: Base.Exception
    captured_exception::T
    salsa_trace::Vector{DependencyKey}
end
DerivedFunctionException(exc::T, trace) where {T} = DerivedFunctionException{T}(exc, trace)

function ExceptionUnwrapping.unwrap_exception(exc::DerivedFunctionException)
    exc.captured_exception
end

function Base.showerror(io::IO, exc::DerivedFunctionException)
    print(io, nameof(typeof(exc)))
    println(io, ": Error encountered while executing derived function:")
    Base.showerror(io, exc.captured_exception)
    println(io, "\n\n------ Salsa Trace -----------------")
    for (idx, dependency_key) in Iterators.reverse(enumerate(exc.salsa_trace))
        println(io, "[$idx] ", dependency_key)  # Uses pretty-printing for Traces defined below
    end
    println(io, "------------------------------------")
end

"""
    DependencyCycleException(problematic_depkey)

Exception thrown when a derived function stacktrace includes a cycle. The exception itself
only contains the problematic call that closed the cycle, the stacktrace itself is needed
for debugging and will get attached to the surrounding `DerivedFunctionException`.
"""
struct DependencyCycleException <: Base.Exception
    problematic_depkey
end

function Base.showerror(io::IO, exc::DependencyCycleException)
    print(io, nameof(typeof(exc)))
    println(io, ": Cycle in derived function stacktrace calling $(exc.problematic_depkey).")
end


text:
function numdiff(f, x0; ϵ=1e-6)
    f0 = f(x0)
    p = size(f0,1)
    q = size(x0,1)
    J = zeros(p,q)
    for j=1:q
        ei = zeros(q)
        ei[j] = 1.0
        J[:,j] = (f(x0+ϵ*ei) - f0)/ϵ
    end
    return J
end

function get_ss_derivatives(model; method=:symbolic)
    m, s, x, p = model.calibration[:exogenous, :states, :controls, :parameters]
    if method==:symbolic
        # it is still incorrect
        g_diff = transition(model,Val{(1,2,3,4)}, m, s, x, m, p)
        f_diff = arbitrage(model, Val{(1,2,3,4,5,6)}, m, s, x, m, s, x, p)
    else
        g_diff = [numdiff(f,u) for (f,u) in [
            (u->transition(model, u, s, x, m, p),m),
            (u->transition(model, m, u, x, m, p),s),
            (u->transition(model, m, s, u, m, p),x),
            (u->transition(model, m, s, x, u, p),m)
        ]]
        f_diff = [numdiff(f,u) for (f,u) in [
            (u->arbitrage(model, u, s, x, m, s, x, p), m),
            (u->arbitrage(model, m, u, x, m, s, x, p), s),
            (u->arbitrage(model, m, s, u, m, s, x, p), x),
            (u->arbitrage(model, m, s, x, u, s, x, p), m),
            (u->arbitrage(model, m, s, x, m, u, x, p), s),
            (u->arbitrage(model, m, s, x, m, s, u, p), x),
        ]]
    end
    g_diff, f_diff
end

perturb(p::IIDExogenous) = (zeros(0), zeros(0, 0))
perturb(p::VAR1) = (p.mu, p.R)

mutable struct PerturbationResult
    dr::BiTaylorExpansion
    generalized_eigenvalues::Union{Vector, Nothing}
    stable::Union{Bool, Nothing}     # biggest e.v. lam of solution is < 1
    determined::Union{Bool, Nothing} # next eigenvalue is > lam + epsilon (MOD solution well defined)
    unique::Union{Bool, Nothing}     # next eigenvalue is > 1
end

function Base.show(io::IO, pbr::PerturbationResult)
    @printf io "Perturbation Results\n"
    @printf io " * Decision Rule type: %s\n" string(typeof(pbr.dr))
    # @printf io " * Blanchard-Kahn: %s\n" blanchard_kahn(pbr)
    if !(typeof(pbr.stable)<:Nothing)
        @printf io " * stable < %s\n" pbr.stable
    end
    if !(typeof(pbr.determined)<:Nothing)
        @printf io " * determined < %s\n" pbr.determined
    end
end

blanchard_kahn(fos::PerturbationResult) = fos.stable && fos.unique

function perturb_first_order(g_s, g_x, f_s, f_x, f_S, f_X)

    eigtol = 1.0+1e-6

    ns = size(g_s, 2)
    nx = size(g_x, 2)

    nv = ns + nx

    A = [Matrix(1.0I,ns,ns) zeros(ns, nx);
         -f_S     -f_X]

    B = [g_s g_x;
         f_s f_x]

    # do orderd QZ decomposition
    gs = schur(A, B)

    genvals = (abs.(gs.α) ./ abs.(gs.β))
    sort!(genvals, rev=true)
    n_keep = ns # number of eigenvalues to keep
    diff = genvals[n_keep+1] - genvals[n_keep]
    eigtol = genvals[n_keep] + diff/2

    select = (abs.(gs.α) .> eigtol*abs.(gs.β))

    ordschur!(gs, select)
    S, T, Q, Z = gs.S, gs.T, gs.Q, gs.Z
    diag_S = diag(S)
    diag_T = diag(T)
    eigval = abs.(diag_S./diag_T)

    # now obtain first order solution
    Z11 = Z[1:ns, 1:ns]
    Z12 = Z[1:ns, ns+1:end]
    Z21 = Z[ns+1:end, 1:ns]
    Z22 = Z[ns+1:end, ns+1:end]
    S11 = S[1:ns, 1:ns]
    T11 = T[1:ns, 1:ns]

    # first order solution
    C = (Z11'\Z21')'
    return C, genvals
end

function get_gf_derivatives(model::AbstractModel)

    g_diff, f_diff = get_ss_derivatives(model)
    _f_m, _f_s, _f_x, _f_M, _f_S, _f_X = f_diff
    _g_m, _g_s, _g_x, _g_M = g_diff

    (M, R) = perturb(model.exogenous)

    f_x = _f_x
    f_X = _f_X
    _m, _s, x, p = model.calibration[:exogenous, :states, :controls, :parameters]

    if size(R, 1)>0
        f_s = [_f_m _f_s]
        f_S = [_f_M _f_S]
        g_s = [R zeros(size(R, 1), size(_g_s, 2)); _g_m _g_s]
        g_x = [zeros(size(_g_m, 1), size(_g_x, 2)); _g_x]
        s = cat(_m, _s; dims=1)
    else
        f_s = _f_s
        f_S = _f_S
        g_s = _g_s
        g_x = _g_x
        s = _s
    end
    return g_s, g_x, f_s, f_x, f_S, f_X
end

"""
TBD
"""
function perturb(model::Model; method=:qz)

    g_s, g_x, f_s, f_x, f_S, f_X = get_gf_derivatives(model)
    nx = size(g_x, 2)

    (M, R) = perturb(model.exogenous)
    _m, _s, x, p = model.calibration[:exogenous, :states, :controls, :parameters]

    if size(R, 1)>0
        s = cat(_m, _s; dims=1)

    else
        s = _s
    end

    n_s = size(g_s, 1)
    if method==:qz
        tol = 1e-6 # minimum distance betweel lam_n and lam_{n+1}
        C, genvals = perturb_first_order(g_s, g_x, f_s, f_x, f_S, f_X)
        sort!(genvals)
        stable = (genvals[n_s]<1)
        deter = (genvals[n_s+1]-genvals[n_s]>tol)
        uni = (genvals[n_s+1]>1)
    elseif method==:linear_time_iteration
        C, nit = linear_time_iteration(g_s, g_x, f_s, f_x, f_S, f_X)
        genvals = nothing
        stable = nothing
        deter = nit>0
        uni = nothing
    end

    if size(R, 1)>0
        C_exo = C[:, 1:length(_m)]
        C_endo = C[:, (length(_m)+1):end]
        dr = BiTaylorExpansion{nx}(_m, _s, x, C_exo, C_endo)
    else
        dr = BiTaylorExpansion{nx}(_m, _s, x, zeros(nx, length(_m)), C)
    end

    PerturbationResult(
        dr,
        genvals,
        stable,
        deter,
        uni
    )

end

function linear_time_iteration(model; maxit=1000, tol_ϵ=1e-8, improve=20, verbose=false)
    g_s, g_x, f_s, f_x, f_S, f_X = Dolo.get_gf_derivatives(model)
    return linear_time_iteration(g_s, g_x, f_s, f_x, f_S, f_X; maxit=maxit)
end

function L(P, Q, X, Y, u)
    return P*u*Q
end

function linear_time_iteration(g_s, g_x, f_s, f_x, f_S, f_X; maxit=1000, tol_ϵ=1e-8, improve=20, verbose=false)

    n_x, n_s = size(f_s)

    K = f_s + f_S*g_s

    X_0 = 1 .+rand(n_x, n_s)
    local X
    for it=1:maxit
        Y = X_0
        B = -K - f_X*Y*g_s
        A = f_x + f_S*g_x + f_X*Y*g_x
        X = A\B
        ϵ = maximum(abs.(X-X_0))
        if ϵ<tol_ϵ
            return (X,it)
        end
        if verbose
            println(it, "\t: ", ϵ)
        end
        if (improve>0)
            S = (X-X_0)
            δ = (X-X_0)
            P = -(f_x + f_S*g_x + f_X*Y*g_x)\f_X
            Q = g_s + g_x*X
            for k=1:improve
                S = δ + L(P, Q, X, Y, S)
            end
            X_0 = X_0 + S
        else
            X_0 = X
        end
    end

    return (X,-1)

end

function ttest(g_s, g_x, f_s, f_x, f_S, f_X, X)
    res = f_s + f_x*X + f_S*(g_s+g_x*X) + f_X*X*(g_s+g_x*X)
    return res
end


text:
module FastKeywords

iskw(ex) = false
iskw(ex::Expr) = ex.head == :kw && length(ex.args) == 2 && (ex.args[1] isa Symbol)

iskw_or_eq(ex) = false
iskw_or_eq(ex::Expr) = (ex.head == :(=) || ex.head == :kw) && length(ex.args) == 2 && (ex.args[1] isa Symbol)

include("KW.jl")
include("fastcall.jl")
include("fastkw.jl")

export @kw, @fastkw, @fastcall

end # module


text:
function regression_model(input_d::Int,
                          output_d::Int,
                          hidden_d::Int,
                          hidden_layers::Int,
                          nonlinearity::Function,
                          initialization)
   if hidden_layers == 0
      return FastChain(FastDense(input_d, output_d))
   elseif hidden_layers == 1
      return FastChain(FastDense(input_d, hidden_d, nonlinearity; initW=initialization),
                       FastDense(hidden_d, output_d; initW=initialization))
   elseif hidden_layers > 1
      layers = [FastDense(input_d, hidden_d, nonlinearity; initW=initialization),
                FastDense(hidden_d, output_d; initW=initialization)]
      for j in 1:hidden_layers-1
         insert!(layers, j+1, FastDense(hidden_d, hidden_d, nonlinearity; initW=initialization))
      end
      return FastChain(layers...)
   end

   error("Hidden layers must be a positive integer or 0")
end

function restructure_parameters(θ::Vector{T},
                                input_d::Int,
                                output_d::Int,
                                hidden_d::Int,
                                hidden_layers::Int) where {T}
     """ Specialized on the type of regression neural networks we use here, and the FastX
   chain of models """
   if hidden_layers == 0
      (output_d * input_d + output_d == length(θ)) || error("dimensions don't match")
      return [θ]
   elseif hidden_layers == 1
      l_1 = hidden_d * input_d + hidden_d
      l_N = output_d * hidden_d + output_d
      l_1 + l_N == length(θ) || error("dimensions don't match")
      return [θ[1:l_1], θ[l_1+1:l_1+l_N]]
   elseif hidden_layers > 1
      l_1 = hidden_d * input_d + hidden_d
      l_j = hidden_d * hidden_d + hidden_d
      l_N = output_d * hidden_d + output_d
      l_1 + (hidden_layers-1) * l_j + l_N == length(θ) || error("dimensions don't match")
      parameters = [θ[1:l_1], θ[end-l_N+1:end]]
      for j in 1:hidden_layers-1
         start = l_1 + (j-1) * l_j + 1
         stop = l_1 + j * l_j
         insert!(parameters, j+1, θ[start:stop])
      end
      return parameters
   end

   error("Hidden layers must be a positive integer or 0")
end

function reshape_all_parameters(θ::Vector{T},
                                input_d::Int,
                                output_d::Int,
                                hidden_d::Int,
                                hidden_layers::Int) where {T}
   _θ = restructure_parameters(θ, input_d, output_d, hidden_d, hidden_layers)
   reshape_all_parameters(_θ, input_d, output_d, hidden_d, hidden_layers)
end

function reshape_all_parameters(θ::Vector{Vector{T}},
                                input_d::Int,
                                output_d::Int,
                                hidden_d::Int,
                                hidden_layers::Int) where {T}
   length(θ) == hidden_layers + 1 || error("dimensions don't match: parameter length $(length(θ)) and expected $(hidden_layers+1) sets of parameters")

   if hidden_layers == 0
      return reshape_parameters(first(θ), input_d, output_d)
   elseif hidden_layers == 1
      return [reshape_parameters(first(θ), input_d, hidden_d),
              reshape_parameters(last(θ), hidden_d, output_d)]
   elseif hidden_layers > 1
      parameters = [reshape_parameters(first(θ), input_d, hidden_d),
                    reshape_parameters(last(θ), hidden_d, output_d)]
      for j in 2:hidden_layers
         insert!(parameters, j, reshape_parameters(θ[j], hidden_d, hidden_d))
      end
      return parameters
   end

end

function reshape_parameters(θ::Vector{T}, input_d::Int, output_d::Int) where {T}
   output_d * input_d + output_d == length(θ) || error("dimensions don't match")
   return (weights = reshape(θ[1:output_d*input_d], output_d, input_d),
           bias = θ[output_d*input_d+1:end])
end

struct CallbackLog{T}
    log_params::Bool
    parameters::Vector{Vector{T}}

    log_loss::Bool
    losses::Vector{T}

    log_preds::Bool
    predictions::Vector{Array{T}}
end

function CallbackLog(;log_params=true, log_loss=true, log_preds=true, T=Any)
    CallbackLog(log_params,
                Vector{Vector{T}}(),
                log_loss,
                Vector{T}(),
                log_preds,
                Vector{Array{T}}())
end

function (cb::CallbackLog)(θ, loss, prediction)

    if cb.log_params
        push!(cb.parameters, copy(θ))
    end

    if cb.log_loss
        push!(cb.losses, loss)
    end

    if cb.log_preds
        push!(cb.predictions, prediction)
    end

    if any(isnan.(loss))
        println("NaN in losses; stop training!")
        return true
    elseif any(isinf.(loss))
        println("Inf in losses; stop training!")
        return true
    end

    return false
end

function repeat_experiment(problem,
                           net_config::OrderedDict,
                           repetitions::Int;
                           longterm=0.0,
                           ε = 0.0,
                           progress=true)

    callbacks = []
    min_losses = []
    min_params = []
    longterm_predictions = []

    real_de, _ = problem[:equation](problem[:parameters], x->nothing)
    prob = ODEProblem(real_de, problem[:u0], problem[:tspan])
    solution = solve(prob, problem[:solver](), saveat=problem[:ts])
    y = Array(solution)

    if !iszero(longterm)
        prob_long = ODEProblem(real_de, problem[:u0], (problem[:tspan][1], longterm))
        solution_long = solve(prob, problem[:solver](), saveat=problem[:ts])
        y_long = Array(solution)
    else
        y_long = []
    end

    Threads.@threads for _ in 1:repetitions

        UA = regression_model(values(net_config)...)
        _ , univ_de = problem[:equation](problem[:parameters], UA)

        θ₀ = initial_params(UA)
        prob_nn = ODEProblem(univ_de, problem[:u0], problem[:tspan], θ₀)

        ỹ = y .+ ε * randn(eltype(y), size(y))

        predict(θ) = Array(concrete_solve(prob_nn, problem[:solver](), problem[:u0], θ, saveat=solution.t))
        loss(θ) = problem[:loss](θ, y, predict)

        callback = CallbackLog(T=Float32)

        res = DiffEqFlux.sciml_train(loss,
                                     θ₀,
                                     problem[:optimizer](),
                                     cb=callback,
                                     maxiters=problem[:max_iter];
                                     progress)
        push!(callbacks, callback)
        push!(min_losses, res.minimum)
        push!(min_params, res.minimizer)

        if !iszero(longterm)
            new_tspan = (problem[:tspan][1], longterm)
            prob_longerm = ODEProblem(univ_de, problem[:u0], new_tspan, res.minimizer)
            longterm_prediction = Array(concrete_solve(
                                                     prob_longerm,
                                                     problem[:solver](),
                                                     problem[:u0],
                                                     res.minimizer,
                                                     saveat=problem[:ts]
                                                     ))
            push!(longterm_predictions, longterm_prediction)
        end

    end

    return Dict([
        "losses" => min_losses,
        "parameters" => min_params,
        "longterm_predictions" => longterm_predictions,
        "longterm_solution" => y_long
    ]), callbacks

end

function grid_experiment(problem, grid_config, repetitions; progress=true)

    result = []

    for net_config in grid_config
        summary = repeat_experiment(problem, net_config, repetitions; progress)
        push!(result, summary[:losses])
    end

    return reshape(result, size(grid_config))
end

# very purpose build data loader to repeat experiments with the same initializations
mutable struct InitializationLoader{T}
    _state::Int

    n_sets::Int
    n_parameters::Int
    grp_id::Int

    data_path::String
    initializations::Vector{Vector{T}}
end

function InitializationLoader(grp::Int)
    data_path = "./data/P$(grp)SelkovsimpleWeightsSolLoss.jld2"
    groupname = "P$(grp)G"

    all_parameters = load(data_path, groupname)
    initializations = [v[1] for v in all_parameters]
    n_sets = length(initializations)
    n_parameters = length(initializations[1])

    T = eltype(initializations[1])

    return InitializationLoader{T}(
        1,
        n_sets,
        n_parameters,
        grp,
        data_path,
        initializations
    )

end

function (il::InitializationLoader)()
    idx = il._state
    ret = il.initializations[idx]
    il._state = idx + 1
    ret
end

# ignore dimensions, in this experiment, we "know" everything.
function (il::InitializationLoader)(dims...)
    idx = il._state

    idx > il._n_sets || error("Too many iterations for number of init param sets")

    ret = il.initializations[idx]
    il._state = idx + 1
    ret
end

function callbacks_to_hdf5(data, fn::String)
    losses = hcat([cb.losses for cb in data]...)
    # Pred is a N_DIMS x T_STEPS x TRAINING_ITERATION x REPTITION
    N_DIMS = size(data[1].predictions[1])[1] # homogeniety of the data is required!
    preds = cat([cat([hcat([pred[i,:] for pred in cb.predictions]...) for cb in data]...; dims=3) for i in 1:N_DIMS]...;
               dims = 4) |> arr -> permutedims(arr, [4,1,2,3])
    params = cat([hcat([param for param in cb.parameters]...) for cb in data]...; dims=3)
    h5open(fn, "w") do file
        file["losses"] = losses
        file["predictions"] = preds
        file["parameters"] = params
    end
end

max_stars_repo_name: pnieters/GeneralizedDynamicsFromData

max_issues_repo_name: pnieters/GeneralizedDynamicsFromData

max_forks_repo_name: pnieters/GeneralizedDynamicsFromData


text:
@testset "adaptive_threshold" begin
    @info "Test: AdaptiveThreshold"

    @testset "API" begin
        img_gray = imresize(testimage("lena_gray_256"); ratio=0.25)
        img = copy(img_gray)

        # AdaptiveThreshold
        @test AdaptiveThreshold(window_size=32) == AdaptiveThreshold(32, 15)
        @test AdaptiveThreshold(img_gray) == AdaptiveThreshold(window_size=8)

        # window_size non-positive integer
        @test_throws ArgumentError AdaptiveThreshold(window_size=-10)
        @test_throws TypeError AdaptiveThreshold(window_size=32.5)
        # percentage ∈ [0, 100]
        @test_throws ArgumentError AdaptiveThreshold(window_size=32, percentage=-10)
        @test_throws ArgumentError AdaptiveThreshold(window_size=32, percentage=150)

        # binarize
        f = AdaptiveThreshold(img, percentage=15)
        binarized_img_1 = binarize(img, f)
        @test img == img_gray # img unchanged
        @test eltype(binarized_img_1) == Gray{N0f8}

        binarized_img_2 = binarize(Gray{Bool}, img, f)
        @test img == img_gray # img unchanged
        @test eltype(binarized_img_2) == Gray{Bool}

        binarized_img_3 = similar(img, Bool)
        binarize!(binarized_img_3, img, f)
        @test img == img_gray # img unchanged
        @test eltype(binarized_img_3) == Bool

        binarized_img_4 = copy(img_gray)
        binarize!(binarized_img_4, f)
        @test eltype(binarized_img_4) == Gray{N0f8}

        @test binarized_img_1 == binarized_img_2
        @test binarized_img_1 == binarized_img_3
        @test binarized_img_1 == binarized_img_4

        for T in generate_test_types([Float32, N0f8, Bool], [Gray])
            @test eltype(binarize(T, img, f)) == T
        end
    end

    @testset "Types" begin
        # Gray
        img_gray = imresize(testimage("lena_gray_256"); ratio=0.25)
        f = AdaptiveThreshold(img_gray, percentage=15)

        type_list = generate_test_types([Float32, N0f8], [Gray])
        for T in type_list
            img = T.(img_gray)
            @test_reference "References/AdaptiveThreshold_Gray.png" Gray.(binarize(img, f)) by=binarization_equality()
        end

        # Color3
        img_color = imresize(testimage("lena_color_256"); ratio=0.25)
        f = AdaptiveThreshold(img_color, percentage=15)

        type_list = generate_test_types([Float32, N0f8], [RGB, Lab])
        for T in type_list
            img = T.(img_gray)
            @test_reference "References/AdaptiveThreshold_Color3.png" Gray.(binarize(img, f)) by=binarization_equality()
        end
    end

    @testset "Numerical" begin
        # Check that the image only has ones or zeros.
        img = imresize(testimage("lena_gray_256"); ratio=0.25)
        f = AdaptiveThreshold(img, percentage=15)
        img₀₁ = binarize(img, f)
        non_zeros = findall(x -> x != 0.0 && x != 1.0, img₀₁)
        @test length(non_zeros) == 0

        # Check that ones and zeros have been assigned to the correct side of the threshold.
        maxval, maxpos = findmax(Gray.(img))
        @test img₀₁[maxpos] == 1
        minval, minpos = findmin(Gray.(img))
        @test img₀₁[minpos] == 0
    end

    @testset "Miscellaneous" begin
        img = testimage("lena_gray_256")

        # deprecations
        @test (@test_deprecated AdaptiveThreshold()) == AdaptiveThreshold(32, 15.0)
        @test (@test_deprecated AdaptiveThreshold(percentage = 15)) == AdaptiveThreshold(32, 15)
        @test (@test_deprecated recommend_size(img)) == 32
    end

end