"""
    _power(d0::UnivariateDistribution, d1::UnivariateDistribution, alpha::Real, tail::Tail)

Return `power` for which `quantile(d0, alpha) == quantile(d1, beta)` and `beta = 1 - power`.
"""
function _power(d0::UnivariateDistribution, d1::UnivariateDistribution, alpha::Real, tail::Tail)
    right_tail = tail == one_tail ? alpha : alpha / 2
    except_right_tail = 1 - right_tail
    critical_value = quantile(d0, except_right_tail)
    beta = cdf(d1, critical_value)
    return 1 - beta
end

"""
    _alpha(d0::UnivariateDistribution, d1::UnivariateDistribution, power::Real, tail::Tail)

Return `alpha` for which `quantile(d0, alpha) == quantile(d1, beta)` and `beta = 1 - power`.
"""
function _alpha(d0::UnivariateDistribution, d1::UnivariateDistribution, power::Real, tail::Tail)
    beta = 1 - power
    critical_value = quantile(d1, beta)
    except_right_tail = cdf(d0, critical_value)
    right_tail = 1 - except_right_tail
    return tail == one_tail ? right_tail : 2 * right_tail
end

_distribution_parameters(T::IndependentSamplesTTest, n) = n - 2 # n1 + n2 - 2
_distribution_parameters(T::PointBiserialTTest, n) = n - 2
_distribution_parameters(T::TTest, n) = n - 1
function _distribution_parameters(T::DeviationFromZeroMultipleRegression, n)
    return (T.n_predictors, n - T.n_predictors - 1)
end
function _distribution_parameters(T::IncreaseMultipleRegression, n)
    return (T.n_tested_predictors, n - T.n_predictors - 1)
end
_distribution_parameters(T::OneWayANOVA, n) = (T.n_groups - 1, n - T.n_groups)
function _distribution_parameters(T::ConstantVectorHotellingTsqTest, n)
    return (T.n_response_variables, n - T.n_response_variables)
end
function _distribution_parameters(T::TwoVectorsHotellingTsqTest, n::AbstractVector)
    return (T.n_response_variables, sum(n) - T.n_response_variables - 1)
end
_distribution_parameters(T::MultifactorFixedEffectsANOVA, n) = (T.df, n - T.n_groups)
_distribution_parameters(T::ConstantVarianceChisqTest, n) = n - 1
_distribution_parameters(T::ChisqTest, n) = T.df

_noncentrality_parameter(T::IndependentSamplesTTest, es, n) = sqrt(n / 2) * es
_noncentrality_parameter(T::PointBiserialTTest, es, n) = sqrt(es^2 / (1 - es^2)) * sqrt(n)
_noncentrality_parameter(T::TTest, es, n) = sqrt(n) * es
_noncentrality_parameter(T::MultifactorFixedEffectsANOVA, es, n) = es^2 * n
function _noncentrality_parameter(T::TwoVectorsHotellingTsqTest, es, n::AbstractVector)
    return es^2 * ((n[1] * n[2]) / (n[1] + n[2]))
end
_noncentrality_parameter(T::FTest, es, n::AbstractVector) = sum(n) * es^2
_noncentrality_parameter(T::FTest, es, n) = n * es^2
_noncentrality_parameter(T::ChisqTest, es, n) = n * es^2

_distribution(T::TTest) = NoncentralT
_distribution(T::ChisqTest) = NoncentralChisq
_distribution(T::FTest) = NoncentralF

_tail(T::TTest) = T.tail
_tail(T::StatisticalTest) = one_tail

"""
    _null_distribution(T::StatisticalTest, v)

Return the distribution for the null hypothesis of test `T` with distribution parameters `v`.
"""
function _null_distribution(T::StatisticalTest, v)
    D = _distribution(T)
    d0 = D(v..., 0)
    return d0
end

"""
    _alternative_distribution(T::StatisticalTest, v; es, n)

Return the distribution for the null hypothesis of test `T` with distribution parameters `v`.
"""
function _alternative_distribution(T::StatisticalTest, v, es, n)
    D = _distribution(T)
    λ = _noncentrality_parameter(T, es, n)
    d1 = D(v..., λ)
    return d1
end

function _alternative_distribution(T::ConstantVarianceChisqTest, v, es, n)
    D = _distribution(T)
    d1 = D(v..., 0)
    d2 = v + es * d1
    return d2
end

"""
    get_power(T::StatisticalTest; es::Real, alpha::Real, n)

Return the power for some test `T` with effect size `es`, required significance level `alpha` and sample size `n`.
"""
function get_power(T::StatisticalTest; es::Real, alpha::Real, n)
    v = _distribution_parameters(T, n)
    d0 = _null_distribution(T, v)
    d1 = _alternative_distribution(T, v, es, n)
    return _power(d0, d1, alpha, _tail(T))
end

"""
    get_alpha(T::StatisticalTest; es::Real, power::Real, n)

Return the significance level for some test `T` with effect size `es`, power `power` and sample size `n`.
"""
function get_alpha(T::StatisticalTest; es::Real, power::Real, n)
    v = _distribution_parameters(T, n)
    d0 = _null_distribution(T, v)
    d1 = _alternative_distribution(T, v, es, n)
    return _alpha(d0, d1, power, _tail(T))
end

"""
    get_es(T::StatisticalTest; alpha::Real, power::Real, n)

Return the minimum effect size for some test `T` with significance level `alpha`, power `power` and sample size `n`.
"""
function get_es(T::StatisticalTest; alpha::Real, power::Real, n)
    f(es) = get_alpha(T; es, power, n) - alpha
    initial_value = 0.5
    return find_zero(f, initial_value)
end

"""
    get_n(T::StatisticalTest; alpha::Real, power::Real, es::Real)

Return minimum sample size `n` for some test `T` with significance level `alpha`, power `power` and effect size `es`.
"""
function get_n(T::StatisticalTest; alpha::Real, power::Real, es::Real)
    f(n) = get_alpha(T; es, power, n) - alpha
    initial_value = 50
    return find_zero(f, initial_value)
end


export singleton, set, card, fullset

const EntropyIndex = UInt

# Set Manipulation
# function subsets(S)
#   if S == 0
#     return []
#   else
#     cum = [S]
#     for i = 0:(n-1)
#       if (S & (1 << i)) != 0
#         append!(cum, subsets(S $ (1 << i)))
#       end
#     end
#     return cum
#   end
# end
#
# function subsetsones(S)
#   ret = zeroentropy(N)
#   ret[subsets(S)] = 1
#   return ret
# end

function Base.setdiff(S::EntropyIndex, I::EntropyIndex)
    S & (~I)
end

# S ∪ T
function Base.union(S::EntropyIndex, T::EntropyIndex)
    S | T
end

# S ∩ T
function Base.intersect(S::EntropyIndex, T::EntropyIndex)
    S & T
end

# S ⊆ T
function issubset(S::EntropyIndex, T::EntropyIndex)
    Base.setdiff(S, T) == zero(S)
end

function singleton(i::Integer)
    EntropyIndex(1) << (i-1)
end

function fullset(n::Integer)
    (EntropyIndex(1) << n) - EntropyIndex(1)
end

"""
    set(i::Integer)

Return the set of digits of the integer `i`. Note that this cannot create a set
with elements other than 1, 2, 3, 4, 5, 6, 7, 8 and 9. For instance, the set
``{2, 4, 10}`` cannot be constructed with this method, use `set([2, 4, 10])`
instead.

## Examples

To create the set ``\\{2, 4\\}``, use `set(24)` or `set(42)`.
"""
function set(i::Integer)
    ret = emptyset()
    while i > 0
        ret = union(ret, singleton(i % 10))
        i = div(i, 10)
    end
    ret
end

"""
    set(I::AbstractArray{<:Integer})

Return the set of elements of `I`.

## Examples

To create the set ``\\{2, 4\\}``, use `set([2, 4])` or `set([4, 2])`.
"""
function set(I::AbstractArray{<:Integer})
    ret = emptyset()
    for i in I
        ret = union(ret, singleton(i))
    end
    ret
end

function myin(i::Signed, I::EntropyIndex)
    (singleton(i) & I) != 0
end

function card(S::EntropyIndex)
    sum = 0
    while S > 0
        sum += S & 1
        S >>= 1
    end
    Signed(sum)
end

function emptyset()
    EntropyIndex(0)
end

function setsto(J::EntropyIndex)
    EntropyIndex(1):J
end

function mymap(map::Vector, S::EntropyIndex, n)
    T = emptyset()
    for i in 1:n
        if myin(i, S)
            T = T ∪ singleton(map[i])
        end
    end
    T
end


@fastmath function ltc_gj!(du,u,p,t, n_in, n_neurons, w_sens, w_syns, w_gj)
  v = u
  stim, C, G_L, E_L, G_sens, s_sens, h_sens, E_sens, G_syns, s_syns, h_syns, E_syns, G_gj = ltc_gj_unpack_p(n_in,n_neurons,p)

  #I_senss = reshape([G_sens[s,n] * NNlib.σ((stim[s] - h_sens[s,n]) * s_sens[s,n]) * (v[n] - E_sens[s,n]) for s in 1:size(G_sens,1), n in 1:size(G_sens,2)], size(G_sens,1), :)
  #I_synss = reshape([G_syns[s,n] * NNlib.σ((v[s] - h_syns[s,n]) * s_syns[s,n]) * (v[n] - E_syns[s,n]) for s in 1:size(G_syns,1), n in 1:size(G_syns,2)], size(G_syns,1), :)

  #I_senss = Array{eltype(stim)}(undef, n_in, n_neurons)
  #I_synss = Array{eltype(stim)}(undef, n_neurons, n_neurons)
  I_senss = zeros(eltype(stim), n_in, n_neurons)
  I_synss = zeros(eltype(stim), n_neurons, n_neurons)
  I_gjs = zeros(eltype(stim), n_neurons, n_neurons)
  @inbounds @simd for n in 1:n_neurons
    for s in 1:n_in
      I_senss[s,n] = w_sens[s,n] * G_sens[s,n] * mysigm((stim[s] - h_sens[s,n]) * s_sens[s,n]) * (v[n] - E_sens[s,n])
    end
    for s in 1:n_neurons
      I_synss[s,n] = w_syns[s,n] * G_syns[s,n] * mysigm((v[s] - h_syns[s,n]) * s_syns[s,n]) * (v[n] - E_syns[s,n])
      I_gjs[s,n] = w_gj[s,n] * G_gj[s,n] * (v[n] - v[s])
    end
  end
  I_sens = reshape(sum(I_senss, dims=1), :)
  I_syns = reshape(sum(I_synss, dims=1), :)
  I_gj = reshape(sum(I_gjs, dims=1), :)

  # I_sens = zeros(eltype(stim), n_neurons)
  # I_syns = zeros(eltype(stim), n_neurons)
  # @inbounds for n in 1:n_neurons
  #   for s in 1:n_in
  #     #w_sens[s,n] == 0 && continue
  #     I_sens[n] += w_sens[s,n] * G_sens[s,n] * mysigm((stim[s] - h_sens[s,n]) * s_sens[s,n]) * (v[n] - E_sens[s,n])
  #   end
  #   # I_sens[n] += sum(G_sens[:,n] .* mysigm.((stim .- h_sens[:,n]) .* s_sens[:,n]) .* (v[n] .- E_sens[:,n]))
  #   for s in 1:n_neurons
  #     #w_sens[s,n] == 0 && continue
  #     I_syns[n] += w_syns[s,n] * G_syns[s,n] * mysigm((v[s] - h_syns[s,n]) * s_syns[s,n]) * (v[n] - E_syns[s,n])
  #   end
  #   # I_syns[n] += sum(G_syns[:,n] .* mysigm.((v .- h_syns[:,n]) .* s_syns[:,n]) .* (v[n] .- E_syns[:,n]))
  # end

  #du .= (-((G_L .* (v .- E_L)) .+ I_sens .+ I_syns) ./ C) #.* dZ
  @.. du = (-((G_L * (v - E_L)) + I_sens + I_syns + I_gj) / C) #.* dZ

  nothing
end

function ltc_gj_init_u0p(n_in, n_neurons; T=Float32)
  v = rand_uniform(T, 0.01, 0.1, n_neurons) #.± 0.1f0

  stim   = fill(T(0), n_in)
  C      = rand_uniform(T,  1.0,     1.002,   n_neurons)
  G_L    = rand_uniform(T,  0.0001,  0.1,     n_neurons)
  E_L    = rand_uniform(T, -0.3,     0.3,     n_neurons) #.± 0.1f0
  G_sens = rand_uniform(T,  0.0001,  0.1,     n_in, n_neurons)
  s_sens = rand_uniform(T,  3,       8,       n_in, n_neurons)
  h_sens = rand_uniform(T,  0.3,     0.8,     n_in, n_neurons)
  E_sens = rand_uniform(T, -0.9,     0.9,     n_in, n_neurons)
  G_syns = rand_uniform(T,  0.0001,  0.1,     n_neurons, n_neurons)
  s_syns = rand_uniform(T,  3,       8,       n_neurons, n_neurons)
  h_syns = rand_uniform(T,  0.3,     0.8,     n_neurons, n_neurons)
  E_syns = rand_uniform(T, -0.9,     0.9,     n_neurons, n_neurons)

  G_gj   = rand_uniform(T,  0.0001,  0.1,     n_neurons, n_neurons)

  u0 = v
  p = vcat( stim,
    C, G_L, E_L,
    vec(G_sens), vec(s_sens), vec(h_sens), vec(E_sens),
    vec(G_syns), vec(s_syns), vec(h_syns), vec(E_syns),
    vec(G_gj)
  )
  return u0, p
end

