import EnergyExpressions: QuantumOperator, NBodyTermFactor, NBodyTerm,
    OrbitalMatrixElement, NBodyMatrixElement
import AngularMomentumAlgebra: radial_integral

function ref_exchange_multipoles(a, b)
    # Table 1⁶ of Condon & Shortley. As noted by in ref [10] of
    #
    # - Cohl, H. S., Rau, A. R. P., Tohline, J. E., Browne, D. A., Cazes,
    #   J. E., & Barnes, E. I. (2001). Useful alternative to the multipole
    #   expansion of1/rpotentials. Physical Review A, 64(5),
    #   052509. http://dx.doi.org/10.1103/physreva.64.052509
    #
    # df, k = 5 is in error, and the numerator has to multiplied by 5
    # and the denominator replaced by 7623.
    reference = Dict((o"ks", o"kp") => ((1, 3, Dict((0,1) => 1,
                                                    (0,0) => 1)),),
                     (o"ks", o"kf") => ((3, 7, Dict((0,3) => 1,
                                                    (0,2) => 1,
                                                    (0,1) => 1,
                                                    (0,0) => 1)),),
                     (o"kp", o"kd") => ((1, 15, Dict((1,2) => 6,
                                                     (1,1) => 3,
                                                     (1,0) => 1,
                                                     (0,2) => 0,
                                                     (0,1) => 3,
                                                     (0,0) => 4,
                                                     (1,-2) => 0,
                                                     (1,-1) => 0)),
                                        (3, 245, Dict((1,2) => 3,
                                                      (1,1) => 9,
                                                      (1,0) => 18,
                                                      (0,2) => 15,
                                                      (0,1) => 24,
                                                      (0,0) => 27,
                                                      (1,-2) => 45,
                                                      (1,-1) => 30))),
                     (o"kd", o"kf") => ((1, 35, Dict((2,3) => 15,
                                                     (2,2) => 5,
                                                     (2,1) => 1,
                                                     (2,0) => 0,
                                                     (1,3) => 0,
                                                     (1,2) => 10,
                                                     (1,1) => 8,
                                                     (1,0) => 3,
                                                     (0,3) => 0,
                                                     (0,2) => 0,
                                                     (0,1) => 6,
                                                     (0,0) => 9,
                                                     (2,-3) => 0,
                                                     (2,-2) => 0,
                                                     (2,-1) => 0,
                                                     (1,-3) => 0,
                                                     (1,-2) => 0,
                                                     (1,-1) => 0)),
                                        (3, 315, Dict((2,3) => 10,
                                                      (2,2) => 20,
                                                      (2,1) => 24,
                                                      (2,0) => 20,
                                                      (1,3) => 25,
                                                      (1,2) => 15,
                                                      (1,1) => 2,
                                                      (1,0) => 2,
                                                      (0,3) => 25,
                                                      (0,2) => 0,
                                                      (0,1) => 9,
                                                      (0,0) => 16,
                                                      (2,-3) => 0,
                                                      (2,-2) => 0,
                                                      (2,-1) => 10,
                                                      (1,-3) => 0,
                                                      (1,-2) => 25,
                                                      (1,-1) => 15)),
                                        (5, 7623, Dict((2,3) => 5,
                                                       (2,2) => 25,
                                                       (2,1) => 75,
                                                       (2,0) => 175,
                                                       (1,3) => 35,
                                                       (1,2) => 120,
                                                       (1,1) => 250,
                                                       (1,0) => 400,
                                                       (0,3) => 140,
                                                       (0,2) => 315,
                                                       (0,1) => 450,
                                                       (0,0) => 500,
                                                       (2,-3) => 1050,
                                                       (2,-2) => 630,
                                                       (2,-1) => 350,
                                                       (1,-3) => 420,
                                                       (1,-2) => 560,
                                                       (1,-1) => 525))),
                     (o"ks", o"ks") => ((0, 1, Dict((0,0) => 1)),),
                     (o"ks", o"kd") => ((2, 5, Dict((0,2) => 1,
                                                    (0,1) => 1,
                                                    (0,0) => 1)),),
                     (o"kp", o"kp") => ((0, 1, Dict((1,1) => 1,
                                                    (1,0) => 0,
                                                    (0,0) => 1,
                                                    (1,-1) => 0)),
                                        (2, 25, Dict((1,1) => 1,
                                                     (1,0) => 3,
                                                     (0,0) => 4,
                                                     (1,-1) => 6))),
                     (o"kp", o"kf") => ((2, 175, Dict((1,3) => 45,
                                                      (1,2) => 30,
                                                      (1,1) => 18,
                                                      (1,0) => 9,
                                                      (0,3) => 0,
                                                      (0,2) => 15,
                                                      (0,1) => 24,
                                                      (0,0) => 27,
                                                      (1,-3) => 0,
                                                      (1,-2) => 0,
                                                      (1,-1) => 3)),
                                        (4, 189, Dict((1,3) => 1,
                                                      (1,2) => 3,
                                                      (1,1) => 6,
                                                      (1,0) => 10,
                                                      (0,3) => 7,
                                                      (0,2) => 12,
                                                      (0,1) => 15,
                                                      (0,0) => 16,
                                                      (1,-3) => 28,
                                                      (1,-2) => 21,
                                                      (1,-1) => 15))),
                     (o"kd", o"kd") => ((0, 1, Dict((2,2) => 1,
                                                    (2,1) => 0,
                                                    (2,0) => 0,
                                                    (1,1) => 1,
                                                    (1,0) => 0,
                                                    (0,0) => 1,
                                                    (2,-2) => 0,
                                                    (2,-1) => 0,
                                                    (1,-1) => 0)),
                                        (2, 49, Dict((2,2) => 4,
                                                     (2,1) => 6,
                                                     (2,0) => 4,
                                                     (1,1) => 1,
                                                     (1,0) => 1,
                                                     (0,0) => 4,
                                                     (2,-2) => 0,
                                                     (2,-1) => 0,
                                                     (1,-1) => 6)),
                                        (4, 441, Dict((2,2) => 1,
                                                      (2,1) => 5,
                                                      (2,0) => 15,
                                                      (1,1) => 16,
                                                      (1,0) => 30,
                                                      (0,0) => 36,
                                                      (2,-2) => 70,
                                                      (2,-1) => 35,
                                                      (1,-1) => 40))),
                     (o"kf", o"kf") => ((0, 1, Dict((3,3) => 1,
                                                    (3,2) => 0,
                                                    (3,1) => 0,
                                                    (3,0) => 0,
                                                    (2,2) => 1,
                                                    (2,1) => 0,
                                                    (2,0) => 0,
                                                    (1,1) => 1,
                                                    (1,0) => 0,
                                                    (0,0) => 1,
                                                    (3,-3) => 0,
                                                    (3,-2) => 0,
                                                    (3,-1) => 0,
                                                    (2,-2) => 0,
                                                    (2,-1) => 0,
                                                    (1,-1) => 0)),
                                        (2, 225, Dict((3,3) => 25,
                                                      (3,2) => 25,
                                                      (3,1) => 10,
                                                      (3,0) => 0,
                                                      (2,2) => 0,
                                                      (2,1) => 15,
                                                      (2,0) => 20,
                                                      (1,1) => 9,
                                                      (1,0) => 2,
                                                      (0,0) => 16,
                                                      (3,-3) => 0,
                                                      (3,-2) => 0,
                                                      (3,-1) => 0,
                                                      (2,-2) => 0,
                                                      (2,-1) => 0,
                                                      (1,-1) => 24)),
                                        (4, 1089, Dict((3,3) => 9,
                                                       (3,2) => 30,
                                                       (3,1) => 54,
                                                       (3,0) => 63,
                                                       (2,2) => 49,
                                                       (2,1) => 32,
                                                       (2,0) => 3,
                                                       (1,1) => 1,
                                                       (1,0) => 15,
                                                       (0,0) => 36,
                                                       (3,-3) => 0,
                                                       (3,-2) => 0,
                                                       (3,-1) => 42,
                                                       (2,-2) => 70,
                                                       (2,-1) => 14,
                                                       (1,-1) => 40)),
                                        (6, 7361*64, Dict((3,3) => 1,
                                                          (3,2) => 7,
                                                          (3,1) => 28,
                                                          (3,0) => 84,
                                                          (2,2) => 36,
                                                          (2,1) => 105,
                                                          (2,0) => 224,
                                                          (1,1) => 225,
                                                          (1,0) => 350,
                                                          (0,0) => 400,
                                                          (3,-3) => 924,
                                                          (3,-2) => 462,
                                                          (3,-1) => 210,
                                                          (2,-2) => 504,
                                                          (2,-1) => 378,
                                                          (1,-1) => 420))))

    g = CoulombInteraction()
    terms = NBodyTerm[]

    push_multipole!(k, coeff) =
        push!(terms, NBodyTerm([radial_integral([a,b], (k,g), [b,a])], coeff))

    ao,bo = Orbital(:k, a.orb.ℓ),Orbital(:k, b.orb.ℓ)
    i,j = ao ≤ bo ? (1,2) : (2,1)
    uv = (ao,bo)[[i,j]]

    m′,m = (a.m[1],b.m[1])[[i,j]]
    ao == bo && abs(m) > abs(m′) && ((m′,m) = (m,m′))
    # The table only stores the upper sign of the projection quantum
    # numbers, since the opposite signs have the same values.
    if iszero(m′)
        m = abs(m)
    else
        m *= sign(m′)
    end
    m′ = abs(m′)

    for (k,D,Ns) in reference[uv]
        N = Ns[(m′,m)]
        iszero(N) && continue
        push_multipole!(k, N/D)
    end

    NBodyMatrixElement(terms)
end

function ref_direct_multipoles(a, b)
    # Table 2⁶ of Condon & Shortley
    reference = Dict((o"kp", o"kp") => ((2, 25, Dict((1,1) => 1,
                                                     (1,0) => -2,
                                                     (0,0) => 4)),),
                     (o"kp", o"kd") => ((2, 35, Dict((1,2) => 2,
                                                     (1,1) => -1,
                                                     (1,0) => -2,
                                                     (0,2) => -4,
                                                     (0,1) => 2,
                                                     (0,0) => 4)),),
                     (o"kp", o"kf") => ((2, 75, Dict((1,3) => 5,
                                                     (1,2) => 0,
                                                     (1,1) => -3,
                                                     (1,0) => -4,
                                                     (0,3) => -10,
                                                     (0,2) => 0,
                                                     (0,1) => 6,
                                                     (0,0) => 8)),),
                     (o"kd", o"kd") => ((2, 49, Dict((2,2) => 4,
                                                     (2,1) => -2,
                                                     (2,0) => -4,
                                                     (1,1) => 1,
                                                     (1,0) => 2,
                                                     (0,0) => 4)),
                                        (4, 441, Dict((2,2) => 1,
                                                      (2,1) => -4,
                                                      (2,0) => 6,
                                                      (1,1) => 16,
                                                      (1,0) => -24,
                                                      (0,0) => 36))),
                     (o"kd", o"kf") => ((2, 105, Dict((2,3) => 10,
                                                      (2,2) => 0,
                                                      (2,1) => -6,
                                                      (2,0) => -8,
                                                      (1,3) => -5,
                                                      (1,2) => 0,
                                                      (1,1) => 3,
                                                      (1,0) => 4,
                                                      (0,3) => -10,
                                                      (0,2) => 0,
                                                      (0,1) => 6,
                                                      (0,0) => 8)),
                                        (4, 693, Dict((2,3) => 3,
                                                      (2,2) => -7,
                                                      (2,1) => 1,
                                                      (2,0) => 6,
                                                      (1,3) => -12,
                                                      (1,2) => 28,
                                                      (1,1) => -4,
                                                      (1,0) => -24,
                                                      (0,3) => 18,
                                                      (0,2) => -42,
                                                      (0,1) => 6,
                                                      (0,0) => 36))),
                     (o"kf", o"kf") => ((2, 225, Dict((3,3) => 25,
                                                      (3,2) => 0,
                                                      (3,1) => -15,
                                                      (3,0) => -20,
                                                      (2,2) => 0,
                                                      (2,1) => 0,
                                                      (2,0) => 0,
                                                      (1,1) => 9,
                                                      (1,0) => 12,
                                                      (0,0) => 16)),
                                        (4, 1089, Dict((3,3) => 9,
                                                       (3,2) => -21,
                                                       (3,1) => 3,
                                                       (3,0) => 18,
                                                       (2,2) => 49,
                                                       (2,1) => -7,
                                                       (2,0) => -42,
                                                       (1,1) => 1,
                                                       (1,0) => 6,
                                                       (0,0) => 36)),
                                        (6, 7361*64, Dict((3,3) => 1,
                                                          (3,2) => -6,
                                                          (3,1) => 15,
                                                          (3,0) => -20,
                                                          (2,2) => 36,
                                                          (2,1) => -90,
                                                          (2,0) => 120,
                                                          (1,1) => 225,
                                                          (1,0) => -300,
                                                          (0,0) => -400))))

    g = CoulombInteraction()
    terms = NBodyTerm[]

    push_multipole!(k, coeff) =
        push!(terms, NBodyTerm([radial_integral([a,b], (k,g), [a,b])], coeff))

    # The monopole is always present in the multipole expansion of the
    # direct interaction.
    push_multipole!(0, 1.0)

    ao,bo = Orbital(:k, a.orb.ℓ),Orbital(:k, b.orb.ℓ)
    i,j = ao ≤ bo ? (1,2) : (2,1)
    uv = (ao,bo)[[i,j]]
    if uv ∈ keys(reference)
        ref = reference[uv]

        m′m = (abs(a.m[1]), abs(b.m[1]))
        m′,m = uv[1] < uv[2] ? m′m[[i,j]] : reverse(minmax(m′m...))

        for (k,D,Ns) in ref
            N = Ns[(m′,m)]
            iszero(N) && continue
            push_multipole!(k, N/D)
        end
    end

    NBodyMatrixElement(terms)
