function AngleToX360(v)
    θ = angle(Complex(v...))
    if θ < 0
        θ = 2*π+θ
    end
    return θ
end

function Angle(u,v)
    α = (u⋅v)/(norm(u)*norm(v))
    if abs(α) > 1
        α=sign(α)*1
    end
    return acos(α)
end

function unwrap(v, inplace=false)
  # currently assuming an array
  unwrapped = inplace ? v : copy(v)
  for i in 2:length(v)
    while unwrapped[i] - unwrapped[i-1] >= pi/2
      unwrapped[i] -= pi
    end
    while unwrapped[i] - unwrapped[i-1] <= -pi/2
      unwrapped[i] += pi
    end
  end
  return unwrapped
end

function Ellipse(n;a=1,b=1,max=2π,θ=0,center=[0.0,0.0])
    x = [a*cos(t)*cos(θ) - b*sin(t)*sin(θ) for t in 0:(2π/n):2π] .+ center[1]
    y = [a*cos(t)*sin(θ) + b*sin(t)*cos(θ) for t in 0:(2π/n):2π] .+ center[2]
    return cat(x,y,dims=2)
end

beetle_length = 13.0
beetle_width = 7.0
dt = 1.0/30.0

function Density(X)
    tess,inds = AlphaShapes.GetDelaunayTriangulation(X,true);
	d = zeros(size(X,1))
	for i in 1:size(X,1)
    	d[i] = AlphaShapes.WeightedDTFELocalDensity(i,tess,inds)
	end
	return d
end

flatten(x) = reshape(x,prod(size(x)))

function logspace(start,stop,n=10)
    s = log10(start)
    S = log10(stop)
    x = []
    for i in linspace(s,S,n)
        push!(x,10^i)
    end
    return x
end

function linspace(l,u,n=10)
    return collect(range(l,stop=u,length=n))
end

function DVBinPlot(;D,V,n_bins=10,beetlescale=true,beetle_length=14.0,xscale=:log10,FPS=30.0,errors=true)
    IFT = 1/FPS
    bin_values = logspace(minimum(D),maximum(D),n_bins)

    Bin_V_SE = [[] for i in 1:n_bins]
    Bin_V = zeros(n_bins)
    number = zeros(n_bins)

    for b in 1:(n_bins-1)
        ind1 = findall(x->x.>=bin_values[b],D)
        v = V[ind1]
        ind = findall(x->x.<bin_values[b+1],D[ind1])
        if size(ind,1) > 0
            Bin_V[b] += sum(v[ind])
            append!(Bin_V_SE[b],v[ind])
            number[b] += size(ind,1)
        end
    end

    if beetlescale
        SE = zeros(n_bins)
        for i in 1:n_bins
            if size(Bin_V_SE[i],1) > 0
                SE[i] = std(Float64.(Bin_V_SE[i]) ./ (beetle_length * IFT)) / sqrt(number[i])
            end
        end
        Bin_V = Bin_V[2:end-1] ./ ((number[2:end-1] .* beetle_length) * IFT)
        bin_values = bin_values[2:end-1].*beetle_length^2.0
        if errors
            scatter(bin_values,Bin_V,xscale=xscale,yerr=SE,label="")
        else
            plot(bin_values,Bin_V,xscale=xscale,label="")
        end
    else
        SE = zeros(n_bins)
        for i in 1:n_bins
            if size(Bin_V_SE[i],1) > 0
                SE[i] = std(Float64.(Bin_V_SE[i])) / sqrt(number[i])
            end
        end
        if errors
            plot(bin_values,Bin_V,xscale=xscale,yerr=SE,label="")
        else
            plot(bin_values,Bin_V,xscale=xscale,label="")
        end
    end
    println(number)
    return bin_values,Bin_V,SE[2:end-1]
end

function MSD(X,t=size(X,1),lags=120)
    l = 0:lags
    D = zeros(size(l,1))
    C = zeros(size(l,1),size(X,1))
    @showprogress for i in 1:t
            for τ in l
                if i+τ > size(X,1)
                    continue
                end
                for j in 1:size(X,2)
                    p = norm(X[i+τ,j,1:2] - X[i,j,1:2])^2
                    if isnan(p) == false
                        D[τ+1] += p
                        C[τ+1,i] += 1
                    end
                end
            end
        end
    return D ./ sum(C,dims=2),collect(l)
end

function msdθ(X,t=size(X,1),lags=30)
    l = 0:lags
    D = zeros(size(l,1))
    A = X[:,:,3]
    C = zeros(size(l,1),t)
    @showprogress for i in 1:t
            for τ in l
                if i+τ > size(A,1)
                    continue
                end
                for j in 1:size(X,2)
                    p = (A[i+τ,j] - A[i,j])^2
                    if isnan(p) == false
                        D[τ+1] += p
                        C[τ+1,i] += 1
                    end
                end
            end
        end
    return D ./ sum(C,dims=2),collect(l)
end

function VelocityHeading(X,t,b)
    h = X[t+1,b,1:2].-X[t,b,1:2]
    return h./norm(h)
end

function Arrow(d,pos,K=1,c="red",linewidth=1)
    v = ones(2).*d*K + pos
    plot!([pos[1],v[1]],[pos[2],v[2]],color=c,label="",linewidth=linewidth,arrow=arrow())
end

function TimeBinDensity(BX,bins,T;length_scale=2.0)
    times = Int.(floor.(range(1,stop=size(T,1),length=bins)))
    collect(times)
    pdfs = zeros(bins-1,length(BX))
    for t in 1:size(times,1)-1
        start = times[t]
        stop = times[t+1]
        Y = T[start:stop,:,1:2]

        D = zeros(size(Y,1),size(Y,2))
        for i in 1:size(Y,1)
            D[i,:] = Density(Y[i,:,1:2]./length_scale)
        end
        D = replace(D,NaN => 0.0)
        x = flatten(D)
        x = x[x.>0]
        x = x[x.<Inf]

        h = fit(Histogram,x,BX)
        h = normalize(h, mode=:density)
        h.weights = h.weights./size(x,1)
        push!(h.weights,0.0)
        pdfs[t,:] = h.weights
    end
    return pdfs
end

function TimeBinDistribution(x,data;bins=10)
    times = Int.(floor.(range(1,stop=size(data,1),length=bins)))
    collect(times)
    pdfs = zeros(bins-1,length(x))
    for t in 1:size(times,1)-1
        start = times[t]
        stop = times[t+1]
        Y = []
        for i in start:stop
            for j in 1:size(data[i],1)
                push!(Y,data[i][j])
            end
        end
        y = replace!(Y,NaN => 0.0)
        y = y[y.>0]
        y = y[y.<Inf]

		h = fit(Histogram,y,x)
		h = normalize(h, mode=:density)
		h.weights = h.weights./size(x,1)
		push!(h.weights,0.0)
		pdfs[t,:] = h.weights
    end
    return pdfs
end

## Collision Filtering

using Distances

