@testset verbose=true "Inequalities                " begin

	n = 16
	m = 5

	# no bounds, lower bound, upper bound, both bounds
	divsize = n ÷ 4
	xl = [-Inf.*ones(divsize); randn(divsize); -Inf.*ones(divsize); randn(divsize)]
	xu = [Inf.*ones(divsize); Inf.*ones(divsize); randn(divsize); randn(divsize)]

	x = [randn(n÷4);
		xl[divsize+1:2*divsize] .+ rand(0:2, divsize);
		xu[2*divsize+1:3*divsize] .- rand(0:2, divsize);
		xl[3*divsize+1:4*divsize] .+ rand(divsize).*(xu[3*divsize+1:4*divsize] .- xl[3*divsize+1:4*divsize])]

	xaug = zeros(2*n)
	xaug[1:n] .= x

	idata = LFPSQP.InequalityData(xl, xu)

	# test aspects of Inequality data
	@testset "InequalityData" begin
		@test all(idata.q .≈ [zeros(3*divsize); ones(divsize)])
		@test all(idata.r .≈ [zeros(divsize);
								xl[divsize+1:2*divsize];
								xu[2*divsize+1:3*divsize];
								xl[3*divsize+1:4*divsize]./2 .+ xu[3*divsize+1:4*divsize]./2])
		@test all(idata.s .≈ [zeros(divsize); -ones(divsize); ones(2*divsize)])
		@test all(idata.t .≈ [zeros(divsize);
								xl[divsize+1:2*divsize];
								xu[2*divsize+1:3*divsize];
								(xu[3*divsize+1:4*divsize] .- xl[3*divsize+1:4*divsize]).^2 ./ 4])

		@test all(idata.isline .== vcat(trues(divsize), falses(3*divsize)))
		@test all(idata.isparabola .== vcat(falses(divsize), trues(2*divsize), falses(divsize)))
	end

	# test y generation satisfies inequalities
	@testset "Initial y" begin
		LFPSQP.generate_initial_y!(xaug, idata)

		xview = view(xaug, 1:n)
		yview = view(xaug, n+1:2*n)

		@test all(isapprox.(idata.q.*(xview .- idata.r).^2 .+ (1.0 .- idata.q.^2).*xview .+
			idata.s.*(yview .- idata.r).^2 .- (1.0 .- idata.s.^2).*yview .- idata.t, 0.0, atol=1e-15))

		cvalaug = zeros(n+m)
		LFPSQP.calculate_h!(cvalaug, xaug, idata)

		@test all(isapprox.(cvalaug, 0.0, atol=1e-15))
	end

	# generate random decomposition
	Jct = randn(n, m)

	∇hx = 2.0*idata.q.*(x .- idata.r) .+ (1.0 .- idata.q.^2)
	∇hy = 2.0*idata.s.*(xaug[n+1:end] .- idata.r) .- (1.0 .- idata.s.^2)

	S = sqrt.(∇hx.^2 .+ ∇hy.^2)
	Dx = ∇hx ./ S
	Dy = ∇hy ./ S

	R = Dx .* Jct

	PJct = zeros(2*n, m)
	PJct[1:n, :] .= (1.0 .- Dx.*Dx) .* Jct
	PJct[n+1:2*n, :] .= -1.0 .* Dy.*Dx .* Jct

	U, Σ, V = svd(PJct)
	Vt = Matrix(V')

	bigA = [diagm(∇hx) Jct; diagm(∇hy) zeros(n, m)]
	bigQ = hcat([diagm(Dx); diagm(Dy)], U)
	bigR = [diagm(S) R; zeros(m, n) diagm(Σ)*Vt]

	idecomp = LFPSQP.InequalityDecomp(U, Σ, Vt, [Array{Float64}(undef, n) for i in 1:3]..., Jct, m)

	# test InequalityDecomposition
	@testset "InequalityDecomposition" begin

		LFPSQP.inequality_gradient!(idecomp, xaug, idata)

		@test all(isapprox.(idecomp.Dx, Dx, atol=1e-15))
		@test all(isapprox.(idecomp.Dy, Dy, atol=1e-15))
		@test all(isapprox.(idecomp.S, S, atol=1e-15))

		@test all(isapprox.(bigA, bigQ*bigR, atol=1e-14))					# decomposition is correct
		@test all(isapprox.(bigQ'*bigQ, Matrix(I, n+m, n+m), atol=1e-14))	# orthogonality
	end

	@testset "Multiplication" begin
		v = randn(n+m)
		w = randn(2*n)

		destv = zeros(2*n)
		destw = zeros(n+m)

		# projections
		mul!(destv, LFPSQP.InequalityDecompProject(idecomp), v)
		@test all(isapprox.(destv, bigQ*v, atol=1e-15))

		mul!(destw, LFPSQP.InequalityDecompProject(idecomp)', w)
		@test all(isapprox.(destw, bigQ'*w, atol=1e-15))

		destv .= 1.0
		mul!(destv, LFPSQP.InequalityDecompProject(idecomp), v, 2.0, 3.0)
		@test all(isapprox.(destv, 2*bigQ*v .+ 3, atol=1e-15))

		# Jacobians
		mul!(destv, idecomp, v)
		@test all(isapprox.(destv, bigA*v, atol=1e-15))

		mul!(destw, idecomp', w)
		@test all(isapprox.(destw, bigA'*w, atol=1e-15))

		destv .= 1.0
		mul!(destv, idecomp, v, 2.0, 3.0)
		@test all(isapprox.(destv, 2*bigA*v .+ 3, atol=3e-15))

		# test lower rank
		idecomp.rank = m - 2

		v = randn(n+m-2)
		destv = zeros(2*n)
		destw = zeros(n+m-2)

		mul!(destv, LFPSQP.InequalityDecompProject(idecomp), v)
		@test all(isapprox.(destv, bigQ[:, 1:n+m-2]*v, atol=1e-15))

		mul!(destw, LFPSQP.InequalityDecompProject(idecomp)', w)
		@test all(isapprox.(destw, bigQ[:, 1:n+m-2]'*w, atol=1e-15))

		destv .= 1.0
		mul!(destv, LFPSQP.InequalityDecompProject(idecomp), v, 2.0, 3.0)
		@test all(isapprox.(destv, 2*bigQ[:, 1:n+m-2]*v .+ 3, atol=1e-15))

		idecomp.rank = m
	end

	@testset "Lagrange multipliers" begin
		d = randn(2*n)
		Qt∇f = zeros(n+m)
		mul!(Qt∇f, LFPSQP.InequalityDecompProject(idecomp)', d)

		λ = zeros(n+m)
		λ_kkt = view(λ, n+1:n+m)
		λy_kkt = view(λ, 1:n)

		LFPSQP.calculate_λ_kkt!(λ_kkt, λy_kkt, Qt∇f, idecomp)

		@test all(isapprox.(λ, bigA \ d, atol=1e-14))
	end

	@testset "Hessian action" begin

		A = randn(n, n)
		C = randn(n, n, m)
		λ_kkt = randn(m)
		λy_kkt = randn(n)

		dest = zeros(2*n)
		v = randn(2*n)

		function hess_lag_vec!(dest, src, x, λ)
			dest .= A*src + sum(λ[i]*C[:,:,i]*src for i in 1:m)
		end

		LFPSQP.augmented_hess_lag_vec!(dest, v, hess_lag_vec!, xaug, λ_kkt, λy_kkt, idata)

		H = A + sum(λ_kkt[i]*C[:,:,i] for i in 1:m)
		bigH = [H+2*diagm(λy_kkt.*idata.q) zeros(n, n); zeros(n,n) 2*diagm(λy_kkt.*idata.s)]

		@test all(isapprox.(dest, bigH*v, atol=1e-14))
	end

	@testset "y retraction" begin
		# take step in tangent plane
		d = randn(2*n)
		tmp_m = zeros(n+m)
		ineqproject = LFPSQP.InequalityDecompProject(idecomp)

		mul!(tmp_m, ineqproject', d)				# tmp_m = Q'd
		mul!(d, ineqproject, tmp_m, -1.0, 1.0)		# d = d - Q*tmp_m

		# test really long step length
		xnewaug = xaug .+ d
		xaugcopy = copy(xaug)
		LFPSQP.y_retract!(xnewaug, xaug, idata)

		# ensure y retraction worked and that inequality constraints are 0
		cvalaug = zeros(n+m)
		LFPSQP.calculate_h!(cvalaug, xnewaug, idata)

		@test all(isapprox.(cvalaug, 0.0, atol=1e-13))
		@test all(xaug .== xaugcopy)					# ensure xaug not overwritten
	end
end


"""
A shape which contains all particles,``x > x_0``, necessary to simulate a plane-wave scattering from an infinite medium, for a reciever at the focal point, for time ``t < D / c``, where ``c`` is the sound speed of the background medium, and ``D`` is some chosen focal distance. The plane-wave travels towards the positive direction of the ``x`` axis.

More precisely, if the focal point is at ``(x_f,y_f)`` then the interior of the shape
is defined as
``(y - y_f)^2 < (D + x_0)^2 - x_f^2 - 2(D + x_0 - x_f)x``  and ``x > min(x_0, x_f)``
where ``D`` is the focal distance.
"""
struct TimeOfFlightPlaneWaveToPoint{T <: AbstractFloat,Dim} <: Shape{T,Dim}
    focal_point::AbstractVector{T}
    focal_distance::T
    minimum_x::T
end

TimeOfFlightPlaneWaveToPoint(focal_point::AbstractVector{T},  focal_distance::T;  minimum_x::T = zero(T)) where T <:AbstractFloat = TimeOfFlightPlaneWaveToPoint{T,length(focal_point)}(focal_point, focal_distance, minimum_x)

name(shape::TimeOfFlightPlaneWaveToPoint) = "Time of flight from planar source to the focal point"

# NOTE: this needs to be redone
function volume(shape::TimeOfFlightPlaneWaveToPoint{T}) where T <: AbstractFloat

    D = shape.focal_distance
    xf = shape.focal_point[1]
    x_min = shape.minimum_x
    x_max = (x_min + xf + D) / T(2)

    a = D + x_min + xf
    b = D + x_min - xf
    return 2/3*(sqrt(b*(a - 2*x_min)^3) - sqrt(b*(a - 2*x_max)^3))
end

import Base.issubset
function issubset(sphere::Sphere{T,Dim}, shape::TimeOfFlightPlaneWaveToPoint{T,Dim}) where {T, Dim}
    return (origin(sphere)[1] > shape.minimum_x) && (norm(origin(sphere) - shape.focal_point) <= shape.focal_distance - (origin(sphere)[1] - shape.minimum_x)  - T(2) * outer_radius(sphere))
end

function bounding_box(shape::TimeOfFlightPlaneWaveToPoint{T,Dim}) where {T,Dim}
    D = shape.focal_distance
    xf = shape.focal_point[1]
    y_f = shape.focal_point[2]
    x_min = shape.minimum_x
    x_max = (x_min + xf + D) / T(2)

    centre = zeros(Dim)
    centre[1] = (x_min + x_max) / T(2)
    centre[2] = y_f

    dimensions = ones(Dim) * T(2) * sqrt((D + x_min)^2 - xf^2 - 2*(D + x_min - xf)*x_min)
    dimensions[1] = x_max - x_min

    return Box(centre, dimensions)
end

function boundary_functions(shape::TimeOfFlightPlaneWaveToPoint{T,2}) where T

    D = shape.focal_distance
    xf = shape.focal_point[1]
    y_f = shape.focal_point[2]
    x_min = shape.minimum_x
    x_max = (x_min + xf + D) / T(2)

    function x(τ)
        check_boundary_coord_range(τ)

        if τ <= 1//3
            return (1 - 3*τ)*x_min + 3*τ*x_max
        elseif τ <= 2//3
            return (3*τ - 1)*x_min + (2 - 3*τ)*x_max
        else
            return x_min
        end
    end

    function y(τ)
        check_boundary_coord_range(τ)

        if τ <= 1//3
            return y_f + sqrt((D + x_min)^2 - xf^2 - 2*(D + x_min - xf)*((1 - 3*τ)*x_min + 3*τ*x_max))
        elseif τ <= 2//3
            return y_f - sqrt((D + x_min)^2 - xf^2 - 2*(D + x_min - xf)*((3*τ - 1)*x_min + 3*(2//3 - τ)*x_max))
        else
            return y_f + (6*τ - 5)*sqrt((D + x_min)^2 - xf^2 - 2*(D + x_min - xf)*x_min)
        end
    end

    return x, y
end


include("cg.jl")
using CG;
include("render.jl");

v1 = Vec4f(1, 0, 0, 0);
v2 = Vec4f(0.0, 42.0, 0.0, 0.0);

println(v1 - v2);

m1 = Mat4f(Vec4f(1.0, 0.0, 0.0, 0.0),
           Vec4f(0.0, 1.0, 0.0, 0.0),
           Vec4f(0.0, 0.0, 1.0, 0.0),
           Vec4f(0.0, 0.0, 0.0, 1.0));

println(2 * (m1 * v1));
println(m1 * v2);

v1 = Vec4f(1, 0, 0, 1);
v2 = Vec4f(0, 0, 0, 1);
v3 = Vec4f(0, 0, 1, 1);

triangle = Object(v1, v2, v3);

house = Object(Vec4f(-1, -1, 0, 1),
Vec4f(1, -1, 0, 1),
Vec4f(-1, 1, 0, 1),
Vec4f(0, 2, 0, 1),
Vec4f(1, 1, 0, 1),
Vec4f(1, 1, 0, 1),
Vec4f(-1, -1, 0, 1),
Vec4f(-1, 1, 0, 1),
Vec4f(1, 1, 0, 1),
Vec4f(1, -1, 0, 1))

#render(translation(Vec4f(1.5, 0, 0, 1)) * house, figAxis=[-2, 2, -2, 2]);
#render(scaling(Vec4f(0.5, 0.5, 0.5, 1)) * house, figAxis=[-2, 2, -2, 2]);
render(roty(pi/4) * house, figAxis=[-2, 2, -2, 2]);


using JuMP, FlexibilityAnalysis, Gurobi

# Set the covariance matrix for the uncertain parameters
means = [4.0; 5.0]
bea = 1
covar = [2. bea;
         bea 3.]

# Setup the model
m = FlexibilityModel(solver = GurobiSolver(OutputFlag = 0))

# Define variables
@randomvariable(m, x, mean = means[1])
@randomvariable(m, y, mean = means[2])

# Define the constraints
@constraint(m, x + y - 14 <= 0.0)
@constraint(m, x - 2y - 2 <= 0.0)
@constraint(m, -x <= 0.0)
@constraint(m, -y <= 0.0)

# Define the uncertainty set
setuncertaintyset(m, :Ellipsoid, covar)

# Solve
status = solve(m, active_constr = true)
if status == :Optimal
    # Retrieve optimized data
    flexibility_index = getflexibilityindex(m)
    x_opt = getvalue(x)
    y_opt = getvalue(y)
    actives = getactiveconstraints(m)

    # Get the confidence level
    conf_lvl = getconfidencelevel(m)

    # Rank the inequality constraints
    rank_data = rankinequalities(m, active_constr = true)
end

# Estimate the value of SF
SF = findstochasticflexibility(m, use_flexibility_index = true)


#!/usr/bin/env julia

using Colors
using Compose
using Base.MathConstants

function golden_rect(n::Int)
    poly_points = [(0, 0), (1, 0), (1, 1), (0, 1)]
    if n == 0 return context() end
    c = compose(context(), polygon(poly_points), fill(LCHab(90, 80, 70-20n)), stroke("black"))
    compose(c, (context(0,0,1/φ,1/φ, rotation=Rotation(-π/2,1,0)),  golden_rect(n-1)))
end

draw(SVG("golden_rect.svg", φ*3inch, 3inch),
    compose(golden_rect(8), linewidth(0.2mm)))


module VectorFieldsOnASphere

using LabelledArrays,OffsetArrays

import LinearAlgebra: norm,normalize,dot
import Base: +,-,*,/,==,≈,collect

using PointsOnASphere
import TwoPointFunctions: cosχ

export CartesianVector,SphericalVector,HelicityVector,SphericalPolarVector,unitvector

abstract type VectorField end

const vector_types = (:CartesianVector,:SphericalVector,:SphericalPolarVector,:HelicityVector)

function cartesian_components(x,y,z)
	v = @LVector Float64 (:x,:y,:z)
	v.x = x
	v.y = y
	v.z = z
	return v
end

function spherical_components(p,z,m)
	v = @LVector ComplexF64 (:p,:z,:m)
	v.p = p
	v.z = z
	v.m = m
	return v
end

function spherical_polar_components(r,θ,ϕ)
	v = @LVector Float64 (:r,:θ,:ϕ)
	v.r = r
	v.θ = θ
	v.ϕ = ϕ
	return v
end

helicity_components(p,z,m) = spherical_components(p,z,m)

# Vector fields evaluated at a point
struct CartesianVector <: VectorField
	components :: AbstractArray
	pt::Point3D
end

struct SphericalVector <: VectorField
	components :: AbstractArray
	pt::Point3D
end

struct SphericalPolarVector <: VectorField
	components :: AbstractArray
	pt::Point3D
end

struct HelicityVector <: VectorField
	components :: AbstractArray
	pt::Point3D
end

const complex_basis = Union{SphericalVector,HelicityVector}
const real_basis = Union{CartesianVector,SphericalPolarVector}
const fixed_basis = Union{CartesianVector,SphericalVector}
const rotating_basis = Union{SphericalPolarVector,HelicityVector}

const CtoS = (@SLArray ComplexF64 (3,3) (:px,:zx,:mx,:py,:zy,:my,:pz,:zz,:mz))(
								-1/√2,	0,	1/√2,
								im/√2,	0,	im/√2,
									0,	1,	0)

const StoC = (@SLArray ComplexF64 (3,3) (:xp,:yp,:zp,:xz,:yz,:zz,:xm,:ym,:zm))(CtoS')
const SPtoH = (@SLArray ComplexF64 (3,3) (:pθ,:zθ,:mθ,:pϕ,:zϕ,:mϕ,:pr,:zr,:mr))(CtoS)
const HtoSP = (@SLArray ComplexF64 (3,3) (:θp,:ϕp,:rp,:θz,:ϕz,:rz,:θm,:ϕm,:rm))(StoC)

function SP_to_C_matrix(v)
	(@SLArray Float64 (3,3) (:xr,:yr,:zr,:xθ,:yθ,:zθ,:xϕ,:yϕ,:zϕ))(
		[sin(v.pt.θ)cos(v.pt.ϕ)	cos(v.pt.θ)cos(v.pt.ϕ)	-sin(v.pt.ϕ)
		 sin(v.pt.θ)sin(v.pt.ϕ)	cos(v.pt.θ)sin(v.pt.ϕ)	cos(v.pt.ϕ)
		 cos(v.pt.θ)			-sin(v.pt.θ)					0])
end

function C_to_SP_matrix(v)
	(@SLArray Float64 (3,3) (:rx,:θx,:ϕx,:ry,:θy,:ϕy,:rz,:θz,:ϕz))(SP_to_C_matrix(v)')
end

# Matrix to transform from cartesian to spherical coordinates
function CartesianToSpherical(v::CartesianVector)
	components = @LVector ComplexF64 (:p, :m, :z)
	vc = v.components
	components.z = vc.z
	components.p = CtoS.px * vc.x + CtoS.py * vc.y
	components.m = CtoS.mx * vc.x + CtoS.my * vc.y
	SphericalVector(components,v.pt)
end

function CartesianToSphericalPolar(v::CartesianVector)
	M = C_to_SP_matrix(v)
	components = @LVector Float64 (:r, :θ, :ϕ)
	vc = v.components
	components.r = M.rx*vc.x + M.ry*vc.y + M.rz*vc.z
	components.θ = M.θx*vc.x + M.θy*vc.y + M.θz*vc.z
	components.ϕ = M.ϕx*vc.x + M.ϕy*vc.y + M.ϕz*vc.z

	SphericalPolarVector(components,v.pt)
end

CartesianToHelicity(v::CartesianVector) = SphericalPolarToHelicity(CartesianToSphericalPolar(v))

function SphericalToCartesian(v::SphericalVector)
	components = @LVector Float64 (:x, :y, :z)
	vc = v.components
	components.x = real(StoC.xp*vc.p + StoC.xm*vc.m)
	components.y = real(StoC.yp*vc.p + StoC.ym*vc.m)
	components.z = real(vc.z)
	CartesianVector(components,v.pt)
end

function SphericalToSphericalPolar(v::SphericalVector)
	CartesianToSphericalPolar(SphericalToCartesian(v))
end

SphericalToHelicity(v::SphericalVector) = SphericalPolarToHelicity(SphericalToSphericalPolar(v))

function SphericalPolarToCartesian(v::SphericalPolarVector)
	M = SP_to_C_matrix(v)
	components = cartesian_components(zero(Float64),zero(Float64),zero(Float64))
	vc = v.components
	components.x = M.xr*vc.r + M.xθ*vc.θ + M.xϕ*vc.ϕ
	components.y = M.yr*vc.r + M.yθ*vc.θ + M.yϕ*vc.ϕ
	components.z = M.zr*vc.r + M.zθ*vc.θ + M.zϕ*vc.ϕ
	CartesianVector(components,v.pt)
end

function SphericalPolarToSpherical(v::SphericalPolarVector)
	CartesianToSpherical(SphericalPolarToCartesian(v))
end

function SphericalPolarToHelicity(v::SphericalPolarVector)
	components = @LVector ComplexF64 (:p, :m, :z)
	vc = v.components
	components.p = SPtoH.pθ*vc.θ + SPtoH.pϕ*vc.ϕ
	components.m = SPtoH.mθ*vc.θ + SPtoH.mϕ*vc.ϕ
	components.z = vc.r
	HelicityVector(components,v.pt)
end

HelicityToCartesian(v::HelicityVector) = SphericalPolarToCartesian(HelicityToSphericalPolar(v))

HelicityToSpherical(v::HelicityVector) = SphericalPolarToSpherical(HelicityToSphericalPolar(v))

function HelicityToSphericalPolar(v::HelicityVector)
	components = @LVector Float64 (:r, :θ, :ϕ)
	vc = v.components
	components.θ = real(HtoSP.θp*vc.p + HtoSP.θm*vc.m)
	components.ϕ = real(HtoSP.ϕp*vc.p + HtoSP.ϕm*vc.m)
	components.r = real(vc.z)
	SphericalPolarVector(components,v.pt)
end

# Assume radius vector by default if coordinates are not specified
SphericalPolarVector(x::Point3D) = SphericalPolarVector(spherical_polar_components(x.r,0,0),x)
# Assume point on unit sphere if distance from origin is not specified
SphericalPolarVector(n::Point2D) = SphericalPolarVector(spherical_polar_components(1,0,0),Point3D(1,n))

CartesianVector(x::SphericalPoint) = SphericalPolarToCartesian(SphericalPolarVector(x))
SphericalVector(x::SphericalPoint) = SphericalPolarToSpherical(SphericalPolarVector(x))
HelicityVector(x::SphericalPoint) = SphericalPolarToHelicity(SphericalPolarVector(x))

CartesianVector(v::CartesianVector) = v
CartesianVector(v::SphericalVector) = SphericalToCartesian(v)
CartesianVector(v::SphericalPolarVector) = SphericalPolarToCartesian(v)
CartesianVector(v::HelicityVector) = HelicityToCartesian(v)

HelicityVector(v::HelicityVector) = v
HelicityVector(v::CartesianVector) = CartesianToHelicity(v)
HelicityVector(v::SphericalVector) = SphericalToHelicity(v)
HelicityVector(v::SphericalPolarVector) = SphericalPolarToHelicity(v)

SphericalVector(v::SphericalVector) = v
SphericalVector(v::CartesianVector) = CartesianToSpherical(v)
SphericalVector(v::SphericalPolarVector) = SphericalPolarToSpherical(v)
SphericalVector(v::HelicityVector) = HelicityToSpherical(v)

SphericalPolarVector(v::SphericalPolarVector) = v
SphericalPolarVector(v::CartesianVector) = CartesianToSphericalPolar(v)
SphericalPolarVector(v::SphericalVector) = SphericalToSphericalPolar(v)
SphericalPolarVector(v::HelicityVector) = HelicityToSphericalPolar(v)

twoD_types = Union{Point2D,Tuple{<:Real,<:Real}}
threeD_types = Union{Point3D,Tuple{<:Real,<:Real,<:Real}}

CartesianVector(x,y,z,pt::threeD_types) = CartesianVector(cartesian_components(x,y,z),Point3D(pt))
SphericalVector(p,z,m,pt::threeD_types) = SphericalVector(spherical_components(p,z,m),Point3D(pt))
SphericalPolarVector(r,θ,ϕ,pt::threeD_types) = SphericalPolarVector(spherical_polar_components(r,θ,ϕ),Point3D(pt))
HelicityVector(p,z,m,pt::threeD_types) = HelicityVector(helicity_components(p,z,m),Point3D(pt))

CartesianVector(x,y,z,n::twoD_types) = CartesianVector(x,y,z,Point3D(1,n))
SphericalVector(p,z,m,n::twoD_types) = SphericalVector(p,z,m,Point3D(1,n))
SphericalPolarVector(r,θ,ϕ,n::twoD_types) = SphericalPolarVector(r,θ,ϕ,Point3D(1,n))
HelicityVector(p,z,m,n::twoD_types) = HelicityVector(p,z,m,Point3D(1,n))

for T in vector_types
	# Iterables
	@eval $(T)(v,n::SphericalPoint) = $(T)(v[1],v[2],v[3],n)
end

(==)(v1::T,v2::T) where T<:VectorField = v1 === v2
(≈)(v1::T,v2::T) where T<:VectorField = (v1.components ≈ v2.components) && (v1.pt === v2.pt)

norm(v::T) where T<:real_basis = norm(v.components)
function norm(v::T) where T<:complex_basis
	real(√( v.components.z^2 - 2v.components.p*v.components.m ))
end

normalize(v::T) where T<:real_basis = v/norm(v)

normalize(v::SphericalVector) = SphericalVector(normalize(CartesianVector(v)))
normalize(v::HelicityVector) = HelicityVector(normalize(SphericalPolarVector(v)))

unitvector(v::VectorField) = normalize(v)

(+)(v1::T,v2::T) where T<:fixed_basis = T(v1.components + v2.components,v1.pt)
(-)(v1::T,v2::T) where T<:fixed_basis = T(v1.components - v2.components,v2.pt)

function (+)(v1::T,v2::T) where T<:rotating_basis
	if v1.pt === v2.pt
		return T(v1.components + v2.components,v1.pt)
	elseif (v1.pt.θ == v2.pt.θ) && (v1.pt.ϕ == v2.pt.ϕ)
		# in this case the points are radially separated,
		# and the basis vectors are aligned
		return T(v1.components + v2.components,v1.pt)
	else
		v1C = CartesianVector(v1)
		v2C = CartesianVector(v2)
		T(v1C + v2C)
	end
end

function (-)(v1::T,v2::T) where T<:rotating_basis
	if v1.pt === v2.pt
		return T(v1.components - v2.components,v2.pt)
	elseif (v1.pt.θ == v2.pt.θ) && (v1.pt.ϕ == v2.pt.ϕ)
		# in this case the points are radially separated,
		# and the basis vectors are aligned
		return T(v1.components - v2.components,v2.pt)
	else
		v1C = CartesianVector(v1)
		v2C = CartesianVector(v2)
		T(v1C - v2C)
	end
end

# Scaling
(*)(v::T,c::Number) where T<:VectorField = T(v.components.*c,v.pt)
(*)(c::Number,v::T) where T<:VectorField = T(v.components.*c,v.pt)
(/)(v::T,c::Number) where T<:VectorField = T(v.components./c,v.pt)

cosχ(v::VectorField,w::VectorField) = cosχ(v.pt,w.pt)

# Inner product of two vectors
function dot(v::VectorField,w::VectorField)
	# Use the fact that a⋅b = |a||b|cos(θ)
	norm(v)*norm(w)*cosχ(v,w)
end

Base.parent(v::VectorField) = v.components
Base.getindex(v::VectorField,inds...) = getindex(parent(v),inds...)

# convert to array
collect(v::real_basis) = collect(v.components)
collect(v::complex_basis) = OffsetArray([v.components.m,v.components.z,v.components.p],-1:1)

# display
function Base.show(io::IO, v::VectorField)
    compact = get(io, :compact, false)

    print(io,round.(v.components,sigdigits=3)," at $(v.pt)")
end

function Base.show(io::IO, ::MIME"text/plain", v::VectorField)
    println(io,"Vector with components $(v.components)")
    print(io,"Defined at $(v.pt)")
end

end # module


using BayesTests
using Base.Test

binomBF = BinomTest(1, 1)
samples = posterior(binomBF; iter = 10000)
credible = hpd(samples)

@test binomBF.BF == 1.0
@test_approx_eq_eps credible[1] .167 .1
@test_approx_eq_eps credible[2] .986 .1

update(binomBF, 3, 3)

corBF = CorrelationTest(1, 1; a = 1)
corBF2 = CorrelationTest(3, 1; a = 1)
@test isnan(corBF.BF)
@test !isnan(corBF2.BF)

x1 = collect(1:10)
x2 = collect(11:20)

ttestBF = TTest(x1, x2)
update(ttestBF, x1, x2)

ttestBF = TTest(x1, x2; typ = :twosample)
update(ttestBF, x1, x2)


##
cd(@__DIR__())
function parse_mask(s, ::Type{T} = UInt64) where {T<:Unsigned}
    mask0 = zero(T)
    bit1 = one(T)
    map(enumerate(reverse(s))) do (i, c)
        c == '0' && return Base.Fix2(&, ~((mask0 | bit1) << (i-1)))
        c == '1' && return Base.Fix2(|, (mask0 | bit1) << (i-1))
        return missing
    end |> skipmissing |> collect
end
maskrx = r"mask = (.*)"
memrx = r"mem\[([0-9]+)\] = ([0-9]+)"
ismask(t) = t[1] === :mask
ismem(t) = t[1] === :mem
function run_program(prog)
    space = extrema(t->t[2], Iterators.filter(ismem, prog))
    mem = zeros(Int, space)
    mask = nothing
    foreach(prog) do cmd
        if ismask(cmd)
            mask = cmd[2]
            return
        end
        if ismem(cmd)
            mem[cmd[2]] = foldl(mask; init = cmd[3]) do v, op2
                op2(v)
            end
            return
        end
    end
    sum(mem)
end

##
prog = map(eachline("input.txt")) do line
    m = match(maskrx, line)
    if !isnothing(m)
        return :mask, parse_mask(m[1])
    end
    m = match(memrx, line)
    if !isnothing(m)
        return :mem, parse(Int, m[1]), parse(Int, m[2])
    end
    return nothing
end;

##
run_program(prog)


@with_kw mutable struct HalfspaceParameters{T <: AbstractFloat} <: AbstractPolicyParameters
    α::T = 0.
    a::Vector{T}
end

struct Halfspace <: AbstractProx
    params::HalfspaceParameters{Float64}
    norm::Float64

    function (::Type{Halfspace})(; kwargs...)
        params = HalfspaceParameters(; kwargs...)
        @unpack a = params
        return new(params,a⋅a)
    end
end

function prox!(halfspace::Halfspace,step::Real,xprev::AbstractVector,gcurr::AbstractVector,xcurr::AbstractVector)
    @unpack α,a = halfspace.params
    x = xprev - step*gcurr
    diff = a⋅x - α
    xcurr .= x - (max(diff,zero(eltype(a)))/halfspace.norm)*a
end


using Test, Random, LinearAlgebra, LuxurySparse

using Yao
using Yao.YaoArrayRegister: invorder
using YaoExtensions
using FFTW

@testset "QFT" begin
    num_bit = 5
    fftblock = qft_circuit(num_bit)
    ifftblock = fftblock'
    reg = rand_state(num_bit)
    rv = copy(statevec(reg))

    @test Matrix(mat(chain(3, qft_circuit(3) |> adjoint, qft_circuit(3)))) ≈ IMatrix(1<<3)

    # test ifft
    reg1 = apply!(copy(reg), ifftblock)

    # permute lines (Manually)
    kv = fft(statevec(reg))/sqrt(length(rv))
    @test statevec(reg1) ≈ invorder(kv)

    # test fft
    reg2 = apply!(invorder!(copy(reg)), fftblock)
    kv = ifft(rv) * sqrt(length(rv))
    @test statevec(reg2) ≈ kv
end

@testset "QFT" begin
    num_bit = 5
    circuit_qft = qft_circuit(num_bit)
    circuit_iqft = circuit_qft'
    block_qft = qft(num_bit)
    block_iqft = qft(num_bit)'
    @test openbox(block_qft) == circuit_qft
    @test openbox(block_iqft) == circuit_iqft
    reg = rand_state(num_bit)

    @test Matrix(chain(3, qft(3)', qft(3))) ≈ IMatrix(1<<3)

    # permute lines (Manually)
    @test apply!(copy(reg), circuit_iqft) ≈ apply!(copy(reg), block_iqft)

    # test fft
    @test apply!(copy(reg), qft) ≈ apply!(copy(reg), block_qft)

    # regression test for nactive
    @test apply!(focus!(copy(reg), 1:3), QFT{3}()) |> isnormalized
end


export getSensMatVecEik, getSensTMatVecEik

function getOp(opCode::Int8,dim::Int64)
if dim==1
	op = mod(opCode,5)-2;
elseif dim==2
	op = mod(opCode,25);
	op = div(op,5)-2;
elseif dim==3
	op = div(opCode,25)-2;
end
end

function getOpCode(op1,op2,op3)
opCode = (op1+2) + 5*(op2+2) + 25*(op3+2);
return opCode;
end

function getSensMatVecEik(b::Array{Float64,1},res::Array{Float64,1},pEik::EikonalParam)
# This is basically a lower triangular solution.
	order = pEik.ordering;
	n = pEik.Mesh.n.+1;
	if pEik.Mesh.dim==2
		rowLength = 5;
		if pEik.HO
			getRowOfSensitivity = getRowOfSensitivity2D;
		else
			getRowOfSensitivity = getRowOfSensitivity2DFO;
		end
		srcCS = loc2cs(pEik.src[1],pEik.src[2],n);
		offset = [1;n[1]];
	else
		rowLength = 7;
		getRowOfSensitivity = getRowOfSensitivity3D;
		srcCS = loc2cs3D(pEik.src[1],pEik.src[2],pEik.src[3],n);
		offset = [1;n[1];n[1]*n[2]];
	end
	G0 = zeros(pEik.Mesh.dim);
	Idxs = zeros(Int64,pEik.Mesh.dim);
	colIdxs = zeros(Int64,rowLength);
	vals 	= zeros(Float64,rowLength);
	############ Invert a lower triangular Matrix ##########################
	res[:] .= 0.0;
	for j=order
		r = b[j];
		rowLength_j = getRowOfSensitivity(pEik,j,vals,colIdxs,srcCS,offset,G0,Idxs,n);
		for i = 1:rowLength_j
			r -= vals[i]*res[colIdxs[i]];
		end
		res[j] = r/vals[1];
	end
	if pEik.Mesh.dim==2
		selfMultiplyWithAnalyticSolution2D(n, pEik.Mesh.h,pEik.src,res);
	else
		selfMultiplyWithAnalyticSolution3D(n, pEik.Mesh.h,pEik.src,res);
	end

	return res;
end

function getSensTMatVecEik(b::Array{Float64,1},res::Array{Float64,1},pEik::EikonalParam)
# This is basically an upper triangular solution. This function ruins b.
	order = pEik.ordering;
	n = pEik.Mesh.n.+1;
	if pEik.Mesh.dim==2
		rowLength = 5;
		if pEik.HO
			getRowOfSensitivity = getRowOfSensitivity2D;
		else
			getRowOfSensitivity = getRowOfSensitivity2DFO;
		end
		srcCS = loc2cs(pEik.src[1],pEik.src[2],n);
		offset = [1;n[1]];
	else
		rowLength = 7;
		getRowOfSensitivity = getRowOfSensitivity3D;
		srcCS = loc2cs3D(pEik.src[1],pEik.src[2],pEik.src[3],n);
		offset = [1;n[1];n[1]*n[2]];
	end
	G0 = zeros(pEik.Mesh.dim);
	Idxs = zeros(Int64,pEik.Mesh.dim);
	colIdxs = zeros(Int64,rowLength);
	vals 	= zeros(Float64,rowLength);

	############ Invert an upper triangular Matrix ##########################
	N = length(b)
	res[:] = b[:]; # up to index i (backwards, starting from n) res holds the result. From that to 1 it holds the residual.
	if pEik.Mesh.dim==2
		selfMultiplyWithAnalyticSolution2D(n, pEik.Mesh.h,pEik.src,res);
	else
		selfMultiplyWithAnalyticSolution3D(n, pEik.Mesh.h,pEik.src,res);
	end

	for k=0:(N-1)
		j = order[N-k];
		rowLength_j = getRowOfSensitivity(pEik,j,vals,colIdxs,srcCS,offset,G0,Idxs,n);
		d = res[colIdxs[1]]/vals[1];
		res[colIdxs[1]] = d;
		for i = 2:rowLength_j
			res[colIdxs[i]] -= vals[i]*d;
		end
	end
	return res;
end

function getRowOfSensitivity3D(pEik::EikonalParam,rowIdx::EikIdxType,vals::Array{Float64},colIdxs::Array{Int64},srcCS::Int64,offsetVec::Array{Int64,1},
								G0::Array{Float64,1},Idxs::Array{Int64,1},n::Array{Int64,1});
# the ordering is : [diag,1st cell dx,2nd cell dx,1st cell dy,2nd cell dy]
# in 2D the length of the rows should be of length 5 at most.

OP = pEik.OP;

h = pEik.Mesh.h;
src = pEik.src;
rowIdx = convert(Int64,rowIdx);
colIdxs[1] = rowIdx;
T1 = pEik.T1;
if srcCS==rowIdx
	vals[1] = 2.0*T1[rowIdx];
	return 1;
end

vals[1] = 0.0;
cs2loc3D(Idxs,rowIdx,n);
@inbounds T0loc = analyticLocal3D(h,src,Idxs[1],Idxs[2],Idxs[3],G0);

k = 1;

for dim = 1:3
	hinvij  = T0loc./h[dim];
	hinv2ij = 0.5*hinvij;
	op = getOp(OP[rowIdx],dim);
	offset = offsetVec[dim];
	if op == 1 # forward 1
		diagDhat  = -hinvij + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] + hinvij*T1[rowIdx+offset]);
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = hinvij*diagSensValue;
		colIdxs[k]  = rowIdx + offset;
	elseif op == -1 # backward 1
		diagDhat  = hinvij + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] - hinvij*T1[rowIdx - offset]);
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = -hinvij*diagSensValue;
		colIdxs[k]  = rowIdx - offset;
	elseif op == -2 # backward 2
		diagDhat  = hinv2ij*3.0 + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] + hinv2ij*(-4.0*T1[rowIdx-offset] + T1[rowIdx-2*offset]));
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = -hinv2ij*4.0*diagSensValue;
		colIdxs[k]  = rowIdx - offset;
		k+=1;
		vals[k]    = hinv2ij*diagSensValue;
		colIdxs[k] = rowIdx - 2*offset;
	elseif op == 2 # forward 2
		diagDhat  = -hinv2ij*3.0 + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] + hinv2ij*(4.0*T1[rowIdx+offset] - T1[rowIdx+2*offset]));
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = hinv2ij*4.0*diagSensValue;
		colIdxs[k]  = rowIdx + offset;
		k+=1;
		vals[k]    = -hinv2ij*diagSensValue;
		colIdxs[k] = rowIdx + 2*offset;
	end
end
return k;
end

function getRowOfSensitivity2D(pEik::EikonalParam,rowIdx::EikIdxType,vals::Array{Float64},colIdxs::Array{Int64},srcCS::Int64,offsetVec::Array{Int64,1},
								G0::Array{Float64,1},Idxs::Array{Int64,1},n::Array{Int64,1});
# the ordering is : [diag,1st cell dx,2nd cell dx,1st cell dy,2nd cell dy]
# in 2D the length of the rows should be of length 5 at most.

OP = pEik.OP;

h = pEik.Mesh.h;
src = pEik.src;

rowIdx = convert(Int64,rowIdx);
colIdxs[1] = rowIdx;
T1 = pEik.T1;
if srcCS==rowIdx
	vals[1]    = 2.0*T1[rowIdx];
	return 1;
end
vals[1] = 0.0;

cs2loc(Idxs,rowIdx,n);
T0loc = analyticLocal(h,src,Idxs[1],Idxs[2],G0);

k = 1;

for dim = 1:2
	hinvij  = T0loc./h[dim];
	hinv2ij = 0.5*hinvij;
	op = getOp(OP[rowIdx],dim);
	offset = offsetVec[dim];
	if op == 1 # forward 1
		diagDhat  = -hinvij + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] + hinvij*T1[rowIdx+offset]);
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = hinvij*diagSensValue;
		colIdxs[k]  = rowIdx + offset;
	elseif op == -1 # backward 1
		diagDhat  = hinvij + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] - hinvij*T1[rowIdx - offset]);
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = -hinvij*diagSensValue;
		colIdxs[k]  = rowIdx - offset;
	elseif op == -2 # backward 2
		diagDhat  = hinv2ij*3.0 + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] + hinv2ij*(-4.0*T1[rowIdx-offset] + T1[rowIdx-2*offset]));
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = -hinv2ij*4.0*diagSensValue;
		colIdxs[k]  = rowIdx - offset;
		k+=1;
		vals[k]    = hinv2ij*diagSensValue;
		colIdxs[k] = rowIdx - 2*offset;
	elseif op == 2 # forward 2
		diagDhat  = -hinv2ij*3.0 + G0[dim];
		diagSensValue = 2.0*(diagDhat*T1[rowIdx] + hinv2ij*(4.0*T1[rowIdx+offset] - T1[rowIdx+2*offset]));
		vals[1]                 += diagDhat*diagSensValue;
		k+=1;
		vals[k]     = hinv2ij*4.0*diagSensValue;
		colIdxs[k]  = rowIdx + offset;
		k+=1;
		vals[k]    = -hinv2ij*diagSensValue;
		colIdxs[k] = rowIdx + 2*offset;
	end
end
return k;
end

function getRowOfSensitivity2DFO(pEik::EikonalParam,rowIdx::EikIdxType,vals::Array{Float64},colIdxs::Array{Int64},srcCS::Int64,offsetVec::Array{Int64,1},
								G0::Array{Float64,1},Idxs::Array{Int64,1},n::Array{Int64,1});
# the ordering is : [diag,1st cell dx,2nd cell dx,1st cell dy,2nd cell dy]
# in 2D the length of the rows should be of length 5 at most.

OP  = pEik.OP;
h   = pEik.Mesh.h;
src = pEik.src;
T1  = pEik.T1;
rowIdx = convert(Int64,rowIdx);

@inbounds colIdxs[1] = rowIdx;

if srcCS==rowIdx
	@inbounds vals[1]    = 2.0*T1[rowIdx];
	return 1;
end
@inbounds vals[1] = 0.0;

cs2loc(Idxs,rowIdx,n);

@inbounds T0loc = analyticLocal(h,src,Idxs[1],Idxs[2],G0);

k = 1;

for dim = 1:2
	@inbounds op = getOp(OP[rowIdx],dim);
	if op != 0
		@inbounds firstNeighbor = rowIdx + op*offsetVec[dim];
		@inbounds hinvij  = op*(T0loc./h[dim]);

		@inbounds diagDhat  = G0[dim]-hinvij;
		@inbounds diagSensValue = 2.0*(diagDhat*T1[rowIdx] + hinvij*T1[firstNeighbor]);
		@inbounds vals[1]     += diagDhat*diagSensValue;
		k+=1;
		@inbounds vals[k]     = hinvij*diagSensValue;
		@inbounds colIdxs[k]  = firstNeighbor;
	end
end
return k;
end


#include("HPDE_options.jl")
##########################################################

mutable struct FVMPDEProblem
    grid::FVMPDEGrid

    """
    Flux function F in x direction

    """
    F::Function

    """
    Flux function G in y direction
    !!! compat
        Will be added in future versions
    """
    G::Function

    """
    Flux function H in z direction
    !!! compat
        Will be added in future versions
    """
    H::Function

    """
    Source function
    !!! compat
        Will be added in future versions
    """
    S::Function

    nvars::Int8

    param::Any

    U0

    U_min
    U_max

    jac
    ode_problem_kwargs
    tspan
    t

    function FVMPDEProblem(grid::FVMPDEGrid, tspan, U0; kwargs...)
        if haskey(kwargs,:jacobian)
            jac = kwargs[:jacobian]
            kwargs = Dict([p for p in pairs(kwargs) if p[1] != :jacobian])
        else
            jac = true
        end
        zeroFunc(U) = zeros(size(U))
        F = kwargs[:F]
        @assert size(U0,1) == grid.nx
        if haskey(kwargs,:G)
            G = kwargs[:G]
            kwargs = Dict([p for p in pairs(kwargs) if p[1] != :G])
            @assert size(U0,2) == grid.ny
            if haskey(kwargs,:H)
                H = kwargs[:H]
                kwargs = Dict([p for p in pairs(kwargs) if p[1] != :H])
                nvars = size(U0,4)
                @assert grid.dimension == _3D
                @assert length(size(U0))-1 == 3
                @assert size(U0,3) == grid.nz
            else
                H = zeroFunc
                nvars = size(U0,3)
                @assert grid.dimension == _2D
                @assert length(size(U0))-1 == 2
            end
        else
            @assert !haskey(kwargs,:H)
            G = zeroFunc
            H = zeroFunc
            nvars = size(U0,2)
            @assert grid.dimension == _1D
            @assert length(size(U0))-1 == 1
        end

        if haskey(kwargs,:source)
            S = kwargs[:source]
            kwargs = Dict([p for p in pairs(kwargs) if p[1] != :source])
        else
            S = zeroFunc
        end
        if jac
            @assert size(F(U0[grid.indices[1],:])) == size(U0[grid.indices[1],:])
            @assert size(G(U0[grid.indices[1],:])) == size(U0[grid.indices[1],:])
            @assert size(H(U0[grid.indices[1],:])) == size(U0[grid.indices[1],:])
        else
            @assert size(F(U0[grid.indices[1],:]),1) == length(U0[grid.indices[1],:])
            @assert size(F(U0[grid.indices[1],:]),2) == length(U0[grid.indices[1],:])
            # TODO: add @assert for G and H
        end
        @assert size(S(U0)) == size(U0)

        #grid.U = reshape(U0,grid.nx,grid.ny,grid.nz,nvars)
        #grid.U = copy(U0)
        if haskey(kwargs,:parameters)
            θ = kwargs[:parameters]
            kwargs = Dict([p for p in pairs(kwargs) if p[1] != :parameters])
        else
            θ = nothing
        end
        if haskey(kwargs,:ul)
            U_min = kwargs[:ul]
            kwargs = Dict([p for p in pairs(kwargs) if p[1] != :ul])
        else
            U_min = -Inf * ones(size(U0))
        end
        if haskey(kwargs,:uh)
            U_max = kwargs[:uh]
            kwargs = Dict([p for p in pairs(kwargs) if p[1] != :uh])
        else
            U_max = Inf * ones(size(U0))
        end

        @assert size(U_min) == size(U0)
        @assert size(U_max) == size(U0)

        #@assert checkDim(D) == size(U0,)

        return new(grid,F, G, H, S, nvars, θ, U0, U_min, U_max,jac,kwargs,tspan,0)
    end

end


"""
    Tetrahedron3

An object `t` of type [`Tetrahedron3`](@ref) is an oriented tetrahedron in the
three-dimensional Euclidean space ``Ε^3``.

It is defined by four vertices ``p_0``, ``p_1``, ``p_2`` and ``p_3``. The
orientation of a tetrahedron is the orientation of its four vertices. That means
it is positive when ``p_3`` is on the positive side of the plane defined by
``p_0``, ``p_1`` and ``p_2``.

The tetrahedron itself splits the space ``Ε^3`` in a *positive* and a *negative*
side.

The boundary of a tetrahedron splits the space in two open regions, a bounded
one and an unbounded one.
"""
Tetrahedron3

"""
    Tetrahedron3(p₀::Point3, p₁::Point3, p₂::Point3, p₃::Point3)

Introduces a tetrahedron `t` with vertices `p₀`, `p₁`, `p₂` and `p₃`.
"""
Tetrahedron3(p₀::Point3, p₁::Point3, p₂::Point3, p₃::Point3)

"""
    ==(t₁::Tetrahedron3, t₂::Tetrahedron3)

Test for equality: two tetrahedra `t₁` and `t₂` are equal, iff `t₁` and `t₂`
have the same orientation and their sets (not sequences) of vertices are equal.
"""
==(t₁::Tetrahedron3, t₂::Tetrahedron3)

"""
    vertex(t::Tetrahedron3, i::Integer)

Returns the `i`ᵗʰ vertex modulo 4 of `t`.
"""
vertex(t::Tetrahedron3, i::Integer)

"""
    is_degenerate(t::Tetrahedron3)

Tetrahedron `t` is degenerate, if the vertices are coplanar.
"""
is_degenerate(t::Tetrahedron3)

"""
    orientation(t::Tetrahedron3)

Returns the orientation of `t`.
"""
orientation(t::Tetrahedron3)

"""
    oriented_side(t::Tetrahedron3, p::Point2)

Returns [`ON_ORIENTED_BOUNDARY`](@ref), or [`POSITIVE_SIDE`](@ref), or the
constant [`ON_NEGATIVE_SIDE`](@ref), determined by the position of point `p`.

!!! info "Precondition"

    `t` is not degenerate.
"""
oriented_side(t::Tetrahedron3, p::Point2)

"""
    bounded_side(t::Tetrahedron3, p::Point2)

Returns [`ON_BOUNDARY`](@ref), [`ON_BOUNDED_SIDE`](@ref), or else
[`ON_UNBOUNDED_SIDE`](@ref), depending on where point `p` is.

!!! info "Precondition"

    `t` is not degenerate.
"""
bounded_side(t::Tetrahedron3, p::Point2)

"""
    has_on_positive_side(t::Tetrahedron3, p::Point2)

Returns `true`, iff `p` lies on the positive side of `t`.
"""
has_on_positive_side(t::Tetrahedron3, p::Point2)

"""
    has_on_negative_side(t::Tetrahedron3, p::Point2)

Returns `true`, iff `p` lies on the negative side of `t`.
"""
has_on_negative_side(t::Tetrahedron3, p::Point2)

"""
    has_on_boundary(t::Tetrahedron3, p::Point2)

Returns `true`, iff `p` lies on the boundary of `t`.
"""
has_on_boundary(t::Tetrahedron3, p::Point2)

"""
    has_on_bounded_side(t::Tetrahedron3, p::Point2)

Returns `true`, iff `p` lies properly inside of `t`.
"""
has_on_bounded_side(t::Tetrahedron3, p::Point2)

"""
    has_on_unbounded_side(t::Tetrahedron3, p::Point2)

Returns `true`, iff `p` lies properly outside of `t`.
"""
has_on_unbounded_side(t::Tetrahedron3, p::Point2)

"""
    volume(t::Tetrahedron3)

Returns the signed volume of `t`.
"""
volume(t::Tetrahedron3)

"""
    bbox(t::Tetrahedron3)

Returns a bounding box containing `t`.
"""
bbox(t::Tetrahedron3)

"""
    transform(t::Tetrahedron3, at::AffTransformation3)

Returns the tetrahedron obtained by applying `at` on the four vertices of `t`.
"""
transform(t::Tetrahedron3, at::AffTransformation3)


"""
    li6(z::ComplexF64)::ComplexF64

Returns the complex 6th order polylogarithm of a complex number `z` of type `ComplexF64`.

Author: Alexander Voigt

License: MIT

# Example
```julia
li6(1.0 + 1.0im)
```
"""
function li6(z::ComplexF64)::ComplexF64
    function clog(z)
        az::Float64 = angle(z)
        return 0.5*log(abs2(z)) + (imag(z) == 0.0 && az < 0.0 ? -az : az)*1.0im
    end

    z6::Float64 = 1.0173430619844491
    bf = (
        1.0                   , -31.0/64.0             ,
        1.5241340877914952e-01, -3.4365555877057613e-02,
        5.7174797239368999e-03, -6.8180453746570645e-04,
        4.9960361948734493e-05, -4.9166051196039048e-07,
       -3.0632975161302164e-07,  1.4414599270849095e-08,
        3.7272438230924107e-09, -3.7300867345487607e-10,
       -5.1246526816085832e-11,  9.0541930956636683e-12,
        6.7381882615512517e-13, -2.1215831150303135e-13,
       -6.8408811719011698e-15,  4.8691178462005581e-15
    )

    if imag(z) == 0.0
        if real(z) == 0.0
            return 0.0 + 0.0im
        end
        if real(z) == 1.0
            return z6 + 0.0im
        end
        if real(z) == -1.0
            return -31.0/32.0*z6 + 0.0im
        end
    end

    nz::Float64  = abs2(z)
    pz::Float64  = angle(z)
    lnz::Float64 = 0.5*log(nz)

    if lnz*lnz + pz*pz < 1.0 # |log(z)| < 1
        v::ComplexF64  = lnz + pz*im
        v2::ComplexF64 = v*v
        c1::Float64 = 1.0369277551433699 # zeta(5)
        c2::Float64 = 0.54116161685556910
        c3::Float64 = 0.20034281719326571
        c4::Float64 = 0.068538919452009435
        c5::ComplexF64 = (137.0/60.0 - clog(-v))/120.0
        c6::Float64 = -1.0/1440.0

        cs = (
            -1.6534391534391534e-05, 2.2964432686654909e-08,
            -9.9413128513657614e-11, 6.6912682653423394e-13,
            -5.7933058574392549e-15
        )

        return z6 + v * c1 +
            v2 * (c2 + v * c3 +
            v2 * (c4 + v * c5 +
            v2 * (c6 +
            v * (cs[1] +
            v2 * (cs[2] +
            v2 * (cs[3] +
            v2 * (cs[4] +
            v2 * (cs[5]))))))))
    end

    (u::ComplexF64, rest::ComplexF64, sgn::Float64) = if nz <= 1.0
        (-clog(1.0 - z), 0.0 + 0.0im, 1.0)
    else # nz > 1.0
        arg::Float64 = pz > 0.0 ? pz - pi : pz + pi
        lmz::ComplexF64 = lnz + arg*im # clog(z)
        lmz2::ComplexF64 = lmz*lmz
        (-clog(1.0 - 1.0/z), -31.0*pi^6/15120.0 + lmz2*(-7.0/720.0*pi^4 + lmz2*(-1.0/144.0*pi^2 - 1.0/720.0*lmz2)), -1.0)
    end

    u2::ComplexF64 = u*u
    u4::ComplexF64 = u2*u2
    u8::ComplexF64 = u4*u4

    rest + sgn * (
       u*bf[1] +
       u2*(bf[2] + u*bf[3]) +
       u4*(bf[4] + u*bf[5] + u2*(bf[6] + u*bf[7])) +
       u8*(bf[8] + u*bf[9] + u2*(bf[10] + u*bf[11]) +
           u4*(bf[12] + u*bf[13] + u2*(bf[14] + u*bf[15]))) +
       u8*u8*(bf[16] + u*bf[17] + u2*bf[18])
    )
end


using Ipopt, ForwardDiff
import Vecchia: ipopt_hessian

include("example_setup.jl")

BLAS.set_num_threads(6)

function allocating_negloglik(kfun, params, pts, vals)
  K = cholesky!(Symmetric([kfun(x,y,params) for x in pts, y in pts]))
  0.5*logdet(K) + 0.5*dot(vals, K\vals)
end

# nll and derivatives:
_enll(p) = allocating_negloglik(kfn, p, pts, sim)
egrad!(p, store) = ForwardDiff.gradient!(store, _enll, p)
ehess(p) = ForwardDiff.hessian(_enll, p)

# Set up and solve the problem in Ipopt:
exact_prob = CreateIpoptProblem(
               2,ones(2).*1.0e-3, ones(2).*100.0,
               0,Float64[], Float64[],
               0,3,
               _enll,
               (args...)->nothing,
               egrad!,
               (args...)->nothing,
               (x,r,c,o,l,v)->ipopt_hessian(x,r,c,o,l,v,ehess,Function[],0))
AddIpoptNumOption(exact_prob, "tol",      1.0e-5)
AddIpoptIntOption(exact_prob, "max_iter", 1000)
exact_prob.x = ones(2)
@time estatus = IpoptSolve(exact_prob)
@assert iszero(estatus) "Optimization for exact likelihood problem failed."

# Here are your MLE and observed information matrix:
exact_mle = copy(exact_prob.x)
exact_obs_inf_mat = ehess(exact_mle)


# ? activate Splines
import Pkg
Pkg.activate("Splines")

using PlotlyJS
using Revise
Revise.revise()

using Splines

@info "Computing"
n, m = 8, 8

# cupula
N1 = [
        cat([0:n-1 zeros(n) 4 .* sin.(range(0, stop=π, length=n))], dims=2)',
        cat([zeros(m) 0:m-1 4 .* sin.(range(0, stop=π, length=n))], dims=2)'
    ]

# angled plane
N2 = [
        cat([1:n zeros(n) zeros(n)], dims=2)',
        cat([zeros(m) 1:m 1:m], dims=2)'
     ]

# flat plane
N3 = [
        cat([0:n-1 ones(n).-1 zeros(n)], dims=2)',
        cat([ones(m).-1 0:m-1 zeros(m)], dims=2)'
     ]

# funny shape
N4 = [
        [[-1, -3, 1] [1, 3, -1] [2, 6, -3] [3, 2, -5] [4, -1, 1] [5, 1, 3] [6, 3, 7]],
        [1:7 1:7 1:7]'
     ]

fig = make_subplots(
    rows=2, cols=2,
    specs=fill(Spec(kind="scene"), 2, 2),
)
rows = [1, 1, 2, 2]
cols = [1, 2, 1, 2]

for (n, (nodes_1, nodes_2)) in enumerate([N1, N2, N3, N4])
    # compute beziers
    b1 = Bezier(nodes_1)
    b2 = Bezier(nodes_2)
    bb = b1 ⊗ b2

    # compute and plot points on the bezier curve
    points = hcat(b1(.5), b2(.5), bb(.5, .5))
    add_trace!(fig,
    plot_nodes(points, color="black", node_size=8), row=rows[n], col=cols[n]
    )

    # plot knots
    for (nodes, color) in zip((b1.nodes, b2.nodes, bb.nodes), ("red", "blue", "black"))
        add_trace!(fig,
            plot_nodes(nodes; color=color, node_size=1), row=rows[n], col=cols[n]
        )
    end

    # plot curves
    for (b, color) in zip((b1, b2), ("red", "blue"))
        add_trace!(fig,
            plot_curve(b.coordinates, color=color), row=rows[n], col=cols[n])
    end

    # plot surface
    add_trace!(
        fig,
        plot_surface(bb.coordinates, opacity=1)[1],
        row=rows[n], col=cols[n]
    )

end

display(fig)
@info "Done"


export compute_stress_rivlin_saunders, compute_maxwell_viscoelasticity_matrices

@doc raw"""
    compute_stress_rivlin_saunders(strain::Union{PyObject, Array{Float64,2}},c1::Union{PyObject, Float64},c2::Union{PyObject, Float64})

Computes the stress using the plane stress incompressible Rivlin Saunders model.
"""
function compute_stress_rivlin_saunders(strain::Union{PyObject, Array{Float64,2}},c1::Union{PyObject, Float64},c2::Union{PyObject, Float64})
    rivlin_saunders_ = load_op_and_grad("$(@__DIR__)/../../deps/CustomOp/build/libDataLib","rivlin_saunders")
    strain,c1,c2 = convert_to_tensor([strain,c1,c2], [Float64,Float64,Float64])
    rivlin_saunders_(strain,c1,c2)
end

@doc raw"""
    compute_maxwell_viscoelasticity_matrices(mu::Union{PyObject, Array{Float64,1}},eta::Union{PyObject, Array{Float64,1}},lambda::Union{PyObject, Array{Float64,1}},dt::Union{PyObject, Float64})

Computes the viscoelasticity matrix in the Maxwell model. `mu`, `lambda` are Lamé parameters, and `eta` is the viscosity parameter.
These three inputs are all length $N$ vectors, where $N$ is the total number of Gauss points.
`dt` is the  time step.

The time advancing formula is

$$\sigma^{n+1} = H(\epsilon^{n+1}-\epsilon^n) + S\sigma^n$$

This function returns `S`, `H`, and both are $N\times 3\times 3$ tensors.
"""
function compute_maxwell_viscoelasticity_matrices(mu::Union{PyObject, Array{Float64,1}},lambda::Union{PyObject, Array{Float64,1}},eta::Union{PyObject, Array{Float64,1}},dt::Union{PyObject, Float64})
    viscoelasticity_matrices_ = load_op_and_grad("$(@__DIR__)/../../deps/CustomOp/build/libDataLib","viscoelasticity_matrices", multiple=true)
    mu,eta,lambda,dt = convert_to_tensor(Any[mu,eta,lambda,dt], [Float64,Float64,Float64,Float64])
    viscoelasticity_matrices_(mu,eta,lambda,dt)
end


function UndoShift(AShot,datainfo)
	### Input:
	### AShot: one shot gather (nt,nrx,nry) with padding zeros
	### datainfo: a Dictionary contains info for one shot gather
	### Output:
	### A shot gather w/o linear shift
	nrx = datainfo["nrx"];
	nry = datainfo["nry"];
	Rgrid = datainfo["Rgrid"];
	dt = datainfo["dt"];

	vel = 1415; #velocity sand(dry) 0.2-1 Km/s, sand(water saturated 1.5-2 Km/s)

	# find the center position of receiver = source position
	CentX = (nrx+1)/2;
	CentY = (nry+1)/2;

	# calculate the shift grid for each receiver and save them in the matrix shiftT
	a = 1:nrx;
	b = 1:nry;
	row = reshape(repeat(a,outer=nry),nrx,nry);
	col = reshape(repeat(b,inner=nry),nrx,nry);
	shiftT = round.(sqrt.((row .-CentX).^2 .+ (col .-CentY).^2) .*(Rgrid/vel/(dt*1e-3)));

	nt1 =  datainfo["nt"] + 300 + 100;
	Shift_Back_Data = zeros(size(AShot));
	for i = 1:nrx
	    for j = 1:nry
		Shift_Back_Data[Int(shiftT[i,j]+1):nt1,i,j] = AShot[1:Int(nt1-shiftT[i,j]),i,j];
	    end
	end

	return Shift_Back_Data[300+1:300+datainfo["nt"],:,:]
end


"""
    tvdiff(data::AbstractVector, iter::Integer, α::Real; kwargs...)

# Arguments

  - `data::AbstractVector`:
    Vector of data to be differentiated.

  - `iter::Integer`:
    Number of iterations to run the main loop.  A stopping
    condition based on the norm of the gradient vector `g`
    below would be an easy modification.
  - `α::Real`:
    Regularization parameter.  This is the main parameter
    to fiddle with.  Start by varying by orders of
    magnitude until reasonable results are obtained.  A
    value to the nearest power of 10 is usally adequate.
    Higher values increase regularization strenght
    and improve conditioning.

## Keywords

  - `u_0::AbstractVector`:
    Initialization of the iteration.  Default value is the
    naive derivative (without scaling), of appropriate
    length (this being different for the two methods).
    Although the solution is theoretically independent of
    the intialization, a poor choice can exacerbate
    conditioning issues when the linear system is solved.

  - `scale::String`:
    Scale of dataset, `\"large\"` or `\"small\"` (case insensitive).
    Default is `\"small\"` .  `\"small\"`  has somewhat better
    boundary behavior, but becomes unwieldly for very large datasets.
    `\"large\"` has simpler numerics but
    is more efficient for large-scale problems.  `\"large\"` is
    more readily modified for higher-order derivatives,
    since the implicit differentiation matrix is square.
  - `ε::Real`:
    Parameter for avoiding division by zero.  Default value
    is `1e-6`.  Results should not be very sensitive to the
    value.  Larger values improve conditioning and
    therefore speed, while smaller values give more
    accurate results with sharper jumps.
  - `dx::Real`:
    Grid spacing, used in the definition of the derivative
    operators.  Default is `1 / length(data)`.
  - `precond::String`:
    Select the preconditioner for the conjugate gradient method.
    Default is `\"none\"`.

      + `scale = \"small\"`:
        While in principle `precond=\"simple\"` should speed things up,
        sometimes the preconditioner can cause convergence problems instead,
        and should be left to `\"none\"`.
      + `scale = \"large\"`:
        The improved preconditioners are one of the main features of the
        algorithm, therefore using the default `\"none\"` is discouraged.
        Currently, `\"diagonal\"`,`\"amg_rs\"`,`\"amg_sa\"` are available.
  - `diff_kernel::String`:
    Kernel to use in the integral to smooth the derivative. By default it is set to
    `\"abs\"`, the absolute value ``|u'|``. However, it can be changed to `\"square\"`,
    the square value ``(u')^2``. The latter produces smoother
    derivatives, whereas the absolute values tends to make them more blocky.
  - `cg_tol::Real`:
    Relative tolerance used in conjugate gradient method. Default is `1e-6`.
  - `cgmaxit::Int`:
    Maximum number of iterations to use in conjugate gradient optimisation.
    Default is `100`.
  - `show_diagn::Bool`:
    Flag whether to display diagnostics at each iteration. Default is `false`.
    Useful for diagnosing preconditioning problems. When tolerance is not met,
    an early iterate being best is more worrying than a large relative residual.

# Output

  - `u`:
    Estimate of the regularized derivative of data with
    `length(u) = length(data)`.
"""
function tvdiff(
    data::AbstractVector,
    iter::Integer,
    α::Real;
    u_0::AbstractVector=[NaN],
    scale::String="small",
    ε::Real=1e-6,
    dx::Real=NaN,
    precond::String="none",
    diff_kernel::String="abs",
    cg_tol::Real=1e-6,
    cg_maxiter::Integer=100,
    show_diagn::Bool=false,
)::AbstractVector
    n = length(data)
    if isnan(dx)
        dx = 1 / (n - 1)
    end

    # Make string inputs case insensitive
    scale = lowercase(scale)
    precond = lowercase(precond)
    diff_kernel = lowercase(diff_kernel)

    # Run tvdiff for selected method
    return tvdiff(
        Val(Symbol(scale)),
        data,
        iter,
        α,
        u_0,
        ε,
        dx,
        cg_tol,
        cg_maxiter,
        precond,
        diff_kernel,
        show_diagn,
    )
end

# Total variance regularized numerical differences for small scale problems.
function tvdiff(
    ::Val{:small},
    data::AbstractVector,
    iter::Integer,
    α::Real,
    u_0::AbstractVector,
    ε::Real,
    dx::Real,
    cg_tol::Real,
    cg_maxiter::Integer,
    precond::String,
    diff_kernel::String,
    show_diagn::Bool,
)::AbstractVector
    n = length(data)

    #= Assert initialization if provided, otherwise set
    default initization to naive derivative =#
    if isequal(u_0, [NaN])
        u_0 = [0; diff(data); 0] / dx
    elseif length(u_0) != (n + 1)
        throw(
            DimensionMismatch(
                "size $(size(u_0)) of u_0 doesn't match size ($(n + 1),) required for scale=\"small\".",
            ),
        )
    end
    u = copy(u_0)

    # Construct differentiation matrix.
    D = spdiagm(n, n + 1, 0 => -ones(n), 1 => ones(n)) / dx
    Dᵀ = transpose(D)

    # Construct antidifferentiation operator and its adjoint.
    function A(x)
        return (cumsum(x) - 0.5 * (x .- x[1]))[2:end] * dx
    end

    function Aᵀ(x)
        return [sum(x) / 2; (sum(x) .- cumsum(x) .- x / 2)] * dx
    end

    # Precompute antidifferentiation adjoint on data
    # Since A([0]) = 0, we need to adjust.
    offset = data[1]
    Aᵀb = Aᵀ(offset .- data)

    for i in 1:iter
        if diff_kernel == "abs"
            # Diagonal matrix of weights, for linearizing E-L equation.
            Q = Diagonal(1 ./ sqrt.((D * u) .^ 2 .+ ε))
            # Linearized diffusion matrix, also approximation of Hessian.
            L = dx * Dᵀ * Q * D
        elseif diff_kernel == "square"
            L = dx * Dᵀ * D
        else
            throw(
                ArgumentError(
                    """unexpected diff_kernel "$(diff_kernel)" for scale="small"."""
                ),
            )
        end

        # Gradient of functional.
        g = Aᵀ(A(u)) + Aᵀb + α * L * u

        # Select preconditioner.
        if precond == "simple"
            P = Diagonal(α * diag(L) .+ 1)
        elseif precond == "none"
            P = I # Identity matrix
        else
            throw(ArgumentError("""unexpected precond "$(precond)" for scale="small"."""))
        end

        # Prepare linear operator for linear equation.
        # Approximation of Hessian of TVR functional at u
        H = LinearMap(u -> Aᵀ(A(u)) + α * L * u, n + 1, n + 1)

        # Solve linear equation.
        s = cg(H, -g; Pl=P, reltol=cg_tol, maxiter=cg_maxiter)

        show_diagn && println(
            "Iteration $(i):\trel. change = $(norm(s) / norm(u)),\tgradient norm = $(norm(g))",
        )

        # Update current solution
        u += s
    end
    return u[1:(end - 1)]
end

# Total variance regularized numerical differences for large scale problems.
function tvdiff(
    ::Val{:large},
    data::AbstractVector,
    iter::Integer,
    α::Real,
    u_0::AbstractVector,
    ε::Real,
    dx::Real,
    cg_tol::Real,
    cg_maxiter::Integer,
    precond::String,
    diff_kernel::String,
    show_diagn::Bool,
)::AbstractVector
    n = length(data)

    # Since Au(0) = 0, we need to adjust.
    data = data .- data[1]

    #= Assert initialization if provided, otherwise set
    default initization to naive derivative =#
    if isequal(u_0, [NaN])
        u_0 = [0; diff(data)] / dx
    elseif length(u_0) != n
        throw(
            DimensionMismatch(
                "size $(size(u_0)) of u_0 doesn't match size ($(n),) required for scale=\"large\".",
            ),
        )
    end
    u = deepcopy(u_0) * dx

    # Construct differentiation matrix.
    D = spdiagm(n, n, 0 => -ones(n - 1), 1 => ones(n - 1)) / dx
    Dᵀ = transpose(D)

    # Construct antidifferentiation operator and its adjoint.
    A(x) = cumsum(x)

    function Aᵀ(x)
        return sum(x) .- [0; cumsum(x[1:(end - 1)])]
    end

    # Precompute antidifferentiation adjoint on data
    Aᵀd = Aᵀ(data)

    for i in 1:iter
        if diff_kernel == "abs"
            # Diagonal matrix of weights, for linearizing E-L equation.
            Q = Diagonal(1 ./ sqrt.((D * u) .^ 2 .+ ε))
            # Linearized diffusion matrix, also approximation of Hessian.
            L = Dᵀ * Q * D
        elseif diff_kernel == "square"
            L = Dᵀ * D
        else
            throw(
                ArgumentError(
                    """unexpected diff_kernel "$(diff_kernel)" for scale="large"."""
                ),
            )
        end

        # Gradient of functional.
        g = Aᵀ(A(u)) - Aᵀd + α * L * u

        # Select preconditioner.
        B = α * L + Diagonal(reverse(cumsum(n:-1:1)))

        if precond == "diagonal"
            P = DiagonalPreconditioner(B)
        elseif precond == "amg_rs"
            # Ruge-Stuben variant
            P = AMGPreconditioner{RugeStuben}(B)
        elseif precond == "amg_sa"
            # Smoothed aggregation
            P = AMGPreconditioner{SmoothedAggregation}(B)
        elseif precond == "none"
            P = I # Identity matrix
        else
            throw(ArgumentError("""unexpected precod "$(precond)" for scale="large"."""))
        end

        # Prepare linear operator for linear equation.
        # Approximation of Hessian of TVR functional at u
        H = LinearMap(u -> Aᵀ(A(u)) + α * L * u, n, n)

        # Solve linear equation.
        s = cg(H, -g; Pl=P, reltol=cg_tol, maxiter=cg_maxiter)

        show_diagn && println(
            "Iteration $(i):\trel. change = $(norm(s) / norm(u)),\tgradient norm = $(norm(g))",
        )

        # Update current solution
        u += s
    end
    return u / dx
end


"""
Quantum Inverse Method

Calculate Hamiltonian from eigen state(s).
"""

@inline inner_product(v1::AbstractVector{<:Number}, v2::AbstractVector{<:Number}) = dot(v1, v2)
@inline inner_product(v1::AbstractMatrix{<:Number}, v2::AbstractMatrix{<:Number}) = dot(v1, v2) / size(v1, 2)

export covmat
function covmat(ol::AbstractVector, v::AbstractVecOrMat{<:Number})
    n = length(ol)
    vs = [oi * v for oi in ol]
    am = [real(inner_product(v, vsi)) for vsi in vs]
    cm = Matrix{Float64}(undef, n, n)
    for i=1:n
        for j=i:n
            cm[i, j] = real(inner_product(vs[i], vs[j])) - am[i] * am[j]
        end
    end
    Hermitian(cm)
end


using OrthogonalPolynomialsQuasi, ContinuumArrays, DifferentialEquations, Plots

###
# Heat
###

S = JacobiWeight(1.0,1.0) .* Jacobi(1.0,1.0)
D = Derivative(axes(S,1))
Δ = -((D*S)'*(D*S))
M = S'S

function evolution(u, (M,A), t)
    M\(A*u)
end
n = 50
u0 = [[1,2,3]; zeros(n-3)]
prob = ODEProblem(evolution,u0,(0.0,3.0),(cholesky(Symmetric(M[1:n,1:n])),Δ[1:n,1:n]))
@time u = solve(prob,TRBDF2()); u(1.0)

###
# Heat natural
###
L = LinearSpline(range(-1,1;length=2))
P = apply(hcat,L,S)
ApplyArray(hcat,Legendre() \ (D*L), Legendre() \ (D*S))

(D*L)'*(D*L)

(D*S)'*(D*S)

((D*S)'*(D*L)).args[3]
(D*L)'*(D*S)
H = (D*L).args[1]
P = Legendre()
(P\H)' * P'P * (P \ (D*S))

###
# Beam
###

# [u, u_t]_t = [0 I; M] * [u, u_t]

S = JacobiWeight(2.0,2.0) .* Jacobi(2.0,2.0)
D = Derivative(axes(S,1))
Δ² = (D^2*S)'*(D^2*S)
M = S'S
function wave(uv, (M,A), t)
    n = size(M,1)
    u,v = uv[1:n],uv[n+1:end]
    [v; M\(A*u)]
end
n = 30
u0 = [1; 2; 3; zeros(n-3)]
prob = ODEProblem(wave,[u0; zeros(n)],(0.0,3.0),(cholesky(Symmetric(M[1:n,1:n])),-(Δ²)[1:n,1:n]))
@time u = solve(prob,TRBDF2()); u(1.0)

g = range(-1,1;length=200)
V = S[g,1:n]
@gif for t in 0.0:0.01:3.0
    plot(g, V*u(t)[1:n]; ylims=(-5,5))
end

prob = ODEProblem(wave,[u0; zeros(n)],(0.0,3.0),(qr(M),-(Δ²)))


using Base.Test
using RandomMatrices

@test Wigner{3} == GaussianHermite{3}

n = 25

for (β, T, N) in [(1, Real, n), (2, Complex, n), (4, Complex, 2n)]
    d = Wigner(β)
    A = rand(d, n)
    @test eltype(A) <: T
    @test size(A) == (N, N)

    At = tridrand(d, n)
    @test eltype(At) <: Real
    @test size(At) == (n, n)

    vals = eigvalrand(d, n)
    @test eltype(vals) <: Real
    @test length(vals) == n

    vd = RandomMatrices.VandermondeDeterminant(vals, β)
    @test isa(vd, Real)

    ed = eigvaljpdf(d, vals)
    @test isa(ed, Real)
end


# Contains evaluation functions
# ----------------------------------
"""
	 Evaluates the FIS

	eval_fis(fis, input_values,defuzz_method = "WTAV")

	 Parameters
	 ----------
	 `fis` is the inference system to evaluate
	 `input_values` is a Vector of inputs
	 `defuzz_method` is the method for defuzzification, see defuzz function definition

"""
function eval_fis(fis::FISMamdani, input_values::Vector{<:AbstractFloat},defuzz_method = "WTAV")

	firing_strengths = AbstractFloat[]
	for rule in fis.rules
		tmp_strengths = AbstractFloat[]
		for i in 1:length(rule.input_mf_names)
			push!(tmp_strengths, fis.input_mfs_dicts[i][rule.input_mf_names[i]].eval(input_values[i]))
		end
			push!(firing_strengths, firing(tmp_strengths,rule.firing_method ))
	end
	defuzz(firing_strengths, fis.rules, fis.output_mfs_dict, defuzz_method)

end

"""
	 Evaluates the FIS

	eval_fis(fis, input_values)

	 Parameters
	 ----------
	 `fis` is the inference system to evaluate
	 `input_values` is a Vector of inputs

"""
function eval_fis(fis::FISSugeno, input_values::Vector{<:AbstractFloat})

	firing_strengths = AbstractFloat[]
	for rule in fis.rules
		tmp_strengths = AbstractFloat[]
		for i in 1:length(rule.input_mf_names)
			push!(tmp_strengths, fis.input_mfs_dicts[i][rule.input_mf_names[i]].eval(input_values[i]))
		end
			push!(firing_strengths, firing(tmp_strengths,rule.firing_method ))
	end

	push!(input_values, 1)
	n_firing_strengths = firing_strengths / sum(firing_strengths)
	out = 0.0
	for i = 1:length(fis.rules)
		out += n_firing_strengths[i] * (fis.rules[i].output_mf' * input_values)[1]
	end
	pop!(input_values)
	return out

end

"""
	 Defuzzifies the output using the given firing strengths

	defuzz(firing_strengths, rules, output_mfs_dict, defuzz_method)

	 Parameters
	 ----------
	 `firing_strengths` is a Vector of firing strengths
	 		one for each output membership function
	 `rules` is a Vector of Rule
	 `output_mfs_dict` is a Dict of output membership functions
	 `defuzz_method` is the method for defuzzification
	 		"MOM" - Mean of Maximum
	 		"WTAV" - Weighted Average
"""
function defuzz(firing_strengths::Vector{AbstractFloat}, rules::Vector{Rule},	output_mfs_dict::Dict{AbstractString, MF}, defuzz_method::AbstractString)

	if defuzz_method == "MOM"
		max_firing_index = indmax(firing_strengths)
		max_fired_mf_name = rules[max_firing_index].output_mf
		output_mfs_dict[max_fired_mf_name].mean_at(maximum(firing_strengths))
	elseif defuzz_method == "WTAV"
		mean_vec = AbstractFloat[]
		for i in 1:length(rules)
			push!(mean_vec, output_mfs_dict[rules[i].output_mf].mean_at(firing_strengths[i]))
		end
		(mean_vec' * firing_strengths)[1] / sum(firing_strengths)
	end

end

"""
	 Docs goes here
"""
function firing(tmp_strengths::Vector{<:AbstractFloat},firing_method::AbstractString)
	if firing_method == "MIN"
		return	minimum_value(tmp_strengths)
	elseif firing_method == "A-PROD"
		return algebraic_product(tmp_strengths)
	elseif firing_method == "B-DIF"
		return bounded_difference(tmp_strengths)
	elseif firing_method == "D-PROD"
		return drastic_product(tmp_strengths)
	elseif firing_method == "E-PROD"
		return einstein_product(tmp_strengths)
	elseif firing_method == "H-PROD"
		return hamacher_product(tmp_strengths)
	elseif firing_method == "MAX"
		return	maximum_value(tmp_strengths)
	elseif firing_method == "A-SUM"
		return algebraic_sum(tmp_strengths)
	elseif firing_method == "B-SUM"
		return bounded_sum(tmp_strengths)
	elseif firing_method == "D-SUM"
		return drastic_sum(tmp_strengths)
	elseif firing_method == "E-SUM"
		return einstein_sum(tmp_strengths)
	elseif firing_method == "H-SUM"
		return hamacher_sum(tmp_strengths)
	end
end


@testset "JJ terms" begin
    @testset "_terms_jw" begin
        function terms_reference(j::HalfInteger, w::Integer)
            w <= 2j+1 || throw(ArgumentError("w=$w too large for $orb orbital"))

            2w ≥ 2j+1 && (w = convert(Int, 2j) + 1 - w)
            w == 0 && return [zero(HalfInt)]
            w == 1 && return [j]

            # Forms full Cartesian product of all mⱼ, not necessarily the most
            # performant method.
            mⱼs = filter(allunique, collect(AtomicLevels.allchoices([-j:j for i = 1:w])))
            MJs = map(x -> reduce(+, x), mⱼs) # TODO: make map(sum, mⱼs) work

            Js = HalfInt[]

            while !isempty(MJs)
                # Identify the maximum MJ and associate it with J.
                MJmax = maximum(MJs)
                N = count(isequal(MJmax), MJs)
                append!(Js, repeat([MJmax], N))
                # Remove the -MJ:MJ series, N times.
                for MJ = -MJmax:MJmax
                    deleteat!(MJs, findall(isequal(MJ), MJs)[1:N])
                end
            end

            # Do we really want unique here?
            sort(unique(Js))
        end

        for twoj = 0:10, w = 1:twoj+1
            j = HalfInteger(twoj//2)
            ts = AtomicLevels._terms_jw(j, w)

            @test issorted(ts, rev=true) # make sure that the array is sorted in descending order
            @test terms_reference(j, w) == sort(unique(ts))
            # Check particle-hole symmetry
            if w != twoj + 1
                @test AtomicLevels._terms_jw(j, twoj+1-w) == ts
            end
        end
    end

    @testset "jj coupling of equivalent electrons" begin
        @test terms(ro"1s", 0) == [0]
        @test terms(ro"1s", 1) == [1//2]
        @test terms(ro"1s", 2) == [0]

        @test terms(ro"3d-", 0) == [0]
        @test terms(ro"3d-", 1) == [3//2]
        @test terms(ro"3d-", 4) == [0]

        @test terms(ro"Xd", 0) == [0]
        @test terms(ro"Xd", 1) == [5//2]
        @test terms(ro"Xd", 6) == [0]

        # Table 4.5, Cowan 1981
        foreach([
            (ro"1s",ro"2p-") => [(0,2) => 0, 1 => 1//2],
            (ro"2p",ro"3d-") => [(0,4) => 0, (1,3) => 3//2, 2 => [0,2]],
            (ro"3d",ro"4f-") => [(0,6) => 0, (1,5) => 5//2, (2,4) => [0,2,4],
                                 3 => [3//2,5//2,9//2]],
            (ro"4f",ro"5g-") => [(0,8) => 0, (1,7) => 7//2, (2,6) => [0,2,4,6],
                                 (3,5) => [3//2,5//2,7//2,9//2,11/2,15//2],
                                 4 => [0,2,2,4,4,5,6,8]],
            (ro"5g",ro"6h-") => [(0,10) => 0, (1,9) => 9//2, (2,8) => [0,2,4,6,8],
                                 (3,7) => [3//2,5//2,7//2,9//2,9//2,11//2,13//2,15//2,17//2,21//2],
                                 (4,6) => [0,0,2,2,3,4,4,4,5,6,6,6,7,8,8,9,10,12],
                                 5 => [1//2,3//2,5//2,5//2,7//2,7//2,9//2,9//2,9//2,11//2,11//2,13//2,13//2,15//2,15//2,17//2,17//2,19//2,21//2,25//2]]
        ]) do (orbs,wsj)
            foreach(orbs) do orb
                foreach(wsj) do (ws,j)
                    js = map(HalfInteger, isa(j, Number) ? [j] : j)
                    foreach(ws) do w
                        @test terms(orb,w) == js
                    end
                end
            end
        end
    end
end


"""
Implementation of the Graphical Lasso (Glasso) for l1-penalized precision matrix estimation

Code credit: https://github.com/bnaul/GraphLasso.jl
"""

function glasso(S::Matrix{Float64}, α::Float64; tol::Float64 = 1e-5,
                    maxit::Int = 1000, penalize_diag::Bool = true, verbose::Bool = false)
    p = size(S, 1)
    adj = abs.(S) .> α
    blocks = connected_components(SimpleGraph(max.(adj, I(p))))
    if(verbose)
        print(blocks)
    end
    Θ = zeros(p, p)
    for block = blocks
        Θ[block,block] .= fit_block(S[block,block], α, tol, maxit, penalize_diag)
    end
    return Θ
end

function fit_block(S::Matrix{Float64}, α::Float64, tol::Float64,
                   maxit::Int, penalize_diag::Bool)
    p = size(S, 1)
    W = copy(S)
    if penalize_diag
        W .+= α * I(p)
    end
    if size(S, 1) == 1
        return inv(W)
    end
    Θ = zeros(p, p)
    i = 0
    β = zeros(p-1, p)
    while i < maxit
        i += 1
        W_old = copy(W)
        for j = 1:p
            inds = collect(1:p)
            splice!(inds, j)
            W11 = W[inds, inds]
            sqrtW11 = sqrt(W11)
            # β[:,j] = fit(LassoPath, sqrtW11, sqrtW11 \ S[inds,j], λ=[α/(p-1)],
            #              standardize=false, intercept=false).coefs
            β[:,j] = glmnet(sqrtW11, sqrtW11 \ S[inds, j], lambda=[α / (p-1)],
                         standardize = false, intercept = false).betas
            W[inds, j] .= W11 * β[:, j]
            W[j, inds] .= W[inds, j]
        end
        if norm(W - W_old) < tol
            break
        end
    end
    if i == maxit
        @warn "Maximum number of iterations reached, graphical lasso failed to converge"
    end
    for j in 1:p
        inds = collect(1:p)
        splice!(inds, j)
        Θ[j, j] = 1 / (W[j, j] - dot(W[inds, j], β[:, j]))
        Θ[inds, j] .= -Θ[j, j] * β[:, j]
        Θ[j, inds] .= Θ[inds, j]
    end
    return Θ
end

# # Testing
# S_iris = cov(Array(dataset("datasets", "iris")[:,1:4]))

# α = 1.0

# W_R = [1.68569351 0.000000 0.2743155 0.01969989
#        0.00000000 1.189979 0.0000000 0.00000000
#        0.27431545 0.000000 4.1162779 0.29560940
#        0.01969989 0.000000 0.2956094 1.58100626]

# Θ_R = [ 0.59973155 0.0000000 -0.03996709  0.00000000
#         0.00000000 0.8403507  0.00000000  0.00000000
#        -0.03996708 0.0000000  0.24890787 -0.04604166
#         0.00000000 0.0000000 -0.04604166  0.64111722]

# W_iris, Θ_iris = graphlasso(S_iris, α; tol=1e-6, penalize_diag=true)

# @test isapprox(W_R, W_iris, atol=1e-5)
# @test isapprox(Θ_R, Θ_iris, atol=1e-5)


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

using DelimitedFiles

name = "double_pendulum2.csv"
file = datadir("experimental", name)
M = readdlm(file, ',')

x = M[:, 5][1:15000]
y = M[:, 6][1:15000]

fig, axs = subplots(1,2; figsize = (10, 5))
axs[1].plot(x)
axs[2].plot(x,y)
axs[1].plot(y)

writedlm(datadir("experimental", "doublependulum.txt"), [x y])

estimate_embedding_plot(x)
estimate_embedding_plot(y)
# Mutual information does NOT have a clear minimum
# the value of 50 is given approximately
# Cao's method provides embedding dimension of 5

_, τ_vals_x = pecuzal_embedding(x .+ 1e-9randn(length(x)); τs = 0:100, w = 5)
_, τ_vals_y = pecuzal_embedding(y .+ 1e-9randn(length(x)); τs = 0:100, w = 5)
# Astonishingly Pecuzal method gives very similar τ for both input timeseries
τx = [0, 51, 25, 39, 12]
τx = [0, 60, 30, 46, 14]

# Let's see how it does for the dual input
A = Dataset(x .+ 1e-9randn(length(x)), y .+ 1e-9randn(length(x)))
_, τ_vals_d, ts_vals = pecuzal_embedding(A; τs = 0:100, w = 5)

# It gives a 6dimension embedding
τd = [0, 61, 31, 47, 31, 0]
jd = [2, 2, 2, 2, 1, 1]

# # PECUZAL method gives 4d embedding with τ = [0, 6, 3, 14]
#
# # %%
# writedlm(datadir("experimental", "roessler_all.txt"), M)
# writedlm(datadir("experimental", "roessler_embed.txt"), x2)
#
# # %% Okay, this will take tremendous amount of computation time,
# # using all 28 timeseries...
# # Y, τ_vals, ts_vals, = pecuzal_embedding(X1)


"""
    NLNS(n_steps, x₀)

A nonlinear, nonstationary process.

## Equations
The systems is given by the following stochastic model:

```math
x(t+1) = x(t) + a(t) a(t-1)
```

with ``a(t) \textasciitilde N(0,1)`` being stationary. The nonlinearity lies in
the ``a(t) a(t-1)`` term.
This is, in essence, a nonlinear random walk. `x₀` sets the initial condition.

This example is from [Lucio et al. (2012)](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202).

## References

Lucio et al., Phys. Rev. E *85*, 056202 (2012), after J. Timmer,
Phys. Rev. E *58*, 5153 (1998). [https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202)

"""
function NLNS(n_steps, x₀)
    a = rand(Uniform(0, 1), n_steps)

    x = Vector{Float64}(n_steps)
    x[1] = x₀
    for i = 2:n_steps
        x[i] = x[i-1] + a[i]*a[i-1]
    end

    x
end

"""
    AR1(n_steps, x₀, k)

Simple AR(1) model with no static transformation from [Lucio et al. (2012)](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202).

## Equations
The system is given by the following map [1]:

```math
x(t+1) = k x(t) + a(t),
```

where ``a(t)`` is a draw from a normal distribution with zero mean and unit
variance. `x₀` sets the initial condition and `k` is the tunable parameter in
the map.

# References

Lucio et al., Phys. Rev. E *85*, 056202 (2012). [https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202)
"""
function AR1(n_steps, x₀, k)
    a = rand(Normal(), n_steps)
    x = Vector{Float64}(n_steps)
    x[1] = x₀
    for i = 2:n_steps
        x[i] = k*x[i-1] + a[i]
    end
    x
end

"""
    NSAR2(n_steps, x₀, x₁)

Cyclostationary AR(2) process from [Lucio et al. (2012)](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202).

## References

Lucio et al., Phys. Rev. E *85*, 056202 (2012), after J. Timmer,
Phys. Rev. E *58*, 5153 (1998). [https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202)

"""
function NSAR2(n_steps, x₀, x₁)
    T₀ = 50.0
    τ = 10.0
    M = 5.5
    Tmod = 250
    σ₀ = 1.0
    a_1_0 = 2*cos(2*π/T₀)*exp(-1/τ)
    a₂ = -exp(-2 / τ)
    a = rand(Normal(), n_steps)
    x = Vector{Float64}(n_steps)
    x[1], x[2] = x₀, x₁

    for i = 3:n_steps
        T_t = T₀ + M*sin(2*π*i / Tmod)
        a₁_t = 2*cos(2*π/T_t) * exp(-1 / τ)

        σ = σ₀^2/(1 - a_1_0^2 - a₂^2 - (2*(a_1_0^2) * a₂)/(1 - a₂)) *
            (1 - a₁_t - a₂^2 - (2*(a₁_t^2) * a₂)/(1 - a₂))
        x[i] = a₁_t * x[i-1] + a₂*x[i-2] + rand(Normal(0, sqrt(σ)))
    end
    x
end

"""
    randomwalk(n_steps, x₀)

Linear random walk (AR(1) process with a unit root) from [Lucio et al. (2012)](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202).
This is an example of a nonstationary linear process.

# References

Lucio et al., Phys. Rev. E *85*, 056202 (2012), after J. Timmer,
Phys. Rev. E *58*, 5153 (1998). [https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202)
"""
function randomwalk(n_steps, x₀)
    a = rand(Normal(), n_steps)

    x = Vector{Float64}(n_steps)
    x[1] = x₀
    for i = 2:n_steps
        x[i] = x[i-1] + a[i]
    end
    x
end

"""
    SNLST(n_steps, x₀, k)

Dynamically linear process transformed by a strongly nonlinear static
transformation (SNLST), from
[Lucio et al., 2012](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202).

## Equations
The system is by the following map:

```math
x(t) = k x(t-1) + a(t)
```

with the transformation ``s(t) = x(t)^3``.

# References

Lucio et al., Phys. Rev. E *85*, 056202 (2012), after J. Timmer,
Phys. Rev. E *58*, 5153 (1998).
[https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.85.056202)

"""
function SNLST(n_steps, x₀, k)
    a = rand(Normal(), n_steps)
    x = Vector{Float64}(n_steps)
    x[1] = x₀
    for i = 2:n_steps
        x[i] = k*x[i-1] + a[i]
    end
    x.^3
end

# Keyword versions of the functions
AR1(;n_steps = 500, x₀ = rand(), k = rand())    = AR1(n_steps, x₀, k)
NLNS(;n_steps = 500, x₀ = rand())               = NLNS(n_steps, x₀)
SNLST(;n_steps = 500, x₀ = rand(), k = rand())  = SNLST(n_steps, x₀, k)
randomwalk(;n_steps = 500, x₀ = rand())         = randomwalk(n_steps, x₀)
NSAR2(;n_steps = 500, x₀ = rand(), x₁ = rand()) = NSAR2(n_steps, x₀, x₁)

export
    AR1,
    NLNS,
    SNLST,
    randomwalk,
    NSAR2


using Fst

#=
a = Wfst()
add_arc!(a, "0", "1", "a", "b", 0.1)
add_arc!(a, "1", "0", "a", "b", 0.2)
add_arc!(a, "1", "2", "b", "b", 0.3)
add_arc!(a, "1", "3", "b", "b", 0.4)
add_arc!(a, "2", "3", "a", "b", 0.5)
add_arc!(a, "3", "3", "a", "a", 0.6)
add_initial_state!(a, "0")
add_final_state!(a, "3")
a.final_weights["3"] = 0.7
println(wfst2dot(a))
create_pdf(a, "a.pdf")

b = Wfst()
add_arc!(b, "0", "1", "b", "b", 0.1)
add_arc!(b, "1", "1", "b", "a", 0.2)
add_arc!(b, "1", "2", "a", "b", 0.3)
add_arc!(b, "1", "3", "a", "b", 0.4)
add_arc!(b, "2", "3", "b", "a", 0.5)
add_initial_state!(b, "0")
add_final_state!(b, "3")
b.final_weights["3"] = 0.6
create_pdf(b, "b.pdf")

c = compose(a, b)
create_pdf(c, "c.pdf")
=#

# Tests for composition involving <epsilon> transitions
#t1 = Wfst()
#add_arc!(t1, 0, 1, "a", "a", 1.0)
#add_arc!(t1, 1, 2, "b", "<eps>", 1.0)
#add_arc!(t1, 2, 3, "c", "<eps>", 1.0)
#add_arc!(t1, 3, 4, "d", "d", 1.0)
#add_initial_state!(t1, 0, 1.0)
#add_final_state!(t1, 4, 1.0)
#create_pdf(t1, "t1.pdf")
#
#t2 = Wfst()
#add_arc!(t2, 0, 1, "a", "d", 1.0)
#add_arc!(t2, 1, 2, "<eps>", "e", 1.0)
#add_arc!(t2, 2, 3, "d", "a", 1.0)
#add_initial_state!(t2, 0, 1.0)
#add_final_state!(t2, 3, 1.0)
#create_pdf(t2, "t2.pdf")
#
#c = compose(t1, t2)
#create_pdf(c, "c.pdf")
#
#c_eps = compose_epsilon(t1,t2)
#create_pdf(c_eps, "c_eps.pdf")

a = read_wfst(readall("fsts/compose_test1/a.txt"), Probability_semiring)
b = read_wfst(readall("fsts/compose_test1/b.txt"), Probability_semiring)
c = compose_basic(a, b)
create_pdf(c, "fsts/compose_test1/c-Probability.pdf")

a = read_wfst(readall("fsts/compose_test1/a.txt"), Tropical_semiring)
b = read_wfst(readall("fsts/compose_test1/b.txt"), Tropical_semiring)
c = compose_basic(a, b)
create_pdf(c, "fsts/compose_test1/c-Tropical.pdf")

a = read_wfst(readall("fsts/compose_test1/a.txt"), Log_semiring)
b = read_wfst(readall("fsts/compose_test1/b.txt"), Log_semiring)
c = compose_basic(a, b)
create_pdf(c, "fsts/compose_test1/c-Log.pdf")


example_ticket_str = ".--------------------------------------------------------.
| ????: 101    ?????: 102   ??????????: 103     ???: 104 |
|                                                        |
| ??: 301  ??: 302             ???????: 303      ??????? |
| ??: 401  ??: 402           ???? ????: 403    ????????? |
'--------------------------------------------------------'"

test_str = "class: 1-3 or 5-7
row: 6-11 or 33-44
seat: 13-40 or 45-50

your ticket:
7,1,14

nearby tickets:
7,3,47
40,4,50
55,2,20
38,6,12"

test_str2 = "class: 0-1 or 4-19
row: 0-5 or 8-19
seat: 0-13 or 16-19

your ticket:
11,12,13

nearby tickets:
3,9,18
15,1,5
5,14,9"

test_list = split(test_str, '\n')

test_list2 = split(test_str2, '\n')

list = readlines("day16_input.txt")

function count_invalid(list)
    rules_end = findfirst(==(""), list) - 1
    rules = map(l -> split(l, r": |-| or ") |> e -> (e[1], parse(Int, e[2]), parse(Int, e[3]), parse(Int, e[4]), parse(Int, e[5])), list[1:rules_end])
    total = 0
    for i in rules_end+6:length(list)
        nums = parse.(Int, split(list[i], ','))
        for n in nums
            if !any(x -> x[2] ≤ n ≤ x[3] || x[4] ≤ n ≤ x[5], rules)
                total += n
            end
        end
    end
    return total
end

@show count_invalid(test_list) == 71
@show count_invalid(list) == 27802

function count_ticket(list)
    num_rules = findfirst(==(""), list) - 1
    rules = map(l -> split(l, r": |-| or ") |> e -> Tuple([e[1], [parse(Int, e[j]) for j in 2:5]]), list[1:num_rules])
    rule_decoder = Dict(map(x -> 2^(x[1]-1) => x[2][1], enumerate(rules)))

    # get valid tickets
    other_tickets = fill(0, length(list) - num_rules - 5, num_rules)
    num_others = 0
    for i in num_rules+6:length(list)
        m = parse.(Int, split(list[i], ','))
        ok = true
        j = 1
        while ok == true && j ≤ num_rules
            if !any(x -> x[2][1] ≤ m[j] ≤ x[2][2] || x[2][3] ≤ m[j] ≤ x[2][4], rules)
                ok = false
            end
            j += 1
        end
        if ok
            num_others += 1
            other_tickets[num_others,:] = m
        end
    end
    other_tickets = other_tickets[1:num_others, :]

    # build array with possible rules (row) for each position (column)
    ruling_off = fill(true, num_rules, num_rules)
    for k in 1:num_rules
        for j in 1:num_rules
            x = rules[j][2]
            for i in 1:num_others
                n = other_tickets[i,k]
                if !(x[1] ≤ n ≤ x[2] || x[3] ≤ n ≤ x[4])
                    ruling_off[j,k] = false
                end
            end
        end
    end

    # find order from array of possible rules
    order = fill(0, num_rules)
    for k in 1:num_rules
        sums = sum(ruling_off, dims=1)
        argmin_sum = argmax(sum(ruling_off, dims=1))[2]
        for j in 1:num_rules
            if 0 < sums[j] ≤ sums[argmin_sum]
                argmin_sum = j
            end
        end
        aux = argmax(ruling_off[:,argmin_sum])
        order[argmin_sum] = 2^(aux-1)
        ruling_off[aux,:] = fill(0, num_rules)
    end

    # get product of values with rules containing "departure"
    product = 1
    ordered_rules = Tuple(rule_decoder[k] for k in order)
    myticket = parse.(Int, split(list[num_rules+3], ','))
    for j in 1:num_rules
        if occursin("departure", ordered_rules[j])
            product *= myticket[j]
        end
    end

    return product, ordered_rules
end

@show count_ticket(test_list2)[2] == ("row", "class", "seat")
@show count_ticket(list)[1] == 279139880759

nothing

#=
```julia
julia> @btime count_invalid(list);
  907.626 μs (1491 allocations: 440.66 KiB)

julia> @btime count_ticket(list);
  15.494 ms (214734 allocations: 3.86 MiB)
```
 =#


# ---
# title: 322. Coin Change
# id: problem322
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/coin-change/description/>
# hidden: true
# ---
#
# You are given coins of different denominations and a total amount of money
# _amount_. Write a function to compute the fewest number of coins that you need
# to make up that amount. If that amount of money cannot be made up by any
# combination of the coins, return `-1`.
#
# You may assume that you have an infinite number of each kind of coin.
#
#
#
# **Example 1:**
#
#
#
#     Input: coins = [1,2,5], amount = 11
#     Output: 3
#     Explanation: 11 = 5 + 5 + 1
#
#
# **Example 2:**
#
#
#
#     Input: coins = [2], amount = 3
#     Output: -1
#
#
# **Example 3:**
#
#
#
#     Input: coins = [1], amount = 0
#     Output: 0
#
#
# **Example 4:**
#
#
#
#     Input: coins = [1], amount = 1
#     Output: 1
#
#
# **Example 5:**
#
#
#
#     Input: coins = [1], amount = 2
#     Output: 2
#
#
#
#
# **Constraints:**
#
#   * `1 <= coins.length <= 12`
#   * `1 <= coins[i] <= 231 - 1`
#   * `0 <= amount <= 104`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


function fib(i::Int64)::Int64
    if i == 0 || i == 1
        return 1
    end
    return (fib((i - 1)) + fib((i - 2)))
end

function main()
    println(join([fib(5)], " "))
end

main()


#=
Popis
  Funkcia vypocita pre kazdy vrchol hodnotu jeho load centrality.

Parametre
  g : datova struktura grafu :: GenericGraph
  dist_key : nazov atributu, v ktorom je ulozena vlastnost vrchola :: String
  normalize : parameter pre aplikovanie normalizacie :: Bool

Navratova hodnota
  loadness : slovnik obsahujuci mapovanie vrcholov a k nim prisluchajucich hodnot
           closeness centrality (index_vrchola => hodnota_loadness_centrality) :: Dict{Int64, Float64}()

Poznamky
  Hodnota centrality moze byt normalizovana nastavenim parametra normalize na True/False.
=#

function centrality_load(g::GenericGraph; normalize::Bool=true, dist_key::AbstractString="weight")

  loadness = [v.index => 0.0 for v=vertices(g)]

  for s in keys(loadness)

    u_loadness = node_loadness(g, s, dist_key)

    for temp in keys(u_loadness)
      loadness[temp] += u_loadness[temp]
    end

  end

  if normalize

    vertex_count = num_vertices(g)

    if vertex_count <= 2
      return loadness
    end

    scale = 1.0 / ((vertex_count - 1) * (vertex_count - 2))

    for v in keys(loadness)
      loadness[v] *= scale
    end

  end

  return loadness

end


abstract type RandBase{T} <: UnaryNodeOp{T} end

mutable struct RandState{RNG<:Random.AbstractRNG} <: NodeEvaluationState
    rng::RNG
end

always_ticks(::RandBase) = true
time_agnostic(::RandBase) = true
value_agnostic(::RandBase) = true
# Note that we should never mutate the random state on the node op itself.
create_operator_evaluation_state(::Tuple{Node}, op::RandBase) = RandState(copy(op.rng))

struct Rand{T<:Real} <: RandBase{T}
    rng::Random.AbstractRNG
end

Base.hash(x::Rand, h::UInt64) = hash(x.rng, hash(:Rand, h))
Base.:(==)(x::Rand{T}, y::Rand{T}) where {T} = x.rng == y.rng

operator!(::Rand{T}, state::RandState) where {T} = rand(state.rng, T)

struct RandArray{T,N} <: RandBase{Array{T,N}}
    rng::Random.AbstractRNG
    dims::Dims{N}
end

Base.hash(x::RandArray, h::UInt64) = hash(x.rng, hash(x.dims, hash(:RandArray, h)))
Base.:(==)(x::RandArray{T}, y::RandArray{T}) where {T} = x.rng == y.rng && x.dims == y.dims

operator!(op::RandArray{T}, state::RandState) where {T} = rand(state.rng, T, op.dims)

"""
    rand([rng=MersenneTwister(),] alignment::Node[, S, dims...])

Generate random numbers aligned to `alignment`, with the given `rng` if provided.

Semantics are otherwise very similar to the usual `Base.rand`:
* If specified, `S` will be the element type, and will default to `Float64` otherwise.
* If specified, `dims` should be a tuple or vararg of integers representing the dimensions
    of an array.

**NB** The values of `alignment` will be ignored.

!!! tip
    If provided, `rng` will be copied before it is used. This is to ensure reproducability
    when evaluating a node multiple times.
"""
Base.rand(alignment::Node) = rand(MersenneTwister(), alignment)
Base.rand(rng::Random.AbstractRNG, alignment::Node) = rand(rng, alignment, Float64)
function Base.rand(rng::Random.AbstractRNG, alignment::Node, ::Type{X}) where {X}
    return obtain_node((alignment,), Rand{X}(copy(rng)))
end

function Base.rand(alignment::Node, dims::Integer...)
    return rand(MersenneTwister(), alignment, Dims(dims))
end
Base.rand(alignment::Node, dims::Dims) = rand(MersenneTwister(), alignment, dims)
function Base.rand(rng::Random.AbstractRNG, alignment::Node, dims::Dims)
    return rand(rng, alignment, Float64, dims)
end
function Base.rand(
    rng::Random.AbstractRNG, alignment::Node, ::Type{X}, dims::Integer...
) where {X}
    return rand(rng, alignment, X, Dims(dims))
end
function Base.rand(
    rng::Random.AbstractRNG, alignment::Node, ::Type{X}, dims::Dims
) where {X}
    return obtain_node((alignment,), RandArray{X,length(dims)}(copy(rng), dims))
end


using Revise
using BNPRegressionGGA2021
using CSV
using DataStructures
using DataFrames
using Distributions
using Gadfly
using Random
using Statistics
const BNP = BNPRegressionGGA2021

# Define the DGP
function simulate_sample(N0, N1)
    dy1 = LogNormal(0.3, 0.7);
    dy2 = MixtureModel(LogNormal, [(0.0, 1.0), (0.3, 1.2)], [0.6, 0.4]);
    X0 = [ones(N0) rand([0, 1], N0, 5)]
    X1 = kron([ones(2, 5) [1, 0]], ones(N1))
    y0 = [X0[i, end] == 1 ? rand(dy2) : rand(dy1) for i in 1:N0]
    c0 = rand(Exponential(10), N0)
    y1 = LinRange(0, 6, N1) |> collect |> x -> repeat(x, size(X1, 1) ÷ N1)
    z0 = deepcopy(y0)
    return dy1, dy2, z0, c0, y0, X0, y1, X1
end

N0, N1, Niter = 500, 2, 1;
df = zeros(Bool, Niter, 6);
Random.seed!(1);
dy1, dy2, z0, c0, y0, X0, y1, X1 = simulate_sample(N0, N1);
mean(y0 .> c0)
grid = LinRange(0, 6, 50) |> collect;
plot(
    layer(x = grid, y = 1 .- cdf.(dy1, grid), Geom.line()),
    layer(x = grid, y = 1 .- cdf.(dy2, grid), Geom.line())
)
plot(
    layer(x = grid, y = pdf.(dy1, grid) ./ (1 .- cdf.(dy1, grid)), Geom.line()),
    layer(x = grid, y = pdf.(dy2, grid) ./ (1 .- cdf.(dy2, grid)), Geom.line())
)
m = BNP.DGPMErlang(; c0, y0, X0, y1, X1);
_, _, _, chaing = BNP.sample!(m; mcmcsize = 10000, burnin = 5000);
mean(chaing)

# Get the MAP of γ for 100 simulated samples
begin
    N0, N1, Niter = 500, 2, 100;
    df = zeros(Bool, Niter, 6);
    Random.seed!(1);
    for iter in 1:Niter
        _, _, _, c0, y0, X0, y1, X1 = simulate_sample(N0, N1);
        m = BNP.DGPMErlang(; c0, y0, X0, y1, X1);
        _, _, _, chaing = BNP.sample!(m; mcmcsize = 10000, burnin = 5000);
        df[iter, :] .= begin
            chaing |>
            DataStructures.counter |>
            collect |>
            x -> sort(x, by = x -> x[2], rev = true) |>
            x -> getindex(x, 1) |>
            x -> getindex(x, 1)
        end
        println(iter)
        println(df[iter, :])
    end
    CSV.write("data/simulation-erlang-nph-gamma-bnp.csv", DataFrame(df, :auto))
end

# Save the 100 simulated samples
N0, N1, Niter = 500, 2, 100;
Random.seed!(1);
df = map(1:Niter) do iter
    _, _, _, c0, y0, X0, y1, X1 = simulate_sample(N0, N1);
    df = DataFrame(X0, :auto)
    df[!, :iter] .= iter
    df[!, :y0] = y0
    df[!, :c0] = c0
    df
end;
df = reduce(vcat, df);
CSV.write("data/simulation-erlang-nph-data.csv", df);


module test_suite

using DomainSets, BasisFunctions, FrameFun, StaticArrays, FastTransforms
using Test, Printf, LinearAlgebra, Random

FE = FrameFun
BA = BasisFunctions

## Settings

# Test fourier extensions for all parameters
Extensive = false

# Show matrix vector product timings
const show_mv_times = false
const verbose = true
total_mv_allocs = 0
total_mv_time = 0.0

# const include_1d_tests = false
# const include_2d_tests = false
# const include_3d_tests = false
# const include_bigfloat_tests = false
const include_1d_tests = true
const include_2d_tests = true
const include_3d_tests = true
const include_bigfloat_tests = true

########
# Auxiliary functions
########

function delimit(s::AbstractString)
    println()
    println("############")
    println("# ",s)
    println("############")
end

show_mv_timings(basis, domain, fun; options...) = show_mv_times && show_mv_timings(ExtensionFrame(domain, basis), fun; options...)

function show_mv_timings(dict::Dictionary, fun; options...)
    A, B = full_discretization(FrameFun.promote_dictionary(dict, fun), fun; options...)
    op = A
    c1 = zeros(eltype(op), size(src(op)))
    c2 = zeros(eltype(op), size(dest(op)))
    apply!(op, c2, c1)
    t = @timed apply!(op, c2, c1)
    if t[3] > 700
        printstyled("\"\tMV product: ", color=:red)
    elseif t[3] == 0
        printstyled("\"\tMV product: ", color=:green)
    else
        printstyled("\"\tMV product: ", color=:blue)
    end
    println(t[3], " bytes allocated, in ", t[2], " seconds")

    global total_mv_allocs += t[3]
    global total_mv_time += t[2]
    return nothing
end

#######
# Testing -- Accuracy
#######

function test_1d_cases()
    delimit("1D")

    ## The tridiagonal solver for Fourier extensions
    solverstyle = TridiagonalProlateStyle()

    n = 61
    T = 1.9
    domain = -1.0..1.0
    basis = Fourier(n) → -T..T

    println()
    println("Testing \t solver = $solverstyle, \n\t\t Domain = $domain, \n\t\t Basis = Fourier,\n\t\t ELT = Float64 ")
    verbose && println("N\t T\t Complex?\t abserror\t time\t\t memory   ")

    # Complex and real functions / Float64
    f(x) = cos(x^2-1//2)-1
    g(x) = 1im*cos(x^2-1//2)-1
    for func in (f,g)
        F = @timed Fun(func, basis, domain; solverstyle=solverstyle, normalizedsampling=false)
        error = abserror(func, F[1])
        if verbose
            print("$n\t $T\t\t")
            func==g ? print("Y\t") : print("N\t")
            @printf("%3.2e\t %3.2e s\t %3.2e bytes \n",error, F[2],F[3])
        end
        @test  (error < sqrt(eps(Float64))*10)
        show_mv_timings(basis, domain, func; solverstyle=solverstyle, normalizedsampling=false)
    end

    @testset "result" for ELT in (Float32,Float64),
            Basis in (Fourier, ChebyshevT),
            domain in [Interval(ELT(-1.5),ELT(0.7)), DomainSets.UnionDomain(Interval(ELT(-1.5),ELT(-0.5)),Interval(ELT(0.5),ELT(1.5)))],
            solverstyle in (AZStyle(), DirectStyle())

        println()
        println("Testing \t solver = $solverstyle, \n\t\t Domain = $domain, \n\t\t Basis = $(repr(Basis(10))),\n\t\t ELT = $ELT ")
        verbose && println("N\t T\t Complex?\t abserror\t time\t\t memory   ")

        for n in (61,)

            # There is some symmetry around T=2, test smaller and larger values
            for T in (ELT(1.9),)
                for func in (f,g)

                    basis = Basis{ELT}(n) → -T..T
                    F = @timed Fun(func, basis, domain; solverstyle=solverstyle)
                    error = abserror(func, F[1])
                    if verbose
                        print("$n\t $T\t\t")
                        func==g ? print("Y\t") : print("N\t")
                        @printf("%3.2e\t %3.2e s\t %3.2e bytes \n",error, F[2],F[3])
                    end
                    @test  (error < sqrt(eps(ELT))*10)
                    show_mv_timings(basis, domain, func; solverstyle=solverstyle)
                end
            end
        end
    end
end

function test_bigfloat()
    f(x) = cos(x.^2) - big(1.0)
    g(x) = big(1.0)im * cos(x.^2) - big(1.0)
    @testset "result" for Basis in (Fourier, ChebyshevT),
            domain in [Interval(big(-3.0)/2, big(7.0)/10)]
        println()
        println("Testing \t solver = QR_solver\n\t\t Domain = $domain, \n\t\t Basis = $(repr(Basis(10))),\n\t\t ELT = BigFloat ")
        verbose && println("N\t T\t Complex?\t abserror\t time\t\t memory   ")
        for T in (big(17.0)/10,)
            for func in (f,g)
                basis = Basis(91) → -T..T
                F = @timed Fun(func, basis, domain; solverstyle = DirectStyle(), directsolver=:qr)
                error = abserror(func, F[1])
                if verbose
                    @printf("91\t %3.2e\t\t",T)
                    func==g ? print("Y\t") : print("N\t")
                    @printf("%3.2e\t %3.2e s\t %3.2e bytes \n",error, F[2],F[3])
                end
                @test error < 1e-20
                show_mv_timings(basis, domain, func; solverstyle = DirectStyle(), directsolver=:qr)
            end
        end
    end
end

function test_2d_cases()
    delimit("2D")

    f(x,y) = cos(0.5*x)+2*sin(0.2*y)-1.0*x*y
    g(x,y) = 1im*cos(0.5*x)+2*sin(0.2*y)-1.0im*x*y
    @testset "result" for Basis in (Fourier, ChebyshevT),
            domain in [Disk(1.2,SA[-0.1,-0.2]), Rectangle(SA[-1.0,-1.5],SA[0.5,0.7])],
            solverstyle in (AZStyle(), DirectStyle())

        println()
        println("Testing \t solver = $solverstyle \n\t\t Domain = $domain, \n\t\t Basis = $(repr(Basis(10)⊗Basis(10))),\n\t\t ELT = Float64 ")
        verbose && println("N\t\t T\t\t Complex?\t abserror\t time\t\t memory   ")

        for n in ((11,11),)
            for T in ((1.7,2.3),)
                basis = (Basis(n[1]) → -T[1]..T[1]) ⊗ (Basis(n[2]) → -T[2]..T[2])
                for func in (f,g)
                    F = @timed Fun(func, basis, domain; solverstyle=solverstyle)
                    error = abserror(func, F[1])
                    if verbose
                        print("$n \t $T \t\t")
                        func==g ? print("Y\t") : print("N\t")
                        @printf("%3.2e\t %3.2e s\t %3.2e bytes \n",error, F[2],F[3])
                    end
                    @test error < 1e-3
                    show_mv_timings(basis, domain, func; solverstyle=solverstyle)
                end

            end
        end
    end
end

function test_3d_cases()
    delimit("3D")

    f(x,y,z) = cos(x)+sin(y)-x*z
    # @testset "result" for Basis in (Fourier, ChebyshevT), D in (Cube((-1.2,-1.0,-0.9),(1.0,0.9,1.2)),FrameFun.tensorproduct(Interval(-1.0,1.0),Disk(1.05)), FrameFun.Ball(1.2,[-0.3,0.25,0.1])), solver in (AZSolver, )
    #             show(solver); println()
    @testset "result" for Basis in (Fourier, ChebyshevT),
                domain in (Rectangle(SA[-1.2,-1.0,-0.9],SA[1.0,0.9,1.2]), Ball(1.2,SA[-0.3,0.25,0.1])),
                solverstyle in (AZStyle(), )

        println()
        println("Testing \t solver = $solverstyle \n\t\t Domain = $domain, \n\t\t Basis = $(repr(Basis(10)⊗Basis(10)⊗Basis(10))),\n\t\t ELT = Float64 ")
        verbose && println("N\t\t T\t\t Complex?\t abserror\t time\t\t memory   ")

        n = (12,12,12)

        for T in ((1.7,1.2,1.3),)
            basis = (Basis(n[1]) → -T[1]..T[1]) ⊗ (Basis(n[2]) → -T[2]..T[2]) ⊗ (Basis(n[3]) → -T[3]..T[3])
            F = @timed Fun(f, basis, domain; solverstyle=solverstyle, threshold=10.0^(3/4*log10(eps(real(codomaintype(basis))))), oversamplingfactor=1.5)
            error = abs(f(0.4,0.4,0.4)-F[1](0.4,0.4,0.4))
            if verbose
                print("$n \t $T \t\t")
                print("N\t")
                @printf("%3.2e\t %3.2e s\t %3.2e bytes \n",error, F[2],F[3])
            end
            @test  error < 1e-1
            show_mv_timings(basis, domain, f; solverstyle=solverstyle, threshold=10.0^(3/4*log10(eps(real(codomaintype(basis))))), oversamplingfactor=1.5)
        end
    end
end

delimit("Algorithm Implementation and Accuracy")
# FFTW.set_num_threads(Sys.CPU_CORES)

if include_1d_tests
    test_1d_cases()
end

if include_bigfloat_tests
    test_bigfloat()
end

if include_2d_tests
    test_2d_cases()
end

if include_3d_tests
    test_3d_cases()
end

# delimit("Random circles")
# dom = FrameFun.randomcircles(10)
# #    b = Fourier(21) ⊗ Fourier(21)
# b = Fourier(31) ⊗ ChebyshevT(31)
# f(x,y) = cos(20*x+22*y)
# @time F = Fun(f,b,dom,verbose=true)

if show_mv_times
    println("Total bytes in MV products:\t$total_mv_allocs")
    println("Total time in MV products:\t$total_mv_time")
end

end


module aoc_03
    using Test
    export read_input, part1, part2
    export test_part1, test_part2, test
    export input, demo

    const input_r = r""

    function read_input(filename = "input.txt")
        lines = readlines(filename)
        # for line in eachline(filename)
        input = [parse_line(line) for line in eachline(filename)]
        permutedims(hcat(input...))
    end

    function parse_line(line)
        Vector{Char}(line) .== '#'
    end

    const input = read_input("input.txt")
    const demo = read_input("demo.txt")

    function part1(input = input, slope = CartesianIndex(1,3))
        pos = CartesianIndex(1,1)
        N = size(input,1)
        M = size(input,2)
        trees = 0
        while pos[1] <= N
            trees += input[ pos ]
            pos += slope
            pos = CartesianIndex( pos[1], mod1( pos[2] , M))
        end
        trees
    end

    function test_part1()
        @test part1( demo ) == 7
        @test part1() == 228
    end

    function part2(input = input)
# Right 1, down 1.
# Right 3, down 1. (This is the slope you already checked.)
# Right 5, down 1.
# Right 7, down 1.
# Right 1, down 2.
        slopes = [
            CartesianIndex(1,1),
            CartesianIndex(1,3),
            CartesianIndex(1,5),
            CartesianIndex(1,7),
            CartesianIndex(2,1)
        ]
        prod( part1.( (input,), slopes ) )
    end

    function test_part2()
        @test part2( demo ) == 336
    end

    function test()
        test_part1()
        test_part2()
    end
end