module LinearMaps

export LinearMap

using LinearAlgebra
using SparseArrays

abstract type LinearMap{T} end

Base.eltype(::LinearMap{T}) where {T} = T

Base.isreal(A::LinearMap) = eltype(A) <: Real
LinearAlgebra.issymmetric(::LinearMap) = false # default assumptions
LinearAlgebra.ishermitian(A::LinearMap{<:Real}) = issymmetric(A)
LinearAlgebra.ishermitian(::LinearMap) = false # default assumptions
LinearAlgebra.isposdef(::LinearMap) = false # default assumptions

Base.ndims(::LinearMap) = 2
Base.size(A::LinearMap, n) = (n==1 || n==2 ? size(A)[n] : error("LinearMap objects have only 2 dimensions"))
Base.length(A::LinearMap) = size(A)[1] * size(A)[2]

Base.:(*)(A::LinearMap, x::AbstractVector) = mul!(similar(x, promote_type(eltype(A), eltype(x)), size(A, 1)), A, x)
function LinearAlgebra.mul!(y::AbstractVector, A::LinearMap{T}, x::AbstractVector, α::Number=one(T), β::Number=zero(T)) where {T}
    length(y) == size(A, 1) || throw(DimensionMismatch("mul!"))
    if α == 1
        β == 0 && (A_mul_B!(y, A, x); return y)
        β == 1 && (y .+= A * x; return y)
        # β != 0, 1
        rmul!(y, β)
        y .+= A * x
        return y
    elseif α == 0
        β == 0 && (fill!(y, zero(eltype(y))); return y)
        β == 1 && return y
        # β != 0, 1
        rmul!(y, β)
        return y
    else # α != 0, 1
        β == 0 && (A_mul_B!(y, A, x); rmul!(y, α); return y)
        β == 1 && (y .+= rmul!(A * x, α); return y)
        # β != 0, 1
        rmul!(y, β)
        y .+= rmul!(A * x, α)
        return y
    end
end
# the following is of interest in, e.g., subspace-iteration methods
function LinearAlgebra.mul!(Y::AbstractMatrix, A::LinearMap{T}, X::AbstractMatrix, α::Number=one(T), β::Number=zero(T)) where {T}
    (size(Y, 1) == size(A, 1) && size(X, 1) == size(A, 2) && size(Y, 2) == size(X, 2)) || throw(DimensionMismatch("mul!"))
    @inbounds @views for i = 1:size(X, 2)
        mul!(Y[:, i], A, X[:, i], α, β)
    end
    # starting from Julia v1.1, we could use the `eachcol` iterator
    # for (Xi, Yi) in zip(eachcol(X), eachcol(Y))
    #     mul!(Yi, A, Xi, α, β)
    # end
    return Y
end

A_mul_B!(y::AbstractVector, A::AbstractMatrix, x::AbstractVector)  = mul!(y, A, x)
At_mul_B!(y::AbstractVector, A::AbstractMatrix, x::AbstractVector) = mul!(y, transpose(A), x)
Ac_mul_B!(y::AbstractVector, A::AbstractMatrix, x::AbstractVector) = mul!(y, adjoint(A), x)

# Matrix: create matrix representation of LinearMap
function Base.Matrix(A::LinearMap)
    M, N = size(A)
    T = eltype(A)
    mat = Matrix{T}(undef, (M, N))
    v = fill(zero(T), N)
    @inbounds for i = 1:N
        v[i] = one(T)
        mul!(view(mat, :, i), A, v)
        v[i] = zero(T)
    end
    return mat
end

Base.Array(A::LinearMap) = Matrix(A)
Base.convert(::Type{Matrix}, A::LinearMap) = Matrix(A)
Base.convert(::Type{Array}, A::LinearMap) = Matrix(A)
Base.convert(::Type{SparseMatrixCSC}, A::LinearMap) = sparse(A)

# sparse: create sparse matrix representation of LinearMap
function SparseArrays.sparse(A::LinearMap{T}) where {T}
    M, N = size(A)
    rowind = Int[]
    nzval = T[]
    colptr = Vector{Int}(undef, N+1)
    v = fill(zero(T), N)
    Av = Vector{T}(undef, M)

    for i = 1:N
        v[i] = one(T)
        mul!(Av, A, v)
        js = findall(!iszero, Av)
        colptr[i] = length(nzval) + 1
        if length(js) > 0
            append!(rowind, js)
            append!(nzval, Av[js])
        end
        v[i] = zero(T)
    end
    colptr[N+1] = length(nzval) + 1

    return SparseMatrixCSC(M, N, colptr, rowind, nzval)
end

include("transpose.jl") # transposing linear maps
include("linearcombination.jl") # defining linear combinations of linear maps
include("composition.jl") # composition of linear maps
include("wrappedmap.jl") # wrap a matrix of linear map in a new type, thereby allowing to alter its properties
include("uniformscalingmap.jl") # the uniform scaling map, to be able to make linear combinations of LinearMap objects and multiples of I
include("functionmap.jl") # using a function as linear map
include("blockmap.jl") # block linear maps

"""
    LinearMap(A; kwargs...)
    LinearMap{T=Float64}(f, [fc,], M::Int, N::Int = M; kwargs...)

Construct a linear map object, either from an existing `LinearMap` or `AbstractMatrix` `A`,
with the purpose of redefining its properties via the keyword arguments `kwargs`, or
from a function or callable object `f`. In the latter case, one also needs to specify
the size of the equivalent matrix representation `(M, N)`, i.e. for functions `f` acting
on length `N` vectors and producing length `M` vectors (with default value `N=M`). Preferably,
also the `eltype` `T` of the corresponding matrix representation needs to be specified, i.e.
whether the action of `f` on a vector will be similar to e.g. multiplying by numbers of type `T`.
If not specified, the devault value `T=Float64` will be assumed. Optionally, a corresponding
function `fc` can be specified that implements the transpose/adjoint of `f`.

The keyword arguments and their default values for functions `f` are
*   issymmetric::Bool = false : whether `A` or `f` acts as a symmetric matrix
*   ishermitian::Bool = issymmetric & T<:Real : whether `A` or `f` acts as a Hermitian matrix
*   isposdef::Bool = false : whether `A` or `f` acts as a positive definite matrix.
For existing linear maps or matrices `A`, the default values will be taken by calling
`issymmetric`, `ishermitian` and `isposdef` on the existing object `A`.

For functions `f`, there is one more keyword arguments
*   ismutating::Bool : flags whether the function acts as a mutating matrix multiplication
    `f(y,x)` where the result vector `y` is the first argument (in case of `true`),
    or as a normal matrix multiplication that is called as `y=f(x)` (in case of `false`).
    The default value is guessed by looking at the number of arguments of the first occurence
    of `f` in the method table.
"""
LinearMap(A::Union{AbstractMatrix, LinearMap}; kwargs...) = WrappedMap(A; kwargs...)
LinearMap(f, M::Int; kwargs...) = LinearMap{Float64}(f, M; kwargs...)
LinearMap(f, M::Int, N::Int; kwargs...) = LinearMap{Float64}(f, M, N; kwargs...)
LinearMap(f, fc, M::Int; kwargs...) = LinearMap{Float64}(f, fc, M; kwargs...)
LinearMap(f, fc, M::Int, N::Int; kwargs...) = LinearMap{Float64}(f, fc, M, N; kwargs...)

LinearMap{T}(A::Union{AbstractMatrix, LinearMap}; kwargs...) where {T} = WrappedMap{T}(A; kwargs...)
LinearMap{T}(f, args...; kwargs...) where {T} = FunctionMap{T}(f, args...; kwargs...)

end # module


struct Pattern{
        A<:StepRangeLen,
        E<:StepRangeLen,
        V<:Matrix
    }
    azs::A
    els::E
    values::V
    max_el::Float64
end

struct Pattern3D{T}
    X::Matrix{T}
    Y::Matrix{T}
    Z::Matrix{T}
    gains::Matrix{T}
    max_el::Float64
end

function Pattern(manifold, reduce_ant_fun = norm; num_az = 360, num_el = 91, max_el = π / 2)
    azs = range(0, step = 2π/num_az, length = num_az)
    els = range(0, stop = max_el, length = num_el)
    values = [reduce_ant_fun(get_steer_vec(manifold, Spherical(1.0, az, π / 2 - el))) for el in els, az in azs]
    Pattern(azs, els, values, max_el)
end

@recipe function f(pattern::Pattern;)
    seriestype := :heatmap
    seriescolor --> :viridis
    colorbar_title --> "Magnitude"
    projection --> :polar
    ylims --> (0, pattern.max_el)
    pattern.azs, pattern.els, pattern.values
end

function Pattern3D(manifold, reduce_ant_fun = norm; num_az = 360, num_el = 181, max_el = π)
    azs = range(0, step = 2π/num_az, length = num_az)
    els = range(0, stop = max_el, length = num_el)
    doas_sph = [Spherical(1.0, az, π / 2 - el) for el in els, az in azs]
    doas_cart = map(doa_sph -> CartesianFromSpherical()(doa_sph), doas_sph)
    gains = [reduce_ant_fun(get_steer_vec(manifold, doa_sph)) for doa_sph in doas_sph]
    scaled_doas_cart = doas_cart .* gains
    X = map(doa -> doa[1], scaled_doas_cart)
    Y = map(doa -> doa[2], scaled_doas_cart)
    Z = map(doa -> doa[3], scaled_doas_cart)
    Pattern3D(X, Y, Z, gains, Float64(max_el))
end

@recipe function f(pattern::Pattern3D;)
    seriestype := :surface
    seriescolor --> :viridis
    pattern.X, pattern.Y, pattern.Z
end


#
# Denoise an S2 valued image demonstrating the effect of using different base points m
#
# L2-TV functional with anisotropic TV
#
# This example is used in Section 6.2 of
#
# > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal Núñez:
# > _Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds_,
# > arXiv: [1908.02022](https://arxiv.org/abs/1908.02022)
#
using CSV, DataFrames, LinearAlgebra, Manopt, Manifolds
#
# Settings
experiment_name = "S2_WhirlImage_CP"
export_orig = false
export_primal = false
export_primal_video = false
export_table = false
asy_render_size = 2
#
# Automatic Settings
export_any = export_orig || export_primal || export_primal_video || export_table
results_folder = joinpath(@__DIR__, "Image_TV")
video_folder = joinpath(results_folder, "video")
(export_any && !isdir(results_folder)) && mkdir(results_folder)
(export_primal_video && !isdir(video_folder)) && mkdir(video_folder)
#
# Experiment Parameters
σ = 0.25
τ = 0.25
γ = 0.2
θ = 1.0
α = 1.5
#
# Data
pixelM = Sphere(2)
f = artificial_S2_whirl_image(64)
#
# load TV model
#
include("Image_TV_commons.jl")

if export_orig
    orig_file = joinpath(results_folder, experiment_name * "-orig.asy")
    asymptote_export_S2_data(orig_file; data=f)
    render_asymptote(orig_file; render=asy_render_size)
end
struct DebugRenderAsy <: DebugAction
    name::String
    folder::String
    DebugRenderAsy(f, n) = new(n, f)
end
function (d::DebugRenderAsy)(::PrimalDualProblem, ::ChambollePockOptions, i)
    if i >= 0
        orig_file = joinpath(d.folder, d.name * "-vid-$(lpad(i[1],7,"0")).asy")
        asymptote_export_S2_data(orig_file; data=f)
        render_asymptote(orig_file; render=asy_render_size)
    end
end
#
# Experiments (entries overwrite defaults from below)
data_mean = mean(pixelM, vec(f), GradientDescentEstimation(); stop_iter=5)
mean_image = fill(data_mean, size(f))
west_image = fill([1.0, 0.0, 0.0], size(f))
#
# Build Experiments
experiments = [
    Dict(:name => "mMean", :m => deepcopy(mean_image)),
    Dict(:name => "mWest", :m => deepcopy(west_image)),
]

#
# Defaults
x0 = deepcopy(f)
m = deepcopy(mean_image)
records = Array{Array{Tuple{Int,Float64,Array},1},1}()
for e in experiments
    name = e[:name]
    print("\n --- Running Experiment $name ---")
    println(
        " Values σ:",
        get(e, :σ, σ),
        " | τ:",
        get(e, :τ, τ),
        " | θ:",
        get(e, :θ, θ),
        " | γ:",
        get(e, :γ, γ),
    )
    #
    # Any Entry in the dictionary overwrites the above default
    @time o = ChambollePock(
        M,
        N,
        cost,
        get(e, :x, x0),
        get(e, :ξ, zero_vector(N, get(e, :n, Λ(get(e, :m, m))))),
        get(e, :m, m),
        get(e, :n, Λ(get(e, :m, m))),
        proxFidelity,
        proxPriorDual,
        AdjDΛ;
        linearized_forward_operator=DΛ,
        primal_stepsize=get(e, :σ, σ),
        dual_stepsize=get(e, :τ, τ),
        relaxation=get(e, :θ, θ),
        acceleration=get(e, :γ, γ),
        relax=:dual,
        debug=[
            :Iteration,
            " | ",
            :Cost,
            "\n",
            export_primal_video ? DebugRenderAsy(video_folder, experiment_name * name) : "",
            100,
            :Stop,
        ],
        record=[:Iteration, :Cost, :Iterate],
        stopping_criterion=StopAfterIteration(get(e, :maxIter, 300)),
        variant=:linearized,
        return_options=true,
    )
    push!(records, get_record(o))
    if export_primal
        result_file = joinpath(results_folder, experiment_name * "-result.asy")
        asymptote_export_S2_data(result_file; data=get_solver_result(o))
        render_asymptote(result_file; render=asy_render_size)
    end
    if export_primal_video
        source = join_path(video_folder, experiment_name * name * "-vid-*.png")
        dest = join_path(results_folder, experiment_name * name * ".mp4")
        # run(`ffmpeg -framerate 15 -pattern_type glob -i $(source) -c:v libx264 -vf pad=ceil\(iw/2\)*2:ceil\(ih/2\)*2 -pix_fmt yuv420p $(dest)`)
    end
end
#
# Finalize - export costs
if export_table
    A = cat(
        first.(records[1]),
        [[r[2] for r in records[i]] for i in 1:length(records)]...;
        dims=2,
    )
    CSV.write(
        joinpath(results_folder, experiment_name * "-Result.csv"),
        DataFrame(A);
        header=["i", [e[:name] for e in experiments]...],
    )
end


module AffineSpaces

using FixedSizeArrays

export AffineSpace, FixedAffineSpace, AffinePoint

"""
    abstract AffineSpace{T}

An `AffineSpace` object represents points in an affine space, which is similar
to a Cartesian space but where no special meaning can be assigned to the origin.
Subtypes of `AffineSpace` should have a field named `pos` with vector-like
abilities describing the coordinates.

The consequences are that points do not have vector-like abilities such as
multiplication by a scalar (scaling the distance relative to the origin), unary
inversion (reflection about the origin) or addition of the coordinates of two
points (coorinates being relative to the origin).

Outside these restriction, one can take affine combinations of points, finding
for instance the midpoint of two points. Affine combinations are built with
the colon operation `p₁:p₂` and indexed with a scalar `(p₁:p₂)[i]`, where
`i = 0.0` indicates the point `p₁`, and `i = 1.0` indicates the point `p₂`, and
other values of `i` trace out the line connecting the two points. Generally,
the affine combination of `n` points can be generated with `(p₁:...:pₙ)[i₁,...,iₙ]`
so long as the `iⱼ`s sum to `1` (or the last `iₙ` is not given, and is
automatically calculated).

Furthermore, standard vectors can be added and subtracted from affine points as
a displacement of the point. The displacement vector between two points is given
by their subtraction.

The goal of these types are to provide some programmer safety to avoid
conceptual mistakes regarding manipulation of points, by removing the meaning
of the possibly-arbitrary origin.
"""
abstract AffineSpace{T}
abstract FixedAffineSpace{N,T} <: AffineSpace

