module KLU

using SparseArrays
using SparseArrays: SparseMatrixCSC
import SparseArrays: nnz

export KLUFactorization, klu

include("../lib/klu_lib.jl")
using .LibKLU

import Base: (\), size, getproperty, setproperty!, propertynames, show

# Convert from 1-based to 0-based indices
function decrement!(A::AbstractArray{T}) where T<:Integer
    for i in eachindex(A); A[i] -= oneunit(T) end
    A
end
decrement(A::AbstractArray{<:Integer}) = decrement!(copy(A))

# Convert from 0-based to 1-based indices
function increment!(A::AbstractArray{T}) where T<:Integer
    for i in eachindex(A); A[i] += oneunit(T) end
    A
end
increment(A::AbstractArray{<:Integer}) = increment!(copy(A))

using .LibKLU:
    SuiteSparse_long,
    klu_l_common,
    klu_common,
    klu_l_defaults,
    klu_defaults,
    klu_l_symbolic,
    klu_symbolic,
    klu_l_free_symbolic,
    klu_free_symbolic,
    klu_l_numeric,
    klu_numeric,
    klu_l_free_numeric,
    klu_zl_free_numeric,
    klu_free_numeric,
    klu_z_free_numeric,
    KLU_OK,
    KLU_SINGULAR,
    KLU_OUT_OF_MEMORY,
    KLU_INVALID,
    KLU_TOO_LARGE,
    klu_l_analyze,
    klu_analyze,
    klu_l_factor,
    klu_factor,
    klu_zl_factor,
    klu_z_factor,
    klu_l_solve,
    klu_zl_solve,
    klu_solve,
    klu_z_solve,
    klu_l_tsolve,
    klu_zl_tsolve,
    klu_tsolve,
    klu_z_tsolve,
    klu_l_extract,
    klu_zl_extract,
    klu_extract,
    klu_z_extract,
    klu_l_sort,
    klu_zl_sort,
    klu_sort,
    klu_z_sort,
    klu_refactor,
    klu_z_refactor,
    klu_l_refactor,
    klu_zl_refactor

using LinearAlgebra

const KLUTypes = Union{Float64, ComplexF64}
const KLUValueTypes = (:Float64, :ComplexF64)

if sizeof(SuiteSparse_long) == 4
    const KLUITypes = Int32
    const KLUIndexTypes = (:Int32,)
else
    const KLUITypes = Union{Int32, Int64}
    const KLUIndexTypes = (:Int32, :Int64)
end

function kluerror(status::Integer)
    if status == KLU_OK
        return
    elseif status == KLU_SINGULAR
        throw(LinearAlgebra.SingularException(0))
    elseif status == KLU_OUT_OF_MEMORY
        throw(OutOfMemoryError())
    elseif status == KLU_INVALID
        throw(ArgumentError("Invalid Status"))
    elseif status == KLU_TOO_LARGE
        throw(OverflowError("Integer overflow has occured"))
    else
        throw(ErrorException("Unknown KLU error code: $status"))
    end
end

kluerror(common::Union{klu_l_common, klu_common}) = kluerror(common.status)

"""
Data structure for parameters of and statistics generated by KLU functions.

# Fields
- `tol::Float64`: Partial pivoting tolerance for diagonal preference
- `btf::Int64`: If `btf != 0` use BTF pre-ordering
- `ordering::Int64`: If `ordering == 0` use AMD to permute, if `ordering == 1` use COLAMD,
if `ordering == 3` use the user provided ordering function.
- `scale::Int64`: If `scale == 1` then `A[:,i] ./= sum(abs.(A[:,i]))`, if `scale == 2` then
`A[:,i] ./= maximum(abs.(A[:,i]))`. If `scale == 0` no scaling is done, and the input is
checked for errors if `scale >= 0`.

See the [KLU User Guide](https://github.com/DrTimothyAldenDavis/SuiteSparse/raw/master/KLU/Doc/KLU_UserGuide.pdf)
for more information.
"""
klu_l_common

"""
Data structure for parameters of and statistics generated by KLU functions.

This is the `Int32` version of [`klu_l_common`](@ref).
"""
klu_common

macro isok(A)
    :(kluerror($(esc(A))))
end

function _klu_name(name, Tv, Ti)
    outname = "klu_" * (Tv === :Float64 ? "" : "z") * (Ti === :Int64 ? "l_" : "_") * name
    return Symbol(replace(outname, "__"=>"_"))
end
function _common(T)
    if T == Int64
        common = klu_l_common()
        ok = klu_l_defaults(Ref(common))
    elseif T == Int32
        common = klu_common()
        ok = klu_defaults(Ref(common))
    else
        throw(ArgumentError("Index type must be Int64 or Int32"))
    end
    if ok == 1
        return common
    else
        throw(ErrorException("Could not initialize common struct."))
    end
end

"""
    KLUFactorization <: Factorization

Matrix factorization type of the KLU factorization of a sparse matrix `A`.
This is the return type of [`klu`](@ref), the corresponding matrix factorization function.

The factors can be obtained from `K::KLUFactorization` via `K.L`, `K.U` and `K.F`
See the [`klu`](@ref) docs for more information.

You typically should not construct this directly, instead use [`klu`](@ref).
"""
mutable struct KLUFactorization{Tv<:KLUTypes, Ti<:KLUITypes} <: LinearAlgebra.Factorization{Tv}
    common::Union{klu_l_common, klu_common}
    _symbolic::Ptr{Cvoid}
    _numeric::Ptr{Cvoid}
    n::Int
    colptr::Vector{Ti}
    rowval::Vector{Ti}
    nzval::Vector{Tv}
    function KLUFactorization(n, colptr, rowval, nzval)
        Ti = eltype(colptr)
        common = _common(Ti)
        obj = new{eltype(nzval), Ti}(common, C_NULL, C_NULL, n, colptr, rowval, nzval)
        function f(klu)
            _free_symbolic(klu)
            _free_numeric(klu)
        end
        return finalizer(f, obj)
    end
end

function _free_symbolic(K::KLUFactorization{Tv, Ti}) where {Ti<:KLUITypes, Tv}
    if Ti == Int64
        klu_l_free_symbolic(Ref(Ptr{klu_l_symbolic}(K._symbolic)), Ref(K.common))
    elseif Ti == Int32
        klu_free_symbolic(Ref(Ptr{klu_symbolic}(K._symbolic)), Ref(K.common))
    end
    K._symbolic = C_NULL
end

for Ti ∈ KLUIndexTypes, Tv ∈ KLUValueTypes
    klufree = _klu_name("free_numeric", Tv, Ti)
    ptr = _klu_name("numeric", :Float64, Ti)
    @eval begin
        function _free_numeric(K::KLUFactorization{$Tv, $Ti})
            $klufree(Ref(Ptr{$ptr}(K._numeric)), Ref(K.common))
            K._numeric = C_NULL
        end
    end
end

function KLUFactorization(A::SparseMatrixCSC{Tv, Ti}) where {Tv<:KLUTypes, Ti<:KLUITypes}
    n = size(A, 1)
    n == size(A, 2) || throw(ArgumentError("KLU only accepts square matrices."))
    # Copying here to match UMFPACK
    # Seems overly defensive, it's probably rare that the user will modify A before the
    # numeric factorization is done.
    return KLUFactorization(n, decrement(A.colptr), decrement(A.rowval), copy(A.nzval))
end

size(K::KLUFactorization) = (K.n, K.n)
function size(K::KLUFactorization, dim::Integer)
    if dim < 1
        throw(ArgumentError("size: dimension $dim out of range"))
    elseif dim == 1 || dim == 2
        return Int(K.n)
    else
        return 1
    end
end

nnz(K::KLUFactorization) = K.lnz + K.unz + K.nzoff

Base.adjoint(K::KLUFactorization) = Adjoint(K)
Base.transpose(K::KLUFactorization) = Transpose(K)

function setproperty!(klu::KLUFactorization, ::Val{:(_symbolic)}, x)
    _free_symbolic(klu)
    setfield!(klu, :(_symbolic), x)
end

function setproperty!(klu::KLUFactorization, ::Val{:(_numeric)}, x)
    _free_numeric(klu)
    setfield!(klu, :(_numeric), x)
end

# Certain sets of inputs must be non-null *together*:
# [Lp, Li, Lx], [Up, Ui, Ux], [Fp, Fi, Fx]
for Tv ∈ KLUValueTypes, Ti ∈ KLUIndexTypes
    extract = _klu_name("extract", Tv, Ti)
    sort = _klu_name("sort", Tv, Ti)
    if Tv === :ComplexF64
        call = :($extract(klu._numeric, klu._symbolic, Lp, Li, Lx, Lz, Up, Ui, Ux, Uz, Fp, Fi, Fx, Fz, P, Q, Rs, R, Ref(klu.common)))
    else
        call = :($extract(klu._numeric, klu._symbolic, Lp, Li, Lx, Up, Ui, Ux, Fp, Fi, Fx, P, Q, Rs, R, Ref(klu.common)))
    end
    @eval begin
        function _extract!(
            klu::KLUFactorization{$Tv, $Ti};
            Lp = C_NULL, Li = C_NULL, Up = C_NULL, Ui = C_NULL, Fp = C_NULL, Fi = C_NULL,
            P = C_NULL, Q = C_NULL, R = C_NULL, Lx = C_NULL, Lz = C_NULL, Ux = C_NULL, Uz = C_NULL,
            Fx = C_NULL, Fz = C_NULL, Rs = C_NULL
        )
            $sort(klu._symbolic, klu._numeric, Ref(klu.common))
            ok = $call
            if ok == 1
                return nothing
            else
                kluerror(klu.common)
            end
        end
    end
end

function Base.propertynames(::KLUFactorization, private::Bool=false)
    publicnames = (:lnz, :unz, :nzoff, :L, :U, :F, :q, :p, :Rs, :symbolic, :numeric,)
    privatenames = (:nblocks, :maxblock,  :(_L), :(_U), :(_F))
    if private
        return (publicnames..., privatenames...)
    else
        return publicnames
    end
end

function getproperty(klu::KLUFactorization{Tv, Ti}, s::Symbol) where {Tv<:KLUTypes, Ti<:KLUITypes}
    # Forwards to the numeric struct:
    if s ∈ [:lnz, :unz, :nzoff]
        klu._numeric == C_NULL && throw(ArgumentError("This KLUFactorization has not yet been factored. Try `klu_factor!`."))
        return getproperty(klu.numeric, s)
    end
    if s ∈ [:nblocks, :maxblock]
        klu._symbolic == C_NULL && throw(ArgumentError("This KLUFactorization has not yet been analyzed. Try `klu_analyze!`."))
        return getproperty(klu.symbolic, s)
    end
    if s === :symbolic
        klu._symbolic == C_NULL && throw(ArgumentError("This KLUFactorization has not yet been analyzed. Try `klu_analyze!`."))
        if Ti == Int64
            return unsafe_load(Ptr{klu_l_symbolic}(klu._symbolic))
        else
            return unsafe_load(Ptr{klu_symbolic}(klu._symbolic))
        end
    end
    if s === :numeric
        klu._numeric == C_NULL && throw(ArgumentError("This KLUFactorization has not yet been factored. Try `klu_factor!`."))
        if Ti == Int64
            return unsafe_load(Ptr{klu_l_numeric}(klu._numeric))
        else
            return unsafe_load(Ptr{klu_numeric}(klu._numeric))
        end
    end
    # Non-overloaded parts:
    if s ∉ [:L, :U, :F, :p, :q, :R, :Rs, :(_L), :(_U), :(_F)]
        return getfield(klu, s)
    end
    # Factor parts:
    if s === :(_L)
        lnz = klu.lnz
        Lp = Vector{Ti}(undef, klu.n + 1)
        Li = Vector{Ti}(undef, lnz)
        Lx = Vector{Float64}(undef, lnz)
        Lz = Tv == Float64 ? C_NULL : Vector{Float64}(undef, lnz)
        _extract!(klu; Lp, Li, Lx, Lz)
        return Lp, Li, Lx, Lz
    elseif s === :(_U)
        unz = klu.unz
        Up = Vector{Ti}(undef, klu.n + 1)
        Ui = Vector{Ti}(undef, unz)
        Ux = Vector{Float64}(undef, unz)
        Uz = Tv == Float64 ? C_NULL : Vector{Float64}(undef, unz)
        _extract!(klu; Up, Ui, Ux, Uz)
        return Up, Ui, Ux, Uz
    elseif s === :(_F)
        fnz = klu.nzoff
        # We often don't have an F, so create the right vectors for an empty SparseMatrixCSC
        if fnz == 0
            Fp = zeros(Ti, klu.n + 1)
            Fi = Vector{Ti}()
            Fx = Vector{Float64}()
            Fz = Tv == Float64 ? C_NULL : Vector{Float64}()
        else
            Fp = Vector{Ti}(undef, klu.n + 1)
            Fi = Vector{Ti}(undef, fnz)
            Fx = Vector{Float64}(undef, fnz)
            Fz = Tv == Float64 ? C_NULL : Vector{Float64}(undef, fnz)
            _extract!(klu; Fp, Fi, Fx, Fz)
            # F is *not* sorted on output, so we'll have to do it here:
            for i ∈ 1:length(Fp) - 1
                Fp[end] -= 1 # deal with last index being last + 1
                # find each segment
                first = Fp[i] + 1
                last = Fp[i+1] + 1
                # sort each column of rowval, nzval, and Fz for complex numbers if necessary
                #by the ascending permutation of rowval.
                Fiview = view(Fi, first:last)
                Fxview = view(Fx, first:last)
                P = sortperm(Fiview)
                Fiview .= Fiview[P]
                Fxview .= Fxview[P]
                if length(Fz) == length(Fx)
                    Fzview = view(Fz, first:last)
                    Fzview .= Fzview[P]
                end
                # reset the last entry
                Fp[end] += 1
            end
        end
        return Fp, Fi, Fx, Fz
    end
    if s ∈ [:q, :p, :R, :Rs]
        if s === :Rs
            out = Vector{Float64}(undef, klu.n)
        elseif s === :R
            out = Vector{Ti}(undef, klu.nblocks + 1)
        else
            out = Vector{Ti}(undef, klu.n)
        end
        # This tuple construction feels hacky, there's a better way I'm sure.
        s === :q && (s = :Q)
        s === :p && (s = :P)
        _extract!(klu; NamedTuple{(s,)}((out,))...)
        if s ∈ [:Q, :P, :R]
            increment!(out)
        end
        return out
    end
    if s ∈ [:L, :U, :F]
        if s === :L
            p, i, x, z = klu._L
        elseif s === :U
            p, i, x, z = klu._U
        elseif s === :F
            p, i, x, z = klu._F
        end
        if Tv == Float64
            return SparseMatrixCSC(klu.n, klu.n, increment!(p), increment!(i), x)
        else
            return SparseMatrixCSC(klu.n, klu.n, increment!(p), increment!(i), Complex.(x, z))
        end
    end