end


#!/usr/bin/env julia

using SpectralElements, OrdinaryDiffEq, Flux

using Plots, LinearAlgebra, SparseArrays

#--------------------------------------#
# setup
#--------------------------------------#
nx1 = 8; Ex = 2;
ny1 = 8; Ey = 2;
#--------------------------------------#
# nodal operators
#--------------------------------------#
zr1,wr1 = SpectralElements.FastGaussQuadrature.gausslobatto(nx1)
zs1,ws1 = SpectralElements.FastGaussQuadrature.gausslobatto(ny1)

Dr1 = derivMat(zr1); Ds1 = derivMat(zs1)

#----------------------------------------------------------------------#
# local operators
#----------------------------------------------------------------------#
wx1 = kron(ones(Ex,1),wr1); wy1 = kron(ones(Ey,1),ws1);

Iex = Matrix(I,Ex,Ex);
Iey = Matrix(I,Ey,Ey);

Dx1 = kron(Iex,Dr1); Dy1 = kron(Iey,Ds1);

#----------------------------------------------------------------------#
# boundary conditions
#----------------------------------------------------------------------#
ifperiodicX = false
ifperiodicY = false

Ix1 = Matrix(I,Ex*nx1,Ex*nx1);
Iy1 = Matrix(I,Ey*ny1,Ey*ny1);

Rx1 = Ix1[2:end-1,:];
Ry1 = Iy1[2:end-1,:];

if(ifperiodicX) Rx1 = Ix1; end
if(ifperiodicY) Ry1 = Iy1; end

M1 = diag(Rx1'*Rx1) * diag(Ry1'*Ry1)';

Ix1 = nothing
Iy1 = nothing
Rx1 = nothing
Ry1 = nothing
#----------------------------------------------------------------------#
# mapping
#----------------------------------------------------------------------#

# Q: global -> local op, Q': local -> global
Qx1 = semq(Ex,nx1,ifperiodicX)
Qy1 = semq(Ey,ny1,ifperiodicY)

# gather scatter op
QQtx1 = Qx1*Qx1';
QQty1 = Qy1*Qy1';

# mult
mult = ones(nx1*Ex,ny1*Ey);
mult = gatherScatter(mult,QQtx1,QQty1);
mult = @. 1 / mult;

mult1 = reshape(mult,:)

#----------------------------------------------------------------------#
# geometry
#----------------------------------------------------------------------#
x1e,_ = semmesh(Ex,nx1); y1e,_ = semmesh(Ey,ny1);
x1,y1 = ndgrid(x1e,y1e);

# deform grid with gordonhall
#x1 = @. 0.5 * (x1 + 1); y1 = @. 0.5 * (y1 + 1);

Jac1,Jaci1,rx1,ry1,sx1,sy1 = jac(x1,y1,Dx1,Dy1);

#----------------------------------------------------------------------#
# case setup
#----------------------------------------------------------------------#
kx=1.0
ky=1.0
kt=1.0
ν =1.0

utrue(x,y,t) = sin(kx*pi*x)*sin.(ky*pi*y)*cos(kt*pi*t)
setIC(x,y)   = utrue(x,y,0.0)

function setForcing(x,y,t)
    ut = utrue(x,y,t)
    f  = ut*((kx^2+ky^2)*pi^2*ν)
    f -= sin(kx*pi*x)*sin(ky*pi*y)*sin(kt*pi*t)*(kt*pi)
end

# callback
function cond(u,t,integrator)
    istep = integrator.iter
    cond = (istep % 1) == 0
    return cond
end

function affect!(integrator)
    istep = integrator.iter
    u = integrator.u
    p = integrator.p
    t = integrator.t
    dt = integrator.dt

    ut = utrue.(xx,yy,t)
    er = norm(ut-u,Inf)
    println("Step=$istep, Time=$t, dt=$dt, er=$er")

    ul  = reshape(Q*Rg'*u ,nx1*Ex,ny1*Ey)
#   dul = reshape(Q*Rg'*du,nx1*Ex,ny1*Ey)
    plt = meshplt(x1,y1,ul)
    plt = plot!(zlims=(-1,1))
    display(plt)
    return
end

cb = DiscreteCallback(cond,affect!,save_positions=(false,false))

#----------------------------------------------------------------------#
# setup
#----------------------------------------------------------------------#

# diagonal mass matrix
B1  = Jac1 .* (wx1*wy1');
Bi1 = 1 ./ B1;

# set up matrices for Laplace op.
G11 = @. ν * B1 * (rx1 * rx1 + ry1 * ry1);
G12 = @. ν * B1 * (rx1 * sx1 + ry1 * sy1);
G22 = @. ν * B1 * (sx1 * sx1 + sy1 * sy1);
#----------------------------------------------------------------------#
# Linear System
#----------------------------------------------------------------------#

Ixl = Matrix(I,Ex*nx1,Ex*nx1);              # Identity on local vectors
Iyl = Matrix(I,Ey*ny1,Ey*ny1);
Ixg = Matrix(I,Ex*(nx1-1)+1,Ex*(nx1-1)+1);  # Identity on global vectors
Iyg = Matrix(I,Ey*(ny1-1)+1,Ey*(ny1-1)+1);
Rxl = Ixl[2:end-1,:]; Rxg = Ixg[2:end-1,:]; # restriction on loc/gl vectors
Ryl = Iyl[2:end-1,:]; Ryg = Iyg[2:end-1,:];
Rl  = kron(Ryl,Rxl)
Rg  = kron(Ryg,Rxg)
Q   = kron(Qy1,Qx1)
Dr  = kron(Iyl,Dx1);
Ds  = kron(Dy1,Ixl);
Drs = [Dr;Ds];
B   = Diagonal(reshape(B1 ,:))
g11 = Diagonal(reshape(G11,:))
g12 = Diagonal(reshape(G12,:))
g22 = Diagonal(reshape(G22,:))
G   = [g11 g12; g12 g22]
A   = Drs' * G * Drs

# local vectors
xl  = reshape(x1,:)
yl  = reshape(y1,:)

# global
xx = Rg * Q' * xl
yy = Rg * Q' * yl

#----------------------------------------------------------------------#
# full rank global system
AA = Rg * Q' * A * Q * Rg'
BB = Rg * Q' * B * Q * Rg'
Bb = Rg * Q' * B

Big = Rg * Q' * reshape(Bi1,:)
BBi = Diagonal(Big)
precondB = BBi \ LinearAlgebra.I

#----------------------------------------------------------------------#
"""
 mass matrix form
"""

function Mdudt!(Mdu,u,p,t)
    f = setForcing.(xl,yl,t)
    Mdu = -AA*u + Bb*f
    return Mdu
end

u0   = setIC.(xx,yy)
dt   = 0.01
tspn = (0.0,1.0)

func = ODEFunction(Mdudt!;mass_matrix=BB)
prob = ODEProblem(func,u0,tspn)
sol  = solve(prob,
#            Rodas5();
#            Rodas5(linsolve=LinSolveGMRES(Pl=precondB));
             Rodas5(linsolve=LinSolveGMRES());
             saveat=0.1,
             callback=cb,
            )

@show sol.retcode
err = sol.u - [utrue.(xx,yy,sol.t[i]) for i=axes(sol.t,1)]
ee  = maximum.(abs.(err[i]) for i=axes(sol.t,1))
#----------------------------------------------------------------------#
"""
 standard form - mass-matrix inversion inside ODEFunciton
"""

function dudt!(du,u,p,t)
    f = setForcing.(xl,yl,t)
    rhs = -AA*u + Bb*f
    pcg!(du,rhs,BB;opM=x->x.*Big,ifv=false)
    return du
end

u0   = setIC.(xx,yy)
dt   = 0.01
tspn = (0.0,1.0)

func = ODEFunction(dudt!)
prob = ODEProblem(func,u0,tspn)
sol  = solve(prob, Rodas5(); saveat=0.1, callback=cb)

@show sol.retcode
err = sol.u - [utrue.(xx,yy,sol.t[i]) for i=axes(sol.t,1)]
ee  = maximum.(abs.(err[i]) for i=axes(sol.t,1))
#----------------------------------------------------------------------#
return ee


using PyPlot
frequencies = range(16, 20000, length=10000)
fs = 96e3
function show_grid()
	minorticks_on()
	grid(which="major", color="gray", linestyle="-")
	grid(which="minor", color="lightgray", linestyle="--")
	xlabel("f in Hz")
end

function magplot(myfilter, frequencies=frequencies)
	plot(frequencies, 20*log10.(abs.(freqz(myfilter, frequencies, fs))), label="Level/dB")
	semilogx()
	show_grid()
end

function phaseplot(myfilter, frequencies=frequencies)
	plot(frequencies, angle.(freqz(myfilter, frequencies, fs)), label="Phase/rad")
	semilogx()
	show_grid()
end

function gdplot(myfilter, frequencies=frequencies)
	deltaf = frequencies.step.hi
	gdfrequencies = range(frequencies[1]+deltaf/2, frequencies[end]-deltaf/2, length=frequencies.len-1)
	phaseresponse = angle.(freqz(myfilter, frequencies, fs))|>unwrap
	gd = -diff(phaseresponse)/deltaf
	plot(gdfrequencies, gd)
	semilogx()
	show_grid()
end

function csvGDplot(filename)
	input = readdlm(filename, ';', skipstart=1)
	deltaf = [input[i+1,1]-input[i,1] for i=1:size(input,1)-1]
	gd = -diff(unwrap(input[:,3]/180*pi))./deltaf
	fAxis = [(input[i+1,1]+input[i,1])/2 for i=1:size(input,1)-1]
	plot(fAxis, gd)
	semilogx()
	show_grid()
end

"""
testplot(filename, myfilter)

plot csv (sep ';', skipstart=1, f/Hz, Level/dB, Phase/deg) level and phase
also calls magplot and phaseplot on given filter
"""
function testplot(filename, myfilter)
	input = readdlm(filename, ';', skipstart=1)
	plot(input[:,1], input[:,2], label="ref Level/dB")
	plot(input[:,1], input[:,3]/180*pi, label="ref Phase/rad")
	magplot(myfilter, input[:,1])
	phaseplot(myfilter, input[:,1])
	legend()
end


using IntervalArithmetic
using Test

let b

@testset "DecoratedInterval tests" begin
    a = DecoratedInterval(@interval(1, 2), com)
    @test decoration(a) == com

    b = sqrt(a)
    @test interval(b) == sqrt(interval(a))
    @test decoration(b) == com

    a = DecoratedInterval(@interval(-1, 1), com)
    b = sqrt(a)
    @test interval(b) == sqrt(Interval(0, 1))
    @test decoration(b) == trv

    d = DecoratedInterval(a, dac)
    @test decoration(d) == dac

    @test decoration(DecoratedInterval(1.1)) == com
    @test decoration(DecoratedInterval(1.1, dac)) == dac

    @test decoration(DecoratedInterval(2, 0.1, com)) == ill
    @test decoration(DecoratedInterval(2, 0.1)) == ill
    @test isnai(DecoratedInterval(2, 0.1))
    @test decoration(@decorated(2, 0.1)) == ill
    @test decoration(DecoratedInterval(big(2), big(1))) == ill
    @test isnai((DecoratedInterval(big(2), big(1))))
    @test isnai(@decorated(big(2), big(1)))

    # Disabling the following tests, because Julia 0.5 has some strange behaviour here
    # @test_throws ArgumentError DecoratedInterval(BigInt(1), 1//10)
    # @test_throws ArgumentError @decorated(BigInt(1), 1//10)

    # Tests related to powers of decorated Intervals
    @test @decorated(2,3) ^ 2 == DecoratedInterval(4, 9)
    @test @decorated(2,3) ^ -2 == DecoratedInterval(1/9,1/4)
    @test @decorated(-3,2) ^ 3 == DecoratedInterval(-27., 8.)
    @test @decorated(-3,-2) ^ -3 == DecoratedInterval(-1/8.,-1/27)
    @test @decorated(0,3) ^ 2 == DecoratedInterval(0, 9)
    @test @decorated(0,3) ^ -2 == DecoratedInterval(1/9, Inf, trv)
    @test @decorated(2,3)^Interval(0.0, 1.0) == DecoratedInterval(1.0,3.0)
    @test @decorated(2,3)^@decorated(0.0, 1.0) == DecoratedInterval(1.0,3.0)
    @test @decorated(0, 2)^Interval(0.0, 1.0) == DecoratedInterval(0.0,2.0, trv)
    @test @decorated(0, 2)^@decorated(0.0, 1.0) == DecoratedInterval(0.0,2.0, trv)
    @test @decorated(-3, 2)^Interval(0.0, 1.0) == DecoratedInterval(0.0,2.0, trv)
    @test @decorated(-3, 2)^@decorated(0.0, 1.0) == DecoratedInterval(0.0,2.0, trv)
    @test @decorated(-3, 2)^Interval(-1.0, 1.0) == DecoratedInterval(0.0,Inf, trv)
    @test @decorated(-3, 2)^@decorated(-1.0, 1.0) == DecoratedInterval(0.0, Inf, trv)
end

@testset "Convert string to DecoratedInterval" begin
    @test convert(DecoratedInterval{Float64}, "[1,2]") ==
                    DecoratedInterval(Interval(1, 2), com)

    @test convert(DecoratedInterval{Float64}, "[1,2]_dac") ==
                        DecoratedInterval(Interval(1, 2), dac)

end

end


"""
bayessupport \\

For calculating the support (number of non-zero probabilities) from a Bayesian point of view  \\
Arguments:  probs = naive probability distribution \\
            n_trials = number of trials \\
Returns:    a bayesian estimate for the support of the probability distribution  \\
"""
function bayessupport(probs::Array{Float64}, n_trials::Int)
  dim = length(probs)
  probs_nonzero = probs[probs .> 0]
  support_naive = length(probs_nonzero)
  support = support_naive
  if support_naive < dim
    support_expected = support_naive - sum((1 - probs_nonzero) .^ n_trials)
    prev_support_delta = dim
    support_delta = abs(support_naive - support_expected)
    xtr = 0.0
    while (support_delta < prev_support_delta) && ((support_naive+xtr)<dim)
      xtr = xtr + 1.0
      support_expected = 0.0
      # occupied bins
      gamma = xtr*(1.0 - (n_trials/(n_trials+support_naive)^(1/n_trials)))
      support_bayes = ((1.0-gamma)/(n_trials+support_naive))*(probs_nonzero*n_trials+1.0)
      suppport_expected = sum(1.0-(1.0-support_bayes)^n_trials)
      # non-occupied bins
      support_bayes = gamma / xtr
      support_expected = support_expected + xtr*(1.0 - (1.0 - support_bayes)^n_trials)
      prev_support_delta = support_delta
      support_delta = abs(support_naive - support_expected)
    end
    support_naive = support_naive + xtr - 1.0
    if support_delta < prev_support_delta
      support_naive += 1.0
    end
  end
return support_naive
end


immutable FDist <: ContinuousUnivariateDistribution
    ndf::Float64
    ddf::Float64
    function FDist(d1::Real, d2::Real)
    	d1 > zero(d1) && d2 > zero(d2) ||
    	    error("Numerator and denominator degrees of freedom must be positive")
  	new(float64(d1), float64(d2))
    end
end

@_jl_dist_2p FDist f

@continuous_distr_support FDist 0.0 Inf

mean(d::FDist) = 2.0 < d.ddf ? d.ddf / (d.ddf - 2.0) : NaN

median(d::FDist) = quantile(d, 0.5)

mode(d::FDist) = d.ndf <= 2. ? 0.0 : (d.ndf - 2.) / d.ndf * d.ddf / (d.ddf + 2.)

var(d::FDist) = d.ddf > 4. ?  2.0 * d.ddf^2 *
		       (d.ndf + d.ddf - 2.0) /
		       (d.ndf * (d.ddf - 2.0)^2 * (d.ddf - 4.0)) : NaN

skewness(d::FDist) = d.ddf > 6. ?  (2.d.ndf + d.ddf - 2.) * sqrt(8.(d.ddf - 4.)) /
           ((d.ddf - 6.) * sqrt(d.ndf * (d.ndf + d.ddf - 2.))) : NaN

function kurtosis(d::FDist)
    d.ddf <= 8. && return NaN
    a = d.ndf * (5. * d.ddf - 22.) * (d.ndf + d.ddf - 2.) +
        (d.ddf - 4.) * (d.ddf - 2.)^2
    b = d.ndf * (d.ddf - 6.) * (d.ddf - 8.) * (d.ndf + d.ddf - 2.)
    12. * a / b
end

entropy(d::FDist) = (log(d.ddf) -log(d.ndf)
                     +lgamma(0.5*d.ndf) +lgamma(0.5*d.ddf) -lgamma(0.5*(d.ndf+d.ddf))
                     +(1.0-0.5*d.ndf)*digamma(0.5*d.ndf) +(-1.0-0.5*d.ddf)*digamma(0.5*d.ddf)
                     +0.5*(d.ndf+d.ddf)*digamma(0.5*(d.ndf+d.ddf)))


immutable SVM <: RegERM
    X::Matrix               # n x m matrix of n m-dimensional training examples
    y::Vector               # 1 x n vector with training classes
    n::Int                  # number of training examples
    m::Int                  # number of features
    kernel::Symbol          # kernel function
    regression_type::Symbol # ordinal, binomial, multinomial
    params::Hyperparameters # hyperparameters (e.g. λ)
end

function SVM(X::Matrix, y::Vector; kernel::Symbol=:linear, λ::Float64=0.1)
    check_arguments(X, y, :binomial)
    SVM(X, y, size(X)..., kernel, :binomial, RegularizationParameters(λ))
end

function SVM(X::Matrix, y::Vector, params::Hyperparameters; kernel::Symbol=:linear)
    check_hyperparameters(params)
    check_arguments(X, y, :binomial)
    SVM(X, y, size(X)..., kernel, :binomial, params)
end

methodname(::SVM) = "Support Vector Machine"
loss(::SVM) = HingeLoss()
regularizer(SVM::SVM, w::Vector) = L2reg(w, SVM.params.λ)


include_implicit_dynamics()

"""
    Cartpole w/ internal friction
"""

struct Cartpole{T}
    dim::Dimensions

    mc::T     # mass of the cart in kg
    mp::T     # mass of the pole (point mass at the end) in kg
    l::T      # length of the pole in m
    g::T      # gravity m/s^2
end

function kinematics(model::Cartpole, q)
    [q[1] + model.l * sin(q[2]); -model.l * cos(q[2])]
end

lagrangian(model::Cartpole, q, q̇) = 0.0

function M_func(model::Cartpole, x)
    H = @SMatrix [model.mc + model.mp model.mp * model.l * cos(x[2]);
				  model.mp * model.l * cos(x[2]) model.mp * model.l^2.0]
    return H
end

function B_func(model::Cartpole, x)
    @SMatrix [1.0;
              0.0]
end

function P_func(model::Cartpole, x)
    @SMatrix [1.0 0.0;
              0.0 1.0]
end

function C_func(model::Cartpole, q, q̇)
    # c_func(model, q, q̇) * q̇ - τ_func(model, q)
    C = @SMatrix [0.0 -1.0 * model.mp * q̇[2] * model.l * sin(q[2]);
	 			  0.0 0.0]
    G = @SVector [0.0,
				  model.mp * model.g * model.l * sin(q[2])]

    return -C * q̇ + G
end

model = Cartpole(Dimensions(2, 1, 0, 4),
    1.0, 0.2, 0.5, 9.81)

model_no_friction = Cartpole(Dimensions(2, 1, 0, 4),
    1.0, 0.2, 0.5, 9.81)

function lagrangian_derivatives(model::Cartpole, q, v)
	D1L = -1.0 * C_func(model, q, v)
    D2L = M_func(model, q) * v
	return D1L, D2L
end

function dynamics(model::Cartpole, h, q0, q1, u1, λ1, q2)
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * u1
        + transpose(P_func(model, q2)) * λ1)
end

function dynamics_no_friction(model::Cartpole, h, q0, q1, u1, λ1, q2)
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * u1)
        # + transpose(P_func(model, q2)) * λ1)
