using Manifolds, Manopt, Test, ManifoldsBase

using Random
Random.seed!(42)
# Test the additional manifold functions
#
@testset "Additional Manifold functions" begin
    @testset "mid point & reflect" begin
        M = Sphere(2)
        p = [1.0, 0.0, 0.0]
        q = [0.0, 1.0, 0.0]

        r = mid_point(M,p,q)
        r2 = similar(r)
        mid_point!(M,r2,p,q)
        r3 = shortest_geodesic(M,p,q,0.5)
        r4 = mid_point(M,p,q,q)
        @test isapprox(M,r,r2)
        @test isapprox(M,r2,r3)
        @test isapprox(M,r3,r4)
        r5 = similar(r4)
        mid_point!(M,r5,p,q,q)
        @test isapprox(M,r4,r5)

        r4 = mid_point(M,p,-p,q)
        r5 = similar(r4)
        mid_point!(M,r5,p,-p,q)
        @test isapprox(M,r4,q)
        @test isapprox(M,r4,r5)

        @test isapprox(M,reflect(M,p,q),-q)
        f = x->x
        @test reflect(M,f,q) == q

        M2 = Euclidean(2)
        p2 = [1.0, 0.0]
        q2 = [0.0, 1.0]
        s = mid_point(M2,p2,q2)
        s2 = similar(s)
        mid_point!(M2,s2,p2,q2)
        @test s==s2
        @test s==(p2+q2)/2
        s = mid_point(M2,p2,q2,s)
        @test s==s2
        s2 = similar(s)
        mid_point!(M2,s2,p2,q2,s)
        @test s==s2

    end
    @testset "random" begin
        Mc = Circle()
        pc = random_point(Mc)
        @test is_manifold_point(Mc, pc, true)
        Xc = random_tangent(Mc,pc)
        @test is_tangent_vector(Mc, pc, Xc, true)

        Me = Euclidean(3)
        pe = random_point(Me)
        @test is_manifold_point(Me, pe, true)
        Xe = random_tangent(Me,pe)
        @test is_tangent_vector(Me, pe, Xe, true)

        Mg = Grassmann(3,2)
        pg = random_point(Mg)
        @test is_manifold_point(Mg, pg, true)
        Xg = random_tangent(Mg,pg)
        @test is_tangent_vector(Mg, pg, Xg, true; atol=10^(-14))

        Mp = ProductManifold(Mg,Me)
        pp = random_point(Mp)
        @test is_manifold_point(Mp, pp, true)
        Xp = random_tangent(Mp,pp)
        @test is_tangent_vector(Mp, pp, Xp,true; atol=10^(-15))

        Mp2 = PowerManifold(Me, NestedPowerRepresentation(),4)
        pp2 = random_point(Mp2)
        @test is_manifold_point(Mp2,pp2)
        Xp2 = random_tangent(Mp2,pp2)
        @test is_tangent_vector(Mp2,pp2,Xp2)

        Mp3 = PowerManifold(Me, ArrayPowerRepresentation(), 4)
        pp3 = random_point(Mp3)
        @test is_manifold_point(Mp3,pp3)
        Xp3 = random_tangent(Mp3,pp3)
        @test is_tangent_vector(Mp3,pp3,Xp3)

        Mr = Rotations(3)
        pr = random_point(Mr)
        @test is_manifold_point(Mr,pr)
        Xr = random_tangent(Mr,pr)
        @test is_tangent_vector(Mr,pr,Xr)

        Mspd = SymmetricPositiveDefinite(3)
        pspd = random_point(Mspd)
        @test is_manifold_point(Mspd,pspd;atol=10^(-14))
        Xspd = random_tangent(Mspd,pspd)
        @test is_tangent_vector(Mspd,pspd,Xspd, true; atol=10^(-15))
        Xspd2 = random_tangent(Mspd,pspd, Val(:Rician))
        @test is_tangent_vector(Mspd,pspd,Xspd2,true;atol=10^(-15))

        Mst = Stiefel(3,2)
        pst = random_point(Mst)
        @test is_manifold_point(Mst, pst, true)

        Msp = Sphere(2)
        psp = random_point(Msp)
        @test is_manifold_point(Msp, psp, true)
        Xsp = random_tangent(Msp,psp)
        @test is_tangent_vector(Msp,psp,Xsp,true;atol=10^(-15))

        Mh = Hyperbolic(2)
        ph = [0.0,0.0,1.0]
        Xh = random_tangent(Mh, ph)
        @test is_tangent_vector(Mh, ph, Xh, true)
    end
end


