using Knet, ArgParse, JLD
include("lm_preprocess.jl")
include("lm_model.jl")

function main(args=ARGS)
    s = ArgParseSettings()
    s.exc_handler = ArgParse.debug_handler
    @add_arg_table s begin
        ("--trainfile"; required=true; help="Infinite stream training file")
        ("--devfile"; help="dev data to test perplexity")
        ("--vocabfile"; required=true; help="Vocabulary file to train a model, each line contains a single word")
        ("--wordsfile"; required=true; help="Words file used to hold all words, each line contains a single word")
        ("--charlines"; arg_type=Int; default=40000; help="Numoflines for character initialization")
        ("--vosave"; help="Words file used to hold all words")
        ("--loadfile"; help="Loadfile for model loading")
        ("--atype"; default=(gpu() >= 0 ? "KnetArray{Float32}" : "Array{Float32}"))
        ("--hiddens"; arg_type=Int; nargs='+'; default=[300]; help="hidden layer configuration")
        ("--embedding"; arg_type=Int; default=350)
        ("--chembedding"; arg_type=Int; default=15)
        ("--batchsize"; arg_type=Int; default=25)
        ("--gclip"; arg_type=Float64; default=5.0)
        ("--savefile"; help="To save the julia model")
    end

    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true)
    atype = eval(parse(o[:atype]))

    for (k, v) in o
        println("$k => $v")
    end

    # word data initialization
    ulimit = 28; maxlines = 500;
    word_vocab_out = create_vocab(o[:vocabfile])
    word_vocab_all = create_vocab(o[:wordsfile])
    i2w_all = Array(AbstractString, length(word_vocab_all))
    for (k, v) in word_vocab_all; i2w_all[v] = k;end;
    dstream = open(o[:trainfile]); s = Dict{Int64, Array{Any, 1}}();
    readstream!(dstream, s, word_vocab_out, word_vocab_all; maxlines=1000, ulimit=ulimit)

    # character level initialization
    char_vocab = create_chvocab(o[:wordsfile], o[:charlines])
    println("Character vocabulary length: $(length(char_vocab))")
    flush(STDOUT)
    if o[:vosave] != nothing
        x = o[:vosave]
        println("saving vocabularies to $x")
        save(o[:vosave], "char_vocab", char_vocab, "word_vocab", word_vocab_out)
    end

    # model initialization
    charhiddens = [o[:embedding]]; wordvsize = length(word_vocab_out); chvsize = length(char_vocab);
    if o[:loadfile] == nothing
        m = initmodel(atype, o[:hiddens], charhiddens, o[:chembedding], wordvsize, chvsize)
    else # to fix broken trainings
        x = load(o[:loadfile])
        m = revconvertmodel(x["model"]) # to convert KnetArray
    end
    char_states = initstate(atype, charhiddens, o[:batchsize])
    states = initstate(atype, o[:hiddens], o[:batchsize])
    opts = oparams(m, Adam; gclip=o[:gclip])

    loss = this_loss = 0

    ids = nextbatch(dstream, s, word_vocab_out, word_vocab_all, o[:batchsize]; maxlines=100, ulimit=ulimit)
    bcount = 1
    println("Training started...")
    flush(STDOUT)
    while ids !=nothing
        this_loss = train(m, char_states, states, ids, i2w_all, char_vocab, opts)

        # running average loss
        if bcount < 100
            loss = (this_loss + (bcount - 1)*loss) / bcount
        else
            loss = loss * 0.99 + this_loss * 0.01
        end

        # checkpoints and save
        if bcount % 4000 == 0
            perp = exp(loss)
            println("Running average perlexity is $perp")
            moc = convertmodel(m)
            save(o[:savefile], "model", moc, "char_vocab", char_vocab, "word_vocab", word_vocab_out)
            flush(STDOUT)
        end

        ids = nextbatch(dstream, s, word_vocab_out, word_vocab_all, o[:batchsize]; maxlines=300, ulimit=ulimit)
        bcount += 1
    end
    convertfile(o[:savefile], "temp.jld")

end

function convertfile(infile, outfile)
    atr(x)=transpose(Array(x)) # LM stores in row-major Array, we use col-major Array
    d = load(infile); m = d["model"]
    save(outfile, "cembed", atr(m[:cembed]), "forw", map(atr,m[:forw]),
         "back",map(atr,m[:back]), "soft",map(atr,m[:soft]), "char",map(atr,m[:char]),
         "char_vocab",d["char_vocab"], "word_vocab",d["word_vocab"],
         "sosword","<s>","eosword","</s>","unkword","<unk>",
         "sowchar",'\x12',"eowchar",'\x13',"unkchar",'\x11')
end
!isinteractive() && main(ARGS)


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

using Chuffed
using MathOptInterface
using ConstraintProgrammingExtensions

using Test

const MOI = MathOptInterface
const MOIB = MOI.Bridges
const MOIT = MOI.Test
const MOIU = MathOptInterface.Utilities
const CP = ConstraintProgrammingExtensions
const COIT = CP.Test

# Adapted copy of AmplNLWriter's optimizer (in the tests)
# https://github.com/jump-dev/AmplNLWriter.jl/blob/327fa0bd46b48d7b2be9bbc08e728070d89d0943/test/MOI_wrapper.jl
function optimizer(T)
    return MOIU.CachingOptimizer(
        MOIU.UniversalFallback(MOIU.Model{T}()),
        MOI.Bridges.full_bridge_optimizer(
            MOIU.CachingOptimizer(
                MOIU.UniversalFallback(MOIU.Model{T}()),
                Chuffed.Optimizer(),
            ),
            T,
        ),
    )
end

@testset "Chuffed" begin
    include("chuffed.jl")
    include("MOI.jl")
end


# Deprecated methods
using NNlib: DenseConvDims
import NNlib: stride, padding, dilation, flipkernel, spatial_dims, kernel_size,
    conv!, ∇conv_filter!, ∇conv_data!,
    maxpool!, meanpool!, ∇maxpool!, ∇meanpool!, PoolDims

using CUDA.CUDNN: scalingParameter, CUDNN_CONVOLUTION, convdims,
                  cudnnConvolutionDescriptor, cudnnConvolutionBwdDataAlgoPerf,
                  cudnnConvolutionForward!, cudnnConvolutionBwdFilterAlgoPerf,
                  cudnnConvolutionBackwardData, cudnnConvolutionBackwardFilter,
                  cudnnConvolutionBackwardBias

const CUDNNFloat = Union{Float16,Float32,Float64}

# Since CUDNN does not support 1D convolution, Conv in Flux will give a CUDNNError if the size is 1-dimensional.
fix1d(x) = x
fix1d(x::DenseCuArray{T, 3}) where T = reshape(x, 1, size(x, 1), size(x, 2), size(x, 3))
fix1d(cdims::DenseConvDims{1,K,C_in,C_out,S,P,D,F}) where {K,C_in,C_out,S,P,D,F} =
    DenseConvDims{2,(1,K...),C_in,C_out,(1,S...),(0,0,P...),(1,D...),F}((1,cdims.I...))
fix1d(pdims::PoolDims{1,K,S,P,D}) where {K,S,P,D,F} =
    PoolDims{2,(1,K...),(1,S...),(0,0,P...),(1,D...)}((1,pdims.I...), pdims.C_in)

# Convolution

function cudnnConvolutionDescriptor(cdims::DenseConvDims, x::DenseCuArray{T}) where T
    cdims, x = fix1d(cdims), fix1d(x)
    mode=(NNlib.flipkernel(cdims) ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION)
    cudnnConvolutionDescriptor(convdims(nnlibPadding(cdims),size(x),0),
                               convdims(NNlib.stride(cdims),size(x),1),
                               convdims(NNlib.dilation(cdims),size(x),1),
                               mode,
                               cudnnDataType(T),
                               math_mode(),
                               CUDNN_DEFAULT_REORDER,
                               Cint(1))
end

function conv!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{T}, cdims::DenseConvDims;
               alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    d = cudnnConvolutionDescriptor(cdims, x)
    cudnnConvolutionForward!(y, w, x, d; alpha, beta, z=y)
end

function NNlib.conv_bias_act!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{T},
                            cdims::DenseConvDims, bias::DenseCuArray{T}, σ=identity;
                            z::DenseCuArray{T}=y, alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "The algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    d = cudnnConvolutionDescriptor(cdims, x)
    # only relu and identity are supported by cudnnConvolutionForward!
    activation = (σ == NNlib.relu ? CUDNN_ACTIVATION_RELU : CUDNN_ACTIVATION_IDENTITY)
    cudnnConvolutionForward!(y, w, x, d; z, bias, activation, alpha, beta)
    if activation === CUDNN_ACTIVATION_IDENTITY && σ ∉ (nothing, identity)
        y = σ.(y)
    end
    return y
end

function ∇conv_data!(dx::DenseCuArray{T}, dy::DenseCuArray{T}, w::DenseCuArray{T},
                     cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "The algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    alpha, beta = scalingParameter(T,alpha), scalingParameter(T,beta);
    xDesc, yDesc, wDesc = cudnnTensorDescriptor(dx), cudnnTensorDescriptor(dy), cudnnFilterDescriptor(w)
    convDesc = cudnnConvolutionDescriptor(cdims, dx)
    p = cudnnConvolutionBwdDataAlgoPerf(wDesc, w, yDesc, dy, convDesc, xDesc, dx)
    with_workspace(p.memory) do workspace
        cudnnConvolutionBackwardData(handle(), alpha, wDesc, w, yDesc, dy, convDesc, p.algo, workspace, sizeof(workspace), beta, xDesc, dx)
    end
    return dx
end

function ∇conv_filter!(dw::DenseCuArray{T}, x::DenseCuArray{T}, dy::DenseCuArray{T},
                       cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "The algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    alpha, beta = scalingParameter(T,alpha), scalingParameter(T,beta);
    xDesc, yDesc, wDesc = cudnnTensorDescriptor(x), cudnnTensorDescriptor(dy), cudnnFilterDescriptor(dw)
    convDesc = cudnnConvolutionDescriptor(cdims, x)
    p = cudnnConvolutionBwdFilterAlgoPerf(xDesc, x, yDesc, dy, convDesc, wDesc, dw);
    with_workspace(p.memory) do workspace
        cudnnConvolutionBackwardFilter(handle(), alpha, xDesc, x, yDesc, dy, convDesc, p.algo, workspace, sizeof(workspace), beta, wDesc, dw);
    end
    return dw
end

function ∇conv_bias!(db::DenseCuArray{T}, dy::DenseCuArray{T}; alpha=1, beta=0) where T<:CUDNNFloat
    alpha,beta = scalingParameter(T,alpha), scalingParameter(T,beta)
    bDesc, yDesc = cudnnTensorDescriptor.((db,dy))
    cudnnConvolutionBackwardBias(handle(), alpha, yDesc, dy, beta, bDesc, db)
    return db
end

# Compatibility shims until users upgrade to new NNlib format
function conv!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{T}; pad=0, stride=1, flipkernel=0, dilation=1, kwargs...) where {T<:CUDNNFloat}
    cdims = DenseConvDims(x, w; padding=pad, stride=stride, flipkernel=(flipkernel!=0), dilation=dilation)
    return conv!(y, x, w, cdims; kwargs...)
end

function ∇conv_filter!(dw::DenseCuArray{T}, dy::DenseCuArray{T}, x::DenseCuArray{T}; pad=0, stride=1, flipkernel=0, dilation=1, kwargs...) where {T<:CUDNNFloat}
    cdims = DenseConvDims(x, dw; padding=pad, stride=stride, flipkernel=(flipkernel!=0), dilation=dilation)
    # NOTE!!!  This compat shim re-arranges the argument order!
    return ∇conv_filter!(dw, x, dy, cdims; kwargs...)
end

function cudnnConvolutionForward(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N},
                                 cdims::DenseConvDims; algo=0, alpha=1, beta=0) where {T,N}
    # @warn "`cudnnConvolutionForward(y,x,w,c::DenseConvDims)` is deprecated, please use one of the methods in `@doc cudnnConvolutionForward!`." maxlog=1
    cudnnConvolutionForward!(y, w, x; alpha, beta, padding=nnlibPadding(cdims), stride=NNlib.stride(cdims), dilation=NNlib.dilation(cdims), mode=(NNlib.flipkernel(cdims) ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION))
