# Shallow ice approximation (SIA) implicit solver for Greenland (steady state)

using JLD, Plots, Printf, LinearAlgebra

# enable plotting & saving by default
if !@isdefined do_visu; do_visu = true end
if !@isdefined do_save; do_save = true end

# finite difference stencil operation support functions
@views av(A)    = 0.25*(A[1:end-1,1:end-1].+A[2:end,1:end-1].+A[1:end-1,2:end].+A[2:end,2:end]) # average
@views av_xa(A) = 0.5.*(A[1:end-1,:].+A[2:end,:]) # average x-dir
@views av_ya(A) = 0.5.*(A[:,1:end-1].+A[:,2:end]) # average y-dir
@views inn(A)   = A[2:end-1,2:end-1] # inner points

@views function iceflow(dx, dy, Zbed, Hice, Mask, grad_b, z_ELA, b_max)
    println("Initialising ice flow model ... ")
    # physics
    s2y      = 3600*24*365.25  # seconds to years
    rho_i    = 910.0           # ice density
    g        = 9.81            # gravity acceleration
    npow     = 3.0             # Glen's power law exponent
    a0       = 1.5e-24         # Glen's law enhancement term
    # numerics
    @assert (dx>0 && dy>0) "dx and dy need to be positive"
    nx, ny   = size(Zbed,1), size(Zbed,2) # numerical grid resolution
    @assert (nx, ny) == size(Zbed) == size(Hice) == size(Mask) "Sizes don't match"
    itMax    = 1e5             # number of iteration (max)
    nout     = 200             # error check frequency
    tolnl    = 1e-6            # nonlinear tolerance
    epsi     = 1e-4            # small number
    damp     = 0.85            # convergence accelerator (this is a tuning parameter, dependent on e.g. grid resolution)
    dtausc   = 1.0/3.0         # iterative dtau scaling
    # derived physics
    a        = 2.0*a0/(npow+2)*(rho_i*g)^npow*s2y
    # derived numerics
    cfl      = max(dx^2,dy^2)/4.1
    # array initialization
    Err      = zeros(nx  , ny  )
    dSdx     = zeros(nx-1, ny  )
    dSdy     = zeros(nx  , ny-1)
    gradS    = zeros(nx-1, ny-1)
    D        = zeros(nx-1, ny-1)
    qHx      = zeros(nx-1, ny-2)
    qHy      = zeros(nx-2, ny-1)
    dtau     = zeros(nx-2, ny-2)
    ResH     = zeros(nx-2, ny-2)
    dHdt     = zeros(nx-2, ny-2)
    Vx       = zeros(nx-1, ny-1)
    Vy       = zeros(nx-1, ny-1)
    M        = zeros(nx  , ny  )
    B        = zeros(nx  , ny  )
    H        = zeros(nx  , ny  )
    S        = zeros(nx  , ny  )
    # initial conditions
    B       .= Zbed
    H       .= Hice
    S       .= B .+ H
    # iteration loop
    println(" starting iteration loop:")
    iter = 1; err = 2*tolnl
    while err>tolnl && iter<itMax
        Err   .= H
        # mass balance
        M     .= min.(grad_b.*(S .- z_ELA), b_max)
        # compute diffusivity
        dSdx  .= diff(S, dims=1)/dx
        dSdy  .= diff(S, dims=2)/dy
        gradS .= sqrt.(av_ya(dSdx).^2 .+ av_xa(dSdy).^2)
        D     .= a*av(H).^(npow+2) .* gradS.^(npow-1)
        # compute flux
        qHx   .= .-av_ya(D).*diff(S[:,2:end-1], dims=1)/dx
        qHy   .= .-av_xa(D).*diff(S[2:end-1,:], dims=2)/dy
        # update ice thickness
        dtau  .= dtausc*min.(10.0, cfl./(epsi .+ av(D)))
        ResH  .= .-(diff(qHx, dims=1)/dx .+ diff(qHy, dims=2)/dy) .+ inn(M)
        dHdt  .= dHdt.*damp .+ ResH
        H[2:end-1,2:end-1] .= max.(0.0, inn(H) .+ dtau.*dHdt)
        # apply mask (a very poor-man's calving law)
        H[Mask.==0] .= 0.0
        # update surface
        S     .= B .+ H
        # error check
        if mod(iter, nout)==0
            Err .= Err .- H
            err = norm(Err)/length(Err)
            @printf(" iter = %d, error = %1.2e \n", iter, err)
            if isnan(err)
                error("""NaNs encountered.  Try a combination of:
                             decreasing `damp` and/or `dtausc`, more smoothing steps""")
            end
        end
        iter += 1
    end
    # compute velocities
    Vx .= -D./(av(H) .+ epsi).*av_ya(dSdx)
    Vy .= -D./(av(H) .+ epsi).*av_xa(dSdy)
    # return as GeoArrays
    return  as_geoarray(H,  Zbed, name=:thickness),
            as_geoarray(S,  Zbed, name=:surface),
            as_geoarray(M,  Zbed, name=:smb),
            as_geoarray(Vx, Zbed, name=:vel_x, staggerd=true),
            as_geoarray(Vy, Zbed, name=:vel_y, staggerd=true)
end
# ------------------------------------------------------------------------------
include(joinpath(@__DIR__, "helpers.jl"))

# load the data
print("Loading the data ... ")
Zbed, Hice, Mask, dx, dy, xc, yc = load_data(; nx=96) # nx=96,160 are included in the repo
                                                      # other numbers will trigger a 2GB download
println("done.")

# apply some smoothing
print("Applying some smoothing ... ")
for is=1:2 # two smoothing steps
    smooth!(Zbed)
    smooth!(Hice)
end
println("done.")

# calculate mass balance coefficients for given spatial grid
grad_b, z_ELA, b_max = mass_balance_constants(xc, yc)

# run the SIA flow model
H, S, M, Vx, Vy = iceflow(dx, dy, Zbed, Hice, Mask, grad_b, z_ELA, b_max)

# visualization and save
nx, ny = size(H)
if do_visu
    !ispath("../output") && mkdir("../output")

    H_v = copy(H); H_v[Mask.==0].=NaN
    Hice_v = copy(Hice); Hice_v[Mask.==0].=NaN
    S_v = copy(S); S_v[Mask.==0].=NaN
    M_v = copy(M); M_v[Mask.==0].=NaN
    V_v = sqrt.(Vx.^2 .+ Vy.^2)

    # outputs
    fontsize  = 7
    opts = (aspect_ratio=1, yaxis=font(fontsize, "Courier"), xaxis=font(fontsize, "Courier"),
            ticks=nothing, framestyle=:box, titlefontsize=fontsize, titlefont="Courier", colorbar_title="",
            xlabel="", ylabel="", xlims=(dims(H_v)[1][1],dims(H_v)[1][end]), ylims=(dims(H_v)[2][end],dims(H_v)[2][1]) )
    p1 = heatmap(S_v; c=:davos, title="Surface elev. [m]", opts...)
    p2 = heatmap(H_v; c=:davos, title="Ice thickness [m]", opts...)
    p3 = heatmap(log10.(V_v); clims=(0.1, 2.0), title="log10(vel) [m/yr]", opts...)
    p4 = heatmap(M_v; c=:devon, title="Mass Bal. rate [m/yr]", opts...)
    p = plot(p1, p2, p3, p4, size=(400,400), dpi=200) #background_color=:transparent, foreground_color=:white
    ## uncomment if you want a pop-up plot pane showing:
    # display(p)
    savefig("../output/iceflow_out1.png")

    # error
    H_diff = Hice_v.-H_v
    fontsize = 7
    p1 = heatmap(Hice_v; c=:davos, title="H data [m]", opts...)
    p2 = heatmap(H_v; c=:davos, title="H model [m]", opts...)
    clim = max(abs.(extrema(H_diff[.!isnan.(H_diff)]))...)
    p3 = heatmap(H_diff; title="H (data-model) [m]",  clims=(-clim,clim), seriescolor=:balance, opts...)
    p = plot(p1, p2, p3, layout=(1, 3), size=(500,160), dpi=200) #background_color=:transparent, foreground_color=:white
    ## uncomment if you want a pop-up plot pane showing:
    # display(p)
    savefig("../output/iceflow_out2.png")
end

if do_save
    save("../output/iceflow_$(nx)x$(ny).jld", "Hice", Hice,
                                              "Mask", Mask,
                                              "H"   , H,
                                              "S"   , S,
                                              "M"   , M,
                                              "Vx"  , Vx,
                                              "Vy"  , Vy,
                                              "xc", xc, "yc", yc)
end

println("... done.")


module BSpline

export
  bspline_basismatrix!, bspline_basismatrix,
  derivative_bspline_basismatrix!, derivative_bspline_basismatrix,
  derivative2_bspline_basismatrix!, derivative2_bspline_basismatrix,
  bspline_basis, derivative_bspline_basis, derivative2_bspline_basis,
  num_basis,
  BSplineT

type BSplineT{T<:FloatingPoint}
  M::Int64                  # B-spline order (2 for linear, 3 for quadratic, etc.)
  knots::Vector{T}          # knot vector
end

BSplineT{T<:FloatingPoint}(M::Int64, knots::Vector{T}) = BSplineT{T}(M, knots)
# BSplineT{T<:FloatingPoint}(M::Int64, knots::LinSpace{T}) = BSplineT{T}(M, [knots;])

function BSplineT{T<:FloatingPoint}(M::Int64, lb::T, ub::T, num_knots::Int64)
  h = (ub - lb) / (num_knots - 1)

  BSplineT{T}(M,
             vcat(lb .- h * [M-1:-1:1;],
                  linspace(lb, ub, num_knots),
                  ub .+ h * [1:M-1;]))
end

num_basis{T<:FloatingPoint}(bs::BSplineT{T}) = length(bs.knots) - bs.M

# B-spline basis function value B(j,n) at x.
#
# Input arguments:
# j:   interval index, 0 =< j < num_basis
#
# x:   value where the basis function is to be evaluated
#
# Output arguments:
# y:
#    B-spline basis function value, nonzero for a knot span of n

function bspline_basis{T<:FloatingPoint}(j::Integer, bs::BSplineT{T}, x::T)
  n = bs.M
  t = bs.knots
  y = zero(T)
  if n > 1
    dn = x - t[j+1]
    dd = t[j+n] - t[j+1]
    if dd != zero(T)  # indeterminate forms 0/0 are deemed to be zero
      b = bspline_basis(j, BSplineT(n-1,t), x)
      y = y + b*(dn/dd)
    end

    dn = t[j+n+1] - x
    dd = t[j+n+1] - t[j+2]
    if dd != zero(T)
      b = bspline_basis(j+1, BSplineT(n-1,t), x)
      y = y + b*(dn/dd)
    end
  else
    y = t[j+1] <= x < t[j+2] ? one(T) : zero(T)
  end
  y
end

# B-spline basis function value matrix B(n) for x.
#
# Input arguments:
# n:
#    B-spline order (2 for linear, 3 for quadratic, etc.)
# t:
#    knot vector
# x :
#    an m-dimensional vector of values where the basis function is to be
#    evaluated
#
# Output arguments:
# out:
#    a matrix of m rows and numel(t)-n columns

function bspline_basismatrix!{T<:FloatingPoint}(
    out::StridedMatrix{T}, bs::BSplineT{T}, x::Vector{T}
    )
  for j = 1 : num_basis(bs)
    for k = 1 : length(x)
      out[k, j] = bspline_basis(j-1, bs, x[k])
    end
  end
  out
end

bspline_basismatrix{T<:FloatingPoint}(bs::BSplineT{T}, x::T) =
  bspline_basismatrix!(Array(T, (1, num_basis(bs))), bs, [x;])

bspline_basismatrix{T<:FloatingPoint}(bs::BSplineT{T}, x::Vector{T}) =
  bspline_basismatrix!(Array(T, (length(x), num_basis(bs))), bs, x)

# derivatives

function derivative_bspline_basis{T<:FloatingPoint}(j::Integer, bs::BSplineT{T}, x::T)
  n = bs.M
  t = bs.knots

  y = zero(T)
  if n > 1
    dn = n - 1
    dd = t[j+n] - t[j+1]
    if dd != zero(T)  # indeterminate forms 0/0 are deemed to be zero
      b = bspline_basis(j, BSplineT(n-1, t), x)
      y = y + b*(dn/dd)
    end

    dn = 1 - n
    dd = t[j+n+1] - t[j+2]
    if dd != zero(T)
      b = bspline_basis(j+1, BSplineT(n-1, t), x)
      y = y + b*(dn/dd)
    end
  else
    y = t[j+1] <= x < t[j+2] ? one(T) : zero(T)
  end
  y
end

function derivative_bspline_basismatrix!{T<:FloatingPoint}(
        out::StridedMatrix{T}, bs::BSplineT{T}, x::Vector{T}
        )
  for j = 1 : num_basis(bs)
    for k = 1 : length(x)
      out[k,j] = derivative_bspline_basis(j-1,bs,x[k])
    end
  end
  out
end

derivative_bspline_basismatrix{T<:FloatingPoint}(bs::BSplineT{T}, x::Vector{T}) =
  derivative_bspline_basismatrix!(Array(T, (length(x), num_basis(bs))), bs, x)

derivative_bspline_basismatrix{T<:FloatingPoint}(bs::BSplineT{T}, x::T) =
  derivative_bspline_basismatrix!(Array(T, (1, num_basis(bs))), bs, [x;])

# second derivative

function derivative2_bspline_basis{T<:FloatingPoint}(j::Integer, bs::BSplineT{T}, x::T)
  n = bs.M
  t = bs.knots
  y = zero(T)
  if n > 1
    dn = n - 1
    dd = t[j+n] - t[j+1]
    if dd != zero(T)  # indeterminate forms 0/0 are deemed to be zero
      b = derivative_bspline_basis(j, BSplineT(n-1, t), x)
      y = y + b*(dn/dd)
    end

    dn = 1 - n
    dd = t[j+n+1] - t[j+2]
    if dd != zero(T)
      b = derivative_bspline_basis(j+1, BSplineT(n-1, t), x)
      y = y + b*(dn/dd)
    end
  else
    y = t[j+1] <= x < t[j+2] ? one(T) : zero(T)
  end
  y
end

function derivative2_bspline_basismatrix!{T<:FloatingPoint}(
            out::StridedMatrix{T}, bs::BSplineT{T}, x::Vector{T}
            )
  for j = 1 : num_basis(bs)
    for k = 1 : length(x)
      out[k,j] = derivative2_bspline_basis(j-1, bs, x[k])
    end
  end
end

derivative_bspline_basismatrix{T<:FloatingPoint}(bs::BSplineT{T}, x::Vector{T}) =
  derivative2_bspline_basismatrix!(Array(T, (length(x), num_basis(bs))), bs, x)

derivative_bspline_basismatrix{T<:FloatingPoint}(bs::BSplineT{T}, x::T) =
  derivative2_bspline_basismatrix!(Array(T, (1, num_basis(bs))), bs, [x;])

end


using Test

using MarketData
using TimeSeries

using MarketTechnicals

@testset "Levels" begin