"""
    ZDT1(D, n_solutions)

ZDT1 returns `(f::function, bounds::Matrix{Float64}, pareto_set::Array{xFgh_indiv})`
where `f` is the objective function and `pareto_set` is an array with optimal Pareto solutions
with `n_solutions`.

### Parameters
- `D` number of variables (dimension)
- `n_solutions` number of pareto solutions.

Main properties:
- convex
"""
function ZDT1(D = 30, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1) )
        return ( [x[1], gx*(1 - sqrt(x[1] / gx))  ], [0.0], [0.0] )
    end

    bounds = Array([zeros(D) ones(D)]')

    x = range(0, 1, length=n_solutions)

    X = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    pareto_set = [ generateChild(x, f(x)) for x in X ]

    return f, bounds, pareto_set
end

"""
    ZDT2(D, n_solutions)

ZDT2 returns `(f::function, bounds::Matrix{Float64}, pareto_set::Array{xFgh_indiv})`
where `f` is the objective function and `pareto_set` is an array with optimal Pareto solutions
with `n_solutions`.

### Parameters
- `D` number of variables (dimension)
- `n_solutions` number of pareto solutions.

Main properties:
- nonconvex
"""
function ZDT2(D = 30, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1) )
        return ( [x[1], gx*(1 - (x[1] / gx)^2)  ], [0.0], [0.0] )
    end
    bounds = Array([zeros(D) ones(D)]')

    x = range(0, 1, length=n_solutions)

    x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    pareto_set = [ generateChild(x, f(x)) for x in x ]

    return f, bounds, pareto_set
end

"""
    ZDT3(D, n_solutions)

ZDT3 returns `(f::function, bounds::Matrix{Float64}, pareto_set::Array{xFgh_indiv})`
where `f` is the objective function and `pareto_set` is an array with optimal Pareto solutions
with `n_solutions`.

### Parameters
- `D` number of variables (dimension)
- `n_solutions` number of pareto solutions.

Main properties:
- convex disconected
"""
function ZDT3(D = 30, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1) )
        a = x[1] / gx
        return ( [x[1], gx*(1 - sqrt(a) - a*sin(10π*x[1]) )  ], [0.0], [0.0] )
    end
    bounds = Array([zeros(D) ones(D)]')

    regions = [ 0 0.0830015349
                0.182228780 0.2577623634
                0.4093136748 0.4538821041
                0.6183967944 0.6525117038
                0.8233317983 0.8518328654]

    n = n_solutions ÷ size(regions, 1)
    x = Float64[]
    for i in 1:size(regions, 1)
        x = vcat(x, range(regions[i,1], regions[i,2], length=n))
    end

    x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    pareto_set = [ generateChild(x, f(x)) for x in x ]

    return f, bounds, get_non_dominated_solutions(pareto_set)

end

"""
    ZDT4(D, n_solutions)

ZDT4 returns `(f::function, bounds::Matrix{Float64}, pareto_set::Array{xFgh_indiv})`
where `f` is the objective function and `pareto_set` is an array with optimal Pareto solutions
with `n_solutions`.

### Parameters
- `D` number of variables (dimension)
- `n_solutions` number of pareto solutions.

Main properties:
- nonconvex
"""
function ZDT4(D = 10, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 10*(length(x)-1) + sum( x[2:end].^2 - 10cos.(4π*x[2:end]))
        return ( [x[1], gx*(1 - sqrt(x[1] / gx))  ], [0.0], [0.0] )
    end
    bounds = Array([-5zeros(D) 5ones(D)]')
    bounds[:,1] = [0, 1.0]

    x = range(0, 1, length=n_solutions)

    x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    pareto_set = [ generateChild(x, f(x)) for x in x ]

    return f, bounds, pareto_set
end

"""
    ZDT6(D, n_solutions)

ZDT6 returns `(f::function, bounds::Matrix{Float64}, pareto_set::Array{xFgh_indiv})`
where `f` is the objective function and `pareto_set` is an array with optimal Pareto solutions
with `n_solutions`.

### Parameters
- `D` number of variables (dimension)
- `n_solutions` number of Pareto solutions.

Main properties:
- nonconvex
- non-uniformly spaced
"""
function ZDT6(D = 10, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1.0) )^(0.25)
        ff1 = 1.0 - exp(-4.0x[1])*sin(6.0π*x[1])^6
        return ( [ ff1 , gx*(1.0 - (ff1 / gx)^2)  ], [0.0], [0.0] )
    end

    bounds = Array([zeros(D) ones(D)]')

    #x = range(0, 1, length=n_solutions)

    #x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    xx = range(0.2807753191, 1, length=100)
    yy = 1 .- (xx).^2
    pareto_set = [ generateChild(zeros(0), ([xx[i], yy[i]], [0.0], [0.0])) for i in 1:length(xx) ]

    return f, bounds, pareto_set
end


using CollisionDetection
using StaticArrays
using JLD
using Base.Test

fn = normpath(joinpath(dirname(@__FILE__),"center_sizes.jld"))
d = load(fn)

tmp = d["ctrs"]
ctrs = [SVector(q...) for q in tmp]
rads = d["rads"]

tree = Octree(ctrs, rads)

# extract all the triangles that (potentially) intersect octant (+,+,+)
pred(i) = all(ctrs[i].+rads[i] .> 0)
bb = SVector(0.5, 0.5, 0.5), 0.5
ids = collect(searchtree(pred, tree, bb))
@test length(ids) == 154

# using MATLAB
# A = [c[i] for i in 1:3, c in ctrs].'
# @mput A rads
# @matlab begin
#     plot3(A[:,1],A[:,2],A[:,3],".")
# end
#
# B = [ctrs[j][i] for i in 1:3, j in ids].'
# @mput B
# @matlab begin
#     hold("on")
#     plot3(B[:,1],B[:,2],B[:,3],"or")
# end


using Dierckx, QuadGK

function int_tabulated(λ, f; int_k=3)
    @assert issorted(λ)
    itp = Spline1D(λ, f, k=int_k, bc="error")
    return quadgk(itp, λ[1], λ[end])
end


@testset "Basis functions" begin
    rₘₐₓ = 10
    ρ = rand()
    N = ceil(Int, rₘₐₓ/ρ + 1/2)

    for B in [StaggeredFiniteDifferences(N, ρ), FiniteDifferences(N, 1.0)]
        x = locs(B)
        χ = B[x, :]
        @test χ isa AbstractSparseMatrix
        @test χ ≈ Diagonal(ones(N))
    end

    @testset "Basis functions in restricted bases" begin
        N = 100
        L = 20.0

        B = FiniteDifferences(range(-L, stop=L, length=N+2)[2:end-1])
        xx = axes(B,1)
        a = xx.domain
        x = range(a.left, stop=a.right, length=max(N,101))
        χ = B[x,:]
        @test χ isa AbstractSparseMatrix

        sell = 1:round(Int, 0.6N)
        selv = round(Int, 0.4N):N

        Bl = B[:,sell]
        Bv = B[:,selv]

        χl = Bl[x, :]
        @test χl isa AbstractSparseMatrix
        @test χl == χ[:,sell]
        χv = Bv[x, :]
        @test χv == χ[:,selv]

        @test Bv[x, 4:8] == χ[:,selv[4:8]]
    end
end


module Drawing

using CoilFunctions, Unitful, Luxor

function drawwire(p::Point, awg::AWG, fillfactor=1.0;u=u"mm")
    f = x->ustrip(u(x))
    setcolor("green")
    circle(p,f(awg.d_insul)/2,:fill)
    setcolor("blue")
    circle(p,f(awg.d_bare)/2,:fill)
    # setcolor("black")
    # circle(p,f(awg.d_insul)/sqrt(fillfactor)/2, :stroke)
    # @info f(awg.d_insul)
end

function drawcoil(awg::AWG,winding::WindingLayers,maxturns::Integer,fillfactor::Real=1.0)
    turns = zero(maxturns)
    d = awg.d_insul/sqrt(fillfactor) |> u"mm" |> ustrip
    for (layer,y) = enumerate(range(d/2;step=d*√(3)/2,length=winding.layers))
        (turnlayers,xoffset) = isodd(layer) ? (winding.odd_layer,d/2) : (winding.even_layer,d)
        for x in range(xoffset;step=d,length=turnlayers)
            turns >= maxturns && break
            drawwire(Point(x,y),awg,fillfactor)
            turns += 1
        end
    end
end
drawcoil(a,fillfactor) = drawcoil(CoilFunctions.AWG_Chart[a.gauge], a.winding, a.turns, fillfactor)

function drawcoil(coil::CoilGeometry)
    y=(coil.od - coil.id)/2 |> u"mm" |> ustrip
    x=ustrip(u"mm"(coil.len))
    gsave()
    setcolor("black")
    box(Point(0,0),Point(x,y), :stroke)
    grestore()
end

function illustratecoil(coil::CoilGeometry, Ω::Unitful.ElectricalResistance, fillfactor=1.0; filename=:png, pad=20,
    resolution=(800,600), minturns=0, maxturns=typemax(Int64))

    x,y = resolution
    a=optimalcoil(coil,Ω,fillfactor)
    Drawing(x+2pad, y+2pad, filename)
    origin(pad,y+pad)
    s = ustrip(x/coil.len)
    scale(s,-s)
    background("white")
    sethue("black")
    drawcoil(coil)
    turns = clamp(a.turns,minturns, maxturns)
    drawcoil(merge(a,(turns=turns,)),a.fillfactor)
    finish()
    @info a
    preview()
end

function illustratecoil(coil::CoilGeometry, v::Unitful.Voltage, p::Unitful.Power, fillfactor=1.0; kwargs...)
    illustratecoil(coil, v^2/p |> u"Ω", fillfactor;kwargs...)
end

function illustratecoil(s::Symbol, v::Unitful.Voltage, fillfactor=1.0; kwargs...)
    coil=CoilFunctions.coil[s]
    Ω=CoilFunctions.Ω(s,v)
    illustratecoil(coil,Ω,fillfactor;kwargs...)
end

end


using Revise
using FwiFlow
using AdFem
using ADCME
using PyPlot
using PyCall
using Statistics
using DelimitedFiles
using ADCMEKit
np = pyimport("numpy")

function visualize_obs(o1)
    close("all")
    plot((0:NT)*Δt, o1[:, 1:m+1])
    xlabel("t")
    ylabel("displacement")
end

function layer_model(m, n)
    z = zeros(4*n*m)
    k = 0
    mm = ones(n)
    mm[n÷3:end] .= 1.5
    for i = 1:n
        for j = 1:m
            for p = 1:2
                for q = 1:2
                    k += 1
                    z[k] = mm[i]
                end
            end
        end
    end
    return z
end

function sin_model(m, n)
end

function visualize_invη(o)
    o = o * 1e12
    Z = zeros(n, m)
    k = 0
    for i = 1:n
        for j = 1:m
            Z[i, j] = mean(o[k+1:k+4])
            k += 4
        end
    end
    pcolormesh(Z)
    axis("scaled")
    xlabel("x")
    ylabel("y")
    gca().invert_yaxis()
    colorbar()
end


# Some simple tests on the entropy/information code

using SubmodularMaximization

grid_size = 10

grid = Grid(grid_size, grid_size)

sensor = RangingSensor()
histogram_filter = Filter(grid)

# Test entropy on uniform prior
@show grid_entropy = entropy(histogram_filter)
real_entropy = -log(2, 1/num_states(grid))
if !isapprox(grid_entropy, real_entropy; rtol=1e-3)
  println("Incorrect entropy for uniform prior")
end


using BesselK, SpecialFunctions, ArbNumerics
include("shared.jl")

function wbesselkxv(v,x)
  try
    return BesselK.adbesselkxv(v,x)
  catch
    return NaN
  end
end

function rbesselkxv(v,x)
  _v = ArbReal(v)
  _x = ArbReal(x)
  xv = _x^_v
  Float64(xv*ArbNumerics.besselk(ArbFloat(v), ArbFloat(x)))
end
abesselkxv(v,x) = (x^v)*SpecialFunctions.besselk(v, x)

const VGRID = range(0.25, 10.0, length=100)
const XGRID = range(0.0,   8.0, length=201)[2:end] # since other impls can't do x=0.

const BASELINE = [rbesselkxv(z[1], z[2]) for z in Iterators.product(VGRID, XGRID)]
const AMOS     = [abesselkxv(z[1], z[2]) for z in Iterators.product(VGRID, XGRID)]
const OURSOL   = [wbesselkxv(z[1], z[2]) for z in Iterators.product(VGRID, XGRID)]
const TOLS_A   = atolfun.(zip(BASELINE, AMOS))
const TOLS_U   = atolfun.(zip(BASELINE, OURSOL))
const TOLS_AU  = rtolfun.(zip(AMOS,     OURSOL))

#=
gnuplot_save_matrix!("../plotdata/atols_amos.csv",      TOLS_A,  VGRID, XGRID)
gnuplot_save_matrix!("../plotdata/atols_ours.csv",      TOLS_U,  VGRID, XGRID)
gnuplot_save_matrix!("../plotdata/rtols_amos_ours.csv", TOLS_AU, VGRID, XGRID)
=#


__precompile__()

module Projections

export project, project!

abstract type ConvexSet end
abstract type NonconvexSet end

"""
    project!(s, v, y)

Overwrite `v` by projection of point `y` to the box set `b`.
"""
function project!(); end

"""
    project(s, y)

Project point `y` to the box set `b`.
"""
function project(); end

include("box.jl")
include("ball.jl")
include("simplex.jl")
include("affine.jl")

end


# fbp2_sino_filter.jl

export fbp2_sino_filter

using FFTW

"""
    sino, Hk, hh, nn = fbp2_sino_filter(how, sino;
    ds=1, dsd=0, extra=0, npad=0, decon1=1, window=:none)

Apply ramp-like filters to sinogram(s) for 2D FBP image reconstruction.
Both parallel-beam and fan-beam tomographic geometries are supported.
This approach of sampling the band-limited ramp avoids the aliasing that
would be caused by sampling the ramp directly in the frequency domain.

in
- `how::Symbol`
  `:arc` (3rd generation CT) or `:flat` (for parallel too)
- `sino::AbstractArray{<:Number}` : `[nb (L)]` sinograms

options
- `ds::RealU` : sample spacing (in distance units, e.g., cm) (default: 1)
- `dsd::RealU` : source-to-detector distance, for `:arc` case only. (default: Inf)
- `extra::Int` : # of extra sinogram radial samples to keep (default: 0)
- `npad::Int` : # of padded samples. (default: 0, means next power of 2)
- `decon1::Int` : deconvolve effect of linear interpolator? (default: 1)
- `window::Symbol` : default: `:none`

out
- `sino::AbstractArray{<:Number}`       filtered sinogram rows
- `Hk::AbstractVector{<:Number}`        apodized ramp filter frequency response
- `hn::AbstractVector{<:Number}`        samples of band-limited ramp filter
- `nn::AbstractVector{<:Number}`        [-np/2,...,np/2-1] vector for convenience

"""
function fbp2_sino_filter(
    sg::SinoGeom,
    sino::AbstractMatrix{<:Number};
    ds::RealU = 1,
    dsd::RealU = Inf,
    extra::Int = 0,
    npad::Int = 0,
    decon1::Int = 1,
    window::Symbol = :none
)

    nb, na = size(sino)
    if npad == 0
        npad = 2^ceil(Int64, log2(2*nb-1)) # padded size
    end

    sino = [sino; zeros(npad-nb,na)] # padded sinogram
    hn, nn = fbp_ramp(sg, npad, ds, dsd)

    reale = (x) -> (@assert x ≈ real(x); real(x))
    Hk = fft(fftshift(hn))
    Hk = reale(Hk)

    Hk = Hk .* fbp2_window(npad, window)
    Hk = ds * Hk # differential for discrete-space convolution vs integral

    #= linear interpolation is like blur with a triangular response,
    so we can compensate for this approximately in frequency domain =#
    if decon1 != 0
        Hk = Hk ./ fftshift(sinc.(nn / npad).^2)
    end

    sino = reale(ifft(fft(sino, 1) .* repeat(Hk, 1, na), 1)) # apply filter
    #NOTE: was fft(sino, [], 1) and ifft(..., [], 1) in matlab
    #NOTE: still assertion error for reale
    # todo: definitely use broadcast or map here.  should be no need to repeat
    # trick: possibly keep extra column(s) for zeros!
    sino = sino[1:(nb+extra),:]
    sino[(nb+1):(nb+extra),:] .= 0
    sino = reshape(sino, nb, na)

    return sino, Hk, hn, nn
end

function fbp2_sino_filter(how::Symbol, sino::AbstractArray{<:Number}; kwargs...)
    return mapslices(sino -> fbp2_sino_filter(how, sino; kwargs...), sino, [1,2])
end


using Test, SciMLBase, SciMLBase.EnsembleAnalysis

EA = SciMLBase.EnsembleAnalysis

# tests for https://github.com/SciML/DifferentialEquations.jl/issues/731
# make sure integer inputs work
@test all(EA.componentwise_mean([[1,1],[2,2],[3,3]]) .≈ [2.0,2.0])
m,v = EA.componentwise_meanvar([[1,1],[2,2],[3,3]])
@test all(m .≈ 2.0)
@test all(v .≈ 1.0)
mx,my,C = EA.componentwise_meancov([[1,1],[2,2],[3,3]], [[3,3],[2,2],[1,1]])
@test all(mx .≈ 2.0)
@test all(my .≈ 2.0)
@test all(C .≈ -1.0)
mx,my,C = EA.componentwise_weighted_meancov([[1,1],[2,2],[3,3]], [[3,3],[2,2],[1,1]], [[1,1],[2,2],[1,1]])
@test all(mx .≈ 2.0)
@test all(my .≈ 2.0)
@test all(C .≈ -.8)


let

    # fitting noisy data to an exponential model
    model(x, p) = @. p[1] * exp(-x * p[2])
    #model(x,p) = [p[1],100*(p[2]-p[1]^2)]
    # some example data
    Random.seed!(12345)
    xdata = range(0, stop=10, length=50_000)
    ydata = model(xdata, [1.0, 2.0]) + 0.01*randn(length(xdata))
    p0 = [.5, .5]

    function jacobian_model(x,p)
        J = Array{Float64}(undef, length(x), length(p))
        @. J[:,1] = exp(-x*p[2])     #dmodel/dp[1]
        @. @views J[:,2] = -x*p[1]*J[:,1]
        J
    end

    #Generating the "automatic" avv
    model_inplace(out, p) =  @. out = p[1] * exp(-xdata * p[2])

    hessians = Array{Float64}(undef, length(xdata)*length(p0), length(p0))
    h! = make_hessian(model_inplace, xdata, p0)
    auto_avv! = Avv(h!, length(p0), length(xdata))

    # a couple notes on the Avv function:
    # - the basic idea is to see the model output as simply a collection of functions: f1...fm
    # - then Avv return an array of size m, where each elements corresponds to
    # v'H(p)v, with H an n*n Hessian matrix of the m-th function, with n the size of p
    function manual_avv!(dir_deriv,p,v)
        v1 = v[1]
        v2 = v[2]
        for i=1:length(xdata)
            #compute all the elements of the H matrix
            h11 = 0
            h12 = (-xdata[i] * exp(-xdata[i] * p[2]))
            #h21 = h12
            h22 = (xdata[i]^2 * p[1] * exp(-xdata[i] * p[2]))

            # manually compute v'Hv. This whole process might seem cumbersome, but
            # allocating temporary matrices quickly becomes REALLY expensive and might even
            # render the use of geodesic acceleration terribly inefficient
            dir_deriv[i] = h11*v1^2 + 2*h12*v1*v2 + h22*v2^2

        end
    end

    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=1); #warmup
    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=1, avv! = manual_avv!,lambda=0, min_step_quality = 0); #lambda = 0 to match Mark's code
    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=10, avv! = auto_avv!,lambda=0, min_step_quality = 0)

    println("--------------\nPerformance of curve_fit vs geo")

    println("\t Non-inplace")
    fit = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=1000)
    @test fit.converged

    println("\t Geodesic")
    fit_geo = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=10, avv! = manual_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo.converged

    println("\t Geodesic - auto avv!")
    fit_geo_auto = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=10, avv! = auto_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo_auto.converged

    @test maximum(abs.(fit.param-[1.0, 2.0])) < 1e-1
    @test maximum(abs.(fit.param-fit_geo.param)) < 1e-6
    @test maximum(abs.(fit.param-fit_geo_auto.param)) < 1e-6

    #with noise
    yvars = 1e-6*rand(length(xdata))
    ydata = model(xdata, [1.0, 2.0]) + sqrt.(yvars) .* randn(length(xdata))

    #warm up
    curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, p0; maxIter=1)
    curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=1, avv! = manual_avv!,lambda=0, min_step_quality = 0)
    curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=1, avv! = auto_avv!,lambda=0, min_step_quality = 0)

    println("--------------\nPerformance of curve_fit vs geo with weights")

    println("\t Non-inplace")
    fit_wt = @time curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, p0; maxIter=100)
    @test fit_wt.converged

    println("\t Geodesic")
    fit_geo_wt = @time curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=100, avv! = manual_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo_wt.converged

    println("\t Geodesic - auto avv!")
    fit_geo_auto_wt = @time curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=100, avv! = auto_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo_auto_wt.converged

    @test maximum(abs.(fit_wt.param-[1.0, 2.0])) < 1e-1
    @test maximum(abs.(fit_wt.param-fit_geo_wt.param)) < 1e-6
    @test maximum(abs.(fit_wt.param-fit_geo_auto_wt.param)) < 1e-6

end


using LinearAlgebra
using CompScienceMeshes
using SauterSchwab3D

const pI   = point(6,5,3)
const pII  = point(5,2,4)
const pIII = point(7,1,0)
const pIV  = point(3,4,0)

const P = simplex(pI,pII,pIII,pIV)
const Q =  simplex(pI,pII,pIII,pIV)

const sing = SauterSchwab3D.singularity_detection(P,Q)

Accuracy2 = 18
ct_ref = CommonVolume6D(sing,SauterSchwab3D._legendre(Accuracy2,0.0,1.0))