"""
    AffinePoint{N,T} <: FixedAffineSpace{T}

A concrete type describing a point in an affine space, where the Cartesian
coordinates of the point are stored in a `FixedSizeArrays.Vec{N,T}`. See
`AffineSpace` for a description of the interface for affine points.
"""
immutable AffinePoint{N,T} <: FixedAffineSpace{T}
    pos::Vec{N,T}
end

Base.length(point::AffineSpace) = length(point.pos)
Base.size(point::AffineSpace) = size(point.pos)
Base.size(point::AffineSpace, d) = size(point.pos, d)
Base.endof(point::AffineSpace) = endof(point.pos)
Base.getindex(point::AffineSpace, i) = point.pos[i]
Base.setindex!(point::AffineSpace, v, i) = setindex!(point.pos, v, i)
Base.start(point::AffineSpace) = start(point.pos)
Base.next(point::AffineSpace, i) = next(point.pos, i)
Base.done(point::AffineSpace, i) = done(point.pos, i)

function Base.show(io::IO, point::AffinePoint)
    print(io, "AffinePoint(")
    showcompact(io, point.pos)
    print(io,")")
end
function Base.showcompact(io::IO, point::AffinePoint)
    showcompact(io, point.pos)
end

# Adding and subtracting vectors from points
Base.(:+){P <: AffineSpace}(point::P, vec::Union{AbstractVector,FixedVector}) = P(point.pos + vec)
Base.(:+){P <: AffineSpace}(vec::Union{AbstractVector,FixedVector}, point::P) = P(vec + point.pos)

Base.(:-){P <: AffineSpace}(point::P, vec::Union{AbstractVector,FixedVector}) = P(point.pos + vec)

# Subtract to give vectors - uses underlying vector structure of the .pos field
Base.(:-)(point1::AffineSpace, point2::AffineSpace) = point1.pos - point2.pos

# Nice messages for disallowed operations
Base.(:+)(point1::AffineSpace, point2::AffineSpace) = error("Cannot add affine points. Consider adding a displacement vector instead.")
Base.(:*)(point::AffineSpace, x::Number) = error("Cannot scale affine points, since scaling is relative to the origin.")
Base.(:*)(x::Number, point::AffineSpace) = error("Cannot scale affine points, since scaling is relative to the origin.")
Base.(:-)(point::AffineSpace) = error("The additive inverse of an affine point is not defined, since affine points cannot be added.")
Base.(:-)(vec::Union{AbstractVector,FixedVector}, point::AffineSpace) = error("Cannot subtract an affine point from a vector.")

# Taking affine combinations of points
immutable AffineSubSpace{N, P<:AffineSpace}
    points::NTuple{N,P}
end

Base.start(ss::AffineSubSpace) = 1
Base.next(ss::AffineSubSpace, n::Int) = (ss.points[n], n+1)
Base.done{N}(ss::AffineSubSpace{N}, n::Int) = n > N

function Base.getindex{N, P}(ss::AffineSubSpace{N, P}, x...)
    lx = length(x)
    if lx == N
        @assert sum(x) ≈ 1 # maybe consider this a boundscheck?
        return P(sum(map(*, x, map(p -> p.pos, ss))))
    elseif lx == N - 1
        x = (x..., 1 - sum(x))
        return P(sum(map(*, x,  map(p -> p.pos, ss))))
    else
        error("Must specify either $N or $(N-1) numbers to define affine combination of $N points - got $lx")
    end
end

# Allow syntax like (point1:point2)[0.5] to get midpoint between point1 and point2
Base.colon(points::AffineSpace...) = AffineSubSpace(points)
Base.colon(ss::AffineSubSpace, points::AffineSpace...) = AffineSubSpace((ss.points..., points...))

# Base.mean as special case of affine combination
Base.mean{N,P <: AffineSpace}(points::NTuple{N,P}) = P(sum(map(p -> p.pos, points)) / N)

end # module


function xorshift32(seed=UInt32(2463534242))
    y = seed
    inner=function counter()
    	y = y $ (y << 13)
    	y = y $ (y >> 17)
    	y = y $ (y << 5)
    	return y
    end
    return inner
end

function uniform(generator,b=0.0,e=1.0)
	inner=function inner()
		value=b+generator()/typemax(UInt32)/(e-b)
		return value
	end
	return inner
end

function u01(generator)
	inner=function inner()
		value=generator()/typemax(UInt32)
		return value
	end
	return inner
end

function main()
	#uniform01=uniform(xorshift32())
	uniform01=u01(xorshift32())
	N=10000000
	count=0
	for i = 1:N
		x=uniform01()
		y=uniform01()
		if x*x+y*y<1.0
			count+=1
		end
	end
	println("Pi=",4.0*count/N)
end

main()


if VERSION > v"0.6.3"
    using Printf
end

include("alias.jl")
include("utils.jl")
include("PGBF.jl")
include("CGBF.jl")
include("overlap.jl")

#=
using PyPlot
const plt = PyPlot

function do_plot_x( bf::PGBF, filesave::String )
    const N = 100
    x = linspace(-5.0,5.0,N)
    y = zeros(N)
    for i = 1:N
      y[i] = evaluate( bf, x[i], 0.0, 0.0 )
    end
    plt.clf()
    plt.plot( x, y, marker="o" )
    plt.grid()
    plt.savefig(filesave, dpi=300)
end
=#

function test_main()
    bf1 = PGBF(0.3, (0.0,0.0,0.0), (1,0,0))
    println(bf1)
    #do_plot_x(bf1,"bf1.png")

    bf2 = PGBF(0.4,(0.1,0.0,0.0), (1,0,0))
    println(bf2)
    #do_plot_x(bf2,"bf2.png")

    S1 = overlap(bf1,bf1)
    @printf("overlap(bf1,bf1) = %18.10f\n", S1)

    S1 = overlap(bf2,bf2)
    @printf("overlap(bf2,bf2) = %18.10f\n", S1)

    S1 = overlap(bf1,bf2)
    @printf("overlap(bf1,bf2) = %18.10f\n", S1)
end

@time test_main()


using Test

using Downhill
using Downhill: argumentvec, gradientvec, step_origin, fnval, fnval_origin
using Downhill: mcholesky!

using LinearAlgebra

using Random

const test_rng = MersenneTwister(123581321)

"""
    rosenbrock!(x, g; b=2)

Return the value of Rosenbrock function in `length(x)` dimensions.
"""
function rosenbrock!(x::AbstractVector, g::AbstractVector; b=2)
    f = zero(eltype(g))
    fill!(g, 0)
    inds = eachindex(x, g)
    for i in 2:last(inds)
        f += (1 - x[i-1])^2 + b * (x[i] - x[i-1]^2)^2
        g[i-1] += 2 * (x[i-1] - 1) + 4 * b * x[i-1] * (x[i-1]^2 - x[i])
        g[i] += 2 * b * (x[i] - x[i-1]^2)
    end
    return f, g
end

OPT_TYPES = (
    SteepestDescent,
    HyperGradDescent,
    FixedRateDescent,
    MomentumDescent,
    NesterovMomentum,
    CGDescent,
    BFGS,
    CholBFGS,
)

@testset "Testing package DescentMethods" begin
# Test that all methods have a constructor
include("constructors.jl")

# Test that all methods have interface functions
include("interface.jl")

# Test linesearch on quadratic functions
include("linesearch.jl")

# Test modified Cholesky decomposition
include("cholesky.jl")

# Test convergence on the Rosenbrock function
include("convergence.jl")
end


using ModelingToolkit
using Test

@parameters t
@variables x(t)
@variables y(t)
@variables z(t)
x1 = Variable(:x, [t])
y1 = Variable(:y, [t])
z1 = Variable(:z, [t])
@test isequal(x1, x)
@test isequal(y1, y)
@test isequal(z1, z)
@test convert(Expr, x) == :x
@test convert(Expr, y) == :y
@test convert(Expr, z) == :z

@parameters begin
    t
    s
end
t1 = Variable(:t; known = true)
s1 = Variable(:s; known = true)
@test isequal(t1, t)
@test isequal(s1, s)
@test convert(Expr, t) == :t
@test convert(Expr, s) == :s
@test convert(Expr, cos(t + sin(s))) == :(cos(t + sin(s)))

@derivatives D'~t
D1 = Differential(t)
@test D1 == D
@test convert(Expr, D) == D

@test isequal(x ≤ y + 1, (x < y + 1) | (x == y + 1))


struct AngularPositionAroundPoint{M,N,S,m} <: TaskMap{M,N,S}
    position_center::SVector{m,S}
end
AngularPositionAroundPoint{M,N}(c::SVector{m,S}) where {M,N,S,m} =
    AngularPositionAroundPoint{M,N,S,dim(M)}(c)

struct AngularPositionAroundPointT{M,N,S,m} <: TaskMapT{M,N,S}
    base_map::AngularPositionAroundPoint{M,N,S,m}
end
AngularPositionAroundPointT{M,N}(c::SVector{m,S}) where {M,N,S,m} =
    AngularPositionAroundPointT{M,N,S,m}(AngularPositionAroundPoint{M,N}(c))

# S1
codomain_coord_rep(::AngularPositionAroundPoint{R2,S1}) = ChartRep()

function task_map_emb_chart(::EmbRep, ::ChartRep, xme, task_map::AngularPositionAroundPoint{R2,S1},
        ::Chart{Angleπ,S1})
    a = xme - task_map.position_center
    d = norm(a)
    if d == 0
        return SA[0.]
    else
        â = a ./ norm(a)
        return SA[atan(â[2], â[1])]
    end
end

function task_map_emb_chart(::EmbRep, ::ChartRep, xme, task_map::AngularPositionAroundPoint{R2,S1},
        C2π::Chart{Angle2π,S1})
    Cπ = Chart{Angleπ,S1}()
    xnπ = task_map_emb_chart(xme, task_map, Cπ)
    chart_transition(xnπ, Cπ, C2π)
end

# S2
codomain_coord_rep(::AngularPositionAroundPoint{R3,S2}) = EmbRep()

function task_map_emb(::EmbRep, ::EmbRep, xme, task_map::AngularPositionAroundPoint{R3,S2})
    Δx = xme - task_map.position_center
    d = norm(Δx)
    if d == 0
        return @SVector zeros(3)
    else
        return Δx/norm(Δx)
    end
end


export mv_avg_2d, delta_h, envelope, mute, mute!, laplacian, offset_map

"""
    mv_avg_2d(x; k=5)

2D moving average with a square window of width k
"""
function mv_avg_2d(x::AbstractArray{T, 2}; k=5) where T
    out = 1f0 * x
    nx, ny = size(x)
    kl, kr = k÷2 + 1, k÷2
    for i=kl:nx-kr, j=kl:ny-kr
        out[i, j] = sum(x[i+k, j+l] for k=-kr:kr, l=-kr:kr)/k^2
    end
    out
end

"""
    envelope(x)

Envelope of a vector or a 2D matrix. The envelope over the first dimension is taken for a 2D matrix (see DSP `hilbert`)
"""
envelope(x::AbstractArray{T, 1}) where T = abs.(hilbert(x))

"""
    envelope(x)

Envelope of a vector or a 1D vector (see DSP `hilbert`)

"""
envelope(x::AbstractArray{T, 2}) where T = abs.(hilbert(x))

"""
    delta_h(ha, h, tol)

Compute the binary mask where `ha` is within `tol` of `h`.
"""
delta_h(ha::AbstractArray{T, 2}, h::Number, tol::Number) where T = Float32.(abs.(h .- ha) .<= tol)

"""
    mute!(shot, offsets;vp=1500, t0=1/10, dt=0.004)

In place direct wave muting of a shot record with water sound speed `vp`, time sampling `dt` and firing time `t0`.
"""
function mute!(shot::AbstractArray{Ts, 2}, offsets::Vector{To}; vp=1500, t0=1/10, dt=.004) where {To, Ts}
    length(offsets) == size(shot, 2) || throw(DimensionMismatch("Number of offsets has to match the number of traces"))
    inds = trunc.(Integer, (offsets ./ vp .+ t0) ./ dt)
    inds = min.(max.(1, inds), size(shot, 1))
    for (rx, i) = enumerate(inds)
        shot[1:i, rx] .= 0f0
    end
end

"""
    mute(shot, offsets;vp=1500, t0=1/10, dt=0.004)

Direct wave muting of a shot record with water sound speed `vp`, time sampling `dt` and firing time `t0`.
"""
function mute(shot::AbstractArray{Ts, 2}, offsets::Vector{To}; vp=1500, t0=1/10, dt=.004) where {To, Ts}
    out = Ts(1) .* shot
    mute!(out, offsets; vp=vp, t0=t0, dt=dt)
    out
end

"""
    laplacian(image; hx=1, hy=1)

2D laplacian of an image with grid spacings (hx, hy)
"""
function laplacian(image::AbstractArray{T, 2}; hx=1, hy=1) where T
    scale = 1/(hx*hy)
    out = 1 .* image
    @views begin
        out[2:end-1, 2:end-1] .= -4 .* image[2:end-1, 2:end-1]
        out[2:end-1, 2:end-1]  .+= image[1:end-2, 2:end-1] + image[2:end-1, 1:end-2]
        out[2:end-1, 2:end-1]  .+= image[3:end, 2:end-1] + image[2:end-1, 3:end]
    end
    return scale .* out
end

"""
    offset_map(rtm, rtmo; scale=0)

Return the regularized least-square division of rtm and rtmo. The regularization consists of the envelope and moving average
followed by the least-square division ['surface_gather'](@ref)
"""
function offset_map(rtm::AbstractArray{T, 2}, rtmo::AbstractArray{T, 2}; scale=0) where T
    rtmn = mv_avg_2d(envelope(rtm))
    rtmo = mv_avg_2d(envelope(rtmo))

    offset_map = rtmn .* rtmo ./ (rtmn .* rtmn .+ eps(Float32)) .- scale
    return offset_map
end


mutable struct Ash2{R1 <: AbstractRange, R2 <: AbstractRange, F1 <: Function, F2 <: Function}
    kernelx::F1
    kernely::F2
    rngx::R1               # rng for x
    rngy::R2               # rng for y
    mx::Int64
    my::Int64
    v::Matrix{Int}      # counts
    z::Matrix{Float64}  # density
    nobs::Int64

    function Ash2(kx::F1, ky::F2, rngx::R1, rngy::R2, mx::Int, my::Int) where
            {F1 <: Function, F2 <: Function, R1 <: AbstractRange, R2 <: AbstractRange}

        v = zeros(Int, length(rngy), length(rngx))
        z = zeros(Float64, length(rngy), length(rngx))
        new{R1, R2, F1, F2}(kx, ky, rngx, rngy, mx, my, v, z, 0)
    end
