using PyPlot, Statistics, NbodyGradient, LinearAlgebra

export State

GNEWT = NbodyGradient.GNEWT

function compute_energy(m::Array{T,1},x::Array{T,2},v::Array{T,2},n::Int64) where {T <: Real}
  KE = 0.0
  for j=1:n
    KE += 0.5*m[j]*(v[1,j]^2+v[2,j]^2+v[3,j]^2)
  end
  PE = 0.0
  for j=1:n-1
    for k=j+1:n
       PE += -GNEWT*m[j]*m[k]/norm(x[:,j] .- x[:,k])
    end
  end
  ang_mom = zeros(3)
  for j=1:n
    ang_mom .+= m[j]*cross(x[:,j],v[:,j])
  end
  return KE,PE,ang_mom
end

n=5
xout = zeros(3,n)
# Positions at time September 5, 1994 at 0h00 in days (from Hairer, Lubich & Wanner
# 2006, Geometric Numerical Integration, 2nd Edition, Springer, pp. 13-14):

xout .= transpose([0.0 0.0 0.0;
  -3.5023653  -3.8169847  -1.5507963;
  9.0755314  -3.0458353  -1.6483708;
  8.3101420  -16.2901086  -7.2521278;
  11.4707666  -25.7294829  -10.8169456])

vout = zeros(3,n)
vout .= transpose([0.0 0.0 0.0;
  0.00565429  -0.00412490  -0.00190589;
  .00168318  .00483525  .00192462;
  .00354178  .00137102  .00055029;
  .00288930  .00114527  .00039677])

# Specify masses, including terrestrial planets in the Sun:
m = [1.00000597682,0.000954786104043,0.000285583733151,
         0.0000437273164546,0.0000517759138449] #,6.58086572e-9];

# Compute the center-of-mass:
vcm = zeros(3);
xcm = zeros(3);
for j=1:n
    vcm .+= m[j]*vout[:,j];
    xcm .+= m[j]*xout[:,j];
end
vcm ./= sum(m);
xcm ./= sum(m);
# Adjust so CoM is stationary
for j=1:n
    vout[:,j] .-= vcm[:];
    xout[:,j] .-= xcm[:];
end

#=
struct CartesianElements{T} <: NbodyGradient.InitialConditions{T}
    x::Matrix{T}
    v::Matrix{T}
    m::Vector{T}
    nbody::Int64
end

function NbodyGradient.State(ic::NbodyGradient.InitialConditions{T}) where T<:AbstractFloat
    n = ic.nbody
    x = copy(ic.x)
    v = copy(ic.v)
    jac_init = zeros(7*n,7*n)
    xerror = zeros(T,size(x))
    verror = zeros(T,size(v))
    jac_step = Matrix{T}(I,7*n,7*n)
    dqdt = zeros(T,7*n)
    dqdt_error = zeros(T,size(dqdt))
    jac_error = zeros(T,size(jac_step))
    pair = zeros(Bool,n,n)

    rij = zeros(T,3)
    a = zeros(T,3,n)
    aij = zeros(T,3)
    x0 = zeros(T,3)
    v0 = zeros(T,3)
    input = zeros(T,8)
    delxv = zeros(T,6)
    rtmp = zeros(T,3)
    return State(x,v,[0.0],copy(ic.m),jac_step,dqdt,jac_init,xerror,verror,dqdt_error,jac_error,n,
    pair,rij,a,aij,x0,v0,input,delxv,rtmp)
end
=#

ic = CartesianIC(xout,vout,m,n,0.0);

# Now, integrate this forward in time:
xerror = zeros(3,n); verror = zeros(3,n);
hgrid = [1.5625,3.125,6.25,12.5,25.0,50.0,100.0,200.0]

power = 20
nstep = 2^power
grad = false
ngrid = length(hgrid)

# 4e6 time steps ~ 2 Myr (takes about 9 minutes to run)
# I've reducd that to ~400,000 to take ~1 minute.

# Set up array to save the state as a function of time:
# Save the potential & kinetic energy, as well as angular momentum:
egrid = zeros(nstep); ang_mom = zeros(3,nstep)
sig_energy = zeros(ngrid); sig_ang_mom = zeros(3,ngrid)
mean_energy = zeros(ngrid); mean_ang_mom = zeros(3,ngrid)
# Time the integration:
for j=1:length(hgrid)
  h = hgrid[j]
  s = State(ic)
  # Specify Kepler step only between planets & star:
  s.pair .= true
  for i=1:s.n
    s.pair[1,i] = false
    s.pair[i,1] = false
  end
  if grad; d = Derivatives(T,s.n); end

  tstart = time()
  # Carry out the integration:
  for i=1:nstep
       if grad
      ahl21!(s,d,h)
    else
      ahl21!(s,h)
    end
    KE_step,PE_step,ang_mom_step=compute_energy(s.m,s.x,s.v,n)
    egrid[i] = KE_step+PE_step
    ang_mom[:,i] = ang_mom_step
  end
  mean_energy[j] = mean(egrid)
  sig_energy[j] = std(egrid)
  mean_ang_mom[1,j] = mean(ang_mom[1,:])
  mean_ang_mom[2,j] = mean(ang_mom[2,:])
  mean_ang_mom[3,j] = mean(ang_mom[3,:])
  sig_ang_mom[1,j] = std(ang_mom[1,:])
  sig_ang_mom[2,j] = std(ang_mom[2,:])
  sig_ang_mom[3,j] = std(ang_mom[3,:])
  telapse = time()- tstart
  println(j," ",telapse," ",sig_energy[j]," ",sig_ang_mom[:,j])
end

# These integrations have already been carried out:  so,
# just plotting the results:

clf()
tstep = hgrid
loglog(tstep,sig_energy,"o",label="Measured scatter in energy")
loglog(tstep,1.403145134013028e-16 .* (tstep ./ tstep[8]).^4,label="Quartic scaling with time step")
xlabel("Time step [days]")
ylabel(L"Energy error [$M_\odot$ AU$^2$ day$^{-2}$]")
loglog([1.5625,200],abs(mean_energy[end])*2.2e-16*[1,1],linestyle=":",label="Double precision floor")
legend()

tight_layout()
savefig("Energy_error_vs_timestep_mixed.pdf",bbox_inches="tight")

clf()

loglog(tstep,sig_ang_mom[1,:],label=L"$L_x$")
loglog(tstep,sig_ang_mom[2,:],label=L"$L_y$")
loglog(tstep,sig_ang_mom[3,:],label=L"$L_z$")
loglog(tstep,mean_ang_mom[1,end]*2.2e-16*ones(8),linestyle=":",color="blue")
loglog(tstep,mean_ang_mom[2,end]*2.2e-16*ones(8),linestyle=":",color="orange")
loglog(tstep,mean_ang_mom[3,end]*2.2e-16*ones(8),linestyle=":",color="green")
loglog(tstep,1e-21 .* tstep, color="k",linestyle="--")
xlabel("Time step [days]")
ylabel(L"Angular momentum error [$M_\odot$ AU$^2$ day$^{-1}$]")
legend()

tight_layout()
savefig("Angular_momentum_error_vs_timestep_mixed.pdf",bbox_inches="tight")


using OrdinaryDiffEq, DiffEqBase, DiffEqProblemLibrary, DiffEqDevTools

f = (t,u) -> u
prob = ODEProblem(f,1/2,(0.0,1.0))
analytic = (t,u0) -> u0*exp(t)