function integrand(x,y)
   ((x-pI)'*(y-pI))*return(exp(-im*1*norm(x-y))/(4pi*norm(x-y)))
end

function INTEGRAND(u,v)
   if u[3]< 0.0 || u[3]>1-u[1]-u[2] || u[2]< 0.0 || u[2]>1-u[1] || u[1]<0.0 || u[1]>1.0
      println("u: ",u)
   end
   if v[3]< 0.0 || v[3]>1-v[1]-v[2] || v[2]< 0.0 || v[2]>1-v[1] || v[1]<0.0 || v[1]>1.0
      println("v: ",v)
   end
   n1 = neighborhood(P,u)
   n2 = neighborhood(Q,v)
   x = cartesian(n1)
   y = cartesian(n2)
   output = integrand(x,y)*jacobian(n1)*jacobian(n2)
   return(output)
end

print("Ref: ")
ref = sauterschwab_parameterized(INTEGRAND, ct_ref)
println(ref)
println()

res_tp =[]
res_sp =[]
res_gm =[]
n1 = []
n2 = []
n3 = []
for i in 2:1:15
   Accuracy = i
   ct = CommonVolume6D(sing,SauterSchwab3D._legendre(Accuracy,0.0,1.0))

   int_tp = sauterschwab_parameterized(INTEGRAND, ct)

   num_pts = 18*length(ct.qps)^6
   push!(n1,num_pts)

   push!(res_tp,int_tp)
end

for i in 2:1:6
   Accuracy = i
   ct_s = CommonVolume6D_S(sing,(SauterSchwab3D._legendre(Accuracy,0.0,1.0),
                          SauterSchwab3D._shunnham2D(Accuracy),
                          SauterSchwab3D._shunnham4D(Accuracy)))

   int_sp = sauterschwab_parameterized(INTEGRAND, ct_s)

   num_pts = 2*length(ct_s.qps[1])^2*length(ct_s.qps[3])+
            16*length(ct_s.qps[3])*length(ct_s.qps[2])
   push!(n2,num_pts)
   push!(res_sp,int_sp)
end

for i in 2:1:12
   Accuracy = i
   ct_gm = CommonVolume6D_S(sing,(SauterSchwab3D._legendre(Accuracy,0.0,1.0),
                          SauterSchwab3D._grundmannMoeller2D(2*Accuracy-1),
                          SauterSchwab3D._grundmannMoeller4D(2*Accuracy-1)))

   int_gm = sauterschwab_parameterized(INTEGRAND, ct_gm)

   num_pts = 2*length(ct_gm.qps[1])^2*length(ct_gm.qps[3])+
            16*length(ct_gm.qps[3])*length(ct_gm.qps[2])
   push!(n3,num_pts)
   push!(res_gm,int_gm)
end

err_tp = norm.(res_tp.-ref)/norm(ref)
err_sp = norm.(res_sp.-ref)/norm(ref)
err_gm = norm.(res_gm.-ref)/norm(ref)

println(err_tp)
println(err_sp)
println(err_gm)

using Plots
plot( yaxis=:log, xaxis=:log, fontfamily="Times" )
plot!(n1,err_tp, label="Gauss Tensor-Product",markershape=:circle)
plot!(n2,err_sp, label="Simplex-Product SH",markershape=:rect)
plot!(n3,err_gm, label="Simplex-Product GM",markershape=:x)
plot!(xlims=(1e2,1e7),ylims=(1e-8,1))
plot!(xlabel="#Quad. Pts./Func. Evals", ylabel="Rel. Error.", title="Identical Tetrahedron 6D",legend=:bottomleft)

using BenchmarkTools

ref = 1.749281958454714 - 4.409866361065269im

ct = CommonVolume6D(sing,SauterSchwab3D._legendre(5,0.0,1.0))

int_tp = sauterschwab_parameterized(INTEGRAND, ct)

num_pts = 18*length(ct.qps)^6
println("#Pts: ",num_pts)
err_tp = norm.(int_tp.-ref)/norm(ref)
println("Rell Err; ",err_tp)
@time for i in 1:100 sauterschwab_parameterized(INTEGRAND, ct) end

ct_s = CommonVolume6D_S(sing,(SauterSchwab3D._legendre(5,0.0,1.0),
                         SauterSchwab3D._shunnham2D(5),
                         SauterSchwab3D._shunnham4D(5)))

int_sp = sauterschwab_parameterized(INTEGRAND, ct_s)

num_pts = 2*length(ct_s.qps[1])^2*length(ct_s.qps[3])+
          16*length(ct_s.qps[3])*length(ct_s.qps[2])
println("#Pts: ",num_pts)
err_sp = norm.(int_sp.-ref)/norm(ref)
println("Rel Err: ",err_sp)
@time for i in 1:100 sauterschwab_parameterized(INTEGRAND, ct_s) end

#=
ct_gm = CommonVolume6D_S(sing,(SauterSchwab3D._legendre(7,0.0,1.0),
                         SauterSchwab3D._grundmannMoeller2D(2*7-1),
                         SauterSchwab3D._grundmannMoeller4D(2*7-1)))

int_gm = sauterschwab_parameterized(INTEGRAND, ct_gm)

num_pts = 2*length(ct_gm.qps[1])^2*length(ct_gm.qps[3])+
          16*length(ct_gm.qps[3])*length(ct_gm.qps[2])
println("#Pts: ",num_pts)
err_gm = norm.(int_gm.-ref)/norm(ref)
println("Rel Err: ",err_gm)

=#


# Personalised Movie Recommender System. JJ.Piao Alessandro Previero

# Import required packages -----------------------------------------------------
using Gurobi, DelimitedFiles, LinearAlgebra, CSV, PyPlot, Random;
using DataFrames, Distributions, MatrixImpute;
plt.style.library["ggplot"]; plt.style.use("ggplot")

# Read in the data -------------------------------------------------------------
movies = CSV.read("movies.csv", header = true)
links = CSV.read("links.csv", header = true)
ratings = CSV.read("ratings.csv", header = true)
tags = CSV.read("tags.csv", header = true)

# Report statistics on the missing entries -------------------------------------
sparseMovies = unstack(ratings, :userId, :movieId, :rating)
colMiss = [count(x->x === missing, sparseMovies[:,i])
          for i =1:size(sparseMovies,2)]./size(sparseMovies,1)*100
rowMiss = [count(x->x === missing, sparseMovies[i,:])
          for i =1:size(sparseMovies,1)]./size(sparseMovies,2)*100

plt.hist(colMiss, bins=30)
xlim(100, 0)
xlabel("Percentage of missing ratings, %")
ylabel("Number of Movies")
gcf()

plt.hist(rowMiss, bins = 30)
xlim(100, 0)
xlabel("Percentage of missing ratings, %")
ylabel("Number of Users")
gcf()

clf()

# PIPELINE
# Output: given a user with some previously seen movies, recommend top 10 movies

# Input: SparseMovie matrix, and MovieFeature matrix

# STEPS
# 1. Perform matrix completion using SparseMovie and MovieFeature
# 2. Develop naïve baseline model, sorting the top 10 scores for each user

# 3. Perform clustering on each user => need to define feature dimenions (cannot
# use individual movies). While varying the number of clusters:

    # 4. Use OCTs to interpret the clusters, until ideal K centroids are defined
# 5.


# Example from Zhou et al., 2016, Stochastic Dual Integer Programming
#

using JuMP, Gurobi, SDDiP

# ===================================================
# P1
m1 = Model(solver=GurobiSolver(OutputFlag=0))
@variables(m1, begin
    x1, Bin
    x2, Bin
end)

@objective(m1, Min, x1 + x2)

# ===================================================
# P2
function Q(z1, z2)
    m2 = Model(solver=GurobiSolver(OutputFlag=0))
    @variables(m2, begin
        y, Int
        0 <= x1 <= 1
        0 <= x2 <= 1
    end)
    @constraints(m2, begin
        y <= 4
        y >= 2.6 - 0.5x2 - 0.25x1
    end)

    c1 = @constraint(m2, x1 == z1)
    c2 = @constraint(m2, x2 == z2)

    @objective(m2, Min, 4y)
    m2, c1, c2

end
# L(π)
function L(π1, π2, z1, z2)
    m3 = Model(solver=GurobiSolver(OutputFlag=0))
    @variables(m3, begin
        y, Int
        0 <= x1 <= 1
        0 <= x2 <= 1
    end)
    @constraints(m3, begin
        y <= 4
        y >= 2.6 - 0.5x2 - 0.25x1
    end)
    @objective(m3, Min, 4y - π1 * x1 - π2 * x2)
    m3
end

# ===================================================
# Solve the first problem
solve(m1)
z1, z2 = getvalue(x1), getvalue(x2)
# ===================================================
# Pass values to the second problem
m1, c1, c2 = Q(z1, z2)
# ===================================================
# Solve the second problem using a Lagrangian solver
method = LevelMethod(initialbound = 100.0, quadsolver=GurobiSolver(OutputFlag=0))
# method = SubgradientMethod(100.0) # gap
lp     = LinearProgramData(m1.obj, [c1; c2], [100.0; 100.0], method=method)
π = [0.0; 0.0]
val, status = lagrangiansolve!(lp, m1, π)

println("With Lagrangian solver")
println("Objective at (0,0): ", getobjectivevalue(m1))
println("Duals are:")
println(m1.linconstrDuals[c1.idx])
println(m1.linconstrDuals[c2.idx])

# Note the primal is degenerate
# ===================================================
# With Benders
println("With Benders")
m2, c1, c2 = Q(z1, z2)
solve(m2, relaxation=true)
println("Objective at (0,0): ", getobjectivevalue(m2))
println("Duals are:")
println(m2.linconstrDuals[c1.idx])
println(m2.linconstrDuals[c2.idx])

# ===================================================
# With strengthened Benders
# z1, z2 = 1, 1
π1, π2 = m2.linconstrDuals[c1.idx], m2.linconstrDuals[c2.idx]
# π1, π2 = 0, -4
m3 = L(π1, π2, z1, z2)
solve(m3)
# println(getobjectivevalue(m3) - (π1 * z1 + π2 * z2))
println("Strengthened benders")
println("Objective at (0,0): ", getobjectivevalue(m3))


# This code was due to Jeffrey Sarnoff. See
# https://github.com/JuliaIntervals/IntervalArithmetic.jl/issues/372

using IntervalArithmetic
import SpecialFunctions
import ArbNumerics
using ArbNumerics: ArbFloat, ArbReal, setinterval, lowerbound, upperbound

# number of bits that suffices to alter the least significant digit
#   use: prevfloat(x, DigitBits), nextfloat(x, DigitBits)
const DigitBits = ceil(Int, log2(10))

# Interval{T} --> ArbReal, thin enclosure
function arbreal(x::Interval{T}) where T
     setprecision(ArbFloat, T)
     lo = prevfloat(ArbFloat(x.lo), DigitBits) # assure least significant
     hi = nextfloat(ArbFloat(x.hi), DigitBits) #   decimal digit changes
     return setinterval(lo, hi)
end

# T, ArbReal --> Interval{T}
#    P is the working precision from ArbNumerics
function tinterval(::Type{T}, x::ArbReal{P}) where {T, P}
     lo, hi = lowerbound(x), upperbound(x)
     return Interval{T}(T(lo), T(hi))
end

function specialfunction(fn, a::Interval{T}) where T
     x = fn(arbreal(a))
     return tinterval(T, x)
end

function specialfunction(fn, a::Interval{T}, b::Interval{T}) where T
     x = fn(arbreal(a), arbreal(b))
     return tinterval(T, x)
end

function specialfunction(fn, a::Interval{T}, b::Interval{T},
                             c::Interval{T}) where T
     x = fn(arbreal(a), arbreal(b), arbreal(c))
     return tinterval(T, x)
end

function specialfunction(fn, a::Interval{T}, b::Interval{T},
                             c::Interval{T}, d::Interval{T}) where T
     x = fn(arbreal(a), arbreal(b), arbreal(c), arbreal(d))
     return tinterval(T, x)
end

Base.setprecision(::Type{ArbFloat}, ::Type{Float64}) =
    setprecision(ArbFloat, 53)
Base.setprecision(::Type{ArbFloat}, ::Type{Float32}) =
    setprecision(ArbFloat, 24)
Base.setprecision(::Type{ArbFloat}, ::Type{BigFloat}) =
    setprecision(ArbFloat, precision(BigFloat))

# Don't import these functions
# functions of 1 arg from Base.Math (for Interval{BigFloat})
#for F in (:acos, :acosh, :acot, :acoth, :acsc, :acsch, :asec,
#          :asech, :asin, :asinh, :atan, :atanh, :cbrt, :cos,
#          :cosc, :cosh, :cospi, :cot, :coth, :csc, :csch,
#          :exp, :expm1, :log, :log10, :log1p, :log2, :sec,
#          :sech, :sin, :sinc, :sinh, :sinpi, :sqrt, :tan, :tanh)
#    @eval Base.$F(x::A) where {T, A<:Interval{BigFloat}} =
#          specialfunction($F,x)
#end

# functions of 2 arg from Base.Math (for Interval{BigFloat})
for F in (:^, :atan, :hypot)
  @eval begin
    Base.$F(x::A, y::A) where {T, A<:Interval{BigFloat}} =
        specialfunction($F, x, y)
    Base.$F(x::A, y::Real) where {T, A<:Interval{BigFloat}} =
        specialfunction($F, x, A(y))
    Base.$F(x::Real, y::A) where {T, A<:Interval{BigFloat}} =
        specialfunction($F, A(x), y)
  end
end

# functions of 1 arg exported by SpecialFunctions and ArbNumerics
for F in (:airyai, :airyaiprime, :airybi, :airybiprime,
          :besselj0, :besselj1, :bessely0, :bessely1,
          #:erf, :erfc, :erfi,
          :gamma, :lgamma, :digamma, :trigamma,
          :eta, :zeta)
    @eval SpecialFunctions.$F(x::A) where {T, A<:Interval{T}} =
          specialfunction($F,x)
end

# functions of 1 arg exported by both with distinct names
for (SF, AF) in ((:cosint, :ci), (:sinint, :si),
                 (:ellipk, :elliptic_k), (:ellipe, :elliptic_e))
    @eval SpecialFunctions.$SF(x::A) where {T, A<:Interval{T}} =
          specialfunction($AF,x)
end

# functions of 1 arg exported by ArbNumerics only
for F in (:chi, :shi, :ei, :xi,  :lambertw)
    @eval ArbNumerics.$F(x::A) where {T, A<:Interval{T}} =
          specialfunction($F,x)
end

# functions of a 2 args exported by ArbNumerics only
for F in (:elliptic_pi, :elliptic_pi2, # the 2 arg versions
          :polylog,
          :hypergeometric_0F1, :regular_hypergeometric_0F1)
    @eval ArbNumerics.$F(x::A, y::A) where {T, A<:Interval{T}} =
          specialfunction($F, x, y)
end

# functions of a 3 args exported by ArbNumerics only
for F in (:elliptic_pi, :elliptic_pi2, # the 3 arg versions
          :hypergeometric_1F1, :regular_hypergeometric_1F1)
    @eval ArbNumerics.$F(a::A, b::A, z::A) where {T, A<:Interval{T}} =
          specialfunction($F, a, b, z)
end

# functions of a 4 args exported by ArbNumerics only
for F in (:hypergeometric_2F1, :regular_hypergeometric_2F1)
    @eval ArbNumerics.$F(a::A, b::A, c::A, z::A) where
                                   {T, A<:Interval{T}} =
          specialfunction($F, a, b, c, z)
end

function __init__()
    # provide padding for prevfloat, nextfloat in arbreal()
    extrabits = ArbNumerics.ExtraBits[] + DigitBits
    ArbNumerics.setextrabits(extrabits)
    return nothing
end


import LinearAlgebra.norm
export norm_inf, norm, norm_1

# deprecate these soon
norm_inf(x::AbstractExpr) = maximum(abs(x))
norm_1(x::AbstractExpr) = sum(abs(x))
norm_fro(x::AbstractExpr) = norm2(vec(x))

# behavior of norm should be consistent with julia:
# * vector norms for vectors
# * operator norms for matrices
function norm(x::AbstractExpr, p::Real=2)
    if length(size(x)) <= 1 || minimum(size(x))==1
        # x is a vector
        if p == 1
            return norm_1(x)
        elseif p == 2
            return norm2(x)
        elseif p == Inf
            return norm_inf(x)
        elseif p > 1
            # TODO: allow tolerance in the rationalize step
            return rationalnorm(x, rationalize(Int, float(p)))
        else
            error("vector p-norms not defined for p < 1")
        end
    else
        # TODO: After the deprecation period, allow this again but make it consistent with
        # LinearAlgebra, i.e. make norm(x, p) for x a matrix the same as norm(vec(x), p).
        Base.depwarn("`norm(x, p)` for matrices will in the future be equivalent to " *
                     "`norm(vec(x), p)`. Use `opnorm(x, p)` for the Julia 0.6 behavior of " *
                     "computing the operator norm for matrices.", :norm)
        return opnorm(x, p)
    end
end

if isdefined(LinearAlgebra, :vecnorm) # deprecated but defined
    import LinearAlgebra: vecnorm
end
Base.@deprecate vecnorm(x::AbstractExpr, p::Real=2) norm(vec(x), p)


using SparseDiffTools, ForwardDiff, FiniteDiff, Zygote, IterativeSolvers
using LinearAlgebra, Test

using Random
Random.seed!(123)

const A = rand(300,300)
f(y,x) = mul!(y,A,x)
f(x) = A*x
x = rand(300)
v = rand(300)
dy = similar(x)
g(x) = sum(abs2,x)
function h(x)
      FiniteDiff.finite_difference_gradient(g,x)
end
function h(dy,x)
      FiniteDiff.finite_difference_gradient!(dy,g,x)
end

cache1 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
cache2 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
@test num_jacvec!(dy, f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
@test num_jacvec!(dy, f, x, v, similar(v), similar(v)) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
@test num_jacvec(f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6

@test auto_jacvec!(dy, f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v
@test auto_jacvec!(dy, f, x, v, cache1, cache2) ≈ ForwardDiff.jacobian(f,similar(x),x)*v
@test auto_jacvec(f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v

@test num_hesvec!(dy, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test num_hesvec!(dy, g, x, v, similar(v), similar(v), similar(v)) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test num_hesvec(g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2

@test numauto_hesvec!(dy, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test numauto_hesvec!(dy, g, x, v, ForwardDiff.GradientConfig(g,x), similar(v), similar(v)) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test numauto_hesvec(g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8

@test autonum_hesvec!(dy, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test autonum_hesvec!(dy, g, x, v, cache1, cache2) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test autonum_hesvec(g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8

@test numback_hesvec!(dy, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test numback_hesvec!(dy, g, x, v, similar(v), similar(v)) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test numback_hesvec(g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8

cache3 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(Nothing,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
cache4 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(Nothing,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
@test autoback_hesvec!(dy, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test autoback_hesvec!(dy, g, x, v, cache3, cache4) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test autoback_hesvec(g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8

@test num_hesvecgrad!(dy, h, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test num_hesvecgrad!(dy, h, x, v, similar(v), similar(v)) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test num_hesvecgrad(h, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2

@test auto_hesvecgrad!(dy, h, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test auto_hesvecgrad!(dy, h, x, v, cache1, cache2) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test auto_hesvecgrad(h, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2

L = JacVec(f,x)
@test L*x ≈ auto_jacvec(f, x, x)
@test L*v ≈ auto_jacvec(f, x, v)
@test mul!(dy,L,v) ≈ auto_jacvec(f, x, v)
L.x .= v
@test mul!(dy,L,v) ≈ auto_jacvec(f, v, v)

L = JacVec(f,x,autodiff=false)
@test L*x ≈ num_jacvec(f, x, x)
@test L*v ≈ num_jacvec(f, x, v)
L.x == x
@test mul!(dy,L,v) ≈ num_jacvec(f, x, v) rtol=1e-6
L.x .= v
@test mul!(dy,L,v) ≈ num_jacvec(f, v, v) rtol=1e-6

### Integration test with IterativeSolvers
out = similar(v)
gmres!(out, L, v)

x = rand(300)
v = rand(300)
L = HesVec(g,x,autodiff=false)
@test L*x ≈ num_hesvec(g, x, x)
@test L*v ≈ num_hesvec(g, x, v)
@test mul!(dy,L,v) ≈ num_hesvec(g, x, v) rtol=1e-2
L.x .= v
@test mul!(dy,L,v) ≈ num_hesvec(g, v, v) rtol=1e-2

L = HesVec(g,x)
@test L*x ≈ numauto_hesvec(g, x, x)
@test L*v ≈ numauto_hesvec(g, x, v)
@test mul!(dy,L,v) ≈ numauto_hesvec(g, x, v) rtol=1e-8
L.x .= v
@test mul!(dy,L,v) ≈ numauto_hesvec(g, v, v) rtol=1e-8

### Integration test with IterativeSolvers
out = similar(v)
gmres!(out, L, v)

x = rand(300)
v = rand(300)
L = HesVecGrad(h,x,autodiff=false)
@test L*x ≈ num_hesvec(g, x, x)
@test L*v ≈ num_hesvec(g, x, v)
@test mul!(dy,L,v) ≈ num_hesvec(g, x, v) rtol=1e-2
L.x .= v
@test mul!(dy,L,v) ≈ num_hesvec(g, v, v) rtol=1e-2

L = HesVecGrad(h,x,autodiff=true)
@test L*x ≈ autonum_hesvec(g, x, x)
@test L*v ≈ numauto_hesvec(g, x, v)
@test mul!(dy,L,v) ≈ numauto_hesvec(g, x, v) rtol=1e-8
L.x .= v
@test mul!(dy,L,v) ≈ numauto_hesvec(g, v, v) rtol=1e-8

### Integration test with IterativeSolvers
out = similar(v)
gmres!(out, L, v)


immutable Exponential{T} <: Model
    a::T
    δ::T
end

Base.show(io::IO, ::Type{Exponential}) = print(io, "Classical Exponential")

function Base.show(io::IO, m::Exponential)
    @printf(
        io,
        "Classical Exponential (a = %.3f, δ = %.3f)",
        m.a,
        m.δ,
    )
end

function initial_parameters(::Type{Exponential})
    return Float64[log(1.0), logit(0.5)]
end

canonical_parameters(m::Exponential) = (m.a, m.δ)

function Exponential(parameters::Vector)
    a = exp(parameters[1])
    δ = invlogit(parameters[2])
    return Exponential(a, δ)
end

@inline function predict(
    m::Exponential,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    a, δ = m.a, m.δ
    u1 = x1 * δ^t1
    u2 = x2 * δ^t2
    z = a * (u2 - u1)
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::Exponential,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    a, δ = m.a, m.δ
    u2 = x2 * δ^t2
    u1 = x1 * δ^t1
    gr = (
        (u2 - u1) * a,
        a * ((x2 * t2 * δ^t2 - x1 * t1 * δ^t1) / δ) * δ * (1 - δ),
    )
    return gr
end


#=
# Normal shocks
In the last notebook, we saw that normal shocks develop in a nozzle when it is
not perfectly expanded. This notebook demonstrates the use of the compressible tools
for dealing with **normal shock waves**.
=#

# ### Set up the module
using Gasdynamics1D
#-
using Plots

#=
### General use of normal shock relations
Normal shock relations are designated with the argument `NormalShock`. For example,
let's look at the ratio of stagnation pressures $p_{02}/p_{01}$ across a normal shock
 with the Mach number $M_1$ entering the shock equal to 2:
=#
p02_over_p01 = StagnationPressureRatio(MachNumber(2),NormalShock)

# and the density $\rho_2/\rho_1$ increases by the factor:
ρ2_over_ρ1 = DensityRatio(MachNumber(2),NormalShock)

# What is the entropy change across this shock?
s2_minus_s1 = Entropy(Entropy(0),MachNumber(2),NormalShock)

# Let's plot the entropy increase as a function of Mach number:
M1 = range(1,3,length=101)
s2 = []
pratio = []
for M in M1
    push!(s2,ustrip(Entropy(Entropy(0),MachNumber(M),NormalShock)))
    push!(pratio,ustrip(PressureRatio(MachNumber(M),NormalShock)))
end
#-
plot(M1,s2,xlim=(1,3),ylim=(0,400),xlabel="Mach number",ylabel="Entropy increase (J/kgK)")
# Not surprisingly, the entropy jump gets bigger as the Mach number increases

#=
### Example 1
For air entering a shock at Mach number 3, 1 atm, and 50 degrees F, what are the
Mach number, pressure, and temperature on the other side of the shock?
=#
p1 = Pressure(1u"atm")
M1 = MachNumber(3)
T1 = Temperature(50u"°F")

# The Mach number exiting the shock is
M2 = MachNumber(M1,NormalShock)

# Pressure exiting the shock (in atm) is
p2 = Pressure(p1*PressureRatio(M1,NormalShock))
value(p2,u"atm")

# and the temperature exiting the shock (in F) is
T2 = Temperature(T1*TemperatureRatio(M1,NormalShock))
value(T2,u"°F")

#=
### Example 2
A blow-down supersonic windtunnel is supplied with air from a large reservoir.
A Pitot tube is placed at the exit plane of a converging-diverging nozzle. The test
section lies at the end of the diverging section. The Mach number in the test
section (M2) is 2 and the pressure is below atmospheric, so a shock is formed at
the exit. The pressure just after the shock (p3) is 14.7 psi. Find the pressure in
the reservoir (p01), the pressure in the throat (pth), the Mach number just after
the shock (M3), the pressure at the Pitot tube (p4), and the temperature at the Pitot
 tube (T4).
=#
M2 = MachNumber(2)
p3 = Pressure(14.7u"psi")

#=
First, let's find the pressure just before the shock and the Mach number
just after the shock
=#
# p3/p2
p3_over_p2 = PressureRatio(M2,NormalShock)

# p03/p02
p03_over_p02 = StagnationPressureRatio(M2,NormalShock)

# M3
M3 = MachNumber(M2,NormalShock)

# We can now calculate the pressure just before the shock, using $p_2 = p_3/(p_3/p_2)$
p2 = Pressure(p3/p3_over_p2)

#=
Now, let's find the conditions in the throat and reservoir, based on what we have
 been given here. We can immediately find the stagnation pressure upstream of the
 shock using the isentropic relation. This stagnation pressure is the same as the
 reservoir pressure.
=#
p02 = StagnationPressure(p2,M2,Isentropic)
p01 = p02

#=
We know that the throat is choked. It must be, because the flow is subsonic before
and supersonic after the throat. Therefore, $M_{th}$ is 1, and we can get the local
pressure ($p_{th}$) from knowing the stagnation pressure and Mach number there:
=#
Mth = MachNumber(1)
p0th = p02
pth = Pressure(p0th,Mth,Isentropic)

#=
Now we can calculate the conditions at the Pitot tube at the exit. At the nose of
the Pitot tube, the flow stagnates. Thus, the pressure and temperature are equal
to the stagnation values. The stagnation pressure at the exit is the same as the
stagnation pressure at point 3 (just after the shock):
=#
p03 = Pressure(p03_over_p02*p02)
#-
value(p03,u"psi")


export Lehmer, random_sample

# MINSTD
const M_BIT_LEN = 31
const M_PLUS_1 = UInt64(2) ^ M_BIT_LEN
const LCG_M = M_PLUS_1 - UInt64(1)  # Mersenne prime, so smaller numbers are always coprime to it
const LCG_A = UInt64(7) ^ 5

"Lehmer RNG, a special case of Linear congruential generator where c = 0."
mutable struct Lehmer
    x::UInt32

    Lehmer(seed::UInt32) = new(seed)
end

function Base.iterate(l::Lehmer, state=l)
    # Fast computes `(a * x) % m` when `m = (2 ^ e - 1)`:
    #
    #     r ⧋ (a * x) % m  s.t. ∃ k: (a * x) = k * m + r
    #
    #     ∃ k: a * x = k * (2 ^ e - 1) + r  ⟹
    #     ∃ k: a * x = k * (2 ^ e) + (r - k)
    #
    # Sometimes, `r - k` is negative, so we have
    #
    #     ∃ k: a * x = (k - 1) * (2 ^ e) + (2 ^ e + r - (k - 1) - 1)
    #
    # Since `r - k << (2 ^ e)`,
    #
    #     k = (a * x) ÷ (2 ^ e)
    #     r = (a * x) % (2 ^ e) + k

    ax = LCG_A * l.x
    k = ax >> 31
    r = (ax & LCG_M) + k
    rr = r + UInt64(1)
    if !iszero(rr & M_PLUS_1)
        r = rr & LCG_M
    end
    l.x = r

    (r, l)
end

UINT32_MAX = typemax(UInt32)

"Randomly samples `k` indices without replacement from 1 to n by Fisher–Yates shuffle."
function random_sample(lehmer::Lehmer, n::UInt32, k::UInt32)::Vector{UInt32}
    samples = collect(UInt32(1):n)

    # Random values from RNG which are greater than `skip_threshold` will be skipped. For dealing
    # with modulo bias.
    for i in 1:k
        n_choices = n - i + 1
        remainder = UINT32_MAX % n_choices
        skip_threshold = UINT32_MAX - (remainder + 1) % n_choices
        j = i + first(Iterators.filter(x -> x <= skip_threshold, lehmer)) % n_choices
        samples[i], samples[j] = samples[j], samples[i]
    end

    samples[1:k]
end

# For testing
function next_lehmer(x::UInt32)::UInt32
    UInt32((LCG_A * x) % LCG_M)
end


@doc raw"""
Axisymmetric tokamak equilibrium in (R,Z,ϕ) coordinates with covariant
components of the vector potential given by
```math
A (R, Z, \phi) = \frac{B_0}{2} \, \bigg( R_0 \, \frac{Z}{R} , \, - R_0 \, \ln \bigg( \frac{R}{R_0} \bigg) , \, - \frac{r^2}{q_0} \bigg)^T ,
```
resulting in the magnetic field with covariant components
```math
B (R, Z, \phi) = \frac{B_0}{q_0} \, \bigg( - \frac{Z}{R} , \, \frac{R - R_0}{R} , \, - q_0 R_0 \bigg)^T ,
```
where $r = \sqrt{ (R - R_0)^2 + Z^2 }$.

Parameters:
 * `R₀`: position of magnetic axis
 * `B₀`: B-field at magnetic axis
 * `q₀`: safety factor at magnetic axis
"""
module AxisymmetricTokamakCylindrical

    using RecipesBase

    import ..ElectromagneticFields
    import ..ElectromagneticFields: AnalyticEquilibrium, code

    export  AxisymmetricTokamakCylindricalEquilibrium

    const DEFAULT_R₀ = 1.0
    const DEFAULT_B₀ = 1.0
    const DEFAULT_q₀ = 2.0

    const ITER_R₀ = 6.2
    const ITER_B₀ = 5.3
    const ITER_q₀ = √2

    struct AxisymmetricTokamakCylindricalEquilibrium{T <: Number} <: AnalyticEquilibrium
        name::String
        R₀::T
        B₀::T
        q₀::T

        function AxisymmetricTokamakCylindricalEquilibrium{T}(R₀::T, B₀::T, q₀::T) where T <: Number
            new("AxisymmetricTokamakCylindricalEquilibrium", R₀, B₀, q₀)
        end
    end

    AxisymmetricTokamakCylindricalEquilibrium(R₀::T=DEFAULT_R₀, B₀::T=DEFAULT_B₀, q₀::T=DEFAULT_q₀) where T <: Number = AxisymmetricTokamakCylindricalEquilibrium{T}(R₀, B₀, q₀)

    function init(R₀=DEFAULT_R₀, B₀=DEFAULT_B₀, q₀=DEFAULT_q₀)
        AxisymmetricTokamakCylindricalEquilibrium(R₀, B₀, q₀)
    end

    function ITER()
        AxisymmetricTokamakCylindricalEquilibrium(ITER_R₀, ITER_B₀, ITER_q₀)
    end

    macro code(R₀=DEFAULT_R₀, B₀=DEFAULT_B₀, q₀=DEFAULT_q₀)
        code(init(R₀, B₀, q₀); escape=true)
    end

    macro code_iter()
        code(ITER(); escape=true)
    end

    function Base.show(io::IO, equ::AxisymmetricTokamakCylindricalEquilibrium)
        print(io, "Axisymmetric Tokamak Equilibrium in (R,Z,ϕ) Coordinates with\n")
        print(io, "  R₀ = ", equ.R₀, "\n")
        print(io, "  B₀ = ", equ.B₀, "\n")
        print(io, "  q₀ = ", equ.q₀)
    end

    R(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = x[1]
    Z(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = x[2]
    ϕ(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = x[3]
    r²(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = (R(x,equ) - equ.R₀)^2 + Z(x,equ)^2
    r(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = sqrt(r²(x, equ))
    X(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = R(x,equ) * cos(ϕ(x,equ))
    Y(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = R(x,equ) * sin(ϕ(x,equ))
    θ(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = atan(Z(x,equ), R(x,equ) - equ.R₀)

    ElectromagneticFields.J(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = R(x,equ)

    ElectromagneticFields.A₁(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = + equ.B₀ * equ.R₀ * Z(x,equ) / R(x,equ) / 2
    ElectromagneticFields.A₂(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = - equ.B₀ * equ.R₀ * log(R(x,equ) / equ.R₀) / 2
    ElectromagneticFields.A₃(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = + equ.B₀ * r²(x,equ) / equ.q₀ / 2

    ElectromagneticFields.x¹(ξ::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = X(ξ,equ)
    ElectromagneticFields.x²(ξ::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = Y(ξ,equ)
    ElectromagneticFields.x³(ξ::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = Z(ξ,equ)

    ElectromagneticFields.ξ¹(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = sqrt(x[1]^2 + x[2]^2)
    ElectromagneticFields.ξ²(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = x[3]
    ElectromagneticFields.ξ³(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = atan(x[2], x[1])

    ElectromagneticFields.g₁₁(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = one(eltype(x))
    ElectromagneticFields.g₂₂(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = one(eltype(x))
    ElectromagneticFields.g₃₃(x::AbstractVector, equ::AxisymmetricTokamakCylindricalEquilibrium) = R(x, equ)^2

    ElectromagneticFields.get_functions(::AxisymmetricTokamakCylindricalEquilibrium) = (X=X, Y=Y, Z=Z, R=R, r=r, θ=θ, ϕ=ϕ, r²=r²)

    function ElectromagneticFields.periodicity(x::AbstractVector, ::AxisymmetricTokamakCylindricalEquilibrium)
        p = zero(x)
        p[3] = 2π
        return p
    end

    @recipe function f(equ::AxisymmetricTokamakCylindricalEquilibrium;
                       nx = 100, ny = 120, levels = 50, size = (400,400),
                       xlims = (  0.5 * equ.R₀,   1.5 * equ.R₀),
                       ylims = (- 0.5 * equ.R₀, + 0.5 * equ.R₀))

        xgrid = LinRange(xlims[1], xlims[2], nx)
        zgrid = LinRange(ylims[1], ylims[2], ny)
        pot   = [ElectromagneticFields.A₃([xgrid[i], zgrid[j], 0.0], equ) / xgrid[i] for i in eachindex(xgrid), j in eachindex(zgrid)]

        seriestype   := :contour
        aspect_ratio := :equal
        size   := size
        xlims  := xlims
        ylims  := ylims
        levels := levels
        legend := :none

        (xgrid, zgrid, pot')
    end

end


using LinearAlgebra, Statistics, Roots, Optim, Plots, Random; Random.seed!(1234)
n = 1000 # number of samples to draw

# the true regression function
r(x) = 2 + 1/50*x*(30-x)
# the true density function
σy = 3/2
f(x,y) = 3/4000 * 1/√(2π*σy^2) * x*(20-x)*exp(-1/(2σy^2)*(y-r(x))^2)

# x values and y values for a grid
xs = 0:1/2^3:20
ys = 0:1/2^3:12

# F(t) is the CDF of X
F(t) = -t^3/4000 + 3t^2/400
# Inverse CDF sampling
function sampleX()
    U = rand()
    find_zero(t->F(t)-U,(0,20),Bisection())
end

function sampleXY(r,σ)
    X = sampleX()
    Y = r(X) + σ*randn()
    (X,Y)
end

samples = [sampleXY(r,σy) for i=1:n]

D(u) = abs(u) < 1 ? 70/81*(1-abs(u)^3)^3 : 0 # tri-cube function
D(λ,u) = 1/λ*D(u/λ) # scaled tri-cube
K(λ,x,y) = D(λ,x) * D(λ,y) # kernel
kde(λ,x,y,samples) = sum(K(λ,x-Xi,y-Yi) for (Xi,Yi) in samples)/length(samples)

# `optimize` takes functions which accept vector arguments, so we treat
# λ as a one-entry vector
L(λ) = sum((f(x,y) - kde(λ,x,y,samples))^2 for x=xs,y=ys)*step(xs)*step(ys)

# minimize L using the BFGS method
λ_best = optimize(λ->L(first(λ)),[1.0],BFGS())

using AsyPlots
#plot = Plots.plot
#heatmap = Plots.heatmap
    # background colors
softseagreen = NamedColor("SoftSeaGreen",(0.96, 0.995, 0.98))
softblue = NamedColor("SoftBlue",(0.92,0.95,0.99))
softyellow = NamedColor("SoftYellow",(0.98, 0.98, 0.9))

densities = [f(x,y) for x=xs,y=ys]
xmax = xs[end]
ymax = ys[end]

# plot accoutrements
extras = Plot([Path([0 10; -0.5 10]),Label("10",(-1.5,10))]) +
    Plot([Path([20 0; 20 -0.66]),Label("20",(20,-1.5))]) +
    Plot(Path([0 0; 0 ymax],arrow=Arrow())) +
    Plot(Path([0 0; 1.05*xmax 0],arrow=Arrow())) +
    Plot(Label("hours studied",(10,-1))) +
    Plot(Label("score",(-1,5);rotation=90);
         width=256,height=128,ignoreaspect=true,bgfill=false)
function sampleXY()
    X = sampleX()
    Y = r(X) + σy*randn()
    if !(0 < Y < 10)
        sampleXY()
    else
        (X,Y)
    end
end
Random.seed!(1234)
n = 1000
samples = [sampleXY() for i=1:n]
samplesfig = Plot([Point(point;color="MidnightBlue",linewidth=2)
                   for point in samples]) + extras
save("figures/exam-samples.pdf",samplesfig)

densityfig = AsyPlots.heatmap(densities,(0,0),(xmax,ymax);
                  colors=[NamedColor("white"),NamedColor("MidnightBlue")],
              smooth=true,alpha=fill(4/5,size(densities)...)) + extras

save("exam-density.pdf",densityfig)

densitywithline = AsyPlots.heatmap(densities,(0,0),(xmax,ymax);
                                        colors=[NamedColor("white"),NamedColor("MidnightBlue")],
                                        smooth=true,alpha=fill(4/5,size(densities)...)) +
                   Plot(Path([(x,r(x)) for x=0:1/2^4:20],color="DarkRed",linewidth=2)) + extras

save("exam-density-line.pdf",densitywithline)

cd(@__DIR__)

pwd()


# fft
a = rand(8) + im*rand(8)
@test norm(ifft(fft(a)) - a) < 1e-8

m4 = [16.    2     3    13;
    5    11    10     8;
    9     7     6    12;
    4    14    15     1]

b = rand(17,14)
b[3:6,9:12] = m4
sm4 = slice(b,3:6,9:12)

pm4 = plan_fft(m4,1)

fft_m4 = fft(m4,1)
fftd2_m4 = fft(m4,2)
ifft_fft_m4 = ifft(fft(m4,1),1)
fftn_m4 = fft(m4)
ifftn_fftn_m4 = ifft(fft(m4))

fft!_m4 = complex(m4); fft!(fft!_m4,1)
fft!d2_m4 = complex(m4); fft!(fft!d2_m4,2)
ifft!_fft_m4 = fft(m4,1); ifft!(ifft!_fft_m4,1)
fft!n_m4 = complex(m4); fft!(fft!n_m4)
ifft!n_fftn_m4 = fft(m4); ifft!(ifft!n_fftn_m4)

pfft_m4 = plan_fft(m4,1)(m4)
pfftd2_m4 = plan_fft(m4,2)(m4)
pifft_fft_m4 = plan_ifft(fft_m4,1)(fft_m4)
pfftn_m4 = plan_fft(m4)(m4)
pifftn_fftn_m4 = plan_ifft(fftn_m4)(fftn_m4)

pfft!_m4 = complex(m4); plan_fft!(pfft!_m4,1)(pfft!_m4)
pfft!d2_m4 = complex(m4); plan_fft!(pfft!d2_m4,2)(pfft!d2_m4)
pifft!_fft_m4 = fft(m4,1); plan_ifft!(pifft!_fft_m4,1)(pifft!_fft_m4)
pfft!n_m4 = complex(m4); plan_fft!(pfft!n_m4)(pfft!n_m4)
pifft!n_fftn_m4 = fft(m4); plan_ifft!(pifft!n_fftn_m4)(pifft!n_fftn_m4)

sfftn_m4 = fft(sm4)
psfftn_m4 = plan_fft(sm4)(sm4)
sfft!n_b = complex128(b)
sfft!n_m4 = slice(sfft!n_b,3:6,9:12); fft!(sfft!n_m4)
psfft!n_b = complex128(b)
psfft!n_m4 = slice(psfft!n_b,3:6,9:12); plan_fft!(psfft!n_m4)(psfft!n_m4)

true_fft_m4 = [
    34.            34.            34.            34.;
     7. - 1.im  -5. + 3.im  -3. + 5.im   1. - 7.im;
    16.           -16.           -16.            16.;
     7. + 1.im  -5. - 3.im  -3. - 5.im   1. + 7.im ]

true_fftn_m4 = [
 136.        0          0         0 ;
   0.       20          8 + 8im   0 - 12im ;
   0.       32 + 32im   0        32 - 32im ;
   0.        0 + 12im   8 - 8im  20 ]

true_fftd2_m4 = [
   34.   13 + 11im    4   13 - 11im ;
   34.   -5 -  3im   -4   -5 +  3im ;
   34.    3 +  5im   -4    3 -  5im ;
   34.  -11 - 13im    4  -11 + 13im ]

for i = 1:length(m4)
    @test_approx_eq fft_m4[i] true_fft_m4[i]
    @test_approx_eq fftd2_m4[i] true_fftd2_m4[i]
    @test_approx_eq ifft_fft_m4[i] m4[i]
    @test_approx_eq fftn_m4[i] true_fftn_m4[i]
    @test_approx_eq ifftn_fftn_m4[i] m4[i]

    @test_approx_eq fft!_m4[i] true_fft_m4[i]
    @test_approx_eq fft!d2_m4[i] true_fftd2_m4[i]
    @test_approx_eq ifft!_fft_m4[i] m4[i]
    @test_approx_eq fft!n_m4[i] true_fftn_m4[i]
    @test_approx_eq ifft!n_fftn_m4[i] m4[i]

    @test_approx_eq pfft_m4[i] true_fft_m4[i]
    @test_approx_eq pfftd2_m4[i] true_fftd2_m4[i]
    @test_approx_eq pifft_fft_m4[i] m4[i]
    @test_approx_eq pfftn_m4[i] true_fftn_m4[i]
    @test_approx_eq pifftn_fftn_m4[i] m4[i]

    @test_approx_eq pfft!_m4[i] true_fft_m4[i]
    @test_approx_eq pfft!d2_m4[i] true_fftd2_m4[i]
    @test_approx_eq pifft!_fft_m4[i] m4[i]
    @test_approx_eq pfft!n_m4[i] true_fftn_m4[i]
    @test_approx_eq pifft!n_fftn_m4[i] m4[i]

    @test_approx_eq sfftn_m4[i] true_fftn_m4[i]
    @test_approx_eq sfft!n_m4[i] true_fftn_m4[i]
    @test_approx_eq psfftn_m4[i] true_fftn_m4[i]
    @test_approx_eq psfft!n_m4[i] true_fftn_m4[i]
end

ifft!(sfft!n_m4)
plan_ifft!(psfft!n_m4)(psfft!n_m4)
@test norm(sfft!n_m4 - m4) < 1e-8
@test norm(psfft!n_m4 - m4) < 1e-8

m3d = float32(reshape(1:5*3*2, 5, 3, 2))
ifft3_fft3_m3d = ifft(fft(m3d))

fftd3_m3d = fft(m3d,3)
ifftd3_fftd3_m3d = ifft(fftd3_m3d,3)

fft!d3_m3d = complex(m3d); fft!(fft!d3_m3d,3)
ifft!d3_fftd3_m3d = copy(fft!d3_m3d); ifft!(ifft!d3_fftd3_m3d,3)

pfftd3_m3d = plan_fft(m3d,3)(m3d)
pifftd3_fftd3_m3d = plan_ifft(fftd3_m3d,3)(fftd3_m3d)

pfft!d3_m3d = complex(m3d); plan_fft!(pfft!d3_m3d,3)(pfft!d3_m3d)
pifft!d3_fftd3_m3d = copy(fft!d3_m3d); plan_ifft!(pifft!d3_fftd3_m3d,3)(pifft!d3_fftd3_m3d)

@test isa(fftd3_m3d, Array{Complex64,3})
@test isa(ifftd3_fftd3_m3d, Array{Complex64,3})
@test isa(fft!d3_m3d, Array{Complex64,3})
@test isa(ifft!d3_fftd3_m3d, Array{Complex64,3})
@test isa(pfftd3_m3d, Array{Complex64,3})
@test isa(pifftd3_fftd3_m3d, Array{Complex64,3})
@test isa(pfft!d3_m3d, Array{Complex64,3})
@test isa(pifft!d3_fftd3_m3d, Array{Complex64,3})

true_fftd3_m3d = Array(Float32, 5, 3, 2)
true_fftd3_m3d[:,:,1] = 17:2:45
true_fftd3_m3d[:,:,2] = -15

for i = 1:length(m3d)
    @test_approx_eq fftd3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq ifftd3_fftd3_m3d[i] m3d[i]
    @test_approx_eq ifft3_fft3_m3d[i] m3d[i]

    @test_approx_eq fft!d3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq ifft!d3_fftd3_m3d[i] m3d[i]

    @test_approx_eq pfftd3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq pifftd3_fftd3_m3d[i] m3d[i]
    @test_approx_eq pfft!d3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq pifft!d3_fftd3_m3d[i] m3d[i]
end

# rfft/rfftn

rfft_m4 = rfft(m4,1)
rfftd2_m4 = rfft(m4,2)
rfftn_m4 = rfft(m4)

prfft_m4 = plan_rfft(m4,1)(m4)
prfftd2_m4 = plan_rfft(m4,2)(m4)
prfftn_m4 = plan_rfft(m4)(m4)

srfftn_m4 = rfft(sm4)
psrfftn_m4 = plan_rfft(sm4)(sm4)

for i = 1:3, j = 1:4
    @test_approx_eq rfft_m4[i,j] true_fft_m4[i,j]
    @test_approx_eq rfftd2_m4[j,i] true_fftd2_m4[j,i]
    @test_approx_eq rfftn_m4[i,j] true_fftn_m4[i,j]

    @test_approx_eq prfft_m4[i,j] true_fft_m4[i,j]
    @test_approx_eq prfftd2_m4[j,i] true_fftd2_m4[j,i]
    @test_approx_eq prfftn_m4[i,j] true_fftn_m4[i,j]

    @test_approx_eq srfftn_m4[i,j] true_fftn_m4[i,j]
    @test_approx_eq psrfftn_m4[i,j] true_fftn_m4[i,j]
end

irfft_rfft_m4 = irfft(rfft_m4,size(m4,1),1)
irfft_rfftd2_m4 = irfft(rfftd2_m4,size(m4,2),2)
irfftn_rfftn_m4 = irfft(rfftn_m4,size(m4,1))

pirfft_rfft_m4 = plan_irfft(rfft_m4,size(m4,1),1)(rfft_m4)
pirfft_rfftd2_m4 = plan_irfft(rfftd2_m4,size(m4,2),2)(rfftd2_m4)
pirfftn_rfftn_m4 = plan_irfft(rfftn_m4,size(m4,1))(rfftn_m4)

for i = 1:length(m4)
    @test_approx_eq irfft_rfft_m4[i] m4[i]
    @test_approx_eq irfft_rfftd2_m4[i] m4[i]
    @test_approx_eq irfftn_rfftn_m4[i] m4[i]

    @test_approx_eq pirfft_rfft_m4[i] m4[i]
    @test_approx_eq pirfft_rfftd2_m4[i] m4[i]
    @test_approx_eq pirfftn_rfftn_m4[i] m4[i]
end

rfftn_m3d = rfft(m3d)
rfftd3_m3d = rfft(m3d,3)
@test size(rfftd3_m3d) == size(fftd3_m3d)
irfft_rfftd3_m3d = irfft(rfftd3_m3d,size(m3d,3),3)
irfftn_rfftn_m3d = irfft(rfftn_m3d,size(m3d,1))
for i = 1:length(m3d)
    @test_approx_eq rfftd3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq irfft_rfftd3_m3d[i] m3d[i]
    @test_approx_eq irfftn_rfftn_m3d[i] m3d[i]
end

fftn_m3d = fft(m3d)
@test size(fftn_m3d) == (5,3,2)
rfftn_m3d = rfft(m3d)
@test size(rfftn_m3d) == (3,3,2)
for i = 1:3, j = 1:3, k = 1:2
    @test_approx_eq rfftn_m3d[i,j,k] fftn_m3d[i,j,k]
end


"""
Compute the power of a test without checking arguments. Useful inside of loops
that need to evaluate the power of many proposed designs.
"""
function analytic_power(
    ::Type{T},
    ns::Union{Real, Tuple},
    δ::Real,
    σs::Union{Real, Tuple},
    α::Real,
    one_sided::Bool,
) where {T}
    null, alt = hypotheses(T, ns, δ, σs)
    if one_sided
        c_r = cquantile(null, α)
        return ccdf(alt, c_r)
    else
        c_l = quantile(null, α / 2)
        c_r = cquantile(null, α / 2)
        return cdf(alt, c_l) + ccdf(alt, c_r)
    end
end

"""
Compute the power of a test. Users must specify:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function power(
    ::Type{T},
    ns::Union{Real, Tuple},
    δ::Real,
    σs::Union{Real, Tuple},
    α::Real = 0.05,
    one_sided::Bool = false,
) where {T}
    check_args(ns, δ, σs, 0.8, α)
    return analytic_power(T, ns, δ, σs, α, one_sided)
end

"""
Compute the smallest effect size measurable using a test. Users must specify:

* `ns`: Per-group sample size(s)
* `p`: Desired power
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function effect_size(
    ::Type{T},
    ns::Union{Real, Tuple},
    p::Real,
    α::Real = 0.05,
    one_sided::Bool = false,
) where {T}
    check_args(ns, 1.0, 1.0, p, α)
    gap(δ::Real) = analytic_power(T, ns, δ, 1.0, α, one_sided) - p
    return fzero(gap, 1e-8, 100.0)
end

"""
Compute the smallest sample size that allows using a test to measure a
specified effect. Users must specify:

* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
* `p`: Desired power
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function sample_size(
    ::Type{T},
    δ::Real,
    σs::Union{Real, Tuple},
    p::Real,
    α::Real = 0.05,
    one_sided::Bool = false,
) where {T}
    check_args(2, δ, σs, p, α)
    gap(n::Real) = analytic_power(T, n, δ, σs, α, one_sided) - p
    return fzero(gap, 2.0, 1.0e12)
end

"""
Compute the type 1 error rate of a test. Users must specify:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function type_1(
    ::Type{T},
    ns::Union{Real, Tuple},
    δ::Real,
    σs::Union{Real, Tuple},
    α::Real = 0.05,
    one_sided::Bool = false,
) where {T}
    check_args(ns, δ, σs, 0.8, α)
    return α
end

"""
Compute the type 2 error rate of a test. Users must specify:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function type_2(
    ::Type{T},
    ns::Union{Real, Tuple},
    δ::Real,
    σs::Union{Real, Tuple},
    α::Real = 0.05,
    one_sided::Bool = false,
) where {T}
    check_args(ns, δ, σs, 0.8, α)
    return 1 - power(T, ns, δ, σs, α, one_sided)
end

"""
Compute the type S error rate of a test. Users must specify:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function type_s(
    ::Type{T},
    ns::Union{Real, Tuple},
    δ::Real,
    σs::Union{Real, Tuple},
    α::Real = 0.05,
) where {T}
    check_args(ns, δ, σs, 0.8, α)
    null, alt = hypotheses(T, ns, δ, σs)
    c_l, c_r = quantile(null, α / 2), cquantile(null, α / 2)
    p_l, p_r = cdf(alt, c_l), ccdf(alt, c_r)
    return p_l / (p_l + p_r)
end

"""
Compute the type M error rate of a test. Users must specify:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function type_m(
    ::Type{T},
    ns::Union{Real, Tuple},
    δ::Real,
    σs::Union{Real, Tuple},
    α::Real = 0.05,
    one_sided::Bool = false,
) where {T}
    check_args(ns, δ, σs, 0.8, α)
    null, alt = hypotheses(T, ns, δ, σs)
    if one_sided
        c_r = cquantile(null, α)
        t_l, t_r = thresholds(alt)
        p_significant = ccdf(alt, c_r)
        p_overestimate = ccdf(alt, max(c_r, t_r))
        return p_overestimate / p_significant
    else
        c_l, c_r = quantile(null, α / 2), cquantile(null, α / 2)
        t_l, t_r = thresholds(alt)
        p_significant = cdf(alt, c_l) + ccdf(alt, c_r)
        p_overestimate = cdf(alt, min(c_l, t_l)) + ccdf(alt, max(c_r, t_r))
        return p_overestimate / p_significant
    end
end

"""
Compute the average exaggeration factor of a test. Users must specify:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
* `α`: Significance level (defaults to 0.05)
* `one_sided`: Is the test one-sided? (defaults to false)
"""
function exaggeration_factor(
    ::Type{T},
    ns::Union{Real, Tuple},
    δ::Real,
    σs::Union{Real, Tuple},
    α::Real = 0.05,
    one_sided::Bool = false,
    bounds::Real = 1e-8,
    tol::Real = 1e-8,
) where {T}
    check_args(ns, δ, σs, 0.8, α)
    null, alt = hypotheses(T, ns, δ, σs)
    if one_sided
        c_r = cquantile(null, α)
        t_l, t_r = thresholds(alt)
        p_significant = ccdf(alt, c_r)
        f0(t) = (t / t_r) * (pdf(alt, t) / p_significant)
        i_r = cquantile(alt, bounds)
        e_r = gauss(f0, c_r, i_r, rtol = tol)[1]
        return e_r
    else
        c_l, c_r = quantile(null, α / 2), cquantile(null, α / 2)
        t_l, t_r = thresholds(alt)
        p_significant = cdf(alt, c_l) + ccdf(alt, c_r)
        f1(t) = (t / t_l) * (pdf(alt, t) / p_significant)
        f2(t) = (t / t_r) * (pdf(alt, t) / p_significant)
        i_l, i_r = quantile(alt, bounds), cquantile(alt, bounds)
        e_l = gauss(f1, i_l, c_l, rtol = tol)[1]
        e_r = gauss(f2, c_r, i_r, rtol = tol)[1]
        return e_l + e_r
    end
end


using ModelingToolkit, MethodOfLines, DomainSets, Test, Symbolics, SymbolicUtils, LinearAlgebra

@parameters x, t
@variables u(..)

Dx(d) = Differential(x)^d

Dt = Differential(t)

t_min= 0.
t_max = 2.0
x_min = 0.
x_max = 20.0

dx = 1.0

domains = [t ∈ Interval(t_min, t_max), x ∈ Interval(x_min, x_max)]

@testset "Test 01: Cartesian derivative discretization" begin
    weights = []
    push!(weights, ([-0.5,0,0.5], [1.,-2.,1.], [-1/2,1.,0.,-1.,1/2]))
    push!(weights, ([1/12, -2/3,0,2/3,-1/12], [-1/12,4/3,-5/2,4/3,-1/12], [1/8,-1.,13/8,0.,-13/8,1.,-1/8]))
    for d in 1:3
        for (i,a) in enumerate([2])
            pde = Dt(u(t,x)) ~ +Dx(d)(u(t,x))
            bcs = [u(0,x) ~ cos(x), u(t,0) ~ exp(-t), u(t,Float64(π)) ~ -exp(-t)]

            @named pdesys = PDESystem(pde,bcs,domains,[t,x],[u(t,x)])

            # Test centered order
            disc = MOLFiniteDifference([x=>dx], t; approx_order=a)

            depvar_ops = map(x->operation(x.val),pdesys.depvars)

            depvars_lhs = MethodOfLines.get_depvars(pde.lhs, depvar_ops)
            depvars_rhs = MethodOfLines.get_depvars(pde.rhs, depvar_ops)
            depvars = collect(depvars_lhs ∪ depvars_rhs)
            # Read the independent variables,
            # ignore if the only argument is [t]
            indvars = first(Set(filter(xs->!isequal(xs, [t]), map(arguments, depvars))))
            x̄ = first(Set(filter(!isempty, map(u->filter(x-> t === nothing || !isequal(x, t.val), arguments(u)), depvars))))

            s = MethodOfLines.DiscreteSpace(domains, depvars, x̄, disc)

            derivweights = MethodOfLines.DifferentialDiscretizer(pdesys, s, disc)

            #@show pde.rhs, operation(pde.rhs), arguments(pde.rhs)
            for II in s.Igrid[s.ū[1]][2:end-1]
                #II = s.Igrid[end-1]
                I1 = MethodOfLines.unitindices(1)[1]

                rules = MethodOfLines.generate_finite_difference_rules(II, s, depvars, pde, derivweights)

                disc_pde=substitute(pde.lhs,rules) ~ substitute(pde.rhs,rules)
                #@show disc_pde
                #@test disc_pde == dot(weights[i][d], s.discvars[depvars[1]][[II + j*I1 for j in MethodOfLines.half_range(length(weights[i][d]))]])
                ufunc(u, Itap, x) = s.discvars[u][Itap]
                if isodd(d)
                    @test isequal(substitute(pde.rhs,rules), MethodOfLines.upwind_difference(d, II, s, derivweights, (1, x), depvars[1], ufunc, true))
                else
                    @test isequal(substitute(pde.rhs,rules), MethodOfLines.central_difference(derivweights.map[Differential(x)^d], II, s, (1, x), depvars[1], ufunc))
                end
            end
        end
    end
end

@testset "Test 01: Nonlinear Diffusion discretization" begin

    pde = Dt(u(t,x)) ~ Dx(1)(u(t,x))
    bcs = [u(0,x) ~ cos(x), u(t,0) ~ exp(-t), u(t,Float64(π)) ~ -exp(-t)]

    @named pdesys = PDESystem(pde,bcs,domains,[t,x],[u(t,x)])

    # Test centered order

    depvar_ops = map(x->operation(x.val),pdesys.depvars)

    depvars_lhs = MethodOfLines.get_depvars(pde.lhs, depvar_ops)
    depvars_rhs = MethodOfLines.get_depvars(pde.rhs, depvar_ops)
    depvars = collect(depvars_lhs ∪ depvars_rhs)
    # Read the independent variables,
    # ignore if the only argument is [t]
    indvars = first(Set(filter(xs->!isequal(xs, [t]), map(arguments, depvars))))
    x̄ = first(Set(filter(!isempty, map(u->filter(x-> t === nothing || !isequal(x, t.val), arguments(u)), depvars))))

    for order in [2]
        disc = MOLFiniteDifference([x=>dx], t; approx_order=order)
        s = MethodOfLines.DiscreteSpace(domains, depvars, x̄, disc)

        derivweights = MethodOfLines.DifferentialDiscretizer(pdesys, s, disc)

        ufunc(u, I, x) = s.discvars[u][I]
        #TODO Test Interpolation of params
        # Test simple case
        for II in s.Igrid[s.ū[1]][2:end-1]
            expr = MethodOfLines.cartesian_nonlinear_laplacian((1~1).lhs, II, derivweights, s, depvars, x, u(t,x))
            expr2 = MethodOfLines.central_difference(derivweights.map[Differential(x)^2], II, s, (1,x), u(t,x), ufunc)
            @test isequal(expr, expr2)
        end
    end
end

@testset "Test 02: Spherical Diffusion discretization" begin

    pde  = Dt(u(t,x)) ~ 1/x^2 * Dx(1)(x^2 * Dx(1)(u(t,x)))

    bcs = [u(0,x) ~ cos(x), u(t,0) ~ exp(-t), u(t,Float64(π)) ~ -exp(-t)]

    @named pdesys = PDESystem(pde,bcs,domains,[t,x],[u(t,x)])

    # Test centered order
    disc = MOLFiniteDifference([x=>dx], t; approx_order=2)

    depvar_ops = map(x->operation(x.val),pdesys.depvars)

    depvars_lhs = MethodOfLines.get_depvars(pde.lhs, depvar_ops)
    depvars_rhs = MethodOfLines.get_depvars(pde.rhs, depvar_ops)
    depvars = collect(depvars_lhs ∪ depvars_rhs)
    # Read the independent variables,
    # ignore if the only argument is [t]
    indvars = first(Set(filter(xs->!isequal(xs, [t]), map(arguments, depvars))))
    x̄ = first(Set(filter(!isempty, map(u->filter(x-> t === nothing || !isequal(x, t.val), arguments(u)), depvars))))

    s = MethodOfLines.DiscreteSpace(domains, depvars, x̄, disc)

    derivweights = MethodOfLines.DifferentialDiscretizer(pdesys, s, disc)

    for II in s.Igrid[s.ū[1]][2:end-1]
        #TODO Test Interpolation of params
        expr = MethodOfLines.spherical_diffusion((1~1).lhs, II, derivweights, s, depvars, x, u(t,x))
        #@show II, expr
    end

end


function ll{T<:Real}(Θᵏ::Vector{T}, D, c, J)
    update_Pᵈ_at_xᵈ!(c.Pᵈ, D.xᵈ, Θᵏ, c.Δz̃ᵈ, c.Δẽᵈ, J)

    ll = zero(T)
    lengthxd = length(D.xᵈ)
    @inbounds for j = 1:J
        for i = 1:lengthxd
            ll -=  D.nxjᵈ[i,j]*log(c.Pᵈ[j][i])
        end
    end
    ll/D.nobs
end

∇ll(Θᵏ, grad, D, c, J) =
                  ∇ll(Θᵏ,            grad, D.x,  J, D.nobs, D.xᵈ, c, D)
function ∇ll{T<:Real}(Θᵏ::Vector{T}, grad,   xᵈ, J, n, Dxd, c, D)
    update_Pᵈ_at_xᵈ!(c.Pᵈ, Dxd, Θᵏ, c.Δz̃ᵈ, c.Δẽᵈ, J)

    z̄ = sum(c.Pᵈ[j].*c.z̃ᵈ[j] for j = 1:J)
    grad .= zero(T)
    @views for j in 1:J
        grad[:] -= sum(D.nxjᵈ[:, j].*(c.z̃ᵈ[j]-z̄),1)'
    end
    grad .= grad/n
end

ll_∇ll(Θᵏ, grad, D, c, J) =
                  ll_∇ll(Θᵏ,           grad, D.x,  J, D.nobs, c, D.nxjᵈ, D.xᵈ)
function ll_∇ll{T<:Real}(x::Vector{T}, grad,   xᵈ, J,   n,    c,   nxjᵈ, Dxd)

    K = length(x)
    lengthxd = length(Dxd)
    update_Pᵈ_at_xᵈ!(c.Pᵈ, Dxd, x, c.Δz̃ᵈ, c.Δẽᵈ, J)
    ll = zero(T)
    @inbounds for j = 1:J
        for i = 1:lengthxd
            ll -=  nxjᵈ[i,j]*log(c.Pᵈ[j][i])
        end
    end

    z̄ = sum(c.Pᵈ[j].*c.z̃ᵈ[j] for j = 1:J)
    grad .= zero(T)
    @views for j in 1:J
        grad[:] -= sum(nxjᵈ[:, j].*(c.z̃ᵈ[j]-z̄),1)'
    end
    grad .= grad/n
    ll/n
end

function ∇²ll{T<:Real}(Θᵏ::Vector{T}, hessian, D, c, J)
    update_Pᵈ_at_xᵈ!(c.Pᵈ, D.xᵈ, Θᵏ, c.Δz̃ᵈ, c.Δẽᵈ, J)
    K = size(c.z̃ᵈ[1],2)

    z̄ = sum(c.Pᵈ[j].*c.z̃ᵈ[j] for j = 1:J)
    nxx = sum(D.nxjᵈ,2)
    hessian .= zero(T)
    for j in 1:J
        hessian .+= (nxx.*c.Pᵈ[j].*(c.z̃ᵈ[j]-z̄))'*(c.z̃ᵈ[j]-z̄)
    end

    hessian .= hessian/D.nobs
end

update_Pᵈ_at_xᵈ!{T}(c, D, Θᵏ::Vector{T}, J) =
          update_Pᵈ_at_xᵈ!(c.Pᵈ, D.xᵈ, Θᵏ,          c.Δz̃ᵈ, c.Δẽᵈ, J)
function update_Pᵈ_at_xᵈ!{T}(Pᵈ,   xᵈ, Θᵏ::Vector{T}, Δz̃ᵈ,   Δẽᵈ, J)
    K = size(Δz̃ᵈ[1],2) # nvar
    nxᵈ = length(xᵈ)
    @inbounds for j = 1:J-1
        Pᵈ[j] .= Δẽᵈ[j]
        zj = Δz̃ᵈ[j]
        for k = 1:K
            @simd for i = 1:nxᵈ
                Pᵈ[j][i] += zj[i, k]*Θᵏ[k] # hoist
            end
        end
        Pᵈ[j] .= exp.(Pᵈ[j])
    end
    denom = 1+sum(Pᵈ[1:J-1])
    Pᵈ[end] .= one(T)
    for j = 1:J-1
        Pᵈ[j] .= Pᵈ[j]./denom
        Pᵈ[end] .-= Pᵈ[j]
    end
end


## Axis transformations
#==========================================================================================#

"""
    body_to_stability_axes(coords, α)

Convert coordinates into stability axes with angle ``\\alpha``.
"""
body_to_stability_axes(coords, α :: T) where T <: Real = RotY{T}(α) * coords

"""
    body_to_stability_axes(coords, α)

Convert coordinates into stability axes with angle ``\\alpha``.
"""
stability_to_body_axes(coords, α :: T) where T <: Real = RotY{T}(-α) * coords

"""
    body_to_wind_axes(coords, α, β)

Convert coordinates from body axes to wind axes with angles ``\\alpha,~ \\beta``.
"""
body_to_wind_axes(coords, α :: T, β :: T) where T <: Real = RotZY{T}(β, α) * coords

"""
    body_to_wind_axes(coords, α, β)

Convert coordinates from wind axes to body axes with angles ``\\alpha,~ \\beta``.
"""
wind_to_body_axes(coords, α :: T, β :: T) where T <: Real = RotZY{T}(-α, -β) * coords

## Line methods
#==========================================================================================#

"""
    body_to_wind_axes(line, freestream)

Transform a Line from body to wind axes in a given `Freestream`.
"""
body_to_wind_axes(line :: Line, α :: T, β :: T) where T <: Real = Line(body_to_wind_axes(line.r1, α, β), body_to_wind_axes(line.r2, α, β))

## Reflections and projections
#==========================================================================================#

"""
    reflect_xz(vector)

Reflect the y-coordinate of a given 3-dimensional vector about the x-z plane.
"""
reflect_xz(vector) = SVector(vector[1], -vector[2], vector[3])

"""
    project_yz(vector)

Project a given 3-dimensional vector into the y-z plane.
"""
project_yz(vector) = SVector(0, vector[2], vector[3])

"""
    reflect_yz(line)

Reflect a Line onto the ``x``-``z`` plane of its reference coordinate system.
"""
reflect_xz(line :: Line) = Line((reflect_xz ∘ p2)(line), (reflect_xz ∘ p1)(line))

"""
    reflect_yz(horseshoe)

Reflect a Horseshoe onto the ``x``-``z`` plane of its reference coordinate system.
"""
reflect_xz(horseshoe :: Horseshoe) = (Horseshoe ∘ reflect_xz ∘ bound_leg)(horseshoe)

"""
    reflect_yz(panel :: Panel3D)

Reflect a Panel3D onto the ``x``-``z`` plane of its reference coordinate system.
"""
reflect_xz(panel :: Panel3D) = Panel3D((reflect_xz ∘ p1)(panel), (reflect_xz ∘ p2)(panel), (reflect_xz ∘ p3)(panel), (reflect_xz ∘ p4)(panel))

"""
    project_yz(line)

Project a Line onto the ``y``-``z`` plane of its reference coordinate system.
"""
project_yz(line :: Line) = Line((project_yz ∘ p1)(line), (project_yz ∘ p2)(line))

"""
    stability_flip(vector)

Reflect the x- and z- coordinates of a given 3-dimensional vector about the y-z and x-y planes respectively.
"""
stability_flip(vector) = SVector(-vector[1], vector[2], -vector[3])


struct Kpca
    X::Array{Float64}
    Kt::Array{Float64}
    T::Array{Float64}
    P::Array{Float64}
    sv::Vector{Float64}
    eig::Vector{Float64}
    D::Array{Float64}
    DKt::Array{Float64}
    vtot::Array{Float64}
    weights::Vector{Float64}
    kern
    dots
end

"""
    kpca(X, Y, weights = ones(size(X, 1)); nlv , kern = "krbf", kwargs...)
Kernel PCA  (Scholkopf et al. 1997, Scholkopf & Smola 2002, Tipping 2001).

* `X` : X-data.
* `weights` : vector (n,).
* `nlv` : Nb. principal components (PCs), or collection of nb. PCs, to consider.
* `kern` : Type of kernel used to compute the Gram matrices.
    Possible values are "krbf" of "kpol" (see respective functions `krbf` and `kpol`).
* `kwargs` : Named arguments to pass in the kernel function.

The method is implemented by SVD factorization of the weighted Gram matrix
D^(1/2) * Phi(`X`) * Phi(`X`)' * D^(1/2), where D is a diagonal matrix of weights for
the observations (rows of X).

The kernel Gram matrices are internally centered.

## References

Scholkopf, B., Smola, A., MÃ¼ller, K.-R., 1997. Kernel principal component analysis,
in: Gerstner, W., Germond, A., Hasler, M., Nicoud, J.-D. (Eds.), Artificial Neural Networks,
ICANN 97, Lecture Notes in Computer Science. Springer, Berlin, Heidelberg,
pp. 583-588. https://doi.org/10.1007/BFb0020217

Scholkopf, B., Smola, A.J., 2002. Learning with kernels: support vector machines, regularization,
optimization, and beyond, Adaptive computation and machine learning. MIT Press, Cambridge, Mass.

Tipping, M.E., 2001. Sparse kernel principal component analysis. Advances in neural information
processing systems, MIT Press. http://papers.nips.cc/paper/1791-sparse-kernel-principal-component-analysis.pdf

"""
function kpca(X, weights = ones(size(X, 1)); nlv, kern = "krbf", kwargs...)
    X = ensure_mat(X)
    n = size(X, 1)
    nlv = min(nlv, n)
    weights = mweights(weights)
    fkern = eval(Meta.parse(kern))
    K = fkern(X, X; kwargs...)     # In the future: fkern!(K, X, X; kwargs...)
    D = Diagonal(weights)
    Kt = K'
    DKt = D * Kt
    vtot = sum(DKt, dims = 1)
    Kc = K .- vtot' .- vtot .+ sum(D * DKt')
    sqrtD = sqrt.(D)
    Kd = sqrtD * Kc * sqrtD
    res = LinearAlgebra.svd(Kd)
    U = res.V[:, 1:nlv]
    eig = res.S
    eig[eig .< 0] .= 0
    sv = sqrt.(eig)
    P = sqrtD * scale(U, sv[1:nlv])     # In the future: scale!
    T = Kc * P       # T = Kc * P = D^(-1/2) * U * Delta
    Kpca(X, Kt, T, P, sv, eig, D, DKt, vtot, weights, kern, kwargs)
end

"""
    summary(object::Kpca, X)
Summarize the maximal (i.e. with maximal nb. PCs) fitted model.
* `object` : The fitted model.
* `X` : The X-data that was used to fit the model.
"""
function Base.summary(object::Kpca)
    nlv = size(object.T, 2)
    TT = object.D * (object.T).^2
    tt = vec(sum(TT, dims = 1))
    sstot = sum(object.eig)
    pvar = tt / sstot
    cumpvar = cumsum(pvar)
    explvar = DataFrame(pc = 1:nlv, var = tt, pvar = pvar, cumpvar = cumpvar)
    (explvar = explvar,)
end

"""
    transform(object::Kpca, X; nlv = nothing)
Compute PCs (scores matrix "T") from a fitted model and X-data.
* `object` : The maximal fitted model.
* `X` : X-data for which PCs are computed.
* `nlv` : Nb. PCs to consider. If nothing, it is the maximum nb. PCs.
"""
function transform(object::Kpca, X; nlv = nothing)
    a = size(object.T, 2)
    isnothing(nlv) ? nlv = a : nlv = min(nlv, a)
    fkern = eval(Meta.parse(object.kern))
    K = fkern(X, object.X; object.dots...)
    DKt = object.D * K'
    vtot = sum(DKt, dims = 1)
    Kc = K .- vtot' .- object.vtot .+ sum(object.D * object.DKt')
    T = Kc * @view(object.P[:, 1:nlv])
    T
end


include("inner_problem.jl")

# qcqp version of relaxation: it's faster than the socp version
# please see the raw.jl file for the dual problem to this relaxation
function portfolios_socp(thePortfolio::SparsePortfolioData)

  n =thePortfolio.f
  m =thePortfolio.m
  f =thePortfolio.n

  socpProblem = Model(solver=MosekSolver(MSK_DPAR_INTPNT_QO_TOL_PFEAS=1e-5, MSK_DPAR_INTPNT_QO_TOL_DFEAS=1e-5, MSK_IPAR_LOG=0, MSK_IPAR_MAX_NUM_WARNINGS=0))

  @variable(socpProblem, α[1:n])
  @variable(socpProblem, λ)
  @variable(socpProblem, t>=0)
  @variable(socpProblem, v[1:f]>=0.0)
  @variable(socpProblem, w[1:f])
  @variable(socpProblem, βl[1:m]>=0.0)
  @variable(socpProblem, βu[1:m]>=0.0)
  @variable(socpProblem, ρ[1:n]>=0.0)

  @constraint(socpProblem, w.>=thePortfolio.X'*α+thePortfolio.A'*(βl-βu)+λ*ones(f)+ρ-thePortfolio.d)
  @constraint(socpProblem, knorm[i=1:f], v[i]+t+ρ[i]*thePortfolio.min_investment[i]>=(thePortfolio.γ[i]/2.0)*w[i]^2)

  @objective(socpProblem, Max, (-0.5*α'*α+thePortfolio.Y'*α+λ+βl'*thePortfolio.l-βu'*thePortfolio.u-sum(v)-thePortfolio.k*t)') #+π'*thePortfolio.min_investment

  status=solve(socpProblem)

  return Dual(getvalue(α), getvalue(λ), getvalue(βl), getvalue(βu), getvalue(ρ), getvalue(w), getobjectivevalue(socpProblem), status), getobjectivevalue(socpProblem)
end

# SOCP version of relaxation. Not in use, as this version is currently slower.
function portfolios_socp2(thePortfolio::SparsePortfolioData)

  n = size(thePortfolio.X, 1)
  m = size(thePortfolio.A, 1)
  f = size(thePortfolio.X, 2)

  socpProblem = Model(solver=MosekSolver( MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-8, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-8))

  @variable(socpProblem, α[1:n])
  @variable(socpProblem, λ)
  @variable(socpProblem, t>=0)
  @variable(socpProblem, v[1:f]>=0.0)
  @variable(socpProblem, w[1:f])
  @variable(socpProblem, βl[1:m]>=0.0)
  @variable(socpProblem, βu[1:m]>=0.0)
  @variable(socpProblem, ρ[1:n]>=0.0)
  @variable(socpProblem, τ>=0.0)

  @constraint(socpProblem, w.>=thePortfolio.X'*α+thePortfolio.A'*(βl-βu)+λ*ones(f)+ρ-thePortfolio.d)
  @constraint(socpProblem, knorm[i=1:f], v[i]+t+ρ[i]*thePortfolio.min_investment[i]+1.0>=norm([sqrt(2.0*thePortfolio.γ[i])*w[i];v[i]+t+ρ[i]*thePortfolio.min_investment[i]-1.0]))
  @constraint(socpProblem, τ+1.0>=norm([2.0*α;τ-1.0]))

  @objective(socpProblem, Max, (-0.5*τ+thePortfolio.Y'*α+λ+βl'*thePortfolio.l-βu'*thePortfolio.u-sum(v)-thePortfolio.k*t)') #+π'*thePortfolio.min_investment

  status=solve(socpProblem)

  return Dual(getvalue(α), getvalue(λ), getvalue(βl), getvalue(βu), getvalue(ρ), getvalue(w), status, getobjectivevalue(socpProblem)), getobjectivevalue(socpProblem)
end


function departarrive_eccentricity_objective(e, a, Ω, i, r̄ₒ, rₛₚ, μ, v̄rel::AbstractVector{<:Real}, c)
    # for a given parking position, and velocity at SoI, and eccentricity, define a departure/arrival orbit
    rₒ = norm(r̄ₒ)
    θₒ = c * orbital_angle(rₒ, a, e)
    θₛₚ = c * orbital_angle(rₛₚ, a, e)
    ω = wrap_angle(angle_in_plane(r̄ₒ, basis_rotation(Ω, i, 0.))-θₒ)

    v̄ₛₚ = orbital_velocity(θₛₚ,a,e,μ,i,Ω,ω)

    return 1-dot(v̄ₛₚ./norm(v̄ₛₚ), v̄rel./norm(v̄rel)), ω
end

function departarrive_true_anomaly(r̄ₒ::AbstractVector{<:Real}, prim::Union{CelestialBody,Star}, v̄rel::AbstractVector{<:Real}, tₛₚ, c)
    # optimize the eccentricity of the departure/arrival orbit against the departure/arrival vector
    μ = prim.μ
    rₒ = norm(r̄ₒ)
    rₛₚ = prim.SoI
    ĥ = normalize(cross(r̄ₒ,v̄rel))
    n̂ = normalize(cross(@SVector([0.,0.,1.]), ĥ))
    if isnan(norm(n̂))
        n̂ = @SVector([1.,0.,0.])
    end
    a =  1 / (2/rₛₚ - sum(abs2,v̄rel)/μ)
    i = wrap_acos(ĥ[3])
    Ω = wrap_angle(copysign(1,n̂[2]) * acos(n̂[1]))

    if a > 0
        lb = 0.
        ub = 1. - eps(Float64)
    else
        lb = 1. + eps(Float64)
        ub = max(1-rₒ/a, rₒ/a-1)
        ub = ub == 1 ? 1. + 2*eps(Float64) : ub
    end

    objectivefun(e) = departarrive_eccentricity_objective(e, a, Ω, i, r̄ₒ, rₛₚ, μ, v̄rel, c)[1]
    res = optimize(objectivefun, lb, ub, Brent())
    e = Optim.minimizer(res)

    eccobj, ω = departarrive_eccentricity_objective(e, a, Ω, i, r̄ₒ, rₛₚ, μ, v̄rel, c)

    θₒ = c * orbital_angle(rₒ, a, e)
    θₛₚ = c * orbital_angle(rₛₚ, a, e)
    T = period(a,μ)
    Mₒ = true_to_mean(θₒ, e)
    Δt = true_to_time(θₛₚ, e, T) - true_to_time(θₒ, e, T)
    tₒ = tₛₚ - Δt

    return Orbit(a,e,i,Ω,ω,Mₒ,tₒ,prim), eccobj, 0.
end

function departarrive_true_anomaly(θ, pkorb::Orbit{<:CelestialBody}, v̄rel::AbstractVector{<:Real}, tₛₚ, c; match_pkorb_Mo=true)
    r̄ₒ = orbital_position(θ, pkorb)
    daorb, eccobj = departarrive_true_anomaly(r̄ₒ, pkorb.primary, v̄rel::AbstractVector{<:Real}, tₛₚ, c)[1:2]

    a, e, i, Ω, ω, Mₒ = daorb.a, daorb.e, daorb.i, daorb.Ω, daorb.ω, daorb.Mo

    tₒ = match_pkorb_Mo ? true_to_time(θ, pkorb, daorb.epoch - pkorb.period/2) : daorb.epoch

    v̄ₒ = time_orbital_velocity(daorb.epoch, daorb)
    v̄ₚₖ = orbital_velocity(θ, pkorb)
    Δv̄ = v̄ₒ .- v̄ₚₖ

    return Orbit(a,e,i,Ω,ω,Mₒ,tₒ,pkorb.primary), eccobj, Δv̄
end

function departarrive_objectivefun(θ, pkorb, v̄rel, tₛₚ, c)
    err, Δv̄ = departarrive_true_anomaly(θ, pkorb, v̄rel, tₛₚ, c)[2:3]
    return norm(Δv̄) * exp(1000*err)        # penalize mismatched direction
end

function departarrive_orbit(pkorb::Orbit{<:CelestialBody}, v̄rel::AbstractVector{<:Real}, tₛₚ; out=true, match_pkorb_Mo=true)
    c = out ? 1 : -1
    # optimize dv against true anomaly, with each true anomaly being optimized for eccentricity
    if pkorb.e < 1           # elliptical case: any true anomaly
        lb = -2π
        ub = 2π
    else                    # hyperbolic case: true anomalies occuring within the SoI
        ub = orbital_angle(pkorb.primary.SoI, pkorb.a, pkorb.e) - eps(Float64)
        lb = -ub
    end

    objectivefun(θ) = departarrive_objectivefun(θ, pkorb, v̄rel, tₛₚ, c)

    res = optimize(objectivefun, lb, ub, Brent())
    θₚₖ = Optim.minimizer(res)

    daorb, eccobj, Δv̄ = departarrive_true_anomaly(θₚₖ, pkorb, v̄rel, tₛₚ, c; match_pkorb_Mo=match_pkorb_Mo)

    if eccobj > 1e-3
        @warn "departure/arrival velocity mismatch from target direction: $(pkorb.primary.name), $eccobj"
    end
    return daorb, eccobj, Δv̄
end

function fast_departarrive_orbit(pkorb::Orbit{<:CelestialBody}, v̄rel::AbstractVector{<:Real}, tₛₚ; out=true,  match_pkorb_Mo=true, tol=0.1, maxit=50)
    c = out ? 1 : -1
    μ = pkorb.primary.μ
    rₛₚ = pkorb.primary.SoI
    vₛₚ = norm(v̄rel)
    err = tol+1
    it = 0
    rₒnext = pkorb.a
    a = 1 / (2/rₛₚ - vₛₚ^2/μ)

    e, θₛₚ = 0., 0.
    r̄ₒ = @SVector zeros(Float64, 3)
    v̄ₒ = @SVector zeros(Float64, 3)
    while abs(err) > tol && it < maxit
        it = it+1
        rₒ = rₒnext
        vₒ = √(vₛₚ^2 + 2(μ/rₒ - μ/rₛₚ))
        e = √(1 + 2((vₒ^2)/2 - μ/rₒ) * rₒ^2 * vₒ^2 / μ^2)
        θₛₚ = c * orbital_angle(rₛₚ, a, e)
        ϕₛₚ = orbital_direction(θₛₚ, e)
        v̄ₛₚ = vₛₚ * @SVector([cos(ϕₛₚ),
                            sin(ϕₛₚ),
                            0])

        r̄ₒ = @SVector([rₒ, 0, 0])
        v̄ₒ = @SVector([0, vₒ, 0])
        v̄relplane = inertial_to_perifocal_bases(v̄rel, pkorb)
        # R = align_vectors(v̄ₛₚ, v̄relplane)               # align velocity at SoI with input direction,
        #                                                # valid only for undefined parking orbit
        ψ₁ = atan(v̄relplane[3], √abs(vₛₚ^2 - v̄ₛₚ[1]^2 - v̄relplane[3]^2))
        # R₁ = MRP([1        0        0;                   # rotate around x-axis to match i
        #           0  cos(ψ₁) -sin(ψ₁);
        #           0  sin(ψ₁)  cos(ψ₁)])
        R₁ = MRP((1.,0.,0.,0.,cos(ψ₁),sin(ψ₁),0.,-sin(ψ₁),cos(ψ₁)))
        R1v̄ₛₚ = R₁*v̄ₛₚ
        ψ₂ = atan(v̄relplane[2], v̄relplane[1]) - atan(R1v̄ₛₚ[2], R1v̄ₛₚ[1])
        # R₂ = MRP([cos(ψ₂) -sin(ψ₂)  0;                   # rotate around z-axis to match ω
        #           sin(ψ₂)  cos(ψ₂)  0;
        #           0        0        1]))
        R₂ = MRP((cos(ψ₂),sin(ψ₂),0.,-sin(ψ₂),cos(ψ₂),0.,0.,0.,1.))

        v̄ₒ = perifocal_to_inertial_bases(R₂*R₁*v̄ₒ, pkorb)
        r̄ₒ = perifocal_to_inertial_bases(R₂*R₁*r̄ₒ, pkorb)

        δθ = angle_in_plane(r̄ₒ, pkorb)
        r̄ₒpkorb = state_vector(δθ, pkorb)[1]

        errprev = err
        err = norm(r̄ₒ) - norm(r̄ₒpkorb)
        if abs(err/errprev) > 0.9
            rₒnext = (norm(r̄ₒpkorb) + rₒ)/2
        else
            rₒnext = norm(r̄ₒpkorb)
        end
    end
    Δt = true_to_time(θₛₚ, e, period(a, μ))
    θₚₖ = angle_in_plane(r̄ₒ, pkorb)
    tₚₖ =  match_pkorb_Mo ? true_to_time(θₚₖ, pkorb, tₛₚ-Δt-pkorb.period/2) : tₛₚ - Δt
    v̄ₚₖ = orbital_velocity(θₚₖ, pkorb)
    Δv̄ = c*(v̄ₒ .- v̄ₚₖ)
    return Orbit(tₚₖ, r̄ₒ, v̄ₒ, pkorb.primary), Δv̄
end

"""
    orb = departure_orbit(pkorb, v̄rel, tₛₚ)

Computes an optimized departure orbit from the parking orbit `pkorb`, the relative velocity
at the patch `v̄rel`, the time at the patch `tₛₚ`.
"""
departure_orbit(pkorb::Orbit{<:CelestialBody}, v̄rel::AbstractVector{<:Real}, tₛₚ; kwargs...) = departarrive_orbit(pkorb, v̄rel, tₛₚ; out=true,  kwargs...)

"""
    orb = arrival_orbit(pkorb, v̄rel, tₛₚ)

Computes an optimized arrival orbit from the parking orbit `pkorb`, the relative velocity
at the patch `v̄rel`, the time at the patch `tₛₚ`.
"""
arrival_orbit(pkorb::Orbit{<:CelestialBody}, v̄rel::AbstractVector{<:Real}, tₛₚ; kwargs...)   = departarrive_orbit(pkorb, v̄rel, tₛₚ; out=false, kwargs...)

"""
    orb = fast_departure_orbit(pkorb, v̄rel, tₛₚ)

Computes a departure orbit from the parking orbit `pkorb`, the relative velocity
at the patch `v̄rel`, the time at the patch `tₛₚ`.

This function is faster than `departure_orbit`, but may return sub-optimal trajectories
for highly elliptical parking orbits.
"""
fast_departure_orbit(pkorb::Orbit{<:CelestialBody}, v̄rel::AbstractVector{<:Real}, tₛₚ; kwargs...) = fast_departarrive_orbit(pkorb, v̄rel, tₛₚ; out=true, kwargs...)

"""
    orb = fast_arrival_orbit(pkorb, v̄rel, tₛₚ)

Computes a arrival orbit from the parking orbit `pkorb`, the relative velocity
at the patch `v̄rel`, the time at the patch `tₛₚ`.

This function is faster than `arrival_orbit`, but may return sub-optimal trajectories
for highly elliptical parking orbits.
"""
fast_arrival_orbit(pkorb::Orbit{<:CelestialBody}, v̄rel::AbstractVector{<:Real}, tₛₚ; kwargs...)   = fast_departarrive_orbit(pkorb, v̄rel, tₛₚ; out=false, kwargs...)

# true anomaly/time at the SoI
patch_angle(daorb::Orbit, c) = c*orbital_angle(daorb.primary.SoI, daorb)
patch_time(daorb::Orbit, c) = true_to_time(patch_angle(daorb, c), daorb)

"""
    θ = ejection_angle(dorb)

Computes the true anomaly at SoI ejection for a departure orbit `dorb`.
"""
ejection_angle(dorb::Orbit)  = patch_angle(dorb,  1)

"""
    t = ejection_time(dorb)

Computes the time at SoI ejection for a departure orbit `dorb`.
"""
ejection_time(dorb::Orbit)   = patch_time(dorb,  1)

"""
    θ = insertion_angle(aorb)

Computes the true anomaly at SoI insertion for a arrival orbit `aorb`.
"""
insertion_angle(aorb::Orbit) = patch_angle(aorb, -1)

"""
    t = insertion_angle(aorb)

Computes the time at SoI insertion for a arrival orbit `aorb`.
"""
insertion_time(aorb::Orbit)  = patch_time(aorb, -1)

patch_position(daorb::Orbit, c) = orbital_position(patch_angle(daorb, c), daorb)
patch_velocity(daorb::Orbit, c) = orbital_velocity(patch_angle(daorb, c), daorb)
patch_state_vector(daorb::Orbit, c) = state_vector(patch_angle(daorb, c), daorb)

"""
    r̄ = ejection_position(dorb)

Computes the position vector at SoI ejection for a departure orbit `dorb`.
"""
ejection_position(dorb::Orbit) = orbital_position(ejection_angle(dorb), dorb)

"""
    v̄ = ejection_velocity(dorb)

Computes the velocity vector at SoI ejection for a departure orbit `dorb`.
"""
ejection_velocity(dorb::Orbit) = orbital_velocity(ejection_angle(dorb), dorb)

"""
    r̄, v̄ = ejection_state_vector(dorb)

Computes the state vector at SoI ejection for a departure orbit `dorb`.
"""
ejection_state_vector(dorb::Orbit) = state_vector(ejection_angle(dorb), dorb)

"""
    r̄ = insertion_position(aorb)

Computes the position vector at SoI insertion for a arrival orbit `aorb`.
"""
insertion_position(aorb::Orbit) = orbital_position(insertion_angle(aorb), aorb)

"""
    v̄= insertion_velocity(aorb)

Computes the velocity vector at SoI insertion for a arrival orbit `aorb`.
"""
insertion_velocity(aorb::Orbit) = orbital_velocity(insertion_angle(aorb), aorb)

"""
    r̄, v̄ = insertion_state_vector(aorb)

Computes the state vector at SoI insertion for a arrival orbit `aorb`.
"""
insertion_state_vector(aorb::Orbit) = state_vector(insertion_angle(aorb), aorb)


using Soss

μdist = @model begin
    a ~ Normal()
    b ~ Normal()
    return a/b
end

σdist = @model begin
    x ~ Normal()
    return x^2
end

m = @model begin
    μ ~ μdist()
    σ ~ σdist()
    x ~ Normal(μ,σ) |> iid(2)
    return x
end

rand(m(),5)

simulate(m(), 5)

dynamicHMC(m() | (x = randn(2),))