function Collisions(R=13.0,data=[X200])
    R = 13.0
    Data = []
    Cs = []
    for X in data
        C = Bool.(zeros(size(X,1),size(X,2),size(X,2)))
        for t in 1:size(X,1)
            d = pairwise(Euclidean(),X[t,:,1:2]').<=R
            for b in 1:size(X,2)
                for j in findall(x->x.==1,d[b,:])
                    C[t,b,j] = 1
                end
                C[t,b,b] = 0
            end
        end

        push!(Cs,C)

        for b in 1:size(C,2)
            flag = sum(C[1,b,:]) > 0
            d = []
            for t in 2:size(C,1)
                if sum(C[t,b,:]) > 0
                    push!(d,[t,X[t,b,:]...,findall(x->x.==1,C[t,b,:])...])
                    flag = 0
                end
                if sum(C[t,b,:]) == 0
                    if flag == 0
                        push!(Data,d)
                        d = []
                    end
                    flag = 1
                end
            end
            push!(Data,d)
        end
    end
    @info length(Data)
    return Cs
end

function CollisionFreeTrajectories(C,X)
    T = []
    d = zeros(size(X,1),size(X,2))
    for i in 1:size(X,1)
        d[i,:] = Density(X[i,:,1:2]./13.0)
    end
    @showprogress for i in 1:size(C,2)
        c = sum(C[:,i,:],dims=2)[:,1]
        flag = c[1]
        for j in 1:size(c,1)
            tau = []
            offset = 0
            while length(findall(x->x.==0,c)) > 0
                s = findfirst(x->x.==0,c)
                e = findfirst(x->x.==1,c)
                if e == nothing
                    e = length(c)
                end
                for k in s:1:e
                    push!(tau,([X[offset+k,i,:]...,d[offset+k,i]],offset+k))
                end
                offset += e
                c = c[(e+1):end]
            end
            if length(tau) > 0
                push!(T,tau)
            end
        end
    end

    # gets length count distribution
    L = []
    O = []
    for i in 1:size(T,1)
        t = T[i]
        if length(t) > 1
            o = [t[1][1]]
            k = t[1][2]
            l = 0
            for i in 2:size(t,1)
                if t[i][2] - k == 1
                    l += 1
                    k += 1
                    push!(o,t[i][1])
                else
                    k = t[i][2]
                    push!(L,l)
                    l = 0
                    o = [t[i][1]]
                    push!(O,o)
                end
            end
        end
    end
    return T,L,O
end

function MSDθ_CF(data,dt=1.0/30.0,lag_max=30)
    C = [[] for i in 1:(lag_max+1)]
    @showprogress for i in 1:length(data)
        d = data[i]
        if length(d) >= 2
            for t in 1:size(d,1)
                for (i,tau) in enumerate(0:lag_max)
                    if 1 < t+tau < size(d,1)
                        c = (d[t+tau][3] - d[t][3])^2.0
                        if isnan(c) == false
                            push!(C[i],c)
                        end
                    end
                end
            end
        end
    end
    return C,collect(0:lag_max).*dt
end

function MSD_CF(data,dt=1.0/30.0,lag_max=30)
    C = [[] for i in 1:(lag_max+1)]
    @showprogress for i in 1:length(data)
        d = data[i]
        if length(d) >= 2
            for t in 1:size(d,1)
                for (i,tau) in enumerate(0:lag_max)
                    if 1 < t+tau < size(d,1)
                        c = (norm(d[t+tau][1:2] .- d[t][1:2]))^2.0
                        if isnan(c) == false
                            push!(C[i],c)
                        end
                    end
                end
            end
        end
    end
    return C,collect(0:lag_max).*dt
end

function NVCorr(data,dt=1.0/30.0,lag_max=30,min_length=2,info=false)
    C = [[] for i in 1:(2*lag_max+1)]
    n = 0
    @showprogress for i in 1:length(data)
        d = data[i]
        if length(d) >= min_length
            n += 1
			# central difference
            for t in 1:size(d,1)
                for (i,tau) in enumerate(-lag_max:1:lag_max)
                    if 1 < t+tau < size(d,1)-1
						# central difference about t+tau
                        h = d[t+tau+1][1:2] .- d[t+tau-1][1:2]
                        h = h ./ norm(h)
                        o = [cos(d[t][3]),sin(d[t][3])]
                        c = h⋅o
                        if isnan(c) == false
                            push!(C[i],c)
                        end
                    end
                end
            end
        end
    end
    if info
        @info n
    end
    return C,collect(-lag_max:1:lag_max).*dt
end


""" Interoperation with mesh files.

This module enables delta sets to be imported from mesh files supported by
MeshIO.jl and for delta sets to be converted to meshes, mainly for the purposes
of plotting. Meshes are represented by the `GeometryBasics.Mesh` type.
"""
module MeshInterop

using FileIO, MeshIO, GeometryBasics

using Catlab.CategoricalAlgebra: copy_parts!
using ..SimplicialSets, ..DiscreteExteriorCalculus
import ..SimplicialSets: EmbeddedDeltaSet2D

# Export meshes
###############

""" Construct a Mesh object from an embedded delta set.
"""
function GeometryBasics.Mesh(ds::EmbeddedDeltaSet2D)
  points = Point{3, Float64}[point(ds)...]
  tris = TriangleFace{Int}[zip(triangle_vertices(ds)...)...]
  Mesh(points, tris)
end

""" Construct a Mesh object from a dual embedded delta set.
"""
function GeometryBasics.Mesh(ds::EmbeddedDeltaDualComplex2D{O,R,P};
                             primal=false) where {O,R,P}
  if primal
    ds′ = EmbeddedDeltaSet2D{O,P}()
    copy_parts!(ds′, ds)
    Mesh(ds′)
  else
    points = Point{3, Float64}[dual_point(ds)...]
    tris = TriangleFace{Int}[zip(dual_triangle_vertices(ds)...)...]
    Mesh(points, tris)
  end
end

# Import meshes
###############

""" Construct EmbeddedDeltaSet2D from Mesh object

This operator should work for any triangular mesh object. Note that it will
not preserve any normal, texture, or other attributes from the Mesh object.

The `force_unique` flag merges all points which are at the same location. This
is necessary to process `.stl` files which references points by location.
"""
function EmbeddedDeltaSet2D(m::GeometryBasics.Mesh; force_unique=false)
  coords = metafree.(coordinates(m))
  ind_map = 1:length(coords)
  if(force_unique)
    indices = unique(i->coords[i], 1:length(coords))
    val2ind = Dict(coords[indices[i]]=>i for i in 1:length(indices))
    ind_map = map(c->val2ind[c], coords)
    coords = coords[indices]
  end
  tris = faces(m)
  s = EmbeddedDeltaSet2D{Bool, eltype(coords)}()
  add_vertices!(s, length(coords), point=coords)
  for tri in tris
    tri = ind_map[convert.(Int64, tri)]
    glue_sorted_triangle!(s, tri...)
  end
  # Properly orient the delta set
  orient!(s)
  s
end

""" Construct EmbeddedDeltaSet2D from mesh file

This operator should work for any file support for import from MeshIO. Note
that it will not preserve any normal, texture, or other data beyond points and
triangles from the mesh file.
"""
function EmbeddedDeltaSet2D(fn::String)
  if(splitext(fn)[2] == ".stl")
    # The .stl format references points by location so we apply the
    # force_unique flag
    EmbeddedDeltaSet2D(FileIO.load(fn); force_unique=true)
  else
    EmbeddedDeltaSet2D(FileIO.load(fn))
  end
end

end


include("centrality.jl")

using Statistics

function main()
    m = 3
    dim = 7
    T = SymTensor3([1, 1, 3, 5],
                   [2, 2, 5, 6],
                   [3, 4, 6, 7],
                   ones(Float64, 4),
                   dim)

    z, conv = Z_evec_dynsys(T)
    println("dynamical systems algorithm converged: $conv")
    z2 = z / norm(z, 2)
    λ1 = mean(apply(T, z2) ./ z2)
    println("tensor z-eigenvalue: $(λ1)")
    println("tensor z-eigenvector (2-norm normalized): $(z2)")

    # Check for stability
    Ustar = Array(qr((I - z2 * z2') * randn(dim, dim - 1)).Q)
    maxval = maximum(z2' * Ustar )
    println("orthogonality test (should be close to 0): $(maxval)")
    diff = Ustar' * Ustar - I
    diffnorm = [norm(vec(diff[:, ind]), 2) for ind in 1:size(diff)[2]]
    println("orthonormality test (should be close to 0): $(diffnorm)")

    C = Ustar' * ((m - 1) * reduce(T, z2) - λ1 * I) * Ustar
    U = triu(C, 1)  # for symmetrization
    C_sym = Matrix(U + U' + Diagonal(C))
    spectrum = eigen(C_sym).values
    println("Spectrum of projected Hessian of the Lagrangian: $spectrum")
end
;


"""
    t_matrix(CapsuleParticle{T,2,Acoustic{T,2},Sphere{T,2}}, Acoustic{T,2}, ω, order)

The T-matrix for a 2D circlular capsule particle in an acoustic medium.
"""
function t_matrix(cap::CapsuleParticle{T,2,Acoustic{T,2},Sphere{T,2}}, medium::Acoustic{T,2}, ω::T, M::Integer)::Diagonal{Complex{T}} where T <: AbstractFloat

    k = ω / medium.c
    k0 = ω / cap.inner.medium.c
    k1 = ω / cap.outer.medium.c
    a0 = outer_radius(cap.inner)
    a1 = outer_radius(cap.outer)
    q  = (medium.ρ*medium.c)/(cap.outer.medium.ρ*cap.outer.medium.c)
    q0 = (cap.inner.medium.ρ*cap.inner.medium.c)/(cap.outer.medium.ρ*cap.outer.medium.c)

    Yn(n::Integer)   = hankelh1(n,k1*a1)*besselj(n,k1*a0) - hankelh1(n,k1*a0)*besselj(n,k1*a1)
    Yddn(n::Integer) = diffhankelh1(n,k1*a1)*diffbesselj(n,k1*a0) - diffhankelh1(n,k1*a0)*diffbesselj(n,k1*a1)

    Ydn(n::Integer,x::Complex{T},y::Complex{T}) = hankelh1(n,x)*diffbesselj(n,y) - diffhankelh1(n,y)*besselj(n,x)

    function Tn(n::Integer)::Complex{T}
        numer = Ydn(n,k*a1,k*a1) * (q0*besselj(n,k0*a0)*Ydn(n,k1*a1,k1*a0) - Yn(n)*diffbesselj(n,k0*a0))
        denom = diffbesselj(n,k0*a0)*
            (q*hankelh1(n,k*a1)*Ydn(n,k1*a0,k1*a1) + diffhankelh1(n,k*a1)*Yn(n)) +
            q0*besselj(n,k0*a0)*
            (q*hankelh1(n,k*a1)*Yddn(n) - diffhankelh1(n,k*a1)*Ydn(n,k1*a1,k1*a0))

         return (numer - denom*besselj(n,k*a1)) / (denom*hankelh1(n,k*a1))
     end

    # Get Tns for positive m
    Tns = map(Tn,0:M)

    return Diagonal(vcat(reverse(Tns), Tns[2:end]))
end

function internal_field(x::AbstractArray{T}, p::CapsuleParticle{T,2,Acoustic{T,2},Sphere{T,2}}, source::RegularSource{T,Acoustic{T,2}}, ω::T, scattering_coefficients::AbstractVector) where T

    Nh = Int((length(scattering_coefficients) - one(T))/T(2.0)) #shorthand

    if iszero(p.outer.medium.c) || isinf(abs((p.outer.medium.c)))
        return zero(Complex{T})
    elseif norm(x - origin(p)) > outer_radius(p)
        @warn "point $x not insider particle $p. Returning zero."
        return zero(Complex{T})
    end

    k = ω / source.medium.c
    k0 = ω / p.inner.medium.c
    k1 = ω / p.outer.medium.c
    a0 = outer_radius(p.inner)
    a1 = outer_radius(p.outer)
    q0 = (p.inner.medium.ρ*p.inner.medium.c) / (p.outer.medium.ρ*p.outer.medium.c)
    q = (source.medium.ρ*source.medium.c) / (p.outer.medium.ρ*p.outer.medium.c)

    Yn(n::Integer) = hankelh1(n,k1*a1)*besselj(n,k1*a0) - hankelh1(n,k1*a0)*besselj(n,k1*a1)
    Yddn(n::Integer) = diffhankelh1(n,k1*a1)*diffbesselj(n,k1*a0) - diffhankelh1(n,k1*a0)*diffbesselj(n,k1*a1)
    Ydn(n::Integer,x::Complex{T},y::Complex{T}) = hankelh1(n,x)*diffbesselj(n,y) - diffhankelh1(n,y)*besselj(n,x)

    denom(n::Integer) = q0 * besselj(n,a0*k0) *
            (q*besselj(n,a1*k)*Yddn(n) - diffbesselj(n,a1*k)*Ydn(n,a1*k1,a0*k1)) + diffbesselj(n,a0*k0) * (q*besselj(n,k*a1)*Ydn(n,a0*k1,a1*k1) + diffbesselj(n,k*a1)*Yn(n))
    forces = scattering_coefficients .* [Ydn(n,k*a1,k*a1)/denom(n) for n = -Nh:Nh];

    if norm(x - origin(p)) <= outer_radius(p.inner)
        coefs = - q0 * forces .* Ydn.(-Nh:Nh, a0*k1, a0*k1)
        basis = regular_basis_function(p.inner, ω)

        return sum(basis(Nh,x-origin(p)) .* coefs)
    else # calculate field in the mid region
        J_coefs = forces .* [
            q0*besselj(n, a0*k0)*diffhankelh1(n, a0*k1) - hankelh1(n, a0*k1)*diffbesselj(n, a0*k0)
        for n = -Nh:Nh]

        H_coefs = forces .* [
            besselj(n, a0*k1)*diffbesselj(n,a0*k0) - q0*besselj(n, a0*k0)*diffbesselj(n,a0*k1)
        for n = -Nh:Nh]

        J_basis = regular_basis_function(p.outer, ω)
        H_basis = outgoing_basis_function(p.outer.medium, ω)
        return sum(
            J_basis(Nh,x-origin(p)) .* J_coefs .+ H_basis(Nh,x-origin(p)) .* H_coefs
        )
    end
end


# providing the necessary arithmetic functions

"""
    +(a::ℤd{T, d}, b::ℤd{T, d}) where {T <: Integer, d}

Adds the values of `a` and `b` modulo `d` and returns a new ℤd{T, d}
"""
+(a::ℤd{T, d}, b::ℤd{T, d}) where {T <: Integer, d} = ℤd{T,d}(a.value + b.value)

"""
    -(a::ℤd{T, d}, b::ℤd{T, d}) where {T <: Integer, d}

Subtracts the value of `b` from `a` modulo `d` and returns a new ℤd{T, d} (always positive).
"""
-(a::ℤd{T, d}, b::ℤd{T, d}) where {T <: Integer, d} = ℤd{T,d}(a.value - b.value)

"""
    -(x::ℤd{T,d}) where {T <: Integer, d}

Negating x in ℤd yields again x.
"""
-(x::ℤd{T,d}) where {T <: Integer, d} = x

"""
    *(a::ℤd{T, d}, b::ℤd{T, d}) where {T <: Integer, d}

Multiplies the values of `a` and `b` modulo `d` and returns a new ℤd{T, d}
"""
*(a::ℤd{T, d}, b::ℤd{T, d}) where {T <: Integer, d} = ℤd{T,d}(a.value * b.value)

"""
    *(a::Integer, b::ℤd{T, d}) where {T <: Integer, d}

Enables ℤ × ℤd → ℤd multiplication.
"""
*(a::Integer, b::ℤd{T, d}) where {T <: Integer, d} = ℤd{T,d}(a * b.value)
*(a::ℤd{T, d}, b::Integer) where {T <: Integer, d} = b * a

"""
    convert(::Type{K}, x::ℤd{T, d}) where {K <: Number, T <: Integer, d}

Converts a given ℤd{T, d} to a general number without `mod d`-logic.

This is usefull for many more complex operations and definitions of functions that live outside ℤd.
Also, some external functions are not reimplemented, therefore they can only be used if ℤd{T, d} is treated like a normal number.
If the function-result should again map to ℤd{T, d}, remember to convert the result back to ℤd{T, d}.

!!! info
    If the external function returns floating point values, it's necessary to convert them back to Integers first using something like
    ```julia
    # rounding a float x to the next integer and converting it back to an integer.
    Int(round(x; digits=0))
    ```
"""
convert(::Type{K}, x::ℤd{T, d}) where {K <: Number, T <: Integer, d} = convert(K, x.value)

"""
Promote the `ℤd{T, d}` type over `T`.
"""
promote_rule(::Type{T}, ::Type{ℤd{T, d}}) where {T <: Integer, d} = ℤd{T,d}

"""
Promote the `ℤ₂` type over `Bool`.
"""
promote_rule(::Type{Bool}, ::Type{ℤ₂}) where {d} = ℤ₂

"""
    Int(x::ℤd{T,d}) where {T <: Integer, d}

Converts a given ℤd{T, d} to a general Integer. Used mostly internal by Julia. See `convert(::Type{K}, x::ℤd{T, d})` for more details on conversion.
"""
Int(x::ℤd{T,d}) where {T <: Integer, d} = x.value

"""
    AbstractFloat(x::ℤd{T,d}) where {T <: Integer, d}

Converts a given ℤd{T, d} to a general Float. Used mostly internal by Julia. See `convert(::Type{K}, x::ℤd{T, d})` for more details on conversion.
"""
AbstractFloat(x::ℤd{T,d}) where {T <: Integer, d} = AbstractFloat(x.value)

"""
    show(io::IO, x::ℤd{T, d}) where {T <: Integer, d}

Pretty-printing for ℤd as Integers.
"""
show(io::IO, x::ℤd{T, d}) where {T <: Integer, d} = print(io, x.value)

"""
    trailing_zeros(x::ℤd{T, d}) where {T <: Integer, d}

Internal Julia function needed in some cases, like when using Sets. Passes the call down to the value of x.
"""
trailing_zeros(x::ℤd{T, d}) where {T <: Integer, d} = trailing_zeros(Int(x.value))

# There may be special cases where we want to simplify the operations, these can be appended below
# The special case handled here is for d = 2

"""
    +(x::ℤ₂, y::ℤ₂)

Special case of addition in ℤ₂ using Booleans.
"""
+(x::ℤ₂, y::ℤ₂) = ℤ₂(x.value != y.value)

"""
    -(x::ℤ₂, y::ℤ₂)

Special case of subtraction in ℤ₂ using Booleans.
"""
-(x::ℤ₂, y::ℤ₂) = ℤ₂(x.value != y.value)

"""
    +(x::ℤ₂, y::ℤ₂)

Special case of multiplication in ℤ₂ using Booleans.
"""
*(x::ℤ₂, y::ℤ₂) = ℤ₂(x.value * y.value)

"""
    show(io::IO, x::ℤ₂)

Special case of pretty-printing in ℤ₂ using Booleans.
"""
show(io::IO, x::ℤ₂) = print(io, Int8(x.value))

"""
    inv(A::Matrix{ℤd{T, d}})

Invert a matrix `A` in ℤd by multiplying the inverse in ℝ with the determinant.
"""
inv(A::Matrix{ℤd{T, d}}) where {T <: Integer, d} = Matrix{ℤd{T, d}}(Int.(round.(inv(Matrix{Int}(A)) * Int(round(det(A); digits=0)); digits=0)))


using SymbolicRegression, SymbolicUtils

X = randn(Float32, 5, 100)
y = 2 * cos.(X[4, :]) + X[1, :] .^ 2 .- 2

options = SymbolicRegression.Options(;
    binary_operators=(+, *, /, -), unary_operators=(cos, exp), npopulations=20
)

hall_of_fame = EquationSearch(X, y; niterations=40, options=options, numprocs=4)

dominating = calculate_pareto_frontier(X, y, hall_of_fame, options)

eqn = node_to_symbolic(dominating[end].tree, options)
println(simplify(eqn * 5 + 3))

println("Complexity\tMSE\tEquation")

for member in dominating
    complexity = compute_complexity(member.tree, options)
    loss = member.loss
    string = string_tree(member.tree, options)

    println("$(complexity)\t$(loss)\t$(string)")
end


using Sigma
using Lens

srand(1234)

# Paramter Estimation
λ = uniform(0,2)
x = mvexponential(λ, 3)
observations = x == [0.083, 0.55, 2.37]
rand(λ, observations)
rand(x, observations)

m = 3
n = 3
A = Sigma.mvuniform(-1,1,m,n);

# Sample matrix whose sum is 3
As = rand(A, sum(A) == 3)
sum(As)

# Sample Orthogonal Matrix
orthogonal_A = rand(A, (A'*A == eye(m,n)))
orthogonal_A * orthogonal_A'
map(round, orthogonal_A * orthogonal_A')

# trace
B = mvexponential(0.5, m, n);
C = mvlogistic(1, 1, m, n);

Bs, Cs = rand((B,C), (trace(B) == trace(C)), 1)
trace(Bs[1])
trace(Cs[1])

# Simple
A = uniform(0,1)
B = exponential(.5)
rand(A, A^2/sin(A) > 0.45 * atan(B))
rand(A, A < -1)

## Non Negative Matrix Factorisation
## ==================================
dist(A,B; epsilon = 1.0) = sum([(A[i] - B[i])^2 for i = 1:length(A)]) < epsilon

function nmf(V, W, H)
  dist(V, W*H)
end

W = mvuniform(0.0,10.0,4,2)
H = mvuniform(0.0,10.0,2,4)
V = rand(4,4)

nmf(V, W, H)


function chebyshev_evaluate(weights::Array{T,1},x::Array{T,1},order::Array{S,1},domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain)

  polynomial_1 = chebyshev_polynomial(order[1],[x1])

  evaluated_polynomial = zero(T)

  for i = 1:order[1]+1

    evaluated_polynomial += polynomial_1[i]*weights[i]

  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,2},x::Array{T,1},order::Array{S,1},domain=[ones(1,2);-ones(1,2)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])

  polynomial_1 = chebyshev_polynomial(order[1],[x1])
  polynomial_2 = chebyshev_polynomial(order[2],[x2])

  evaluated_polynomial = zero(T)

  for j = 1:order[2]+1
    for i = 1:order[1]+1

      evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*weights[i,j]

    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,3},x::Array{T,1},order::Array{S,1},domain=[ones(1,3);-ones(1,3)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])

  polynomial_1 = chebyshev_polynomial(order[1],[x1])
  polynomial_2 = chebyshev_polynomial(order[2],[x2])
  polynomial_3 = chebyshev_polynomial(order[3],[x3])

  evaluated_polynomial = zero(T)

  for k = 1:order[3]+1
    for j = 1:order[2]+1
      for i = 1:order[1]+1

        evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*weights[i,j,k]

      end
    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,4},x::Array{T,1},order::Array{S,1},domain=[ones(1,4);-ones(1,4)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])
  x4 = normalize_node(x[4],domain[:,4])

  polynomial_1 = chebyshev_polynomial(order[1],[x1])
  polynomial_2 = chebyshev_polynomial(order[2],[x2])
  polynomial_3 = chebyshev_polynomial(order[3],[x3])
  polynomial_4 = chebyshev_polynomial(order[4],[x4])

  evaluated_polynomial = zero(T)

  for l = 1:order[4]+1
    for k = 1:order[3]+1
      for j = 1:order[2]+1
        for i = 1:order[1]+1

          evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*polynomial_4[l]*weights[i,j,k,l]

  			 end
			end
    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,5},x::Array{T,1},order::Array{S,1},domain=[ones(1,5);-ones(1,5)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])
  x4 = normalize_node(x[4],domain[:,4])
  x5 = normalize_node(x[5],domain[:,5])

  polynomial_1 = chebyshev_polynomial(order[1],[x1])
  polynomial_2 = chebyshev_polynomial(order[2],[x2])
  polynomial_3 = chebyshev_polynomial(order[3],[x3])
  polynomial_4 = chebyshev_polynomial(order[4],[x4])
  polynomial_5 = chebyshev_polynomial(order[5],[x5])

  evaluated_polynomial = zero(T)

	for m = 1:order[5]+1
	  for l = 1:order[4]+1
      for k = 1:order[3]+1
        for j = 1:order[2]+1
          for i = 1:order[1]+1

            evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*polynomial_4[l]*polynomial_5[m]*weights[i,j,k,l,m]

					end
				end
      end
    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,6},x::Array{T,1},order::Array{S,1},domain=[ones(1,6);-ones(1,6)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])
  x4 = normalize_node(x[4],domain[:,4])
  x5 = normalize_node(x[5],domain[:,5])
  x6 = normalize_node(x[6],domain[:,6])

  polynomial_1 = chebyshev_polynomial(order[1],[x1])
  polynomial_2 = chebyshev_polynomial(order[2],[x2])
  polynomial_3 = chebyshev_polynomial(order[3],[x3])
  polynomial_4 = chebyshev_polynomial(order[4],[x4])
  polynomial_5 = chebyshev_polynomial(order[5],[x5])
  polynomial_6 = chebyshev_polynomial(order[6],[x6])

  evaluated_polynomial = zero(T)

	for n = 1:order[6]+1
	  for m = 1:order[5]+1
	    for l = 1:order[4]+1
        for k = 1:order[3]+1
          for j = 1:order[2]+1
            for i = 1:order[1]+1

              evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*polynomial_4[l]*polynomial_5[m]*polynomial_6[n]*weights[i,j,k,l,m,n]

            end
          end
				end
      end
    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,1},x::Array{T,1},order::S,domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain)

  polynomial_1 = chebyshev_polynomial(order,[x1])

  evaluated_polynomial = zero(T)

  for i = 1:order+1

    evaluated_polynomial += polynomial_1[i]*weights[i]

  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,2},x::Array{T,1},order::S,domain=[ones(1,2);-ones(1,2)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])

  polynomial_1 = chebyshev_polynomial(order,[x1])
  polynomial_2 = chebyshev_polynomial(order,[x2])

  evaluated_polynomial = zero(T)

  for j = 1:order+1
    for i = 1:order+1

      if (i+j <= order+2)

        evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*weights[i,j]

      end

    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,3},x::Array{T,1},order::S,domain=[ones(1,3);-ones(1,3)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])

  polynomial_1 = chebyshev_polynomial(order,[x1])
  polynomial_2 = chebyshev_polynomial(order,[x2])
  polynomial_3 = chebyshev_polynomial(order,[x3])

  evaluated_polynomial = zero(T)

  for k = 1:order+1
    for j = 1:order+1
      for i = 1:order+1

        if (i+j+k <= order+3)

          evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*weights[i,j,k]

        end

      end
    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,4},x::Array{T,1},order::S,domain=[ones(1,4);-ones(1,4)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])
  x4 = normalize_node(x[4],domain[:,4])

  polynomial_1 = chebyshev_polynomial(order,[x1])
  polynomial_2 = chebyshev_polynomial(order,[x2])
  polynomial_3 = chebyshev_polynomial(order,[x3])
  polynomial_4 = chebyshev_polynomial(order,[x4])

  evaluated_polynomial = zero(T)

  for l = 1:order+1
    for k = 1:order+1
      for j = 1:order+1
        for i = 1:order+1

          if (i+j+k+l <= order+4)

            evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*polynomial_4[l]*weights[i,j,k,l]

          end

  			end
			end
    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,5},x::Array{T,1},order::S,domain=[ones(1,5);-ones(1,5)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])
  x4 = normalize_node(x[4],domain[:,4])
  x5 = normalize_node(x[5],domain[:,5])

  polynomial_1 = chebyshev_polynomial(order,[x1])
  polynomial_2 = chebyshev_polynomial(order,[x2])
  polynomial_3 = chebyshev_polynomial(order,[x3])
  polynomial_4 = chebyshev_polynomial(order,[x4])
  polynomial_5 = chebyshev_polynomial(order,[x5])

  evaluated_polynomial = zero(T)

  for m = 1:order+1
    for l = 1:order+1
      for k = 1:order+1
	      for j = 1:order+1
          for i = 1:order+1

            if (i+j+k+l+m <= order+5)

              evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*polynomial_4[l]*polynomial_5[m]*weights[i,j,k,l,m]

            end

					end
				end
      end
    end
  end

  return evaluated_polynomial