end
function Base.show(io::IO, o::Ash2)
    println(io, "Ash2")
    f, l, s = round.((first(o.rngx), last(o.rngx), step(o.rngx)), digits=4)
    print(io, "X:")
    println(io, "  > edges  | $f : $s : $l")
    println(io, "    > kernel | $(o.kernelx)")
    println(io, "    > m      | $(o.mx)")
    f, l, s = round.((first(o.rngy), last(o.rngy), step(o.rngy)), digits=4)
    print(io, "Y:")
    println(io, "  > edges  | $f : $s : $l")
    println(io, "    > kernel | $(o.kernely)")
    println(io, "    > m      | $(o.my)")
    print(io, UnicodePlots.spy(reverse(o.z, dims=1)).graphics)
end

function Base.:(==)(o::Ash2, o2::Ash2)
    fns = fieldnames(typeof(o))
    all(getfield.(Ref(o), fns) .== getfield.(Ref(o2), fns))
end

function _histogram!(o::Ash2, x::AbstractVector, y::AbstractVector)
    length(x) == length(y) || throw(ArgumentError("data lengths differ"))
    δx, δy = step(o.rngx), step(o.rngy)
    ax, ay = first(o.rngx), first(o.rngy)
    nx, ny = length(o.rngx), length(o.rngy)
    for i in eachindex(x)
        kx = floor(Int, (x[i] - ax) / δx + 1.5)
        ky = floor(Int, (y[i] - ay) / δy + 1.5)
        if (1 ≤ kx ≤ nx) && (1 ≤ ky ≤ ny)
            o.v[ky, kx] += 1
        end
    end
    o.nobs += length(x)
    return o
end

function _ash!(o::Ash2)
    δx, δy = step(o.rngx), step(o.rngy)
    nx, ny = length(o.rngx), length(o.rngy)
    mx, my = o.mx, o.my
    kernelx, kernely = o.kernelx, o.kernely
    w = zeros(2 * my - 1, 2 * mx - 1)
    for i = (1 - my):(my - 1), j = (1 - mx):(mx - 1)
        w[i + my, j + mx] = kernely(i / my) * kernelx(j / mx)
    end
    v, z = o.v, o.z
    fill!(z, 0.0)
    for k = 1:ny
       for l = 1:nx
           if v[k, l] == 0
                l += 1
            else
                for i = max(1, k - my + 1):min(ny, k + my - 1)
                    for j = max(1, l - mx + 1):min(nx, l + mx - 1)
                        z[i, j] += v[k, l] * w[my + i - k, mx + j - l]
                    end
                end
            end
        end
    end
    # make density integrate to 1
    denom = 1.0 / (sum(o.z) * δx * δy)
    for i in eachindex(z)
        @inbounds z[i] *= denom
    end
    o
end

function ash(x, y; nbin=500, nbinx=nbin, nbiny=nbin,
        rngx::AbstractRange = extendrange(x, nbinx), rngy::AbstractRange = extendrange(y, nbiny),
        mx = 5, my = 5, kernelx = Kernels.biweight, kernely = Kernels.biweight)
    o = Ash2(kernelx, kernely, rngx, rngy, mx, my)
    _histogram!(o, x, y)
    _ash!(o)
    return o
end

function ash!(o::Ash2; mx = o.mx, my = o.my, kernelx = o.kernelx, kernely = o.kernely)
    o.mx, o.my, o.kernelx, o.kernely = mx, my, kernelx, kernely
    _ash!(o)
end
function ash!(o::Ash2, x, y; mx = o.mx, my = o.my,
        kernelx = o.kernelx, kernely = o.kernely)
    o.mx, o.my, o.kernelx, o.kernely = mx, my, kernelx, kernely
    _histogram!(o, x, y)
    _ash!(o)
end

"return ranges and density of biviariate ASH"
xyz(o::Ash2) = (o.rngx, o.rngy, copy(o.z))
nobs(o::Ash2) = o.nobs
nout(o::Ash2) = nobs(o) - sum(o.v)
function Statistics.mean(o::Ash2)
    meanx = mean(o.rngx, fweights(vec(sum(o.v, dims=1))))
    meany = mean(o.rngy, fweights(vec(sum(o.v, dims=2))))
    [meanx; meany]
end
function Statistics.var(o::Ash2)
    varx = var(o.rngx, fweights(vec(sum(o.v, dims=1))); corrected=true)
    vary = var(o.rngy, fweights(vec(sum(o.v, dims=2))); corrected=true)
    [varx; vary]
end
Statistics.std(o::Ash2) = sqrt.(var(o))

RecipesBase.@recipe function f(o::Ash2; hist=false)
    seriestype --> :heatmap
    o.rngx, o.rngy, (hist ? o.v : o.z)
end


"""
    UnitVector{T}(N, k)

Helper vector type of length `N` where the `k`th element is `one(T)`
and all the others `zero(T)`.
"""
struct UnitVector{T} <: AbstractVector{T}
    N::Int
    k::Int
end

Base.size(e::UnitVector) = (e.N,)
Base.getindex(e::UnitVector{T}, i::Int) where T = i == e.k ? one(T) : zero(T)
Base.show(io::IO, e::UnitVector{T}) where T = write(io, "ê{$T}($(e.k)|$(e.N))")


using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot
using ProgressMeter, Statistics
using JLD2

include("koch.jl")

function run_exs_takens(k, F, N, eps_step)

    name = savename("takens_epsilonmax_regularized", @dict(k, N,F,eps_step),"jld2")
    isfile(datadir("takens", name)) && return

    systems = ["Koch", "Henon", "Lorenz", "L96_D5","L96_D6","L96_D7","L96_D8","L96_D9"]
    systs = [pointskoch(points, k), Systems.henon(), Systems.lorenz(),Systems.lorenz96(5, F=F),Systems.lorenz96(6, F=F),Systems.lorenz96(7, F=F),Systems.lorenz96(8, F=F),Systems.lorenz96(9, F=F)]

    Delta_T = Dict()

    for (system, ds) in zip(systems, systs)

        println("System: $system")

        X = nothing

        if system == "Koch"
            points = [[0.0; 0.0], [1.0; 0.0]]
            koch = pointskoch(points, k)
            # making sure we have the same amount of points for the Koch curve
            pt = rand(1:length(koch), 40000)
            X = regularize(Dataset(koch[pt]))
        else
            dt = system == "Henon" ? 1 : 0.1
            T = N * dt
            ts = trajectory(ds, T; dt = dt, Ttr = 100)
            X = regularize(ts)
        end

        εs = 0.2:eps_step:1
        α, αmin, αmax = zeros(length(εs)), zeros(length(εs)), zeros(length(εs))

        @showprogress for (i,ε) in enumerate(εs)
            try
                α[i], αmax[i], αmin[i] = takens_best_estimate(X, ε)
            catch
                # An error above means that the estimator diverged
                α[i], αmax[i], αmin[i] = NaN, NaN, NaN
            end
        end
        result = @strdict α αmin αmax

        Delta_T[system] = result
    end
    @tagsave(datadir("takens", name), Delta_T)
end

function plot_ex_takens(F, N, eps_step,k)

    name = savename("takens_epsilonmax_regularized", @dict(N,F,eps_step,k),"jld2")
    plotname = savename("takens_epsilonmax_regularized_all", @dict(N,F, eps_step,k),"png")

    systems = ["Koch", "Henon", "L96_D5","L96_D7","L96_D9"]
    labels = Dict("Koch"=> "Koch", "Henon"=> "Hénon", "L96_D5"=>"Lorenz-96, \$D = 5\$","L96_D7"=>"Lorenz-96, \$D = 7\$","L96_D9"=>"Lorenz-96, \$D = 9\$")

    @load datadir("takens", name) Koch Henon L96_D5 L96_D7 L96_D9

    system_results = [Koch, Henon, L96_D5, L96_D7, L96_D9]

    @load datadir("lorenz96_Dmax=13_F=$(Int(F))_T=100000_ics=10.jld2") K

    K2 = [log(4)/log(3), 1.26] # analytial solution for Koch, Δ^(L) for Hénon

    Ks = [K2...,mean.(K[[1,3,5]])...]

    εs = 0.2:eps_step:1

    fig, axs = plt.subplots(2, sharex=true,  gridspec_kw=Dict("hspace"=> .1),constrained_layout=false)

    mylinestyles = [":", "--", "-.",":", "--"]

    for (k,(system,system_result)) in enumerate(zip(systems, system_results))

        ax = (system == "Henon" || system == "Koch") ? axs[2] : axs[1]

        ax.fill_between(εs, system_result["αmin"], system_result["αmax"], color = COLORS[k], alpha = .5)
        ax.plot(εs,system_result["α"], color = COLORS[k],linestyle = mylinestyles[k],label = labels[system])

        ax.plot(εs, ones(length(εs))*Ks[k], color = COLORS[k], linestyle = "-", lw = 3)
        ax.set_xlim(.2,1)
    end

    handles, labels = axs[1].get_legend_handles_labels()
    handles2, labels2 = axs[2].get_legend_handles_labels()
    xlabel(L"$\varepsilon_{\mathrm{max}}$")
    fig.text(0.04, 0.55, L"$\Delta^{(T)}$", va="center", rotation="vertical", fontsize=30)
    fig.tight_layout()
    fig.subplots_adjust(bottom=0.13, top = .84)
    fig.legend([handles...,handles2...],[labels...,labels2...], loc = "upper center", ncol = 3)

    tight_layout()
    wsave(plotsdir("takens", plotname), fig)
    close("all")
end

N = 40000
F = 8.
k = 8
run_exs_takens(k,F,N,.01)
plot_ex_takens(F, N,.01,k)


type SkewTDist <: ContinuousUnivariateDistribution
    ξ::Float64  # Location
    ω::Float64  # Scale
    α::Float64  # Skewness
    df::Float64 # Degrees of freedom
end

skewtdistpdf(α::Real, df::Real, z::Real) = 2.0* _t₁(z, df) * _T₁(α * z * sqrt((df + 1)/(z^2 + df)), df + 1)

function pdf(dist::SkewTDist, x::Real)
    z = (x - dist.ξ)/dist.ω
    skewtdistpdf(dist.α, dist.df, z)/dist.ω
end

function rand(dist::SkewTDist)
    chisqd = Chisq(dist.df)
    w = rand(chisqd)/dist.df
    Ω = Array(Float64, 1,1)
    Ω[1,1] = dist.ω^2
    sndist = MvSkewNormal(Ω, [dist.α])
    x = rand(sndist)[1]
    return dist.ξ + x/sqrt(w)
    ## Ω = Array(Float64, 1,1)
    ## Ω[1,1] = dist.ω^2
    ## return rand(MvSkewTDist([dist.ξ], Ω, [dist.α], dist.df))[1]
end

function cdf(dist::SkewTDist, x::Real)
    quadgk(t->pdf(dist,t), -Inf, x)[1]
end

function quantile(dist::SkewTDist, β::Float64)
    if dist.α < 0
        β = 1 - β
    end
    a = tdistinvcdf(dist.df, β)
    b = sqrt(fdistinvcdf(1.0, dist.df, β))
    qz = fzero(x->cdf(dist, x) - β, a, b)
end

minimum(dist::SkewTDist) = -Inf
maximum(dist::SkewTDist) = Inf

μ(dist::SkewTDist) = δ(dist.α) * sqrt(dist.df/π) * exp(lgamma(0.5*(dist.df-1.0)) - lgamma(0.5*dist.df))
mean(dist::SkewTDist) = dist.ξ + dist.ω * μ(dist)
var(dist::SkewTDist) = (dist.ω^2) * (dist.df/(dist.df - 2.0)) - dist.ω^2 * μ(dist)^2
dof(dist::SkewTDist) = dist.df


using MOOSE

include("CoupledConvection.jl")

# Create the Mesh
mesh = buildSquare(0, 1, 0, 1, 20, 20)

# Create the System to hold the equations
diffusion_system = System{Float64}(mesh)

# Add some variables
u = addVariable!(diffusion_system, "u")
v = addVariable!(diffusion_system, "v")

# Apply the Laplacian operator to the "u" variable
addKernel!(diffusion_system, Diffusion(u))

# Apply the Coupled Convection operator
addKernel!(diffusion_system, CoupledConvection(u, v))

# Apply the Laplacian operator to the "v" variable
addKernel!(diffusion_system, Diffusion(v))

# u = 0 on the Left
addBC!(diffusion_system, DirichletBC(u, [4], 0.0))

# u = 1 on the Right
addBC!(diffusion_system, DirichletBC(u, [2], 1.0))

# v = 0 on the Left
addBC!(diffusion_system, DirichletBC(v, [4], 0.0))

# v = 1 on the Right
addBC!(diffusion_system, DirichletBC(v, [2], 1.0))

# Initialize the system of equations
initialize!(diffusion_system)

# Create a solver and solve
solver = JuliaDenseNonlinearImplicitSolver(diffusion_system)
solve!(solver, nl_max_its=5)

# Output
out = VTKOutput()
output(out, solver, "coupled_convection_out")


"""
**Laplace likelihood**

Laplace likelihood for regression: ``\\frac{1}{2\\beta}\\exp\\left(-\\frac{|y-f|}{\\beta}\\right)``
see [wiki page](https://en.wikipedia.org/wiki/Laplace_distribution)

```julia
LaplaceLikelihood(β::T=1.0)  #  Laplace likelihood with scale β
```

---

For the analytical solution, it is augmented via:
```math
p(y|f,\\omega) = \\mathcal{N}(y|f,\\omega^{-1})
```
where ``\\omega \\sim \\text{Exp}\\left(\\omega \\mid \\frac{1}{2 \\beta^2}\\right)``, and Exp is the [Exponential distribution](https://en.wikipedia.org/wiki/Exponential_distribution)
We approximate ``q(\\omega) = \\mathcal{GIG}\\left(\\omega \\mid a,b,p\\right)
"""
struct LaplaceLikelihood{T<:Real} <: RegressionLikelihood{T}
    β::LatentArray{T}
    a::LatentArray{T}
    p::LatentArray{T}
    b::LatentArray{Vector{T}}
    θ::LatentArray{Vector{T}} #Expected value of ω
    function LaplaceLikelihood{T}(β::AbstractVector{T}) where {T<:Real}
        new{T}(β,β.^-2,0.5*ones(size(β)))
    end
    function LaplaceLikelihood{T}(β::AbstractVector{T},b::AbstractVector{<:AbstractVector{T}},θ::AbstractVector{<:AbstractVector{T}}) where {T<:Real}
        new{T}(β,β.^(-2),0.5*ones(size(β)),b,θ)
    end
end

function LaplaceLikelihood(β::T=1.0) where {T<:Real}
    LaplaceLikelihood{T}([β])
end

function init_likelihood(likelihood::LaplaceLikelihood{T},inference::Inference{T},nLatent::Int,nSamplesUsed::Int,nFeatures::Int) where T
    if inference isa AnalyticVI || inference isa GibbsSampling
        LaplaceLikelihood{T}(
        likelihood.β,
        [abs2.(T.(rand(T,nSamplesUsed))) for _ in 1:nLatent],
        [zeros(T,nSamplesUsed) for _ in 1:nLatent])
    else
        LaplaceLikelihood{T}(likelihood.β)
    end
end

function pdf(l::LaplaceLikelihood,y::Real,f::Real)
    Distributions.pdf(Laplace(f,l.β[1]),y) #WARNING multioutput invalid
end

function Base.show(io::IO,model::LaplaceLikelihood{T}) where T
    print(io,"Laplace likelihood")