end

function residual(model::Cartpole, z, θ, κ)

    nq = model.dim.q
    nu = model.dim.u
    nc = model.dim.c

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]
    μ_slider = θ[2nq + nu + 1 .+ (1:1)]
    μ_angle = θ[2nq + nu + 1 + 1 .+ (1:1)]

    q2 = z[1:nq]
    c1 = z[nq .+ (1:nc)]
    c2 = z[nq + nc .+ (1:nc)]

    vT1 = (q2[1] - q1[1]) / h
    vT2 = (q2[2] - q1[2]) / h

    β1 = c1[1:2]
    β2 = c2[1:2]

    η1 = c1[2 .+ (1:2)]
    η2 = c2[2 .+ (1:2)]

    λ1 = [β1[2]; β2[2]]

    [
     dynamics(model, h, q0, q1, u1, λ1, q2);
     η1[2] - vT1[1];
     # β1[1] - (μ_slider[1] * 1.0); # static normal
     β1[1] .- μ_slider[1] * (model.mp + model.mc) * model.g * h[1];
     second_order_cone_product(η1, β1) - κ .* [1.0; 0.0];
     η2[2] - vT2[1];
     # β2[1] - (μ_angle[1] * 1.0); # static "normal"
     β2[1] .- μ_angle[1] * (model.mp * model.g * model.l) * h[1];
     second_order_cone_product(η2, β2) - κ .* [1.0; 0.0];
    ]
end

function residual_no_friction(model::Cartpole, z, θ, κ)

    nq = model.dim.q
    nu = model.dim.u
    nc = model.dim.c

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]

    return dynamics_no_friction(model, h, q0, q1, u1, zeros(nc), q2);
end

nq = model.dim.q
nu = model.dim.u
nc = model.dim.c
nz = nq + nc + nc
nz_no_friction = nq
nθ = nq + nq + nu + 1 + 2

# Declare variables
@variables z[1:nz]
@variables z_no_friction[1:nz_no_friction]
@variables θ[1:nθ]
@variables κ[1:1]

# Residual
r = vec(residual(model, z, θ, κ))
r = Symbolics.simplify.(r)
rz = Symbolics.jacobian(r, z, simplify = true)
rθ = Symbolics.jacobian(r, θ, simplify = true)

# Build function
r_func = eval(build_function(r, z, θ, κ)[2])
rz_func = eval(build_function(rz, z, θ)[2])
rθ_func = eval(build_function(rθ, z, θ)[2])

rz_array = similar(rz, Float64)
rθ_array = similar(rθ, Float64)

@save joinpath(@__DIR__, "dynamics/residual.jl") r_func rz_func rθ_func rz_array rθ_array
@load joinpath(@__DIR__, "dynamics/residual.jl") r_func rz_func rθ_func rz_array rθ_array

# Residual
r_no_friction = residual_no_friction(model, z_no_friction, θ, κ)
r_no_friction = Symbolics.simplify.(r_no_friction)
rz_no_friction = Symbolics.jacobian(r_no_friction, z_no_friction, simplify = true)
rθ_no_friction = Symbolics.jacobian(r_no_friction, θ, simplify = true)

# Build function
r_no_friction_func = eval(build_function(r_no_friction, z_no_friction, θ, κ)[2])
rz_no_friction_func = eval(build_function(rz_no_friction, z_no_friction, θ)[2])
rθ_no_friction_func = eval(build_function(rθ_no_friction, z_no_friction, θ)[2])

rz_no_friction_array = similar(rz_no_friction, Float64)
rθ_no_friction_array = similar(rθ_no_friction, Float64)

@save joinpath(@__DIR__, "dynamics/residual_no_friction.jl") r_no_friction_func rz_no_friction_func rθ_no_friction_func rz_no_friction_array rθ_no_friction_array
@load joinpath(@__DIR__, "dynamics/residual_no_friction.jl") r_no_friction_func rz_no_friction_func rθ_no_friction_func rz_no_friction_array rθ_no_friction_array

idx_ineq = collect(1:0)
idx_soc = [collect([nq .+ (1:2)]...), collect([nq + 2 .+ (1:2)]...), collect([nq + nc .+ (1:2)]...), collect([nq + nc + 2 .+ (1:2)]...)]

z_subset_init = [1.0; 0.1; 1.0; 0.1; 1.0; 0.1; 1.0; 0.1]

q0 = [0.0; 0.0]
q1 = [0.0; 0.0]
u1 = [0.0]
h = 0.1

z0 = copy([q1; ones(nc + nc)])
θ0 = [q0; q1; u1; h; 0.1; 0.1]

# options
opts = InteriorPointOptions(
   κ_init = 0.1,
   κ_tol = 1.0e-4,
   r_tol = 1.0e-8,
   diff_sol = true)

# solver
ip = interior_point(z0, θ0,
   r! = r_func, rz! = rz_func,
   rz = similar(rz, Float64),
   rθ! = rθ_func,
   rθ = similar(rθ, Float64),
   # idx_ineq = idx_ineq,
   idx_soc = idx_soc,
   opts = opts)

# simulate
T = 100
q_hist = [q0, q1]

