text:
module Peaks

using Compat

export argmaxima, argminima, findmaxima, findminima, findnextmaxima, findnextminima,
    peakproms, peakproms!, peakwidths, peakwidths!

include("minmax.jl")
include("peakprom.jl")
include("peakwidth.jl")

end # module Peaks


text:
using ReTest

include(joinpath(@__DIR__, "main.jl"))

# You can run `Pkg.test("Talkon", test_args = ["foo", "bar"])` or just
# `Pkg.test(test_args = ["foo", "bar"])` to select only specific tests. If no `test_args`
# is given or you are running usual `> ] test` command, then all tests are executed.
# Strings are used as regexps and you can prepend "-" char before filter match to exclude specific subset of tests, for example `Pkg.test("BIBOAdminTools, test_args = ["-foo.*"])` execute all tests except those which starts with `foo`.
if isempty(ARGS)
    TalkonTest.runtests()
else
    TalkonTest.runtests(map(arg -> startswith(arg, "-") ? not(Regex(arg[2:end])) : Regex(arg), ARGS))
end

ext: jl


text:
using SearchSortedNearest
using Test

@testset "SearchSortedNearest.jl" begin
    @testset "Range checks" begin
        for rng in [1:10, (1:10) .+ 0.4, (1:10) .+ 0.5, 1.0:10.0], item in rng
            @test searchsortednearest(rng, item) == floor(Int, item)
        end
        @test searchsortednearest(1:10, 0) == 1
        @test searchsortednearest(1:10, 11) == 10
        @test searchsortednearest(10:-1:1, 0, rev=true) == 10
        @test searchsortednearest(10:-1:1, 11, rev=true) == 1
    end
    @testset "by" begin
        struct StructWithoutOrder end
        rng = map(x -> x => StructWithoutOrder(), 1:10)
        for item in rng
            @test searchsortednearest(rng, item, by=first) == first(item)
            @test searchsortednearest(reverse(rng), item, by=first, rev=true) == 11 - first(item)
        end
        a = ["a", "aa", "aaa", "aaaa"]
        x = "aaa"
        @test searchsortedfirst(a, x, by=length) == 3
    end
    @testset "lt" begin
        a = 1:10.0
        x = NaN
        VERSION > v"1.5.4" && @test searchsortednearest(a, x, lt=isless) == 10
        @test searchsortednearest(a, x, lt=<) == 1
    end
    @testset "distance" begin
        a = ["abc", "abd", "abe", "abz"]
        x = "abf"
        @test searchsortednearest(a, x, distance=(a,b) -> sum(map(-, a, b))) == 3
    end
    @testset "random" begin
        for _ in 1:100
            a, x = sort!(randn(100)), randn()
            @test searchsortednearest(a, x) == argmin(abs.(a .- x))
            @test searchsortednearest(reverse(a), x, rev=true) == argmin(abs.(reverse(a) .- x))
        end
    end
end


text:
function newton_fixpoint(f::F,x::T,atol) where {F,T}
    fx, gx = f(x)
    abs(fx) < atol && return x
    Δx = fx/gx
    x =x - Δx
    return x
end

"""
    ad_newton(f,x0;atol = oneunit(T) * (eps(one(T)))^(4/5),rtol = )

    Performs newton root finding (via Roots.jl), but the derivatives are calculated with ForwardDiff.jl
    kwargs are passed to `Roots.find_zero`
"""
function ad_newton(f::F,x0::T;
    rtol= (eps(one(T)))^(4/5),
    atol=rtol*oneunit(T),
    max_iters = 100) where {F,T}
    fdf(x) =  f∂f(f,x)
    f0(x) = newton_fixpoint(fdf,x,atol)
    return fixpoint(f0,x0;rtol,atol,max_iters)
end

function newton(f::F,x0::T;
    rtol= (eps(one(T)))^(4/5),
    atol=rtol*oneunit(T),
    max_iters = 100) where {F,T}
    f0(x) = newton_fixpoint(f,x,atol)
    return fixpoint(f0,x0;rtol,atol,max_iters)
end

function halley_fixpoint(f::F,x::T,atol) where {F,T}
    ff,gg,hh = f(x)
    abs(ff) < atol && return x
    Δx = ff/gg/(1-ff*hh/(2*gg^2))
    x =x - Δx
    return x
end

function halley(f::F,x0::T;
    rtol= (eps(one(T)))^(4/5),
    atol=rtol*oneunit(T),
    max_iters = 100) where {F,T}
    f0(x) = halley_fixpoint(f,x,atol)
    return fixpoint(f0,x0;rtol,atol,max_iters)
end


text:
module TestID

import Historic
using Test

function test_taskid()
    global task1 = @task nothing
    global task2 = @task nothing
    @test Historic.taskid(task1) === Historic.taskid(task1)
    @test Historic.taskid(task1) !== Historic.taskid(task2)
end

function test_objid()
    global ref1 = Ref(0)
    global ref2 = Ref(0)
    @test Historic.objid(ref1) === Historic.objid(ref1)
    @test Historic.objid(ref1) !== Historic.objid(ref2)
    @test Historic.objid(1) === Historic.objid(1)
end

end  # module


text:
module TestModule

using Test
using MLStyle

macro test_macro_throws(errortype, m)
    :(
        @test_throws $errortype try
                @eval $m
            catch err
                while err isa LoadError
                    err = err.error
                end
                throw(err)
            end
    )
end

MODULE = TestModule

@use GADT
include("uncomp.jl")
include("when.jl")
include("MQuery/test.jl")
include("modules/cond.jl")
include("modules/ast.jl")
include("as_record.jl")
include("adt.jl")
include("active_patterns.jl")
include("exception.jl")
include("render.jl")
include("pervasive.jl")
include("expr_template.jl")
include("gallery/simple.jl")
include("match.jl")
include("pattern.jl")
include("dot_expression.jl")
include("typelevel.jl")
include("untyped_lam.jl")
include("nothing.jl")
include("lambda.jl")

end


text:
using BinaryBandit

horizon = 60

BinaryBandit.DP_2_action_lin( horizon , 64 )
BinaryBandit.DP_2_action_lin( horizon , 32 )
BinaryBandit.DP_2_action_lin( horizon , 16 )

BinaryBandit.DP_2_action_lin_with_administration( horizon , 0 , 64 )
BinaryBandit.DP_2_action_lin_with_administration( horizon , 0 , 32 )
BinaryBandit.DP_2_action_lin_with_administration( horizon , 0 , 16 )

BinaryBandit.DP_2_action( horizon , 64 )
BinaryBandit.DP_2_action( horizon , 32 )
BinaryBandit.DP_2_action( horizon , 16 )

BinaryBandit.DP_2_action_with_administration( horizon , 0 , 64 )
BinaryBandit.DP_2_action_with_administration( horizon , 0 , 32 )
BinaryBandit.DP_2_action_with_administration( horizon , 0 , 16 )

BinaryBandit.DP_2_policy_bin_lin( horizon )
BinaryBandit.DP_2_policy_bin_lin_with_administration( horizon , 0 )

BinaryBandit.DP_2_policy_bin( horizon )
BinaryBandit.DP_2_policy_bin_with_administration( horizon , 0 )

BinaryBandit.DP_2_policy_lin_new( horizon )

BinaryBandit.DP_2_policy_lin( horizon )
BinaryBandit.DP_2_policy_lin_with_administration( horizon , 0 )

BinaryBandit.DP_2_policy( horizon )
BinaryBandit.DP_2_policy_with_administration( horizon , 0 )

success_probability_arm_1 = 0.3
success_probability_arm_2 = 0.5

BinaryBandit.DP_2_NS( horizon , success_probability_arm_1 , success_probability_arm_2 )

max_stars_repo_name: PeterJacko/BinaryBandit

max_issues_repo_name: PeterJacko/BinaryBandit

max_forks_repo_name: PeterJacko/BinaryBandit


text:
"""
    StatsModelComparisons

This module implements several model comparison methods such as PSIS, WAIC and DIC.

Included functions
------------------
psisloo
    Pareto smoothed importance sampling leave-one-out log predictive densities.

psislw
    Pareto smoothed importance sampling.

gpdfitnew
    Estimate the paramaters for the Generalized Pareto Distribution (GPD).

gpinv
    Inverse Generalised Pareto distribution function.

waic
    ...

dic
    ...

deviance
    ...

aic
    Akaike Information Criterion

bic
    Bayesian Information Criterion

References
----------
Aki Vehtari, Andrew Gelman and Jonah Gabry (2015). Efficient implementation
of leave-one-out cross-validation and WAIC for evaluating fitted Bayesian
models. arXiv preprint arXiv:1507.04544.

Aki Vehtari and Andrew Gelman (2015). Pareto smoothed importance sampling.
arXiv preprint arXiv:1507.02646.
"""
module StatsModelComparisons

using RecipesBase
using StatsFuns
using Statistics

include("psisloo.jl")
include("psislw.jl")
include("gpdfitnew.jl")
include("gpinv.jl")
include("waic.jl")
include("pk_utilities.jl")
include("dic.jl")
include("legacy.jl")

export deviance, dic, psisloo, gpdfitnew, gpinv
export var2, waic, aic, bic
export psislw
export pk_qualify
end


text:
using ClimateMachine

ClimateMachine.init()

using ClimateMachine.Ocean.OceanProblems: InitialConditions, InitialValueProblem

@testset "$(@__FILE__)" begin
    U = 0.1
    L = 0.2
    a = 0.3
    U = 0.4

    Ψ(x, L) = exp(-x^2 / 2 * L^2) # a Gaussian

    uᵢ(x, y, z) = +U * y / L * Ψ(x, L)
    vᵢ(x, y, z) = -U * x / L * Ψ(x, L)
    ηᵢ(x, y, z) = a * Ψ(x, L)
    θᵢ(x, y, z) = 20.0 + 1e-3 * z

    initial_conditions = InitialConditions(u = uᵢ, v = vᵢ, η = ηᵢ, θ = θᵢ)

    problem = InitialValueProblem{Float64}(
        dimensions = (π, 42, 1.1),
        initial_conditions = initial_conditions,
    )

    @test problem.Lˣ == Float64(π)
    @test problem.Lʸ == 42.0
    @test problem.H == 1.1

    @test problem.initial_conditions.u === uᵢ
    @test problem.initial_conditions.v === vᵢ
    @test problem.initial_conditions.η === ηᵢ
    @test problem.initial_conditions.θ === θᵢ
end


text:
"""
    SimpleSDMPredictor(::Type{WorldClim}, ::Type{BioClim}, layer::Integer=1; resolution::Float64=10.0, left=nothing, right=nothing, bottom=nothing, top=nothing)

Download and prepare a WorldClim 2.1 bioclimatic variable, and returns it as an
`SimpleSDMPredictor`. Layers are called by their number, from 1 to 19. The list
of available layers is given in a table below.

The keywords are `resolution`, which must be a floating point value, and either
`2.5`, `5.0`, or `10.0`, as well as optionally `left`, `right`, `bottom`, and
`top`, which will allow to only load parts of the data.

Internally, this function will download the main zip file for the required
resolution from the WordlClim website, extract it, and parse the required
layers.

It is recommended to *keep* the content of the `path` folder, as it will
eliminate the need to download and/or extract the tiff files. For example,
calling `wordlclim(1:19)` will download and extract everything, and future calls
will be much faster.

| Variable | Description                                                |
| ------   | ------                                                     |
| 1        | Annual Mean Temperature                                    |
| 2        | Mean Diurnal Range (Mean of monthly (max temp - min temp)) |
| 3        | Isothermality (BIO2/BIO7) (* 100)                          |
| 4        | Temperature Seasonality (standard deviation *100)          |
| 5        | Max Temperature of Warmest Month                           |
| 6        | Min Temperature of Coldest Month                           |
| 7        | Temperature Annual Range (BIO5-BIO6)                       |
| 8        | Mean Temperature of Wettest Quarter                        |
| 9        | Mean Temperature of Driest Quarter                         |
| 10       | Mean Temperature of Warmest Quarter                        |
| 11       | Mean Temperature of Coldest Quarter                        |
| 12       | Annual Precipitation                                       |
| 13       | Precipitation of Wettest Month                             |
| 14       | Precipitation of Driest Month                              |
| 15       | Precipitation Seasonality (Coefficient of Variation)       |
| 16       | Precipitation of Wettest Quarter                           |
| 17       | Precipitation of Driest Quarter                            |
| 18       | Precipitation of Warmest Quarter                           |
| 19       | Precipitation of Coldest Quarter                           |

Original data: https://www.worldclim.org/data/worldclim21.html
"""
function SimpleSDMPredictor(::Type{WorldClim}, ::Type{BioClim}, layer::Integer=1; resolution::Float64=10.0, kwargs...)
    @assert resolution in [0.5, 2.5, 5.0, 10.0]
    file = _get_raster(WorldClim, BioClim, layer, resolution)
    return geotiff(SimpleSDMPredictor, file; kwargs...)
end

function SimpleSDMPredictor(::Type{WorldClim}, ::Type{BioClim}, layers::AbstractArray; kwargs...)
    @assert eltype(layers) <: Integer
    return [SimpleSDMPredictor(WorldClim, BioClim, l; kwargs...) for l in layers]
end