end

function compute_proba(l::LaplaceLikelihood{T},μ::AbstractVector{T},σ²::AbstractVector{T}) where {T<:Real}
    return μ,max.(σ²,0.0).+ 2*l.β[1]^2
end

## Local Updates ##

function local_updates!(model::Union{VGP{T,<:LaplaceLikelihood,<:AnalyticVI},VStP{T,<:LaplaceLikelihood,<:AnalyticVI}}) where {T}
    model.likelihood.b .= broadcast((Σ,μ,y)->(Σ+abs2.(μ-y)),diag.(model.Σ),model.μ,model.y)
    model.likelihood.θ .= broadcast((a,b)->sqrt(a)./sqrt.(b),model.likelihood.a,model.likelihood.b)
end

function local_updates!(model::SVGP{T,<:LaplaceLikelihood,<:AnalyticVI}) where {T}
    model.likelihood.b .= broadcast((K̃,κ,Σ,μ,y)->(K̃ + opt_diag(κ*Σ,κ) + abs2.(κ*μ-y)),model.K̃,model.κ,model.Σ,model.μ,model.inference.y)
    model.likelihood.θ .= broadcast((a,b)->sqrt(a)./sqrt.(b),model.likelihood.a,model.likelihood.b)
end

function sample_local!(model::VGP{T,<:LaplaceLikelihood,<:GibbsSampling}) where {T}
    model.likelihood.b .= broadcast((μ,y,β)->rand.(GeneralizedInverseGaussian.(1/β^2,abs2.(μ-y),0.5)),model.μ,model.inference.y,model.likelihood.β)
    model.likelihood.θ .= broadcast(b->1.0./b,model.likelihood.b)
    return nothing
end

@inline ∇E_μ(model::AbstractGP{T,<:LaplaceLikelihood,<:GibbsorVI}) where {T} =  hadamard.(model.likelihood.θ,model.inference.y)
@inline ∇E_μ(model::AbstractGP{T,<:LaplaceLikelihood,<:GibbsorVI},i::Int) where {T} =  model.likelihood.θ[i].*model.inference.y[i]

@inline ∇E_Σ(model::AbstractGP{T,<:LaplaceLikelihood,<:GibbsorVI}) where {T} = 0.5*model.likelihood.θ
@inline ∇E_Σ(model::AbstractGP{T,<:LaplaceLikelihood,<:GibbsorVI},i::Int) where {T} = 0.5*model.likelihood.θ[i]

## ELBO ##
function ELBO(model::AbstractGP{T,<:LaplaceLikelihood,<:AnalyticVI}) where {T}
    return expecLogLikelihood(model) - GIGExpKL(model) - GaussianKL(model)
end

function expecLogLikelihood(model::VGP{T,LaplaceLikelihood{T},AnalyticVI{T}}) where {T}
    tot = -0.5*model.nLatent*model.nSample*log(twoπ)
    tot += 0.5.*sum(broadcast(θ->sum(log.(θ)),model.likelihood.θ))
    tot += -0.5.*sum(broadcast((θ,Σ,μ,y)->dot(θ,(Σ+abs2.(μ)-2.0*μ.*y-abs2.(y))),model.likelihood.θ,diag.(model.Σ),model.μ,model.y))
    return tot
end

function expecLogLikelihood(model::SVGP{T,LaplaceLikelihood{T},AnalyticVI{T}}) where {T}
    tot = -0.5*model.nLatent*model.inference.nSamplesUsed*log(twoπ)
    tot += 0.5.*sum(broadcast(θ->sum(log.(θ)),model.likelihood.θ))
    tot += -0.5.*sum(broadcast((θ,K̃,κ,Σ,μ,y)->dot(θ,(K̃+opt_diag(κ*Σ,κ)+abs2.(κ*μ)-2.0*(κ*μ).*y-abs2.(y))),model.likelihood.θ,model.K̃,model.κ,model.Σ,model.μ,model.inference.y))
    return model.inference.ρ*tot
end

function GIGExpKL(model::AbstractGP{T,<:LaplaceLikelihood}) where {T}
    GIGEntropy(model)-expecExponentialGIG(model)
end

function GIGEntropy(model::AbstractGP{T,<:LaplaceLikelihood}) where {T}
    model.inference.ρ*sum(broadcast(GIGEntropy,model.likelihood.a,model.likelihood.b,model.likelihood.p))
end

function expecExponentialGIG(model::AbstractGP{T,<:LaplaceLikelihood}) where {T}
    sum(broadcast((β,a,b)->sum(-log(2*β^2).-0.5*(a.*sqrt.(b)+b.*sqrt(a))./(a.*b*β^2)),model.likelihood.β,model.likelihood.a,model.likelihood.b))
end

## PDF and Log PDF Gradients ##

function grad_quad(likelihood::LaplaceLikelihood{T},y::Real,μ::Real,σ²::Real,inference::Inference) where {T<:Real}
    nodes = inference.nodes*sqrt2*sqrt(σ²) .+ μ
    Edlogpdf = dot(inference.weights,grad_log_pdf.(likelihood,y,nodes))
    Ed²logpdf =  (1/sqrt(twoπ*σ²))/(likelihood.β[1]^2)
    return -Edlogpdf::T, Ed²logpdf::T
end

@inline grad_log_pdf(l::LaplaceLikelihood{T},y::Real,f::Real) where {T<:Real} = sign(y-f)./l.β[1]

function gradpdf(l::LaplaceLikelihood{T},y::Real,f::Real) where {T<:Real}
    grad_log_pdf_μ(l,y,f)*pdf(l,y,f)
end

function hessiandiagpdf(l::LaplaceLikelihood{T},y::Real,f::Real) where {T<:Real}
    pdf(l,y,f)/(l.β[1]^2)
end

@inline hessian_log_pdf(l::LaplaceLikelihood{T},y::Real,f::Real) where {T<:Real} = zero(T)


using Flux
struct SirenDense
    weight
    bias
    ω0
end

Flux.@functor SirenDense

function SirenDense(din, dout; c=6, ω0=1, isfirst=false)
    w_std = isfirst ? (1 / din) : (sqrt(c / din) / ω0)
    dist = Uniform(-w_std, w_std)
    weight = Float32.(rand(dist, dout, din))
    bias = Float32.(rand(dist, dout))
    SirenDense(weight, bias, ω0)
end

function (a::SirenDense)(x)
  sin.(a.ω0 * (a.weight*x .+ a.bias))
end

struct ModulatedSiren
    sirens
    modulator
end

Flux.trainable(π::ModulatedSiren) = (Flux.trainable(π.sirens)..., Flux.trainable(π.modulator)...)

function (a::ModulatedSiren)(x, z)
    slayers = layers(a.sirens)
    mlayers = layers(a.modulator)
    @assert length(slayers) == 4
    @assert length(slayers) == length(mlayers) + 1

    h = z

    h = mlayers[1](h)
    x = slayers[1](x) .* h
    h = vcat(h, z)

    h = mlayers[2](h)
    x = slayers[2](x) .* h
    h = vcat(h, z)

    h = mlayers[3](h)
    x = slayers[3](x) .* h
    # h = vcat(h, z)
    #
    # h = mlayers[4](h)
    # x = slayers[4](x) .* h
    # h = vcat(h, z)
    #
    # h = mlayers[5](h)
    # x = slayers[5](x) .* h

    slayers[end](x)
end


function rand(::AbstractRNG, ::SamplerTrivial{CloseOpen01{T}}) where {P, T<:ArbFloat{P}}
    oldprec = precision(BigFloat)
    setprecision(BigFloat, P+9)
    denom = BigInt(2)^P
    bigfloat = BigFloat(rand(BigInt(0):denom))
    bigfloat = bigfloat / BigFloat(denom)
    result = ArbFloat{P}(bigfloat)
    setprecision(BigFloat, oldprec)
    return result
end

@inline function rand(::AbstractRNG, ::SamplerTrivial{CloseOpen01{ArbFloat}})
    P = workingprecision(ArbFloat)
    return rand(ArbFloat{P})
end

function rand(::AbstractRNG, ::SamplerType{T}) where {P, T<:ArbReal{P}}
    arbfloat = rand(ArbFloat{P})
    return ArbReal{P}(arbfloat)
end

@inline function rand(::AbstractRNG, ::SamplerType{ArbReal})
    P = workingprecision(ArbReal)
    return rand(ArbReal{P})
end

function rand(::AbstractRNG, ::SamplerType{T}) where {P, T<:ArbComplex{P}}
    re = ArbReal{P}(rand(ArbFloat{P}))
    im = ArbReal{P}(rand(ArbFloat{P}))
    return ArbComplex{P}(re, im)
end

@inline function rand(::AbstractRNG, ::SamplerType{ArbComplex})
    P = workingprecision(ArbFloat)
    return rand(ArbComplex{P})
end

function randn(rng::AbstractRNG, ::Type{ArbFloat{N}}) where {N}
    urand1, urand2 = rand(ArbFloat{N}, 2)
    urand1 = urand1 + urand1 - 1
    urand2 = urand2 + urand2 - 1
    s = urand1*urand1 + urand2*urand2

    while s >= 1 || s === 0
        urand1, urand2 = rand(ArbFloat{N}, 2)
        urand1 = urand1 + urand1 - 1
        urand2 = urand2 + urand2 - 1
        s = urand1*urand1 + urand2*urand2
    end

    s = sqrt( -log(s) / s )
    return (urand1 + urand2) * s
end

randn(rng::AbstractRNG, ::Type{ArbFloat})  = randn(rng, ArbFloat{workingprecision(ArbFloat)})

function randn(rng::AbstractRNG, ::Type{ArbReal{N}}) where {N}
    urand1, urand2 = rand(ArbReal{N}, 2)
    urand1 = urand1 + urand1 - 1
    urand2 = urand2 + urand2 - 1
    s = urand1*urand1 + urand2*urand2

    while s >= 1 || s === 0
        urand1, urand2 = rand(ArbReal{N}, 2)
        urand1 = urand1 + urand1 - 1
        urand2 = urand2 + urand2 - 1
        s = urand1*urand1 + urand2*urand2
    end

    s = sqrt( -log(s) / s )
    return (urand1 + urand2) * s
end

randn(rng::AbstractRNG, ::Type{ArbReal})  = randn(rng, ArbReal{workingprecision(ArbReal)})


using Revise; using DynamicalBilliards, BenchmarkTools

bt = billiard_sinai()
p = MagneticParticle(0.1,0.1, π/8, 0.5)

@btime next_collision($p, $bt)

@btime bounce!($p, $bt)

# before sincos():
julia> @btime bounce!($p, $bt)
  278.529 ns (0 allocations: 0 bytes)
(1, 0.44691013238163385, [0.707091, 0.640047], [-0.353179, 0.935556])

# after sincos():
julia> @btime bounce!($p, $bt)
  57.476 ns (0 allocations: 0 bytes)
(0, Inf, [0.415228, 1.10629], [0.957396, -0.288778])


using DifferentialEquations
using Plots
using LinearAlgebra
using MatrixEquations

const M = 1
const m = 0.7
const l = 0.15
const D_x = 0.01
const D_θ = 0.01
const g = 9.80665

function fxu(x::Vector{T}, u::T) where T
    #dx = Vector{T}(undef, 4)
    dx1 = x[2]
    dx2 = (3*D_θ*cos(x[3])*x[4]/(4*M*l + 3*l*m*sin(x[3])^2 + l*m)
    - 4*D_x*l*x[2]/(4*M*l + 3*l*m*sin(x[3])^2 + l*m)
    - 3*g*l*m*sin(2*x[3])/(8*M*l + 6*l*m*sin(x[3])^2 + 2*l*m)
    + 4*l^2*m*sin(x[3])*x[4]^2/(4*M*l + 3*l*m*sin(x[3])^2 + l*m)
    + 4*l*u/(4*M*l + 3*l*m*sin(x[3])^2 + l*m))
    println("dx2 = ", dx2)
    dx3 = x[4]
    dx4 = (-3*D_θ*M*x[4]/(4*M*l^2*m + 3*l^2*m^2*sin(x[3])^2 + l^2*m^2)
    - 3*D_θ*m*x[4]/(4*M*l^2*m + 3*l^2*m^2*sin(x[3])^2 + l^2*m^2)
    + 3*D_x*l*m*cos(x[3])*x[2]/(4*M*l^2*m + 3*l^2*m^2*sin(x[3])^2 + l^2*m^2)
    + 3*M*g*l*m*sin(x[3])/(4*M*l^2*m + 3*l^2*m^2*sin(x[3])^2 + l^2*m^2)
    + 3*g*l*m^2*sin(x[3])/(4*M*l^2*m + 3*l^2*m^2*sin(x[3])^2 + l^2*m^2)
    - 3*l^2*m^2*sin(2*x[3])*x[4]^2/(8*M*l^2*m + 6*l^2*m^2*sin(x[3])^2 + 2*l^2*m^2)
    - 3*l*m*u*cos(x[3])/(4*M*l^2*m + 3*l^2*m^2*sin(x[3])^2 + l^2*m^2))
    dx = [dx1, dx2, dx3, dx4]
    println("x = ", x)
    println("dx = ", dx)
    return dx
end

function A(x::Vector{T}) where T
    out = Matrix{T}(undef, 4, 4)
    out[1,1] = 0
    out[1,2] = 1
    out[1,3] = 0
    out[1,4] = 0
    out[2,1] = 0
    out[2,2] = -4*D_x/(4*M + 3*m*sin(x[3])^2 + m)
    out[2,3] = (-6*m*(3*D_θ*cos(x[3])*x[4] - 4*D_x*l*x[2] - 3*g*l*m*sin(2*x[3])/2
    + 4*l^2*m*sin(x[3])*x[4]^2)*sin(x[3])*cos(x[3])/(l*(4*M
    + 3*m*sin(x[3])^2 + m)^2) + (-3*D_θ*sin(x[3])*x[4] - 3*g*l*m*cos(2*x[3])
    + 4*l^2*m*cos(x[3])*x[4]^2)/(l*(4*M + 3*m*sin(x[3])^2 + m)))
    out[2,4] = (3*D_θ*cos(x[3]) + 8*l^2*m*sin(x[3])*x[4])/(l*(4*M + 3*m*sin(x[3])^2 + m))
    out[3,1] = 0
    out[3,2] = 0
    out[3,3] = 0
    out[3,4] = 1
    out[4,1] = 0
    out[4,2] = 3*D_x*cos(x[3])/(l*(4*M + 3*m*sin(x[3])^2 + m))
    out[4,3] = (-18*(-D_θ*M*x[4] - D_θ*m*x[4] + D_x*l*m*cos(x[3])*x[2] + M*g*l*m*sin(x[3])
    + g*l*m^2*sin(x[3]) - l^2*m^2*sin(2*x[3])*x[4]^2/2)*sin(x[3])*cos(x[3])/(l^2*(4*M + 3*m*sin(x[3])^2 + m)^2)
    + 3*(-D_x*l*m*sin(x[3])*x[2] + M*g*l*m*cos(x[3]) + g*l*m^2*cos(x[3])
    - l^2*m^2*cos(2*x[3])*x[4]^2)/(l^2*m*(4*M + 3*m*sin(x[3])^2 + m)))
    out[4,4] = 3*(-D_θ*M - D_θ*m - l^2*m^2*sin(2*x[3])*x[4])/(l^2*m*(4*M + 3*m*sin(x[3])^2 + m))
    out
