#=== Cube ===#

struct Cube{P} <: AbstractCell
    origin::AbstractVector{P}
    extent::AbstractVector{P}
    function Cube(o::AbstractVector{P}, x::AbstractVector{P}) where{P}
        @assert length(x) <= length(o) "Too many extent elements"
        new{P}(o, x)
    end
end
Cube(x::AbstractVector{P}) where {P} = Cube(zeros(P, length(x)), x)
Cube(::Type{P}) where {P} = Cube(P[], P[])
show(io::IO, c::Cube) = show(io, "Cube[$(c.origin) + $(c.extent)]")

# Public methods

dim(c::Cube) = length(findall(!iszero, c.origin - c.extent))

==(a::Cube, b::Cube) = hash(a) == hash(b)

hash(c::Cube) = hash(vcat(c.origin, c.extent))

function faces(c::Cube{P}, ε=zero(P)) where {P}
    d = dim(c)
    d == 0 && return Cube[]
    kₙ = findall(!iszero, c.origin - c.extent) # on-degenerate intervals
    (Cube(copy(c.origin), [i == j ? ε : c.extent[j] for j in 1:length(c.origin)]) for i in kₙ)
end

function boundary(::Type{R}, c::Cube) where {R}
    d = dim(c)
    ch = Chain(d-1, R)
    d == 0 && return ch
    o = one(R)
    for (i,(λ⁰, λ¹)) in enumerate(zip(faces(c,zero(R)), faces(c, one(R))))
        # println("λ⁰:", λ⁰," λ¹:", λ¹)
        push!(ch, hash(λ⁰)=>( isodd(i) ? -o : o))
        push!(ch, hash(λ¹)=>( isodd(i) ? o : -o))
    end
    return simplify(ch)
end

# Misc. methods

function volume(c::Cube)
    sizes = similar(c.origin,0)
    for (i,ex) in enumerate(c.extent)
        if ex != 0.
            expoint = copy(c.origin)
            expoint[i] += ex
            l = norm(expoint - c.origin)
            push!(sizes, l)
        end
    end
    return prod(sizes)
end


export SuperpositionMeasure

"""
    struct SuperpositionMeasure{X,NT} <: AbstractMeasure
        components :: NT
    end

Superposition of measures is analogous to mixture distributions, but (because
measures need not be normalized) requires no scaling.

The superposition of two measures μ and ν can be more concisely written as μ + ν.
"""
struct SuperpositionMeasure{X,NT} <: AbstractMeasure
    components :: NT
end

SuperpositionMeasure(ms :: AbstractMeasure...) where {X} = SuperpositionMeasure{X,length(ms)}(ms)

# SuperpositionMeasure(m::NTuple{N, Measure{X}}) where {N,X} = SuperpositionMeasure(m...)

# Base.length(::SuperpositionMeasure{X,N}) where {X,N} = N

function Base.:+(μ::SuperpositionMeasure{X,N1}, ν::SuperpositionMeasure{X,N2}) where {X,N1,N2}
    components = (μ.components..., ν.components...)
    SuperpositionMeasure{X, N1+N2}(components)
end

function Base.:+(μ::AbstractMeasure, ν::SuperpositionMeasure{X,N}) where {X,N}
    components = (μ, ν.components...)
    SuperpositionMeasure{X,N+1}(components)
end

function Base.:+(μ::SuperpositionMeasure{X,N}, ν::AbstractMeasure) where {X,N}
    components = (μ.components..., ν)
    SuperpositionMeasure{X,N+1}(components)
end

function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure) where {X}
    components = (μ, ν)
    SuperpositionMeasure{X,2}(components)
end

# TODO: Fix `rand` method (this one is wrong)
# function Base.rand(μ::SuperpositionMeasure{X,N}) where {X,N}
#     return rand(rand(μ.components))
# end


module RecurrenceAnalysis

using Distances

export embed,
       distancematrix,
       recurrencematrix,
       crossrecurrencematrix,
       jointrecurrencematrix,
       recurrencerate,
       determinism,
       avgdiag,
       maxdiag,
       divergence,
       entropy,
       trend,
       laminarity,
       trappingtime,
       maxvert,
       autocorrelation,
       ami,
       gmi,
       fnn,
       afnn,
       ffnn,
       sorteddistances

# get distance metric of the Distance packages
function getmetric(normtype::AbstractString)
    normtype = lowercase(normtype)
    metrics = Dict(
        "euclidean"=>Euclidean(),
        "max"=>Chebyshev(),
        "inf"=>Chebyshev()
        )
    !haskey(metrics,normtype) && error("incorrect norm type. Accepted values are \""
        *join(keys(metrics),"\", \"", "\" or \"") * "\".")
    metrics[normtype]
end

"""
    embed(x, m, delay)

Embed a time series in `m` dimensions with a given delay.
"""
function embed(x::AbstractVecOrMat, m::Integer, delay::Integer)
    dims = size(x)
    n = dims[1]
    nm = n-delay*(m-1)
    (nm < 2) && warning("the emedded time series has length < 2")
    ix = (1:nm) .+ (0:delay:delay*(m-1))'
    embed_indices(x, ix)
end

embed(x, m, delay) = isinteger(m) && isinteger(delay) ?
    embed(x, m, delay) : error("embedding dimension and delay must be integer values")

embed_indices(x::AbstractVector, indices) = x[indices]