end

function chebyshev_evaluate(weights::Array{T,6},x::Array{T,1},order::S,domain=[ones(1,6);-ones(1,6)]) where {T<:AbstractFloat,S<:Integer}

  x1 = normalize_node(x[1],domain[:,1])
  x2 = normalize_node(x[2],domain[:,2])
  x3 = normalize_node(x[3],domain[:,3])
  x4 = normalize_node(x[4],domain[:,4])
  x5 = normalize_node(x[5],domain[:,5])
  x6 = normalize_node(x[6],domain[:,6])

  polynomial_1 = chebyshev_polynomial(order,[x1])
  polynomial_2 = chebyshev_polynomial(order,[x2])
  polynomial_3 = chebyshev_polynomial(order,[x3])
  polynomial_4 = chebyshev_polynomial(order,[x4])
  polynomial_5 = chebyshev_polynomial(order,[x5])
  polynomial_6 = chebyshev_polynomial(order,[x6])

  evaluated_polynomial = zero(T)

  for n = 1:order+1
    for m = 1:order+1
      for l = 1:order+1
        for k = 1:order+1
	        for j = 1:order+1
            for i = 1:order+1

              if (i+j+k+l+m+n <= order+6)

                evaluated_polynomial += polynomial_1[i]*polynomial_2[j]*polynomial_3[k]*polynomial_4[l]*polynomial_5[m]*polynomial_6[n]*weights[i,j,k,l,m,n]

              end

            end
					end
				end
      end
    end
  end

  return evaluated_polynomial

end


include("computation/derivatives.jl")
using Plots

f = x-> x[1]*x[2] + x[2]^3 #derivative
df = x -> [x[2];  x[1]+3*x[2]^2] #Jacobian, pass in as column vector
d2f = x ->  [0 1; 1 6*x[2]] #Hessian

(h, A) = checkDerivative(f, df, [1.0;3.0], d2f, 3, 6)
polyfit(h, A, 1)
plt1 = Plots.plot(h, A, title = "Finite Difference Derivative Checker", xlabel = "log of h", ylabel = "log of error",fontfamily=font(48, "Courier") , reuse = false)
#plot(polyfit(h, A, 1), reuse = true)
println("derivative of p(z0|z)")
println(polyfit(h, A, 1))
display(plt1)


@diffusion_process SIRAux begin
    :parameters
    (α, β, σ1, σ2, t, T) --> Float64
    (u, v) --> ℝ{2}

    :additional
    linear --> true
    constdiff --> true
    statespace -->BoundedStateSpace(((1,2), (0.0,0.0)), ((1,2), (1.0,1.0)))
end

function B(t, P::SIRAux)
    @SMatrix[
        (P.α*(1 - P.v[1] - P.v[2]) - P.β)  0.0;
        P.β   0.0
    ]
end

β(t, P::SIRAux) = ℝ{2}(0.0, 0.0)

function σ(t, P::SIRAux)
    @SMatrix Float64[
        (-P.σ1*(1 - v[1] - v[2])*v[1])  -P.σ2*v[1];
        0.0   P.σ2*v[1]
    ]
end


function GeometricAverageRate(CallPutFlag, S, X, T, r, b, v)

  # b = r, give the BS stock option model
  # b = r - q, gives Merton 73, stock option with dividend
  # b = 0, give black 76
  # b = 0 && r = 0, gives Asay for margined futures options
  # b = r- r_f, gives garman kohlhagen currency options model

#require("CND.jl")
d1::Float64
d2::Float64

b = (1/2) * (b - v^2/6)
v = v/sqrt(3)

d1 = (log(S / X) + (b + v^2 / 2) * T) / (v * sqrt(T))
d2 = d1 - v * sqrt(T)

if CallPutFlag == "c"

        GBS = S * exp((b - r) * T) * CND(d1) - X * exp(-r * T) * CND(d2)

elseif CallPutFlag == "p"

        GBS = X * exp(-r * T) * CND(-d2) - S * exp((b - r) * T) * CND(-d1)

end

end


function pseudocharge(gridpos, Ls, atoms, YukawaK, epsil0)
# Create the pseudo-charge for Coulomb interaction
# $\rho_a = \sum_{N_{\text{atoms}}} - Z_j/\sqrt{2 \pi \sigma_i}$

Nsglb = length(gridpos);
Natoms = atoms.Natoms;

Z     = atoms.Z;
sigma = atoms.sigma;
R     = atoms.R;

# Calculate the total pseudo-charge
# NOTE: The pseudo-charge should not extend over twice the domain size!
rhoa = zeros(Nsglb);

    for j=1:Natoms
        d = R[j] .- gridpos;
        d = d - round.(Integer, d/Ls)*Ls;
        # Note: Z has minus sign in front of it!
        rhoa = rhoa - Z[j]./sqrt(2*pi*sigma[j]^2) .* (exp.(-0.5*(d./sigma[j]).^2 ));
    end

drhoa = zeros(Nsglb, Natoms);

    for j=1:Natoms
        d = R[j] .- gridpos;
        d = d - round.(Integer, d/Ls)*Ls;
        drhoa[:,j] = drhoa[:,j] - Z[j] ./sqrt(2*pi*sigma[j]^2)./sigma[j]^2 .* (exp.(-0.5*(d./sigma[j]).^2)).*d;
    end
    return rhoa, drhoa
end


"""
Parent Centric Crossover (PCX).

See
    Deb, K., Anand, A., and Joshi, D., "A Computationally Efficient Evolutionary Algorithm for Real-Parameter Optimization," Evolutionary Computation, vol. 10, no. 4, pp. 371-395, 2002.
"""
struct ParentCentricCrossover{NP} <: CrossoverOperator{NP,1}
    ζ::Float64 # sd for the orthogonal directions defined by 2nd,3rd etc parents
    η::Float64 # sd for the direction of the 1st parent

    function ParentCentricCrossover{NP}(ζ::Number, η::Number) where NP
        ζ > 0 || throw(ArgumentError("ζ must be positive"))
        η > 0 || throw(ArgumentError("η must be non-negative"))
        new{NP}(ζ, η)
    end
    ParentCentricCrossover{NP}(params::Parameters) where NP =
        new{NP}(params[:PCX_ζ], params[:PCX_η])