end

function LinearAlgebra.issuccess(K::KLUFactorization)
    return K.common.status == KLU_OK && K._numeric != C_NULL
end
function show(io::IO, mime::MIME{Symbol("text/plain")}, K::KLUFactorization)
    if issuccess(K)
        summary(io, K); println(io)
        println(io, "L factor:")
        show(io, mime, K.L)
        println(io, "\nU factor:")
        show(io, mime, K.U)
        F = K.F
        if F !== nothing
            println(io, "\nF factor:")
            show(io, mime, K.F)
        end
    else
        throw(ArgumentError("Failed factorization of type $(typeof(K)). Try `klu_factor!(K)`."))
    end
end

function klu_analyze!(K::KLUFactorization{Tv, Ti}) where {Tv, Ti<:KLUITypes}
    if K._symbolic != C_NULL return K end
    if Ti == Int64
        sym = klu_l_analyze(K.n, K.colptr, K.rowval, Ref(K.common))
    else
        sym = klu_analyze(K.n, K.colptr, K.rowval, Ref(K.common))
    end
    if sym == C_NULL
        kluerror(K.common)
    else
        K._symbolic = sym
    end
    return K
end

# User provided permutation vectors:
function klu_analyze!(K::KLUFactorization{Tv, Ti}, P::Vector{Ti}, Q::Vector{Ti}) where {Tv, Ti<:KLUITypes}
    if K._symbolic != C_NULL return K end
    if Ti == Int64
        sym = klu_l_analyze_given(K.n, K.colptr, K.rowval, P, Q, Ref(K.common))
    else
        sym = klu_analyze_given(K.n, K.colptr, K.rowval, P, Q, Ref(K.common))
    end
    if sym == C_NULL
        kluerror(K.common)
    else
        K._symbolic = sym
    end
    return K
end

for Tv ∈ KLUValueTypes, Ti ∈ KLUIndexTypes
    factor = _klu_name("factor", Tv, Ti)
    @eval begin
        function klu_factor!(K::KLUFactorization{$Tv, $Ti})
            K._symbolic == C_NULL  && klu_analyze!(K)
            num = $factor(K.colptr, K.rowval, K.nzval, K._symbolic, Ref(K.common))
            if num == C_NULL
                kluerror(K.common)
            else
                K._numeric = num
            end
            return K
        end
    end
end

"""
    klu_factor!(K::KLUFactorization)

Factor `K` into components `K.L`, `K.U`, and `K.F`.
This function will perform both the symbolic and numeric steps of factoriation on an
existing `KLUFactorization` instance.

The `K.common` struct can be used to modify certain options and parameters, see the
[KLU documentation](https://github.com/DrTimothyAldenDavis/SuiteSparse/raw/master/KLU/Doc/KLU_UserGuide.pdf)
or [`klu_common`](@ref) for more information.
"""
klu_factor!

"""
    klu(A::SparseMatrixCSC) -> K::KLUFactorization
    klu(n, colptr::Vector{Ti}, rowval::Vector{Ti}, nzval::Vector{Tv}) -> K::KLUFactorization

Compute the LU factorization of a sparse matrix `A` using KLU.

For sparse `A` with real or complex element type, the return type of `K` is
`KLUFactorization{Tv, Ti}`, with `Tv` = `Float64` or `ComplexF64`
respectively and `Ti` is an integer type (`Int32` or `Int64`).

The individual components of the factorization `K` can be accessed by indexing:

| Component | Description                                                      |
|:----------|:-----------------------------------------------------------------|
| `L`       | `L` (lower triangular) part of `LU` of the diagonal blocks       |
| `U`       | `U` (upper triangular) part of `LU` of the diagonal blocks       |
| `F`       | `F` (upper triangular) part of `LU + F`, the off-diagonal blocks |
| `p`       | right permutation `Vector`                                       |
| `q`       | left permutation `Vector`                                        |
| `Rs`      | `Vector` of scaling factors                                      |

The relation between `K` and `A` is

`K.L * K.U + K.F  == K.Rs `\\` A[K.p, K.q]`

`K` further supports the following functions:

- `LinearAlgebra.\\`

!!! note
    `klu(A::SparseMatrixCSC)` uses the KLU library that is part of
    SuiteSparse. As this library only supports sparse matrices with [`Float64`](@ref) or
    `ComplexF64` elements, `lu` converts `A` into a copy that is of type
    `SparseMatrixCSC{Float64}` or `SparseMatrixCSC{ComplexF64}` as appropriate.
"""
function klu(n, colptr::Vector{Ti}, rowval::Vector{Ti}, nzval::Vector{Tv}) where {Ti<:KLUITypes, Tv<:AbstractFloat}
    if Tv != Float64
        nzval = convert(Vector{Float64}, nzval)
    end
    K = KLUFactorization(n, colptr, rowval, nzval)
    return klu_factor!(K)
end

function klu(n, colptr::Vector{Ti}, rowval::Vector{Ti}, nzval::Vector{Tv}) where {Ti<:KLUITypes, Tv<:Complex}
    if Tv != ComplexF64
        nzval = convert(Vector{ComplexF64}, nzval)
    end
    K = KLUFactorization(n, colptr, rowval, nzval)
    return klu_factor!(K)
end

function klu(A::SparseMatrixCSC{Tv, Ti}) where {Tv<:Union{AbstractFloat, Complex}, Ti<:KLUITypes}
    n = size(A, 1)
    n == size(A, 2) || throw(DimensionMismatch())
    return klu(n, decrement(A.colptr), decrement(A.rowval), A.nzval)
end

for Tv ∈ KLUValueTypes, Ti ∈ KLUIndexTypes
    refactor = _klu_name("refactor", Tv, Ti)
    @eval begin
        function klu!(K::KLUFactorization{$Tv, $Ti}, nzval::Vector{$Tv})
            length(nzval) != length(K.nzval)  && throw(DimensionMismatch())
            K.nzval = nzval
            ok = $refactor(K.colptr, K.rowval, K.nzval, K._symbolic, K._numeric, Ref(K.common))
            if ok == 1
                return K
            else
                kluerror(K.common)
            end
        end
    end
end

function klu!(K::KLUFactorization{ComplexF64}, nzval::Vector{U}) where {U<:Complex}
    return klu!(K, convert(Vector{ComplexF64}, nzval))
end

function klu!(K::KLUFactorization{Float64}, nzval::Vector{U}) where {U<:AbstractFloat}
    return klu!(K, convert(Vector{Float64}, nzval))
end

function klu!(K::KLUFactorization{U}, S::SparseMatrixCSC{U}) where {U}
    size(K) == size(S) || throw(ArgumentError("Sizes of K and S must match."))
    return klu!(K, S.nzval)
end
#B is the modified argument here. To match with the math it should be (klu, B). But convention is
# modified first. Thoughts?
for Tv ∈ KLUValueTypes, Ti ∈ KLUIndexTypes
    solve = _klu_name("solve", Tv, Ti)
    @eval begin
        function solve!(klu::KLUFactorization{$Tv, $Ti}, B::StridedVecOrMat{$Tv})
            stride(B, 1) == 1 || throw(ArgumentError("B must have unit strides"))
            klu._numeric == C_NULL && klu_factor!(klu)
            size(B, 1) == size(klu, 1) || throw(DimensionMismatch())
            isok = $solve(klu._symbolic, klu._numeric, size(B, 1), size(B, 2), B, Ref(klu.common))
            isok == 0 && kluerror(klu.common)
            return B
        end
    end
end

for Tv ∈ KLUValueTypes, Ti ∈ KLUIndexTypes
    tsolve = _klu_name("tsolve", Tv, Ti)
    if Tv === :ComplexF64
        call = :($tsolve(klu._symbolic, klu._numeric, size(B, 1), size(B, 2), B, conj, Ref(klu.common)))
    else
        call = :($tsolve(klu._symbolic, klu._numeric, size(B, 1), size(B, 2), B, Ref(klu.common)))
    end
    @eval begin
        function solve!(klu::Adjoint{$Tv, KLUFactorization{$Tv, $Ti}}, B::StridedVecOrMat{$Tv})
            conj = 1
            klu = parent(klu)
            stride(B, 1) == 1 || throw(ArgumentError("B must have unit strides"))
            klu._numeric == C_NULL && klu_factor!(klu)
            size(B, 1) == size(klu, 1) || throw(DimensionMismatch())
            isok = $call
            isok == 0 && kluerror(klu.common)
            return B
        end
        function solve!(klu::Transpose{$Tv, KLUFactorization{$Tv, $Ti}}, B::StridedVecOrMat{$Tv})
            conj = 0
            klu = parent(klu)
            stride(B, 1) == 1 || throw(ArgumentError("B must have unit strides"))
            klu._numeric == C_NULL && klu_factor!(klu)
            size(B, 1) == size(klu, 1) || throw(DimensionMismatch())
            isok = $call
            isok == 0 && kluerror(klu.common)
            return B
        end
    end
end

function solve(klu, B)
    X = copy(B)
    return solve!(klu, X)
end
LinearAlgebra.ldiv!(klu::KLUFactorization{Tv}, B::StridedVecOrMat{Tv}) where {Tv<:KLUTypes} =
    solve!(klu, B)
LinearAlgebra.ldiv!(klu::LinearAlgebra.AdjOrTrans{Tv, KLUFactorization{Tv, Ti}}, B::StridedVecOrMat{Tv}) where {Tv, Ti} =
    solve!(klu, B)
function LinearAlgebra.ldiv!(klu::KLUFactorization{<:AbstractFloat}, B::StridedVecOrMat{<:Complex})
    imagX = solve(klu, imag(B))
    realX = solve(klu, real(B))
    map!(complex, B, realX, imagX)
end

function LinearAlgebra.ldiv!(klu::LinearAlgebra.AdjOrTrans{Tv, KLUFactorization{Tv, Ti}}, B::StridedVecOrMat{<:Complex}) where {Tv<:AbstractFloat, Ti}
    imagX = solve(klu, imag(B))
    realX = solve(klu, real(B))
    map!(complex, B, realX, imagX)
end

end


