module LinQuadInterpMod

import FEM: AbstractInterpolator, Point2, FENode2, Vertex4, inv2x2t!, det2x2,
       dNmatrix, Jmatrix, dNdxmatrix, Nvec, get_area

export LinQuadInterp

immutable LinQuadInterp <: AbstractInterpolator
    N::Vector{Float64}
    dN::Matrix{Float64}
    dNdx::Matrix{Float64}
    J::Matrix{Float64}
end

function LinQuadInterp()
    N = zeros(4)
    dN = zeros(4,2)
    dNdx = Array(Float64, 4, 2)
    J = Array(Float64, 2, 2)
    LinQuadInterp(N, dN, dNdx, J)
end
get_allocated_N(i::LinQuadInterp) = i.N
get_allocated_dN(i::LinQuadInterp) = i.dN
get_allocated_dNdx(i::LinQuadInterp) = i.dNdx
get_allocated_J(i::LinQuadInterp) = i.J

# Shape functions in local coords
function Nvec(interp::LinQuadInterp, loc_coords::Point2)

    ξ = loc_coords[1]
    η = loc_coords[2]

    N = get_allocated_N(interp)

    N[1] = (1 + ξ) * (1 + η) * 0.25
    N[2] = (1 - ξ) * (1 + η) * 0.25
    N[3] = (1 - ξ) * (1 - η) * 0.25
    N[4] = (1 + ξ) * (1 - η) * 0.25

    return N
end

function dNmatrix(interp::LinQuadInterp, loc_coords::Point2)
    ξ = loc_coords[1]
    η = loc_coords[2]

    dN = get_allocated_dN(interp)

    dN[1,1] =  (1 + η) * 0.25
    dN[1,2] =  (1 + ξ) * 0.25

    dN[2,1] = -(1 + η) * 0.25
    dN[2,2] =  (1 - ξ) * 0.25

    dN[3,1] = -(1 - η) * 0.25
    dN[3,2] = -(1 - ξ) * 0.25

    dN[4,1] =  (1 - η) * 0.25
    dN[4,2] = -(1 + ξ) * 0.25

    return dN
end

function Jmatrix(interp::LinQuadInterp,
                 vertices::Vertex4, nodes::Vector{FENode2},
                 dN::Matrix{Float64})

    J = get_allocated_J(interp)
    fill!(J, 0.0)
    for row in 1:size(dN, 1)
        x = nodes[vertices[row]].coords.x
        y = nodes[vertices[row]].coords.y

        J[1, 1] += dN[row, 1] * x
        J[1, 2] += dN[row, 1] * y
        J[2, 1] += dN[row, 2] * x
        J[2, 2] += dN[row, 2] * y
    end
    return J
end

function dNdxmatrix(interp::LinQuadInterp, local_coords::Point2,
                    vertices::Vertex4, nodes::Vector{FENode2})
        dN = dNmatrix(interp, local_coords)
        J = Jmatrix(interp, vertices, nodes, dN)
        A_mul_B!(interp.dNdx, dN, inv2x2t!(J))
        return interp.dNdx
end

end # module


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -3.0943792528344187 <= q <= 2.261396294653902)

                     add_NL_constraint(m, :(log(1 + exp(-0.6000573802834048 + -0.4460703798694121*$(x[1]) + -0.8919582737388767*$(x[2]))) + log(1 + exp(0.4850728208540551 + 0.8239401438043505*$(x[1]) + -0.080858577965925*$(x[2]))) + log(1 + exp(-0.3015069196609086 + 0.14649274768979348*$(x[1]) + 0.2885676506758026*$(x[2]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using ApproxFun, SpectralMeasures, Plots
gr(linewidth=3,legend=true,xlims=(0,100))

D = freejacobioperator() - I

## Discrete Diffusion equations
u0 = [zeros(20);ones(5);zeros(20);1.2*ones(3)]

# J1 = D
J1 = D
x1,U1 = eigen(J1)
plot(pad(u0,100),ylims=(0,1.5),marker=2,label="t=0",xlabel="k",ylabel="u_k",title="Pure diffusion on the half-line")
  for t in [1,10,100,1000]
    ut1 = U1\(exp(x1*t)*(U1*u0))
    plot!(pad(ut1.coefficients,100),ylims=(0,1.5),marker=2,label="t=$(t)")
  end
  plot!()
Plots.savefig("purediffusion1.pdf")

J1n = full(J1[1:70,1:70])
plot(pad(u0,100),ylims=(0,1.5),marker=2,label="t=0",xlabel="k",ylabel="u_k",title="Pure diffusion on the half-line (n=70 finite section method)")
  for t in [1,10,100,1000]
    ut1 = expm(J1n*t)*pad(u0,70)
    push!(areasfs,sum(ut1))
    push!(energiesfs,sum(ut1.^2))
    plot!(pad(ut1,100),ylims=(0,1.5),marker=2,label="t=$(t)")
  end
  plot!()
Plots.savefig("purediffusion1finitesection.pdf")

# J2 = D + diag(v) where v is a reaction term
react = [zeros(35);-1.;-1.;-1]
J2 = D + SymTriOperator(react,[0.])
x2,U2 = eigen(J2)
plot(pad(u0,100),marker=2,label="t=0",ylims=(0,1.5),xlabel="k",ylabel="u_k",title="Diffusion plus reaction on the half-line")
  for t in [1,10,100,1000]
    ut2 = U2\(exp(x2*t)*(U2*u0))
    plot!(pad(ut2.coefficients,100),ylims=(0,1.5),marker=2,label="t=$(t)")
  end
  plot!()
plot!(pad(-react,100),label="Reaction potential",color=:black,ylims=(0,1.5),linewidth=2)
Plots.savefig("reactiondiffusion.pdf")
eigs = discreteeigs(J2)
plot(U2[1,1:100],marker=2,title="Stable eigenmodes of reaction-diffusion Jacobi operator",xlabel="k",ylabel="u_k",label="v_1")
  plot!(U2[2,1:100],marker=2,label="v_2")
  plot!(pad(react,100),color=:black,label="Reaction potential",yticks=-1.5:0.5:1.5)
Plots.savefig("stableeigenmodes.pdf")

# Pure fractional diffusion α = 0.65
J4 = D
x4,U4 = eigen(J4)
α = 0.65
plot(pad(u0,100),ylims=(0,1.5),marker=2,label="t=0",xlabel="k",ylabel="u_k",title="Pure fractional diffusion on the half-line, \\alpha = 0.65")
  for t in [1,10,100,1000]
    h4 = Fun(x->exp(-(abs(x)^α)*t),space(x4))
    ut4 = U4\(h4*(U4*u0))
    plot!(pad(ut4.coefficients,100),ylims=(0,1.5),marker=2,label="t=$(t)")
  end
  plot!()
Plots.savefig("purefracdiffusion1.pdf")

# Pure fractional diffusion α = 0.85
J5 = D
x5,U5 = eigen(J5)
α = 0.85
plot(pad(u0,100),ylims=(0,1.5),marker=2,label="t=0",xlabel="k",ylabel="u_k",title="Pure fractional diffusion on the half-line, \\alpha = 0.85")
  for t in [1,10,100,1000]
    h5 = Fun(x->exp(-(abs(x)^α)*t),space(x5))
    ut5 = U5\(h5*(U5*u0))
    plot!(pad(ut5.coefficients,100),ylims=(0,1.5),marker=2,label="t=$(t)")
  end
  plot!()
Plots.savefig("purefracdiffusion085.pdf")

# Pure fractional diffusion α = .25
J6 = D
x6,U6 = eigen(J6)
α = .25
plot(pad(u0,100),ylims=(0,1.5),marker=2,label="t=0",xlabel="k",ylabel="u_k",title="Pure fractional diffusion on the half-line, \\alpha = 0.25")
  for t in [1,10,100,1000]
    h6 = Fun(x->exp(-(abs(x)^α)*t),space(x6))
    ut6 = U6\(h6*(U6*u0))
    plot!(pad(ut6.coefficients,100),ylims=(0,1.5),marker=2,label="t=$(t)")
  end
  plot!()
Plots.savefig("purefracdiffusion025.pdf")

# Fractional diffusion with reaction α = 0.85
react = [zeros(20);1.;1.;1.]
J7 = D + SymTriOperator(react,[0.])
discreteeigs(J7)
x7,U7 = eigen(J7)
α = 0.85
plot(pad(u0,100),ylims=(0,1.5),marker=2,label="t=0",xlabel="k",ylabel="u_k",title="Fractional diffusion with reaction on the half-line, \\alpha = 0.85")
#for t in [1;10;100;1000]
t = 1
h7 = Fun(x->exp(-(abs(x)^α)*t),space(x7)[3])
pt1 = space(x7)[1].points[1]
pt2 = space(x7)[2].points[1]
h7 = Fun(space(x7),[exp((pt1^α)*t);exp((pt2^α)*t);coefficients(h7)])

h7 = Fun(x->exp(-(abs(x)^α)*t),space(x7))

ut7 = U7\(h7*(U7*u0))
plot!(pad(ut7.coefficients,100),ylims=(0,1.5),marker=2,label="t=$(t)")
#end
plot!()
plot!(pad(react,100),color=:black,label="Reaction potential")
Plots.savefig("fracreactdiffusion085withpot.pdf")

SymTriPertToeplitz(V8+1.0,[.5],1.0,-.5)

# Discrete Schrodinger equations
V8 = [zeros(40);2*ones(5)]
  J8 = -D + SymTriOperator(V8,[0.])
  x8,U8 = eigen(J8)
  λ8 = discreteeigs(J8)
  n8 = length(λ8)
  Eigv8 = full(U8[1:n8,1:100])

Eigv8[1,40:100]

n8

V9 = [zeros(40);1*ones(3)]
  J9 = D + SymTriOperator(V9,[0.])
  x9,U9 = eigen(J9)
  λ9 = discreteeigs(J9)
  n9 = length(λ9)
  Eigv9 = full(U9[1:n9,1:100])

if n8 > 0
  plot(Eigv8[1,:]/4+λ8[1]-2,marker=2,title="Tunneling barrier potential and eigenstates",xlabel="k",ylabel="u_k",label="v_1")
  if n8 > 1
    for k in 2:n8
      plot!(Eigv8[k,:]/4+λ8[k]-2,marker=2,label="v_$(k)")
    end
  end
  plot!(pad(V8,100),color=:black,label="Potential")
end
Plots.savefig("tunnelbarriereig.pdf")

if n9 > 0
  plot(Eigv9[1,:]/4+λ9[1],marker=2,title="Double well potential and eigenstates (far)",xlabel="k",ylabel="u_k",label="v_1",ylims=(-.075,.4))
  if n9 > 1
    for k in 2:n9
      plot!((-1)^(k+1)*Eigv9[k,:]/4+λ9[k],marker=2,label="v_$(k)")
    end
  end
  plot!(.2-pad(V9,100),color=:black,label="Potential")
end
Plots.savefig("doublewelleigfar.pdf")

V8 = [zeros(40);2*ones(5)]
  J8 = -D + SymTriOperator(V8,[0.])
  n = 400
  uinit = pad!(exp(-(-28:.75:28).^2),n)
  Jn = full(J8[1:n,1:n])
t = 20
  ut8 = expm(im*Jn*t)*uinit
  plot(abs(ut8[1:100]),ylims=(-.5,1.5),marker=2,label="t=$(t) (abs)")
  plot!(pad(V8,100),ylims=(-.5,1.5),color=:black,label="Potential",xlabel="k",ylabel="u_k",title="No quantum tunneling on the half-line")

d,Q = eigen(Jn)

for t in [0;5;10;20;30;40]
  ut8 = expm(im*Jn*t)*uinit
  plot(abs(ut8[1:100]),ylims=(-.5,1.5),marker=2,label="t=$(t) (abs)")
  plot!(pad(V8,100),ylims=(-.5,1.5),color=:black,label="Potential",xlabel="k",ylabel="u_k",title="No quantum tunneling on the half-line")
  Plots.savefig("noquantumtunnelingt=$(t).pdf")
end
plot!()

V8 = [zeros(10);-.5*ones(6);zeros(10);-.5*ones(6)]
  J8 = D + SymTriOperator(V8,[0.])
  n = 400
  uinit = pad!([zeros(15);.45],n)
  Jn = full(J8[1:n,1:n])
  t = 50
  ut8 = expm(im*Jn*t)*uinit
  plot(abs(ut8[1:100]),ylims=(-.5,1),marker=2,label="t=$(t)")
  plot!(pad(V8,100),ylims=(-.5,1),color=:black,label="Potential",xlabel="k",ylabel="u_k",title="Quantum tunneling on the half-line")
Plots.savefig("doublewellsimclose.pdf")

n = 1000
uinit = pad!([zeros(11);.05;.09;.09;.05],n)
  Jn = full(J9[1:n,1:n])
  plot(pad(V9,100),ylims=(0,1.5),marker=2,label="t=0",xlabel="k",ylabel="u_k",title="Quantum tunneling on the half-line")
  t = 1000000
  ut9 = expm(im*Jn*t)*uinit
  plot!(abs(ut9[1:100]),ylims=(0,1.5),marker=2,label="t=$(t)")
Plots.savefig("doublewellsimfar.pdf")

u0 = [zeros(10);1.;1.;1.;1.]
plot(pad(u0,100),marker=2)
for t = [5;10;15;20;25]
  ut = SchroEvo(u0U\(exp(im*x*t)*(U*u0))
  plot!(pad(abs(ut.coefficients),100),marker=2)
end
plot!()


#=
   The Generous King

   From
   Gunnar Blom, Lars Holst, Dennis Sandell:
   "Problems and Snapshots from the World of Probability"
   Page 11f, Problem 1.8 The generous king

   A) This is the first part (run1)
      """
      The king of a country with m inhabitants is generous enough to give
      each citizen a gold coin as a Christmas gift.
      ...
      All m citizen line up in a random order in front of the Royal Palace.
      The king comes out on to a balcony and tosses k gold coins in the air.
      [The probabilities of head/tails is 0.5.]
      The first person in the queue receives as a gift all coins showing
      tails and returns the other to the king. The kind tosses the remaining
      coins; the second person obtains those which come up tails and returns
      the rest, and so on.
      The procedure stops at the m'th toss or earlier if there are no coins
      left.
      """

   B) The second part is to select k so that the last m citizen will
      receive a gold coin with at least probability 1/2.
      However, the model (run2) does not give the same expectation as the
      theoretical in the book and is perhaps not the best approach to
      this problem...

  * First part: 5 citizen and 6 coins

    parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
        Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

           len    3.6846    1.1191     0.0112    0.0150   6288.7212    0.9999     1530.1025
   lastCitizen    0.1984    0.4350     0.0044    0.0054   6080.2801    0.9999     1479.3869

   Distributions of variable len (num:0)
   5.00000 =>    3192  (0.319200)
   3.00000 =>    2825  (0.282500)
   4.00000 =>    2299  (0.229900)
   2.00000 =>    1531  (0.153100)
   1.00000 =>     153  (0.015300)

   Distributions of variable lastCitizen (num:0)
   0.00000 =>    8162  (0.816200)
   1.00000 =>    1697  (0.169700)
   2.00000 =>     136  (0.013600)
   3.00000 =>       5  (0.000500)

   * Second part (for 8 citizen)
     Note: the value of k1 and the theoretical value are too distinct, so this second
           model is probably not correct. One can note that k is about the expected
           value of the prior of k (Poisson(8*8)  ~ 64) which indicates that something
           is strange...

     Theoretical: 88.722839111673

     Summary Statistics
      parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
          Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

              k1   64.7986    7.8484     0.0785    0.1921   1892.3737    0.9999      434.3295
     lastCitizen    1.1325    0.3725     0.0037    0.0086   1993.1598    1.0018      457.4615

     Distributions of variable lastCitizen (num:0)
     1.00000 =>    8779  (0.877900)
     2.00000 =>    1132  (0.113200)
     3.00000 =>      74  (0.007400)
     4.00000 =>      15  (0.001500)

  Cf ~/webppl/the_generous_king.wppl

=#
using Turing, StatsPlots, Distributions
include("jl_utils.jl")

# First part of the problem
@model function the_generous_king1(m,k)

    function toss(a,m,k)
        if m == 0 || k == 0
            return a
        else
            numTails = rand(Binomial(k,0.5))
            return toss(push!(a,numTails),m-1,k-numTails)
        end
    end

    a = toss([],m,k)

    len ~ Dirac(length(a))

    # How many coins did the last citizen in the queue got?
    lastCitizen ~ Dirac(length(a) >= m ? a[end] : 0)

end

println("Part 1: 5 citizen and 6 coins")
model = the_generous_king1(5,6)
# model = the_generous_king1(15,36)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns,:len)
show_var_dist_pct(chns,:lastCitizen)

# Second part:
#

# Theoretical probability
theoreticalProbPart2 = m -> log(2)*2^(m-1)

# Calculate k such that the last m citizen has a 100% chance
# of winning a gold coin.
#
@model function the_generous_king2(m)

    # Prior for k.
    # Note: The mean of k is quite sensitive to the prior
    #
    k1 ~ Poisson(m*m)
    # k1 ~ truncated(Poisson(m*m),0,Inf)
    # k1 ~ DiscreteUniform(1, m*m)
    # k1 ~ Geometric(1/(m*m))

    # This is the same as in part1.
    function toss(a,m,k)
        if m == 0 || k == 0
            return a
        else
            numTails = rand(Binomial(k,0.5))
            return toss(push!(a,numTails),m-1,k-numTails)
        end
    end

    a = toss([],m,k1);

    # How many coins did the last citizen in the queue got?
    lastCitizen ~ Dirac(length(a) >= m ? a[end] : 0)

    # Ensure the probability of 1/2 of last citizen getting a coin
    # if flip() == true
    #     true ~ Dirac(lastCitizen == 1)
    # end
    true ~ Dirac(lastCitizen >= 0.5)

end

m = 8
println("\nPart 2: $m citizen")

println("Theoretical: ", theoreticalProbPart2(m))

model = the_generous_king2(m)
# model = the_generous_king1(15,36)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns,:len)
show_var_dist_pct(chns,:lastCitizen)


using MultivectorAlgebras
using Base.Test: @testset

@testset "1-dimensional" begin include("one-dimensional_tests.jl") end
@testset "2-dimensional" begin include("two-dimensional_tests.jl") end


"A circular or spherical specie of particle with homogenious material properties."
mutable struct Specie{T<:Real}
  ρ::T # density
  r::T # radius
  c::Complex{T} # sound speed
  num_density::T # number density
end

"A homogenious background material."
mutable struct Medium{T}
  ρ::T # density
  c::Complex{T} # sound speed
end

"Returns the volume fraction of the specie."
volume_fraction(sp::Specie; dim = 2) = (dim == 2 ? 1 : 4/3) * sp.r^dim * pi * sp.num_density

"Returns pressure wave speed, when β is the adiabatic bulk modulus/"
p_speed(ρ,β,μ) = sqrt(β+4*μ/3)/sqrt(ρ)

function Specie(ρ::T, r::T, c=one(Complex{T}); volfrac::T=0.1*one(T)) where T <: AbstractFloat
  Specie{T}(ρ,r,c,volfrac/(T(pi)*r^2.0))
end

function Specie(T=Float64;ρ = 1.0, r = 1.0, c= 1.0+0.0im, volfrac=0.1)
  if T <: Int T = Float64 end
  Specie{T}(T(ρ),T(r),Complex{T}(c),T(volfrac/(pi*r^2.0)))
end
# function Specie(T=Float64; ρ = zero(T), r = one(T), c=one(Complex{T}), volfrac=0.1*one(T))
#   Specie{T}(T(ρ),T(r),c,T(volfrac/(T(pi)*r^2.0)))
# end

# function Specie(T; ρ = zero(Float64), r = one(Float64), c=one(Complex{Float64}), volfrac=0.1*one(Float64))
#   Specie{Float64}(Float64(ρ),Float64(r),c,Float64(volfrac/(Float64(pi)*r^2.0)))
# end

Medium(;ρ=1.0, c=1.0+0.0im) = Medium{typeof(ρ)}(ρ,Complex{typeof(ρ)}(c))
# Medium(;ρ::T = 1.0, c::Union{R,Complex{T}} = 1.0+0.0im) = Medium(ρ,Complex{T}(c))


__precompile__()

module QuaternionArgs
using Quaternions
import Base: convert, promote_rule, show, real, imag, conj, abs, abs2, inv, +, -, /, *, rand
import Quaternions: Quaternion, imagi, imagj, imagk
import LinearAlgebra: normalize
import Random: AbstractRNG, SamplerType
import Base: isinteger, isfinite, isnan, isinf, iszero, isone

export QuaternionArg, QuaternionArgF64, QuaternionArgF32, QuaternionArgF16
export amp, phase1, phase2, phase3, normalize

struct QuaternionArg{T<:AbstractFloat} <: Number
    q::T
    phi::T
    theta::T
    psi::T
end

QuaternionArg(q::AbstractFloat,phi::AbstractFloat,theta::AbstractFloat,psi::AbstractFloat) = QuaternionArg(promote(q,phi,theta,psi)...)
QuaternionArg(q::Integer,phi::Integer,theta::Integer,psi::Integer) = QuaternionArg{Float64}(promote(q,phi,theta,psi)...)
QuaternionArg(x::AbstractFloat) = QuaternionArg(x,zero(x),zero(x),zero(x))
QuaternionArg(x::Integer) = QuaternionArg(float(x))

function Quaternion(qarg::QuaternionArg{T}) where T <: AbstractFloat
    q0=qarg.q*(cos(qarg.phi)*cos(qarg.theta)*cos(qarg.psi) + sin(qarg.phi)*sin(qarg.theta)*sin(qarg.psi))
    q1=qarg.q*(sin(qarg.phi)*cos(qarg.theta)*cos(qarg.psi) - cos(qarg.phi)*sin(qarg.theta)*sin(qarg.psi))
    q2=qarg.q*(cos(qarg.phi)*sin(qarg.theta)*cos(qarg.psi) - sin(qarg.phi)*cos(qarg.theta)*sin(qarg.psi))
    q3=qarg.q*(cos(qarg.phi)*cos(qarg.theta)*sin(qarg.psi) + sin(qarg.phi)*sin(qarg.theta)*cos(qarg.psi))

    return Quaternion(q0,q1,q2,q3)
end

function QuaternionArg(x::Quaternion{T}) where T<:Real
    QuaternionArg(convert(QuaternionF64,x))
end

function QuaternionArg(x::Quaternion{T}) where T<:AbstractFloat
    q = abs(x)
    if q==zero(typeof(q))
        return QuaternionArg(zero(typeof(q)))
    end
    x = x/q

    val = 2*(x.q1*x.q2 - x.q0*x.q3)
    # error adjustment for satisfying the domain of asin
    if val > one(typeof(q))
        val = one(typeof(q))
    elseif val < -one(typeof(q))
        val = -one(typeof(q))
    end
    psi = -asin(val)/2

    if (psi != T(pi/4)) && (psi != -T(pi/4))
        phi  = argi(x*beta((conj(x))))/2
        theta= argj(alpha(conj(x))*x)/2
    else
        phi = T(0)
        theta= argj(gamma(conj(x))*x)/2
    end

    t = Quaternion(QuaternionArg(one(T),phi,theta,psi));
    if isapprox(t,-x)
        phi = phi - sign(phi)*T(pi)
    end

    return QuaternionArg(q, phi, theta, psi)
end

const QuaternionArgF64 = QuaternionArg{Float64}
const QuaternionArgF32 = QuaternionArg{Float32}
const QuaternionArgF16 = QuaternionArg{Float16}

convert(::Type{QuaternionArg}, x::Real) = QuaternionArg(x)
convert(::Type{QuaternionArg{T}}, x::Real) where T<:AbstractFloat = QuaternionArg(x)
convert(::Type{QuaternionArg{T}}, q::QuaternionArg{S}) where {T<:AbstractFloat,S<:AbstractFloat} = QuaternionArg{T}(T(q.q), T(q.phi), T(q.theta), T(q.psi))
convert(::Type{QuaternionArg{T}}, q::QuaternionArg{T}) where T<:AbstractFloat = q
convert(::Type{QuaternionArg{T}}, q::Quaternion{T}) where T<:AbstractFloat = QuaternionArg(q)
convert(::Type{Quaternion{T}}, q::QuaternionArg{T}) where T<:Real = Quaternion(q)

promote_rule(::Type{QuaternionArg{T}}, ::Type{S}) where {T<:AbstractFloat,S<:Real} = QuaternionArg{promote_type(T,S)}
promote_rule(::Type{QuaternionArg{T}}, ::Type{QuaternionArg{S}}) where {T<:AbstractFloat,S<:Real} = QuaternionArg{promote_type(T,S)}
promote_rule(::Type{QuaternionArg{T}}, ::Type{Quaternion{S}}) where {T<:AbstractFloat,S<:Real} = Quaternion{promote_type(T,S)}

quaternionArg(q,phi,theta,psi) = QuaternionArg(q,phi,theta,psi)
quaternionArg(x) = QuaternionArg(x)
quaternionArg(q::QuaternionArg) = q

function show(io::IO, z::QuaternionArg)
    pm(z) = z < 0 ? "-$(-z)" : "+$z"
    print(io, z.q, " (phi=",pm(z.phi)," theta=", pm(z.theta), " psi=", pm(z.psi), ")")
end

QuaternionArg(x::AbstractArray{T}) where T<:QuaternionArg = x
QuaternionArg(x::AbstractArray) = copy!(similar(x,typeof(quaternion(one(eltype(x))))), x)
QuaternionArg(x::AbstractArray{T}) where T<:Quaternion = map(QuaternionArg,x)
Quaternion(x::AbstractArray{T}) where T<:QuaternionArg = map(Quaternion,x)

amp(z::QuaternionArg) = z.q
phase1(z::QuaternionArg) = z.phi
phase2(z::QuaternionArg) = z.theta
phase3(z::QuaternionArg) = z.psi
normalize(z::QuaternionArg) = QuaternionArg(one(z.q),z.phi,z.theta,z.psi)

phase1(x::AbstractVector{T}) where T<:Real = zero(x)
phase2(x::AbstractVector{T}) where T<:Real = zero(x)
phase3(x::AbstractVector{T}) where T<:Real = zero(x)

for fn in (:amp,:phase1,:phase2,:phase3,:normalize)
    @eval begin
        ($fn)(A::AbstractArray) = map(($fn),A)
    end
end

conj(z::QuaternionArg) = QuaternionArg(z.q, -z.phi, -z.theta, -z.psi)
abs(z::QuaternionArg) = z.q
abs2(z::QuaternionArg) = z.q*z.q
inv(z::QuaternionArg) = QuaternionArg(inv(Quaternion(z)))

(-)(z::QuaternionArg) = QuaternionArg(-z.q, z.phi, z.theta, z.psi)
(/)(z::QuaternionArg, x::Real) = QuaternionArg(z.q/x, z.phi, z.theta, z.psi)
(+)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)+Quaternion(w))
(-)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)-Quaternion(w))
(*)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)*Quaternion(w))
(/)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)/Quaternion(w))