end

function B(x::Vector{T}) where T
    out = Vector{T}(undef, 4)
    out[1] = 0
    out[2] = 4*l/(4*M*l + 3*l*m*(sin(x[3])^2) + l*m)
    out[3] = 0
    out[4] = -3*l*m*cos(x[3])/(4*M*l^2*m + 3*l^2*m^2*(sin(x[3])^2) + l^2*m^2)
    out
end

struct PID{T}
    Kp::T
    Ki::T
    Kd::T
end

function input(p::PID{T}, x::Vector{T}) where T
    0
end

function dx!(dx::Vector{T}, x::Vector{T}, p::PID{T}, t::T) where T
    fxu!(dx, x, input(p, x))
end

function run_simuration(p::PID{T}, t_span::Tuple{T, T}, Δt::T, x₀::Vector{T}) where T
    prob = ODEProblem(dx!, x₀, t_span, p)
    sol = solve(prob)
    sol
end

mutable struct LQR{T}
    Q::Matrix{T}
    R::Matrix{T}
    F::Matrix{T}
end

function init!(p::LQR{T}, A::Matrix{T}, B::Vector{T}) where T
    S = zero(B)
    P, _, _ = arec(A, B, p.R, p.Q, S)
    p.F = inv(p.R) * B' * P
end

function input(p::LQR{T}, x::Vector{T}) where T
    (p.F * x)[1,1]
end

function dx(x::Vector{T}, p::LQR{T}, t::T) where T
    println("t = ", t)
    u = input(p, x)
    println("u = ", u)
    fxu(x, u)
end

function run_simuration(p::LQR{T}, t_span::Tuple{T, T}, Δt::T, x₀::Vector{T}) where T
    init!(p, A(zeros(T, 4)), B(zeros(T, 4)))
    prob = ODEProblem(dx, x₀, t_span, p)
    sol = solve(prob)
    sol
end

x₀ = [0, 0, π/6, 0]
p = LQR(diagm([10.0, 0, 100, 0]), Matrix{Float64}(I, 1, 1).*0.01, Matrix{Float64}(undef, 4, 4))
sol = run_simuration(p, (0.0, 5.0), 0.01, x₀)
plot(sol)


"""
Truncation in a box between `a` and `b`.
"""
struct BoxTruncationRegion <: TruncationRegion
    a::Vector{Float64}
    b::Vector{Float64}
end

intruncationregion(r::BoxTruncationRegion, x::AbstractArray) = all(r.a .<= x) && all(x .<= r.b)

"""
An Elliptical Truncation region.
"""
struct EllipticalTruncationRegion <: TruncationRegion
    H::PDMat
    h::Vector{Float64}
    c::Float64
end

intruncationregion(r::EllipticalTruncationRegion, x::AbstractArray) = (x-r.h)'*r.H*(x-r.h) <= r.c

abstract type PolytopeTruncationRegion <: TruncationRegion end


using Plots
using Distributions
pyplot()

function gen_rand_uniform_mult(min1, max1, min2, max2, n)
    rand(Uniform(min1, max1), n) .* rand(Uniform(min2, max2), n)
end

histogram(gen_rand_uniform_mult(0, 10, 35, 40, 10000))


using DrWatson
@quickactivate
using ArgParse
using GenerativeAD
using GenerativeAD.Models
import StatsBase: fit!, predict
using StatsBase
using BSON
# because of vae and 2stage
using DataFrames
using CSV
using ValueHistories
using Flux
using ConditionalDists
using GenerativeModels
import GenerativeModels: VAE
using Distributions
using DistributionsAD
using Distances
using Statistics

s = ArgParseSettings()
@add_arg_table! s begin
   "max_seed"
		required = true
		arg_type = Int
		help = "seed"
	"dataset"
		required = true
		arg_type = String
		help = "dataset"
	"contamination"
    	arg_type = Float64
    	help = "contamination rate of training data"
    	default = 0.0
end
parsed_args = parse_args(ARGS, s)
@unpack dataset, max_seed, contamination = parsed_args

# we select the encoding of the model with lowest loss
tab_name = "vae_LOSS_tabular"
mi = 2

#######################################################################################
################ THIS PART IS TO BE PROVIDED FOR EACH MODEL SEPARATELY ################
# get criterion
sp = split(tab_name, "_")
enc = sp[1]
criterion = lowercase(sp[2])

modelname = "$(enc)_ocsvm"

# default hyperparams
function median_l2_dist(X)
	dists = pairwise(Euclidean(), X)
	# take only the upper diagonal
	ds = []
	for i in 1:size(dists,1)
		for j in (i+1):size(dists,1)
			push!(ds, dists[i,j])
		end
	end
	ml=median(ds)
end
function set_params(data)
	D, N = size(data[1][1])
	gamma = 1/median_l2_dist(data[1][1])
	return (gamma=gamma, kernel="rbf", nu=0.5)
end

function fit(data, parameters, aux_info)
	# construct model - constructor should only accept kwargs
	model = GenerativeAD.Models.OCSVM(;parameters...)

	# sumbsample and fit train data
	tr_data = data[1][1]
	M,N = size(tr_data)
	tr_data = tr_data[:, sample(1:N, min(10000, N), replace=false)]
	try
		global info, fit_t, _, _, _ = @timed fit!(model, data[1][1])
	catch e
		# return an empty array if fit fails so nothing is computed
		return (fit_t = NaN,), []
	end

	# construct return information - put e.g. the model structure here for generative models
	training_info = (
		fit_t = fit_t,
		model = nothing
		)

	# now return the different scoring functions
	training_info, [(x->predict(model, x), merge(parameters, aux_info))]
end

####################################################################
################ THIS PART IS COMMON FOR ALL MODELS ################
# set a maximum for parameter sampling retries
cont_string = (contamination == 0.0) ? "" : "_contamination-$contamination"
for seed in 1:max_seed
	savepath = datadir("experiments/tabular_clean_val_default/$(modelname)/$(dataset)/seed=$(seed)")
	mkpath(savepath)

	aux_info = (model_index=mi, criterion=criterion)
	# get data
	data = GenerativeAD.load_data(dataset, seed=seed, contamination=contamination)
	output, encode_t, _, _, _ = @timed GenerativeAD.Models.load_encoding(tab_name, data, dataset=dataset, seed=seed, model_index=mi)
	data, encoding_name, encoder_params, encoder_fit_t = output

	# set parameters
	parameters = set_params(data)

	@info "Trying to fit $modelname on $dataset with parameters $(parameters)..."
	# check if a combination of parameters and seed alread exists
	# fit
	training_info, results = fit(data, parameters, aux_info)
	# here define what additional info should be saved together with parameters, scores, labels and predict times
	save_entries = merge(training_info, (modelname = modelname,
										 seed = seed,
										 dataset = dataset,
										 encoder=encoding_name,
										 encoder_fit_t = encoder_fit_t,
										 encoder_params=encoder_params,
										 encode_t = encode_t,
										 model_index=mi,
										 criterion=criterion,
										 contamination = contamination))

	# now loop over all anomaly score funs
	for result in results
		GenerativeAD.experiment(result..., data, savepath; save_entries...)
	end
end


function log_beta(::Type{Beta}, arg0::jdouble, arg1::jdouble)
    return jcall(Beta, "logBeta", jdouble, (jdouble, jdouble), arg0, arg1)
end

function regularized_beta(::Type{Beta}, arg0::jdouble, arg1::jdouble, arg2::jdouble)
    return jcall(Beta, "regularizedBeta", jdouble, (jdouble, jdouble, jdouble), arg0, arg1, arg2)
end

function regularized_beta(::Type{Beta}, arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return jcall(Beta, "regularizedBeta", jdouble, (jdouble, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function regularized_beta(::Type{Beta}, arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jdouble, arg4::jint)
    return jcall(Beta, "regularizedBeta", jdouble, (jdouble, jdouble, jdouble, jdouble, jint), arg0, arg1, arg2, arg3, arg4)
end

function regularized_beta(::Type{Beta}, arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jint)
    return jcall(Beta, "regularizedBeta", jdouble, (jdouble, jdouble, jdouble, jint), arg0, arg1, arg2, arg3)
end


workspace()
#using EAGOSemiInfinite
#using EAGOGlobalSolver
using IntervalArithmetic
using EAGOBranchBound
#using EAGOSmoothMcCormickGrad
#using StaticArrays

function FDI_invApprox(x::Interval)
  Ap = 0.00005
  g = 9.8
  low = -(x^2)/(2*g*Ap^2)
  high = (x^2)/(2*g*Ap^2)
  if x.lo>0
    return high
  elseif x.hi<0
    return low
  else
    return Interval(0,max(low.hi,high.hi))
  end
end

function FDI_invApprox(x::Float64)
  Ap = 0.00005
  g = 9.8
  if x>=0
    return (x^2)/(2*g*Ap^2)
  else x<0
    return -(x^2)/(2*g*Ap^2)
  end
end

function xbnds(u,p)

         # parameters
         Ap = 0.00005
         g = 9.8
         z = []
         # bound z via McCormick operators
         z1 = (((u[1]/p[3])/Ap)^2)/(2.0*g)
         z3 = (1.0/(2.0*g))*((u[1]+u[2])/(p[4]*Ap))^2
         z2 = (-((u[2]-p[4]*Ap*sqrt(2*g*z3))/(p[5])/Ap)^2)/(2.0*g)
         z4 = (((u[1]*p[1]/p[3])/Ap)^2)/(2.0*g)
         z6 = (1.0/(2.0*g))*((p[1]*u[1]+u[2])/(p[4]*Ap))^2
         z5 = (-((u[2]-p[4]*Ap*sqrt(2*g*z6))/(p[5])/Ap)^2)/(2.0*g)
         z7 = (((u[1]/p[3])/Ap)^2)/(2.0*g)
         z9 = (1.0/(2.0*g))*((u[1]+u[2])/(p[4]*(Ap+3.14159*p[2]^2)))^2
         z8 = FDI_invApprox((u[2]-p[4]*(Ap+3.14159*p[2]^2)*sqrt(2.0*g*z9))/(p[5]))
         push!(z,z1,z2,z3,z4,z5,z6,z7,z8,z9)

         # rotate coordinates back to X & intersects

         x = []
         x1 = min(max(z1-z2+z3,0.0),0.75)
         x2 = min(max(z3,0.0),0.75)
         x3 = min(max(z3-z2,0.0),0.75)
         x4 = min(max(z4-z5+z6,0.0),0.75)
         x5 = min(max(z6,0.0),0.75)
         x6 = min(max(z6-z5,0.0),0.75)
         x7 = min(max(z7-z8+z9,0.0),0.75)
         x8 = min(max(z9,0.0),0.75)
         x9 = min(max(z9-z8,0.0),0.75)
         push!(x,x1,x2,x3,x4,x5,x6,x7,x8,x9)

         return x
end

# sets up SIP constraint with imbedded x(p) calculation
function SIP_g(un,p)
  u = un[1:2]
  xbnd = xbnds(u,p)
  r_xbnds = [zero(xbnd[1]) for i=1:3,j=1:3]
  r_xbnds[:,1] = xbnd[1:3]
  r_xbnds[:,2] = xbnd[4:6]
  r_xbnds[:,3] = xbnd[7:9]
  out = un[3]
  for i=1:3
    for j=1:2
      for k=(j+1):3
        out = out - (r_xbnds[i,j]-r_xbnds[i,k])^2
      end
    end
  end
  return out
end

# sets up SIP objective
function SIP_f(un)
  return -un[3]
end
# sets up SIP
ex_SIP_X = [(1E-5)..(1E-4),(1E-5)..(1E-4),(-10.0)..(10.0)]
ex_SIP_P = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),(0.65)..(0.95),(0.85)..(1.15)]

#Intv = SIP_g(ex_SIP_X,ex_SIP_P)
#Trial1 = [1E-5,1E-5,0.0]
#Intv1 = SIP_g(Trial1,ex_SIP_P)
#Trial2 = [1E-5,1E-4,0.0]
#Intv2 = SIP_g(Trial2,ex_SIP_P)
#Trial3 = [1E-4,1E-5,0.0]
#Intv3 = SIP_g(Trial3,ex_SIP_P)
#Trial4 = [1E-4,1E-4,0.0]
#Intv4 = SIP_g(Trial4,ex_SIP_P)
#Trial5 = [(1E-5+1E-4)/2,(1E-5+1E-4)/2,0.0]
#Intv5 = SIP_g(Trial5,ex_SIP_P)
#Trial6 = [(1E-5+1E-4)/3,(1E-5+1E-4)/3,0.0]
#Intv6 = SIP_g(Trial6,ex_SIP_P)
#Trial7 = [2*(1E-5+1E-4)/3,2*(1E-5+1E-4)/3,0.0]
#Intv7 = SIP_g(Trial7,ex_SIP_P)
#Trial8 = [(1E-5+1E-4)/2,(1E-5+1E-4)/3,0.0]
#Intv8 = SIP_g(Trial8,ex_SIP_P)
#Trial9 = [(1E-5+1E-4)/2,2*(1E-5+1E-4)/3,0.0]
#Intv9 = SIP_g(Trial9,ex_SIP_P)
#Trial10 = [(1E-5+1E-4)/3,(1E-5+1E-4)/2,0.0]
#Intv10 = SIP_g(Trial10,ex_SIP_P)
Trial11 = [2*(1E-5+1E-4)/3,(1E-5+1E-4)/2,0.0]
Intv11 = SIP_g(Trial11,ex_SIP_P)

ex_SIP_P1a = [0.66,(0.0005)..(0.005),(0.85)..(1.15),(0.65)..(0.95),(0.85)..(1.15)]
ex_SIP_P2a = [(0.54)..(0.66),0.005,(0.85)..(1.15),(0.65)..(0.95),(0.85)..(1.15)]
ex_SIP_P3a = [(0.54)..(0.66),(0.0005)..(0.005),1.15,(0.65)..(0.95),(0.85)..(1.15)]
ex_SIP_P4a = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),0.95,(0.85)..(1.15)]
ex_SIP_P5a = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),(0.65)..(0.95),1.15]
Intv11a1 = SIP_g(Trial11,ex_SIP_P1a)
Intv11b1 = SIP_g(Trial11,ex_SIP_P2a)
Intv11c1 = SIP_g(Trial11,ex_SIP_P3a)
Intv11d1 = SIP_g(Trial11,ex_SIP_P4a)
Intv11e1 = SIP_g(Trial11,ex_SIP_P5a)

ex_SIP_P1b = [0.54,(0.0005)..(0.005),(0.85)..(1.15),(0.65)..(0.95),(0.85)..(1.15)]
ex_SIP_P2b = [(0.54)..(0.66),0.0005,(0.85)..(1.15),(0.65)..(0.95),(0.85)..(1.15)]
ex_SIP_P3b = [(0.54)..(0.66),(0.0005)..(0.005),0.85,(0.65)..(0.95),(0.85)..(1.15)]
ex_SIP_P4b = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),0.65,(0.85)..(1.15)]
ex_SIP_P5b = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),(0.65)..(0.95),0.85]
Intv11a2 = SIP_g(Trial11,ex_SIP_P1b)
Intv11b2 = SIP_g(Trial11,ex_SIP_P2b)
Intv11c2 = SIP_g(Trial11,ex_SIP_P3b)
Intv11d2 = SIP_g(Trial11,ex_SIP_P4b)
Intv11e2 = SIP_g(Trial11,ex_SIP_P5b)
ex_SIP_P4ca = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),0.70,(0.85)..(1.15)]
ex_SIP_P4c = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),0.75,(0.85)..(1.15)]
ex_SIP_P4da = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),0.80,(0.85)..(1.15)]
ex_SIP_P4d = [(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),0.85,(0.85)..(1.15)]
Intv11d2a = SIP_g(Trial11,ex_SIP_P4ca)
Intv11d3 = SIP_g(Trial11,ex_SIP_P4c)
Intv11d3a = SIP_g(Trial11,ex_SIP_P4da)
Intv11d4 = SIP_g(Trial11,ex_SIP_P4d)