for t = 1:T-1
    u1 = [t == 1 ? 1.0 : t == 2 ? 0.0 : 0.0]
    ip.z .= copy([q_hist[end]; z_subset_init])
    ip.θ .= copy([q_hist[end-1]; q_hist[end]; u1; h; 0.1; 0.05])
    status = interior_point_solve!(ip)

    if status
        push!(q_hist, ip.z[1:nq])
    else
        println("dynamics failure t = $t")
        println("res norm = $(norm(ip.r, Inf))")
        # break
    end
end

include(joinpath(@__DIR__, "visuals.jl"))
include(joinpath(pwd(), "models/visualize.jl"))

vis = Visualizer()
# open(vis)
render(vis)
default_background!(vis)
settransform!(vis["/Cameras/default"],
        compose(Translation(0.0, -95.0, -1.0), LinearMap(RotY(0.0 * π) * RotZ(-π / 2.0))))
setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 1)

visualize!(vis, model, q_hist, Δt = h)


using Test
using GeoParams

@testset "DiffusionCreepLaws" begin

# This tests the MaterialParameters structure
CharUnits_GEO   =   GEO_units(viscosity=1Pa*s, length=1m);

# Define a linear viscous creep law ---------------------------------
x1      =   DiffusionCreep()
@test Value(x1.n) == 1.0
@test Value(x1.p) == -3.0
@test Value(x1.A) == 1.5MPa^-1.0*s^-1*m^(3.0)

x2      =   DiffusionCreep(n=3, p=-3.0)
@test Value(x2.A) == 1.5MPa^-3.0*s^-1*m^(3.0)

# perform a computation with the dislocation creep laws
    # Calculate EpsII, using a set of pre-defined values
CharDim = GEO_units(length=1000km, viscosity=1e19Pa*s, stress=100MPa, temperature=1000C)
EpsII   = GeoUnit(1.0s^-1.0)
EpsII_nd= nondimensionalize(EpsII,CharDim)
TauII   = GeoUnit(0.3MPa)
TauII_nd= nondimensionalize(TauII,CharDim)
P       = GeoUnit(1.0e9Pa)
P_nd    = nondimensionalize(P,CharDim)
T       = GeoUnit(1400C)
T_nd    = nondimensionalize(T,CharDim)
f       = GeoUnit(1000NoUnits)
f_nd    = nondimensionalize(f,CharDim)
d       = GeoUnit(10mm)
d_nd    = nondimensionalize(d,CharDim)
c       = CreepLawVariables(P=P,T=T,f=f,d=d)
#Phase   = SetMaterialParams(Name="Viscous Matrix", Phase=1,
#                                     Density   = ConstantDensity(),
#                                     CreepLaws = DiffusionCreep(n=1.0NoUnits, r=1.0NoUnits, p=-3.0NoUnits, A=1.0e6MPa^(-1.0)*s^(-1.0)*µm^(3.0), E=335000.0J/mol, V=4.0e-6m^(3.0)/mol, Apparatus="Invariant"), CharDim = CharDim)
#EpsII.val = ComputeCreepLaw_EpsII(TauII,Phase.CreepLaws[1],c)
# Tested value is Dimensional so need to non-dimensionalize it back to compare
#@test EpsII.val  ≈ nondimensionalize(7.823165072337786e-15s^(-1.0), CharDim) rtol = 1e-12
    # Check that once inverted, we get back the TauII that we used to calculate EpsII
#NewTau = ComputeCreepLaw_TauII(EpsII,Phase.CreepLaws[1],c)
#@show NewTau
#@test NewTau ≈ TauII.val

# Given stress
#@test ComputeCreepLaw_EpsII(1e6Pa, x1, CreepLawParams())==5e-13/s                # dimensional input

# Given strainrate
#@test ComputeCreepLaw_TauII(1e-13/s, x1, CreepLawParams())==1e18*2*1e-13Pa       # dimensional input
# -------------------------------------------------------------------

end


using Seaborn
using Pandas
using PyPlot
using PyCall

@pyimport numpy

#for some strange reason loading doesn't work the first time
tips = nothing
try
    tips = load_dataset("tips");
catch
    tips = load_dataset("tips");
end

head(tips)

g = regplot(x="total_bill", y="tip", data=tips)
title("Total Bill vs. Tip") #current active figure PyPlot
# alternatively
# g[:figure][:axes][1][:set_title]("Total Bill vs. Tip")

f, (ax1, ax2,ax3)  = subplots(1, 3, sharey=true)
regplot(x="size", y="tip", data=tips, ax=ax1);
regplot(x="size", y="tip", data=tips, x_jitter=.05, ax=ax2);
regplot(x="size", y="tip", data=tips, x_estimator=numpy.mean, ax=ax3);

anscombe = load_dataset("anscombe");
head(anscombe)

lmplot(x="x", y="y", data=query(anscombe, "dataset == 'II'"),
           ci=nothing, scatter_kws=Dict("s"=> 80));

lmplot(x="x", y="y", data=query(anscombe, "dataset == 'II'"),
           ci=nothing, scatter_kws=Dict("s"=> 80), order =2);

jointplot(x="total_bill", y="tip", data=tips, kind="reg")
savefig("./test.svg")

lmplot(x="total_bill", y="tip", hue="smoker", data=tips);

lmplot(x="total_bill", y="tip", col="day", data=tips,
           aspect=.5);


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 70b737a2-a9b2-11eb-1f6f-a330281aaaeb
using Omega

# ╔═╡ 305de14b-4732-4abf-ac08-c542b1bf0203
using CSV

# ╔═╡ f8ac57f8-680a-4246-b2a2-92714c38c167
using DataFrames

# ╔═╡ 4d5a4967-d8c6-4bf9-8c16-c1486d39b334
using Plots

# ╔═╡ 0d6a1f13-1b0c-46c0-8885-23530fd9a02c
using Query

# ╔═╡ e8195cc0-45c8-4edd-b3bc-425f0f902645
using VegaLite

# ╔═╡ 601b0f0d-5dab-43c2-9440-5254c802a3dc
using Distributions

# ╔═╡ 46fcaa08-ba11-41f4-aa65-9cc5ac9f44ec
md"# Inferring Degrees of Racial Descrimination through Bayesian Inference from Police Stopping Data"

# ╔═╡ 3f4880ba-65f2-4dd0-8eed-2d4f811fb8b0
md""" ## The Data
The data here is of police stops in San Diego
"""

# ╔═╡ 71b12af6-1886-4af9-bd3b-a7f06d89f83b
datapath = joinpath(dirname(pathof(Omega)), "..", "examples")

# ╔═╡ 847d6099-c655-4a92-a39a-ec7ae9982848
pd = DataFrame(CSV.File(joinpath(datapath, "sandiegopolicestops.csv")))

# ╔═╡ 471dea65-1eb1-415f-89bc-0612d7d88ed2
x = @from i in pd begin
    @where i.race_category == "Black"
    @select {i.block, i.land_mark}
    @collect DataFrame
end;

# ╔═╡ 1a4e1fb2-9c59-4cca-b925-ba06947f82ce
md"""# Stopping Test

The basic idea behind the threshold test is the following.

Suppose there is a disparity between the the rate at which members of minority groups are being stopped compared to non-minority groups.

Alone, this evidence is not sufficient to prove bias against the minority groups because there may be other causal factors that lead to the same conclusion.  For example, ethnic minorities tend to live in cities, which have a higher density of police offers.

One indication that disparities in stopping are indeed due to discrimination are if there are disparities in the rates at which the stops were justified.
A justified stop for example is one where an illegal weapon was found on the participant.

However the presence of a disparity is still insufficient to indicate bias.
This can be demonstrated through a counterexample: it's possible for a police offer to adopt a completely random (and hence indifferent to race) criteria by which they choose to stop someone, but there still to be differences in the justified stopping rates.
"""

# ╔═╡ e8d79dfc-13b3-4ae1-bc19-0d4eadd60cdc
md"""## A Bayesian Model

We will model an encounter between a police officer and a civillian.  The basic setup is as follows:

- A police-officer encounters a civillian
- Based on the civillian's race and the location, the officer infers a belief $p$ which represents the probability that the civillian is carrying a weapon
- The police-offer actually executes a search if $p$ exceeds some threshold $t_r$ which is race specific.

We use Bayesian inference to infer the thresholds for different races.  Any major disparity of $t_r$ between races indicates discrimination.
"""

# ╔═╡ 71ab4e3a-3888-480e-aa7e-82cc771f6c55
md"First we extract the race categories from the dataset"

# ╔═╡ cd2e8516-68f8-4c72-baf6-11cbbfd721ec
#race_categories = (@from i in pd begin
#    @select i.race_category
#	end) |> @unique() |> collect

# ╔═╡ 71b4adbc-4f84-4d9a-a158-4a3667881a4a
race_categories = ["Black", "White"]

# ╔═╡ ad52aaa5-fbb5-4747-9554-e536b81f1736
md"Let's add a little helper function to get the race index from the string"

# ╔═╡ 4a0aa2ef-654c-4bec-8d88-3ac5af3f0a94
race_id(race) = findfirst(x -> x == race, race_categories);

# ╔═╡ 8a174224-9e12-43aa-96be-fb0e3e951c2a
n_races = length(race_categories)

# ╔═╡ c92c75a7-2a72-4520-a03b-39116e7106ac
md"The following function `risk` constructs the parameters for the risk distribution as a function of race, as indicated by its id `i` (which spans from 1 to the total number of races `n_races`)"

# ╔═╡ a6678b03-f493-447e-83db-b63a03028744
risk(i) =
  (ϕ_r = (i, 1, 1) ~ Normal(0, 1),
   ϕ_d = (i, 2, 1) ~ Normal(0, 1),
   λ_r = (i, 3, 1) ~ Normal(0, 1),
   λ_d = (i, 4, 1) ~ Normal(0, 1))

# ╔═╡ ae577e52-8ced-46b9-a1ea-ee1098428de4
md"For each race category we assume there is an independent risk distribution"

# ╔═╡ 825b4d05-78f7-4619-a9c8-54f016dfd7f4
risks = map(risk, 1:n_races);

# ╔═╡ de8c1f27-6dc3-477f-b5e1-9e89d608bf92
md"In addition, for each race category there is a threshold value $t_r$, above which, the model suggests that a person will be searched"

# ╔═╡ d59e8d13-cce1-4ef2-b14c-27604fd42f3f
thresholds = [(i, 2) ~ Beta(1, 2) for i = 1:n_races];

# ╔═╡ 09fd183f-8677-4922-b86e-bf786b667bbe
invlogit(x) = (e = exp(x); e / (e + 1))

# ╔═╡ 5a63da51-65fe-4c8d-8821-78a4644c509a
md"Finally we can construct the stopping model.  Models in Omega are normal Julia functions which take as input ω (which we can think of as an random number generator object, or if we are more mathematically inclined as an element in probability space."

# ╔═╡ 1b73ba1c-3e7b-4594-a8fb-9e1835ebf3cc
function stop(id, ω)
	race_ = ω |> ((id, 3) ~ DiscreteUniform(1, n_races))

	# Get risk parameters
	riskp = risks[race_]

	# Construct beta distribution priors
	ϕ_rd = invlogit(riskp.ϕ_r(ω) + riskp.ϕ_d(ω))
	λ_rd = exp(riskp.λ_r(ω) + riskp.λ_d(ω))

	# Draw $p$ -- belief weapon-carrying from Beta
	p = ω |> ((id, 4) ~ Beta(ϕ_rd, λ_rd))

	# Draw threshold (as function of race)
	thresh = thresholds[race_](ω)

	# Officer stops if p exceeds threshold
	race_, Int(p > thresh)
end

# ╔═╡ 5afbcc6c-0012-4d10-b1ec-b194899396ba
n = 10

# ╔═╡ 3438de0c-692d-4376-b45e-85de49d95129
md"`stop` models a single data point.  To model multiple data points we next construct `data_size` (conditionally) independent draws from this model"

# ╔═╡ a251f256-3c95-46ac-be85-26aa08ca0eda
stopmany = Mv(1:n, stop)

# ╔═╡ 6bc7da07-7a60-4916-9d2f-3bf4d7838fcf
md"Here's a single sample from `stopmany`"

# ╔═╡ ecb05030-bb6e-40ea-aae8-718eb85a08b7
randsample(stopmany)

# ╔═╡ 92308112-eb7f-42bb-b303-e37dbbc4c0d4
md"We can show the unconditional distribution of the stops vs non-stops";

# ╔═╡ e659f421-19d5-43fe-a9bf-c845494efbe6
histogram(randsample(stopmany));

# ╔═╡ 599a49f2-0b86-4661-8bc8-a8473fe68241
md"""### Conditioning

Finally we can condition the model with data to infer the posterior distribution over $t$.

First, let's construct the actual proposition we want to condition on
"""

# ╔═╡ 1be486c7-cd84-4843-a388-aaadb39881ea
black_threshold = thresholds[race_id("Black")]

# ╔═╡ be0f16df-608f-43fa-accb-5af34c5608cb
d1 = @from i in pd begin
     @where i.search_conducted == 0
	 @where i.race_category in race_categories
     @select (race_id(i.race_category), i.search_conducted)
	 @collect