rand(r::AbstractRNG, ::SamplerType{QuaternionArg{T}}) where T<:AbstractFloat =
     quaternionArg(one(T),T(2.0)*T(pi)*(rand(T)-T(0.5)),one(T)*T(pi)*(rand(T)-T(0.5)),T(0.5)*T(pi)*(rand(T)-T(0.5)))

real(z::QuaternionArg)  = Quaternion(z).q0
imagi(z::QuaternionArg) = Quaternion(z).q1
imagj(z::QuaternionArg) = Quaternion(z).q2
imagk(z::QuaternionArg) = Quaternion(z).q3

argi(z::Quaternion) = atan(z.q1,z.q0)
argj(z::Quaternion) = atan(z.q2,z.q0)
argk(z::Quaternion) = atan(z.q3,z.q0)

alpha(z::Quaternion) = Quaternion(z.q0,  z.q1, -z.q2, -z.q3)
beta(z::Quaternion)  = Quaternion(z.q0, -z.q1,  z.q2, -z.q3)
gamma(z::Quaternion) = Quaternion(z.q0, -z.q1, -z.q2,  z.q3)

isreal(q::QuaternionArg) = isreal(Quaternion(q))
isinteger(q::QuaternionArg) = isinteger(Quaternion(q))
isfinite(q::QuaternionArg) = isfinite(Quaternion(q))
isnan(q::QuaternionArg) = isnan(Quaternion(q))
isinf(q::QuaternionArg) = isinf(Quaternion(q))
iszero(q::QuaternionArg) = iszero(Quaternion(q))
isone(q::QuaternionArg) = isone(Quaternion(q))

end


using PQBaseCamp
using DataFrames
using CSV
using Distributions

# load up a sample data file -
path_to_data_file = "$(pwd())/test/data/SPY.csv"
df = CSV.read(path_to_data_file, DataFrame)

# build a computation model -
model = LogReturnComputationModel()
model.ticker = "SPY"
model.data = df
model.map = :timestamp => :close
Δ_table = Δ(model)

# estimate a distribution from this data -
D = 𝒟(Laplace, Δ_table; colkey = :Δ)


module FETermsTests

using Test
using Gridap.Arrays
using Gridap.Algebra
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.Integration
using Gridap.Fields
using Gridap.FESpaces
using Gridap.CellData
using LinearAlgebra

domain =(0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)

order = 1
grid_topology = get_grid_topology(model)
polytopes = get_polytopes(grid_topology)
reffes = [LagrangianRefFE(Float64,p,order) for p in polytopes]

u_sol(x) = x[1]+x[2]
f(x) = 0

dirichlet_tags = "boundary"
V = GradConformingFESpace(reffes,model,dirichlet_tags)
U = TrialFESpace(V,u_sol)

trian = get_triangulation(model)
degree = 2
quad = CellQuadrature(trian,degree)

strian = SkeletonTriangulation(model)
degree = 2
squad = CellQuadrature(strian,degree)

assem = SparseMatrixAssembler(U,V)

uhd = zero(U)
v = get_cell_basis(V)
u = get_cell_basis(U)
uh = interpolate(u_sol,U)

a(u,v) = ∇(v)⊙∇(u)
l(v) = f*v

j(u,du,v) = a(du,v)
r(u,v) = a(u,v) - l(v)

z(u,v) = jump(v)*jump(u)
w(v) = jump(v)*f