@testset "floor pivots" begin
    # values verified by various website calculators
    @test isapprox(values(floorpivots(ohlc)[:r3])[1]   , 123.310, atol=.01)
    @test isapprox(values(floorpivots(ohlc)[:r2])[1]   , 119.52, atol=.01)
    @test isapprox(values(floorpivots(ohlc)[:r1])[1]   , 115.73, atol=.01)
    @test isapprox(values(floorpivots(ohlc)[:pivot])[1], 108.71, atol=.01)
    @test isapprox(values(floorpivots(ohlc)[:s1])[1]   , 104.92, atol=.01)
    @test isapprox(values(floorpivots(ohlc)[:s2])[1]   , 97.900, atol=.01)
    @test isapprox(values(floorpivots(ohlc)[:s3])[1]   , 94.110, atol=.01)
    @test timestamp(floorpivots(ohlc))[end] == Date(2001,12,31)
end

@testset "woodiespivots" begin
    # @test_approx_eq  97.37500000000001 value(wr4)[2]   # values NEED to be verified with online calculators
    # @test_approx_eq   88.62500000000001 value(ws4)[2]
    @test isapprox(values(woodiespivots(ohlc)[:r3])[1]   , 124.465, atol=.01)
    @test isapprox(values(woodiespivots(ohlc)[:r2])[1]   , 118.480, atol=.01)
    @test isapprox(values(woodiespivots(ohlc)[:r1])[1]   , 113.655, atol=.01)
    @test isapprox(values(woodiespivots(ohlc)[:pivot])[1], 107.670, atol=.01)
    @test isapprox(values(woodiespivots(ohlc)[:s1])[1]   , 102.845, atol=.01)
    @test isapprox(values(woodiespivots(ohlc)[:s2])[1]   , 96.8625, atol=.01)
    @test isapprox(values(woodiespivots(ohlc)[:s3])[1]   , 92.035, atol=.01)
    @test timestamp(woodiespivots(ohlc))[end] == Date(2001,12,31)
end

end  # @testset "Levels"


#=
Visual illustration of the central limit theorem in 3d

@author : Spencer Lyon <spencer.lyon@nyu.edu>
          Victoria Gregory <victoria.gregory@nyu.edu>

References
----------
Based off the original python  file clt3d.py
=#
using Plots
pyplot()
using Distributions
using KernelDensity

beta_dist = Beta(2.0, 2.0)

function gen_x_draws(k)
    bdraws = rand(beta_dist, 3, k)

    # == Transform rows, so each represents a different distribution == #
    bdraws[1, :] -= 0.5
    bdraws[2, :] += 0.6
    bdraws[3, :] -= 1.1

    # == Set X[i] = bdraws[j, i], where j is a random draw from {1, 2, 3} == #
    js = rand(1:3, k)
    X = Array{Float64}(k)
    for i=1:k
        X[i]=  bdraws[js[i], i]
    end

    # == Rescale, so that the random variable is zero mean == #
    m, sigma = mean(X), std(X)
    return (X .- m) ./ sigma
end

nmax = 5
reps = 100000
ns = 1:nmax

# == Form a matrix Z such that each column is reps independent draws of X == #
Z = Array{Float64}(reps, nmax)
for i=ns
    Z[:, i] = gen_x_draws(reps)
end

# == Take cumulative sum across columns
S = cumsum(Z, 2)

# == Multiply j-th column by sqrt j == #
Y = S .* (1. ./ sqrt(ns))'

# == Plot == #
a, b = -3, 3
gs = 100
xs = linspace(a, b, gs)

x_vec = []
y_vec = []
z_vec = []
colors = []
for n=ns
    kde_est = kde(Y[:, n])
    _xs, ys = kde_est.x, kde_est.density
    push!(x_vec, collect(_xs))
    push!(y_vec, ys)
    push!(z_vec, collect(n*ones( length(_xs))))
    push!(colors, RGBA(0, 0, 0, 1-(n-1)/nmax))
end

plot(x_vec, z_vec, y_vec, color=reshape(colors,1,length(colors)), legend=:none)
plot!(xlims=(a,b), xticks=[-3; 0; 3])
plot!(ylims=(1, nmax), yticks=ns, ylabel="n")
plot!(zlims=(0, 0.4), zticks=[0.2; 0.4])


using ClassifierMetrics
using DataFrames
using CSV
using Test
using LinearAlgebra
using Random
using Plots
Random.seed!(1234)

#utils
add_noise(label::Bool, λ=0.0) = label ? 1 - λ*rand() : λ*rand()

@testset "ClassifierMetrics" begin
  @testset "Binrary Metrics" begin
    include("binarymetrics.jl")
  end
  @testset "testing ROC" begin
    include("roc.jl")
  end
  @testset "testing PR" begin
    include("pr.jl")
  end
  @testset "testing DET" begin
    include("det.jl")
  end
  @testset "testing Plots" begin
    include("plots.jl")
  end
end


# ---
# title: Custom median filters
# cover: assets/median.gif
# author: Johnny Chen
# date: 2020-09-23
# ---

# With median filter as an example, this demo shows how you could construct a custom kernel and
# pass it to [`mapwindow`](@ref) for common stencil operations.

# Unlike other image processing toolboxes, `ImageFiltering.jl` does not provide you an out-of-box
# function for median filters. It instead provides the more general [`mapwindow`](@ref) function,
# which allows you to apply any function to the window around each pixel.

using ImageFiltering, ImageCore, ImageShow # or you could just `using Images`
using TestImages
using Statistics
using Random #hide
Random.seed!(0) #hide

img = Float64[isodd(i) - isodd(j) for i = 1:5, j = 1:5]
img[3, 3] = 1000
img #hide #md

#-

patch_size = (3, 3)
imgm = mapwindow(median, img, patch_size)

# `mapwindow` provides a high-level interface to loop over the image and call a function (`median`
# in this demo) on each patch/window. It's performing an operation that is nearly equivalent to the
# loop below, albeit more efficiently:

## For simplicity, border condition is not included here.
imgm = zeros(axes(img))
R = CartesianIndices(img)
I_first, I_last = first(R), last(R)
Δ = CartesianIndex(patch_size .÷ 2)
for I in R
    patch = max(I_first, I-Δ):min(I_last, I+Δ)
    imgm[I] = median(img[patch])
end
imgm

# Compared to this hand-written loop, `mapwindow` offers additional flexibility in handling the
# borders and allows you to process only a subset of the windows. For some functions (e.g., `min`,
# `max`, and `extrema`), `mapwindow` uses a custom implementation that is far more efficient than
# naively computing the function "freshly" on each window.

# When the input array has `NaN` or some unwanted pixel values, a pre-filtering process is needed to
# exclude them first. This can be done quite easily by compositing a given kernel operation and a
# `filter` operation. Let's still take `median` filter as an example.

img[2,2] = NaN
imgm = mapwindow(median, img, (3,3))

# `NaN` has polluted the output, which is usually not wanted. This can be fixed quite easily by
# compositing a new filter kernel. This way, we only compute the median value on the non-NaN subset
# of each patch/window.

_median(patch) = median(filter(x->!isnan(x), patch))
imgm = mapwindow(_median, img, patch_size)

# ## Impulse (Salt and Pepper) noise removal

# Median filters are quite robust to outliers (unusual values), this property can be used to remove
# salt&pepper noise. An image with `n`-level salt&papper noise is defined as: each pixel `p` has
# `n/2` probabilty to be filled by `0`(the minimal gamut value), `n/2` probabilty to be filled by
# `1`(the maximal gamut value), and `1-n` probablity to be undamaged.

img = testimage("cameraman")

n = 0.2 # noise level
noisy_img = map(img) do p
    sp = rand()
    if sp < n/2
        eltype(img)(gamutmin(eltype(img))...)
    elseif sp < n
        eltype(img)(gamutmax(eltype(img))...)
    else
        p
    end
end

denoised_img = mapwindow(median!, noisy_img, (5, 5))

mosaicview(img, noisy_img, denoised_img; nrow=1)

## save covers #src
using ImageMagick #src
mkpath("assets") #src
ImageMagick.save("assets/median.gif", cat(noisy_img, denoised_img; dims=3); fps=1) #src


# Advent of code 2021, day 1

function day01()
    part = [0, 0]
    input1 = map(s -> parse(Int, s), readlines("AoCdata/AoC_2021_day1.txt"))

    increases(arr) = count(arr[i + 1] > arr[i] for i in 1:(length(arr) - 1))
    part[1] = increases(input1) # 1477

    function slidingincreases(arr)
        return count(
            sum(arr[(i + 1):(i + 3)]) > sum(arr[i:(i + 2)]) for i in 1:(length(arr) - 3)
        )
    end
    part[2] = slidingincreases(input1) # 1523
    return part
end

part = day01()

println("Part 1: ", part[1])
println("Part 2: ", part[2])


@testset "affine.jl" begin

    map = NLS_Fit.Map_Affine(1=>2,3=>5)

    @test parameter_size(map) == 2

    θ = Float32[1, 1]
    X_hat = Int[1,2,3]
    X = eval_map(map,X_hat,θ)

    @test eltype(X) == Float32
    @test X ≈ Float32[2,3.5,5]

end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ de098278-8e74-11eb-13e3-49c1b6e06e7d
using Pkg; Pkg.activate("MLJ_env", shared=true)

# ╔═╡ 10d6d90e-8d43-11eb-216c-8f94a3946fab
begin
	push!(LOAD_PATH, "./src")
	using YaDeepLearning
end

# ╔═╡ ac463e7a-8b59-11eb-229e-db560e17c5f5
begin
	using Test
	using Random
	using Distributions
	using PlutoUI
	using Plots
	using Flux.Data.MNIST
	using Images
end

# ╔═╡ 8c80e072-8b59-11eb-3c21-a18fe43c4536
md"""
## Deep Learning (DL) (cont'ed)

ref. from book **"Data Science from Scratch"**, Chap 19

$(html"
© Pascal, April 2021
")
"""

# ╔═╡ 3972ef40-8d4f-11eb-3369-6dd94e7cfc13
Layers.Linear <: AbstractLayers.AbstractLayer, Layers.Linear <: AbstractLayers.AL

# ╔═╡ 6b3c5460-8dab-11eb-3046-331119a0dc9b
# on using: using vs import and provide extension
# cf. https://towardsdatascience.com/julias-big-problem-with-namespace-996d2e9ed71e

# ╔═╡ e7373726-8b59-11eb-2a2b-b5138e4f5268
PlutoUI.TableOfContents(indent=true, depth=4, aside=true)

# ╔═╡ f5ee64b2-8b59-11eb-2751-0778efd589cd
html"""

"""

# ╔═╡ 81290d1c-8ce2-11eb-3340-337957fd81b7
md"""
#### Dropout

Droput is a technique for regularizing neural network (to prevent overfitting). It works as follows:
  - at training time we randomly turn off a set of neurons (by replacing their output with 0) given a fixed probability. This seems to prevent overfitting by forcing the neural network to not rely  on any particular neuron.

  - at evaluation time, we do not drop the neurons, but we scale down the outputs (uniformly) by using the same fraction (used at training time).
"""

# ╔═╡ 8ff1bb20-8ce2-11eb-1de6-fd84daec8930
md"""
Let us define a helper function on initialize the binary mask we need to apply dorpout at training time.
"""

# ╔═╡ d3ee2138-8ce2-11eb-0b29-659a3be01512
function prand(dt::DataType, shape::NTuple{N, Integer};
               keep_proba::Float64=.5)::BitArray where {N}
	rand(dt, shape) .≥ (1. - keep_proba)
end

# ╔═╡ 13a0940e-8ce4-11eb-231c-f331a607203c
begin
	using .TensorDT: Tensor
	import .AbstractLayers: AbstractLayer, forward, backward, parms, ∇parms

	mutable struct Dropout <: AbstractLayer
		keep_prob::Float64
		train::Bool
		mask::BitArray
		_type::Symbol
		#
		Dropout(;keep_prob=.5, train=true) =
			new(keep_prob, train, BitArray(undef), :regularization)
	end

	function forward(self::Dropout, input::Tensor)::Tensor
		if self.train
			self.mask = prand(eltype(input), size(input);
				keep_proba=self.keep_prob)
			input .* self.mask
		else
			input * self.keep_prob # (1. - )
		end
	end

	function backward(self::Dropout, ∇p::Tensor)::Tensor
		!self.train &&
			throws(RuntimeError("Do NOT all backward when NOT in training mode"))
		∇p .* self.mask
	end

	parms(_self::Dropout) = []
	∇parms(_self::Dropout) = []
end

# ╔═╡ d3d074e4-8ce2-11eb-0307-d932886e4b10
md"""
Then let us define our dropout  layer.
"""

# ╔═╡ 80ddf0f4-8ce2-11eb-3046-331119a0dc9b
html"""

"""

# ╔═╡ 80a71408-8ce2-11eb-3978-75a4a2df9116
md"""
#### MNIST
"""

# ╔═╡ d1ed1e36-8ce5-11eb-34ee-4db60fb0db8a
const NROWS, NCOLS, NLABELS = 28, 28, 10

# ╔═╡ d2b198b6-8ce2-11eb-170e-0f17904c9f2c
begin
  ## Load labels and images from Flux.Data.MNIST
  train_labels, train_imgs = MNIST.labels(:train), MNIST.images(:train);
  test_labels, test_imgs = MNIST.labels(:test), MNIST.images(:test)
  size(train_labels), size(train_imgs), size(test_labels), size(test_imgs)
end

# ╔═╡ d37f37b8-8ce8-11eb-2c00-3f98ca407f41
size(train_imgs[1]), size(test_imgs[1])

# ╔═╡ 5bf7729e-8dfd-11eb-070f-9b7ec0746bd7
typeof(train_imgs[1])

# ╔═╡ 394093fa-8cea-11eb-0071-eff3045a012b
plot(train_imgs[1]), train_labels[1]

# ╔═╡ 435c4444-8dfd-11eb-24e8-5543f905f199
# sum((train_imgs[10])),
minimum(Float32.(train_imgs[1]))

# ╔═╡ e95983ba-8ceb-11eb-38fd-ed92cdcf754c
mosaicview(train_imgs[1:100]; nrow=5, ncol=20, npad=5)

# ╔═╡ d3a749a2-8cec-11eb-1f06-b568f244b576
md"""
Each image is  28x28 pixels, but our linear layer can only deal with one-dimensional input, therefore we need to flatten them.

We will also recenter both train and test sets by using the mean calculated over the training set.
"""

# ╔═╡ 31d3226a-8cf4-11eb-0897-39989ba76b58
begin
	## reshape (flatten) and cast from Gray{Normed{UInt8,8}} to Float32
	_train_imgs = [Float32.(img) for img ∈ (reshape.(train_imgs, NROWS * NCOLS))]
	_test_imgs = [Float32.(img) for img ∈ (reshape.(test_imgs, NROWS * NCOLS))]

	## calc. avg
	img_avg = Float32.(sum(_train_imgs) / length(_train_imgs) / (NROWS * NCOLS))
	@assert size(_train_imgs[1]) == size(img_avg) == (NROWS * NCOLS,)

	## Re-center, re-scale
	λ = img -> Float32.(img .- img_avg)
	ntrain_imgs = λ.(_train_imgs); _train_imgs = nothing
	ntest_imgs = λ.(_test_imgs); _test_imgs = nothing

  	## ntrain_imgs ≡ Array{Array{Float32,1},1}
	@assert size(ntrain_imgs) == size(train_imgs)
	@assert size(ntest_imgs) == size(test_imgs)
end

# ╔═╡ 6e7e7896-8cf8-11eb-062e-99492ec8cff8
md"""
Let us write a one hot encoder for our 10 outputs (digits from 0 to 9)
"""

# ╔═╡ 70f707c0-8cf6-11eb-11c2-73d7b28f7a0c
function one_hot_encoder(label::Integer; nlabels=NLABELS, DT=Float64)::Tensor
	# @assert label < nlabels
	a = zeros(DT, nlabels)
	a[label + 1] = one(DT)
	a