@testset "compacctable.jl" begin
    @testset "cleanTable(variables::Vector{<:AbstractIntVar}, table::Matrix{Int})::Matrix{Int}" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(1, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        z = SeaPearl.IntVar(2, 3, "Z", trailer)
        vec = Vector{SeaPearl.AbstractIntVar}([x, y, z])
        table = [
            1 2 2;
            3 3 3;
            1 2 3
        ]
        res = SeaPearl.cleanTable(vec, table)
        @test all(res .== [2 2; 3 3; 2 3])
    end
    @testset "buildSupport(variables::Vector{<:AbstractIntVar}, table::Matrix{Int})::Dict{Pair{Int,Int},BitVector}" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(1, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        z = SeaPearl.IntVar(2, 3, "Z", trailer)
        vec = Vector{SeaPearl.AbstractIntVar}([x, y, z])
        table = [
            1 2 2;
            3 3 3;
            1 2 3
        ]
        table = SeaPearl.cleanTable(vec, table)
        support = SeaPearl.buildSupport(vec, table)
        @test support[2 => 3] == [1,1]
        @test support[1 => 1] == [0,0]
    end
    @testset "cleanSupports!(supports::Dict{Pair{Int,Int},BitVector}, variables::Vector{<:AbstractIntVar})::Nothing" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(1, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        z = SeaPearl.IntVar(2, 3, "Z", trailer)
        vec = [x, y, z]
        table = [
            1 2 2;
            3 3 3;
            1 2 3
        ]
        table = SeaPearl.cleanTable(vec, table)
        support = SeaPearl.buildSupport(vec,table)
        @test all([(i => j) in keys(support) for i = 1:3, j = 2:3])
        @test (1 => 1) in keys(support)
        SeaPearl.cleanSupports!(support, vec)
        @test !((1 => 1) in keys(support))
        @test !((1 => 3) in keys(support))
        @test !((2 => 2) in keys(support))
        @test !((3 => 1) in keys(support))
        @test length(keys(support)) == 4
    end
    @testset "TableConstraint(variables::Vector{<:AbstractIntVar}, table::Matrix{Int}, supports::Dict{Pair{Int,Int},BitVector}, trailer)" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(1, 3, "x", trailer)
        y = SeaPearl.IntVar(1, 3, "y", trailer)
        z = SeaPearl.IntVar(1, 3, "Z", trailer)
        vec1 = [x, y, z]
        vec2 = [z, y, x]
        table = [
            1 2 2;
            3 3 3;
            1 2 3
        ]
        supports = SeaPearl.buildSupport(vec1, table)
        SeaPearl.cleanSupports!(supports, vec1)
        table = SeaPearl.cleanTable(vec1, table)
        constraint1 = SeaPearl.TableConstraint(vec1, table, supports, trailer)
        constraint2 = SeaPearl.TableConstraint(vec2, table, supports, trailer)
        @test size(constraint1.table, 2)==3
        @test size(constraint2.table, 2)==3
        @test (constraint1.supports[3=>2] == [0,1,0])
        @test (constraint2.supports[3=>2] == [0,1,0])
        @test(isempty(constraint1.modifiedVariables))
        @test(isempty(constraint2.modifiedVariables))
        @test(constraint1.residues[3=>2]==1)
        @test(constraint2.residues[3=>2]==1)
    end
    @testset "TableConstraint(variables::Vector{<:AbstractIntVar}, table::Matrix{Int}, trailer)" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(1, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        z = SeaPearl.IntVar(2, 3, "Z", trailer)
        vec = [x, y, z]
        table = [
            1 2 2;
            3 3 3;
            1 2 3
        ]
        constraint = SeaPearl.TableConstraint(vec, table, trailer)
        @test size(constraint.table, 2)==2
        @test (constraint.supports[3=>2] == [1,0])
        @test(isempty(constraint.modifiedVariables))
        @test(constraint.residues[3=>2]==1)
    end
    @testset "propagate!(constraint::TableConstraint, toPropagate::Set{Constraint}, prunedDomains::CPModification)" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(1, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        z = SeaPearl.IntVar(2, 3, "z", trailer)
        a = SeaPearl.IntVar(2, 4, "a", trailer)
        SeaPearl.remove!(a.domain, 3)
        b = SeaPearl.IntVar(3, 6, "b", trailer)
        c = SeaPearl.IntVar(5, 7, "c", trailer)

        table = [
            1 2 2;
            3 3 3;
            1 2 3;
            2 5 4;
            3 7 6;
            5 7 7
        ]
        vec = [x, y, z, a, b, c]
        table = SeaPearl.cleanTable(vec,table)

        constraint = SeaPearl.TableConstraint(vec, table, trailer)

        toPropagate = Set{SeaPearl.Constraint}([constraint])
        modif = SeaPearl.CPModification()
        res = SeaPearl.propagate!(constraint, toPropagate, modif)

        @test !(constraint in toPropagate)

        @test length(x.domain) == 1
        @test length(y.domain) == 1
        @test length(z.domain) == 1
        @test length(a.domain) == 1
        @test length(b.domain) == 1
        @test length(c.domain) == 1
        @test 2 in x.domain
        @test 3 in y.domain
        @test 3 in z.domain
        @test 4 in a.domain
        @test 6 in b.domain
        @test 7 in c.domain
    end

    @testset "compactTable fullPropagation" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(1, 3, "x", trailer)
        y = SeaPearl.IntVar(1, 3, "y", trailer)
        z = SeaPearl.IntVar(1, 3, "z", trailer)
        a = SeaPearl.IntVar(1, 3, "a", trailer)
        b = SeaPearl.IntVar(1, 3, "b", trailer)
        c = SeaPearl.IntVar(1, 3, "c", trailer)

        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addVariable!(model, z)
        SeaPearl.addVariable!(model, a)
        SeaPearl.addVariable!(model, b)
        SeaPearl.addVariable!(model, c)
        table1 = [
            1 2 3 1 2 3;
            1 1 2 2 3 3;
            1 1 1 3 3 3;
            3 2 1 3 2 1
        ]
        table2 = [
            1 2 3;
            1 2 3;
            1 2 3
        ]

        constraint1 = SeaPearl.TableConstraint([x, y, z, a], table1, trailer)
        constraint2 = SeaPearl.TableConstraint([a, b, c], table2, trailer)

        SeaPearl.addConstraint!(model, constraint1)
        SeaPearl.addConstraint!(model, constraint2)

        variableSelection = SeaPearl.MinDomainVariableSelection{false}()
        SeaPearl.solve!(model; variableHeuristic=variableSelection)

        @test length(model.statistics.solutions) == 6
        @test Dict("c" => 2, "x" => 2, "b" => 2, "z" => 3, "a" => 2, "y" => 3) in model.statistics.solutions
        @test Dict("c" => 1, "x" => 3, "b" => 1, "z" => 3, "a" => 1, "y" => 3) in model.statistics.solutions
        @test Dict("c" => 3, "x" => 1, "b" => 3, "z" => 3, "a" => 3, "y" => 2) in model.statistics.solutions
        @test Dict("c" => 1, "x" => 3, "b" => 1, "z" => 1, "a" => 1, "y" => 2) in model.statistics.solutions
        @test Dict("c" => 3, "x" => 1, "b" => 3, "z" => 1, "a" => 3, "y" => 1) in model.statistics.solutions
        @test Dict("c" => 2, "x" => 2, "b" => 2, "z" => 1, "a" => 2, "y" => 1) in model.statistics.solutions
    end
end


struct CrossValidation
    indices
    partition::Integer
    samples::Array
    folds::Integer
    use_all::Bool
end

function CrossValidation(folds::Integer, samples::AbstractArray;shuffle::Bool=true,use_all::Bool=true)
    ids = LinearIndices(samples) |> collect
    shuffle!(ids)
    partition = length(ids) ÷ folds
    CrossValidation(ids, partition, samples, folds, use_all)
end

function (x::CrossValidation)(fold::Integer;)
    (fold < 1 || fold > x.folds) && throw(ArgumentError("Please check 1 ≤ fold ≤ $(x.folds)"))
    last = x.use_all ? length(x.indices) : x.folds * x.partition
    select_start = (fold - 1) * x.partition + 1
    select_end = fold * x.partition
    train = x.samples[x.indices[1:select_start - 1]]
    append!(train, x.samples[x.indices[select_end + 1:last]])
    test = x.samples[x.indices[select_start:select_end]]
    return (train = train, test = test)
end

struct ValidateResult{T <: AbstractFloat}
    dir::String
    scores::Matrix{T}
    labels::Matrix{Bool}
    mask::Matrix{Bool}
    drugs::Array{String}
    diseases::Array{String}
end

show(io::IO,vr::ValidateResult) = print(io, "Validate Results($(vr.dir))")

function ValidateResult(dir::String)
    labels = readdlm(joinpath(dir, "RD.csv"), ',', Bool)
    drugs = readdlm(joinpath(dir, "drugs.csv"), ',', String)
    diseases = readdlm(joinpath(dir, "diseases.csv"), ',', String)
    mask = zeros(size(labels))
    scores = zeros(Float32, size(labels))
    folds = map(x->joinpath(abspath(dir),x),readdir(dir))
    filter!(isdir,folds)
    for f in folds
        mask = mask .+ readdlm(joinpath(dir, f, "mask.csv"), ',', Bool)
        scores = scores .+ readdlm(joinpath(dir, f, "result.csv"), ',', Float32)
    end
    scores = scores ./ mask
    ValidateResult(dir, scores, labels, map(>(0), mask), drugs, diseases)
end

function analyze(res::ValidateResult;union_all::Bool=false)
    if union_all
        cms = threshold_scan(res.labels[res.mask], res.scores[res.mask])
    else
        cms_all = []
        for (name, l, s, m) in zip(res.drugs, eachrow(res.labels), eachrow(res.scores), eachrow(res.mask))
            sum(m)==0 && continue
            tmp = threshold_scan(l[m], s[m])
            push!(cms_all, tmp)
        end
        cms=[]
        for i=1:length(cms_all[1])
            cm=ConfusionMatrix(zeros(eltype(cms_all[1][1]),size(cms_all[1][1])))
            for each in cms_all
                length(each)<i && continue
                cm=cm+each[i]
            end
            push!(cms,cm)
        end
    end
    cms
end

function roc(r::ValidateResult;union_all::Bool=false)
    cms=analyze(r,union_all=union_all)
    TPR = tpr.(cms, 2)
    FPR = fpr.(cms, 2)
    FPR,TPR
end

function pr(r::ValidateResult,union_all::Bool=false)
    cms=analyze(r,union_all=union_all)
    TPR = tpr.(cms, 2)
    FPR = fpr.(cms, 2)
    FPR,TPR
end

function analyze_by_drugs(result::ValidateResult)
    df = DataFrame("Drug Name" => String[], "AUPR" => Float64[], "AUROC" => Float64[])
    for (name, l, s, m) in zip(result.drugs, eachrow(result.labels), eachrow(result.scores), eachrow(result.mask))
        sum(m)==0 && continue
        cms = threshold_scan(l[m], s[m])
        TPR = tpr.(cms, 2)
        FPR = fpr.(cms, 2)
        P = precision.(cms, 2)
        R = recall.(cms, 2)
        aupr = auc(R, P)
        auroc = auc(FPR, TPR)
        push!(df, (name, aupr, auroc))
    end
    push!(df, ("Average of $(size(df, 1)) drugs", mean(df.AUPR), mean(df.AUROC)))
    df
end

function pr_roc(res::ValidateResult;union_all::Bool=false)
    cms=analyze(res;union_all=union_all)
    TPR = tpr.(cms, 2)
    FPR = fpr.(cms, 2)
    P = precision.(cms, 2)
    R = recall.(cms, 2)
    R,P,FPR,TPR
end

function plot(res::ValidateResult;union_all::Bool=false,)
    R,P,FPR,TPR=pr_roc(res,union_all=union_all)
    aupr = auc(R, P)
    auroc = auc(FPR, TPR)
    left = plot(R, P, title="P-R(auc=$(round(aupr, digits=3)))")
    right = plot(FPR, TPR, title="ROC(auc=$(round(auroc, digits=3)))")
    plot(left, right, layout=(1, 2)) |> display
end

export CrossValidation,ValidateResult,analyze_by_drugs,analyze,pr,roc,pr_roc


using LoopVectorization
using StatsBase
using Tullio

const LIKELY_ERROR_CAUSES = """
1. Incorrect inputs -- check your program for bugs. If you provided an `r_eff` argument,
double check it is correct.
2. Your chains failed to converge. Check diagnostics.
3. You do not have enough posterior samples (ESS < ~25).
"""
const MIN_TAIL_LEN = 5  # Minimum size of a tail for PSIS to give sensible answers
const SAMPLE_SOURCES = ["mcmc", "vi", "other"]

export psis, PsisLoo, PsisLooMethod, Psis

###########################
###### RESULT STRUCT ######
###########################

"""
    Psis{R<:Real, AT<:AbstractArray{R, 3}, VT<:AbstractVector{R}}

A struct containing the results of Pareto-smoothed importance sampling.

# Fields

  - `weights`: A vector of smoothed, truncated, and normalized importance sampling weights.
  - `pareto_k`: Estimates of the shape parameter `k` of the generalized Pareto distribution.
  - `ess`: Estimated effective sample size for each LOO evaluation, based on the variance of
    the weights.
  - `sup_ess`: Estimated effective sample size for each LOO evaluation, based on the
    supremum norm, i.e. the size of the largest weight. More likely than `ess` to warn when
    importance sampling has failed. However, it can have a high variance.
  - `r_eff`: The relative efficiency of the MCMC chain, i.e. ESS / posterior sample size.
  - `tail_len`: Vector indicating how large the "tail" is for each observation.
  - `posterior_sample_size`: How many draws from an MCMC chain were used for PSIS.
  - `data_size`: How many data points were used for PSIS.
"""
struct Psis{
    RealType <: Real,
    AT <: AbstractArray{RealType, 3},
    VT <: AbstractVector{RealType},
}
    weights::AT
    pareto_k::VT
    ess::VT
    sup_ess::VT
    r_eff::VT
    tail_len::Vector{Int}
    posterior_sample_size::Int
    data_size::Int
end

function Base.show(io::IO, ::MIME"text/plain", psis_object::Psis)
    table = hcat(psis_object.pareto_k, psis_object.ess, psis_object.sup_ess)
    post_samples = psis_object.posterior_sample_size
    data_size = psis_object.data_size
    println("Results of PSIS with $post_samples posterior samples and $data_size cases.")
    _throw_pareto_k_warning(psis_object.pareto_k)
    return pretty_table(
        table;
        compact_printing=false,
        header=[:pareto_k, :ess, :sup_ess],
        formatters=ft_printf("%5.2f"),
        alignment=:r,
    )
end

###########################
####### PSIS FNCTNS #######
###########################

"""
    psis(
        log_ratios::AbstractArray{T<:Real},
        r_eff::AbstractVector;
        source::String="mcmc"
    ) -> Psis

Implements Pareto-smoothed importance sampling (PSIS).

# Arguments
## Positional Arguments
  - `log_ratios::AbstractArray`: A 2d or 3d array of (unnormalized) importance ratios on the
    log scale. Indices must be ordered as `[data, step, chain]`. The chain index can be left
    off if there is only one chain, or if keyword argument `chain_index` is provided.
  - $R_EFF_DOC

## Keyword Arguments

  - $CHAIN_INDEX_DOC
  - `source::String="mcmc"`: A string or symbol describing the source of the sample being
    used. If `"mcmc"`, adjusts ESS for autocorrelation. Otherwise, samples are assumed to be
    independent. Currently permitted values are $SAMPLE_SOURCES.

See also: [`relative_eff`]@ref, [`psis_loo`]@ref, [`psis_ess`]@ref.
"""
function psis(
    log_ratios::AbstractArray{<:Real, 3};
    r_eff::AbstractVector{<:Real}=similar(log_ratios, 0),
    source::Union{AbstractString, Symbol}="mcmc"
)

    source = lowercase(String(source))
    dims = size(log_ratios)

    data_size = dims[1]
    post_sample_size = dims[2] * dims[3]

    # Reshape to matrix (easier to deal with)
    log_ratios = reshape(log_ratios, data_size, post_sample_size)
    r_eff = _generate_r_eff(log_ratios, dims, r_eff, source)
    weights = similar(log_ratios)
    # Shift ratios by maximum to prevent overflow
    @tturbo @. weights = exp(log_ratios - $maximum(log_ratios; dims=2))

    _check_input_validity_psis(reshape(log_ratios, dims), r_eff)

    tail_length = Vector{Int}(undef, data_size)
    ξ = similar(r_eff)
    @inbounds Threads.@threads for i in eachindex(tail_length)
        tail_length[i] = _def_tail_length(post_sample_size, r_eff[i])
        ξ[i] = @views ParetoSmooth._do_psis_i!(weights[i, :], tail_length[i])
    end

    @tullio norm_const[i] := weights[i, j]
    @tturbo weights .= weights ./ norm_const
    ess = psis_ess(weights, r_eff)
    inf_ess = sup_ess(weights, r_eff)

    weights = reshape(weights, dims)

    return Psis(
        weights,
        ξ,
        ess,
        inf_ess,
        r_eff,
        tail_length,
        post_sample_size,
        data_size
    )
end

function psis(
    log_ratios::AbstractMatrix{<:Real};
    chain_index::AbstractVector{<:Integer}=_assume_one_chain(log_ratios),
    kwargs...,
)
    new_log_ratios = _convert_to_array(log_ratios, chain_index)
    return psis(new_log_ratios; kwargs...)
end

"""
    _do_psis_i!(is_ratios::AbstractVector{Real}, tail_length::Integer) -> T

Do PSIS on a single vector, smoothing its tail values.

# Arguments

  - `is_ratios::AbstractVector{Real}`: A vector of importance sampling ratios,
    scaled to have a maximum of 1.

# Returns

  - `T<:Real`: ξ, the shape parameter for the GPD; big numbers indicate thick tails.
"""
function _do_psis_i!(is_ratios::AbstractVector{T}, tail_length::Integer) where {T <: Real}

    len = length(is_ratios)
    tail_start = len - tail_length + 1  # index of smallest tail value

    # sort is_ratios and also get results of sortperm() at the same time
    ratio_index = collect(zip(is_ratios, Base.OneTo(len)))
    partialsort!(ratio_index, (tail_start-1):len; by=first)
    is_ratios .= first.(ratio_index)
    @views tail = is_ratios[tail_start:len]
    _check_tail(tail)

    # Get value just before the tail starts:
    cutoff = is_ratios[tail_start - 1]
    ξ = _psis_smooth_tail!(tail, cutoff)

    # truncate at max of raw weights (1 after scaling)
    clamp!(is_ratios, 0, 1)
    # unsort the ratios to their original position:
    invpermute!(is_ratios, last.(ratio_index))

    return ξ::T
end

"""
    _def_tail_length(log_ratios::AbstractVector, r_eff::Real) -> tail_len::Integer

Define the tail length as in Vehtari et al. (2019).
"""
function _def_tail_length(length::Int, r_eff::Real)
    return Int(ceil(min(length / 5, 3 * sqrt(length / r_eff))))
end

"""
    _psis_smooth_tail!(tail::AbstractVector{T}, cutoff::T) where {T<:Real} -> ξ::T

Takes an *already sorted* vector of observations from the tail and smooths it *in place*
with PSIS before returning shape parameter `ξ`.
"""
function _psis_smooth_tail!(tail::AbstractVector{T}, cutoff::T) where {T <: Real}
    len = length(tail)
    if any(isinf.(tail))
        return ξ = Inf
    else
        @turbo @. tail = tail - cutoff

        # save time not sorting since tail is already sorted
        ξ, σ = gpdfit(tail)
        @turbo @. tail = gpd_quantile(($(1:len) - 0.5) / len, ξ, σ) + cutoff
    end
    return ξ
end

##########################
#### HELPER FUNCTIONS ####
##########################

"""
Generate the relative effective sample size if not provided by the user.
"""
function _generate_r_eff(
    weights::AbstractArray{R},
    dims::Base.AbstractVecOrTuple,
    r_eff::T,
    source::String,
)::T where {R<:Real, T<:AbstractVector{R}}
    output::T = similar(r_eff, dims[1])
    if isempty(r_eff)
        if source == "mcmc"
            @info "Adjusting for autocorrelation. If the posterior samples are not " *
                  "autocorrelated, specify the source of the posterior sample using the " *
                  "keyword argument `source`. MCMC samples are always autocorrelated; VI " *
                  "samples are not."
            output .= relative_eff(reshape(weights, dims))
        elseif source ∈ SAMPLE_SOURCES
            @info "Samples have not been adjusted for autocorrelation. If the posterior " *
                  "samples are autocorrelated, as in MCMC methods, ESS estimates will be " *
                  "upward-biased, and standard error estimates will be downward-biased. " *
                  "MCMC samples are always autocorrelated; VI samples are not."
            return output .= ones(R, dims[1])
        else
            throw(
                ArgumentError(
                    "$source is not a valid source. Valid sources are $SAMPLE_SOURCES."
                ),
            )
            return output .= ones(R, dims[1])
        end
    else
        return r_eff
    end
    if any(_invalid_number, r_eff)
        throw(
            ArgumentError(
                "PSIS-LOO has encountered an error calculating ESS values for your " *
                "Markov chains. Please check your inputs. $LIKELY_ERROR_CAUSES"
            )
        )
    end
    return output::T
end

"""
Make sure all inputs to `psis` are valid.
"""
function _check_input_validity_psis(
    log_ratios::AbstractArray{T, 3}, r_eff::AbstractVector{T}
) where {T <: Real}
    if any(_invalid_number, log_ratios)
        throw(DomainError("Invalid input for `log_ratios` (contains NaN  or inf values)."))
    elseif isempty(log_ratios)
        throw(ArgumentError("Invalid input for `log_ratios` (array is empty)."))
    elseif length(r_eff) ≠ size(log_ratios, 1)
        throw(ArgumentError("Size of `r_eff` does not equal the number of data points."))
    end
    return nothing
end

"""
Check the tail to make sure a GPD fit is possible.
"""
function _check_tail(tail::AbstractVector{T}) where {T <: Real}
    if maximum(tail) ≈ minimum(tail)
        throw(
            ArgumentError(
                "Unable to fit generalized Pareto distribution: all tail values are the " *
                "same. Likely causes are: \n$LIKELY_ERROR_CAUSES",
            ),
        )
    elseif length(tail) < MIN_TAIL_LEN
        throw(
            ArgumentError(
                "Unable to fit generalized Pareto distribution: tail length was too " *
                "short. Likely causese are: \n$LIKELY_ERROR_CAUSES",
            ),
        )
    end
    return nothing
end

"""
Check if an input is invalid.
"""
function _invalid_number(x::Real)
    return isinf(x) || isnan(x)
end


module SIMDTest

using Base.Test
using ForwardDiff: Dual, valtype

const DUALS = (Dual(1., 2., 3., 4.),
               Dual(1., 2., 3., 4., 5.),
               Dual(Dual(1., 2.), Dual(3., 4.)))

function simd_sum(x::Vector{T}) where T
    s = zero(T)
    @simd for i in eachindex(x)
        @inbounds s = s + x[i]
    end
    return s
end

for D in map(typeof, DUALS)
    plus_bitcode = sprint(io -> code_llvm(io, +, (D, D)))
    @test contains(plus_bitcode, "fadd <4 x double>")

    minus_bitcode = sprint(io -> code_llvm(io, -, (D, D)))
    @test contains(minus_bitcode, "fsub <4 x double>")

    times_bitcode = sprint(io -> code_llvm(io, *, (D, D)))
    @test ismatch(r"fadd \<.*?x double\>", times_bitcode)
    @test ismatch(r"fmul \<.*?x double\>", times_bitcode)

    div_bitcode = sprint(io -> code_llvm(io, /, (D, D)))
    @test ismatch(r"fadd \<.*?x double\>", div_bitcode)
    @test ismatch(r"fmul \<.*?x double\>", div_bitcode)

    exp_bitcode = sprint(io -> code_llvm(io, ^, (D, D)))
    @test ismatch(r"fadd \<.*?x double\>", exp_bitcode)
    if !(valtype(D) <: Dual)
        # see https://github.com/JuliaDiff/ForwardDiff.jl/issues/167
        @test ismatch(r"fmul \<.*?x double\>", exp_bitcode)

        # see https://github.com/JuliaDiff/ForwardDiff.jl/pull/201
        sum_bitcode = sprint(io -> code_llvm(io, simd_sum, (Vector{D},)))
        @test ismatch(r"fadd \<.*?x double\>", sum_bitcode)
    end
end

end # module


let
    lrec = ListRecord(1.0, Float64, Bool, Int)
    append!(lrec.frames, [RecordFrame(1,2), RecordFrame(3,4)])
    append!(lrec.states, [RecordState(1.0,1),RecordState(2.0,2),RecordState(3.0,1),RecordState(4.0,3)])
    lrec.defs[1] = true
    lrec.defs[2] = true
    lrec.defs[3] = false

    sparsemat, id_lookup = get_sparse_lookup(lrec)
    @test sparsemat[1,1] == 1.0
    @test sparsemat[2,1] == 3.0
    @test sparsemat[1,2] == 2.0
    @test sparsemat[2,2] == 0.0
    @test sparsemat[1,3] == 0.0
    @test sparsemat[2,3] == 4.0
    @test id_lookup == Dict(2=>2,3=>3,1=>1)

    qrec = convert(QueueRecord{Entity{Float64, Bool, Int}}, lrec)
    @test qrec[-1][1] == Entity(1.0,true,1)
    @test qrec[-1][2] == Entity(2.0,true,2)
    @test qrec[ 0][1] == Entity(3.0,true,1)
    @test qrec[ 0][2] == Entity(4.0,false,3)
    @test nframes(qrec) == 2
end

let
    qrec = QueueRecord(Entity{Float64, Bool, Int}, 2, 0.1)
    @test nframes(qrec) == 0

    update!(qrec, Frame([Entity(1.0,true,1), Entity(2.0,true,2)]))
    update!(qrec, Frame([Entity(3.0,true,1), Entity(4.0,false,3)]))

    @test nframes(qrec) == 2

    lrec = convert(ListRecord{Float64, Bool, Int}, qrec)
    @test nframes(lrec) == 2
    @test lrec.frames == [RecordFrame(1,2), RecordFrame(3,4)]
    @test lrec.states == [RecordState(1.0,1),RecordState(2.0,2),RecordState(3.0,1),RecordState(4.0,3)]
    @test lrec.defs[1] == true
    @test lrec.defs[2] == true
    @test lrec.defs[3] == false
    @test length(lrec.defs) == 3
end


function sumnondiagonal(x::Matrix{T}) where {T}
  @assert size(x,1) == size(x,2)
  s = zero(T)
  @inbounds for j in 1:size(x,2)
    for i in 1:size(x,1)
      s += (i == j) ? zero(T) : x[i,j]
    end
  end
  s
end

function backsumnondiagonal(x::AbstractMatrix,Δ::T) where {T}
  @assert size(x,1) == size(x,2)
  s = zero(x)
  @inbounds for j in 1:size(x,2)
    for i in 1:size(x,1)
      s[i,j] = (i == j) ? zero(T) : Δ
    end
  end
  s
end

sumnondiagonal(x::Flux.Tracker.TrackedMatrix) = Flux.Tracker.track(sumnondiagonal,x)
Flux.Tracker.@grad function sumnondiagonal(x)
  return(sumnondiagonal(Flux.data(x)), Δ -> (backsumnondiagonal(Flux.data(x) , Δ),))
end


export get_lbi_timescale, get_lbi!

"""
	get_lbi_timescale(lineage, segment)

Return the timescale used for computing LBI values, in branch length units.
 """
function get_lbi_timescale(lineage, segment)
	if typeof(lbi_integration_time) == Year
		ct = lbi_integration_time.value
	elseif typeof(lbi_integration_time) == Month
		ct = lbi_integration_time.value / 12
	elseif typeof(lbi_integration_time) == Day
		ct = lbi_integration_time.value / 365
	else
		ct = lbi_integration_time
	end
	τ = substitution_rate[lineage, segment] * ct
	return τ
end

"""
	get_lbi!(t::Tree{LBIData}, fp::FluPop, strainnames::Array{String}, datemin, datemax, τ, lbi_field)

Compute lbi of `strainsnames` according to tree `t`. Leaves of tree posterior to `datemax` or anterior to `datemin` are set to dead. Tree is modified in the process (life-state and lbi value of nodes).
Store the resulting lbi value in `fp.strains[lbi_field]`
"""
function get_lbi!(t::Tree{LBIData}, fp::FluPop, strainnames::Array{<:AbstractString}, datemin, datemax, τ, lbi_field, verbose=false)
	# Setting life-state of tree nodes
	for node in values(t.leaves)
		if haskey(fp.strains, node.label)
			node.data.alive = datemin <= fp.strains[node.label][:date] < datemax
		else
			node.data.alive = false
		end
	end
	TreeTools.set_live_nodes!(t, set_leaves=false)
	# Computing LBI
	TreeTools.lbi!(t, τ, normalize=true) && verbose && begin
		println("LBI computation failed")
	    println("-- ", length(strainnames), " strains")
	    println("-- Date range $datemin - $datemax")
	end
	for (i,s) in enumerate(strainnames)
		if haskey(t.lleaves, s)
			fp.strains[s].data[lbi_field] = t.lleaves[s].data.lbi
		else
			fp.strains[s].data[lbi_field] = missing
		end
	end
	nothing
end

"""
	function get_lbi!(fp::FluPop, t::Tree{LBIData}, datestyle=:all_anterior;
				lineage="h3n2",
				segment="ha",
				τ=get_tau(lineage,segment),
				lbi_field=:lbi)

Compute lbi of strains in `fp` based on tree `t` and date binning of `fp.datebin`. Tree is modified in the process (life-state of nodes).
The `datestyle` arguments modifies the life-state of nodes when LBI is computed. If `:all_anterior`, all nodes anterior to the current datebin are set to live. If `:datebin`, only nodes in the current datebin are set to live.
If a strain does not exist in the tree, its lbi will be set to `missing`.

"""
function get_lbi!(fp::FluPop, t::Tree{LBIData}, datestyle=:all_anterior;
				lineage="h3n2",
				segment="ha",
				τ=get_lbi_timescale(lineage,segment),
				lbi_field=:lbi,
				verbose=false)
	# Find `datemin` as a function of `datestyle`
	if datestyle == :all_anterior
		datemin = let val = findmin(collect(keys(fp.datebin)))[1][1]
			(x) -> val
		end
	elseif datestyle == :datebin
		datemin = (x) -> minimum(x)
	else
		error("Accepted values for `datestyle`: `:all_anterior` or `:datebin`")
	end
	# Compute LBI for all datebins
	for (datebin,strains) in sort(OrderedDict(fp.datebin), by=x->x[1])
		if length(strains)>1
			labels = [x[:strain] for x in strains]
			get_lbi!(t, fp, labels, datemin(datebin), max(datebin...), τ, lbi_field, verbose)
		else
			for s in strains
				s.data[lbi_field] = missing
			end
		end
	end
	# Setting lbi to missing for strains not in a datebin
	for s in values(fp.strains)
		if !haskey(s.data, lbi_field)
			s.data[lbi_field] = missing
		end
	end
end

"""
	local_lbi_maximas(t::Tree{LBIData}, fp::FluPop, τ)

The idea is to find local maximas of the LBI for each date bin of `fp`. In order to determine if `n` is a local maximum, one needs to know the LBI of its ancestor `a`. Thus, the lbi of `a` has to be recomputed for every datebin, when new leaves come live and old ones die.
## Notes
- `fp` should be date-binned already.
- Strains in `fp` and `t` should match **exactly**.
"""
function local_lbi_maximas!(t::Tree{LBIData}, fp::FluPop, τ=get_lbi_timescale("h3n2","ha"))
	out = Dict()
	# Setting nodes alive in date order
	for (db, strains) in sort(OrderedDict(fp.datebin), by=x->x[1])
		for n in values(t.lnodes)
			n.data.alive = false
		end
		for s in strains
			t.lleaves[s[:strain]].data.alive = true
		end
		TreeTools.set_live_nodes!(t, set_leaves=false)
		TreeTools.lbi!(t, τ, normalize=true)
		# Local maximas for the current tree
		llm = local_lbi_maximas(t)
		out[db] = [(label=llm[i], lbi=t.lnodes[llm[i]].data.lbi) for i in 1:length(llm)]
	end
	return out
end
"""
	local_lbi_maximas(t::Tree{LBIData})

Find all nodes that are a local maxima of the LBI. Does not compute the LBI.
"""
function local_lbi_maximas(t::Tree{LBIData})
	out = String[]
	for (l,n) in t.lnodes
		if n.isroot || n.data.lbi > n.anc.data.lbi
			if mapreduce(x-> n.data.lbi > x.data.lbi, *, n.child, init=true)
				push!(out, n.label)
			end
		end
	end
	return out
end

# if false

# """
# 	local_lbi_maximas(t::Tree{LBIData}, sp::StrainPop, τ)

# The idea is to find local maximas of the LBI for each date bin of `sp`. In order to determine if `n` is a local maximum, one needs to know the LBI of its ancestor `a`. Thus, the lbi of `a` has to be recomputed for every datebin, when new leaves come live and old ones die.
# ## Notes
# - `sp` should be date-binned already.
# - Strains in `sp` and `t` should match **exactly**.
# """
# function local_lbi_maximas!(t::Tree{LBIData}, sp::StrainPop, τ)
# 	out = Dict()
# 	# Setting nodes alive in date order
# 	for (db, strains) in sort(sp.datebin, by=x->x[1])
# 		for n in values(t.lnodes)
# 			n.data.alive = false
# 		end
# 		for s in strains
# 			t.lleaves[s.strain].data.alive = true
# 		end
# 		TreeTools.set_live_nodes!(t, set_leaves=false)
# 		TreeTools.lbi!(t, τ, normalize=true)
# 		# Local maximas for the current tree
# 		llm = local_lbi_maximas(t)
# 		out[db] = [(label=llm[i], lbi=t.lnodes[llm[i]].data.lbi) for i in 1:length(llm)]
# 	end
# 	return out
# end
# """
# 	local_lbi_maximas(t::Tree{LBIData})

# Find all nodes that are a local maxima of the LBI. Does not recompute LBI.
# """
# function local_lbi_maximas(t::Tree{LBIData})
# 	out = String[]
# 	for (l,n) in t.lnodes
# 		if n.isroot || n.data.lbi > n.anc.data.lbi
# 			if mapreduce(x-> n.data.lbi > x.data.lbi, *, n.child, init=true)
# 				push!(out, n.label)
# 			end
# 		end
# 	end
# 	return out
# end

# end


import Base:
  # types
  Array,
  CartesianIndices,
  Vector,
  # symbols
  +,
  -,
  *,
  /,
  ==,
  <,
  >,
  # functions
  adjoint,
  axes,
  complex,
  convert,
  copy,
  copyto!,
  deepcopy,
  deleteat!,
  eachindex,
  eltype,
  fill!,
  getindex,
  hash,
  isapprox,
  isempty,
  isless,
  iterate,
  keys,
  lastindex,
  length,
  map,
  map!,
  ndims,
  push!,
  resize!,
  setindex!,
  show,
  similar,
  size,
  summary,
  truncate,
  zero

import Base.Broadcast:
  # types
  AbstractArrayStyle,
  Broadcasted,
  BroadcastStyle,
  DefaultArrayStyle,
  Style,
  # functions
  _broadcast_getindex,
  broadcasted,
  broadcastable,
  instantiate

import ITensors.ContractionSequenceOptimization:
  optimized_contraction_sequence

import HDF5:
  read,
  write

import LinearAlgebra:
  axpby!,
  axpy!,
  dot,
  eigen,
  exp,
  factorize,
  lmul!,
  mul!,
  norm,
  normalize!,
  qr,
  rmul!,
  svd,
  tr

import NDTensors:
  # Modules
  Strided, # to control threading
  # Methods
  array,
  blockdim,
  blockoffsets,
  contract,
  dense,
  dim,
  dims,
  disable_tblis,
  eachnzblock,
  enable_tblis,
  ind,
  inds,
  insertblock!,
  matrix,
  #maxdim,
  mindim,
  nblocks,
  nnz,
  nnzblocks,
  nzblock,
  nzblocks,
  outer,
  permuteblocks,
  polar,
  scale!,
  setblockdim!,
  sim,
  store,
  sum,
  tensor,
  truncate!,
  using_tblis,
  vector,
  # Deprecated
  addblock!

import Random:
  randn!


##############################################################################
##
## The lightest type that can (i) print table (ii) predict etc
##
##############################################################################

abstract AbstractRegressionResult <: RegressionModel

# fields
coef(x::AbstractRegressionResult) = x.coef
coefnames(x::AbstractRegressionResult) = x.coefnames
vcov(x::AbstractRegressionResult) = x.vcov
nobs(x::AbstractRegressionResult) = x.nobs
df_residual(x::AbstractRegressionResult) = x.df_residual
function confint(x::AbstractRegressionResult)
    scale = quantile(TDist(x.df_residual), 1 - (1-0.95)/2)
    se = stderr(x)
    hcat(x.coef -  scale * se, x.coef + scale * se)
end

# predict, residuals, modelresponse
function predict(x::AbstractRegressionResult, df::AbstractDataFrame)
    rf = deepcopy(x.formula)
    secondstage!(rf)

    newTerms = dropresponse(Terms(rf))
    mf = ModelFrame(newTerms, df)
    newX = ModelMatrix(mf).m

    out = DataArray(Float64, size(df, 1))
    out[mf.msng] = newX * x.coef
end

function residuals(x::AbstractRegressionResult, df::AbstractDataFrame)
    rf = deepcopy(x.formula)
    secondstage!(rf)

    mf = ModelFrame(Terms(rf), df)
    newX = ModelMatrix(mf).m
    out = DataArray(Float64, size(df, 1))
    out[mf.msng] = model_response(mf) -  newX * x.coef
end

function model_response(x::AbstractRegressionResult, df::AbstractDataFrame)
    rf = deepcopy(x.formula)
    secondstage!(rf)
    mf = ModelFrame(Terms(rf), df)
    model_response(mf)
end

# Display Results
function title(::AbstractRegressionResult)
    error("function title has no general method for AbstractRegressionResult")
end

function top(::AbstractRegressionResult)
    error("function top has no general method for AbstractRegressionResult")
end

function coeftable(x::AbstractRegressionResult)
    ctitle = title(x)
    ctop = top(x)
    cc = coef(x)
    se = stderr(x)
    coefnms = coefnames(x)
    conf_int = confint(x)
    # put (intercept) last
    if !isempty(coefnms) && ((coefnms[1] == Symbol("(Intercept)")) || (coefnms[1] == "(Intercept)"))
        newindex = vcat(2:length(cc), 1)
        cc = cc[newindex]
        se = se[newindex]
        conf_int = conf_int[newindex, :]
        coefnms = coefnms[newindex]
    end
    tt = cc ./ se
    CoefTable2(
        hcat(cc, se, tt, ccdf(FDist(1, df_residual(x)), abs2(tt)), conf_int[:, 1], conf_int[:, 2]),
        ["Estimate","Std.Error","t value", "Pr(>|t|)", "Lower 95%", "Upper 95%" ],
        ["$(coefnms[i])" for i = 1:length(cc)], 4, ctitle, ctop)
end

function Base.show(io::IO, x::AbstractRegressionResult)
    show(io, coeftable(x))
end

## Coeftalble2 is a modified Coeftable allowing for a top String matrix displayed before the coefficients.
## Pull request: https://github.com/JuliaStats/StatsBase.jl/pull/119

type CoefTable2
    mat::Matrix
    colnms::Vector
    rownms::Vector
    pvalcol::Integer
    title::AbstractString
    top::Matrix{AbstractString}
    function CoefTable2(mat::Matrix,colnms::Vector,rownms::Vector,pvalcol::Int=0,
                        title::AbstractString = "", top::Matrix = Any[])
        nr,nc = size(mat)
        0 <= pvalcol <= nc || error("pvalcol = $pvalcol should be in 0,...,$nc]")
        length(colnms) in [0,nc] || error("colnms should have length 0 or $nc")
        length(rownms) in [0,nr] || error("rownms should have length 0 or $nr")
        length(top) == 0 || (size(top, 2) == 2 || error("top should have 2 columns"))
        new(mat,colnms,rownms,pvalcol, title, top)
    end
end

## format numbers in the p-value column
function format_scientific(pv::Number)
    return @sprintf("%.3f", pv)
end

function show(io::IO, ct::CoefTable2)
    mat = ct.mat; nr,nc = size(mat); rownms = ct.rownms; colnms = ct.colnms;
    pvc = ct.pvalcol; title = ct.title;   top = ct.top
    if length(rownms) == 0
        rownms = AbstractString[lpad("[$i]",floor(Integer, log10(nr))+3) for i in 1:nr]
    end
    rnwidth = max(4,maximum([length(nm) for nm in rownms]) + 1)
    rownms = [rpad(nm,rnwidth) for nm in rownms]
    widths = [length(cn)::Int for cn in colnms]
    str = [sprint(showcompact,mat[i,j]) for i in 1:nr, j in 1:nc]
    if pvc != 0                         # format the p-values column
        for i in 1:nr
            str[i, pvc] = format_scientific(mat[i, pvc])
        end
    end
    for j in 1:nc
        for i in 1:nr
            lij = length(str[i, j])
            if lij > widths[j]
                widths[j] = lij
            end
        end
    end
    widths .+= 1
    totalwidth = sum(widths) + rnwidth
    if length(title) > 0
        halfwidth = div(totalwidth - length(title), 2)
        println(io, " " ^ halfwidth * string(title) * " " ^ halfwidth)
    end
    if length(top) > 0
        for i in 1:size(top, 1)
            top[i, 1] = top[i, 1] * ":"
        end
        println(io, "=" ^totalwidth)
        halfwidth = div(totalwidth, 2) - 1
        interwidth = 2 +  mod(totalwidth, 2)
        for i in 1:(div(size(top, 1) - 1, 2)+1)
            print(io, top[2*i-1, 1])
            print(io, lpad(top[2*i-1, 2], halfwidth - length(top[2*i-1, 1])))
            print(io, " " ^interwidth)
            if size(top, 1) >= 2*i
                print(io, top[2*i, 1])
                print(io, lpad(top[2*i, 2], halfwidth - length(top[2*i, 1])))
            end
            println(io)
        end
    end
    println("=" ^totalwidth)
    println(io," " ^ rnwidth *
            join([lpad(string(colnms[i]), widths[i]) for i = 1:nc], ""))
    println("-" ^totalwidth)
    for i in 1:nr
        print(io, rownms[i])
        for j in 1:nc
            print(io, lpad(str[i,j],widths[j]))
        end
        println(io)
    end
    println("=" ^totalwidth)
end

##############################################################################
##
## Subtypes of Regression Result
##
##############################################################################

type RegressionResult <: AbstractRegressionResult
    coef::Vector{Float64}   # Vector of coefficients
    vcov::Matrix{Float64}   # Covariance matrix

    esample::BitVector      # Is the row of the original dataframe part of the estimation sample?
    augmentdf::DataFrame

    coefnames::Vector       # Name of coefficients
    yname::Symbol           # Name of dependent variable
    formula::Formula        # Original formula

    nobs::Int64             # Number of observations
    df_residual::Int64      # degrees of freedoms

    r2::Float64             # R squared
    r2_a::Float64           # R squared adjusted
    F::Float64              # F statistics
    p::Float64              # p value for the F statistics
end
title(::RegressionResult) =  "Linear Model"
top(x::RegressionResult) = [
            "Number of obs" sprint(showcompact, nobs(x));
            "Degrees of freedom" sprint(showcompact, nobs(x) - df_residual(x));
            "R2" format_scientific(x.r2);
            "R2 Adjusted" format_scientific(x.r2_a);
            "F Statistic" sprint(showcompact, x.F);
            "p-value" format_scientific(x.p);
            ]

type RegressionResultIV <: AbstractRegressionResult
    coef::Vector{Float64}   # Vector of coefficients
    vcov::Matrix{Float64}   # Covariance matrix

    esample::BitVector      # Is the row of the original dataframe part of the estimation sample?
    augmentdf::DataFrame

    coefnames::Vector       # Name of coefficients
    yname::Symbol           # Name of dependent variable
    formula::Formula        # Original formula

    nobs::Int64             # Number of observations
    df_residual::Int64      # degrees of freedoms

    r2::Float64             # R squared
    r2_a::Float64           # R squared adjusted
    F::Float64              # F statistics
    p::Float64              # p value for the F statistics

    F_kp::Float64           # First Stage F statistics KP
    p_kp::Float64           # First Stage p value KP
end

title(::RegressionResultIV) = "IV Model"
top(x::RegressionResultIV) = [
            "Number of obs" sprint(showcompact, nobs(x));
            "Degrees of freedom" sprint(showcompact, nobs(x) - df_residual(x));
            "R2" format_scientific(x.r2);
            "R2 Adjusted" format_scientific(x.r2_a);
            "F-Statistic" sprint(showcompact, x.F);
            "p-value" format_scientific(x.p);
            "First Stage F-stat (KP)" sprint(showcompact, x.F_kp);
            "First State p-val (KP)" format_scientific(x.p_kp);
            ]

type RegressionResultFE <: AbstractRegressionResult
    coef::Vector{Float64}   # Vector of coefficients
    vcov::Matrix{Float64}   # Covariance matrix

    esample::BitVector      # Is the row of the original dataframe part of the estimation sample?
    augmentdf::DataFrame

    coefnames::Vector       # Name of coefficients
    yname::Symbol           # Name of dependent variable
    formula::Formula        # Original formula

    nobs::Int64             # Number of observations
    df_residual::Int64      # degrees of freedoms

    r2::Float64             # R squared
    r2_a::Float64           # R squared adjusted
    r2_within::Float64
    F::Float64              # F statistics
    p::Float64              # p value for the F statistics

    iterations::Int         # Number of iterations
    converged::Bool         # Has the demeaning algorithm converged?
end
function predict(::RegressionResultFE, ::AbstractDataFrame)
    error("predict is not defined by default for fixed effect models.  Run reg with the the option save = true")
end
function residuals(::RegressionResultFE, ::AbstractDataFrame)
    error("predict is not defined by default for fixed effect models.  Run reg with the the option save = true")
end
title(::RegressionResultFE) = "Fixed Effect Model"
top(x::RegressionResultFE) = [
            "Number of obs" sprint(showcompact, nobs(x));
            "Degrees of freedom" sprint(showcompact, nobs(x) - df_residual(x));
            "R2" format_scientific(x.r2);
            "R2 within" format_scientific(x.r2_within);
            "F-Statistic" sprint(showcompact, x.F);
            "p-value" format_scientific(x.p);
            "Iterations" sprint(showcompact, x.iterations);
            "Converged" sprint(showcompact, x.converged)
            ]

type RegressionResultFEIV <: AbstractRegressionResult
    coef::Vector{Float64}   # Vector of coefficients
    vcov::Matrix{Float64}   # Covariance matrix

    esample::BitVector      # Is the row of the original dataframe part of the estimation sample?
    augmentdf::DataFrame

    coefnames::Vector       # Name of coefficients
    yname::Symbol           # Name of dependent variable
    formula::Formula        # Original formula

    nobs::Int64             # Number of observations
    df_residual::Int64      # degrees of freedoms

    r2::Float64             # R squared
    r2_a::Float64           # R squared adjusted
    r2_within::Float64
    F::Float64              # F statistics
    p::Float64              # p value for the F statistics

    F_kp::Float64           # First Stage F statistics KP
    p_kp::Float64           # First Stage p value KP

    iterations::Int         # Number of iterations
    converged::Bool         # Has the demeaning algorithm converged?
end
function predict(::RegressionResultFEIV, ::AbstractDataFrame)
    error("predict is not defined for fixed effect models. Run reg with the the option savefe = true")
end
function residuals(::RegressionResultFEIV, ::AbstractDataFrame)
    error("residuals is not defined for fixed effect models. Use the function partial_out")
end
title(::RegressionResultFEIV) = "Fixed effect IV Model"
top(x::RegressionResultFEIV) = [
            "Number of obs" sprint(showcompact, nobs(x));
            "Degrees of freedom" sprint(showcompact, nobs(x) - df_residual(x));
            "R2" format_scientific(x.r2);
            "R2 within" format_scientific(x.r2_within);
            "F Statistic" sprint(showcompact, x.F);
            "p-value" format_scientific(x.p);
            "First Stage F-stat (KP)" sprint(showcompact, x.F_kp);
            "First State p-val (KP)" format_scientific(x.p_kp);
            "Iterations" sprint(showcompact, x.iterations);
            "Converged" sprint(showcompact, x.converged)
            ]


#=
Generalized versions of the JuMPeR interface
Delegates to correct cutting plane algorithm via suppFcn
=#
typealias DDUSOracles Union{UMOracle, UIOracle, FBOracle, LCXOracle, UCSOracle}

# JuMPeR alerting us that we're handling this contraint
registerConstraint(w::DDUSOracles, rm::Model, ind::Int, prefs) =
    !get(prefs, :prefer_cuts, true) && error("Only cutting plane supported")

# JuMPeR wants us to generate a constraint for every uncertain constraint in inds
function generateCut(w::DDUSOracles, m::Model, rm::Model, inds::Vector{Int}, active=false)
    new_cons = Any[]
    rd = JuMPeR.getRobust(rm)
    for ind in inds
        con = JuMPeR.get_uncertain_constraint(rm, ind)
        cut_sense, xs, lhs_const = JuMPeR.build_cut_objective(rm, con, m.colVal)
        zstar, ustar = suppFcn(xs, w, cut_sense)
        lhs_of_cut = zstar + lhs_const

        # SUBJECT TO CHANGE: active cut detection
        if active
            push!(rd.activecuts[ind],
                JuMPeR.cut_to_scen(ustar,
                    JuMPeR.check_cut_status(con, lhs_of_cut, w.cut_tol) == :Active))
            continue
        end

        # Check violation
        if JuMPeR.check_cut_status(con, lhs_of_cut, w.cut_tol) != :Violate
            w.debug_printcut && JuMPeR.debug_printcut(rm ,m,w,lhs_of_cut,con,nothing)
            continue  # No violation, no new cut
        end

        # Create and add the new constraint
        new_con = JuMPeR.build_certain_constraint(m, con, ustar)
        w.debug_printcut && JuMPeR.debug_printcut(rm, m, w, lhs_of_cut, con, new_con)
        push!(new_cons, new_con)
    end
    return new_cons
end

# JuMPeR asking us for any reformulations we might want to make - we make none
generateReform(w::DDUSOracles, m::Model, rm::Model, inds::Vector{Int}) = 0


"""
David F. Anderson, Masanori Koyama; An asymptotic relationship between coupling
methods for stochastically modeled population processes. IMA J Numer Anal 2015;
35 (4): 1757-1778. doi: 10.1093/imanum/dru044
"""
SplitCoupledJumpProblem(prob::DiffEqBase.AbstractJumpProblem,prob_control::DiffEqBase.AbstractJumpProblem,aggregator::AbstractAggregatorAlgorithm,coupling_map::Vector{Tuple{Int,Int}})= JumpProblem(cat_problems(prob.prob,prob_control.prob),aggregator,build_split_jumps(prob,prob_control,coupling_map)...)

# make new problem by joining initial_data
function cat_problems(prob::DiscreteProblem,prob_control::DiscreteProblem)
  u0_coupled = CoupledArray(prob.u0,prob_control.u0,true)
  DiscreteProblem(u0_coupled,prob.tspan,prob.p)
end

function cat_problems(prob::DiffEqBase.AbstractODEProblem,prob_control::DiffEqBase.AbstractODEProblem)
  l = length(prob.u0) # add l_c = length(prob_control.u0)
  new_f = function (du,u,p,t)
    prob.f(@view(du[1:l]),u.u,p,t)
    prob_control.f(@view(du[l+1:2*l]),u.u_control,p,t)
  end
  u0_coupled = CoupledArray(prob.u0,prob_control.u0,true)
  ODEProblem(new_f,u0_coupled,prob.tspan,prob.p)
end

function cat_problems(prob::DiscreteProblem,prob_control::DiffEqBase.AbstractODEProblem)
  l = length(prob.u0) # add l_c = length(prob_control.u0)
  if !(typeof(prob.f) <: typeof(DiffEqBase.DISCRETE_INPLACE_DEFAULT))
    @warn("Coupling to DiscreteProblem with nontrivial f. Note that, unless scale_by_time=true, the meaning of f will change when using an ODE/SDE/DDE/DAE solver.")
  end
  new_f = function (du,u,p,t)
    prob.f(@view(du[1:l]),u.u,p,t)
    prob_control.f(@view(du[l+1:2*l]),u.u_control,p,t)
  end
  u0_coupled = CoupledArray(prob.u0,prob_control.u0,true)
  ODEProblem(new_f,u0_coupled,prob.tspan,prob.p)
end

function cat_problems(prob::DiffEqBase.AbstractSDEProblem,prob_control::DiffEqBase.AbstractSDEProblem)
  l = length(prob.u0)
  new_f = function (du,u,p,t)
    prob.f(@view(du[1:l]),u.u,p,t)
    prob_control.f(@view(du[l+1:2*l]),u.u_control,p,t)
  end
  new_g = function (du,u,p,t)
    prob.g(@view(du[1:l]),u.u,p,t)
    prob_control.g(@view(du[l+1:2*l]),u.u_control,p,t)
  end
  u0_coupled = CoupledArray(prob.u0,prob_control.u0,true)
  SDEProblem(new_f,new_g,u0_coupled,prob.tspan,prob.p)
end

function cat_problems(prob::DiffEqBase.AbstractSDEProblem,prob_control::DiffEqBase.AbstractODEProblem)
  l = length(prob.u0)
  new_f = function (du,u,p,t)
    prob.f(@view(du[1:l]),u.u,p,t)
    prob_control.f(@view(du[l+1:2*l]),u.u_control,p,t)
  end
  new_g = function (du,u,p,t)
    prob.g(@view(du[1:l]),u.u,p,t)
    for i in l+1:2*l
      du[i] = 0.
    end
  end
  u0_coupled = CoupledArray(prob.u0,prob_control.u0,true)
  SDEProblem(new_f,new_g,u0_coupled,prob.tspan,prob.p)
end

function cat_problems(prob::DiffEqBase.AbstractSDEProblem,prob_control::DiscreteProblem)
  l = length(prob.u0)
  if !(typeof(prob_control.f) <: typeof(DiffEqBase.DISCRETE_INPLACE_DEFAULT))
    @warn("Coupling to DiscreteProblem with nontrivial f. Note that, unless scale_by_time=true, the meaning of f will change when using an ODE/SDE/DDE/DAE solver.")
  end
  new_f = function (du,u,p,t)
    prob.f(@view(du[1:l]),u.u,p,t)
    prob_control.f(@view(du[l+1:2*l]),u.u_control,p,t)
  end
  new_g = function (du,u,p,t)
    prob.g(@view(du[1:l]),u.u,p,t)
    for i in l+1:2*l
      du[i] = 0.
    end
  end
  u0_coupled = CoupledArray(prob.u0,prob_control.u0,true)
  SDEProblem(new_f,new_g,u0_coupled,prob.tspan)
end

cat_problems(prob_control::DiffEqBase.AbstractODEProblem,prob::DiscreteProblem) = cat_problems(prob,prob_control)
cat_problems(prob_control::DiscreteProblem,prob::DiffEqBase.AbstractSDEProblem) = cat_problems(prob,prob_control)
cat_problems(prob_control::DiffEqBase.AbstractODEProblem,prob::DiffEqBase.AbstractSDEProblem) = cat_problems(prob,prob_control)

# this only depends on the jumps in prob, not prob.prob
function build_split_jumps(prob::DiffEqBase.AbstractJumpProblem,prob_control::DiffEqBase.AbstractJumpProblem,coupling_map::Vector{Tuple{Int,Int}})
  num_jumps = length(prob.discrete_jump_aggregation.rates)
  num_jumps_control = length(prob_control.discrete_jump_aggregation.rates)
  jumps = []
  # overallocates, will fix later
  uncoupled = deleteat!(Vector(1:num_jumps),[c[1] for c in coupling_map ])
  uncoupled_control = deleteat!(Vector(1:num_jumps_control),[c[2] for c in coupling_map ])
  for c in uncoupled   # make uncoupled jumps in prob
    new_rate = prob.discrete_jump_aggregation.rates[c]
    new_affect! = prob.discrete_jump_aggregation.affects![c]
    push!(jumps,ConstantRateJump(new_rate,new_affect!))
  end
  for c in uncoupled_control  # make uncoupled jumps in prob_control
    rate = prob_control.discrete_jump_aggregation.rates[c]
    new_rate = (u,p,t)->rate(u.u_control,p,t)
    affect! = prob_control.discrete_jump_aggregation.affects![c]
    new_affect! = function (integrator)
        flip_u!(integrator.u)
        affect!(integrator)
        flip_u!(integrator.u)
    end
    push!(jumps,ConstantRateJump(new_rate,new_affect!))
  end

  for c in coupling_map # make coupled jumps. 3 new jumps for each old one
    rate = prob.discrete_jump_aggregation.rates[c[1]]
    rate_control = prob_control.discrete_jump_aggregation.rates[c[2]]
    affect! = prob.discrete_jump_aggregation.affects![c[1]]
    affect_control! = prob_control.discrete_jump_aggregation.affects![c[2]]
    # shared jump
    new_affect! = function (integrator)
        affect!(integrator)
        flip_u!(integrator.u)
        affect_control!(integrator)
        flip_u!(integrator.u)
    end
    new_rate = (u,p,t)->min(rate(u.u,p,t),rate_control(u.u_control,p,t))
    push!(jumps,ConstantRateJump(new_rate,new_affect!))
     # only prob
    new_affect! = affect!
    new_rate = (u,p,t)->rate(u.u,p,t)-min(rate(u.u,p,t),rate_control(u.u_control,p,t))
    push!(jumps,ConstantRateJump(new_rate,new_affect!))
    # only prob_control
    new_affect! = function (integrator)
        flip_u!(integrator.u)
        affect!(integrator)
        flip_u!(integrator.u)
    end
    new_rate = (u,p,t)->rate_control(u.u_control,p,t)-min(rate(u.u,p,t),rate_control(u.u_control,p,t))
    push!(jumps,ConstantRateJump(new_rate,new_affect!))
  end
  jumps
end


include("square_metric.jl")
include("norm2_metric.jl")
include("inner_prod_metric.jl")


include("SSA.jl")

function addSim(sys::System, tauVar::TauVar, n::Int64)
    ens = newEnsemble(sys, n)
    npoints = length(sys.data.dd)
    if sys.routine.ssa == :Direct
        for i in 2:npoints
            if i == 2 && sys.routine.shoot
                reShoot!(sys, ens, i-1)
            end
            _sim!(sys, ens, i)
            for k in eachindex(ens)
                for xx in ens[k].x
                    push!(ens[k].xa,xx)
                end
                push!(ens[k].ta,sys.times[i])
            end
        end
    elseif sys.routine.ssa == :Tau
        for i in 2:npoints
            if i == 2 && sys.routine.shoot
                reShoot!(sys, ens, i-1)
            end
            _sim!(sys, ens, i, tauVar)
            for k in eachindex(ens)
                for xx in ens[k].x
                    push!(ens[k].xa,xx)
                end
                push!(ens[k].ta,sys.times[i])
            end
        end
    end
    ens
end

function sim(sys::System, θ::Vector{Float64}, σ2::Vector{Float64}, n::Int64)
    sys.state.θ = θ
    sys.state.σ2 = diagm(σ2)
    tauVar = preAllocTauVar(sys)
    ens = Ensemble()
    append!(ens, addSim(sys, tauVar, n))
    ens
end


using PGFPlots
import JSON

case = ARGS[1]
case_kind = ARGS[2]
string_split = split(case, "_")
case_number = parse(Int, string_split[1])
case_type = string_split[2]

k_max = 0
if case_number == 24
    if case_kind == "det"
        k_max = 38
    else
        k_max = 20
    end
else
    if case_kind == "det"
        k_max = 20
    else
        k_max = 10
    end
end

k_values = collect(2:k_max)

dc_load_shed = zeros(length(k_values))
dc_expected_load_shed = zeros(length(k_values))
dc_times = zeros(length(k_values))
dc_iterations = zeros(length(k_values))
dc_ac_load_shed = zeros(length(k_values))
dc_expected_ac_load_shed = zeros(length(k_values))

soc1_load_shed = zeros(length(k_values))
soc1_expected_load_shed = zeros(length(k_values))
soc1_times = zeros(length(k_values))
soc1_iterations = zeros(length(k_values))
soc1_ac_load_shed = zeros(length(k_values))
soc1_expected_ac_load_shed = zeros(length(k_values))

soc2_load_shed = zeros(length(k_values))
soc2_times = zeros(length(k_values))
soc2_iterations = zeros(length(k_values))
soc2_ac_load_shed = zeros(length(k_values))

for k in 2:k_max
    if case_type == "normal"
        file = string(case_kind, "_", k, "_", case_number, ".json")
    else
        file = string(case_kind, "_", k, "_", case_number, "_", case_type, ".json")
    end
    solution = JSON.parsefile(file)["solution"]
    dc_load_shed[k-1] = solution[1]["load_shed"]
    dc_expected_load_shed[k-1] = solution[1]["expected_load_shed"]
    dc_ac_load_shed[k-1] = (typeof(solution[1]["ac_load_shed"]) == Void) ? dc_load_shed[k-1] : solution[1]["ac_load_shed"]
    dc_expected_ac_load_shed[k-1] = (typeof(solution[1]["expected_ac_load_shed"]) == Void) ? dc_load_shed[k-1] : solution[1]["expected_ac_load_shed"]
    dc_times[k-1] = solution[1]["time"]
    dc_iterations[k-1] = solution[1]["iterations"]

    soc1_load_shed[k-1] = solution[2]["load_shed"]
    soc1_expected_load_shed[k-1] = solution[2]["expected_load_shed"]
    soc1_ac_load_shed[k-1] = (typeof(solution[2]["ac_load_shed"]) == Void) ? soc1_load_shed[k-1] : solution[2]["ac_load_shed"]
    soc1_expected_ac_load_shed[k-1] = (typeof(solution[2]["expected_ac_load_shed"]) == Void) ? soc1_load_shed[k-1] : solution[2]["expected_ac_load_shed"]
    soc1_times[k-1] = solution[2]["time"]
    soc1_iterations[k-1] = solution[2]["iterations"]

    #=
    soc2_load_shed[k-1] = solution[3]["load_shed"]
    soc2_ac_load_shed[k-1] = (typeof(solution[3]["ac_load_shed"]) == Void) ? soc2_load_shed[k-1] : solution[3]["ac_load_shed"]
    soc2_times[k-1] = solution[3]["time"]
    soc2_iterations[k-1] = solution[3]["iterations"]
    =#
end

plt = Axis([
            Plots.Linear(k_values, dc_load_shed*100, markSize=0.5, legendentry=L"DC"),
            Plots.Linear(k_values, soc1_load_shed*100, markSize=0.5, legendentry=L"SOC"),
            Plots.Linear(k_values, dc_ac_load_shed*100, markSize=0.5, legendentry=L"DC(AC)"),
            Plots.Linear(k_values, soc1_ac_load_shed*100, markSize=0.5, legendentry=L"SOC(AC)"),
            #Plots.Linear(k_values, soc2_load_shed, markSize=0.5, legendentry=L"soc-ac")
           ],
           xlabel="k value", title="objective values", xmin=0, xmax=k_values[length(k_values)]+1, legendPos="south east")

save("objective_$(case_number)_$(case_type)_$(case_kind).pdf", plt)

plt = Axis([
            Plots.Linear(k_values, dc_expected_load_shed*100, markSize=0.5, legendentry=L"DC"),
            Plots.Linear(k_values, soc1_expected_load_shed*100, markSize=0.5, legendentry=L"SOC"),
            Plots.Linear(k_values, dc_expected_ac_load_shed*100, markSize=0.5, legendentry=L"DC(AC)"),
            Plots.Linear(k_values, soc1_expected_ac_load_shed*100, markSize=0.5, legendentry=L"SOC(AC)"),
            #Plots.Linear(k_values, soc2_load_shed, markSize=0.5, legendentry=L"soc-ac")
           ],
           xlabel="k value", title="expected values", xmin=0, xmax=k_values[length(k_values)]+1, legendPos="north east")

save("expected_objective_$(case_number)_$(case_type)_$(case_kind).pdf", plt)

plt = Axis([
            Plots.Linear(k_values, dc_iterations, markSize=0.5, legendentry=L"DC"),
            Plots.Linear(k_values, soc1_iterations, markSize=0.5, legendentry=L"SOC"),
            #Plots.Linear(k_values, soc2_iterations, markSize=0.5, legendentry=L"soc-ac")
           ],
           xlabel="k value", title="iteration count", xmin=0, xmax=k_values[length(k_values)]+1, legendPos="north east")

save("iteration_$(case_number)_$(case_type)_$(case_kind).pdf", plt)

plt = Axis([
            Plots.Linear(k_values, dc_times, markSize=0.5, legendentry=L"DC"),
            Plots.Linear(k_values, soc1_times, markSize=0.5, legendentry=L"SOC"),
            #Plots.Linear(k_values, soc2_times, markSize=0.5, legendentry=L"soc-ac")
           ],
           xlabel="k value", title="time taken (seconds)", xmin=0, xmax=k_values[length(k_values)]+1, legendPos="north east")

save("time_$(case_number)_$(case_type)_$(case_kind).pdf", plt)


using MRFingerprintingRecon
using Test

@testset "Greens_Interpolation" begin
    include("cmaps.jl")
end


module MTD

## distributed packages
using Combinatorics
using Dates
using Distributions
using StatsBase

## personal packages
using SparseProbVec

include("mtd_inference.jl")
include("mtdg_inference.jl")
include("mmtd_inference.jl")
include("gof.jl")
include("general.jl")
include("reduce.jl")

end # module


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

@recipe function f(samples::DensitySampleVector;
                params=collect(1:5),
                mean=false,
                std_dev=false,
                globalmode=false,
                localmode=false,
                diagonal = Dict(),
                upper = Dict(),
                lower = Dict(),
                param_labels = [])

    parsel = params
    parsel = parsel[parsel .<= Base.size(samples.params[1], 1)]

    if Base.size(param_labels, 1) == 0
        param_labels = [latexstring("\\theta_$i") for i in parsel]
        param_labels_y = [latexstring("p(\\theta_$i)") for i in parsel]
    else
        param_labels_y = [latexstring("p("*param_labels[i]*")") for i in 1:length(param_labels)]
        param_labels = [latexstring(param_labels[i]) for i in 1:length(param_labels)]
    end

    nparams = length(parsel)
    layout --> nparams^2
    size --> (1000, 600)

    for i in 1:nparams

        # diagonal
        @series begin
            subplot := i + (i-1)*nparams

            seriestype --> get(diagonal, "seriestype", :smallest_intervals)
            bins --> get(diagonal, "bins", 200)
            colors --> get(diagonal, "colors", standard_colors)
            intervals --> get(diagonal, "intervals", standard_confidence_vals)
            legend --> get(diagonal, "legend", false)
            mean --> get(diagonal, "mean", mean)
            std_dev --> get(diagonal, "std_dev", std_dev)
            globalmode --> get(diagonal, "globalmode", globalmode)
            localmode --> get(diagonal, "localmode", localmode)
            xlims --> get(diagonal, "xlims", :auto)
            ylims --> get(diagonal, "ylims", :auto)
            xguide --> param_labels[i]
            yguide --> param_labels_y[i]

            samples, (parsel[i])
        end

        # upper right plots
        for j in i+1:nparams

            @series begin
                subplot := j + (i-1)*nparams

                seriestype --> get(upper, "seriestype", :histogram)
                bins --> get(upper, "bins", 200)
                colors --> get(upper, "colors", standard_colors)
                intervals --> get(upper, "intervals", standard_confidence_vals)
                legend --> get(upper, "legend", false)
                mean --> get(upper, "mean", mean)
                std_dev --> get(upper, "std_dev", std_dev)
                globalmode --> get(upper, "globalmode", globalmode)
                localmode --> get(upper, "localmode", localmode)
                xlims --> get(upper, "xlims", :auto)
                ylims --> get(upper, "ylims", :auto)
                xguide --> param_labels[i]
                yguide --> param_labels[j]

                samples, (parsel[i], parsel[j])
            end

            # lower left plots
            @series begin
                subplot := i + (j-1)*nparams

                seriestype --> get(lower, "seriestype", :smallest_intervals)
                bins --> get(lower, "bins", 200)
                colors --> get(lower, "colors", standard_colors)
                intervals --> get(lower, "intervals", standard_confidence_vals)
                legend --> get(lower, "legend", false)
                mean --> get(lower, "mean", mean)
                std_dev --> get(lower, "std_dev", std_dev)
                globalmode --> get(lower, "globalmode", globalmode)
                localmode --> get(lower, "localmode", localmode)
                xlims --> get(lower, "xlims", :auto)
                ylims --> get(lower, "ylims", :auto)
                xguide --> param_labels[i]
                yguide --> param_labels[j]

                samples, (parsel[i], parsel[j])
            end

        end
    end

end


include("../src/Graph.jl")
include("../src/Tools.jl")
include("../src/Algorithm.jl")
include("../src/Sampling.jl")
include("../src/Greedy.jl")

using StatsBase
using JSON
using JLD2

"""
Set Seed Nodes & change innate opinions
Set parameters.
k = 50
epsilon : s + epsilon
delta : edge probability parameter
beta: threshold
json struct js
"""

epsilon = 0.1
delta = 0.5
beta = 0.5
tags = []

jsdict = Dict{String, Any}()

# Read Graph
buf = split(ARGS[1], ',')
fileName = string("../data/all/", buf[1], ".jld2")

println(fileName)

(name, n, m, V, E) = load(fileName, "name", "n", "m", "V", "E")
G = Graph(n, m, V, E)

merge!(jsdict, Dict("name"=>name, "n"=>n, "m"=>m))

if ARGS[2] == "1"
    #Weighted Cascade
    push!(tags, "WC")
    E_p = load(fileName, "E_WC")
elseif ARGS[2] == "2"
    #Trivalency
    push!(tags, "TM") # Trivalency Model
    E_p = load(fileName, "E_TM")
end
G2 = Graph(n, m, V, E_p)

# Before Making changes
if ARGS[3] == "1"
    # Uniform distribution
    push!(tags, "Uniform")
    s = load(fileName, "s_Uni")
elseif ARGS[3] == "2"
    # Exponential distribution
    push!(tags, "Exponential")
    s = load(fileName, "s_Exp")
elseif ARGS[3] == "3"
    # Power-law distribution
    push!(tags, "Power-law")
    s = load(fileName, "s_Pow")
end

# Different parameters to select k
if ARGS[4] == "1"
    #push!(tags, "0.005")
    k = 1
elseif ARGS[4] == "3"
    #push!(tags, "0.015")
    k = 3
elseif ARGS[4] == "5"
    #push!(tags, "0.02")
    k = 5
end

println("n is ", n)
println("m is ", m)
println("k is ", k)

merge!(jsdict, Dict("k"=>k))
merge!(jsdict, Dict("model"=>tags[1], "s_dist"=>tags[2]))

name_greedy = ["aci", "ad", "ap", "aidc"]
name = ["aci", "ad", "ap", "aidc", "aidc"]

##################### Above are configs
#####################
# The only difference for marketing setting and backfire setting is that Δs is different

if ARGS[5] == "1"
    Δs = delta_s_marketing(s, epsilon)
elseif ARGS[5] == "2"
    Δs = delta_s_backfire(s, epsilon, beta)
end

##################### Below are related to algorithms
######################

#println("s is ", Δs)
Δaci, Δad, Δap, Δaidc, Δaidc_upp = assign_node_weights(G, s, Δs) # linear part weight
lin_qua = [Δaci, Δad, Δap, Δaidc]
lin = [Δaci, Δad, Δap, Δaidc, Δaidc_upp]

Δci, Δd, Δp, Δidc = assign_pair_weights(G, Δs) # quadratic part weight
qua = [Δci, Δd, Δp, Δidc]

lin_labels = ["max_aci_", "max_ad_", "max_ap_", "max_aidc_", "max_aidc_upp_"] # max lin part
qua_labels = ["greedy_aci_", "greedy_ad_", "greedy_ap_", "greedy_aidc_"]

# The initial of different measures
_, init_aci, init_ad, init_ap, init_aidc = Approx(G, s)
init_result_greedy = [init_aci, init_ad, init_ap, init_aidc, init_aidc]
#
init_labels = "init_".*name_greedy
for i in 1:length(name_greedy)
   merge!(jsdict, Dict(init_labels[i]=>init_result_greedy[i]))
end

#println("init aci is ", init_aci)

######################### For the pairs

ris_me, T_me, seeds_me, R_me, R_sample = [], [], [], [], []

for i in 1:length(qua)
    (ris, seeds_qua, R_size, T_qua, R) = ris_measure_greedy(G2, k, lin_qua[i], qua[i], delta) ## RR to select seed nodes
    push!(ris_me, ris)
    push!(T_me, T_qua)
    push!(R_me, R_size)
    push!(R_sample, R)
    push!(seeds_me, seeds_qua)
end

for i in 1:length(qua)
    nn = qua_labels[i].*name_greedy
    lin_nn = qua_labels[i]*"rr" ## init rr_algorithm
    lin_t = qua_labels[i]*"time"
    lin_size = qua_labels[i]*"r_size"

    merge!(jsdict, Dict(nn[i]=>(ris_me[i] + init_result_greedy[i])))
    merge!(jsdict, Dict(lin_t=>T_me[i]))
    merge!(jsdict, Dict(lin_size=>R_me[i]))
end

################### For maximize linear part

ris_me, seeds_me, T_me, R_me = [], [], [], []
# Calculation
for i = 1:length(lin)
    flag = 1
    (ris, seeds_lin, R_size, T_lin) = ris_measure_linear(G2, k, lin[i], flag, delta) ## RR to select seed nodes
    push!(ris_me, ris)
    push!(T_me, T_lin)
    push!(R_me, R_size)
    push!(seeds_me, seeds_lin)
end

for i = 1:length(lin)
    nn = lin_labels[i].*name
    lin_nn = lin_labels[i]*"lin" # only the linear part
    lin_t = lin_labels[i]*"time"
    lin_size = lin_labels[i]*"r_size"

    if i < 5
        result = evaluation_ris(R_sample[i], n, lin_qua[i], qua[i], seeds_me[i])
    else
        result = evaluation_ris(R_sample[4], n, lin_qua[4], qua[4], seeds_me[i])
    end

    merge!(jsdict, Dict(nn[i]=>(result + init_result_greedy[i])))
    merge!(jsdict, Dict(lin_t=>T_me[i]))
    merge!(jsdict, Dict(lin_size=>R_me[i]))
    merge!(jsdict, Dict(lin_nn=>ris_me[i] + init_result_greedy[i]))
end

####################### IO

js = JSON.json(jsdict)

println(jsdict)

if ARGS[5] == "1"
    open("out/greedy_marketing.json", "a+") do ff
        write(ff, js)
        write(ff, "\n")
    end
elseif ARGS[5] == "2"
    open("out/greedy_backfire.json", "a+") do ff
        write(ff, js)
        write(ff, "\n")
    end
end


module Qux

using Hyperspecialize

struct Wobble <: Real
  w::Int
end

f(::Real) = false
g(::Real) = false
@concretize TypicalTag []
@replicable f(::@hyperspecialize(TypicalTag)) = true
@replicable g(::@hyperspecialize(TypicalTag)) = true

h(::Real, ::Wobble) = false

export h, Wobble

end #module


include("a.jl")
include("g.jl")

include("a_p.jl")
include("epsilon_p.jl")
include("kappa_p.jl")

include("kappa_tau.jl")

include("alpha_i_star.jl")
include("alpha_n_star.jl")
include("alpha_r_star.jl")

include("epsilon_b.jl")
include("r_cs.jl")
include("h_cs.jl")
include("b.jl")
include("c.jl")
include("d.jl")

include("perimeter.jl")
include("g_p.jl")

include("kappa_x.jl")
include("delta_x.jl")

include("cross_sectional_area.jl")
include("surface_area.jl")
include("volume.jl")


"""
    most_common_colors(img::String)
img: path to local image file
n: how many common colors to return
"""
function most_common_colors(img::String, n::Int)

    colorthief = pyimport("colorthief")
    color_thief = colorthief.ColorThief(img)
    cols = color_thief.get_palette(color_count=n)

	function getone(x)
		RGB(cols[x][1]/255, cols[x][2]/255,  cols[x][3]/255)
	end

	output = []
	for i in 1:n-1
		push!(output,getone(i))
	end

    return output

end


function coverage_robot(inputs::Array{Float64})
    # TODO: visualize why this doesn't work
    # wall or covered spot in front
    # green = RoboGrid.color_to_float("green")
    # if inputs[3] == 1.0 || inputs[4] == green
    #     return rand([3, 4])
    # end
    # # paint green in self spot
    # if inputs[10] != "green"
    #     return 6
    # end
    # 1
    rand(1:4)
end

no_terminate(e::Episode) = false

@testset "Exit reward" begin
    e = Episode(Grid(), reward=RoboGrid.exit_reward)
    total_reward = run!(e, coverage_robot; terminate=no_terminate)
    println("Found exit in ", e.meta["step"], " steps, ", e.meta["time"], " s, ",
            e.meta["memory"], " bytes")
    @test total_reward == 1.0
end

@testset "Food reward" begin
    g = Grid()
    foodind = findfirst([RoboGrid.object(c) == "food" for c in g.cells])
    g.exits[1] = LinearIndices(g.cells)[foodind]
    e = Episode(g, reward=RoboGrid.food_reward)
    total_reward = run!(e, coverage_robot; terminate=no_terminate)
    println("Found food in ", e.meta["step"], " steps, ", e.meta["time"], " s, ",
            e.meta["memory"], " bytes")
    @test total_reward == 1.0
    @test sum(map(c->RoboGrid.object(c) == "food", e.grid.cells)) == 0
end

@testset "Coverage reward" begin
    e = Episode(Grid(), reward=RoboGrid.coverage_reward)
    total_reward = run!(e, coverage_robot; terminate=no_terminate)
    println("Found exit in ", e.meta["step"], " steps, ", e.meta["time"], " s, ",
            e.meta["memory"], " bytes")
    visited = sum(e.meta["visited"])
    println("Covered $visited cells for $total_reward")
    @test total_reward >= 1
    @test visited == total_reward
end


using Metatheory
using Test
using Metatheory.Library
using Metatheory.EGraphs
using Metatheory.EGraphs.Schedulers

t = @theory begin
    f(a...) |> (println.(a); 42)
end

dump(t)

g = EGraph(:(f(1,2,3)))
saturate!(g, t)


import GeometryBasics, GeometryTypes
using ..TopOptProblems: getdh

"""
map Ferrite cell type to VTKDataTypes cell type
"""
const ferrite_to_vtk = Dict( Triangle => 5,
                            QuadraticTriangle => 22,
                            Quadrilateral => 9,
                            QuadraticQuadrilateral => 23,
                            Tetrahedron => 10,
                            QuadraticTetrahedron => 24,
                            Hexahedron => 12,
                            QuadraticHexahedron => 25
                          )

"""
Converting a Ferrite grid to a VTKUnstructuredData from [VTKDataTypes](https://github.com/mohamed82008/VTKDataTypes.jl).
"""
function VTKDataTypes.VTKUnstructuredData(grid::Ferrite.Grid{dim, <:Ferrite.Cell{dim,N,M}, T}) where {dim, N, M, T}
    celltype = ferrite_to_vtk[eltype(grid.cells)]
    celltypes = [celltype for i in 1:length(grid.cells)]
    connectivity = copy(reinterpret(NTuple{N, Int}, grid.cells))
    node_coords = copy(reshape(reinterpret(Float64, grid.nodes), dim, length(grid.nodes)))
    return VTKUnstructuredData(node_coords, celltypes, connectivity)
end
function VTKDataTypes.VTKUnstructuredData(problem::AbstractTopOptProblem)
    return VTKUnstructuredData(getdh(problem).grid)
end
VTKDataTypes.GLMesh(grid::Ferrite.Grid; kwargs...) = GLMesh(VTKUnstructuredData(grid); kwargs...)
VTKDataTypes.GLMesh(problem::AbstractTopOptProblem; kwargs...) = GLMesh(VTKUnstructuredData(problem); kwargs...)

```
workaround taken from https://github.com/JuliaPlots/Makie.jl/issues/647
TODO: should directly convert VTKDataTypes.VTKUnstructuredData to GeometryBasics.Mesh
Do not want to spend more time on this now...
```
function GeometryBasics.Mesh(glmesh::GeometryTypes.GLNormalVertexcolorMesh)
    newverts = reinterpret(GeometryBasics.Point{3, Float32}, glmesh.vertices)
    newfaces = reinterpret(GeometryBasics.NgonFace{3, GeometryBasics.OffsetInteger{-1, UInt32}}, glmesh.faces)
    newnormals = reinterpret(GeometryBasics.Vec{3, Float32}, glmesh.normals)
    return GeometryBasics.Mesh(GeometryBasics.meta(newverts; normals = newnormals, color = glmesh.color), newfaces)
end

"""
Get mesh of the topopt problem with a given topology indicator vector
"""
function GeometryBasics.Mesh(problem::AbstractTopOptProblem, topology::Array{T,1}; kwargs...) where {T}
    mesh = VTKUnstructuredData(problem)
    topology = round.(topology)
    inds = findall(isequal(0), topology)
    deleteat!(mesh.cell_connectivity, inds)
    deleteat!(mesh.cell_types, inds)
    topology = topology[setdiff(1:length(topology), inds)]
    mesh.cell_data["topology"] = topology
    glmesh = GLMesh(mesh, color = "topology"; kwargs...)
    return GeometryBasics.Mesh(glmesh)
end


cd(@__DIR__)
allowed = ["case", "space", "flux", "collision"]
D = KitBase.read_dict("config.txt", allowed)

D = KitBase.read_dict("config.txt")
ks, ctr, face, simTime = KitBase.initialize(D)

ks, ctr, face, simTime = KitBase.initialize("config.txt")
KitBase.write_jld(ks, ctr)

using KitBase.Plots
plot_line(ks, ctr)
#plot(ks, ctr, legend=:none)


@testset "object" begin
    r, c = (100, 500)
    data = rand(0:10, r, c)
    obs = DataFrame(index=1:c, A=rand(c), B=rand(c))
    var = DataFrame(index=1:r, C=rand(r), D=rand(r))
    @test_throws AssertionError Profile(data, obs, var)

    prof = Profile(data, var, obs)
    prof.layers[:a] = rand(r, c)
    prof.layers[:b] = rand(c, r)
    @test obsnames(prof) == ["index", "A", "B"]
    @test varnames(prof) == ["index", "C", "D"]
    @test nrow(prof) == r
    @test ncol(prof) == c
    @test nvar(prof) == r
    @test nobs(prof) == c
    @test maximum(prof) == 10
    @test minimum(prof) == 0
    @test size(prof) == (r, c)
    @test axes(prof) == (Base.OneTo(r), Base.OneTo(c))
    @test_throws AssertionError prof.obs = var
    @test_throws AssertionError prof.var = obs
    @test size(prof.layers[:a]) == (r, c)
    @test size(prof.layers[:b]) == (c, r)
    @test vec(get_gene_expr(prof, 50)) == data[50, :]
    @test vec(get_gene_expr(prof, 50, :a)) == prof.layers[:a][50, :]
    @test isempty(get_gene_expr(prof, r+11, :a))

    prof2 = copy(prof)
    @test prof2 !== prof
    @test prof.data == prof2.data
    @test prof.var == prof2.var
    @test prof.obs == prof2.obs

    prof2 = filter(:C => x -> x > 0, prof)
    @test prof2 !== prof
    @test prof2.var == prof.var[prof.var.C .> 0, :]
    @test prof2.data == prof.data[prof.var.C .> 0, :]
    @test prof2.obs == prof.obs
    @test prof2.layers[:a] == prof.layers[:a][prof.var.C .> 0, :]
    @test prof2.layers[:b] == prof.layers[:b]

    filter!(:C => x -> x > 0, prof2)
    @test prof2.var == prof.var[prof.var.C .> 0, :]
    @test prof2.data == prof.data[prof.var.C .> 0, :]
    @test prof2.obs == prof.obs
    @test prof2.layers[:a] == prof.layers[:a][prof.var.C .> 0, :]
    @test prof2.layers[:b] == prof.layers[:b]

    prof2 = prof[1:50, :]
    @test prof2.data == prof.data[1:50, :]

    idx = rand([false,true], r)
    prof3 = prof[idx, :]
    @test prof3.data == prof.data[idx, :]
    @test prof3.var == prof.var[idx, :]
    @test prof3.obs == prof.obs

    idx2 = rand([false,true], c)
    prof4 = prof[:, idx2]
    @test prof4.data == prof.data[:, idx2]
    @test prof4.var == prof.var
    @test prof4.obs == prof.obs[idx2, :]
end