function embed_indices(x::AbstractMatrix, indices)
    dx = size(x)
    dxm = size(indices)
    ix_rep = repeat(indices, inner=[1,dx[2]])
    ix_rep += repeat(dx[1]*(0:dx[2]-1)', outer=[dxm...])
    x[ix_rep]
end

# Recurrence matrix creation

"""
    distancematrix(x, metric="max")
    distancematrix(x, y, metric="max")

Create a distance matrix from one or two (possibly embedded) time series.

Available metrics for the distances are `"max"` (default), `"inf"` (same), and `"euclidean"`.
"""
function distancematrix(x::AbstractVecOrMat, metric::AbstractString="max")
    dist = getmetric(metric)
    pairwise(dist, x')
end

function distancematrix(x::AbstractVecOrMat, y::AbstractVecOrMat, metric::AbstractString="max")
    dist = getmetric(metric)
    pairwise(dist, x', y')
end

"""
    recurrencematrix(x, radius; <keyword arguments>)

Create a recurrence matrix from an embedded time series.

# Arguments
* `x::Any` : embedded time series.
* `radius::Any` : threshold parameter to classify distances as recurrences.
* `scale::Any` : function of the distance matrix, or fixed number to scale the distances
between points. Typical choices are `maximum` (default) to scale distances into
the unit inteval, or `mean`. Use `1` to keep the distances unscaled.
* `metric::String` : metric of the norm, as in `distancematrix`.
"""
function recurrencematrix(x, radius; scale=maximum, kwargs...)
    kwargs = Dict(kwargs)
    argsdm = haskey(kwargs,:metric) ? (x, kwargs[:metric]) : (x,)
    dm = distancematrix(argsdm...)
    (typeof(scale) <: Function) && (scale = scale(dm))
    dm /= scale
    sparse(dm .< radius)
end

"""
    crossrecurrencematrix(x, y, radius; <keyword arguments>)

Create a cross recurrence matrix from two embeded time series.

See `?recurrencematrix` for details.
"""
function crossrecurrencematrix(x, y, radius; scale=maximum, kwargs...)
    kwargs = Dict(kwargs)
    argsdm = haskey(kwargs,:metric) ? (x, y, kwargs[:metric]) : (x, y)
    dm = distancematrix(argsdm...)
    (typeof(scale) <: Function) && (scale = scale(dm))
    dm /= scale
    sparse(dm .< radius)
end

"""
    crossrecurrencematrix(x, y, radius; <keyword arguments>)

Create a joint recurrence matrix from two embeded time series.

See `?recurrencematrix` for details.
"""
function jointrecurrencematrix(x, y, radius; kwargs...)
    rm1 = recurrencematrix(x, radius, kwargs...)
    rm2 = recurrencematrix(y, radius, kwargs...)
    n = minimum([size(rm1)..., size(rm2)...])
    rm1[1:n, 1:n] .* rm2[1:n, 1:n]
end

include("rqa.jl")
include("delay.jl")
include("dimension.jl")
include("radius.jl")

end


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

"""
wedge.jl defines methods for the wedge(x, y) function
"""

# --- Exports

export wedge, ∧

# --- Operator aliases

∧(M::AbstractMultivector, N::AbstractMultivector) = wedge(M, N)

∧(M::AbstractMultivector, x::Real) = wedge(M, x)
∧(x::Real, M::AbstractMultivector) = wedge(x, M)

∧(M::AbstractMultivector, v::Vector{<:Real}) = wedge(M, v)
∧(v::Vector{<:Real}, M::AbstractMultivector) = wedge(v, M)

∧(v::Vector{<:Real}, w::Vector{<:Real}) = wedge(v, w)

# --- Method definitions

"""
    wedge(M::AbstractMultivector, N::AbstractMultivector)::AbstractMultivector
    ∧(M::AbstractMultivector, N::AbstractMultivector)::AbstractMultivector

Compute the outer product of `M` and `N`.
"""
function wedge end

# ------ Specializations involving an AbstractMultivector instance

# M::AbstractMultivector, B::AbstractBlade
# B::AbstractBlade, M::AbstractMultivector
wedge(M::AbstractMultivector, B::AbstractBlade) =
    Multivector(map(C -> wedge(C, B), blades(M)))

wedge(B::AbstractBlade, M::AbstractMultivector) =
    Multivector(map(C -> wedge(B, C), blades(M)))

# M::AbstractMultivector, v::Vector
# v::Vector, M::AbstractMultivector
wedge(M::AbstractMultivector, v::Vector{<:Real}) =
    Multivector(map(C -> wedge(C, v), blades(M)))

wedge(v::Vector{<:Real}, M::AbstractMultivector) =
    Multivector(map(C -> wedge(v, C), blades(M)))

# M::AbstractMultivector, B::AbstractScalar
# B::AbstractScalar, M::AbstractMultivector
wedge(B::AbstractMultivector, C::AbstractScalar) = B * C
wedge(B::AbstractScalar, C::AbstractMultivector) = B * C

# M::AbstractMultivector, B::One
# B::One, M::AbstractMultivector
wedge(M::AbstractMultivector, B::One) = M
wedge(B::One, M::AbstractMultivector) = M

# M::AbstractMultivector, B::Zero
# B::Zero, M::AbstractMultivector
wedge(B::Zero, M::AbstractMultivector) = B
wedge(M::AbstractMultivector, B::Zero) = B

# M::AbstractMultivector, x::Real
# x::Real, M::AbstractMultivector
wedge(M::AbstractMultivector, x::Real) = x * M
wedge(x::Real, M::AbstractMultivector) = x * M

# ------ Specializations involving an AbstractBlade instance

# B::AbstractBlade, C::AbstractScalar
# B::AbstractScalar, C::AbstractBlade
wedge(B::AbstractBlade, C::AbstractScalar) = B * C
wedge(B::AbstractScalar, C::AbstractBlade) = B * C

# B::AbstractBlade, C::One
# B::One, C::AbstractBlade
wedge(B::AbstractBlade, C::One) = B
wedge(B::One, C::AbstractBlade) = C

# B::AbstractBlade, C::Zero
# B::Zero, C::AbstractBlade
wedge(B::AbstractBlade, C::Zero) = zero(B)
wedge(B::Zero, C::AbstractBlade) = zero(C)

# B::AbstractBlade, x::Real
# x::Real, B::AbstractBlade
wedge(B::AbstractBlade, x::Real) = B * x
wedge(x::Real, B::AbstractBlade) = x * B

# ------ Specializations involving a Blade instance

# B::Blade, C::Blade
function wedge(B::Blade, C::Blade)
    # --- Check arguments

    assert_dim_equal(B, C)

    if grade(B) + grade(C) > dim(B)
        return zero(B)
    end

    # --- Construct new blade

    volume(B) * volume(C) * Blade(hcat(basis(B), basis(C)))
end

# B::Blade, C::Pseudoscalar
# B::Pseudoscalar, C::Blade
function wedge(B::Blade, C::Pseudoscalar)
    assert_dim_equal(B, C)
    zero(B)
end

wedge(B::Pseudoscalar, C::Blade) = wedge(C, B)

# B::Blade, v::Vector
# v::Vector, B::Blade
function wedge(B::Blade, v::Vector{<:Real})
    assert_dim_equal(B, v)

    # Note: volume(B) is incorporated into the norm of `v`
    Blade(hcat(basis(B), volume(B) * v))
end

function wedge(v::Vector{<:Real}, B::Blade)
    assert_dim_equal(v, B)

    # Note: volume(B) is incorporated into the norm of `v`
    Blade(hcat(volume(B) * v, basis(B)))
end

# v::Vector, w::Vector
wedge(v::Vector{<:Real}, w::Vector{<:Real}) = Blade(hcat(v, w))

# ------ Specializations involving a Pseudoscalar instance

# B::Pseudoscalar, C::Pseudoscalar
function wedge(B::Pseudoscalar, C::Pseudoscalar)
    assert_dim_equal(B, C)
    zero(B)
end

# B::Pseudoscalar, v::Vector
# v::Vector, B::Pseudoscalar
function wedge(B::Pseudoscalar, v::Vector{<:Real})
    assert_dim_equal(B, v)
    zero(B)
end

wedge(v::Vector{<:Real}, B::Pseudoscalar) = B ∧ v

# ------ Specializations involving an AbstractScalar instance

# B::AbstractScalar, C::One
# B::One, C::AbstractScalar
wedge(B::AbstractScalar, C::One) = B
wedge(B::One, C::AbstractScalar) = C

# B::AbstractScalar, C::Zero
# B::Zero, C::AbstractScalar
wedge(B::AbstractScalar, C::Zero) = C
wedge(B::Zero, C::AbstractScalar) = B

# B::AbstractScalar, x::Real
# x::Real, B::AbstractScalar
wedge(B::AbstractScalar, x::Real) = B * x
wedge(x::Real, B::AbstractScalar) = x * B

# B::AbstractScalar, v::Vector
# v::Vector, B::AbstractScalar
wedge(B::AbstractScalar, v::Vector{<:Real}) = value(B) * Blade(v)
wedge(v::Vector{<:Real}, B::AbstractScalar) = wedge(B, v)

# ------ Specializations involving a Scalar instance

# B::Scalar, C::Scalar
wedge(B::Scalar, C::Scalar) = B * C

# ------ Specializations involving a One instance

# B::One, C::One
wedge(B::One, C::One) = B

# B::One, C::Zero
# B::Zero, C::One
wedge(B::One, C::Zero) = C
wedge(B::Zero, C::One) = B

# ------ Specializations involving a Zero instance

# B::Zero, C::Zero
wedge(B::Zero, C::Zero) = B


@doc raw""" ```julia
struct ALSolver <: ConstrainedSolver{T}
```
Augmented Lagrangian (AL) is a standard tool for constrained optimization. For a trajectory optimization problem of the form:
```math
\begin{aligned}
  \min_{x_{0:N},u_{0:N-1}} \quad & \ell_f(x_N) + \sum_{k=0}^{N-1} \ell_k(x_k, u_k, dt) \\
  \textrm{s.t.}            \quad & x_{k+1} = f(x_k, u_k), \\
                                 & g_k(x_k,u_k) \leq 0, \\
                                 & h_k(x_k,u_k) = 0.
\end{aligned}
```
AL methods form the following augmented Lagrangian function:
```math
\begin{aligned}
    \ell_f(x_N) + &λ_N^T c_N(x_N) + c_N(x_N)^T I_{\mu_N} c_N(x_N) \\
           & + \sum_{k=0}^{N-1} \ell_k(x_k,u_k,dt) + λ_k^T c_k(x_k,u_k) + c_k(x_k,u_k)^T I_{\mu_k} c_k(x_k,u_k)
\end{aligned}
```
This function is then minimized with respect to the primal variables using any unconstrained minimization solver (e.g. iLQR).
    After a local minima is found, the AL method updates the Lagrange multipliers λ and the penalty terms μ and repeats the unconstrained minimization.
    AL methods have superlinear convergence as long as the penalty term μ is updated each iteration.
"""
struct ALSolver{T,S} <: ConstrainedSolver{T}
    opts::SolverOptions{T}
    stats::SolverStats{T}
    ilqr::S
end

function ALSolver(
        prob::Problem{T},
        opts::SolverOptions=SolverOptions(),
        stats::SolverStats=SolverStats(parent=solvername(ALSolver));
        use_static=usestaticdefault(get_model(prob)[1]),
        kwarg_opts...
    ) where {T}
    set_options!(opts; kwarg_opts...)

    # Build Augmented Lagrangian Objective
    alobj = ALObjective{T}(prob.obj)
    nx,nu = dims(prob)
    prob_al = Problem(prob.model, alobj, ConstraintList(nx,nu),
        prob.x0, prob.xf, prob.Z, prob.N, TO.get_initial_time(prob), TO.get_final_time(prob))

    # Instantiate the iLQR solver
    ilqr = iLQRSolver(prob_al, opts, stats, use_static=use_static)
    initialize!(alobj.conset, prob.constraints, ilqr.Z, ilqr.opts, alobj.alcost, ilqr.Efull)
    # settraj!(alobj.conset, get_trajectory(ilqr))

    # Build the solver
    solver = ALSolver(opts, stats, ilqr)
    reset!(solver)
    return solver
end

# Getters
for method in (:(RD.dims), :(RD.state_dim), :(RD.errstate_dim), :(RD.control_dim),
              :(TO.get_trajectory), :(TO.get_objective), :(TO.get_model),
              :(TO.get_initial_state), :getlogger)
    @eval $method(solver::ALSolver) = $method(solver.ilqr)
end
solvername(::Type{<:ALSolver}) = :AugmentedLagrangian
get_ilqr(solver::ALSolver) = solver.ilqr
TO.get_constraints(solver::ALSolver) = solver.ilqr.obj.conset
stats(solver::ALSolver) = solver.stats
options(solver::ALSolver) = solver.opts

# Methods
function reset!(solver::ALSolver)
    # reset_solver!(solver)
    opts = options(solver)::SolverOptions
    reset!(stats(solver), opts.iterations, solvername(solver))
    reset!(solver.ilqr)

    # Reset constraints
    conset = get_constraints(solver)
    reset!(conset)
end

function max_violation(solver::ALSolver)
    evaluate_constraints!(get_constraints(solver), get_trajectory(solver))
    max_violation(get_constraints(solver))
end


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("../../test/useful_routines.jl")

function time_step_size(levelsetspeed, dx; CFL = 0.5)
    s = maximum(abs.(levelsetspeed))
    return CFL * dx / s
end

function step_interface(
    levelset,
    levelsetspeed,
    dt,
    paddedmesh,
    tol,
    boundingradius,
)
    mesh = CutCellDG.background_mesh(levelset)
    cutmesh = CutCellDG.CutMesh(mesh, levelset)

    refseedpoints, refseedcellids =
        CutCellDG.seed_zero_levelset(2, levelset, cutmesh)
    spatialseedpoints =
        CutCellDG.map_to_spatial(refseedpoints, refseedcellids, mesh)

    paddedlevelset = CutCellDG.BoundaryPaddedLevelSet(
        paddedmesh,
        spatialseedpoints,
        refseedcellids,
        levelset,
        tol,
        boundingradius,
    )

    return CutCellDG.step_first_order_levelset(
        paddedlevelset,
        levelsetspeed,
        dt,
    )
end

function signed_distance_coefficients(levelset, mesh, tol, boundingradius)
    cutmesh = CutCellDG.CutMesh(mesh, levelset)
    refseedpoints, refseedcellids =
        CutCellDG.seed_zero_levelset(2, levelset, cutmesh)
    spatialseedpoints =
        CutCellDG.map_to_spatial(refseedpoints, refseedcellids, mesh)
    nodalcoordinates = CutCellDG.nodal_coordinates(mesh)

    signeddistance = CutCellDG.distance_to_zero_levelset(
        nodalcoordinates,
        spatialseedpoints,
        refseedcellids,
        levelset,
        tol,
        boundingradius,
    )

    return signeddistance
end

function run_time_steps(levelset, levelsetspeed, paddedmesh, dt, numiter)
    mesh = CutCellDG.background_mesh(levelset)
    dx = CutCellDG.element_size(mesh)
    polyorder = CutCellDG.order(levelset)
    tol = 1e-2(minimum(dx))^(polyorder + 1)
    boundingradius = 1.5maximum(dx)

    numnodes = CutCellDG.number_of_nodes(mesh)
    newcoeffs = zeros(numnodes)
    itercoeffs = zeros(numnodes, numiter + 1)
    itercoeffs[:, 1] = CutCellDG.coefficients(levelset)

    for iter = 1:numiter
        newcoeffs = step_interface(
            levelset,
            levelsetspeed,
            dt,
            paddedmesh,
            tol,
            boundingradius,
        )
        CutCellDG.update_coefficients!(levelset, newcoeffs)
        dist = signed_distance_coefficients(levelset, mesh, tol, boundingradius)
        CutCellDG.update_coefficients!(levelset, dist)
        newcoeffs = dist
    end
    return newcoeffs
end

function grid_range(mesh)
    x0 = CutCellDG.reference_corner(mesh)
    w = CutCellDG.mesh_widths(mesh)
    nfmside = CutCellDG.nodes_per_mesh_side(mesh)

    x = range(x0[1], stop = x0[1] + w[1], length = nfmside[1])
    y = range(x0[2], stop = x0[2] + w[2], length = nfmside[2])

    return x, y
end

function error_for_nelmts(nelmts)
    x0 = [0.0, 0.0]
    L, W = 1.0, 1.0
    polyorder = 2
    numqp = required_quadrature_order(polyorder)

    xc = [0.5, 0.5]
    radius = 0.45
    speed = -1.0
    stoptime = 0.25

    basis = LagrangeTensorProductBasis(2, polyorder)
    dim,nf = size(interpolation_points(basis))
    quad = tensor_product_quadrature(2, numqp)
    mesh = CutCellDG.CGMesh(x0, [L, W], [nelmts, nelmts], nf)
    paddedmesh = CutCellDG.BoundaryPaddedMesh(mesh, 1)
    numnodes = CutCellDG.number_of_nodes(mesh)
    dx = minimum(CutCellDG.grid_size(paddedmesh))

    levelset = CutCellDG.LevelSet(
        x -> -circle_distance_function(x, xc, radius),
        mesh,
        basis,
    )
    levelsetspeed = speed * ones(numnodes)
    dt = time_step_size(levelsetspeed, dx)
    @assert isinteger(stoptime / dt)
    numiter = round(Int, stoptime / dt)
    itercoeffs =
        run_time_steps(levelset, levelsetspeed, paddedmesh, dt, numiter)

    finalradius = radius - stoptime
    err = uniform_mesh_L2_error(
        itercoeffs',
        x -> -circle_distance_function(x, xc, finalradius)[1],
        basis,
        quad,
        mesh,
    )[1]

    normalizer = integral_norm_on_uniform_mesh(
        x -> -circle_distance_function(x, xc, finalradius)[1],
        quad,
        mesh,
        1,
    )[1]
    normalizederr = err / normalizer

    return normalizederr
end

powers = 2:5
nelmts = 2 .^powers
err = error_for_nelmts.(nelmts)
dx = 1.0 ./ nelmts
rate = convergence_rate(dx,err)

@test all(rate .> 0.9)


# A collection of extensions to the DomainSets package.

using DomainSets: inverse_map, forward_map, MappedDomain

##########################################################################
### Distances and Normals
##########################################################################
export normal, distance
# Function that returns the normal vector when evaluated at the domain boundary
# This vector is of unit length and points to the OUTSIDE.
function normal(d::Domain, x)
    if hascanonicaldomain(d)
        normal(convert(MappedDomain, d), x)
    else
        error("Normal not available for this domain type")
    end
end

# Auxiliary function that returns distance from the boundary in some metric
function distance(x, d::Domain)
    if hascanonicaldomain(d)
        distance(x, convert(MappedDomain, d))
    else
        error("Domain distance not available for this domain type")
    end
end

distance(x, ::UnitSimplex) = min(minimum(x),1-sum(x))

distance(x, d::AbstractInterval) = min(abs(maximum(d)-x),abs(x-minimum(d)))

distance(x, d::UnitBall) = 1-norm(x)

distance(x, d::UnitSphere) = 1-norm(x)

distance(x,d::UnionDomain) = indomain(x,d) ? sum(map(di->max(0,distance(x,di)),components(d))) : maximum(map(di->distance(x,di),components(d)))

normal(d::UnionDomain, x) = normal(components(d)[findmin(map(di->abs(distance(x,di)),components(d)))[2]], x)

distance(x,d::IntersectDomain) = minimum(map(di->distance(x,di),components(d)))

normal(d::IntersectDomain, x) = normal(components(d)[findmin(map(di->distance(x,di),components(d)))[2]], x)

distance(x,d::SetdiffDomain) = indomain(x,d) ? min(abs(distance(x,d.domains[1])),abs(distance(x,d.domains[2]))) : -1*min(abs(distance(x,d.domains[1])),abs(distance(x,d.domains[2])))

normal(d::SetdiffDomain, x) = abs(distance(x,d.domains[1]))<abs(distance(x,d.domains[2])) ? normal(d.domains[1], x) : -1*normal(d.domains[2], x)

distance(x, t::ProductDomain) = minimum(map(distance, x, components(t)))

distance(x, d::MappedDomain) = distance(inverse_map(d)(x),superdomain(d))

function normal(d::MappedDomain, x)
    x = applymap(forward_map(d),normal(superdomain(d),inverse_map(d)(x)))
    x0 = inverse(inverse_map(d))(zeros(size(x)))
   (x-x0)/norm(x-x0)
end
function normal(t::ProductDomain, x)
    index = findmin(map(distance, x, components(t)))[2]
    [(i==index)*normal(component(t,i),x[i]) for i =1:length(components(t))]
end


"""
    ARCHLMTest <: HypothesisTest
Engle's (1982) LM test for autoregressive conditional heteroskedasticity.
"""
struct ARCHLMTest{T<:Real} <: HypothesisTest
    n::Int         # number of observations
    p::Int         # number of lags
    LM::T    # test statistic
end

"""
    ARCHLMTest(am::UnivariateARCHModel, p=max(o, p, q, ...))
Conduct Engle's (1982) LM test for autoregressive conditional heteroskedasticity with
p lags in the test regression.
"""
ARCHLMTest(am::UnivariateARCHModel, p=presample(typeof(am.spec))) = ARCHLMTest(residuals(am), p)

"""
    ARCHLMTest(u::Vector, p::Integer)
Conduct Engle's (1982) LM test for autoregressive conditional heteroskedasticity with
p lags in the test regression.
"""
function ARCHLMTest(u::Vector{T}, p::Integer) where T<:Real
    @assert p>0
    n = length(u)
    u2 = u.^2
    X = zeros(T, (n-p, p+1))
    X[:, 1] .= one(eltype(u))
    for i in 1:p
        X[:, i+1] = u2[p-i+1:n-i]
    end
    y = u2[p+1:n]
    B = X \ y
    e = y - X*B
    ybar = y .- mean(y)
    LM = n * (1 - (e'e)/(ybar'ybar)) #T*R^2
    ARCHLMTest(n, p, LM)
end

testname(::ARCHLMTest) = "ARCH LM test for conditional heteroskedasticity"
population_param_of_interest(x::ARCHLMTest) = ("T⋅R² in auxiliary regression", 0, x.LM)
function show_params(io::IO, x::ARCHLMTest, ident)
    println(io, ident, "sample size:                    ", x.n)
    println(io, ident, "number of lags:                 ", x.p)
    println(io, ident, "LM statistic:                   ", x.LM)
end

pvalue(x::ARCHLMTest) = pvalue(Chisq(x.p), x.LM; tail=:right)

"""
    DQTest <: HypothesisTest
Engle and Manganelli's (2004) out-of-sample dynamic quantile test.
"""
struct DQTest{T<:Real} <: HypothesisTest
    n::Int         # number of observations
    p::Int         # number of lags
    level::T       # VaR level
    DQ::T          # test statistic
end

"""
    DQTest(data, vars, level, p=1)
Conduct Engle and Manganelli's (2004) out-of-sample dynamic quantile test with
p lags in the test regression. `vars` shoud be a vector of out-of-sample Value at Risk
predictions at level `level`.
"""
function DQTest(data::Vector{T}, vars::Vector{T}, level::AbstractFloat, p::Integer=1) where T<:Real
    @assert p>0
    @assert length(data) == length(vars)
    n = length(data)
    hit = (data .< -vars).*1 .- level
    y = hit[p+1:n]
    X = zeros(T, (n-p, p+2))
    X[:, 1] .= one(T)
    for i in 1:p
        X[:, i+1] = hit[p-i+1:n-i]
    end
    X[:, p+2] = vars[p+1:n]
    B = X \ y
    DQ = B' * (X'*X) *B/(level*(1-level)) # y'X * inv(X'X) * X'y / (level*(1-level)); note 2 typos in the paper
    DQTest(n, p, level, DQ)
end

testname(::DQTest) = "Engle and Manganelli's (2004) DQ test (out of sample)"
population_param_of_interest(x::DQTest) = ("Wald statistic in auxiliary regression", 0, x.DQ)
function show_params(io::IO, x::DQTest, ident)
    println(io, ident, "sample size:                    ", x.n)
    println(io, ident, "number of lags:                 ", x.p)
    println(io, ident, "VaR level:                      ", x.level)
    println(io, ident, "DQ statistic:                   ", x.DQ)
end

pvalue(x::DQTest) = pvalue(Chisq(x.p+2), x.DQ; tail=:right)


export dcs_lp, charikar_lp, dcs_lp_extract, extract_primal_solution, charikar_JuMP_LP
export LP_METHODS

using JuMP
using Gurobi
using Convex
LP_METHODS=["DCS"; "MIN"; "AVG"]

function dcs_convert{T}(G::Array{SparseMatrixCSC{T, Int64}, 1})
  m = length(G)
  H = SparseMatrixCSC{Float64, Int64}[]
  for j=1:m
    push!(H, similar(G[j], Float64))
    ii, jj = findn(G[j])
    for k=1:length(ii)
       H[j][ii[k], jj[k] ]  = 1.0
    end
  end
  return H
end

function old_dcs_lp{T}(H::Array{SparseMatrixCSC{T, Int64}, 1})
  G = dcs_convert(H)
  m = length(G)
  n = size(G[1], 1)
  X = Convex.Variable[]
  for j=1:m
    push!(X, Variable(n, n,Positive()))
  end
  Y = Variable(n, Positive())
  t = Variable()
  z = sum(Y)
  constraints = Convex.Constraint[]
  push!(constraints, z <= 1.0)

  for j=1:m
    push!(constraints, trace(G[j] * X[j]) >= t)
    ii, jj = findn(G[j])
    for k=1:length(ii)
      crow = ii[k]
      ccol = jj[k]
      push!(constraints, X[j][crow, ccol] <= Y[crow])
      push!(constraints, X[j][crow, ccol] <= Y[ccol])
    end
  end

  problem  = maximize(t, constraints);
  solve!(problem);
  # println("Optval: ", problem.optval)
  return problem, X, Y, t
end

function dcs_lp{T}(G::Array{SparseMatrixCSC{T, Int64}, 1})

  m = length(G)
  n = size(G[1], 1)
  nz = Int64[]
  @inbounds begin
    iv = {}
    jv = {}
    gv = {}
    for j=1:m
      ii, jj= findn(G[j])
      push!(iv, ii)
      push!(jv, jj)
      push!(nz, length(ii))
      ss = similar(ii, Float64)
      for k=1:length(ii)
        ss[k] = G[j][ii[k], jj[k] ]
      end
      push!(gv, ss)
    end
    ni = deepcopy(nz)
    insert!(ni, 1, 1)
    ni = cumsum(ni)

    nv = sum(nz)
    X = Convex.Variable(nv, 1 , Positive())

    Y = Variable(n, Positive())
    t = Variable()
    w = sum(Y)
    constraints = Convex.Constraint[]
    push!(constraints, w <= 1.0)

    for j=1:m
      cb = ni[j]
      ce = ni[j+1]-1
      # println("gv: $(length(gv[j])) ni: [$cb:$ce]")
      # @assert length(gv[j]) == length([cb:ce;])
      # @assert cb >= 1
      # @assert ce <= length(X)
      push!(constraints, sum(X[cb:ce] .* gv[j] ) >= t)
      # @assert length(gv[j]) == nz[j]
      for k=1:nz[j]
        crow = iv[j][k]
        ccol = jv[j][k]
        ci = ni[j] - 1 + k
        # @assert (ci >= cb) && (ci <= ce)
        push!(constraints, X[ci] <= Y[crow])
        push!(constraints, X[ci] <= Y[ccol])
        # println("j: $j X[$ci] <= Y[$crow], Y[$ccol]")
      end
    end
  end
  problem  = maximize(t, constraints);
  solve!(problem);

  Z = SparseMatrixCSC{Float64, Int64}[]
  typeof(Z)
  for j=1:m
    push!(Z, spzeros(Float64, n, n))
  end
  for j=1:m
    for k=1:nz[j]
      crow = iv[j][k]
      ccol = jv[j][k]
      ci = ni[j] - 1 + k
      # @assert crow >= 1 && crow <= n
      # @assert ccol >= 1 && ccol <= n
      # @assert G[j][crow, ccol] != 0
      # println(X.value[ci])
      Z[j][crow, ccol] = X.value[ci]
    end
  end

  # println("Optval: ", problem.optval)
  return problem, Z, Y.value, t.value
end

function faster_dcs_lp{T}(G::Array{SparseMatrixCSC{T, Int64}, 1})
  m = length(G)
  n = size(G[1], 1)
  nz = Int64[]
  @inbounds begin
    for j=1:m
      cval = nnz(G[j])
      push!(nz, cval)
    end

    ni = deepcopy(nz)
    insert!(ni, 1, 1)
    ni = cumsum(ni)

    X = Convex.Variable(sum(nz), 1 , Positive())
    Y = Variable(n, Positive())
    t = Variable()
    z = sum(Y)
    constraints = Convex.Constraint[]
    push!(constraints, z <= 1.0)

    for j=1:m
      cb = ni[j]
      ce = (ni[j+1]-1)
      gi = find(G[j])
      gval = map((x) -> convert(Float64, x), G[j][gi])
      push!(constraints, sum(X[cb:ce] .* gval) >= t)
      for k=1:nz[j]
        ccol = iceil(gi[k]/n)
        crow = (gi[k] % n)
        if crow == 0
          crow = n
        end
        ci = ni[j] - 1 + k

        push!(constraints, X[ci] <= Y[crow])
        push!(constraints, X[ci] <= Y[ccol])
        # println("j: $j X[$ci] <= Y[$crow], Y[$ccol]")
      end
    end
  end
  problem  = maximize(t, constraints);
  solve!(problem);
  return problem, X, Y, t
end

# function DCSP(G::Graph[])
#   m = length(G)
#   n = size(G[1], 1)
#   R = Model(solver=GurobiSolver(OutputFlag=1))
#   @defVar(R, Y[1:n] >= 0)
#   for i=1:m
#     [gi, gj] = findn(G[i])
#     ne = length(gi)
#     @defVar(R, X[i=i, j=1:ne] >= 0)

# end

function charikar_JuMP_LP(G::Graph)
  n = size(G, 1)
  tic()
  ii, jj = findn(G)
  ne = length(ii)
  println("graph with $n nodes and $ne edges. Setting up Charikar LP ... ")
  R = Model(solver=GurobiSolver(OutputFlag=0))
  @defVar(R, X[1:ne] >= 0)
  @defVar(R, Y[1:n] >= 0)
  @addConstraint(R, gamma, sum{Y[i], i=1:n} <= 1)

  @addConstraint(R, alpha[k=1:ne], X[k] <= Y[ ii[k] ] )
  @addConstraint(R, beta[k=1:ne], X[k] <= Y[ jj[k] ] )
#   for k=1:ne
#     crow = ii[k]
#     ccol = jj[k]
#     @addConstraint(R, X[crow, ccol] <= Y[crow])
#     @addConstraint(R, X[crow, ccol] <= Y[ccol])
#   end
  @setObjective(R, Max, sum{X[k], k=1:ne })
  t_elapsed = toc();
  println("Completed setup, starting to solve LP with Gurobi ")
  status = solve(R);
  if status == :Optimal
    println("solved LP successfully, returning primal variables")
  end

  t = R.objVal
  X1 = spzeros(n, n)
  Y1 = spzeros(n, 1)
  for k=1:ne
    X1[ii[k], jj[k]]= getValue(X[k])
  end
  for j=1:n
    Y1[j] = getValue(Y[j])
  end
  return X1, Y1, t, status
end

function charikar_lp{T}(G::SparseMatrixCSC{T, Int64})

  set_default_solver(GurobiSolver())
  n = size(G, 1)
  X = Variable(n, n, Positive())
  Y = Variable(n, Positive())
  constraints = Convex.Constraint[]

  push!(constraints, sum(Y) <= 1)

  ii, jj = findn(G)
  for k=1:length(ii)
    crow = ii[k]
    ccol = jj[k]
    push!(constraints, X[crow, ccol] <= Y[crow])
    push!(constraints, X[crow, ccol] <= Y[ccol])
  end

  problem = maximize(trace(G * X), constraints);
  solve!(problem);
  t= problem.optval
  return problem, X.value, Y.value, t
end

function extract_primal_solution(Y)
  V = [j > 0.0 for j in Y]
  S = find(V)
  return S
end

# Uses LP primal solution to find dense common subgraph
#
#
function dcs_lp_extract(G::GraphVec; method::String = "DCS" )
  m = length(G)
  @assert m >= 1
  n = size(G{1}, 1)

  if method == "DCS"
    p, X, Y, t = dcs_lp(G)
  elseif method == "MIN"
    H = dcs_intersection_graph(G)
    p, X, Y, t =  charikar_lp(H)
  elseif method == "AVG"
    H = dcs_avg_graph(G)
    p, X, Y, t = charikar_lp(H)
  else
    error("Undefined method: $method") #  ('dcs', 'charikar', 'weighted')"
  end
  S = extract_primal_solution(Y)
  d = common_subgraph_density(S)
  return (S, d, t)
end

# function dcs_charikar_lp{T}(H::Array{SparseMatrixCSC{T, Int64}, 1})
#   m = length(H)
#   n = size(H[1], 1)
#   G = spzeros(T, n, n)

#   ii, jj = findn(H[1])
#   for k=1:length(ii)
#     cr = ii[k]
#     cc = jj[k]
#     G[cr, cc] = H[1][cr, cc]
#     for j=2:m
#       G[cr, cc] = min(G[cr, cc], H[j][cr, cc])
#     end
#   end
#   return charikar_lp(G)

  # end


mutable struct ClearDistances
    cardinal::SVector{4, Int}
    diagonal::SVector{4, Int}
end

function find_distances(f::Floor, obstacles::Set{Coord}, s::LTState)
    card = MVector{4, Int}(undef)
    diag = MVector{4, Int}(undef)
    for i in 1:4
        d = 1
        while !opaque(f, obstacles, s, s.robot+d*CARDINALS[i])
            d += 1
        end
        card[i] = d-1
    end
    for i in 1:4
        d = 1
        while !opaque(f, obstacles, s, s.robot+d*DIAGONALS[i])
            d += 1
        end
        diag[i] = d-1
    end
    return ClearDistances(card, diag)
end

mutable struct LTDistanceCache
    floor::Floor
    distances::Vector{ClearDistances}
end

function LTDistanceCache(f::Floor, obstacles::Set{Coord})
    dists = Array{ClearDistances}(undef,n_pos(f)^2)
    visited = falses(n_pos(f)^2)
    for i in 1:f.n_cols, j in 1:f.n_rows, k in 1:f.n_cols, l in 1:f.n_rows
        s = LTState(Coord(i,j), Coord(k,l), false)
        ii = stateindex(f, s)
        visited[ii] = true
        dists[ii] = find_distances(f, obstacles, s)
    end
    @assert all(visited)
    push!(dists, ClearDistances(zeros(4), zeros(4)))
    return LTDistanceCache(f, dists)
end

Base.getindex(c::LTDistanceCache, s::LTState) = c.distances[stateindex(c.floor, s)]

function n_clear_cells(d::ClearDistances, dir::Int)
    if dir <= 4
        return d.cardinal[dir]
    else
        return d.diagonal[dir-4]
    end
end

struct ReadingCDF
    cardcdf::Matrix{Float64}
    diagcdf::Matrix{Float64}
end

# reading CDF
function ReadingCDF(f::Floor,
                    std::Float64,
                    shortonly::Bool=false,
                    maxread::Int=ceil(Int, max_diag(f)+4*std))
    maxclear = max(f.n_rows, f.n_cols) - 1
    cardcdf = Array{Float64}(undef, maxclear + 1, maxread + 1)
    diagcdf = Array{Float64}(undef, maxclear + 1, maxread + 1)

    for c in 0:maxclear
        for r in 0:maxread
            cardcdf[c+1, r+1] = (1+erf((r+1.0-c)/(std*sqrt(2))))/2
        end
    end

    for c in 0:maxclear
        for r in 0:maxread
            diagcdf[c+1, r+1] = (1+erf((r+1.0-sqrt(2)*c)/(std*sqrt(2))))/2
        end
    end

    for c in 0:maxclear
        @assert abs(sum(diff(cardcdf[c, :])) - 1.0) < 1e-5
        @assert abs(sum(diff(diagcdf[c, :])) - 1.0) < 1e-5
    end

    return ReadingCDF(cardcdf, diagcdf)
end

function cdf(c::ReadingCDF, dir::Int, clear::Int, reading::Int)
    if dir <= 4 # cardinal
        if reading > size(c.cardcdf, 2)
            return 1.0
        else
            return c.cardcdf[clear + 1, reading + 1]
        end
    else # diagonal
        if reading > size(c.diagcdf, 2)
            return 1.0
        else
            return c.diagcdf[clear + 1, reading + 1]
        end
    end
end


using JLD
using SharedArrays
using Distributed
include("../examples/solenoid.jl")
function obj_fun(u)
	return u[1]^2 + u[2]^2
end
function obj_fun_erg_avg(s, nSteps)
	nRunup = 500
	u = rand(3)
	J = 0.
	for i = 1:nRunup
		u = next(u,s)
	end
	for i = 1:nSteps
		J += obj_fun(u)/nSteps
		u = next(u,s)
	end
	return J
end
function get_JavgN(sind)
	s = [1.0, 4.0, sind]
	n_rep = 160
	n_pts = 8
	J = zeros(n_pts)
	J_proc = SharedArray{Float64}(n_rep)
	J_proc .= 0.
	Ni = 1200
	N = zeros(Int64, n_pts)
	for i = 1:n_pts
		N[i] = Ni*n_rep
		@show Ni
		t = @distributed for n=1:n_rep
			J_proc[n] = obj_fun_erg_avg(s,Ni)/n_rep
		end
		wait(t)
		J[i] = sum(J_proc)
		Ni = Ni*5
	end
	save("../data/finiteDifference/solenoid/r2_$(sind).jld","s", sind,"J", J, "N", N)
end


# Noncentral hypergeometric distribution
# TODO: this distribution needs clean-up and testing

abstract NoncentralHypergeometric <: DiscreteUnivariateDistribution

### handling support

function insupport(d::NoncentralHypergeometric, x::Real)
    isinteger(x) && minimum(d) <= x <= maximum(d)
end

@distr_support NoncentralHypergeometric max(d.n - d.nf, 0) min(d.ns, d.n)

# Functions

function quantile(d::NoncentralHypergeometric, q::Float64)
    if !(zero(q) <= q <= one(q))
        NaN
    else
        range = support(d)
        if q > 0.5
            q = 1-q
            range = reverse(range)
        end

        qsum, i = 0.0, 0
        while qsum < q
            i += 1
            qsum += pdf(d, range[i])
        end
        range[i]
    end
end

params(d::NoncentralHypergeometric) = (d.ns, d.nf, d.n, d.ω)

## Fisher's noncentral hypergeometric distribution

immutable FisherNoncentralHypergeometric <: NoncentralHypergeometric
    ns::Int    # number of successes in population
    nf::Int    # number of failures in population
    n::Int     # sample size
    ω::Float64 # odds ratio

    function FisherNoncentralHypergeometric(ns::Real, nf::Real, n::Real, ω::Float64)
        @check_args(FisherNoncentralHypergeometric, ns >= zero(ns) && nf >= zero(nf))
        @check_args(FisherNoncentralHypergeometric, zero(n) < n < ns + nf)
        @check_args(FisherNoncentralHypergeometric, ω > zero(ω))
        new(ns, nf, n, ω)
    end
end

# Properties
binomial{T<:Integer}(n::T, k::T) = Base.binomial(BigInt(n), BigInt(k))
binomial{T<:Integer}(n::T, kR::OrdinalRange{T}) = [binomial(n,k) for k in kR]
function _P(d::FisherNoncentralHypergeometric, k::Int)
    y = support(d)
    p = binomial(d.ns, y) .* binomial(d.nf, d.n-y) .* d.ω.^y .* y.^k
    sum(p)
end

function _mode(d::FisherNoncentralHypergeometric)
    A = d.ω - 1
    B = d.n - d.nf - (d.ns + d.n + 2)*d.ω
    C = (d.ns + 1)*(d.n + 1)*d.ω
    -2*C / (B - sqrt(B^2-4*A*C))
end

mean(d::FisherNoncentralHypergeometric) =_P(d,1) / _P(d,0)
var(d::FisherNoncentralHypergeometric) = _P(d,2)/_P(d,0) - (_P(d,1) / _P(d,0))^2
mode(d::FisherNoncentralHypergeometric) = floor(Int, _mode(d))

logpdf(d::FisherNoncentralHypergeometric, k::Int) =
    Float64(log(binomial(d.ns, k)) + log(binomial(d.nf, d.n-k)) + k*log(d.ω) - log(_P(d,0)))

pdf(d::FisherNoncentralHypergeometric, k::Int) = exp(logpdf(d, k))

## Wallenius' noncentral hypergeometric distribution

immutable WalleniusNoncentralHypergeometric <: NoncentralHypergeometric
    ns::Int    # number of successes in population
    nf::Int    # number of failures in population
    n::Int     # sample size
    ω::Float64 # odds ratio
    function WalleniusNoncentralHypergeometric(ns::Real, nf::Real, n::Real, ω::Float64)
        @check_args(WalleniusNoncentralHypergeometric, ns >= zero(ns) && nf >= zero(nf))
        @check_args(WalleniusNoncentralHypergeometric, zero(n) < n < ns + nf)
        @check_args(WalleniusNoncentralHypergeometric, ω > zero(ω))
        new(ns, nf, n, ω)
    end
end

# Properties
mean(d::WalleniusNoncentralHypergeometric)=sum(support(d).*pdf(d,support(d)))
var(d::WalleniusNoncentralHypergeometric) = sum((support(d)-mean(d)).^2.*pdf(d,support(d)))
mode(d::WalleniusNoncentralHypergeometric) = support(d)[indmax(pdf(d,support(d)))]

function pdf(d::WalleniusNoncentralHypergeometric, k::Int)
    D = d.ω*(d.ns-k)+(d.nf-d.n+k)
    f(t) = (1-t^(d.ω/D))^k * (1-t^(1/D))^(d.n-k)
    I,_ = quadgk(f,0,1)
    Float64(binomial(d.ns,k)*binomial(d.nf,d.n-k)*I)
end

logpdf(d::WalleniusNoncentralHypergeometric, k::Int) = log(pdf(d, k))


@testset "cl" begin
    for n in vcat(collect(1:16), [1000, 1001, 1_000_000])
        data = open(readdlm, joinpath(@__DIR__, "data", "Cl$(n).txt"))

        for r in 1:size(data, 1)
            row      = data[r, :]
            x        = row[1]
            expected = row[2]
            sgn      = -(-1)^n

            cl  = ClausenFunctions.cl(n, x)
            clm = sgn*ClausenFunctions.cl(n, -x)

            @test cl == clm
            @test cl ≈ expected atol=1e-14
        end

    end

    @test_throws DomainError ClausenFunctions.cl(0, 1.0)
    @test_throws DomainError ClausenFunctions.cl(-1, 1.0)
    @test_throws DomainError ClausenFunctions.cl(-2, 1.0)
end


# check the Bernoulli numbers and polynomials
#    this is mostly just a check that code is correct
#    most errors are arising because of errors in the functions being used
using Polylogarithms
using PyPlot
using SpecialFunctions
import Base: parse
using Printf

Q = Polylogarithms.Q
c1 = Polylogarithms.c
c2 = Polylogarithms.c_crandall
b = Polylogarithms.b_crandall
f = Polylogarithms.f_crandall
g = Polylogarithms.g_crandall

# check f
@printf("Testing f\n")
for k=0:5
    @printf(" k = %d\n", k)
    @printf("   f(k,1) = %f, -H_k= %f\n", f(k,1), -harmonic(k) )
    @printf("   f(k,2) = %f, (H_k^2 + H_{k,2})/2= %f\n", f(k,2), (harmonic(k)^2 + harmonic(k,2))/2 )
end
@printf("\n")

# # check g against numerical derivates of Gamma at 1
@printf("Testing g\n")
function G(t::Integer, x::Float64; dx::Float64 = 1.0e-3)
    if t==0
        return 1
    elseif t==1
        return ( SpecialFunctions.gamma(x+dx) - SpecialFunctions.gamma(x) )/dx
    elseif t>1
        return ( G(t-1,x+dx) - G(t-1,x) ) /dx
    else
        throw(DomainError(t))
    end
end
for t=0:5
    @printf(" t = %d, ", t)
    @printf(" G^{t}(1) = %.7f,  g(t) = %.7f\n",  G(t,1.0), g(t) )
end
@printf("\n")

# check c
z = complex(-1.0/2.0)
μ = log(z)
ℒ = log(-μ)
for k=0:5
    @printf(" k = %d\n", k)
    for j=0:2
        @printf("  j = %d,  ", j)
        @printf("  c_explicit = %.12f + %.12fi, ", real(c1(k, j, ℒ)), imag(c1(k, j, ℒ)) )
        @printf("  c_crandall = %.12f + %.12fi, ", real(c2(k, j, ℒ)), imag(c2(k, j, ℒ)) )
        @printf("\n")
   end
   @printf("\n")
end

c2(1, 3, ℒ)


using BatchedRoutines, Test

@testset "Testing batched_scal! with $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
    A = rand(elty, 100)
    B = rand(elty, 10, 10, 100)
    expect_B = copy(B)
    batched_scal!(A, B)

    for k in 1:100
        expect_B[:, :, k] .= A[k] * expect_B[:, :, k]
    end

    @test expect_B ≈ B
end

@testset "Testing batched_scal with $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
    A = rand(elty, 100)
    B = rand(elty, 10, 10, 100)
    expect_B = copy(B)
    test_B = batched_scal(A, B)

    for k in 1:100
        expect_B[:, :, k] .= A[k] * expect_B[:, :, k]
    end

    @test expect_B ≈ test_B
end

@testset "Testing batched_gemm! with $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]

    alpha, beta = rand(elty), rand(elty)
    A = rand(elty, 10, 10, 100)
    B = rand(elty, 10, 10, 100)
    C = rand(elty, 10, 10, 100)
    test_C = copy(C)

    batched_gemm!('N', 'N', alpha, A, B, beta, C)

    for k in 1:100
        test_C[:, :, k] = alpha * (A[:, :, k]) * B[:, :, k] + beta * test_C[:, :, k]
    end

    @test C ≈ test_C
end

@testset "Testing batched_gemm with $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
    alpha = rand(elty)
    A = rand(elty, 10, 10, 100)
    B = rand(elty, 10, 10, 100)

    C = batched_gemm('N', 'N', alpha, A, B)

    test_C = zeros(elty, 10, 10, 100)
    for k in 1:100
        test_C[:, :, k] = alpha * (A[:, :, k]) * B[:, :, k]
    end

    @test test_C ≈ C
end


using Revise, StableRNGs, Pathfinder, Distributions, ForwardDiff, Plots, Optim, LineSearches, NonconvexIpopt, Random

function run_pathfinder(; nruns = 10, ndraws = 1000_000, ndraws_per_run = ndraws ÷ nruns, ndists = 2, rng = StableRNG(123), means = 100 * randn(rng, ndists), initx = [1000.0], power = 3.0, deflation_ub = 1000, deflation_radius = 0.0, plot_title = "")
    function mixturef(means, stds = ones(length(means)))
        return MixtureModel(Normal.(means, stds))
    end
    mixture = mixturef(means)
    logπ = x -> logpdf(mixture, x[1])
    ∇logπ = x -> ForwardDiff.gradient(logπ, x)
    plt = Plots.plot(logπ, minimum(means) - 5, maximum(means) + 5, label = "target pdf", legend = legend = :outertopright, title = plot_title)
    @show means

    # x₀s = [30 * randn(rng, length(initx)) for _ in 1:nruns]
    x₀s = [initx for _ in 1:nruns]

    # Deflation
    optimizer1 = IpoptAlg()
    q1, ϕ1, component_ids1 = multipathfinder(
        logπ, ∇logπ, x₀s, ndraws; ndraws_elbo=100, ndraws_per_run=ndraws_per_run, optimizer=optimizer1,
        rng, power, deflation_ub, deflation_radius,
    )

    logπ2 = x -> logpdf(q1, [x])
    Plots.plot!(plt, logπ2, minimum(means) - 5, maximum(means) + 5, label = "approx")

    return q1, ϕ1, component_ids1, plt
end

RESULT_DIR = joinpath(@__DIR__, "results");
for ndists in [2, 4, 6, 8]
    nruns = 100
    ndraws = 1000
    rng = StableRNG(1234)
    means = 20 * randn(rng, ndists)
    initx = [0.0]
    deflation_radius = 0.0

    q1, ϕ1, component_ids1, plt = run_pathfinder(; nruns, ndraws, ndists, rng, means, initx, deflation_radius)
    savefig(plt, joinpath(RESULT_DIR, "pathfinder_$ndists.pdf"))
end

# # No deflation

# optimizer2 = Optim.LBFGS(; m=6, linesearch=LineSearches.MoreThuente())
# q2, ϕ2, component_ids1 = multipathfinder(
#     logπ, ∇logπ, x₀s, ndraws; ndraws_elbo=100, ndraws_per_run=ndraws_per_run, optimizer=optimizer2,
#     rng=rng,
# )
# logπ3 = x -> logpdf(q2, [x])
# Plots.plot!(logπ3, minimum(means) - 5, maximum(means) + 5)


import LinearAlgebra.BLAS: @blasfunc
import LinearAlgebra: BlasInt
import LinearAlgebra.LAPACK: stegr!
const liblapack = Base.liblapack_name

mutable struct StegrWork{T<:Real}
    jobz::Char
    range::Char
    dv::Vector{T}
    ev::Vector{T}
    vl::Real
    vu::Real
    il::BlasInt
    iu::BlasInt
    abstol::Vector{T}
    m::Vector{BlasInt}
    w::Vector{T}
    Z::Matrix{T}
    isuppz::Vector{BlasInt}
    work::Vector{T}
    lwork::BlasInt
    iwork::Vector{BlasInt}
    liwork::BlasInt
    info::Vector{BlasInt}
end

for (stegr,elty) in ((:dstegr_,:Float64),
                     (:sstegr_,:Float32))
    @eval begin
        function stegr!(n::BlasInt, sw::StegrWork{$elty})
            ldz = stride(sw.Z, 2)
            ccall((@blasfunc($stegr), liblapack), Cvoid,
                  (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ptr{$elty},
                   Ptr{$elty}, Ref{$elty}, Ref{$elty}, Ref{BlasInt},
                   Ref{BlasInt}, Ptr{$elty}, Ptr{BlasInt}, Ptr{$elty},
                   Ptr{$elty}, Ref{BlasInt}, Ptr{BlasInt}, Ptr{$elty},
                   Ref{BlasInt}, Ptr{BlasInt}, Ref{BlasInt}, Ptr{BlasInt}),
                  sw.jobz, sw.range, n,
                  sw.dv, sw.ev,
                  sw.vl, sw.vu, sw.il, sw.iu,
                  sw.abstol, sw.m,
                  sw.w, sw.Z, ldz,
                  sw.isuppz, sw.work, sw.lwork, sw.iwork, sw.liwork,
                  sw.info)
        end
    end
end

"""
    stegr!(α, β, sw)

Diagonalize the real-symmetric tridiagonal matrix with `α` on the
diagonal and `β` on the super-/subdiagonal, using the workspaces
allocated in `sw`.
"""
function stegr!(α::AbstractVector{T}, β::AbstractVector{T}, sw::StegrWork{T}) where T
    # @assert length(sw.dv) >= length(α)
    # @assert length(sw.ev) >= length(β)
    copyto!(sw.dv, α)
    copyto!(sw.ev, β)
    stegr!(BlasInt(length(α)), sw)
end

"""
    StegrWork(T, n)

Allocate work arrays for diagonalization of real-symmetric tridiagonal
matrices of sizes up to `n`×`n`.
"""
function StegrWork(::Type{T}, n::BlasInt,
                   jobz::Char = 'V', range::Char = 'A') where T
    dv = Array{T}(undef, n)
    ev = Array{T}(undef, n)
    abstol = Array{T}(undef, 1)
    m = Vector{BlasInt}(undef, 1)
    w = Array{T}(undef, n)
    ldz = jobz == 'N' ? 1 : n
    Z = Array{T}(undef, ldz, n)
    isuppz = Array{BlasInt}(undef, 2n)
    work = Array{T}(undef, 1)
    lwork = -one(BlasInt)
    iwork = Array{BlasInt}(undef, 1)
    liwork = -one(BlasInt)
    info = Array{BlasInt}(undef, 1)
    sw = StegrWork(jobz, range,
                   dv, ev,
                   0.0, 0.0,
                   BlasInt(0), BlasInt(0),
                   abstol, m,
                   w, Z,
                   isuppz,
                   work, lwork,
                   iwork, liwork,
                   info)
    # Call stegr! once to query for necessary workspace sizes.
    stegr!(n, sw)
    sw.lwork = BlasInt(sw.work[1])
    sw.work = Array{T}(undef, sw.lwork)
    sw.liwork = sw.iwork[1]
    sw.iwork = Array{BlasInt}(undef, sw.liwork)
    sw
end

mutable struct StegrCache{T,R<:Real} <: HermitianSubspaceCache{T}
    v::Vector{T} # Subspace-propagated vector
    w::Vector{T}
    sw::StegrWork{R}
    StegrCache(::Type{T}, n::Integer) where T = new{T,real(T)}(
        Vector{T}(undef, n), Vector{T}(undef, n),
        StegrWork(real(T), BlasInt(n)))
end

StegrCache(Ks::KrylovSubspace{B,T,U}) where {B,T,U} =
    StegrCache(T, Ks.maxiter)

"""
    expT!(α, β, t, cache)

Calculate the subspace exponential `exp(t*T)` for a tridiagonal
subspace matrix `T` with `α` on the diagonal and `β` on the
super-/subdiagonal, diagonalizing via `stegr!`.
"""
function expT!(α::AbstractVector{R}, β::AbstractVector{R}, t::Number,
               cache::StegrCache{T,R}) where {T,R<:Real}
    stegr!(α, β, cache.sw)
    sel = 1:length(α)
    @inbounds for i = sel
        cache.w[i] = exp(t*cache.sw.w[i])*cache.sw.Z[1,i]
    end
    mul!(@view(cache.v[sel]), @view(cache.sw.Z[sel,sel]), @view(cache.w[sel]))
end


export nImgIn, nImgOut, nFeatIn, nFeatOut, nTheta, getOp, initTheta, AbstractConvKernel

abstract type AbstractConvKernel{T} <: AbstractMeganetElement{T} end

## All convKernel types are assumed to have fields nImage (size of the image) and sK (size of the Convolution Kernel)

function nImgIn(this::AbstractConvKernel)
    return [this.nImg[1]; this.nImg[2]; this.sK[3]]
end

function nImgOut(this::AbstractConvKernel)
   return [this.nImg[1]; this.nImg[2]; this.sK[4]]
end

function nFeatIn(this::AbstractConvKernel)
    return prod(nImgIn(this));
end
function nFeatOut(this::AbstractConvKernel)
   return prod(nImgOut(this));
end

function nTheta(this::AbstractConvKernel)
   return prod(this.sK);
end

function getOp(this::AbstractConvKernel{T},theta::Array{T}) where {T <: Number}

    m = prod(nImgOut(this))
    n = prod(nImgIn(this))

    A = LinearOperator{T}(m,n,false,false,
                        v -> Amv(this,theta,v),
                        Nullable{Function}(),
                        w -> ATmv(this,theta,w))
    return A
end

function initTheta(this::AbstractConvKernel{T}) where {T <: Number}

    sd    = T(0.01);
    theta = sd*randn(T,prod(this.sK));
    #id1 = find(theta>2*sd);
    #theta(id1[:]) = randn(numel(id1),1);

    #id2 = find(theta< -2*sd);
    #theta(id2(:)) = randn(numel(id2),1);
    #theta = max(min(2*sd, theta),-2*sd);
    return theta
end


export pairwise_fst

# generate a wrapper struct so we can nicely print the results
struct PairwiseFST
    results::DataFrame
    method::String
end

# pretty-printing of FST results
function Base.show(io::IO, data::PairwiseFST)
    show(
        io,
        data.results,
        show_row_number=false,
        rowlabel = Symbol(" "),
        eltypes = false,
        row_names = names(data.results),
        title = "Pairwise FST: " * data.method
    )
end

function FST_global(data::AbstractDataFrame)
    # observed/expected het per locus per pop
    het_df = DataFrames.combine(
        groupby(data, [:locus, :population]),
        :genotype => nonmissing => :n,
        :genotype => hetero_o => :het_obs,
        :genotype => hetero_e => :het_exp,
        :genotype => allele_freq => :alleles
    )
    # collapse down to retrieve averages and counts
    n_df = DataFrames.combine(
        groupby(het_df, :locus),
        :n => count_nonzeros => :count,
        :n => (n -> count_nonzeros(n) / reciprocal_sum(n)) => :mn,
        [:het_obs, :het_exp, :n] => ((o,e,n) -> mean(skipmissing(gene_diversity_nei87.(e, o, count_nonzeros(n) / reciprocal_sum(n))))) => :HS,
        :het_obs => (o -> mean(skipmissing(o)))=> :Het_obs,
        :alleles => (alleles ->  sum(values(avg_allele_freq(alleles, 2))))=> :avg_freq
        )

    Ht = @inbounds 1.0 .- n_df.avg_freq .+ (n_df.HS ./ n_df.mn ./ n_df.count) - (n_df.Het_obs ./ 2.0 ./ n_df.mn ./ n_df.count)
    DST = @inbounds Ht .- n_df.HS
    DST′ = @inbounds n_df.count ./ (n_df.count .- 1) .* DST
    HT′ = @inbounds n_df.HS .+ DST′

    round(safemean(DST′) / safemean(HT′), digits = 4)
end

function _pairwise_Nei(data::PopData)
    idx_pdata = groupby(data.loci, :population)
    pops = getindex.(keys(idx_pdata), :population)
    npops = length(idx_pdata)
    results = zeros(npops, npops)
    for i in 2:npops
        for j in 1:(i-1)
            results[i,j] = FST_global(DataFrame(idx_pdata[[j,i]]))
        end
    end
    return PairwiseFST(DataFrame(results, Symbol.(pops)), "Nei 1987")
end

function _pairwise_WeirCockerham(data::PopData)
    idx_pdata = groupby(data.loci, :population)
    pops = getindex.(keys(idx_pdata), :population)
    npops = length(idx_pdata)
    results = zeros(npops, npops)
    for i in 2:npops
        for j in 1:(i-1)
            results[i,j] = weircockerham1984(DataFrame(idx_pdata[[j,i]]))
        end
    end
    return PairwiseFST(DataFrame(results, Symbol.(pops)), "Weir-Cockerham")
end

#TODO add methods and complete docstring
"""
    pairwise_fst(data::PopData; method::String)
Calculate pairwise FST between populations in a `PopData` object.
#### Methods:
- `"Nei87"`: Nei (1987) method
- `"WC84"` : Weir-Cockerham (1984) method (default)
"""
function pairwise_fst(data::PopData; method::String = "WC")
    occursin("nei", lowercase(method)) && return _pairwise_Nei(data)
    occursin("wc", lowercase(method)) && return _pairwise_WeirCockerham(data)
end

function FST_global(data::PopData)
    # observed/expected het per locus per pop
    het_df = DataFrames.combine(
        groupby(data, [:locus, :population]),
        :genotype => nonmissing => :n,
        :genotype => hetero_o => :het_obs,
        :genotype => (i -> hetero_e(i)) => :het_exp,
        :genotype => allele_freq => :alleles
    )
    # collapse down to retrieve averages and counts
    n_df = DataFrames.combine(
        groupby(het_df, :locus),
        :n => (n -> sum(.!iszero.(n)))=> :count,
        :n => (n -> sum(.!iszero.(n)) ./ sum(reciprocal.(n))) => :mn,
        [:het_obs, :het_exp, :n] => ((o,e,n) -> mean(skipmissing(gene_diversity_nei87.(e,o,sum(.!iszero.(n)) ./ sum(reciprocal.(n)))))) => :HS,
        :het_obs => (o -> mean(skipmissing(o)))=> :Het_obs,
        :alleles => (alleles ->  sum(values(avg_allele_freq(alleles, 2))))=> :avg_freq
        )

    Ht = @inbounds 1.0 .- n_df.avg_freq .+ (n_df.HS ./ n_df.mn ./ n_df.count) - (n_df.Het_obs ./ 2.0 ./ n_df.mn ./ n_df.count)
    DST = @inbounds Ht .- n_df.HS
    DST′ = @inbounds n_df.count ./ (n_df.count .- 1) .* DST
    HT′ = @inbounds n_df.HS .+ DST′

    HT = safemean(HT)
    DST = safemean(DST)
    DST′ = safemean(DST′)
    HT′ = safemean(HT′)

    return Dict{Symbol, Float64}(
        :FST => DST / HT,
        :FST′ => DST′ / HT′
    )
end

function FST_locus(data::PopData)
    # observed/expected het per locus per pop
    het_df = DataFrames.combine(
        groupby(data.loci, [:locus, :population]),
        :genotype => nonmissing => :n,
        :genotype => hetero_o => :het_obs,
        :genotype => (i -> hetero_e(i)) => :het_exp,
        :genotype => allele_freq => :alleles
    )
    # collapse down to retrieve averages and counts
    n_df = DataFrames.combine(
        [:het_obs, :het_exp, :n, :alleles] => (o,e,n,alleles) -> (
            count = sum(.!iszero.(n)),
            mn = sum(.!iszero.(n)) ./ sum(reciprocal.(n)),
            HS = mean(skipmissing(gene_diversity_nei87.(e,o,sum(.!iszero.(n)) ./ sum(reciprocal.(n))))),
            Het_obs = mean(skipmissing(o)),
            avg_freq = sum(values(avg_allele_freq(alleles)).^2)
            ),
        groupby(het_df, :locus)
    )

    HT = 1.0 .- n_df.avg_freq .+ (n_df.HS ./ n_df.mn ./ n_df.count) - (n_df.Het_obs ./ 2.0 ./ n_df.mn ./ n_df.count)
    DST = HT .- n_df.HS
    DST′ = n_df.count ./ (n_df.count .- 1) .* DST
    HT′ = n_df.HS .+ DST′

    return Dict{Symbol, Vector{Float64}}(
        :FST => DST ./ HT,
        :FST′ => DST′ ./ HT′
        )
end

function f_stat_p(data::PopData; nperm::Int = 1000)
    observed = FST_global(data)
    popnames = data.meta.population
    perm_dict = Dict{Symbol,Vector{Float64}}(
        :FST => Vector{Float64}(undef, nperm-1),
        :FST′ => Vector{Float64}(undef, nperm-1)
        )
    p = Progress(nperm-1, dt = 1, color = :blue)
    @inbounds @sync for n in 1:nperm-1
        Base.Threads.@spawn begin
            tmp = FST_global(permute_samples!(copy(data.loci), popnames))
            perm_dict[:FST][n] = tmp[:FST]
            perm_dict[:FST′][n] = tmp[:FST′]
            next!(p)
        end
    end
    @inbounds for (k,v) in perm_dict
        observed[k] = (1 + sum(observed[k] .<= v))/nperm
    end

    return observed
end


using Plots

function makeCosPlots()
    x = 0:0.01:1.00
    yn2, yn4, yn8, yn16 = [], [], [], []
    for i in x
        push!(yn2, cosx(i, 2))
    end
    for i in x
        push!(yn4, cosx(i, 4))
    end
    for i in x
        push!(yn8, cosx(i, 8))
    end
    for i in x
        push!(yn16, cosx(i, 16))
    end
    p = plot(x, yn2)
    plot!(p, x, yn4)
    plot!(p, x, yn8)
    plot!(p, x, yn16)
    display(p)
end


const TIVN = Tuple{Int, Union{Vector, Nothing}}

function read_input(input::Vector)::Matrix{Int}
  n = length(input)
  dist = fill(zero(Int), n, n)

  for (rix, row) ∈ input |> enumerate
    for (cjx, col) ∈ split(row, r"\s+") |> enumerate
      dist[rix, cjx] = parse(Int, col)
    end
  end

  dist
end

function max_path(d::Matrix{Int}; show_path=false)::TIVN
  """
  Start from penultimate row downto 1st row using max. from
  2 possible cells (from the row below)
  """
  for prow_ix ∈ size(d)[1] - 1:-1:1   # penultimate row index
    for cjx ∈ 1:prow_ix
      d[prow_ix, cjx] += max(d[prow_ix + 1, cjx], d[prow_ix + 1, cjx + 1])
    end
  end

  show_path ? (d[1, 1], find_path(d)) : (d[1, 1], nothing)
end

function find_path(d::Matrix{Int})::Vector
  path = [(1, 1)]

  for row ∈ 2:size(d)[1]
    push!(path, coord_max(d, path[end]))
  end

  path
end

function coord_max(d, coord)
  (x, y) = coord
  d[x + 1, y] > d[x + 1, y + 1] ? (x + 1, y) : (x + 1, y + 1)
end

function calc_max_path(input::AbstractString; show_path=false)::TIVN
  vstr = if occursin(".txt", input)
    # file
    open(input) do f
      readlines(f)
    end
  else
    split(input, "\n")
  end

  vstr .|>
    strip |>
    a -> filter(s -> length(s) > 0, a) .|>
    String |>
    read_input |>
    d -> max_path(d; show_path)
end


__precompile__()

module NeuralNetDiffEq
#dependencies

using Reexport
@reexport using DiffEqBase
using Knet, Compat, ForwardDiff
import DiffEqBase: solve

# Abstract Types
@compat abstract type NeuralNetDiffEqAlgorithm <: AbstractODEAlgorithm end

# DAE Algorithms
immutable nnode <: NeuralNetDiffEqAlgorithm
    hl_width::Int
end

nnode(;hl_width=10) = nnode(hl_width)

export nnode

#constant functions
#z(P,i,x) = P[1][i]*x + P[2][i]
sig_der(x) = sigm(x)*(1-sigm(x))

function solve(
    prob::AbstractODEProblem,
    alg::NeuralNetDiffEqAlgorithm;
    dt = nothing,
    timeseries_errors = true,
    iterations = 50)

    u0 = prob.u0
    tspan = prob.tspan
    f = prob.f
    t0 = tspan[1]

    if dt == nothing
        error("dt must be set.")
    end

    #types and dimensions
    uElType = eltype(u0)
    tType = typeof(tspan[1])
    outdim = length(u0)

    #hidden layer
    hl_width = alg.hl_width

    #The phi trial solution
    phi(P,t) = u0 + (t-t0)*predict(P,t)

    #train points generation
    dtrn = generate_data(tspan[1],tspan[2],dt,atype=tType)

    #iterations
    _maxiters = iterations

    #initialization of weights and bias
    w = init_params(uElType,hl_width)
    # println(w[1])
    # println(w[2])
    # println(w[3])

    #initialization of optimization parameters (Adam by default for now)
    lr_ = 0.1
    beta1_ = 0.9
    beta2_ = 0.95
    eps_ = 1e-6
    prms = Any[]

    for i=1:length(w)
    prm = Adam(lr=lr_, beta1=beta1_, beta2=beta2_, eps=eps_)
    push!(prms, prm)
    end

    #iters = 1000

    #reporting the accuracy
    #report(epoch)=println((:epoch,epoch,:trn,accuracy(w,dtrn,predict),:tst,accuracy(w,dtst,predict)))
    #report(epoch)=println((:epoch,epoch,:trn,loss_trial(w,dtrn)))
    #report(0)
    #P_tuned = train(w,prms,dtrn,regFlag; epochs=100, iters=1000)
    #@time for epoch=1:epochs
    #@time train(w, prms, dtrn, f, phi, hl_width; maxiters=_maxiters)
        #report(epoch)
    #end

    @time for iters=1:_maxiters
            train(w, prms, dtrn, f, phi, hl_width; maxiters=1)
            loss = test(w,dtrn,f,phi,hl_width)
            if mod(iters,100) == 0
                println((:iteration,iters,:loss,loss))
            end
            #gradcheck(loss_trial, w, dtrn, f, phi, hl_width...; gcheck=10, verbose=true)
            #check_Phi_gradient(w,dtrn,hl_width)
            if loss < 10^(-8.0)
                break
            end
        end

    # for t in log; println(t); end
    # return w

    #solutions at timepoints
    u = [phi(w,x) for x in dtrn]

    build_solution(prob,alg,dtrn,u,
               timeseries_errors = timeseries_errors,
               retcode = :Success)

end #solve

include("training_utils.jl")
#include("interface.jl")
end # module


module NeuralNetwork

typealias AA AbstractArray

include("utils.jl")
include("mlp.jl")

export ConvolutionLayer, LogisticRegression, fit!, predict
export onehot, maxindices

type ConvolutionLayer
    filterw::Int
    filterh::Int
    stepw::Int
    steph::Int
    weights::Matrix
    bias::Matrix

    function ConvolutionLayer(filter_size::Int, step_size::Int,
                              weights = zeros(filter_size, filter_size))
        filterw, filterh = to_tuple_if_int(filter_size)
        stepw, steph = to_tuple_if_int(step_size)
        ConvolutionLayer(filterw, filterh, stepw, steph, weights)
    end

    function ConvolutionLayer(filterw::Int, filterh::Int,
                              stepw::Int, steph::Int,
                              weights = zeros(filterw, filterh))
        new(filterw, filterh, stepw, steph, weights)
    end
end

function forward(layer::ConvolutionLayer, image::Matrix)
    """
    Forward computation for an image
    """
    filterw, filterh = layer.filterw, layer.filterh
    startpointsx = 1:layer.stepw:(size(image, 1)-filterw+1)
    startpointsy = 1:layer.steph:(size(image, 2)-filterh+1)

    output = zeros(length(startpointsx), length(startpointsy))
    for (ox, x) in enumerate(startpointsx)
        for (oy, y) in enumerate(startpointsy)
            subimage = image[x:x+filterw-1, y:y+filterh-1]
            output[ox, oy] = sum(subimage .* layer.weights)
        end
    end
    output += layer.bias
    return output
end

function backward(layer::ConvolutionLayer)
end

end # module


"""

    Interval unions sets of defined by unions of disjoint intervals.
    This file includes constructors, arithmetic (including intervals and scalars)
    and complement functions

    Empty sets and intersecting intervals are appropriately handled in the constructor:

    julia> a = interval(0,2) ∪ interval(3,4)
    [0, 2] ∪ [3, 4]

    julia> b = interval(1,2) ∪ interval(4,5) ∪ ∅
    [1, 2] ∪ [4, 5]

    julia> c = a * b
    [0, 10] ∪ [12, 20]

    julia> complement(c)
    [-∞, 0] ∪ [10, 12] ∪ [20, ∞]

"""

###
#   IntervalUnion constructor. Consists of a vector of intervals
###
struct IntervalUnion{T<:Real} <: AbstractInterval{T}
    v :: Vector{Interval{T}}
end

###
#   Outer constructors
###
function intervalUnion(v)
    x = IntervalUnion(v)
    sort!(x.v)
    remove_empties!(x)
    condense!(x)
    closeGaps!(x)
    return x
end

intervalUnion(num :: Real) = IntervalUnion([interval(num)])

intervalUnion(lo :: Real, hi :: Real) = IntervalUnion([interval(lo,hi)])

intervalUnion(x :: Interval) = intervalUnion([x])
∪(x :: Interval) = intervalUnion(x)

∪(x :: Interval, y :: Interval) = intervalUnion([x; y])
∪(x :: Array{Interval{T}}) where T <:Real = intervalUnion(x)

intervalUnion(x :: Interval, y :: IntervalUnion) = intervalUnion([x; y.v])
∪(x :: Interval, y :: IntervalUnion) = intervalUnion(x,y)

intervalUnion(x :: IntervalUnion, y :: Interval) = intervalUnion([x.v; y])
∪(x :: IntervalUnion, y :: Interval) = intervalUnion(x,y)

intervalUnion(x :: IntervalUnion, y :: IntervalUnion) = intervalUnion([x.v; y.v])
∪(x :: IntervalUnion, y :: IntervalUnion) = intervalUnion(x,y)

# MultiInterval can act like a vector
getindex(x :: IntervalUnion, ind :: Integer) = getindex(x.v,ind)
getindex(x :: IntervalUnion, ind :: Array{ <: Integer}) = getindex(x.v,ind)

# Remove ∅ from IntervalUnion
function remove_empties!(x :: IntervalUnion)
    deleteat!(x.v, (i for i in eachindex(x.v) if isempty(x.v[i])))
    x
end

function closeGaps!(x :: IntervalUnion, maxInts = MAXINTS[1])

    while length(x.v) > maxInts     # Global

        # Complement code
        v = sort(x.v)
        vLo = left.(v)
        vHi = right.(v)
        vLo[1] == -∞ ? popfirst!(vLo) : pushfirst!(vHi, -∞)
        vHi[end] == ∞ ? pop!(vHi) : push!(vLo, ∞)
        xc = interval.(vHi,vLo)

        # Close smallest width of complement
        widths = diam.(xc)
        d, i = findmin(widths)
        merge = hull(x[i-1], x[i])
        popat!(x.v, i)
        popat!(x.v, i-1)
        push!(x.v, merge)
        sort!(x.v)
    end
    return x
end

function condense!(f::Function, x :: IntervalUnion)
    v = x.v
    sort!(v)
    i = 1
    while i < length(v)
        these = (i - 1) .+ findall(collect(!f(v[i] ∩ v[j]) for j in i:length(v)))
        if length(these) > 1
            v[i] = hull(v[these])
            deleteat!(v, @view(these[2:end]))
        end
        i += 1
    end
    return x
end

condense(f, x) = condense!(f, copy(x))

# join intervals that intersect.
condense!(x::IntervalUnion) = condense!(isempty, x)
# join intervals that intersect.
condense(x::IntervalUnion) = condense(isempty, x)
# Join intervals which intersect in more than one point
condense_weak(x::IntervalUnion) = condense(i->isempty(i) || isthin(i), x)

iscondensed(f::Function, x :: IntervalUnion) = all(f(intersect(x.v[i], x.v[j])) for i in eachindex(x.v) for j in 1:i-1)

iscondensed(x::IntervalUnion) = iscondensed(isempty, x)
iscondensed_weak(x::IntervalUnion) = iscondensed(i->isempty(i) || isthin(i), x)


@testset "Misc/MPolyAbsFact" begin

  f = Hecke.MPolyFact.example(wildanger_field(3, 13)[1], 3, 5)
  @test length(factor_absolute(f)) >= 2

  f = Hecke.MPolyFact.example(cyclotomic_field(4)[1], 3, 5)
  @test length(factor_absolute(f)) >= 2

  @test length(factor(f, ComplexField(40)).fac) >= 2
  @test length(factor(f, ArbField(40)).fac) >= 1
end


# ## Building

# ## Model

module building #jl

using ReachabilityAnalysis, MAT
using ReachabilityModels: @relpath

file = matopen(@relpath "building.mat")

## system matrix
A = read(file, "A")

## input matrix
B = read(file, "B")

## state domain
X = Universe(48)

## input domain
U = BallInf([0.5], 0.3)

function model(X0)
    S = @system(x' = Ax + Bu, x ∈ X, u ∈ U)
    return IVP(S, X0)
end

end  #jl

# ## References


module DiffEq

using LinearAlgebra
using ..Op:setup
using ..EdgeDetect:find
using ..Chebyshev:make_cheb

function solve(;N = 20, bc1 = [:a,0,0.0], bc2 = [:b,0,0.0], ranges=zeros(0,2), contin = Bool[], A0=0.0, A1=0.0, A2=0.0, B=0.0, a = -1.0, b = 1.0, dosplit=true)

#    if split && !(typeof(A0) <: Number) && length(ranges) == 0
#        ranges, contin = find(A0, a=a,b=b)
#    elseif split && !(typeof(B) <: Number) && length(ranges) == 0
#        ranges, contin = find(B, a=a,b=b)
#    else
#        ranges = zeros(0,2)
#        contin = Bool[]
#    end

    H, S, rhs =  setup(N=N, a=a,b=b,bc1=bc1, bc2=bc2, ranges=ranges, contin=contin, A0=A0, A1=A1,A2=A2,B=B,dosplit=dosplit)
    u = H\rhs

    C = make_cheb(u, a = a, b = b)

    return C

end

end #end module


include("../utility.jl")
using Turing, Test

N = 10
setchunksize(N)
# alg = HMCDA(2000, 0.65, 1.5)
alg = HMC(1000, 0.2, 4)

@model vdemo() = begin
  x = Vector{Real}(undef, N)
  for i = 1:N
    x[i] ~ Normal(0, sqrt(4))
  end
end

t_loop = @elapsed res = sample(vdemo(), alg)

# Test for vectorize UnivariateDistribution
@model vdemo() = begin
  x = Vector{Real}(undef, N)
  x ~ [Normal(0, 2)]
end

t_vec = @elapsed res = sample(vdemo(), alg)

@model vdemo() = begin
  x ~ MvNormal(zeros(N), 2 * ones(N))
end

t_mv = @elapsed res = sample(vdemo(), alg)

println("Time for")
println("  Loop : $t_loop")
println("  Vec  : $t_vec")
println("  Mv   : $t_mv")

# Transformed test
@model vdemo() = begin
  x = Vector{Real}(undef, N)
  x ~ [InverseGamma(2, 3)]
end

sample(vdemo(), alg)


using GeometricFlux
using GraphSignals
using Flux
using Flux: onehotbatch, onecold, logitcrossentropy, throttle
using Flux: @epochs
using JLD2
using Statistics
using SparseArrays
using Graphs.SimpleGraphs
using CUDA

@load "data/cora_features.jld2" features
@load "data/cora_labels.jld2" labels
@load "data/cora_graph.jld2" g

num_nodes = 2708
num_features = 1433
hidden = 16
target_catg = 7
epochs = 200
λ = 5e-4

## Preprocessing data
train_X = Matrix{Float32}(features) |> gpu  # dim: num_features * num_nodes
train_y = Matrix{Float32}(labels) |> gpu  # dim: target_catg * num_nodes
fg = FeaturedGraph(g)  # pass to gpu together in model layers

## Model
model = Chain(GCNConv(fg, num_features=>hidden, relu),
              Dropout(0.5),
              GCNConv(fg, hidden=>target_catg),
              ) |> gpu;
# do not show model architecture, showing CuSparseMatrix will trigger errors

## Loss
l2norm(x) = sum(abs2, x)
# cross entropy with first layer L2 regularization
loss(x, y) = logitcrossentropy(model(x), y) + λ*sum(l2norm, Flux.params(model[1]))
accuracy(x, y) = mean(onecold(softmax(cpu(model(x)))) .== onecold(cpu(y)))

## Training
ps = Flux.params(model)
train_data = [(train_X, train_y)]
opt = ADAM(0.01)
evalcb() = @show(accuracy(train_X, train_y))

@epochs epochs Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))