sol =solve(prob,Euler();dt=1//2^(4))

sol2 =solve(prob,Vern9();dt=1//2^(10))

prob2 = ODETestProblem(f,1/2,analytic,(0.0,1.0))
sol3 =solve(prob_ode_linear,Euler();dt=1//2^(4))

errsol1 = appxtrue(sol,sol2)

sol4 =solve(prob,Euler();dt=1//2^(4))
test_sol = TestSolution(sol2)
errsol2 = appxtrue(sol4,test_sol)

sol5 =solve(prob,Euler();dt=1//2^(4))
test_sol = TestSolution(sol2.t,sol2[end])
errsol3 = appxtrue(sol5,test_sol)

@test errsol1.errors[:L2] ≈ 0.018865798306718855 && errsol1.errors[:L2] ≈ errsol2.errors[:L2]


using JLD2, Printf, OceanTurb

@use_pyplot_utils

# For this example we use the same "simple flux" scenario used in kpp_examples.ipynb.

function simple_flux_model(; N=10, H=100, Tz=0.01, Qb=1e-8, Qu=0, parameters=KPP.Parameters())

    model = KPP.Model(N=N, H=H, parameters=parameters, stepper=:BackwardEuler)

    # Initial condition
    T₀(z) = 20 + Tz*z

    # Set T to the function T0(z)
    model.solution.T = T₀

    # Fluxes
    Qθ = Qb / (model.constants.α * model.constants.g)
    model.bcs.U.top = FluxBoundaryCondition(Qu)
    model.bcs.T.top = FluxBoundaryCondition(Qθ)
    model.bcs.T.bottom = GradientBoundaryCondition(Tz)

    return model
end

## # Saving data
#
# To facilitate data, we define a few functions that save useful things both on model
# initialization (init_data), and while the model is running (save_data).
#
# We use JLD2.jl.

function save_data(path, model)
    iteration = iter(model)
    jldopen(path, "a+") do file
        file["timeseries/t/$iteration"] = time(model)
        file["timeseries/U/$iteration"] = collect(data(model.solution.U))
        file["timeseries/T/$iteration"] = collect(data(model.solution.T))
    end
end

function init_data(path, model)
    Qu = OceanTurb.getbc(model, model.bcs.U.top)
    Qb = OceanTurb.getbc(model, model.bcs.T.top) * model.constants.α * model.constants.g
    Tz = OceanTurb.getbc(model, model.bcs.T.bottom)

    jldopen(path, "a+") do file
        for gridfield in (:N, :H)
            file["grid/$gridfield"] = getproperty(model.grid, gridfield)
        end

        file["bcs/Qu"] = Qu
        file["bcs/Qb"] = Qb
        file["bcs/Tz"] = Tz
    end

    return nothing
end

# Next, we generate a model with some benign parameters and run, saving data every `dout` seconds.

dt = 10 # time step in seconds
dout = 10*minute # interval between output
tfinal = 4*hour # final time

ntot = Int(tfinal/dt)

nint = Int(dout/dt)
nout = Int(ntot/nint)

filepath = joinpath(".", "test_free_convection.jld2")

model = simple_flux_model(
     N = 100,
     H = 100,
    Tz = 4.08e-4,
    Qb = 5e-8,
    Qu = 0
)

isfile(filepath) && rm(filepath)
init_data(filepath, model)

for i = 1:nout
    @printf "\nrunning from iter %d" iter(model)
    @time time_step!(model, dt, nint)

    @printf "    ...saving at iter %d" iter(model)
    @time save_data(filepath, model)
end

# # Opening and analyzing data
#
# Running the model created a jld file, which we can open and look at:

file = jldopen(filepath, "r")

@show file

N = file["grid/N"]
H = file["grid/H"]
@show iters = parse.(Int, keys(file["timeseries/t"]))

close(file)

# # Plotting data
#
# Now we can plot the data (and otherwise analyze it if we wish).

grid = UniformGrid(N, H)

function get_t_and_T(path, i)
    file = jldopen(path, "r")
    t = file["timeseries/t/$i"]
    T = file["timeseries/T/$i"]
    close(file)
    t, T
end

fig, axs = subplots()
cornerspines()
xlabel(L"T")
ylabel(L"z")

for (ii, iter) in enumerate(iters)
    t, T = get_t_and_T(filepath, iter)
    plotlabel = ii % 4 == 0 ? @sprintf("\$ t = %.2f \$ hours", t/hour) : ""
    plot(T, grid.zc, label=plotlabel)
end

legend()


println("\nTesting extraction from Terms\n")

# testing extracting stuff from terms

m,n1 = 5,3
x1 = Variable(n1)
A = randn(m,n1)
# single term, single variable
cf = ls(A*x1)
xAll = StructuredOptimization.extract_variables(cf)
@test xAll[1] == x1
L = StructuredOptimization.extract_operators(xAll,cf)
@test typeof(L) <: MatrixOp
La = StructuredOptimization.extract_affines(xAll,cf)
@test typeof(La) <: MatrixOp
f = StructuredOptimization.extract_functions(cf)
@test typeof(f) <: SqrNormL2

# multiple terms, single variable
b1 = randn(n1)
cf = ls(A*x1) + 2.5*norm(x1+b1,1)
xAll = StructuredOptimization.extract_variables(cf)
@test xAll[1] == x1
V = StructuredOptimization.extract_operators(xAll,cf)
@test typeof(V) <: VCAT
@test typeof(V[1]) <: MatrixOp
@test typeof(V[2]) <: Eye
V2 = StructuredOptimization.extract_affines(xAll,cf)
@test typeof(V2) <: VCAT
@test typeof(V2[1]) <: MatrixOp
@test typeof(V2[2]) <: AffineAdd{T} where {T <: Eye}
f = StructuredOptimization.extract_functions(cf)
@test typeof(f) <: SeparableSum
@test typeof(f.fs[1]) <: SqrNormL2
@test typeof(f.fs[2]) <: Postcompose
x = randn(n1)
@test norm(f.fs[2](x) - 2.5*norm(x+b1,1)) < 1e-12

# single term, multiple variables
x2 = Variable(m)
cf = ls(A*x1+x2+20)
xAll = StructuredOptimization.extract_variables(cf)
xAll = (x2,x1) # change the order on pourpose
H = StructuredOptimization.extract_operators(xAll,cf)
@test typeof(H) <: HCAT
@test typeof(H[1]) <: Eye
@test typeof(H[2]) <: MatrixOp
H2 = StructuredOptimization.extract_affines(xAll,cf)
@test typeof(H2[1]) <: AffineAdd{T} where {T <: Eye}
@test typeof(H2[2]) <: AffineAdd{T} where {T <: MatrixOp}
f = StructuredOptimization.extract_functions(cf)
@test typeof(f) <: PrecomposeDiagonal

### multiple terms, multiple variables
n1,n2,n3,n4,n5 = 3,3,4,4,7
A = randn(n5,n1)
x1,x2,x3,x4,x5 = Variable(randn(n1)),Variable(randn(n2)),Variable(randn(n3)),Variable(randn(n4)),Variable(randn(n5))

cf = ls(x1+x2)
xAll = StructuredOptimization.extract_variables(cf)
@test xAll == (x1,x2)

cf = ls(x1+x2)+ls(x1)
xAll = StructuredOptimization.extract_variables(cf)
@test xAll == (x1,x2)

cf = ls(x1+x2)+ls(x3+x4)+ls(x5)+ls(x5+A*x2)+ls(x1)+ls(x5)
xAll = StructuredOptimization.extract_variables(cf)
@test xAll == (x1,x2,x3,x4,x5)

V = StructuredOptimization.extract_operators(xAll,cf)

@test typeof(V[1][1]) <: Eye
@test typeof(V[1][2]) <: Eye
@test typeof(V[1][3]) <: Zeros
@test typeof(V[1][4]) <: Zeros
@test typeof(V[1][5]) <: Zeros

@test typeof(V[2][1]) <: Zeros
@test typeof(V[2][2]) <: Zeros
@test typeof(V[2][3]) <: Eye
@test typeof(V[2][4]) <: Eye
@test typeof(V[2][5]) <: Zeros

@test typeof(V[3][1]) <: Zeros
@test typeof(V[3][2]) <: Zeros
@test typeof(V[3][3]) <: Zeros
@test typeof(V[3][4]) <: Zeros
@test typeof(V[3][5]) <: Eye

@test typeof(V[4][1]) <: Zeros
@test typeof(V[4][2]) <: MatrixOp
@test typeof(V[4][3]) <: Zeros
@test typeof(V[4][4]) <: Zeros
@test typeof(V[4][5]) <: Eye

@test typeof(V[5][1]) <: Eye
@test typeof(V[5][2]) <: Zeros
@test typeof(V[5][3]) <: Zeros
@test typeof(V[5][4]) <: Zeros
@test typeof(V[5][5]) <: Zeros

@test typeof(V[6][1]) <: Zeros
@test typeof(V[6][2]) <: Zeros
@test typeof(V[6][3]) <: Zeros
@test typeof(V[6][4]) <: Zeros
@test typeof(V[6][5]) <: Eye

println("\nTesting splitting Terms\n")

x = Variable(5)
y = Variable(5)
cf = ls(x)+10*norm(x,2)+ls(x+y)

f, g = StructuredOptimization.split_smooth(cf)
@test f[1] == cf[1]
@test f[2] == cf[3]
@test g[1] == cf[2]

cf = ls(x)
f, g = StructuredOptimization.split_smooth((cf,))
@test f == (cf,)
@test g == ()

cf = norm(x,1)+norm(y,2)+norm(randn(5,5)*x+y,Inf)
xAll = StructuredOptimization.extract_variables(cf)
AAc, nonAAc = StructuredOptimization.split_AAc_diagonal(cf)
@test AAc[1] == cf[1]
@test AAc[2] == cf[2]
@test nonAAc[1] == cf[3]

cf = ls(sigmoid(x)) + ls(x)
fq, fs = StructuredOptimization.split_quadratic(cf)
@test fs[1] == cf[1]
@test fq[1] == cf[2]

println("\nTesting extracting Proximable functions\n")
# testing is_proximable
@test StructuredOptimization.is_proximable(AAc) == true
@test StructuredOptimization.is_proximable(nonAAc) == false

cf = norm(x[1:2],1)+norm(x[3:5])
xAll = StructuredOptimization.extract_variables(cf)

@test all(StructuredOptimization.is_AAc_diagonal.(cf)) == true
@test StructuredOptimization.is_proximable(cf) == true

cf = norm(x[1:2],1)+norm(x[3:5])+norm(x,Inf)
xAll = StructuredOptimization.extract_variables(cf)

@test all(StructuredOptimization.is_AAc_diagonal.(cf)) == true
@test StructuredOptimization.is_proximable(cf) == false

# testing extract_proximable
# single variable, single term
x = Variable(randn(5))
b = randn(5)
cf = 10*norm(x-b,1)
xAll = StructuredOptimization.extract_variables(cf)
@test StructuredOptimization.is_proximable(cf) == true

f = StructuredOptimization.extract_proximable(xAll,cf)
@test norm(f(~x) - 10*norm(~x-b,1)) < 1e-12

# single variable, single term, diagonal term
x = Variable(randn(5))
b = randn(5)
d = randn(5)
cf = 10*norm(d.*x-b,1)
xAll = StructuredOptimization.extract_variables(cf)
@test StructuredOptimization.is_proximable(cf) == true

f = StructuredOptimization.extract_proximable(xAll,cf)
@test norm(f(~x) - 10*norm(d.*~x-b,1)) < 1e-12

# single variable, single term, tight frame term
x = Variable(randn(5))
b = randn(5)
d = randn(5)
cf = 10*norm(dct(x)-b,1)
xAll = StructuredOptimization.extract_variables(cf)
@test StructuredOptimization.is_proximable(cf) == true

f = StructuredOptimization.extract_proximable(xAll,cf)
@test norm(f(~x) - 10*norm(dct(~x)-b,1)) < 1e-12

# single variable, single term, tight frame term, fft
# TODO this not working (probably fix needed in ProxOp)
#x = Variable(randn(5))
#b = randn(5)
#d = randn(5)
#cf = 10*norm(fft(x)-b,1)
#xAll = StructuredOptimization.extract_variables(cf)
#@test StructuredOptimization.is_proximable(cf) == true
#
#f = StructuredOptimization.extract_proximable(xAll,cf)
#@test norm(f(~x) - 10*norm(fft(~x)-b,1)) < 1e-12

# single variable, multiple terms with GetIndex
x = Variable(randn(5))
b = randn(2)
cf = 10*norm(x[1:2]-b,1)+norm(x[3:5],2)
xAll = StructuredOptimization.extract_variables(cf)
@test StructuredOptimization.is_proximable(cf) == true
f = StructuredOptimization.extract_proximable(xAll,cf)
@test norm(f(~x) - sum([10*norm((~x)[1:2]-b,1);norm((~x)[3:5],2)])) < 1e-12

# single variable, multiple terms with GetIndex composed with dct
x = Variable(randn(5))
b = randn(2)
cf = 10*norm(x[1:2]-b,1)+norm(dct(x[3:5]),2)
xAll = StructuredOptimization.extract_variables(cf)
@test StructuredOptimization.is_proximable(cf) == true
f = StructuredOptimization.extract_proximable(xAll,cf)
@test norm(f(~x) - sum([10*norm((~x)[1:2]-b,1);norm(dct((~x)[3:5]),2)])) < 1e-12

# multiple variables, multiple terms
x1 = Variable(randn(5))
b1 = randn(5)
x2 = Variable(randn(3))
b2 = randn(3)

cf = 10*norm(x2-b2,1)+norm(x1+b1,2)
xAll = (x1,x2)
@test StructuredOptimization.is_proximable(cf) == true
f = StructuredOptimization.extract_proximable(xAll,cf)
@test norm(f.fs[1](~x1)-norm(~x1+b1,2) ) < 1e-12
@test norm(f.fs[2](~x2)-10*norm(~x2-b2,1) ) < 1e-12

x1 = Variable(randn(5))
b1 = randn(5)
x2 = Variable(randn(5))
b2 = randn(5)

# TODO fix this?
#cf = 10*norm(x2+x1+b2,1)
#xAll = (x1,x2)
#@test StructuredOptimization.is_proximable(cf) == true
#f = StructuredOptimization.extract_proximable(xAll,cf)
# TODO fix this! in ProxOp?
# @test norm(f((~x1,~x2))-10*norm(~x2+~x1+b2,1) ) < 1e-12

# multiple variables, missing terms
x1 = Variable(randn(5))
b1 = randn(5)
x2 = Variable(randn(3))
b2 = randn(3)

cf = 10*norm(x2-b2,1)
xAll = (x1,x2)
@test StructuredOptimization.is_proximable(cf) == true
f = StructuredOptimization.extract_proximable(xAll,cf)
@test f.fs[1](~x1) == 0.
@test norm(f.fs[2](~x2)-10*norm(~x2-b2,1) ) < 1e-12

# multiple variables, multiple terms, with GetIndex
x1 = Variable(randn(5))
b1 = randn(5)
x2 = Variable(randn(3))
b2 = randn(3)

cf = norm(x1[3:5]+b1[3:5],1)+10*norm(x2-b2,1)+norm(x1[1:2]+b1[1:2],2)
xAll = (x1,x2)
@test StructuredOptimization.is_proximable(cf) == true
f = StructuredOptimization.extract_proximable(xAll,cf)
@test norm(f.fs[1](~x1)-norm((~x1)[1:2]+b1[1:2],2)-norm((~x1)[3:5]+b1[3:5],1) ) < 1e-12
@test norm(f.fs[2](~x2)-10*norm(~x2-b2,1) ) < 1e-12


# roessler chaotic phase diff
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random, OrdinaryDiffEq, LinearAlgebra

a = 0.2
b = 0.2
c = 5.7
amu = 0.02
ω1 = 1.0 + amu
ω2 = 1.0 - amu
k1 = 0.01
k2 = k1

y_init = [0.1, 0.2, 0., 0.11, 0.19, 0.1]

croesslers = Systems.coupled_roessler(y_init; ω1, ω2, a, b, c, k1, k2)
diffeq = (alg = Vern9(), reltol = 1e-9, abstol = 1e-9)
Δt = 0.01
T = 1000
t = 0:Δt:T
L = length(t)÷10

ks = [0.03, 0.04] # estimated from the Lyapunov exponents plot
bbox2 = Dict(:boxstyle => "round,pad=0.3", :facecolor=>"white", :alpha => 0.8)

fig, axs = subplots(1,2; figsize = (0.66figx, figy))

function phase(x, y)
    f = atan.(y, x)
    for i in 2:length(f)
        if f[i] - f[i-1] < 0
            f[i:end] .+= 2π
        end
    end
    return f
end

for (i, k) in enumerate(ks)
    set_parameter!(croesslers, 6, k)
    set_parameter!(croesslers, 7, k)
    tr = trajectory(croesslers, T; Ttr = 200, Δt, diffeq...)
    x1 = tr[:, 1]
    x2 = tr[:, 2]
    y1 = tr[:, 4]
    y2 = tr[:, 5]
    f1 = phase(x1, x2)
    f2 = phase(y1, y2)
    axs[2].plot(t, f1 .- f2, label = "\$k=$(k)\$")

    if i == 1
        c = COLORS[1]
        axs[1].plot(x1[1:L], x2[1:L]; color = "0.5", linewidth = 2)
        j = 250
        axs[1].quiver(zeros(1), zeros(1), [x1[j+2]], [x2[j+2]],
            color = c,
            scale = 1., scale_units = "xy", zorder = 9,
            width = 0.01
        )
        axs[1].plot([x1[j+1]], [x2[j+1]], marker = "o", zorder = 99, ms = 8,
        ls = "None", mec = c, mfc = :white,
            # scale = 1., scale_units = "xy", zorder = 9,
            # width = 0.02
        )
        θ = atan(x2[j], x1[j])
        r = norm([x1[j], x2[j]])
        a = matplotlib.patches.Wedge((0,0), 0.3r;
        theta1 = 0, theta2 = 180θ/π, color = c, alpha = 0.5)
        axs[1].add_artist(a)
        axs[1].text(0.3x1[j-10], 0.3x2[j-10], "\$\\alpha\$"; color = c)
    end
end

axs[1].set_xlabel("\$x_1\$"; labelpad = -20)
axs[1].set_ylabel("\$x_2\$"; labelpad = -20)
axs[1].set_xticks([-10, 0, 10])
axs[1].set_xticklabels([-10, "", 10])
axs[2].set_xlabel("\$t\$"; labelpad = -20)
axs[2].set_ylabel("\$\\Delta \\alpha\$")
axs[2].set_xticks([0,t[end]])
axs[2].set_xlim(0,t[end])
axs[2].legend()
add_identifiers!(fig)
fig.tight_layout(;pad=0.33)
wsave(plotsdir("9", "coupled_roessler_phasediff"), fig)


include("utils.jl")
include("sized_abstract_array.jl")

# starting with tests of simple manifolds
include("euclidean.jl")
include("sphere.jl")
include("rotations.jl")
include("symmetric_positive_definite.jl")
include("cholesky_space.jl")
include("product_manifold.jl")
include("power_manifold.jl")
include("vector_bundle.jl")

include("metric.jl")


using AdvancedHMC
using Distributions, Random
using ForwardDiff, StatsPlots

include(string(pwd(),"/src/simulation.jl"))
include(string(pwd(),"/gali_bayesian.jl"))
include(string(pwd(),"/mcmc/priors_v2.jl"))
include(string(pwd(),"/mcmc/renormalization_manual.jl"))

yy,shocks = simulate_dsge(GAMMA_0,GAMMA_1,PSI,PI,500)

ell(par) = dens([1/3;par],yy[:,2])
ell_grad(par) = dens_and_grad([1/3;par],yy[:,2])

ff(par) = log_like_dsge(par,yy[:,2])[1]

from_unit(x) = -log(1/x-1)
from_pos(x) = log(x)
from_one_to_inf(x) = log(x-1)

par = [1/3;0.99;6;2/3;1;1;1;1.5;0.5/4;0.5]

ll1,dd1 = log_like_dsge(par,yy[:,2])#,reduce=false)

aa = [from_unit(0.99),from_pos(6),from_unit(2/3),from_pos(1),from_pos(1),from_pos(1),from_one_to_inf(1.5),from_pos(0.5/4),from_unit(0.5)]

ll,dd = ell_grad(aa)

#m = inv(diagm(diag(dd*dd'))/500)

n_samples, n_adapts = 2_000, 1_000

# Define a Hamiltonian system
metric = DiagEuclideanMetric(9)
hamiltonian = Hamiltonian(metric, ell,ell_grad)

initial = aa

# Define a leapfrog solver, with initial step size chosen heuristically
initial_ϵ = find_good_stepsize(hamiltonian,initial)
integrator = JitteredLeapfrog(initial_ϵ,1.0)

# Define an HMC sampler, with the following components
#   - multinomial sampling scheme,
#   - generalised No-U-Turn criteria, and
#   - windowed adaption for step-size and diagonal mass matrix
proposal = NUTS{MultinomialTS,ClassicNoUTurn}(integrator) #StaticTrajectory(integrator,100)#
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.65, integrator))#StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))

# Run the sampler to draw samples from the specified Gaussian, where
#   - `samples` will store the samples
#   - `stats` will store diagnostic statistics for each sample
samples, stats = sample(hamiltonian, proposal, initial, n_samples, adaptor, n_adapts; progress=true)

tab = zeros(2000,9)

for j = 1:size(samples,1)
    tab[j,:] = samples[j]
end

true_pars  = [0.99,6,2/3,1,1,1,1.5,0.5/4,0.5]

StatsPlots.density(to_unit.(tab[1:2000,1]), legend = :topleft, label = "Distribution")
vline!([true_pars[1]],label = "True Value")
StatsPlots.density(to_positive.(tab[1:2000,2]), label = "Distribution")
vline!([true_pars[2]],label = "True Value")
StatsPlots.density(to_unit.(tab[1:2000,3]), label = "Distribution")
vline!([true_pars[3]],label = "True Value")
StatsPlots.density(to_positive.(tab[1:2000,4]), label = "Distribution")
vline!([true_pars[4]],label = "True Value")
StatsPlots.density(to_positive.(tab[1:2000,5]), label = "Distribution")
vline!([true_pars[5]],label = "True Value")
StatsPlots.density(to_positive.(tab[1:2000,6]), label = "Distribution")
vline!([true_pars[6]],label = "True Value")
StatsPlots.density(to_one_inf.(tab[1:2000,7]), label = "Distribution")
vline!([true_pars[7]],label = "True Value")
StatsPlots.density(to_positive.(tab[1:2000,8]), label = "Distribution")
vline!([true_pars[8]],label = "True Value")
StatsPlots.density(to_unit.(tab[1:2000,9]), label = "Distribution")
vline!([true_pars[9]],label = "True Value")


module phasespace

export qfunc, wigner

using ..bases, ..states, ..operators, ..operators_dense, ..fock

"""
    qfunc(a, α)
    qfunc(a, x, y)
    qfunc(a, xvec, yvec)

Husimi Q representation ``⟨α|ρ|α⟩/π`` for the given state or operator `a`. The
function can either be evaluated on one point α or on a grid specified by
the vectors `xvec` and `yvec`. Note that conversion from `x` and `y` to `α` is
done via the relation ``α = \\frac{1}{\\sqrt{2}}(x + i y)``.
"""
function qfunc(rho::Operator, alpha::Number)
    b = basis(rho)
    @assert isa(b, FockBasis)
    _qfunc_operator(rho, convert(Complex128, alpha), Ket(b), Ket(b))
end

function qfunc(rho::Operator, xvec::Vector{Float64}, yvec::Vector{Float64})
    b = basis(rho)
    @assert isa(b, FockBasis)
    Nx = length(xvec)
    Ny = length(yvec)
    tmp1 = Ket(b)
    tmp2 = Ket(b)
    result = Matrix{Complex128}(Nx, Ny)
    @inbounds for j=1:Ny, i=1:Nx
        result[i, j] = _qfunc_operator(rho, complex(xvec[i], yvec[j])/sqrt(2), tmp1, tmp2)
    end
    result
end

function qfunc(psi::Ket, alpha::Number)
    b = basis(psi)
    @assert isa(b, FockBasis)
    _alpha = convert(Complex128, alpha)
    _conj_alpha = conj(_alpha)
    N = length(psi.basis)
    s = psi.data[N]/sqrt(N-1)
    @inbounds for i=1:N-2
        s = (psi.data[N-i] + s*_conj_alpha)/sqrt(N-i-1)
    end
    s = psi.data[1] + s*_conj_alpha
    return abs2(s)*exp(-abs2(_alpha))/pi
end

function qfunc(psi::Ket, xvec::Vector{Float64}, yvec::Vector{Float64})
    b = basis(psi)
    @assert isa(b, FockBasis)
    Nx = length(xvec)
    Ny = length(yvec)
    points = Nx*Ny
    N = length(b)::Int
    _conj_alpha = [complex(x, -y)/sqrt(2) for x=xvec, y=yvec]
    q = fill(psi.data[N]/sqrt(N-1), size(_conj_alpha))
    @inbounds for n=1:N-2
        f0_ = 1/sqrt(N-n-1)
        x = psi.data[N-n]
        for i=1:points
            q[i] = (x + q[i]*_conj_alpha[i])*f0_
        end
    end
    result = similar(q, Float64)
    x = psi.data[1]
    @inbounds for i=1:points
        result[i] = abs2(x + q[i]*_conj_alpha[i])*exp(-abs2(_conj_alpha[i]))/pi
    end
    result
end

function qfunc(state::Union{Ket, Operator}, x::Number, y::Number)
    qfunc(state, Complex128(x, y)/sqrt(2))
end

function _qfunc_operator(rho::Operator, alpha::Complex128, tmp1::Ket, tmp2::Ket)
    coherentstate(basis(rho), alpha, tmp1)
    operators.gemv!(complex(1.), rho, tmp1, complex(0.), tmp2)
    a = dot(tmp1.data, tmp2.data)
    return a/pi
end

"""
    wigner(a, α)
    wigner(a, x, y)
    wigner(a, xvec, yvec)

Wigner function for the given state or operator `a`. The
function can either be evaluated on one point α or on a grid specified by
the vectors `xvec` and `yvec`. Note that conversion from `x` and `y` to `α` is
done via the relation ``α = \\frac{1}{\\sqrt{2}}(x + i y)``.
"""
function wigner(rho::DenseOperator, x::Number, y::Number)
    b = basis(rho)
    @assert isa(b, FockBasis)
    N = b.N::Int
    _2α = complex(convert(Float64, x), convert(Float64, y))*sqrt(2)
    abs2_2α = abs2(_2α)
    w = complex(0.)
    coefficient = complex(0.)
    @inbounds for L=N:-1:1
        coefficient = 2*_clenshaw(L, abs2_2α, rho.data)
        w = coefficient + w*_2α/sqrt(L+1)
    end
    coefficient = _clenshaw(0, abs2_2α, rho.data)
    w = coefficient + w*_2α
    exp(-abs2_2α/2)/pi*real(w)
end

function wigner(rho::DenseOperator, xvec::Vector{Float64}, yvec::Vector{Float64})
    b = basis(rho)
    @assert isa(b, FockBasis)
    N = b.N::Int
    _2α = [complex(x, y)*sqrt(2) for x=xvec, y=yvec]
    abs2_2α = abs2.(_2α)
    w = zeros(_2α)
    b0 = similar(_2α)
    b1 = similar(_2α)
    b2 = similar(_2α)
    @inbounds for L=N:-1:1
        _clenshaw_grid(L, rho.data, abs2_2α, _2α, w, b0, b1, b2, 2)
    end
    _clenshaw_grid(0, rho.data, abs2_2α, _2α, w, b0, b1, b2, 1)
    @inbounds for i=eachindex(w)
        abs2_2α[i] = exp(-abs2_2α[i]/2)/pi.*real(w[i])
    end
    abs2_2α
end

wigner(psi::Ket, x, y) = wigner(dm(psi), x, y)
wigner(state, alpha::Number) = wigner(state, real(alpha)*sqrt(2), imag(alpha)*sqrt(2))

function _clenshaw_grid(L::Int, ρ::Matrix{Complex128},
                abs2_2α::Matrix{Float64}, _2α::Matrix{Complex128}, w::Matrix{Complex128},
                b0::Matrix{Complex128}, b1::Matrix{Complex128}, b2::Matrix{Complex128}, scale::Int)
    n = size(ρ, 1)-L-1
    points = length(w)
    if n==0
        f = scale*ρ[1, L+1]
        @inbounds for i=1:points
            w[i] = f + w[i]*_2α[i]/sqrt(L+1)
        end
    elseif n==1
        f1 = 1/sqrt(L+1)
        @inbounds for i=1:points
            w[i] = scale*(ρ[1, L+1] - ρ[2, L+2]*(L+1-abs2_2α[i])*f1) + w[i]*_2α[i]*f1
        end
    else
        f0 = sqrt(float((n+L-1)*(n-1)))
        f1 = sqrt(float((n+L)*n))
        f0_ = 1/f0
        f1_ = 1/f1
        fill!(b1, ρ[n+1, L+n+1])
        @inbounds for i=1:points
            b0[i] = ρ[n, L+n] - (2*n-1+L-abs2_2α[i])*f1_*b1[i]
        end
        @inbounds for k=n-2:-1:1
            b1, b2, b0 = b0, b1, b2
            x = ρ[k+1, L+k+1]
            a1 = -(2*k+1+L)
            a2 = -f0*f1_
            @inbounds for i=1:points
                b0[i] = x + (a1+abs2_2α[i])*f0_*b1[i] + a2*b2[i]
            end
            f1 , f1_ = f0, f0_
            f0 = sqrt((k+L)*k)
            f0_ = 1/f0
        end
        @inbounds for i=1:points
            w[i] = scale*(ρ[1, L+1] - (L+1-abs2_2α[i])*f0_*b0[i] - f0*f1_*b1[i]) + w[i]*_2α[i]*f0_
        end
    end
end

function _clenshaw(L::Int, abs2_2α::Float64, ρ::Matrix{Complex128})
    n = size(ρ, 1)-L-1
    if n==0
        return ρ[1, L+1]
    elseif n==1
        ϕ1 = -(L+1-abs2_2α)/sqrt(L+1)
        return ρ[1, L+1] + ρ[2, L+2]*ϕ1
    else
        f0 = sqrt(float((n+L-1)*(n-1)))
        f1 = sqrt(float((n+L)*n))
        f0_ = 1/f0
        f1_ = 1/f1
        b2 = complex(0.)
        b1 = ρ[n+1, L+n+1]
        b0 = ρ[n, L+n] - (2*n-1+L-abs2_2α)*f1_*b1
        @inbounds for k=n-2:-1:1
            b1, b2 = b0, b1
            b0 = ρ[k+1, L+k+1] - (2*k+1+L-abs2_2α)*f0_*b1 - f0*f1_*b2
            f1, f1_ = f0, f0_
            f0 = sqrt((k+L)*k)
            f0_ = 1/f0
        end
        return ρ[1, L+1] - (L+1-abs2_2α)*f0_*b0 - f0*f1_*b1
    end
end

end #module


function G_to_R( Ns::Array{Int,1}, fG )
    out = reshape( ifft( reshape(fG,Ns[1],Ns[2],Ns[3]) ), size(fG) )
end

function c_G_to_R(Ns::Array{Int,1}, fG)
    out = zeros(Complex128,Ns[1]*Ns[2]*Ns[3])
    ccall( (:fftw_inv_fft3d, "../libs/fft3d.so"), Void,
           (Ptr{Complex128}, Ptr{Complex128}, Int64, Int64, Int64),
            fG, out, Ns[3],Ns[2],Ns[1] )
    return out
end

function R_to_G( Ns::Array{Int,1}, fR )
    #
    out = reshape( fft( reshape(fR,Ns[1],Ns[2],Ns[3]) ),
          size(fR) )
    #
end

function c_R_to_G( Ns::Array{Int,1}, fR::Array{Complex128,1} )
    #
    out = zeros(Complex128,Ns[1]*Ns[2]*Ns[3])
    ccall( (:fftw_fw_fft3d, "../libs/fft3d.so"), Void,
           (Ptr{Complex128}, Ptr{Complex128}, Int64, Int64, Int64),
            fR, out, Ns[3],Ns[2],Ns[1] )
    return out
    #
end

#
# Input array is of type Float64. This input is converted
# to Complex128 before supplied to `fftw_fw_fft3d`
#
function c_R_to_G( Ns::Array{Int,1}, fR::Array{Float64,1} )
    out = zeros(Complex128,Ns[1]*Ns[2]*Ns[3])
    ccall( (:fftw_fw_fft3d, "../libs/fft3d.so"), Void,
           (Ptr{Complex128}, Ptr{Complex128}, Int64, Int64, Int64),
            convert(Array{Complex128,1},fR), out, Ns[3],Ns[2],Ns[1] )
    return out
end


using RobotZoo
using RobotDynamics
using StaticArrays
using BenchmarkTools
using FiniteDiff
using LinearAlgebra
using Test
using Random
Random.seed!(1)
const RD = RobotDynamics

## Double Integrator
D = 2
linmodel = RobotZoo.LinearModels.DoubleIntegrator(D)
model = RobotZoo.DoubleIntegrator(D)
x,u = rand(model)
@test RD.dynamics(model, x, u) ≈ RD.dynamics(linmodel, x, u)

h = 0.1
dmodel_exp = RobotZoo.DiscreteLinearModel(linmodel, h)
dmodel_lin = RobotZoo.LinearModels.DiscreteDoubleIntegrator(h, D)
@test RD.discrete_dynamics(dmodel_exp, x, u, 0.0, h) ≈
    RD.discrete_dynamics(dmodel_lin, x, u, 0.0, h)

## FlexibleSatellite
model = RobotZoo.LinearModels.FlexibleSatellite()
@test RD.dims(model) == (12,3,12)


export neighborhood,parametric,cartesian,jacobian,normal,shapefuns,shapeders,get_shape_curl
mutable struct subd_point{T}
    cart::Vector{T}
    paracoords::Vector{T}
    J::Array{T,2}
    detJ::T
    normal::SVector{3,T}
    shapes::Matrix{T}
    shapeders::Matrix{T}
end
Base.getindex(p::subd_point{T}, i::Int) where {T} = p.cart[i]
function neighborhood(chart::subd_chart{T}, u) where {T}
    vertices = chart.vertices
    Nv = chart.N
    val = Nv - 6
    shapefunc = shape_function(u[1],u[2],val)
    shape_der = shape_function_der(u[1],u[2],val)
    coords = zeros(3)
    t1 = zeros(3)
    t2 = zeros(3)
    for id = 1:3
        for iv = 1:Nv
            coords[id] += vertices[iv][id]*shapefunc[iv]
            t1[id] += vertices[iv][id]*shape_der[iv,1]
            t2[id] += vertices[iv][id]*shape_der[iv,2]
        end
    end
    J = [t1 t2]'
    norm = [t1[2] * t2[3] - t1[3] * t2[2];
            t1[3] * t2[1] - t1[1] * t2[3];
            t1[1] * t2[2] - t1[2] * t2[1]]
    detJ = sqrt(norm[1]*norm[1] + norm[2]*norm[2] + norm[3]*norm[3])
    norm = norm./detJ
    return subd_point{T}(coords,u,J,detJ,norm,shapefunc,shape_der)
end

function parametric(d::subd_point)
    return d.paracoords
end

function cartesian(d::subd_point)
    v = d.cart
    return v
end

function jacobian(d::subd_point)
    return d.detJ
end

function normal(d::subd_point)
    return d.normal
end

function shapefuns(d::subd_point)
    return d.shapes
end

function shapeder(d::subd_point)
    return d.shapeders
end

function JacobMatrix33(d::subd_point)
    Jacob = d.J
    norm = d.normal
    return [Jacob;norm']
end

function Jacob_inv(d::subd_point)
    J = JacobMatrix33(d)
    invJ = inv(J)
    return invJ[:,1:2]
end

function get_shape_curl(d::subd_point)
    norm = d.normal
    sders = d.shapeders
    invJ = Jacob_inv(d)
    grad = invJ*sders'
    ncol = size(grad,2)
    curl = zeros(size(grad))
    for i = 1 : ncol
        curl[:,i] = cross(norm,grad[:,i])
    end
    return curl'
end


f(x) = begin
        x2 = x ^ 2
        x * Main.muladd(x2, Main.muladd(x2, 0.008333333333333333, -0.16666666666666666), 1.0)
    end


#== ChhabraJensen.jl  by Tehran Davis, 2020

This is an Julia-based translation of the Matlab script written by
Authors: Lucas França, Yujiang Wang, José Miranda
https://lucasfr.github.io/, http://xaphire.de/

# --------------------------------------------------------

#INPUT:

This function assumes that your time series is all positive values (sign transform it if needed)

#OUTPUT:
#alpha and falpha scattered against each other give the multifractal spectrum
#qValues and Dq scattered against each other give the generalised dimensions spectrum
#Rsqr_alpha, Rsqr_falpha, and Rsqr_Dq are the R^2 values for each of the values in alpha, falpha, and Dq respectively. It is recommended to run the following code outside of the function:

==#

using DataFrames, CSV, GLM

#timeseries = CSV.read("multifractal_toolbox/exampleTS.txt",header=false)[:,1]
#qValues = [-5:1:5;]
#scales = [3:1:11;]

function ChhabraJensen(timeseries, qValues, scales)

    timeseries_trim = length(timeseries) % 2^scales[length(scales)]
    timeseries = timeseries[1:(length(timeseries)-timeseries_trim)]

    # start function here

    nq=length(qValues)
    ns=length(scales)
    Ma=zeros(nq,ns)
    Mf=zeros(nq,ns)
    Md=zeros(nq,ns)

    muScale= -log10.(2 .^scales)

    alpha=zeros(nq)
    falpha=zeros(nq)
    Dq=zeros(nq)

    Rsqr_alpha=zeros(nq)
    Rsqr_falpha=zeros(nq)
    Rsqr_Dq=zeros(nq)

    for i=1:nq

        q=qValues[i]

        for j=1:ns

            #determine how many windows we will have at this scale
            window=2^scales[j]

            #break the time series into windows & sum
            timeseriesReshaped=reshape(timeseries,:,window)
            timeseriesSummed=sum(timeseries)

            #calculate p
            ps=sum(timeseriesReshaped, dims = 1)
            p=ps./timeseriesSummed

            Nor=sum(p.^q)

            #calculation of Md
            #not accounting for q between 0 and 1
            if q<=1 && q>0
                Md[i,j]=sum(p.*log10.(p))/Nor
            else
                Md[i,j]=log10(Nor)
            end

            #Ma & Mf
            mu=(p.^q)/Nor
            Ma[i,j]=sum(mu.*log10.(p))
            Mf[i,j]=sum(mu.*log10.(mu))
        end

        Ma_mdl = lm(@formula(Ma ~ muScale), DataFrame(Ma = Ma[i,:], muScale = muScale))
        Mf_mdl = lm(@formula(Mf ~ muScale), DataFrame(Mf = Mf[i,:], muScale = muScale))
        Md_mdl = lm(@formula(Md ~ muScale), DataFrame(Md = Md[i,:], muScale = muScale))

        b_Ma = coef(Ma_mdl)[2]
        b_Mf = coef(Mf_mdl)[2]
        b_Md = coef(Md_mdl)[2]

        alpha[i]=b_Ma
        falpha[i]=b_Mf

        if (q<=1 && q>0)
            Dq[i,]=b_Md
        else
            Dq[i]=b_Md/(q-1)
        end

        Rsqr_alpha[i] = r2(Ma_mdl)
        Rsqr_falpha[i] = r2(Mf_mdl)
        Rsqr_Dq[i] = r2(Md_mdl)
    end #looping through scales

    output = Dict("alpha" => alpha,
                    "Rsqr_alpha" => Rsqr_alpha,
                    "falpha" => falpha,
                    "Rsqr_falpha" => Rsqr_falpha,
                    "Dq" => Dq,
                    "Rsqr_Dq" => Rsqr_Dq,
                    "q" => qValues,
                    "mfw" => maximum(alpha)-minimum(alpha)
                    )

    return output

end


"""
## glob_to_loc!

This subroutine transforms the global end reactions and
moments into the local system (2- or 3-d). Called from hinge!().

### Function
```julia
glob_to_loc!(loc, glob, gamma, coord)
```

### Arguments
```julia
* loc::Vector{Float64}       : Local force and momemts (Updated)
* glob::Vector{Float64}      : Globale forces and moments
* gamma::Float64             : Element orientation angle (3D)
* coord::Matrix{Float64}     : Nodal coordinates
```
"""
function glob_to_loc!(loc::Vector{Float64}, glob::Vector{Float64},
  gamma::Float64, coord::LinearAlgebra.Adjoint{Float64,Array{Float64,2}})
  #
  # This subroutine transforms the global end reactions and
  # moments into the local system (2- or 3-d).
  #
  # Updates loc[:]
  #

  ndim = size(coord, 2)
  if ndim == 2
    x1 = coord[1, 1]
    y1 = coord[1, 2]
    x2 = coord[2, 1]
    y2 = coord[2, 2]
    ell = sqrt((x2-x1)^2+(y2-y1)^2)
    cg = (x2-x1)/ell
    sg = (y2-y1)/ell
    loc[1] = cg*glob[1]+sg*glob[2]
    loc[2] = cg*glob[2]-sg*glob[1]
    loc[3] = glob[3]
    loc[4] = cg*glob[4]+sg*glob[5]
    loc[5] = cg*glob[5]-sg*glob[4]
    loc[6] = glob[6]
  elseif ndim == 3
    t = zeros(12,12)
    r0 = zeros(3, 3)
    x1 = coord[1, 1]
    y1 = coord[1, 2]
    z1 = coord[1, 3]
    x2 = coord[2, 1]
    y2 = coord[2, 2]
    z2 = coord[2, 3]
    xl = x2 - x1
    yl = y2 - y1
    zl = z2 - z1
    ell = sqrt(xl*xl+yl*yl+zl*zl)
    gamrad = gamma*pi/180.0
    cg = cos(gamrad)
    sg = sin(gamrad)
    den = ell*sqrt(xl*xl+zl*zl)
    if den !== 0.0
      r0[1,1] = xl/ell
      r0[1,2] = yl/ell
      r0[1,3] = zl/ell
      r0[2,1] = (-xl*yl*cg-ell*zl*sg)/den
      r0[2,2] = den*cg/(ell*ell)
      r0[2,3] = (-yl*zl*cg+ell*xl*sg)/den
      r0[3,1] = (xl*yl*sg-ell*zl*cg)/den
      r0[3,2] = -den*sg/(ell*ell)
      r0[3,3] = (yl*zl*sg+ell*xl*cg)/den
    else
      r0[1,1] = 0.0
      r0[1,3] = 0.0
      r0[2,2] = 0.0
      r0[3,2] = 0.0
      r0[1,2] = 1.0
      r0[2,1] = -cg
      r0[3,3] = cg
      r0[2,3] = sg
      r0[3,1] = sg
    end
    for i in 1:3
      for j in 1:3
        x = r0[i,j]
        for  k in 0:3:9
          t[i+k,j+k] = x
        end
      end
    end
    for i in 1:12
      sum = 0.0
      for j in 1:12
        sum = sum+t[i,j]*glob[j]
      end
      loc[i] = sum
    end
  end
end


module OptimizeQCQP

using LinearAlgebra
using SparseArrays
using Arpack

export minimize, Problem

include("linalg.jl")

function minimize(Q0::AbstractMatrix{T}, g0::AbstractVector{T},
                  ::Nothing, ::Nothing, Q1::Union{AbstractMatrix{T},UniformScaling},
                  ::Nothing, ::Nothing, cu::Real; tol=T(1e-4)) where T
    # Solve the problem:
    #    minimize x'*Q0*x/2 + g0'*x
    #    subject to x'*Q1*x/2 <= cu

    # Pick a safe value of λ that makes Q0 + λ*Q1 positive-definite.
    # Keep track of the smallest known such value in λp
    λ, C = posdefλ(Q0, Q1)
    λp = λ

    ## Construct an initial estimate
    x = C \ (-g0)
    xQx = x'*(Q1*x)
    if xQx > 2*cu
        # Since λ only guarantees positivity and does not even try to satisfy the cu
        # constraint, it's possible that this λ is far too small. Try to estimate the right
        # value in terms of a correction Δλ where
        #          x = (C + Δλ*Q1) \ (-g)
        # and then picking Δλ so that x satisfies the constraint.
        # We can expand that inverse in two ways:
        # (1) If C dominates Δλ*Q1, use the expansion inv(C + Δλ*Q1) ≈ inv(C) - Δλ*inv(C)*Q1*inv(C)
        Δx = - (C \ (Q1 * x))  # coefficient of Δλ in correction to x
        # Solve for Δλ from (x + Δλ*Δx)'*Q1*(x + Δλ*Δx) == 2*cu
        a = Δx'*(Q1*Δx)       # coefficient of Δλ^2
        b = 2*x'*(Q1*Δx)      # coefficient of Δλ (guaranteed negative)
        c = xQx - 2*cu        # const (coefficient of Δλ^0)
        s2 = b^2 - 4*a*c
        if s2 >= 0
            ΔλC = 2*c/(-b + sqrt(s2))   # choose the smaller of the two roots
        else
            # it doesn't completely reach the criterion, so choose the minimum
            # (this is likely a sign that solution (2) will be selected)
            ΔλC = -b/(2*a)
        end
        @assert ΔλC >= 0
        xC = x + ΔλC*Δx   # this is our estimate of the solution, given this Δλ
        # (2) If Δλ*Q1 dominates C, then just use x ≈ (1/Δλ) * (Q1 \ (-g))  (dropping C altogether)
        v1 = Q1 \ (-g0)
        Δλ1 = sqrt(-(g0'*v1)/(2*cu))   # initial alternative estimate of Δλ, discounting higher-order Δλ
        x1 = v1 ./ Δλ1         # the alternate estimate of x
        # Pick between them by assessing accuracy in matching
        #     (C + Δλ*Q1) * x = -g0
        errC = Q0*xC .+ (λ + ΔλC).*(Q1*xC) .+ g0
        err1 = Q0*x1 .+ (λ + Δλ1).*(Q1*x1) .+ g0
        absprod(x, y) = abs(x)*abs(y)
        errmagC = sum(absprod(x,y) for (x,y) in zip(errC, xC))
        errmag1 = sum(absprod(x,y) for (x,y) in zip(err1, x1))
        # λ = (s2 < 0 || errmagC > errmag1) ? λ + Δλ1 : λ + ΔλC
        # Use interpolation, weighted by the error
        λ += (errmag1*ΔλC + errmagC*Δλ1)/(errmag1 + errmagC)
        C = cholesky(Q0 + λ*Q1)
        x = C \ (-g0)
        xQx = x'*(Q1*x)
    end

    ## Newton's method iterative improvement
    # Algorithm 4.3 in Nodecal & Wright, generalized for Q1 != I
    # Also written to make it globally convergent
    iter = 0
    while iter < 10
        λp = min(λ, λp)
        xQxold = xQx
        q = Q1*x
        dxdλ = - (C \ q)
        Δ = cu - xQx/2
        Δλ = Δ/cu * (xQx / (2*(q'*dxdλ)))
        λold = λ
        # Global convergence iteration
        itergc = 0
        while itergc < 10
            λ = max(λold + Δλ, zero(T))
            # posdef iteration
            C = cholesky(Q0 + λ*Q1; check=false)
            iterc = 0
            while !issuccess(C) && iterc < 10
                iterc += 1
                λ = (λ + iterc*λp)/(iterc+1)   # move towards a known-safe value
                C = cholesky(Q0 + λ*Q1; check=false)
            end
            if !issuccess(C)
                λ = (λold + λp)/2
                C = cholesky(Q0 + λ*Q1)
            end
            x = C \ (-g0)
            xQx = x'*(Q1*x)
            (xQx <= 2*cu || abs(xQx - 2*cu) < abs(xQxold - 2*cu)) && break
            itergc += 1
            Δλ /= itergc+1
        end
        abs(xQx - xQxold) < tol*(xQx + xQxold) && break
        iter += 1
    end
    return x, λ
end

# Junk from attempt to incorporate higher-order correction
        # (2) If Δλ*Q1 dominates C, use the expansion
        #        inv(C + Δλ*Q1) ≈ inv(Δλ*Q1) - inv(Δλ*Q1)*C*inv(Δλ*Q1)
        #                       = (1/Δλ - λ/Δλ^2)*inv(Q1) - (1/Δλ^2)*inv(Q1)*Q0*inv(Q1)
        # v2 = Q1 \ (Q0 * v1)
        # v1Qv1, v2Qv2, v1Qv2 = v1'*Q1*v1, v2'*Q1*v2, v1'*Q1*v2
        # let λ=λ, v1Qv1=v1Qv1, v2Qv2=v2Qv2, v1Qv2=v1Qv2, cu=cu
        #     froot(Δλ) = (1/Δλ^2) * ((1-λ/Δλ)^2*v1Qv1 - 2*(1-λ/Δλ)/Δλ*v1Qv2 + 1/Δλ^2*v2Qv2) - 2*cu
        #     Δλ1 = bisectroot(froot, zero(Δλ1), Δλ1)
        # end
        # x1 = (1/Δλ1-λ/Δλ1^2).*v1 - v2./Δλ1^2         # the alternate estimate of x

# function bisectroot(f, a, b; tol=1e-3, maxiter=10, maxseek=20)
#     fa, fb = f(a), f(b)
#     f0 = min(abs(fa), abs(fb))
#     fa == 0 && return a
#     fb == 0 && return b
#     iter = 0
#     while fa*fb > 0 && iter < maxseek
#         a, fa = b, fb
#         b *= 2
#         fb = f(b)
#         fb == 0 && return b
#         iter += 1
#     end
#     iter = 0
#     while iter < maxiter
#         c = (a+b)/2
#         fc = f(c)
#         (fc == 0 || (abs(fc - fa) + abs(fc - fb)) < tol*f0) && return c
#         if sign(fc) == sign(fa)
#             a, fa = c, fc
#         else
#             b, fb = c, fc
#         end
#         iter += 1
#     end
#     return b
# end

# Solve the problem:
#    (approximately) minimize x'*Q0*x/2 + g0'*x
#    subject to cl[i] <= x'*Qi*x/2 + gi'*x <= cu[i]       for i = 1,...,m
#                l[k] <= x[k] <= u[k]                     for k = 1,...,n
#
# Note that equality constraints are handled by setting `cl[i] = cu[i]` and/or
# `l[k] = u[k]`.
#
# For positive-semidefinite Q0 and Qi, the problem is convex, otherwise it is NP-hard
# (hence the "approximate" minimization).

struct Problem{T<:Real,Q0T<:AbstractMatrix{T},g0T<:AbstractVector{T}}
    n::Int
    l::Vector{T}
    u::Vector{T}
    cl::Vector{T}
    cu::Vector{T}
    Q0::Q0T
    g0::g0T
    Qidense::Vector{Matrix{T}}
    Qisparse::Vector{SparseMatrixCSC{T,Int}}
    gidense::Vector{Vector{T}}
    gisparse::Vector{SparseVector{T,Int}}
    issparseQi::Vector{Bool}
    issparsegi::Vector{Bool}
    isposdefQi::Vector{Bool}
end

function Problem(Q0::AbstractMatrix{T}, g0::AbstractVector{T},
                 lower=fill(-T(Inf), length(g0)), upper=fill(T(Inf), length(g0))) where T
    n = length(g0)
    indices(g0) == (Base.OneTo(n),) || throw(DimensionMismatch("g0 indices must start with 1, got $(indices(g0))"))
    indices(Q0) == (Base.OneTo(n), Base.OneTo(n))  || throw(DimensionMismatch("Q0 indices must match g0, got $(indices(g0)) and $(indices(Q0))"))
    indices(lower) == (Base.OneTo(n),) || throw(DimensionMismatch("lower's indices must match g0, got $(indices(lower)) and $(indices(g0))"))
    indices(upper) == (Base.OneTo(n),) || throw(DimensionMismatch("upper's indices must match g0, got $(indices(upper)) and $(indices(g0))"))
    Qidense = Matrix{T}[]
    Qisparse = SparseMatrixCSC{T,Int}[]
    gidense = Vector{T}[]
    gisparse = SparseVector{T,Int}[]
    cl, cu = T[], T[]
    issparseQi = Bool[]
    issparsegi = Bool[]
    isposdefQi = Bool[]
    return Problem{T,typeof(Q0),typeof(g0)}(n, lower, upper, cl, cu,
        Q0, g0, Qidense, Qisparse, gidense, gisparse,
        issparseQi, issparsegi, isposdefQi)
end

iszerom(Q::Matrix)          = false
iszerom(Q::SparseMatrixCSC) = isempty(Q.nzval)

zerom(::Type{T}, n) where T = sparse(Int[], Int[], T[], n, n)
zerov(::Type{T}, n) where T = sparse(Int[], T[], n)

function addconstraint!(P::Problem{T}, Qi::AbstractMatrix, gi::AbstractVector, cl::Real, cu::Real;
                        isposdef::Bool=false) where T
    n = P.n
    indices(gi) == (Base.OneTo(n),) || throw(DimensionMismatch("gi indices must match 1:$n, got $(indices(gi))"))
    indices(Qi) == (Base.OneTo(n), Base.OneTo(n))  || throw(DimensionMismatch("Qi indices must match (1:$n, 1:$n), got $(indices(Qi))"))
    Qsp = issparse(Qi)
    gsp = issparse(gi)
    Qsp ? push!(P.Qisparse, Qi) : push!(P.Qidense, Qi)
    gsp ? push!(P.gisparse, gi) : push!(P.gidense, gi)
    push!(P.cl, cl)
    push!(P.cu, cu)
    push!(P.issparseQi, Qsp)
    push!(P.issparsegi, gsp)
    push!(P.isposdefQi, isposdef)
    return P
end
function addconstraint!(P::Problem{T}, Qi::Real, gi::Real, cl::Real, cu::Real; kwargs...) where T
    @assert Qi == 0 && gi == 0
    n = P.n
    return addconstraint!(P, zerom(T, n), zerov(T, n), cl, cu; kwargs...)
end
function addconstraint!(P::Problem{T}, Qi::Real, gi::AbstractVector, cl::Real, cu::Real; kwargs...) where T
    @assert Qi == 0
    n = P.n
    return addconstraint!(P, zerom(T, n), gi, cl, cu; kwargs...)
end
function addconstraint!(P::Problem{T}, Qi::AbstractMatrix, gi::Real, cl::Real, cu::Real; kwargs...) where T
    @assert gi == 0
    n = P.n
    return addconstraint!(P, Qi, zerov(T, n), cl, cu; kwargs...)
end

function minimize(P::Problem; usefactorization::Bool=!issparse(P.Q0))
    if usefactorization
        return minimize_factorization(P)
    else
        error("not yet implemented")
    end
end

function minimize_factorization(P::Problem{T}, x::AbstractVector, μ) where T
    m, n = length(P.issparseQi), P.n
    indices(x) == (Base.OneTo(n),) || throw(DimensionMismatch("x must have indices 1:$n, got $(linearindices(x))"))
    λx, λc, sx, sc, constrval = initialize_auxvars(P, x)
    error("unfinished")
    # if isempty(λx) && m == 1
    #     # No bounds constraints
end

function initialize_auxvars(P::Problem{T}, x) where T
    m, n = length(P.issparseQi), P.n
    # How many inequality constraints are "real"? (x[k] <= Inf does not count as "real")
    nx = count(isfinite, P.l)  + count(isfinite, P.u)
    nc = count(isfinite, P.cl) + count(isfinite, P.cu)
    # Lagrange multipliers
    λx = zeros(T, nx)
    λc = zeros(T, nc)
    # Slack variables
    sx = zeros(T, nx)
    sc = zeros(T, nc)
    # Temporaries
    constrval = fill(T(NaN), m)
    # Initialize the slacks
    si = 0
    for k = 1:n
        if isfinite(P.l[k])
            sx[si+=1] = min(zero(T), x[k] - P.l[k])
        end
    end
    for k = 1:n
        if isfinite(P.u[k])
            sx[si+=1] = min(zero(T), P.u[k] - x[k])
        end
    end
    si = 0
    for k = 1:m
        if isfinite(P.cl[k])
            sc[si+=1] = min(zero(T), evalconstr!(constrval, P, k, x) - P.cl[k])
        end
    end
    for k = 1:m
        if isfinite(P.cu[k])
            sc[si+=1] = min(zero(T), P.cu[k] - evalconstr!(constrval, P, k, x))
        end
    end
    return λx, λc, sx, sc, constrval
end

end # module


# tests transformations from GeoData <=> Cartesian
using Test
using GeophysicalModelGenerator

# Create 3D volume with some fake data
Lon,Lat,Depth   =   LonLatDepthGrid(5:25,20:50,(-1300:100:0)km);
Data_set3D      =   GeoData(Lon,Lat,Depth,(Depthdata=Depth*2 + Lon*km,LonData=Lon))

proj            =   ProjectionPoint(Lon=20,Lat=35)

# Convert this 3D dataset to a Cartesian dataset (the grid will not be orthogonal)
Data_set3D_Cart =   Convert2CartData(Data_set3D, proj)
@test sum(abs.(Data_set3D_Cart.x.val)) ≈ 5.293469089428514e6km

# Create Cartesian grid
X,Y,Z           =   XYZGrid(-500:100:500,-900:200:900,(-500:100:0)km);
Data_Cart       =   CartData(X,Y,Z,(Z=Z,))

# Project values of Data_set3D to the cartesian data
Data_Cart       =   ProjectCartData(Data_Cart, Data_set3D, proj)
@test sum(Data_Cart.fields.Depthdata) ≈ -316834.28716859705km
@test sum(Data_Cart.fields.LonData) ≈ 13165.712831402916

# Next, 3D surface (like topography)
Lon,Lat,Depth   =   LonLatDepthGrid(5:25,20:50,0);
Depth           =   cos.(Lon/5).*sin.(Lat)*10
Data_surf       =   GeoData(Lon,Lat,Depth,(Z=Depth,))
Data_surf_Cart  =   Convert2CartData(Data_surf, proj)

# Cartesian surface
X,Y,Z           =   XYZGrid(-500:10:500,-900:20:900,0);
Data_Cart       =   CartData(X,Y,Z,(Z=Z,))

Data_Cart       =   ProjectCartData(Data_Cart, Data_surf, proj)
@test sum(Data_Cart.z.val) ≈ 1858.2487019158766km
@test sum(Data_Cart.fields.Z) ≈ 1858.2487019158766

# Cartesian surface when UTM data is used
WE,SN,depth     =   XYZGrid(420000:1000:430000, 4510000:1000:4520000, 0);

Data_surfUTM    =   UTMData(WE, SN, depth, 33, true, (Depth = WE,));
Data_Cart       =   CartData(X,Y,Z,(Z=Z,))
Data_Cart       =   ProjectCartData(Data_Cart, Data_surfUTM, proj)

@test sum(Data_Cart.z.val) ≈ 0.0km
@test sum(Data_Cart.fields.Depth) ≈ 3.9046959539921126e9


using JuMP, Ipopt

@testset "dat export import consistency check" begin
    # min x - y
    # st  x + x^2 + x*y + y^2 <= 1
    #     -2 <= x, y <= 2
    # solution: x+y = -1/3
    # optimal objective -1-4/sqrt(3)

    qcqp = Problem()

    x = MPC.Variable("x", Complex)
    y = MPC.Variable("y", Real)
    b = MPC.Variable("b", Bool)

    set_objective!(qcqp, x - y)

    add_constraint!(qcqp, "ctr1", (x + x^2 + x*y + y^2) << 1)
    add_constraint!(qcqp, "ctr2", (x^3 + (1.0+3im)*conj(x^2)*x*y + y^2) << (1 + 2im))
    add_constraint!(qcqp, "ctr3", -3im << (5*b + 5im*y) << 1)
    add_constraint!(qcqp, "x_bounds", -2 << x << 2)
    add_constraint!(qcqp, "y_bounds", -2 << y << 2)

    point = Point([x], [3])
    dat_exportpath = joinpath(Pkg.dir("MathProgComplex"), "tmp_datexport")
    mkpath(splitdir(dat_exportpath)[1])

    export_to_dat(qcqp, dat_exportpath, filename="POP.dat", point=point)

    qcqp2, point2 = import_from_dat(joinpath(dat_exportpath, "POP.dat"))

    rm(joinpath(dat_exportpath, "POP.dat"))

    @test qcqp.variables == qcqp2.variables
    @test qcqp.objective == qcqp2.objective

    @test Set(keys(qcqp.constraints)) == Set(keys(qcqp2.constraints))

    for (ctrname, ctr) in qcqp.constraints
        @test qcqp.constraints[ctrname].ub == qcqp2.constraints[ctrname].ub
        @test qcqp.constraints[ctrname].lb == qcqp2.constraints[ctrname].lb
        @test qcqp.constraints[ctrname].p == qcqp2.constraints[ctrname].p
    end

    @test point == point2

    rm(dat_exportpath, recursive=true)
end

@testset "WB5.dat import and Ipopt solve" begin
    instancepath = joinpath(Pkg.dir("MathProgComplex"), "test", "instances")
    WB5_cplx, initpt = import_from_dat(joinpath(instancepath, "WB5.dat"))

    WB5 = pb_cplx2real(WB5_cplx)

    mysolver = IpoptSolver(print_level = 0)
    m, variables_jump, ctr_jump, ctr_exp = get_JuMP_cartesian_model(WB5, mysolver)
    solve(m)

    sol = get_JuMP_solution(m, variables_jump, WB5)

    @test get_objective(WB5, sol) ≈ 1146.478 atol=1e-3
end


# # Von Mises Plasticity with Mixed Hardening
#
#
# ### Load packages
using Tensors, ConstLab
using ConstLab: MixedControl, Plastic, PlasticState

# ### Define the strain control
ncycles = 5
timesteps = range(0.0, 2pi*ncycles; length=200*ncycles)
ε_amp = 0.005 # Strain amplitute for ε_11
ctrl = MixedControl(
    ## Strain
    t -> SymmetricTensor{2,3}((i,j) -> i == j == 1 ? ε_amp*sin(t) : 0.0),
    ## Controlled strain components
    SymmetricTensor{2,3}((true, false, false, false, false, false)),
    ## Stress
    t -> zero(SymmetricTensor{2,3}),
    ## Controlled stress components
    SymmetricTensor{2,3}((false, true, true, true, true, true))
);

# ### Set material parameters
E = 200.0e9
ν =  0.3
G = E/(2*(1+ν))
K = E/(3*(1-2*ν))
σ_y = 200.0e6
H = 0.1*E
κ∞ = 0.2*σ_y
α∞ = 0.1*σ_y
r = 0.5

model = Plastic(G, K, σ_y, H, κ∞, α∞, r);

# ### Set initial conditions
s0 = PlasticState{3}();

# ### Do integration
res = ConstLab.integrate(model, ctrl, timesteps, s0; solver_params=Dict(:method => :newton));

# ### Plot results
using PGFPlotsX, LaTeXStrings

p = @pgf Axis({xlabel=L"\varepsilon_{11}", ylabel=L"\sigma_{11}"},
    PlotInc({mark="none"},
        Coordinates([r.ε[1,1] for r in res],
                    [r.σ[1,1] for r in res])
    ),
)


# checks Fresnel's transmission coefficient numerically
n_incident = 1.
n_transmit = 1.45
ref_local_field = get_local_field(Float64[0.5], refractive_indexes=ones(3) * n_incident)[3]
interface_localfield = get_local_field(Float64[0.5],  refractive_indexes=Float64[n_transmit, n_incident, n_incident])[3]

transmission_numerical = abs2(interface_localfield)/abs2(ref_local_field)/n_incident*n_transmit
transmission_fresnel = ((2n_incident)/(n_incident+n_transmit))^2/n_incident*n_transmit

@test abs(transmission_fresnel - transmission_numerical)/abs(transmission_fresnel)< 1e-2


module UniDerivative
"""
Compute derivative of a univariate function.
"""

abstract type DerivativeMethod end

struct CentralDiff <: DerivativeMethod
    h
    CentralDiff(h = cbrt(eps())) = new(h)
end

function df(param::CentralDiff, f, x::Real)::Real
    (f(x + param.h / 2) - f(x - param.h / 2)) / param.h
end

struct ComplexDiff <: DerivativeMethod
    h
    ComplexDiff(h = 1e-20) = new(h)
end

function df(param::ComplexDiff, f, x::Real)::Real
    imag(f(x + param.h * im)) / param.h
end

end # module


"""
    Sphere([origin=zeros(),] radius)

A [`Shape`](@ref) where boundary is a fixed distance from the origin. In 2D this is a circle, in 3D the usual sphere, and in higher dimensions if difficult to visualise.
"""
struct Sphere{T,Dim} <: Shape{T,Dim}
    origin::SVector{Dim,T}
    radius::T
end

# Alternate constructors, where type is inferred naturally
Sphere(origin::NTuple{Dim}, radius::T) where {T,Dim} = Sphere{T,Dim}(origin, radius)
Sphere(origin::AbstractVector, radius::T) where {T} = Sphere{T,length(origin)}(origin, radius)

Sphere(Dim, radius::T) where {T} = Sphere{T,Dim}(zeros(T,Dim), radius)

Circle(origin::Union{AbstractVector{T},NTuple{2,T}}, radius::T) where T <: AbstractFloat = Sphere{T,2}(origin, radius::T)
Circle(radius::T) where T <: AbstractFloat = Sphere(2, radius::T)
Sphere(radius::T) where {T} = Sphere{T,3}(zeros(T,3), radius)

name(shape::Sphere) = "Sphere"
name(shape::Sphere{T,2}) where T = "Circle"
Symmetry(shape::Sphere{T,Dim}) where {T,Dim} = RadialSymmetry{Dim}()

outer_radius(sphere::Sphere) = sphere.radius
volume(shape::Sphere{T,3}) where T = 4//3 * π * shape.radius^3
volume(shape::Sphere{T,2}) where T = π * shape.radius^2

# bounding_box(sphere::Sphere{T,3}; kws...) where T = bounding_box(Circle(sphere; kws...))
function bounding_box(sphere::Sphere{T,Dim}) where {T,Dim}
    return Box(origin(sphere), T(2)*sphere.radius .* ones(T,Dim))
end

import Base.in
function in(x::AbstractVector, sphere::Sphere)::Bool
    norm(origin(sphere) .- x) <= sphere.radius
end

import Base.issubset
function issubset(inner_sphere::Sphere{T,Dim}, outer_sphere::Sphere{T,Dim}) where {T,Dim}
    norm(origin(outer_sphere) - origin(inner_sphere)) <= outer_sphere.radius - inner_sphere.radius
end

function issubset(sphere::Sphere, box::Box)
    sphere_box = bounding_box(sphere)
    return issubset(sphere_box,box)
end

function issubset(box::Box, sphere::Sphere)
    issubset(Sphere(origin(box),outer_radius(box)), sphere)
end

import Base.(==)
function ==(c1::Sphere, c2::Sphere)
    c1.origin == c2.origin &&
    c1.radius == c2.radius
end

import Base.isequal
function isequal(c1::Sphere, c2::Sphere)
    isequal(c1.origin, c2.origin) &&
    isequal(c1.radius, c2.radius)
end

function iscongruent(s1::Sphere, s2::Sphere)
    s1.radius == s2.radius
end

function congruent(s::Sphere, x)
    Sphere(x, s.radius)
end

function Circle(sphere::Sphere{T}; y = sphere.origin[2]) where T
    if abs(y - sphere.origin[2]) > sphere.radius
        return EmptyShape(sphere)
    else
        r = sqrt(sphere.radius^2 - (sphere.origin[2] - y)^2)
        return Sphere{T,2}(sphere.origin[[1,3]], r)
    end
end

# boundary_functions(sphere::Sphere; kws...) = boundary_functions(Circle(sphere; kws...))

function boundary_functions(sphere::Sphere{T,3}) where T

    function x(t,s=T(0))
        check_boundary_coord_range(t)
        check_boundary_coord_range(s)

        sphere.radius * sin(t * π) * cos(s * 2π) + origin(sphere)[1]
    end

    function y(t,s=T(0))
        check_boundary_coord_range(t)
        check_boundary_coord_range(s)

        sphere.radius * sin(t * π) * sin(s * 2π) + origin(sphere)[2]
    end

    function z(t,s=T(0))
        check_boundary_coord_range(t)
        check_boundary_coord_range(s)

        sphere.radius * cos(t * π) + origin(sphere)[3]
    end

    return x, y, z
end

function boundary_functions(circle::Sphere{T,2}) where T

    function x(t)
        check_boundary_coord_range(t)
        circle.radius * cos(2π * t) + origin(circle)[1]
    end

    function y(t)
        check_boundary_coord_range(t)
        circle.radius * sin(2π * t) + origin(circle)[2]
    end

    return x, y
end


# The type hierarchy is as follows:
# abstract Ball
# |-> abstract UnitBall: radius is 1
#     |-> StaticUnitBall: dimension is part of type
#     |-> DynamicUnitBall: dimension is specified by int field
# |-> GenericBall: stores center and radius. Here, the dimension can be
#     obtained from the center, so there is only one type.
# There are aliases for SVector{N,T} and Vector{T}:
#   EuclideanBall, VectorBall,
#   EuclideanUnitBall (StaticUnitBall), VectorUnitBall (DynamicUnitBall).

"""
Supertype of balls for which elements satisfy `norm(x) < radius(ball)` (open ball)
or `norm(x) <= radius(ball)` (closed ball).
"""
abstract type Ball{T,C} <: Domain{T} end

"A ball in a fixed N-dimensional Euclidean space."
const EuclideanBall{N,T,C} = Ball{SVector{N,T},C}

"A ball with vector elements of variable length."
const VectorBall{T,C} = Ball{Vector{T},C}

iscompatiblepair(x::AbstractVector, d::VectorBall) = length(x) == dimension(d)

const ClosedBall{T} = Ball{T,:closed}
const OpenBall{T} = Ball{T,:open}

isclosedset(::ClosedBall) = true
isclosedset(::OpenBall) = false
isopenset(d::Ball) = !isclosedset(d)

# Convenience constructors for the abstract type
Ball() = UnitBall()
Ball{T}() where {T} = UnitBall{T}()
Ball{T,C}() where {T,C} = UnitBall{T,C}()
Ball(radius::Number) = GenericBall(radius)
Ball{T}(radius::Number) where {T} = GenericBall{T}(radius)
Ball{T,C}(radius::Number) where {T,C} = GenericBall{T,C}(radius)
Ball(radius::Number, center) = GenericBall(radius, center)
Ball{T}(radius::Number, center) where {T} = GenericBall{T}(radius, center)
Ball{T,C}(radius::Number, center) where {T,C} = GenericBall{T,C}(radius, center)

indomain(x, d::OpenBall) = norm(x-center(d)) < radius(d)
indomain(x, d::ClosedBall) = norm(x-center(d)) <= radius(d)
approx_indomain(x, d::OpenBall, tolerance) = norm(x-center(d)) < radius(d)+tolerance
approx_indomain(x, d::ClosedBall, tolerance) = norm(x-center(d)) <= radius(d)+tolerance

# A closed ball always contains at least its center point.
isempty(d::ClosedBall) = false
isempty(d::OpenBall) = radius(d) == 0

==(d1::Ball, d2::Ball) = isclosedset(d1)==isclosedset(d2) &&
    radius(d1)==radius(d2) && center(d1)==center(d2)

function issubset(d1::Ball, d2::Ball)
    if dimension(d1) == dimension(d2)
        if center(d1) == center(d2)
            if radius(d1) < radius(d2)
                true
            elseif radius(d1) == radius(d2)
                isclosedset(d2) || isopenset(d1)
            else
                false
            end
        else
            dist = norm(center(d1)-center(d2))
            if dist + radius(d1) < radius(d2)
                true
            else
                false
            end
        end
    else
        false
    end
end

# We choose the center of the ball here. Concrete types should implement 'center'
point_in_domain(d::Ball) = center(d)

"The unit ball."
abstract type UnitBall{T,C} <: Ball{T,C} end

radius(::UnitBall) = 1
center(d::UnitBall{T}) where {T<:StaticTypes} = zero(T)

UnitBall(n::Int) = DynamicUnitBall(n)
UnitBall(::Val{N} = Val(3)) where {N} = EuclideanUnitBall{N}()

UnitBall{T}(n::Int) where {T <: StaticTypes} = StaticUnitBall{T}(n)
UnitBall{T}(::Val{N}) where {N,T} = StaticUnitBall{T}(Val(N))
UnitBall{T}() where {T <: StaticTypes} = StaticUnitBall{T}()
UnitBall{T}(n::Int) where {T} = DynamicUnitBall{T}(n)

UnitBall{T,C}(n::Int) where {T <: StaticTypes,C} = StaticUnitBall{T,C}(n)
UnitBall{T,C}(::Val{N}) where {N,T,C} = StaticUnitBall{T,C}(Val(N))
UnitBall{T,C}() where {T <: StaticTypes,C} = StaticUnitBall{T,C}()
UnitBall{T,C}(n::Int) where {T,C} = DynamicUnitBall{T,C}(n)

const ClosedUnitBall{T} = UnitBall{T,:closed}
const OpenUnitBall{T} = UnitBall{T,:open}

indomain(x, d::OpenUnitBall) = norm(x) < 1
indomain(x, d::ClosedUnitBall) = norm(x) <= 1
approx_indomain(x, d::OpenUnitBall, tolerance) = norm(x) < 1+tolerance
approx_indomain(x, d::ClosedUnitBall, tolerance) = norm(x) <= 1+tolerance

==(d1::UnitBall, d2::UnitBall) = isclosedset(d1)==isclosedset(d2) &&
    dimension(d1) == dimension(d2)

issubset(d1::UnitBall, d2::UnitBall) =
    dimension(d1) == dimension(d2) && (isclosedset(d2) || isopenset(d1))

convert(::Type{SublevelSet}, d::UnitBall{T,C}) where {T,C} =
    SublevelSet{T,C}(norm, radius(d))
convert(::Type{SublevelSet{T}}, d::UnitBall{S,C}) where {T,S,C} =
    SublevelSet{T,C}(norm, radius(d))

"The unit ball with fixed dimension(s) specified by the element type."
struct StaticUnitBall{T,C} <: UnitBall{T,C}
end

StaticUnitBall() = StaticUnitBall{Float64}()
StaticUnitBall(::Val{N}) where {N} = StaticUnitBall{SVector{N,Float64}}()

StaticUnitBall{T}() where {T} = StaticUnitBall{T,:closed}()

StaticUnitBall{T}(n::Int) where {T} =
    (@assert n == euclideandimension(T); StaticUnitBall{T}())
StaticUnitBall{T}(::Val{N}) where {N,T} =
    (@assert N == euclideandimension(T); StaticUnitBall{T}())

StaticUnitBall{T,C}(n::Int) where {T,C} =
    (@assert n == euclideandimension(T); StaticUnitBall{T,C}())
StaticUnitBall{T,C}(::Val{N}) where {N,T,C} =
    (@assert N == euclideandimension(T); StaticUnitBall{T,C}())

"The unit ball in a fixed N-dimensional space."
const EuclideanUnitBall{N,T,C} = StaticUnitBall{SVector{N,T},C}

EuclideanUnitBall{N}() where {N} = EuclideanUnitBall{N,Float64}()

similardomain(d::StaticUnitBall{S,C}, ::Type{T}) where {S,C,T<:StaticTypes} =
    StaticUnitBall{T,C}()
similardomain(d::StaticUnitBall{S,C}, ::Type{T}) where {S,C,T} =
    DynamicUnitBall{T,C}(dimension(d))

const UnitDisk{T} = UnitBall{SVector{2,T},:closed}
UnitDisk() = UnitDisk{Float64}()

## A StaticUnitBall{<:Number} equals the interval [-1,1]  (open or closed)
convert(::Type{Interval}, d::StaticUnitBall{T,:closed}) where {T <: Number} =
    ChebyshevInterval{T}()
convert(::Type{Interval}, d::StaticUnitBall{T,:open}) where {T <: Number} =
    OpenInterval{T}(-1, 1)

canonicaldomain(::Equal, d::StaticUnitBall{T}) where {T<:Number} = convert(Interval, d)

"The unit ball with variable dimension stored in a data field."
struct DynamicUnitBall{T,C} <: UnitBall{T,C}
    dimension   ::  Int

    DynamicUnitBall{T,C}(n::Int) where {T,C} = new(n)
    DynamicUnitBall{T,C}(n::Int) where {T<:StaticTypes,C} =
        (@assert n == euclideandimension(T); new(n))
end

DynamicUnitBall(n::Int) = DynamicUnitBall{Vector{Float64}}(n)
DynamicUnitBall{T}(n::Int) where {T} = DynamicUnitBall{T,:closed}(n)

dimension(d::DynamicUnitBall) = d.dimension

center(d::DynamicUnitBall{T}) where {T} = zeros(eltype(T), dimension(d))

"The unit ball with vector elements of a given dimension."
const VectorUnitBall{T,C} = DynamicUnitBall{Vector{T},C}

VectorUnitBall(n::Int = 3) = VectorUnitBall{Float64}(n)
VectorUnitDisk() = VectorUnitBall(2)

similardomain(d::DynamicUnitBall{S,C}, ::Type{T}) where {S,C,T} =
    DynamicUnitBall{T,C}(dimension(d))
similardomain(d::DynamicUnitBall{S,C}, ::Type{T}) where {S,C,T<:StaticTypes} =
    StaticUnitBall{T,C}()

"A `GenericBall` is a ball with a given radius and center."
struct GenericBall{T,C,S} <: Ball{T,C}
    radius  ::  S
    center  ::  T
end

GenericBall() = GenericBall(1.0)
GenericBall(radius::S) where {S<:Number} = GenericBall(radius, zero(SVector{3,float(S)}))
GenericBall(radius::S, center::T) where {S<:Number,U,T<:Vector{U}} =
    GenericBall{Vector{promote_type(S,U)}}(radius, center)
GenericBall(radius::S, center::T) where {S<:Number,N,U,T<:SVector{N,U}} =
    GenericBall{SVector{N,promote_type(S,U)}}(radius, center)
GenericBall(radius::S, center::T) where {S<:Number,T<:AbstractVector} =
    GenericBall{Vector{promote_type(S,eltype(T))}}(radius, center)
GenericBall(radius::S, center::T) where {S<:Number,T<:Number} =
    GenericBall{promote_type(S,T)}(radius, center)
GenericBall{T}(radius) where {T} = GenericBall{T,:closed}(radius)
GenericBall{T}(radius, center) where {T} = GenericBall{T,:closed}(radius, center)
GenericBall{T,C}(radius) where {T,C} = GenericBall{T,C}(radius, zero(T))
GenericBall{T,C}(radius, center) where {T,C} = GenericBall{T,C,eltype(T)}(radius, center)

radius(d::GenericBall) = d.radius
center(d::GenericBall) = d.center

dimension(d::GenericBall) = length(d.center)

canonicaldomain(d::GenericBall{T,C}) where {T,C} = UnitBall{T,C}(dimension(d))
mapfrom_canonical(d::GenericBall) = AffineMap(radius(d), center(d))

boundingbox(d::GenericBall) =
    map_boundingbox(boundingbox(canonicaldomain(d)), mapfrom_canonical(d))

# Preserve the `Ball` type under affine maps which preserve shape
map_domain(m::GenericAffineMap{T,S}, d::Ball{U,C}) where {T<:AbstractVector,S<:Number,U<:AbstractVector,C} =
    Ball{U,C}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::AffineMap{T}, d::Ball{U,C}) where {T<:Number,U<:Number,C} =
    Ball{promote_type(U,T),C}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::GenericLinearMap{T,S}, d::Ball{U,C}) where {S<:Number,T<:AbstractVector,U<:AbstractVector,C} =
    Ball{promote_type(T,U),C}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Ball{U,C}) where {T<:Number,U<:Number,C} =
    Ball{promote_type(T,U),C}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Ball{U,C}) where {T<:Number,U,C} =
    Ball{U,C}(m.A * radius(d), m.A * center(d))
map_domain(m::Translation{T}, d::Ball{S,C}) where {T<:AbstractVector,S<:AbstractVector,C} =
    Ball{S,C}(radius(d), center(d) + m.b)
map_domain(m::Translation{T}, d::Ball{S,C}) where {T<:Number,S<:Number,C} =
    Ball{S,C}(radius(d), center(d) + m.b)

show(io::IO, d::EuclideanUnitBall{3,Float64,:closed}) = print(io, "UnitBall()")
show(io::IO, d::EuclideanUnitBall{N,Float64,:closed}) where {N} = print(io, "UnitBall(Val($(N)))")
show(io::IO, d::EuclideanUnitBall{3,Float64,:open}) = print(io, "UnitBall()  (open)")
show(io::IO, d::EuclideanUnitBall{N,Float64,:open}) where {N} = print(io, "UnitBall(Val($(N)))  (open)")
Display.object_parentheses(d::EuclideanUnitBall{N,Float64,:open}) where {N} = true
show(io::IO, d::EuclideanUnitBall{2,Float64,:closed}) = print(io, "UnitDisk()")
show(io::IO, d::UnitDisk{T}) where {T} = print(io, "UnitDisk{$(T)}()")
show(io::IO, d::UnitBall{T}) where {T<:Number} = print(io, "UnitBall{$(T)}()")
show(io::IO, d::UnitBall{T,:open}) where {T<:Number} = print(io, "UnitBall{$(T)}()  (open)")
Display.object_parentheses(d::UnitBall{T,:open}) where {T<:Number} = true
show(io::IO, d::VectorUnitBall{Float64,:closed}) = print(io, "UnitBall($(dimension(d)))")
show(io::IO, d::VectorUnitBall{Float64,:open}) = print(io, "UnitBall($(dimension(d)))  (open)")
show(io::IO, d::Ball) = print(io, "Ball($(radius(d)), $(center(d)))")

# The type hierarchy of spheres parallels that of Ball above:
# abstract Sphere
# |-> abstract UnitSphere: radius is 1
#     |-> StaticUnitSphere: dimension is part of type
#     |-> DynamicUnitSphere: dimension is specified by int field
# |-> GenericSphere: stores center and radius. Here, the dimension can be
#     obtained from the center, so there is only one type.
# There are aliases for SVector{N,T} and Vector{T}.

"Supertype of spherical domains for which elements satisfy `norm(x) == radius(sphere)`."
abstract type Sphere{T} <: Domain{T} end

indomain(x, d::Sphere) = norm(x-center(d)) == radius(d)
approx_indomain(x, d::Sphere, tolerance) =
    radius(d)-tolerance <= norm(x-center(d)) <= radius(d)+tolerance

# Convenience constructors for the abstract type
Sphere() = UnitSphere()
Sphere{T}() where {T} = UnitSphere{T}()
Sphere(radius::Number) = GenericSphere(radius)
Sphere{T}(radius::Number) where {T} = GenericSphere{T}(radius)
Sphere(radius::Number, center) = GenericSphere(radius, center)
Sphere{T}(radius::Number, center) where {T} = GenericSphere{T}(radius, center)

isempty(::Sphere) = false

isclosedset(::Sphere) = true
isopenset(::Sphere) = false

point_in_domain(d::Sphere) = center(d) + unitvector(d)

==(d1::Sphere, d2::Sphere) = radius(d1)==radius(d2) && center(d1)==center(d2)

"A hypersphere in a fixed N-dimensional Euclidean space."
const EuclideanSphere{N,T} = Sphere{SVector{N,T}}

"A sphere with vector elements of variable length."
const VectorSphere{T} = Sphere{Vector{T}}

"The unit sphere."
abstract type UnitSphere{T} <: Sphere{T} end

radius(::UnitSphere) = 1
center(d::UnitSphere{T}) where {T<:StaticTypes} = zero(T)

indomain(x, d::UnitSphere) = norm(x) == 1
approx_indomain(x, d::UnitSphere, tolerance) =
    1-tolerance <= norm(x) <= 1+tolerance

UnitSphere(n::Int) = DynamicUnitSphere(n)
UnitSphere(::Val{N} = Val(3)) where {N} = EuclideanUnitSphere{N}()

UnitSphere{T}(n::Int) where {T <: StaticTypes} = StaticUnitSphere{T}(n)
UnitSphere{T}(::Val{N}) where {N,T} = StaticUnitSphere{T}(Val(N))
UnitSphere{T}() where {T <: StaticTypes} = StaticUnitSphere{T}()
UnitSphere{T}(n::Int) where {T} = DynamicUnitSphere{T}(n)

==(d1::UnitSphere, d2::UnitSphere) = dimension(d1)==dimension(d2)

issubset1(d1::UnitSphere, d2::UnitBall) =
    dimension(d1) == dimension(d2) && isclosedset(d2)

convert(::Type{LevelSet}, d::UnitSphere{T}) where {T} = LevelSet{T}(norm, radius(d))
convert(::Type{LevelSet{T}}, d::UnitSphere) where {T} = LevelSet{T}(norm, radius(d))

"The unit sphere with fixed dimension(s) specified by the element type."
struct StaticUnitSphere{T} <: UnitSphere{T}
end

StaticUnitSphere() = StaticUnitSphere{Float64}()
StaticUnitSphere(::Val{N}) where {N} = StaticUnitSphere{SVector{N,Float64}}()

StaticUnitSphere{T}(n::Int) where {T} =
    (@assert n == euclideandimension(T); StaticUnitSphere{T}())
StaticUnitSphere{T}(::Val{N}) where {N,T} =
    (@assert N == euclideandimension(T); StaticUnitSphere{T}())

similardomain(d::StaticUnitSphere, ::Type{T}) where {T<:StaticTypes} =
    StaticUnitSphere{T}()
similardomain(d::StaticUnitSphere, ::Type{T}) where {T} =
    DynamicUnitSphere{T}(dimension(d))

"The unit sphere in a fixed N-dimensional Euclidean space."
const EuclideanUnitSphere{N,T} = StaticUnitSphere{SVector{N,T}}

EuclideanUnitSphere{N}() where {N} = EuclideanUnitSphere{N,Float64}()

"The unit circle in 2D."
const UnitCircle{T} = UnitSphere{SVector{2,T}}
UnitCircle() = UnitCircle{Float64}()

"The unit sphere with variable dimension."
struct DynamicUnitSphere{T} <: UnitSphere{T}
    dimension   ::  Int

    DynamicUnitSphere{T}(n::Int) where {T} = new(n)
    DynamicUnitSphere{T}(n::Int) where {T<:StaticTypes} =
        (@assert n == euclideandimension(T); new(n))
end

DynamicUnitSphere(n::Int) = DynamicUnitSphere{Vector{Float64}}(n)

dimension(d::DynamicUnitSphere) = d.dimension

center(d::DynamicUnitSphere{T}) where {T} = zeros(eltype(T), dimension(d))

"The unit sphere with vector elements of a given dimension."
const VectorUnitSphere{T} = DynamicUnitSphere{Vector{T}}

VectorUnitSphere(n::Int = 3) = VectorUnitSphere{Float64}(n)
VectorUnitCircle() = VectorUnitSphere(2)

similardomain(d::DynamicUnitSphere, ::Type{T}) where {T} =
    DynamicUnitSphere{T}(d.dimension)
similardomain(d::DynamicUnitSphere, ::Type{T}) where {T <: StaticTypes} =
    StaticUnitSphere{T}()

"A `GenericSphere` is a sphere with a given radius and center."
struct GenericSphere{T,S} <: Sphere{T}
    radius  ::  S
    center  ::  T
end

GenericSphere() = GenericSphere(1.0)
GenericSphere(radius::S) where {S<:Number} =
    GenericSphere(radius, zero(SVector{3,float(S)}))
GenericSphere(radius::S, center::T) where {S<:Number,U,T<:Vector{U}} =
    GenericSphere{Vector{promote_type(S,U)}}(radius, center)
GenericSphere(radius::S, center::T) where {S<:Number,N,U,T<:SVector{N,U}} =
    GenericSphere{SVector{N,promote_type(S,U)}}(radius, center)
GenericSphere(radius::S, center::T) where {S<:Number,T<:AbstractVector} =
    GenericSphere{Vector{promote_type(S,eltype(T))}}(radius, center)
GenericSphere{T}(radius) where {T} = GenericSphere{T}(radius, zero(T))
GenericSphere{T}(radius, center) where {T} =
    GenericSphere{T,eltype(T)}(radius, center)

radius(d::GenericSphere) = d.radius
center(d::GenericSphere) = d.center

dimension(d::GenericSphere) = length(d.center)

canonicaldomain(d::GenericSphere{T}) where {T} = UnitSphere{T}(dimension(d))
mapfrom_canonical(d::GenericSphere) = AffineMap(radius(d), center(d))

boundingbox(d::GenericSphere) =
    map_boundingbox(boundingbox(canonicaldomain(d)), mapfrom_canonical(d))

# Preserve the `Sphere` type under affine maps which preserve shape
map_domain(m::GenericAffineMap{T,S}, d::Sphere{U}) where {T<:AbstractVector,S<:Number,U<:AbstractVector} =
    Sphere{U}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::AffineMap{T}, d::Sphere{U}) where {T<:Number,U<:Number} =
    Sphere{promote_type(U,T)}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::GenericLinearMap{T,S}, d::Sphere{U}) where {S<:Number,T<:AbstractVector,U<:AbstractVector} =
    Sphere{promote_type(T,U)}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Sphere{U}) where {T<:Number,U<:Number} =
    Sphere{promote_type(T,U)}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Sphere{U}) where {T<:Number,U} =
    Sphere{U}(m.A * radius(d), m.A * center(d))
map_domain(m::Translation{T}, d::Sphere{S}) where {T<:AbstractVector,S<:AbstractVector} =
    Sphere{S}(radius(d), center(d) + m.b)
map_domain(m::Translation{T}, d::Sphere{S}) where {T<:Number,S<:Number} =
    Sphere{S}(radius(d), center(d) + m.b)

show(io::IO, d::EuclideanUnitSphere{3,Float64}) = print(io, "UnitSphere()")
show(io::IO, d::EuclideanUnitSphere{N,Float64}) where {N} = print(io, "UnitSphere(Val($(N)))")
show(io::IO, d::EuclideanUnitSphere{2,Float64}) = print(io, "UnitCircle()")
show(io::IO, d::UnitCircle{T}) where {T} = print(io, "UnitCircle{$(T)}()")
show(io::IO, d::VectorUnitSphere{Float64}) = print(io, "UnitSphere($(dimension(d)))")
show(io::IO, d::UnitSphere{T}) where {T<:Number} = print(io, "UnitSphere{$(T)}()")
show(io::IO, d::Sphere) = print(io, "Sphere($(radius(d)), $(center(d)))")

boundary(d::UnitBall{T}) where {T} = UnitSphere{T}(dimension(d))
boundary(d::Ball{T}) where {T} = Sphere{T}(radius(d), center(d))

interior(d::StaticUnitBall{T}) where {T} = StaticUnitBall{T,:open}()
interior(d::DynamicUnitBall{T}) where {T} = DynamicUnitBall{T,:open}(dimension(d))
interior(d::Ball{T}) where {T} = Ball{T,:open}(radius(d), center(d))

closure(d::StaticUnitBall{T}) where {T} = StaticUnitBall{T,:closed}()
closure(d::DynamicUnitBall{T}) where {T} = DynamicUnitBall{T,:closed}(dimension(d))
closure(d::Ball{T}) where {T} = Ball{T,:closed}(radius(d), center(d))

boundingbox(d::UnitBall{T}) where {T<:Number} = ChebyshevInterval{T}()
boundingbox(d::UnitBall{SVector{N,T}}) where {N,T} =
    ChebyshevProductDomain{N,T}()
boundingbox(d::UnitBall{T}) where {T} =
    Rectangle{T}(-ones(eltype(T), dimension(d)), ones(eltype(T), dimension(d)))

boundingbox(d::UnitSphere{T}) where {T<:Number} = ChebyshevInterval{T}()
boundingbox(d::UnitSphere{SVector{N,T}}) where {N,T} =
    ChebyshevProductDomain{N,T}()
boundingbox(d::UnitSphere{T}) where {T} =
    Rectangle{T}(-ones(eltype(T), dimension(d)), ones(eltype(T), dimension(d)))

################
# Applications
################

"Create a cylinder with given radius and length."
cylinder(::Type{T} = Float64) where {T} = UnitDisk{T}() × UnitInterval{T}()
cylinder(radius, length) = cylinder(promote(radius, length)...)
cylinder(radius::T, length::T) where {T} = (radius .* UnitDisk{T}()) × (0..length)

"Create an ellipse curve with semi-axes lengths `a` and `b` respectively."
ellipse(a::Number, b::Number) = ellipse(promote(a,b)...)
ellipse(a::T, b::T) where {T <: Number} = LinearMap(a, b).(UnitCircle{T}())

"Create an ellipse-shaped domain with semi-axes lengths `a` and `b` respectively."
ellipse_shape(a::Number, b::Number) = ellipse_shape(promote(a,b)...)
ellipse_shape(a::T, b::T) where {T <: Number} = LinearMap(a, b).(UnitDisk{T}())

"""
The map `[cos(2πt), sin(2πt)]` from `[0,1]` to the unit circle in `ℝ^2`.
"""
struct UnitCircleMap{T} <: Map{T} end

UnitCircleMap() = UnitCircleMap{Float64}()

mapsize(m::UnitCircleMap) = (2,)

applymap(m::UnitCircleMap{T}, t) where {T} = SVector(cos(2*T(pi)*t), sin(2*T(pi)*t))

function jacobian(m::UnitCircleMap{T}, t) where {T}
    a = 2*T(pi)
    SVector(-a*sin(a*t), a*cos(a*t))
end

# we know that the differential volume is 2*pi
diffvolume(m::UnitCircleMap{T}) where {T} = ConstantMap{T}(2*T(pi))

"""
`AngleMap` is a left inverse of `UnitCircleMap`. A 2D vector `x` is projected onto
the intersection point with the unit circle of the line connecting `x` to the
origin. The angle of this point, scaled to the interval `[0,1)`, is the result.
"""
struct AngleMap{T} <: Map{SVector{2,T}}
end

AngleMap() = AngleMap{Float64}()

function applymap(m::AngleMap{T}, x) where {T}
    twopi = 2*convert(T, pi)
    θ = atan(x[2],x[1])
    if θ < 0
        # atan2 returns an angle in (-π,π], convert to [0,2π) using periodicity.
        θ += twopi
    end
    # And divide by 2π to scale to [0,1)
    θ / twopi
end

mapsize(m::AngleMap) = (1,2)

function jacobian(m::AngleMap{T}, t) where {T}
    x = t[1]; y = t[2]
    twopi = 2*convert(T, pi)
    v = SVector(one(T) / (1+(y/x)^2) * (-y/x^2) * 1/twopi,
            one(T) / (1+(y/x)^2) * one(T)/x * 1/twopi)
    transpose(v)
end

leftinverse(m::UnitCircleMap{T}) where {T} = AngleMap{T}()
leftinverse(m::UnitCircleMap, x) = leftinverse(m)(x)
rightinverse(m::AngleMap{T}) where {T} = UnitCircleMap{T}()
rightinverse(m::AngleMap, x) = rightinverse(m)(x)

canonicaldomain(::Parameterization, d::UnitCircle{T}) where {T} = UnitInterval{T}()
mapfrom_canonical(::Parameterization, d::UnitCircle{T}) where {T} = UnitCircleMap{T}()

## The complex plane

const ComplexUnitCircle{T} = StaticUnitSphere{Complex{T}}
const ComplexUnitDisk{T,C} = StaticUnitBall{Complex{T},C}

ComplexUnitCircle() = ComplexUnitCircle{Float64}()
ComplexUnitDisk() = ComplexUnitDisk{Float64}()
ComplexUnitDisk{T}() where {T} = ComplexUnitDisk{T,:closed}()

show(io::IO, d::ComplexUnitCircle{Float64}) = print(io, "ComplexUnitCircle()")
show(io::IO, d::ComplexUnitDisk{Float64,:closed}) = print(io, "ComplexUnitDisk()")
show(io::IO, d::ComplexUnitDisk{Float64,:open}) = print(io, "ComplexUnitDisk()  (open)")
show(io::IO, d::ComplexUnitDisk{T,:closed}) where {T} = print(io, "ComplexUnitDisk{$(T)}()")
show(io::IO, d::ComplexUnitDisk{T,:open}) where {T} = print(io, "ComplexUnitDisk{$(T)}()  (open)")

canonicaldomain(::Parameterization, d::ComplexUnitCircle{T}) where {T} =
    UnitInterval{T}()
mapfrom_canonical(::Parameterization, d::ComplexUnitCircle{T}) where {T} =
    VectorToComplex{T}() ∘ UnitCircleMap{T}()


"""
Module Setup2D

Module to aid in setting up reference operators, meshes, geometric terms
Also provides functions to assemble global DG-SBP operators

"""

module SetupDG

# non-DG modules
using LinearAlgebra # for diagm
using SparseArrays # for sparse, droptol
using UnPack # for easy setting/getting in mutable structs

# matlab-like modules
using CommonUtils # for I matrix in geometricFactors
using Basis1D

# triangular routines
import Basis2DTri
import UniformTriMesh # for face vertices

# quadrilateral routines
import Basis2DQuad
import UniformQuadMesh # for face vertices

# hex routines
import Basis3DHex
import UniformHexMesh # for face vertices

# initialization of mesh/reference element data
export init_reference_interval, init_reference_tri
export init_reference_quad, init_reference_hex
export init_mesh
export MeshData, RefElemData

# do we realy need mutable here?  I think we can remove it
# TODO: add types to all of these, siwtch to immutable struct
mutable struct MeshData

    VX;VY;VZ # vertex coordinates
    K::Int # num elems
    EToV # mesh vertex array
    FToF # face connectivity

    x; y; z # physical points
    xf;yf;zf
    xq;yq;zq;wJq # phys quad points, Jacobian-scaled weights

    mapM; mapP; mapB # connectivity between face nodes

    # volume geofacs
    rxJ; sxJ; txJ
    ryJ; syJ; tyJ
    rzJ; szJ; tzJ; J

    # surface geofacs
    nxJ; nyJ; nzJ; sJ

    MeshData() = new() # empty initializer

end

mutable struct RefElemData

    Nfaces; fv # face vertex tuple list

    V1 # low order interp nodes and matrix

    # probably won't use nodes, but might as well keep them around
    r; s; t        # interpolation nodes
    rq; sq; tq; wq # volume quadrature
    rf; sf; tf; wf # surface quadrature
    rp; sp; tp     # plotting nodes

    VDM         # Vandermonde matrix
    Dr; Ds; Dt  # differentiation matrices
    Vq; Vf      # quadrature interpolation matrices
    M; Pq       # mass matrix, L2 projection matrix
    LIFT        # quadrature-based lift matrix
    Vp          # interp to equispaced nodes

    nrJ; nsJ; ntJ # reference normals

    RefElemData() = new() # empty initializer
end

function init_reference_interval(N)
    # initialize a new reference element data struct
    rd = RefElemData()

    # Construct matrices on reference elements
    r,_ = gauss_lobatto_quad(0,0,N)
    VDM = vandermonde_1D(N, r)
    Dr = grad_vandermonde_1D(N, r)/VDM
    @pack! rd = r,VDM,Dr

    V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1])
    @pack! rd = V1

    rq,wq = gauss_quad(0,0,N+1)
    Vq = vandermonde_1D(N, rq)/VDM
    M = Vq'*diagm(wq)*Vq
    Pq = M\(Vq'*diagm(wq))
    @pack! rd = rq,wq,Vq,M,Pq

    rf = [-1;1]
    nrJ = [-1;1]
    Vf = vandermonde_1D(N,rf)/VDM
    LIFT = M\(transpose(Vf)) # lift matrix
    @pack! rd = rf,nrJ,Vf,LIFT

    # plotting nodes
    rp = LinRange(-1,1,50)
    Vp = vandermonde_1D(N,rp)/VDM
    @pack! rd = rp,Vp

    return rd

end

function init_reference_tri(N)

    # initialize a new reference element data struct
    rd = RefElemData()

    fv = UniformTriMesh.tri_face_vertices() # set faces for triangle
    Nfaces = length(fv)
    @pack! rd = fv, Nfaces

    # Construct matrices on reference elements
    r, s = Basis2DTri.nodes_2D(N)
    VDM = Basis2DTri.vandermonde_2D(N, r, s)
    Vr, Vs = Basis2DTri.grad_vandermonde_2D(N, r, s)
    Dr = Vr/VDM
    Ds = Vs/VDM
    @pack! rd = r,s,VDM,Dr,Ds

    # low order interpolation nodes
    r1,s1 = Basis2DTri.nodes_2D(1)
    V1 = Basis2DTri.vandermonde_2D(1,r,s)/Basis2DTri.vandermonde_2D(1,r1,s1)
    @pack! rd = V1

    #Nodes on faces, and face node coordinate
    r1D, w1D = gauss_quad(0,0,N)
    Nfp = length(r1D) # number of points per face
    e = ones(Nfp) # vector of all ones
    z = zeros(Nfp) # vector of all zeros
    rf = [r1D; -r1D; -e];
    sf = [-e; r1D; -r1D];
    wf = vec(repeat(w1D,3,1));
    nrJ = [z; e; -e]
    nsJ = [-e; e; z]
    @pack! rd = rf,sf,wf,nrJ,nsJ

    rq,sq,wq = Basis2DTri.quad_nodes_2D(2*N)
    Vq = Basis2DTri.vandermonde_2D(N,rq,sq)/VDM
    M = transpose(Vq)*diagm(wq)*Vq
    Pq = M\(transpose(Vq)*diagm(wq))
    @pack! rd = rq,sq,wq,Vq,M,Pq

    Vf = Basis2DTri.vandermonde_2D(N,rf,sf)/VDM # interpolates from nodes to face nodes
    LIFT = M\(transpose(Vf)*diagm(wf)) # lift matrix used in rhs evaluation
    @pack! rd = Vf,LIFT

    # plotting nodes
    rp, sp = Basis2DTri.equi_nodes_2D(15)
    Vp = Basis2DTri.vandermonde_2D(N,rp,sp)/VDM
    @pack! rd = rp,sp,Vp

    return rd
end

# default to full quadrature nodes
# if quad_nodes_1D=tuple of (r1D,w1D) is supplied, use those nodes
function init_reference_quad(N,quad_nodes_1D = gauss_quad(0,0,N))

    # initialize a new reference element data struct
    rd = RefElemData()

    fv = UniformQuadMesh.quad_face_vertices() # set faces for triangle
    Nfaces = length(fv)
    @pack! rd = fv, Nfaces

    # Construct matrices on reference elements
    r, s = Basis2DQuad.nodes_2D(N)
    VDM = Basis2DQuad.vandermonde_2D(N, r, s)
    Vr, Vs = Basis2DQuad.grad_vandermonde_2D(N, r, s)
    Dr = Vr/VDM
    Ds = Vs/VDM
    @pack! rd = r,s,VDM

    # low order interpolation nodes
    r1,s1 = Basis2DQuad.nodes_2D(1)
    V1 = Basis2DQuad.vandermonde_2D(1,r,s)/Basis2DQuad.vandermonde_2D(1,r1,s1)
    @pack! rd = V1

    #Nodes on faces, and face node coordinate
    # r1D,w1D = quad_nodes_1D(0,0,N)
    # r1D,w1D = gauss_lobatto_quad(0,0,N)
    # r1D,w1D = gauss_quad(0,0,N)
    r1D,w1D = quad_nodes_1D
    Nfp = length(r1D)
    e = ones(size(r1D))
    z = zeros(size(r1D))
    rf = [r1D; e; -r1D; -e]
    sf = [-e; r1D; e; -r1D]
    wf = vec(repeat(w1D,Nfaces,1));
    nrJ = [z; e; z; -e]
    nsJ = [-e; z; e; z]
    @pack! rd = rf,sf,wf,nrJ,nsJ

    # quadrature nodes - build from 1D nodes.
    # can also use "rq,sq,wq = Basis2DQuad.quad_nodes_2D(2*N)"
    rq,sq = (x->x[:]).(meshgrid(r1D))
    wr,ws = meshgrid(w1D)
    wq = wr[:] .* ws[:]
    Vq = Basis2DQuad.vandermonde_2D(N,rq,sq)/VDM
    M = transpose(Vq)*diagm(wq)*Vq
    Pq = M\(transpose(Vq)*diagm(wq))
    @pack! rd = rq,sq,wq,Vq,M,Pq

    Vf = Basis2DQuad.vandermonde_2D(N,rf,sf)/VDM # interpolates from nodes to face nodes
    LIFT = M\(transpose(Vf)*diagm(wf)) # lift matrix used in rhs evaluation

    # expose kronecker product sparsity
    Dr = droptol!(sparse(Dr), 1e-10)
    Ds = droptol!(sparse(Ds), 1e-10)
    Vf = droptol!(sparse(Vf),1e-10)
    LIFT = droptol!(sparse(LIFT),1e-10)
    @pack! rd = Dr,Ds,Vf,LIFT

    # plotting nodes
    rp, sp = Basis2DQuad.equi_nodes_2D(15)
    Vp = Basis2DQuad.vandermonde_2D(N,rp,sp)/VDM
    @pack! rd = rp,sp,Vp

    return rd
end

# dispatch to 2D or 3D version if tuple called
function init_mesh(VXYZ,EToV,rd::RefElemData)
    return init_mesh(VXYZ...,EToV,rd)
end

function init_mesh(VX,VY,EToV,rd::RefElemData)

    # initialize a new mesh data struct
    md = MeshData()

    @unpack fv = rd
    FToF = connect_mesh(EToV,fv)
    Nfaces,K = size(FToF)
    @pack! md = FToF,K,VX,VY,EToV

    #Construct global coordinates
    @unpack V1 = rd
    x = V1*VX[transpose(EToV)]
    y = V1*VY[transpose(EToV)]
    @pack! md = x,y

    #Compute connectivity maps: uP = exterior value used in DG numerical fluxes
    @unpack r,s,Vf = rd
    xf,yf = (x->Vf*x).((x,y))
    mapM,mapP,mapB = build_node_maps((xf,yf),FToF)
    Nfp = convert(Int,size(Vf,1)/Nfaces)
    mapM = reshape(mapM,Nfp*Nfaces,K)
    mapP = reshape(mapP,Nfp*Nfaces,K)
    @pack! md = xf,yf,mapM,mapP,mapB

    #Compute geometric factors and surface normals
    @unpack Dr,Ds = rd
    rxJ, sxJ, ryJ, syJ, J = geometric_factors(x, y, Dr, Ds)
    @pack! md = rxJ, sxJ, ryJ, syJ, J

    @unpack Vq,wq = rd
    xq,yq = (x->Vq*x).((x,y))
    wJq = diagm(wq)*(Vq*J)
    @pack! md = xq,yq,wJq

    #physical normals are computed via G*nhatJ, G = matrix of geometric terms
    @unpack nrJ,nsJ = rd
    nxJ = (Vf*rxJ).*nrJ + (Vf*sxJ).*nsJ;
    nyJ = (Vf*ryJ).*nrJ + (Vf*syJ).*nsJ;
    sJ = @. sqrt(nxJ^2 + nyJ^2)
    @pack! md = nxJ,nyJ,sJ

    return md
end

"========== 3D routines ============="

function init_reference_hex(N,quad_nodes_1D=gauss_quad(0,0,N))

    # initialize a new reference element data struct
    rd = RefElemData()

    fv = UniformHexMesh.hex_face_vertices() # set faces for triangle
    Nfaces = length(fv)
    @pack! rd = fv, Nfaces

    # Construct matrices on reference elements
    r,s,t = Basis3DHex.nodes_3D(N)
    VDM = Basis3DHex.vandermonde_3D(N,r,s,t)
    Vr,Vs,Vt = Basis3DHex.grad_vandermonde_3D(N,r,s,t)
    Dr,Ds,Dt = (A->A/VDM).(Basis3DHex.grad_vandermonde_3D(N,r,s,t))
    @pack! rd = r,s,t,VDM

    # low order interpolation nodes
    r1,s1,t1 = Basis3DHex.nodes_3D(1)
    V1 = Basis3DHex.vandermonde_3D(1,r,s,t)/Basis3DHex.vandermonde_3D(1,r1,s1,t1)
    @pack! rd = V1

    #Nodes on faces, and face node coordinate
    r1D,w1D = quad_nodes_1D
    rquad,squad = (x->x[:]).(meshgrid(r1D,r1D))
    wr,ws = (x->x[:]).(meshgrid(w1D,w1D))
    wquad = wr.*ws
    e = ones(size(rquad))
    zz = zeros(size(rquad))
    rf = [-e; e; rquad; rquad; rquad; rquad]
    sf = [rquad; rquad; -e; e; squad; squad]
    tf = [squad; squad; squad; squad; -e; e]
    wf = vec(repeat(wquad,Nfaces,1));
    nrJ = [-e; e; zz;zz; zz;zz]
    nsJ = [zz;zz; -e; e; zz;zz]
    ntJ = [zz;zz; zz;zz; -e; e]

    @pack! rd = rf,sf,tf,wf,nrJ,nsJ,ntJ

    # quadrature nodes - build from 1D nodes.
    rq,sq,tq = meshgrid(r1D,r1D,r1D)
    wr,ws,wt = meshgrid(w1D,w1D,w1D)
    wq = wr.*ws.*wt
    Vq = Basis3DHex.vandermonde_3D(N,rq,sq,tq)/VDM
    M = transpose(Vq)*diagm(wq)*Vq
    Pq = M\(transpose(Vq)*diagm(wq))
    @pack! rd = rq,sq,tq,wq,Vq,M,Pq

    Vf = Basis3DHex.vandermonde_3D(N,rf,sf,tf)/VDM
    LIFT = M\(transpose(Vf)*diagm(wf))

    # expose kronecker product sparsity
    Dr = droptol!(sparse(Dr),1e-12)
    Ds = droptol!(sparse(Ds),1e-12)
    Dt = droptol!(sparse(Dt),1e-12)
    Vf = droptol!(sparse(Vf),1e-12)
    LIFT = droptol!(sparse(LIFT),1e-12)
    @pack! rd = Dr,Ds,Dt,Vf,LIFT

    # plotting nodes
    rp,sp,tp = Basis3DHex.equi_nodes_3D(15)
    Vp = Basis3DHex.vandermonde_3D(N,rp,sp,tp)/VDM
    @pack! rd = rp,sp,tp,Vp

    return rd
end

function init_mesh(VX,VY,VZ,EToV,rd::RefElemData)

    # initialize a new mesh data struct
    md = MeshData()

    @unpack fv = rd
    FToF = connect_mesh(EToV,fv)
    Nfaces,K = size(FToF)
    @pack! md = FToF,K,VX,VY,VZ,EToV

    #Construct global coordinates
    @unpack V1 = rd
    x = V1*VX[transpose(EToV)]
    y = V1*VY[transpose(EToV)]
    z = V1*VZ[transpose(EToV)]
    @pack! md = x,y,z

    #Compute connectivity maps: uP = exterior value used in DG numerical fluxes
    @unpack r,s,t,Vf = rd
    xf,yf,zf = (x->Vf*x).((x,y,z))
    mapM,mapP,mapB = build_node_maps((xf,yf,zf),FToF)
    Nfp = convert(Int,size(Vf,1)/Nfaces)
    mapM = reshape(mapM,Nfp*Nfaces,K)
    mapP = reshape(mapP,Nfp*Nfaces,K)
    @pack! md = xf,yf,zf,mapM,mapP,mapB

    #Compute geometric factors and surface normals
    @unpack Dr,Ds,Dt = rd
    rxJ,sxJ,txJ,ryJ,syJ,tyJ,rzJ,szJ,tzJ,J = geometric_factors(x,y,z,Dr,Ds,Dt)
    @pack! md = rxJ,sxJ,txJ,ryJ,syJ,tyJ,rzJ,szJ,tzJ,J

    @unpack Vq,wq = rd
    xq,yq,zq = (x->Vq*x).((x,y,z))
    wJq = diagm(wq)*(Vq*J)
    @pack! md = xq,yq,zq,wJq

    #physical normals are computed via G*nhatJ, G = matrix of geometric terms
    @unpack nrJ,nsJ,ntJ = rd
    nxJ = nrJ.*(Vf*rxJ) + nsJ.*(Vf*sxJ) + ntJ.*(Vf*txJ)
    nyJ = nrJ.*(Vf*ryJ) + nsJ.*(Vf*syJ) + ntJ.*(Vf*tyJ)
    nzJ = nrJ.*(Vf*rzJ) + nsJ.*(Vf*szJ) + ntJ.*(Vf*tzJ)
    sJ = @. sqrt(nxJ.^2 + nyJ.^2 + nzJ.^2)
    @pack! md = nxJ,nyJ,nzJ,sJ

    return md
end

end


using LinearAlgebra
using SparseArrays
using DifferentialEquations
using Statistics
using Sundials
using DataStructures
using Plots
using Plots.PlotMeasures
include("MC_functions.jl")

model_name = "US_Bank_Callcenter"
p = Dict("forwards" => true)
S = 470
t = 1440.0
u_0 = bd_stat(US_Bank_Q(0))
prob = ODEProblem(US_Bank_dudt, u_0, (0.0, 1440.0),p)
sol = solve(prob, CVODE_BDF(), reltol=1e-14, abstol=1e-14, saveat=5.0)

NS_perf_mes = zeros(288)
rhos = zeros(288)
lzs = zeros(288)
for i in range(1,stop=288)
    if ~ is_stochastic_vector(sol[i])
        sol[i][sol[i].<0].=0
        if ~ is_stochastic_vector(sol[i])
            print("Not stochastic after shift.")
        end
    end
    NS_perf_mes[i] = sol[i]' * US_Bank_Number_Waiting_r(sol.t[i])
    rhos[i] = US_Bank_Lambda(sol.t[i])/(1/3*US_Bank_Servers(sol.t[i]))
    lzs[i] = US_Bank_Lambda(sol.t[i])
end

perf_mes = zeros(288)
for i in range(1,stop=288)
    # perf_mes[i] = get_stat_dist_Q(US_Bank_Q(sol.t[i]))' * US_Bank_Number_Waiting_r(sol.t[i])
    print(i)
    perf_mes[i] = bd_stat_mode(US_Bank_Q(sol.t[i]))' * US_Bank_Number_Waiting_r(sol.t[i])
end

rel_error = (perf_mes.-NS_perf_mes)./NS_perf_mes

plot(sol.t[1:288], rel_error,
    title="PSA versus NS Relative Error", legend=:topleft,
    label="Abs (PSA-NS)/NS", xlabel="Time (Minutes)",
    ylabel="Abs (PSA-NS)/NS",right_margin=20mm)
plot!(NaN.*[1:288], label="rho(t)", grid=false,
    legend=:topleft,linecolor=:green)
plot!(twinx(), sol.t[1:288], rhos, label="rho(t)", legend=false,linecolor=:green, ylabel="rho(t)")
savefig(string(model_name,"PSA_vs_NS_relative_error_USBank.png"))

plot(sol.t[1:288], perf_mes, title="PSA and NSA Actual",
    label="PSA", xlabel="Time (Minutes)",yscale =:log,
        ylabel="Expected Number Waiting",
        right_margin=20mm, legend =:topleft)
plot!(sol.t[1:288], NS_perf_mes, label="NS", xlabel="Time (Minutes)",right_margin=20mm)
plot!(NaN.*[1:288], label="rho(t)", grid=false, legend=:left,linecolor=:green)
plot!(twinx(), sol.t[1:288], rhos, label="rho(t)", legend=false,linecolor=:green, ylabel="rho(t)")
savefig(string(model_name,"PSA_vs_NS_Actual_Values"))

# Also plot the different in measures in 2-norm.

# Same thing but for a simple sin M/M/s queue
# We look at the expected delay on such a thing.

model_name = "Queue_Saturated"

p = Dict("forwards" => true, "S"=> 600, "s"=>330, "mu"=> .333)
p["sin_base"] = .8*p["mu"]*p["s"]
p["sin_rel_amp"] = .1
p["phase_shift"] = 0
p["sin_rate"] = 2*pi/1440.0

u_0 = get_stat_dist_Q(Finite_Server_Q(0,p))
prob = ODEProblem(Finite_Server_dudt, u_0, (0.0, 1440.0),p)
sol = solve(prob, CVODE_BDF(), reltol=1e-12, abstol=1e-12, saveat=5.0)

NS_perf_mes = zeros(288)
rhos = zeros(288)
for i in range(1,stop=288)
    if ~ is_stochastic_vector(sol[i])
        sol[i][sol[i].<0].=0
        if ~ is_stochastic_vector(sol[i])
            print("Not stochastic after shift.")
        end
    end
    NS_perf_mes[i] = sol[i]' * Finite_Server_Number_Waiting_r(p)
    rhos[i] = lambda_func(sol.t[i], p)/(p["mu"]*p["s"])
end

perf_mes = zeros(288)
for i in range(1,stop=288)
    perf_mes[i] = (bd_stat(Finite_Server_Q(sol.t[i],p)))' * Finite_Server_Number_Waiting_r(p)
end

rel_error = (perf_mes.-NS_perf_mes)./NS_perf_mes

plot(sol.t[1:288], rel_error, title="PSA versus NS Relative Error", label="(PSA-NS)/NS",
        xlabel="Time (Minutes)", ylabel="(PSA-NS)/NS",right_margin=20mm)
plot!(NaN.*[1:288], label="rho(t)", grid=false, legend=:left,linecolor=:green)
plot!(twinx(), sol.t[1:288], rhos, label="rho(t)", legend=false,linecolor=:green, ylabel="rho(t)")
savefig(string(model_name,"PSA_vs_NS_relative_error_large_rho_USBankLevel.png"))

plot(sol.t[1:288], perf_mes, title="PSA and NS Actual",
 label="PSA", xlabel="Time (Minutes)", ylabel="Expected Number Waiting",
 right_margin=20mm)
plot!(sol.t[1:288], NS_perf_mes, label="NS")
plot!(NaN.*[1:288], label="rho(t)", grid=false,
    legend=:topleft, linecolor=:green)
plot!(twinx(), sol.t[1:288], rhos,
        label="rho(t)", legend=false,linecolor=:green, ylabel="rho(t)")
savefig(string(model_name,"PSA_vs_NS_Actual_Values_large_rho_USBankLevel.png"))


Random.seed!(0)

# facts("qrdelcol") do
    m = 100
    A = randn(m,m)
    Q, R = qr(A)
    for i in 100:-1:1
        k = rand(1:i)
        global A = A[:,1:i .!= k]
        global R = qrdelcol(R, k)
        @assert norm( R'*R - A'*A ) < 1e-5
    end
# end


#!/usr/bin/env julia

export pe5

pe5()=lcm(1:20)

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe5()|>print
end


"""A package for graph computations related to graph Laplacians

Graphs are represented by sparse adjacency matrices, etc.
"""
module Laplacians

  using LinearAlgebra, SparseArrays, Statistics

  include("approxCholTypes.jl")
  include("approxChol.jl")
  include("graphOps.jl")
  include("graphUtils.jl")

  export adj, lap, extendMatrix, approxchol_lap_pc, LDLsolver!, LDLinv

end # module Laplacians.jl


function default_in_partition(sites::Tuple{Int}, p::Integer, nparts::Integer)
  return p == mod1(sites[1], nparts)
end

# i ≥ j
_f(i, j) = (i - 1) * i ÷ 2 + j
function default_in_partition(sites::NTuple{2,Int}, p, nparts)
  i, j = sites
  if i ≤ j
    return p == mod1(_f(j, i), nparts)
  end
  return p == mod1(_f(i, j), nparts)
end

function default_in_partition_sorted(sites::NTuple{4,Int}, p, nparts)
  i, j, k, l = sites
  if j == k
    return default_in_partition((j,), p, nparts)
  end
  return default_in_partition((i, j), p, nparts)
end

function default_in_partition(sites::NTuple{4,Int}, p, nparts)
  return default_in_partition_sorted(sort(sites), p, nparts)
end


include("../functions/Functions.jl")
include("../evaluate/Evaluate.jl")
include("../opt/pgd.jl")
include("../opt/agd.jl")
include("../opt/GD-const.jl")
include("../data/data.jl")
include("../plot/Plotting.jl")

using Plotting

X, y = getOldData()

x0 = zeros(size(X,2),1)

lambda = 0.000001

iter = 50
step = 0.01

f = f_logreg(X,y,x0,lambda)
g = g_logreg(X,y,x0,0)

g_l2 = g_logreg(X,y,x0,lambda)

prox_operator = SqrNormL2(lambda/2)

@time (x_pgd,state,val_array_pgd,stop_array_pgd) = pgd(X,y,x0,f,g,1e-8,lambda,iter,step,size(X,1),prox_operator)
@time (x_agd,state,val_array_agd,stop_array_agd) = agd(X,y,x0,f,g_l2,1e-8,lambda,iter,step,size(X,1),prox_operator)
@time (x_gd,state,val_array_gd,stop_array_gd) = gd(X,y,x0,f,g_l2,1e-8,lambda,iter,step,size(X,1))

plotTripleArray(val_array_pgd,val_array_agd,val_array_gd, "GD-FVAL-COMPARE.svg")
plotTripleArray(stop_array_pgd,stop_array_agd,stop_array_gd, "GD-STOP-COMPARE.svg")