end

function cudnnConvolutionBiasActivationForward(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N}, z::DenseCuArray{T,N}, bias::DenseCuArray{T,N},
                                               cdims::DenseConvDims; algo=0, alpha1=1, alpha2=1,
                                               activationMode=CUDNN_ACTIVATION_RELU, activationCoeff=0.0, activationReluNanOpt=CUDNN_NOT_PROPAGATE_NAN) where {T,N}
    # @warn "`cudnnConvolutionBiasActivationForward` is deprecated, please use one of the methods in `@doc cudnnConvolutionForward!`." maxlog=1
    cudnnConvolutionForward!(y, w, x; bias, activation=activationMode, z, alpha=alpha1, beta=alpha2, padding=nnlibPadding(cdims), stride=NNlib.stride(cdims), dilation=NNlib.dilation(cdims), mode=(NNlib.flipkernel(cdims) ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION))
end


#=
Copyright 2021 BlackRock, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

@testset "Settings" begin
    # should not error
    Settings(; ρ=[1], σ=[1])

    # proper consntruction
    @test_throws AssertionError Settings(; ρ=[1], σ=[1], α=3)
    @test_throws AssertionError Settings(; ρ=[-1], σ=[1])
    @test_throws AssertionError Settings(; ρ=[1], σ=[-1])
    @test_throws AssertionError Settings(; ρ=[1], σ=[1], max_iters=-1)
    @test_throws AssertionError Settings(; ρ=[1], σ=[1], obj_tol=-1)
    @test_throws AssertionError Settings(; ρ=[1], σ=[1], res_tol=-1)
    @test_throws AssertionError Settings(; ρ=[1], σ=[1], non_improvement_iters=-1)

    s1 = Settings(3, 3)
    s2 = Settings(; ρ=[1, 1, 1], σ=[1, 1, 1])

    @test s1.ρ == s2.ρ
    @test s1.σ == s2.σ
end

@testset "Params" begin
    # should not error
    P = [1.0 0.0; 0.0 2.0]
    q = [1.0, 2.0]
    A = [1.0 2.0; 0.0 1.0; 1.0 2.0]
    b = [1.0, 2.0, 3.0]
    g = repeat([zero(PiecewiseQuadratic)], 2)
    AdmmParams(P, q, A, b, g)

    @test_throws AssertionError AdmmParams([1.0 0.0 3.0; 0.0 2.0 2.0], q, A, b, g)
    @test_throws AssertionError AdmmParams(P, [1.0, 2.0, 3.0], A, b, g)
    @test_throws AssertionError AdmmParams(P, q, A, [1.0, 2.0], g)
    @test_throws AssertionError AdmmParams(P, q, A, b,
                                           repeat([zero(PiecewiseQuadratic)], 3))
end

@testset "Vars" begin
    # should not error
    Vars(5, 1)
    Vars(zeros(5), zeros(1))
    Vars(zeros(5), zeros(1), zeros(5), zeros(1), zeros(5), zeros(1))

    @test_throws AssertionError Vars(zeros(5), zeros(2), zeros(5), zeros(1), zeros(5),
                                     zeros(1))
    @test_throws AssertionError Vars(zeros(5), zeros(1), zeros(2), zeros(1), zeros(5),
                                     zeros(1))
    @test_throws AssertionError Vars(zeros(5), zeros(1), zeros(5), zeros(2), zeros(5),
                                     zeros(1))
    @test_throws AssertionError Vars(zeros(5), zeros(1), zeros(5), zeros(1), zeros(2),
                                     zeros(1))
    @test_throws AssertionError Vars(zeros(5), zeros(1), zeros(5), zeros(1), zeros(5),
                                     zeros(2))
end

@testset "Validity" begin
    P = [1.0 0.0; 0.0 2.0]
    q = [1.0, 2.0]
    A = [1.0 2.0; 0.0 1.0; 1.0 2.0]
    b = [1.0, 2.0, 3.0]
    g = repeat([zero(PiecewiseQuadratic)], 2)
    params = AdmmParams(P, q, A, b, g)

    vars = Vars(2, 3)

    settings = Settings(; ρ=[1.0, 2.0, 3.0], σ=[1.0, 2.0])

    # should not error
    assert_valid(params, settings, vars)
end


# infdecomp.jl
module InfDecomp

using InfDecomp_Base
using ExpCone
using GradDesc

using MathProgBase

features_list  = [:Grad,:Jac,:JacVec,:Hess]::Vector{Symbol} # not doing :HessVec right now

# ---------------------------------------------
# M a t h P r o g B a s e    I n t e r f a c e
# ---------------------------------------------

import MathProgBase.initialize
import MathProgBase.features_available
import MathProgBase.eval_f
import MathProgBase.eval_g
import MathProgBase.eval_grad_f
import MathProgBase.jac_structure
import MathProgBase.hesslag_structure
import MathProgBase.eval_jac_g
import MathProgBase.eval_jac_prod
import MathProgBase.eval_jac_prod_t
import MathProgBase.eval_hesslag_prod
import MathProgBase.eval_hesslag
import MathProgBase.isobjlinear
import MathProgBase.isobjquadratic
import MathProgBase.isconstrlinear
# import MathProgBase.obj_expr
# import MathProgBase.constr_expr
import MathProgBase.getreducedcosts
import MathProgBase.getconstrduals
import MathProgBase.getsolution
import MathProgBase.status
import MathProgBase.getsolvetime

# ------------
# B a s i c s
# ------------

# initialize()
function initialize(e::My_Eval, requested_features::Vector{Symbol})
    println("here!")
    for feat in requested_features
        if feat ∉ features_list
            error("infdecomp.jl: initialize():\n-   JuliaOpt:MathProgBase is asking for a feature ($feat) that I don't have.-   Maybe use another solver?")
        end
    end
end

# features_available()
features_available(::My_Eval) = features_list

# Properties:
isobjlinear(::My_Eval)           = false
isobjquadratic(::My_Eval)        = false
isconstrlinear(::My_Eval, ::Integer) = true

#global count_hesseval = 0
# ------------------------------------------
# E v a l u a t i o n :   0 t h   o r d e r
# ------------------------------------------
function eval_f{TFloat_2,TFloat}(e::My_Eval, x::Vector{TFloat_2},dummy::TFloat=Float64(0)) :: TFloat
    local condent::Float64
    if e.TmpFloat==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        condent = condEntropy(e,x,BigFloat(0.))
        setprecision(BigFloat,prc)
    else
        condent = condEntropy(e,x,Float64(0))
    end
    e.count_eval_f += 1
    return -condent
    ;
end #^ eval_f()

# eval_g --- eval of constraint into g
function eval_g(e::My_Eval, g::Vector{Float64}, x::Vector{Float64})  :: Void
    g .= reshape( reshape(x,1,e.n)*e.Gt , e.m, ) .- e.rhs
    e.count_eval_g += 1
    return nothing
    ;
end # eval_g()

# ------------------------------------------
# E v a l u a t i o n :   1 s t   o r d e r
# ------------------------------------------

####################################################################################################
# G L O B A L   V A R I A B L E
global__ill_sol   = "What?!???"
ill_sol__copy_sol = "Oh, no!!!"
function ill_sol__do_copy_sol(x) :: Void
    if global__ill_sol == nothing
        global global__ill_sol = x[:]
    elseif typeof(global__ill_sol)==typeof(x)
        global global__ill_sol .= x
    else
        @show "FUCK!!!!!"
    end
    return nothing
end
function  ill_sol__dont_copy_sol(x) :: Void
    return nothing
end

function set_copy_sol_behaviour(doit::Bool)
    if doit
        global global__ill_sol   = nothing
        global ill_sol__copy_sol = ill_sol__do_copy_sol
    else
        global global__ill_sol   = nothing
        global ill_sol__copy_sol = ill_sol__dont_copy_sol
    end
    ;
end
####################################################################################################

# eval_grad_f --- eval gradient of objective function
function eval_grad_f(e::My_Eval, g::Vector{Float64}, x::Vector{Float64}) :: Void
    if e.TmpFloat==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        ∇f(e,g,x,BigFloat(0.))
        setprecision(BigFloat,prc)
    else
        ∇f(e,g,x,Float64(0))
    end
    e.count_eval_grad_f += 1

    # Crap:
    ill_sol__copy_sol(x)

    return nothing
    ;
end # eval_grad_f()

# Constraint Jacobian
# jac_structure() --- zero-nonzero pattern of constraint Jacobian
jac_structure(e::My_Eval) :: Tuple{Vector{Int64},Vector{Int64}}   =  ( e.Gt_L , e.Gt_K )
# Note: Gt is transposed, so K and L are swapped [K: rows of G^T; L: columns of G^T]

# eval_jac_g() --- constraint Jacobian   -> J
function eval_jac_g(e::My_Eval, J::Vector{Float64}, x::Vector{Float64}) :: Void
    J .= e.Gt.nzval
    e.count_eval_jac_g += 1
    return nothing
    ;
end # eval_jac_g()

# eval_jac_prod() --- constraint_Jacobian * w   -> y
function eval_jac_prod(e::My_Eval, y::Vector{Float64}, x::Vector{Float64}, w::Vector{Float64}) :: Void
    y .= reshape( reshape(w,1,e.n)*e.Gt , e.m, )
    return nothing
    ;
end # eval_jac_prod()

# eval_jac_prod_t() --- constraint_Jacobian^T * w  -> y
function eval_jac_prod_t{T<:AbstractFloat}(e::My_Eval, y::Vector{T}, x::Vector{T}, w::Vector{T}) :: Void
    y .= e.Gt*w
    return nothing
end

# ------------------------------------------
# E v a l u a t i o n :   2 n d   o r d e r
# ------------------------------------------

# Lagrangian:
# L(x, (σ,μ) ) = σ f(x) + μ' G(x)
# Since G(.) is linear, it's Hessian is 0 anyway, so it won't bother us.

# hesslag_structure() --- zero-nonzero pattern of Hessian [wrt x] of the Lagrangian
function hesslag_structure(e::My_Eval)  :: Tuple{Vector{Int64},Vector{Int64}}
    K = Vector{Int64}()
    L = Vector{Int64}()
    counter = 0
    for y = 1:e.n_y
        for z = 1:e.n_z
            # Start with the diagonal
            for x = 1:e.n_x
                i = e.varidx[x,y,z]
                if i>0
                    counter += 1
                    push!(K,i)
                    push!(L,i)
                end
            end
            # Now off-diagonal.
            # H is  treated as a symmetric matrix, but:
            # if both (k,l) & (l,k) are present, their values will be added!
            for x = 1:e.n_x
                for u = 1:(x-1)
                    i_x = e.varidx[x,y,z]
                    i_u = e.varidx[u,y,z]
                    if i_x>0 && i_u>0
                        counter += 1
                        push!(K,i_x)
                        push!(L,i_u)
                    end
                end
            end
        end# for z
    end #^ for y
    return (K,L)
    ;
end # hesslag_structure()

function Hess{TFloat}(e::My_Eval, H::Vector{Float64}, p::Vector{Float64}, σ::Float64, dummy::TFloat) :: Void
    counter = 0
    for y = 1:e.n_y
        for z = 1:e.n_z
            # make marginal P(*yz)
            P_yz  ::TFloat = TFloat(0.)
            for x = 1:e.n_x
                i = e.varidx[x,y,z]
                if i>0
                    P_yz += p[i]
                end
            end

            # now: for all pairs x,u we have:
            # if x ≠  u:   -1/P_yz;
            # if x == u:   ( P_yz - P(xyz) )/(  P_yz * P(xyz) )

            # Start with the diagonal
            for x = 1:e.n_x
                i = e.varidx[x,y,z]
                if i>0
                    counter += 1
                    P_xyz = p[i]
                    H[counter] = Float64(   (P_xyz == 0 ) ?  1.e50  : TFloat(σ)*( P_yz - P_xyz )/(  P_yz * P_xyz )  )
                end
            end
            # Now off-diagonal.
            # H is  treated as a symmetric matrix, but:
            # if both (k,l) & (l,k) are present, their values will be added!
            for x = 1:e.n_x
                for u = 1:(x-1)
                    i_x = e.varidx[x,y,z]
                    i_u = e.varidx[u,y,z]
                    if i_x>0 && i_u>0
                        counter += 1
                        H[counter] = -TFloat(σ)/P_yz
                    end
                end
            end
        end#for z
    end# for y
    return nothing
    ;