disc1 = [0.6,0.0005,1.0,0.65,1.00]
ex_SIP_Xt = [(1E-5)..(1E-4),(1E-5)..(1E-4),(0.0)..(0.0)]
Xtest = SIP_g(ex_SIP_Xt,disc1)
m = BnBModel([(1E-5)..(1E-4),(1E-5)..(1E-4),(-10.0)..(10.0)])
s = BnBSolver()
set_to_default!(s)
s.BnB_atol = 1E-2

# solve optimization problem via interval extension
function test_lower(X,k,pos,opt_inner,UBD)
  FInt = -X[3]
  disc1 = [0.6,0.005,1.0,0.65,1.00]
  GInt = SIP_g(X,disc1)
  feas = true
  for i=1:length(GInt)
    if (GInt[i].lo>0.0)
      feas = false
      break
    end
  end
  val = FInt.lo
  pnt = mid.(X)
  return pnt, val, feas, X, [FInt]
end

function test_upper(X,k,pos,opt_inner,temp)
  feas = true
  val = temp[1].hi
  pnt = mid.(X)
  return pnt, val, feas, X
end

s.Lower_Prob = test_lower
s.Upper_Prob = test_upper
solveBnB!(s,m)

mi = BnBModel([(0.54)..(0.66),(0.0005)..(0.005),(0.85)..(1.15),(0.65)..(0.95),(0.85)..(1.15)])
si = BnBSolver()
set_to_default!(si)
si.BnB_atol = 5E-2

function inner_lower(X,k,pos,opt_inner,UBD)
  disc1 = [5.5E-5,5.5E-5,4.99]
  FInt = -SIP_g(disc1,X)
  feas = true
  val = FInt.lo
  pnt = mid.(X)
  return pnt, val, feas, X, [FInt]
end

function inner_upper(X,k,pos,opt_inner,temp)
  feas = true
  val = temp[1].hi
  pnt = mid.(X)
  return pnt, val, feas, X
end

si.Lower_Prob = inner_lower
si.Upper_Prob = inner_upper
#solveBnB!(si,mi)

m1 = BnBModel([(1E-5)..(1E-4),(1E-5)..(1E-4),(-10.0)..(10.0)])
s1 = BnBSolver()
set_to_default!(s1)
s1.BnB_atol = 1E-2

# solve optimization problem via interval extension
function test_lower(X,k,pos,opt_inner,UBD)
  FInt = -X[3]
  disc1 = [0.6,0.005,1.0,0.65,1.00]
  GInt = SIP_g(X,disc1) +0.5
  feas = true
  for i=1:length(GInt)
    if (GInt[i].lo>0.0)
      feas = false
      break
    end
  end
  val = FInt.lo
  pnt = mid.(X)
  return pnt, val, feas, X, [FInt]
end

function test_upper(X,k,pos,opt_inner,temp)
  feas = true
  val = temp[1].hi
  pnt = mid.(X)
  return pnt, val, feas, X
end

s1.Lower_Prob = test_lower
s1.Upper_Prob = test_upper
solveBnB!(s1,m1)

#=
SIPopt = SIP_opts()
set_to_default!(SIPopt)
SIPopt.Verbosity = "Full" #"Normal"

# solves sample LBP
vals = 1
count = 1
p_sto = zeros(vals^5,5)
for i=1:vals
  for j=1:vals
    for k=1:vals
      for p=1:vals
        for q=1:vals
          t1 = ex_SIP_P[1].lo + (ex_SIP_P[1].hi - ex_SIP_P[1].lo)*(i-1)/(vals-1)
          t2 = ex_SIP_P[2].lo + (ex_SIP_P[2].hi - ex_SIP_P[2].lo)*(j-1)/(vals-1)
          t3 = ex_SIP_P[3].lo + (ex_SIP_P[3].hi - ex_SIP_P[3].lo)*(k-1)/(vals-1)
          t4 = ex_SIP_P[4].lo + (ex_SIP_P[4].hi - ex_SIP_P[4].lo)*(p-1)/(vals-1)
          t5 = ex_SIP_P[5].lo + (ex_SIP_P[5].hi - ex_SIP_P[5].lo)*(q-1)/(vals-1)
          p_sto[count,:] = [t1,t2,t3,t4,t5]
          count += 1
        end
      end
    end
  end
end

bnd_z2 = ex_SIP_X[2]-ex_SIP_P[4]*0.00005*sqrt((ex_SIP_P[1]*ex_SIP_X[1]+ex_SIP_X[2])/(ex_SIP_P[5]*(ex_SIP_P[4]*0.00005)^2))
bnd_z5 = ex_SIP_X[2]-ex_SIP_P[4]*0.00005*sqrt((ex_SIP_P[1]*ex_SIP_X[1]+ex_SIP_X[2])/(ex_SIP_P[5]*(ex_SIP_P[4]*0.00005)^2))
bnd_z8 = (ex_SIP_X[2]-ex_SIP_P[4]*(0.00005+3.14159*ex_SIP_P[2]^2)*sqrt((((ex_SIP_X[1]+ex_SIP_X[2])/(ex_SIP_P[4]*(0.00005+3.14159*ex_SIP_P[2]^2)))^2)))/(ex_SIP_P[5])

SIPopt.tol = 1E-1
SIPopt.eps_g0 = 1.0
SIPopt.P_LBD = [p_sto[i,:] for i=1:vals^5]
SIPopt.P_UBD = [p_sto[i,:] for i=1:vals^5]
SIPopt.LBP_Opt.DAG_depth = -1
SIPopt.LBP_Opt.f = SIP_f
SIPopt.LBP_Opt.g = SIP_g
SIPopt.LBP_Opt.X = ex_SIP_X
SIPopt.LBP_Opt.g = x -> BndProb_reform(x,[],SIP_g,[],0.0)
#SIPopt.LBP_Opt.LBD_func_relax = "Diff2-MV-OFF"
#SIPopt.LBP_Opt.LBD_problem_relax = "NLP2"
#SIPopt.LBP_Opt.LBD_problem_solver = "Ipopt"
SIPopt.LBP_Opt.LBD_func_relax = "Interval"
SIPopt.LBP_Opt.LBD_problem_relax = "Interval"
SIPopt.LBP_Opt.LBD_problem_solver = "Interval"
SIPopt.LBP_Opt.UBD_func_relax = "Interval"
SIPopt.LBP_Opt.UBD_problem_relax = "Interval"
SIPopt.LBP_Opt.UBD_problem_solver = "Interval"
SIPopt.LBP_Opt.BnB_options.Verbosity = "Normal" #"Full"
SIPopt.LBP_Opt.BnB_options.BnB_tol = 1E-2
#LBDg, xbar, feas, tLBP1,tLBP2 = EAGO_Global_Explicit(SIPopt.LBP_Opt)
#time_lower = @time EAGO_Global_Explicit(SIPopt.LBP_Opt)
#time_lower = @time EAGO_Global_Explicit(SIPopt.LBP_Opt)

xbar = mid.(ex_SIP_X)
SIPopt.LLP_Opt.DAG_depth = -1
SIPopt.LLP_Opt.f = p -> -SIP_g(xbar,p)
SIPopt.LLP_Opt.X = ex_SIP_P
#SIPopt.LLP_Opt.LBD_func_relax = "Diff2-MV-OFF"
#SIPopt.LLP_Opt.LBD_problem_relax = "NLP2"
#SIPopt.LLP_Opt.LBD_problem_solver = "Ipopt"
SIPopt.LLP_Opt.LBD_func_relax = "Interval"
SIPopt.LLP_Opt.LBD_problem_relax = "Interval"
SIPopt.LLP_Opt.LBD_problem_solver = "Interval"
SIPopt.LLP_Opt.UBD_func_relax = "Interval"
SIPopt.LLP_Opt.UBD_problem_relax = "Interval"
SIPopt.LLP_Opt.UBD_problem_solver = "Interval"
SIPopt.LLP_Opt.BnB_options.Verbosity = "Normal"
SIPopt.LLP_Opt.BnB_options.BnB_tol = 1E-2
#INNg1, pbar,feas,tLLP1,tLLP2 = EAGO_Global_Explicit(SIPopt.LLP_Opt)
#time_inner = @time EAGO_Global_Explicit(SIPopt.LLP_Opt)
#time_inner = @time EAGO_Global_Explicit(SIPopt.LLP_Opt)

SIPopt.UBP_Opt.DAG_depth = -1
SIPopt.UBP_Opt.f = SIP_f
SIPopt.UBP_Opt.g = SIP_g
SIPopt.UBP_Opt.X = ex_SIP_X
SIPopt.UBP_Opt.g = x -> BndProb_reform(x,[],SIP_g,[],1.0)
#SIPopt.UBP_Opt.LBD_func_relax = "Diff2-MV-OFF"
#SIPopt.UBP_Opt.LBD_problem_relax = "NLP2"
#SIPopt.UBP_Opt.LBD_problem_solver = "Ipopt"
SIPopt.UBP_Opt.LBD_func_relax = "Interval"
SIPopt.UBP_Opt.LBD_problem_relax = "Interval"
SIPopt.UBP_Opt.LBD_problem_solver = "Interval"
SIPopt.UBP_Opt.UBD_func_relax = "Interval"
SIPopt.UBP_Opt.UBD_problem_relax = "Interval"
SIPopt.UBP_Opt.UBD_problem_solver = "Interval"
SIPopt.UBP_Opt.BnB_options.Verbosity = "Normal"
SIPopt.UBP_Opt.BnB_options.BnB_tol = 1E-2
#UBD_temp, xbar, feas,tUBP1,tUBP2 = EAGO_Global_Explicit(SIPopt.UBP_Opt)
#time_upper = @time EAGO_Global_Explicit(SIPopt.UBP_Opt)
#time_upper = @time EAGO_Global_Explicit(SIPopt.UBP_Opt)
#pbar = [0.66, 0.00275, 1.13595, 0.95, 1.15]
#ubar = [5.5e-05, 0.000100001] × [5.5e-05, 0.000100001] × [9.9781, 9.97815]

xbnd_range1 = xbnds(ex_SIP_X[1:2],ex_SIP_P)
SIP_range = SIP_g(ex_SIP_X,ex_SIP_P)
ex_out = Explicit_SIP_Solve(SIP_f,[],SIP_g,ex_SIP_X,ex_SIP_P,SIPopt)
=#


export  PlaneStress

@doc raw"""
    PlaneStress

Creates a plane stress element

- `H`: Linear elasticity matrix, $3\times3$
- `E`: Young's modulus
- `ν`: Poisson's ratio
- `ρ`: density
- `σ0`: stress at the **last** time step
- `σ0_`: (for internal use), stress to be updated in `commitHistory`
- `ε0`: strain at the **last** time step
- `ε0_`: (for internal use), strain to be updated in `commitHistory`

# Example
```julia
prop = Dict("name"=> "PlaneStress", "rho"=> 0.0876584, "E"=>0.07180760098, "nu"=>0.4)
mat = PlaneStress(prop)
```
"""
mutable struct PlaneStress
    H::Array{Float64}
    E::Float64
    ν::Float64
    ρ::Float64
    σ0::Array{Float64} # stress at last time step
    σ0_::Array{Float64} # σ0 to be updated in `commitHistory`
    ε0::Array{Float64}
    ε0_::Array{Float64}
end

"""
    PlaneStress(prop::Dict{String, Any})

`prop` should contain at least the following three fields: `E`, `nu`, `rho`
"""
function PlaneStress(prop::Dict{String, Any})
    E = prop["E"]; ν = prop["nu"]; ρ = prop["rho"]
    H = zeros(3,3)

    H[1,1] = E/(1. -ν*ν)
    H[1,2] = H[1,1]*ν
    H[2,1] = H[1,2]
    H[2,2] = H[1,1]
    H[3,3] = E/(2.0*(1.0+ν))
    σ0 = zeros(3); σ0_ = zeros(3); ε0 = zeros(3); ε0_ = zeros(3)
    PlaneStress(H, E, ν, ρ, σ0, σ0_,ε0,ε0_)
end

function getStress(self::PlaneStress, strain::Array{Float64}, Dstrain::Array{Float64}, Δt::Float64 = 0.0)
    sigma = self.H * strain
    self.σ0_ = copy(sigma)
    self.ε0_ = copy(strain)
    return sigma, self.H
end

function getTangent(self::PlaneStress)
    self.H
end

function commitHistory(self::PlaneStress)
    self.σ0 = self.σ0_
    self.ε0 = self.ε0_
end


# Practical Machine learning
# Bayesian learning - Naive Bayes example
# Chapter 9

import StatsBase.fit
import StatsBase.predict
using Distributions

#### core naive Bayes functions ####

function ensure_data_size(X, y)
    @assert(size(X, 2) == length(y),
            "Number of observations in X ($(size(X, 2))) is not equal to " *
            "number of class labels in y ($(length(y)))")
end

function logprob_c{C}(m::NBModel, c::C)
    return m.c_counts[c] / m.n_obs
end

# predict log probabilities for all classes
function predict_logprobs{V<:Number}(m::NBModel, x::Vector{V})
    C = eltype(m.c_counts)[1]
    logprobs = Dict{C, Float64}()
    for c in keys(m.c_counts)
        logprobs[c] = logprob_c(m, c) + logprob_x_given_c(m, x, c)
    end
    return keys(logprobs), values(logprobs)
end