"""
    SimpleSDMPredictor(::Type{WorldClim}, ::Type{BioClim}, mod::CMIP6, fut::SharedSocioeconomicPathway, layer::Integer=1; year="2021-2040", resolution::Float64=10.0, kwargs...)

Future biolcim data, where year is in "2021-2040", "2041-2060", "2061-2080", "2081-2100" .
"""
function SimpleSDMPredictor(::Type{WorldClim}, ::Type{BioClim}, mod::CMIP6, fut::SharedSocioeconomicPathway, layer::Integer=1; year="2021-2040", resolution::Float64=10.0, kwargs...)
    @assert year in ["2021-2040", "2041-2060", "2061-2080", "2081-2100"]
    @assert resolution in [2.5, 5.0, 10.0]
    file = _get_raster(WorldClim, BioClim, mod, fut, resolution, year)
    return geotiff(SimpleSDMPredictor, file, layer; kwargs...)
end

function SimpleSDMPredictor(::Type{WorldClim}, ::Type{BioClim}, mod::CMIP6, fut::SharedSocioeconomicPathway, layers::AbstractArray; kwargs...)
    @assert eltype(layers) <: Integer
    return [SimpleSDMPredictor(WorldClim, BioClim, mod, fut, l; kwargs...) for l in layers]
end


text:
## activate project environment
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

## load packages
using DelimitedFiles, LinearAlgebra, Statistics
using BenchmarkTools
using StaticArrays
using StrideArrays
using Trixi

struct PrimitiveVariablesCompressibleEuler2D{T<:Real} <: FieldVector{4, T}
  rho::T
  v1 ::T
  v2 ::T
  p  ::T
end

@inline function Trixi.cons2prim(u::PrimitiveVariablesCompressibleEuler2D,
                                 equations::CompressibleEulerEquations2D)
  return u
end

struct PrimitiveVariablesCompressibleEuler3D{T<:Real} <: FieldVector{5, T}
  rho::T
  v1 ::T
  v2 ::T
  v3 ::T
  p  ::T
end

@inline function Trixi.cons2prim(u::PrimitiveVariablesCompressibleEuler3D,
                                 equations::CompressibleEulerEquations3D)
  return u
end

@inline function Trixi.get_node_vars(u::AbstractArray{<:StaticVector},
                                     equations, dg::DG, indices...)
  return u[indices...]
end

function cons2prim!(prim::AbstractVector{<:StaticVector},
                    u::AbstractArray{<:Real, 2},
                    equations, solver)
  for i in eachindex(prim)
    u_local = Trixi.get_node_vars(u, equations, solver, i)
    prim_local = cons2prim(u_local, equations)
    prim[i] = prim_local
  end
end

function cons2prim!(_prim::AbstractVector{<:PrimitiveVariablesCompressibleEuler2D},
                    u::AbstractArray{<:Real, 2},
                    equations::CompressibleEulerEquations2D, solver)
  prim = reinterpret(reshape, eltype(eltype(_prim)), _prim)

  Trixi.@turbo for i in Trixi.indices((prim, u), 2)
    rho    = u[1, i]
    rho_v1 = u[2, i]
    rho_v2 = u[3, i]
    rho_e  = u[4, i]

    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    p = (equations.gamma - 1) * (rho_e - 0.5 * (rho_v1 * v1 + rho_v2 * v2))

    prim[1, i] = rho
    prim[2, i] = v1
    prim[3, i] = v2
    prim[4, i] = p
  end
end

function cons2prim!(_prim::AbstractVector{<:PrimitiveVariablesCompressibleEuler3D},
                    u::AbstractArray{<:Real, 2},
                    equations::CompressibleEulerEquations3D, solver)
  prim = reinterpret(reshape, eltype(eltype(_prim)), _prim)

  Trixi.@turbo for i in Trixi.indices((prim, u), 2)
    rho    = u[1, i]
    rho_v1 = u[2, i]
    rho_v2 = u[3, i]
    rho_v3 = u[4, i]
    rho_e  = u[5, i]

    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    v3 = rho_v3 / rho
    p = (equations.gamma - 1) * (rho_e - 0.5 * (rho_v1 * v1 + rho_v2 * v2 + rho_v3 * v3))

    prim[1, i] = rho
    prim[2, i] = v1
    prim[3, i] = v2
    prim[4, i] = v3
    prim[5, i] = p
  end
end

## set up benchmark code
"""
    initial_condition_isentropic_vortex(x, t, equations)

The classical isentropic vortex test case of
- Chi-Wang Shu (1997)
  Essentially Non-Oscillatory and Weighted Essentially Non-Oscillatory
  Schemes for Hyperbolic Conservation Laws.
  [NASA/CR-97-206253](https://ntrs.nasa.gov/citations/19980007543)
"""
function initial_condition_isentropic_vortex(x, t, equations::CompressibleEulerEquations2D)
  ϱ0 = 1.0               # background density
  v0 = SVector(1.0, 1.0) # background velocity
  p0 = 10.0              # background pressure
  ε = 20.0               # vortex strength
  L = 10.0               # size of the domain per coordinate direction

  T0 = p0 / ϱ0           # background temperature
  γ = equations.gamma    # ideal gas constant

  x0 = v0 * t            # current center of the vortex
  dx = vortex_center.(x - x0, L)
  r2 = sum(abs2, dx)

  # perturbed primitive variables
  T = T0 - (γ - 1) * ε^2 / (8 * γ * π^2) * exp(1 - r2)
  v = v0 + ε / (2 * π) * exp(0.5 * (1 - r2)) * SVector(-dx[2], dx[1])
  ϱ = ϱ0 * (T / T0)^(1 / (γ - 1))
  p = ϱ * T

  return prim2cons(SVector(ϱ, v..., p), equations)
end

function initial_condition_isentropic_vortex(x, t, equations::CompressibleEulerEquations3D)
  ϱ0 = 1.0               # background density
  v0 = SVector(1.0, 1.0, 0.0) # background velocity
  p0 = 10.0              # background pressure
  ε = 20.0               # vortex strength
  L = 10.0               # size of the domain per coordinate direction

  T0 = p0 / ϱ0           # background temperature
  γ = equations.gamma    # ideal gas constant

  x0 = v0 * t            # current center of the vortex
  dx = vortex_center.(x - x0, L)
  r2 = sum(abs2, dx)

  # perturbed primitive variables
  T = T0 - (γ - 1) * ε^2 / (8 * γ * π^2) * exp(1 - r2)
  v = v0 + ε / (2 * π) * exp(0.5 * (1 - r2)) * SVector(-dx[2], dx[1], 0.0)
  ϱ = ϱ0 * (T / T0)^(1 / (γ - 1))
  p = ϱ * T

  return prim2cons(SVector(ϱ, v..., p), equations)
end

vortex_center(x, L) = mod(x + L/2, L) - L/2

function split_form_prim_kernel!(du, u_prim, u_cons,
                                 element, mesh,
                                 nonconservative_terms, equations,
                                 volume_flux::VF, dg, cache) where {VF}

  cons2prim!(vec(u_prim), reshape(u_cons, nvariables(equations), :),
             equations, dg)

  Trixi.split_form_kernel!(du, u_prim, element, mesh,
                           nonconservative_terms, equations,
                           volume_flux, dg, cache)

  return nothing
end