using Revise
using AdFem
using PyPlot
using SparseArrays

cmesh = CrackMesh(50,50,1/50,15)
mmesh = cmesh.mesh

NT = 500
Δt = 2/NT

Ca = spzeros(2cmesh.nslip, 2mmesh.ndof)
ls = zeros(2cmesh.nslip)
for i = 1:cmesh.nslip
    u = cmesh.upper[i+1]
    l = cmesh.lower[i+1]
    Ca[i, u] = 1.0
    Ca[i, l] = -1.0
    ls[i] = 0.0

    u = cmesh.upper[i+1] + mmesh.ndof
    l = cmesh.lower[i+1] + mmesh.ndof
    Ca[i+cmesh.nslip, u] = 1.0
    Ca[i+cmesh.nslip, l] = -1.0
    ls[i+cmesh.nslip] = 0.0
end
Ca = constant(Ca)

upbd = bcnode((x,y)->y>1.0-0.001, mmesh)
downbd = bcnode((x,y)->y<0.001, mmesh)

ν = 0.1
E = 1.0
D = compute_plane_strain_matrix(E, ν)
K = constant(compute_fem_stiffness_matrix(D, mmesh))

rhs = zeros(2mmesh.ndof)
K, rhs = impose_Dirichlet_boundary_conditions(
    K, rhs,
    [upbd; downbd; upbd .+ mmesh.ndof; downbd .+ mmesh.ndof],
    [0.001*ones(length(upbd)); zeros(length(downbd)+length(upbd)+length(downbd))]
)