end;

# ╔═╡ 1b72954e-0a25-40a2-9638-5d6f75ae77bf
d2 = @from i in pd begin
     @where i.search_conducted == 1
     @where i.race_category in race_categories
     @select (race_id(i.race_category), i.search_conducted)
	 @collect
end;

# ╔═╡ 2d70f306-abda-4805-ab8d-0cde05decdc7
data_slice =
Tuple{Int, Int64}[d1[1:div(n,2)]; d2[1:div(n,2)]]

# ╔═╡ 2a8a54d0-5460-4679-96b7-fbde5c82cf68
data_condition(ω) = stopmany(ω) == data_slice;

# ╔═╡ c52be3dd-6123-418d-837f-b111ce908fdd
md"Construct the posterior distribution"

# ╔═╡ abaa6f63-c89b-4b54-afc3-27316352e9bc
black_threshold_posterior = black_threshold |ᶜ data_condition;

# ╔═╡ 0d8bef6f-92cb-449c-a77a-fd6650a9e586
md"Sample from the posterior distribution"

# ╔═╡ 7aa560da-08b9-4a61-96f0-b24ddb7a6da0
samples = randsample(black_threshold_posterior, 10000; alg = RejectionSample)

# ╔═╡ 4d10e6c7-0824-4258-822c-624e132cd309
histogram(samples)

# ╔═╡ Cell order:
# ╟─46fcaa08-ba11-41f4-aa65-9cc5ac9f44ec
# ╠═70b737a2-a9b2-11eb-1f6f-a330281aaaeb
# ╠═305de14b-4732-4abf-ac08-c542b1bf0203
# ╠═f8ac57f8-680a-4246-b2a2-92714c38c167
# ╠═4d5a4967-d8c6-4bf9-8c16-c1486d39b334
# ╠═0d6a1f13-1b0c-46c0-8885-23530fd9a02c
# ╠═e8195cc0-45c8-4edd-b3bc-425f0f902645
# ╠═601b0f0d-5dab-43c2-9440-5254c802a3dc
# ╟─3f4880ba-65f2-4dd0-8eed-2d4f811fb8b0
# ╟─71b12af6-1886-4af9-bd3b-a7f06d89f83b
# ╠═847d6099-c655-4a92-a39a-ec7ae9982848
# ╟─471dea65-1eb1-415f-89bc-0612d7d88ed2
# ╟─1a4e1fb2-9c59-4cca-b925-ba06947f82ce
# ╟─e8d79dfc-13b3-4ae1-bc19-0d4eadd60cdc
# ╟─71ab4e3a-3888-480e-aa7e-82cc771f6c55
# ╟─cd2e8516-68f8-4c72-baf6-11cbbfd721ec
# ╟─71b4adbc-4f84-4d9a-a158-4a3667881a4a
# ╟─ad52aaa5-fbb5-4747-9554-e536b81f1736
# ╠═4a0aa2ef-654c-4bec-8d88-3ac5af3f0a94
# ╠═8a174224-9e12-43aa-96be-fb0e3e951c2a
# ╟─c92c75a7-2a72-4520-a03b-39116e7106ac
# ╠═a6678b03-f493-447e-83db-b63a03028744
# ╟─ae577e52-8ced-46b9-a1ea-ee1098428de4
# ╠═825b4d05-78f7-4619-a9c8-54f016dfd7f4
# ╟─de8c1f27-6dc3-477f-b5e1-9e89d608bf92
# ╠═d59e8d13-cce1-4ef2-b14c-27604fd42f3f
# ╠═09fd183f-8677-4922-b86e-bf786b667bbe
# ╟─5a63da51-65fe-4c8d-8821-78a4644c509a
# ╠═1b73ba1c-3e7b-4594-a8fb-9e1835ebf3cc
# ╠═5afbcc6c-0012-4d10-b1ec-b194899396ba
# ╟─3438de0c-692d-4376-b45e-85de49d95129
# ╠═a251f256-3c95-46ac-be85-26aa08ca0eda
# ╟─6bc7da07-7a60-4916-9d2f-3bf4d7838fcf
# ╠═ecb05030-bb6e-40ea-aae8-718eb85a08b7
# ╟─92308112-eb7f-42bb-b303-e37dbbc4c0d4
# ╟─e659f421-19d5-43fe-a9bf-c845494efbe6
# ╟─599a49f2-0b86-4661-8bc8-a8473fe68241
# ╠═1be486c7-cd84-4843-a388-aaadb39881ea
# ╟─be0f16df-608f-43fa-accb-5af34c5608cb
# ╟─1b72954e-0a25-40a2-9638-5d6f75ae77bf
# ╟─2d70f306-abda-4805-ab8d-0cde05decdc7
# ╠═2a8a54d0-5460-4679-96b7-fbde5c82cf68
# ╟─c52be3dd-6123-418d-837f-b111ce908fdd
# ╠═abaa6f63-c89b-4b54-afc3-27316352e9bc
# ╟─0d8bef6f-92cb-449c-a77a-fd6650a9e586
# ╠═7aa560da-08b9-4a61-96f0-b24ddb7a6da0
# ╠═4d10e6c7-0824-4258-822c-624e132cd309


using Dates, BusinessDays, Roots, AOPE
# MIT License
# Copyright (c) 2019: Antonio Saragga Seabra

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
    yca(cleanprice,annualcoupons,redemptionvalue,startdate,enddate)

**Yield Corrente Ajustada** ou yield simples.

The simple yield is computed as:
```math
\\frac{\\text{annual total value of coupons}+(\\text{redemption value}-\\text{clean price})/L}{\\text{clean price}}
\\text{where}
L = \\frac{\\text{# days in the period excluding 29 Feb.}}{365}
```
"""
function yca(cleanprice, annualcoupons, redemptionvalue, startdate::Date, enddate::Date)
    L = (Dates.value(enddate - startdate) - ndays29feb(startdate,enddate))/365
    return (annualcoupons + (redemptionvalue - cleanprice)/L)/cleanprice
end

"""
    ytm(price,cf,m,dc)
        price: scalar
        cf: K vector of cash flows
        m:  K vector of times for the cash flows
        dc: type of composition: Discrete/Continuous

Computes the **Yield-to-Maturity** of a bond. The redemption date may be: maturity, call date, put date

The ytr is computed as:
```math
ytr = y: \\text{price} = \\sum_{i=1}^k \\frac{cf_i}{(1+y)^{m_i}}
ytr = y: \\text{price} = \\sum_{i=1}^k cf_i*\\exp{-y*m_i}}
```
"""
function ytm(price,cf,m, dc)
    return Roots.find_zero(y->spreço(cf, m, y, dc)-price,0.01)
end

function ytmisda(P, g, k, h, n, f1, f2, C, v)
    return Roots.fzero(y->spreço(g, k, h, n, f1, f2, C, v, y)-P,0.01)
end

"""
zspread(price, cf, m, r, dc)

Computes the z-spread of a bond given benchmark interest rates and the dirty price (P) of the bond
    P:  liquidation price
    r:  scalar or K vector of benchmark interest rates
    cf: scalar or K vector of cash flows
    m:  K vector of times for the cash flows
    dc: Discrete/Continuous composition
returns
    scalar, the z-spread
"""
function zspread(P, cf, m, r, dc)
    return Roots.find_zero(z->spreço(cf, m, r, z, dc)-P,0.01)
end


include("weeds.jl")

#run the inverse analysis repeatedly for a set of parameters
sigma, lambda, mu, deltah, sqrtnumobs, alpha, beta, gamma, iterations = 1.0, 50.0, 0.0, 5.0, 16, 1.0, 1.0, 1.0, 100#here are the parameters we will run the inverse analysis with
@time results = pmap(x->inversemodel(sigma, lambda, mu, deltah, sqrtnumobs, alpha, beta, gamma, iterations), 1:100)#run the inverse analysis 100 times (in parallel if julia is run like, e.g., "julia -p 8" where 8 processors will be used)

#postprocess the results
eig_residuals = map(x->x[1], results)
eig_rmse = sum(map(x->x .^ 2, eig_residuals)) / length(eig_residuals)
logK_rmse = map(x->x[2], results)
of_vals = map(x->x[3], results)

#make a couple plots based on the results
fig, ax = PyPlot.subplots()
ax.semilogy(eig_rmse, ".")
ax.set_xlabel("Eigenvector Index")
ax.set_ylabel("Eigenvector Coefficient RMSE")
PyPlot.savefig("eigenvector_coefficient_fit.pdf")
display(fig)
println()
PyPlot.close(fig)

fig, ax = PyPlot.subplots()
for i = 1:length(logK_rmse)
	ax.loglog(of_vals[i], logK_rmse[i], ".", alpha=0.1)
end
ax.set_xlabel("Objective Function")
ax.set_ylabel("L2 Error in logK")
PyPlot.savefig("logK_vs_objective_function.pdf")
display(fig)
println()
PyPlot.close(fig)


# Ported from: http://www.math.unl.edu/~tshores1/Public/Math496S06/MatlabTools/
#                                                Examples/chap7/examp1/bvls.m
# TODO: take a critical look at speed some time
#
# x=bvls(A,b,l,u)
#
# Solves a bounded variables least squares problem
#
#   min || Ax-b ||
#          l <= x <= u
#
# The algorithm is based on the bounded variables least squares algorithm
# of Stark and Parker.  See
#
#  P.B. Stark and R. L. Parker, "Bounded-Variable Least-Squares: An Algorithm
#  and Applications", Computational Statistics 10:129-141, 1995.
#

function bvls(A::AbstractMatrix, b::AbstractVector, l::AbstractVector, u::AbstractVector)
    #
    # Get the size of A for future reference.
    #
    m, n = size(A)
    #
    # Initialize a bunch of variables.  This speeds things up by avoiding
    # realloation of storage each time an entry is added to the end of a vector.
    #
    oopslist = falses(n)
    state = fill(0, n)
    x = zeros(n)
    atbound = falses(n)
    between = falses(n)
    criti = 0
    crits = 0
    myeps = 1.0e-10
    #
    # setup an initial solution with all vars locked at a lower or upper bound.
    # set any free vars to 0.
    #
    for i in 1:n
        if ((u[i] >= +Inf) && (l[i] <= -Inf))
            x[i] = 0
            state[i] = 0
            between[i] = true
        elseif (u[i] >= +Inf)
            x[i] = l[i]
            state[i] = 1
            atbound[i] = true
        elseif (l[i] <= -Inf)
            x[i] = u[i]
            state[i] = 2
            atbound[i] = true
        elseif (abs(l[i]) <= abs(u[i]))
            x[i] = l[i]
            state[i] = 1
            atbound[i] = true
        else
            x[i] = u[i]
            state[i] = 2
            atbound[i] = true
        end
    end
    #
    # The main loop.  Stop after 10*n iterations if nothing else.
    #
    iter = 0
    while (iter < 10*n)
        iter += 1
        #
        # optimality test.
        #
        grad = A'*(A*x-b)
        #
        # We ignore any variable that has already been tried and failed.
        #
        grad[oopslist] = 0.
        #
        # Check for optimality.
        #
        done = true
        for i in 1:n
            if ((abs(grad[i]) > (1+norm(b))*myeps) && (state[i] == 0)) ||
                   ((grad[i] < 0) && (state[i] == 1)) ||
                   ((grad[i] > 0) && (state[i] == 2))
               done = false
               break
           end
        end
        done && return x
        #
        # Not optimal, so we need to free up some vars at bounds.
        #
        newi = 0
        newg = 0.0
        for i in 1:length(atbound)
            if atbound[i]
                #
                # Don't free up a variable that was just locked at its bound.
                #
                i == criti && continue
                #
                # Look for the locked variable with the biggest gradient.
                #
                if (grad[i] > 0) && (state[i] == 2) && (abs(grad[i]) > newg)
                    newi = i
                    newg = abs(grad[i])
                end
                if (grad[i] < 0) && (state[i] == 1) && (abs(grad[i]) > newg)
                    newi = i
                    newg = abs(grad[i])
                end
            end
        end
        #
        # Free the locked variable with the biggest gradient if there is one.
        #
        if (newi != 0)
            atbound[newi] = false
            state[newi] = 0
            between[newi] = true
        end
        #
        # Make sure the projected problem is nontrivial.
        #
        if (length(between)==0)
          # println("Empty projected problem")
          continue
        end
        #
        # Construct the new projected problem.
        #
        Aproj = A[:, between]
        An = A[:, atbound]
        if any(atbound)
            bproj = b - An*x[atbound]
        else
            bproj = b
        end
        #
        # Solve the projected problem.
        #
        z = Aproj\bproj
        xnew = copy(x)
        xnew[between] = z

        if newi != 0 && ((xnew[newi] <= l[newi] && x[newi] == l[newi]) ||
                (xnew[newi] >= u[newi] && x[newi] == u[newi]))
            #
            # Ooops- the freed variable wants to go beyond its bound.  Lock it down
            # and look for some other variable.
            #
            oopslist[newi] = true
            if xnew[newi] <= l[newi] && state[newi] == 1
                state[newi] = 1
                x[newi] = l[newi]
            end
            if xnew[newi] >= u[newi] && state[newi] == 2
                state[newi] = 2
                x[newi] = u[newi]
            end
            atbound[newi] = true
            between[newi] = false
            continue
        end
        #
        # We've got a good variable freed up.  Reset the oopslist.
        #
        oopslist = falses(n)
        #
        # Move as far as possible towards the optimal solution to the projected
        # problem.
        #
        alpha = 1.
        for i in 1:length(atbound)
            if between[i]
                if (xnew[i] > u[i])
                    newalpha = min(alpha, (u[i]-x[i])/(xnew[i]-x[i]))
                    if (newalpha < alpha)
                        criti = i
                        crits = 2
                        alpha = newalpha
                    end
                end
                if (xnew[i] < l[i])
                    newalpha = min(alpha, (l[i]-x[i])/(xnew[i]-x[i]))
                    if (newalpha < alpha)
                        criti = i
                        crits = 1
                        alpha = newalpha
                    end
                end
            end
        end
        #
        # Take the step.
        #
        x = x + alpha*(xnew-x)
        #
        # Update the state of variables.
        #
        if (alpha < 1)
            between[criti] = false
            atbound[criti] = true
            state[criti] = crits
        end

        for i in 1:n
            if (x[i] >= u[i])
                x[i] = u[i]
                state[i] = 2
                between[i] = false
                atbound[i] = true
            end
            if (x[i] <= l[i])
                x[i] = l[i]
                state[i] = 1
                between[i] = false
                atbound[i] = true
            end
        end
    end
end

function bvls{m,n}(A::SMatrix{m,n}, b::SVector{m}, l::SVector{n}, u::SVector{n})
    SVector{n}(bvls(Matrix(A), Vector(b), Vector(l), Vector(u)))
end


const CARRYCAP_PARAM      = Param(100000.0; bounds=(0.0, 10.0))
const INTRINSICRATE_PARAM = Param(0.1,      bounds=(0.0, 10.0))
const THRESHOLD_PARAM     = Param(0.5;      bounds=(0.0, 1.0))

"""
    GrowthRule <: CellRule