end

# ╔═╡ 70d84470-8cf6-11eb-23b6-c7ba506d8552
begin
  @test one_hot_encoder(3) == [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.]
  @test one_hot_encoder(2; nlabels=5) == [0., 0., 1., 0., 0.]
end

# ╔═╡ c451451a-8cf7-11eb-183b-e358c9d618e0
begin
	ntrain_labels = one_hot_encoder.(train_labels)
	ntest_labels = one_hot_encoder.(test_labels);
	ntrain_labels[1:3]
end

# ╔═╡ ebd531e6-8cf7-11eb-228f-d530e7a48151
md"""
Let adjust our (previoulsy defined) train_loop function.
"""

# ╔═╡ 30bdc788-8d01-11eb-2f8f-9fd79593ddb8
begin

const AL, ALoss, NAOpt = AbstractLayers.AL, Loss.ALoss, AbstractOptimizers.NAOpt

function train_loop(model::AL, imgs::AA, labels::AA, loss_fn::ALoss;
    epochs=1, opt_fn::NAOpt=nothing, verbose=false)
	#
	rec_avg_loss, rec_acc = [], []
	ncorr, totloss = 0, 0.0
	#
	for e ∈ 1:epochs, t in 1:length(imgs)
		ŷ = Layers.forward(model, imgs[t])
		y = labels[t]
		## NOTE argmax => CartesianIndex
		argmax(ŷ)[1] == argmax(y) && (ncorr += 1)
		totloss += Loss.loss(loss_fn, ŷ, y)
		#
		if !isnothing(opt_fn)
			∇p = Loss.∇loss(loss_fn, ŷ, y)
			Layers.backward(model, ∇p)
			Optimizers.astep(opt_fn, model)
		end
		avg_loss = totloss / (e * t)
		acc = ncorr / (e * t)
		verbose && t % 5000 == 0 && (@show "$(e)/$(t)", avg_loss, acc)
		push!(rec_avg_loss, avg_loss)
		push!(rec_acc, acc)
	end
	(rec_avg_loss, rec_acc)
end

end

# ╔═╡ 84946000-8d02-11eb-3160-571efee8fb0b
md"""
Let us create a baseline to train a mluticlass logistic regression (which is a single linear layer followed by a softmax). One pass through the 60000 training examples should be enough to learn the model.
"""

# ╔═╡ d05f1c40-8d17-11eb-2724-4d989c4c6b92
begin
	lr_model = Layers.Linear(NROWS * NCOLS, NLABELS)
	lr_loss = Loss.SoftmaxXEntropy()
	lr_opt = Optimizers.MomentumGD(;η=.01, α=.99)
end

# ╔═╡ 848fd812-8d02-11eb-01d6-75965b08bcc5
with_terminal() do
	# Train
	global (tr1_avg_loss, tr1_acc) = train_loop(lr_model, ntrain_imgs, ntrain_labels,
		lr_loss; opt_fn=lr_opt, verbose=true)

	# Test model
	global (te1_avg_loss, te1_acc) = train_loop(lr_model, ntrain_imgs, ntrain_labels,		lr_loss)
	# rouglhy 90% of accuracy
end

# ╔═╡ 7d35f85e-8e07-11eb-228c-81d28a444b52
begin
	ntr, nte = 60_000, 10_000
	lg1 = grid(2, 1, heights=[0.5, 0.5])
	a = plot(tr1_avg_loss[1:ntr], label="train avg loss", leg=:bottomright)
	b = plot(tr1_acc[1:ntr], label="train acc.")
	plot(a, b, layout=lg1)
end

# ╔═╡ 6e7c18b4-8e09-11eb-0238-2515d39a89cd
begin
	lg2 = grid(2, 1, heights=[0.5, 0.5])
	a₂ = plot(te1_avg_loss[1:nte], label="test avg loss", leg=:bottomright)
	b₂ = plot(te1_acc[1:nte], label="test acc.")
	plot(a₂, b₂, layout=lg2)
end

# ╔═╡ 40d46802-8e0f-11eb-2e47-53096e24dbd8
md"""
And now let us define and train our neural network
"""

# ╔═╡ 80dff352-8d02-11eb-06b1-5f5e325046f5
begin
	dropout₁ = Dropout(;keep_prob=0.9)
	dropout₂ = Dropout(;keep_prob=0.8)

	m_model = Layers.Sequential([
		Layers.Linear(NROWS * NCOLS, 3 * NLABELS),
		dropout₁,
		Activations.Tanh(),
		Layers.Linear(3 * NLABELS, NLABELS),
		dropout₂,
		Activations.Tanh(),
		Layers.Linear(NLABELS, NLABELS)
	])
	m_loss = Loss.SoftmaxXEntropy()
	m_opt = Optimizers.MomentumGD(;η=.01, α=.99)
	with_terminal() do
		## Train
		global (tr_avg_loss, tr_acc) = train_loop(m_model, ntrain_imgs,
			ntrain_labels, m_loss; epochs=5, opt_fn=m_opt)

		## Test
		dropout₁.train = dropout₂.train = false
		global (te_avg_loss, te_acc) = train_loop(m_model, ntrain_imgs,
			ntrain_labels, m_loss, verbose=true)
	end
end

# ╔═╡ f1af4256-8e02-11eb-0238-2515d39a89cd
begin
	## 2 epochs for training...
	x₁ = plot(tr_avg_loss[1:ntr], label="train avg loss", leg=:bottomright)
	x₂ = plot(tr_avg_loss[ntr+1:end], label="train avg loss", leg=:bottomright)

	y₁ = plot(tr_acc[1:ntr], label="train acc.", leg=:bottomright)
	y₂ = plot(tr_acc[ntr+1:end], label="train acc.")
	plot(x₁, x₂, y₁, y₂, layout=(2, 2))
end

# ╔═╡ 43a52692-8e10-11eb-2043-8f0be195f58a
begin
	x₃ = plot(te_avg_loss[1:ntr], label="test1 avg loss", leg=:bottomright)
	y₃ = plot(te_acc[1:ntr], label="test2 acc.", leg=:bottomright)
	plot(x₃, y₃, layout=(2, 1))
end

# ╔═╡ Cell order:
# ╠═8c80e072-8b59-11eb-3c21-a18fe43c4536
# ╠═de098278-8e74-11eb-13e3-49c1b6e06e7d
# ╠═ac463e7a-8b59-11eb-229e-db560e17c5f5
# ╠═10d6d90e-8d43-11eb-216c-8f94a3946fab
# ╠═3972ef40-8d4f-11eb-3369-6dd94e7cfc13
# ╠═6b3c5460-8dab-11eb-3046-331119a0dc9b
# ╟─e7373726-8b59-11eb-2a2b-b5138e4f5268
# ╟─f5ee64b2-8b59-11eb-2751-0778efd589cd
# ╟─81290d1c-8ce2-11eb-3340-337957fd81b7
# ╟─8ff1bb20-8ce2-11eb-1de6-fd84daec8930
# ╠═d3ee2138-8ce2-11eb-0b29-659a3be01512
# ╟─d3d074e4-8ce2-11eb-0307-d932886e4b10
# ╠═13a0940e-8ce4-11eb-231c-f331a607203c
# ╟─80ddf0f4-8ce2-11eb-3046-331119a0dc9b
# ╟─80a71408-8ce2-11eb-3978-75a4a2df9116
# ╠═d1ed1e36-8ce5-11eb-34ee-4db60fb0db8a
# ╠═d2b198b6-8ce2-11eb-170e-0f17904c9f2c
# ╠═d37f37b8-8ce8-11eb-2c00-3f98ca407f41
# ╠═5bf7729e-8dfd-11eb-070f-9b7ec0746bd7
# ╠═394093fa-8cea-11eb-0071-eff3045a012b
# ╠═435c4444-8dfd-11eb-24e8-5543f905f199
# ╠═e95983ba-8ceb-11eb-38fd-ed92cdcf754c
# ╟─d3a749a2-8cec-11eb-1f06-b568f244b576
# ╠═31d3226a-8cf4-11eb-0897-39989ba76b58
# ╟─6e7e7896-8cf8-11eb-062e-99492ec8cff8
# ╠═70f707c0-8cf6-11eb-11c2-73d7b28f7a0c
# ╠═70d84470-8cf6-11eb-23b6-c7ba506d8552
# ╠═c451451a-8cf7-11eb-183b-e358c9d618e0
# ╟─ebd531e6-8cf7-11eb-228f-d530e7a48151
# ╠═30bdc788-8d01-11eb-2f8f-9fd79593ddb8
# ╟─84946000-8d02-11eb-3160-571efee8fb0b
# ╠═d05f1c40-8d17-11eb-2724-4d989c4c6b92
# ╠═848fd812-8d02-11eb-01d6-75965b08bcc5
# ╠═7d35f85e-8e07-11eb-228c-81d28a444b52
# ╠═6e7c18b4-8e09-11eb-0238-2515d39a89cd
# ╟─40d46802-8e0f-11eb-2e47-53096e24dbd8
# ╠═80dff352-8d02-11eb-06b1-5f5e325046f5
# ╠═f1af4256-8e02-11eb-0238-2515d39a89cd
# ╠═43a52692-8e10-11eb-2043-8f0be195f58a


# Code to test basic functionality and plotting
using RobustControl, PyPlot

A = [0 1;-1 -1];
B = reshape([0;1],2,1);
C = RobustControl.eye(2);
D = reshape([0;0],2,1);;
s = RobustControl.StateSpace(A,B,C);
rr,ii,om,d = RobustControl.rifd(s);
z = RobustControl.TransmissionZeros(s,2,1);
p = RobustControl.Poles(s);
d = RobustControl.Damping(s);
om = RobustControl.NaturalFrequencies(s);

# Frequency response with auto frequency grid.
G,om = RobustControl.FrequencyResponse(s,1,1);

# Frequency response with specified frequency grid.
om1 = RobustControl.logspace(1E-2,1E2,200);
G1 = RobustControl.FrequencyResponse(s,2,1,om1);

# Test plots
figure(1);clf(); pygui(true);
subplot(2,1,1);
RobustControl.BodeMagPlot(G,om); title("Bode Mag Plot");
RobustControl.BodeMagPlot(G1,om1);
grid(which="both",color=[1,1,1]*0.9,linewidth=0.35)
legend(["G1","G2"]);
tight_layout();

# Another way to get Bode plots -- just magnitude.
subplot(2,1,2);
RobustControl.BodeMagPlot(s,1,1);
RobustControl.BodeMagPlot(s,2,1);
grid(which="both",color=[1,1,1]*0.9,linewidth=0.35)
legend(["G1","G2"]);
tight_layout();

# 3rd way -- get both mag and phase in one plot
figure(2); clf();
RobustControl.BodePlot(s,1,1);
RobustControl.BodePlot(s,2,1);
subplot(211);grid(which="both",color=[1,1,1]*0.9,linewidth=0.35)
subplot(212);grid(which="both",color=[1,1,1]*0.9,linewidth=0.35)
subplot(211);legend(["G1","G2"]);
title("Bode Plot");
tight_layout();

# Nyquist Plot
figure(3);clf();
subplot(2,1,1);RobustControl.NyquistPlot(s,1,1); title("Nyquist Plot"); # With auto frequency grid
subplot(2,1,2);RobustControl.NyquistPlot(s,1,1,om1); # With specified frequency grid

# System Responses
figure(4); clf();
subplot(2,2,1);
t,y = RobustControl.ImpulseResponse(s,1,1);
plot(t,y); grid();
xlabel("Time (sec)");
#ylabel("Impulse Response");
title("Impulse Response");tight_layout();

subplot(2,2,2);
t,y = RobustControl.StepResponse(s,1,1);
plot(t,y); grid();
xlabel("Time (sec)");
#ylabel("Step Response");
title("Step Response");tight_layout();

subplot(2,2,3);
t = Vector(0:0.01:50);
u = ones(length(t)) + 0.1*randn(Float64,length(t)); # Step + Disturbance
x0 = [0.0;0.0];
y = RobustControl.ForcedResponse(s,1,1,t,u,x0);
plot(t,y); grid();
xlabel("Time (sec)");
#ylabel("Time Response to Step+Disturbance");
title("Forced Response");tight_layout();

subplot(2,2,4);
t = Vector(0:0.01:10);
x0 = [1.0;0.0];
t1,y1 = RobustControl.InitialResponse(s,x0); # Auto time grid
y2 = RobustControl.InitialResponse(s,x0,t); # Specified time grid
l1=plot(t1,y1,"b",linewidth=1);
l2=plot(t,y2,"r",linestyle=":",linewidth=2);
xlabel("Time (sec)");
#ylabel("Initial Response");
legend([l1[1],l2[1]],["(auto time grid)","(specified time grid)"]);
title("Initial Response");
grid(); tight_layout();


using DataStructures
function SimplePeel(A)
    """
    GreedyPeel with parameter p
    """
    n = size(A,1)
    d = vec(round.(Int64,sum(A,dims = 1)))
    numer = sum(d)

    zz = length(findall(x->x==0,d))
    # save graph as adjacency list
    cp = A.colptr
    row = A.rowval
    AdjList = Vector{Vector{Int64}}()
    for i = 1:n
        push!(AdjList,row[cp[i]:cp[i+1]-1])
    end

    D = MutableBinaryMinHeap(d)

    # greedily peel
    ranking = zeros(Int64,n)
    Objs = zeros(n+1)
    Objs[1] = (numer/n)
    next = 2
    denom = n-1
    degrees = zeros(Int64,n)

    # if there are zero values by chance, get rid of them
    for t = 1:zz
        deg, i = top_with_handle(D)
        pop!(D)
        Objs[next] = 0
        denom -= 1
        @assert(deg == 0)
        degrees[next] = 0
        ranking[next-1] = i
        next += 1
    end

    while ~isempty(D)
        deg, i = top_with_handle(D)
        ranking[next-1] = i
        degrees[next-1] = deg
        pop!(D)
        Ni = AdjList[i]
        di = d[i]
        if di > 0
            # Decrease the degree of all neighbors
            d[i] = 0
            for nb = Ni
                removej!(AdjList[nb],i)             # This is no longer a neighboring node, we're removed it
                update!(D,nb,D[nb]-1)  # subtract degree, update the heap
                d[nb] -= 1
            end
        end
        numer -= 2*di
        Objs[next] = numer/denom
        denom -= 1
        next +=1
    end
    Objs = Objs[1:end-1]
    density, w = findmax(Objs)
    S = ranking[w:end]
    degeneracy, w2 = findmax(degrees)
    maxcore = ranking[w2:end]

    return S, density, degeneracy, maxcore, Objs, ranking, degrees
end