A = [K Ca';Ca spzero(2cmesh.nslip)]
rhs = [rhs; constant(ls)]
u = A\rhs

u = u[1:2mmesh.ndof]

sess = Session(); init(sess)

U = run(sess, u)

# close("all")
# visualize_displacement(100*U, mmesh)
# savefig("mesh.png")

Se = compute_von_mises_stress_term(D, U, mmesh)
close("all")
visualize_scalar_on_gauss_points(Se, mmesh)
savefig("stress.png")


"""
    SyntheticDiD

Synthetic Differences-in-Differences (SDID) estimator as described in Arkhangelsky et al. (2021).
SDID is defined as

    μ̂, α̂, β̂, τ̂ = argmin ∑ₙ∑ₜ(Yᵢₜ - μ - αᵢ - βₜ - Wᵢₜτ)²ω̂ᵢλ̂ₜ

where i = 1, 2, ..., N denotes the observation units in the panel, t = 1, ..., T denotes the time
periods, Y is an (N×T) matrix of outcomes, W is an (N×T) indicator matrix of treatment status, ω is
a unit weight and λ is a time period weight.

"""
struct SyntheticDiD{T1}# <: SCM where T1
    tp::T1
    ω̂::Vector{Float64}
    λ̂::Vector{Float64}
    ŷ::Vector{Float64}
    τ̂::Vector{Float64}
    se_τ̂::Vector{Float64}
end

function SyntheticDiD(x::BalancedPanel{SingleUnitTreatment{Continuous}})
    ω̂ = zeros(length(x.is) - 1)
    λ̂ = zeros(length_T₀(x))
    ŷ = zeros(size(x.Y, 2))
    τ̂ = zeros(1)
    se_τ̂ = zeros(1)
    SyntheticDiD(x, ω̂, λ̂, ŷ, τ̂, se_τ̂)
end

function fit!(x::SyntheticDiD{BalancedPanel{SingleUnitTreatment{Continuous}}})

    (;Y, W) = x.tp
    y₁₀, y₁₁, y₀₀, y₀₁ = decompose_y(x.tp)
    T₀ = length_T₀(x.tp)
    T₁ = length_T₁(x.tp)
    Nₜᵣ = 1 # length(treated_ids(x)) - has to be one given type of BalancedPanel in this method
    Nₖₒ = size(Y, 1) - Nₜᵣ

    ### Algorithm presented in Arkhangelsky et al. (2021), v4 arXiv

    ## (1) Compute regularization parameter ζ
    # First difference of outcomes for control units before treatment time
    Δᵢₜ = Y[Not(treated_ids(x.tp)), 2:T₀] .- Y[Not(treated_ids(x.tp)), 1:T₀-1]
    Δ̄ = mean(Δᵢₜ) # Average first difference

    σ̂ = √(sum((Δᵢₜ .- Δ̄)'*(Δᵢₜ .- Δ̄)) / (Nₖₒ*(T₀ - 1)))

    ζ = (Nₜᵣ*T₁)^(1/4) * σ̂

    ## (2) Compute unit weights ω̂
    #!# Note - paper seems to generalise to multiple treated units already, here using one unit for
    # now but should be 1/Nₖₒ * sum(Y[treated_ids(x), t]) at each time period
    #!# Weights produced here are very non-sparse due to ζ*norm regularization - double check!
    function ℓᵢ(ω)
        ω₀ = ω[1]; ω = @view ω[2:end]
        (ω₀ .+ y₀₀'*ω .- y₁₀)'*(ω₀ .+ y₀₀'*ω .- y₁₀) + 1e12*(1.0 .- sum(ω))^2 + ζ^2*T₀*norm(ω)
    end

    # Bounds: ω₀ ∈ ℝ, ω ∈ Ω = {ω ∈ ℝ₊ᴺ: ∑wᵢ = 1} - essentially 0 ≤ ωᵢ ≤ 1 with sum constraint and ω₀
    # unbounded
    #!# again this will differ for multiple treated units
    lower = [-Inf; fill(0.0, Nₖₒ)]
    upper = [Inf; fill(1.0, Nₖₒ)]
    initial = [0.0; [1/Nₖₒ for _ ∈ 1:Nₖₒ]]

    res = optimize(ℓᵢ, lower, upper, initial)

    ω̂₀ = res.minimizer[1]  #!# this is probably actually unused?
    ω̂ = res.minimizer[2:end]
    x.ω̂ .= ω̂

    ## (3) Compute time weights λ̂
    # Note small regularization added as per footnote 3 of the paper
    function ℓₜ(λ)
        λ₀ = λ[1]; λ = @view λ[2:end]
        ȳ₀₁ = vec(sum(y₀₁, dims = 2) ./ size(y₀₁, 2)) # Average post-treatment outcome for control
        (λ₀ .+ y₀₀*λ .- ȳ₀₁)'*(λ₀ .+ y₀₀*λ .- ȳ₀₁) + 1e12*(1.0 .- sum(λ))^2 + ζ^2*Nₖₒ*norm(λ)
    end

    lower = [-Inf; fill(0.0, T₀)]
    upper = [Inf; fill(1.0, T₀)]
    initial = [0.0; [1/T₀ for _ ∈ 1:T₀]]

    res = optimize(ℓₜ, lower, upper, initial)

    λ̂₀ = res.minimizer[1]  #!# this is probably actually unused?
    λ̂ = res.minimizer[2:end]

    x.λ̂ .= λ̂

    ## (4) Run regression
    # Need to estimate fixed effect model with time/unit fixed effects and
    #!# To do: for covariates we want to partial them out before this regression
    # Add unit weights - 1 for treated unit, ω̂ᵢ for untreated
    ω_df = DataFrame(Symbol(x.tp.id_var) => [treated_labels(x.tp); x.tp.is[x.tp.is .!= treated_labels(x.tp)]], :ω̂ => [1.0; ω̂])
    df = leftjoin(x.tp.df, ω_df, on = x.tp.id_var)

    # Add time weights
    λ_df = DataFrame(Symbol(x.tp.t_var) => x.tp.ts, :λ̂ => [fill(1/T₁, T₁); λ̂])
    df = leftjoin(df, λ_df, on = x.tp.t_var)

    # Total weights - ωλ
    df.sdid_weight = df.ω̂ .* df.λ̂

    # Add treatment dummy based on treatment matrix
    df.W = reshape(x.tp.W', length(x.tp.ts)*length(x.tp.is))

    sdid_reg = reg(df,
        term(x.tp.outcome_var) ~ term("W") + fe(term(x.tp.id_var)) + fe(term(x.tp.t_var)),
        weights = :sdid_weight, Vcov.robust())

    # Treatment effect is estimated as the coefficient on treatment dummy
    τ̂ = coef(sdid_reg)
    x.τ̂ .= τ̂
    x.se_τ̂ .= stderror(sdid_reg)

    # Counterfactual outcome is actual outcome plus treatment effect
    ŷ₁₁ = y₁₁ .+ τ̂
    x.ŷ .= [y₁₀; ŷ₁₁]
end

#!# TO DO - show method for this type


module VarianceTests
using Base.Test;
using FactCheck;
using ValidatedNumerics;
using IntervalStatistic;
using Distributions;

srand(10)

contain(value) = (i) -> value in i

facts("estimate variance of standard (0, 1) distribution") do
    d = Normal()
    length = 100
    values = rand(d, length)
    mu, sigma = params(d)
    average = mean(values)

    context("by confidence probability") do
        result = IntervalStatistic.var(values, IntervalStatistic.Variance.byConfidenceProbability(0.95))
        println(result)
        @fact result --> contain(sigma^2)
    end

    context("by point variance") do
        variance = var(values)
        result = IntervalStatistic.var(values, IntervalStatistic.Variance.byPointVariance(0.95, variance))
        println(result)
        @fact result --> contain(variance)
    end
end

facts("estimate variance of normanl mu=3 sigma=0.1 distribution") do
    d = Normal(3, 0.1)
    length = 10
    values = rand(d, length)
    mu, sigma = params(d)
    average = mean(values)

    context("by confidence probability") do
        result = IntervalStatistic.var(values, IntervalStatistic.Variance.byConfidenceProbability(0.95))
        println(result)
        @fact result --> contain(sigma^2)
    end

    context("by point variance") do
        variance = var(values)
        result = IntervalStatistic.var(values, IntervalStatistic.Variance.byPointVariance(0.95, variance))
        println(result)
        @fact result --> contain(variance)
    end
end

facts("estimate variance of standard (0, 1) distribution with 10 samples") do
    d = Normal()
    length = 10
    values = rand(d, length)
    mu, sigma = params(d)
    average = mean(values)

    context("by mean abs deviation") do
        result = IntervalStatistic.var(values, IntervalStatistic.Variance.byMeanAbsoluteDeviation(0.95))
        println(result)
        @fact result --> contain(sigma^2)
    end
end

facts("estimate variance of normanl mu=3 sigma=0.1 distribution with 10 samples") do
    d = Normal(3, 0.1)
    length = 10
    values = rand(d, length)
    mu, sigma = params(d)
    average = mean(values)

   context("by mean abs deviation") do
        result = IntervalStatistic.var(values, IntervalStatistic.Variance.byMeanAbsoluteDeviation(0.95))
        println(result)
        @fact result --> contain(sigma^2)
    end
end

end


using LightGraphs: AbstractGraph

## Linear algebra API for AbstractGraph

function degrees(g::AbstractGraph, T::DataType=eltype(g); dir::Symbol=:out)
    adj = adjacency_matrix(g, T; dir=dir)
    degrees(adj, T; dir=dir)
end

function degree_matrix(g::AbstractGraph, T::DataType=eltype(g); dir::Symbol=:out)
    adj = adjacency_matrix(g, T; dir=dir)
    degree_matrix(adj, T; dir=dir)
end

function inv_sqrt_degree_matrix(g::AbstractGraph, T::DataType=eltype(g); dir::Symbol=:out)
    adj = adjacency_matrix(g, T; dir=dir)
    inv_sqrt_degree_matrix(adj, T; dir=dir)
end

function laplacian_matrix(g::AbstractGraph, T::DataType=eltype(g); dir::Symbol=:out)
    adj = adjacency_matrix(g, T; dir=dir)
    laplacian_matrix(adj, T; dir=dir)
end

function normalized_laplacian(g::AbstractGraph, T::DataType=eltype(g); selfloop::Bool=false)
    adj = adjacency_matrix(g, T)
    selfloop && (adj += I)
    normalized_laplacian(adj, T)
end

function scaled_laplacian(g::AbstractGraph, T::DataType=eltype(g))
    adj = adjacency_matrix(g, T)
    scaled_laplacian(adj, T)
end


# The halfpsace functions are private methods, and all take a coordinate and a surface to determine which halfspace the coordinate is in.

# Plane halfspace determination
function halfspace(c::Coord, plane::Plane)
    d::Float64 = -dot(plane.normal, plane.point)
    half::Float64 = dot(plane.normal, c) + d
    if half <= 0
        return -1
    else
        return 1
    end
end

# Cone halfspace determination
function halfspace(c::Coord, cone::Cone)
	p_minus_c::Coord = c - cone.tip
    half::Float64 = dot(p_minus_c,cone.axis)^2 - dot(p_minus_c,p_minus_c) * (cos(cone.theta))^2
    if half <= 0
        return -1
    else
        return 1
    end
end

# Sphere halfspace determination
function halfspace(c::Coord, sphere::Sphere)
    half::Float64 = (c.x - sphere.center.x)^2 + (c.y - sphere.center.y)^2 + (c.z - sphere.center.z)^2 - sphere.radius^2
    if half <= 0
        return -1
    else
        return 1
    end
end

# Infinite cylinder halfspace
function halfspace(c::Coord, cyl::InfCylinder)
    tmp::Coord = cross((c-cyl.center), cyl.normal)
    half::Float64 = dot(tmp, tmp) - cyl.radius^2
    if half <= 0
        return -1
    else
        return 1
    end
end

function ^(a::Region, b::Region)
    if halfspace(_p, a.surface) == a.halfspace
        if halfspace(_p, b.surface) == b.halfspace
            return true
        end
    end
    return false
end

function ^(a::Region, b::Bool)
    if halfspace(_p, a.surface) == a.halfspace
        if b == true
            return true
        end
    end
    return false
end

function ^(b::Bool, a::Region)
    if halfspace(_p, a.surface) == a.halfspace
        if b == true
            return true
        end
    end
    return false
end

function |(a::Region, b::Region)
    if halfspace(_p, a.surface) == a.halfspace
        return true
    end
    if halfspace(_p, b.surface) == b.halfspace
        return true
    end
    return false
end

function |(a::Region, b::Bool)
    if halfspace(_p, a.surface) == a.halfspace
        return true
    end
    if b == true
        return true
    end
    return false
end

function |(b::Bool, a::Region)
    if halfspace(_p, a.surface) == a.halfspace
        return true
    end
    if b == true
        return true
    end
    return false
end

function ~(a::Region)
    b::Region = Region(a.surface, a.halfspace)
    if a.halfspace == -1
        b.halfspace = 1
    else
        b.halfspace = -1
    end
    return b
end


function fdc(order::Int64)

    # This function calculate FD coefficient for spatial finite difference
    # Currently we define the coefficient as the one for the 4th order
    # search "Finite difference coefficient" online for different order

    if order == 4
        FDC = [1/24, -9/8, 9/8, -1/24] # Finite-Difference Coefficient
        else error("please choose choose 4")
    end
    return FDC
end


push!(LOAD_PATH, pwd()*"/../../src/")
push!(LOAD_PATH, pwd()*"/src/")
using Distributions
using PyPlot
using BVB

lower = -20.
upper = 20.

true_d = Cauchy(0.0, 2)
ll(x) = logpdf(true_d, x)

init_approx_d = Normal(0, 40)
x = linspace(-20,20,500)

max_iter = 30

# include("run_laplacian_gradient_boosting.jl")
include("run_laplacian_gradient_boosting_exclusive.jl")
save_trace("./cauchy/", x, init_approx_d, approx_d, true_d,
                    iter_vec, α_vec,
                    eval_vec, -elbo_vec)
show()

println("Done.")


using ReservoirComputing

#model parameters
const approx_res_size = 30
const radius = 1.2
const activation = tanh
const degree = 6
const sigma = 0.1
const beta = 0.0
const alpha = 1.0
nla_type = NLADefault()
const in_size = 3
const out_size = 3
const extended_states = false

const train_len = 50
const predict_len = 12
data = ones(Float64, in_size, 100)
train = data[:, 1:1+train_len-1]
test = data[:, train_len:train_len+predict_len-1]

#test non linear algos
nla = [NLAT1(), NLAT2(), NLAT3()]
for t in nla
    nla_type = t
    esn = ESN(approx_res_size,
        train,
        degree,
        radius,
        activation = activation,
        sigma = sigma,
        alpha = alpha,
        nla_type = nla_type,
        extended_states = extended_states)

    W_out = ESNtrain(esn, beta)
    @test size(W_out) == (out_size, esn.res_size)
    output = ESNpredict(esn, predict_len, W_out)
    @test size(output) == (out_size, predict_len)

end


using DifferentialEquations

# G = 6.67430e-11 #m^3/(kg*s^2)
# m1 = 5.9723e+24 # kg
# m2 = 1.9884e+30 # kg m2 = 333*m1

G = 1.0
m1 = 0.9
m2 = 0.1
M = m1 + m2
mu = m1 * m2 / M

function kepler(t, u, du)
    du[1] = u[3]
    du[2] = u[4]

    norm_r = sprt(u[3]^2 + u[4]^2)^3

    du[3] = -u[1] / norm_r
    du[4] = -u[2] / norm_r
end

u0 = [1.0;0.0]
tspan = (0.0, 100.0)
problem = ODEProblem(kepler, u0, tspan)
solution = solve(problem, Euler(), reltol=1e-6, saveat=0.1, save_everystep=true)


 @testset "Nodes initialization algorithm n=3 entry = 3 output = 2" begin
    M = 1 # Constante para la función rampa
    # Bien definido para tamaño n = 2 y salida de dimensión 1
    f_regression(x,y,z)=[x*y-z,x]
    data_set_size =  6
    entry_dimension = 3
    output_dimension = 2
    # Número de neuronas
    n = data_set_size # Debe de ser mayor que 1 para que no de error
    X_train= rand(Float32, data_set_size, entry_dimension)
    Y_train::Matrix =  mapreduce(permutedims, vcat, map(x->f_regression(x...), eachrow(X_train)))

    h = nn_from_data(X_train, Y_train, n, M)

    # veamos que el tamaño de la salida es la adecuada
    @test size(h.W1) == (n,entry_dimension+1)
    @test size(h.W2) == (output_dimension,n)

    # Si ha sido bien construida:
    # Evaluar la red neuronal en los datos con los que se construyó
    # debería de resultar el valor de Y_train respectivo
    evaluar(x)=forward_propagation(h,
     RampFunction,x)

    for i in 1:n
        @test evaluar(X_train[i,:]) ≈ Y_train[i,:]
    end

end


export DP

# Stick breaking construction.

@pp function pickstick(sticks, j)
    flip(sticks(j)) ? j : pickstick(sticks, j+1)
end

@pp function makesticks(alpha)
    local sticks = mem(index -> ~Beta(1.0, alpha))
    () -> pickstick(sticks,1)
end

@pp function DP(thunk, alpha)
    local atoms = mem(index -> thunk())
    local sticks = makesticks(alpha)
    () -> atoms(sticks())
end


using JuMP
using Ipopt

function Solve_Dist_SOC(buses, lines, generators)
info("Starting DistSOC Model")

bus_set = collect(keys(buses))
line_set = collect(keys(lines))
gen_set = collect(keys(generators))
gen_bus = []
for g in keys(generators)
    push!(gen_bus, generators[g].bus_idx)
end
root_bus = 0
for b in keys(buses)
    if buses[b].is_root
        root_bus = b
    end
end

lines_to = Dict()
for l in keys(lines)
    lines_to[lines[l].to_node] = lines[l]
end

v_root = 1

m = Model(solver=IpoptSolver(print_level=3))

@variable(m, v[bus_set] >= 0) #variable for voltage square
@variable(m, fp[bus_set]) #variable for active power flow
@variable(m, fq[bus_set]) # variable for reactive power flow
@variable(m, gp[bus_set]) #variable for active power generation
@variable(m, gq[bus_set]) #variable for reactive power generation
@variable(m, l[bus_set] >= 0) # variable for quadratic current

@objective(m, Min, sum((gp[b])*buses[b].generator.cost for b in gen_bus))

for b in gen_bus
# All buses with generation
    @constraint(m, gp[b] <= buses[b].generator.g_P_max) #Upper limit constraint for active power generation
    @constraint(m, gp[b] >= 0) #Lower limit constraint for active power generation
    @constraint(m, gq[b] <= buses[b].generator.g_Q_max) #Upper limit constraint for reactive power generation
    @constraint(m, gq[b] >= -buses[b].generator.g_Q_max) #Lower limit constraint for reactive power generation
end

for b in setdiff(bus_set, gen_bus)
# All buses without generation
    @constraint(m, gp[b] == 0)
    @constraint(m, gq[b] == 0)
end

for b in bus_set
# All buses
    @constraint(m, buses[b].d_P - gp[b] + sum((fp[k] + l[k]*lines_to[k].r) for k in buses[b].children) == fp[b])
    @constraint(m, buses[b].d_Q - gq[b] + sum((fq[k] + l[k]*lines_to[k].x) for k in buses[b].children) == fq[b])

    @constraint(m, v[b] <= buses[b].v_max)
    @constraint(m, v[b] >= buses[b].v_min)
end

for b in setdiff(bus_set, [root_bus])
# All buses without root
    b_ancestor = buses[b].ancestor[1]
    R = lines_to[b].r
    X = lines_to[b].x
    @constraint(m, v[b] == v[b_ancestor] - 2*(R*fp[b] + X*fq[b]) - l[b]*(R+X))
    @constraint(m, fp[b]^2 + fq[b]^2 <= lines_to[b].s_max^2)
    @constraint(m, (fp[b] + l[b]*R)^2 + (fq[b] + l[b]*X)^2 <= lines_to[b].s_max^2)
    @NLconstraint(m, (fp[b]^2 + fq[b]^2)/l[b] <= v[b])
end

#Voltage and flow constraints for root node
@constraint(m, v[root_bus] == v_root)
@constraint(m, l[root_bus] == 0)
@constraint(m, fp[root_bus] == 0)
@constraint(m, fq[root_bus] == 0)

solve(m)

# Prepare Results
bus_results = DataFrame(bus = Any[], d_P = Any[], d_Q = Any[], g_P = Any[], g_Q = Any[], v_squared = Any[])
line_results = DataFrame(line = Any[], from = Any[], to = Any[], f_P = Any[], f_Q = Any[], a_squared = Any[])
for b in bus_set
    res = [b, buses[b].d_P, buses[b].d_Q, getvalue(gp[b]), getvalue(gq[b]), getvalue(v[b])]
    push!(bus_results, res)

    # Calc square current on line
    a_res = 0
    fp_res = getvalue(fp[b])
    fq_res = getvalue(fq[b])
    if b != root_bus
        a_res = getvalue(l[b])
        lres = [lines_to[b].index, buses[b].ancestor[1], b, fp_res, fq_res, a_res]
        push!(line_results, lres)
    end
end

sort!(bus_results, cols=:bus)
sort!(line_results, cols=:to)

return getobjectivevalue(m), bus_results, line_results

end


using StaticArrays
using BenchmarkTools

combination(x) = combination_(x...)
combination_(x) = x
combination_(x,y) = (x,y,(x,y))
combination_(x,y...) = (x,combination(y)...,map(z->(x,z...),combination(y))...)

combination((1))
combination((2,3))
combination((2,3,4))
combination((1,2,3,4))

for i in 1:15
    print(count_ones(i)," ")
end

function one_positions_array(x::Int64, bits=64)
    out = Vector{Int}([])
    sizehint!(out,count_ones(x))
    for position in 1:bits-leading_zeros(x)
        if x&1 == 1
            push!(out,position)
        end
        x>>>=1
    end
    out
end

function one_positions_mvector(x::Int64, bits=64)
    out = zeros(MVector{count_ones(x),Int64})
    i = 1
    for position in 1:bits-leading_zeros(x)
        if x&1 == 1
            out[i]=position
            i+=1
        end
        x>>>=1
    end
    out
end

function one_positions_array2(x::Int64, bits=64)
    out = zeros(Int64,count_ones(x))
    i = 1
    for position in 1:bits-leading_zeros(x)
        if x&1 == 1
            out[i]=position
            i+=1
        end
        x>>>=1
    end
    out
end

@benchmark one_positions_array(x) setup=(x=rand(Int))
@benchmark one_positions_mvector(x) setup=(x=rand(Int))
@benchmark one_positions_array2(x) setup=(x=rand(Int))

for i in 1:20
    println(i," => ",one_positions(i))
end
for i in 1:20
    print(i," => ")
    for p in OnePositions(i)
        print(p," ")
    end
    println()
end

@benchmark sum(one_positions_array2(x)) setup=(x=rand(Int))
@benchmark sum(OnePositions(x)) setup=(x=rand(Int))

struct Strategy{T}
    groups :: Vector{T}
end

mutable struct Solution{T<:Integer}
    s :: T
    n :: Int
    a :: Array{T,1}
end
solution(s,n) = Solution{Int}(s,n,[])
function canonical!(solution::Solution)
    sort!(solution.a)
end
function isvalid(solution::Solution)
    s = solution.s
    n = solution.n
    a = solution.a
    s<3 && return false
    n<16 && return false
    n>40 && return false
    length(a)<2 && return false
    length(a)>(s-1) && return false
    for x in a
        x<1 && return false
        x>(s-1) && return false
    end
    true
end

function Base.string(sol::Solution)
    out = string(sol.s,"^",sol.n,"=>{",sol.a[1])
    for value in sol.a[2:end]
        out = out*string(",",value)
    end
    out = out*"}"
    out
end

function err(sol::Solution, acc=BigInt(0))
    acc = acc + sol.s^sol.n
    for a in sol.a
        acc = acc - a^sol.n
    end
    acc
end

function random_search(s, n, iterations=1e5,
                             best_e=s^(BigInt(n)),
                             powers=powers_tuple_old(s,n))
    s_to_n = s^(BigInt(n))
    best_a = Vector{Int}()
    all_a = Set(1:s-1)
    gotone = false
    for i in 1:iterations
        x = copy(all_a)
        a = Vector{Int}()
        e = s_to_n
        while true
            length(x)==0 && break
            new_a = rand(x)
            pop!(x,new_a)
            new_e = e - powers[s-new_a]
            abs(new_e)>abs(e) && break
            push!(a,new_a)
            e = new_e
        end
        if abs(e)<abs(best_e)
            gotone = true
            best_e = e
            best_a = a
        end
        if best_e == 0
            break # can't get better
        end
    end
    if gotone
        return Solution(s,n,sort(best_a,rev=true)),best_e
    else
        return nothing
    end
end

function improve!(solution::Solution, strategy::Strategy;
                  powers=powers_tuple(solution.s,solution.n),
                  initial_error=err(solution),
                  s_to_n=solution.s^(BigInt(solution.n))
                  )
    current_error = BigInt(initial_error)
    new_error = BigInt(initial_error)
    current_solution = solution
    while true
        for group in strategy.groups
            new_error = best_group!(current_solution, group,
                                    powers=powers,
                                    initial_error=current_error,
                                    s_to_n = s_to_n)
        end
        new_error>=current_error && break
        current_error = new_error
    end
    solution = current_solution
    current_error
end

function best_group!(solution::Solution, group;
                     powers=powers_tuple(solution.s,solution.n),
                     initial_error=err(solution),
                     s_to_n=solution.s^(BigInt(solution.n))
                    )
    lhs = BigInt(s_to_n)
    a_not_in_group = setdiff(solution.a,group)
    for i in a_not_in_group
        lhs -= powers[i]
    end
    best_e = BigInt(initial_error)
    best_c = Array{Int,1}[]
    for c in combinations(group)
        e = BigInt(lhs)
        for i in c
            e=e-powers[i]
            e<0 && abs(e)>abs(best_e) && break
        end
        if abs(e)<=abs(best_e)
            best_e = e
            best_c = c
        end
        best_e==0 && break
    end
    best_a = sort!(union(best_c,a_not_in_group))
    solution.a = best_a
    best_e
end

powers_tuple(s,n) = ntuple(x->x^BigInt(n),s-1)
powers_tuple_old(s,n) = ntuple(x->(s-x)^BigInt(n),s-1)
function best(s,n,powers=powers_tuple_old(s,n))
    s_to_n = s^(BigInt(n))
    best_e = s_to_n
    all = collect(s-1:-1:1)
    e = s_to_n
    for i in all
        e = e - powers[s-i]
    end
    if e>=0
        return Solution(s,n,all),e
    end
    best_c = Array{Int,1}[]
    for c = combinations(s-1:-1:1)
        e = s_to_n
        for i in c
            e = e - powers[s-i]
            e<0 && abs(e)>abs(best_e) && break # teminate if hopeless
        end
        if abs(e)<abs(best_e)
            best_e = e
            best_c = c
        end
        if best_e == 0
            break # can't get better
        end
    end
    return Solution(s,n,best_c),best_e
end

function max_all_a(n)
    s=2
    while true
        powers=powers_tuple_old(s,n)
        all = collect(s-1:-1:1)
        e = s^(BigInt(n))
        for i in all
            e = e - powers[s-i]
        end
        if e<0
            return s-1
        end
        s+=1
    end
end

function print_best(io::IO,s,n,powers=powers_tuple_old(s,n))
    sol,e = best(s,n,powers)
    println(io,string(sol),",e=",string(e))
end

function print_best_to_file(s_start,s_stop,n,file="data/n$(n)best.txt")
    open(file,"a") do io
        for s in s_start:s_stop
            print_best(io,s,n)
            flush(io)
        end
    end
end

function ns()
    n = 2:17
    s = [4,5,7,8,10,11,12,14,15,17,18,20,21,23,24,26]
    n,s
end

end # module


using DiffEqFlux, CUDA, MLDatasets, OrdinaryDiffEq, Printf, Test
using Flux.Losses: logitcrossentropy
using Flux.Data: DataLoader
using MLDataUtils: LabelEnc, convertlabel, stratifiedobs

CUDA.allowscalar(false)

function loadmnist(batchsize = bs, train_split = 0.9)
    # Use MLDataUtils LabelEnc for natural onehot conversion
    onehot(labels_raw) = convertlabel(LabelEnc.OneOfK, labels_raw,
                                      LabelEnc.NativeLabels(collect(0:9)))
    # Load MNIST
    imgs, labels_raw = MNIST.traindata()
    # Process images into (H,W,C,BS) batches
    x_data = Float32.(reshape(imgs, size(imgs,1), size(imgs,2), 1, size(imgs,3)))
    y_data = onehot(labels_raw)
    (x_train, y_train), (x_test, y_test) = stratifiedobs((x_data, y_data),
                                                         p = train_split)
    return (
        # Use Flux's DataLoader to automatically minibatch and shuffle the data
        DataLoader(gpu.(collect.((x_train, y_train))); batchsize = batchsize,
                   shuffle = true),
        # Don't shuffle the test data
        DataLoader(gpu.(collect.((x_test, y_test))); batchsize = batchsize,
                   shuffle = false)
    )
end

# Main
const bs = 128
const train_split = 0.9
train_dataloader, test_dataloader = loadmnist(bs, train_split)

down = Chain(Conv((3, 3), 1=>64, relu, stride = 1), GroupNorm(64, 64),
             Conv((4, 4), 64=>64, relu, stride = 2, pad=1), GroupNorm(64, 64),
             Conv((4, 4), 64=>64, stride = 2, pad = 1)) |>gpu

dudt = Chain(Conv((3, 3), 64=>64, tanh, stride=1, pad=1),
             Conv((3, 3), 64=>64, tanh, stride=1, pad=1)) |>gpu

fc = Chain(GroupNorm(64, 64), x -> relu.(x), MeanPool((6, 6)),
           x -> reshape(x, (64, :)), Dense(64,10)) |> gpu

nn_ode = NeuralODE(dudt, (0.f0, 1.f0), Tsit5(),
                   save_everystep = false,
                   reltol = 1e-3, abstol = 1e-3,
                   save_start = false) |> gpu

function DiffEqArray_to_Array(x)
    xarr = gpu(x)
    return xarr[:,:,:,:,1]
end

# Build our over-all model topology
model = Chain(down,                 # (28, 28, 1, BS) -> (6, 6, 64, BS)
              nn_ode,               # (6, 6, 64, BS) -> (6, 6, 64, BS, 1)
              DiffEqArray_to_Array, # (6, 6, 64, BS, 1) -> (6, 6, 64, BS)
              fc)                   # (6, 6, 64, BS) -> (10, BS)

# To understand the intermediate NN-ODE layer, we can examine it's dimensionality
img, lab = train_dataloader.data[1][:, :, :, 1:1], train_dataloader.data[2][:, 1:1]

x_d = down(img)

# We can see that we can compute the forward pass through the NN topology
# featuring an NNODE layer.
x_m = model(img)

classify(x) = argmax.(eachcol(x))

function accuracy(model, data; n_batches = 10)
    total_correct = 0
    total = 0
    for (i, (x, y)) in enumerate(data)
        # Only evaluate accuracy for n_batches
        i > n_batches && break
        target_class = classify(cpu(y))
        predicted_class = classify(cpu(model(x)))
        total_correct += sum(target_class .== predicted_class)
        total += length(target_class)
    end
    return total_correct / total
end

# burn in accuracy
accuracy(model, train_dataloader)

loss(x, y) = logitcrossentropy(model(x), y)

# burn in loss
loss(img, lab)

opt = ADAM(0.05)
iter = 0

cb() = begin
    global iter += 1
    # Monitor that the weights do infact update
    # Every 10 training iterations show accuracy
    if iter % 10 == 1
        train_accuracy = accuracy(model, train_dataloader) * 100
        test_accuracy = accuracy(model, test_dataloader;
                                 n_batches = length(test_dataloader)) * 100
        @printf("Iter: %3d || Train Accuracy: %2.3f || Test Accuracy: %2.3f\n",
                iter, train_accuracy, test_accuracy)
    end
end

Flux.train!(loss, Flux.params(down, nn_ode.p, fc), train_dataloader, opt, cb = cb)
@test accuracy(model, test_dataloader; n_batches = length(test_dataloader)) > 0.8