Abstract supertype for growth dynamics rules.
"""
abstract type GrowthRule{R,W} <: CellRule{R,W} end

"""
    ExponentialGrowth <: CellRule

    ExponentialGrowth(; rate, timestep, [nsteps_type])
    ExponentialGrowth{R}(; rate, timestep, [nsteps_type])
    ExponentialGrowth{R,W}(; rate, timestep, [nsteps_type])

Exponential growth of population size N based on an intrinsic growth rate ``r``, using the
exact solution between timesteps ``t`` and ``t-1``:

```math
N_t = N_{t-1}e^{r t}
```

# Keywords

- `rate`: Intrinsic growth rate. May be a `Number`, an `Aux` array or another `Grid`.
- `timestep`: Time step for the growth rate calculation, in a type compatible with the
    simulation `tspan`.
- `nsteps_type`: Specify the floating point type to use when `nsteps` is generated from the
    timestep, if it is required for type-stability or performance. The default is `Float64`.

Pass grid `Symbol`s to `R` or both `R` and `W` type parameters to use to specific grids.
"""
struct ExponentialGrowth{R,W,GR,TS,S} <: GrowthRule{R,W}
    rate::GR
    timestep::TS
    nsteps::S
end
function ExponentialGrowth{R,W}(;
    rate=INTRINSICRATE_PARAM,
    timestep,
    nsteps_type=Float64,
) where {R,W}
    ExponentialGrowth{R,W}(rate, timestep, zero(nsteps_type))
end

modifyrule(rule::ExponentialGrowth, data) = precalc_nsteps(rule, data)

@inline function applyrule(data, rule::ExponentialGrowth, N, I)
    N > zero(N) || return zero(N)
    rt = get(data, rule.rate, I...) * rule.nsteps

    return @fastmath N * exp(rt)
end

"""
    LogisticGrowth <: GrowthRule

    LogisticGrowth(; rate, carrycap, timestep, [nsteps_type])
    LogisticGrowth{R}(; rate, carrycap, timestep, [nsteps_type])
    LogisticGrowth{R,W}(; rate, carrycap, timestep, [nsteps_type])

Logistic growth rate of population size ``N`` based on an intrinsic growth rate ``r`` and
carry capacity ``K``, using the exact solution between timesteps ``t+1`` and ``t``:

```math
N_{t+1} = (N_t K) / (N_t + (K - N_t) e^{-rt})
```
Saturation only applies with positive growth.

# Keywords

These may be a `Number`, an `Aux` array or another `Grid`:

- `rate`: Intrinsic growth rate.
- `carrycap`: Carrying capacity.
- `timestep`: Time step for the growth rate, in a type compatible with the simulation `tspan`.
- `nsteps_type`: Specify the floating point type to use when `nsteps` is generated from the
    timestep, if it is required for type-stability or performance. The default is `Float64`.

Pass grid `Symbol`s to `R` or both `R` and `W` type parameters to use to specific grids.
"""
struct LogisticGrowth{R,W,GR,CC,TS,S} <: GrowthRule{R,W}
    rate::GR
    carrycap::CC
    timestep::TS
    nsteps::S
end
function LogisticGrowth{R,W}(;
    rate=INTRINSICRATE_PARAM,
    carrycap=CARRYCAP_PARAM,
    timestep,
    nsteps_type=Float64,
) where {R,W}
    LogisticGrowth{R,W}(rate, carrycap, timestep, zero(nsteps_type))
end

modifyrule(rule::LogisticGrowth, data) = precalc_nsteps(rule, data)

@inline function applyrule(data, rule::LogisticGrowth, N, I)
    N > zero(N) || return zero(N)

    rt = get(data, rule.rate, I...) * rule.nsteps
    k = get(data, rule.carrycap, I...)

    if rt > zero(rt)
        return @fastmath (N * k) / (N + (k - N) * exp(-rt))
    else
        return @fastmath N * exp(rt)
    end
end

"""
    ThresholdGrowth <: CellRule

    ThresholdGrowth(; rate, threshold)
    ThresholdGrowth{R}(; rate, threshold)
    ThresholdGrowth{R,W}(; rate, threshold)

Simple threshold mask. Values below a certain threshold are replaced with zero.

# Keywords

These may be a `Number`, an `Aux` array or another `Grid`.

- `rate`: Intrinsic growth rate.
- `threshold`: Minimum viability threshold below which population falls to zero.

Pass grid `Symbol`s to `R` or both `R` and `W` type parameters to use to specific grids.
"""
struct ThresholdGrowth{R,W,GR,Th} <: GrowthRule{R,W}
    rate::GR
    threshold::Th
end
function ThresholdGrowth{R,W}(;
    rate=INTRINSICRATE_PARAM,
    threshold=THRESHOLD_PARAM,
) where {R,W}
    ThresholdGrowth{R,W}(rate, threshold)
end

@inline function applyrule(data, rule::ThresholdGrowth, N, I)
    r = get(data, rule.rate, I...)
    t = get(data, rule.threshold, I...)

    return r >= t ? N : zero(N)
end

precalc_nsteps(rule, data) = precalc_nsteps(rule.timestep, rule, data)
precalc_nsteps(rulestep::DatePeriod, rule, data) =
    @set rule.nsteps = typeof(rule.nsteps)(currenttimestep(data) / Millisecond(rulestep))
precalc_nsteps(rulestep, rule, data) =
    @set rule.nsteps = typeof(rule.nsteps)(timestep(data) / rulestep)


using RecurrentNN
reload("RecurrentNN.jl")
using Base.Test

# graph output test
m1 = RecurrentNN.NNMatrix(3,2)
m1.w[1,1] = 1.; m1.w[1,2] = 2.
m1.w[2,1] = 3.; m1.w[2,2] = 4.
m1.w[3,1] = 5.; m1.w[3,2] = 6.

m2 = RecurrentNN.NNMatrix(2,3)
m2.w[1,1] = 2.; m2.w[1,2] = 3.; m2.w[1,3] = 4.
m2.w[2,1] = 5.; m2.w[2,2] = 6.; m2.w[2,3] = 7.

# add test
g =  RecurrentNN.Graph()
m3 = RecurrentNN.add!(g,m1,m1)
m3.dw[1,1] = .1; m3.dw[1,2] = .2
m3.dw[2,1] = .3; m3.dw[2,2] = .4
m3.dw[3,1] = .5; m3.dw[3,2] = .6

g.backprop[1]()
# set previous chained gradient
@test m1.dw[1,1] == 0.2
@test m1.dw[1,2] == 0.4
@test m1.dw[2,1] == 0.6
@test m1.dw[2,2] == 0.8
@test m1.dw[3,1] == 1.0
@test m1.dw[3,2] == 1.2

# mul test
m1.dw[:] = 0. # reset gradient matrices
m2.dw[:] = 0. # reset  gradient matrices
g =  RecurrentNN.Graph()
m3 = RecurrentNN.mul!(g,m1,m2)
@test m3.w[1,1] == 12.
@test m3.w[1,2] == 15.
@test m3.w[1,3] == 18.
@test m3.w[2,1] == 26.
@test m3.w[2,2] == 33.
@test m3.w[2,3] == 40.
@test m3.w[3,1] == 40.
@test m3.w[3,2] == 51.
@test m3.w[3,3] == 62.

m3.dw[1,1] = .1; m3.dw[1,2] = .2; m3.dw[1,3] = .3
m3.dw[2,1] = .4; m3.dw[2,2] = .5; m3.dw[2,3] = .6
m3.dw[3,1] = .7; m3.dw[3,2] = .8; m3.dw[3,3] = .9
g.backprop[1]()

# m1 gradient tests
@test m1.dw[1,1] == 2.
@test m1.dw[1,2] == 3.8000000000000003
@test m1.dw[2,1] == 4.699999999999999
@test m1.dw[2,2] == 9.2
@test m1.dw[3,1] == 7.4
@test m1.dw[3,2] == 14.600000000000001

# m2 gradient tests
@test m2.dw[1,1] == 4.800000000000001
@test m2.dw[1,2] == 5.7
@test m2.dw[1,3] == 6.6
@test m2.dw[2,1] == 5.999999999999999
@test m2.dw[2,2] == 7.200000000000001
@test m2.dw[2,3] == 8.4

# reul() tests
m4 = RecurrentNN.NNMatrix(3,2)
m4.w[1,1] = 1.; m4.w[1,2] =-2.
m4.w[2,1] =-3.; m4.w[2,2] = 4.
m4.w[3,1] = 5.; m4.w[3,2] =-6.
g =  RecurrentNN.Graph()
m5 = RecurrentNN.relu!(g,m4)
@test m5.w[1,1] == 1.
@test m5.w[1,2] == 0.
@test m5.w[2,1] == 0.
@test m5.w[2,2] == 4.
@test m5.w[3,1] == 5.
@test m5.w[3,2] == 0.

m5.dw[1,1] =-.1; m5.dw[1,2] = .2
m5.dw[2,1] = .3; m5.dw[2,2] = .4
m5.dw[3,1] = .5; m5.dw[3,2] = .6

g.backprop[1]()
@test m4.dw[1,1] == -0.1
@test m4.dw[1,2] == 0.
@test m4.dw[2,1] == 0.
@test m4.dw[2,2] == 0.4
@test m4.dw[3,1] == 0.5
@test m4.dw[3,2] == 0.

# rowpluck!() tests
m4 = RecurrentNN.NNMatrix(3,2)
m4.w[1,1] = 1.; m4.w[1,2] =-2.
m4.w[2,1] =-3.; m4.w[2,2] = 4.
m4.w[3,1] = 5.; m4.w[3,2] =-6.

g =  RecurrentNN.Graph()
m5 = RecurrentNN.rowpluck!(g,m4,2)
@test m5.w[1,1] == -3.
@test m5.w[2,1] == 4.

m5.dw[1,1] =-.1; m5.dw[2,1] = .2
g.backprop[1]()
m4.dw
@test m4.dw[1,1] == 0
@test m4.dw[1,2] == 0
@test m4.dw[2,1] ==-0.1
@test m4.dw[2,2] == 0.2
@test m4.dw[3,1] == 0.
@test m4.dw[3,2] == 0.

# softmax tests
m6 = RecurrentNN.NNMatrix(5,1)
m6.w[1,1] = 0.3; m6.w[2,1] =0.1; m6.w[3,1] =0.6; m6.w[4,1] = 0.002; m6.w[5,1] = 0.00001
sm = RecurrentNN.softmax(m6)
# @test_approx_eq
@test_approx_eq sm.w[1,1] 0.21494050089813527
@test_approx_eq sm.w[2,1] 0.17597839816728894
@test_approx_eq sm.w[3,1] 0.2901393282421457
@test_approx_eq sm.w[4,1] 0.1595506217827435
@test_approx_eq sm.w[5,1] 0.15939115090968653


# Exponential decay at a point
#
# We set up a problem in which a tracer obeys
#
# ∂c/∂t = - c
#
# at a single point. This leads to exponential decay,
#
# c(t) = exp(-t)
#
# in the case that c(0) = 1. We use this problem to test
# the order of convergence of the time-stepper.

module PointExponentialDecay

using Printf
using Statistics
using Oceananigans

include("analysis.jl")

# Simple Exponential decay
c(t) = exp(-t)
Fᶜ = Relaxation(rate=1)

function run_test(; Δt, stop_iteration, timestepper, architecture = CPU())

    grid = RegularRectilinearGrid(size=(1, 1, 1), x=(0, 1), y=(0, 1), z=(0, 1))

    model = NonhydrostaticModel(architecture = architecture,
                                 timestepper = timestepper,
                                        grid = grid,
                                    coriolis = nothing,
                                    buoyancy = nothing,
                                     tracers = :c,
                                     forcing = (c=Fᶜ,))

    set!(model, c=1)
    simulation = Simulation(model, Δt=Δt, stop_iteration=stop_iteration, iteration_interval=stop_iteration)

    @info "Running point exponential decay test for Δt = $Δt, stop_iteration = $stop_iteration..."
    run!(simulation)

    # Calculate errors
    c_simulation = model.tracers.c[1, 1, 1]
    c_analytical = c(model.clock.time)
    L₁_error = abs(c_simulation - c_analytical)

    return (simulation = c_simulation,
            analytical = c_analytical,
                    L₁ = L₁_error)
end

end # module


type BatesProcess<: AbstractHestonProcess
  hestonProcess::HestonProcess
  lambda::Float64
  nu::Float64
  delta::Float64
  m::Float64
end

function BatesProcess(riskFreeRate::YieldTermStructure, dividendYield::YieldTermStructure, s0::Quote, v0::Float64,
                      kappa::Float64, theta::Float64, sigma::Float64, rho::Float64, lambda::Float64, nu::Float64,
                      delta::Float64, d::AbstractHestonDiscretization = FullTruncation())

  hp = HestonProcess(riskFreeRate, dividendYield, s0, v0, kappa, theta, sigma, rho, d)
  return BatesProcess(hp, lambda, nu, delta, expm1(nu + 0.5 * delta * delta))
end

get_risk_free_rate(bp::BatesProcess) = bp.hestonProcess.riskFreeRate
get_dividend_yield(bp::BatesProcess) = bp.hestonProcess.dividendYield
get_s0(bp::BatesProcess) = bp.hestonProcess.s0
get_v0(bp::BatesProcess) = bp.hestonProcess.v0
get_kappa(bp::BatesProcess) = bp.hestonProcess.kappa
get_theta(bp::BatesProcess) = bp.hestonProcess.theta
get_sigma(bp::BatesProcess) = bp.hestonProcess.sigma
get_rho(bp::BatesProcess) = bp.hestonProcess.rho
get_nu(bp::BatesProcess) = bp.nu
get_lambda(bp::BatesProcess) = bp.lambda
get_delta(bp::BatesProcess) = bp.delta
get_m(bp::BatesProcess) = bp.m


#Version Julia "1.1.1"
using DataFrames
using MathProgBase
using CSV
using LinearAlgebra
################################################################################
## Setting-up directory
tempdir1=@__DIR__
repdir=tempdir1[1:findfirst("ReplicationAK",tempdir1)[end]]
diroutput=repdir*"/Output_all"
dirdata=repdir*"/Data_all"

## Data
# Sample size
const n=154
# Number of time periods
const T=50
# Number of goods
const K=3
# Read csv files
dum0=CSV.read(dirdata*"/rationalitydata3goods.csv",DataFrame)
# break the dataset into subdatasets by individual
splitdum0=groupby(dum0,:id)
@eval  splitp=$splitdum0
#Initialize array of effective prices
rho=zeros(n,T,K)
#Initialize array of consumption
cve=zeros(n,T,K)
# Fill the arrays
# Columns 10:12 correspond to prices
# Columns 4:6 correspond to consumption bundles
for i=1:n
    dum0=convert(Array,splitp[i])
    rho[i,:,:]=dum0[1:T,10:12]
    cve[i,:,:]=dum0[1:T,4:6]
end

################################################################################
## Deterministic
ind=1
p=rho[ind,:,:]
q=cve[ind,:,:]

function rdmatrix(;p=p,q=q)
    rdmatrixfill=zeros(T,T)
    for t=1:T, s=1:T
      if (p[t,:]'*q[t,:])[1]>=(p[t,:]'*q[s,:])[1]
        rdmatrixfill[t,s]=1
      end
    end
    rdmatrixfill
end

function tclosure(;m=m)
    n0=size(m)[2]
    for k=1:n0, i=1:n0, j=1:n0
      if  m[i,k]==1 && m[k,j]==1
        m[i,j]=1
      end
    end
    m
end

function garpok(;p=p,q=q)
    garpmatrixfill=zeros(T,T)
    m=tclosure(m=rdmatrix(p=p,q=q))
    for t=1:T
        for s=1:T
            if m[t,s]==1 && ((p[s,:]'*q[s,:])[1]>(p[s,:]'*q[t,:])[1])
                garpmatrixfill[t,s]=1
            end
        end
    end

    maximum(garpmatrixfill)
    ## output of zero is a pass, no violation of GARP
end

function testgarp(;rho=rho,cve=cve)
    garpresults=zeros(n)
    for ind=1:n
      garpresults[ind]=garpok(p=rho[ind,:,:],q=cve[ind,:,:])
    end
    garpresults
end
garpresults=testgarp(rho=rho,cve=cve)
rate=1-sum(garpresults)/n

## Pass Rate
ratev=["pass-garp" rate]
DFsolv=convert(DataFrame,ratev)
CSV.write(diroutput*"/2App_dt_rr.csv",DFsolv)


#=--------------------------------------------------------------------

geometry - Functions for some basic geometric operations

Part of the ImageProjectiveGeometry Module

Copyright (c) 2016 Peter Kovesi
pk@peterkovesi.com

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

The Software is provided "as is", without warranty of any kind.

PK June 2016

---------------------------------------------------------------------=#
using LinearAlgebra

export ray2raydist, circleintersectionpts
export rectintersect, pointinconvexpoly, isconvexpoly
export convexpolyintersect

#--------------------------------------------------------------------------
"""
ray2raydist -  Minimum distance between two 3D rays