t_linear = LinearFETerm(z,strian,squad)

t_source = FESource(w,strian,squad)

# AffineFETerm

t_affine = AffineFETerm(a,l,trian,quad)

matdata = collect_cell_matrix(u,v,[t_affine, t_linear])
vecdata = collect_cell_vector(uhd,v,[t_affine, t_linear, t_source])
A = assemble_matrix(assem,matdata)
b = assemble_vector(assem,vecdata)
x = A \ b
@test x ≈ get_free_values(uh)

data = collect_cell_matrix_and_vector(uhd,u,v,[t_affine, t_linear, t_source])
A, b = allocate_matrix_and_vector(assem,data)
assemble_matrix_and_vector!(A,b,assem,data)
x = A \ b
@test x ≈ get_free_values(uh)
A, b = assemble_matrix_and_vector(assem,data)
x = A \ b
@test x ≈ get_free_values(uh)

matdata = collect_cell_jacobian(uh,u,v,[t_affine])
vecdata = collect_cell_residual(uh,v,[t_affine])
A = assemble_matrix(assem,matdata)
b = assemble_vector(assem,vecdata)
x = A \ -b
@test (x.+1) ≈ ones(length(x))

# LinearFETerm and FESource

t_linear = LinearFETerm(a,trian,quad)
t_source = FESource(l,trian,quad)

matdata = collect_cell_matrix(u,v,[t_linear, t_source])
vecdata = collect_cell_vector(uhd,v,[t_linear, t_source])
A = assemble_matrix(assem,matdata)
b = assemble_vector(assem,vecdata)
x = A \ b
@test x ≈ get_free_values(uh)

matdata = collect_cell_jacobian(uh,u,v,[t_linear, t_source])
vecdata = collect_cell_residual(uh,v,[t_linear, t_source])
A = assemble_matrix(assem,matdata)
b = assemble_vector(assem,vecdata)
x = A \ -b
@test (x.+1) ≈ ones(length(x))

# NonlinearFETerm

t_nonlinear = FETerm(r,j,trian,quad)

matdata = collect_cell_jacobian(uh,u,v,[t_nonlinear])
vecdata = collect_cell_residual(uh,v,[t_nonlinear])
A = assemble_matrix(assem,matdata)
b = assemble_vector(assem,vecdata)
x = A \ -b
@test (x.+1) ≈ ones(length(x))

data = collect_cell_jacobian_and_residual(uh,u,v,[t_nonlinear])
A, b = allocate_matrix_and_vector(assem,data)
assemble_matrix_and_vector!(A,b,assem,data)
x = A \ -b
@test (x.+1) ≈ ones(length(x))

# AffineFETermFromCellMatVec

q = get_coordinates(quad)
w_q = get_weights(quad)
ϕ = get_cell_map(trian)
jac = ∇(ϕ)
jac_q = evaluate(jac,q)
x_q = evaluate(ϕ,q)

function matvecfun(u,v)
  cellmat = integrate(∇(u)⋅∇(v),trian,quad)
  cellvec = integrate(v*f,trian,quad)
  pair_arrays(cellmat,cellvec)
end

t_matvec_Ω = AffineFETermFromCellMatVec(matvecfun,trian)

matdata = collect_cell_matrix(u,v,[t_matvec_Ω,])
vecdata = collect_cell_vector(uhd,v,[t_matvec_Ω,])
A = assemble_matrix(assem,matdata)
b = assemble_vector(assem,vecdata)
x = A \ b
@test x ≈ get_free_values(uh)

data = collect_cell_matrix_and_vector(uhd,u,v,[t_matvec_Ω,])
A, b = allocate_matrix_and_vector(assem,data)
assemble_matrix_and_vector!(A,b,assem,data)
x = A \ b
@test x ≈ get_free_values(uh)

function jacresfun(uh,du,v)
  cellmat = integrate(∇(du)⋅∇(v),trian,quad)
  cellvec = integrate(∇(uh)⋅∇(v)-v*f,trian,quad)
  pair_arrays(cellmat,cellvec)
end

t_jacres_Ω = FETermFromCellJacRes(jacresfun,trian)

matdata = collect_cell_jacobian(uh,u,v,[t_jacres_Ω])
vecdata = collect_cell_residual(uh,v,[t_jacres_Ω])
A = assemble_matrix(assem,matdata)
b = assemble_vector(assem,vecdata)
x = A \ -b
@test (x.+1) ≈ ones(length(x))

data = collect_cell_jacobian_and_residual(uh,u,v,[t_jacres_Ω])
A, b = allocate_matrix_and_vector(assem,data)
assemble_matrix_and_vector!(A,b,assem,data)
x = A \ -b
@test (x.+1) ≈ ones(length(x))

end # module


using DifferentialEquations, Plots
A = 1
B = 1
C = 4
x0 = 1
function f(du, u, p, t)
    du[1] = -0.5 * u[1]^2 + 0.5 * C^2 * u[2] / B
    du[2] = -A + u[1] * u[2]
end

function bc1!(residual, u, p, t)
    residual[1] = u[1][1] - x0
    residual[2] = u[end][2]
end
tspan = (0,1)
bvp1 = BVProblem(f, bc1!, [0, 0], tspan)
sol1 = solve(bvp1, GeneralMIRK4(), dt = 0.05)
plot(sol1)


import SpecialFunctions

function flux_Gxp(Gxp, nx, ny, u1, u2, rho, pr)
    tx = ny
    ty = -nx

    ut = u1*tx + u2*ty
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S1 = ut*sqrt(beta)
    B1 = 0.5*exp(-S1*S1)/sqrt(pi*beta)
    A1pos = 0.5*(1 + SpecialFunctions.erf(S1))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un

    Gxp[1] = (rho*(ut*A1pos + B1))

    temp1 = pr_by_rho + ut*ut
    temp2 = temp1*A1pos + ut*B1
    Gxp[2] = (rho*temp2)

    temp1 = ut*un*A1pos + un*B1
    Gxp[3] = (rho*temp1)

    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*ut*temp1*A1pos
    temp1 = (6*pr_by_rho) + u_sqr
    Gxp[4] = (rho*(temp2 + 0.5*temp1*B1))
    return nothing
end

function flux_Gxn(Gxn, nx, ny, u1, u2, rho, pr)

    tx = ny
    ty = -nx

    ut = u1*tx + u2*ty
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S1 = ut*sqrt(beta)
    B1 = 0.5*exp(-S1*S1)/sqrt(pi*beta)
    A1neg = 0.5*(1 - SpecialFunctions.erf(S1))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un

    Gxn[1] = (rho*(ut*A1neg - B1))

    temp1 = pr_by_rho + ut*ut
    temp2 = temp1*A1neg - ut*B1
    Gxn[2] = (rho*temp2)

    temp1 = ut*un*A1neg - un*B1
    Gxn[3] = (rho*temp1)

    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*ut*temp1*A1neg
    temp1 = (6*pr_by_rho) + u_sqr
    Gxn[4] = (rho*(temp2 - 0.5*temp1*B1))
    return nothing
end

function flux_Gyp(Gyp, nx, ny, u1, u2, rho, pr)
    tx = ny
    ty = -nx

    ut = u1*tx + u2*ty
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S2 = un*sqrt(beta)
    B2 = 0.5*exp(-S2*S2)/sqrt(pi*beta)
    A2pos = 0.5*(1 + SpecialFunctions.erf(S2))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un

    Gyp[1] = (rho*(un*A2pos + B2))

    temp1 = pr_by_rho + un*un
    temp2 = temp1*A2pos + un*B2

    temp1 = ut*un*A2pos + ut*B2
    Gyp[2] = (rho*temp1)

    Gyp[3] = (rho*temp2)

    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*un*temp1*A2pos
    temp1 = (6*pr_by_rho) + u_sqr
    Gyp[4] = (rho*(temp2 + 0.5*temp1*B2))
    return nothing
end

function flux_Gyn(Gyn, nx, ny, u1, u2, rho, pr)
    tx = ny
    ty = -nx

    ut = u1*tx + u2*ty
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S2 = un*sqrt(beta)
    B2 = 0.5*exp(-S2*S2)/sqrt(pi*beta)
    A2neg = 0.5*(1 - SpecialFunctions.erf(S2))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un

    Gyn[1] = (rho*(un*A2neg - B2))

    temp1 = pr_by_rho + un*un
    temp2 = temp1*A2neg - un*B2

    temp1 = ut*un*A2neg - ut*B2
    Gyn[2] = (rho*temp1)

    Gyn[3] = (rho*temp2)

    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*un*temp1*A2neg
    temp1 = (6*pr_by_rho) + u_sqr
    Gyn[4] = (rho*(temp2 - 0.5*temp1*B2))
    return nothing
end

function flux_Gx(Gx, nx, ny, u1, u2, rho, pr)
    tx = ny
    ty = -nx

    ut = u1*tx + u2*ty
    un = u1*nx + u2*ny

    Gx[1] = rho*ut

    Gx[2] = pr + rho*ut*ut

    Gx[3] = rho*ut*un

    temp1 = 0.5*(ut*ut + un*un)
    rho_e = 2.5*pr + rho*temp1
    Gx[4] = (pr + rho_e)*ut
    return  nothing
end

function flux_Gy(Gy, nx, ny, u1, u2, rho, pr)
    tx = ny
    ty = -nx

    ut = u1*tx + u2*ty
    un = u1*nx + u2*ny

    Gy[1] = rho*un

    Gy[2] = rho*ut*un

    Gy[3] = pr + rho*un*un

    temp1 = 0.5*(ut*ut + un*un)
    rho_e = 2.5*pr + rho*temp1
    Gy[4] = (pr + rho_e)*un
    return nothing
end

# print(flux_Gy(flux_Gxn(1, 3, 2, 2, 2, 1),5,2,3,4,3,2))


using Optim
using Statistics

"""
    k_on, k_off, k_bleach, p_miss, nemitters =
        estimateparams(smd::SMD, clusterdata::Vector{Matrix{Float32}})

Estimate rate parameters from the clusters in `smd` and `clusterdata`.

# Description
The rate parameters `k_on`, `k_off`, and `k_bleach`, the probability of missing
a localization for a visible emitter `p_miss`, and the underlying number of
emitters `nemitters` are estimated from the pre-clusters in
`smd`/`clusterdata`.  This is done by assuming that pre-clusters are, on
average, representative of a single blinking event of a single emitter.
"""
function estimateparams(smd::SMD, clusterdata::Vector{Matrix{Float32}})
    # Compute some quantities needed from each of the preclusters.
    nclusters = size(clusterdata)[1]
    clusterdurations, nobservations = computeclusterinfo(clusterdata)

    # Estimate the sum of the off rate and the bleaching rate, as well as the
    # probability of missing a localization.
    k_offpbleach = -log(1 - 1/Statistics.mean(clusterdurations))
    k_offpbleach = isinf(k_offpbleach) ?  1.0 :  k_offpbleach
    p_miss = 1 - Statistics.mean(nobservations./clusterdurations)

    # Compute other parameters by fitting the observed number of localizations
    # over time.
    nlocperframe = counts(smd.framenum)
    nlocperframe = nlocperframe[nlocperframe .!= 0]
    frames = convert(Array{Float32}, unique(smd.framenum))
    nloccumulative = cumsum(nlocperframe)
    k_off = nclusters / nloccumulative[end]
    k_bleach = maximum([10^-5  k_offpbleach - k_off])
    k(k_on) = k_on + k_offpbleach
    l1(k_on) = k_on * k_bleach / k(k_on)
    l2(k_on) = k(k_on) - l1(k_on)
    model(x) = ceil(x[1]) * (1.0.-p_miss) * (x[2]/k(x[2])) *
        ((1/l1(x[2]))*(1.0.-exp.(-l1(x[2])*(frames.-1.0))) .-
        (1/l2(x[2]))*(1.0.-exp.(-l2(x[2])*(frames.-1.0))))
    costfunction(x) = Statistics.mean((nloccumulative .- model(x)).^2)
    lowerbound = [maximum(nlocperframe) 10^-5]
    upperbound = [nclusters nloccumulative[end]/frames[end]]
    x1_guess = ceil(nclusters * k_bleach)
    x1_guess = (x1_guess>lowerbound[1]) & (x1_guess<upperbound[1]) ?
        x1_guess : (upperbound[1]-lowerbound[1])/2.0
    initialguess = [x1_guess 1.0/frames[end]]
    inner_optimizer = NelderMead()
    optimresults = optimize(costfunction, lowerbound, upperbound, initialguess,
        Fminbox(inner_optimizer))
    xhat = optimresults.minimizer
    nemitters = xhat[1]
    k_on = xhat[2]

    return k_on, k_off, k_bleach, p_miss, nemitters
end


using LinearAlgebra

mutable struct DiscriminativeKalmanFilter
    A::Matrix{Float64}  # from eq. (2.1b)
    Γ::Matrix{Float64}  # from eq. (2.1b)
    S::Matrix{Float64}  # from eq. (2.1a)
    f::Function  # from eq. (2.2)
    Q::Function  # from eq. (2.2)
    μₜ::Vector{Float64}  # from eq. (2.6)
    Σₜ::Matrix{Float64}  # from eq. (2.6)
    d::Int  # as in eq. (2)
end

"""
calculates the first 2 lines of eq. (2.7) in-place
"""
function stateUpdate!(DKF::DiscriminativeKalmanFilter)
    DKF.μₜ = DKF.A * DKF.μₜ
    DKF.Σₜ = DKF.A * DKF.Σₜ * transpose(DKF.A) + DKF.Γ
end

"""
calculates the last 2 lines of eq. (2.7)
"""
function measurementUpdate!(DKF::DiscriminativeKalmanFilter, newMeasurement::Vector{Float64})
    Qxₜ = DKF.Q(newMeasurement)
    fxₜ = DKF.f(newMeasurement)
    if minimum(eigvals(inv(Qxₜ) - inv(DKF.S))) <= 1e-6
        newPosteriorCovInv = inv(DKF.Σₜ) + inv(Qxₜ)
    else
        newPosteriorCovInv = inv(DKF.Σₜ) + inv(Qxₜ) - inv(DKF.S)
    end
    DKF.μₜ = newPosteriorCovInv \ (DKF.Σₜ \ DKF.μₜ + Qxₜ \ fxₜ)
    DKF.Σₜ = inv(newPosteriorCovInv)
end

"""
performs stateUpdate() and measurementUpdate(newMeasurement);
returns new posterior mean as in eq. (2.7)
"""
function predict!(DKF::DiscriminativeKalmanFilter, newMeasurement::Vector{Float64})
    stateUpdate!(DKF)
    measurementUpdate!(DKF, newMeasurement)
    return DKF.μₜ
end


###################################################
## geodesy.jl
## fork of some code in geodesy.jl
###################################################

### Ellipsoid
# Specify datum for translation between LLA and other coordinate systems
struct Ellipsoid
    a::Float64        # Semi-major axis
    b::Float64        # Semi-minor axis
    e²::Float64       # Eccentricity squared
    e′²::Float64      # Second eccentricity squared
end

function Ellipsoid(; a::AbstractString="", b::AbstractString="", f_inv::AbstractString="")
    if isempty(a) || isempty(b) == isempty(f_inv)
        throw(ArgumentError("Specify parameter 'a' and either 'b' or 'f_inv'"))
    end
    if isempty(b)
        _ellipsoid_af(parse(BigFloat,a), parse(BigFloat,f_inv))
    else
        _ellipsoid_ab(parse(BigFloat,a), parse(BigFloat,b))
    end