end # eval_hesslag()

# eval_hesslag() --- Hessian [wrt x] of the Lagrangian
function eval_hesslag(e::My_Eval, H::Vector{Float64}, x::Vector{Float64}, σ::Float64, μ::Vector{Float64}) :: Void
    if e.TmpFloat==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        Hess(e,H,x,σ,BigFloat(0.))
        setprecision(BigFloat,prc)
        e.count_hesseval += 1
    else
        Hess(e,H,x,σ,Float64(0))
        e.count_hesseval += 1
    end
    return nothing
    ;
end # eval_hesslag()

# eval_hesslag() --- ( Hessian [wrt x] of the Lagrangian ) * v
# function eval_hesslag_prod{T<:AbstractFloat}(e::My_Eval, h::Vector{T}, x::Vector{T}, v::Vector{T}, σ::T, μ::Vector{T}) :: Void
#     h .= σ .* Hf(x)*v
#     return nothing
# end

#------------------
# M A R G I N A L S
#------------------

function marginal_X{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: Array{TFloat_2,1}
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
    end
    marg_x::Array{TFloat_2,1}  = zeros(TFloat_2,e.n_x)
    if TFloat_2==BigFloat
        setprecision(BigFloat,prc)
    end
    for x in 1:e.n_x
        for y in 1:e.n_y
            for z in 1:e.n_z
                i = e.varidx[x,y,z]
                if i>0
                    marg_x[x] += p[i]
                end
            end
        end
    end
    return marg_x
    ;
end#^ marginal_X

function marginal_Y{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: Array{TFloat_2,1}
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
    end
    marg_y::Array{TFloat_2,1}  = zeros(TFloat_2,e.n_y)
    if TFloat_2==BigFloat
        setprecision(BigFloat,prc)
    end
    for y in 1:e.n_y
        for x in 1:e.n_x
            for z in 1:e.n_z
                i = e.varidx[x,y,z]
                if i>0
                    marg_y[y] += p[i]
                end
            end
        end
    end
    return marg_y
    ;
end#^ marginal_Y

function marginal_Z{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: Array{TFloat_2,1}
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
    end
    marg_z::Array{TFloat_2,1}  = zeros(TFloat_2,e.n_z)
    if TFloat_2==BigFloat
        setprecision(BigFloat,prc)
    end
    for z in 1:e.n_z
        for x in 1:e.n_x
            for y in 1:e.n_y
                i = e.varidx[x,y,z]
                if i>0
                    marg_z[z] += p[i]
                end
            end
        end
    end
    return marg_z
    ;
end#^ marginal_Z

function marginal_XY{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: Array{TFloat_2,2}
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        setprecision(BigFloat,prc)
    end
    marg_xy::Array{TFloat_2,2}  = zeros(TFloat_2,e.n_x,e.n_y)
    if TFloat_2==BigFloat
        setprecision(BigFloat,prc)
    end
    for x in 1:e.n_x
        for y in 1:e.n_y
            for z in 1:e.n_z
                i = e.varidx[x,y,z]
                if i>0
                    marg_xy[x,y] += p[i]
                end
            end
        end
    end
    return marg_xy
    ;
end#^ marginals_XY

function marginal_XZ{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: Array{TFloat_2,2}
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
    end
    marg_xz::Array{TFloat_2,2}  = zeros(TFloat_2,e.n_x,e.n_z)
    if TFloat_2==BigFloat
        setprecision(BigFloat,prc)
    end
    for x in 1:e.n_x
        for z in 1:e.n_z
            for y in 1:e.n_y
                i = e.varidx[x,y,z]
                if i>0
                    marg_xz[x,z] += p[i]
                end
            end
        end
    end
    return marg_xz
    ;
end#^ marginal_XZ

function marginal_YZ{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: Array{TFloat_2,2}
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
    end
    marg_yz::Array{TFloat_2,2}  = zeros(TFloat_2,e.n_y,e.n_z)
    if TFloat_2==BigFloat
        setprecision(BigFloat,prc)
    end
    for y in 1:e.n_y
        for z in 1:e.n_z
            for x in 1:e.n_x
                i = e.varidx[x,y,z]
                if i>0
                    marg_yz[y,z] += p[i]
                end
            end
        end
    end
    return marg_yz
    ;
end#^ marginal_YZ

#------------------------------------------
# I N F O R M A T I O N - F U N C T I O N S
#------------------------------------------

# entropy_X --- Shannon entropy of X --- H(X)

function entropy_X{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: TFloat_2
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        s::BigFloat = BigFloat(0.)
        setprecision(BigFloat,prc)
    else
        s = Float64(0)
    end
    marg_x = marginal_X(e, p, TFloat)
    for x in 1:e.n_x
        s  +=  (  ( marg_[x] ≤ 0 )  ?   TFloat(0.)   :   marg_x[x]*log(marg_x[x])  )
    end
    return s
    ;
end#^ entropy_X

# I_X_Y --- Mutual information of  X & Y --- I(X;Y)

function I_X_Y{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: TFloat_2
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        s::BigFloat = BigFloat(0.)
        setprecision(BigFloat,prc)
    else
        s = Float64(0)
    end
    marg_x  = marginal_X(e, p, TFloat)
    marg_y  = marginal_Y(e, p, TFloat)
    marg_xy = marginal_XY(e, p, TFloat)
    for x in 1:e.n_x
        for y in 1:e.n_y
            s  +=  (  (marg_xy[x,y]  ≤ 0 || marg_x[x] ≤ 0 || marg_y[y] ≤ 0)  ?   TFloat_2(0.)   :   marg_xy[x,y]*log( marg_xy[x,y] / (marg_y[y] * marg_x[x]) )  )
        end
    end# ^y
    return s/log(2)
    ;
end#^ I_X_Y

# I_X_Y__Z --- Mutual information of X & Y|Z --- I(X;Y|Z)

function I_X_Y__Z{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: TFloat_2
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        s::BigFloat = BigFloat(0.)
        setprecision(BigFloat,prc)
    else
        s = Float64(0)
    end
    marg_z  = marginal_Z(e, p, TFloat)
    marg_xz = marginal_XZ(e, p, TFloat)
    marg_yz = marginal_YZ(e, p, TFloat)
    for x in 1:e.n_x
        for y in 1:e.n_y
            for z in 1:e.n_z
                i = e.varidx[x,y,z]
                if i>0
                    s  +=  (  (marg_yz[y,z]  ≤ 0 || marg_z[z] ≤ 0 || marg_xz[x,z] ≤ 0 || p[i] ≤ 0 )  ?   TFloat_2(0.)   :   p[i]*log( (p[i] * marg_z[z]) / (marg_xz[x,z] * marg_yz[y,z]) )  )
                end
            end
        end
    end
    return s/log(2)
    ;
end#^ I_X_Y__Z

# I_X_Z__Y --- Mutual information of X & Z|Y --- I(X;Z|Y)

function I_X_Z__Y{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: TFloat_2
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        s::BigFloat = BigFloat(0.)
        setprecision(BigFloat,prc)
    else
        s = Float64(0)
    end
    marg_y  = marginal_Y(e, p, TFloat)
    marg_xy = marginal_XY(e, p, TFloat)
    marg_yz = marginal_YZ(e, p, TFloat)
    for x in 1:e.n_x
        for y in 1:e.n_y
            for z in 1:e.n_z
                i = e.varidx[x,y,z]
                if i>0
                    s  +=  (  (marg_yz[y,z]  ≤ 0 || marg_y[y] ≤ 0 || marg_xy[x,y] ≤ 0 || p[i] ≤ 0 )  ?   TFloat_2(0.)   :   p[i]*log( (p[i] * marg_y[y]) / (marg_xy[x,y] * marg_yz[y,z]) )  )
                end
            end
        end
    end
    return s/log(2)
end#^ I_X_Y__Z

# SI --- Shared Information of Y & Z --- SI(Y;Z)

function SI{TFloat_2,TFloat}(e::My_Eval, p::Vector{TFloat_2}, dummy::TFloat)   :: TFloat_2
    if TFloat_2==BigFloat
        prc = precision(BigFloat)
        setprecision(BigFloat,e.bigfloat_nbits)
        s::BigFloat = BigFloat(0.)
        setprecision(BigFloat,prc)
    else
        s = Float64(0)
    end
    return I_X_Y(e, p, TFloat) - I_X_Y__Z(e, p, TFloat)
    ;
end#^ SI

# ----------------
# U S I N G   I T
# ----------------

numo_vars(e::My_Eval)                   :: Int64            = e.n
numo_constraints(e::My_Eval)            :: Int64            = e.m
constraints_lowerbounds_vec(e::My_Eval) :: Vector{Float64}  = zeros(Float64,e.m)
constraints_upperbounds_vec(e::My_Eval) :: Vector{Float64}  = zeros(Float64,e.m)
vars_lowerbounds_vec(e::My_Eval)        :: Vector{Float64}  = zeros(Float64,e.n)
vars_upperbounds_vec(e::My_Eval)        :: Vector{Float64}  = [Inf  for j=1:e.n]

using Base.Test

function do_it{T1,T2,T3}(pdf::Dict{Tuple{T1,T2,T3},Float64}, solver, tmpFloat::DataType=BigFloat ;  warmstart::Bool=false, model_type=:IPM)
    print("Starting optimization now (",now(),")!\n")

    # global count_hesseval = 0
    const sd,myeval = create_stuff(pdf,tmpFloat)

    stats = nothing

    if model_type ∈ [:ECOS_L,:SCS_L]
        # Conic Programming, large model
        stats = ExpCone.model_L(myeval,solver)
        return sd,myeval,stats

    elseif model_type ∈ [:ECOS_S,:SCS_S]
        # Conic Programming, small model
        stats = ExpCone.model_S(myeval,solver)
        return sd,myeval,stats
    elseif model_type == :SCS_D
        # Conic Programming, Dual problem
        stats = ExpCone.model_D(myeval,solver)
        return sd,myeval,stats

    elseif model_type == :My_GradDesc
        stats = my_gradient_descent(myeval;
                                    max_iter       = 1000,
                                    eps_grad       = 1.e-20,
                                    eps_steplength = 1.e-20,
                                    stepfactor     = .1)
        return sd,myeval,stats

    else # not Conic Program
        const lb = constraints_lowerbounds_vec(myeval)
        const ub = constraints_upperbounds_vec(myeval)
        const l = vars_lowerbounds_vec(myeval)
        const u = vars_upperbounds_vec(myeval)
        const model = MathProgBase.NonlinearModel(solver)
        MathProgBase.loadproblem!(model, myeval.n, myeval.m, l, u, lb, ub, :Min, myeval)

        if warmstart
            MathProgBase.setwarmstart!(model,ones(Float64,myeval.n))
        end

        MathProgBase.optimize!(model)
        # stat = MathProgBase.status(model)

        return sd,myeval,model;
    end

end #^ do_it()

#-----------------------------------------------
# Solution & Statistics
#-----------------------------------------------
type Solution_and_Stats
    instance_name       :: String
    solver              :: Symbol
    var_num             :: Int64
    x_sz                :: Int64
    y_sz                :: Int64
    z_sz                :: Int64
    status              :: Symbol
    obj_val             :: BigFloat
    q_nonneg_viol       :: BigFloat
    q_min_entry         :: BigFloat
    marginals_1         :: BigFloat
    marginals_2         :: BigFloat
    marginals_Inf       :: BigFloat
    mu_nonneg_viol      :: BigFloat
    complementarity_max :: BigFloat
    complementarity_sum :: BigFloat
    # MI_X_YZ             :: BigFloat
    CI                  :: BigFloat
    SI                  :: BigFloat
    UI_Y                :: BigFloat
    UI_Z                :: BigFloat
    num_eval_f          :: Int64
    num_eval_grad_f     :: Int64
    num_eval_g          :: Int64
    num_eval_jac_g      :: Int64
    num_hessevals       :: Int64
    opt_time            :: BigFloat
    # entropy_X           :: BigFloat
end #^ type Solution_and_Stats

function check_feasibility(name, model, myeval, solver) :: Solution_and_Stats
    if solver ∈ [:ECOS_L,:SCS_L,:ECOS_S,:SCS_S,:SCS_D]
        # Conic Program
        stats = model
        model = stats.model

        fstat = Solution_and_Stats(name,solver, 0,0,0,0, status(model) ,  big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),0,0,0,0,0,big(0.))
        fstat.x_sz = myeval.n_x
        fstat.y_sz = myeval.n_y
        fstat.z_sz = myeval.n_z

        fstat.var_num = MathProgBase.numvar(model)
        if status(model) ∈ [:Solve_Succeeded,:Optimal,:NearOptimal,:Suboptimal,:UserLimit,:FeasibleApproximate,:Error]
            # fstat.obj_val = stats.optimum
            fourtuple = information_quantities(myeval,fstat.obj_val,stats.q)
            (fstat.CI, fstat.SI, fstat.UI_Y, fstat.UI_Z) = fourtuple

            fstat.obj_val = eval_f(myeval,stats.q,Float64(0))

            fstat.q_nonneg_viol = max(-minimum(stats.q),0)
            fstat.q_min_entry   = max(minimum(stats.q),0)

            equation = (stats.q'*myeval.Gt)' - myeval.rhs
            fstat.marginals_1   = norm(equation,1)
            fstat.marginals_2   = norm(equation,2)
            fstat.marginals_Inf = norm(equation,Inf)
        end
        fstat.opt_time = stats.time

        return fstat

    elseif solver == :My_GradDesc
        # First order method: Projected Gradient Descent
        stats = model

        fstat = Solution_and_Stats(name,solver, 0,0,0,0, stats.status ,  big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),0,0,0,0,0,big(0.))
        fstat.x_sz = myeval.n_x
        fstat.y_sz = myeval.n_y
        fstat.z_sz = myeval.n_z

        fstat.var_num = myeval.n
        if stats.status ∈ [:grad0 :step0 :iter]
            # fstat.obj_val = stats.optimum
            fourtuple = information_quantities(myeval,fstat.obj_val,stats.q)
            (fstat.CI, fstat.SI, fstat.UI_Y, fstat.UI_Z) = fourtuple

            fstat.obj_val = eval_f(myeval,stats.q,Float64(0))

            fstat.q_nonneg_viol = max(-minimum(stats.q),0)
            fstat.q_min_entry   = max(minimum(stats.q),0)

            equation = (stats.q'*myeval.Gt)' - myeval.rhs
            fstat.marginals_1   = norm(equation,1)
            fstat.marginals_2   = norm(equation,2)
            fstat.marginals_Inf = norm(equation,Inf)
        end
        fstat.opt_time = stats.time

        return fstat

    else
        # Other solvers mainly Interior Point solvers
        fstat = Solution_and_Stats(name,solver, 0,0,0,0, status(model) ,  big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),big(0.),0,0,0,0,0,big(0.))
        fstat.x_sz = myeval.n_x
        fstat.y_sz = myeval.n_y
        fstat.z_sz = myeval.n_z

        fstat.var_num = myeval.n
        if status(model)==:Solve_Succeeded || status(model)==:Optimal || status(model)==:NearOptimal || status(model)==:Suboptimal || status(model)==:KnitroError || ( status(model)==:UserLimit && solver !=:Mosek ) || status(model)==:FeasibleApproximate
            q = Vector{Float64}( getsolution(model) )
            grad = zeros(Float64, myeval.n)
            InfDecomp.∇f(myeval, grad, q, Float64(.0))

            lambda = Vector{Float64}( getconstrduals(model) )

            mu = grad - myeval.Gt*lambda

            fstat.mu_nonneg_viol = max(-minimum(mu),0)

            fstat.complementarity_max = maximum( abs.(mu) .* abs.(q) )
            fstat.complementarity_sum = sum( abs.(mu) .* abs.(q) )
        else
            if global__ill_sol != nothing
                q = Vector{Float64}( global__ill_sol )
            end
            fstat.mu_nonneg_viol = -10
            fstat.complementarity_max = -10
            fstat.complementarity_sum = -10
        end# ^if status

        fstat.obj_val = eval_f(myeval,q,Float64(0))

        fstat.q_nonneg_viol = max(-minimum(q),0)
        fstat.q_min_entry   = max(minimum(q),0)

        equation = (q'*myeval.Gt)' - myeval.rhs
        fstat.marginals_1   = norm(equation,1)
        fstat.marginals_2   = norm(equation,2)
        fstat.marginals_Inf = norm(equation,Inf)
        # fstat.entropy_X   = entropy_X(myeval,q,Float64(0))
        # fstat.MI_X_YZ     = fstat.entropy_X + fstat.obj_val

        fstat.CI, fstat.SI, fstat.UI_Y, fstat.UI_Z =
            information_quantities(myeval,fstat.obj_val,q)

        fstat.num_eval_f        = myeval.count_eval_f
        fstat.num_eval_grad_f   = myeval.count_eval_grad_f
        fstat.num_eval_g        = myeval.count_eval_g
        fstat.num_eval_jac_g    = myeval.count_eval_jac_g
        fstat.num_hessevals = myeval.count_hesseval
        if solver != :Ipopt
            fstat.opt_time = getsolvetime(model)
        end
        #  end #^if status
        return fstat
    end
    ;
end #^ check_feasibility()

function information_quantities(myeval,optimum,q)
    p = Float64[]
    for x = 1:myeval.n_x
        for y = 1:myeval.n_y
            for z = 1:myeval.n_z
                if myeval.marg_xy[x,y] > 0  &&  myeval.marg_xz[x,z] > 0
                    push!(p,myeval.prb_xyz[x,y,z])
                end
            end# z
        end# y
    end# x

    ci = -(eval_f(myeval,p,Float64(0)) + optimum)/log(2)
    si = SI(myeval,q,Float64(0))
    ui_Y = I_X_Y__Z(myeval,q,Float64(0))
    ui_Z = I_X_Z__Y(myeval,q,Float64(0))

    return ci,si, ui_Y, ui_Z
end

end #^ module

; # EOF


module TestPatterns
using Test

using ReTest: and, or, not, interpolated

@testset "patterns: ==" begin
    basics = [and(), or(), not(0), interpolated, 0, r""]
    for a = basics, b = basics
        if a === b
            @test a == b
        else
            @test a != b
            for f in (and, or)
                @test f(a) == f(a)
                @test f(a) != f(b)
                @test f(a, b) == f(a, b)
                @test f(a, b) != f(b, a)
                @test f(and(a, b), or(b, a)) == f(and(a, b), or(b, a))
                @test f(and(a, b), or(a, b)) != f(and(a, b), or(b, a))
                @test f(not(a), and(a, b)) == f(not(a), and(a, b))
                @test f(and(a, b), not(a)) != f(not(a), and(a, b))
                @test f(not(a), r"a", and(a, b)) == f(not(a), r"a", and(a, b))
                @test f(not(a), r"a", and(a, b)) != f(not(a), r"b", and(a, b))
            end
            @test not(a) == not(a)
            @test not(a) != not(b)
            @test not(not(a)) == not(not(a))
        end
    end
end

end # TestPatterns


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Helper structs to dispatch known decision variable construction
struct NoSpecifiedConstraint <: MOI.AbstractSet end
Base.copy(set::NoSpecifiedConstraint) = set

struct DecisionSet <: MOI.AbstractScalarSet
    stage::Int
    constraint::Union{MOI.AbstractSet, JuMP.AbstractVectorSet}
    is_recourse::Bool

    function DecisionSet(stage::Integer; constraint::Union{MOI.AbstractSet, JuMP.AbstractVectorSet} = NoSpecifiedConstraint(), is_recourse::Bool = false)
        return new(stage, constraint, is_recourse)
    end
end

struct KnownSet <: MOI.AbstractScalarSet
    stage::Int
end

struct SingleDecisionSet{T} <: MOI.AbstractScalarSet
    stage::Int
    decision::Decision{T}
    constraint::MOI.AbstractSet
    is_recourse::Bool
end
Base.copy(set::SingleDecisionSet) = reuse(set, set.decision)

set_constraint(set::SingleDecisionSet) = set.constraint

struct MultipleDecisionSet{T} <: MOI.AbstractVectorSet
    stage::Int
    decisions::Vector{Decision{T}}
    constraint::MOI.AbstractSet
    is_recourse::Bool
end
MOI.dimension(set::MultipleDecisionSet) = length(set.decisions)
Base.copy(set::MultipleDecisionSet) = reuse(set, set.decisions)

MOIU.variable_function_type(::Type{<:SingleDecisionSet}) = SingleDecision
MOIU.variable_function_type(::Type{<:MultipleDecisionSet}) = VectorOfDecisions

is_decision_type(::Type{SingleDecisionSet}) = true
is_decision_type(::Type{MultipleDecisionSet}) = false

function JuMP.in_set_string(print_mode, set::SingleDecisionSet)
    decision_str(set) = set.is_recourse ? "RecourseDecisions" : "Decisions"
    if state(set.decision) == Taken
        return string(JuMP._math_symbol(print_mode, :in), " $(decision_str(set))(value = $(set.decision.value))")
    end
    if state(set.decision) == Known
        return string(JuMP._math_symbol(print_mode, :in), " Known(value = $(set.decision.value))")
    end
    if set.constraint isa NoSpecifiedConstraint
        return string(JuMP._math_symbol(print_mode, :in), " $(decision_str(set))")
    else
        return string(JuMP._math_symbol(print_mode, :in), " $(decision_str(set))($(JuMP.in_set_string(print_mode, set.constraint)))")
    end
end

function JuMP.in_set_string(print_mode, set::MultipleDecisionSet)
    decision_str(set) = set.is_recourse ? "RecourseDecisions" : "Decisions"
    if all(d -> state(d) == Taken, set.decisions)
        return string(JuMP._math_symbol(print_mode, :in), " $(decision_str(set))(values = $([d.value for d in set.decisions]))")
    end
    if all(d -> state(d) == Known, set.decisions)
        return string(JuMP._math_symbol(print_mode, :in), " Known(values = $([d.value for d in set.decisions]))")
    end
    if set.constraint isa NoSpecifiedConstraint
        return string(JuMP._math_symbol(print_mode, :in), " $(decision_str(set))")
    else
        return string(JuMP._math_symbol(print_mode, :in), " $(decision_str(set))($(JuMP.in_set_string(print_mode, set.constraint)))")
    end
end

function reuse(set::SingleDecisionSet, decision::Decision)
    return SingleDecisionSet(set.stage, decision, copy(set.constraint), set.is_recourse)
end

function reuse(set::MultipleDecisionSet, decisions::Vector{<:Decision})
    return MultipleDecisionSet(set.stage, decisions, copy(set.constraint), set.is_recourse)
end

function decision_set(variable::JuMP.ScalarVariable, set::DecisionSet)
    return SingleDecisionSet(set.stage, Decision(variable.info, Float64), set.constraint, set.is_recourse)
end

function decision_set(variables::Vector{<:JuMP.ScalarVariable}, set::DecisionSet)
    decisions = map(variables) do variable
        Decision(variable.info, Float64)
    end
    if set.constraint isa JuMP.AbstractVectorSet
        return MultipleDecisionSet(set.stage,
                                   decisions,
                                   JuMP.moi_set(set.constraint, length(variables)),
                                   set.is_recourse)
    else
        return MultipleDecisionSet(set.stage,
                                   decisions,
                                   set.constraint,
                                   set.is_recourse)
    end
end

function decision_set(variables::Matrix{<:JuMP.ScalarVariable}, set::DecisionSet)
    decisions = map(variables) do variable
        Decision(variable.info, Float64)
    end
    if set.constraint isa JuMP.AbstractVectorSet
        return MultipleDecisionSet(set.stage,
                                   decisions,
                                   JuMP.moi_set(set.constraint, length(variables)),
                                   set.is_recourse)
    else
        return MultipleDecisionSet(set.stage,
                                   decisions,
                                   set.constraint,
                                   set.is_recourse)
    end
end

function decision_set(variable::JuMP.ScalarVariable, set::KnownSet)
    return SingleDecisionSet(set.stage, KnownDecision(variable.info, Float64), NoSpecifiedConstraint(), false)
end

function decision_set(variables, set::KnownSet)
    decisions = map(variables) do variable
        KnownDecision(variable.info, Float64)
    end
    return MultipleDecisionSet(set.stage,
                               decisions,
                               NoSpecifiedConstraint(),
                               false)
end


push!(LOAD_PATH,"./src/")

using Conda, PyCall

module sail_route

include("route/domain.jl")
include("weather/load_weather.jl")
include("scenarios/comparison.jl")
include("performance/aerodynamics.jl")
include("performance/polar.jl")

end


export @D, @DD, @GRAD, @pot

# ===========================================================================
#     implement some fun little macros for easier access
#     to the potentials
# ===========================================================================

# REMARK on @pot:
# ---------------
# Julia 0.4 version:
# call(pp::Potential, varargs...) = evaluate(pp, varargs...)
# call(pp::Potential, ::Type{Val{:D}}, varargs...) = evaluate_d(pp, varargs...)
# call(pp::Potential, ::Type{Val{:DD}}, varargs...) = evaluate_dd(pp, varargs...)
# call(pp::Potential, ::Type{Val{:GRAD}}, varargs...) = grad(pp, varargs...)
# unfortunately, in 0.5 `call` doesn't take anabstractargument anymore,
# which means that we need to specify for every potential how to
# create this syntactic sugar. This is what `@pot` is for.

"""
Annotate a type with `@pot` to setup the syntax sugar
for `evaluate, evaluate_d, evaluate_dd, grad`.

## Usage:

For example, the declaration
```julia
"documentation for `LennardJones`"
mutable struct LennardJones <: PairPotential
   r0::Float64
end
@pot LennardJones
```
creates the following aliases:
```julia
lj = LennardJones(1.0)
lj(args...) = evaluate(lj, args...)
@D lj(args...) = evaluate_d(lj, args...)
@DD lj(args...) = evaluate_dd(lj, args...)
@GRAD lj(args...) = grad(lj, args...)
```

Usage of `@pot` is not restricted to pair potentials, but can be applied to
*any* type.
"""
macro pot(fsig)
   @assert fsig isa Symbol
   sym = esc(:x)
   tsym = esc(fsig)
   quote
      @inline ($sym::$tsym)(args...) = evaluate($sym, args...)
      @inline ($sym::$tsym)(::Type{Val{:D}}, args...) = evaluate_d($sym, args...)
      @inline ($sym::$tsym)(::Type{Val{:DD}}, args...) = evaluate_dd($sym, args...)
      @inline ($sym::$tsym)(::Type{Val{:GRAD}}, args...) = grad($sym, args...)
   end
end

# --------------------------------------------------------------------------

# next create macros that translate
"""
`@D`: Use to evaluate the derivative of a potential. E.g., to compute the
Lennard-Jones potential,
```julia
lj = LennardJones()
r = 1.0 + rand(10)
ϕ = lj(r)
ϕ' = @D lj(r)
```
see also `@DD`.
"""
macro D(fsig::Expr)
    @assert fsig.head == :call
    insert!(fsig.args, 2, Val{:D})
    for n = 1:length(fsig.args)
        fsig.args[n] = esc(fsig.args[n])
    end
    return fsig
end

"`@DD` : analogous to `@D`"
macro DD(fsig::Expr)
    @assert fsig.head == :call
    for n = 1:length(fsig.args)
        fsig.args[n] = esc(fsig.args[n])
    end
    insert!(fsig.args, 2, Val{:DD})
    return fsig
end

"`@GRAD` : analogous to `@D`, but escapes to `grad`"
macro GRAD(fsig::Expr)
    @assert fsig.head == :call
    for n = 1:length(fsig.args)
        fsig.args[n] = esc(fsig.args[n])
    end
    insert!(fsig.args, 2, Val{:GRAD})
    return fsig
end


# Based on https://github.com/fchollet/keras/raw/master/keras/datasets/imdb.py
# Also see https://github.com/fchollet/keras/raw/master/examples/imdb_lstm.py
# Also see https://github.com/ilkarman/DeepLearningFrameworks/raw/master/common/utils.py

using Pkg; for p in ("PyCall","JSON","JLD2"); haskey(Pkg.installed(),p) || Pkg.add(p); end
using PyCall,JSON,JLD2,Random
np = pyimport("numpy")

"""

    imdb()

Load the IMDB Movie reviews sentiment classification dataset from
https://keras.io/datasets and return (xtrn,ytrn,xtst,ytst,dict) tuple.

# Keyword Arguments:
- url=https://s3.amazonaws.com/text-datasets: where to download the data (imdb.npz) from.
- dir=Pkg.dir("Knet/data"): where to cache the data.
- maxval=nothing: max number of token values to include. Words are ranked by how often they occur (in the training set) and only the most frequent words are kept. nothing means keep all, equivalent to maxval = vocabSize + pad + stoken.
- maxlen=nothing: truncate sequences after this length. nothing means do not truncate.
- seed=0: random seed for sample shuffling. Use system seed if 0.
- pad=true: whether to pad short sequences (padding is done at the beginning of sequences). pad_token = maxval.
- stoken=true: whether to add a start token to the beginning of each sequence. start_token = maxval - pad.
- oov=true: whether to replace words >= oov_token with oov_token (the alternative is to skip them). oov_token = maxval - pad - stoken.

"""
function imdb(;
              url = "https://s3.amazonaws.com/text-datasets",
              dir = joinpath(@__DIR__, "imdb"),
              data="imdb.npz",
              dict="imdb_word_index.json",
              jld2="imdb.jld2",
              maxval=nothing,
              maxlen=nothing,
              seed=0, oov=true, stoken=true, pad=true
              )
    global _imdb_xtrn,_imdb_ytrn,_imdb_xtst,_imdb_ytst,_imdb_dict
    if !(@isdefined _imdb_xtrn)
        isdir(dir) || mkpath(dir)
        jld2path = joinpath(dir,jld2)
        if !isfile(jld2path)
            @info("Downloading IMDB...")
            datapath = joinpath(dir,data)
            dictpath = joinpath(dir,dict)
            isfile(datapath) || download("$url/$data",datapath)
            isfile(dictpath) || download("$url/$dict",dictpath)
            d = np.load(datapath, allow_pickle=true)
            _imdb_xtrn = map(a->np.asarray(a,dtype=np.int32), get(d, "x_train"))
            _imdb_ytrn = Array{Int8}(get(d, "y_train") .+ 1)
            _imdb_xtst = map(a->np.asarray(a,dtype=np.int32), get(d, "x_test"))
            _imdb_ytst = Array{Int8}(get(d, "y_test") .+ 1)
            _imdb_dict = Dict{String,Int32}(JSON.parsefile(dictpath))
            JLD2.@save jld2path _imdb_xtrn _imdb_ytrn _imdb_xtst _imdb_ytst _imdb_dict
            #rm(datapath)
            #rm(dictpath)
        end
        @info("Loading IMDB...")
        JLD2.@load jld2path _imdb_xtrn _imdb_ytrn _imdb_xtst _imdb_ytst _imdb_dict
    end
    if seed != 0; Random.seed!(seed); end
    xs = [_imdb_xtrn;_imdb_xtst]
    if maxlen == nothing; maxlen = maximum(map(length,xs)); end
    if maxval == nothing; maxval = maximum(map(maximum,xs)) + pad + stoken; end
    if pad; pad_token = maxval; maxval -= 1; end
    if stoken; start_token = maxval; maxval -= 1; end
    if oov; oov_token = maxval; end
    function _imdb_helper(x,y)
        rp = randperm(length(x))
        newy = y[rp]
        newx = similar(x)
        for i in 1:length(x)
            xi = x[rp[i]]
            if oov
                xi = map(w->(w<=oov_token ? w : oov_token), xi)
            else
                xi = filter(w->(w<=oov_token), xi)
            end
            if stoken
                xi = [ start_token; xi ]
            end
            if length(xi) > maxlen
                xi = xi[end-maxlen+1:end]
            end
            if pad && length(xi) < maxlen
                xi = append!(repeat([pad_token], maxlen-length(xi)), xi)
            end
            newx[i] = xi
        end
        newx,newy
    end
    xtrn,ytrn = _imdb_helper(_imdb_xtrn,_imdb_ytrn)
    xtst,ytst = _imdb_helper(_imdb_xtst,_imdb_ytst)
    return xtrn,ytrn,xtst,ytst,_imdb_dict
end


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

struct DensityWithDiff{D<:AbstractDensity,VJP<:DifferentiationAlgorithm} <: AbstractDensity
    vjpalg::VJP
    density::D
end

@inline Base.parent(density::DensityWithDiff) = density.density

vjp_algorithm(density::DensityWithDiff) = density.vjpalg

Base.:(|)(vjpalg::DifferentiationAlgorithm, density::AbstractDensity) = DensityWithDiff(vjpalg, density)
Base.:(|)(vjpalg::DifferentiationAlgorithm, density::DensityWithDiff) = DensityWithDiff(vjpalg, parent(density))
Base.:(|)(vjpalg::DifferentiationAlgorithm, density::PosteriorDensity) = PosteriorDensity(vjpalg | density.likelihood, vjpalg | density.prior)

function Base.show(io::IO, density::DensityWithDiff)
    vjpalg = vjp_algorithm(density)
    Base.show(io, vjpalg)
    print(io, "|")
    Base.show(io, parent(density))
end

var_bounds(density::DensityWithDiff) = var_bounds(density.density)

ValueShapes.varshape(density::DensityWithDiff) = varshape(density.density)

ValueShapes.unshaped(density::DensityWithDiff) = DensityWithDiff(density.vjpalg, unshaped(density.density))

@inline DensityInterface.logdensityof(density::DensityWithDiff, v::Any) = logdensityof(density.density, v)

@inline checked_logdensityof(density::DensityWithDiff, v::Any) = checked_logdensityof(density.density, v)

@inline get_initsrc_from_target(target::DensityWithDiff) = get_initsrc_from_target(target.density)

function bat_transform_impl(target::NoDensityTransform, density::DensityWithDiff, algorithm::DensityIdentityTransform)
    (result = density, trafo = identity)
end

function bat_transform_impl(target::AbstractDensityTransformTarget, density::DensityWithDiff, algorithm::TransformAlgorithm)
    transformed_density, trafo = bat_transform_impl(target, parent(density), algorithm)
    (result = DensityWithDiff(vjp_algorithm(density), transformed_density), trafo = trafo)
end


##################################
### Particles
##################################

"""
    Particles(type::Symbol, properties::Dict{String})

Particles of a type (typically `:stars`, `:dm`, `:gas`, or `:bh`) with their properties in a `Dict`.
The `Dict` has `String` keys, which are set to all uppercase by default, with exception
of `iM` and `Zs`.

The properties are accessible via `particles.prop` (case insensitive).

# Examples
```jldoctest
julia> p = GadgetGalaxies.Particles(:stars, Dict("ID"=>[1, 2], "MASS"=>[1e11, 2e11]));

julia> p.id
2-element Vector{Int64}:
 1
 2

julia> p.pos = [1 2; 3 4; 5 6]
3×2 Matrix{Int64}:
 1  2
 3  4
 5  6

julia> p.properties["POS"] === p.pos
true
```
"""
struct Particles
    type::Symbol
    properties::Dict{String,Any}
end

function Base.getproperty(obj::Particles, sym::Symbol)
    if sym in fieldnames(Particles)
        return getfield(obj, sym)
    else
        return obj.properties[sym |> String |> uppercase |> _particle_property_exceptions]
    end
end

function Base.setproperty!(obj::Particles, sym::Symbol, val)
    if sym in fieldnames(Particles)
        setfield!(obj, sym, val)
    else
        obj.properties[sym |> String |> uppercase |> _particle_property_exceptions] = val
    end
end

Base.getindex(obj::Particles, str::String) = obj.properties[str]
Base.getindex(obj::Particles, sym::Symbol) = Base.getproperty(obj, sym)
Base.setindex!(obj::Particles, val, str::String) = Base.setindex!(obj.properties, val, str)
Base.setindex!(obj::Particles, val, sym::Symbol) = Base.setproperty!(obj, sym, val)
Base.keys(obj::Particles) = keys(obj.properties) .|> lowercase .|> Symbol
Base.haskey(obj::Particles, key) = haskey(obj.properties, key)
Base.haskey(obj::Particles, key::Symbol) = key in keys(obj)
Base.values(obj::Particles) = values(obj.properties)
Base.propertynames(obj::Particles) = [:type; keys(obj)]
Base.copy(obj::Particles) = Particles(obj.type, copy(obj.properties))

function Base.show(io::IO, ::MIME"text/plain", obj::Particles)
    printstyled(io, String(obj.type); bold=true)
    len = 0
    for val in values(obj)
        if typeof(val) <: AbstractArray && size(val, ndims(val)) > len
            len = size(val, ndims(val))
        end
    end
    println(io, ": $len Particles")
    print(io, " ")
    props = join(keys(obj.properties) .|> _particle_property_exceptions |> sort, " ")
    println(io, props)
end

function _particle_property_exceptions(key::String)
    if key == "IM"
        return "iM"
    elseif key == "ZS"
        return "Zs"
    else
        return key
    end
end

const _particle_type_id = Dict(:gas => 0, :dm => 1, :stars => 4, :bh => 5)

"""
    particle_type_id(type::Symbol)

Returns the Gadget particle type from a particle `Symbol`
(currently `:gas`, `:dm`, `:stars`, and `:bh`).
"""
particle_type_id(type::Symbol) = _particle_type_id[type]

##################################
### Galaxy
##################################

abstract type AbstractGalaxy end

"""
    Galaxy(snapshot::Snapshot,
           isub::Integer,
           subid::Union{HaloID,Nothing},
           particles::Dict{Symbol,Particles})

Galaxy of a given `snapshot` ([`Snapshot`](@ref)) with the zero-based subhalo index `isub`
with [`Particles`](@ref).
Including `subid` (containing information on the exact position in the subfind files) makes
accessing properties of the galaxy faster than when only having `isub`.

The particles are accessible via `galaxy.particletype` or `galaxy[:particletype]`
(see [`Particles`](@ref) for typical types).

# Examples
```jldoctest
julia> p = GadgetGalaxies.Particles(:stars, Dict("ID"=>[1, 2], "MASS"=>[1e11, 2e11]));

julia> g = Galaxy(Snapshot("box", 13), 1532, nothing, Dict(:stars=>p));

julia> g.stars === p
true
```
"""
struct Galaxy <: AbstractGalaxy
    snapshot::Snapshot
    isub::Integer
    subid::Union{HaloID,Nothing}
    particles::Dict{Symbol,Particles}
end

"""
    Galaxy(snapshot::Snapshot, isub[, get_id=true])

Convenience method for creating a `Galaxy` with only a `snapshot` and a subfind id.
Will try to read "MSUB" to extract the  if `get_id` is `true`.

[`Particles`](@ref) can still be added after initializing the galaxy by calling
for example `galaxy.stars = [...]`.
"""
function Galaxy(snapshot::Snapshot, isub::Integer, get_id::Bool)
    if get_id
        subbase = string(snapshot.subbase)
        subfile = GadgetIO.select_file(subbase, 0)

        n_files = read_header(subfile).num_files
        _, subid = read_halo_prop_and_id(subbase, isub, "MSUB", n_files; verbose=false)
    else
        subid = nothing
    end

    return Galaxy(snapshot, isub, subid, Dict{Symbol,Particles}())
end

Galaxy(snapshot::Snapshot, isub::Integer) = Galaxy(snapshot, isub, true)

"""
    GalaxyGroup(snapshot::Snapshot,
           igroup::Integer,
           groupid::Union{HaloID,Nothing},
           particles::Dict{Symbol,Particles})

Group of a given `snapshot` ([`Snapshot`](@ref)) with the zero-based group index `igroup`
with [`Particles`](@ref).
Including `groupid` (containing information on the exact position in the subfind files) makes
accessing properties of the galaxy faster than when only having `igroup`.

The particles are accessible via `group.particletype` or `group[:particletype]`
(see [`Particles`](@ref) for typical types).

# Examples
```jldoctest
julia> p = GadgetGalaxies.Particles(:stars, Dict("ID"=>[1, 2], "MASS"=>[1e11, 2e11]));

julia> g = GalaxyGroup(Snapshot("box", 13), 164, nothing, Dict(:stars=>p));

julia> g.stars === p
true
```
"""
struct GalaxyGroup <: AbstractGalaxy
    snapshot::Snapshot
    igroup::Integer
    groupid::Union{HaloID,Nothing}
    particles::Dict{Symbol,Particles}
end

"""
    GalaxyGroup(snapshot::Snapshot, igroup[, get_id=true])

Convenience method for creating a `GalaxyGroup` with only a `snapshot` and a subfind id.
Will try to read "MTOT" to extract the  if `get_id` is `true`.

[`Particles`](@ref) can still be added after initializing the group by calling for example
`group.stars = [...]`.
"""
function GalaxyGroup(snapshot::Snapshot, igroup::Integer, get_id::Bool)
    if get_id
        subbase = string(snapshot.subbase)
        subfile = GadgetIO.select_file(subbase, 0)

        n_files = read_header(subfile).num_files
        _, groupid = read_halo_prop_and_id(subbase, igroup, "MTOT", n_files; verbose=false)
    else
        groupid = nothing
    end

    return GalaxyGroup(snapshot, igroup, groupid, Dict{Symbol,Particles}())
end

GalaxyGroup(snapshot::Snapshot, igroup::Integer) = GalaxyGroup(snapshot, igroup, true)

function Base.getproperty(obj::T, sym::Symbol) where {T<:AbstractGalaxy}
    if sym in fieldnames(T)
        return getfield(obj, sym)
    else
        return obj.particles[sym]
    end
end

function Base.setproperty!(obj::T, sym::Symbol, val) where {T<:AbstractGalaxy}
    if sym in fieldnames(T)
        setfield!(obj, sym, val)
    else
        obj.particles[sym] = val
    end
end

Base.getindex(obj::AbstractGalaxy, sym::Symbol) = Base.getproperty(obj, sym)
Base.setindex!(obj::AbstractGalaxy, val, sym::Symbol) = Base.setproperty!(obj, sym, val)
Base.keys(obj::AbstractGalaxy) = keys(obj.particles) |> collect
Base.haskey(obj::AbstractGalaxy, key) = haskey(obj.particles, key)
Base.values(obj::AbstractGalaxy) = values(obj.particles)
Base.propertynames(obj::Galaxy) = [[:snapshot, :isub, :subid]; keys(obj)]
Base.propertynames(obj::GalaxyGroup) = [[:snapshot, :igroup, :groupid]; keys(obj)]

function Base.copy(obj::T) where {T<:AbstractGalaxy}
    particles = Dict(key => copy(p) for (key, p) in pairs(obj.particles))
    T(obj.snapshot, getind(obj), getid(obj), particles)
end

function Base.show(io::IO, ::MIME"text/plain", obj::AbstractGalaxy)
    printstyled(io, typeof(obj); bold=true)
    try
        z = round(read_redshift(obj.snapshot); digits=2)
        println(io, " at z=$z")
    catch
        println()
    end
    println(io, " $(getindname(obj)) $(getind(obj))")
    for key in sort(keys(obj))
        show(io, "text/plain", obj[key])
    end
end

getindname(g::Galaxy) = "isub"
getindname(g::GalaxyGroup) = "igroup"

getind(g::Galaxy) = g.isub
getind(g::GalaxyGroup) = g.igroup

getid(g::Galaxy) = g.subid
getid(g::GalaxyGroup) = g.groupid

get_subfind_type(g::Galaxy) = 2
get_subfind_type(g::GalaxyGroup) = 1


# Counting
# ========

@generated function _temporal_counts!(counts::ContingencyTable{T,N,A}, weights,
                           seqs::Vararg{AbstractVector{Residue},N}) where {T,N,A}
    quote
        @assert N == length(seqs) "Number of residue arrays and table dimension doesn't match."
        # seq_1 = seqs[1]
        # seq_2 = ...
        @nextract $N seq d -> seqs[d]
        len = length(seq_1)
        # @assert len == length(seq_1) "Residue arrays have different lengths"
        # @assert len == length(seq_2) ...
        @nexprs $N d -> @assert len == length(seq_d) "Residue arrays have different lengths."
        if isa(weights, AbstractArray)
            @assert len == length(weights) "Residue array and weights sizes doesn't match."
        end
        _cleanup_temporal!(counts)
        temporal = counts.temporal
        @inbounds @simd for index in 1:length(seq_1)
            # temporal[Int(seq_1[index]), Int(seq_2... += getweight(weights, index)
            @nref($N, temporal, d -> Int(seq_d[index])) += getweight(weights, index)
        end
        counts
    end
end

"""
It populates a `ContingencyTable` (first argument) using the frequencies in the sequences
(last positional arguments). The dimension of the table must match the number of sequences
and all the sequences must have the same length. You must indicate the used weights and
pseudocounts as second and third positional arguments respectively. You can use
`NoPseudofrequencies()` and `NoClustering()` to avoid the use of sequence weighting and
pseudocounts, respectively.
"""
function count!(table::ContingencyTable{T,N,A},
                weights,
                pseudocounts::Pseudocount,
                seqs::Vararg{AbstractVector{Residue},N}) where {T,N,A}
    _temporal_counts!(table, weights, seqs...)
    apply_pseudocount!(table, pseudocounts)
    _update!(table)
    table
end

function count!(table::Counts{T,N,A}, args...) where {T,N,A}
    count!(getcontingencytable(table), args...)
    table
end

# Default counters
# ================

function _count(alphabet::A,
                weights, pseudocounts,
                seqs::Vararg{AbstractVector{Residue},N}
                )::ContingencyTable{Float64,N,A} where {N,A <: ResidueAlphabet}
    table = ContingencyTable(Float64, Val{N}, alphabet)::ContingencyTable{Float64,N,A}
    count!(table, weights, pseudocounts, seqs...)
    table
end

"""
It returns a `ContingencyTable` wrapped in a `Counts` type with the frequencies of residues
in the sequences that takes as arguments. The dimension of the table is equal to the number
of sequences. You can use the keyword arguments `alphabet`, `weights` and `pseudocounts`
to indicate the alphabet of the table (default to `UngappedAlphabet()`), a clustering
result (default to `NoClustering()`) and the pseudocounts (default to `NoPseudocount()`)
to be used during the estimation of the frequencies.
"""
function Base.count(seqs::Vararg{AbstractVector{Residue},N};
                    alphabet::ResidueAlphabet = UngappedAlphabet(),
                    weights = NoClustering(),
                    pseudocounts::Pseudocount = NoPseudocount()) where N
    Counts(_count(alphabet, weights, pseudocounts, seqs...))
end

# Probabilities
# =============

"""
It populates a `ContingencyTable` (first argument) using the probabilities in the sequences
(last positional arguments). The dimension of the table must match the number of sequences
and all the sequences must have the same length. You must indicate the used weights,
pseudocounts and pseudofrequencies as second, third and fourth positional arguments
respectively. You can use `NoClustering()`, `NoPseudocount()` and `NoPseudofrequencies()`
to avoid the use of sequence weighting, pseudocounts and pseudofrequencies, respectively.
"""
function probabilities!(table::ContingencyTable{T,N,A},
                        weights,
                        pseudocounts::Pseudocount,
                        pseudofrequencies::Pseudofrequencies,
                        seqs::Vararg{AbstractVector{Residue},N}) where {T,N,A}
    count!(table, weights, pseudocounts, seqs...)
    normalize!(table)
    apply_pseudofrequencies!(table, pseudofrequencies)
    table
end

function probabilities!(table::Probabilities{T,N,A}, args...) where {T,N,A}
    probabilities!(getcontingencytable(table), args...)
    table
end

# Default probabilities
# =====================

function _probabilities(alphabet::A,
                        weights, pseudocounts, pseudofrequencies,
                        seqs::Vararg{AbstractVector{Residue},N}
                        )::ContingencyTable{Float64,N,A} where {N,A <: ResidueAlphabet}
    table = ContingencyTable(Float64, Val{N}, alphabet)::ContingencyTable{Float64,N,A}
    probabilities!(table, weights, pseudocounts, pseudofrequencies, seqs...)
    table
end

"""
It returns a `ContingencyTable` wrapped in a `Probabilities` type with the frequencies of
residues in the sequences that takes as arguments. The dimension of the table is equal to
the number of sequences. You can use the keyword arguments `alphabet`, `weights`,
`pseudocounts` and `pseudofrequencies` to indicate the alphabet of the table
(default to `UngappedAlphabet()`), a clustering result (default to `NoClustering()`),
the pseudocounts (default to `NoPseudocount()`) and the pseudofrequencies
(default to `NoPseudofrequencies()`) to be used during the estimation of the probabilities.
"""
function probabilities(seqs::Vararg{AbstractVector{Residue},N};
                       alphabet::ResidueAlphabet = UngappedAlphabet(),
                       weights = NoClustering(),
                       pseudocounts::Pseudocount = NoPseudocount(),
                       pseudofrequencies::Pseudofrequencies = NoPseudofrequencies()) where N
    Probabilities(_probabilities(alphabet,weights,pseudocounts,pseudofrequencies,seqs...))
end


using Resonance

metadata = CSV.read("data/wrangled.csv", DataFrame)
@rsubset!(metadata, !ismissing(:subject),
                    !ismissing(:sample),
                    )
@rsubset!(metadata, !startswith(:sample, "C"),
                    !startswith(:sample, "z"),
                    )

volumes = metadata[:, r"^(left|right|sample)"i][:, 2:end]
dropmissing!(volumes)

spec = metaphlan_profiles(filter(f-> contains(f, "profile"), readdir("/grace/echo/analysis/biobakery3/links/metaphlan", join=true)), :species)

sns = map(samplenames(spec)) do s
    replace(s, r"_S\d+_profile"=>"")
end

idx = unique([findfirst(==(x), sns) for x in sns])

spec = CommunityProfile(abundances(spec)[:, idx], features(spec), map(samplenames(spec)[idx]) do s
    MicrobiomeSample(replace(s, r"_S\d+_profile"=>""))
end)

overlap = sort(samplenames(spec) ∩ volumes.sample)

brainspec = spec[:, overlap]

volumes = @chain volumes begin
    @rsubset(:sample in overlap)
    @orderby(:sample)
end

permutedims(volumes, "sample")

CSV.write("data/halla_volumes.tsv", permutedims(volumes, :sample), delim='\t')
CSV.write("data/halla_species.tsv", brainspec, delim='\t')

##

run(`halla -x data/halla_species.tsv -y data/halla_volumes.tsv --alla -o data/halla_brain/`)

##

metabs = CSV.read("data/metabolites.csv", DataFrame)

##

sample_meta = airtable_metadata()
@rsubset!(sample_meta, in(:sid_old, names(metabs)))

rename!(metabs, Dict(row.sid_old => row.sample for row in eachrow(sample_meta)))
rename!(sample_meta, "sample"=>"etoh_sample")

sample_meta = leftjoin(sample_meta[:, [:subject, :timepoint, :etoh_sample]], metadata, on=[:subject,:timepoint])

@rsubset!(sample_meta, !ismissing(:sample) && in(:sample, samplenames(spec)))
metabspec = let cmp = spec[:, string.(sample_meta.sample)]
    CommunityProfile(abundances(cmp), features(cmp), MicrobiomeSample.(string.(sample_meta.etoh_sample)))
end

CSV.write("data/halla_metabolites.tsv", select(metabs, ["Metabolite", sample_meta.etoh_sample...]), delim='\t')
CSV.write("data/halla_metab_species.tsv", metabspec, delim='\t')

##

run(`halla -x data/halla_metab_species.tsv -y data/halla_metabolites.tsv --alla -o data/halla_metab/`)

##

using CairoMakie

fig = Figure(resolution=(1200,1200))

ax, sc = scatter(fig[1,1], volumes."Right-Cerebral-Cortex", volumes."Left-Cerebral-Cortex")

fig

CSV.write("data/volumes_desc.csv", describe(volumes))


##########################################################################
# Copyright 2017 Samuel Ridler.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##########################################################################

# get travel mode for given time and priority
# currentTime is sim.time, startTime is time that travel starts (>= sim.time)
function getTravelMode!(travel::Travel, priority::Priority, currentTime::Float; startTime::Float = currentTime)
	@assert(currentTime != nullTime)
	@assert(priority != nullPriority)
	@assert(currentTime <= startTime)

	travel.recentSetsStartTimesIndex = getTravelSetsStartTimesIndex(travel, currentTime)

	i = getTravelSetsStartTimesIndex(travel, startTime)
	travelSetIndex = travel.setsTimeOrder[i]
	travelModeIndex = travel.modeLookup[travelSetIndex, Int(priority)]

	return travel.modes[travelModeIndex]
end

# # get travel mode index for given time and priority
# function getTravelModeIndex!(travel::Travel, priority::Priority, startTime::Float)
	# travelMode = getTravelMode!(travel, priority, startTime)
	# return travelMode.index
# end

# get travel setsStartTimes index for given travel start time
function getTravelSetsStartTimesIndex(travel::Travel, startTime::Float)
	# shorthand:
	setsStartTimes = travel.setsStartTimes
	i = travel.recentSetsStartTimesIndex
	n = length(setsStartTimes)

	@assert(setsStartTimes[i] <= startTime) # otherwise, have gone back in time?

	# use linear search to find position in setsStartTimes
	while i < n && setsStartTimes[i+1] <= startTime
		i += 1
	end

	return i
end

# find nearest hospital to call location, given the travel priority
function nearestHospitalToCall!(sim::Simulation, call::Call, priority::Priority)
	travelMode = getTravelMode!(sim.travel, priority, sim.time)
	hospitalIndex = travelMode.fNetTravel.fNodeNearestHospitalIndex[call.nearestNodeIndex]
	return hospitalIndex
end


_nnz(Q::AbstractMatrix) = length(Q)
_nnz(Q::SparseMatrixCSC) = nnz(Q)
function write(io::IO,
    A::AbstractMatrix,       # the constraint matrix
    collb::Vector,  # vector of variable lower bounds
    colub::Vector,  # vector of variable upper bounds
    c::Vector,      # vector containing variable objective coefficients
    rowlb::Vector,  # constraint lower bounds
    rowub::Vector,  # constraint upper bounds
    sense::Symbol,           # model sense
    colcat::Vector{Symbol},  # constraint types
    sos::Vector{SOS},        # SOS information
    Q::AbstractMatrix,      #  Quadratic objectives 0.5 * x' Q x
    modelname::AbstractString="MPSWriter_jl",  # MPS model name
    colnames::Vector{String} = ["V$i" for i in 1:length(c)],
    rownames::Vector{String} = ["C$i" for i in 1:length(rowub)]
)
    for (i, cname) in enumerate(colnames)
        if !verifyname(cname)
            colnames[i] = correctname(cname)
        end
    end
    for (i, rname) in enumerate(rownames)
        if !verifyname(rname)
            rownames[i] = correctname(rname)
        end
    end
    if _nnz(Q) > 0
        error("LP writer does not support Quadratic objective")
    end
    if sense != :Max && sense != :Min
        error("sense must be either :Min or :Max. Currently sense =$(sense).")
    end

    if sense == :Max
        println(io,"Maximize")
    else
        println(io,"Minimize")
    end

    print_objective!(io, c, colnames)
    print_constraints!(io, A, rowlb, rowub, colnames, rownames)
    print_sos!(io, sos, colnames)
    print_bounds!(io, collb, colub, colnames)
    print_category!(io, colcat, colnames)

    println(io, "End")
end

function print_objective!(io, c, colnames)
    print(io, "obj: ")
    sp_c = sparse(c)
    is_first = true
    for (col, val) in zip(sp_c.nzind, sp_c.nzval)
        print_variable_coefficient!(io, val, colnames[col], is_first)
        is_first = false
    end
    println(io, "")
end

function print_variable_coefficient!(io, val, name, is_first)
    if is_first
        print_shortest(io, val)
    else
        print(io, val < 0?" - ":" + ")
        print_shortest(io, abs(val))
    end
    print(io, " $(name)")
end

function getrowsense{T1 <: Real, T2<: Real}(rowlb::Vector{T1}, rowub::Vector{T2})
    @assert length(rowlb) == length(rowub)
    row_sense = Array{Symbol}(length(rowub))
    hasranged = false
    for r=1:length(rowlb)
        @assert rowlb[r] <= rowub[r]
    	if (rowlb[r] == -Inf && rowub[r] != Inf) || (rowlb[r] == typemin(eltype(rowlb)) && rowub[r] != typemax(eltype(rowub)))
    		row_sense[r] = :(<=) # LE constraint
    	elseif (rowlb[r] != -Inf && rowub[r] == Inf)  || (rowlb[r] != typemin(eltype(rowlb)) && rowub[r] == typemax(eltype(rowub)))
    		row_sense[r] = :(>=) # GE constraint
    	elseif rowlb[r] == rowub[r]
    		row_sense[r] = :(==) # Eq constraint
        elseif (rowlb[r] == -Inf && rowub[r] == Inf)
            error("Cannot have a constraint with no bounds")
    	else
            row_sense[r] = :ranged
            hasranged = true
    	end
    end
    row_sense, hasranged
end

function print_constraints!(io, A, rowlb, rowub, colnames, rownames)
    println(io, "Subject To")
    row_sense, hasranged = getrowsense(rowlb, rowub)
    if hasranged
        error("LP Writer does not support ranged constraints")
    end

    sA = sparse(A)' # so we can iterate over rows instead of columns

    cols, vals = rowvals(sA), nonzeros(sA)
    for row in 1:length(rowlb)
        is_first = true
        print(io, "$(rownames[row]): ")
        for j in nzrange(sA, row)
            print_variable_coefficient!(io, vals[j], colnames[cols[j]], is_first)
            is_first = false
        end
        if row_sense[row] == :(>=)
            println(io, " $(row_sense[row]) $(rowlb[row])")
        else
            println(io, " $(row_sense[row]) $(rowub[row])")
        end
    end
end

function print_bounds!(io, collb, colub, colnames)
    println(io, "Bounds")
    for (i, (lb, ub, name)) in enumerate(zip(collb, colub, colnames))
        if lb == -Inf && ub == +Inf
            println(io, "$(name) free")
        else
            if lb == -Inf
                print(io, "-inf")
            else
                print_shortest(io, lb)
            end
             print(io, " <= $(name) <= ")
            if ub == Inf
                print(io, "+inf")
            else
                print_shortest(io, ub)
            end
            println(io)
        end
    end
end

function print_category!(io, colcat, colnames)
    if any(colcat .== :SemiCont) || any(colcat .== :SemiInt)
        error("The LP file writer does not currently support semicontinuous or semi-integer variables")
    end
    println(io, "General")# Integer
    for (cat, name) in zip(colcat, colnames)
        if cat == :Int
            println(io, "$(name)")
        end
    end
    println(io, "Binary")
    for (cat, name) in zip(colcat, colnames)
        if cat == :Bin
            println(io, "$(name)")
        end
    end
end

function print_sos!(io, sosconstraints::Vector{SOS}, colnames)
    for (i, sos) in enumerate(sosconstraints)
        print_sos!(io, string("csos", i), sos, colnames)
    end
end
function print_sos!(io, rowname::String, sos::SOS, colnames)
    # format
    # c5: S1:: col_1:1 y:2 z:3
    print(io, rowname, ": ")
    print(io, "S", string(sos[1]), "::")
    for (index, weight) in zip(sos[2], sos[3])
        print(io, " ", colnames[index], ":")
        print_shortest(io, weight)
    end
    println(io)
end


import Oceananigans.Grids: required_halo_size
using Oceananigans.Utils: prettysummary

"""
    struct ScalarBiharmonicDiffusivity{F, N, K} <: AbstractScalarBiharmonicDiffusivity{F}

Holds viscosity and diffusivities for models with prescribed isotropic diffusivities.
"""
struct ScalarBiharmonicDiffusivity{F, N, K} <: AbstractScalarBiharmonicDiffusivity{F}
    ν :: N
    κ :: K

    function ScalarBiharmonicDiffusivity{F}(ν::N, κ::K) where {F, N, K}
        return new{F, N, K}(ν, κ)
    end
end

# Aliases that allow specify the floating type, assuming that the discretization is Explicit in time
          ScalarBiharmonicDiffusivity(FT::DataType;         kwargs...) = ScalarBiharmonicDiffusivity(ThreeDimensionalFormulation(), FT; kwargs...)
  VerticalScalarBiharmonicDiffusivity(FT::DataType=Float64; kwargs...) = ScalarBiharmonicDiffusivity(VerticalFormulation(), FT; kwargs...)
HorizontalScalarBiharmonicDiffusivity(FT::DataType=Float64; kwargs...) = ScalarBiharmonicDiffusivity(HorizontalFormulation(), FT; kwargs...)

required_halo_size(::ScalarBiharmonicDiffusivity) = 2

"""
    ScalarBiharmonicDiffusivity([formulation=ThreeDimensionalFormulation(), FT=Float64;]
                                ν=0, κ=0,
                                discrete_form = false)

Return a scalar biharmonic diffusivity turbulence closure with viscosity coefficient `ν` and tracer
diffusivities `κ` for each tracer field in `tracers`. If a single `κ` is provided, it is applied to
all tracers. Otherwise `κ` must be a `NamedTuple` with values for every tracer individually.

Arguments
=========

* `formulation`:
  - `HorizontalFormulation()` for diffusivity applied in the horizontal direction(s)
  - `VerticalFormulation()` for diffusivity applied in the vertical direction,
  - `ThreeDimensionalFormulation()` (default) for diffusivity applied isotropically to all directions

* `FT`: the float datatype (default: `Float64`)

Keyword arguments
=================

  - `ν`: Viscosity. `Number`, `AbstractArray`, or `Function(x, y, z, t)`.

  - `κ`: Diffusivity. `Number`, `AbstractArray`, or `Function(x, y, z, t)`, or
         `NamedTuple` of diffusivities with entries for each tracer.

  - `discrete_form`: `Boolean`.
"""
function ScalarBiharmonicDiffusivity(formulation=ThreeDimensionalFormulation(), FT=Float64;
                                     ν=0, κ=0,
                                     discrete_form = false)

    ν = convert_diffusivity(FT, ν; discrete_form)
    κ = convert_diffusivity(FT, κ; discrete_form)
    return ScalarBiharmonicDiffusivity{typeof(formulation)}(ν, κ)
end

function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F}) where {F}
    κ = tracer_diffusivities(tracers, closure.κ)
    return ScalarBiharmonicDiffusivity{F}(closure.ν, κ)
end

@inline viscosity(closure::ScalarBiharmonicDiffusivity, K) = closure.ν
@inline diffusivity(closure::ScalarBiharmonicDiffusivity, K, ::Val{id}) where id = closure.κ[id]

calculate_diffusivities!(diffusivities, closure::ScalarBiharmonicDiffusivity, args...) = nothing

function Base.summary(closure::ScalarBiharmonicDiffusivity)
    F = summary(formulation(closure))
    return string("ScalarBiharmonicDiffusivity{$F}(ν=", prettysummary(closure.ν), ", κ=", prettysummary(closure.κ), ")")
end

Base.show(io::IO, closure::ScalarBiharmonicDiffusivity) = print(io, summary(closure))


#
# oiformat1.jl --
#
# Define 1st revision of OI-FITS format.
#
#------------------------------------------------------------------------------
#
# This file is part of OIFITS.jl which is licensed under the MIT "Expat"
# License:
#
# Copyright (C) 2015: Éric Thiébaut.
#
#------------------------------------------------------------------------------

# OI-FITS FORMAT DESCRIPTION TABLES
#
# The format of the OI-FITS data block is described by a vector of strings
# like:
#
#   ["KEYWORD FORMAT UNITS DESCR",
#     ...,
#     ...,
#    "---------------------------",
#    "COLUMN  FORMAT UNITS DESCR",
#     ...,
#     ...,
#     ...]
#
# where:
#
#   KEYWORD = keyword for HDU header.
#   COLUMN = column name for table (TTYPE).
#   FORMAT = nL where n is an integer and L a letter;
#       for keywords, 0 means optional and 1 means required;
#       for columns, a negative number means abs(n)*NWAVE.
#   UNITS = default units.
#   DESCR = description/comment.
#
# There may be any number of keyword definitions and any number of column
# definitions, the two parts are separated by a dash line like
# "--------------".
#

# OI_TARGET definition (1st revision):
add_def("OI_TARGET", 1,
        ["OI_REVN    1I -      revision number of the table definition",
         "------------------------------------------------------------",
         "TARGET_ID  1I -      index number",
         "TARGET    16A -      target name",
         "RAEP0      1D deg    RA at mean equinox",
         "DECEP0     1D deg    DEC at mean equinox",
         "EQUINOX    1E yr     equinox",
         "RA_ERR     1D deg    error in RA at mean equinox",
         "DEC_ERR    1D deg    error in DEC at mean equino",
         "SYSVEL     1D m/s    systemic radial velocity",
         "VELTYP     8A -      reference for radial velocity",
         "VELDEF     8A -      definition of radial velocity",
         "PMRA       1D deg/yr proper motion in RA",
         "PMDEC      1D deg/yr proper motion in DEC",
         "PMRA_ERR   1D deg/yr error of proper motion in RA",
         "PMDEC_ERR  1D deg/yr error of proper motion in DEC",
         "PARALLAX   1E deg    parallax",
         "PARA_ERR   1E deg    error in parallax",
         "SPECTYP   16A -      spectral type"])

# OI_ARRAY definition (1st revision):
add_def("OI_ARRAY", 1,
        ["OI_REVN    1I - revision number of the table definition",
         "ARRNAME    1A - array name for cross-referencing",
         "FRAME      1A - coordinate frame",
         "ARRAYX     1D m array center X-coordinate",
         "ARRAYY     1D m array center Y-coordinate",
         "ARRAYZ     1D m array center Z-coordinate",
         "------------------------------------------------------------",
         "TEL_NAME  16A - telescope name",
         "STA_NAME  16A - station name",
         "STA_INDEX  1I - station index",
         "DIAMETER   1E m element diameter",
         "STAXYZ     3D m station coordinates relative to array center"])

# OI_WAVELENGTH definition (1st revision):
add_def("OI_WAVELENGTH", 1,
        ["OI_REVN    1I - revision number of the table definition",
         "INSNAME    1A - name of detector for cross-referencing",
         "------------------------------------------------------------",
         "EFF_WAVE   1E m effective wavelength of channel",
         "EFF_BAND   1E m effective bandpass of channel"])

# OI_VIS definition (1st revision):
add_def("OI_VIS", 1,
        ["OI_REVN    1I -   revision number of the table definition",
         "DATE-OBS   1A -   UTC start date of observations",
         "ARRNAME    0A -   name of corresponding array",
         "INSNAME    1A -   name of corresponding detector",
         "------------------------------------------------------------",
         "TARGET_ID  1I -   target number as index into OI_TARGET table",
         "TIME       1D s   UTC time of observation",
         "MJD        1D day modified Julian Day",
         "INT_TIME   1D s   integration time",
         "VISAMP    -1D -   visibility amplitude",
         "VISAMPERR -1D -   error in visibility amplitude",
         "VISPHI    -1D deg visibility phase",
         "VISPHIERR -1D deg error in visibility phase",
         "UCOORD     1D m   U coordinate of the data",
         "VCOORD     1D m   V coordinate of the data",
         "STA_INDEX  2I -   station numbers contributing to the data",
         "FLAG      -1L -   flag"])

# OI_VIS2 definition (1st revision):
add_def("OI_VIS2", 1,
        ["OI_REVN    1I -   revision number of the table definition",
         "DATE-OBS   1A -   UTC start date of observations",
         "ARRNAME    0A -   name of corresponding array",
         "INSNAME    1A -   name of corresponding detector",
         "------------------------------------------------------------",
         "TARGET_ID  1I -   target number as index into OI_TARGET table",
         "TIME       1D s   UTC time of observation",
         "MJD        1D day modified Julian Day",
         "INT_TIME   1D s   integration time",
         "VIS2DATA  -1D -   squared visibility",
         "VIS2ERR   -1D -   error in squared visibility",
         "UCOORD     1D m   U coordinate of the data",
         "VCOORD     1D m   V coordinate of the data",
         "STA_INDEX  2I -   station numbers contributing to the data",
         "FLAG      -1L -   flag"])

# OI_T3 definition (1st revision):
add_def("OI_T3", 1,
        ["OI_REVN    1I -   revision number of the table definition",
         "DATE-OBS   1A -   UTC start date of observations",
         "ARRNAME    0A -   name of corresponding array",
         "INSNAME    1A -   name of corresponding detector",
         "------------------------------------------------------------",
         "TARGET_ID  1I -   target number as index into OI_TARGET table",
         "TIME       1D s   UTC time of observation",
         "MJD        1D day modified Julian Day",
         "INT_TIME   1D s   integration time",
         "T3AMP     -1D -   triple product amplitude",
         "T3AMPERR  -1D -   error in triple product amplitude",
         "T3PHI     -1D deg triple product phase",
         "T3PHIERR  -1D deg error in triple product phase",
         "U1COORD    1D m   U coordinate of baseline AB of the triangle",
         "V1COORD    1D m   V coordinate of baseline AB of the triangle",
         "U2COORD    1D m   U coordinate of baseline BC of the triangle",
         "V2COORD    1D m   V coordinate of baseline BC of the triangle",
         "STA_INDEX  3I -   station numbers contributing to the data",
         "FLAG      -1L -   flag"])

# OI_SPECTRUM definition (1st revision):
add_def("OI_SPECTRUM", 1,
        ["OI_REVN    1I -   revision number of the table definition",
         "DATE-OBS   1A -   UTC start date of observations",
         "INSNAME    1A -   name of corresponding detector",
        #"FOV        1D -   area of sky over which flux is integrated",
         "------------------------------------------------------------",
         "TARGET_ID  1I -   target number as index into OI_TARGET table",
         "MJD        1D day modified Julian Day",
         "INT_TIME   1D s   integration time",
         "FLUXDATA  -1D -   flux",
         "FLUXERR   -1D -   flux error"])


"""
    Cubic() <: Kernel
"""
struct Cubic <: Kernel
    # TODO
end


module TestJultran
using Jultran
using FactCheck

facts("matches values computed from pytran") do
    κ = compute_κs(:co2, [600], 260, 1.e4)[1]
    println(κ)
    @fact abs(κ - 11.2) < 0.01 => true
end

end