```
Usage: d = ray2raydist(p1, v1, p2, v2)

Arguments:
      p1, p2 - 3D points that lie on rays 1 and 2.
      v1, v2 - 3D vectors defining the direction of each ray.

Returns:
           d - The minimum distance between the rays.
```

Each ray is defined by a point on the ray and a vector giving the direction
of the ray.  Thus a point on ray 1 will be given by  p1 + alpha*v1  where
alpha is some scalar.
"""
function ray2raydist(p1::Vector, v1::Vector, p2::Vector, v2::Vector)

    # Get vector perpendicular to both rays
    n = cross(v1, v2)
    nnorm = sqrt(dot(n,n))

    # Check if lines are parallel. If so, form a vector perpendicular to v1
    # that is within the plane formed by the parallel rays.
    if nnorm < eps()
        n = cross(v1, p1 .- p2) # Vector perpendicular to plane formed by pair
                                # of rays.
        n = cross(v1, n)        # Vector perpendicular to v1 within the plane
                                # formed by the 2 rays.
        nnorm = sqrt(dot(n,n))

        if nnorm < eps()      # p1 and p2 must be coincident
            return  0.0
        end
    end

    n /= nnorm         # Unit vector

    # The vector joining the two closest points on the rays is:
    #    d*n = p2 + beta*v2 - (p1 + alpha*v1)
    # for some unknown values of alpha and beta.
    #
    # Take dot product of n on both sides
    #    d*n.n = p2.n + beta*v2.n - p1.n - alpha*v1.n
    #
    # as n is prependicular to v1 and v2 this reduces to
    #   d = (p2 - p1).n
    return abs(dot((p2 .- p1), n))
end

#--------------------------------------------------------------------------
"""
circleintersectionpts - Finds intersections of two circles.

Function computes the intersection points between two circles given
their centres and radii.

```
Usage: (i1, i2) = circleintersectionpts(c1, r1, c2, r2, lr)
             i1 = circleintersectionpts(c1, r1, c2, r2, lr)

Arguments:
        c1, c2 - 2-vectors specifying the centres of the two circles.
        r1, r2 - Radii of the two circles.
            lr - Optional string specifying what solutions are wanted:
                "l" for the solution to the left of the line from c1 to c2
                "r" for the solution to the right of the line from c1 to c2
                "lr" if both solutions are wanted (default).

Returns:
     (i1, i2) - Tuple of 2D intersection points if both solutions requested.
          i1  - Single 2D intersection point if a single solution requested.
```

If no solution exists empty vectors are returned for i1, i2.  If the
two circles are identical there are an infinite number of solutions.
In this case we simply choose to subtract/add r1 in the x direction to
c1 to obtain the 'left' and 'right' solutions.
"""
function circleintersectionpts(c1i::Vector, r1::Real, c2i::Vector, r2::Real, lr="lr")

    c1 = float(c1i) # Ensure c1 and c2 are floats for type stability
    c2 = float(c2i)

    maxmag = maximum(abs.([c1; c2; r1; r2]))  # Maximum value in data input
    tol = (maxmag+1)*eps()              # Tolerance used for deciding whether values are equal

    bv = c2 .- c1        # Baseline vector from c1 to c2
    b = norm(bv)         # The distance between the centres

    # Trap case of baseline of zero length.  If r1 == r2 we have a
    # valid geometric situation, but there are an infinite number of
    # solutions.  Here we simply choose to subtract/add r1 in the x
    # direction to c1 to obtain the 'left' and 'right' solutions.
    if b < tol && abs(r1-r2) < tol
        i1 = c1 .- [r1, 0]
        i2 = c1 .+ [r1, 0]
    else
        bv = bv/b             # Normalised baseline vector.
        bvp = [-bv[2], bv[1]] # Vector perpendicular to baseline

        # Trap the degenerate cases where one of the radii are zero, or nearly zero
        if r1 < tol && abs(b - r2) < tol
            i1 = c1
            i2 = c1
        elseif r2 < tol && abs(b - r1) < tol
            i1 = c2
            i2 = c2

        # Check triangle inequality
        elseif b > (r1+r2) || r1 > (b+r2) || r2 > (b+r1)
            i1 = zeros(eltype(c1), 0)
            i2 = zeros(eltype(c1), 0)

        else   # Normal solution
            cosR2 = (b^2 + r1^2 - r2^2)/(2*b*r1)
            sinR2 = sqrt(1-cosR2^2)

            i1 = c1 .+ r1*cosR2*bv .+ r1*sinR2*bvp   # 'left' solution
            i2 = c1 .+ r1*cosR2*bv .- r1*sinR2*bvp   # and 'right solution
        end
    end

    # Return requested solution(s)
    if lr == "lr"
        return (i1, i2)

    elseif lr == "l"
        return i1

    elseif lr == "r"
        return i2

    else
        error("illegal left/right solution request")
    end

end

#--------------------------------------------------------------------------
"""
rectintersect - Test if rectangles intersect

```
Usage:  intersect = rectintersect(r1, r2)

Arguments:  r1, r2 - The two rectangles to be tested where the rectangles
                     are defined using 4-vectors with the elements
                     defining  [left, bottom, width, height]

Returns:             Boolean result
```

See also: convexpolyintersect()
"""
function rectintersect(r1::Vector, r2::Vector)

    intersect = (r1[1] <= r2[1]+r2[3] &&
                 r2[1] <= r1[1]+r1[3] &&
                 r1[2] <= r2[2]+r2[4] &&
                 r2[2] <= r1[2]+r1[4])
end

#--------------------------------------------------------------------------
"""
pointinconvexpoly - Determine if a 2D point is within a convex polygon

```
Usage:  v = pointinconvexpoly(p, poly)

Arguments:   p - 2-vector specifying the point.
          poly - Convex polygon defined as a series of vertices in
                 sequence, clockwise or anticlockwise around the polygon as
                 a 2 x N array.

Returns:    v - +1 if within the polygon
                -1 if outside
                 0 if on the boundary
```

Note:  There is no check to see if the polygon is indeed convex, use
isconvexpoly() if needed

See also: isconvexpoly()
"""
function pointinconvexpoly(p::Vector, poly::Array)

    (dim, N) = size(poly)

    if length(p) != 2 || dim != 2
        error("Data must be 2D")
    end

    # Strategy: Determine whether, in traveling from p to a vertex and
    # then to the following vertex, we turn clockwise or anticlockwise.
    c = zeros(Int, N)
    for n = 1:N-1
        c[n] = clockwise(p, view(poly,:,n), view(poly,:,n+1))
    end
    c[N] = clockwise(p, view(poly,:,N), view(poly,:,1))

    # If for every vertex we turn consistently clockwise, or
    # consistently anticlockwise we are inside the polygon.  If for
    # one of the vertices we did not turn clockwise or anticlockwise
    # then we must be on the boundary.
    if all(c .>= 0) || all(c .<= 0)
        if any(c .== 0)  # We are on the boundary
            v = 0
        else             # We are inside
            v = 1
        end
    else                 # We are outside
        v = -1
    end

    return v
end

#----------------------------------------------------------------------
# Determine whether, in traveling from p1 to p2 to p3 we turn clockwise or
# anticlockwise.  Returns +1 for clockwise, -1 for anticlockwise, and 0 for
# p1, p2, p3 in a straight line.

function clockwise(p1::AbstractArray, p2::AbstractArray, p3::AbstractArray)

    # Form vectors p1->p2 and p2->p3 with z component = 0, form cross product
    # if the resulting z value is -ve the vectors turn clockwise, if +ve
    # anticlockwise, and if 0 the points are in a line.

    return -sign((p2[1]-p1[1])*(p3[2]-p2[2]) - (p2[2]-p1[2])*(p3[1]-p2[1]))

end

#----------------------------------------------------------------------
"""
convexpolyintersect - Returns true if convex polygons intersect.