end

function _ellipsoid_ab(a::BigFloat, b::BigFloat)
    e² = (a^2 - b^2) / a^2
    e′² = (a^2 - b^2) / b^2

    Ellipsoid(a, b, e², e′²)
end
function _ellipsoid_af(a::BigFloat, f_inv::BigFloat)
    b = a * (1 - inv(f_inv))

    _ellipsoid_ab(a, b)
end

### World Geodetic Coordinate System of 1984 (WGS 84)
# Standardized coordinate system for Earth
# Global ellipsoidal reference surface
const WGS84  = Ellipsoid(a = "6378137.0", f_inv = "298.257223563")

### Point in Latitude-Longitude-Altitude (LLA) coordinates
# Used to store node data in OpenStreetMap XML files
struct LLA
    lat::Float64
    lon::Float64
    alt::Float64
end
LLA(lat, lon) = LLA(lat, lon, 0.0)

### Point in Earth-Centered-Earth-Fixed (ECEF) coordinates
# Global cartesian coordinate system rotating with the Earth
struct ECEF
    x::Float64
    y::Float64
    z::Float64
end

### Point in East-North-Up (ENU) coordinates
# Local cartesian coordinate system
# Linearized about a reference point
struct ENU
    east::Float64
    north::Float64
    up::Float64
end
ENU(x, y) = ENU(x, y, 0.0)

# Given a reference point for linarization
function ENU(lla::LLA, lla_ref::LLA, datum::Ellipsoid = WGS84)
    ecef = ECEF(lla, datum)
    return ENU(ecef, lla_ref, datum)
end

###############################
### LLA to ECEF coordinates ###
###############################

function ECEF(lla::LLA, datum::Ellipsoid = WGS84)
    ϕdeg, λdeg, h = lla.lat, lla.lon, lla.alt
    d = datum

    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)
    sinλ, cosλ = sind(λdeg), cosd(λdeg)

    N = d.a / sqrt(1 - d.e² * sinϕ^2)  # Radius of curvature (meters)

    x = (N + h) * cosϕ * cosλ
    y = (N + h) * cosϕ * sinλ
    z = (N * (1 - d.e²) + h) * sinϕ

    return ECEF(x, y, z)
end

function ENU(ecef::ECEF, lla_ref::LLA, datum::Ellipsoid = WGS84)
    ϕdeg, λdeg = lla_ref.lat, lla_ref.lon

    ecef_ref = ECEF(lla_ref, datum)
    ∂x = ecef.x - ecef_ref.x
    ∂y = ecef.y - ecef_ref.y
    ∂z = ecef.z - ecef_ref.z

    # Compute rotation matrix
    sinλ, cosλ = sind(λdeg), cosd(λdeg)
    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)

    # R = [     -sinλ       cosλ  0.0
    #      -cosλ*sinϕ -sinλ*sinϕ cosϕ
    #       cosλ*cosϕ  sinλ*cosϕ sinϕ]
    #
    # east, north, up = R * [∂x, ∂y, ∂z]
    east  = ∂x * -sinλ      + ∂y * cosλ       + ∂z * 0.0
    north = ∂x * -cosλ*sinϕ + ∂y * -sinλ*sinϕ + ∂z * cosϕ
    up    = ∂x * cosλ*cosϕ  + ∂y * sinλ*cosϕ  + ∂z * sinϕ

    return ENU(east, north, up)
end


# recoupling function
# two options for the type:
#   - either all arguments are Float64 or energy variables are complex
function change_basis_3from1(σ1, cosθ1, ϕ1, cosθ23, ϕ23,
        m1sq, m2sq, m3sq, s)
        # calculate σ3 in (23) frame
        m1 = sqrt(m1sq); sqrts = sqrt(s)
        sqrtλσ1 = sqrt((sqrts-m1)^2-σ1)*sqrt((sqrts+m1)^2-σ1)/(2*sqrt(σ1));
        σ3 = m1sq + m2sq +
        # 2*(  E2 * E1 -
            2*( (σ1 + m2sq - m3sq)/(2*sqrt(σ1)) * (s-m1sq-σ1)/(2*sqrt(σ1)) -
            # |p2|*cos(theta23) * (-|p1|)
            sqrt(Kallen(σ1, m2sq, m3sq)/(4*σ1))*cosθ23 * (-sqrtλσ1) );
        # caluclate p3 in (23) frame
        p23bu = sqrt(Kallen(σ1, m2sq, m3sq)/(4*σ1));
        p3_in23 = [(σ1 + m3sq - m2sq)/sqrt(4*σ1),
                   -p23bu*sqrt(1-cosθ23^2)*cos(ϕ23),
                   -p23bu*sqrt(1-cosθ23^2)*sin(ϕ23),
                   -p23bu*cosθ23];
        # boost to lab frame from (23) frame
        γ1 = (s + σ1 - m1sq)/sqrt(4*s*σ1);
        β1 = sqrt(1.0-1.0/γ1^2);
        p3_b = [γ1*(p3_in23[1]+β1*p3_in23[4]),
                p3_in23[2],
                p3_in23[3],
                γ1*(β1*p3_in23[1]+p3_in23[4])];

        ct1 = cosθ1; st1 = sqrt(1.0-cosθ1^2); cp1 = cos(ϕ1); sp1 = sin(ϕ1);
        # Rz(phi1) * Ry(theta1) * p3_boost
        p3_rot = [p3_b[1],
                  cp1*ct1* p3_b[2] + (-sp1)* p3_b[3] + cp1*st1* p3_b[4],
                  sp1*ct1* p3_b[2] +   cp1 * p3_b[3] + sp1*st1* p3_b[4],
                     -st1* p3_b[2] +     0 * p3_b[3] +     ct1* p3_b[4]];

        cosθ3 = -p3_rot[4]/sqrt(p3_b[2]^2+p3_b[3]^2+p3_b[4]^2);
        ϕ3 = (p3_rot[2] != zero(p3_rot[2])) ? atan(-p3_rot[3], -p3_rot[2]) : rand()*one(p3_rot[2]);

        cosθ12_n = m2sq + m3sq + 2* (σ3+m2sq-m1sq)/(2*sqrt(σ3)) * (s-m3sq-σ3)/(2*sqrt(σ3)) - σ1;
        m3 = sqrt(m3sq)
        sqrtλσ3 = sqrt((sqrts-m3)^2-σ3)*sqrt((sqrts+m3)^2-σ3)
        cosθ12_d = 2 * sqrt(Kallen(σ3, m1sq, m2sq)/(4*(σ3))) * sqrtλσ3 / (2*sqrt(σ3));
        cosθ12 = cosθ12_d ≈ 0.0+0.0im ? 2.0*rand()-1.00*one(σ1) : cosθ12_n/cosθ12_d;

        n1 = [0.,
             -sqrt(1.0-cosθ1^2)*cos(ϕ1),
             -sqrt(1.0-cosθ1^2)*sin(ϕ1),
             -cosθ1];
        ct3 = cosθ3; st3 = sqrt(1.0-cosθ3^2); cp3 = cos(ϕ3); sp3 = sin(ϕ3);
        # Rz(phi23) * Ry(theta23) * p3_boost
        n1_rot = [n1[1],
                  cp3*ct3* n1[2] + sp3*ct3* n1[3] + (-st3)* n1[4],
                   (-sp3)* n1[2] +    cp3 * n1[3] +   0.0 * n1[4],
                  cp3*st3* n1[2] + sp3*st3* n1[3] +   ct3 * n1[4]];
        # println(n1_rot[3], " ", n1_rot[2])
        ϕ12 = (n1_rot[2] != zero(n1_rot[2])) ? atan(n1_rot[3], n1_rot[2]) : rand()*one(n1_rot[2]);
        return σ3,cosθ3,ϕ3,cosθ12,ϕ12
end


#Function which compute the Energy and Force of the system

function force_pair(p1,p2,r,Data)

 r6 = (r^2)^3
 r12 = (r^2)^6
 r13 = r12*r
 r7 = r6*r

 dudr1 = -12*(Data.sig12/r13)
 dudr2 = 6*(Data.sig6/r7)

 drdx = -(p1[1]-p2[1])/r
 drdy = -(p1[2]-p2[2])/r

  upair = Data.eps4*(Data.sig12/r12 - Data.sig6/r6)
  fx = (Data.eps4*(dudr1*drdx + dudr2*drdx), Data.eps4*(dudr1*drdy + dudr2*drdy))

 return upair,fx
end


using JAC.AngularMomentum: clebschgordan, wigner3j, wigner6j, wigner9j

@testset "@racahsum" begin
    # Clebsch–Gordan coefficients
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, 5/2, M), M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, 5/2,-M), M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J, M), J, M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J,-M), J, M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J, M), M, J) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J,-M), M, J) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(j₁, m₁, 0, 0, 5/2, 3/2), m₁, j₁) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(j₁,-m₁, 0, 0, 5/2, 3/2), m₁, j₁) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, j₂, m₂, 5/2, 3/2), j₂, m₂) ≈
        sum(clebschgordan(5/2, 3/2, j₂, 0, 5/2, 3/2) for j₂=0:5)
    @test @racahsum(clebschgordan(5/2, 3/2, j₂,-m₂, 5/2, 3/2), j₂, m₂) ≈
        sum(clebschgordan(5/2, 3/2, j₂, 0, 5/2, 3/2) for j₂=0:5)
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2), J) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(1/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, M), J, M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(1/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J,-M), J, M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(1/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, M), J=HalfInt(3/2):HalfInt(9/2), M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(3/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J,-M), J=HalfInt(3/2):HalfInt(9/2), M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(3/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, M), J=HalfInt(1):HalfInt(4), M) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J,-M), J=HalfInt(1):HalfInt(4), M) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(5/2-m) * clebschgordan(5/2, m, 5/2, -m, 0, 0), m) ≈ √6
    @test @racahsum((-1)^Integer(5/2-m) * clebschgordan(5/2, m, 5/2, -m, 2, 0), m) ≈ 0 atol=1e-15
    # Orthogonality relations for Clebsch–Gordan coefficients
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J, M)*clebschgordan(2, 2, 3/2, -3/2, J, M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 1, 3/2, -1/2, J, M)*clebschgordan(2, 1, 3/2, -1/2, J, M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J, M)*clebschgordan(2, 1, 3/2, -1/2, J, M), J, M) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J,-M)*clebschgordan(2, 2, 3/2, -3/2, J,-M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 1, 3/2, -1/2, J,-M)*clebschgordan(2, 1, 3/2, -1/2, J,-M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J,-M)*clebschgordan(2, 1, 3/2, -1/2, J,-M), J, M) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 2, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 3, 1)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 1), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2,-m₁, 3/2, m₂, 2, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2,-m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2,-m₂, 3, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2,-m₂, 3, 1)*clebschgordan(5/2,-m₁, 3/2,-m₂, 3, 1), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2,-m₁, 3/2, m₂, 3, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2,-m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2,-m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2,-m₂, 2, 0)*clebschgordan(5/2,-m₁, 3/2,-m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    # 3-j symbol
    @test @racahsum(wigner3j(2, j, 4, 0, 0, 0), j) ≈ sum(wigner3j(2, j, 4, 0, 0, 0) for j=2:2:6)
    @test @racahsum(wigner3j(2, j, 4, m, 0, m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum(wigner3j(2, j, 4,-m, 0, m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum(wigner3j(2, j, 4, m, 0,-m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum(wigner3j(2, j, 4,-m, 0,-m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum((-1)^Integer(3/2-m)*wigner3j(3/2, 3/2, 0,  m, -m, 0), m) ≈ 2
    @test @racahsum((-1)^Integer(3/2+m)*wigner3j(3/2, 3/2, 0, -m,  m, 0), m) ≈ 2
    @test @racahsum((-1)^Integer(3/2-m)*wigner3j(3/2, 3/2, 1,  m, -m, 0), m) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(3/2+m)*wigner3j(3/2, 3/2, 1, -m,  m, 0), m) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(4-m)*wigner3j(4, 4, 0,  m, -m, 0), m) ≈ 3
    @test @racahsum((-1)^Integer(4+m)*wigner3j(4, 4, 0, -m,  m, 0), m) ≈ 3
    @test @racahsum((-1)^Integer(4-m)*wigner3j(4, 4, 1,  m, -m, 0), m) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(4+m)*wigner3j(4, 4, 1, -m,  m, 0), m) ≈ 0 atol=1e-15
    # Orthogonality relations for 3-j symbols
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁,  m₂, 0)*wigner3j(2, 6, 4,  m₁,  m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁,  m₂, 0)*wigner3j(2, 6, 4, -m₁,  m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁, -m₂, 0)*wigner3j(2, 6, 4,  m₁, -m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁, -m₂, 0)*wigner3j(2, 6, 4, -m₁, -m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁,  m₂, 0)*wigner3j(2, 6, 5,  m₁,  m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁,  m₂, 0)*wigner3j(2, 6, 5, -m₁,  m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁, -m₂, 0)*wigner3j(2, 6, 5,  m₁, -m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁, -m₂, 0)*wigner3j(2, 6, 5, -m₁, -m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁,  m₂, 0)*wigner3j(2, 6, 4,  m₁,  m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁,  m₂, 0)*wigner3j(2, 6, 4, -m₁,  m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁, -m₂, 0)*wigner3j(2, 6, 4,  m₁, -m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁, -m₂, 0)*wigner3j(2, 6, 4, -m₁, -m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 6 * @racahsum(wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2), m₁, m₂) ≈ 1
    @test 4 * @racahsum(wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2), m₁, m₂) ≈ 1
    @test 6 * @racahsum(wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2), m₁, m₂) ≈ 0 atol=1e-15
    @test 6 * @racahsum(wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 5/2, m₁, m₂,  1/2), m₁, m₂) ≈ 0 atol=1e-15
    @test 4 * @racahsum(wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 3/2, m₁, m₂, -3/2), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(2, 6, j₃, 0, 0, m₃)*wigner3j(2, 6, j₃, 0, 0, m₃), j₃, m₃) ≈ 1
    @test @racahsum((2j₃+1)*wigner3j(2, 6, j₃, 0, 0, m₃)*wigner3j(2, 6, j₃, 0, 1, m₃), j₃, m₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(2, 6, j₃, 0, 0, m₃)*wigner3j(2, 6, j₃, 1, 0, m₃), j₃, m₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(3/2, 1, j₃, 1/2, 0, m₃)*wigner3j(3/2, 1, j₃,  1/2, 0, m₃), j₃, m₃) ≈ 1
    @test @racahsum((2j₃+1)*wigner3j(3/2, 1, j₃, 1/2, 0, m₃)*wigner3j(3/2, 1, j₃,  1/2, 1, m₃), j₃, m₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(3/2, 1, j₃, 1/2, 0, m₃)*wigner3j(3/2, 1, j₃, -1/2, 0, m₃), j₃, m₃) ≈ 0 atol=1e-15
    # 6-j symbol as sum over 3-j symbols
    @test @racahsum((-1)^Integer(18-m₁-m₂-m₃-m₄-m₅-m₆) * wigner3j(3, 3, 3, -m₁, -m₂, -m₃) *
                    wigner3j(3, 3, 3, m₁, -m₅, m₆) * wigner3j(3, 3, 3, m₄, m₂, -m₆) *
                    wigner3j(3, 3, 3, -m₄, m₅, m₃), m₁, m₂, m₃, m₄, m₅, m₆) ≈ wigner6j(3, 3, 3, 3, 3, 3)
    @test @racahsum((-1)^Integer(30-m₁-m₂-m₃-m₄-m₅-m₆) * wigner3j(5, 5, 5, -m₁, -m₂, -m₃) *
                    wigner3j(5, 5, 5, m₁, -m₅, m₆) * wigner3j(5, 5, 5, m₄, m₂, -m₆) *
                    wigner3j(5, 5, 5, -m₄, m₅, m₃), m₁, m₂, m₃, m₄, m₅, m₆) ≈ wigner6j(5, 5, 5, 5, 5, 5)
    @test @racahsum((-1)^Integer(9-m₁-m₂-m₃-m₄-m₅-m₆) * wigner3j(1, 3/2, 5/2, -m₁, -m₂, -m₃) *
                    wigner3j(1, 5/2, 3/2, m₁, -m₅, m₆) * wigner3j(0, 3/2, 3/2, m₄, m₂, -m₆) *
                    wigner3j(0, 5/2, 5/2, -m₄, m₅, m₃), m₁, m₂, m₃, m₄, m₅, m₆) ≈ wigner6j(1, 3/2, 5/2, 0, 5/2, 3/2)
    # 6-j symbol
    @test @racahsum(wigner6j(1,1,1,1,1,j), j) ≈ wigner6j(1,1,1,1,1,0) + wigner6j(1,1,1,1,1,1) + wigner6j(1,1,1,1,1,2)
    @test @racahsum(j*wigner6j(1,1,1,1,1,j), j) ≈ wigner6j(1,1,1,1,1,1) + 2wigner6j(1,1,1,1,1,2)
    @test @racahsum(wigner6j(1,3/2,j,0,5/2,3/2), j) ≈ wigner6j(1,3/2,5/2,0,5/2,3/2)
    # Orthogonality relation for 6-j symbols
    @test @racahsum((2j₃+1)*wigner6j(1,3/2,j₃,0,5/2,3/2)*wigner6j(1,3/2,j₃,0,5/2,3/2), j₃) ≈ 1/4
    @test @racahsum((2j₃+1)*wigner6j(1,3/2,j₃,0,5/2,3/2)*wigner6j(1,3/2,j₃,0,5/2,5/2), j₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner6j(1,1,j₃,1,1,1)*wigner6j(1,1,j₃,1,1,0), j₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner6j(1,1,j₃,1,1,1)*wigner6j(1,1,j₃,1,1,1), j₃) ≈ 1/3
    @test @racahsum((2j₃+1)*wigner6j(1,1,j₃,1,1,1)*wigner6j(1,1,j₃,1,1,2), j₃) ≈ 0 atol=1e-15
    # 9-j symbol as sum over 6-j symbols
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1,1,1,1,0,j)*wigner6j(1,1,1,1,j,1)*wigner6j(1,1,0,j,1,1), j) ≈
        wigner9j(1,1,1,1,1,1,1,1,0)
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1,5/2,3/2,3/2,0,j)*wigner6j(3/2,0,3/2,5/2,j,5/2)*
                    wigner6j(5/2,5/2,0,j,1,3/2), j) ≈ wigner9j(1,3/2,5/2,5/2,0,5/2,3/2,3/2,0)
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1,5/2,3/2,3/2,0,j)*wigner6j(3/2,0,3/2,5/2,j,5/2)*
                    wigner6j(5/2,5/2,0,j,1,3/2), j) ≈ wigner9j(1,3/2,5/2,5/2,0,5/2,3/2,3/2,0)
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1/2,3,5/2,5/2,1,j)*wigner6j(1/2,3,5/2,3,j,2)*
                    wigner6j(1,2,1,j,1/2,1/2), j) ≈ wigner9j(1/2,1/2,1,3,3,2,5/2,5/2,1)
    # 9-j symbol
    @test @racahsum(wigner9j(1,1,1,1,1,1,1,1,j), j) ≈ sum(wigner9j(1,1,1,1,1,1,1,1,j) for j=0:2)
    @test @racahsum(j^2*wigner9j(1,1,1,1,1,1,1,1,j), j) ≈ wigner9j(1,1,1,1,1,1,1,1,1) + 4wigner9j(1,1,1,1,1,1,1,1,2)
    @test @racahsum(wigner9j(1,3/2,5/2,5/2,j,5/2,3/2,3/2,0), j) ≈ sum(wigner9j(1,3/2,5/2,5/2,j,5/2,3/2,3/2,0) for j=0:3)
    @test @racahsum(wigner9j(1/2,1/2,1,j,3,j,5/2,5/2,1), j) ≈ wigner9j(1/2,1/2,1,2,3,2,5/2,5/2,1)
    # Orthogonality relation for 9-j symbols
    @test @racahsum((2j₇+1)*(2j₈+1)*wigner9j(1/2,1/2,1,3,3,2,j₇,j₈,1)*wigner9j(1/2,1/2,1,3,3,2,j₇,j₈,1), j₇, j₈) ≈ 1/15
    @test @racahsum((2j₇+1)*(2j₈+1)*wigner9j(1/2,1/2,0,3,3,1,j₇,j₈,1)*wigner9j(1/2,1/2,0,3,3,1,j₇,j₈,1), j₇, j₈) ≈ 1/3
    @test @racahsum((2j₇+1)*(2j₈+1)*wigner9j(1/2,1/2,0,3,3,1,j₇,j₈,1)*wigner9j(1/2,1/2,1,3,3,2,j₇,j₈,1), j₇, j₈) ≈ 0 atol=1e-15
    # If summation bounds cannot be determined and were not specified explicitly, sum over 0:1/2:30
    @test @racahsum(wigner6j(j,0,j,0,j,0), j) ≈ sum(wigner6j(j,0,j,0,j,0) for j=0:HalfInt(1/2):30)
    @test @racahsum(wigner6j(j,0,j,0,j,0), j=0:10) ≈ sum(wigner6j(j,0,j,0,j,0) for j=0:10)