# predict log probabilities for all classes
function predict_logprobs{V<:Number}(m::NBModel, X::Matrix{V})
    C = eltype(m.c_counts)[1]
    logprobs_per_class = Dict{C, Vector{Float64}}()
    for c in keys(m.c_counts)
        logprobs_per_class[c] = logprob_c(m, c) + logprob_x_given_c(m, X, c)
    end
    return (collect(keys(logprobs_per_class)),
            hcat(collect(values(logprobs_per_class))...)')
end

# preditct logprobs, return tuples of predicted class and its logprob
function predict_proba{V<:Number}(m::NBModel, X::Matrix{V})
    C = eltype(m.c_counts)[1]
    classes, logprobs = predict_logprobs(m, X)
    predictions = Array((C, Float64), size(X, 2))
    for j=1:size(X, 2)
        maxprob_idx = indmax(logprobs[:, j])
        c = classes[maxprob_idx]
        logprob = logprobs[maxprob_idx, j]
        predictions[j] = (c, logprob)
    end
    return predictions
end

function predict{V<:Number}(m::NBModel, X::Matrix{V})
    return [k for (k,v) in predict_proba(m, X)]
end

#####  Multinomial Naive Bayes  #####

function fit{C}(m::MultinomialNB, X::Matrix{Int64}, y::Vector{C})
    ensure_data_size(X, y)
    for j=1:size(X, 2)
        c = y[j]
        m.c_counts[c] += 1
        m.x_counts[c] .+= X[:, j]
        m.x_totals += X[:, j]
        m.n_obs += 1
    end
    return m
end

# Calculate log P(x|C)
function logprob_x_given_c{C}(m::MultinomialNB, x::Vector{Int64}, c::C)
    x_priors_for_c = m.x_counts[c] ./ m.x_totals
    x_probs_given_c = x_priors_for_c .^ x
    logprob = sum(log(x_probs_given_c))
    return logprob
end

# Calculate log P(x|C)
function logprob_x_given_c{C}(m::MultinomialNB, X::Matrix{Int64}, c::C)
    x_priors_for_c = m.x_counts[c] ./ m.x_totals
    x_probs_given_c = x_priors_for_c .^ X
    logprob = sum(log(x_probs_given_c), 1)
    return squeeze(logprob, 1)
end

######  Gaussian Naive Bayes  #######

function fit{C}(m::GaussianNB, X::Matrix{Float64}, y::Vector{C})
    ensure_data_size(X, y)
    # updatestats(m.dstats, X)
    # m.gaussian = MvNormal(mean(m.dstats), cov(m.dstats))
    # m.n_obs = m.dstats.n_obs
    n_vars = size(X, 1)
    for j=1:size(X, 2)
        c = y[j]
        m.c_counts[c] += 1
        updatestats(m.c_stats[c], reshape(X[:, j], n_vars, 1))
        # m.x_counts[c] .+= X[:, j]
        # m.x_totals += X[:, j]
        m.n_obs += 1
    end
    # precompute distributions for each class
    for c in keys(m.c_counts)
        m.gaussians[c] = MvNormal(mean(m.c_stats[c]), cov(m.c_stats[c]))
    end
    return m
end

# Calculate log P(x|C)
function logprob_x_given_c{C}(m::GaussianNB, x::Vector{Float64}, c::C)
    return logpdf(m.gaussians[c], x)
end

# Calculate log P(x|C)
function logprob_x_given_c{C}(m::GaussianNB, X::Matrix{Float64}, c::C)
    ## x_priors_for_c = m.x_counts[c] ./ m.x_totals
    ## x_probs_given_c = x_priors_for_c .^ x
    ## logprob = sum(log(x_probs_given_c))
    ## return logprob
    return logpdf(m.gaussians[c], X)
end


include("ArgLexMin.jl")
include("Dominated.jl")
include("Explore.jl")
include("MakeModel.jl")
include("MinimalT1.jl")
include("MinimalT2.jl")
include("MinMaxModel.jl")
include("MIMMPSolver.jl")
include("DBMIMMPSolver.jl")
include("CBMIMMPSolver.jl")
include("ShaoRSolver.jl")
include("SelectVariable.jl")
include("SelectNode.jl")
include("GLBUpdater.jl")
include("MakeP_bar.jl")
include("VariableType.jl")
include("WeakRemoverModel.jl")
include("Preprocess2.jl")
include("WSM.jl")
include("Vert.jl")
include("FindIntersectingPoint.jl")
include("WeightedMinMaxModel.jl")
include("ConModel.jl")


function op_V_loc( pw::PWGrid, V_loc, psi::Array{ComplexF64,2} )
    #
    Ns = pw.Ns
    Ω  = pw.Ω
    Npoints = prod(Ns)
    Ncols   = size(psi)[2]

    ctmp = zeros(ComplexF64, Npoints, Ncols)
    idx = pw.gvecw.idx_gw2r
    for ic = 1:Ncols
        ctmp[idx,ic] = psi[:,ic]
    end

    # get values of psi in real space grid via forward transform
    ctmp = G_to_R( Ns, ctmp )

    cVpsi = R_to_G( Ns, Diagprod(V_loc, ctmp) )
    return cVpsi[idx,:]
end

function op_V_loc( pw::PWGrid, V_loc, psi::Array{ComplexF64,1} )
    #
    Ns = pw.Ns
    Ω  = pw.Ω
    Npoints = prod(Ns)

    ctmp = zeros(ComplexF64, Npoints)
    idx = pw.gvecw.idx_gw2r
    ctmp[idx] = psi[:]

    # get values of psi in real space grid via forward transform
    ctmp = G_to_R( Ns, ctmp )

    cVpsi = R_to_G( Ns, V_loc .* ctmp )
    return cVpsi[idx]
end

# B is usually consists of more than one-column
function Diagprod( a, B::Array{ComplexF64,2} )
    Nstates = size(B)[2]
    Npoints = size(B)[1]
    out = zeros( ComplexF64, size(B) )
    for ist = 1:Nstates
        for ip = 1:Npoints
            out[ip,ist] = a[ip]*B[ip,ist]
        end
    end
    return out
end


using ObaraSaika
using Test
using PyCall
using QuadGK

import ObaraSaika: ERI, compute_integrals, boys, compute_normalisation

@testset "ObaraSaika.jl" begin
    @testset "Normalisation" begin
        f(x) = exp(-1.5 * abs(x)^2)
        integral = 4π * quadgk(x -> f(x) * f(x) * x * x, 0, 100)[1]
        @test compute_normalisation(1.5, [0]) ≈ 1 / sqrt(integral)
    end

    @testset "Boys" begin
        boystest(m, t) = quadgk(x -> x^(2m) * exp(-t * x^2), 0, 1)[1]
        @test boystest(0, 1.5)   ≈ boys(0, 1.5)
        @test boystest(0, 0.375) ≈ boys(0, 0.375)
        @test boystest(1, 1.5)   ≈ boys(1, 1.5)
        @test boystest(1, 0.375) ≈ boys(1, 0.375)
    end

    @testset "H 1s" begin
        bas = pyimport("pyscf.gto").parse("H S\n1.5 1.0\nH S\n0.9 1.0")
        mol  = pyimport("pyscf.gto").M(basis=bas, atom="H 0 0 0", spin=1, unit="Bohr")
        ref = mol.intor("int2e")

        #                    c    zeta  P         am
        basis_functions = [[(1.0, 1.5, [0, 0, 0], 0)],
                           [(1.0, 0.9, [0, 0, 0], 0)]]
        res = compute_integrals(ERI(), basis_functions)
        @test res ≈ ref
    end

    @testset "H₂ 1s" begin
        bas = pyimport("pyscf.gto").parse("H S\n1.0 1.0\n")
        mol  = pyimport("pyscf.gto").M(basis=bas, atom="H 0 0 0; H 0 1 0", unit="Bohr")
        ref = mol.intor("int2e")

        #                    c    zeta  P         am
        basis_functions = [[(1.0, 1.0, [0, 0, 0], 0)],
                           [(1.0, 1.0, [0, 1, 0], 0)]]
        res = compute_integrals(ERI(), basis_functions)
        @test res ≈ ref
    end

    @testset "H sto-3g" begin
        mol  = pyimport("pyscf.gto").M(basis="sto-3g", atom="H 0 0 0", spin=1, unit="Bohr")
        ref = mol.intor("int2e")

        #                    c           zeta         P         am
        basis_functions = [[(0.15432897, 3.42525091, [0, 0, 0], 0),
                            (0.53532814, 0.62391373, [0, 0, 0], 0),
                            (0.44463454, 0.16885540, [0, 0, 0], 0)]]
        res = compute_integrals(ERI(), basis_functions)
        @test res ≈ ref atol=1e-7  # TODO That's not a great agreement
    end

    @testset "H₂ sto-3g" begin
        mol  = pyimport("pyscf.gto").M(basis="sto-3g", atom="H 0 0 0; H 0 1 0", unit="Bohr")
        ref = mol.intor("int2e")

        #                    c           zeta         P         am
        basis_functions = [[(0.15432897, 3.42525091, [0, 0, 0], 0),
                            (0.53532814, 0.62391373, [0, 0, 0], 0),
                            (0.44463454, 0.16885540, [0, 0, 0], 0)],
                           [(0.15432897, 3.42525091, [0, 1, 0], 0),
                            (0.53532814, 0.62391373, [0, 1, 0], 0),
                            (0.44463454, 0.16885540, [0, 1, 0], 0)]]
        res = compute_integrals(ERI(), basis_functions)
        @test res ≈ ref atol=1e-7  # TODO That's not a great agreement
    end

    @testset "Be 3-21G" begin
        mol  = pyimport("pyscf.gto").M(basis="3-21g", atom="Be 0 0 0", unit="Bohr")
        ref = mol.intor("int2e")
        #                    c           zeta         P         am
        basis_functions = [[( 0.0644263, 71.8876000, [0, 0, 0], 0),
                            ( 0.3660960, 10.7289000, [0, 0, 0], 0),
                            ( 0.6959340,  2.2220500, [0, 0, 0], 0)],
                           [(-0.4210640,  1.2954800, [0, 0, 0], 0),
                            ( 1.2240700,  0.2688810, [0, 0, 0], 0)],
                           [(-0.4210640,  1.2954800, [0, 0, 0], 1),
                            ( 1.2240700,  0.2688810, [0, 0, 0], 1)],
                           [( 1.0000000,  0.0773500, [0, 0, 0], 0)],
                           [( 1.0000000,  0.0773500, [0, 0, 0], 1)]]
        res = compute_integrals(ERI(), basis_functions)
        @test res ≈ ref atol=1e-7  # TODO That's not a great agreement
    end
end


using SMC, Base.Test
srand(125)

## Linear Gaussian

dx, dy = 5, 3
A = randn(dx,dx); A /= 0.9norm(A)
B = randn(dy,dx); B /= 1.1norm(B)
Q = randn(dx,dx); Q *= Q'; Q += 0.5*eye(dx); Q += Q'; Q /= 5
R = randn(dy,dy); R *= R'; R += 0.5*eye(dy); R += R'; R /= 20

lg  = LinearGaussian(A,B,Q,R)
hmm = HMM(lg)
x0  = randn(dx)

K = 50
N = 150
(states, observations) = generate(lg, x0, K)

srand(155)
@time (psf, ess) = particlefilter(hmm, observations, N, bootstrapprop(lg))

@test length(psf)==K

pfm  = mean(psf)
pfmm = zeros(dx,K)
for k in 1:K
    pfmm[:,k] = pfm[k]
end

@test norm(pfmm-states)/norm(states) < 0.4
println("PF    : $(norm(pfmm-states)/norm(states))")

srand(521)
@time psffbs  = particlesmoother_ffbs(hmm, psf)

@test length(psffbs)==K

psm  = mean(psffbs)
psmm = zeros(dx,K)
for k in 1:K
    psmm[:,k] = psm[k]
end

@test norm(psmm-states)/norm(states) < 0.25
println("PSFFBS: $(norm(psmm-states)/norm(states))")

srand(521)
@time (psbbis, ess) = particlesmoother_bbis(hmm, observations,
                                        psf, bootstrapprop(lg))

psm3  = mean(psbbis)
psmm3 = zeros(dx,K)
for k in 1:K
    psmm3[:,k] = psm3[k]
end

@test norm(psmm3-states)/norm(states) < 0.23
println("PSBISQ: $(norm(psmm3-states)/norm(states))")

srand(521)
@time (pslbbis, ess) = particlesmoother_lbbis(hmm, observations,
                                              psf, bootstrapprop(lg))

psm4  = mean(pslbbis)
psmm4 = zeros(dx,K)
for k in 1:K
    psmm4[:,k] = psm4[k]
end

@test norm(psmm4-states)/norm(states) < 0.23
println("PSBISL: $(norm(psmm4-states)/norm(states))")

srand(521)
@time (psllbbis, ess) = particlesmoother_llbbis(hmm, observations, psf, 25,
                                                bootstrapprop(lg))

psm5  = mean(psllbbis)
psmm5 = zeros(dx,K)
for k in 1:K
    psmm5[:,k] = psm5[k]
end

@test norm(psmm5-states)/norm(states) < 0.23
println("PSBISLL: $(norm(psmm5-states)/norm(states))")


using WAV
using ChirpSignal

fs = 16000.0; # sampling frequency
fl = 500.0;   # lower bound frequency
fh = 8000.0;  # upper bound frequency
T  = 5.0 ;    # total time of chirp signal

wav = chirp(T, fs, fl, fh; method="linear");
wavwrite(wav, "./doc/chirpLinear.wav", Fs=fs, nbits=32)

wav = chirp(T, fs, fl, fh; method="quadratic");
wavwrite(wav, "./doc/chirpQuadratic.wav", Fs=fs, nbits=32)

wav = chirp(T, fs, fl, fh; method="logarithmic");
wavwrite(wav, "./doc/chirpLogarithmic.wav", Fs=fs, nbits=32)

wav = chirp(T, fs, fl, fh; method="exponential");
wavwrite(wav, "./doc/chirpExponential.wav", Fs=fs, nbits=32)

fun(t) = 2000*t
wav = chirp(2.0, 16000.0, fun)
wavwrite(wav, "./doc/0_4000hz_linear.wav", Fs=fs, nbits=32)

wav = chirp(2.0, 16000.0, t -> 2000*t^2)
wavwrite(wav, "./doc/0_8000hz_quadratic.wav", Fs=fs, nbits=32)


#=

   https://reference.wolfram.com/language/ref/BinomialDistribution.html
   """
   Two players roll dice. If the total of both numbers is less than 10, the second player
   is paid 4 cents; otherwise the first player is paid 9 cents. Is the game fair?:

   ...

   [The game is not fair: mean scores are 1.5 and 3.33333.]

   """

   Model 1: Is the game fair?
   mean L: 10.0
   mean prob1: 0.16558055555555556
   mean prob2: 0.8344194444444445
   mean prob1Times9: 1.490225
   mean prob2Times4: 3.337677777777778
   mean total: -17.7128

   Model 2: What L would be fair?
   mean L: 7.0457
   mean prob1: 0.5409611111111111
   mean prob2: 0.4590388888888889
   mean prob1Times9: 4.86865
   mean prob2Times4: 1.8361555555555555
   mean total: 6.9705

   Note: Model 2 (checkFair==true) is NOT correct.
   In this Turing model, prob1Times9 and prob2Time2 are too different
   (4.8 vs 1.8) to be considered the same! I'm not sure why the constraint
        true ~ Dirac(prob1Times9 == prob2Times4)
   don't work as expected....

   According to the WebPPL model, the optimal L is 9 (not 7)
   and prob1Times9 ~ prob2Times4.
   Here are the expected values from the WebPPL model:
    [ 'total', -5.998999999999998 ],
    [ 'prob1', 0.3333333333333333 ],
    [ 'prob2', 0.75 ],
    [ 'prob1Times9', 3 ],
    [ 'prob2Times4', 3 ],
    [ 'L', 8.9286 ] ]

   Cf ~/webppl/binomial_dice2.wppl

=#

using Turing
include("jl_utils.jl")

@model function binomial_coin(checkFair=false)
    n = 36

    # Limit (optimal if checkFair==true)
    L ~ DiscreteUniform(2,12)

    player1 ~ filldist(DiscreteUniform(1,6), n)
    player2 ~ filldist(DiscreteUniform(1,6), n)

    #  Outcome (seen as player1)
    outcome = tzeros(n)
    for i in 1:n
        if player1[i] + player2[i] < L
            outcome[i] ~ Dirac(-4)
        else
            outcome[i] ~ Dirac(9)
        end
    end

    # Probabilities given the limit L
    prob1 ~ Dirac(sum([player1[i] + player2[i] >= L for i in 1:n]) / n)
    prob2 ~ Dirac(sum([player1[i] + player2[i] < L for i in 1:n]) / n)

    # The total for L games
    total ~ Dirac(sum(outcome[i] for i in 1:L))

    prob1Times9 ~ Dirac(prob1*9)
    prob2Times4 ~ Dirac(prob2*4)

    if checkFair
        # What is L for a fair game?
        # This don't give the expected answer
        # true ~ Dirac(prob1Times9 == prob2Times4)
        true ~ Dirac(L == 8.9)
    else
        # The given
        true ~ Dirac(L == 10)
    end