function GenPeel(A,p)
    """
    GreedyPeel with parameter p
    """
    if p == 1
        S, density, degeneracy, maxcore, Objs, ranking, degrees = SimplePeel(A)
        return S, density, Objs, ranking
    end

    n = size(A,1)
    d = sum(A,dims = 1)
    numer = sum(d.^p)

    # converting to dictionary helps later
    d2 = Dict{Int64,Int64}()
    for i = 1:n
        d2[i] = d[i]
    end
    d = d2

    # save graph as adjacency list
    cp = A.colptr
    row = A.rowval
    tic = time()
    AdjList = Vector{Vector{Int64}}()
    for i = 1:n
        push!(AdjList,row[cp[i]:cp[i+1]-1])
    end

    # Create initial minheap
    D = zeros(n)
    for i = 1:n
        D[i] = d[i]^p
        for k = AdjList[i]
            D[i] += d[k]^p - (d[k]-1)^p
        end
    end
    H = MutableBinaryMinHeap(D)

    # greedily peel
    ranking = zeros(Int64,n)
    Objs = zeros(n+1)
    Objs[1] = (numer/n)^(1/p)
    next = 2
    denom = n-1
    while ~isempty(H)
        # global numer, denom,next
        v, i = top_with_handle(H)
        # push!(order,i)
        ranking[next-1] = i
        pop!(H)
        Ni = AdjList[i]
        AdjList[i] = Vector{Int64}()
        di = d[i]
        if di > 0
            dichange = di^p - (di-1)^p
            d[i] = 0
            numer -= di^p
            # Remove this from its neighbors list
            for nb = Ni
                removej!(AdjList[nb],i)             # This is no longer a neighboring node, we're removed it
                db = d[nb]
                dbchange = -db^p + (db-1)^p
                Old = H[nb]                         # previous value of D[nb]
                New = Old + dbchange - dichange     # new value now that node i is gone
                update!(H,nb,New)                   # update the heap
                # D[nb] = New
                numer += dbchange
                d[nb] -= 1                          # decrease degree
            end

            # At this point, we have made all changes due to node i disappearing.
            # Now we need to account for changes in D/H
            # that are due to changes in the degree for i's neighbors
            for b = Ni
                db = d[b]
                if db > 0
                    dbchange = 2*db^p - (db-1)^p -(db+1)^p
                    Nb = AdjList[b]
                    for k = Nb
                        update!(H,k,H[k]+dbchange)
                    end
                end
            end

        end

        # if mod(next,100) == 0
        #     println("$next $numer")
        # end
        # if numer < 0
        #     @show numer
        #     break
        # end
        numer = max(numer,0)                # there can be rounding errors at the very end
        Objs[next] = (numer/denom)^(1/p)
        denom -= 1
        next +=1
    end
    Objs = Objs[1:end-1]
    density, w = findmax(Objs)
    S = ranking[w:end]

    return S, density, Objs, ranking

end

# remove element j from vector v, assuming that j appears exactly once
function removej!(v::Vector{Int64},j)
    for i = 1:length(v)
        if v[i] == j
            # v = v[1:i-1,i+1:end]
            deleteat!(v,i)
            return
        end
    end
end

function pdensityobj(A,S,p)

    if length(S) == 0
        return 0
    else
        As = A[S,S]
        d = sum(As,dims = 1)

        return (sum(d.^p)/length(S))^(1/p)
    end

end