```
Usage:  intersect = convexpolyintersect(p1, p2)

Arguments:  p1, p2 - The two convex polygons to be tested where the
                     polygons are defined as 2xN arrays of vertex
                     coordinates in sequence around the polygon.

Returns:             Boolean result.
```

Note: There is no check to see if the polygons are indeed convex, use
isconvexpoly() if needed

See also: isconvexpoly(), rectintersect()
"""
function convexpolyintersect(p1::Array, p2::Array)

    # Function to test if two 1D segments overlap.
    # Each segment is defined by a 2 element vector with 1st element
    # being the minimum x value and 2nd element the maximum x value.
    function  overlap(s1::Vector, s2::Vector)
        return  s2[2] >= s1[1] && s1[2] >= s2[1]
    end

    if size(p1,1) != 2 || size(p2,1) != 2
        error("Polygon data must be in 2xN arrays")
    end

    # Build a sequence of vertices that trace the edges of both
    # polygons.  We prepend the last vertex to each polygon list to
    # form a complete cycle.  We also concatenate the two polygon
    # arrays together, this adds an extra 'edge' that connects the
    # last vertex of the 1st polygon with the last vertex of the 2nd
    # polygon.  This allows us to run a single loop of tests below.
    p = [p1[:,end] p1  p2[:,end] p2]

    for n = 2:size(p,2)
        # Generate a 2x2 rotation matrix describing a frame with the y axis
        # aligned with the polygon side.
        y = p[:,n] - p[:,n-1]
        normy = sqrt(y[1]^2 + y[2]^2)
        if normy < eps()  # repeated vertex
            continue
        end
        y = y./normy
        x = [y[2], -y[1]]
        R = [x y]'

        # Rotate all vertices to frame defined by this side
        p1r = R*p1
        p2r = R*p2

        # Check whether x coordinates of rotated p1 vertices overlap
        # with x coordinates of rotated p2 vertices.  If they do not
        # overlap we have found a separating axis indicating that the
        # polygons do not intersect.
        if !overlap([minimum(p1r[1,:]), maximum(p1r[1,:])],
                    [minimum(p2r[1,:]), maximum(p2r[1,:])])
            return false
        end
    end

    return true        # No separating axis found
end

#--------------------------------------------------------------------------
"""
isconvexpoly - Test if 2D polygon is convex.

```
Usage: v = isconvexpoly(p)

Argument:  p - 2D polygon are defined as 2xN array of vertices.

Returns:   true if polygon is convex.

```

See also: convexpolyintersect(), pointinconvexpoly()
"""
function isconvexpoly(poly::Array)

    (dim, N) = size(poly)

    if dim != 2
        error("Polygon data must be a 2xN array")
    end

    # Strategy: Determine turning direction as one travels through
    # each set of three successive vertices.  If, for all vertices,
    # the turning direction is always clockwise or straight, or always
    # anti-clockwise or straight then the polygon is convex.

    c = zeros(Int, N)
    for n = 1:N-2
        c[n] = clockwise(view(poly,:,n), view(poly,:,n+1), view(poly,:,n+2))
    end
    c[N-1] = clockwise(view(poly,:,N-1), view(poly,:,N), view(poly,:,1))
    c[N] = clockwise(view(poly,:,N), view(poly,:,1), view(poly,:,2))

    if all(c .>= 0) || all(c .<= 0)
        return true
    else
        return false
    end
end


@testset verbose=true "𝔇" begin

    @testset "Compare H/D indexing ($T)" for T in [Float64, Float32]
        # Here, we check that we can pass in either an "H wedge" array to be used with
        # WignerHindex, or a full 𝔇 array used with WignerDindex, and obtain the same
        # H recurrence results
        ℓₘₐₓ = 8
        for m′ₘₐₓ in 0:ℓₘₐₓ
            expiβ = cis(rand(0:eps(T):π))
            expiβNaNCheck = complex(NaNCheck{T}(expiβ.re), NaNCheck{T}(expiβ.im))
            NCTN = NaNCheck{T}(NaN)
            Hw = fill(NCTN, WignerHsize(ℓₘₐₓ, m′ₘₐₓ))
            H!(Hw, expiβNaNCheck, ℓₘₐₓ, m′ₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
            𝔇 = fill(NCTN, WignerDsize(ℓₘₐₓ, m′ₘₐₓ))
            H!(
                𝔇, expiβNaNCheck, ℓₘₐₓ, m′ₘₐₓ,
                H_recursion_coefficients(ℓₘₐₓ, T), WignerDindex
            )
            for n in 0:ℓₘₐₓ
                for m′ in -min(n, m′ₘₐₓ):min(n, m′ₘₐₓ)
                    for m in abs(m′):n
                        Hnm′m = Hw[WignerHindex(n, m′, m, m′ₘₐₓ)]
                        𝔇nm′m = 𝔇[WignerDindex(n, m′, m, m′ₘₐₓ)]
                        @test Hnm′m == 𝔇nm′m
                    end
                end
            end
        end
    end

    @testset "Compare 𝔇 to formulaic d ($T)" for T in [BigFloat, Float64, Float32]
        # Now, we're ready to check that d_{n}^{m′,m}(β) matches the expected values
        # for a range of β values
        for ℓₘₐₓ in 0:4
            H_rec_coeffs = H_recursion_coefficients(ℓₘₐₓ, T)
            𝔇 = Array{Complex{T}}(undef, WignerDsize(ℓₘₐₓ, ℓₘₐₓ))
            expimα = Array{Complex{T}}(undef, ℓₘₐₓ+1)
            expimγ = Array{Complex{T}}(undef, ℓₘₐₓ+1)
            expiα = complex(one(T))
            expiγ = complex(one(T))
            for β in βrange(T)
                expiβ = cis(β)
                R = from_euler_angles(zero(T), β, zero(T))
                D!(𝔇, R, ℓₘₐₓ, H_rec_coeffs, expimα, expimγ)
                for n in 0:ℓₘₐₓ
                    for m′ in -n:n
                        for m in -n:n
                            𝔇_formula = ExplicitWignerMatrices.D_formula(
                                n, m′, m, expiα, expiβ, expiγ
                            )
                            𝔇_recurrence = 𝔇[WignerDindex(n, m′, m)]
                            @test 𝔇_formula ≈ 𝔇_recurrence atol=200eps(T) rtol=200eps(T)
                        end
                    end
                end
            end
        end
    end

    @testset "Compare 𝔇 to formulaic 𝔇 ($T)" for T in [BigFloat, Float64, Float32]
        # Now, we're ready to check that 𝔇_{n}^{m′,m}(β) matches the expected values
        # for a range of α, β, γ values
        Random.seed!(123)
        ℓₘₐₓ = T===BigFloat ? 4 : 8
        H_rec_coeffs = H_recursion_coefficients(ℓₘₐₓ, T)
        𝔇 = Array{Complex{T}}(undef, WignerDsize(ℓₘₐₓ, ℓₘₐₓ))
        expimα = Array{Complex{T}}(undef, ℓₘₐₓ+1)
        expimγ = Array{Complex{T}}(undef, ℓₘₐₓ+1)
        @showprogress "Compare 𝔇 to formulaic 𝔇 ($T)" for α in αrange(T, 5)
            for β in βrange(T, 5)
                for γ in γrange(T, 5)
                    R = from_euler_angles(α, β, γ)
                    expiα, expiβ, expiγ = to_euler_phases(R)
                    D!(𝔇, R, ℓₘₐₓ, H_rec_coeffs, expimα, expimγ)
                    for n in 0:ℓₘₐₓ
                        for m′ in -n:n
                            for m in -n:n
                                𝔇_formula = ExplicitWignerMatrices.D_formula(
                                    n, m′, m, expiα, expiβ, expiγ
                                )
                                𝔇_recurrence = 𝔇[WignerDindex(n, m′, m)]
                                @test 𝔇_formula ≈ 𝔇_recurrence atol=30eps(T) rtol=30eps(T)
                            end
                        end
                    end
                end
            end
        end
    end

    @testset "Group characters $T" for T in [BigFloat, Float64, Float32]
        # χʲ(β) ≔ Σₘ dʲₘₘ(β) ≡ Σₘ 𝔇ʲₘₘ(exp(v̂ β/2)) = sin((2j+1)β/2) / sin(β/2)
        # Here, v̂ is any unit vector; group characters are constant on conjugacy classes and
        # conjugacy classes of SO(3) are rotations through the same angle about any axis.
        ℓₘₐₓ = T===BigFloat ? 10 : 20
        m′ₘₐₓ = ℓₘₐₓ
        H_rec_coeffs = H_recursion_coefficients(ℓₘₐₓ, T)
        d = Array{T}(undef, WignerDsize(ℓₘₐₓ, m′ₘₐₓ))
        𝔇 = Array{Complex{T}}(undef, WignerDsize(ℓₘₐₓ, m′ₘₐₓ))
        expimα = Array{Complex{T}}(undef, ℓₘₐₓ+1)
        expimγ = Array{Complex{T}}(undef, ℓₘₐₓ+1)
        @showprogress "Group characters $T" for β in βrange(T)
            expiβ = cis(β)
            d!(d, expiβ, ℓₘₐₓ, H_rec_coeffs)
            for j in 0:ℓₘₐₓ
                sin_ratio = sin((2j+1)*β/2) / sin(β/2)
                if abs(β) < 10eps(T)
                    sin_ratio = T(2j+1)
                elseif abs(β-π) < 10eps(T)
                    sin_ratio = T(-1)^j
                end
                χʲ = sum(d[WignerDindex(j, m, m)] for m in -j:j)
                @test χʲ ≈ sin_ratio atol=500eps(T) rtol=500eps(T)
                for v̂ in v̂range(T)
                    R = exp(β/2 * v̂)
                    D!(𝔇, R, ℓₘₐₓ, H_rec_coeffs, expimα, expimγ)
                    χʲ = sum(𝔇[WignerDindex(j, m, m)] for m in -j:j)
                    @test χʲ ≈ sin_ratio atol=500eps(T) rtol=500eps(T)
                end
            end
        end
    end

    @testset "Representation property ($T)" for T in [Float64, Float32, BigFloat]
        # For each l, 𝔇ˡₙ,ₘ(R₁ R₂) = Σₚ 𝔇ˡₙ,ₚ(R₁) 𝔇ˡₚ,ₘ(R₂)
        tol = 3eps(T)
        ℓₘₐₓ = 10
        𝔇₁ = Array{Complex{T}}(undef, WignerDsize(ℓₘₐₓ))
        𝔇₂ = Array{Complex{T}}(undef, WignerDsize(ℓₘₐₓ))
        𝔇₁₂ = Array{Complex{T}}(undef, WignerDsize(ℓₘₐₓ))
        H_rec_coeffs = H_recursion_coefficients(ℓₘₐₓ, T)
        expimα = Array{Complex{T}}(undef, ℓₘₐₓ+1)
        expimγ = Array{Complex{T}}(undef, ℓₘₐₓ+1)
        @showprogress "Representation property ($T)" for R₁ in Rrange(T)
            for R₂ in Rrange(T)
                D!(𝔇₁, R₁, ℓₘₐₓ, H_rec_coeffs, expimα, expimγ)
                D!(𝔇₂, R₂, ℓₘₐₓ, H_rec_coeffs, expimα, expimγ)
                D!(𝔇₁₂, R₁*R₂, ℓₘₐₓ, H_rec_coeffs, expimα, expimγ)
                for ℓ in 0:ℓₘₐₓ
                    i = WignerDindex(ℓ, -ℓ, -ℓ)
                    j = WignerDindex(ℓ, ℓ, ℓ)
                    𝔇₁ˡ = transpose(reshape(𝔇₁[i:j], 2ℓ+1, 2ℓ+1))
                    𝔇₂ˡ = transpose(reshape(𝔇₂[i:j], 2ℓ+1, 2ℓ+1))
                    𝔇₁₂ˡ = transpose(reshape(𝔇₁₂[i:j], 2ℓ+1, 2ℓ+1))
                    @test 𝔇₁ˡ * 𝔇₂ˡ ≈ 𝔇₁₂ˡ atol=(2ℓ+1)^2*tol rtol=(2ℓ+1)^2*tol
                end
            end
        end
    end

end