end

println("Model 1: Is the game fair?")
model = binomial_coin(false)

num_chains = 4
# chs = sample(model, Prior(), 1000)
# chs = sample(model, MH(), 10_000)
# chs = sample(model, PG(5), 10_000)
# chs = sample(model, IS(), 10_000)
chs = sample(model, SMC(), 10_000)
# chs = sample(model, SMC(), MCMCThreads(), 10_000, num_chains)

# chs = sample(model, SGLD(), 10_000)
# chs = sample(model,NUTS(), 10_000)
# chs = sample(model,HMC(0.01,5), 10_000)

# display(chs)
# display(plot(chs))

# probs = [:L, :prob1,:prob2,:prob1B,:prob2B,:prob1Times9,:prob2Times4,:total]
probs = [:L, :prob1,:prob2,:prob1Times9,:prob2Times4,:total]
for p in probs
    chsp = chs[p]
    println("mean ", p,": ", mean(chsp))
end

println("\nModel 2: What L would be fair?")
model = binomial_coin(true)

num_chains = 4
# chs = sample(model, Prior(), 1000)
# chs = sample(model, MH(), 10_000)
# chs = sample(model, PG(15), 10_000)
# chs = sample(model, IS(), 10_000)
chs = sample(model, SMC(), 10_000)
# chs = sample(model, SMC(), MCMCThreads(), 10_000, num_chains)

# chs = sample(model, SGLD(), 10_000)
# chs = sample(model,NUTS(), 10_000)
# chs = sample(model,HMC(0.01,5), 10_000)

# display(chs)
# display(plot(chs))

for p in probs
    chsp = chs[p]
    println("mean ", p,": ", mean(chsp))
end


using Test
using TopOpt
using TopOpt.TopOptProblems: getE
using Base.Iterators
# if get(ENV, "CI", nothing) != "true"
#     import Makie
#     @eval using TopOpt.TrussTopOptProblems.TrussVisualization: visualize
# end

ins_dir = joinpath(@__DIR__, "instances", "ground_meshes");

@testset "Tim problem to solve" for (problem_dim, lc_ind) in product(["2d", "3d"], [0, 1])
    # problem_dim = "2d"
    # file_name = "tim_2d.json"
    # lc_ind = 1
    file_name = "tim_$(problem_dim).json"
    problem_file = joinpath(ins_dir, file_name)

    node_points, elements, mats, crosssecs, fixities, load_cases = parse_truss_json(problem_file);
    loads = load_cases[string(lc_ind)]

    problem = TrussProblem(Val{:Linear}, node_points, elements, loads, fixities, mats, crosssecs);

    ndim, nnodes, ncells = length(node_points[1]), length(node_points), length(elements)
    @test getE(problem) == [m.E for m in mats]

    V = 0.3 # volume fraction
    xmin = 0.001 # minimum density
    rmin = 4.0; # density filter radius

    penalty = TopOpt.PowerPenalty(1.0) # 1
    solver = FEASolver(Displacement, Direct, problem, xmin = xmin,
        penalty = penalty);
    ## call solver to trigger assemble!
    solver()

    # ##############################
    # #! buckling
    # import TrussTopOpt.TrussTopOptProblems: default_quad_order
    # einfo = ElementFEAInfo(problem, TrussTopOpt.TrussTopOptProblems.default_quad_order(problem), Val{:Static})
    # ginfo = GlobalFEAInfo(problem)

    # using TrussTopOpt.TrussTopOptProblems: buckling, get_Kσs
    # buckling(problem, solver.globalinfo, solver.elementinfo)

    # ##############################

    # TODO TopOpt.LogBarrier
    # TODO linear_elasticity, du/dx
    # * Compliance
    comp = TopOpt.Compliance(problem, solver)
    obj = comp
    volfrac = TopOpt.Volume(problem, solver)
    constr = x -> volfrac(x) - V

    mma_options = options = Nonconvex.MMAOptions(
        maxiter = 3000, tol = Nonconvex.Tolerance(kkt = 0.001),
    )
    convcriteria = Nonconvex.KKTCriteria()
    x0 = fill(V, length(solver.vars))
    optimizer = Optimizer(
        obj, constr, x0, Nonconvex.MMA87(),
        options = mma_options, convcriteria = convcriteria,
    )
    simp = SIMP(optimizer, solver, penalty.p)

    # ? 1.0 might induce an infeasible solution, which gives the optimizer a hard time to escape
    # from infeasible regions and return a result
    result = simp(x0);

    println("="^10)
    println("tim-$(problem_dim) - LC $(lc_ind) - #elements $(ncells), #dof: $(ncells*ndim): opt iter $(simp.optimizer.workspace.iter)")
    println("$(result.convstate)")

    solver()

    # if get(ENV, "CI", nothing) != "true"
    #     fig = visualize(
    #         problem, solver.u; topology = result.topology, vector_arrowsize = 0.1,
    #         vector_linewidth=0.8, default_exagg_scale=ndim == 3 ? 1.0 : 0.01,
    #         exagg_range = ndim == 3 ? 10.0 : 0.1,
    #     )
    #     Makie.display(fig)
    # end
end # end testset


using YaoBlocks, SymEngine, LuxurySparse, LinearAlgebra

Base.promote_rule(::Type{Bool}, ::Type{Basic}) = Basic
Base.conj(x::Basic) = real(x) - imag(x)

const SymReal = Union{Basic, SymEngine.BasicRealNumber}
YaoBlocks.RotationGate(block::GT, theta::T) where {N, T <: SymReal, GT<:AbstractBlock{N}} = RotationGate{N, T, GT}(block, theta)

YaoBlocks.phase(θ::SymReal) = PhaseGate(θ)
YaoBlocks.shift(θ::SymReal) = ShiftGate(θ)

YaoBlocks.mat(::Type{Basic}, ::HGate) = 1/sqrt(Basic(2)) * Basic[1 1;1 -1]
YaoBlocks.mat(::Type{Basic}, ::XGate) = Basic[0 1;1 0]
YaoBlocks.mat(::Type{Basic}, ::YGate) = Basic[0 -1im;1im 0]
YaoBlocks.mat(::Type{Basic}, ::ZGate) = Basic[1 0;0 -1]

YaoBlocks.mat(gate::ShiftGate{<:SymReal}) =
    Diagonal([1.0, exp(im * gate.theta)])
YaoBlocks.mat(gate::PhaseGate{<:SymReal}) =
    exp(im * gate.theta) * IMatrix{2}()
function YaoBlocks.mat(R::RotationGate{N, <:SymReal}) where N
    I = IMatrix{1<<N}()
    return I * cos(R.theta / 2) - im * sin(R.theta / 2) * mat(Basic,R.block)
end

YaoBlocks.PSwap{N}(locs::Tuple{Int, Int}, θ::SymReal) where N = YaoBlocks.PutBlock{N}(rot(ConstGate.SWAPGate(), θ), locs)

YaoBlocks.pswap(n::Int, i::Int, j::Int, α::SymReal) = PSwap{n}((i,j), α)
YaoBlocks.pswap(i::Int, j::Int, α::SymReal) = n->pswap(n,i,j,α)


function Vert(PP :: Array{VerticeType,1}, NewCut,CutNum)
    const EPS = 0;
    V_Plus = [];
    v_Plus = 0;
    V_Minus = [];
    v_Minus = 0;
    for i = 1:length(PP)
        if (sum(NewCut[1:end-1].*PP[i].Point) + NewCut[end] > EPS)
            push!(V_Plus, i)
            v_Plus = v_Plus + 1;
        else
            push!(V_Minus, i)
            v_Minus = v_Minus + 1;
        end
    end
    # println(V_Plus)
    # println(V_Minus)
# ---------------------------------------------------------------------------
    # Step I
    if v_Minus == 0
        error("The cut makes it infeasible!!!!")
        return []
    elseif v_Plus == 0
        println("The cut is redundant!!!!")
##############################################################
        return []
##########################################################3##
        return PP
    end

    if v_Plus <= v_Minus
        V1 = V_Plus;
        V2 = V_Minus;
    else
        V2 = V_Plus;
        V1 = V_Minus;
    end
# ---------------------------------------------------------------------------
    # Step II

    VS = [];
    counter = 1;
    for u in V1
        for v in PP[u].adj
            if v in V2
                W = FindIntersectingPoint(PP[u].Point,PP[v].Point,NewCut)
                NewPoint = VerticeType([],[],[],-1)
                NewPoint.Point = W;
                if v_Plus <= v_Minus
                    PP[v].adj = filter(i -> (i != u), PP[v].adj);
                    push!(PP[v].adj,length(PP)+counter);
                    counter = counter + 1;
                    push!(NewPoint.adj,v);
                else
                    PP[u].adj = filter(i -> (i != v), PP[u].adj);
                    push!(PP[u].adj,length(PP)+counter)
                    counter = counter + 1;
                    push!(NewPoint.adj,u)
                end
                temp = intersect(PP[u].J,PP[v].J)
                for ii in temp
                    push!(NewPoint.J,ii)
                end
                push!(NewPoint.J,CutNum)
                if isempty(VS)
                    VS = [NewPoint];
                else
                    push!(VS,NewPoint);
                end
            end
        end
    end
# ---------------------------------------------------------------------------
    IndTrans = Dict();
    for i = 1:v_Minus
        IndTrans[V_Minus[i]]=i;
        for j = 1:length(VS)
            VS[j].adj[VS[j].adj .== V_Minus[i]] =i;
        end
    end

    for i in V_Minus
        for j = 1:length(PP[i].adj)
            if PP[i].adj[j] > length(PP)
                PP[i].adj[j] = PP[i].adj[j]- v_Plus;
            else
                PP[i].adj[j] = IndTrans[PP[i].adj[j]];
            end
        end
    end
    PP = [PP[i] for i in V_Minus];

    # # Step III
    for u = 1:length(VS), v = 1:length(VS)
        if length(intersect(VS[u].J,VS[v].J)) == length(PP[1].Point)-1
            if !(length(V_Minus)+v in VS[u].adj)
                push!(VS[u].adj,length(V_Minus)+v)
            end
            if !(length(V_Minus)+u in VS[v].adj)
                push!(VS[v].adj,length(V_Minus)+u)
            end
        end
    end

    for i = 1:length(VS)
        VS[i].T = prod(VS[i].Point);
        push!(PP,VS[i])
    end
    return PP;

end


export SymSemiseparableChol

struct SymSemiseparableChol <: SymSemiseparableCholesky
    n::Int64
    p::Int64
    U::AbstractArray
    W::AbstractArray
end

# Constuctors
function SymSemiseparableChol(U::AbstractArray, W::AbstractArray)
	if size(U,1) == size(W,1) && size(U,2) == size(W,2)
		return SymSemiseparableChol(size(U,1),size(U,2),U,W)
	else
		error("Dimension mismatch between the generators U and W")
	end
end

# Mappings
mul!(y::AbstractArray, L::SymSemiseparableChol, 		         b::AbstractArray) =
	ss_tri_mul!(y, L.U,   L.W,   b)
mul!(y::AbstractArray, L::AdjointOperator{SymSemiseparableChol}, b::AbstractArray) =
	ssa_tri_mul!(y, L.A.U, L.A.W, b)
inv!(y::AbstractArray, L::SymSemiseparableChol, 		    	 b::AbstractArray) =
	ss_forward!(y, L.U,   L.W,   b);
inv!(y::AbstractArray, L::AdjointOperator{SymSemiseparableChol}, b::AbstractArray) =
	ssa_backward!(y, L.A.U, L.A.W, b);
newlogdet(L::SymSemiseparableChol) = ss_logdet(L.U, L.W)
newlogdet(L::AdjointOperator{SymSemiseparableChol}) = ss_logdet(L.A.U, L.A.W)

#### Inverse of a SymSemiseparableChol using ####
function inv(L::SymSemiseparableChol)
	return L'\(L\Diagonal(ones(L.n)))
end
function inv(L::SymSemiseparableChol, b::AbstractArray)
	return L'\(L\b)
end

########################################################################
#### Cholesky factoriaztion of:                                     ####
####    Extended generator representable {p}-semiseperable matrices ####
########################################################################

# Creating W, s.t. L = tril(UW')
function ss_create_w(U,V)
    n,p = size(U);
    wTw = zeros(p,p);
    W = zeros(n,p)
    for j = 1:n
        tmpu = U[j,:];
        tmp  = V[j,:] - wTw*tmpu;
        w = tmp/sqrt(abs(tmpu'*tmp))
        W[j,:] = w;
        wTw = wTw + w*w';
    end
    return W
end
#### Triangular product ####
function ss_tri_mul!(Y::AbstractArray,U::AbstractArray,
                     W::AbstractArray,X::AbstractArray)
     n, m = size(U);
     mx = size(X,2);
     Wbar = zeros(m,mx);
     for i = 1:n
         tmpW   = W[i,:]
         tmpU   = U[i,:]
         tmpX   = X[i:i,:]
         Wbar  += tmpW .* tmpX;
         Y[i,:] = Wbar'*tmpU;
     end
end
#### Adjoint triangular product ####
function ssa_tri_mul!(Y::AbstractArray,U::AbstractArray,
                      W::AbstractArray,X::AbstractArray)
     n, m = size(U);
	 mx = size(X,2);
     Ubar = zeros(m,mx);
     Ubar = Ubar + U'*X;
     for i = 1:n
         tmpW = W[i,:]
         tmpU = U[i,:]
         tmpX = X[i:i,:]
         Y[i,:] = Ubar'*tmpW;
         Ubar  -= tmpU .* tmpX
     end
end
#### Forward substitution (solve Lx = b) ####
function ss_forward!(X::AbstractArray, U::AbstractArray,
                     W::AbstractArray, B::AbstractArray)
    n, m = size(U)
    mx = size(B,2)
    Wbar = zeros(m, mx);
    for i = 1:n
        tmpU = U[i,:]
        tmpW = W[i,:]
        X[i,:] = (B[i:i,:] - tmpU'*Wbar)./(tmpU'*tmpW);
        Wbar += tmpW .* X[i:i,:];
    end
end
#### Backward substitution (solve L'x = b) ####
function ssa_backward!(X::AbstractArray, U::AbstractArray,
                       W::AbstractArray, B::AbstractArray)
    n, m = size(U)
    mx = size(B,2)
    Ubar = zeros(m,mx);
    for i = n:-1:1
        tmpU = U[i,:];
        tmpW = W[i,:];
        X[i,:] = (B[i:i,:] - tmpW'*Ubar)/(tmpU'*tmpW);
        Ubar += tmpU .* X[i:i,:];
    end
end
#### Logarithm of determinant ####
function ss_logdet(U::AbstractArray, W::AbstractArray)
    n = size(U, 1);
    a = 0;
    for i = 1:n
        a += log(dot(U[i,:],W[i,:]));
    end
    return a
end