# Function for running Riemannian manifold Hamiltonian Monte Carlo (RMHMC)
function rmhmc(model::Model, opts::RmhmcOpts)
  mcmc = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  z = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  proposedGradLogPosterior = Array(Float64, model.nPars)
  invGDerivG = Array(Float64, model.nPars, model.nPars, model.nPars)
  traceInvGDerivG = Array(Float64, model.nPars)
  momentumTerm = Array(Float64, model.nPars)
  proposed, accepted = 0., 0.

  println("Running burn-in iterations...")

  currentPars = model.randPrior()
  currentLogPosterior = model.logPosterior(currentPars)
  currentGradLogPosterior = model.gradLogPosterior(currentPars)
  leapStep = opts.setLeapStep(1, 0., opts.mcmc.n, opts.mcmc.nBurnin, 0.)

  for i = 1:opts.mcmc.n
    proposed += 1

    proposedPars = copy(currentPars)
    G = model.tensor(proposedPars)
    invG = inv(G)
    cholG = chol(G)
    momentum = cholG'*randn(model.nPars)
    currentLogDetG = 0.5*(log(2)+model.nPars*log(pi)+2*sum(log(diag(cholG))))
    currentHamiltonian  =
      -currentLogPosterior+currentLogDetG+(momentum'*(invG*momentum))/2
    derivG = model.derivTensor(proposedPars)
    for j = 1:model.nPars
      invGDerivG[:, :, j] = invG*derivG[:, :, j]
      traceInvGDerivG[j] = trace(invGDerivG[:, :, j])
    end
    timeStep = (randn() > 0.5 ? 1. : -1.)
    randomSteps = ceil(rand()*opts.nLeaps)
    for j = 1:randomSteps
      leapGradLogPosterior  = model.gradLogPosterior(proposedPars)
      leapMomentum = copy(momentum)
      for k = 1:opts.nNewton
        invGMomentum = invG*leapMomentum
        for r = 1:model.nPars
          momentumTerm[r] =
            (0.5*(leapMomentum'*invGDerivG[:, :, r]*invGMomentum))[1]
        end
        leapMomentum = (momentum
          +timeStep*(leapStep/2)*(leapGradLogPosterior-0.5*traceInvGDerivG
          +momentumTerm))
      end
      momentum = copy(leapMomentum)
      leapInvGMomentum = invG*momentum
      leapParameters = copy(proposedPars)
      for k = 1:opts.nNewton
        G = model.tensor(leapParameters)
        invGMomentum = G\momentum
        leapParameters =
          proposedPars+(timeStep*(leapStep/2))*(leapInvGMomentum+invGMomentum)
      end
      proposedPars = copy(leapParameters)
      G = model.tensor(proposedPars)
      invG = inv(G)
      derivG = model.derivTensor(proposedPars)
      for k = 1:model.nPars
        invGDerivG[:, :, k] = invG*derivG[:, :, k]
        traceInvGDerivG[k] = trace(invGDerivG[:, :, k])
      end
      invGMomentum = invG*momentum
      for k = 1:model.nPars
        momentumTerm[k] = (0.5*(momentum'*invGDerivG[:, :, k]*invGMomentum))[1]
      end
      proposedGradLogPosterior = model.gradLogPosterior(proposedPars)
      momentum = (momentum
        +timeStep*(leapStep/2)*(proposedGradLogPosterior-0.5*traceInvGDerivG
        +momentumTerm))
    end
    proposedLogPosterior = model.logPosterior(proposedPars)
    proposedLogDet = 0.5*(log(2)+model.nPars*log(pi)+2*sum(log(diag(chol(G)))))
    proposedHamiltonian =
      -proposedLogPosterior+proposedLogDet+(momentum'*(invG*momentum))/2

    ratio = (currentHamiltonian-proposedHamiltonian)[1]

    if ratio > 0 || (ratio > log(rand()))
      accepted += 1

      currentPars = copy(proposedPars)
      currentLogPosterior = copy(proposedLogPosterior)
      currentGradLogPosterior = copy(proposedGradLogPosterior)
    end

    if i > opts.mcmc.nBurnin
      mcmc[i-opts.mcmc.nBurnin, :] = currentPars
      z[i-opts.mcmc.nBurnin, :] = -currentGradLogPosterior/2
    end

    if mod(i, opts.mcmc.monitorRate) == 0
      acceptanceRatio = accepted/proposed

    leapStep = opts.setLeapStep(i, acceptanceRatio, opts.mcmc.n,
      opts.mcmc.nBurnin, leapStep)

      println("Iteration $i of $(opts.mcmc.n): ", round(100*acceptanceRatio, 2),
        " % acceptance ratio")

      proposed, accepted = 0., 0.
    end

    if i == opts.mcmc.nBurnin
      println("Burn-in completed...\n\nRunning post burn-in MCMC...");
    end
  end

  return mcmc, z
end


using Revise

using DataFrames, Statistics, Distributions, StatsBase,
  LinearAlgebra, BenchmarkTools, Random, Optim, Zygote, UnPack

MFloat64 = Union{Missing, Float64}

Random.seed!(11)

# a simple AR(1) model
struct AR1{T<:Real, Tσ²<:Union{Nothing,Real}}
  ϕ::T
  b::T
  μ::T
  σ²::Tσ² #note- this is σ²e, not the overall σ² of the series
end

AR1(;ϕ, μ=missing, b=missing, σ²=nothing) = AR1(
  ϕ, coalesce(b, μ*(1-ϕ)), coalesce(μ, b/(1-ϕ)),σ²)

centerar1(ar1::AR1) = AR1(ar1.ϕ, 0.0, 0.0, ar1.σ²)

function estimatear1ols(x::Vector{T}) where T
  x̃ = @view x[2:end]
  Lx̃ = @view x[1:(end-1)]

  ϕ = cov(x̃, Lx̃)/(var(Lx̃))
  b = mean(x̃) - ϕ * mean(Lx̃)
  μ = b/(1-ϕ)

  ε = similar(x)
  ε[1] = x[1] - μ
  ε[2:end] = x̃ .- ϕ .* Lx̃
  σ² = var(ε)

  @assert b ≈ mean(x̃) - ϕ * mean(Lx̃)
  @assert ϕ ≈ cov(x̃ .- mean(x̃), Lx̃ .- mean(Lx̃))/var(
    Lx̃ .- mean(Lx̃))

  return AR1(;ϕ,b,σ²)
end

function estimatear1mle(x; tol=10^-10)

  #start by centering
  N=length(x)
  μ = mean(x)
  x̃ = @view(x[2:end]) .- μ
  Lx̃ = @view(x[1:(end-1)]) .- μ

  x₁ = x[1]
  #computes the next iteration of ϕ given a previous ϕ
  cσ²(ϕ) = (x₁^2*(1-ϕ^2) + sum((x̃ .- ϕ*Lx̃).^2))/N
  function objϕ(ϕ)
    σ² = cσ²(ϕ)
    return (sum((x̃ .- ϕ * Lx̃) .* Lx̃) - (σ²/(1-ϕ^2) - x₁^2)*ϕ)^2
  end

  ϕ = optimize(objϕ, -1.01, 1.01, Brent()) |> Optim.minimizer

  return AR1(;ϕ, μ, σ² = cσ²(ϕ))
end

function estimatear1mlecheck(x; tol=10^-10)

  #start by centering
  N=length(x)
  μ = mean(x)
  x̃ = @view(x[2:end]) .- μ
  Lx̃ = @view(x[1:(end-1)]) .- μ

  x₁ = x[1]

  l(ϕ, σ²) = (-N/2*log(2*π*σ²)
    + 1/2*log(1-ϕ^2)
    - x₁^2*(1-ϕ^2)/(2*σ²)
    - sum((x̃ .- ϕ .* Lx̃).^2)/(2*σ²)) * -1

  l(Φ) = l(Φ[1], Φ[2])

  ϕ₀ = cov(x̃, Lx̃)/var(Lx̃) #start by guessing the least squares solution
  ε₀= x̃ - ϕ₀*Lx̃
  σ²₀ = var(ε₀)

  ϕ, σ² = optimize(l,  [-1.0,0.0], [1.0,Inf], [ϕ₀, σ²₀], Fminbox(LBFGS()),
    autodiff = :forward) |> Optim.minimizer

  return AR1(;ϕ, μ, σ²)
end

AR1(x::AbstractVector; estimationmethod=estimatear1mle) = estimationmethod(x)

function checkmath2(x)
  if length(x) > 250
    x = x[1:250] |> deepcopy
    @warn "length(x) > 250 so truncating length to 250. This is only for the purposes
    of computing the full (not log) likelihood"
  end

  N=length(x)
  μ = mean(x)
  x̃ = @view(x[2:end]) .- μ
  Lx̃ = @view(x[1:(end-1)]) .- μ

  ϕ₀ = 0.75
  σ²₀ = 2.2
  x₁ = x̃[1]

  L(ϕ, σ²) =  (1-ϕ^2)^(1/2)/σ²^(1/2)*pdf(Normal(), x₁/(σ²/(1-ϕ^2))^(1/2))*prod(
    ((x̃ᵢ,Lx̃ᵢ)->1/σ²^(1/2)*pdf(Normal(), (x̃ᵢ-ϕ*Lx̃ᵢ)/σ²^(1/2))).(x̃,Lx̃))

  l(ϕ, σ²) = (-N/2*log(2*π*σ²)
    + 1/2*log(1-ϕ^2)
    - x₁^2*(1-ϕ^2)/(2*σ²)
    - sum((x̃ .- ϕ .* Lx̃).^2)/(2*σ²))

  @assert log(L(ϕ₀, σ²₀)) ≈ l(ϕ₀, σ²₀,) "log(L(ϕ, σ²)): $(log(L(ϕ₀, σ²₀)))
    l(ϕ, σ²): $(l(ϕ₀, σ²₀))"

  #now verify the derivitive and conditional wrt sigma^2
  lσ²(ϕ=ϕ₀, σ²=σ²₀,) = (-N/(2σ²)
    + x₁^2*(1-ϕ^2)/(2σ²^2)
    + sum((x̃ .- ϕ .* Lx̃).^2)/(2σ²^2))
  lσ²z(ϕ=ϕ₀, σ²=σ²₀,) = gradient(σ²->l(ϕ, σ²,), σ²)[1]
  lσ²n() =  (l(ϕ₀, σ²₀+1e-6,)-l(ϕ₀, σ²₀,))/1e-6
  @assert lσ²() ≈ lσ²z() "lσ²(...): $(lσ²())
    lσ²z(...): $(lσ²z())"

  cσ²(ϕ) = (x₁^2*(1-ϕ^2) + sum((x̃ .- ϕ*Lx̃).^2))/N
  @assert 1 + lσ²z(ϕ₀, cσ²(ϕ₀)) ≈ 1.0 "cσ²(ϕ₀): $(cσ²(ϕ₀))
    lσ²z(ϕ₀, cσ²(ϕ₀)): $(lσ²z(ϕ₀, cσ²(ϕ₀)))"

  #finally, verify the derivitive wrt ϕ
  lϕ(ϕ=ϕ₀, σ²=σ²₀) = (-ϕ/(1-ϕ^2) +
    x₁^2*ϕ/σ² +
    sum((x̃ .- ϕ * Lx̃) .* Lx̃)/σ²)
  lϕz(ϕ=ϕ₀, σ²=σ²₀,) = gradient(ϕ->l(ϕ, σ²), ϕ)[1]
  @assert lϕ() ≈ lϕz() "lϕ(...): $(lϕ())
    lϕz(...): $(lϕz())"

  ϕcerr(ϕ, σ²=σ²₀)= sum((x̃ .- ϕ * Lx̃) .* Lx̃) - (σ²/(1-ϕ^2) - x₁^2)*ϕ
  cϕ(σ²=σ²₀,) = optimize(ϕ->ϕcerr(ϕ)^2, -1.01, 1.01, Brent(), abs_tol=10^-14) |> Optim.minimizer
  @assert isapprox(lϕz(cϕ()), 0.0,atol=1e-5) "cϕ(ϕ₀): 2$(cϕ())
    lϕz(ϕcerr()): $(lϕz(ϕcerr()))"
  #lμ()
end

#simulates an AR1 series
function simulatear1!(ar1::AR1, x::AbstractVector)
  @unpack ϕ, b, σ², μ = ar1

  σ² === nothing && throw(" σ² must not be nothing")

  #first simulate the centered distribution
  #we will do this all in place
  σ=σ²^0.5
  rand!(Normal(0,σ), x)
  x[1] *= 1/(1-ϕ^2)^0.5 # the first entry is simulated from the unconditional distribution
  for t ∈ 2:length(x)
    x[t] += ϕ * x[t-1]
  end
  x .+= μ

  return x
end

#which approach is faster?
function simulatear1(ar1::AR1, T::Int)
  @unpack ϕ, σ², μ = ar1

  σ² === nothing && throw(" σ² must not be nothing")
  #T::Int = length(x)

  #first simulate the centered distribution
  #we will do this all in place
  σ=σ²^0.5
  x = rand(Normal(), T) .* σ
  x[1] *= 1/(1-ϕ^2)^0.5
  for t ∈ 2:T
    x[t] += ϕ * x[t-1]
  end
  x .+= μ

  return x
end

simulatear1!(ar1::AR1{T,Tσ²}, N::Int) where {T,Tσ²} = simulatear1!(ar1, Vector{T}(undef, N))

compareoncear1xy!(xar1::AR1, x::AbstractVector, yar1::AR1,  y::AbstractVector,
  comparefunc::Tcomparefunc
  ) where Tcomparefunc = comparefunc(simulatear1!(xar1,x), simulatear1!(yar1,y))

#this runs the simulations- basically a parametric bootstrap under the null
#statfunc is computed for each pass, while aggfunc is computed on the statistics
#at the end of the pass
function comparear1undernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4
    ) where {Tx, Tstatfunc, Taggfunc}
  xar1 = AR1(x) |> centerar1
  yar1 = AR1(y) |> centerar1

  results = Dict(s => Vector{MFloat64}(undef, Nsimulations) for s ∈ keys(statfunc(x,y)))

  #this assumes stationary results are meaningless
  if checkstationarity && max(abs(xar1.ϕ), abs(yar1.ϕ)) ≥ 1.0-stationaritytol
    xar1 = AR1(NaN, NaN, xar1.μ, xar1.σ²)
    yar1 = AR1(NaN, NaN, xar1.μ, xar1.σ²)
  end

  xs = [similar(x) for t ∈ 1:Threads.nthreads()]
  ys = [similar(y) for t ∈ 1:Threads.nthreads()]

  #draw a sample and compute the aggregate function, sotring the reuslts in an array
  Threads.@threads for n ∈ 1:Nsimulations
    t::Int = Threads.threadid()
    res = compareoncear1xy!(xar1, xs[t], yar1, ys[t], statfunc)

    #this structure allows us to compute multiple statistics
    for s ∈ keys(res)
      @assert results[s][n] ≡ missing
      results[s][n] = res[s]
    end
  end

  return (; Dict(s=>aggfunc(skipmissing(results[s]), s) for s ∈ keys(results))...)
end

function comparearima11undernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4) where {Tx, Tstatfunc, Taggfunc}

    #difference everything
    dx = x[2:end] .- x[1:(end-1)]
    dy = x[2:end] .- y[1:(end-1)]

    dxar1 = AR1(dx) |> centerar1
    dyar1 = AR1(dy) |> centerar1

    #make integration funcitons
    #ix(δx) = cumsum([x[1]; δx])
    #iy(δy) = cumsum([y[1]; δy])

    #dstatfunc(δx, δy) = statfunc(δx |> ix, δy |> iy)

    results = Dict(s => Vector{MFloat64}(undef, Nsimulations) for s ∈ keys(statfunc(dx,dy)))

    #this assumes stationary results are meaningless
    if checkstationarity && max(abs(dxar1.ϕ), abs(dyar1.ϕ)) ≥ 1.0-stationaritytol
      dxar1 = AR1(NaN, NaN, dxar1.μ, dxar1.σ²)
      dyar1 = AR1(NaN, NaN, dxar1.μ, dxar1.σ²)
    end

    dxs = [similar(dx) for t ∈ 1:Threads.nthreads()]
    dys = [similar(dy) for t ∈ 1:Threads.nthreads()]

    #draw a sample and compute the aggregate function, sotring the reuslts in an array
    Threads.@threads for n ∈ 1:Nsimulations
      t::Int = Threads.threadid()
      res = compareoncear1xy!(dxar1, dxs[t], dyar1, dys[t], statfunc)

      #this structure allows us to compute multiple statistics
      for s ∈ keys(res)
        @assert results[s][n] ≡ missing
        results[s][n] = res[s]
      end
    end

    return (; Dict(s=>aggfunc(skipmissing(results[s]), s) for s ∈ keys(results))...)
end

nonparametricnullp(v::AbstractVector{T}, x) where T<:Real = (
  x ≥ 0 ? sum(v .> x)/length(v) : sum(v .< x)/length(v))
nonparametricnullp(v, x) = nonparametricnullp(v |> skipmissing |> collect, x)

function testar1(;N)
  ϕ = 0.8
  ε = rand(Normal(), N) + rand(Normal(1,1), N)
  x = missings(Float64, N)
  x[1] = ε[1]

  global x = x
  for t ∈ 2:N
    x[t] = ϕ*x[t-1] + ε[t]
  end
  @info "ϕ=$ϕ, b=$(mean(ε)), μ=$(mean(x))"

  checkmath2(x)

  #=estolsar1 = @btime AR1($x, estimationmethod=estimatear1ols)
  @info "ols conditional: $estolsar1"

  estmlear1 = @btime AR1($x, estimationmethod=estimatear1mle)
  @info "mle: $estmlear1"

  estmle2ar1 = @btime AR1($x, estimationmethod=estimatear1mlecheck)
  @info "mle2: $estmle2ar1"=#

  actar1=AR1(;ϕ, b=0.5, σ²=2)
  xsim = simulatear1(actar1, 10^6)
  @assert isapprox(mean(xsim), 2.5, atol=0.01)  "mean(xsim): $(mean(xsim))"
  @assert isapprox(var(xsim), 2/(1-ϕ^2), atol=0.01) "var(xsim): $(var(xsim))"
  ar1sim = AR1(xsim)
  @info "ar1sim: $(ar1sim)"
  @assert isapprox(ar1sim.ϕ, ϕ, atol=0.01)
  @assert isapprox(ar1sim.b, 0.5, atol=0.01)
  @assert isapprox(ar1sim.μ, 2.5, atol=0.01)
  @assert isapprox(ar1sim.σ², 2, atol=0.01)

  xsim! = simulatear1!(actar1, Vector{Float64}(undef,10^6))
  @assert isapprox(mean(xsim), 2.5, atol=0.01)  "mean(xsim): $(mean(xsim))"
  @assert isapprox(var(xsim), 2/(1-ϕ^2), atol=0.01) "var(xsim): $(var(xsim))"
  ar1sim = AR1(xsim!)
  @info "ar1sim: $(ar1sim)"
  @assert isapprox(ar1sim.ϕ, ϕ, atol=0.01)
  @assert isapprox(ar1sim.b, 0.5, atol=0.01)
  @assert isapprox(ar1sim.μ, 2.5, atol=0.01)
  @assert isapprox(ar1sim.σ², 2, atol=0.01)

  Nsimulations=10^4
  dmeantest = 0.2
  #recall the compare func will pass the name of the relevant stat func when it is called
  aggfunc(v, s) = nonparametricnullp(v, Dict(:cor=>0.1, :dmean=>dmeantest)[s])
  statfunc(x,y) = Dict(:cor=>cor(x,y), :dmean=>mean(x)-mean(y))

  @info "std(x): $(std(x))"
  @info "expected std(x): $((2/(1-ϕ^2))^0.5)"
  Sdmean = (2*2/(1-ϕ^2)/length(x))^0.5
  Epdmean = cdf(Normal(), - dmeantest/Sdmean)
  res = comparear1undernull(x, deepcopy(x), statfunc, aggfunc; Nsimulations)
  @info "compareundernull result: $res (Epdmean: $Epdmean)"
  y = deepcopy(x)

  @btime comparear1undernull($x, $y, $statfunc, $aggfunc; Nsimulations=$Nsimulations)
  throw("simulation complete")

  ix = cumsum(x .+ 0.2)
  iy = cumsum(x .+ 0.2)
  res = comparear1undernull(ix, deepcopy(ix), statfunc, aggfunc; Nsimulations)
  @info "compareundernull dmean: $res"
  res = comparearima11undernull(ix, deepcopy(ix), statfunc, aggfunc; Nsimulations)
  @info "compareundernull dmean: $res"

end

@time testar1(N=1000)


# sample from the Muller potential using EM
using Plots
using Printf
using Random
using LinearAlgebra
using ForwardDiff
using StaticArrays

push!(LOAD_PATH,"../src/")

using JuBasicMD

include("potentials.jl")

β = 0.1;
x₀ = [-0.75, 1.5];
Δt = 1e-3;
n_iters = 10^6;

V = x->Muller(x);
cfg = ForwardDiff.GradientConfig(V, x₀);
gradV! = (gradV, x)-> ForwardDiff.gradient!(gradV, V, x, cfg);

sampler = EM(gradV!, β, Δt);

Random.seed!(100);
X₀ = copy(x₀);
sample_trajectory!(X₀, sampler, options=Options(n_iters=n_iters));
@printf("In Place X after %d iterations: (%g, %g)\n",n_iters, X₀[1], X₀[2])

Random.seed!(100);
Xvals = sample_trajectory(x₀, sampler, options=Options(n_iters=n_iters,n_save_iters=n_iters));
X = Xvals[end];
@printf("X after %d iterations: (%g, %g)\n",n_iters, X[1], X[2])

Random.seed!(100);
X_vals = sample_trajectory(x₀, sampler, options=Options(n_iters=n_iters));
histogram2d([X[1] for X in X_vals], [X[2] for X in X_vals],normalize=true,color=:viridis)
xlims!(-1.5,1.5)
ylims!(-0.5, 2.0)
xlabel!("x")
ylabel!("y")


# Script to generate figures in Figure 2.
using Plots, SparseArrays, JLD2, LinearAlgebra, MultiscaleGraphSignalTransforms

# First define a utility function "Glevel", which generate a Graph Signal that
# are piecewise constants whose values are actual region indices "k".
function Glevel(G::GraphSig, GP::GraphPart, j::Int64)
    f = zeros(size(G.f))
    for k in 1:size(GP.rs,1)
        a = GP.rs[k,j+1]
        b = GP.rs[k+1,j+1] - 1
        if b == -1
            break
        end
        f[a:b] .= k*1.0
    end
    Gsub = deepcopy(G)
    Gsub.f[GP.ind] = f
    return Gsub
end

# Set up the resolution and display size
gr(dpi=200, size=(800,600))

# Load the Toronto street network (vehicular volume counts as its graph signal)
JLD2.@load "Toronto_new.jld2"
tmp1 = toronto["G"]
G=GraphSig(tmp1["W"],xy=tmp1["xy"],f=tmp1["f"],name =tmp1["name"],
    plotspecs = tmp1["plotspecs"])

# Assign correct edge weights via 1/Euclidean distance
G = Adj2InvEuc(G)
# Perform the hierarchical graph partitioning using the Fiedler vectors
GP = partition_tree_fiedler(G) # Lrw by default
# Compute the expansion coefficients of the full GHWT c2f dictionary
dmatrix = ghwt_analysis!(G, GP=GP)
N = length(G.f)

# Fig. 2a: Level 1 first since Level 0 is not interesting
j = 1
GraphSig_Plot(Glevel(G,GP,j), linewidth = 1., markersize = 4.,
    markercolor = :viridis, markerstrokealpha =0., notitle = true, nocolorbar = true)
plot!(showaxis = false, ticks = false)
savefig("toronto_j1.pdf")
savefig("toronto_j1.png")

# Fig. 2b: Level 2
j = 2
GraphSig_Plot(Glevel(G,GP,j), linewidth = 1., markersize = 4.,
    markercolor = :viridis, markerstrokealpha =0., notitle = true, nocolorbar = true)
plot!(showaxis = false, ticks = false)
savefig("toronto_j2.pdf")
savefig("toronto_j2.png")

# Fig. 2c: Level 3
j = 3
GraphSig_Plot(Glevel(G,GP,j), linewidth = 1., markersize = 4.,
    markercolor = :viridis, markerstrokealpha =0., notitle = true, nocolorbar = true)
plot!(showaxis = false, ticks = false)
savefig("toronto_j3.pdf")
savefig("toronto_j3.png")


"""
    fgsm_s(R, chars)

Compute the Sᵢⱼ's in FGSM, without taking the product over the choices of the signs of the square roots. Return the smallest absolute value of the FGSM relations.
"""
function fgsm_s(R, chars)
    z = zeros(5);
    theta_constants = [[theta(z, R, char=c) for c in chars_s] for chars_s in chars];
    p = [sqrt(prod(theta_constant)) for theta_constant in theta_constants];
    S = collect(Iterators.flatten([p[1] + i2*p[2] + i3*p[3] + i4*p[4] + i5*p[5] + i6*p[6] for i2 in [1,-1], i3 in [1,-1], i4 in [1,-1], i5 in [1,-1], i6 in [1,-1]]));
    min_fgsm = minimum(abs.(S));
    return min_fgsm;
end

"""
    fgsm(τ, chars=fgsm_chars())

Compute the FGSM relations in genus 5. Return the largest absolute value of the relation.
"""
function fgsm(τ::Array{<:Number}, chars=fgsm_chars())
    R = RiemannMatrix(τ, nderivs=0);
    return fgsm(R, chars);
end

"""
    fgsm(R, chars=fgsm_chars())

Compute the FGSM relations in genus 5. Return the largest absolute value of the relation.
"""
function fgsm(R::RiemannMatrix, chars=fgsm_chars())
    max_fgsm = maximum([fgsm_s(R, c) for c in chars]);
    return max_fgsm;
end

"""
    random_nonfgsm(tol=0.1, trials=100)

Find a random genus 5 matrix in the Siegel upper half space which is not in the FGSM locus, such that the largest FGSM relation has absolute value at least tol, using input number of trials.
"""
function random_nonfgsm(tol::Real=0.1, trials::Integer=100)
    chars = fgsm_chars();
    t = 0; # stores value of largest FGSM relation encountered
    i = 0; # counter for number of trials
    max_matrix = rand(5,5); # stores candidate matrix
    while t < tol && i < trials
        τ = random_siegel(5);
        s = fgsm(τ, chars);
        if s > t
            t = s;
            max_matrix = τ;
        end
        i += 1;
    end
    return [max_matrix, t];
end

"""
    fgsm_chars()

Compute the characteristics used in the FGSM relations in genus 5.
"""
function fgsm_chars()
    chars_34 = [[[[0,0,0,0,0], [0,0,0,0,0]],
                 [[0,0,0,0,0], [1,1,1,1,1]],
                 [[0,0,1,1,0], [0,1,0,0,0]],
                 [[0,0,1,1,0], [1,0,1,1,1]],
                 [[1,1,0,0,0], [0,0,0,1,0]],
                 [[1,1,0,0,0], [1,1,1,0,1]],
                 [[1,1,1,1,0], [0,1,0,1,0]],
                 [[1,1,1,1,0], [1,0,1,0,1]]],
                [[[1,0,1,0,0], [0,0,0,0,0]],
                 [[1,0,1,0,0], [1,1,1,1,1]],
                 [[1,0,0,1,0], [0,1,0,0,0]],
                 [[1,0,0,1,0], [1,0,1,1,1]],
                 [[0,1,1,0,0], [0,0,0,1,0]],
                 [[0,1,1,0,0], [1,1,1,0,1]],
                 [[0,1,0,1,0], [0,1,0,1,0]],
                 [[0,1,0,1,0], [1,0,1,0,1]]],
                [[[0,0,0,0,0], [0,0,1,1,0]],
                 [[0,0,0,0,0], [1,1,0,0,1]],
                 [[0,0,1,1,0], [0,1,1,1,0]],
                 [[0,0,1,1,0], [1,0,0,0,1]],
                 [[1,1,0,0,0], [0,0,1,0,0]],
                 [[1,1,0,0,0], [1,1,0,1,1]],
                 [[1,1,1,1,0], [0,1,1,0,0]],
                 [[1,1,1,1,0], [1,0,0,1,1]]],
                [[[1,0,0,0,1], [0,0,0,0,0]],
                 [[1,0,0,0,1], [1,1,1,1,1]],
                 [[1,0,1,1,1], [0,1,0,0,0]],
                 [[1,0,1,1,1], [1,0,1,1,1]],
                 [[0,1,0,0,1], [0,0,0,1,0]],
                 [[0,1,0,0,1], [1,1,1,0,1]],
                 [[0,1,1,1,1], [0,1,0,1,0]],
                 [[0,1,1,1,1], [1,0,1,0,1]]],
                [[[0,0,1,0,1], [0,0,0,0,0]],
                 [[0,0,1,0,1], [1,1,1,1,1]],
                 [[0,0,0,1,1], [0,1,0,0,0]],
                 [[0,0,0,1,1], [1,0,1,1,1]],
                 [[1,1,1,0,1], [0,0,0,1,0]],
                 [[1,1,1,0,1], [1,1,1,0,1]],
                 [[1,1,0,1,1], [0,1,0,1,0]],
                 [[1,1,0,1,1], [1,0,1,0,1]]],
                [[[1,0,0,0,1], [0,0,1,1,0]],
                 [[1,0,0,0,1], [1,1,0,0,1]],
                 [[1,0,1,1,1], [0,1,1,1,0]],
                 [[1,0,1,1,1], [1,0,0,0,1]],
                 [[0,1,0,0,1], [0,0,1,0,0]],
                 [[0,1,0,0,1], [1,1,0,1,1]],
                 [[0,1,1,1,1], [0,1,1,0,0]],
                 [[0,1,1,1,1], [1,0,0,1,1]]]];
    chars_35 = [[[[0,0,0,0,0], [0,0,0,0,0]],
                 [[0,0,0,0,0], [1,1,1,1,1]],
                 [[0,0,1,0,1], [0,1,0,0,0]],
                 [[0,0,1,0,1], [1,0,1,1,1]],
                 [[1,1,0,0,0], [0,0,0,0,1]],
                 [[1,1,0,0,0], [1,1,1,1,0]],
                 [[1,1,1,0,1], [0,1,0,0,1]],
                 [[1,1,1,0,1], [1,0,1,1,0]]],
                [[[1,0,1,0,0], [0,0,0,0,0]],
                 [[1,0,1,0,0], [1,1,1,1,1]],
                 [[1,0,0,0,1], [0,1,0,0,0]],
                 [[1,0,0,0,1], [1,0,1,1,1]],
                 [[0,1,1,0,0], [0,0,0,0,1]],
                 [[0,1,1,0,0], [1,1,1,1,0]],
                 [[0,1,0,0,1], [0,1,0,0,1]],
                 [[0,1,0,0,1], [1,0,1,1,0]]],
                [[[0,0,0,0,0], [0,0,1,0,1]],
                 [[0,0,0,0,0], [1,1,0,1,0]],
                 [[0,0,1,0,1], [0,1,1,0,1]],
                 [[0,0,1,0,1], [1,0,0,1,0]],
                 [[1,1,0,0,0], [0,0,1,0,0]],
                 [[1,1,0,0,0], [1,1,0,1,1]],
                 [[1,1,1,0,1], [0,1,1,0,0]],
                 [[1,1,1,0,1], [1,0,0,1,1]]],
                [[[1,0,0,1,0], [0,0,0,0,0]],
                 [[1,0,0,1,0], [1,1,1,1,1]],
                 [[1,0,1,1,1], [0,1,0,0,0]],
                 [[1,0,1,1,1], [1,0,1,1,1]],
                 [[0,1,0,1,0], [0,0,0,0,1]],
                 [[0,1,0,1,0], [1,1,1,1,0]],
                 [[0,1,1,1,1], [0,1,0,0,1]],
                 [[0,1,1,1,1], [1,0,1,1,0]]],
                [[[0,0,1,1,0], [0,0,0,0,0]],
                 [[0,0,1,1,0], [1,1,1,1,1]],
                 [[0,0,0,1,1], [0,1,0,0,0]],
                 [[0,0,0,1,1], [1,0,1,1,1]],
                 [[1,1,1,1,0], [0,0,0,0,1]],
                 [[1,1,1,1,0], [1,1,1,1,0]],
                 [[1,1,0,1,1], [0,1,0,0,1]],
                 [[1,1,0,1,1], [1,0,1,1,0]]],
                [[[1,0,0,1,0], [0,0,1,0,1]],
                 [[1,0,0,1,0], [1,1,0,1,0]],
                 [[1,0,1,1,1], [0,1,1,0,1]],
                 [[1,0,1,1,1], [1,0,0,1,0]],
                 [[0,1,0,1,0], [0,0,1,0,0]],
                 [[0,1,0,1,0], [1,1,0,1,1]],
                 [[0,1,1,1,1], [0,1,1,0,0]],
                 [[0,1,1,1,1], [1,0,0,1,1]]]];
    chars_45 = [[[[0,0,0,0,0], [0,0,0,0,0]],
                 [[0,0,0,0,0], [1,1,1,1,1]],
                 [[0,0,0,1,1], [0,1,0,0,0]],
                 [[0,0,0,1,1], [1,0,1,1,1]],
                 [[1,1,0,0,0], [0,0,0,0,1]],
                 [[1,1,0,0,0], [1,1,1,1,0]],
                 [[1,1,0,1,1], [0,1,0,0,1]],
                 [[1,1,0,1,1], [1,0,1,1,0]]],
                [[[1,0,0,1,0], [0,0,0,0,0]],
                 [[1,0,0,1,0], [1,1,1,1,1]],
                 [[1,0,0,0,1], [0,1,0,0,0]],
                 [[1,0,0,0,1], [1,0,1,1,1]],
                 [[0,1,0,1,0], [0,0,0,0,1]],
                 [[0,1,0,1,0], [1,1,1,1,0]],
                 [[0,1,0,0,1], [0,1,0,0,1]],
                 [[0,1,0,0,1], [1,0,1,1,0]]],
                [[[0,0,0,0,0], [0,0,0,1,1]],
                 [[0,0,0,0,0], [1,1,1,0,0]],
                 [[0,0,0,1,1], [0,1,0,1,1]],
                 [[0,0,0,1,1], [1,0,1,0,0]],
                 [[1,1,0,0,0], [0,0,0,1,0]],
                 [[1,1,0,0,0], [1,1,1,0,1]],
                 [[1,1,0,1,1], [0,1,0,1,0]],
                 [[1,1,0,1,1], [1,0,1,0,1]]],
                [[[1,0,1,0,0], [0,0,0,0,0]],
                 [[1,0,1,0,0], [1,1,1,1,1]],
                 [[1,0,1,1,1], [0,1,0,0,0]],
                 [[1,0,1,1,1], [1,0,1,1,1]],
                 [[0,1,1,0,0], [0,0,0,0,1]],
                 [[0,1,1,0,0], [1,1,1,1,0]],
                 [[0,1,1,1,1], [0,1,0,0,1]],
                 [[0,1,1,1,1], [1,0,1,1,0]]],
                [[[0,0,1,1,0], [0,0,0,0,0]],
                 [[0,0,1,1,0], [1,1,1,1,1]],
                 [[0,0,1,0,1], [0,1,0,0,0]],
                 [[0,0,1,0,1], [1,0,1,1,1]],
                 [[1,1,1,1,0], [0,0,0,0,1]],
                 [[1,1,1,1,0], [1,1,1,1,0]],
                 [[1,1,1,0,1], [0,1,0,0,1]],
                 [[1,1,1,0,1], [1,0,1,1,0]]],
                [[[1,0,1,0,0], [0,0,0,1,1]],
                 [[1,0,1,0,0], [1,1,1,0,0]],
                 [[1,0,1,1,1], [0,1,0,1,1]],
                 [[1,0,1,1,1], [1,0,1,0,0]],
                 [[0,1,1,0,0], [0,0,0,1,0]],
                 [[0,1,1,0,0], [1,1,1,0,1]],
                 [[0,1,1,1,1], [0,1,0,1,0]],
                 [[0,1,1,1,1], [1,0,1,0,1]]]];
    return [chars_34, chars_35, chars_45];
end


using PyPlot, Images
using PyCall

function trapz(a)
  return sum(a) - 0.5*a[1] - 0.5*a[end]
end

function lid_method(img_name="../DICimages/DICmandrill.png")

  raw_img = PyPlot.imread(img_name);
  img = convert(Image{Images.Gray},raw_img[:,:,1:3]);
  M = convert(Array,img);
  M = float(M);
  M= M'
  M = M - mean(mean(M))
  intM = [ trapz( M[1:i,j] ) for i in 1:size(M,1), j in 1:size(M,2)]

  maxi = maximum(intM)
  mini = minimum(intM)
  clamped = 1 +  (-intM + maxi) * 1 / ( mini-maxi);
  clamped = convert(Array{Float64,2},clamped)
  out_name = "../images/LID" * basename(img_name)[4:end]
  imwrite(grayim(clamped),out_name)

  @printf("file outputted to /images/")

  return
end

function exdecay_method(img_name="../DICimages/DICmandrill.png"; s=10.0)

    raw_img = PyPlot.imread(img_name);
    img = convert(Image{Images.Gray},raw_img[:,:,1:3]);
    M = convert(Array,img);
    M = float(M);
    M= M'
    M = M - mean(mean(M))
    lx, ly = size(M)
    xrange =[ -lx+1:1:lx-1]

    lM = [ trapz( M[1:i,j]   .* exp(xrange[lx+1-i:lx]/s ) ) - trapz( M[i:end,j] .* exp(-1*xrange[lx:end+1-i]/s ) )  for i in 1:lx, j in 1:ly]
    #rM = [ trapz( M[i:end,j] .* exp(-1*xrange[lx:end+1-i]/s ) ) for i in 1:lx, j in 1:ly]

    intM =lM

    maxi = maximum(intM)
    mini = minimum(intM)
    clamped = 1 +  (-intM + maxi) * 1 / ( mini-maxi);
    clamped = convert(Array{Float64,2},clamped)
    out_name = "../images/exp" * basename(img_name)[4:end]
    imwrite(grayim(clamped),out_name)

    @printf("file outputted to /images/")

    return
  end


@testset "Orderings" begin

    import KnuthBendix.set_inversion!
    import Base.Order.lt

    A = Alphabet(['a', 'b', 'c', 'd'])
    set_inversion!(A, 'a', 'b')
    set_inversion!(A, 'c', 'd')

    lenlexord = LenLex(A)

    @test lenlexord isa Base.Order.Ordering

    u1 = Word([1,2])
    u3 = Word([1,3])
    u4 = Word([1,2,3])
    u5 = Word([1,4,2])

    @test lt(lenlexord, u1, u3) == true
    @test lt(lenlexord, u3, u1) == false
    @test lt(lenlexord, u3, u4) == true
    @test lt(lenlexord, u4, u5) == true
    @test lt(lenlexord, u5, u4) == false
    @test lt(lenlexord, u1, u1) == false

    wo = WreathOrder(A)
    @test wo isa KnuthBendix.WordOrdering

    w1 = Word([3,1,4,2,3])
    w5 = Word([1,2,3,2,4,1,2,4])
    w4 = Word([1,3,1,2,4,2,1])
    w6 = Word([1,2,3,2,1,4,2,4])
    w2 = Word([1,3,4,3,2])
    w3 = Word([1,3,2,1,4,1,2])
    ε = one(w1)

    a = [w1, w5, w4, w6, w2, w3, ε]

    @test sort(a, order = wo) == [ε, w1, w2, w3, w4, w5, w6]

    rpo = RecursivePathOrder(A)

    w1 = Word([1])
    w2 = Word([2])
    w14 = Word([1,4])
    w214 = Word([2,1,4])
    w41 = Word([4,1])
    w241 = Word([2,4,1])
    w1224 = Word([1,2,3,2,1,4,2,4])
    w32141 = Word([3,2,1,4,1])

    a = [w14, w1, w1224, w214, ε]
    b = [w32141, w214, w1, ε, w241, w2, w41]

    @test sort(a, order = rpo) == [ε, w1, w14, w214, w1224]
    @test sort(b, order = rpo) == [ε, w1, w2, w214, w41, w241, w32141]

end


import Base: axpy!, scale!, ctranspose, dot, size
import ArrayViews: view

export AbstractSketch, IdentitySketch, SymmetricSketch, AsymmetricSketch,
	additive_update!, reconstruct, dot

abstract AbstractSketch

# Identity

type IdentitySketch<:AbstractSketch
	X
end

function additive_update!(s::IdentitySketch, tilde_X, eta=1)
	axpy!(eta, tilde_X, s.X)
end

function cgd_update!(s::IdentitySketch, tilde_X, eta)
	scale!(s.X, 1-eta)
	axpy!(eta, tilde_X, s.X)
end

function reconstruct(s::IdentitySketch)
	return s.X
end

# Symmetric

type SymmetricSketch<:AbstractSketch
	r::Int
	Omega::Array{Float64,2}
	Psi::Array{Float64,2}
	Y::Array{Float64,2}
	W::Array{Float64,2}
end

function SymmetricSketch(n::Int, r::Int)
	k = 2r + 1
	l = 4r + 3 # = 2k+1
	Omega = randn(n,k)
	Psi   = randn(n,l)
	Y     = zeros(n,k)
	W     = zeros(n,l)
	return SymmetricSketch(r, Omega, Psi, Y, W)
end

function additive_update!(s::SymmetricSketch, vvt, eta=1)
	# Y = Y + eta*v*v'*Omega
	axpy!(eta, vvt*s.Omega, s.Y)
	# W = W + eta*v*v'*Psi
	axpy!(eta, vvt*s.Psi, s.W)
end

function cgd_update!(s::SymmetricSketch, vvt, eta)
	# Y = (1-eta)Y + eta*v*v'*Omega
	scale!(s.Y, 1-eta)
	axpy!(eta, vvt*s.Omega, s.Y)
	# W = (1-eta)W + eta*v*v'*Psi
	scale!(s.W, 1-eta)
	axpy!(eta, vvt*s.Psi, s.W)
end

function reconstruct(s::SymmetricSketch)
	# Q = orth(s.Y)
	Q,_ = qr(s.Y)
	B = (Q's.W) / (Q's.Psi)
	d,V = eigs((B+B')/2, nev=s.r, which=:LM)
	U = Q*V
	d = pos(d)
	return LowRankOperator(U, spdiagm(d)*U') # reconstruction as square matrix is U*diagm(d)*U'
end

# Asymmetric

type AsymmetricSketch<:AbstractSketch
	r::Int
	Omega::Array{Float64,2}
	Psi::Array{Float64,2}
	Y::Array{Float64,2}
	W::Array{Float64,2}
end

function AsymmetricSketch(m::Int, n::Int, r::Int)
	k = 2r + 1
	l = 4r + 3 # == 2k+1
	Omega = randn(n,k)
	Psi   = randn(m,l)
	Y     = zeros(m,k)
	W     = zeros(n,l)
	return AsymmetricSketch(r, Omega, Psi, Y, W)
end

function additive_update!(s::AsymmetricSketch, uvt, eta=1)
	# Y = Y + eta*v*v'*Omega
	axpy!(eta, uvt*s.Omega, s.Y)
	# W = W + eta*v*v'*Psi
	axpy!(eta, uvt'*s.Psi, s.W)
end

function cgd_update!(s::AsymmetricSketch, uvt, eta)
	# Y = (1-eta)Y + eta*v*v'*Omega
	scale!(s.Y, 1-eta)
	axpy!(eta, uvt*s.Omega, s.Y)
	# W = (1-eta)W + eta*v*v'*Psi
	scale!(s.W, 1-eta)
	axpy!(eta, uvt'*s.Psi, s.W)
end

function reconstruct(s::AsymmetricSketch, r::Int=s.r)
	# Q = orth(s.Y)
	k = 2r + 1
	l = 4r + 3
	Q,_ = qr(view(s.Y,:,1:k))
	B = view(s.W,:,1:l) / (Q'*view(s.Psi,:,1:l)) # Q's.Psi is k x l, its pinv is l x k, so B is n x k
	U,s,V = mysvds(B, nsv=r) # U is n x r
	return LowRankOperator(Q*V, spdiagm(s)*U') # reconstruction as square matrix is Q*V*diagm(s)*U'
end
# function reconstruct(s::AsymmetricSketch)
# 	# Q = orth(s.Y)
# 	Q,_ = qr(s.Y)
# 	B = s.W * pinv(Q'*s.Psi) # B is k x l
# 	U,s,V,_ = svds(B, nsv=s.r)
# 	return LowRankOperator(Q*U, spdiagm(s), V) # reconstruction as square matrix is Q*U*diagm(s)*V'
# end

### utilities
pos(a::Number) = max(a,0)
function pos(a::Array)
	for i in length(a)
		a[i] = pos(a[i])
	end
	a
end

# not yet efficient
function dot(M::SparseMatrixCSC{Float64,Int64}, A::Array{Float64,2})
	dot(M.nzval, A[find(M)])
end


using HMatrices
using LoopVectorization
using StaticArrays
using LinearAlgebra

const EPS = 1e-8

# some simple implementation of laplace and helmholtz single layer matrices.
function laplace_matrix(X, Y)
    f = (x, y) -> begin
        EPS = 1e-8 # fudge factor to avoid division by zero
        d = norm(x - y) + EPS
        inv(4π * d)
    end
    KernelMatrix(f, X, Y)
    end

function helmholtz_matrix(X, Y, k)
    f = (x, y) -> begin
        EPS = 1e-8 # fudge factor to avoid division by zero
        d = norm(x - y) + EPS
        exp(im * k * d) * inv(4π * d)
    end
    KernelMatrix(f, X, Y)
end

function elastostatic_matrix(X, Y, μ, λ)
    f = (x, y) -> begin
        ν = λ / (2 * (μ + λ))
        r = x - y
        d = norm(r) + EPS
        RRT = r * transpose(r) # r ⊗ rᵗ
        return 1 / (16π * μ * (1 - ν) * d) * ((3 - 4 * ν) * I + RRT / d^2)
    end
    KernelMatrix(f, X, Y)
end

function elastosdynamic_matrix(X, Y, μ, λ, ω, ρ)
    f = (x, y) -> begin
        c1 = sqrt((λ + 2μ) / ρ)
        c2 = sqrt(μ / ρ)
        r = x - y
        d   = norm(r) + EPS
        RRT = r * transpose(r) # r ⊗ rᵗ
        s   = -im * ω
        z1  = s * d / c1
        z2  = s * d / c2
        α   = 4
        ψ   = exp(-z2) / d + (1 + z2) / (z2^2) * exp(-z2) / d - c2^2 / c1^2 * (1 + z1) / (z1^2) * exp(-z1) / d
        chi = 3 * ψ - 2 * exp(-z2) / d - c2^2 / c1^2 * exp(-z1) / d
        return 1 / (α * π * μ) * (ψ * I - chi * RRT / d^2) + x * transpose(y)
    end
    KernelMatrix(f, X, Y)
end

"""
    struct LaplaceMatrixVec{T,Td} <: AbstractKernelMatrix{T}

Vectorized version of [`laplace_matrix`](@ref).
"""
struct LaplaceMatrixVec{T,Td} <: AbstractKernelMatrix{T}
    X::Matrix{Td}
    Y::Matrix{Td}
    function LaplaceMatrixVec{T}(X::Matrix{Td}, Y::Matrix{Td}) where {T,Td}
        @assert size(X, 2) == size(Y, 2) == 3
        new{T,Td}(X, Y)
    end
end

Base.size(K::LaplaceMatrixVec) = size(K.X, 1), size(K.Y, 1)

# constructor based on Vector of StaticVector
function LaplaceMatrixVec{T}(_X::Vector{SVector{3,Td}}, _Y::Vector{SVector{3,Td}}) where {T,Td}
    X = reshape(reinterpret(Td, _X), 3, :) |> transpose |> collect
    Y = reshape(reinterpret(Td, _Y), 3, :) |> transpose |> collect
    LaplaceMatrixVec{T}(X, Y)
end
LaplaceMatrixVec(args...) = LaplaceMatrixVec{Float64}(args...) # default to Float64

function Base.getindex(K::LaplaceMatrixVec{T}, i::Int, j::Int)::T where {T}
    d2 = (K.X[i,1] - K.Y[j,1])^2 + (K.X[i,2] - K.Y[j,2])^2 + (K.X[i,3] - K.Y[j,3])^2
    d = sqrt(d2) + EPS
    return inv(4π * d)
end
function Base.getindex(K::LaplaceMatrixVec, I::UnitRange, J::UnitRange)
    T = eltype(K)
    m = length(I)
    n = length(J)
    Xv = view(K.X, I, :)
    Yv = view(K.Y, J, :)
    out = Matrix{T}(undef, m, n)
    @avx for j in 1:n
        for i in 1:m
            d2 = (Xv[i,1] - Yv[j,1])^2
            d2 += (Xv[i,2] - Yv[j,2])^2
            d2 += (Xv[i,3] - Yv[j,3])^2
            d = sqrt(d2) + EPS
            out[i,j] = inv(4 * π * d)
        end
    end
    return out
end
function Base.getindex(K::LaplaceMatrixVec, I::UnitRange, j::Int)
    return vec(K[I,j:j])
end
function Base.getindex(adjK::Adjoint{<:Any,LaplaceMatrixVec}, I::UnitRange, j::Int)
    K = parent(adjK)
    vec(K[j:j,I])
end

"""
    struct HelmholtzMatrix{T,Td,Tk} <: AbstractKernelMatrix{T}

Vectorized version of [`helmholtz_matrix`](@ref).
"""
struct HelmholtzMatrixVec{T,Td,Tk} <: AbstractKernelMatrix{T}
    X::Matrix{Td}
    Y::Matrix{Td}
    k::Tk
    function HelmholtzMatrixVec{T}(X::Matrix{Td}, Y::Matrix{Td}, k::Tk) where {T,Td,Tk}
        @assert size(X, 2) == size(Y, 2) == 3
        new{T,Td,Tk}(X, Y, k)
    end
end
HelmholtzMatrixVec(args...) = HelmholtzMatrixVec{ComplexF64}(args...)

Base.size(K::HelmholtzMatrixVec) = size(K.X, 1), size(K.Y, 1)

function HelmholtzMatrixVec{T}(_X::Vector{SVector{3,Td}}, _Y::Vector{SVector{3,Td}}, k) where {T,Td}
    X = reshape(reinterpret(Td, _X), 3, :) |> transpose |> collect
    Y = reshape(reinterpret(Td, _Y), 3, :) |> transpose |> collect
    HelmholtzMatrixVec{T}(X, Y, k)
end

function Base.getindex(K::HelmholtzMatrixVec{T}, i::Int, j::Int)::T where {T}
    d2 = (K.X[i,1] - K.Y[j,1])^2 + (K.X[i,2] - K.Y[j,2])^2 + (K.X[i,3] - K.Y[j,3])^2
    d  = sqrt(d2) + EPS
    return inv(4π * d) * exp(im * K.k * d)
end
function Base.getindex(K::HelmholtzMatrixVec{Complex{T}}, I::UnitRange, J::UnitRange) where {T}
    k = K.k
    m = length(I)
    n = length(J)
    Xv = view(K.X, I, :)
    Yv = view(K.Y, J, :)
    # since LoopVectorization does not (yet) support Complex{T} types, we will
    # reinterpret the output as a Matrix{T}, then use views. This can probably be
    # done better.
    out    = Matrix{Complex{T}}(undef, m, n)
    out_T  = reinterpret(T, out)
    out_r = @views out_T[1:2:end,:]
    out_i = @views out_T[2:2:end,:]
    @avx for j in 1:n
        for i in 1:m
            d2 = (Xv[i,1] - Yv[j,1])^2
            d2 += (Xv[i,2] - Yv[j,2])^2
            d2 += (Xv[i,3] - Yv[j,3])^2
            d  = sqrt(d2) + EPS
            s, c = sincos(k * d)
            zr = inv(4π * d) * c
            zi = inv(4π * d) * s
            out_r[i,j] = zr
            out_i[i,j] = zi
        end
    end
    return out
end
function Base.getindex(K::HelmholtzMatrixVec, I::UnitRange, j::Int)
    vec(K[I,j:j])
end
function Base.getindex(adjK::Adjoint{<:Any,<:HelmholtzMatrixVec}, I::UnitRange, j::Int)
    K = parent(adjK)
    conj!(K[j:j,I]) |> vec
end

function points_on_sphere(npts, R=1)
    theta = π * rand(npts)
    phi   = 2 * π * rand(npts)
    x     = @. sin(theta) * cos(phi)
    y     = @. R * sin(theta) * sin(phi)
    z     = @. R * cos(theta)
    data  = vcat(x', y', z')
    pts = reinterpret(SVector{3,Float64}, vec(data)) |> collect
    return pts
end

function points_on_cylinder(radius,n,shift=SVector(0,0,0))
    step     = 1.75*π*radius/sqrt(n)
    result          = Vector{SVector{3,Float64}}(undef,n)
    length          = 2*π*radius
    pointsPerCircle = length/step
    angleStep       = 2*π/pointsPerCircle
    for i=0:n-1
        x = radius * cos(angleStep*i)
        y = radius * sin(angleStep*i)
        z = step*i/pointsPerCircle
        result[i+1] = shift + SVector(x,y,z)
    end
    return result
end


using Statistics
using Printf

using Images
using Noise

include("hopfield.jl")

using .HopfieldNetwork

const base_d = "data/EnglishHnd/English/Hnd/Img/"
const pre_d = "Sample"

function get_img_files(s_char::AbstractChar)
    n = s_char <= '9' ? s_char - '0' + 1 : s_char - 'A' + 11
    readdir(joinpath(base_d, pre_d*(@sprintf("%03d", n))), join=true)
end

function get_img_hnd(
        out_t::typeof(AbstractFloat),
        s_char::AbstractChar,
        (h, w)::Tuple{Integer, Integer},
)::Vector{Matrix{out_t}}
    Matrix{out_t}[out_t.(Gray.(imresize(load(f), (h, w)))) for f ∈ get_img_files(s_char)]
end

function get_data(chars::StepRange{Char, Int}, (h, w)::Tuple{Integer, Integer}; wnoise::Bool=false)::Matrix{Float64}
    rawdata = Matrix{Float64}[]
    for c in chars
        imgs = get_img_hnd(Float64, c, (h, w))
        best_img = imgs[argmin(mean.(imgs))]
        if wnoise
            best_img = salt_pepper(best_img)
        end
        push!(rawdata, best_img)
    end
    data = Vector{Float64}[]
    for d in rawdata
        push!(data, vcat(d...))
    end
    data = collect(hcat(data...)')

    rep(x) = iszero(x) ? 1.0 : -1.0
    data = rep.(data)
    data
end

function run_test(chars::StepRange{Char, Int}, (h, w)::Tuple{Integer, Integer}; wnoise::Bool=false)::AbstractFloat
    traindata = get_data(chars, (h, w))
    testdata = get_data(chars, (h, w), wnoise=wnoise)
    hn = SimpleHN{Float64}(traindata, testdata)
    1 - predict_error(hn)
end

function hamming_dis(chars::StepRange{Char, Int}, (h, w)::Tuple{Integer, Integer})::Matrix{Float64}
    data = get_data(chars, (h, w))
    res = zeros(size(data, 1), size(data, 1))
    for (i1, x1) ∈ enumerate(eachrow(data)), (i2, x2) ∈ enumerate(eachrow(data))
        res[i1, i2] = sum(x1 .!= x2)
    end
    res
end

let quz = "question-1-2"
    @show run_test('0':'4', (10, 10))
    @show run_test('A':'E', (10, 10))
    println("$(quz) done.")
end

let quz = "question-1-3"
    println("0--4")
    h0_4 = hamming_dis('0':'4', (10, 10))
    display(h0_4)
    println()
    @show mean(h0_4)

    println("A--E")
    hA_E = hamming_dis('A':'E', (10, 10))
    display(hA_E)
    println()
    @show mean(hA_E)
    println("$(quz) done.")
end

let quz = "question-1-4"
    @show run_test('0':'4', (10, 10), wnoise=true)
    @show run_test('A':'E', (10, 10), wnoise=true)
    println("$(quz) done.")
end

let quz = "question-2"
    @show run_test('0':'4', (30, 30))
    @show run_test('A':'E', (30, 30))
    println("$(quz) done.")
end

let quz = "question-3"
    @show run_test('0':'9', (30, 30))
    @show run_test('A':'O', (30, 30))
    println("$(quz) done.")
end


# The original license on the t-tests from HypothesisTests.jl
# is as follows. We're including it here because we're mostly copying the
# documentation of the functions.
#
# t.jl
# Various forms of t-tests
#
# Copyright (C) 2012   Simon Kornblith
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import HypothesisTests.OneSampleTTest
import HypothesisTests.EqualVarianceTTest
import HypothesisTests.UnequalVarianceTTest

"""
    EqualVarianceTTest(d1::AbstractUncertainValue, d2::AbstractUncertainValue,
        n::Int = 1000; μ0::Real = 0) -> EqualVarianceTTest

Consider two samples `s1` and `s2`, each consisting of `n` random draws from the
distributions furnishing `d1` and `d2`, respectively.

This function performs a two-sample t-test of the null hypothesis that `s1` and
`s2` come from distributions with equal means and variances against the
alternative hypothesis that the distributions have different means but equal
variances.
"""
function EqualVarianceTTest(d1::AbstractUncertainValue,
        d2::AbstractUncertainValue,
        n::Int = 1000;
        μ0::Real = 0)

    x = resample(d1, n)
    y = resample(d2, n)
    EqualVarianceTTest(x, y, μ0)
end

"""
    EqualVarianceTTestPooled(d1::UncertainDataset, d2::UncertainDataset,
        n::Int = 1000; μ0::Real = 0) -> EqualVarianceTTest

Consider two samples `s1[i]` and `s2[i]`, each consisting of `n` random draws
from the distributions furnishing the uncertain values `d1[i]` and `d2[i]`,
respectively. Gather all `s1[i]` in a pooled sample `S1`, and all `s2[i]` in
a pooled sample `S2`.

Perform a two-sample t-test of the null hypothesis that `S1` and `S2` come from
distributions with equal means and variances against the alternative hypothesis
that the distributions have different means but equal variances.
"""
function EqualVarianceTTestPooled(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000;
        μ0::Real = 0)

    x = vcat(resample(d1, n)...,)
    y = vcat(resample(d2, n)...,)

    EqualVarianceTTest(x, y, μ0)
end

"""
    EqualVarianceTTestElementWise(d1::UncertainDataset, d2::UncertainDataset,
        n::Int = 1000; μ0::Real = 0) -> Vector{EqualVarianceTTest}

Consider two samples `s1[i]` and `s2[i]`, each consisting of `n` random draws
from the distributions furnishing the uncertain values `d1[i]` and `d2[i]`,
respectively. This function performs an elementwise `EqualVarianceTTest`
on the pairs `(s1[i], s2[i])`. Specifically:

Performs an pairwise two-sample t-test of the null hypothesis that `s1[i]` and
`s2[i]` come from distributions with equal means and variances against the
alternative hypothesis that the distributions have different means but equal
variances.
"""
function EqualVarianceTTestElementWise(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000;
        μ0::Real = 0)

    N = length(d1)
    [EqualVarianceTTest(resample(d1[i], n), resample(d2[i], n), μ0) for i = 1:N]
end

"""
    UnequalVarianceTTest(d1::AbstractUncertainValue, d2::AbstractUncertainValue,
        n::Int = 1000; μ0::Real = 0) -> UnequalVarianceTTest

Consider two samples `s1` and `s2`, each consisting of `n` random draws from the
distributions furnishing `d1` and `d2`, respectively.

Perform an unequal variance two-sample t-test of the null hypothesis that `s1`
and `s2` come from distributions with equal means against the alternative
hypothesis that the distributions have different means.
"""
function UnequalVarianceTTest(d1::AbstractUncertainValue,
        d2::AbstractUncertainValue,
        n::Int = 1000;
        μ0::Real = 0)

    x = resample(d1, n)
    y = resample(d2, n)
    UnequalVarianceTTest(x, y, μ0)
end

"""
    UnequalVarianceTTestPooled(d1::UncertainDataset, d2::UncertainDataset,
        n::Int = 1000; μ0::Real = 0) -> UnequalVarianceTTest

Consider two samples `s1[i]` and `s2[i]`, each consisting of `n` random draws
from the distributions furnishing the uncertain values `d1[i]` and `d2[i]`,
respectively. Gather all `s1[i]` in a pooled sample `S1`, and all `s2[i]` in
a pooled sample `S2`.

This function performs an unequal variance two-sample t-test of the null
hypothesis that `S1` and `S2` come from distributions with equal means against
the alternative hypothesis that the distributions have different means.

"""
function UnequalVarianceTTestPooled(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000;
        μ0::Real = 0)

    x = vcat(resample(d1, n)...,)
    y = vcat(resample(d2, n)...,)

    UnequalVarianceTTest(x, y, μ0)
end

"""
    UnequalVarianceTTestElementWise(d1::UncertainDataset, d2::UncertainDataset,
        n::Int = 1000; μ0::Real = 0) -> Vector{UnequalVarianceTTest}

Consider two samples `s1[i]` and `s2[i]`, each consisting of `n` random draws
from the distributions furnishing the uncertain values `d1[i]` and `d2[i]`,
respectively. This function performs an elementwise `EqualVarianceTTest`
on the pairs `(s1[i], s2[i])`. Specifically:

Performs an pairwise unequal variance two-sample t-test of the null hypothesis
that `s1[i]` and `s2[i]` come from distributions with equal means against the
alternative hypothesis that the distributions have different means.

This test is sometimes known as Welch's t-test. It differs from the equal
variance t-test in that it computes the number of degrees of freedom of the
test using the Welch-Satterthwaite equation:

```math
    ν_{χ'} ≈ \\frac{\\left(\\sum_{i=1}^n k_i s_i^2\\right)^2}{\\sum_{i=1}^n
        \\frac{(k_i s_i^2)^2}{ν_i}}
```
"""
function UnequalVarianceTTestElementWise(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000;
        μ0::Real = 0)

    N = length(d1)
    [UnequalVarianceTTest(resample(d1[i], n), resample(d2[i], n), μ0) for i = 1:N]
end

"""
    OneSampleTTest(d::AbstractUncertainValue, n::Int = 1000;
        μ0::Real = 0) -> OneSampleTTest

Perform a one sample t-test of the null hypothesis that the uncertain value has
a distribution with mean `μ0` against the alternative hypothesis that its
distribution does not have mean `μ0`. `n` indicates the number of draws during
resampling.
"""
function OneSampleTTest(d::AbstractUncertainValue, n::Int = 1000; μ0::Real = 0)
    x = resample(d, n)
    OneSampleTTest(x, μ0)
end

"""
    OneSampleTTest(d1::AbstractUncertainValue,
                    d2::AbstractUncertainValue,
                    n::Int = 1000; μ0::Real=0) -> OneSampleTTest

Perform a paired sample t-test of the null hypothesis that the differences
between pairs of uncertain values in `d1` and `d2` come from a distribution
with mean `μ0` against the alternative hypothesis that the distribution does
not have mean `μ0`.
"""
function OneSampleTTest(d1::AbstractUncertainValue,
                        d2::AbstractUncertainValue,
                        n::Int = 1000; μ0::Real = 0)
    x = resample(d1, n)
    y = resample(d2, n)
    OneSampleTTest(x, y, μ0)
end

"""
    OneSampleTTestPooled(d::UncertainDataset, n::Int = 1000;
        μ0::Real=0) -> OneSampleTTest

First, sample `n` draws of each uncertain value in `d1`, then pooling the draws
together. Then, perform a one sample t-test of the null hypothesis that the
uncertain values have a pooled distribution with mean `μ0` against the
alternative hypothesis that its pooled distribution does not have mean `μ0`.
`n` indicates the number of draws during resampling.
"""
function OneSampleTTestPooled(d::UncertainDataset, n::Int = 1000; μ0::Real = 0)
    x = vcat([resample(uncertainval, n) for uncertainval in d]...,)
    OneSampleTTest(x, μ0)
end

"""
    OneSampleTTestElementWise(d::UncertainDataset, n::Int = 1000;
        μ0::Real = 0) -> Vector{OneSampleTTest}

First, sample `n` draws of each uncertain value in `d`, then pooling the draws
together. Then, perform a one sample t-test of the null hypothesis that the
uncertain values have a pooled distribution with mean `μ0` against the
alternative hypothesis that its pooled distribution
does not have mean `μ0`. `n` indicates the number of draws during resampling.
"""
function OneSampleTTestElementWise(d::UncertainDataset, n::Int = 1000; μ0::Real = 0)
    [OneSampleTTest(resample(uncertainval, n), μ0) for uncertainval in d]
end

"""
    OneSampleTTestPooled(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000; μ0::Real = 0) -> OneSampleTTest

First, sample `n` draws of each uncertain value in each dataset, pooling the
draws from the elements of `d1` and the draws from the elements of `d2` separately.
Then, perform a paired sample t-test of the null hypothesis that the differences
between pairs of uncertain values in `d1` and `d2` come from a distribution
with mean `μ0` against the alternative hypothesis that the distribution does
not have mean `μ0`.
"""
function OneSampleTTestPooled(d1::UncertainDataset,
                                d2::UncertainDataset,
                                n::Int = 1000; μ0::Real = 0)
    x = vcat(resample(d1, n)...,)
    y = vcat(resample(d2, n)...,)
    OneSampleTTest(x, y, μ0)
end

"""
    OneSampleTTestElementWise(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000; μ0::Real = 0) -> Vector{OneSampleTTest}

Perform a one sample t-test of the null hypothesis that the uncertain value has
a distribution with mean `μ0` against the alternative hypothesis that its
distribution does not have mean `μ0` for uncertain value in `d`.

`n` indicates the number of draws during resampling.
"""
function OneSampleTTestElementWise(d1::UncertainDataset,
                                    d2::UncertainDataset,
                                    n::Int = 1000; μ0::Real = 0)
    N = length(d1)
    [OneSampleTTest(resample(d1[i], n), resample(d2[i], n), μ0) for i in 1:N]
end

export
EqualVarianceTTest,
EqualVarianceTTestPooled,
EqualVarianceTTestElementWise,

UnequalVarianceTTest,
UnequalVarianceTTestPooled,
UnequalVarianceTTestElementWise,

OneSampleTTest,
OneSampleTTestPooled,
OneSampleTTestElementWise


using CVChannel
using Test

"""
This script shows the multiplicativity of the Siddhu channel with itself.
"""

println("In this script we look at the parallel communication value of the Siddhu channel with itself.")
println("Since we know cv(N) is 2 always, all we need to do is check that cvPPT(N ⊗ N) = 4 always.")
@testset "Verify that cv is 2 and cv(N ⊗ N) is multiplicative" begin
    scan_range =[0:0.01:0.5;]
    is_mult = true
    ctr = 1
    println("Now starting to scan over s.")
    for s_id in scan_range
        ctr % 10 == 0 ? println("Now on ", ctr, " of ", length(scan_range), " points.") : nothing
        #To save time I don't check the single copy cv
        # parallel channel uses
        sidchan(X) = siddhuChannel(X,s_id)
        sid_chan= Choi(sidchan,3,3)
        par_cv, = pptCV(parChoi(sid_chan, sid_chan), :dual)
        !isapprox(par_cv,4,atol=3e-6) ? is_mult = false : nothing
        ctr += 1
    end
    @test is_mult
end


@testset "mul" begin
    m = rand(10, 10)
    n = rand(10, 100)
    @test isapprox(Matrix(mul(GBMatrix(m), GBMatrix(n))), m * n, atol=8e-15)
    m = GBMatrix([1,3,5,7], [7,5,3,1], [1,2,3,4])
    n = GBMatrix{Int8}(7, 1)
    n[1:2:7, 1] = [1, 10, 20, 30]
    o = mul(m, n)
    @test size(o, 1) == 7
    @test eltype(o) == Int64
    @test o[7, 1] == 4 && o[5, 1] == 30
    o = GBMatrix(ones(Int64, 7, 1))
    mask = GBMatrix(ones(Bool, 7, 1))
    mask[3,1] = false
    @test mul!(o, m, n; mask, accum=BinaryOps.PLUS) ==
        GBMatrix([31,1,1,1,31,1,5])

    m = GBMatrix([[1,2,3] [4,5,6]])
    n = GBVector([10,20,30])
    @test_throws DimensionMismatch m * n
    @test m' * n == Matrix(m)' * Vector(n)
    @test n' * m == Vector(n)' * Matrix(m)
    @test n' * n == GBVector([Vector(n)' * Vector(n)])
end


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

using Random
Random.seed!(123)

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

cache1 = ForwardDiff.Dual{SparseDiffTools.DeivVecTag}.(x, v)
cache2 = ForwardDiff.Dual{SparseDiffTools.DeivVecTag}.(x, v)
@test num_jacvec!(du, f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
@test num_jacvec!(du, 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!(du, f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v
@test auto_jacvec!(du, 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!(du, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test num_hesvec!(du, 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!(du, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test numauto_hesvec!(du, 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!(du, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test autonum_hesvec!(du, 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!(du, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test numback_hesvec!(du, 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{Nothing}.(x, v)
cache4 = ForwardDiff.Dual{Nothing}.(x, v)
@test autoback_hesvec!(du, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
@test autoback_hesvec!(du, 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!(du, h, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test num_hesvecgrad!(du, 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!(du, h, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-2
@test auto_hesvecgrad!(du, 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!(du,L,v) ≈ auto_jacvec(f, x, v)
L.u .= v
@test mul!(du,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.u == x
@test mul!(du,L,v) ≈ num_jacvec(f, x, v) rtol=1e-6
L.u .= v
@test mul!(du,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!(du,L,v) ≈ num_hesvec(g, x, v) rtol=1e-2
L.u .= v
@test mul!(du,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!(du,L,v) ≈ numauto_hesvec(g, x, v) rtol=1e-8
L.u .= v
@test mul!(du,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!(du,L,v) ≈ num_hesvec(g, x, v) rtol=1e-2
L.u .= v
@test mul!(du,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!(du,L,v) ≈ numauto_hesvec(g, x, v) rtol=1e-8
L.u .= v
@test mul!(du,L,v) ≈ numauto_hesvec(g, v, v) rtol=1e-8

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


abstract type PairwisePotential end

struct PseudoHS{T<:AbstractFloat} <: PairwisePotential
    λ::T
    σ::T
    a::T
    b::T
    temp::T
end

function _pseudoHS(λ::T, σ::T, temp::T) where {T<:AbstractFloat}
    b = λ / (λ - 1.0)
    b = convert(T, b)
    a = λ * b^(λ - 1.0)
    a = convert(T, a)
    return PseudoHS{T}(λ, σ, a, b, temp)
end

PseudoHS(λ::T, σ::T, temp::T) where {T<:AbstractFloat} = _pseudoHS(λ, σ, temp)
function PseudoHS(λ::T, temp::T) where {T<:AbstractFloat}
    return _pseudoHS(λ, convert(T, 1.0), temp)
end

@inline function _potential_func(
    x,
    λ::T,
    a::T,
    temp::T,
) where {T<:AbstractFloat}
    energy = (a / temp) * (x^-λ - x^-(λ - 1.0))
    return energy + (1.0 / temp)
end

@inline function _force_func(x, λ::T, a::T, temp::T) where {T<:AbstractFloat}
    force = (λ * x^-(λ + 1.0)) - ((λ - 1.0) * x^-λ)
    return force * (a / temp)
end  # function _force_func

function _apply!(x, λ, a, b, temp)
    if x < b
        energy = _potential_func(x, λ, a, temp)
        energy = oftype(b, energy)
        force = _force_func(x, λ, a, temp)
        force = oftype(b, force)
    else
        energy = zero(b)
        force = zero(b)
    end
    return energy, force
end

function apply!(p::PseudoHS, x::AbstractFloat)
    _apply!(x, p.λ, p.a, p.b, p.temp)
end
# TODO: Work on a cleaner interface for custom potentials


module GaussianSlab
function get_fn(argd::Dict)
    center = [float(x) for x in argd["center"]]
    sigma  = float(argd["sigma"])
    k = float(argd["k"])
    upper = float(argd["upper"])
    lower = float(argd["lower"])
    raw_direction = [float(x) for x in argd["direction"]]
    direction = raw_direction/norm(raw_direction)
    function fn_compute(coordinates)
        d = dot(coordinates - center, direction)
        if lower <= d <= upper
            return k * exp(-(d)^2 /(2 *sigma))
        else
            return 0.0
        end
    end

    return fn_compute
end

export get_fn
end


"""
    statscholesky(xtx::Symmetric{T}, tol::Real=-1) where {T<:AbstractFloat}

Return a `CholeskyPivoted` object created from `xtx` where the pivoting scheme
retains the original order unless singularity is detected.  Columns that are
(computationally) linearly dependent on columns to their left are moved to the
right hand side in a left circular shift.
"""
function statscholesky(xtx::Symmetric{T}, tol::Real = -1) where {T<:AbstractFloat}
    n = size(xtx, 2)
    chpiv = cholesky(xtx, Val(true), tol = T(-1), check = false)
    chunp = cholesky(xtx, check = false)
    r = chpiv.rank
    piv = [1:n;]
    if r < n
        nleft = n
        while r < nleft
            k = chunp.info
            if k < nleft
                piv = piv[[1:k-1; k+1:n; k]]
                chunp = cholesky!(Symmetric(xtx[piv, piv]), check = false)
            end
            nleft -= 1
        end
    end
    for j = (r+1):n   # an MKL <-> OpenBLAS difference
        for i = (r+1):j
            chunp.factors[i, j] = zero(T)
        end
    end
    CholeskyPivoted(chunp.factors, chunp.uplo, piv, r, tol, chpiv.info)
end