end


tile_size = 50
board_size = 4 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]

chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
img = chessboard

###
using Random, PyPlot
using ScatteredInterpolation
using Dierckx

# params:
point_count = 100

flow = gen_tiled_flow(size(img), 30, 40)

rand_inds = rand(CartesianIndices(flow), point_count)

flow[rand_inds[1]]

## ScatteredInterpolation:
using BenchmarkTools, ScatteredInterpolation

points = reshape([ind[i] for i=1:length(rand_inds[1]) for ind=rand_inds], length(rand_inds), length(rand_inds[1]))'
samples = [flow[I] for I in rand_inds]

idx = rand_inds
@benchmark reduce((x,y)->hcat(x, [y.I...]), idx[2:end]; init=[idx[1].I...])

@benchmark reshape([ind[i] for i=1:length(rand_inds[1]) for ind=rand_inds], length(rand_inds), length(rand_inds[1]))'

# fun
flow_size = size(flow)
x1 = 1:flow_size[1]
y1 = 1:flow_size[2]

X = repeat(x1, flow_size[1])[:]
Y = repeat(y1', flow_size[1])[:]
gridPoints = [X Y]'

itp = interpolate(method, points, samples);
interpolated = ScatteredInterpolation.evaluate(itp, gridPoints)
gridded = reshape(interpolated, flow_size)

# fun end

# methods: RBFs = [Gaussian(2), Multiquadratic(2), InverseQuadratic(2), InverseMultiquadratic(2)]
function test_scatter(points, samples, smoothness; method)
    flow_size = size(flow)
    x1 = 1:flow_size[1]
    y1 = 1:flow_size[2]

    X = repeat(x1, flow_size[1])[:]
    Y = repeat(y1', flow_size[1])[:]
    gridPoints = [X Y]'

    itp = interpolate(method, points, samples);
    interpolated = ScatteredInterpolation.evaluate(itp, gridPoints)
    gridded = reshape(interpolated, flow_size)
    return gridded
end

@btime gridded_scatter = test_scatter($points, $samples, $0)
@benchmark gridded_scatter = test_scatter($points, $samples, $0, method=Shepard())
@benchmark gridded_scatter = test_scatter($points, $samples, $0)
gridded_scatter = test_scatter(points, samples, 0, method=Multiquadratic(2))
gridded_scatter = test_scatter(points, samples, 0, method=Shepard())

showflow(gridded_scatter, figtitle="scatter")

showflow(flow, figtitle="orig")

import LAP_julia: helpers.mse, helpers.angle_rmse, helpers.angle_mae, helpers.vec_len

mse(gridded_scatter, flow)
vec_len(mse(gridded_scatter, flow))
angle_rmse(gridded_scatter, flow)
angle_mae(gridded_scatter, flow)

sum(flow - gridded_scatter)/length(flow)

showflow(flow)
x = [ind[1] for ind in rand_inds]
y = [ind[2] for ind in rand_inds]
scatter(x, y, marker=:x)
gcf()

## Dierckx
using Dierckx

x = float([ind[1] for ind in rand_inds])
y = float([ind[2] for ind in rand_inds])
z1 = [real(flow[ind]) for ind in rand_inds]
z2 = [imag(flow[ind]) for ind in rand_inds]

function test_dier(x, y, z1, z2)
    storage = 1000
    spl1 = Spline2D(x, y, z1; w=ones(length(x)), kx=1, ky=1, s=storage)
    real1 = evalgrid(spl1, 1:size(flow)[1], 1:size(flow)[2])

    spl2 = Spline2D(x, y, z2; w=ones(length(x)), kx=1, ky=1, s=storage)
    imag1 = evalgrid(spl2, 1:size(flow)[1], 1:size(flow)[2])

    gridded = real1 .+ imag1 .* im
    return gridded
end

@benchmark gridded_dier = test_dier($x, $y, $z1, $z2)
gridded_dier = test_dier(x, y, z1, z2)

showflow(gridded_dier, figtitle="dier")
showflow(flow)

### IrregularInterpolate

# using PyPlot
#
# f(x, y) = (x+y-1)^2
#
# x = float([ind[1] for ind in rand_inds]) ./ 200
# y = float([ind[2] for ind in rand_inds]) ./ 200
# a = f.(x, y)
#
# p = Main.IrregularInterpolate.PolyharmonicSpline(2, [x y], a)  # 3 is the order
#
# inds = CartesianIndices(flow)
# x2 = float(vec([I[1] for I in inds])) ./ 200
# y2 = float(vec([I[2] for I in inds])) ./ 200
# data = Main.IrregularInterpolate.interpolate(p, [x2 y2])
#
# gridded = reshape(data, size(flow))
#
# imgshow(gridded)
#
#
# figure()
# plot_surface([I[1] for I in inds], [I[2] for I in inds], gridded)
# gcf()

###

using PyPlot

x = float([ind[1] for ind in rand_inds])
y = float([ind[2] for ind in rand_inds])
z1 = [real(flow[ind]) for ind in rand_inds]
z2 = [imag(flow[ind]) for ind in rand_inds]

function test_code(x, y, z1, z2)
    p1 = Main.IrregularInterpolate.PolyharmonicSpline(2, [x y], z1)  # 3 is the order
    p2 = Main.IrregularInterpolate.PolyharmonicSpline(2, [x y], z2)  # 3 is the order

    inds = CartesianIndices(flow)
    x2 = float(vec([I[1] for I in inds]))
    y2 = float(vec([I[2] for I in inds]))
    data_real = Main.IrregularInterpolate.interpolate(p1, [x2 y2])
    data_imag = Main.IrregularInterpolate.interpolate(p2, [x2 y2])

    gridded = reshape((data_real .+ data_imag .* im), size(flow))
    return gridded
end

@benchmark gridded_code = test_code($x, $y, $z1, $z2)
gridded_code = test_code(x, y, z1, z2)

showflow(gridded_code)

x = [ind[1] for ind in rand_inds]
y = [ind[2] for ind in rand_inds]
scatter(x, y, marker=:x)
gcf()

showflow(flow)

figure()
plot_surface([I[1] for I in inds], [I[2] for I in inds], imag(gridded), alpha=0.5)
gcf()

z1 = [real(point) for point in z]
z2 = [imag(point) for point in z]

### compare results

# prepare data:
flow = gen_tiled_flow(size(img), 30, 60)

# set params:
point_count = 10
rand_inds = rand(CartesianIndices(flow), point_count) # random points

# prepare inputs
points = reshape([ind[i] for i=1:length(rand_inds[1]) for ind=rand_inds], length(rand_inds), length(rand_inds[1]))'
samples = [flow[I] for I in rand_inds]
x = float([ind[1] for ind in rand_inds])
y = float([ind[2] for ind in rand_inds])
z1 = [real(flow[ind]) for ind in rand_inds]
z2 = [imag(flow[ind]) for ind in rand_inds]

# scatter
@benchmark gridded_scatter = test_scatter($points, $samples, $0)
gridded_scatter = test_scatter(points, samples, 0)
# Dierckx
@benchmark gridded_dier = test_dier($x, $y, $z1, $z2)
gridded_dier = test_dier(x, y, z1, z2)
# code
b = @benchmark gridded_code = test_code($x, $y, $z1, $z2)
gridded_code = test_code(x, y, z1, z2)

# show flows:
showflow(flow, figtitle="orig")
showflow(gridded_scatter, figtitle="scatter")
showflow(gridded_dier, figtitle="dier")
showflow(gridded_code, figtitle="code")

# add points
s = scatter(x, y, marker=:x)
gcf()

###
using DataFrames

Methods = ["Dier",
           "Scatter - Thin plate spline",
           "Scatter - Multiquadratic"]

Modes = ["Random_points",
         "Semigridded",
         "Gridded"]

Point_counts = [9,
                16,
                25,
                49,
                81,
                100,
                144]

###
# gridded
test_flow = CartesianIndices(ones(10,10))

flow_size = size(test_flow)
point_count = 10
points_in_dim = round(Int64, sqrt(point_count))
point_spacing = floor.(flow_size ./ points_in_dim)

function gen_rand_points(flow, point_count, mode; max_skew_const=4)
    if mode == "Random_points"
        ret_inds = rand(CartesianIndices(flow), point_count)
    elseif mode == "Gridded"
        points_in_dim = round(Int64, sqrt(point_count))
        point_spacing = floor.(size(flow) ./ points_in_dim)

        ret_inds = Array{CartesianIndex{2}, 1}(undef, 0)
        for ind in CartesianIndices(flow)
            if all((Tuple(ind)) .% point_spacing .== 0)
                add_ind = CartesianIndex(Tuple(ind) .- floor.(Int64, point_spacing ./ 2))
                push!(ret_inds, add_ind)
            end
        end
    elseif mode == "Semigridded"
        points_in_dim = round(Int64, sqrt(point_count))
        point_spacing = floor.(size(flow) ./ points_in_dim)

        ret_inds = Array{CartesianIndex{2}, 1}(undef, 0)
        print(size(ret_inds))
        for ind in CartesianIndices(flow)
            if all((Tuple(ind)) .% point_spacing .== 0)
                skew = (rand((-floor(Int64, point_spacing[1]/max_skew_const):floor(Int64, point_spacing[1]/max_skew_const))),
                        rand((-floor(Int64, point_spacing[1]/max_skew_const):floor(Int64, point_spacing[2]/max_skew_const))))
                add_ind = CartesianIndex(Tuple(ind) .- floor.(Int64, point_spacing ./ 2) .+ skew)
                @assert add_ind in CartesianIndices(flow)

                push!(ret_inds, add_ind)
            end
        end
    end
    return ret_inds
end

points = gen_rand_points(ones(100, 100), 40, "Semigridded")

imgshow(ones(100, 100))
x = [ind[1] for ind in points]
y = [ind[2] for ind in points]
scatter(x, y, marker=:x)
gcf()

# make Semigridded

###

scatter_interp_bench = DataFrame(Methods = )

df = DataFrame(A = 1:2, a = ["ss", "sdfd"], B = [bench, bench], B2 = [showflow(flow, ret="PyObject"), showflow(flow, ret="PyObject")], C =[ 1, 4])

display(df)

showflow(flow, figtitle="orig", ret="PyObject")

import PyCall.PyObject

scatter_interp_bench = DataFrame(
                                Index = Int[],
                                Method = String[],
                                Mode = String[],
                                Point_count = Int[],
                                Median_speed = Float64[],
                                Magnitude_of_MSE = Float64[],
                                Angle_RMS = Float64[],
                                Angle_mean = Float64[],
                                Benchmark = BenchmarkTools.Trial[],
                                Truth_flow = PyObject[],
                                Method_flow = PyObject[],
                                Scatter_points = PyObject[]
                                )

df = scatter_interp_bench
df = DataFrame(A = Int[], B = String[])

push!(df, (1, "shit"))

showtable(df)

insert!(df, 3, undef, :aaArea)

df.aaArea = [s, s]

# test
A = rand(50000, 30)
d = DataFrame(A)

tmp = d[d.x1 .> 0.7, :]
tmp === d[d.x1 .> 0.7, :]

arr = [12.3545235, 12.323, 2.123412]

mat = [12; 12]

round.(arr, digits=2)

vcat([3, 4, 5], arr)

## test try catch
arr = []

for k in 1:12
    try
        tmp = 12
        if k == 4
            @assert true == false
        end
    catch e
        push!(arr, (k, e))
        continue;
    end
    println(tmp)
    println(k)
end

try
    @assert true == false
catch e
    @warn e
    typeof(e)
end

function f(x::Type, y) where Type <: Real
    x .+ y
end

function g2(x::Type, y::mat) where {Type <:Real, mat <: Matrix{<:Complex{<:Real}}}
    x .+ y
end

@code_warntype f(12, fl)

g2(12, 12)

fl = gen_tiled_flow((12,12), 12)

typeof(fl)

point_count_df = DataFrame(Point_c = Int[],
                           Method = String[],
                           Median_speed_mean = Float64[],
                           Magnitude_of_MSE_mean = Float64[],
                           Angle_RMS_mean = Float64[],
                           Angle_Mean_mean = Float64[])

for points in Point_counts
    for method in Methods
        tmp_df = df[((df.Point_count == points) .& (df.Method .== method)), :]
        tmp_df = select(tmp_df, :Median_speed, :Magnitude_of_MSE, :Angle_RMS, :Angle_mean)
        test_count = length(tmp_df.Median_speed)
        tmp_df = describe(tmp_df)
        retrieved_data = select(tmp_df, :mean, copycols=false)
        push!(poin_count_df, [points method retrieved_data])
    end
end

bench = @benchmark 12+1

df = DataFrame(a = Matrix[])

push!(df, Dict(:a => flow))

df[:, :a][1]

typeof(flow)

df = DataFrame()

flow = gen_tiled_flow((100,100), 12)
plot = showflow(flow)

@save "save_df.JLD2" df flow


# TODO: check trees are on the same leave set

"""
    RF(tree1::T, tree2::T)::Int64 where T <:AbstractNode

Calculate the Robinson-Foulds distance between the two trees.
In its current form the function assumes the trees have identical leave sets.

Returns result of algorithm as integer.

* `tree1` : tree used to determine RF distance.

* `tree2` : tree used to determine RF distance.
"""
function RF(tree1::T, tree2::T)::Int64 where T <: GeneralNode
    initialize_tree!(tree1, height=false)
    initialize_tree!(tree2, height=false)
    r, _ = RF_int(tree1, tree2)
    return r
end

function RF_int(tree1::T, tree2::T)::Tuple{Int64, Int64} where T <: GeneralNode
    bt3 = get_bipartitions(tree1)
    bt4 = get_bipartitions(tree2)
    length(bt3)+length(bt4) - 2* length(intersect(bt3, bt4)), length(bt3)+length(bt4)
end

"""
    RF_weighted(tree1::T, tree2::T)::Float64 where T <:AbstractNode

Calculate the weighted Robinson-Foulds distance between the two trees.
The raw Robinson-Foulds distance is weighted by the maximum distance between the trees.
In its current form the function assumes the trees have identical leave sets.

Returns result of algorithm as integer.

* `tree1` : tree used to determine RF distance.

* `tree2` : tree used to determine RF distance.
"""
function RF_weighted(tree1::T, tree2::T)::Float64 where T <: GeneralNode
    initialize_tree!(tree1, height=false)
    initialize_tree!(tree2, height=false)
    rf, l1 = RF_int(tree1, tree2)
    rf / l1
end

"""
    get_bipartitions(tree::T)::Vector{Tuple} where T <:AbstractNode

Get a vector of all bipartions of `tree`.

Returns a vector containing Tuples of sets representing the bipartions.
"""
function get_bipartitions(tree::T)::Vector{Tuple} where T <:AbstractNode
    po_vect= post_order(tree)[1:end-1]
    bt = Vector{Tuple}(undef, length(po_vect))
    all_leaves = [n.name for n in get_leaves(tree)]
    for ind in eachindex(po_vect)
        elem = po_vect[ind]::T
        outset = String[]
        inset = sort([i.name for i in get_leaves(elem)])
        outset = setdiff(all_leaves, inset)
        @inbounds bt[ind] = (join(sort(inset),","), join(sort(outset),","))
    end # for
    bt
end

"""
    BHV_bounds(tree1::T, tree2::T)::Tuple{Float64, Float64} where T <:AbstractNode

This function calculates the lower and upper bounds of the geodesic in the
Billera-Holmes-Vogtman space.

Returns tuple of floats.
"""
function BHV_bounds(tree1::T, tree2::T)::Tuple{Float64, Float64} where T <:AbstractNode
    bp_t1 = get_bipartitions(tree1)
    bp_t2 = get_bipartitions(tree2)

    T1minusT2 = 0.0
    T2minusT1 = 0.0
    T1andT2 = 0.0

    for bp in bp_t1
        found1 = find_lca(tree1, String.(split(bp[1], ",")))
        if found1.root
            found1 = find_lca(tree1, String.(split(bp[2], ",")))
        end

        ind = findfirst(isequal(bp), bp_t2)
        if !isnothing(ind)
            found2 = find_lca(tree2, String.(split(bp[1], ",")))
            if found2.root
                found2 = find_lca(tree2, String.(split(bp[2], ",")))
            end
            T1andT2 += (found1.inc_length - found2.inc_length)^2
            deleteat!(bp_t2, ind)
        else
            T1minusT2 += found1.inc_length^2
        end
    end

    # only nodes which are present in the second tree are left in bp_t2
    for bp in bp_t2
        found2 = find_lca(tree2, String.(split(bp[1], ",")))
        if found2.root
            found2 = find_lca(tree2, String.(split(bp[2], ",")))
        end
        T2minusT1 += found2.inc_length^2
    end

    res_low::Float64 = T1minusT2+T2minusT1+T1andT2
    res_high::Float64 = (sqrt(T1minusT2)+sqrt(T2minusT1))^2+T1andT2

    sqrt(res_low), sqrt(res_high)
end


using Plots, LaTeXStrings, Measures; pyplot()

f(x,y) = x^2 + y^2
f0(x) = f(x,0)
f2(x) = f(x,2)

xVals, yVals = -5:0.1:5 , -5:0.1:5
plot(xVals, [f0.(xVals), f2.(xVals)],
	c=[:blue :red], xlims=(-5,5), legend=:top,
	ylims=(-5,25), ylabel=L"f(x,\cdot)", label=[L"f(x,0)" L"f(x,2)"])
p1 = annotate!(0, -0.2, text("(0,0) The minimum\n of f(x,0)", :left, :top, 10))

z = [ f(x,y) for y in yVals, x in xVals ]
p2 = surface(xVals, yVals, z, c=cgrad([:blue, :red]),legend=:none,
	ylabel="y", zlabel=L"f(x,y)")

M = z[1:10,1:10]
p3 = heatmap(M, c=cgrad([:blue, :red]), yflip=true, ylabel="y",
	xticks=([1:10;], xVals), yticks=([1:10;], yVals))

plot(p1, p2, p3, layout=(1,3), size=(1200,400), xlabel="x", margin=5mm)


"""
        grey(decisionMat, weights, fs, zeta)

    Perform GRA (Grey Relational Analysis) for a given decision matrix and weights.

# Arguments:
 - `decisionMat::Array{Float,2}`: n × m matrix of decision matrix.
 - `weights::Array{Float64, 1}`: m-vector of weights for criteria.
 - `fs::Array{Function, 1}`: m-vector of functions that are either maximize or minimize for each single criterion.
 - `zeta::Float64`: zeta parameter for the algorithm. The default is 0.5.

# Description
Applies GRA (Grey Relational Analysis).

# Output
- `::GameResult`: GreyResult object that holds many values including ordering of strategies or candidates and best index.

# Examples
```julia-repl
julia> # Decision matrix of 5 candidates and 6 criteria.
julia> df = DataFrame(
               :K1 => [105000.0, 120000, 150000, 115000, 135000],
               :K2 => [105.0, 110, 120, 105, 115],
               :K3 => [10.0, 15, 12, 20, 15],
               :K4 => [4.0, 4, 3, 4, 5],
               :K5 => [300.0, 500, 550, 600, 400],
               :K6 => [10.0, 8, 12, 9, 9]
        )
5×6 DataFrame
 Row │ K1        K2       K3       K4       K5       K6
     │ Float64   Float64  Float64  Float64  Float64  Float64
─────┼───────────────────────────────────────────────────────
   1 │ 105000.0    105.0     10.0      4.0    300.0     10.0
   2 │ 120000.0    110.0     15.0      4.0    500.0      8.0
   3 │ 150000.0    120.0     12.0      3.0    550.0     12.0
   4 │ 115000.0    105.0     20.0      4.0    600.0      9.0
   5 │ 135000.0    115.0     15.0      5.0    400.0      9.0

julia> # Direction of optimization for each single criterion
julia> functionlist = [minimum, maximum, minimum, maximum, maximum, minimum];

julia> # Weights
julia> w = [0.05, 0.20, 0.10, 0.15, 0.10, 0.40];

julia> result = grey(df, w, functionlist);

julia> result.scores
5-element Array{Float64,1}:
 0.525
 0.7007142857142857
 0.5464285714285715
 0.5762820512820512
 0.650952380952381

julia> result.bestIndex
2
```

# References
İşletmeciler, Mühendisler ve Yöneticiler için Operasyonel, Yönetsel ve Stratejik Problemlerin
Çözümünde Çok Kriterli Karar verme Yöntemleri, Editörler: Bahadır Fatih Yıldırım ve Emrah Önder,
Dora, 2. Basım, 2015, ISBN: 978-605-9929-44-8
"""
function grey(decisionMat::DataFrame, weights::Array{Float64,1}, fs::Array{Function,1}, zeta::Float64=0.5)::GreyResult

    #mat = convert(Matrix, decisionMat)
    mat = Matrix(decisionMat)

    nrows, ncols = size(mat)

    w = unitize(weights)

    referenceRow = apply_columns(fs, mat)

    normalizedReferenceRow = zeros(Float64, ncols)

    normalizedMat = similar(mat)
    for col in 1:ncols
        mmax = maximum(mat[:, col])
        mmin = minimum(mat[:, col])
        mrange = mmax - mmin
        for row in 1:nrows
            if fs[col] == maximum
                normalizedMat[row, col] = (mat[row, col] - mmin) / mrange
            elseif fs[col] == minimum
                normalizedMat[row, col] = (mmax - mat[row, col]) / mrange
            else
                @error fs[col]
                error("Function not defined")
            end
        end
        if fs[col] == maximum
            normalizedReferenceRow[col] = (referenceRow[col] - mmin) / mrange
        elseif fs[col] == minimum
            normalizedReferenceRow[col] = (mmax - referenceRow[col]) / mrange
        end
    end

    absoluteValueMat = similar(mat)
    for row in 1:nrows
        absoluteValueMat[row,:] = normalizedReferenceRow .- normalizedMat[row, :]  .|> abs
    end

    deltamin = absoluteValueMat |> minimum
    deltamax = absoluteValueMat |> maximum
    zetadeltamax = zeta * deltamax

    greytable = similar(absoluteValueMat)
    for row in 1:nrows
        for col in 1:ncols
            greytable[row, col] = (deltamin + zetadeltamax) / (absoluteValueMat[row, col] + zetadeltamax)
        end
    end

    weightedsums = zeros(Float64, nrows)
    for i in 1:nrows
        weightedsums[i] = w .* greytable[i, :] |> sum
    end

    orderings = sortperm(weightedsums)
    bestIndex = orderings |> last

    result = GreyResult(
        referenceRow,
        normalizedMat,
        absoluteValueMat,
        greytable,
        weightedsums,
        orderings,
        bestIndex
    )

    return result
end


# this file not included in package

struct PSF_gauss2D <: PSF
    σ::Float32
end

function genmodel!(s::StateFlatBg, sz::Int32, psf::PSF_gauss2D, model::Array{Float32,2})
    for ii = 1:sz
        for jj = 1:sz
            model[ii,jj] = s.bg
            for nn = 1:s.n
                model[ii,jj] += s.photons[nn] / (2 * π * psf.σ^2) *
                exp(-(ii - s.y[nn])^2 / (2 * psf.σ^2)) *
                exp(-(jj - s.x[nn])^2 / (2 * psf.σ^2))
            end
        end
    end
end

function genmodel_gauss2D_CUDA!(s_n::Int32,s_x, s_y,
    s_photons,s_bg::Float32, sz::Int32, σ::Float32, model)
     #note that the 2d array is linearized and using 1-based indexing in kernel
     ii = blockIdx().x
     jj = threadIdx().x

     idx=(ii-1)*sz+jj
     model[idx] = s_bg + 1f-4
     for nn = 1:s_n
         model[idx] += s_photons[nn] / (2 * π * σ^2) *
                 exp(-(ii - s_y[nn])^2 / (2 * σ^2)) *
                 exp(-(jj - s_x[nn])^2 / (2 * σ^2))
     end
     return nothing
end

function genmodel!(s::StateFlatBg, sz::Int32, psf::PSF_gauss2D, model::CuArray{Float32,2})
    @cuda threads=sz blocks=sz genmodel_gauss2D_CUDA!(s.n,s.x,s.y,s.photons,s.bg, sz, psf.σ, model)
end

function max2int(psf::PSF_gauss2D)
    return 2*pi*psf.σ
end


using ExprRules
using AbstractTrees

using BenchmarkTools
using Random

include("cas_theory.jl")
include("cas_simplify.jl")

grammar = @grammar begin
    Real = x | y | z | a | b | c     # symbol
    Real = -Real
    Real = Real * Real | Real + Real | Real - Real | Real / Real           # julia expression
    Real = _(Base.rand(1.0:100.0))  # special syntax, eval argument of _() at derivation time
    Real = sin(Real) | cos(Real)        # multiple rules on a single line
    Real = |(1:100)                 # same as Real = 4 | 5 | 6
end

Random.seed!(rand(UInt))

for i ∈ 1:100
    rulenode = rand(RuleNode, grammar, :Real, 10)
    println(get_executable(rulenode, grammar))
end

@profview :(4a)         == simplify(:(2a + a + a))
@profview :(a*b*c)      == simplify(:(a * c * b))
@profview :(2x)         == simplify(:(1 * x * 2))
@profview :((a*b)^2)    == simplify(:((a*b)^2))
@profview :((a*b)^6)    == simplify(:((a^2*b^2)^3))
@profview :(a+b+d)      == simplify(:(a + b + (0*c) + d))
@profview :(a+b)        == simplify(:(a + b + (c*0) + d - d))
@profview :(a)          == simplify(:((a + d) - d))
@profview :(a + b + d)  == simplify(:(a + b * c^0 + d))
@profview :(a * b * x ^ (d+y))  == @simplify a * x^y * b * x^d
@profview :(a * b * x ^ 74103)  == @simplify a * x^(12 + 3) * b * x^(42^3)
@profview 1 == @simplify (x+y)^(a*0) / (y+x)^0
@profview 2 == @simplify cos(x)^2 + 1 + sin(x)^2
@profview 2 == @simplify cos(y)^2 + 1 + sin(y)^2
@profview 2 == @simplify sin(y)^2 + cos(y)^2 + 1
@profview :(y + sec(x)^2 ) == @simplify 1 + y + tan(x)^2
@profview :(y + csc(x)^2 ) == @simplify 1 + y + cot(x)^2

@profview :(2x^3) == simplify(:( x * ∂(x^2, x) * x))


#=
module NonsmoothFwdAD
=====================
A quick implementation of:

- the vector forward mode of automatic differentiation (AD) for generalized
  derivative evaluation, developed in the article:
  KA Khan and PI Barton (2015), https://doi.org/10.1080/10556788.2015.1025400

- the "compass difference" rule for generalized derivatives of scalar-valued bivariate
  functions, developed in the article:
  KA Khan and Y Yuan (2020), https://doi.org/10.46298/jnsao-2020-6061

This implementation applies generalized differentiation rules via operator overloading,
and is modeled on the standard "smooth" vector forward AD mode implementation described
in Chapter 6 of "Evaluating Derivatives (2nd ed.)" by Griewank and Walther (2008).
The "AFloat" struct here is analogous to Griewank and Walther's "adouble" class.

The following operators have been overloaded:
x+y, -x, x-y, x*y, inv(x), x/y, x^p, exp, log, sin, cos, abs, max, min, hypot
(where "p" is a fixed integer)

Additional operators may be overloaded in the same way; a new unary operation would be
overloaded just like "log" or "abs", and a new binary operation would be overloaded
just like "x*y" or "hypot".

Written by Kamil Khan on February 5, 2022
Edited by Maha Chaudhry in April 2022
=#
module NonsmoothFwdAD

using LinearAlgebra

export AFloat

export eval_ld_derivative,
    eval_dir_derivative,
    eval_gen_derivative,
    eval_gen_gradient,
    eval_compass_difference

# default tolerances; feel free to change.
const DEFAULT_ZTOL = 1e-8   # used to decide if we're at the kink of "abs" or "hypot"

struct AFloat
    val::Float64
    dot::Vector{Float64}
    ztol::Float64  # used in "abs" and "hypot" to decide if quantities are 0
end

# outer constructors for when this.dot and/or this.ztol aren't provided explicitly
AFloat(val::Float64, dot::Vector{Float64}) = AFloat(val, dot, DEFAULT_ZTOL)
AFloat(val::Float64, n::Int, ztol::Float64 = DEFAULT_ZTOL) = AFloat(val, zeros(n), ztol)

# display e.g. with "println"
Base.show(io::IO, x::AFloat) = print(io, "(", x.val, "; ", x.dot, ")")

# define promotion from Float64 to AFloat.
# Base.convert can't be used, because it doesn't know the intended dimension
# of the new "dot" vector.
Base.promote(uA::AFloat, uB::Float64) = (uA, AFloat(uB, length(uA.dot), uA.ztol))
Base.promote(uA::Float64, uB::AFloat) = reverse(promote(uB, uA))
Base.promote(uA::AFloat, uB::AFloat) = (uA, uB)

# for u::AFloat, define "u[1]" to mean "u". Helps handle vector/scalar outputs.
function Base.getindex(u::AFloat, i::Int)
    return (i == 1) ? u : throw(DomainError("i: must be 1"))
end

## define high-level generalized differentiation operations, given a mathematical
## function f composed from supported elemental operations, and written as though
## its input is a Vector

# compute:
#  y = the function value f(x), as a Vector, and
#  yDot = the LD-derivative f'(x; xDot)
#  yDot is the same type as xDot, which is either a Matrix or a Vector{Vector}
function eval_ld_derivative(
    f::Function,
    x::Vector{Float64},
    xDot::Matrix{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    # use operator overloading to compute f(x) and f'(x; xDot)
    yLD = eval_AFloat_vec_output(f, x, xDot, ztol = ztol)

    # recover outputs from AFloats
    y = [vLD.val for vLD in yLD]
    yDot = reduce(vcat, [vLD.dot for vLD in yLD]')

    return y, yDot
end

function eval_ld_derivative(
    f::Function,
    x::Vector{Float64},
    xDot::Vector{Vector{Float64}};
    ztol::Float64 = DEFAULT_ZTOL
)
    # use operator overloading to compute f(x) and f'(x; xDot)
    yLD = eval_AFloat_vec_output(f, x, xDot, ztol = ztol)

    # recover outputs from AFloats
    y = [vLD.val for vLD in yLD]
    yDot = [vLD.dot for vLD in yLD]

    return y, yDot
end

function eval_AFloat_vec_output(
    f::Function,
    x::Vector{Float64},
    xDot::Matrix{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    # express inputs as AFloats
    xLD = [AFloat(v, Vector(vDot), ztol) for (v, vDot) in zip(x, eachrow(xDot))]

    # use operator overloading to compute f(x) and f'(x; xDot)
    yLD = f(xLD)
    if !(yLD isa Vector)
        yLD = [yLD]
    end

    return yLD
end

function eval_AFloat_vec_output(
    f::Function,
    x::Vector{Float64},
    xDot::Vector{Vector{Float64}};
    ztol::Float64 = DEFAULT_ZTOL
)
    # express inputs as AFloats
    xLD = map(AFloat, x, xDot, ztol)

    # use operator overloading to compute f(x) and f'(x; xDot)
    yLD = f(xLD)
    if !(yLD isa Vector)
        yLD = [yLD]
    end

    return yLD
end

# compute:
#   y = the function value f(x), as a vector, and
#   yDot = the directional derivative f'(x; xDot::Vector), as a vector
function eval_dir_derivative(
    f::Function,
    x::Vector{Float64},
    xDot::Vector{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    xDotMatrix = reshape(xDot, :, 1)
    (y, yDotMatrix) = eval_ld_derivative(f, x, xDotMatrix, ztol = ztol)
    yDot = vec(yDotMatrix)
    return y, yDot
end

# compute:
#   y = the function value f(x), and
#   yDeriv = the lexicographic derivative D_L f(x; xDot)
#   (xDot defaults to I if not provided, skipping the final linear solve.)
function eval_gen_derivative(
    f::Function,
    x::Vector{Float64},
    xDot::Matrix{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    (y, yDot) = eval_ld_derivative(f, x, xDot, ztol = ztol)
    yDeriv = yDot / xDot
    return y, yDeriv
end

function eval_gen_derivative(
    f::Function,
    x::Vector{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    (y, yDot) = eval_ld_derivative(f, x, Matrix{Float64}(I(length(x))), ztol = ztol)
    return y, yDot
end

# for scalar-valued f (ordinarily returning a Float64), compute:
#   y = the function value f(x), and
#   yGrad = the lexicographic gradient grad_L f(x; xDot),
#           which is a vector (and the transpose of the lex. derivative)
#   (xDot defaults to I if not provided, skipping the final linear solve.)
function eval_gen_gradient(
    f::Function,
    x::Vector{Float64},
    xDot::Matrix{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    # use operator overloading to compute f(x) and f'(x; xDot)
    yLD = eval_AFloat_vec_output(f, x, xDot, ztol = ztol)

    # compute generalized gradient element
    return yLD[1].val, (yLD[1].dot'/ xDot)'
end

function eval_gen_gradient(
    f::Function,
    x::Vector{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    # use operator overloading to compute f(x) and f'(x; xDot)
    yLD = eval_AFloat_vec_output(f, x, Matrix{Float64}(I(length(x))), ztol = ztol)

    return yLD[1].val, yLD[1].dot
end

# for a scalar-valued function f, compute:
#   y = the function value f(x), and
#   yCompass = the compass difference of f at x.
# x must be either a scalar or vector, and yCompass is the same type as x.
function eval_compass_difference(
    f::Function,
    x::Vector{Float64};
    ztol::Float64 = DEFAULT_ZTOL
)
    y = f(x)

    (length(y) == 1) ||
        throw(DomainError("f; this function is not scalar-valued"))

    fVec(u) = [f(u)[1]] # account for f returning either a Float64 or Vector{Float64}

    yCompass = copy(x)
    coordVec = zeros(length(x))
    for i in eachindex(x)
        coordVec[i] = 1.0
        _, yDotPlus = eval_dir_derivative(fVec, x, coordVec, ztol = ztol)
        _, yDotMinus = eval_dir_derivative(fVec, x, -coordVec, ztol = ztol)
        yCompass[i] = 0.5*(yDotPlus[1] - yDotMinus[1])
        coordVec[i] = 0.0
    end
    return y, yCompass
end

function eval_compass_difference(
    f::Function,
    x::Float64;
    ztol::Float64 = DEFAULT_ZTOL
)
    (y, yCompassVec) = eval_compass_difference(f, [x], ztol = ztol)
    return y, yCompassVec[1]
end

## define generalized differentiation rules for the simplest elemental operations

# macro to define (::AFloat, ::Float64) and (::Float64, ::AFloat) variants
# of a defined bivariate operation(::AFloat, ::AFloat).
macro define_mixed_input_variants(op)
    return quote
        Base.$op(uA::AFloat, uB::Float64) = $op(promote(uA, uB)...)
        Base.$op(uA::Float64, uB::AFloat) = $op(promote(uA, uB)...)
    end
end

# addition
function Base.:+(uA::AFloat, uB::AFloat)
    vVal = uA.val + uB.val
    vDot = uA.dot + uB.dot
    return AFloat(vVal, vDot, uA.ztol)
end
@define_mixed_input_variants +

# negative and subtraction
Base.:-(u::AFloat) = AFloat(-u.val, -u.dot, u.ztol)

Base.:-(uA::AFloat, uB::AFloat) = uA + (-uB)
@define_mixed_input_variants -

# multiplication
function Base.:*(uA::AFloat, uB::AFloat)
    vVal = uA.val * uB.val
    vDot = (uB.val * uA.dot) + (uA.val * uB.dot)
    return AFloat(vVal, vDot, uA.ztol)
end
@define_mixed_input_variants *

# reciprocal and division
function Base.inv(u::AFloat)
    vVal = inv(u.val)
    vDot = -u.dot / ((u.val)^2)
    return AFloat(vVal, vDot, u.ztol)
end

Base.:/(uA::AFloat, uB::AFloat) = uA * inv(uB)
@define_mixed_input_variants /

# integer powers
function Base.:^(u::AFloat, p::Int)
    if p == 0
        return 0.0*u + 1.0
    else
        vVal = (u.val)^p
        vDot = p * (u.val)^(p-1) * u.dot
        return AFloat(vVal, vDot, u.ztol)
    end
end

# exponential, logarithm, sine, and cosine
Base.exp(u::AFloat) = AFloat(exp(u.val), exp(u.val) * u.dot, u.ztol)

Base.log(u::AFloat) = AFloat(log(u.val), u.dot / u.val, u.ztol)

Base.sin(u::AFloat) = AFloat(sin(u.val), cos(u.val) * u.dot, u.ztol)

Base.cos(u::AFloat) = AFloat(cos(u.val), -sin(u.val) * u.dot, u.ztol)

# absolute value.
# When evaluating zDot, if any quantity "q" has abs(q) < u.ztol,
# then "q" is considered to be 0.
function Base.abs(u::AFloat)
    vVal = abs(u.val)
    if vVal > u.ztol
        s = sign(u.val)
    else
        s = 0.0
        for d in u.dot
            if abs(d) > u.ztol
                s = sign(d)
                break
            end
        end
    end
    vDot = s * u.dot
    return AFloat(vVal, vDot, u.ztol)
end

# bivariate max and min.
# Uses the identity max(uA,uB) = 0.5*(uA+uB+abs(uA-uB)), and similarly for min
Base.max(uA::AFloat, uB::AFloat) = 0.5*(uA + uB + abs(uA - uB))
@define_mixed_input_variants max

Base.min(uA::AFloat, uB::AFloat) = 0.5*(uA + uB - abs(uA - uB))
@define_mixed_input_variants min

# sqrt(uA^2 + uB^2); this is LinearAlgebra.hypot in Julia.
# uA.ztol is used as in abs(::AFloat)
function Base.hypot(uA::AFloat, uB::AFloat)
    vVal = hypot(uA.val, uB.val)
    if vVal > uA.ztol
        sA = uA.val / vVal
        sB = uB.val / vVal
    else
        sA = 0.0
        sB = 0.0
        for (dA, dB) in zip(uA.dot, uB.dot)
            sNorm = hypot(dA, dB)
            if sNorm > uA.ztol
                sA = dA / sNorm
                sB = dB / sNorm
                break
            end
        end
    end
    vDot = sA * uA.dot + sB * uB.dot
    return AFloat(vVal, vDot, uA.ztol)
end
@define_mixed_input_variants hypot

## overload comparisons to permit simple if-statements, but these conditions' values
## should not change under small perturbations in inputs.
Base.:<(uA::AFloat, uB::AFloat) = (uA.val < uB.val)
@define_mixed_input_variants <

Base.:>(uA::AFloat, uB::AFloat) = (uA.val > uB.val)
@define_mixed_input_variants >

Base.:>=(uA::AFloat, uB::AFloat) = (uA.val >= uB.val)
@define_mixed_input_variants >=

Base.:<=(uA::AFloat, uB::AFloat) = (uA.val <= uB.val)
@define_mixed_input_variants <=

end # module


validate(square) = 1 <= square <= 64 || throw(DomainError("Invalid value: $square"))

"""Calculate the number of grains on square `square`."""
function on_square(square)
    validate(square)
    BigInt(1) * BigInt(2)^(square - 1)
    # exp(log(BigInt(1)) + log(2) * (square - 1))
end

"""Calculate the total number of grains after square `square`."""
function total_after(square)
    validate(square)
    sum(on_square(i) for i in 1:square)
end


"""
usage:
$ julia -p auto -L buffon.jl
$ buffon(convert(Int,1e9))
$ buffon_parallel(convert(Int,1e9))
"""

using BenchmarkTools

function buffon(n_trial)
    hit_counter = zero(Bool)
    for i in 1:n_trial
        position = rand()
        φ = (rand() * π) - π / 2
        x_right = position + cos(φ) / 2
        x_left = position - cos(φ) / 2
        is_hit = (x_right >= 1 || x_left <= 0) ? 1 : 0
        hit_counter += is_hit
    end
    approx_π = n_trial/hit_counter * 2
end

using Distributed

function buffon_parallel(n_trial)
    hit_counter = @distributed (+) for i in 1:n_trial
        position = rand()
        φ = (rand() * π) - π / 2
        x_right = position + cos(φ) / 2
        x_left = position - cos(φ) / 2
        (x_right >= 1 || x_left <= 0) ? 1 : 0
    end
    approx_π = n_trial/hit_counter * 2
end

function main()
    @benchmark buffon(convert(Int,1e9))
    @benchmark buffon_parallel(convert(Int,1e9))
end

#main()


using Gadfly

plot((x,y) -> x*exp(-(x-int(x))^2-y^2), -8., 8, -2., 2)


isdefined(:FourierOptics) || include("../src/FourierOptics.jl")

module TestFourierOptics

using Base.Test

using FourierOptics
using FourierOptics.CoordinateTransforms
using FourierOptics.Regions

Base.isapprox(x::NTuple{2,Real}, y::NTuple{2,Real}; kwds...) =
    (isapprox(x[1], y[1]; kwds...) && isapprox(x[2], y[2]; kwds...))

Base.isapprox(x::NTuple{4,Real}, y::NTuple{4,Real}; kwds...) =
    (isapprox(x[1], y[1]; kwds...) && isapprox(x[2], y[2]; kwds...) &&
     isapprox(x[3], y[3]; kwds...) && isapprox(x[4], y[4]; kwds...))

@testset "Coordinate Transforms" begin
    stp, x0, y0 = 1.2, -2.7, 3.1
    t = (0.7, -1.8)
    p = (0.9, 11.0)
    α = 0.17
    I = CoordinateTransform()
    A = CoordinateTransform(stp, x0, y0)
    B = inv(A)
    @test origin(A) ≈ (x0, y0)
    @test step(A) ≈ stp
    @test A ≈ (x0,y0) + stp*I
    @test A∘B ≈ I  atol=4*eps(Float64)
    @test B∘A ≈ I  atol=4*eps(Float64)
    @test A/A ≈ I  atol=4*eps(Float64)
    @test A\A ≈ I  atol=4*eps(Float64)
    B = t + A
    @test B(p) ≈ t .+ A(p)
    @test (t - A)(p) ≈ t .- A(p)
    @test (A + t)(p) ≈ A(p .+ t)
    @test (A - t)(p) ≈ A(p .- t)
    @test (α*A)(p) ≈ α.*A(p)
    @test (A*α)(p) ≈ A(α.*p)
end

@testset "Regions" begin
    stp, xc, yc = 1.2, -2.7, 3.1
    n1, n2 = 27, 52
    c = 0.7, -1.1
    grdxmin = xc - (n1 - 1)*stp/2
    grdxmax = xc + (n1 - 1)*stp/2
    grdymin = yc - (n2 - 1)*stp/2
    grdymax = yc + (n2 - 1)*stp/2
    boxxmin = xc - n1*stp/2
    boxxmax = xc + n1*stp/2
    boxymin = yc - n2*stp/2
    boxymax = yc + n2*stp/2
    X = linspace(grdxmin, grdxmax, n1)
    Y = linspace(grdymin, grdymax, n2)
    R1 = Region(X, Y)
    @test center(R1) ≈ (xc, yc)
    @test step(R1) ≈ stp
    @test size(R1) ≈ (n1, n2)
    @test extrema(R1) ≈ (grdxmin, grdxmax, grdymin, grdymax)
    @test boundingbox(R1) ≈ (boxxmin, boxxmax, boxymin, boxymax)
    R2 = recenter(R1)
    @test center(R2) ≈ (0,0)
    R2 = recenter(R1, c)
    @test center(R2) ≈ c
    @test center(R2 - c) ≈ (0,0)
end

end


function fast_pca!(X::Matrix{T}, λ::Vector{T}, P::Matrix{T}, n::Int) where T<:AbstractFloat
    r = sum(λ .> 0)
    s = n - r

    if r == 0
        X .= zeros(T, n, n)
    elseif r == n
        return nothing
    elseif r == 1
        X .= (λ[1] * λ[1]) * (P[:,1] * P[:,1]')
    elseif r ≤ s
        P₁   = @view P[:, 1:r]
        λ₁   = sqrt.(λ[1:r])
        P₁λ₁ = P₁ .* λ₁'
        X   .= P₁λ₁ * P₁λ₁'
    else
        P₂   = @view P[:, (r+1):n]
        λ₂   = sqrt.(-λ[(r+1):n])
        P₂λ₂ = P₂ .* λ₂'
        X   .= X .+ P₂λ₂ * P₂λ₂'
    end
    nothing
end

"""
    cor_fastPD!(R::Matrix{<:AbstractFloat}[, τ=1e-6])

Same as [`cor_fastPD`](@ref), but saves space by overwriting the input `R`
instead of creating a copy.

See also: [`cor_fastPD`](@ref), [`cor_nearPD`](@ref)

# Examples
```jldoctest
julia> import LinearAlgebra: isposdef

julia> r = [1.00 0.82 0.56 0.44; 0.82 1.00 0.28 0.85; 0.56 0.28 1.00 0.22; 0.44 0.85 0.22 1.00]
4×4 Array{Float64,2}:
 1.0   0.82  0.56  0.44
 0.82  1.0   0.28  0.85
 0.56  0.28  1.0   0.22
 0.44  0.85  0.22  1.0

julia> isposdef(r)
false

julia> cor_fastPD!(r)

julia> isposdef(r)
true
```
"""
function cor_fastPD!(R::Matrix{<:AbstractFloat}, τ=1e-6)
    n  = size(R, 1)
    τ  = max(eps(eltype(R)), τ)

    R .= Symmetric(R, :U)
    R[diagind(R)] .= (one(eltype(R)) - τ)

    λ, P = eigen(R)
    λ   .= reverse(λ)
    P   .= reverse(P, dims=2)

    fast_pca!(R, λ, P, n)

    R[diagind(R)] .+= τ
    R .= cov2cor(R)
    nothing
end

"""
    cor_fastPD(R::Matrix{<:AbstractFloat}[, τ=1e-6])

Return a positive definite correlation matrix that is close to `R`. `τ` is a
tuning parameter that controls the minimum eigenvalue of the resulting matrix.
`τ` can be set to zero if only a positive semidefinite matrix is needed.

See also: [`cor_fastPD!`](@ref), [`cor_nearPD`](@ref)

# Examples
```jldoctest
julia> import LinearAlgebra: isposdef

julia> r = [1.00 0.82 0.56 0.44; 0.82 1.00 0.28 0.85; 0.56 0.28 1.00 0.22; 0.44 0.85 0.22 1.00]
4×4 Array{Float64,2}:
 1.0   0.82  0.56  0.44
 0.82  1.0   0.28  0.85
 0.56  0.28  1.0   0.22
 0.44  0.85  0.22  1.0

julia> isposdef(r)
false

julia> p = cor_fastPD(r)
4×4 Array{Float64,2}:
 1.0       0.817095  0.559306  0.440514
 0.817095  1.0       0.280196  0.847352
 0.559306  0.280196  1.0       0.219582
 0.440514  0.847352  0.219582  1.0

julia> isposdef(p)
true
```
"""
function cor_fastPD(R::Matrix{<:AbstractFloat}, τ=1e-6)
    X = copy(R)
    cor_fastPD!(X, τ)
    X
end
function cor_fastPD(R::Matrix{Float16}, τ=1e-6)
    X = Matrix{Float32}(R)
    cor_fastPD!(X, τ)
    Matrix{Float16}(X)
end


# ======================
# Univariate functions
# ======================

# (*) can be checked by hand
DAISY["sine"] = Dict("dom" => a,
    "numvars" => 1,
    "f" => (x -> sin(x)),
    "ref" => Interval(-1.0, 0.977530117665097) # (*)
    )

DAISY["bspline0"] = Dict("dom" => a,
    "numvars" => 1,
    "f" => (x -> (1 - x) * (1 - x) * (1 - x) / 6.0),
    "ref" => Interval(0.36616666666666675, 27.729166666666668) # (*)
    )

DAISY["bspline1"] = Dict("dom" => a,
   "numvars" => 1,
    "f" => (x -> (3*x*x*x - 6*x*x + 4) / 6.0),
    "ref" => Interval(-65.14583333333333, 0.5631666666666667) # (*)
    )

DAISY["bspline2"] = Dict("dom" => a,
    "numvars" => 1,
    "f" => (x -> (-3*x*x*x  + 3*x*x + 3*x + 1) / 6.0),
    "ref" => Interval(0.07516666666666667, 53.604166666666664) # (*)
    )

DAISY["bspline3"] = Dict("dom" => a,
    "numvars" => 1,
    "f" => (x -> -x*x*x / 6.0),
    "ref" => Interval(0.0045, 15.1875) # (*)
    )


export centerofmass, findmid

# find center of mass in 2D image
function centerofmass(img)
    x_pos = []
    y_pos = []
    for i = 1:size(img, 3)
        tmp = img[:, :, i]
        x_data = mean(tmp, dims = 2)
        push!(x_pos, findmid(x_data))
        y_data = mean(tmp, dims = 1)
        push!(y_pos, findmid(y_data))
    end
    return x_pos, y_pos
end

# find weighted center in axis
function findmid(data)
    weighted_sum = 0
    int_sum = 0
    for i = 1:length(data)
        weighted_sum += data[i] * i
        int_sum += data[i]
    end
    return weighted_sum / int_sum
end


@testset "Show" begin
    Mod.@polyvar x y z
    @test sprint(show, (x*y^2 + x + 1 + y)) == "xy² + x + y + 1"
    @test sprint(show, (x + 1 + y) / x^2) == "(x + y + 1) / (x²)"
    @test sprint(show, (x - y - x + y) / (x^2 - x)) == "(0) / (x² - x)"
    cp = CustomPoly(1 + x)
    @test sprint(show, cp) == "$(typeof(cp))(x + 1)"
    # Test taken from TypedPolynomials
    @test sprint(show, x) == "x"
    @test sprint(show, x^0) == "1"
    @test sprint(show, x^1) == "x"
    @test sprint(show, x^2) == "x²"
    @test sprint(show, 1x^2) == "x²"
    @test sprint(show, 5x) == "5x"
    @test sprint(show, x * y) == "xy"
    @test sprint(show, y * x) == "xy"
    @test sprint(show, 5 + y + x) == "x + y + 5"
    @test sprint(show, y + 5 + x) == "x + y + 5"
    @test sprint(show, x + x^2) == "x² + x"
    @test sprint(show, x^2 + x) == "x² + x"
    @test sprint(show, x^2 - 3.0x) == "x² - 3.0x"
    @test sprint(show, -x^3) == "-x³"
    @test sprint(show, -x^3 + y^2 - x) == "-x³ + y² - x"
    @test sprint(show, -2.0x^2) == "-2.0x²"
    @test sprint(show, (1.0 + 3.1im) * x*z) == "(1.0 + 3.1im)xz"
    @test sprint(show, -(1.0 + 3.1im) * z*x) == "(-1.0 - 3.1im)xz"
    @test sprint(show, x^2 + (1.0 + 3.1im) * x) == "x² + (1.0 + 3.1im)x"
    @test sprint(show, x^2 - (1.0 + 3.1im) * x) == "x² + (-1.0 - 3.1im)x"

    Mod.@polyvar x[0:9]
    @test sprint(show, sum(i*x[i]^i for i=1:10)) == "10x₉¹⁰ + 9x₈⁹ + 8x₇⁸ + 7x₆⁷ + 6x₅⁶ + 5x₄⁵ + 4x₃⁴ + 3x₂³ + 2x₁² + x₀"
    @test sprint(show, "text/latex", sum(i*x[i]^i for i=1:10)) == "\$\$ 10x_{9}^{10} + 9x_{8}^{9} + 8x_{7}^{8} + 7x_{6}^{7} + 6x_{5}^{6} + 5x_{4}^{5} + 4x_{3}^{4} + 3x_{2}^{3} + 2x_{1}^{2} + x_{0} \$\$"
    @test sprint(show, "text/latex", (x[2] + 1) / x[3]^2) == "\$\$ \\frac{x_{1} + 1}{x_{2}^{2}} \$\$"

    Mod.@polyvar x[1:11]
    @test sprint(show, "text/latex", x[10]) == "\$\$ x_{10} \$\$"
    @test sprint(show, x[10]) == "x₁₀"

    @test sprint(print, 2x[1]^2+3x[3]+1+x[4]) == "2*x[1]^2 + 3*x[3] + x[4] + 1"
end


function read_dots_folds(input_file)
    coords = []
    folds = []
    for line in readlines(input_file)
        if ',' in line
            x,y = map(x->parse(Int,x), split(line, ','))
            push!(coords, [x,y])
        elseif '=' in line
            axis, value = split(split(line)[3], '=')
            value = parse(Int, value)
            push!(folds, [axis, value])
        end
    end
    return coords, folds
end

function fold(input_coords, fold_axis, fold_value)
    println("folding ",fold_axis, " ", fold_value)
    output_coords = []
    for coord in input_coords
        if fold_axis == "y" #fold up
            if coord[2] < fold_value
                push!(output_coords, coord)
            else
                folded_coord = [coord[1], 2*fold_value - coord[2]]
                push!(output_coords, folded_coord)
            end
        else # fold left
            if coord[1] < fold_value
                push!(output_coords, coord)
            else
                folded_coord = [2*fold_value - coord[1], coord[2]]
                push!(output_coords, folded_coord)
            end
        end
    end
    # deduplicate by putting into a set
    return Set(output_coords)
end

function test_folds()
    test_coords = [[3,4],[11,18]]
    @assert fold(test_coords, "x", 10) == Set([[3,4],[9,18]])
    @assert fold(test_coords, "y", 10) == Set([[3,4],[11,2]])
    new_dots = fold(test_coords, "y", 10)
    @assert fold(new_dots, "x", 6) == Set([[3,4],[1,2]])
end

test_folds()

function do_all_folds(coords, folds)
    for f in folds
        coords = fold(coords, f[1], f[2])
        println("number of coords ", length(coords))
    end
    return coords
end

function draw_grid(coords)
    xmax = maximum(x->x[1], coords)+1
    ymax = maximum(x->x[2], coords)+1
    grid = ones(String, ymax, xmax)

    # now draw on the coords
    for coord in coords
        grid[coord[2]+1,coord[1]+1] = "#"
    end
    return grid

end

#
input_dots, folds = read_dots_folds("input.txt")
#
output_dots = do_all_folds(input_dots, folds)
#
println("number of output dots ", length(output_dots))
#
draw_grid(output_dots)