function ltc_gj_bounds(n_in, n_neurons; T=Float32)
  lb = vcat(
    [1.0 for _ in 1:n_neurons],               # C
    [0 for _ in 1:n_neurons],                 # G_L
    [-1 for _ in 1:n_neurons],                # E_L
    [0 for _ in 1:n_in*n_neurons],            # G_sens
    [2 for _ in 1:n_in*n_neurons],            # s_sens
    [-0.3 for _ in 1:n_in*n_neurons],          # h_sens
    [-1 for _ in 1:n_in*n_neurons],           # E_sens
    [0 for _ in 1:n_neurons*n_neurons],       # G_syns
    [2 for _ in 1:n_neurons*n_neurons],       # s_syns
    [-0.3 for _ in 1:n_neurons*n_neurons],     # h_syns
    [-1 for _ in 1:n_neurons*n_neurons],      # E_syns
    [0 for _ in 1:n_neurons*n_neurons],       # G_gj

    [-0.4 for _ in 1:n_neurons])             # v

  ub = vcat(
    [2.0 for _ in 1:n_neurons],
    [0.4 for _ in 1:n_neurons],
    [1 for _ in 1:n_neurons],
    [1 for _ in 1:n_in*n_neurons],
    [9 for _ in 1:n_in*n_neurons],
    [0.9 for _ in 1:n_in*n_neurons],
    [1 for _ in 1:n_in*n_neurons],
    [1 for _ in 1:n_neurons*n_neurons],
    [9 for _ in 1:n_neurons*n_neurons],
    [0.9 for _ in 1:n_neurons*n_neurons],
    [1 for _ in 1:n_neurons*n_neurons],
    [1 for _ in 1:n_neurons*n_neurons],       # G_gj

    [0.8 for _ in 1:n_neurons]) # v)

  T.(lb), T.(ub)
end

function ltc_gj_unpack_p(n_in,n_neurons,p)
  stiml = n_in
  Cl = n_neurons
  G_Ll = n_neurons
  E_Ll = n_neurons
  sensl = n_in*n_neurons
  synsl = n_neurons*n_neurons
  G_sensl = sensl
  s_sensl = sensl
  h_sensl = sensl
  E_sensl = sensl
  G_synsl = synsl
  s_synsl = synsl
  h_synsl = synsl
  E_synsl = synsl

  G_gjl = synsl

  s = 1
  stim = @view p[s:s+stiml-1]
  s += stiml
  C = @view p[s:s+Cl-1]
  s += Cl
  G_L = @view p[s:s+G_Ll-1]
  s += G_Ll
  E_L = @view p[s:s+E_Ll-1]
  s += E_Ll
  G_sens = @view p[reshape(s:s+G_sensl-1, n_in, n_neurons)]
  s += G_sensl
  s_sens = @view p[reshape(s:s+s_sensl-1, n_in, n_neurons)]
  s += s_sensl
  h_sens = @view p[reshape(s:s+h_sensl-1, n_in, n_neurons)]
  s += h_sensl
  E_sens = @view p[reshape(s:s+E_sensl-1, n_in, n_neurons)]
  s += E_sensl

  G_syns = @view p[reshape(s:s+G_synsl-1, n_neurons, n_neurons)]
  s += G_synsl
  s_syns = @view p[reshape(s:s+s_synsl-1, n_neurons, n_neurons)]
  s += s_synsl
  h_syns = @view p[reshape(s:s+h_synsl-1, n_neurons, n_neurons)]
  s += h_synsl
  E_syns = @view p[reshape(s:s+E_synsl-1, n_neurons, n_neurons)]
  s += E_synsl

  G_gj = @view p[reshape(s:s+G_gjl-1, n_neurons, n_neurons)]
  s += G_gjl

  return stim, C, G_L, E_L, G_sens, s_sens, h_sens, E_sens, G_syns, s_syns, h_syns, E_syns, G_gj
end

function LTCGJ(wiring, solver, sensealg;
  T=Float32, tspan=T.((0, 1)),# n_sens=n_neurons, n_out=n_neurons,
  #w_sens=ones(T, n_in, n_neurons),
  #w_syns=ones(T, n_neurons, n_neurons),
  #w_gj   = ones(T, n_neurons, n_neurons),
  mtkize=false, gen_jac=false, kwargs...)

  s_in = wiring.s_in
  n_neurons = wiring.n_total

  w_gj = ones(T, n_neurons, n_neurons)
  new_matrices = Dict(
    :w_gj => w_gj
  )
  w_sens = wiring.matrices[:w_sens]
  w_syns = wiring.matrices[:w_syns]
  wiring.matrices = merge(wiring.matrices, new_matrices)

  lb, ub = ltc_gj_bounds(s_in, n_neurons; T)
  u0, p = ltc_gj_init_u0p(s_in, n_neurons; T)
  dudt!(du,u,p,t) = ltc_gj!(du,u,p,t, s_in, n_neurons, w_sens, w_syns, w_gj)

  rnncell = BCTRNNCell(wiring, solver, sensealg, dudt!, u0, tspan, p, lb, ub; mtkize, gen_jac, kwargs...)
  MyRecur(rnncell)
end


func(x,y) = x*y
dfunc(x,y) = ForwardDiff.derivative(x->func(x,y),x)


# We define the Either type simply by the interactions between Identity and Const

# MonoidAlternative
# =================

# Every Const is neutral towards Identity, we choose Const(nothing) here, the default one, following the Option implementation.
TypeClasses.neutral(::Type{<:Identity}) = Option()
# Option()::Const{Nothing}, so this should help for the Option case
TypeClasses.neutral(::Type{Const{Nothing}}) = Option()

# orelse and combine behave identical for Either, both return the Identity element
# - for combine semantics we follow the general approach of """Any semigroup S may be turned into a monoid simply by adjoining an element e not in S and defining e • s = s = s • e for all s ∈ S""" https://en.wikipedia.org/wiki/Monoid)
TypeClasses.combine(a::Identity, b::Const) = a
TypeClasses.combine(a::Const, b::Identity) = b
# - for orelse semantic Identity denotes success
TypeClasses.orelse(a::Identity, b::Const) = a
TypeClasses.orelse(a::Const, b::Identity) = b

# for generic Either, both Const and Identity can be combined on their own, but combining Const with Identity gives Identity.
# Hence, to comply with monoid laws, the neutral element is `Cons(neutral(L))`
TypeClasses.neutral(::Type{Either{L, R}}) where {L, R} = Const(neutral(L))
TypeClasses.neutral(::Type{Either{L, <:UR}}) where {L, UR} = Const(neutral(L))
# we don't provide the fallback of Option if given a generic Either type, as the monoid laws would get broken
# TypeClasses.neutral(::Type{Either}) = Const(nothing)

TypeClasses.neutral(::Type{Option{T}}) where T = Const(nothing)
TypeClasses.neutral(::Type{Option{<:UT}}) where UT = Const(nothing)
TypeClasses.neutral(::Type{Option}) = Const(nothing)

# FunctorApplicativeMonad
# =======================

TypeClasses.ap(f::Const, x::Identity) = f
TypeClasses.ap(f::Identity, x::Const) = x

# you can think of Identity as the neutral element for monadic composition: adding Identity as an additional layer, and then flattening the layers, nothing is changed in total.
TypeClasses.pure(::Type{Either{L, R}}, a) where {L, R} = Identity(a)  # includes Option, as Const{Nothing} => L=Nothing
TypeClasses.pure(::Type{Either{L, <:UR}}, a) where {L, UR} = Identity(a)
TypeClasses.pure(::Type{Either{<:UL, R}}, a) where {UL, R} = Identity(a)  # includes Try as Const{<:Exception} => UL = Exception
TypeClasses.pure(::Type{Either{<:UL, <:UR}}, a) where {UL, UR} = Identity(a)  # includes Either

# we cannot overload this generically, because `Base.map(f, ::Vector...)` would get overwritten as well (even without warning surprisingly)
# hence we do it individually for Either
Base.map(f, a::Either, b::Either, more::Either...) = mapn(f, a, b, more...)

# FlipTypes
# =========

# all covered by Identity and Const


function solve(data)
    volume(v) = prod(max(0, v[i+1] - v[i] + 1) for i=1:2:5)
    intersection(v1, v2) = [(i % 2 == 1 ? max : min)(v1[i], v2[i]) for i=1:6]
    cubes = []
    for (onoff, nums) in data
        function remove_volume(original, remove)
            # Function to remove one cube from another, by replacing it with up to 6 other cubes
            intersecting = intersection(original, remove)
            volume(intersecting) == 0 && return [original]
            original_divided = []
            add(v) = volume(v) != 0 && push!(original_divided, v)
            x, X, y, Y, z, Z = original
            xi, Xi, yi, Yi, zi, Zi = intersecting
            add([x, xi-1, y, Y, z, Z])
            add([Xi+1, X, y, Y, z, Z])
            add([xi, Xi, y, yi-1, z, Z])
            add([xi, Xi, Yi+1, Y, z, Z])
            add([xi, Xi, yi, Yi, z, zi-1])
            add([xi, Xi, yi, Yi, Zi+1, Z])
            return original_divided
        end
        # Remove the current cube from every previous cube (doesn't matter if on or off)
        !isempty(cubes) && (cubes = reduce(vcat, remove_volume(cube, nums) for cube in cubes))
        # Add the cube if it is turned on
        onoff && push!(cubes, nums)
    end
    # Sum all cubes, since none are intersecting after the previous operations
    return cubes .|> volume |> sum
end
function main()
    extract_ints(str) = parse.(Int, m.match for m in eachmatch(r"-?\d+", str))
    data = split.(readlines()) .|> sp -> (sp[1] == "on", extract_ints(sp[2]))
    data_part1 = filter(nums -> all(i ∈ -50:50 for i=nums[2]), data)

    println(solve(data_part1))
    println(solve(data))
end
main()


struct GeneralConstraint{T}
    evaluate::Any
    jacobian::Any
    hessian::Any
    num_variables::Int
    num_parameter::Int
    num_constraint::Int
    num_jacobian::Int
    num_hessian::Int
    jacobian_sparsity::Vector{Vector{Int}}
    hessian_sparsity::Vector{Vector{Int}}
    evaluate_cache::Vector{T}
    jacobian_cache::Vector{T}
    hessian_cache::Vector{T}
    indices_inequality::Vector{Int}
end

function GeneralConstraint(f::Function, num_variables::Int, num_parameter::Int;
    indices_inequality=collect(1:0),
    evaluate_hessian=false)

    #TODO: option to load/save methods
    @variables z[1:num_variables], w[1:num_parameter]
    evaluate = f(z, w)
    jacobian = Symbolics.sparsejacobian(evaluate, z)
    evaluate_func = eval(Symbolics.build_function(evaluate, z, w)[2])
    jacobian_func = eval(Symbolics.build_function(jacobian.nzval, z, w)[2])
    num_constraint = length(evaluate)
    num_jacobian = length(jacobian.nzval)
    jacobian_sparsity = [findnz(jacobian)[1:2]...]
    if evaluate_hessian
        @variables λ[1:num_constraint]
        lag_con = dot(λ, evaluate)
        hessian = Symbolics.sparsehessian(lag_con, z)
        hessian_func = eval(Symbolics.build_function(hessian.nzval, z, w, λ)[2])
        hessian_sparsity = [findnz(hessian)[1:2]...]
        num_hessian = length(hessian.nzval)
    else
        hessian_func = Expr(:null)
        hessian_sparsity = [Int[]]
        num_hessian = 0
    end

    return GeneralConstraint(
        evaluate_func,
        jacobian_func,
        hessian_func,
        num_variables,
        num_parameter,
        num_constraint,
        num_jacobian,
        num_hessian,
        jacobian_sparsity,
        hessian_sparsity,
        zeros(num_constraint),
        zeros(num_jacobian),
        zeros(num_hessian),
        indices_inequality)
end

function GeneralConstraint()
    return GeneralConstraint(
        (constraint, variables, parameters) -> nothing,
        (jacobian,   variables, parameters) -> nothing,
        (hessian,    variables, parameters) -> nothing,
        0, 0, 0, 0, 0,
        [Int[], Int[]],
        [Int[], Int[]],
        Float64[], Float64[], Float64[],
        collect(1:0))
end

function constraints!(violations, indices, general::GeneralConstraint{T}, variables, parameters) where T
    general.evaluate(general.evaluate_cache, variables, parameters)
    @views violations[indices] .= general.evaluate_cache
    fill!(general.evaluate_cache, 0.0) # TODO: confirm this is necessary
end

function jacobian!(jacobians, indices, general::GeneralConstraint{T}, variables, parameters) where T
    general.jacobian(general.jacobian_cache, variables, parameters)
    @views jacobians[indices] .= general.jacobian_cache
    fill!(general.jacobian_cache, 0.0) # TODO: confirm this is necessary
end

function hessian_lagrangian!(hessian, indices, general::GeneralConstraint{T}, variables, parameters, duals) where T
    if !isempty(general.hessian_cache)
        general.hessian(general.hessian_cache, variables, duals)
        @views hessian[indices] .+= general.hessian_cache
        fill!(general.hessian_cache, 0.0) # TODO: confirm this is necessary
    end
end

function sparsity_jacobian(general::GeneralConstraint{T}, num_variables::Int;
    row_shift=0) where T

    row = Int[]
    col = Int[]

    push!(row, (general.jacobian_sparsity[1] .+ row_shift)...)
    push!(col,  general.jacobian_sparsity[2]...)

    return collect(zip(row, col))
end

function sparsity_hessian(general::GeneralConstraint{T}, num_variables::Int) where T
    row = Int[]
    col = Int[]

    if !isempty(general.hessian_sparsity[1])
        shift = 0
        push!(row, (general.hessian_sparsity[1] .+ shift)...)
        push!(col, (general.hessian_sparsity[2] .+ shift)...)
    end

    return collect(zip(row, col))
end

num_constraint(general::GeneralConstraint{T}) where T = general.num_constraint
num_jacobian(general::GeneralConstraint{T}) where T = general.num_jacobian
# num_hessian(general::GeneralConstraint{T}) where T = general.num_hessian

constraint_indices(general::GeneralConstraint{T}; shift=0) where T = collect(shift .+ (1:general.num_constraint))

jacobian_indices(general::GeneralConstraint{T}; shift=0) where T = collect(shift .+ (1:general.num_jacobian))

function hessian_indices(general::GeneralConstraint{T}, key::Vector{Tuple{Int,Int}}, num_variables::Int) where T
    if !isempty(general.hessian_sparsity[1])
        row = Int[]
        col = Int[]
        shift = 0
        push!(row, (general.hessian_sparsity[1] .+ shift)...)
        push!(col, (general.hessian_sparsity[2] .+ shift)...)
        rc = collect(zip(row, col))
        indices = [findfirst(x -> x == i, key) for i in rc]
    else
        indices = Vector{Int}[]
    end
    return indices
end


abstract type CoordinateTransform end