end

const PCX_DefaultOptions = ParamsDict(
    :PCX_ζ => 0.1,
    :PCX_η => 0.5
)

function apply!(xover::ParentCentricCrossover,
                target::AbstractIndividual, targetIndex::Int,
                pop, parentIndices)
    @assert length(parentIndices) == numparents(xover)

    center = mean(viewer(pop, parentIndices), dims=2)
    # project other parents vectors to the subspace
    # orthogonal to the 1st (central) parent
    parent_1 = viewer(pop, parentIndices[1]) .- dropdims(center, dims=2)
    other_parents_centered = viewer(pop, parentIndices[2:end]) .- center
    if (parent_1_sqnorm = sum(abs2, parent_1)) > 0.0
        dotprods = transpose(parent_1) * other_parents_centered
        dotprods ./= parent_1_sqnorm
        other_parents_centered .-= dotprods .* parent_1
    end
    sd = mean(sqrt.(sum(abs2, other_parents_centered, dims=2)))
    if sd > 1E-8
        svdf = svd(other_parents_centered, full=false)
        svals_norm = norm(svdf.S)
        svals = svdf.S * sd * xover.ζ / svals_norm
        svals .*= randn(length(svals))
        mul!(target, svdf.U, svals)
    else # degenerated
        fill!(target, zero(eltype(target)))
    end

    selParent = viewer(pop, parentIndices[1])
    selScale = randn() * xover.η
    @inbounds target .+= selParent .+ selScale .* parent_1
    return target
end


using Colors, Distributions, Bijectors, Zygote, StatsPlots, MCMCChains
using StatsFuns: logsumexp, logit, logistic
using Revise, AdvancedHMC

########
# Data #
########
const DATA = [RGB(0.5, 0.0, 0.0), RGB(0.6, 0.0, 0.0), RGB(0.4, 0.0, 0.0)]
const DATA_MIXTURE = vcat(DATA, [RGB(0.5, 0.0, 1.0), RGB(0.6, 0.0, 1.0), RGB(0.4, 0.0, 1.0)])

#######################
# Color distributions #
#######################

struct UniformRGB end
Distributions.rand(::UniformRGB) = RGB{Float64}(rand(3)...)
Distributions.logpdf(::UniformRGB, ::RGB{Float64}) = zero(Float64)

# struct BetaRGB
#     α
#     β
# end
# Distributions.rand(b::BetaRGB) =
#     RGB(map(rand, [Beta(b.α[i], b.β[i]) for i = 1:3])...)
# function Distributions.logpdf(b::BetaRGB, x::RGB{T}) where {T<:AbstractFloat}
#     betas = [Beta(b.α[i], b.β[i]) for i = 1:3]
#     xs = [x.r, x.g, x.b]
#     return sum(map(i -> logpdf(betas[i], xs[i]), 1:3))
# end

function Bijectors.logpdf_with_trans(d, c::RGB{Float64})
    θ = [c.r, c.g, c.b]
    return logpdf(d, c) + sum(log.(θ .* (one(Float64) .- θ)))
end

struct NormalRGB
    μ::RGB{Float64}
    σ::Float64
end
Distributions.logpdf(d::NormalRGB, x::RGB{Float64}) = logpdf(Normal(zero(Float64), d.σ), colordiff(d.μ, x)^2)

###################
# Sampling helper #
###################

function sample_colors(θ_init, logjoint, grad; stepsize=1e-3, n_samples=2_000, n_adapts=1_000)
    # Dimension
    dim = length(θ_init)

    # Define metric space, Hamiltonian and sampling method
    metric = DiagEuclideanMetric(dim)
    h = Hamiltonian(metric, logjoint, grad)
    # prop = NUTS(Leapfrog(find_good_eps(h, θ_init)), 5, 1000.0)
    prop = StaticTrajectory(Leapfrog(stepsize), 20)
    # adaptor = StanNUTSAdaptor(n_adapts, Preconditioner(metric), NesterovDualAveraging(0.8, prop.integrator.ϵ))

    # Sampling
    # samples = sample(h, prop, θ_init, n_samples, adaptor, n_adapts; progress=true)
    samples = sample(h, prop, θ_init, n_samples; progress=true)
    return samples
end

############################
# Model 1: single Gaussian #
############################
# Log-joint and it's gradient
function logπ(θ::Vector{Float64}) :: Float64
    c = RGB{Float64}(θ...)
    logprior = logpdf(UniformRGB(), c)
    loglikelihood = sum([logpdf(NormalRGB(t, 5one(Float64)), c) for t in  DATA])
    return logprior + loglikelihood