# 2D version
function run_benchmarks_2d(polydeg, numerical_flux)
  println("Benchmarking ", numerical_flux, " in 2D with polydeg=", polydeg)

  # setup basic code
  equations = CompressibleEulerEquations2D(1.4)

  coordinates_min = (-5.0, -5.0)
  coordinates_max = ( 5.0,  5.0)
  mesh = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=0, n_cells_max=100_000, periodicity=true)

  solver = DGSEM(polydeg=polydeg, surface_flux=numerical_flux,
                 volume_integral=VolumeIntegralFluxDifferencing(numerical_flux))

  semi = SemidiscretizationHyperbolic(mesh, equations,
                                      initial_condition_isentropic_vortex, solver)

  u_ode  = compute_coefficients(0.0, semi)
  du_ode = similar(u_ode)
  GC.@preserve u_ode du_ode begin
    u  = Trixi.wrap_array(u_ode,  semi)
    du = Trixi.wrap_array(du_ode, semi)
    u_prim = StrideArray(undef,
                         PrimitiveVariablesCompressibleEuler2D{eltype(u)},
                         Base.tail(StrideArrays.ArrayInterface.size(u))...)
    element = 1

    # benchmark flux differencing volume terms based on conservative variables
    fill!(du, zero(eltype(du)))
    Trixi.split_form_kernel!(
      du, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_cons = copy(du)

    sleep(0.5)
    result_cons = @benchmark Trixi.split_form_kernel!(
      $du, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_cons)

    # benchmark flux differencing volume terms based on primitive variables
    fill!(du, zero(eltype(du)))
    split_form_prim_kernel!(
      du, u_prim, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_prim = copy(du)
    @assert du_cons ≈ du_prim

    sleep(0.5)
    result_prim = @benchmark split_form_prim_kernel!(
      $du, $u_prim, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_prim)
  end
  println()

  return result_cons, result_prim
end

# 3D version
function run_benchmarks_3d(polydeg, numerical_flux)
  println("Benchmarking ", numerical_flux, " in 3D with polydeg=", polydeg)

  # setup basic code
  equations = CompressibleEulerEquations3D(1.4)

  coordinates_min = (-5.0, -5.0, -5.0)
  coordinates_max = ( 5.0,  5.0,  5.0)
  mesh = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=0, n_cells_max=100_000, periodicity=true)

  solver = DGSEM(polydeg=polydeg, surface_flux=numerical_flux,
                 volume_integral=VolumeIntegralFluxDifferencing(numerical_flux))

  semi = SemidiscretizationHyperbolic(mesh, equations,
                                      initial_condition_isentropic_vortex, solver)

  u_ode  = compute_coefficients(0.0, semi)
  du_ode = similar(u_ode)
  GC.@preserve u_ode du_ode begin
    u  = Trixi.wrap_array(u_ode,  semi)
    du = Trixi.wrap_array(du_ode, semi)
    u_prim = StrideArray(undef,
                         PrimitiveVariablesCompressibleEuler3D{eltype(u)},
                         Base.tail(StrideArrays.ArrayInterface.size(u))...)
    element = 1

    # benchmark flux differencing volume terms based on conservative variables
    fill!(du, zero(eltype(du)))
    Trixi.split_form_kernel!(
      du, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_cons = copy(du)

    sleep(0.5)
    result_cons = @benchmark Trixi.split_form_kernel!(
      $du, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_cons)

    # benchmark flux differencing volume terms based on primitive variables
    fill!(du, zero(eltype(du)))
    split_form_prim_kernel!(
      du, u_prim, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_prim = copy(du)
    @assert du_cons ≈ du_prim

    sleep(0.5)
    result_prim = @benchmark split_form_prim_kernel!(
      $du, $u_prim, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_prim)
  end
  println()

  return result_cons, result_prim
end

# general code
Base.ndims(::typeof(run_benchmarks_2d)) = 2
Base.ndims(::typeof(run_benchmarks_3d)) = 3

function run_benchmarks(run_benchmarks_inner, polydegs, numerical_flux)
  results = zeros(length(polydegs), 1)
  results .= polydegs
  header = "# Polydeg"

  # flux differencing
  println("#"^80)
  println("flux differencing\n")
  means_cons = zeros(length(polydegs))
  stds_cons  = zero(means_cons)
  means_prim = zero(means_cons)
  stds_prim  = zero(means_cons)
  for (i, polydeg) in enumerate(polydegs)
    result_cons, result_prim = run_benchmarks_inner(polydeg, numerical_flux)
    means_cons[i] = time(mean(result_cons)) * 1.0e-9
    stds_cons[i]  = time(std(result_cons))  * 1.0e-9
    means_prim[i] = time(mean(result_prim)) * 1.0e-9
    stds_prim[i]  = time(std(result_prim))  * 1.0e-9
  end
  results = hcat(results, means_cons, stds_cons, means_prim, stds_prim)
  header = header * "\tConservative-variables-mean\tstd" *
                    "\tPrimitive-variables-mean\tstd"

  # write results to a file
  open(joinpath(@__DIR__, "primitive_$(ndims(run_benchmarks_inner))D_$(numerical_flux).dat"), "w") do io
    println(io, header)
    writedlm(io, results)
  end

  return results
end

## run benchmarks
polydegs = 3:15
run_benchmarks(run_benchmarks_2d, polydegs, flux_shima_etal)
run_benchmarks(run_benchmarks_2d, polydegs, flux_ranocha)

run_benchmarks(run_benchmarks_3d, polydegs, flux_shima_etal)
run_benchmarks(run_benchmarks_3d, polydegs, flux_ranocha)


text:
#  Compute r = A * x + s - b in-place
function primal_kkt_condition!(r::AbstractVector{T}, A::AbstractMatrix{T}, x::AbstractVector{T}, s::SplitVector{T}, b::AbstractVector{T}) where {T <: AbstractFloat}
	# r = A * x
	mul!(r, A, x)
	@. r += s
	@. r -= b
	return nothing
end

# Compute r = P * x + q - A' * μ in-place
function dual_kkt_condition!(r::AbstractVector{T}, r_temp::AbstractVector{T}, P::AbstractMatrix{T}, x::AbstractVector{T}, q::AbstractVector{T}, A::AbstractMatrix{T}, μ::AbstractVector{T}) where {T <: AbstractFloat}
	mul!(r, P, x)
	@. r += q

	mul!(r_temp, A', μ)
	@. r -= r_temp
	return nothing
end

# Einv * x
unscale_primal!(x::AbstractVector{T}, Einv::AbstractMatrix{T}) where {T} = lmul!(Einv, x)

# cinv * Dinv * x
function unscale_dual!(x::AbstractVector{T}, Dinv::AbstractMatrix{T}, cinv) where {T}
		lmul!(Dinv, x)
		@. x *= cinv
		return nothing
end

function calculate_residuals!(ws::COSMO.Workspace,  IGNORESCALING_FLAG::Bool = false)
	# reuse vectors from temporary workspace
	r_prim = ws.utility_vars.vec_m
	r_dual = ws.utility_vars.vec_n
	r_temp = ws.utility_vars.vec_n2

	# r_prim = A * x + s - b
	primal_kkt_condition!(r_prim, ws.p.A, ws.vars.x, ws.vars.s, ws.p.b)

	# ∇f0 + ∑ νi ∇hi(x*) == 0 condition
	# r_dual = P * x + q - A' * μ
	dual_kkt_condition!(r_dual, r_temp, ws.p.P, ws.vars.x, ws.p.q, ws.p.A, ws.vars.μ)

	if (ws.settings.scaling != 0 && !IGNORESCALING_FLAG)
		# unscale primal residual: r_prim = Einv * ( A * x + s - b)
		unscale_primal!(r_prim, ws.sm.Einv)

		# unscale dual residual: r_dual = cinv * Dinv * (P * x + q - A' * μ)
		unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
	end

	return norm(r_prim, Inf), norm(r_dual, Inf)

end

function max_res_component_norm(ws::COSMO.Workspace, IGNORESCALING_FLAG::Bool = false)
		r_prim = ws.utility_vars.vec_m
		r_dual = ws.utility_vars.vec_n

		# Determine if the components inside max{ } have to be unscaled
		unscale = (ws.settings.scaling != 0 && !IGNORESCALING_FLAG)

		# Calculate: max { ||Einv * Ax ||_Inf, ||Einv * s||_Inf, ||Einv * b||_Inf }
		# Einv * A * x
		mul!(r_prim, ws.p.A, ws.vars.x)
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = norm(r_prim, Inf)

		# Einv * s
		@. r_prim = ws.vars.s
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = max.(max_norm_prim, norm(r_prim, Inf))

		# Einv * b
		@. r_prim = ws.p.b
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = max.(max_norm_prim, norm(r_prim, Inf))

		# Calculate: max { ||cinv * Dinv * P * x ||_Inf, ||cinv * Dinv * q||_Inf, ||cinv * Dinv * A' * μ||_Inf }
		# ||cinv * Dinv * P * x ||_Inf
		mul!(r_dual, ws.p.P, ws.vars.x)
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = norm(r_dual, Inf)

		# ||cinv * Dinv * q||_Inf
		@. r_dual = ws.p.q
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = max.(max_norm_dual, norm(r_dual, Inf))

		# ||cinv * Dinv * A' * μ||_Inf
		mul!(r_dual, ws.p.A', ws.vars.μ)
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = max.(max_norm_dual, norm(r_dual, Inf))

	return max_norm_prim, max_norm_dual
end

function isapprox_feasible(residual, max_norm, eps_abs, eps_rel)
	ϵ = eps_abs + eps_rel * max_norm
	return residual < ϵ
end

function isapprox_primal_feasible(r::ResultInfo, eps_abs, eps_rel)
    return isapprox_feasible(r.r_prim, r.max_norm_prim, eps_abs, eps_rel)
end

function isapprox_dual_feasible(r::ResultInfo, eps_abs, eps_rel)
    return isapprox_feasible(r.r_dual, r.max_norm_dual, eps_abs, eps_rel)
end

function is_primal_feasible(r::ResultInfo, settings::Settings)
	return isapprox_primal_feasible(r, settings.eps_abs, settings.eps_rel)
end

function is_dual_feasible(r::ResultInfo, settings::Settings)
	return isapprox_dual_feasible(r, settings.eps_abs, settings.eps_rel)
end

function is_primal_nearly_feasible(r::ResultInfo, settings::Settings)
	return isapprox_primal_feasible(r, settings.nearly_ratio * settings.eps_abs, settings.nearly_ratio * settings.eps_rel)
end

function is_dual_nearly_feasible(r::ResultInfo, settings::Settings)
	return isapprox_dual_feasible(r, settings.nearly_ratio * settings.eps_abs, settings.nearly_ratio * settings.eps_rel)
end

function has_converged(ws::COSMO.Workspace{T}, r::ResultInfo) where {T <: AbstractFloat}
	# check activation of accelerator
	COSMO.check_activation!(ws, ws.activation_reason, r)

	# if an optimal objective value was specified for the problem check if current solution is within specified accuracy
	obj_true_FLAG = true
	settings = ws.settings
	if !isnan(settings.obj_true)
		current_cost = calculate_cost!(ws.utility_vars.n, ws.vars.x, ws.p.P, ws.p.q, ws.sm.cinv[])
		obj_true_FLAG = abs(settings.obj_true - current_cost) <= settings.obj_true_tol
	end

	return is_primal_feasible(r, settings) && is_dual_feasible(r, settings) && obj_true_FLAG
end

# cost = cinv *( 1/2 x' * P * x + q' x)
function calculate_cost!(temp::AbstractVector{T}, x::AbstractVector{T}, P::SparseMatrixCSC{T, Int}, q::AbstractVector{T}, cinv::T = one(T)) where {T <: AbstractFloat}
	# P x
	mul!(temp, P, x)
	return cinv * (T(0.5) * dot(temp, x) + dot(q, x))
end

function calculate_result_info!(ws::COSMO.Workspace)
    r_prim, r_dual = calculate_residuals!(ws)
    max_norm_prim, max_norm_dual = max_res_component_norm(ws)
    return ResultInfo(r_prim, r_dual, max_norm_prim, max_norm_dual, ws.rho_updates)
end


text:
module XfromProjections

using SparseArrays
using LinearAlgebra
using PyCall
using Suppressor
using LinearOperators
using Logging

# analytic
include("analytic/filter_proj.jl")
export filter_proj, bp_slices
# include("analytic/gridrec.jl")
# export recon2d_gridrec

# iterative
include("iterative/util_convexopt.jl")
include("iterative/tv_primaldual.jl")
include("iterative/tv_primaldual_flow.jl")
include("iterative/ctv_primaldual.jl")
include("iterative/sirt.jl")
export recon2d_sirt!, recon2d_stack_sirt!, _compute_sum_rows_cols
export recon2d_tv_primaldual!, recon2d_stack_tv_primaldual!, recon2d_ctv_primaldual!, recon2d_tv_primaldual_flow
export get_flows, compute_warping_operator

end

max_stars_repo_name: JuliaTomo/DynamicTomographyWithFlow

max_issues_repo_name: JuliaTomo/DynamicTomographyWithFlow

max_forks_repo_name: JuliaTomo/DynamicTomographyWithFlow


text:
function CholeskyDecomposer(arg0::jdouble, arg1::jdouble)
    return CholeskyDecomposer((jdouble, jdouble), arg0, arg1)
end

function decompose(obj::CholeskyDecomposer, arg0::RealMatrix)
    return jcall(obj, "decompose", DecompositionSolver, (RealMatrix,), arg0)
end


text:
using OmegaCore, Distributions, LinearAlgebra, Test, OmegaTest, OmegaDistributions

function test_mem()
  x = dimsnth(Normal(0, 1), (3, 3))
  h(x) = (println("call!"); svd(x).S)
  y(ω) = h(x(ω))
  y_ = Variable(y)
  vars(ω) = (y_(ω), y_(ω)*10, y_(ω)*20)
  ω = defω()
  memvars = mem(vars)
  ω = defω()
  memvars(ω)
end

function test_mem()
  x = dimsnth(Normal(0, 1), (1000, 1000))
  h(x) = (println("call!"); svd(x).S)
  y(ω) = h(x(ω))
  y_ = Variable(y)
  vars(ω) = (y_(ω), y_(ω)*10, y_(ω)*20)
  ω = defω()
  memvars = mem(vars)
  # @show vars(ω)
  # @show memvars(ω)
  @test vars(ω) == memvars(ω)
  # @test isinferred(vars, ω)
  # @test isinferred(memvars, ω)
  ω = defω()
  _, t1 = @timed vars(ω)
  _, t2 = @timed memvars(ω)
  # @test t1 == t2/3
end

@testset "Mem" begin
  test_mem()
end

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
@testset "test_conversion" begin

# Easy second order system
sys1 = ss([-1 0;1 1],[1;0],[1 1],0)
G1 = tf([1,0],[1,0,-1])
H1 = zpk([0.0], [-1, 1.0], 1.0)

w = 10.0 .^ (-2:2:50)

@test ss(sys1) == sys1
@test freqresp(ss(G1), w) ≈ freqresp(sys1, w) rtol=1e-15
@test freqresp(ss(H1), w) ≈ freqresp(sys1, w) rtol=1e-15

# Test range in freqresp
@test freqresp(ss(G1), collect(1:2:10)) == freqresp(ss(G1), collect(1:2:10))

@test tf(sys1) ≈ G1 rtol=1e-15
@test tf(G1) == G1
@test tf(H1) == G1

@test zpk(sys1) ≈ H1 rtol=1e-15
@test zpk(G1) == H1
@test zpk(H1) == H1

# With direct term second order system
sys2 = ss([-2.5 0;1 1.5],[1;3],[1 2],2.5)
G2 = tf([2.5, 9.5, 6.125],[1, 1, -3.75])
H2 = zpk([-2.97703296142690, -0.82296703857310], [1.5, -2.5], 2.5)

w = 10.0 .^ (-2:2:50)

@test ss(sys2) == sys2
@test freqresp(ss(G2), w) ≈ freqresp(sys2, w) rtol=1e-15
@test freqresp(ss(H2), w) ≈ freqresp(sys2, w) rtol=1e-15

@test tf(sys2) ≈ G2 rtol=1e-15
@test tf(G2) == G2
@test tf(H2) ≈ G2 rtol=1e-15

@test zpk(sys2) ≈ H2 rtol=1e-15
@test zpk(G2) ≈ H2 rtol=1e-15
@test zpk(H2) == H2

# Non-proper system
G3 = tf([4.0, 1, 5],[2, 3])
H3 = zpk([(-1+im*sqrt(79))/8, (-1-im*sqrt(79))/8], [-3/2], 2)

@test tf(G3) == G3
@test tf(H3) ≈ G3 rtol=1e-2

@test zpk(G3) ≈ H3 rtol=1e-15
@test zpk(H3) == H3

# Test complex 1
A = [1.0 + im 1; 0 -2-3im]
B = [0;2]
C = [1 0]
D = zeros(1,1);
sys = ss(A,B,C,D)
sys2 = convert(TransferFunction, sys)
w = 10.0 .^ (-2:2:50)
@test freqresp(sys, w) ≈ freqresp(sys2, w)
csort = v -> sort(v, lt = (x,y) -> abs(x) < abs(y))
@test csort(pole(zpk(sys2))) ≈ [1+im, -2.0-3im]

# Test complex 2
sys4 = ss([-1.0+im 0;1 1],[1;0],[1 im],im)
tf(sys4)
G4 = tf([1.0im,2,-2],[1,-1im,-1+im])
H4 = zpk([im*(1 + sqrt(1+2im)), im*(1 - sqrt(1+2im))], [-1+im,1], 1.0im)

w = 10.0 .^ (-2:2:50)

@test ss(sys4) == sys4
@test freqresp(ss(G4), w) ≈ freqresp(sys4, w) rtol=1e-15
@test freqresp(ss(H4), w) ≈ freqresp(sys4, w) rtol=1e-15

@test tf(sys4) ≈ G4 rtol=1e-15
@test tf(G4) == G4
@test tf(H4) ≈ G4 rtol=1e-15

@test zpk(sys4) ≈ H4 rtol=1e-15
@test zpk(G4) ≈ H4 rtol=1e-15
@test zpk(H4) == H4

# Test conversion of Int system with direct term
w = [0.5, 1, 1.5, 2.0]
G = tf([2, 3],[5, 1])
@test freqresp(ss(G), w) ≈ freqresp(G, w)

sys2 = ss([-1 0;1 1],[1;0],[1 1],1)
G2 = tf([1,1,-1],[1,0,-1])
sys3 = ss([0 1;1 0],[0;1],[1 1],0)
G3 = tf([1,1],[1,0,-1])

#Test ss2tf
@test tf(sys2) ≈ G2 atol=1e-15
@test tf(sys3) ≈ G3 atol=1e-15

#Test TransferFunction -> StateSpace
#ss is not unique so go back and forth
@test G2 ≈ tf(ss(G2)) atol=1e-15
@test G3 ≈ tf(ss(G3)) atol=1e-15

### test mimo
@test hinfnorm(zpk([sys1 sys2])-zpk([ss(sys1) ss(sys2)]))[1] < 1e-12

## Test some BigFloat
SSBigFloat = StateSpace{BigFloat,Array{BigFloat,2}}
ZpkBigFloat = TransferFunction{ControlSystems.SisoZpk{BigFloat,Complex{BigFloat}}}
RationalBigFloat = TransferFunction{ControlSystems.SisoRational{BigFloat}}

s = tf("s")
f = zpk(1.0*(2s+3)/((5s+7)*(11s+13)))
fb = BigFloat(1.0)*f

# Reasonable defaults
# @test fb isa RationalBigFloat # Do we want this? We get TransferFunction{ControlSystems.SisoZpk{BigFloat,Complex{Float64}}}
@test ss(fb) isa SSBigFloat
@test tf(fb) isa RationalBigFloat
# Cant compute poles:
# @test zpk(tf(fb)) isa SSBigFloat

@test convert(SSBigFloat, fb) isa SSBigFloat
@test convert(ZpkBigFloat, fb) isa ZpkBigFloat
@test convert(RationalBigFloat, fb) isa RationalBigFloat

## Test some rather trivial conversions of numeric types

b = 1.5
D22 = [1.0 2.0; 3.0 4.0]

@test convert(StateSpace{Float64,Matrix{Float64}}, D22) == ss(D22)
@test convert(StateSpace{Float64,Matrix{Float64}}, b) == ss(b)

@test convert(TransferFunction{ControlSystems.SisoRational{Float64}}, D22) == tf(D22)
@test convert(TransferFunction{ControlSystems.SisoRational{Float64}}, b) == tf(b)

@test convert(TransferFunction{ControlSystems.SisoZpk{Float64,ComplexF64}}, D22) == zpk(D22)
@test convert(TransferFunction{ControlSystems.SisoZpk{Float64,ComplexF64}}, b) == zpk(b)

# Error, not proper
@test_throws ErrorException ss(tf([1,0],[1]))

end


text:
### Network models

struct CopperPlatePowerModel <: PM.AbstractActivePowerModel end
struct AreaBalancePowerModel <: PM.AbstractActivePowerModel end
abstract type AbstractPTDFModel <: PM.AbstractDCPModel end
struct StandardPTDFModel <: AbstractPTDFModel end
struct PTDFPowerModel <: AbstractPTDFModel end

#================================================
    # exact non-convex models
    ACPPowerModel, ACRPowerModel, ACTPowerModel

    # linear approximations
    DCPPowerModel, NFAPowerModel

    # quadratic approximations
    DCPLLPowerModel, LPACCPowerModel

    # quadratic relaxations
    SOCWRPowerModel, SOCWRConicPowerModel,
    SOCBFPowerModel, SOCBFConicPowerModel,
    QCRMPowerModel, QCLSPowerModel,

    # sdp relaxations
    SDPWRMPowerModel, SparseSDPWRMPowerModel
================================================#

##### Exact Non-Convex Models #####
import PowerModels: ACPPowerModel

import PowerModels: ACRPowerModel

import PowerModels: ACTPowerModel

##### Linear Approximations #####
import PowerModels: DCPPowerModel

import PowerModels: NFAPowerModel

##### Quadratic Approximations #####
import PowerModels: DCPLLPowerModel

import PowerModels: LPACCPowerModel

##### Quadratic Relaxations #####
import PowerModels: SOCWRPowerModel

import PowerModels: SOCWRConicPowerModel

import PowerModels: QCRMPowerModel

import PowerModels: QCLSPowerModel


text:
using LinearAlgebra

@testset "Test QXTensor constructors and basic usage" begin
    # create a QXTensor with no data
    a = QXTensor([Index(2)])
    @assert size(a) == (2,)

    # test hyper edge search
    data = Diagonal(ones(4))
    indices = [Index(4), Index(4)]
    a = QXTensor(indices, data)
    @test size(a) == (4, 4)
    @test hyperindices(a) == [indices]

    data = reshape(QXTns.Gates.cx(), (2, 2, 2 ,2))
    indices = [Index(2), Index(2), Index(2), Index(2)]
    a = QXTensor(indices, data)
    @test hyperindices(a) == [indices[[2, 4]]]
    @test hyperindices(a, all_indices=true) == [[indices[1]], indices[[2, 4]], [indices[3]]]

    data = reshape(Diagonal(ones(4)), (2, 2, 2 ,2))
    indices = [Index(2), Index(2), Index(2), Index(2)]
    a = QXTensor(indices, data)
    @test hyperindices(a) == [indices[[1, 3]], indices[[2, 4]]]
end

@testset "Test contraction of tensors with hyper indices" begin
    # we create two sets of indices and identify subsets of indices that are hyper indices
    # we then contract over these sets to ensure the results set has the expected hyper indices
    as = [Index(2), Index(2), Index(2), Index(2)]
    a_hyper_indices = [[1, 3]]
    a = QXTensor(as, a_hyper_indices)
    bs = [as[3], as[4], Index(2), Index(2)]
    b_hyper_indices = [[1, 3]]
    b = QXTensor(bs, b_hyper_indices)

    c = contract_tensors(a, b)
    @test Set(hyperindices(c)...) == Set([as[1], bs[3]])

    # now an example with two sets of hyper indices
    as = [Index(2), Index(2), Index(2), Index(2)]
    a_hyper_indices = [[1, 3], [2, 4]]
    a = QXTensor(as, a_hyper_indices)
    bs = [as[3], as[4], Index(2), Index(2)]
    b_hyper_indices = [[1, 3], [2, 4]]
    b = QXTensor(bs, b_hyper_indices)

    c = contract_tensors(a, b)
    @test Set(hyperindices(c)[1]) == Set([as[1], bs[3]])
    @test Set(hyperindices(c)[2]) == Set([as[2], bs[4]])

    # next an example where the first tensor will have a group of hyper indices remaining
    # but the second tensor won't
    as = [Index(2), Index(2), Index(2), Index(5), Index(5)]
    a_hyper_indices = [[1, 2, 3]]
    a = QXTensor(as, a_hyper_indices)
    bs = [as[1], Index(4), Index(5), as[4]]
    b_hyper_indices = [[3, 4]]
    b = QXTensor(bs, b_hyper_indices)

    c = contract_tensors(a, b)
    @test Set(hyperindices(c)[1]) == Set([as[2], as[3]])

    # next an example where the first tensor has two groups linked with group from second tensor
    as = [Index(2), Index(2), Index(2), Index(2)]
    a_hyper_indices = [[1, 2], [3, 4]]
    a = QXTensor(as, a_hyper_indices)
    bs = [as[1], as[4]]
    b_hyper_indices = [[1, 2]]
    b = QXTensor(bs, b_hyper_indices)

    c = contract_tensors(a, b)
    @test Set(hyperindices(c)[1]) == Set([as[2], as[3]])
end

@testset "Test tensor_data when considering hyperedges" begin
    # test a diagonal
    data = Diagonal(ones(4))
    indices = [Index(4), Index(4)]
    a = QXTensor(indices, data)
    @test tensor_data(a, consider_hyperindices=true) == ones(4)

    # test a rank 4 tensor with 2 sets of hyper edges
    data = reshape(Diagonal(ones(4)), (2, 2, 2 ,2))
    indices = [Index(2), Index(2), Index(2), Index(2)]
    a = QXTensor(indices, data)
    @test tensor_data(a, consider_hyperindices=true) == ones(2, 2)

    # test a 2x2x2 tensor with single group of hyper edges include all ranks
    data = zeros(2, 2, 2)
    for i in 1:2 data[i, i, i] = 1 end
    a = QXTensor([Index(2), Index(2), Index(2)], data)
    @test tensor_data(a, consider_hyperindices=true) == ones(2)
end


text:
# Input data
data = map(x -> parse(Int64, x), readlines("Day01.input"));

# Part 1
getans =
    (offset) ->
        count(map(((i, v),) -> (i > offset && v - data[i-offset] > 0), enumerate(data)));

ans = getans(1);
println("Part 1: $ans")

# Part 2
ans = getans(3);
println("Part 2: $ans")


text:
const STRUCTSYMBOL = VERSION < v"0.7-" ? :type : :struct
isstructdef(ex) = Meta.isexpr(ex, STRUCTSYMBOL)

function splitstructdef(ex)
    ex = MacroTools.striplines(ex)
    ex = MacroTools.flatten(ex)
    d = Dict{Symbol, Any}()
    if @capture(ex, struct header_ body__ end)
        d[:mutable] = false
    elseif @capture(ex, mutable struct header_ body__ end)
        d[:mutable] = true
    else
        parse_error(ex)
    end

    if @capture header nameparam_ <: super_
        nothing
    elseif @capture header nameparam_
        super = :Any
    else
        parse_error(ex)
    end
    d[:supertype] = super
    if @capture nameparam name_{param__}
        nothing
    elseif @capture nameparam name_
        param = []
    else
        parse_error(ex)
    end
    d[:name] = name
    d[:params] = param
    d[:fields] = []
    d[:constructors] = []
    for item in body
        if @capture item field_::T_
            push!(d[:fields], (field, T))
        elseif item isa Symbol
            push!(d[:fields], (item, Any))
        else
            push!(d[:constructors], item)
        end
    end
    d
end

function combinestructdef(d)::Expr
    name = d[:name]
    parameters = d[:params]
    nameparam = isempty(parameters) ? name : :($name{$(parameters...)})
    header = :($nameparam <: $(d[:supertype]))
    fields = map(d[:fields]) do field
        fieldname, typ = field
        :($fieldname::$typ)
    end
    body = quote
        $(fields...)
        $(d[:constructors]...)
    end

    Expr(STRUCTSYMBOL, d[:mutable], header, body)
end

function combinefield(x)
    fieldname, T = x
    :($fieldname::$T)
end


text:
using PDMats
using Test

_randPDMat(T, n) = (X = randn(T, n, n); PDMat(X * X'))
_randPDiagMat(T, n) = PDiagMat(rand(T, n))
_randScalMat(T, n) = ScalMat(n, rand(T))
function _pd_compare(A::AbstractPDMat, B::AbstractPDMat)
    @test dim(A) == dim(B)
    @test Matrix(A) ≈ Matrix(B)
    @test cholesky(A).L ≈ cholesky(B).L
    @test cholesky(A).U ≈ cholesky(B).U
    nothing
end
function _pd_kron_compare(A::AbstractPDMat, B::AbstractPDMat)
    PDAkB1 = kron(A, B)
    PDAkB2 = PDMat( kron( Matrix(A), Matrix(B) ) )
    _pd_compare(PDAkB1, PDAkB2)
    nothing
end

n = 4
m = 7

for T in [Float32, Float64]
    _pd_kron_compare( _randPDMat(T, n),    _randPDMat(T, m) )
    _pd_kron_compare( _randPDiagMat(T, n), _randPDiagMat(T, m) )
    _pd_kron_compare( _randScalMat(T, n),  _randScalMat(T, m) )
    _pd_kron_compare( _randPDMat(T, n),    _randPDiagMat(T, m) )
    _pd_kron_compare( _randPDiagMat(T, m), _randPDMat(T, n) )
    _pd_kron_compare( _randPDMat(T, n),    _randScalMat(T, m) )
    _pd_kron_compare( _randScalMat(T, m),  _randPDMat(T, n) )
    _pd_kron_compare( _randPDiagMat(T, n), _randScalMat(T, m) )
    _pd_kron_compare( _randScalMat(T, m),  _randPDiagMat(T, n) )
end


text:
# ODE SIR/SICR models
# variables are concentrations, X/N
# Z = 1-S

# SICRLQM ODE
# Cases are infective, mitigation applied
# 7+1
function sicrqm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,mbase,mlocation = p
    beta *= mbase + (1-mbase)/(1 + exp(mrate*(t - mlocation)))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# SICR ODE
# Cases are noninfective
# 4+1 parameters
function sicr!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I)*(1-Z)
end

# SICRq ODE
# Cases are quarantined with effectiveness q
# 5+1 parameters
function sicrq!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# SICRf ODE
# Cases are quarantined with effectiveness q
# 5+1 parameters
function sicrf!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,f = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I)*(1-Z)
end

# SICRqf ODE
# Cases are quarantined with factor q
# Cases recover or die with factor f
# 6+1 parameters
function sicrqf!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,f = p
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# Mitigation applied
#6+1
function sicrm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I)*(1-Z)
end

# Mitigation applied
# 7+1
function sicrfm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,f,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I)*(1-Z)

end

# 7+1
function sicrqm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

# 8+1
function sicrqfm!(du,u,p,t)
	C,D,R,I,Z = u
	beta,sigmac,sigma,sigmar,q,f,mc,ml = p
    beta *= mc + (1-mc)/(1 + exp(t - 35. - ml))
	du[1] = dC = sigmac*I - sigma*C
	du[2] = dD = (sigma - sigmar)*C
	du[3] = dR = sigmar*C
	du[4] = dI = beta*(I+q*C)*(1-Z) - sigmac*I - f*sigma*I
	du[5] = dZ = beta*(I+q*C)*(1-Z)
end

function sicrlqfm!(du,u,p,t)
	I,C = u
	beta,sigmac,sigma,q,f,mbase,mlocation,mrate,cmax,c50 = p
    beta *= mbase + (1-mbase)/(1 + exp(mrate*(t - mlocation)))
    sigmac *= 1 + cmax*t/(c50+t)
    du[1] = dI = beta*(I+q*C) - sigmac*I - f*sigma*I
	du[2] = dC = sigmac*I - sigma*C
end

function linearrd!(du,u,p,t)
	I,C = u
	beta,sigmac,sigmar,sigmad,q,f,mbase,mlocation,mrate,cmax,c50 = p
    beta *= mbase + (1-mbase)/(1 + exp(mrate*(t - mlocation)))
    sigmac *= 1 + cmax*t/(c50+t)
    du[1] = dI = beta*(I+q*C) - sigmac*I - f*(sigmar+sigmad)*I
	du[2] = dC = sigmac*I - (sigmar+sigmad)*C
end

max_stars_repo_name: cannin/covid-sicr-test

max_issues_repo_name: cannin/covid-sicr-test

max_forks_repo_name: cannin/covid-sicr-test


text:
import JSON2

function validate_serialization(sys::IS.SystemData; time_series_read_only = false)
    #path, io = mktemp()
    # For some reason files aren't getting deleted when written to /tmp. Using current dir.
    path = "test_system_serialization.json"
    @info "Serializing to $path"

    try
        if isfile(path)
            rm(path)
        end
        IS.prepare_for_serialization!(sys, path)
        IS.to_json(sys, path)
    catch
        rm(path)
        rethrow()
    end

    # Make sure the code supports the files changing directories.
    test_dir = mktempdir()
    path = mv(path, joinpath(test_dir, path))

    t_file = splitext(basename(path))[1] * "_" * IS.TIME_SERIES_STORAGE_FILE
    mv(t_file, joinpath(test_dir, t_file))
    v_file = splitext(basename(path))[1] * "_" * IS.VALIDATION_DESCRIPTOR_FILE
    mv(v_file, joinpath(test_dir, v_file))

    ts_file = open(path) do file
        JSON2.read(file).time_series_storage_file
    end
    sys2 = IS.SystemData(path; time_series_read_only = time_series_read_only)
    return sys2, IS.compare_values(sys, sys2)
end

@testset "Test JSON serialization of system data" begin
    for in_memory in (true, false)
        sys = create_system_data_shared_forecasts(; time_series_in_memory = in_memory)
        _, result = validate_serialization(sys)
        @test result
    end
end

@testset "Test prepare_for_serialization" begin
    sys = create_system_data_shared_forecasts()
    directory = joinpath("dir1", "dir2")
    IS.prepare_for_serialization!(sys, joinpath(directory, "sys.json"))
    @test IS.get_ext(sys.internal)["serialization_directory"] == directory
end

@testset "Test JSON serialization of with read-only time series" begin
    sys = create_system_data_shared_forecasts(; time_series_in_memory = false)
    sys2, result = validate_serialization(sys; time_series_read_only = true)
    @test result
    component = collect(IS.get_components(IS.TestComponent, sys2))[1]
    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)
    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    forecast = IS.Deterministic("get_val", ta)
    @test_throws ErrorException IS.add_forecast!(sys2, component, forecast)
    @test_throws ErrorException IS.clear_forecasts!(sys2)
    forecast = collect(IS.iterate_forecasts(component))[1]
    @test_throws ErrorException IS.remove_forecast!(
        IS.Deterministic,
        sys2,
        component,
        IS.get_initial_time(forecast),
        IS.get_label(forecast),
    )
end

@testset "Test JSON serialization of with mutable time series" begin
    sys = create_system_data_shared_forecasts(; time_series_in_memory = false)
    sys2, result = validate_serialization(sys; time_series_read_only = false)
    @test result
    IS.clear_forecasts!(sys2)

    sys2, result = validate_serialization(sys; time_series_read_only = false)
    @test result
    component = collect(IS.iterate_components_with_forecasts(sys2.components))[1]
    forecast = collect(IS.iterate_forecasts(component))[1]
    IS.remove_forecast!(
        IS.Deterministic,
        sys2,
        component,
        IS.get_initial_time(forecast),
        IS.get_label(forecast),
    )
    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)
    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    forecast = IS.Deterministic("get_val", ta)
    IS.add_forecast!(sys2, component, forecast)
end


text:
# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# McCormick.jl
# A McCormick operator library in Julia
# See https://github.com/PSORLab/McCormick.jl
#############################################################################
# src/forward_operators/convex_increasing.jl
# Contains definitions of exp, exp2, exp10, expm1.
#############################################################################

for f in (:exp, :exp2, :exp10, :expm1)
    function ($f)(::Union{NS,MV}, x::MCNoGrad)
        X = x.Intv;    xL = x.Intv.lo;   xU = x.Intv.hi
        z = ($f)(X);   zL = z.lo;        zU = z.hi
        mcc = mid3v(x.cv, x.cc, xU)
        mcv = mid3v(x.cv, x.cc, xL)
        zcc = zU > zL ? (zL*(xU - mcc) + zU*(mcc - xL))/(xU - xL) : zU
        zcv = ($f)(mcv)
        (zL > zcv) && (zcv = zL;)
        (zU < zcc) && (zcc = zU;)
        return MCNoGrad(zcv, zcc, z, x.cnst)
    end
    function ($f)(::Diff, x::MCNoGrad)
        X = x.Intv;    xL = x.Intv.lo;   xU = x.Intv.hi
        z = ($f)(X);   zL = z.lo;        zU = z.hi
        mcc = mid3v(x.cv, x.cc, xU)
        mcv = mid3v(x.cv, x.cc, xL)
        zcc = zU > zL ? (zL*(xU - mcc) + zU*(mcc - xL))/(xU - xL) : zU
        zcv = ($f)(mcv)
       return MCNoGrad(zcv, zcc, z, x.cnst)
    end
    f_kernel = Symbol(String(opMC)*"_kernel")
    df = diffrule(:Base, opMC, :mcv) # Replace with cv ruleset
    function ($f_kernel)(::Union{NS,MV}, x::MCNoGrad, z::Interval{Float64})
        xL = x.Intv.lo;   xU = x.Intv.hi;   zL = z.lo;   zU = z.hi
        mcc, cci = mid3(x.cc, x.cv, xU)
        mcv, cvi = mid3(x.cc, x.cv, xL)
        zcc = zU > zL ? (zL*(xU - mcc) + zU*(mcc - xL))/(xU - xL) : zU
        zcv = ($f)(mcv)
        if zL > zcv
            cvi = 3
            dcv = 0.0
        else
            dcv = $df
        end
        if zU < zcc
            cci = 3
            dcc = 0.0
        else
            dcc = (zU - zL)/(xU - xL)
        end
        return MCNoGrad(u, o, z, x.cnst), cvi, cci, dcv, dcc
    end
    ddf = diffrule(:Base, opMC, :mcv) # Replace with cv ruleset
    function ($f_kernel)(::Diff, x::MCNoGrad, z::Interval{Float64})
        xL = x.Intv.lo;   xU = x.Intv.hi;   zL = z.lo;   zU = z.hi
        mcc, cci = mid3(x.cc, x.cv, xU)
        mcv, cvi = mid3(x.cc, x.cv, xL)
        zcc = zU > zL ? (zL*(xU - mcc) + zU*(mcc - xL))/(xU - xL) : zU
        zcv = ($f)(mcv)
        dcv = $ddf
        dcc = (zU - zL)/(xU - xL)
       return MCNoGrad(u, o, z, x.cnst), cvi, cci, dcv, dcc
    end
end

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
import ForwardDiff
import MacroTools
import LinearAlgebra

struct InvolutionDSLProgram
    fn!::Function
end

# See this math writeup for an understanding of how this code works:
# docs/tex/mcmc.pdf

struct FirstPassState

    "trace containing the input model choice map ``t``"
    trace

    "the input proposal choice map ``u``"
    u::ChoiceMap

    "subset of the output model choice map ``t'``"
    constraints::ChoiceMap

    "output proposal choice map ``u'``"
    u_back::ChoiceMap

    t_cont_reads::Dict
    u_cont_reads::Dict
    t_cont_writes::Dict
    u_cont_writes::Dict
    t_copy_reads::DynamicSelection
    u_copy_reads::DynamicSelection
end

function FirstPassState(trace, u::ChoiceMap)
    FirstPassState(
        trace, u, choicemap(), choicemap(),
        Dict(), Dict(), Dict(), Dict(),
        DynamicSelection(), DynamicSelection())
end

struct JacobianPassState{T<:Real}
    trace
    u::ChoiceMap
    input_arr::AbstractArray{T}
    output_arr::Array{T}
    t_key_to_index::Dict
    u_key_to_index::Dict
    cont_constraints_key_to_index::Dict
    cont_u_back_key_to_index::Dict
end

const inv_state = gensym("inv_state")

"""
    @involution function f(...)
        ..
    end

Write a program in the [Involution DSL](@ref).
"""
macro involution(ex)
    ex = MacroTools.longdef(ex)
    MacroTools.@capture(ex, function f_(args__) body_ end) || error("expected syntax: function f(..) .. end")

    fn! = gensym("$(esc(f))_fn!")

    quote

    # mutates the state
    function $fn!($(esc(inv_state))::Union{FirstPassState,JacobianPassState}, $(map(esc, args)...))
        $(esc(body))
        nothing
    end

    Core.@__doc__ $(esc(f)) = InvolutionDSLProgram($fn!)

    end # quote

end # macro involution()

# read discrete from model

macro read_discrete_from_model(addr)
    quote read_discrete_from_model($(esc(inv_state)), $(esc(addr))) end
end

function read_discrete_from_model(state::Union{FirstPassState,JacobianPassState}, addr)
    state.trace[addr]
end

# read discrete from proposal

macro read_discrete_from_proposal(addr)
    quote read_discrete_from_proposal($(esc(inv_state)), $(esc(addr))) end
end

function read_discrete_from_proposal(state::Union{FirstPassState,JacobianPassState}, addr)
    state.u[addr]
end

# write_discrete_to_proposal

macro write_discrete_to_proposal(addr, value)
    quote write_discrete_to_proposal($(esc(inv_state)), $(esc(addr)), $(esc(value))) end
end

function write_discrete_to_proposal(state::FirstPassState, addr, value)
    state.u_back[addr] = value
    value
end

function write_discrete_to_proposal(state::JacobianPassState, addr, value)
    value
end

# write_discrete_to_model

macro write_discrete_to_model(addr, value)
    quote write_discrete_to_model($(esc(inv_state)), $(esc(addr)), $(esc(value))) end
end

function write_discrete_to_model(state::FirstPassState, addr, value)
    state.constraints[addr] = value
    value
end

function write_discrete_to_model(state::JacobianPassState, addr, value)
    value
end

# read continuous from proposal

macro read_continuous_from_proposal(addr)
    quote read_continuous_from_proposal($(esc(inv_state)), $(esc(addr))) end
end

function read_continuous_from_proposal(state::FirstPassState, addr)
    state.u_cont_reads[addr] = state.u[addr]
    state.u[addr]
end

function read_continuous_from_proposal(state::JacobianPassState, addr)
    if haskey(state.u_key_to_index, addr)
        state.input_arr[state.u_key_to_index[addr]]
    else
        state.u[addr]
    end
end

# read continuous from model

macro read_continuous_from_model(addr)
    quote read_continuous_from_model($(esc(inv_state)), $(esc(addr))) end
end

function read_continuous_from_model(state::FirstPassState, addr)
    state.t_cont_reads[addr] = state.trace[addr]
    state.trace[addr]
end

function read_continuous_from_model(state::JacobianPassState, addr)
    if haskey(state.t_key_to_index, addr)
        state.input_arr[state.t_key_to_index[addr]]
    else
        state.trace[addr]
    end
end

# write_continuous_to_proposal

macro write_continuous_to_proposal(addr, value)
    quote write_continuous_to_proposal($(esc(inv_state)), $(esc(addr)), $(esc(value))) end
end

function write_continuous_to_proposal(state::FirstPassState, addr, value)
    has_value(state.u_back, addr) && error("Proposal address $addr already written to")
    state.u_back[addr] = value
    state.u_cont_writes[addr] = value
    value
end

function write_continuous_to_proposal(state::JacobianPassState, addr, value)
    state.output_arr[state.cont_u_back_key_to_index[addr]] = value
end

# write_continuous_to_model

macro write_continuous_to_model(addr, value)
    quote write_continuous_to_model($(esc(inv_state)), $(esc(addr)), $(esc(value))) end
end

function write_continuous_to_model(state::FirstPassState, addr, value)
    has_value(state.constraints, addr) && error("Model address $addr already written to")
    state.constraints[addr] = value
    state.t_cont_writes[addr] = value
    value
end

function write_continuous_to_model(state::JacobianPassState, addr, value)
    state.output_arr[state.cont_constraints_key_to_index[addr]] = value
end

# copy_model_to_model

macro copy_model_to_model(from_addr, to_addr)
    quote copy_model_to_model($(esc(inv_state)), $(esc(from_addr)), $(esc(to_addr))) end
end

function copy_model_to_model(state::FirstPassState, from_addr, to_addr)
    trace_choices = get_choices(state.trace)
    push!(state.t_copy_reads, from_addr)
    if has_value(trace_choices, from_addr)
        state.constraints[to_addr] = state.trace[from_addr]
    else
        set_submap!(state.constraints, to_addr, get_submap(trace_choices, from_addr))
    end
    nothing
end

function copy_model_to_model(state::JacobianPassState, from_addr, to_addr)
    nothing
end

# copy_model_to_proposal

macro copy_model_to_proposal(from_addr, to_addr)
    quote copy_model_to_proposal($(esc(inv_state)), $(esc(from_addr)), $(esc(to_addr))) end
end

function copy_model_to_proposal(state::FirstPassState, from_addr, to_addr)
    trace_choices = get_choices(state.trace)
    push!(state.t_copy_reads, from_addr)
    if has_value(trace_choices, from_addr)
        state.u_back[to_addr] = state.trace[from_addr]
    else
        set_submap!(state.u_back, to_addr, get_submap(trace_choices, from_addr))
    end
    nothing
end

function copy_model_to_proposal(state::JacobianPassState, from_addr, to_addr)
    nothing
end

# copy_proposal_to_proposal

macro copy_proposal_to_proposal(from_addr, to_addr)
    quote copy_proposal_to_proposal($(esc(inv_state)), $(esc(from_addr)), $(esc(to_addr))) end
end

function copy_proposal_to_proposal(state::FirstPassState, from_addr, to_addr)
    push!(state.u_copy_reads, from_addr)
    if has_value(state.u, from_addr)
        state.u_back[to_addr] = state.u[from_addr]
    else
        set_submap!(state.u_back, to_addr, get_submap(state.u, from_addr))
    end
    nothing
end

function copy_proposal_to_proposal(state::JacobianPassState, from_addr, to_addr)
    nothing
end

# copy_proposal_to_model

macro copy_proposal_to_model(from_addr, to_addr)
    quote copy_proposal_to_model($(esc(inv_state)), $(esc(from_addr)), $(esc(to_addr))) end
end

function copy_proposal_to_model(state::FirstPassState, from_addr, to_addr)
    push!(state.u_copy_reads, from_addr)
    if has_value(state.u, from_addr)
        state.constraints[to_addr] = state.u[from_addr]
    else
        set_submap!(state.constraints, to_addr, get_submap(state.u, from_addr))
    end
    nothing
end

function copy_proposal_to_model(state::JacobianPassState, from_addr, to_addr)
    nothing
end

# call another involution function (NOTE: only the top-level function is actually technically an involution)

macro invcall(ex)
    MacroTools.@capture(ex, f_(args__)) || error("expected syntax: f(..)")
    quote $(esc(f)).fn!($(esc(inv_state)), $(map(esc, args)...)) end
end

# apply

function apply_involution(involution::InvolutionDSLProgram, trace, u, proposal_args, proposal_retval)

    # take the first pass through the involution, mutating first_pass_state
    first_pass_state = FirstPassState(trace, u)
    involution.fn!(first_pass_state, get_args(trace), proposal_args, proposal_retval)

    # update model trace
    (new_trace, model_weight, _, discard) = update(
        trace, get_args(trace), map((_) -> NoChange(), get_args(trace)), first_pass_state.constraints)

    # create input array and mappings input addresses that are needed for Jacobian
    # exclude addresses that were moved to another address
    input_arr = Vector{Float64}()
    next_input_index = 1

    t_key_to_index = Dict()
    for (addr, v) in first_pass_state.t_cont_reads
        if addr in first_pass_state.t_copy_reads
            continue
        end
        if !has_value(discard, addr)
            continue
        end
        t_key_to_index[addr] = next_input_index
        next_input_index += 1
        push!(input_arr, v)
    end

    u_key_to_index = Dict()
    for (addr, v) in first_pass_state.u_cont_reads
        if addr in first_pass_state.u_copy_reads
            continue
        end
        u_key_to_index[addr] = next_input_index
        next_input_index += 1
        push!(input_arr, v)
    end

    # create mappings for output addresses that are needed for Jacobian
    next_output_index = 1

    cont_constraints_key_to_index = Dict()
    for (addr, v) in first_pass_state.t_cont_writes
        cont_constraints_key_to_index[addr] = next_output_index
        next_output_index += 1
    end

    cont_u_back_key_to_index = Dict()
    for (addr, v) in first_pass_state.u_cont_writes
        cont_u_back_key_to_index[addr] = next_output_index
        next_output_index += 1
    end

    # this function is the partial application of the continuous part of the
    # involution, with inputs corresponding to a particular superset of the
    # columns of the reduced Jacobian matrix
    function f_array(input_arr::AbstractArray{T}) where {T <: Real}

        # closing over:
        # - trace, u
        # - u_key_to_index, t_key_to_index, cont_constraints_key_to_index, cont_u_back_key_to_index
        # - proposal_args, proposal_retval

        n_output = length(cont_constraints_key_to_index) + length(cont_u_back_key_to_index)
        output_arr = Vector{T}(undef, n_output)

        jacobian_pass_state = JacobianPassState(
            trace, u, input_arr, output_arr,
            t_key_to_index, u_key_to_index,
            cont_constraints_key_to_index,
            cont_u_back_key_to_index)

        # mutates the state
        involution.fn!(
            jacobian_pass_state, get_args(trace), proposal_args, proposal_retval)

        # return the output array
        output_arr
    end

    (new_trace, model_weight, first_pass_state.u_back, input_arr, f_array)
end

# callable

function (involution::InvolutionDSLProgram)(trace, u::ChoiceMap, proposal_retval, proposal_args; check=false)

    # apply the involution, and get metadata back about it, including the
    # function f_array from arrays to arrays, that is the restriction of the
    # involution to just a subset of the relevant continuous random choices
    (new_trace, model_weight, u_back, input_arr, f_array) = apply_involution(
        involution, trace, u, proposal_args, proposal_retval)

    # compute Jacobian matrix of f_array, where columns are inputs, rows are outputs
    J = ForwardDiff.jacobian(f_array, input_arr)
    @assert size(J)[2] == length(input_arr)
    num_outputs = size(J)[1]
    if size(J) != (num_outputs, num_outputs)
        error("Jacobian was not square (size was $(size(J)); the function may not be an involution")
    end

    # log absolute value of Jacobian determinant
    correction = LinearAlgebra.logabsdet(J)[1]
    if isinf(correction)
        @error "Weight correction is infinite; the function may not be an involution"
    end

    (new_trace, u_back, model_weight + correction)
end

export @involution
export @read_discrete_from_proposal, @read_discrete_from_model
export @write_discrete_to_proposal, @write_discrete_to_model
export @read_continuous_from_proposal, @read_continuous_from_model
export @write_continuous_to_proposal, @write_continuous_to_model
export @copy_model_to_model, @copy_model_to_proposal
export @copy_proposal_to_proposal, @copy_proposal_to_model
export @invcall


text:
# Inspired by https://discourse.julialang.org/t/slow-custom-digits-iterator/31367
# Expanded to accept more kinds of integer
struct DigitsIterator{T<:Integer}; n::T; end
@inline Base.iterate(it::DigitsIterator) = it.n == 0 ? (0,0) : iterate(it,it.n)
@inline Base.iterate(it::DigitsIterator,el) = el == 0 ? nothing : reverse(divrem(el,10))
Base.length(it::DigitsIterator) = ndigits(it.n)
Base.eltype(::Type{DigitsIterator{T}}) where T = T

function isarmstrong(n::Integer)
    ds = DigitsIterator(n) # digits(n) also works but allocates and is slower
    n == sum(x->x^length(ds),ds)
end


text:
using Plots
using RecipesBase

# function primarily to create an individual sliceplorer sub-plot
#@recipe function f(::Type{Val{:sliceplorerdim}}, x, y, z)
@recipe function f(spDim::SliceplorerDim)
  #spDim = y;
  # separate out the x and y values from the data structure
  xs = [[v′.pt[1] for v′ in v.sample] for v in spDim.dim]
  ys = [[v′.pt[2] for v′ in v.sample] for v in spDim.dim]

  # plot properties
  # FIXME: set line width to something nice
  seriestype := :line # sort by x
  legend := false
  linealpha --> 0.4
  linecolor --> :black

  (hcat(xs...), hcat(ys...))
end

# TODO: add option to lay out everything vertically, etc
@recipe function f(sp::SliceSet1D)

  # global options for the subplots
  legend := false
  link := :both
  layout := length(sp.problemSpec) # number of dimensions
  #seriestype := :sliceplorerdim

  # FIXME: adjust positioning of title
  title := [k for j in 1:1, k in keys(sp.samples.samples)]
  titlelocation := :right
  titlefont := font(8)

  # FIXME: might need to be a map (i.e. returns something)
  for (i, spDim) in enumerate(values(sp.samples.samples))
    @series begin
      subplot := i
      spDim
    end
  end
end

# TODO: is it worth splitting this by suplot like the other dataset?
@recipe function f(hsp::Vector{SliceSet2D})

  # need to sort by dims otherwise layout and splitting doesn't work
  slices = sort(hsp.slices, by=x -> (x.d1, x.d2));

  # universal settings
  linecolor := :black
  linealpha --> 0.4
  seriestype := :path

  # layout is a SPLOM-style layout
  dims = length(hsp.problemSpec);
  layoutRows = map(1:(dims-1)) do d
    hcat(fill(nothing, (1, )), fill(∘, (1, dims-d)))
  end
  # FIXME: might need to search/replace the underscore...
  layout := @layout vcat(layoutRows)

  # cut up the segment list by dimension
  # and generate a subplot per segment
  # FIXME: seems like there could be a better algorithm
  subplot = 1
  iRow = 1
  while iRow <= length(slices)
    dimSet = (slices[iRow].d1, slices[iRow].d2);
    i = iRow + 1
    while dimSet == (slices[i].d1, slices[i].d2) && i <= length(slices)
      i = i + 1;
    end

    @series begin
      slices[iRow:i]
      subplot := subplot
    end

    iRow = i;
  end

end


text:
function get_evaluations(obj::BaseOptimizer)
    return jcall(obj, "getEvaluations", jint, ())
end

function optimize(obj::BaseMultivariateOptimizer, arg0::Vector{OptimizationData})
    return jcall(obj, "optimize", Object, (Vector{OptimizationData},), arg0)
end


text:
module LatentGaussianMixtureModel

using StatsModels
using Distributions
using StatsBase, StatsFuns
using FastGaussQuadrature
using Yeppp
using Printf
using Distributed
using LinearAlgebra
using Statistics
using SharedArrays

import GaussianMixtureTest: gmm, pn, decidepenalty
import StatsBase: RealArray, RealVector, RealMatrix, IntegerArray, IntegerVector, IntegerMatrix, IntUnitRange
import StatsFuns: logπ

import StatsBase: coef, confint, deviance, loglikelihood, nobs, stderror, vcov, predict, fit!, model_response, RegressionModel, coeftable

export LGMModel, ModelTable,
    initialize!, latentgmm, EMtest, asymptoticdistribution,
    ranefmixture, infomatrix, multiplefit!, ranef!,
    FDR, detect, latexprint,
    fit!, coef, confint, deviance, loglikelihood, nobs, stderror, vcov, predict, model_response, coeftable

import Base.show
using DataFrames

include("tableprint.jl")
include("model.jl")

#include("MCEM.jl")

include("arithmetic.jl")
include("utility.jl")
include("EM.jl")
#include("FractionalImputation.jl")

PKGVERSION=v"0.4.1"
export PKGVERSION
end # module


text:
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Descriptions
# ==============================================================================
#
#   Function calls to create the precompilation statements using SnoopCompiler.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

function precompilation_input()
    # Input: Array
    # ==========================================================================

    data = Any[1    false      1.0     0x01 ;
               2     true      2.0     0x02 ;
               3    false      3.0     0x03 ;
               4     true      4.0     0x04 ;
               5    false      5.0     0x05 ;
               6     true      6.0     0x06 ;]

    pretty_table(data)

    # Pre-compile the most common types.
    #
    # The functions that parse the cells (`_parse_cell_text`) could not be
    # despecialized because it degraded too much the performance. Hence, there
    # will be a small precompilation for each different cell type.
    types = [Int8(1),
             Int16(1),
             Int32(1),
             Int64(1),
             UInt8(1),
             UInt16(1),
             UInt32(1),
             UInt64(1),
             Float16(1),
             Float32(1),
             Float64(1),
             Bool(1),
             "S",
             'C',
             Int8[1, 2, 3],
             Int16[1, 2, 3],
             Int32[1, 2, 3],
             Int64[1, 2, 3],
             UInt8[1, 2, 3],
             UInt16[1, 2, 3],
             UInt32[1, 2, 3],
             UInt64[1, 2, 3],
             Float16[1, 2, 3],
             Float32[1, 2, 3],
             Float64[1, 2, 3],
             Bool[0, 1],
             ["S", "S"],
             ['C', 'C'],
             Int8[1 2; 3 4],
             Int16[1 2; 3 4],
             Int32[1 2; 3 4],
             Int64[1 2; 3 4],
             UInt8[1 2; 3 4],
             UInt16[1 2; 3 4],
             UInt32[1 2; 3 4],
             UInt64[1 2; 3 4],
             Float16[1 2; 3 4],
             Float32[1 2; 3 4],
             Float64[1 2; 3 4],
             Bool[0 1; 0 1],
             ["S" "S"; "S" "S"],
             ['C' 'C'; 'C' 'C']]

    pretty_table(types;
                 alignment = :l,
                 crop = :none)

    # Input: Tables.jl
    # ==========================================================================

    # A named tuple is compliant with Table.jl.
    table = (a = 1:1:10, b = ["S" for i = 1:10], c = ['C' for i = 1:10])

    # This example is created based on DataFrames.jl options.
    pretty_table(table,
                 alignment                   = [:l, :c, :r],
                 crop                        = :both,
                 crop_num_lines_at_beginning = 2,
                 display_size                = (15, 34),
                 ellipsis_line_skip          = 3,
                 header_alignment            = :l,
                 hlines                      = [:header],
                 newline_at_end              = false,
                 row_name_alignment          = :r,
                 row_name_crayon             = Crayon(),
                 row_name_column_title       = "Name",
                 row_names                   = ["row" for i = 1:10],
                 row_number_alignment        = :r,
                 row_number_column_title     = "Row",
                 show_row_number             = true,
                 title                       = "Test table",
                 vcrop_mode                  = :middle,
                 vlines                      = [1])
    println()

    # Input: Dictionary
    # ==========================================================================

    dict = Dict(:a => (1, 1),
                :b => (2, 2),
                :c => (3, 3))

    pretty_table(dict, sortkeys = true)

    # Filters
    # ==========================================================================

    pretty_table(data, filters_row = ((data, i) -> i % 2 == 0,))
    pretty_table(data, filters_col = ((data, i) -> i % 2 == 0,))
    pretty_table(
        data,
        filters_col = ((data, i) -> i % 2 == 0,),
        filters_row = ((data, i) -> i % 2 == 0,)
    )

    # Input: Data with URLTextCell
    # ==========================================================================

    custom_cells = [
        1 "Ronan Arraes Jardim Chagas" URLTextCell("Ronan Arraes Jardim Chagas", "https://ronanarraes.com")
        2 "Google" URLTextCell("Google", "https://google.com")
        3 "Apple" URLTextCell("Apple", "https://apple.com")
        4 "Emojis!" URLTextCell("😃"^20, "https://emojipedia.org/github/")
    ]

    pretty_table(custom_cells)

    # Input: Data with AnsiCellText
    # ==========================================================================

    b = crayon"blue bold"
    y = crayon"yellow bold"
    g = crayon"green bold"

    ansi_table = [
        AnsiTextCell("$(g)This $(y)is $(b)awesome!")
        AnsiTextCell("$(g)😃😃 $(y)is $(b)awesome!")
        AnsiTextCell("$(g)σ𝛕θ⍺ $(y)is $(b)awesome!")
    ]

    pretty_table(ansi_table)
end


text:
# memory_function.jl

"""
χ(Ω::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the memory function χ(Ω) of the polaron at finite temperatures (equation (35a) in FHIP 1962) for a given frequency Ω. v and w are the variational polaron parameters that minimise the free energy, for the supplied α Frohlich coupling.
"""
function χ(Ω, β, α, v, w)

    # FHIP1962, page 1011, eqn (47c).
    R = (v^2 - w^2) / (w^2 * v)

    # FHIP1962, page 1009, eqn (35c).
    D(x) = w^2 / v^2 * (R * (1 - cos(v * x)) * coth(β * v / 2) + x^2 / β - 1im * (R * sin(v * x) + x))

    # FHIP1962, page 1009, eqn (36).
    S(x) = 2 * α / (3 * √π) * (exp(1im * x) + 2 * cos(x) / (exp(β) - 1)) / (D(x))^(3 / 2)

    # FHIP1962, page 1009, eqn (35a).
    integrand(x) = (1 - exp(-1im * Ω * x)) * imag(S(x))
    QuadGK.quadgk(x -> integrand(x), 0.0, Inf)[1]
end

"""
χ(Ω::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the memory function χ(Ω) of the polaron at zero-temperatures (equation (35a) in FHIP 1962) for a given frequency Ω. v and w are the variational polaron parameters that minimise the free energy, for the supplied α Frohlich coupling.
"""
function χ(Ω, α, v, w)

    # FHIP1962, page 1011, eqn (47c).
    R = (v^2 - w^2) / (w^2 * v)

    # FHIP1962, page 1009, eqn (35c) with β → ∞.
    D(x) = w^2 / v^2 * (R * (1 - cos(v * x)) - 1im * (R * sin(v * x) + x))

    # FHIP1962, page 1009, eqn (36) with β → ∞.
    S(x) = 2 * α / (3 * √π) * exp(1im * x) / (D(x))^(3 / 2)

    # FHIP1962, page 1009, eqn (35a). Set upper limit < ∞ but very large so cos argument finite.
    integrand(x) = (1 - exp(-1im * Ω * x)) * imag(S(x))
    QuadGK.quadgk(x -> integrand(x), 0.0, 1e200)[1]
end

"""
χ_dc(β::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the memory function lim(Ω → 0){χ(Ω) / Ω} of the polaron at finite temperatures (equation (35a) in FHIP 1962) at zero frequency. v and w are the variational polaron parameters that minimise the free energy, for the supplied α Frohlich coupling.
"""
function χ_dc(β, α, v, w)

    # FHIP1962, page 1011, eqn (47c).
    R = (v^2 - w^2) / (w^2 * v)

    # FHIP1962, page 1009, eqn (35c).
    D(x) = w^2 / v^2 * (R * (1 - cos(v * x)) * coth(β * v / 2) + x^2 / β - 1im * (R * sin(v * x) + x))

    # FHIP1962, page 1009, eqn (36).
    S(x) = 2 * α / (3 * √π) * (exp(1im * x) + 2 * cos(x) / (exp(β) - 1)) / (D(x))^(3 / 2)

    # Set frequency small enough to mimic Ω = 0 without generating numerical instabilities in integral.
    Ω = 1e-200

    # FHIP1962, page 1009, eqn (35a). Readily divided by Ω to get sinc function in integrand rather than sine which would give 0 all the time.
    integrand(x) = (1 - exp(-1im * Ω * x)) * imag(S(x)) / Ω
    QuadGK.quadgk(x -> integrand(x), 0.0, Inf)[1]
end


text:
include("fit.jl")
include("data.jl")

function example_iso1(epsi=0.5)
    data = [100, 1, 0, 0, 3, 5, 0, 99, 0, 0]
    sorted_data = sort(data)
    noisy_data = DP.laplace_mechanism(sorted_data, sensitivity=1, epsilon=epsi)
    postprocessed1 = sort(noisy_data)
    dpdata = DP.isotonic(noisy_data)
    println("Sorted: ", sorted_data)
    println("Noisy: ", round.(noisy_data, digits=3))
    println("Sorting the noisy data: ", round.(postprocessed1, digits=3))
    println("Isotonic regression: ", round.(dpdata, digits=3))
    # this is not very sorted now
    error_noisy_data = sum(x -> x^2, sorted_data - noisy_data)
    error_sorted = sum(x -> x^2, sorted_data - postprocessed1)
    error_isotonic = sum(x -> x^2, sorted_data - dpdata)
    println("Error of noisy data: ", error_noisy_data)
    println("Error of sorted noisy data: ", error_sorted)
    println("Error of isotonic: ", error_isotonic)

end

mean(x::Array) = sum(x)/length(x)

function example_iso2(epsi=0.5)
    data = getnet()
    sorted_data = sort(data)
    noisy_data = DP.laplace_mechanism(sorted_data, sensitivity=1, epsilon=epsi)
    dpdata = DP.isotonic(noisy_data)
    println("Noisy Data error: ", mean((sorted_data - noisy_data).^2))
    println("Postprocessed error: ", mean((sorted_data - dpdata).^2))
end

function example_marginal(epsi=3.0)
    datashape = (3,4,5)
    data = reshape(1:60, datashape)
    q1 = DP.Marginal((1,))
    q2 = DP.Marginal((2,))
    q3 = DP.Marginal((3,))

    answer1 = DP.answer(q1, data)
    answer2 = DP.answer(q2, data)
    answer3 = DP.answer(q3, data)

    dp1 = DP.laplace_mechanism(answer1, sensitivity=DP.sens(q1), epsilon=epsi/3)
    dp2 = DP.laplace_mechanism(answer2, sensitivity=DP.sens(q2), epsilon=epsi/3)
    dp3 = DP.laplace_mechanism(answer3, sensitivity=DP.sens(q3), epsilon=epsi/3)

    dphist = DP.fit(datashape, [(q1, dp1, 1.0), (q2, dp2, 1.0), (q3, dp3, 1.0)])

    for (i, q) in enumerate([q1, q2, q3])
        println("####################")
        println("Marginal ", i)
        println("####################")
        println("True:")
        show(DP.answer(q, data))
        println()
        println("DP:")
        show(round.(DP.answer(q, dphist), digits=3))
        println()
    end
end

sqerror(x::Number, y::Number) = (x-y)^2
sqerror(x::Array, y::Array) = sum(a -> a^2, x-y)

function evaluate_strategy(workload, strategy, d, data, epsilon)
    noisy_answers = DP.protect(strategy, data, epsilon)
    noisy_data = DP.fit(size(data), [(strategy, noisy_answers, 1.0)])
    mean([sqerror(DP.answer(q, data), DP.answer(q, noisy_data)) for q in workload])
end

function evaluate_mwem(workload, T, data, epsilon)
    noisy_data = DP.mwem(workload, T, data, epsilon)
    mean([sqerror(DP.answer(q, data), DP.answer(q, noisy_data)) for q in workload])
end

function example_range(epsilon=2, T=3, amount=100000)
    data = getfrank()[1:200]
    d = length(data)

    rangequeries = DP.all_range_queries(d)
    idquery = DP.IDQuery()
    error_id = mean([evaluate_strategy(rangequeries, idquery, d, data, epsilon) for _ in amount])

    tree = DP.tree_queries(d, numsplits=100)
    error_tree = mean([evaluate_strategy(rangequeries, tree, d, data, epsilon) for _ in amount])

    #error_mwem = mean([evaluate_mwem(rangequeries, T, data, epsilon) for _ in amount])

    println("#### Results: ####")
    println("ID Query mean squared error: ", error_id)
    println("Range Query mean squared error: ", error_tree)
    #println("MWEM mean squared error: ", error_mwem)
end

max_stars_repo_name: dkifer/dpcourse

max_issues_repo_name: dkifer/dpcourse

max_forks_repo_name: dkifer/dpcourse


text:
# Autogenerated wrapper script for Libtiff_jll for x86_64-unknown-freebsd
export libtiff

using JpegTurbo_jll
using Zlib_jll
using Zstd_jll
JLLWrappers.@generate_wrapper_header("Libtiff")
JLLWrappers.@declare_library_product(libtiff, "libtiff.so.5")
function __init__()
    JLLWrappers.@generate_init_header(JpegTurbo_jll, Zlib_jll, Zstd_jll)
    JLLWrappers.@init_library_product(
        libtiff,
        "lib/libtiff.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
## Assumptions
## http://docs.sympy.org/0.7.2/modules/assumptions/index.html

## TODO: Deprecate Qeven, etc. in favor of Q.even

"""
refine: http://docs.sympy.org/dev/modules/assumptions/refine.html
"""
refine(ex, assumpts...) = sympy_meth(:refine, ex, assumpts...)
export refine

"""

ask. Returns true, false or nothing

ask(Qinteger(x*y), And(Qinteger(x), Qinteger(y)))
## really slow isprime:
filter(x -> ask(Qprime(x)), [1:1000])

"""
ask(x::Sym, args...) = sympy_meth(:ask, x, args...)
export ask

## should we support & and | for (sym,sym) pairs? Not sure
## dependso on what happens to | (x==0) and ex == x usage for ==
## for now, we can combine terms logically with And, Or, Not...
## these are in logic module

## simple methods (x, args) -> y (y coercion happens via PyCall)
logic_sympy_methods = (
                     :And, :Or, :Not,
                     :Xor, :Nand, :Nor, :Implies,
                     :Equivalent,
                     :satisfiable
                     )

## This is a bit messed up, as we use Qeven in place of Q.even, ...
## This is now deprecated. See below
Q_predicates = (:even, :odd, :prime, :nonzero,
                :complex, :extended_real, :imaginary, :infinitesimal,
                :integer, :irrational, :real,
                :positive, :negative,
                :bounded, :commutative)
Q_nms = [@compat(Symbol("Q" * string(i))) for i in Q_predicates]

for (fn, meth) in zip(Q_nms, Q_predicates)
    nm = string(meth)
    @eval begin
        @doc """
`$($nm)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($nm)
""" ->
        ($fn)(x::SymbolicObject) = PyCall.pyeval("f(x)", Sym, f=SymPy.sympy[:Q][($nm)], x=SymPy.project(x))::Sym
    end
    eval(Expr(:export, fn))
end

@deprecate Qeven(x) Q.even(x)
@deprecate Qodd(x) Q.odd(x)
@deprecate Qprime(x) Q.prime(x)
@deprecate Qnonzero(x) Q.nonzero(x)
@deprecate Qcomplex(x) Q.complex(x)
@deprecate Qextended_real(x) Q.extended_real(x)
@deprecate Qimaginary(x) Q.imaginary(x)
@deprecate Qinfinitesimal(x) Q.infinitesimal(x)
@deprecate Qinteger(x) Q.integer(x)
@deprecate Qirrational(x) Q.irrational(x)
@deprecate Qreal(x) Q.real(x)
@deprecate Qpositive(x) Q.positive(x)
@deprecate Qnegative(x) Q.negative(x)
@deprecate Qbounded(x) Q.bounded(x)
@deprecate Qcommutative(x) Q.commutative(x)

## We make a module Q to hold the assumptions
## this follows this page http://docs.sympy.org/0.7.5/_modules/sympy/assumptions/ask.html
module Q
import SymPy
import PyCall

##http://docs.sympy.org/dev/_modules/sympy/assumptions/ask.html#ask
Q_predicates = (:antihermitian,
                :bounded, :finite, # bounded deprecated
                :commutative,
#                :complex,
                :composite,
                :even,
                :extended_real,
                :hermitian,
                :imaginary,
                :infinitesimal,
                :infinity, :infinite, # :infinity deprecated
                :integer,
                :irrational,
                :rational,
                :algebraic,
                :transcendental,
                :negative,
                :nonzero, :zero,
                :positive,
                :prime,
                :real,
                :odd,
                :is_true,
                :nonpositive,
                :nonnegative,
                :symmetric,
                :invertible,
                :singular,
                :orthogonal,
                :unitary,
                :normal,
                :positive_definite,
                :upper_triangular,
                :lower_triangular,
                :diagonal,
                :triangular,
                :unit_triangular,
                :fullrank,
                :square,
                :real_elements,
                :complex_elements,
                :integer_elements)

for meth in Q_predicates
        nm = string(meth)
        @eval begin
            @doc """
`$($nm)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($nm)
""" ->
            ($meth)(x) = PyCall.pyeval("f(x)", SymPy.Sym, f=SymPy.sympy[:Q][($nm)], x=SymPy.project(x))::SymPy.Sym
        end
    end
end

export Q


text:
using BasisFunctions, BasisFunctions.Test, DomainSets, StaticArrays,
    DoubleFloats, GaussQuadrature, SpecialFunctions

using Test

import BasisFunctions.Test:
    supports_approximation,
    supports_interpolation,
    suitable_function,
    suitable_interpolation_grid,
    dimension_tuple

BF = BasisFunctions

domaintypes = (Float64, LargeFloat)

GaussQuadrature.maxiterations[Double64] = 50
SpecialFunctions.logabsgamma(x::Double64) = Double64.(logabsgamma(BigFloat(x)))
Base.rationalize(x::SVector{N,Double64}) where {N} = SVector{N,Rational{Int}}([rationalize(x_i) for x_i in x])

include("test_dictionaries_util.jl")
include("test_dictionaries_derived.jl")
include("test_dictionaries_discrete.jl")
include("test_dictionaries_tensor.jl")
include("test_dictionaries_mapped.jl")

@testset "test custom dictionary" begin
    mydict = BasisFunctions.MyDictionary()
    test_generic_dict_interface(mydict)
end

test_dictionaries = [Fourier,
    ChebyshevT,
    ChebyshevU,
    Legendre,
    Laguerre,
    Hermite,
    Jacobi,
    CosineSeries,
    SineSeries,
    Monomials,
    RationalFunctions]

for T in domaintypes
    delimit("1D dictionaries ($(T))")
    for DICT in test_dictionaries
        n = 9
        basis = DICT{T}(n)
        @testset "$(string(basis))" begin
            @test length(basis) == n
            @test domaintype(basis) == T
            test_generic_dict_interface(basis)
        end
    end
    # also try a Fourier series with an even length
    @testset begin test_generic_dict_interface(Fourier{T}(8)) end

    delimit("derived dictionaries ($(T))")
    test_derived_dicts(T)

    delimit("Tensor specific tests ($(T))")
    @testset "test iteration" begin
        test_tensor_sets(T)
    end

    delimit("Tensor product set interfaces ($(T))")
    @testset "$(repr(basis))" for basis in
                ( Fourier{T}(5) ⊗ Fourier{T}(7), # Two odd-length Fourier series
                  Fourier{T}(6) ⊗ ChebyshevT{T}(4), # combination of Fourier and Chebyshev
                  Fourier{T}(5) ⊗ Fourier{T}(6), # Odd and even-length Fourier series
                  ChebyshevT{T}(5) ⊗ ChebyshevT{T}(6), # Two Chebyshev sets
                  (Fourier{T}(5) → 2..3) ⊗ (Fourier{T}(7) → 4..5), # Two mapped Fourier series
                  (ChebyshevT{T}(5) → 2..3) ⊗ (ChebyshevT{T}(6) → 4..5), # Two mapped Chebyshev series
                  Fourier{T}(2)^2 ∘ AffineMap(rand(T,2,2),rand(T,2))) # A 2d-mapped dict
        test_generic_dict_interface(basis)
    end

    delimit("Discrete sets ($(T))")
    @testset "discrete sets" begin
        test_discrete_sets(T)
    end
end

delimit("Mapped dictionaries")
test_mapped_dicts()


text:
"""
Reference Sequence

Reference sequence is a sequence of A/C/G/T/N. In the internals, it compresses
`N` positions and consumes less than three bits per base. Unlike `BioSequence`
in Bio.jl, reference sequences are immutable and hence no modifyting
operators are provided.
"""
immutable ReferenceSequence <: Sequence
    data::Vector{UInt64}
    nmask::NMask
    part::UnitRange{Int}
end

function ReferenceSequence{T<:Integer}(seq::ReferenceSequence, part::UnitRange{T})
    ReferenceSequence(seq.data, seq.nmask, part)
end

function Base.convert{A<:DNAAlphabet}(::Type{ReferenceSequence}, seq::BioSequence{A})
    data = Vector{UInt64}(cld(length(seq), 32))
    nmask = falses(length(seq))
    i = 1
    for j in 1:endof(data)
        x = UInt64(0)
        r = 0
        while r < 64 && i ≤ endof(seq)
            nt = seq[i]
            if nt ≤ DNA_T
                x |= convert(UInt64, nt) << r
            elseif nt == DNA_N
                nmask[i] = true
            else
                throw(ArgumentError("found a invalid symbol $(seq[i]) ∉ {A,C,G,T,N} at $i"))
            end
            i += 1
            r += 2
        end
        data[j] = x
    end
    return ReferenceSequence(data, NMask(nmask), 1:length(seq))
end

function Base.convert(::Type{ReferenceSequence}, seq::AbstractString)
    return ReferenceSequence(DNASequence(seq))
end

function Base.convert(::Type{DNASequence}, seq::ReferenceSequence)
    bioseq = DNASequence(length(seq))
    for i in 1:endof(seq)
        bioseq[i] = seq[i]
    end
    return bioseq
end

function Base.convert{S<:AbstractString}(::Type{S}, seq::ReferenceSequence)
    return S([Char(nt) for nt in seq])
end

Base.endof(seq::ReferenceSequence) = length(seq.part)
Base.length(seq::ReferenceSequence) = length(seq.part)
Base.isempty(seq::ReferenceSequence) = length(seq) == 0

function Base.checkbounds(seq::ReferenceSequence, i::Integer)
    if 1 ≤ i ≤ endof(seq)
        return true
    end
    throw(BoundsError(i))
end

function Base.checkbounds(seq::ReferenceSequence, part::UnitRange)
    if 1 ≤ part.start && part.stop ≤ endof(seq)
        return true
    end
    throw(BoundsError(part))
end

@inline function Base.getindex(seq::ReferenceSequence, i::Integer)
    checkbounds(seq, i)
    return unsafe_getindex(seq, i)
end

@inline function unsafe_getindex(seq::ReferenceSequence, i::Integer)
    j = Int(i) + seq.part.start - 2
    d, r = divrem32(j)
    @inbounds begin
        if seq.nmask[j+1]
            return DNA_N
        else
            return DNANucleotide((seq.data[d+1] >> 2r) & 0b11)
        end
    end
end

function Base.getindex{T<:Integer}(seq::ReferenceSequence, part::UnitRange{T})
    checkbounds(seq, part)
    return ReferenceSequence(seq, part)
end

divrem32(i) = i >> 5, i & 0b11111

Base.eltype(::Type{ReferenceSequence}) = DNANucleotide

# iterator
Base.start(seq::ReferenceSequence) = 1
Base.done(seq::ReferenceSequence, i) = i > endof(seq)
Base.next(seq::ReferenceSequence, i) = unsafe_getindex(seq, i), i + 1

# Comparison
# ----------

function Base.(:(==))(seq1::ReferenceSequence, seq2::ReferenceSequence)
    if length(seq1) != length(seq2)
        return false
    end
    for (x, y) in zip(seq1, seq2)
        if x != y
            return false
        end
    end
    return true
end

function Base.(:(==)){A<:DNAAlphabet}(seq1::ReferenceSequence, seq2::BioSequence{A})
    if length(seq1) != length(seq2)
        return false
    end
    for (x, y) in zip(seq1, seq2)
        if x != y
            return false
        end
    end
    return true
end

Base.(:(==)){A<:DNAAlphabet}(seq1::BioSequence{A}, seq2::ReferenceSequence) =
    seq2 == seq1

function Base.cmp(seq1::ReferenceSequence, seq2::ReferenceSequence)
    for (x, y) in zip(seq1, seq2)
        c = cmp(x, y)
        if c != 0
            return c
        end
    end
    return cmp(length(seq1), length(seq2))
end

function Base.cmp{A<:DNAAlphabet}(seq1::ReferenceSequence, seq2::BioSequence{A})
    for (x, y) in zip(seq1, seq2)
        c = cmp(x, y)
        if c != 0
            return c
        end
    end
    return cmp(length(seq1), length(seq2))
end

Base.cmp{A<:DNAAlphabet}(seq1::BioSequence{A}, seq2::ReferenceSequence) =
    -cmp(seq2, seq1)

Base.isless(seq1::ReferenceSequence, seq2::ReferenceSequence) =
    cmp(seq1, seq2) == -1
Base.isless{A<:DNAAlphabet}(seq1::ReferenceSequence, seq2::BioSequence{A}) =
    cmp(seq1, seq2) == -1
Base.isless{A<:DNAAlphabet}(seq1::BioSequence{A}, seq2::ReferenceSequence) =
    cmp(seq1, seq2) == -1

# Finders
# -------

function Base.findnext(seq::ReferenceSequence, val, start::Integer)
    checkbounds(seq, start)
    v = convert(DNANucleotide, val)
    if v == DNA_N
        return findnextn(seq.nmask, start)
    else
        for i in Int(start):endof(seq)
            x = unsafe_getindex(seq, i)
            if x == val
                return i
            end
        end
    end
    return 0
end

function Base.findprev(seq::ReferenceSequence, val, start::Integer)
    checkbounds(seq, start)
    v = convert(DNANucleotide, val)
    if v == DNA_N
        return findprevn(seq.nmask, start)
    else
        for i in Int(start):-1:1
            x = unsafe_getindex(seq, i)
            if x == val
                return i
            end
        end
    end
    return 0
end

Base.findfirst(seq::ReferenceSequence, val) = findnext(seq, val, 1)
Base.findlast(seq::ReferenceSequence, val)  = findprev(seq, val, endof(seq))

# Printers
# --------

Base.summary(seq::ReferenceSequence) = string(length(seq), "nt Reference Sequence")

function Base.show(io::IO, seq::ReferenceSequence)
    println(io, summary(seq), ':')
    showcompact(io, seq)
end

function Base.showcompact(io::IO, seq::ReferenceSequence)
    width = Base.tty_size()[2]
    if length(seq) > width
        h = div(width, 2)
        for i in 1:h-2
            print(io, seq[i])
        end
        print(io, '…')
        for i in endof(seq)-h+1:endof(seq)
            print(io, seq[i])
        end
    else
        for i in 1:endof(seq)
            print(io, seq[i])
        end
    end
end

function Base.print(io::IO, seq::ReferenceSequence)
    for nt in seq
        print(io, nt)
    end
end

# TODO: implement following functions for incremantal construction
# Base.push!(seq::ReferenceSequence, x)
# Base.append!(seq::ReferenceSequence, other)


text:
module AssemblersTests

using Test

end # module


text:
module RLEVectors

using Base.Broadcast
using Requires
using Statistics
using StatsBase
using DataFrames
using Reexport
@reexport using Missings

# RLEVector type
export RLEVector, FloatRle, IntegerRle, BoolRle, StringRle, RLEVectorList, rfirst, rwidth, rlast, rvalue, nrun, similar, collect, similar, starts, widths, widths!, ends, values
import Base: show, length, size, first, last, step, convert, similar, collect, isequal, values, copy

# collections
import Base.==
import Base: eltype, vcat, pop!, push!, popfirst!, pushfirst!, insert!, deleteat!, splice!, resize!, empty!, lastindex, intersect, append!
export       eltype, vcat, pop!, push!, popfirst!, pushfirst!, insert!, deleteat!, splice!, resize!, growat!, empty!, lastindex, intersect, append!
export deleterun!, decrement_run!

# indexing
import Base.Broadcast: BroadcastStyle, Broadcasted
import Base: getindex, setindex!
import Base: iterate
export getindex, setindex!, ind2run, setrun!, ind2runcontext, RLERangesIterator, eachrange, tapply, iterate

# describe
import Base: isempty
export endtype

# group_generics
import Base: broadcast, map
import Base: in, indexin, findmin, findmax
import Base: range, any, all, sum, eltype, unique, minimum, maximum, extrema, first, last
export range, any, all, eltype, sum, unique, minimum, maximum, extrema, first, last
export in, indexin, findmin, findmax
import Statistics: mean, median
export mean, median

# math
import StatsBase: mode, countmap
export mode, countmap

# ranges
export disjoin, disjoin_length, ree, ree!, inverse_ree, numruns, rangeMeans

# utils
export rep

# sorting
import Base.Order: Ordering
import Base.Sort: QuickSortAlg
import Base: sort, sort!, issorted, reverse, reverse!, sortperm, Algorithm
export       sort, sort!, issorted, reverse, reverse!, sortperm, permute_runs

# DataFrames
import DataFrames: AbstractDataFrame, DataFrame, Index, index, nrow, ncol
export RLEDataFrame, nrow, ncol, columns, index, names
export rowSums, rowMeans, rowMedians, colSums, colMeans, colMedians

# Missings
import Missings: allowmissing, disallowmissing
export allowmissing, disallowmissing

### Includes
include("utils.jl")
include("runs.jl")
include("RLEVector-type.jl")
include("RLEDataFrame-type.jl")
include("ranges.jl")
include("describe.jl")
include("indexing.jl")
include("group_generics.jl")
include("collections_api.jl")
include("math.jl")
include("sorting.jl")
include("missing.jl")
include("precompile.jl")
_precompile_()

function __init__()
	 @require RCall="6f49c342-dc21-5d91-9882-a32aef131414" include("rcall.jl")
end

end # module RLEVectors


text:
export AbstractOptBuffer, OptBuffer, Wrapper

"""
    AbstractOptBuffer

Abstract type for buffers used to implement optimizations routines (including stopping
    criteria etc.)
"""
abstract type AbstractOptBuffer end

"""
    OptBuffer

Abstract type for structs designed to store the data required for core optimization logic.
"""
abstract type OptBuffer <: AbstractOptBuffer end

"""
    Wrapper

Abstract type for wrappers around core buffers meant for auxiliary purposes (stopping
    conditions, logging etc.)
"""
abstract type Wrapper <: AbstractOptBuffer end

const OptLogLevel = LogLevel(-10)
const LSLogLevel = LogLevel(-20)

# Default getters of `OptBuffer` fields

"""
    gradientvec(M::AbstractOptBuffer)

Return the gradient vector at the end of the optimization step.
"""
@inline gradientvec(M::OptBuffer) = M.g

"""
    argumentvec(M::AbstractOptBuffer)

Return the argument vector at the end of the optimization step.
"""
@inline argumentvec(M::OptBuffer) = M.x

"""
    step_origin(M::AbstractOptBuffer)

Return the argument vector at the start of the optimization step.
"""
@inline step_origin(M::OptBuffer) = M.xpre

"""
    fnval(M::AbstractOptBuffer)

Return the objective function value at the end of the optimization step.
"""
@inline fnval(M::OptBuffer) = M.y

"""
    fnval_origin(M::AbstractOptBuffer)

Return the objective function value at the start of the optimization step.
"""
@inline fnval_origin(M::OptBuffer) = M.ypre