# Based on the pseudo-code on page 156 in Loschchilov's PhD Thesis from 2013.
mutable struct AdaptiveEncoding <: CoordinateTransform
  n::Int
  sqrtn::Float64
  initialized::Bool
  c_p::Float64
  c_mu::Float64
  c_1::Float64
  p::Array{Float64, 2}
  C::Array{Float64, 2}
  B::Array{Float64, 2}
  m::Array{Float64, 2}
  m_prev::Array{Float64, 2}

  AdaptiveEncoding(mu, n) = begin
    c_p = 1.0 / sqrt(n)
    c_mu = c_1 = 0.5 / n
    p = zeros(n, 1)
    C = Matrix{Float64}(I, n, n)
    B = Matrix{Float64}(I, n, n)
    new(n, sqrt(n), false, c_p, c_mu, c_1,
      p, C, B, zeros(n, 1), zeros(n, 1))
  end
end

# Update the coordinate transform given a new set of points (x) and their
# utilities/fitnesses (us). Based on
# http://en.wikipedia.org/wiki/CMA-ES#Example_code_in_MATLAB.2FOctave
function update_transform!(ae::AdaptiveEncoding, x::Array{Float64, 2}, us::Array{Float64,2})
  # Init m on the first call.
  if !ae.initialized
    ae.m = x * us
    ae.initialized = true
    # The other variables were already updated in the constructor...
    return ae.B
  end

  ae.m_prev = ae.m
  ae.m = x * us
  Binv = inv(ae.B)

  diff_m = ae.m - ae.m_prev
  z0 = (ae.sqrtn / norm(Binv * diff_m)) * diff_m

  # Update evolution path.
  ae.p = (1.0 - ae.c_p) * ae.p + sqrt(ae.c_p * (2 - ae.c_p)) * z0

  # Skip the rank-mu update when you need to save time.
  # Only do this randomly, at times, then, to save time?
  xis_minus_m_prev = broadcast(-, xi, m_prev)
  zis = (ae.sqrtn / norm(Binv * xis_minus_m_prev)) * xis_minus_m_prev
  C_mu = zeros(ae.n, ae.n)
  # Should be some simpler way to achieve this... Maybe C_mu = zis * diag(us) * zis'
  for i in 1:ae.n
    z = zis[:,i]
    C_mu += (us[i] * (z * z'))
  end

  # Now update the covar matrix. Skip the c_mu parts if not doing the rank-u update...
  ae.C = (1 - ae.c_1 - ae.c_mu) * ae.C + (ae.c_1 * ae.p * ae.p') + ae.c_mu * C_mu

  # Should we enforce symmetry in C?
  # tc = triu(ae.C)
  # ae.C = tc + tc' - Matrix{Float64}(I, ae.C, ae.C) * diag(ae.C)

  # Now B is the square root of the covar matrix, but faster to do eig decomposition.
  eigvalues, eigvectors = eig(ae.C)
  diagD = diag(eigvectors)
  ae.B = eigvalues .* repmat(sqrt(diagD)', ae.n, 1)

  return ae.B
end


using JLD2, Attitude, MATLAB, SatelliteDynamics, Dierckx
# include ksfunctions.jl
include(joinpath(dirname(dirname(@__FILE__)),"ks_functions.jl"))
@load "rate_control_transfers/30_v3_day_transfer.jld2" X U
dscale = 1e7 # m
tscale = 20000 # s
uscale = 10000.0
dt = 3e-2

t_days,t_hist = get_time_transfer(X,dt,tscale)

r_eci_hist = dscale*mat_from_vec([x_from_u(X[i][1:4]) for i = 1:length(X)])
v_eci_hist = (dscale/tscale)*mat_from_vec([xdot_from_u(X[i][1:4],X[i][5:8]) for i = 1:length(X)])

# function discretize_t(t_hist,n)
#     # new time vec
#     new_t = [0.0]
#     for i = 1:(length(t_hist)-1)
#
#         # time spacing betwen two knot points
#         δt = t_hist[i+1] - t_hist[i]
#
#         # add intermediate points between the knot points
#         for j = 1:n
#             push!(new_t,δt*j/n + t_hist[i])
#         end
#     end
#     return new_t
# end
new_t = discretize_t(t_hist,4)
mat"
$rnew = spline($t_hist,$r_eci_hist,$new_t)/1000;
"
mat"
rnew = spline($t_hist,$r_eci_hist,$new_t)/1000;
rnew = round(rnew,3,'Significant')
figure
hold on
plot3(rnew(1,:),rnew(2,:),rnew(3,:),'k')
[X,Y,Z] = sphere(10)
X = round(X,3,'Significant')
Y = round(Y,3,'Significant')
Z = round(Z,3,'Significant')
Re = 6010
h = surf(Re*X,Re*Y,Re*Z)
axis equal
view(-69,14)
hold off
h.EdgeAlpha = 0
h.FaceColor = 'blue'
view([58 2])
%axis equal
hold off
addpath('/Users/kevintracy/devel/Low-Thrust-TrajOpt/matlab2tikz/src')
%matlab2tikz('full_30_traj.tex')
"
# TODO: Spline this in a data efficient way
mat"
figure
hold on
plot3($r_eci_hist(1,:),$r_eci_hist(2,:),$r_eci_hist(3,:))
axis equal
hold off
"

# ANIMATION STUFF
# output_to_matlab = (new_time_vec = new_t, x_hist_t_correct = rnew)
# file = matopen("matlab_video_data.mat", "w")
# write(file, "output_to_matlab", output_to_matlab)
# close(file)
# mat"""
# [x,y,z] = sphere(20);
# new_time_vec = $new_t;
# x_hist_t_correct = $rnew;
# a_el_1 = [0 90];
# a_el_2 = [-90 0];
#
#
# f1 = figure('units','normalized','outerposition',[0 0 1 1]);
#
# iter = 0;
#
# len = length(1:200:length(x_hist_t_correct));
#
# %v = VideoWriter('newfile.avi');
# %open(v);
#
# [x,y,z] = sphere(20);
#
# $for i = 1:200:length(x_hist_t_correct)
# for i = 1:5
#     i
#     iter = iter+1;
#     figure(f1)
#     %title('delta T' = ',num2str(floor(new_time_vec(i))),' Days'])
#     hold on
#     cla
#     hold on
#
#     plot3(x_hist_t_correct(1,1:i),x_hist_t_correct(2,1:i),x_hist_t_correct(3,1:i))
#     surf(6378*x,6378*y,6378*z,'FaceColor','blue','EdgeColor', 'black')
#     plot3(x_hist_t_correct(1,i),x_hist_t_correct(2,i),x_hist_t_correct(3,i),'o')
#     axis equal
#     axis(dscale*[-6.5 4.5 -4.5 4.5 -1.8 1.8])
#     %view((i/500),10)
#     view(-60,5)
#     title('GTO-GEO Trajectory Optimization')
#     xlabel('ECI I (m)')
#     ylabel('ECI J (m)')
#     zlabel('ECI K (m)')
#     %view(0 - 90*i/(50*len),90 - 90*i/(50*len))
#
#     drawnow
#      %frame= getframe(gcf);
#      %writeVideo(v,frame)
#
#
#
# end
#
# %close(v);
#
# """

Xm = mat_from_vec(X)
U_real = Xm[12:14,:]

mat"
figure
hold on
title('Thrust Vector')
plot($t_hist,$U_real')
ylabel('Thrust Acceleration (0.0001m/s^2)')
xlabel('Time (days)')
hold off
xlim([0,$t_hist(end)])
legend('u_x','u_y','u_z')
%saveas(gcf,'long_traj.png')
"

U_real_vec = [X[i][12:14] for i = 1:length(X)]
Unorm = [norm(U_real_vec[i]) for i = 1:length(X)]

mat"
figure
hold on
plot($t_hist,$Unorm/$uscale)
xlim([$t_hist(4),$t_hist(end)])
xlabel('Time (days)')
ylabel('Thrust Magnitude (m/s^2)')
hold off
"

angle_hist = get_rtn_in_out(U_real,r_eci_hist,v_eci_hist)

angles = mat_from_vec(angle_hist)

mat"
figure
hold on
plot(round($t_hist,4,'Significant'),round(rad2deg($angles'),4,'Significant'),'linewidth',1.2)
xlabel('Time (days)')
xlim([0 $t_hist(end)])
ylabel('Angle (degrees)')
legend('In Plane','Out of Plane','Location','southwest')
hold off
addpath('/Users/kevintracy/devel/Low-Thrust-TrajOpt/matlab2tikz/src')
%matlab2tikz('thrust_angles_30.tex')
"

# picture in picture
mat"
figure
hold on
plot(round($t_hist,4,'Significant'),round(rad2deg($angles'),4,'Significant'),'linewidth',1.2)
xlabel('Time (days)')
xlim([0 $t_hist(end)])
ylabel('Angle (degrees)')
legend('In Plane','Out of Plane','Location','southwest')
a2 = axes()
a2.Position = [0.1200 0.2600 0.2 0.2]; % xlocation, ylocation, xsize, ysize
plot(a2,$t_hist(1:500),$angles(:,1:500)'); axis tight
hold off
addpath('/Users/kevintracy/devel/Low-Thrust-TrajOpt/matlab2tikz/src')
%matlab2tikz('thrust_angles_30.tex')
"

# now let's get angle from velocity vector
angle_from_v = [acos(dot(normalize(v_eci_hist[:,i]),normalize(U_real[:,i]))) for i = 1:size(U_real,2)]

mat"
figure
hold on
plot(rad2deg($angle_from_v))
hold off
"
println("Time in days:  $t_days")

# now let's get orbital elements
const SD = SatelliteDynamics
oe_hist = [SD.sCARTtoOSC([r_eci_hist[:,i];v_eci_hist[:,i]]) for i = 1:size(r_eci_hist,2)]

a_hist = [oe_hist[i][1] for i = 1:length(oe_hist)]
e_hist = [oe_hist[i][2] for i = 1:length(oe_hist)]
i_hist = [rad2deg(oe_hist[i][3]) for i = 1:length(oe_hist)]

mat"
figure
hold on
subplot(3,1,1)
hold on
plot($t_hist,$a_hist/1000)
h = plot($t_hist, 42e3*ones(length($t_hist),1),'r');
legend([h(1)],'GEO SMA','Location','SouthEast')
ylabel('Semi-Major Axis (km)')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
ylim([20e3,46e3])
subplot(3,1,2)
hold on
plot($t_hist,$e_hist)
ylabel('Eccentricity')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
subplot(3,1,3)
hold on
plot($t_hist,$i_hist)
ylabel('Inclination (deg)')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
hold off
"

angle_hist = get_rtn_in_out(U_real,r_eci_hist,v_eci_hist)
angles = mat_from_vec(angle_hist)

thirty = (t_hist = t_hist, a_hist = a_hist, e_hist = e_hist, i_hist = i_hist, Unorm = Unorm,angles = angles)

using JLD2
@save "plotting_data/plotting_data_30.jld2" thirty


"""
An upper arrowhead matrix:

A = [d_1           e_1]
    [    d_2       e_2]
    [       d_3    e_3]
    [           ⋱   ⋮ ]
    [               f ]
"""
struct UpperArrow{T} <: AbstractMatrix{T}
    d::Vector{T}
    e::Vector{T}
    f::T
end

size(A::UpperArrow) = length(A.d) + 1, length(A.d) + 1

function getindex(A::UpperArrow{T}, i::Integer, j::Integer) where T
    n = size(A, 1)
    if i == j
        if i == n
            return A.f
        else
            return A.d[i]
        end
    elseif j == n
        return A.e[i]
    else
        return zero(T)
    end
end

issymmetric(A::UpperArrow) = false

+(A::UpperArrow, B::UpperArrow) = UpperArrow(A.d+B.d, A.e+B.e, A.f+B.f)
-(A::UpperArrow, B::UpperArrow) = UpperArrow(A.d-B.d, A.e-B.e, A.f-B.f)

function chol(A::SymArrow{T}) where T
    a, b, c = A.a, A.b, A.c
    d = sqrt.(a)
    e = b./d
    f = sqrt(c - norm(e)^2)
    UpperArrow(d, e, f)
end

function inv(A::UpperArrow{T}) where T
    d, e, f = A.d, A.e, A.f
    d1 = inv.(d)
    f1 = inv(f)
    e1 = scale!(-f1, d1.*e)
    UpperArrow(d1, e1, f1)
end

function similarity(A::SymArrow{T}, B::SymArrow{T}) where T
    n = size(A, 1)
    R = chol(B)
    Ri = inv(R)
    a, b, c = A.a, A.b, A.c
    d, e, f = Ri.d, Ri.e, Ri.f
    a1 = Vector{T}(n - 1)
    b1 = Vector{T}(n - 1)
    c1 = c*f^2
    twof = 2f
    @inbounds @simd for i = 1:n-1
        a1[i] = a[i]*d[i]^2
        ae = a[i]*e[i]
        b1[i] = d[i]*(ae + b[i]*f)
        c1 += e[i]*(ae + b[i]*twof)
    end
    SymArrow(a1, b1, c1)
end

eigvals(A::SymArrow{T}, B::SymArrow{T}) where T = eigvals(similarity(A, B))


# project.jl

using SPECTrecon: SPECTplan
using SPECTrecon: project!
using MAT
using LinearAlgebra: norm

function cal_nrmse(x, xtrue)
    len = size(xtrue, 3)
    err = zeros(eltype(xtrue), len)
    for l = 1:len
        err[l] = norm(vec(x[:,:,l]) - vec(xtrue[:,:,l])) / norm(vec(xtrue[:,:,l]))
    end
    nrmse = sum(err) / len
    return nrmse
end

function project_error()
    T = Float32
    path = "../data/"
    file = matopen(path*"mumap208.mat")
    mumap = read(file, "mumap208")
    close(file)

    file = matopen(path*"psf_208.mat")
    psfs = read(file, "psf_208")
    close(file)

    file = matopen(path*"xtrue.mat")
    xtrue = convert(Array{Float32, 3}, read(file, "xtrue"))
    close(file)

    file = matopen(path*"proj_jeff_newmumap.mat")
    proj_jeff = read(file, "proj_jeff")
    close(file)

    dy = T(4.7952)
    nview = size(psfs, 4)

    (nx, ny, nz) = size(xtrue)

    for interpmeth in (:one, :two)
        for mode in (:fast, :mem)
            println(string(interpmeth)*", "*string(mode))
            plan = SPECTplan(mumap, psfs, dy; interpmeth, mode)
            views = zeros(T, nx, nz, nview)
            project!(views, xtrue, plan)
            nrmse = cal_nrmse(views, proj_jeff)
            @show nrmse
        end
    end

    # 1d interp: 2e-7 nrmse
    # 2d interp: 4e-3 nrmse
end

# shows accuracy
project_error()
nothing
# plot(nrmse1d)


# Note that we tried and failed to automate the detection of "train" mode looking at the type
# of argument.  The argument is of type Value only during training and only if it is a value
# influenced by model weights.  However people typically apply dropout to the input which is
# not a Value.  So we are going back to no automation, make sure to supply p=0 during testing
# to stop dropout.

# In 1.1.3 trying to automate again, this time using the @diff context: By default drop if
# AutoGrad is recording.

"""
    dropout(x, p; drop, seed)

Given an array `x` and probability `0<=p<=1` return an array `y` in which each element is 0
with probability `p` or `x[i]/(1-p)` with probability `1-p`. Just return `x` if `p==0`, or
`drop=false`. By default `drop=true` in a `@diff` context, `drop=false` otherwise.  Specify a
non-zero `seed::Number` to set the random number seed for reproducible results. See
[(Srivastava et al. 2014)](http://www.jmlr.org/papers/v15/srivastava14a.html) for a reference.

"""
function dropout(x,p; seed=0, drop=training())
    if !drop
        x
    elseif 0 < p < 1
        if seed != 0; Knet.seed!(seed); end
        dropout!(p,x,similar(x))
    elseif p == 0
        x
    elseif p == 1
        zero(x)
    else
        error("Dropout probability not in [0:1]: $p")
    end
end

# Use a more meaningful name than `recording` for Knet.
"`training()` returns `true` only inside a `@diff` context, e.g. during a training iteration of a model."
training() = AutoGrad.recording()

function dropback(dy,y,x,p)
    if 0 < p < 1
        dropback!(p,x,y,dy,similar(x))
    elseif p == 0
        dy
    elseif p == 1
        zero(x)
    else
        error("Dropout probability not in [0:1]: $p")
    end
end

# Turn dropout into an AutoGrad primitive
@primitive dropout(x,p;seed=0,drop=training()),dy,y dropback(value.((dy,y,x,p))...)

# GPU implementation
for S in (32,64)
    T = Symbol("Float$S")
    forw = Symbol("dropout_$S")
    back = Symbol("dropback_$S")
    @eval begin
        function dropout!(p::Number, x::KnetArray{$T}, y::KnetArray{$T})
            rand!(y)
            @knet8($forw,(Cint,$T,Ptr{$T},Ptr{$T}),length(y),$T(p),x,y)
            return y
        end
        function dropback!(p::Number, x::KnetArray{$T}, y::KnetArray{$T}, dy::KnetArray{$T}, dx::KnetArray{$T})
            @knet8($back,(Cint,$T,Ptr{$T},Ptr{$T},Ptr{$T},Ptr{$T}),length(dx),$T(p),x,y,dy,dx)
            return dx
        end
    end
end

# CPU implementation
function dropout!(p,x,y)
    rand!(y)
    p = convert(eltype(y),p)
    q = 1-p
    @inbounds for i=1:length(y)
        if y[i] > p
            y[i] = x[i] / q
        else
            y[i] = 0
        end
    end
    return y
end

function dropback!(p,x,y,dy,dx)
    p = convert(eltype(y),p)
    q = 1-p
    @inbounds for i=1:length(dx)
        if y[i] == 0
            dx[i] = 0
        else
            dx[i] = dy[i] / q
        end
    end
    return dx
end


module AttractorNetworksBase
using LinearAlgebra
using OrdinaryDiffEq
using Distributions, Random
# numerical defaults for the constructors
const default_values = Dict(
    :muw=>3.0,
    :Wmeanf => [ 2.5  -1.3
                 2.4   -1.  ],
    :Wmeanf_alt => [ 1.08  -0.8
                       0.8   -2.25 ],
    :external_current => 7.0,
    :tau_ei => (0.02,0.01),
    :gainalpha => 0.02 )

abstract type IOFunction{R} end
Base.Broadcast.broadcastable(g::IOFunction)=Ref(g)
Base.copy(g::IOFunction) = g

struct IOQuad{T} <: IOFunction{T}
    α::T
end
@inline function (gf::IOQuad{T})(x::T) where T
     return gf.α*(log(1.0+exp(x)))^2
end
@inline function ioprime(x::T,gf::IOQuad{T}) where T
    ex= exp(x)
    return gf.α * 2.0ex*log(1. + ex)/(1. + ex)
end
@inline function ioprimeprime(x::T,gf::IOQuad{T}) where T
    ex= exp(x)
    oex= 1. + ex
    return 2.0*gf.α*ex*(log(oex)+ex) / (oex*oex)
end
@inline function ioinv(y::T,gf::IOQuad{T}) where T
    return log(exp( sqrt(y/gf.α) )-1.0)
end

struct IOId{T} <:IOFunction{T}
end
@inline function (gf::IOId{T})(x::T) where T
     return x
end
@inline function ioprime(x::T,gf::IOId{T}) where T
    return one(T)
end
@inline function ioprimeprime(x::T,gf::IOId{T}) where T
    return zero(T)
end
@inline function ioinv(x::T,gf::IOId{T}) where T
    return x
end

# struct IORelu{T} <:IOFunction{T}
# end
#

"""
        populations(ne::I,ni::I) where I<:Integer -> (pe::BitArray,pi::BitArray)

For `ne` excitatory an `ni` inhibitory neurons, returns the positions of each in the form of BitArrays. The convention is that E neurons come first.  The indexes of the connections in `W` *from* poulation `pa` *to* population `pb` can then be selected as: `view(W,pb,pa)`
"""
function populations(ne::I,ni::I) where I<:Integer
    ntot = ne+ni
    pe = BitArray(falses(ntot))
    pi = copy(pe)
    pe[1:ne] .= true
    pi[ne+1:end] .= true
    return (pe,pi)
end
"""
        populations(ne::I,ni::I, ne_aux::I,ni_aux::I) where I<:Integer
            -> ((pe,pi,pe_pri,pe_aux,pi_pri,pi_aux)

Like `populations(ne,ni)`, but now E and I populations can have a subset of auxiliary neurons, with `ne = ne_pri+ne_aux`. Returns the indexing for each population in the form of BitArrays. The convetion is that excitatory neurons go first, and principal neurons go first.
"""
function populations(ne::I,ni::I, ne_aux::I,ni_aux::I) where I<:Integer
    ne_pri = ne - ne_aux
    ni_pri = ni - ni_aux
    @assert (ne_pri >= 0 ) && (ni_pri >= 0)
    ntot = ne+ni
    pe,pi = populations(ne,ni)
    pe_pri,pe_aux,pi_pri,pi_aux = (BitArray(falses(ntot)) for _ in 1:4)
    pe_pri[1:ne_pri] .= true
    pe_aux[ne_pri+1:ne_pri+ne_aux] .= true
    pi_pri[ne+1:ne+ni_pri] .= true
    pi_aux[ne+ni_pri+1:end] .= true
    return (pe,pi,pe_pri,pe_aux,pi_pri,pi_aux)
end

"""
        diagtozero!(M::AbstractMatrix{T}) where T
Replaces the matrix diagonal with zeros
"""
function diagtozero!(M::AbstractMatrix{T}) where T<:Real
    ms = minimum(size(M))
    for i in 1:ms
        @inbounds M[i,i] = zero(T)
    end
    return M
end

"""
        norm_sum_rows!(mat)
Rescales the matrix by row so that the sum of each row is 1.0
"""
function norm_sum_rows!(mat::AbstractMatrix{T}) where T<:Real
    normf=abs.(inv.(sum(mat,dims=2)))
    broadcast!(*,mat,mat,normf)
    return mat
end

 # two populations means Dale!
function make_wmat(ne::I,ni::I,wmeanf::M ; noautapses=true) where
            {I<: Integer,M <: AbstractMatrix{<:Real}}
    nw = size(wmeanf,1)
    @assert (nw == 2) || (nw==ne+ni)
    if nw == ne+ni #nothing to do
        return wmeanf
    end
    # else, expand the 2D
    Wμ = wmeanf
    d = Exponential(default_values[:muw])
    W=rand(d,ne+ni,ne+ni)
    noautapses && diagtozero!(W)
    # now normalize by block
    pe,pi = populations(ne,ni)
    let wee =  view(W,pe,pe)
        norm_sum_rows!(wee)
        wee .*= Wμ[1,1]
    end
    let wii =  view(W,pi,pi)
        norm_sum_rows!(wii)
        wii .*= Wμ[2,2]
    end
    let wei =  view(W,pe,pi) # inhibitory to excitatory
        norm_sum_rows!(wei)
        wei .*= Wμ[1,2]
    end
    let wie =  view(W,pi,pe) # excitatory to inhibitory
        norm_sum_rows!(wie)
        wie .*= Wμ[2,1]
    end
    return W
end

abstract type RecurrentNetwork{R} end

struct BaseNetwork{R} <: RecurrentNetwork{R}
    weights::Matrix{R}
    iofunction::IOFunction{R}
    membrane_taus::Vector{R}
    external_input::Vector{R}
end
struct AttractorNetwork{R} <: RecurrentNetwork{R}
    weights::Matrix{R}
    iofunction::IOFunction{R}
    attractors_u::Matrix{R}
    membrane_taus::Vector{R}
    external_input::Vector{R}
end

n_neurons(rn::RecurrentNetwork) = size(rn.weights,1)
n_attractors(rn::AttractorNetwork) = size(rn.attractors_u,2)
Base.ndims(rn::RecurrentNetwork)=n_neurons(rn)

#Base.Broadcast.broadcastable(g::RecurrentNetwork)=Ref(g)

Base.copy(ntw::BaseNetwork{R}) where R =
    BaseNetwork{R}( (copy(getfield(ntw,n))
        for n in fieldnames(BaseNetwork) )...)
Base.copy(ntw::AttractorNetwork{R}) where R =
    AttractorNetwork{R}( (copy(getfield(ntw,n))
                for n in fieldnames(AttractorNetwork) )...)

@inline function iofunction!(dest::AbstractVector{R},u::AbstractVector{R},
    nt::RecurrentNetwork{R}) where R<:Real
  for i in eachindex(dest)
      dest[i]=nt.iofunction(u[i])
  end
  return dest
end
@inline function ioprime!(dest::Vector{R},u::AbstractVector{R},
    nt::RecurrentNetwork{R}) where R<:Real
  for i in eachindex(dest)
      dest[i]=ioprime(u[i],nt.iofunction)
  end
  return dest
end
@inline function ioprime(u::AbstractVector{R},nt::RecurrentNetwork{R}) where R<:Real
    return ioprime!(similar(u),u,nt)
end
@inline function ioprimeprime!(dest::AbstractVector{R},
    u::AbstractVector{R},nt::RecurrentNetwork{R}) where R<:Real
  for i in eachindex(dest)
      dest[i]=ioprimeprime(u[i],nt.iofunction)
  end
  return dest
end
@inline function ioprimeprime(u::AbstractVector{R},nt::RecurrentNetwork{R}) where R<:Real
    return ioprimeprime!(similar(u),u,nt)
end

@inline function ioinv!(dest::Vector{R},u::AbstractVector{R},
    nt::RecurrentNetwork{R}) where R<:Real
  for i in eachindex(dest)
      dest[i]=ioinv(u[i],nt.iofunction)
  end
  return dest
end
@inline function ioinv(u::AbstractVector{R},nt::RecurrentNetwork{R}) where R<:Real
    return ioinv!(similar(u),u,nt)
end

function normalize_membrane_taus!(v::Vector{R},ntw::RecurrentNetwork{R}) where R
    for (i,τ) in enumerate(ntw.membrane_taus)
        v[i] /= τ
    end
    return v
end

function BaseNetwork(ne::I,ni::I;
        gfun::Union{Nothing,G}=nothing,
        taus::Union{Nothing,V}=nothing,
        external_current::Union{Nothing,V}=nothing,
        W::Union{Nothing,M}=nothing,
        noautapses = true) where
                {I<:Integer, G<:IOFunction,M<:AbstractMatrix,V<:AbstractVector}

   W = if isnothing(W)
       make_wmat(ne,ni,default_values[:Wmeanf];noautapses=noautapses)
   # if full D , just use it
   elseif size(W,1)==ne+ni
       W
   # if 2D use as mean field
   elseif (size(W,1) == 2 ) && (size(W,2) == 2 )
       make_wmat(ne,ni,W;noautapses=noautapses)
   else
       @warn "something might be wrong with the weight matrix"
       W
   end
   # now the taus
   taus =
       if isnothing(taus)
           τe,τi =default_values[:tau_ei]
           vcat(fill(τe,ne),fill(τi,ni))
       elseif length(taus)==2
           τe,τi =taus
           vcat(fill(τe,ne),fill(τi,ni))
       else
           taus
    end
    # currents
    h  = something(external_current,
            fill(default_values[:external_current],ne+ni) )
    # gain
    gfun = something(gfun, IOQuad(default_values[:gainalpha]))
   return BaseNetwork(W,gfun,taus,h)
end
function AttractorNetwork(bn::BaseNetwork{R},attr::Matrix{R}) where R
    return AttractorNetwork{R}(bn.weights,bn.iofunction,attr,
        bn.membrane_taus,bn.external_input)
end

function velocity!(dest::AbstractVector{R},u::AbstractVector{R},
        gu::AbstractVector{R},rn::RecurrentNetwork{R}) where R<:Real
    copy!(dest, rn.external_input)
    dest .-= u  #  v_out <-  - u +  h
    LinearAlgebra.BLAS.gemv!('N',1.0,rn.weights,gu,1.0,dest) # W*g(u) - u + h
    normalize_membrane_taus!(dest,rn) # ( W*g(u) - u + h) / taus
    return dest
end
velocity(u,rn) = velocity!(similar(u),u,rn.iofunction.(u), rn)

# Jacobian stuff ! Compute and derivatives!
# let's define (and test) the derivatives of the jacobian here!

struct JAlloc{R}
    dgu::Vector{R}
    ddgu::Vector{R}
    inv_taus::Vector{R}
end
function JAlloc(ntw::RN) where {R<:Real,RN<:RecurrentNetwork{R}}
    dgu = similar(ntw.membrane_taus)
    ddgu = similar(ntw.membrane_taus)
    JAlloc{R}(dgu,ddgu,inv.(ntw.membrane_taus))
end

function jacobian!(Jdest::AbstractMatrix{R},
        Wgrad::Union{Nothing,Matrix{R}},ugrad::Union{Nothing,Matrix{R}},
        u::AbstractVector{R},ntw::RecurrentNetwork{R},
        Jallocs::Union{Nothing,JAlloc{R}}=nothing) where R
  Jallocs=something(Jallocs,JAlloc(ntw))
  ioprime!(Jallocs.dgu,u,ntw)
  broadcast!(*,Jdest, ntw.weights, transpose(Jallocs.dgu)) # multiply columnwise
  for i in 1:n_neurons(ntw)
    Jdest[i,i] -= 1.0 #subtract diagonal
  end
  #normalize by taus, rowwise
  broadcast!(/,Jdest,Jdest,ntw.membrane_taus)
  # GRADIENTS ! with respect to  W first
  if !isnothing(Wgrad)
    broadcast!(*,Wgrad,Jallocs.inv_taus,transpose(Jallocs.dgu))
  end
  # now with respect to u
  if !isnothing(ugrad)
    ioprimeprime!(Jallocs.ddgu,u,ntw)
    #broadcast!(/,Jallocs.weights,transpose(dgu),ntw.membrane_taus)
    broadcast!(*,ugrad,Jallocs.inv_taus,
      ntw.weights, transpose(Jallocs.ddgu))
  end
  return Jdest
end

"""
        jacobian(u,rn::RecurrentNetwork)
Jacobian matrix of the
"""
@inline function jacobian(u::AbstractVector{R},rn::RecurrentNetwork{R}) where R
    return jacobian!(similar(rn.weights),nothing,nothing,u,rn)
end

function spectral_abscissa(u,rn::RecurrentNetwork)
    J=jacobian(u,rn)
    return maximum(real.(eigvals(J)))
end

# Dynamics
function run_network(ntw::RecurrentNetwork{R},r_start::Vector{R},
    t_end::Real; verbose::Bool=false,stepsize=0.05) where R<:Real
  u0=ioinv(r_start,ntw)
  iofun_alloc=similar(r_start)
  f(du,u,p,t) = velocity!(du,u,iofunction!(iofun_alloc,u,ntw),ntw)
  prob = ODEProblem(f,u0,(0.,t_end))
  solv = solve(prob,Tsit5();verbose=verbose,saveat=stepsize)
  ret_u = hcat(solv.u...)
  ret_r = ntw.iofunction.(ret_u)
  return solv.t,ret_u,ret_r
end

function run_network_noise(ntw::RecurrentNetwork{R},
    r_start::Vector{R},noiselevel::Real,t_end::Real;
    verbose::Bool=false,stepsize=0.05) where R<:Real
  u0=ioinv(r_start,ntw)
  iofun_alloc=similar(r_start)
  f(du,u,p,t) = velocity!(du,u,iofunction!(iofun_alloc,u,ntw),ntw)
  σ_f(du,u,p,t) = fill!(du,noiselevel)
  prob = SDEProblem(f,σ_f,u0,(0.,t_end))
  solv =  solve(prob,EM();verbose=verbose,dt=stepsize)
  ret_u = hcat(solv.u...)
  ret_r = ntw.iofunction.(ret_u)
  return solv.t,ret_u,ret_r
end

"""
        run_network_to_convergence(rn::RecurrentNetwork,r_start::Vector{<:Real} ;
                t_end=80. , veltol=1E-1)
Runs the network as described in [`run_network_nonoise`](@ref), but stops as soon as
`norm(v) / n < veltol` where `v` is the velocity at time `t`.
If this condition is not satisfied (no convergence to attractor), it runs until `t_end` and prints a warning.
# Arguments
- `rn::RecurrentNetwork`
- `r_start::Vector` : initial conditions
- `t_end::Real`: the maximum time considered
- `veltol::Real` : the norm (divided by num. dimensions) for velocity at convergence
# Outputs
- `u_end`::Vector : the final state at convergence
- `r_end`::Vector : the final state at convergence as rate
"""
function run_network_to_convergence(ntw::RecurrentNetwork{R},
     r_start::Vector{R};t_end::Real=50. , veltol::Real=1E-1) where R
    function  condition(u,t,integrator)
        v = get_du(integrator)
        return norm(v) / length(v) < veltol
    end
    function affect!(integrator)
        savevalues!(integrator)
        return terminate!(integrator)
    end
    u0=ioinv(r_start,ntw)
    cb=DiscreteCallback(condition,affect!)
    iofun_alloc=similar(r_start)
    f(du,u,p,t) = velocity!(du,u,iofunction!(iofun_alloc,u,ntw),ntw)
    prob = ODEProblem(f,u0,(0.,t_end))
    out = solve(prob,Tsit5();verbose=false,callback=cb)
    u_out = out.u[end]
    t_out = out.t[end]
    if isapprox(t_out,t_end; atol=0.05)
        vel = velocity(u_out,ntw)
        @warn "no convergence after max time $t_end"
        @info "the norm (divided by n) of the velocity is $(norm(vel)/length(vel)) "
    end
    ret_r = ntw.iofunction.(u_out)
    return u_out,ret_r
end

# here I make the AttractorNetwork, the one with attractors

"""
        lognorm_reparametrize(m,std) -> distr::LogNormal
# parameters
  + `m`   sample mean
  + `std` sample std
"""
function lognorm_reparametrize(m,std)
    vm2= (std/m)^2
    μ = log(m / sqrt(1. + vm2))
    σ = sqrt(log( 1. + vm2))
    return LogNormal(μ,σ)
end

function make_attractors(ntot,natt,gainf::IOFunction{R};
        mu_r=5.0,std_r=2.0) where R
    distr = lognorm_reparametrize(mu_r,std_r)
    attr_r = rand(distr,(ntot,natt))
    attr_u = ioinv.(attr_r,gainf)
    return (attr_r,attr_u)
end
make_attractors(ntot,natt,ntw::RecurrentNetwork;mu_r=5.0,std_r=2.0) =
    make_attractors(ntot,natt,ntw.iofunction;mu_r=mu_r,std_r=std_r)

function AttractorNetwork(natt::I,ne::I,ni::I;
        gfun::Union{Nothing,G}=nothing,
        taus::Union{Nothing,V}=nothing,
        external_current::Union{Nothing,V}=nothing,
        W::Union{Nothing,M}=nothing,
        noautapses::Bool = true ,
        mu_attr::Real=5.0,
        std_attr::Real=2.0) where
                {I<:Integer, G<:IOFunction,M<:AbstractMatrix,V<:AbstractVector}
  base=BaseNetwork(ne,ni;gfun=gfun,taus=taus,external_current=external_current,
    W=W,noautapses=noautapses)
  _,attr_u = make_attractors(ne+ni,natt,base;mu_r=mu_attr,std_r=std_attr)
  return AttractorNetwork(base,attr_u)
end

end # module


# Factorize two different matrices (one real, one complex) and solve for multiple rhs and free memory
using Distributed
if nworkers()<3
    addprocs(3)
end
@everywhere using MUMPSjInv
using Test

include("getDivGrad.jl");
A = getDivGrad(24,20,22);
A2 = getDivGrad(30,31,33);
n = size(A,1);
n2 = size(A2,1);
A  = A + im*spdiagm(0 => rand(n));
nrhs = 10;
rhs = randn(n,nrhs) + im*randn(n,nrhs);
rhs2 = randn(n2,nrhs);

@test nworkers() >= 3
println("Factorize complex matrix on worker 1")
@time F1 = remotecall_fetch(factorMUMPS,workers()[1], A, 1)
dump(F1)

println("Factorize real matrix on worker 2")
@time F2 = remotecall_fetch(factorMUMPS, workers()[2], A2, 1)
dump(F2)

println("Solve complex system on worker 1")
@time x = remotecall_fetch(applyMUMPS, workers()[3], F1, rhs)
err = zeros(nrhs)
for i=1:nrhs
    err[i] =  norm(A * x[:, i]-rhs[:, i]) / norm(rhs[:, i])
end
@test maximum(err) < 1e-14

println("Solve real system on worker 2")
@time x2 = remotecall_fetch(applyMUMPS, workers()[2], F2, rhs2)
err = zeros(nrhs)
for i = 1:nrhs
    err[i] = norm(A2 * x2[:, i] - rhs2[:, i]) / norm(rhs2[:, i])
end
@test maximum(err) < 1e-14

println("Free memory on worker 1")
destroyMUMPS(F1)

println("Free memory on worker 2")
destroyMUMPS(F2)


# parameter settings: \nu, \lambda, \sigma, \sigma_e)=(0.001, 0.02, 0.1, 0.001,
# parameter settings: sigma_epsilon = 0.001
# start value: x_0 = 5

using Pkg
using DataFrames
import Statistics.mean
import Statistics.var
using LinearAlgebra
using CSV
using PyPlot

# load functions
include(pwd()*"/src/SDEMEM OU neuron data/ou_sdemem.jl")
include(pwd()*"/src/SDEMEM OU neuron data/mcmc.jl")

data = CSV.read("data/SDEMEM OU neuron data/OUsim.dat"; header=0)

y = data[2:end,1]

dt = 1

time_grid = 1:dt:length(y)

PyPlot.figure()
PyPlot.plot(time_grid, y)

M_subjects = 1

ϕ = zeros(M_subjects,3) # lambda nu sigma

ϕ[1,:] = log.([0.001, 0.02, 0.1])

log_σ_ϵ = log(0.001)

nbr_particles = 1*ones(Int64, M_subjects)
ρ = 0.0

# set number of obs for each data set
T_vec = zeros(Int64, M_subjects)
for i = 1:M_subjects; T_vec[i] = length(y); end

u_prop_old = [randn(nbr_particles[1], T_vec[1]+1)]
u_resample_old = [randn(T_vec[1], 2)]

for i in 2:M_subjects;
    append!(u_prop_old, [randn(nbr_particles[i],T_vec[i]+1)])
    append!(u_resample_old, [randn(T_vec[i],2)])
end

m = 1

dt
# Plot paths

loglik_est, x_hat_kalman = @time kalman_filter_test(y, exp(log_σ_ϵ),  ϕ[m,:], dt, true)

PyPlot.figure()
PyPlot.plot(time_grid, y, "g")
PyPlot.plot(time_grid, x_hat_kalman, "r--")

nbr_particles_bootstrap = 10000*ones(Int64, M_subjects)

# set number of obs for each data set
T_vec = zeros(Int64, M_subjects)
for i = 1:M_subjects; T_vec[i] = length(y); end

u_prop_old_bootstrap = [randn(nbr_particles_bootstrap[1], T_vec[1]+1)]
u_resample_old_bootstrap = [randn(T_vec[1], 2)]

for i in 2:M_subjects;
    append!(u_prop_old_bootstrap, [randn(nbr_particles_bootstrap[i],T_vec[i]+1)])
    append!(u_resample_old_bootstrap, [randn(T_vec[i],2)])
end

loglik, x,nbr_resamlping = @time particlefilter_test(y, exp(log_σ_ϵ), ϕ[m,:], dt, u_prop_old_bootstrap[m], u_resample_old_bootstrap[m], nbr_particles_bootstrap[m], true, true)

PyPlot.figure()
PyPlot.plot(time_grid, x[1:100,:]', "r--")
PyPlot.plot(time_grid, y, "g")

loglik, x,nbr_resamlping = diffusionbridgeparticlefilter_test(y, exp(log_σ_ϵ), ϕ[m,:], dt, u_prop_old[m], u_resample_old[m], nbr_particles[m], true, true)

PyPlot.figure()
PyPlot.plot(time_grid, y, "g")
PyPlot.plot(time_grid, x[:], "r--")

# run multiple loglik calcs

runs = 100

loglik_matrix = zeros(3, runs)

for i in 1:runs

    T_vec = zeros(Int64, M_subjects)
    for i = 1:M_subjects; T_vec[i] = length(y); end

    u_prop_old = [randn(nbr_particles[1], T_vec[1]+1)]
    u_resample_old = [randn(T_vec[1], 2)]

    for i in 2:M_subjects;
        append!(u_prop_old, [randn(nbr_particles[i],T_vec[i]+1)])
        append!(u_resample_old, [randn(T_vec[i],2)])
    end

    nbr_particles_bootstrap = 10000*ones(Int64, M_subjects)

    # set number of obs for each data set
    T_vec = zeros(Int64, M_subjects)
    for i = 1:M_subjects; T_vec[i] = length(y); end

    u_prop_old_bootstrap = [randn(nbr_particles_bootstrap[1], T_vec[1]+1)]
    u_resample_old_bootstrap = [randn(T_vec[1], 2)]

    for i in 2:M_subjects;
        append!(u_prop_old_bootstrap, [randn(nbr_particles_bootstrap[i],T_vec[i]+1)])
        append!(u_resample_old_bootstrap, [randn(T_vec[i],2)])
    end

    m = 1

    loglik_matrix[1,i], ~ = kalman_filter_test(y, exp(log_σ_ϵ),  ϕ[m,:], dt, true)
    loglik_matrix[2,i], ~ = particlefilter_test(y, exp(log_σ_ϵ), ϕ[m,:], dt, u_prop_old_bootstrap[m], u_resample_old_bootstrap[m], nbr_particles_bootstrap[m], true, true)
    loglik_matrix[3,i], ~ = diffusionbridgeparticlefilter_test(y, exp(log_σ_ϵ), ϕ[m,:], dt, u_prop_old[m], u_resample_old[m], nbr_particles[m], true, true)

end

println("Mean:")
println(round.(mean(loglik_matrix, dims = 2),digits=2))

println("Std:")
println(round.(std(loglik_matrix, dims = 2),digits=2))

src/SDEMEM OU neuron data/test_particle_filters_with_simulated_data.jl


"""
   inverse

Inverse a permutation.
If n > 0, fill with zeros. Can also use `invperm`
"""
function inverse(p::Vector{Int}, n::Int = 0)
    if n <= 0
        n = length(p)
    end
    q = zeros(Int, n)
    for (i, j) in enumerate(p)
        q[j] = i
    end
    q
end

"""
    concat(a, b)

TODO
"""
concat(a::Vector{T}, i::Vector{Int}) where T =
    T[a[j] for j in i if j > 0]


using Random

function is_sorted(a::Vector{Float64})
    for i = 1:length(a)-1
        if (a[i] > a[i + 1])
            return false
        end
    end
    return true
end

function bogo_sort(a::Vector{Float64})
    while(!is_sorted(a))
        shuffle!(a)
    end
end

function main()
    a = [1, 3, 2, 4]
    bogo_sort(a)
    println(a)
end

main()


using ArgParse

s=ArgParseSettings()
@add_arg_table s begin
    "--part2"
        help="Do part 2"
        action=:store_true
    "--test"
        help="Run test data"
        action=:store_true
end
parsed_args = parse_args(ARGS, s)

mutable struct Board
    squares::Dict
    columns::Array
    rows::Array
    done::Bool
end

function mark_board(board, num)
    if !board.done && haskey(board.squares, num)
        (row, col) = pop!(board.squares, num)
        board.rows[row] += 1 # Crossed off one more number in this row
        board.columns[col] += 1 # Likewise colum
        board.done = (board.rows[row] == 5 || board.columns[col]==5)
        return board.done
    end
    false
end

function find_winner(numbers, boards)
    for num in numbers
        for board in boards
            if mark_board(board, num)
               return num, board
            end
        end
    end
end

function find_loser(numbers, boards)
    number_ofboards = length(boards)
    for num in numbers
        for board in boards
            if mark_board(board, num)
                number_ofboards -= 1
                if number_ofboards == 0
                    return num, board
               end
            end
        end
    end
end

test_suffix = parsed_args["test"] ? "_test" : ""
lines = readlines("input/day_4$test_suffix.txt")
numbers = split(lines[1], ",")
boards = []
for boardIdx in 2:6:length(lines)-1
    squares = []
    for row in 1:5
        row_nums = split(lines[boardIdx+row])
        for (col, val) in enumerate(row_nums)
            push!(squares, (val,(row, col)))
        end
    end
    push!(boards,Board(Dict(squares), zeros(5), zeros(5), false))
end
if parsed_args["part2"]
    (last, the_one) = find_loser(numbers, boards)
else
    (last, the_one) = find_winner(numbers, boards)
end
unmarked = sum(s -> parse(Int32, s), keys(the_one.squares))
answer = unmarked * parse(Int32, last)
println("The answer is: $answer")


using Base.Test
using PValueAdjust

pvalues = [0.05, 0.05, 0.05]
newPvalues = padjust(pvalues, Bonferroni)
@test_approx_eq newPvalues pvalues * 3.0

@test all(padjust([1,1,1], Bonferroni) .<= 1)

@test_throws DomainError padjust([-1], Bonferroni)
@test_throws DomainError padjust([2], Bonferroni)


using TDD
using Test

@testset "TDD.jl" begin
    # Write your tests here.

    #1.

    undirected_graph = [
        [2, 3], #nodes reacheable from node 1
        [1],    #nodes reacheable from node 2
        [1],    #nodes reacheable from node 3
        [5],    #nodes reacheable from node 4
        [4]     #nodes reacheable from node 5
    ]

    A_undirected_graph = Bool[               #A short for "Adjacency Matrix Format"
                        1 1 1 0 0;
                        1 1 0 0 0;
                        1 0 1 0 0;
                        0 0 0 1 1;
                        0 0 0 1 1;
                        ]

    #based on Figure 1 from homework, node order 2 --> 1 --> 3; 4 --> 5
    directed_graph = [
        [3],    #nodes reacheable from node 1
        [1],    #nodes reacheable from node 2
        [],     #nodes reacheable from node 3
        [5],    #nodes reacheable from node 4
        []      #nodes reacheable from node 5
    ]

    A_directed_graph = Bool[
                        1 0 1 0 0;
                        1 1 0 0 0;
                        0 0 1 0 0;
                        0 0 0 1 1;
                        0 0 0 0 1;
                        ]

    #new test crazy case - a unidirectional train directed graph
    # 8 --> 7 --> 6 --> 5 --> 4 --> 3 --> 2 --> 1
    train_graph = [
        [],  #nodes reacheable from node 1
        [1], #nodes reacheable from node 2
        [2],
        [3],
        [4],
        [5],
        [6],
        [7]  #nodes reacheable from node 8
    ]

    A_train_graph = Bool[
                    1 0 0 0 0 0 0 0;
                    1 1 0 0 0 0 0 0;
                    0 1 1 0 0 0 0 0;
                    0 0 1 1 0 0 0 0;
                    0 0 0 1 1 0 0 0;
                    0 0 0 0 1 1 0 0;
                    0 0 0 0 0 1 1 0;
                    0 0 0 0 0 0 1 1;
                    ]

    @test DCneighbors(undirected_graph, 1) == Set([3, 2, 1]) #Sets are unordered
    @test DCneighbors(undirected_graph, 2) == Set([2, 1])
    @test DCneighbors(undirected_graph, 3) == Set([3, 1])
    @test DCneighbors(undirected_graph, 4) == Set([4, 5])
    @test DCneighbors(undirected_graph, 5) == Set([5, 4])

    @test DCneighbors(directed_graph, 1) == Set([1, 3])
    @test DCneighbors(directed_graph, 2) == Set([2, 1])
    @test DCneighbors(directed_graph, 3) == Set([3])
    @test DCneighbors(directed_graph, 4) == Set([4, 5])
    @test DCneighbors(directed_graph, 5) == Set([5])

    @test DCneighbors(train_graph, 1) == Set([1])
    @test DCneighbors(train_graph, 2) == Set([2,1])
    @test DCneighbors(train_graph, 3) == Set([3,2])
    @test DCneighbors(train_graph, 4) == Set([4,3])
    @test DCneighbors(train_graph, 5) == Set([5,4])
    @test DCneighbors(train_graph, 6) == Set([6,5])
    @test DCneighbors(train_graph, 7) == Set([7,6])
    @test DCneighbors(train_graph, 8) == Set([8,7])

    @test_throws ErrorException("node does not exist") DCneighbors(undirected_graph, 6)

    #2

    @test ICneighbors(undirected_graph, 1) == Set([1, 2, 3])
    @test ICneighbors(undirected_graph, 2) == Set([1, 2, 3])
    @test ICneighbors(undirected_graph, 3) == Set([1, 2, 3])
    @test ICneighbors(undirected_graph, 4) == Set([4, 5])
    @test ICneighbors(undirected_graph, 5) == Set([4, 5])

    @test ICneighbors(directed_graph, 1) == Set([1, 2, 3])
    @test ICneighbors(directed_graph, 2) == Set([1, 2, 3])
    @test ICneighbors(directed_graph, 3) == Set([1, 2, 3])
    @test ICneighbors(directed_graph, 4) == Set([4, 5])
    @test ICneighbors(directed_graph, 5) == Set([4, 5])

    @test_throws ErrorException("node does not exist") ICneighbors(undirected_graph, 6)

    #new tests
    @test ICneighbors(train_graph, 1) == Set([1, 2, 3,4,5,6,7,8])
    @test ICneighbors(train_graph, 2) == Set([1, 2, 3,4,5,6,7,8])
    @test ICneighbors(train_graph, 3) == Set([1, 2, 3,4,5,6,7,8])
    @test ICneighbors(train_graph, 4) == Set([1, 2, 3,4,5,6,7,8])
    @test ICneighbors(train_graph, 5) == Set([1, 2, 3,4,5,6,7,8])
    @test ICneighbors(train_graph, 6) == Set([1, 2, 3,4,5,6,7,8])
    @test ICneighbors(train_graph, 7) == Set([1, 2, 3,4,5,6,7,8])
    @test ICneighbors(train_graph, 8) == Set([1, 2, 3,4,5,6,7,8])

    #3

    @test AllComponents(directed_graph) == Set( [ Set([1, 2, 3]), Set([4, 5]) ] ) #for some reason need to wrap Set argument in []
    @test AllComponents(undirected_graph) == Set( [ Set([1, 2, 3]), Set([4, 5]) ] )
    @test AllComponents(train_graph) == Set( [ Set([1, 2, 3,4,5,6,7,8]) ] )

    #4 tests for adjacency matrix format -- see matrices above, tests below (copied, changed matrices input)

        @test DCneighbors(A_undirected_graph, 1) == Set([3, 2, 1])                              #Sets are unordered
        @test DCneighbors(A_undirected_graph, 2) == Set([2, 1]) #passes erreously
        @test DCneighbors(A_undirected_graph, 3) == Set([3, 1]) #passes erreously
        @test DCneighbors(A_undirected_graph, 4) == Set([4, 5])
        @test DCneighbors(A_undirected_graph, 5) == Set([5, 4])

        @test DCneighbors(A_directed_graph, 1) == Set([1, 3])
        @test DCneighbors(A_directed_graph, 2) == Set([2, 1]) #passes erreously
        @test DCneighbors(A_directed_graph, 3) == Set([3])
        @test DCneighbors(A_directed_graph, 4) == Set([4, 5])
        @test DCneighbors(A_directed_graph, 5) == Set([5])

        @test DCneighbors(A_train_graph, 1) == Set([1]) #passes erreously
        @test DCneighbors(A_train_graph, 2) == Set([2,1]) #passes erreously
        @test DCneighbors(A_train_graph, 3) == Set([3,2])
        @test DCneighbors(A_train_graph, 4) == Set([4,3])
        @test DCneighbors(A_train_graph, 5) == Set([5,4])
        @test DCneighbors(A_train_graph, 6) == Set([6,5])
        @test DCneighbors(A_train_graph, 7) == Set([7,6])
        @test DCneighbors(A_train_graph, 8) == Set([8,7])

        @test_throws ErrorException("node does not exist") DCneighbors(undirected_graph, 6) #passes erreously (though expected? this is technicallythetruth)

        #4-2

        @test ICneighbors(A_undirected_graph, 1) == Set([1, 2, 3])
        @test ICneighbors(A_undirected_graph, 2) == Set([1, 2, 3])
        @test ICneighbors(A_undirected_graph, 3) == Set([1, 2, 3])
        @test ICneighbors(A_undirected_graph, 4) == Set([4, 5])
        @test ICneighbors(A_undirected_graph, 5) == Set([4, 5])

        @test ICneighbors(A_directed_graph, 1) == Set([1, 2, 3])
        @test ICneighbors(A_directed_graph, 2) == Set([1, 2, 3])
        @test ICneighbors(A_directed_graph, 3) == Set([1, 2, 3])
        @test ICneighbors(A_directed_graph, 4) == Set([4, 5])
        @test ICneighbors(A_directed_graph, 5) == Set([4, 5])

        @test_throws ErrorException("node does not exist") ICneighbors(undirected_graph, 6) #passes erreously (though expected? this is technicallythetruth)

        #new tests
        @test ICneighbors(A_train_graph, 1) == Set([1, 2, 3,4,5,6,7,8])
        @test ICneighbors(A_train_graph, 2) == Set([1, 2, 3,4,5,6,7,8])
        @test ICneighbors(A_train_graph, 3) == Set([1, 2, 3,4,5,6,7,8])
        @test ICneighbors(A_train_graph, 4) == Set([1, 2, 3,4,5,6,7,8])
        @test ICneighbors(A_train_graph, 5) == Set([1, 2, 3,4,5,6,7,8])
        @test ICneighbors(A_train_graph, 6) == Set([1, 2, 3,4,5,6,7,8])
        @test ICneighbors(A_train_graph, 7) == Set([1, 2, 3,4,5,6,7,8])
        @test ICneighbors(A_train_graph, 8) == Set([1, 2, 3,4,5,6,7,8])

        #4-3

        @test AllComponents(A_directed_graph) == Set( [ Set([1, 2, 3]), Set([4, 5]) ] ) #for some reason need to wrap Set argument in []
        @test AllComponents(A_undirected_graph) == Set( [ Set([1, 2, 3]), Set([4, 5]) ] )
        @test AllComponents(A_train_graph) == Set( [ Set([1, 2, 3,4,5,6,7,8]) ] )

        #note at this point, 7 tests for Adjacency Matrix input pass, even though I have not implemented this...
        #need to go back figure out which tests erreously passed
        #I assume the error exceptions. But also which others?
        #Yep erroneous passes labelled in comments see above. Now gotta figure out why
        #relevant bugs are in DCneighbors
        #Ok so ICneighbors if input some Bool_matrix and a node, it does Bool_matrix[node] which just returns true/false, so makes sense all those tests fail right now
        #AllComponents depends on ICneighbors so also makes sense those tests fail
        #But DCneighbors does something different
        #If I run DCneighbors(A_undirected_graph, somenode) for example it returns some Set{Int64} which may or may not pass the test
        #Because it is just doing `Set(vcat(node, graph_input[node]))` which returns a Set{Int64}([node, 0 or 1 depending on if indexed value from Bool_matrix]) which may or may not be the right answer
        #not gonna fix DCneighbors right now so that it correctly fails the tests, afraid gonna mess something up
        #just keep this in mind. When I implement the functions with multiple dispatch specifying method for input Bool matrix this will end up being moot

end


"""
	asympFit(x,columns)

Function to estimate the asmpytotic value of α(x).

# Arguments
 - `alphaValues::Array{Float64,1}`: α(x) array.
# Returns
 - `Array{Float64,2}`: Array of array containing asymptotic values, lower confidence interval, higher confidence interval
"""
function asympFit(alphaValues::Array{Float64,1})

	# Model
	asympModel(x,p) = @. p[1] + p[2]*exp(-x*p[3])

	# Fit values
	fitted1    = curve_fit(asympModel,collect(1:size(alphaValues,1)),alphaValues,[-1.0,-1.0,1.0];lower=[-1.0,-1.0,1.0],upper=[1.0, 1.0, 10.0])
	fitted2    = curve_fit(asympModel,collect(1:size(alphaValues,1)),alphaValues,fitted1.param)
	asymp      = asympModel(size(alphaValues,1),fitted2.param)

	ciLow, ciHigh   = try
		confidence_interval(fitted2)[1][1],LsqFit.confidence_interval(fitted2)[1][2]
	catch err
		(0.0,0.0)
	end

	return asymp,[ciLow,ciHigh],fitted2.param
	# return asymp
end

"""
	imputedMK(sfs,divergence,m,cutoff)

Function to estimate the asmpytotic value of α(x).

# Arguments
 - `alphaValues::Array{Float64,1}`: α(x) array.
# Returns
 - `Array{Float64,2}`: Array of array containing asymptotic values, lower confidence interval, higher confidence interval
"""
function impMK(;sfs::Array,divergence::Array,m::T,cutoff::Float64=0.15) where {T<:Union{Nothing,Array}}

	output = OrderedDict{String,Float64}()

	pn = sum(sfs[:,2])
	ps = sum(sfs[:,3])
	dn = divergence[1]
	ds = divergence[2]

	deleterious = 0
	### Estimating slightly deleterious with pn/ps ratio
	fltLow = (sfs[:, 1] .<= cutoff)
	pnLow   = sum(sfs[fltLow,2])
	psLow   = sum(sfs[fltLow,3])

	fltInter = (sfs[:, 1] .>= cutoff) .& (sfs[:, 1] .<= 1)
	pnInter = sum(sfs[fltInter,2])
	psInter = sum(sfs[fltInter,3])

	ratioPs       = psLow / psInter
	deleterious   = pnLow - (pnInter * ratioPs)

	if deleterious > pn
		deleterious = 0
		pnNeutral = round(pn - deleterious,digits=3)
	else
		pnNeutral = round(pn - deleterious,digits=3)
	end
	# output['alpha'] = 1 - (((pn - deleterious) / ps) * (ds / dn))
	output["alpha"] = round(1 - ((pnNeutral/ps) * (ds / dn)),digits=3)
	#  method = :minlike same results R, python two.sides
	output["pvalue"] = pvalue(FisherExactTest(Int(ps),Int(ceil(pnNeutral)),Int(ds),Int(dn)))

	if (!isnothing(m))

		mn = m[1]; ms = m[2]
		# ## Estimation of b: weakly deleterious
		output["b"] = (deleterious / ps) * (ms / mn)

		## Estimation of f: neutral sites
		output["f"] = (ms * pnNeutral) / (mn * ps)

		## Estimation of d, strongly deleterious sites
		output["d"] = 1 - (output["f"] + output["b"])

		ka      = dn / mn
		ks       = ds / ms
		output["omega"]    = ka/ ks

		# Omega A and Omega D
		output["omegaA"] = output["omega"] * output["alpha"]
		output["omegaD"] = output["omega"] - output["omegaA"]
	end

	return output
end

"""
	fwwMK(sfs,divergence,m,cutoff)

Function to estimate the asmpytotic value of α(x).

# Arguments
 - `alphaValues::Array{Float64,1}`: α(x) array.
# Returns
 - `Array{Float64,2}`: Array of array containing asymptotic values, lower confidence interval, higher confidence interval
"""
function fwwMK(;sfs::Array,divergence::Array,m::T,cutoff::Float64=0.15) where {T<:Union{Nothing,Array}}

	output = OrderedDict{String,Float64}()

	ps = sum(sfs[:,2])
	pn = sum(sfs[:,3])
	dn = divergence[1]
	ds = divergence[2]

	deleterious = 0
	### Estimating slightly deleterious with pn/ps ratio
	fltInter = (sfs[:, 1] .>= cutoff) .& (sfs[:, 1] .<= 1)
	pnInter = sum(sfs[fltInter,2])
	psInter = sum(sfs[fltInter,3])

	# output['alpha'] = 1 - (((pn - deleterious) / ps) * (ds / dn))
	output["alpha"] = round(1 - ((pnInter/psInter) * (ds / dn)),digits=3)
	#  method = :minlike same results R, python two.sides
	output["pvalue"] = pvalue(FisherExactTest(Int(psInter),Int(ceil(pnInter)),Int(ds),Int(dn)))

	if (!isnothing(m))

		mn = m[1]; ms = m[2]
		ka      = dn / mn
		ks       = ds / ms
		output["omega"]    = ka/ ks

		# Omega A and Omega D
		output["omegaA"] = output["omega"] * output["alpha"]
		output["omegaD"] = output["omega"] - output["omegaA"]
	end

	return output
end

"""
	standardMK(x,columns)

Function to estimate the original α value

# Arguments
 - `alphaValues::Array{Float64,1}`: α(x) array.
# Returns
 - `Array{Float64,2}`: Array of array containing asymptotic values, lower confidence interval, higher confidence interval
"""
function standardMK(;sfs::Array,divergence::Array,m::T=nothing) where {T<:Union{Nothing,Array}}

	output = OrderedDict{String,Float64}()

	pn = sum(sfs[:,2])
	ps = sum(sfs[:,3])
	dn = divergence[1]
	ds = divergence[2]

	output["alpha"] = round(1 - ((pn/ps) * (ds / dn)),digits=3)
	#  method = :mnnlike same results R, python two.sides
	output["pvalue"] = pvalue(FisherExactTest(Int(ps),Int(ceil(pn)),Int(ds),Int(dn)))

	if (!isnothing(m))

		mn = m[1]; ms = m[2]

		ka      = dn / mn
		ks       = ds / ms
		output["omega"]    = ka/ ks

		# Omega A and Omega D
		output["omegaA"] = output["omega"] * output["alpha"]
		output["omegaD"] = output["omega"] - output["omegaA"]
	end

	return output
end


"""
    check_errs(f, ∇f, Ȳ::AS, X::AS, V::AS, ε_abs::Real=1e-7, ε_rel::Real=1e-5)::Bool
"""
check_errs(f, ∇f, Ȳ::AS, X::AS, V::AS, ε_abs::Real=1e-7, ε_rel::Real=1e-5)::Bool =
    assert_approx_equal(
        central_fdm(2, 1)(ϵ->sum(Ȳ .* f(X + ϵ * V))),
        sum(∇f(f(X), Ȳ, X) .* V),
        ε_abs, ε_rel)

"""
    check_errs(N::Int, f, ∇f, Ȳ, X, V, ε_abs=1e-7, ε_rel=1e-5)::Bool

Call check_errs `N` times with arguments generated by 0-ary functions `Ȳ`, `X` and `V`.
"""
check_errs(N::Int, f, ∇f, Ȳ, X, V, ε_abs=1e-7, ε_rel=1e-5)::Bool =
    all(map(n->check_errs(f, ∇f, Ȳ(), X(), V(), ε_abs, ε_rel), 1:N))

# Utility to create the closures required for unit testing.
unary_ȲD(f) = (f, (Y, Ȳ, X)->∇(f, Val{1}, (), Y, Ȳ, X))
unary_ȲD_inplace(f, X̄0) = (f, (Y, Ȳ, X)->∇(copy(X̄0), f, Val{1}, (), Y, Ȳ, X) - X̄0)
function binary_ȲD(f, arg::Int, G)
    G_ = G()
    if arg == 1
        return X->f(X, G_), (Y, Ȳ, X)->∇(f, Val{1}, (), Y, Ȳ, X, G_)
    else
        return X->f(G_, X), (Y, Ȳ, X)->∇(f, Val{2}, (), Y, Ȳ, G_, X)
    end
end
function binary_ȲD_inplace(f, arg::Int, G, X̄0)
    G_ = G()
    if arg == 1
        return X->f(X, G_), (Y, Ȳ, X)->∇(copy(X̄0), f, Val{1}, (), Y, Ȳ, X, G_) - X̄0
    else
        return X->f(G_, X), (Y, Ȳ, X)->∇(copy(X̄0), f, Val{2}, (), Y, Ȳ, G_, X) - X̄0
    end
end


@testset "Model matrix" begin

    using StatsBase: StatisticalModel

    using SparseArrays, DataFrames, Tables

    sparsetype = SparseMatrixCSC{Float64,Int}

    d = DataFrame(y = 1:4,
                  x1 = 5:8,
                  x2 = 9:12,
                  x3 = 13:16,
                  x4 = 17:20)
    d.x1p = categorical(d.x1)

    d_orig = deepcopy(d)

    x1 = [5.:8;]
    x2 = [9.:12;]
    x3 = [13.:16;]
    x4 = [17.:20;]
    f = @formula(y ~ 1 + x1 + x2)
    mf = ModelFrame(f, d)
    @test coefnames(mf) == ["(Intercept)","x1","x2"]
    @test response(mf) == [1:4;]
    mm = ModelMatrix(mf)
    smm = ModelMatrix{sparsetype}(mf)
    @test mm.m[:,1] == ones(4)
    @test mm.m[:,2:3] == [x1 x2]
    @test mm.m == smm.m

    @test isa(mm.m, Matrix{Float64})
    @test isa(smm.m, sparsetype)

    f_implint = @formula(y ~ x1 + x2)
    @test ModelMatrix(ModelFrame(f_implint, d)).m == mm.m

    @test ModelMatrix(ModelFrame(f_implint, d, model=Nothing)).m == hcat(x1, x2)

    #test_group("expanding a nominal array into a design matrix of indicators for each dummy variable")

    mf = ModelFrame(@formula(y ~ 1 + x1p), d)
    mm = ModelMatrix(mf)

    @test mm.m[:,2] == [0, 1., 0, 0]
    @test mm.m[:,3] == [0, 0, 1., 0]
    @test mm.m[:,4] == [0, 0, 0, 1.]
    @test coefnames(mf)[2:end] == ["x1p: 6", "x1p: 7", "x1p: 8"]
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    #test_group("Creating a model matrix using full formulas: y => x1 + x2, etc")

    d = deepcopy(d_orig)
    f = @formula(y ~ 1 + x1 & x2)
    mf = ModelFrame(f, d)
    mm = ModelMatrix(mf)
    @test mm.m == [ones(4) x1.*x2]
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    f = @formula(y ~ 1 + x1 * x2)
    mf = ModelFrame(f, d)
    mm = ModelMatrix(mf)
    @test mm.m == [ones(4) x1 x2 x1.*x2]
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    d.x1 = categorical(x1)
    x1e = [[0, 1, 0, 0] [0, 0, 1, 0] [0, 0, 0, 1]]
    f = @formula(y ~ 1 + x1 * x2)
    mf = ModelFrame(f, d)
    mm = ModelMatrix(mf)
    @test mm.m == [ones(4) x1e x2 [0, 10, 0, 0] [0, 0, 11, 0] [0, 0, 0, 12]]
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    # additional tests from Tom
    y = [1., 2, 3, 4]
    mf = ModelFrame(@formula(y ~ 1 + x2), d)
    mm = ModelMatrix(mf)
    @test mm.m == [ones(4) x2]
    @test mm.m == ModelMatrix{sparsetype}(mf).m
    @test response(mf) == y''

    d = deepcopy(d_orig)
    d.x1 = CategoricalArray{Union{Missing, Float64}}(d.x1)

    f = @formula(y ~ 1 + x2 + x3 + x3*x2)
    mm = ModelMatrix(ModelFrame(f, d))
    @test mm.m == [ones(4) x2 x3 x2.*x3]
    mm = ModelMatrix(ModelFrame(@formula(y ~ 1 + x3*x2 + x2 + x3), d))
    @test mm.m == [ones(4) x3 x2 x2.*x3]
    mm = ModelMatrix(ModelFrame(@formula(y ~ 1 + x1 + x2 + x3 + x4), d))
    @test mm.m[:,2] == [0, 1., 0, 0]
    @test mm.m[:,3] == [0, 0, 1., 0]
    @test mm.m[:,4] == [0, 0, 0, 1.]
    @test mm.m[:,5] == x2
    @test mm.m[:,6] == x3
    @test mm.m[:,7] == x4

    mm = ModelMatrix(ModelFrame(@formula(y ~ 1 + x2 + x3 + x4), d))
    @test mm.m == [ones(4) x2 x3 x4]
    mm = ModelMatrix(ModelFrame(@formula(y ~ 1 + x2 + x2), d))
    @test mm.m == [ones(4) x2]
    mm = ModelMatrix(ModelFrame(@formula(y ~ 1 + x2*x3 + x2&x3), d))
    @test mm.m == [ones(4) x2 x3 x2.*x3]
    mm = ModelMatrix(ModelFrame(@formula(y ~ 1 + x2*x3*x4), d))
    @test mm.m == [ones(4) x2 x3 x4 x2.*x3 x2.*x4 x3.*x4 x2.*x3.*x4]
    mm = ModelMatrix(ModelFrame(@formula(y ~ 1 + x2&x3 + x2*x3), d))
    @test mm.m == [ones(4) x2 x3 x2.*x3]

    f = @formula(y ~ 1 + x2 & x3 & x4)
    mf = ModelFrame(f, d)
    mm = ModelMatrix(mf)
    @test mm.m == [ones(4) x2.*x3.*x4]
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    f = @formula(y ~ 1 + x1 & x2 & x3)
    mf = ModelFrame(f, d)
    mm = ModelMatrix(mf)
    @test mm.m[:, 2:end] == Matrix(Diagonal(x2.*x3))
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    ## Distributive property of :& over :+
    d = deepcopy(d_orig)
    f = @formula(y ~ 1 + (x1+x2) & (x3+x4))
    mf = ModelFrame(f, d)
    mm = ModelMatrix(mf)
    @test mm.m == hcat(ones(4), x1.*x3, x1.*x4, x2.*x3, x2.*x4)
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    ## Condensing nested :+ calls
    f = @formula(y ~ 1 + x1 + (x2 + (x3 + x4)))
    @test ModelMatrix(ModelFrame(f, d)).m == hcat(ones(4), x1, x2, x3, x4)

    ## Extra levels in categorical column
    mf_full = ModelFrame(@formula(y ~ 1 + x1p), d)
    mm_full = ModelMatrix(mf_full)
    @test size(mm_full) == (4,4)

    mf_sub = ModelFrame(@formula(y ~ 1 + x1p), d[2:4, :])
    mm_sub = ModelMatrix(mf_sub)
    ## should have only three rows, and only three columns (intercept plus two
    ## levels of factor) - it also catches when some levels do not occur in the data
    @test size(mm_sub) == (3,3)

    ## Missing data
    d.x1m = [5, 6, missing, 7]
    mf = ModelFrame(@formula(y ~ 1 + x1m), d)
    mm = ModelMatrix(mf)
    @test mm.m[:, 2] == d[completecases(d), :x1m]
    @test mm.m == ModelMatrix{sparsetype}(mf).m

    ## Same variable on left and right side
    mf = ModelFrame(@formula(x1 ~ x1), d)
    mm = ModelMatrix(mf)
    mm.m == response(mf)

    ## Promote non-redundant categorical terms to full rank
    @testset "non-redundant categorical terms" begin
        d = DataFrame(x = repeat([:a, :b], outer = 4),
                      y = repeat([:c, :d], inner = 2, outer = 2),
                      z = repeat([:e, :f], inner = 4))
        cs = Dict([Symbol(name) => EffectsCoding() for name in names(d)])
        d.n = 1.:8

        ## No intercept
        mf = ModelFrame(@formula(n ~ 0 + x), d, contrasts=cs)
        mm = ModelMatrix(mf)
        @test mm.m == [1 0
                       0 1
                       1 0
                       0 1
                       1 0
                       0 1
                       1 0
                       0 1]
        @test mm.m == ModelMatrix{sparsetype}(mf).m
        @test coefnames(mf) == ["x: a", "x: b"]

        ## promotion blocked when we block default model=StatisticalModel
        mf = ModelFrame(@formula(n ~ 0 + x), d, model=Nothing, contrasts=cs)
        mm = ModelMatrix(mf)
        @test all(mm.m .== ifelse.(d.x .== :a, -1, 1))
        @test coefnames(mf) == ["x: b"]

        ## No first-order term for interaction
        mf = ModelFrame(@formula(n ~ 1 + x + x&y), d, contrasts=cs)
        mm = ModelMatrix(mf)
        @test mm.m[:, 2:end] == [-1 -1  0
                                 1  0 -1
                                 -1  1  0
                                 1  0  1
                                 -1 -1  0
                                 1  0 -1
                                 -1  1  0
                                 1  0  1]
        @test mm.m == ModelMatrix{sparsetype}(mf).m
        @test coefnames(mf) == ["(Intercept)", "x: b", "x: a & y: d", "x: b & y: d"]

        ## When both terms of interaction are non-redundant:
        mf = ModelFrame(@formula(n ~ 0 + x&y), d, contrasts=cs)
        mm = ModelMatrix(mf)
        @test mm.m == [1 0 0 0
                       0 1 0 0
                       0 0 1 0
                       0 0 0 1
                       1 0 0 0
                       0 1 0 0
                       0 0 1 0
                       0 0 0 1]
        @test mm.m == ModelMatrix{sparsetype}(mf).m
        @test coefnames(mf) == ["x: a & y: c", "x: b & y: c",
                                "x: a & y: d", "x: b & y: d"]

        # only a three-way interaction: every term is promoted.
        mf = ModelFrame(@formula(n ~ 0 + x&y&z), d, contrasts=cs)
        mm = ModelMatrix(mf)
        @test mm.m == Matrix(1.0I, 8, 8)
        @test mm.m == ModelMatrix{sparsetype}(mf).m

        # two two-way interactions, with no lower-order term. both are promoted in
        # first (both x and y), but only the old term (x) in the second (because
        # dropping x gives z which isn't found elsewhere, but dropping z gives x
        # which is found (implicitly) in the promoted interaction x&y).
        mf = ModelFrame(@formula(n ~ 0 + x&y + x&z), d, contrasts=cs)
        mm = ModelMatrix(mf)
        @test mm.m == [1 0 0 0 -1  0
                       0 1 0 0  0 -1
                       0 0 1 0 -1  0
                       0 0 0 1  0 -1
                       1 0 0 0  1  0
                       0 1 0 0  0  1
                       0 0 1 0  1  0
                       0 0 0 1  0  1]
        @test mm.m == ModelMatrix{sparsetype}(mf).m
        @test coefnames(mf) == ["x: a & y: c", "x: b & y: c",
                                "x: a & y: d", "x: b & y: d",
                                "x: a & z: f", "x: b & z: f"]

        # ...and adding a three-way interaction, only the shared term (x) is promoted.
        # this is because dropping x gives y&z which isn't present, but dropping y or z
        # gives x&z or x&z respectively, which are both present.
        mf = ModelFrame(@formula(n ~ 0 + x&y + x&z + x&y&z), d, contrasts=cs)
        mm = ModelMatrix(mf)
        @test mm.m == [1 0 0 0 -1  0  1  0
                       0 1 0 0  0 -1  0  1
                       0 0 1 0 -1  0 -1  0
                       0 0 0 1  0 -1  0 -1
                       1 0 0 0  1  0 -1  0
                       0 1 0 0  0  1  0 -1
                       0 0 1 0  1  0  1  0
                       0 0 0 1  0  1  0  1]
        @test mm.m == ModelMatrix{sparsetype}(mf).m
        @test coefnames(mf) == ["x: a & y: c", "x: b & y: c",
                                "x: a & y: d", "x: b & y: d",
                                "x: a & z: f", "x: b & z: f",
                                "x: a & y: d & z: f", "x: b & y: d & z: f"]

        # two two-way interactions, with common lower-order term. the common term x is
        # promoted in both (along with lower-order term), because in every case, when
        # x is dropped, the remaining terms (1, y, and z) aren't present elsewhere.
        mf = ModelFrame(@formula(n ~ 0 + x + x&y + x&z), d, contrasts=cs)
        mm = ModelMatrix(mf)
        @test mm.m == [1 0 -1  0 -1  0
                       0 1  0 -1  0 -1
                       1 0  1  0 -1  0
                       0 1  0  1  0 -1
                       1 0 -1  0  1  0
                       0 1  0 -1  0  1
                       1 0  1  0  1  0
                       0 1  0  1  0  1]
        @test mm.m == ModelMatrix{sparsetype}(mf).m
        @test coefnames(mf) == ["x: a", "x: b",
                                "x: a & y: d", "x: b & y: d",
                                "x: a & z: f", "x: b & z: f"]

        ## FAILS: When both terms are non-redundant and intercept is PRESENT
        ## (not fully redundant). Ideally, would drop last column. Might make sense
        ## to warn about this, and suggest recoding x and y into a single variable.
        mf = ModelFrame(@formula(n ~ 1 + x&y), d[1:4, :], contrasts=cs)
        @test_broken ModelMatrix(mf).m == [1 1 0 0
                                           1 0 1 0
                                           1 0 0 1
                                           1 0 0 0]
        @test_broken coefnames(mf) == ["x: a & y: c", "x: b & y: c",
                                       "x: a & y: d", "x: b & y: d"]

        ## note that R also does not detect this automatically. it's left to glm et al.
        ## to detect numerically when the model matrix is rank deficient, which is hard
        ## to do correctly.
        # > d = data.frame(x = factor(c(1, 2, 1, 2)), y = factor(c(3, 3, 4, 4)))
        # > model.matrix(~ 1 + x:y, d)
        #   (Intercept) x1:y3 x2:y3 x1:y4 x2:y4
        # 1           1     1     0     0     0
        # 2           1     0     1     0     0
        # 3           1     0     0     1     0
        # 4           1     0     0     0     1

    end

    @testset "arbitrary functions in formulae" begin
        d = deepcopy(d_orig)
        mf = ModelFrame(@formula(y ~ log(x1)), d, model=Nothing)
        @test coefnames(mf) == ["log(x1)"]
        mm = ModelMatrix(mf)
        @test all(mm.m .== log.(x1))

        # | is not special in base formula:
        d = DataFrame(x = [1,2,3], y = [4,5,6])
        mf = ModelFrame(@formula(y ~ 1 + (1 | x)), d)
        @test coefnames(mf) == ["(Intercept)", "1 | x"]

        mf = ModelFrame(@formula(y ~ 0 + (1 | x)), d)
        @test all(ModelMatrix(mf).m .== float.(1 .| d.x))
        @test coefnames(mf) == ["1 | x"]
    end

    # Ensure X is not a view on df column
    d = DataFrame(x = [1.0,2.0,3.0], y = [4.0,5.0,6.0])
    mf = ModelFrame(@formula(y ~ 0 + x), d)
    X = ModelMatrix(mf).m
    X[1] = 0.0
    @test mf.data[:x][1] === 1.0

    # Ensure string columns are supported
    d1 = DataFrame(A = 1:4, B = categorical(["M", "F", "F", "M"]))
    d2 = DataFrame(A = 1:4, B = ["M", "F", "F", "M"])
    d3 = DataFrame(Any[1:4, ["M", "F", "F", "M"]], [:A, :B])

    M1 = ModelMatrix(ModelFrame(@formula(A ~ B), d1))
    M2 = ModelMatrix(ModelFrame(@formula(A ~ B), d2))
    M3 = ModelMatrix(ModelFrame(@formula(A ~ B), d3))

    @test (M1.m, M1.assign) == (M2.m, M2.assign) == (M3.m, M3.assign)

    @testset "row-wise model matrix construction" begin
        d = DataFrame(r = rand(8),
                      w = rand(8),
                      x = repeat([:a, :b], outer = 4),
                      y = repeat([:c, :d], inner = 2, outer = 2),
                      z = repeat([:e, :f], inner = 4))

        f = apply_schema(@formula(r ~ 1 + w*x*y*z), schema(d))
        modelmatrix(f, d)
        @test reduce(vcat, last.(modelcols.(Ref(f), Tables.rowtable(d)))') == modelmatrix(f,d)
    end

    @testset "modelmatrix and response set schema if needed" begin
        d = DataFrame(r = rand(8),
                      w = rand(8),
                      x = repeat([:a, :b], outer = 4),
                      y = repeat([:c, :d], inner = 2, outer = 2),
                      z = repeat([:e, :f], inner = 4))

        f = @formula(r ~ 1 + w*x*y*z)

        mm1 = modelmatrix(f, d)
        mm2 = modelmatrix(apply_schema(f, schema(d)), d)
        @test mm1 == mm2

        r1 = response(f, d)
        r2 = response(apply_schema(f, schema(d)), d)
        @test r1 == r2
    end

    @testset "setcontrasts!" begin
        @testset "#95" begin
            tbl = (Y = randn(8),
                   A = repeat(['N','Y'], outer=4),
                   B=repeat(['-','+'], inner=2, outer=2),
                   C=repeat(['L','H'], inner=4))

            contrasts = Dict(:A=>HelmertCoding(), :B=>HelmertCoding(), :C=>HelmertCoding())

            mf = ModelFrame(@formula(Y ~ 1 + A*B*C), tbl)
            mf_helm = ModelFrame(@formula(Y ~ 1 + A*B*C), tbl, contrasts = contrasts)

            @test size(modelmatrix(mf)) == size(modelmatrix(mf_helm))

            mf_helm2 = setcontrasts!(ModelFrame(@formula(Y ~ 1 + A*B*C), tbl), contrasts)

            @test size(modelmatrix(mf)) == size(modelmatrix(mf_helm2))
            @test modelmatrix(mf_helm) == modelmatrix(mf_helm2)

        end
    end

    @testset "#136" begin
        t = (x = rand(100), y = randn(100));
        f = @formula(y ~ x)
        @test_throws ArgumentError modelcols(f, t)
    end

    @testset "#185 - interactions of scalar terms for row tables" begin
        t = (a = rand(10), b = rand(10), c = rand(10))
        f = apply_schema(@formula(0 ~ a&b&c), schema(t))
        @test vec(modelcols(f.rhs, t)) == modelcols.(Ref(f.rhs), Tables.rowtable(t))
    end

end


vector(x::Vector{T}) where {T} = x
vector(x::Array{T,N}) where {T,N} = vec(x)

function vector(x::NTuple{N,T}) where {N,T}
     result = Vector{T}(undef, N)
     for i=1:N
         result[i] = x[i]
     end
     return result
end

"""
     coordvecs(seq_of_coord, ::Type{T}=Float64)

```
vec_of_coord_tuples = [(1, 2), (2, 4)]
coordvecs(vec_of_coord_tuples) == [[1.0, 2.0], [2.0, 4.0]]

vec_of_coord_vectors = [[1, 2], [2, 4]]
coordvecs(vec_of_coord_vectors) == [[1.0, 2.0], [2.0, 4.0]]

tuple_of_coord_tuples = ((1, 2), (2, 4))
coordvecs(tuple_of_coord_tuples) == ([1.0, 2.0], [2.0, 4.0])
coordvecs(tuple_of_coord_tuples, Float32) == ([1.0f0, 2.0f0], [2.0f0, 4.0f0])
```
"""
coordvecs(coordseq::Vector{Vector{T}}) where {T<:Integer} = coordvecs(coordseq, Float64)
coordvecs(coordseq::Vector{NTuple{N,T}}) where {N, T<:Integer} = coordvecs(coordseq, Float64)
coordvecs(coordseq::NTuple{N1, NTuple{N2,T}}) where {N1, N2, T<:Integer} = coordvecs(coordseq, Float64)
coordvecs(coordseq::NTuple{N, Vector{T}}) where {N, T<:Integer} = coordvecs(coordseq, Float64)

coordvecs(coordseq::Vector{Vector{T}}) where {T<:AbstractFloat} = coordvecs(coordseq, T)
coordvecs(coordseq::Vector{NTuple{N,T}}) where {N, T<:AbstractFloat} = coordvecs(coordseq, T)
coordvecs(coordseq::NTuple{N1, NTuple{N2,T}}) where {N1, N2, T<:AbstractFloat} = coordvecs(coordseq, T)
coordvecs(coordseq::NTuple{N, Vector{T}}) where {N, T<:AbstractFloat} = coordvecs(coordseq, T)

coordvecs(coordseq::Vector{Vector{T1}}, ::Type{T2}) where {T1<:Real, T2<:AbstractFloat} = map(x->T2.(x), coordseq)
coordvecs(coordseq::Vector{NTuple{N,T1}}, ::Type{T2}) where {N, T1<:Real, T2<:AbstractFloat} = map(x->[T2.(x)...,], coordseq)
coordvecs(coordseq::NTuple{N1, NTuple{N2,T1}}, ::Type{T2}) where {N1, N2, T1<:Real, T2<:AbstractFloat} = map(x->[T2.(x)...,], coordseq)
coordvecs(coordseq::NTuple{N, Vector{T1}}, ::Type{T2}) where {N, T1<:Real, T2<:AbstractFloat} = map(x->[T2.(x)...,], coordseq)

vectorofcoordvecs(coordseq) = vector(coordvecs(coordseq))
matrixfromcoords(coordseq) = permutedims(reduce(hcat, vectorofcoordvecs(coordseq)))


f(x, y) = 2x + y           # call-like form on left-hand side
g = x -> x * x             # arrow notation
@assert g(f(2, 9)) == 169