end
# @code_warntype logπ(rand(3))
∂logπ∂θ(θ::Vector{Float64}) :: Vector{Float64} = vec(collect(gradient(logπ, θ)[1]))
# @code_warntype ∂logπ∂θ(rand(6))
samples = sample_colors(rand(3), logπ, ∂logπ∂θ)
plot(
    Chains(
        reshape(hcat(samples...)', (2_000, 3, 1)),
        ["c.r", "c.g", "c.b"]
    )
)

colors = map(s -> RGB(s...), samples)
mean_samples = mean(samples)
mean_color = RGB(mean_samples...)

###################################
# Model 1: mixture of 2 Gaussians #
###################################
# Log-joint and it's gradient
logπ_mixture(θlink::Vector{Float64}) = logπ_mixture(θlink[1:3], θlink[4:6])
# @Mike: I cannot define `logπ_mixture(θ::Vector{Float64})` directly
#        and construct colors by `c1 = RGB(θ[1:3]...)`, which gives me
#        error when calling `gradient`.
function logπ_mixture(θ1link::Vector{Float64}, θ2link::Vector{Float64}) :: Float64
    p = ones(Float64) / 2
    σ = 20one(Float64)
    c1 = RGB{Float64}(logistic.(θ1link)...)
    c2 = RGB{Float64}(logistic.(θ2link)...)
    logprior = logpdf_with_trans(UniformRGB(), c1) + logpdf_with_trans(UniformRGB(), c2)
    loglikelihood = sum(
        [logsumexp([logpdf(NormalRGB(t, σ), c1) - log(2one(Float64)),
                    logpdf(NormalRGB(t, σ), c2) - log(2one(Float64))]) for t in DATA_MIXTURE]
    )
    return logprior + loglikelihood
end
# @code_warntype logπ(rand(6))
∂logπ∂θ_mixture(θlink::Vector{Float64}) :: Vector{Float64} = ∂logπ∂θ_mixture(θlink[1:3], θlink[4:6])
∂logπ∂θ_mixture(θ1link::Vector{Float64}, θ2link::Vector{Float64}) :: Vector{Float64} =
    vcat(map(collect, gradient(logπ_mixture, θ1link, θ2link))...)
# @code_warntype ∂logπ∂θ_mixture(rand(6))
sampleslink_mixture = sample_colors(randn(6), logπ_mixture, ∂logπ∂θ_mixture; stepsize=1e-2)
samples_mixture = map(s -> logistic.(s), sampleslink_mixture)
plot(
    Chains(
        reshape(hcat(samples_mixture...)', (2_000, 6, 1)),
        ["c1.r", "c1.g", "c1.b", "c2.r", "c2.g", "c2.b"]
    )
)

colors_mixture = map(s -> [RGB(s[1:3]...), RGB(s[4:6]...)], samples_mixture)
map(cs->cs[1], colors_mixture)
map(cs->cs[2], colors_mixture)
mean_samples_mixture = mean(samples_mixture)
mean_color = [RGB(mean_samples_mixture[1:3]...), RGB(mean_samples_mixture[4:6]...)]


using SpatialFields
using Base.Test
import MultiPoly

function test_linear_fit()
    x, y = MultiPoly.generators(MultiPoly.MPoly{Float64}, :x, :y)

    true_polynomial = 3 + 4x - 7y

    points = SVector{2, Float64}[[5; -2], [-1; 10], [-1; 11]]
    data = Float64[MultiPoly.evaluate(true_polynomial, p...) for p in points]

    fit_polynomial = SpatialFields.linear_fit(points, data)

    for p in points
        @test isapprox(MultiPoly.evaluate(fit_polynomial, p...), MultiPoly.evaluate(true_polynomial, p...))
    end
    @test isapprox(fit_polynomial[0,0], true_polynomial[0,0])
    @test isapprox(fit_polynomial[1,0], true_polynomial[1,0])
    @test isapprox(fit_polynomial[0,1], true_polynomial[0,1])
end

test_linear_fit()

function test_vector_linear_fit()
    x, y = MultiPoly.generators(MultiPoly.MPoly{Float64}, :x, :y)

    true_polynomials = [3 + 4x - 7y; -1 - 8x + 2y]
    points = SVector{2, Float64}[[5; -2], [-1;10], [-1; 11]]
    data = SVector{2, Float64}[[MultiPoly.evaluate(poly, p...) for poly in true_polynomials] for p in points]

    fit_polynomials = SpatialFields.linear_fit(points, data)

    for i = 1:length(true_polynomials)
        for p in points
            @test isapprox(MultiPoly.evaluate(fit_polynomials[i], p...), MultiPoly.evaluate(true_polynomials[i], p...))
        end
        @test isapprox(fit_polynomials[i][0,0], true_polynomials[i][0,0])
        @test isapprox(fit_polynomials[i][1,0], true_polynomials[i][1,0])
        @test isapprox(fit_polynomials[i][0,1], true_polynomials[i][0,1])
    end
end

test_vector_linear_fit()


using EconPDEs, Distributions

mutable struct DiTellaModel
  # Utility Function
  γ::Float64
  ψ::Float64
  ρ::Float64
  τ::Float64

  # Technology
  A::Float64
  σ::Float64

  # MoralHazard
  ϕ::Float64

  # Idiosyncratic
  νbar::Float64
  κν::Float64
  σνbar::Float64
end

function DiTellaModel(;γ = 5.0, ψ = 1.5, ρ = 0.05, τ = 0.4, A = 200.0, σ = 0.03, ϕ = 0.2, νbar = 0.24, κν = 0.22, σνbar = -0.13)
  DiTellaModel(γ, ψ, ρ, τ, A, σ, ϕ, νbar, κν, σνbar)
end

function initialize_stategrid(m::DiTellaModel; xn = 80, νn = 10)
  γ = m.γ ; ψ = m.ψ ; ρ = m.ρ ; τ = m.τ ; A = m.A ; σ = m.σ ; ϕ = m.ϕ ; νbar = m.νbar ; κν = m.κν ; σνbar = m.σνbar
  distribution = Gamma(2 * κν * νbar / σνbar^2, σνbar^2 / (2 * κν))
  νmin = quantile(distribution, 0.001)
  νmax = quantile(distribution, 0.999)
  OrderedDict(:x => range(0.01, stop = 0.99, length = xn), :ν => range(νmin, stop = νmax, length = νn))
end

function initialize_y(m::DiTellaModel, stategrid::OrderedDict)
  x = fill(1.0, length(stategrid[:x]), length(stategrid[:ν]))
  OrderedDict(:pA => x, :pB => x, :p => x)
end

function (m::DiTellaModel)(state::NamedTuple, y::NamedTuple)
  γ = m.γ ; ψ = m.ψ ; ρ = m.ρ ; τ = m.τ ; A = m.A ; σ = m.σ ; ϕ = m.ϕ ; νbar = m.νbar ; κν = m.κν ; σνbar = m.σνbar
  x, ν = state.x, state.ν
  pA, pAx, pAν, pAxx, pAxν, pAνν, pB, pBx, pBν, pBxx, pBxν, pBνν, p, px, pν, pxx, pxν, pνν = y.pA, y.pAx, y.pAν, y.pAxx, y.pAxν, y.pAνν, y.pB, y.pBx, y.pBν, y.pBxx, y.pBxν, y.pBνν, y.p, y.px, y.pν, y.pxx, y.pxν, y.pνν

  # drift and volatility of state variable ν
  g = p / (2 * A)
  i = A * g^2
  μν = κν * (νbar - ν)
  σν = σνbar * sqrt(ν)

  # Market price of risk κ
  σX = x * (1 - x) * (1 - γ) / (γ * (ψ - 1)) * (pAν / pA - pBν / pB) * σν / (1 - x * (1 - x) * (1 - γ) / (γ * (ψ - 1)) * (pAx / pA - pBx / pB))
  σpA = pAx / pA * σX + pAν / pA * σν
  σpB = pBx / pB * σX + pBν / pB * σν
  σp = px / p * σX + pν / p * σν
  κ = (σp + σ - (1 - γ) / (γ * (ψ - 1)) * (x * σpA + (1 - x) * σpB)) / (1 / γ)
  κν = γ * ϕ * ν / x
  σA = κ / γ + (1 - γ) / (γ * (ψ - 1)) * σpA
  νA = κν / γ
  σB = κ / γ + (1 - γ) / (γ * (ψ - 1)) * σpB

  # Interest rate r
  μX = x * (1 - x) * ((σA * κ + νA * κν - 1 / pA - τ) - (σB * κ -  1 / pB + τ * x / (1 - x)) - (σA - σB) * (σ + σp))
  μpA = pAx / pA * μX + pAν / pA * μν + 0.5 * pAxx / pA * σX^2 + 0.5 * pAνν / pA * σν^2 + pAxν / pA * σX * σν
  μpB = pBx / pB * μX + pBν / pB * μν + 0.5 * pBxx / pB * σX^2 + 0.5 * pBνν / pB * σν^2 + pBxν / pB * σX * σν
  μp = px / p * μX + pν / p * μν + 0.5 * pxx / p * σX^2 + 0.5 * pνν / p * σν^2 + pxν / p * σX * σν
  r = (1 - i) / p + g + μp + σ * σp - κ * (σ + σp) - γ / x * (ϕ * ν)^2

  # Market Pricing
  pAt = pA * (1 / pA  + (ψ - 1) * τ / (1 - γ) * ((pA / pB)^((1 - γ) / (1 - ψ)) - 1) - ψ * ρ + (ψ - 1) * (r + κ * σA + κν * νA) + μpA - (ψ - 1) * γ / 2 * (σA^2 + νA^2) + (2 - ψ - γ) / (2 * (ψ - 1)) * σpA^2 + (1 - γ) * σpA * σA)
  pBt = pB * (1 / pB - ψ * ρ + (ψ - 1) * (r + κ * σB) + μpB - (ψ - 1) * γ / 2 * σB^2 + (2 - ψ - γ) / (2 * (ψ - 1)) * σpB^2 + (1 - γ) * σpB * σB)
  # algebraic constraint
  pt = p * ((1 - i) / p - x / pA - (1 - x) / pB)

  return (pAt, pBt, pt), (μX, μν), (μX = μX, μν = μν, p = p, pA = pA, pB = pB, κ = κ, r = r, σX = σX)
end

m = DiTellaModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
y, result, distance = pdesolve(m, stategrid, y0)
y, result, distance = pdesolve(m, stategrid, y0; is_algebraic = OrderedDict(:pA => false, :pB => false, :p => true))


function earthquake_exercise_setup()
    # explosion detector (using spiral coordinate system)
    # define the coordinate system:
    S=2000 # number of points on the spiral
    rate=25 # angular rate of spiral
    sd=0.2 # standard deviation of the sensor Gaussian noise

    x=zeros(S); y=zeros(S)
    for s=1:S
        theta=rate*2*pi*s/S;  r=s/S
        x[s]=r*cos(theta); y[s]=r*sin(theta)
    end

    # define the locations of the detection stations on the surface
    # Also define what value on each sensor would be generated by an explostion at internal location s
    N=30 # number of stations
    x_sensor=zeros(N); y_sensor=zeros(N)
    v=zeros(S,N)
    for sensor=1:N
        theta_sensor=2*pi*sensor/N
        x_sensor[sensor]=cos(theta_sensor); y_sensor[sensor]=sin(theta_sensor)
        for s=1:S
            v[s,sensor]=value(x[s],y[s],x_sensor[sensor],y_sensor[sensor]) # explosion value for some value function
        end
    end
end

Earthqauke - Assignment 1/.ipynb_checkpoints/earthquakeExerciseSetup-checkpoint.jl


#Importing Neccessary Packages
using DifferentialEquations,ModelingToolkit #Solving ODEs
using Turing,Distributions #Running statistics on MCMC
using CSV, DataFrames #Handle the data
using StatsPlots,ProgressMeter #Plotting and Monitoring

#Not sure why but include assumes you're in same folder as file(?)
include("../HelperFunctions.jl")
include("../MCMCRun.jl")

if pwd()!="C:\\Users\\Portable\\Documents\\GitHub\\TuringODE\\ViralSensingModel"
  cd("./ViralSensingModel")
end

#General workflow for MCMC parameter fitting

#=
1. Set up the ODE
2. Import and process the data
3. Transform the ODE output to match the data
4. Run the MCMC algorithm
5. Postprocessing
=#

###############################################################
                    # 1. Set up the ODE
###############################################################
#Some Model Information needed for MCMC (parameters and States)
parNames = [:k12, :k13,:k14,
            :k21, :tau2,
            :k31,:k32, :k33,
            :k41,:k42,
            :k51,
            :k61,
            :k71, :k72]

varNames = ["IFN","IFNe","STATP","IRF7","IRF7P","Infected","Virus"]
#Set the number of parameters and states
parNum = length(parNames)
varNum = length(varNames)
modelInfo = ModelInformation(parNames,varNames)

## Define any constants for the model
const k11 = 0.0
const n=3
const tau3 = 1.238 #Cambridge 2011
const tau4 = 0.347 #Sharova 2009
const tau5 = 6.93 #Prakash 2006
#Can be defined here or just as a number in the equations

function Model!(dy,y,par,t)
  #IFN, ODE 1 parameters
  #k11=0 #PR8, RIGI is assumed antagonized
  k12=par[1]
  k13=par[2]
  k14=par[3]
  #IFNe, ODE 2 parameters
  k21=par[4]
  tau2=par[5]
  #STATP, ODE 3 parameters
  k31=par[6]
  k32=par[7]
  k33=par[8]
  #IRF7, ODE 4 parameters
  k41=par[9]
  k42=par[10]
  #IRF7P, ODE 5 parameters
  k51=par[11]
  #Infected cells, ODE 6 parameters
  k61=par[12]
  #Viral replication, ODE 7 paramters
  k71=par[13]
  k72=par[14]

  #ODE System
  v=y[7]/(10^4) #scale virus effects
  dy[1]=k11*v+(k12*(v^n))/(k13+(v^n))+k14*y[5]-k21*y[1] #IFN in cells
  dy[2]=k21*y[1]-y[2]*tau2 #IFN in environment
  dy[3]=(k31*y[2])/(k32+k33*y[2])-y[3]*tau3 #STATP
  dy[4]=k41*y[3]+k42*y[5]-y[4]*tau4 #IRF7
  dy[5]=k51*y[4]-y[5]*tau5 #IRF7P
  dy[6]=-k61*y[6] #Productive infected cells
  dy[7]=(k71*y[6]*(v^n))/(y[2]*(v^n))-k72*y[7] #Virus count

end

#Define information for ODE model
p=rand(parNum) #Parameter values
u0 = [7.94, 0.01, 12.2, 14.15, 0.01, 1, 7.5E-2] #Initial Conditions
tspan = (0.25,24.0) #Time (start, end)

#Contruct the ODE Problem
prob = ODEProblem(Model!,u0,tspan,p)
alg = Rodas5()  #ODE solver

#Testing the ODEs
sol = solve(prob,alg)
plot(sol,layout=varNum,legend=false, framestyle=:box,title=[v for i=1:1, v in varNames])
xlabel!("")

###############################################################
                    # 2. Import the data
###############################################################
#Read in the CSV
data = CSV.read("PR8.csv", missingstring= "-")

#What if we normalize the virus?
#vMin = 7.5E-2 #Initial Condition
#vMax = maximum(skipmissing(data.Virus))
#virusNorm(x) = @. (x - vMin)/(vMax - vMin)
#data.Virus = virusNorm(data.Virus)

control = CSV.read("Control.csv", missingstring= "-")

#Convert the data into a workable form
dataTransform = ConvertData(data)

###############################################################
        # 3. Transform the ODE output to match the data
###############################################################

#What states are being measured?
const allMeasuredStates = [1,3,4,7] #IFN, STATp, IRF7, Virus

#This function takes in the solved ODE, and outputs the observable species
function ObserveTransform(sol,measuredTime,currentStates)
  #Create a container to hold the solution at the desired time points
  #Get indices of currently measured states
  measuredIdx = allMeasuredStates[currentStates]
  obversedSpecies = sol(measuredTime)[measuredIdx]

  #Create a vector of control data
  timeFilter = x->x==measuredTime

if 7 ∉ measuredIdx #if not virus
  #Get the control values and calculate a lfc
  controlVec = convert(Vector,control[findfirst(timeFilter,control.Time),3:end])
  obversedSpecies = @. log2(maximum(obversedSpecies,0.0)+1.0) - log2(controlVec)
elseif 7 ∈ measuredIdx && length(measuredIdx)==1 #Only Virus measured
  obversedSpecies = obversedSpecies
else #Both virus and species are measured
  #Get the control values and calculate a lfc, but skip the virus
  controlVec = convert(Vector,control[findfirst(timeFilter,control.Time),3:end])
  obversedSpecies[1:end-1] = @. log2(maximum(obversedSpecies[1:end-1],0.0)+1.0) - log2(controlVec)
end
  return obversedSpecies
end

###############################################################
                    # 4. Run the MCMC algorithm
###############################################################

#Provide any prior knowledge for the parameters
#priors = fill(FlatPos(0.0),length(prob.p))
priors = fill(Uniform(0.0,100.0),length(prob.p))
priors[5]=Uniform(0,1) #tau2, IFNe degradation
priors[11]=Uniform(0,1) #k61, Infected cell clearance
priors[13]=Uniform(0,1) #k72, viral clearance
#How many MCMC sample do you want
mcmcSamples = 1000

#Gather all the information to one structure
sampleProblem = MCMCSetup(modelInfo,prob,alg,dataTransform,mcmcSamples,priors)

#Run the MCMC
result = MCMCRun(sampleProblem)

#Save the parameter chains and information about the MCMC run
chainParameters = DataFrame(result,:parameters)
CSV.write("./Parameters.csv",chainParameters)

chainInternals = DataFrame(result,:internals)
CSV.write("./Internals.csv",chainInternals)
###############################################################
                    # 5. Postprocessing
###############################################################

#Make a plot of the chains and density plots
chainsPlot = plot(result)
savefig(chainsPlot,"./Figures/ChainsODE.pdf")

#Corner plot for correlations
#corPlot = autocorplot(result)
#savefig(corPlot,"./Figures/AutoCorr.pdf")

#Running average Plot
runAvePlot = meanplot(result)
savefig(runAvePlot,"./Figures/RunAve.pdf")

#Running average Plot
#cornerPlot = corner(result)
#savefig(cornerPlot,"./Figures/Corner.pdf")

#Create a plot fitting "best" parameter values
@df data plot(:Time,[:IFN,:STATP,:IRF7,:Virus],markershape=:auto,
              layout=length(allMeasuredStates),legend=false)

#Retrieve the chains and information about the chains
#Get best parameter set
bestPar = BestParSet(result,parNames) #ignore the std

#Rerun the problem with these parameters
#newProb = remake(prob, p=bestPar)
newProb = ODEProblem(Model!,u0,[0.25,18],bestPar)
newSol = solve(newProb,alg)

plotBestSol = zeros(length(newSol.t),varNum)
controlLFC = log2.(mean(convert(Matrix,control[:,3:end]),dims=1))

for t=1:length(newSol.t)
  for s = 1:varNum
    if s ∈ allMeasuredStates[1:end-1]
      idx = findfirst(x->x==s,allMeasuredStates)
      #@show s,t,idx
      plotBestSol[t,s] = log2(max(newSol[s,t],0.0)+1.0) - controlLFC[idx]
    else
      plotBestSol[t,s] = newSol[s,t]
    end
  end
end

plot!(newSol.t,plotBestSol[:,allMeasuredStates],layout=length(allMeasuredStates),legend=false,title=[v for i=1:1, v in varNames[allMeasuredStates]])
plot!(newSol,vars=allMeasuredStates,layout=5)
savefig("./Figures/DataFit.pdf")

plot(newSol,layout=varNum,legend=false, framestyle=:box,title=[v for i=1:1, v in varNames])
xlabel!("")
savefig("./Figures/AllODEs.pdf")


# TODO: Add ci()
# TODO: Add one-sided hypothesis tests

immutable BinomialTest <: HypothesisTest
    test_name::ASCIIString
    # Number of Successes
    statistic::Int
    # Test parameter: p
    parameter::Float64
    p_value::Float64
    confidence_interval::Vector{Float64}
    estimate::Float64
    null_value::Float64
    alternative_hypothesis::ASCIIString
    method::ASCIIString
end

function p_value(::Type{BinomialTest},
                 x::Integer,
                 n::Integer,
                 p::Real)
    d = Binomial(n, p)
    expectation = mean(d)

    if x == expectation
        return 1.0
    end

    if (x != 0 && p == 0.0) || (x != n && p == 1.0)
        return 0.0
    end

    mirror = expectation + (expectation - x) - 1.0

    pv = pmf(d, x)

    if x < expectation
        # TODO: Perform a binary search here
        if pmf(d, ifloor(mirror)) > pv
            lower_bound = ifloor(mirror)
        else
            lower_bound = iceil(expectation)
        end
        for i in lower_bound:n
            if pmf(d, i) <= pv
                return cdf(d, x) + ccdf(d, i - 1)
            end
        end
        return cdf(d, x)
    else
        # TODO: Perform a binary search here
        if pmf(d, iceil(mirror)) > pv
            upper_bound = iceil(mirror)
        else
            upper_bound = ifloor(expectation)
        end
        for i in upper_bound:-1:0
            if pmf(d, i) <= pv
                return cdf(d, i) + ccdf(d, x - 1)
            end
        end
        return ccdf(d, x - 1)
    end
end

function p_value(::Type{BinomialTest},
                 s::Real,
                 n::Real,
                 p::Real)
    p_value(BinomialTest, int(s), int(n), p)
end

function p_value{T <: Real}(::Type{BinomialTest},
                            x::Vector{T},
                            p::Real)
    p_value(BinomialTest, int(sum(x)), length(x), p)
end

function Base.run(::Type{BinomialTest},
                  s::Int,
                  n::Int;
                  p::Float64 = 0.5)
    return BinomialTest("Binomial Test",
                        s,
                        p,
                        p_value(BinomialTest, s, n, p),
                        [NaN, NaN],
                        s / n,
                        p,
                        "Two-sided",
                        "Exact binomial test")
end

function Base.run(::Type{BinomialTest},
                  s::Real,
                  n::Real;
                  p::Real = 0.5)
    Base.run(BinomialTest, int(s), int(n), p = float64(p))
end

function Base.run{T <: Real}(::Type{BinomialTest},
                             x::Vector{T};
                             p::Real = 0.5)
    Base.run(BinomialTest, int(sum(x)), length(x), p = float64(p))
end


@testset "Continous Weight" begin

    T = Float64

    ## known set of modified moments
    P = ShiftedLegendre{T}
    f(x) = -log(x)
    q0 = ContinuousWeight(P,f,[1], :x)
    WW = typeof(q0)
    function SpecialPolynomials.modified_moment(w::W, j) where {W <: WW}
        iszero(j) && return 1.0
        (-1)^j * gamma(j+1)^2/(j*(j+1)*gamma(2j+1))
    end

    q0 = ContinuousWeight(P, f, [1], :x)
    q1 = ContinuousWeight(P, f, [0,1], :x)
    q2 = ContinuousWeight(P, f, [0,0,1], :x)
    q3 = ContinuousWeight(P, f, [0,0,0,1], :x)
    q4 = ContinuousWeight(P, f, [0,0,0,0,1], :x)
    q5 = ContinuousWeight(P, f, [0,0,0,0,0,1], :x)
    qs = convert.(Polynomial, (q0,q1,q2,q3,q4,q5))
    a, b = first(domain(q0)), last(domain(q0))
    for i in eachindex(qs)
        for j in eachindex(qs)
            j <= i  && continue
            @test abs(SpecialPolynomials._quadgk(x -> qs[i](x) * qs[j](x) * f(x),  first(domain(q0)), last(domain(q0)))) <= 100*sqrt(eps(T))
        end
    end

    ## non-known set
    T = Float64
    P = Chebyshev{T}
    f(x) = SpecialPolynomials.weight_function(ChebyshevU{T})(x)
    q0 = ContinuousWeight(P, f, [1], :x)
    q1 = ContinuousWeight(P, f, [0,1], :x)
    q2 = ContinuousWeight(P, f, [0,0,1], :x)
    q3 = ContinuousWeight(P, f, [0,0,0,1], :x)
    q4 = ContinuousWeight(P, f, [0,0,0,0,1], :x)
    q5 = ContinuousWeight(P, f, [0,0,0,0,0,1], :x)
    qs = convert.(Polynomial, (q0,q1,q2,q3,q4,q5))
    a, b = first(domain(q0)), last(domain(q0))
    for qi in qs
        for qj in qs
            j <= i && continue
            @test abs(SpecialPolynomials._quadgk(x -> qi(x) * qj(x) * f(x), a, b)) <= 100*sqrt(eps(T))
        end
    end

end


"""
    LogBayesQuad(k::Kernel; n_candidates=100, l=1.0, σ::Real=1.0)

Tool for running Bayesian Quadrature by assuming a GP on the log integrand instead
of the integrand.

## Arguments
- `n_candidates` : will be sampled around the first samples to perform a linear approximation of the result.
- `l` : Lengthscale of the kernel
- `σ` : variance of the kernel

## Reference
- **Active Learning of Model Evidence Using Bayesian Quadrature** - Osborne et al. - NIPS 2012.
"""
struct LogBayesQuad{TK,Tl,Tσ} <: AbstractBQ{TK}
    kernel::TK
    n_candidates::Int
    l::Tl
    σ::Tσ
end

function LogBayesQuad(k::Kernel; n_candidates=100, l=1.0, σ::Real=1.0)
    k, (l, σ) = get_kernel_params(k; l, σ)
    check_kernel_parameters(l, σ)
    return LogBayesQuad(k, n_candidates, l, σ)
end

scale(bquad::LogBayesQuad) = bquad.σ

function quadrature(
    bquad::LogBayesQuad{<:SqExponentialKernel}, model::AbstractBQModel{<:MvNormal}, samples
)
    isempty(samples) && error("The collection of samples is empty")
    nsamples = length(samples)
    logf = logintegrand(model).(samples) # Evaluate integrand on samples
    normalize = false
    if normalize
        max_logf = maximum(logf) # Find maximum of all given samples
        logf .-= max_logf # Normalize log integrand
    end
    f = exp.(logf) # Compute integrand on samples

    x_c = sample_candidates(bquad, samples, bquad.n_candidates) # Sample candidates around the samples
    ke = KernelEmbedding(bquad, p_0(model))

    gp = create_gp(bquad, samples)
    f_c_0 = mean.(predict(gp, f, x_c)) # Predict integrand on x_c
    logf_c_0 = mean.(predict(gp, logf, x_c)) # Predict log-integrand on x_c
    Δ_c = exp.(logf_c_0) - f_c_0 # Compute difference of predictions

    z = kernel_mean(ke, samples) # Compute mean for the basic BQ
    K = kernelpdmat(kernel(bquad), samples) # and the kernel matrix

    z_c = kernel_mean(ke, x_c) # Compute mean for the ΔlogBQ
    K_c = kernelpdmat(kernel(bquad), x_c) # and the kernel matrix for the candidates

    m_evidence = evaluate_mean(z, K, f) # Compute m(Z|samples)
    m_correction = evaluate_mean(z_c, K_c, Δ_c) #

    C = kernel_variance(ke) # Compute the kernel variance

    var_evidence = evaluate_var(z, K, C)
    var_correction = evaluate_var(z_c, K_c, C)

    if normalize
        m = exp(log(m_evidence + m_correction) + max_logf)
        v = exp(log(var_evidence + var_correction) + 2 * max_logf)
    else
        m = m_evidence + m_correction
        v = var_evidence + var_correction
    end
    return Normal(m, sqrt(v))
end

function predict(gp::AbstractGPs.FiniteGP, f, x_c)
    f_post = posterior(gp, f) # Create a posterior given the observations f
    return mean(f_post(x_c)) # Return the predictive mean on the candidates
end

function create_gp(bquad::LogBayesQuad, samples)
    gp = GP(kernel(bquad)) # Create a GP prior
    return gp(samples) # Project it on the samples
end

"""
    sample_candidates(bquad::LogBayesQuad, samples, n_c)

Sample new candidates around the existing `samples` within an hyper-ellipse, i.e
in the space where |x_i|^n + |y_i|^n = 1.
The number of candidates is given by `bquad.n_candidates`
This is needed to estimate the correction term for using a log transformation.
See Osborne et al. 2012 and [`LogBayesQuad`](@ref)
"""
function sample_candidates(bquad::LogBayesQuad, samples, n_c)
    n_samples = length(samples)
    n_dim = length(first(samples))
    T = eltype(first(samples))
    directions = [project_on_ellipse(rand(T, n_dim) .- 0.5) for _ in 1:n_c] # Sample withing the hyper ellipse
    directions = scale_directions(bquad.l, directions) # Scale the hyperellipse
    p_indices = rand(1:n_samples, n_c) # Get n_c indices from the samples
    return directions .+= samples[p_indices] # Create candidates around the existing samples
end

function project_on_ellipse(x)
    return x / norm(x, 1)
end

scale_directions(l::Real, directions) = l .* directions
scale_directions(l::AbstractVector, directions) = hadamard(Ref(l), directions)
scale_directions(l::AbstractMatrix, directions) = Ref(l) .* directions

hadamard(x, y) = x .* y


# Test using the Random data
@testset "Random" begin

    @testset "Negative Moran" begin
        geom = reggeomlattice(10, 10)
        W = polyneigh(geom)
        X = rand(StableRNG(1234567), 100)
        Wmat = Matrix(W)
        rho = -0.9

        Y = X + rho * Wmat * X
        mnegative = moran(Y, W, rng = StableRNG(1234567))

        @test score(mnegative) ≈ -0.197453 atol = 1e-5
        @test mean(mnegative) ≈ -0.010786 atol = 1e-5
        @test std(mnegative) ≈ 0.053380 atol = 1e-5
        @test zscore(mnegative) ≈ -3.496951 atol = 1e-5
        @test pvalue(mnegative) ≈ 0.0001 atol = 1e-5
    end

end


@testset "PathTracker" begin
    @testset "Constructor + Options" begin
        @polyvar x y z
        f = [x^2 - 2, x + y - 1]
        tracker, starts = pathtracker_startsolutions(f; system = FPSystem)
        @test tracker isa PathTracker{Vector{ComplexF64}}
        test_show_juno(tracker)
        @test !isempty(sprint(show, tracker))
        test_show_juno(tracker.options)
        @test !isempty(sprint(show, tracker.options))
        test_show_juno(tracker.state)
        @test !isempty(sprint(show, tracker.state))
        @test Base.broadcastable(tracker) isa Ref{typeof(tracker)}
        tracker = pathtracker(f; system = FPSystem, precision_strategy = :adaptive_never)
        @test tracker.options.precision_strategy == HC.PREC_STRATEGY_NEVER
        tracker = pathtracker(f; system = FPSystem, precision_strategy = :adaptive_always)
        @test tracker.options.precision_strategy == HC.PREC_STRATEGY_ALWAYS
        @test_throws ArgumentError pathtracker(
            f;
            system = FPSystem,
            precision_strategy = :NOT_AN_OPTION,
        )

        @test_throws ArgumentError pathtracker(
            f;
            not_an_option = true,
            this_option = "nope",
        )

        @test_throws ArgumentError pathtracker(f, patch = OrthogonalPatch())
        @test_throws ArgumentError pathtracker(f, system_scaling = :lalala)
    end

    # Simple
    @testset "Simple" begin
        @polyvar x y z
        f = [x^2 - 2, x + y - 1]
        tracker, starts = pathtracker_startsolutions(f; system = FPSystem)

        for xᵢ in starts
            @test is_success(track!(tracker, xᵢ))
            s = solution(tracker)
            @test abs(s[1]) ≈ sqrt(2) atol = 1e-8
            @test sum(s) ≈ 1 atol = 1e-8
            @test isnothing(winding_number(tracker))
            @test !isnan(tracker.state.solution_cond)
            @test !isnan(tracker.state.solution_accuracy)
            @test tracker.state.solution_accuracy < 1e-14
            @test tracker.state.solution_residual < 1e-14
        end

        # test iterator
        init!(tracker, first(starts))
        for _ in tracker
        end
        @test is_success(status(tracker))

        tracker, starts =
            pathtracker_startsolutions(f; projective_tracking = true, system = FPSystem)
        @test tracker isa PathTracker{PVector{ComplexF64,1}}
        for x in starts
            @test is_success(track!(tracker, x))
            s = solution(tracker)
            @test s isa Vector{ComplexF64}
            @test abs(s[1]) ≈ sqrt(2) atol = 1e-8
            @test sum(s) ≈ 1 atol = 1e-8
            @test isnothing(winding_number(tracker))
            @test !isnan(tracker.state.solution_cond)
            @test !isnan(tracker.state.solution_accuracy)
            @test !isnan(tracker.state.solution_residual)
        end

        tracker, starts = pathtracker_startsolutions(homogenize(f, z); system = FPSystem)
        @test tracker isa PathTracker{PVector{ComplexF64,1}}
        for x in starts
            @test is_success(track!(tracker, x))
            s = solution(tracker)
            @test s isa PVector{ComplexF64,1}
        end

        g = equations(katsura(5))
        tracker, starts = pathtracker_startsolutions(g; system = FPSystem)
        @test all(s -> is_success(track!(tracker, s)), starts)

        # start solution is also target solution
        @polyvar x
        tracker = pathtracker([x - 1]; system = FPSystem)
        @test is_success(track(tracker, [1]))
    end

    # Univariate singular
    @testset "Univariate singular" begin
        @polyvar x
        print("Univariate singular: ")
        for n = 2:12
            print(n, ",")
            f = [(x - 3)^n]
            g = [x^n - 1]
            S = [[cis(i * 2π / n)] for i = 0:(n-1)]

            tracker = pathtracker(g, f, S; seed = 842121, system = FPSystem)
            @test all(S) do s
                is_success(track!(tracker, s)) && (
                    winding_number(tracker) == n || isnothing(winding_number(tracker))
                ) && isapprox(solution(tracker)[1], 3.0; atol = 1e-6) &&
                     !isnan(tracker.state.solution_cond) &&
                     tracker.state.solution_accuracy < 1e-5
            end
        end
        println("")
    end

    # Wilkinson
    @testset "Wilkinson" begin
        @polyvar x
        print("Wilkinson: ")
        for n = 1:18
            print(n, ",")
            g = [x^n - 1]
            f = [prod(x - i for i = 1:n)]
            S = [[cis(i * 2π / n)] for i = 0:(n-1)]
            tracker = pathtracker(g, f, S; seed = 512321, system = FPSystem)
            @test all(s -> is_success(track!(tracker, s)), S)
            # check that we actually get the correct results
            @test sort(round.(
                Int,
                real.(first.(solution.(track.(tracker, S)))),
            )) == collect(1:n)
        end
        println("")
    end

    # Bivariate singular
    @testset "Bivariate singular" begin
        f = equations(griewank_osborne())
        tracker, starts = pathtracker_startsolutions(f; seed = 78373, system = FPSystem)
        @test all(starts) do s
            retcode = track!(tracker, s)
            (is_success(retcode) && winding_number(tracker) == 3) || is_at_infinity(retcode)
        end

        # Bivariate projetive singular
        @polyvar x z y
        # This has two roots of multiplicity 6 at the hyperplane z=0.
        # But the winding numbers are only 3 at each singularity
        F = [
            0.75 * x^4 + 1.5 * x^2 * y^2 - 2.5 * x^2 * z^2 + 0.75 * y^4 -
                2.5 * y^2 * z^2 + 0.75 * z^4,
            10 * x^2 * z + 10 * y^2 * z - 6 * z^3,
        ]
        @test pathtracker(F) isa PathTracker{PVector{ComplexF64,1}}
        tracker, starts = pathtracker_startsolutions(F; seed = 29831, system = FPSystem)
        @test all(s -> is_success(track!(tracker, s)), starts)
        @test all(starts) do s
            is_success(track!(tracker, s)) &&
            winding_number(tracker) == 3 && tracker.state.solution_cond > 1e10
        end
        # affine
        F̂ = subs.(F, Ref(y => 1))
        tracker, starts = pathtracker_startsolutions(F̂; seed = 29831, system = FPSystem)
        @test all(starts) do s
            is_success(track!(tracker, s)) &&
            winding_number(tracker) == 3 && tracker.state.solution_cond > 1e10
        end

        # Fix z == 1 --> all solutions at infinity
        G = subs.(F, Ref(z => 1))
        tracker, starts = pathtracker_startsolutions(G; seed = 29831, system = FPSystem)
        @test all(starts) do s
            is_at_infinity(track!(tracker, s))
        end
    end

    @testset "Non-Singular bad-conditioned" begin
        f = equations(PolynomialTestSystems.bacillus_subtilis())
        tracker, starts = pathtracker_startsolutions(f; seed = 78373, system = FPSystem)
        @test count(s -> is_success(track!(tracker, s)), starts) == 44
    end

    @testset "Diverging path without ill-conditioning" begin
        f = equations(cyclic(7))
        tracker, starts = pathtracker_startsolutions(f; seed = 78373, system = FPSystem)
        S = collect(starts)
        @test is_at_infinity(track!(tracker, S[7]))
        @test is_at_infinity(track!(tracker, S[9]))
        @test is_at_infinity(track!(tracker, S[5016]))
    end

    @testset "Cyclic 7 correct root count" begin
        f = equations(cyclic(7))
        tracker, starts = pathtracker_startsolutions(f; seed = 78373, system = FPSystem)
        results = map(s -> track!(tracker, s), starts)
        @test count(is_success, results) == 924

        successfull_paths =
            [1 4 6 10 13 14 16 18 43 44 46 59 63 66 69 70 73 78 85 93 94 98 104 106 111 119 121 122 134 141 148 152 155 157 160 163 171 173 182 185 189 200 204 206 208 211 215 218 225 227 230 231 236 244 245 246 256 258 259 261 264 272 277 280 284 287 298 305 318 319 326 336 340 342 347 358 367 380 384 397 400 405 410 411 416 420 423 424 430 431 439 442 447 470 476 483 486 497 510 511 550 552 558 566 569 575 583 585 587 589 591 599 601 602 605 610 615 622 627 630 634 638 639 645 650 651 653 656 658 669 673 679 681 683 689 691 693 696 707 709 715 719 722 723 724 733 746 755 757 765 767 770 773 786 796 805 807 808 811 816 818 832 837 840 846 850 857 858 864 873 874 875 880 882 900 905 909 910 911 918 919 920 923 925 929 935 936 942 943 947 955 960 972 979 982 990 1008 1010 1022 1031 1035 1037 1047 1059 1060 1062 1070 1075 1076 1082 1092 1096 1100 1102 1103 1109 1111 1115 1121 1124 1129 1147 1148 1152 1166 1168 1178 1189 1197 1198 1200 1201 1215 1216 1217 1223 1229 1235 1237 1242 1247 1251 1258 1261 1262 1272 1273 1277 1282 1293 1313 1318 1333 1336 1352 1362 1376 1377 1406 1408 1409 1419 1421 1422 1423 1424 1425 1442 1444 1446 1453 1454 1460 1462 1465 1473 1490 1503 1509 1520 1522 1523 1526 1527 1530 1531 1534 1549 1552 1558 1559 1561 1565 1579 1585 1594 1600 1621 1629 1631 1637 1640 1642 1644 1645 1649 1655 1656 1662 1665 1669 1671 1679 1705 1706 1712 1731 1734 1736 1737 1749 1750 1751 1753 1757 1759 1769 1784 1790 1795 1799 1800 1801 1802 1803 1807 1809 1817 1819 1825 1832 1833 1842 1847 1848 1852 1853 1856 1861 1865 1867 1871 1876 1885 1890 1896 1900 1903 1914 1925 1926 1927 1930 1933 1937 1947 1951 1954 1956 1962 1968 1971 1974 1980 1982 1988 1997 2000 2001 2014 2015 2017 2020 2030 2033 2045 2047 2054 2058 2062 2064 2066 2070 2080 2085 2099 2102 2103 2104 2111 2113 2122 2133 2136 2137 2139 2140 2145 2148 2149 2156 2159 2161 2163 2166 2171 2174 2182 2185 2190 2193 2195 2201 2202 2224 2225 2226 2229 2230 2233 2243 2260 2262 2267 2268 2272 2273 2292 2293 2295 2300 2301 2302 2306 2310 2314 2321 2326 2334 2338 2345 2356 2357 2358 2362 2363 2368 2369 2388 2391 2400 2405 2407 2416 2426 2430 2434 2435 2451 2454 2456 2463 2465 2473 2476 2477 2493 2495 2506 2509 2521 2527 2528 2551 2560 2562 2564 2566 2572 2574 2576 2603 2608 2613 2615 2624 2633 2635 2645 2650 2653 2658 2660 2668 2675 2682 2688 2692 2693 2696 2697 2699 2740 2741 2742 2767 2768 2777 2778 2784 2785 2799 2806 2808 2818 2822 2825 2836 2840 2852 2865 2875 2882 2894 2900 2902 2909 2910 2915 2925 2927 2935 2937 2938 2953 2956 2958 2974 2982 2985 2987 2991 2996 3000 3012 3013 3021 3024 3025 3026 3030 3032 3033 3042 3046 3064 3068 3077 3085 3091 3092 3093 3096 3102 3107 3110 3114 3119 3128 3135 3148 3150 3151 3154 3161 3164 3167 3169 3183 3187 3188 3190 3194 3201 3204 3206 3207 3212 3224 3230 3237 3242 3245 3250 3256 3260 3262 3263 3268 3271 3275 3277 3281 3283 3285 3288 3294 3305 3311 3319 3321 3324 3327 3335 3341 3344 3357 3367 3368 3371 3384 3388 3389 3397 3409 3412 3429 3434 3437 3444 3458 3464 3468 3473 3479 3492 3501 3505 3511 3513 3515 3517 3527 3539 3540 3549 3550 3553 3560 3565 3567 3577 3578 3584 3588 3600 3607 3608 3612 3614 3632 3659 3667 3670 3671 3674 3679 3680 3684 3686 3693 3700 3709 3712 3713 3729 3731 3735 3736 3737 3738 3749 3751 3752 3783 3798 3799 3804 3808 3815 3819 3826 3846 3847 3851 3855 3860 3864 3865 3871 3874 3880 3881 3886 3905 3915 3916 3917 3924 3941 3950 3956 3958 3963 3976 3978 3982 3986 3987 3989 3990 4002 4014 4015 4016 4018 4020 4023 4025 4032 4036 4045 4055 4064 4065 4068 4073 4078 4084 4087 4088 4097 4100 4107 4110 4116 4117 4135 4138 4143 4148 4150 4153 4159 4164 4168 4189 4193 4194 4199 4200 4201 4206 4208 4213 4219 4228 4238 4241 4243 4252 4254 4258 4259 4263 4264 4294 4303 4305 4314 4315 4324 4329 4359 4360 4363 4370 4379 4381 4383 4387 4390 4400 4403 4415 4436 4441 4452 4453 4461 4465 4466 4471 4478 4479 4480 4485 4504 4510 4526 4534 4550 4552 4553 4556 4566 4570 4576 4580 4584 4597 4600 4603 4606 4608 4613 4615 4616 4617 4618 4632 4639 4650 4663 4679 4683 4689 4691 4698 4700 4709 4718 4719 4720 4721 4724 4725 4728 4729 4734 4738 4739 4741 4744 4750 4753 4754 4759 4762 4768 4775 4776 4780 4781 4785 4788 4804 4807 4808 4814 4818 4821 4823 4826 4828 4830 4838 4843 4844 4864 4870 4872 4873 4889 4895 4909 4916 4918 4920 4922 4934 4936 4939 4949 4954 4962 4972 4973 4975 4976 4983 4989 4994 4996 4997 5003 5006 5008 5013 5015 5023 5025 5029 5035]

        @test all(successfull_paths' .== findall(is_success, results))
    end

    @testset "PathResult" begin
        # simple affine
        @polyvar x y
        f = [x^2 - 2, x + y - 1]
        tracker, starts = pathtracker_startsolutions(f; system = FPSystem)
        for (i, sᵢ) in enumerate(starts)
            result = track(tracker, sᵢ; path_number = i, details = :extensive)
            @test is_success(result)
            s = solution(tracker)
            @test abs(s[1]) ≈ sqrt(2) atol = 1e-8
            @test sum(s) ≈ 1 atol = 1e-8

            @test accuracy(result) < 1e-8
            @test residual(result) < 1e-8

            @test isnothing(winding_number(result))
            @test isnothing(multiplicity(result))
            @test condition_jacobian(result) < 1e3
            @test cond(result) == condition_jacobian(result)
            @test path_number(result) == i
            @test is_affine(result)
            @test !is_projective(result)
            @test is_finite(result)
            @test !is_singular(result)
            @test is_nonsingular(result)
            @test start_solution(result) == sᵢ
            @test is_nonsingular(result, 1e10)
            @test is_real(result)
            @test isreal(result)
            @test !is_failed(result)

            test_show_juno(result)
            @test !isempty(sprint(show, result))
        end

        # simple projective
        tracker, starts = pathtracker_startsolutions(homogenize(f); system = FPSystem)
        @test tracker isa PathTracker{PVector{ComplexF64,1}}
        for (i, sᵢ) in enumerate(starts)
            result = track(tracker, sᵢ; path_number = i, details = :extensive)
            @test is_success(result)
            @test is_projective(result)
            test_show_juno(result)
            @test !isempty(sprint(show, result))
        end

        # 1 singular solution (multiplicity 3) and 3 paths going to infinity. Checked earlier.
        f = equations(griewank_osborne())
        tracker, starts = pathtracker_startsolutions(f; seed = 78373, system = FPSystem)
        for sᵢ in starts
            result = track(tracker, sᵢ; details = :extensive)

            test_show_juno(result)
            @test !isempty(sprint(show, result))

            if is_at_infinity(result)
                @test result.t > 0
                @test !is_finite(result)
            elseif is_success(result)
                @test solution(result) ≈ [0.0, 0.0] atol = 1e-6
                @test is_real(result)
                @test winding_number(result) == 3
                @test accuracy(result) < 1e-6
                @test condition_jacobian(result) > 1e10
                @test is_singular(result)
                @test is_finite(result)
                @test isnothing(path_number(result))
            else
                @test false
            end
        end
    end

    @testset "Status codes" begin
        @test HC.path_tracker_status(HC.CoreTrackerStatus.success) == HC.PathTrackerStatus.success
        @test HC.path_tracker_status(HC.CoreTrackerStatus.terminated_invalid_startvalue) == HC.PathTrackerStatus.terminated_invalid_startvalue
        @test HC.path_tracker_status(HC.CoreTrackerStatus.terminated_maximal_iterations) == HC.PathTrackerStatus.terminated_max_iters
        @test HC.path_tracker_status(HC.CoreTrackerStatus.terminated_step_size_too_small) == HC.PathTrackerStatus.terminated_step_size_too_small
        @test HC.path_tracker_status(HC.CoreTrackerStatus.terminated_ill_conditioned) == HC.PathTrackerStatus.terminated_ill_conditioned
        @test HC.path_tracker_status(HC.CoreTrackerStatus.tracking) == HC.PathTrackerStatus.tracking
    end

    @testset "track with parameters change" begin
        @polyvar x a y b
        F = [x^2 - a, x * y - a + b]
        p = [a, b]

        tracker, starts = pathtracker_startsolutions(
            F,
            [1.0, 1.0 + 0.0 * im],
            parameters = p,
            p₁ = [5, 5], # wrong params
            p₀ = [10, 10],
            affine_tracking = true,
        )
        p₁ = [2, 0]
        p₀ = [2, 4]
        res =
            track(tracker, starts[1]; start_parameters = [1, 0], target_parameters = [2, 4])

        @test is_success(res)
        @test isa(solution(res), Vector{ComplexF64})
        @test length(solution(res)) == 2
    end

    @testset "Overdetermined" begin
        @polyvar x y z

        p₁ = (x^2 + y^2 + z^2 - 1) * (x - 0.5)
        p₂ = (x^2 + y^2 + z^2 - 1) * (y - 0.5)
        p₃ = (z - x^2 - 2) * (x^2 + y^2 + z^2 - 1) * (z - 0.5)
        L₁ = randn(ComplexF64, 2, 4) * [x, y, z, 1]
        L₂ = randn(ComplexF64, 2, 4) * [x, y, z, 1]

        s = let
            tracker, starts =
                    pathtracker_startsolutions([x^2 + y^2 + z^2 - 1; L₁]; system = FPSystem)
            solution(track(tracker, first(starts)))
        end

        tracker = pathtracker([[p₁, p₂, p₃]; L₁], [[p₁, p₂, p₃]; L₂], s; system = FPSystem)
        @test is_success(track(tracker, s))

        tracker, starts = pathtracker_startsolutions(
            [[p₁, p₂, p₃]; L₁],
            [[p₁, p₂, p₃]; L₂],
            s;
            system = FPSystem, projective_tracking = true,
        )
        @test is_success(track(tracker, s))
    end

    @testset "Homotopy input" begin
        @polyvar x a y b
        E = [[2 1 0; 0 0 0], [1 0; 1 0]]
        start = [[1.0 + 0im, -3.0, 2.0], [2.0 + 0im, -2.0]]
        target = [randn(ComplexF64, 3), randn(ComplexF64, 2)]
        H = CoefficientHomotopy(E, start, target)
        tracker = pathtracker(H, [[1, 1]])
        @test is_success(track!(tracker, [1, 1]))
    end
end


using QuasiCopula, LinearAlgebra, GLM
using Random, Distributions, DataFrames, ToeplitzMatrices
using Test, BenchmarkTools

BLAS.set_num_threads(1)
Threads.nthreads()

p = 3   # number of fixed effects, including intercept
m = 1    # number of variance components
# true parameter values
Random.seed!(12345)
βtrue = rand(Uniform(-2, 2), p)
σ2true = [0.5]
ρtrue = [0.5]
τtrue = 100.0
σ2 = inv(τtrue)
σ = sqrt(σ2)

function get_V(ρ, n)
    vec = zeros(n)
    vec[1] = 1.0
    for i in 2:n
        vec[i] = ρ
    end
    V = ToeplitzMatrices.SymmetricToeplitz(vec)
    V
end

#simulation parameters
samplesize = 10000
ni = 25

V = get_V(ρtrue[1], ni)

# true Gamma
Γ = σ2true[1] * V

T = Float64
gcs = Vector{GaussianCopulaCSObs{T}}(undef, samplesize)

# for reproducibility I will simulate all the design matrices here
Random.seed!(12345)
X_samplesize = [randn(ni, p - 1) for i in 1:samplesize]

for i in 1:samplesize
    X = [ones(ni) X_samplesize[i]]
    μ = X * βtrue
    vecd = Vector{ContinuousUnivariateDistribution}(undef, ni)
    for i in 1:ni
        vecd[i] = Normal(μ[i], σ)
    end
    nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)
    # simuate single vector y
    y = Vector{Float64}(undef, ni)
    res = Vector{Float64}(undef, ni)
    rand(nonmixed_multivariate_dist, y, res)
    gcs[i] = GaussianCopulaCSObs(y, X)
end

gcm = GaussianCopulaCSModel(gcs)
# precompile
println("precompiling Gaussian CS fit")
gcm2 = deepcopy(gcm);
QuasiCopula.fit!(gcm2, IpoptSolver(print_level = 0, max_iter = 20));

loglikelihood!(gcm, true, true)
fittime = @elapsed QuasiCopula.fit!(gcm, IpoptSolver(print_level = 5, max_iter = 100, tol = 10^-8, limited_memory_max_history = 50, accept_after_max_steps = 4, hessian_approximation = "limited-memory"))
@show fittime
@show gcm.β
@show gcm.σ2
@show gcm.ρ
@show gcm.∇β
@show gcm.∇σ2
@show gcm.∇ρ

@show gcm.τ
@show gcm.∇τ

loglikelihood!(gcm, true, true)
vcov!(gcm)
@show QuasiCopula.confint(gcm)

mseβ, mseρ, mseσ2, mseτ = MSE(gcm, βtrue, τtrue[1], ρtrue, σ2true)
@show mseβ
@show mseτ
@show mseσ2
@show mseρ

@test mseβ < 0.01
@test mseσ2 < 1
@test mseρ < 0.01
@test mseτ < 0.01


using Distributions, Random, Test
using EconLH.ProductionFunctionsLH
using EconLH

function pv_test()
	@testset "present value" begin
		R = 1.05;
		T = 3;

		function f(x)
			return 0.7 .* x
		end
		pv = EconLH.present_value(f, R, T);
		@test length(pv) == 1
		@test isa(pv, Float64)
		@test pv ≈ f(1) + f(2) ./ R + f(3) ./ R^2
	end
end

@testset "econLH" begin
	pv_test();
	include("extreme_value_decision_test.jl");
	include("prod_fct_test.jl");
	include("ces_test.jl")
end

# ----------


# Arvind Prasadan
# 2019 January
# DMF Package
# Tests for functions in package

push!(LOAD_PATH, "../src/")
using DMF
using LinearAlgebra
using Statistics
using StatsBase
using Test

# gen_arma_sequence, no arma
tol = 1e-2
n = 10000
ar = []
ma = []
arma_std = 1.0
trials = 200
running_sum = 0.0
for tr = 1:1:trials
    x = gen_arma_sequence(n, ar, ma, arma_std)
    global running_sum += std(x) / trials
end
@test isapprox(arma_std, running_sum; atol = tol)

# gen_arma_sequence, AR(1)
tol = 1e-2
n = 10000
ar = 0.5
ma = []
arma_std = 1.0
trials = 200
running_sum = 0.0
for tr = 1:1:trials
    x = gen_arma_sequence(n, ar, ma, arma_std)
    global running_sum += std(x) / trials
end
@test isapprox(arma_std / sqrt(1.0 - ar^2.0), running_sum; atol = tol)

# gen_cos_sequence, element
tol = 1e-8
n = 1000
f = 2.0
fs = 1.0
phase = 0.1
x, t = gen_cos_sequence(n, f, fs, phase)
@test isapprox(x[2], cos((f / fs) + phase); atol = tol)
@test isapprox(t[2], 1.0 / fs; atol = tol)

# autocorrelation, standard
tol = 1e-8
x = [1; 1; 1; 1]
gamma = autocorrelation(x; circular = false, bias_correct = false)
@test isapprox(gamma[1], 0.0; atol = tol)
@test isapprox(gamma[2], 0.0; atol = tol)

# autocorrelation, bias correction
tol = 1e-8
x = [1; 1; 1; 1]
gamma = autocorrelation(x; circular = false, bias_correct = true)
@test isapprox(gamma[1], 0.0; atol = tol)
@test isapprox(gamma[2], 0.0; atol = tol)

# autocorrelation, circular
tol = 1e-8
x = [1; 1; 1; 1]
gamma = autocorrelation(x; circular = true, bias_correct = true)
@test isapprox(gamma[1], 0.0; atol = tol)
@test isapprox(gamma[2], 0.0; atol = tol)

# autocorr_mat
tol = 1e-2
p = 5
n = 10000
trials = 200
running_sum = zeros(1, p, p)
for tr = 1:1:trials
    global running_sum += autocorr_mat(randn(p, n), 1)[1] / trials
end
@test isapprox(running_sum[1, :, :], zeros(p, p); atol = tol)

# eigenvector_error
tol = 1e-8
p = 10
Q = randn(p, 2)
Q_hat = [Q[:, 2] Q[:, 1] randn(p)]
@test isapprox(eigenvector_error(Q, Q_hat), 0.0; atol = tol)

# For Source Separation Tests

p = 100
k = 2
n = 1000
Q = qr(randn(p, k)).Q
Q = Q[:, 1:k]
S = [gen_cos_sequence(n, 2.0)[1] gen_cos_sequence(n, 0.25)[1]]
S = mapslices(normalize, S; dims = 1)
D = Diagonal(1:k)
X = Q * D * S'

# SOBI_Wrapper
tol = 1e-2
Q_hat, S_hat = SOBI_Wrapper(X, 1)
@test isapprox(eigenvector_error(Q, Q_hat), 0.0; atol = tol)
@test isapprox(eigenvector_error(S, S_hat), 0.0; atol = tol)

# DMF: Standard
tol = 1e-3
w, Q_hat, C_hat, A = dmf(X; tsvd = false, os = false, hil = false, nsv = k, C_nsv = k, demean = false, lag = 1)
@test isapprox(eigenvector_error(Q, Q_hat), 0.0; atol = tol)
@test isapprox(eigenvector_error(S, C_hat), 0.0; atol = tol)
@test size(C_hat, 2) == k

# DMF: Hilbert
tol = 1e-3
w, Q_hat, C_hat, A = dmf(X; tsvd = false, os = false, hil = true, nsv = k, C_nsv = k, demean = false, lag = 1)
@test isapprox(eigenvector_error(Q, Q_hat), 0.0; atol = tol)
@test isapprox(eigenvector_error(S, C_hat), 0.0; atol = tol)
@test size(C_hat, 2) == k

# DMF: tSVD
tol = 1e-3
w, Q_hat, C_hat, A = dmf(X; tsvd = true, os = false, hil = false, nsv = k, C_nsv = k, demean = false, lag = 1)
@test isapprox(eigenvector_error(Q, Q_hat), 0.0; atol = tol)
@test isapprox(eigenvector_error(S, C_hat), 0.0; atol = tol)
@test size(C_hat, 2) == k

# DMF: OptShrink
tol = 1e-3
w, Q_hat, C_hat, A = dmf(X; tsvd = false, os = true, hil = false, nsv = k, C_nsv = k, demean = false, lag = 1)
@test isapprox(eigenvector_error(Q, Q_hat), 0.0; atol = tol)
@test isapprox(eigenvector_error(S, C_hat), 0.0; atol = tol)
@test size(C_hat, 2) == k

# DMF: De-Mean
tol = 1e-3
w, Q_hat, C_hat, A = dmf(X; tsvd = false, os = false, hil = false, nsv = k, C_nsv = k, demean = true, lag = 1)
@test isapprox(eigenvector_error(Q, Q_hat), 0.0; atol = tol)
@test isapprox(eigenvector_error(S, C_hat), 0.0; atol = tol)
@test size(C_hat, 2) == k

# OptShrink, rank = p or n
tol = 1e-8
p = 10
n = 100
X = randn(p, n)
Xh, w = optshrink(X, p)
@test isapprox(X, Xh; atol = tol)
Xh, w = optshrink(X, n)
@test isapprox(X, Xh; atol = tol)
Xh, w = optshrink(X, n + p)
@test isapprox(X, Xh; atol = tol)

# OptShrink
tol = 1e-8
X = ones(p, n)
Xh, w = optshrink(X, 1)
@test isapprox(w[1], norm(ones(p * n), 2.0); atol = tol)
@test isapprox(X, Xh; atol = tol)


function onedLinForm( Ff, test, wg )
#= """
  onedLinForm - routine to compute \int{ f*test }

  Author: Jeff Borggaard, Virginia Tech
          part of FEMfunctions.jl

  Licensing:
     This code is distributed under the MIT license.

  Usage:
  ```julia
    F = onedLinForm( Ff, test, w_g )
  ```

  Arguments:
  - `Ff`: Function values at the Gauss points
  - `test`: Matrix of test functions evaluated at the Gauss points
            (dim: n_gauss, n_dof)
  - `w_g`: Column vector of Gauss weights

  Output arguments:
  - `F`: the element's contribution to a right-hand-side
""" =#

  F = test'*(wg.*Ff)

  return F

end


@testset "7.5 Inverse hyperbolic secant" begin
    include("7.5.1 u (a+b arcsech(c x))^n.jl")
    include("7.5.2 Inverse hyperbolic secant functions.jl")
end

test/7 Inverse hyperbolic functions/7.5 Inverse hyperbolic secant/.jl


using LinearAlgebra
using Random
using SparseArrays
using Test

using Jutils.Elements
using Jutils.Functions
using Jutils.Integration
using Jutils.Mesh
using Jutils.Transforms
using Jutils.Topologies

const lineelt = Element(Simplex{1}(), 1)
const squareelt = Element(Tensor([Simplex{1}(), Simplex{1}()]), 1)

ev(func, pt, elt) = callable(func)(pt, elt)

@testset "Transforms" begin include("Transforms.jl") end

@testset "Functions" begin include("Functions.jl") end
@testset "Gradients" begin include("Gradients.jl") end

@testset "Optimization" begin include("Optimization.jl") end

@testset "Integration" begin include("Integration.jl") end


module Day2

import DelimitedFiles

function read_input(filepath)
    return DelimitedFiles.readdlm(filepath, ' ', '\n')
end

function calculate_position_and_depth(course)
    position = 0
    depth = 0
    for row = 1:(size(course)[1])
        direction = course[row, 1]
        magnitude = course[row, 2]
        if direction == "forward"
            position += magnitude
        elseif direction == "up"
            depth -= magnitude
        elseif direction == "down"
            depth += magnitude
        end
    end
    return (position, depth)
end

function part1_answer()
    input_filepath = joinpath(@__DIR__, "input.txt")
    position, depth = calculate_position_and_depth(read_input(input_filepath))
    return position * depth
end

function calculate_position_and_depth_with_aim(course)
    position = 0
    depth = 0
    aim = 0
    for row = 1:(size(course)[1])
        direction = course[row, 1]
        magnitude = course[row, 2]
        if direction == "forward"
            position += magnitude
            depth += aim*magnitude
        elseif direction == "up"
            aim -= magnitude
        elseif direction == "down"
            aim += magnitude
        end
    end
    return (position, depth)
end

function part2_answer()
    input_filepath = joinpath(@__DIR__, "input.txt")
    position, depth = calculate_position_and_depth_with_aim(read_input(input_filepath))
    return position * depth
end

end


"""
## bmat_nonaxi!

This subroutine forms the strain-displacement matrix for
axisymmetric solids subjected to non-axisymmetric loading.

### Method
```julia
bmat_nonaxi!(bee, radius, coord, deriv, fun, iflag, lth)
```

### Arguments
```julia
* bee::Matrix{Float64}         : Bee matrix (Updated)
* radius::Float64              : r coordinate of the Gauss point
* coord::Matrix{Float64}       : Nodal coordinate matrix
* deriv::Matrix{Float64}       : Derivative
* fun::Vector{Float64}         : Shape function
* iflag::Int                 : 1 = symmetric, -1 = anti-symmetric
* lth::Int                   : Harmonic
```
"""
function bmat_nonaxi!(bee::Matrix{Float64}, radius::Float64,
  coord::LinearAlgebra.Adjoint{Float64,Array{Float64,2}}, deriv::Matrix{Float64},
  fun::Vector{Float64}, iflag::Int, lth::Int)
  #
  # This subroutine forms the strain-displacement matrix for
  # axisymmetric solids subjected to non-axisymmetric loading.
  #
  # Updates bee, radius
  #
 nod = size(deriv, 2)
 radius = sum(fun.*coord[:,1])
 for m in 1:nod
   n = 3 * m
   k = n-1
   l = k-1
   bee[1,l] = deriv[1, m]
   bee[2,k] = deriv[2, m]
   bee[3,l] = fun[m]/radius
   bee[3,n] = iflag*lth*bee[3, l]
   bee[4,l] = deriv[2, m]
   bee[4,k] = deriv[1, m]
   bee[5,k] = -iflag*lth*fun[m]/radius
   bee[5,n] = deriv[2, m]
   bee[6,l] = bee[5,k]
   bee[6,n] = deriv[1, m]-fun[m]/radius
 end
 (radius, bee)
end