export SolitaryWaveWhitham

"""
    `SolitaryWaveWhitham(param; kwargs...)`

Computes the Whitham solitary wave with prescribed velocity.

# Argument
- `param :: NamedTuple`: parameters of the problem containing velocity `c` and dimensionless parameters `ϵ` and `μ`, and mesh size `L` and number of collocation points `N`;

## Keywords (optional)
- `guess :: Vector{Real}`: initial guess for the surface deformation (if not provided, the exact formula for KdV is used);
- `x₀ :: Real`: center of solitary wave (if guess is not provided);
- `iterative :: Bool`: inverts Jacobian through GMRES if `true`, LU decomposition if `false`;
- `verbose :: Bool`: prints numerical errors at each step if `true`;
- `max_iter :: Int`: maximum number of iterations of the Newton algorithm;
- `tol :: Real`: general tolerance (default is `1e-10`);
- `ktol :: Real`: tolerance of the Krasny filter (default is `0`, i.e. no filtering);
- `gtol :: Real`: relative tolerance of the GMRES algorithm;
- `dealias :: Int`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing);
- `q :: Real`: Newton algorithm modified with
`u_{n+1}=q*u_{n+1}+(1-q)*u_n`
(default is `1`);
- `α :: Real`: adds `α` times spectral projection onto the Kernel to the Jacobian;
- `KdV :: Bool`: if `true` computes the KdV (instead of Whitham) solitary wave.
# Return values
`u :: Vector{Float64}` the solution
"""
function SolitaryWaveWhitham(
                param :: NamedTuple;
                guess = zeros(0) :: Vector{Float64},
                x₀ = 0 :: Real,
                iterative = false :: Bool,
                verbose = true :: Bool,
                max_iter = 20 :: Int,
                tol = 1e-14 :: Real,
                gtol = 1e-14 :: Real,
                ktol = 0 :: Real,
                dealias = 0 :: Int,
                q=1 :: Real,
                α=0 :: Real,
                KdV = false :: Bool)

        mesh=Mesh(param)
        c = param.c
        ϵ = param.ϵ
        μ = param.μ

        if guess == zeros(0)
                @info "Using the exact formula for the KdV solitary wave as initial guess"
                guess = 2*(c-1)/ϵ*sech.(sqrt(3*2*(c-1)/μ)/2*mesh.x.-x₀).^2
        end

        k  = mesh.k
        Dx =  1im * k
        if KdV == true
                @info string("Solving the KdV solitary wave with velocity c=",param.c)
                F₁ = 1 .-μ/6*k.^2
        else
                @info string("Solving the Whitham solitary wave with velocity c=",param.c)
                F₁ = sqrt.(tanh.(sqrt(μ)*abs.(k))./(sqrt(μ)*abs.(k)))
                F₁[1]=1
        end

        if dealias == 0
                Π = ones(size(k))
        else
                Π = abs.(k) .< maximum(k) * (1-dealias/(dealias+2))
        end
        krasny(k) = (abs.(k).> ktol ).*k
        krasny!(k) = k[abs.(k).< ktol ].=0

        function filter( v :: Vector{Float64} )
                real.(ifft(krasny(Π.*fft(v))))
        end
        function filter( v :: Vector{Complex{Float64}} )
                ifft(krasny(Π.*fft(v)))
        end

        function F( u :: Vector{Float64} )
            -c*u+real.(ifft(F₁.*fft(u)))+3*ϵ/4*u.^2
        end
        function Fabs( u :: Vector{Float64} )
            c*abs.(u)+abs.(ifft(F₁.*fft(u)))+3*ϵ/4*u.^2
        end

        if iterative == false
                k = mesh.k
                x = mesh.x
                x₀ = mesh.x[1]
                FFT = exp.(-1im*k*(x.-x₀)');
                IFFT = exp.(1im*k*(x.-x₀)')/length(x);
                M = real.(IFFT*(Diagonal( F₁)*FFT))
                function JacF( v , dxv  )
                    Symmetric(M+Diagonal(3*ϵ/2*v .-c)+ α*dxv*dxv')
                end
        else
                function JacFfast( v ,dxv )
                        dF(φ) = -c*φ+real.(ifft(F₁.*fft(φ)))+3*ϵ/2*v.*φ+ α*dot(dxv,φ)*dxv
                        return LinearMap(dF, length(v); issymmetric=true, ismutating=false)
                end
        end

        flag=0
        iter = 0
        u = filter(guess)
        du = similar(u)
        fu = similar(u)
        dxu = similar(u)
        for i in range(0, stop=max_iter)
                dxu .= real.(ifft(Dx.*fft(u)))
                dxu ./= norm(dxu,2)
                fu .= F(u)
                relerr = norm(fu,Inf)/norm(Fabs(u),Inf)
                abserr = norm(fu,Inf)
                if relerr < tol
    			@info string("Converged : relative error ",relerr," in ",i," steps\n")
    			break
    		elseif verbose == true
                        print(string("absolute error at step ",i,": ",abserr,"\n"))
                        print(string("relative error at step ",i,": ",relerr,"\n"))

    		end
                if i == max_iter
                        @warn string("The algorithm did not converge after ",i," steps: final relative error is ",relerr,"\n")
                        break
                end

                if iterative == false
                        du .= JacF(u,dxu) \ fu
                else
                        du .= gmres( JacFfast(u,dxu) , fu ; verbose = verbose, tol = gtol )
                end
    		u .-= q*filter(du)
        end
        return u
end


# by @walra356
using CairoMakie
CairoMakie.activate!() #HIDE
let
    supertitle = "heatmap of camera chip exposure with x and y sections"
  # define transformation from pixel coordinates to physical coordinates
    edges(p,Δx,x0) = collect(p .* Δx) .-(x0+0.5Δx)
  # parameters camera chip (pixel coordinates)
    dimX = 40; dimY = 30                       # pixel dimensions
    (px, py) = (1:dimX, 1:dimY)                  # pixel positions
    (ix, iy) = (10, 3)                        # selected pixel
  # physical coordinates
    (Δx, Δy) = (4.0, 5,0)                       # pixel lattice periods (μm)
    (Ox, Oy) = (edges(ix,Δx,0), edges(iy,Δy,0)) # manual offset (μm)
    (x, y) = (edges(px,Δx,Ox), edges(py,Δy,Oy)) # pixel positions (μm)
  # create exposure matrix σ (model alternative for measurement file)
    valmax = 500.0                              # maximum value exposure matrix
    exposure(i, j , ix, iy, wx, wy) = 0.9valmax *
                        (exp(-(((i-ix)/wx)^2 + ((Δy/Δx)*(j-iy)/wy)^2)) + 0.1)
    σ = round.(Int, [exposure(i, j, ix, iy, 6, 6) for i=1:dimX, j=1:dimY])
    σ1 = σ[px,py]  # raw matrix
    sh = σ[:,iy]   # horizontal section (pixel row iy from bottom)
    sv = σ[ix,:]   # vertical section (pixel column ix from left)
    footnote = "Origin of coordinate system was shifted by ($Ox μm, $Oy μm)
                w.r.t. left-bottom edge of camera chip
                \nPixel lattice periods: (Δx = $Δx μm, Δy = $Δy μm);
                pixel position of light spot: (ix, y0) = ($(ix), $(iy))"

  # start actual plot
    theme = Theme(fontsize = 12, colormap = :gist_earth, resolution = (750,850))
    set_theme!(theme)
    fig = Figure()
  # collect attributes
    fsize = fig.scene.attributes.fontsize.val
    fresx = fig.scene.attributes.resolution.val[1]
    fresy = fig.scene.attributes.resolution.val[2]
    attr = (xlabelsize = 6fsize/5, ylabelsize = 6fsize/5, titlesize = 7fsize/5,
            xautolimitmargin = (0,0), yautolimitmargin = (0,0), )
    attr1 = (attr..., title = "physical scale",
        xticks = [x[1]-Δx/2, x[ix], 50, x[end]+Δx/2],
        yticks = [y[1]-Δy/2, y[iy], 50, y[end]+Δy/2],
        xlabel = "x (μm)", ylabel = "y (μm)",
        aspect = (Δx * size(σ1,1))/(Δy * size(σ1,2)), )
    attrc = (label = "exposure (counts)", ticksize=6fsize/5, tickalign = 1,
        width = 15, height = Relative(0.88), ) # tweaked
    attrh = (attr..., title = "horizontal section",
        xlabel = "x (μm)", ylabel = "exposure (counts)",
        xticks = [x[1]-Δx/2, x[ix], 50, x[end]+Δx/2],
        yticks = [0, 250, valmax],
        limits = ((x[1]-Δx/2, x[end]+Δx/2),(0, valmax)),
        height = fresy/5, aspect = 2.05, ) # tweaked
    attrv = (attr..., xlabel = "exposure (counts)",
        xticks = [0, 250, valmax],
        yticks = [y[1]-Δy/2, y[iy], 50, y[end]+Δy/2],
        xreversed = true, yaxisposition = :right,
        limits = ((0,valmax),(y[1]-Δy/2, y[end]+Δy/2)),
        width = fresy/5, aspect = 0.58, ) # tweaked
  # create axes and add contents
    ax1 = Axis(fig; attr1...)
    axh = Axis(fig; attrh...)
    axv = Axis(fig; attrv...)
    hm1 = heatmap!(ax1, x, y, σ1)
    axc = Colorbar(fig, hm1; attrc...)
    lfn = Label(fig, text = footnote, textsize = 6fsize/5)
    lst = Label(fig, text = supertitle, textsize = 2fsize)
    lvs = Label(fig, text = "vertical section", padding = (0, 5, 0, 0),
        rotation = pi/2, textsize = 7fsize/5)
    sch = scatter!(axh, x, sh, direction = :x, color = :red, markersize = 5, )
    scv = scatter!(axv, sv, y, color = :red, markersize = 5 )
    lnh = lines!(ax1, x , fill(0, length(x)), color = :white, linewidth = 1)
    lnv = lines!(ax1, fill(0, length(y)), y, color = :white, linewidth = 1)
  # create layout and show figure
    fig[1,2] = axh
    fig[2,1] = axv
    fig[2,1, Left()] = lvs
    fig[2,2] = ax1
    fig[2,3] = axc
    fig[3,:] = lfn
    fig[0,:] = lst
    fig
    save(joinpath(@__DIR__, "output", "heatmapScaleSections.png"), fig, px_per_unit = 2.0) # HIDE
end
using Pkg # HIDE
Pkg.status(["CairoMakie"]) # HIDE


module StableDQMC

using LinearAlgebra, Pkg
using Requires

include("helpers.jl")
include("helpers_sdw.jl")
include("qr_udt.jl")
include("svd.jl")

# Slice matrix chain B_M .... B_1
export calc_Bchain, calc_Bchain_qr, calc_Bchain_svd
export calc_tdgf_qr, calc_tdgf_svd
export calc_Bchain_svd_faster

# QR / UDT
export UDT, inv!, fact_mult, Matrix!
export udt!, udt

export udt_inv_one_plus
export inv_one_plus!
export inv_one_plus

export udt_inv_sum
export inv_sum!
export inv_sum

export udt_inv_one_plus_loh
export inv_one_plus_loh!
export inv_one_plus_loh

export udt_inv_sum_loh
export inv_sum_loh!
export inv_sum_loh

# SVD / UDV
export gesdd!
export gesdd
export gesvd!
export gesvd

export svd_inv_one_plus
export inv_one_plus!
export inv_one_plus

export svd_inv_one_plus_loh
export inv_one_plus_loh!
export inv_one_plus_loh

export svd_inv_sum
export inv_sum!
export inv_sum

export svd_inv_sum_loh
export inv_sum_loh!
export inv_sum_loh

# Optional JacobiSVD + GenericSVD
function addJacobiSVD()
    pkg"add https://github.com/RalphAS/JacobiSVD.jl"
    @eval import JacobiSVD
end
rmJacobiSVD() = pkg"rm JacobiSVD"

function addGenericSVD()
    pkg"add GenericSVD"
    @eval import GenericSVD
end
rmGenericSVD() = pkg"rm GenericSVD"

function __init__()
    @require JacobiSVD="2ca068c6-2156-5cf0-8317-c67edf277a2c" include("svd_jacobi.jl")
    @require GenericSVD="01680d73-4ee2-5a08-a1aa-533608c188bb" include("svd_generic.jl")
end

end # module


# SPDX-License-Identifier: X11
# 2020-11-27
# Sum of gcd of Tuples (Easy) (300pt)

function main()
  K = parse(Int, readline())

  sum = 0
  for i ∈ 1:K
    for j ∈ 1:K
      for k ∈ 1:K
        sum += gcd(i,j,k)
      end
    end
  end

  println(sum)
end

main()


using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using Measures
using ProgressMeter
using Plots
theme(:dark)
using Revise
push!(LOAD_PATH, pwd()*"/main")
using eqws

filter_hash = fink_filter_hash(1,8)
a = rand(256,256)
d = eqws_compute_convmap(a,filter_hash)

heatmap(d[57])

d2 = sum.(d)

d1 = eqws_compute(a,filter_hash)

count(d1 .≈ d2)

maximum(abs.(d1.-d2)) < 1e-15

rod_test = rod_image(1,1,30,35,8)

d1 = eqws_compute(rod_test,filter_hash)
d = eqws_compute_convmap(rod_test,filter_hash)
d2 = sum.(d)
count(d1 .≈ d2)
test2 = (d1 .≈ d2)
argmin(test2)

d1[57]
d2[57]

d1[57]
d2[57]

maximum(abs.(d1.-d2)) < 1e-16

scatter(d1 .≈ d2)

b = [im_rd_0_1[:,:,i]  for i in 1:size(im_rd_0_1)[3]]

b[1]

filter_hash["S1_iso_mat"]*b

filter_hash["S1_iso_mat"]

filter_hash["S1_iso_mat"]*b

function test_compute(i)
    a = rand(256,256)
    return eqws_compute_convmap(a,filter_hash)
end

using Distributed

test = pmap(test_compute,1:10)

DHC_out = hcat(test...)

DHC_out[3,:]

49*69000


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("../useful_routines.jl")

p = [1.0 2.0 3.0]

bp = CutCellDG.extend_to_face(p, 1)
testbp = vcat(p, -ones(3)')
@test allapprox(bp, testbp)

rp = CutCellDG.extend_to_face(p, 2)
testrp = vcat(ones(3)', p)
@test allapprox(rp, testrp)

tp = CutCellDG.extend_to_face(p, 3)
testtp = vcat(p, ones(3)')
@test allapprox(tp, testtp)

lp = CutCellDG.extend_to_face(p, 4)
testlp = vcat(-ones(3)', p)
@test allapprox(testlp, lp)

quad1d = tensor_product_quadrature(1, 4)
points = quad1d.points
bp = vcat(points, -ones(4)')
rp = vcat(ones(4)', points)
tp = vcat(points, ones(4)')
lp = vcat(-ones(4)', points)
bq = CutCellDG.extend_to_face(quad1d, 1)
rq = CutCellDG.extend_to_face(quad1d, 2)
tq = CutCellDG.extend_to_face(quad1d, 3)
lq = CutCellDG.extend_to_face(quad1d, 4)
@test allapprox(bq.points, bp)
@test allapprox(rq.points, rp)
@test allapprox(tq.points, tp)
@test allapprox(lq.points, lp)
@test allapprox(bq.weights, quad1d.weights)
@test allapprox(rq.weights, quad1d.weights)
@test allapprox(tq.weights, quad1d.weights)
@test allapprox(lq.weights, quad1d.weights)

facequads = CutCellDG.face_quadratures(4)
@test allapprox(facequads[1].points, bp)
@test allapprox(facequads[2].points, rp)
@test allapprox(facequads[3].points, tp)
@test allapprox(facequads[4].points, lp)
@test allapprox(facequads[1].weights, quad1d.weights)
@test allapprox(facequads[2].weights, quad1d.weights)
@test allapprox(facequads[3].weights, quad1d.weights)
@test allapprox(facequads[4].weights, quad1d.weights)

polyorder = 1
numqp = 2
levelsetbasis = HermiteTensorProductBasis(2)

elasticitybasis = LagrangeTensorProductBasis(2, polyorder)
dim, levelsetnf = size(interpolation_points(levelsetbasis))
points = interpolation_points(elasticitybasis)
levelset = InterpolatingPolynomial(1, levelsetbasis)
x0 = [0.0, 0.0]
meshwidths = [2.0, 1.0]
nelements = [2, 1]
dgmesh = CutCellDG.DGMesh(x0, meshwidths, nelements, points)
cgmesh = CutCellDG.CGMesh(x0, meshwidths, nelements, levelsetnf)

normal = [1.0, 0.0]
xI = [0.5, 0.0]

levelset = CutCellDG.LevelSet(
    x -> plane_distance_function(x, normal, xI),
    cgmesh,
    levelsetbasis,
)

cutmesh = CutCellDG.CutMesh(dgmesh, levelset)
facequads = CutCellDG.FaceQuadratures(cutmesh, levelset, numqp)
idx = [5, 9, 6, 10, 7, 11, 8, 12, 1, 0, 2, 0, 3, 0, 4, 0]
testfacetoquad = reshape(idx, 2, 4, 2)
@test allequal(testfacetoquad, facequads.facetoquad)


# This is an example which uses much of the functionality provided by
# Interplanetary.jl. It is based on a homework assignment, but that is
# no guarentee that it's correct.
#
# The objective is to explore optimal trajectories from Earth to
# asteroid 2015PDC - a fictional asteroid created for a hypothetical
# asteroid impact scenario held at the 2015 Planetary Defense
# Conference (see http://neo.jpl.nasa.gov/pdc15/). This is
# accomplished via a Pork Chop Plot - a contour plot showing the value
# of some objective function (mass at arrival, delta V, etc.) over a
# grid of Earth departure dates and time of flight values.
#
# This particular example searches for the minimum-delta V solution. I
# also included some of my experiments with using Julia's built-in
# parallel computing capabilities, because it seemed like a good place
# for parallelization.

using DelimitedFiles, Dates, LinearAlgebra

function meshgrid(x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat}
    return (repeat(x, 1, length(y)),
             repeat(y',length(x),1))
end

"""
This function creates a Pork Chop Plot given a range of departure
dates, time of flight for a given grid point (trajectory), and the
figure of merit value for a given trajectory.
"""
function createPCC( departure_date_range::Vector{T},
                   TOF_range::Vector{T},
                   FOM_array::Matrix{T}; levels = [],
                   FOM_label::String = "") where {T<:AbstractFloat}
    (depart_grid, TOF_grid) = meshgrid(departure_date_range, TOF_range)

    if isempty(levels)
        cplot = contourf(depart_grid, TOF_grid, FOM_array)
    else
        cplot = contourf(depart_grid, TOF_grid, FOM_array, levels)
    end

    # formatting
    ylocs,ylabels = yticks()
    xlocs,xlabels = xticks()
    cb = colorbar()
    cb.set_label(FOM_label)

    xticks(xlocs,Date.(Dates.julian2datetime.(mjd2jd.(xlocs))),rotation=30)
    #yticks(ylocs,Date(julian2datetime(mjd2jd(ylocs))))
    xlabel("Departure date")
    ylabel("TOF [days]")
    # need to set axis limits
end

# read ephemeris files for Earth and Asteroid 2015PDC
Earth_ephem = readdlm("data/Earth_2015-04-01_to_2022-09-03_step_1d.orb")
ast_2015_PDC_ephem = readdlm("data/ast_2015_PDC_2015-04-01_to_2022-09-03_step_1d.orb")

#@everywhere begin # need to add cores before running this: addprocs(n)
using Interplanetary, PyPlot
departure_lb_mjd = jd2mjd( Dates.datetime2julian( DateTime("2015-10-01") ) )
departure_ub_mjd = jd2mjd( Dates.datetime2julian( DateTime("2022-02-26") ) )

# set up ranges for grid search
departure_step = 5. #days
TOF_step = 5. #days

departure_date_range = collect(departure_lb_mjd:departure_step:departure_ub_mjd)
TOF_range = collect(90.:TOF_step:1500.)

# now loop over all combinations

dV_total = Array{Float64}(undef,length(departure_date_range),length(TOF_range)) # for non-parallel
#dV_total = SharedArray{Float64}(length(departure_date_range),length(TOF_range))

rPark = RE + 185. #parking orbit radius
#end

#@time @sync @parallel for i = 1:length(departure_date_range)
for i = 1:length(departure_date_range)
    for j = 1:length(TOF_range)
        dep_date = departure_date_range[i]
        arr_date = dep_date + TOF_range[j]
        if arr_date <= ast_2015_PDC_ephem[end,1]
            (dv1,dv2,c3,rla,dla) = lambert_transfer(muEarth,
                                        TOF_range[j]*86400.,
                                        mjd2jd(departure_date_range[i]),
                                        rPark,
                                        vec(Earth_ephem[Earth_ephem[:,1].==dep_date,2:7]),
                                        vec(ast_2015_PDC_ephem[ast_2015_PDC_ephem[:,1].==arr_date,2:7]))
            dV_total[i,j] = dv1+dv2
        end#if
    end#for
end#for

# create the contour plot
# sdata() makes a SharedArray into a regular array
createPCC(departure_date_range,
          TOF_range,sdata(dV_total),
          levels = collect(5.:2:30),
          FOM_label = "Delta V [km/s]")

# find the minimum dV solution
# linear index of minimum, excluding zeros:
ind = findall(dV_total .== minimum(dV_total[dV_total.>0]))[1]

subs = CartesianIndices(dV_total)[ind] # convert to array subscripts
opt_arrival_date = departure_date_range[subs[1][1]] + TOF_range[subs[2][1]]

# plot the optimal solution on the pork chop plot
plot(departure_date_range[subs[1][1]],TOF_range[subs[2][1]],"r*",markersize=10)

# rerun the Lambert targeter for the optimal solution
(dv1,dv2,c3,rla,dla) = lambert_transfer(muEarth,
                            TOF_range[subs[2][1]]*86400.,
                            mjd2jd(departure_date_range[subs[1][1]]),
                            rPark,
                            vec(Earth_ephem[Earth_ephem[:,1].==departure_date_range[subs[1][1]],2:7]),
                            vec(ast_2015_PDC_ephem[ast_2015_PDC_ephem[:,1].==opt_arrival_date,2:7]))

# print results to screen
println("Optimal solution:")
println("Departure date: ", Date.(Dates.julian2datetime.(mjd2jd(departure_date_range[subs[1][1]]))))
println("TOF (days): ",TOF_range[subs[2][1]])
println("Earth departure ΔV [km/s]: ", dv1)
println("Earth departure C3 [km^2/s^2]: ", c3)
println("DLA (deg): ", rad2deg(dla))
println("Arrival ΔV [km/s]: ", dv2)
println("Total solution ΔV [km/s]: ", dV_total[subs][1])


using Test,
    QSimulator,
    Distributed

function test_parallel_dissipators_1()
    sys = CompositeQSystem();

    qA = Qubit("qA",8.0);
    qB = Qubit("qB",8.0);

    g=.1

    sys += qA;
    sys += qB;
    sys += FlipFlop(qA,qB,g);

    q0 = [1 0; 0 0];
    q1 = [0 0; 0 1];

    obs = kron(q1,q0);

    Rabi_rate = 2g;
    max_time = 3/Rabi_rate;
    times = range(0.0; stop = max_time, length = 1000);

    result =
        pmap(times) do time
          local u
          local out
            out = zeros(3)
            u = unitary_propagator(sys, .01, .0, time);
            out[1] = real(trace(obs*u*kron(q0,q0)*u'));
            out[2] = real(trace(obs*u*kron(q0,q1)*u'));
            out[3] = real(trace(obs*u*kron(q1,q0)*u'));
            out
        end;

    pred1 = zeros(length(times))
    pred2 = map(t->1/2*cos(2pi*2g*t+pi)+1/2,times);
    pred3 = map(t->1/2*cos(2pi*2g*t)+1/2,times);

    r1 = Float64[r[1] for r in result];
    r2 = Float64[r[2] for r in result];
    r3 = Float64[r[3] for r in result];

    @test isapprox(maximum(abs(r1-pred1)), 0.0, atol = 1e-2)
    @test isapprox(maximum(abs(r2-pred2)), 0.0, atol = 1e-2)
    @test isapprox(maximum(abs(r3-pred3)), 0.0, atol = 1e-2)
end

function test_parallel_dissipators_2()
    sys3 = CompositeQSystem();

    qA = Qubit("qA",8.0);

    rrate = .1
    relaxA = Cooling("T1",.1,qA)

    sys3 += qA
    sys3 += relaxA

    max_time = 3/rrate
    times = range(0.0, stop = max_time, length = 1000);

    rhos = [exp(time*generator(sys3))*vec([0 0; 0 1]) for time in times];

    p = Float64[ (vec([0 0; 0 1])'*rho)[1,1] for rho in rhos ];
    pred = map(t->exp(-2pi*t*rate(relaxA)),times);

    @test isapprox(maximum(abs(p-pred)), 0.0, atol = 1e-2)
end

# @testset "parallel dissipators 1" begin
#     test_parallel_dissipators_1()
# end

# @testset "parallel dissipators 2" begin
#     test_parallel_dissipators_2()
# end


using Revise
using ApproxFun, LinearAlgebra, Parameters, Setfield

using PseudoArcLengthContinuation, Plots
const PALC = PseudoArcLengthContinuation

####################################################################################################
# specific methods for ApproxFun
import Base: eltype, similar, copyto!, length
import LinearAlgebra: mul!, rmul!, axpy!, axpby!, dot, norm

similar(x::ApproxFun.Fun, T) = (copy(x))
similar(x::ApproxFun.Fun) = copy(x)
mul!(w::ApproxFun.Fun, v::ApproxFun.Fun, α) = (w .= α * v)

eltype(x::ApproxFun.Fun) = eltype(x.coefficients)
length(x::ApproxFun.Fun) = length(x.coefficients)

dot(x::ApproxFun.Fun, y::ApproxFun.Fun) = sum(x * y)

# do not put y .= a .* x .+ y, this puts a lot of coefficients!
axpy!(a::Float64, x::ApproxFun.Fun, y::ApproxFun.Fun) = (y .= a * x + y)
axpby!(a::Float64, x::ApproxFun.Fun, b::Float64, y::ApproxFun.Fun) = (y .= a * x + b * y)
rmul!(y::ApproxFun.Fun, b::Float64) = (y.coefficients .*= b; y)
rmul!(y::ApproxFun.Fun, b::Bool) = b == true ? y : (y.coefficients .*= 0; y)

# copyto!(x::ApproxFun.Fun, y::ApproxFun.Fun) = ( copyto!(x.coefficients, y.coefficients);x)
copyto!(x::ApproxFun.Fun, y::ApproxFun.Fun) = ( (x.coefficients = copy(y.coefficients);x))

####################################################################################################

N(x; a = 0.5, b = 0.01) = 1 + (x + a * x^2) / (1 + b * x^2)
dN(x; a = 0.5, b = 0.01) = (1 - b * x^2 + 2 * a * x)/(1 + b * x^2)^2

function F_chan(u, alpha, beta = 0.01)
	return [Fun(u(0.), domain(u)) - beta,
			Fun(u(1.), domain(u)) - beta,
			Δ * u + alpha * N(u, b = beta)]
end

function dF_chan(u, v, alpha, beta = 0.01)
	return [Fun(v(0.), domain(u)),
			Fun(v(1.), domain(u)),
			Δ * v + alpha * dN(u, b = beta) * v]
end

function Jac_chan(u, alpha, beta = 0.01)
	return [Evaluation(u.space, 0.),
			Evaluation(u.space, 1.),
			Δ + alpha * dN(u, b = beta)]
end

function finalise_solution(z, tau, step, contResult)
	printstyled(color=:red,"--> AF length = ", (z, tau) .|> length ,"\n")
	chop!(z.u, 1e-14);chop!(tau.u, 1e-14)
	true
end

sol = Fun( x -> x * (1-x), Interval(0.0, 1.0))
const Δ = Derivative(sol.space, 2);

optnew = NewtonPar(tol = 1e-12, verbose = true)
	out, _, flag = @time PALC.newton(
		u -> F_chan(u, 3.0),
		u -> Jac_chan(u, 3.0),
		sol, optnew, normN = x -> norm(x, Inf64))
	# Plots.plot(out, label="Solution")

optcont = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 4.1, plotEveryNsteps = 10, newtonOptions = NewtonPar(tol = 1e-8, maxIter = 20, verbose = true), maxSteps = 300)

	br, _ = @time continuation(
		(x, p) ->   F_chan(x, p),
		(x, p) -> Jac_chan(x, p),
		out, 3.0, optcont,
		plot = true,
		plotSolution = (x; kwargs...) -> plot!(x; label = "l = $(length(x))", kwargs...),
		verbosity = 2,
		# printsolution = x -> norm(x, Inf64),
		normC = x -> norm(x, Inf64))
####################################################################################################
# other dot product
# dot(x::ApproxFun.Fun, y::ApproxFun.Fun) = sum(x * y) * length(x) # gives 0.1

optcont = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 4.1, plotEveryNsteps = 10, newtonOptions = NewtonPar(tol = 1e-8, maxIter = 20, verbose = true), maxSteps = 300, theta = 0.2)

	br, _ = @time continuation(
		(x, p) ->   F_chan(x, p),
		(x, p) -> Jac_chan(x, p),
		out, 3.0, optcont;
		dotPALC = (x, y) -> dot(x, y),
		plot = true,
		# finaliseSolution = finalise_solution,
		plotSolution = (x; kwargs...) -> plot!(x; label = "l = $(length(x))", kwargs...),
		verbosity = 2,
		# printsolution = x -> norm(x, Inf64),
		normC = x -> norm(x, Inf64))
####################################################################################################
# tangent predictor with Bordered system
br, _ = @time continuation(
	(x, p) ->   F_chan(x, p),
	(x, p) -> Jac_chan(x, p),
	out, 3.0, optcont,
	tangentAlgo = BorderedPred(),
	plot = true,
	finaliseSolution = finalise_solution,
	plotSolution = (x;kwargs...)-> plot!(x; label = "l = $(length(x))", kwargs...))
####################################################################################################
# tangent predictor with Bordered system
# optcont = @set optcont.newtonOptions.verbose = true
indfold = 2
outfold, _, flag = @time newtonFold(
		(x, α) -> F_chan(x, α),
		(x, p) -> Jac_chan(x, p),
		br, indfold, #index of the fold point
		optcont.newtonOptions)
	flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold[end], ", β = 0.01, from ", br.bifpoint[indfold][3],"\n")
#################################################################################################### Continuation of the Fold Point using minimally augmented
indfold = 2

outfold, _, flag = @time newtonFold(
	(x, p) -> F_chan(x, p),
	(x, p) -> Jac_chan(x, p),
	(x, p) -> Jac_chan(x, p),
	br, indfold, #index of the fold point
	optcont.newtonOptions)

br.bifpoint[2].


module QueensAlwaysPuzzleMe

using Distributed, BenchmarkTools

# There is also a simple (i.e. not distributed) version:
# https://github.com/PeterLuschny/IntegerSequences.jl/blob/master/demos/SimpleQueens.jl
# For some background see https://wp.me/paipV7-E and https://oeis.org/A319284

function solve!(profile, level, size, start, cols, diag4, diag1)

    if level > 0
        for i in start:size
            save = cols & (1 << i) +
            diag1 & (1 << (i + level)) +
            diag4 & (1 << (32 + i - level))

            if save == 0

                cols  = xor(cols,  1 << i)
                diag1 = xor(diag1, 1 << (i + level))
                diag4 = xor(diag4, 1 << (32 + i - level))

                solve!(profile, level - 1, size, 0, cols, diag4, diag1)

                cols  = xor(cols,  1 << i)
                diag1 = xor(diag1, 1 << (i + level))
                diag4 = xor(diag4, 1 << (32 + i - level))

                profile[level + 1] += 1

                cols == 0 && return
            end
        end
    else
        for i in 0:size
            save = cols & (1<<i) + diag1 & (1<<i) + diag4 & (1<<(32+i))
            save == 0 && (profile[1] += 1)
        end
    end
end

function distributed_search(n::Int, k::Int)

    profile = zeros(Int, n + 1)
    cols = diag4 = diag1 = Int(0)
    solve!(profile, n - 1, n - 1, k, cols, diag4, diag1)
    return profile
end

function queens(n::Int)

    n == 0 && return [1]

    profile = Distributed.@distributed (+) for k = 0:n-1
        distributed_search(n, k)
    end

    profile[n+1] = 1  # add the root
    [profile[n-i+1] for i = 0:n]
end

function ProfileQueens(up_to)
    for n in 0:up_to
        profile = queens(n)
        #show_profile(n, profile)
    end
end

"""
julia> show_profile(7, queens(7))
    size:      7
    profile:   [1, 7, 30, 76, 140, 164, 94, 40]
    nodes:     552
    solutions: 40
"""
function show_profile(n, p)
    println("size:      ", n)
    println("profile:   ", p)
    println("nodes:     ", sum(p))
    println("solutions: ", p[n+1])
    println()
end

"""
julia> BechmarkQueens()
    Number of cpu-threads : 4
    Number of threads used: 4
    95.699 ms (589 allocations: 48.88 KiB)
"""
function BechmarkQueens()
    println("Number of cpu-threads : ", Base.Sys.CPU_THREADS)
    println("Number of threads used: ", Threads.nthreads())
    println()

    @btime ProfileQueens(12)
end

BechmarkQueens()

end # module


struct BathymetryModel{Td<:AbstractMvNormal, Tω<:AbstractFisheryDomain} <: Any
    D::Td
    Ω::Tω
end
## Might be worth making a type system for spatial kernels...
function BathymetryModel(Ω::GriddedFisheryDomain{Tf, Ti},
                         μ::AbstractVecOrMat{Tf},
                         kern::K) where {Tf<:Real, Ti, K<:AbstractCovarianceKernel}
    length(μ) == length(Ω) || throw(DimensionMismatch("Need mean for every grid cell."))
    D = MvNormal(vec(μ), cov(kern, Ω))
    BathymetryModel(D, Ω)
end

"""
    Bathymetry
        bathymetry::Matrix
        Ω::AbstractFisheryDomain

A realization of a `BathymetryModel`. Typically created using the
`rand(BM::BathyModel)` method.
"""
struct Bathymetry{T<:Real, Tω<:AbstractFisheryDomain} <: Any
    bathymetry::Matrix{T}
    Ω::Tω

    function Bathymetry(bathymetry::Matrix{T}, Ω::Tω) where {T<:Real, Tω<:AbstractFisheryDomain}
        all(size(Ω) .== size(bathymetry)) ||
            throw(DimensionMismatch("Bathymetry must conform to domain grid"))
        new{T, Tω}(bathymetry, Ω)
    end
end

function rand(rng::Random.AbstractRNG, BM::BathymetryModel)
    bathy = rand(BM.D)
    bathy .-= minimum(bathy)
    Bathymetry(reshape(bathy, size(BM.Ω)...), BM.Ω)
end
rand(BM::BathymetryModel) = rand(Random.GLOBAL_RNG, BM)

getindex(B::Bathymetry, i, j) = B.bathymetry[i, j]
getindex(B::Bathymetry, i) = B.bathymetry[i]


using Random
using SparseArrays

# From the adjacency matrix, build an adjacency list for the graph
function ConstructAdj(C::SparseMatrixCSC,n::Int64)
    rp = C.rowval
    ci = C.colptr
    Neighbs = Vector{Vector{Int64}}()
    d = zeros(Int64,n)
    for i = 1:n
        # chop up the rp vector and put it in Neighbs
        push!(Neighbs,rp[ci[i]:ci[i+1]-1])
        d[i] = ci[i+1]-ci[i]
    end

    # d is the number of neighbors. This is the unweighted degree,
    # but note importantly that if the original graph is weighted this is
    # not the same as the degree vector d we will sometimes use
    return Neighbs, d
end

function LamCCobj(A::SparseMatrixCSC,c,w,lam)

    w_volA = sum(w)
    # d = sum(A,dims = 1)
    # volA = sum(d)
    obj = (lam*(w_volA)^2-lam*sum(w.^2))/2

    for i = 1:maximum(c)
        S = findall(x->x==i,c)
        AS = A[:,S]
        vol = sum(AS.nzval);
        SAS = AS[S,:]
        edges = sum(SAS.nzval);
        cut = vol-edges
        w_volS = sum(w[S])
        obj += 1/2*(cut - lam*w_volS*(w_volA-w_volS))
    end
    return obj
end

# Take a clustering vector that may have gaps in the numbers, and re-name clusters
function renumber(c::Vector{Int64},Clusters::Vector{Vector{Int}})

    n = length(c)
    map = unique(c)
    cnew = zeros(Int64,n)

    Clusters = Clusters[map]

    # Rename the clusters
    for i = 1:n
        newClus = findfirst(x->x == c[i],map)
        cnew[i] = newClus
        push!(Clusters[newClus],i)
    end

    return cnew, Clusters

end

function renumber(c::Vector{Int64})

    n = length(c)
    map = sort(unique(c))
    cnew = zeros(Int64,n)

    # Rename the clusters
    for i = 1:n
        newClus = findfirst(x->x == c[i],map)
        cnew[i] = newClus
    end

    return cnew

end

function Many_Louvain(A::SparseMatrixCSC{Float64,Int64},w::Vector{Float64},lam::Float64,numtimes::Int64,maxits::Int64=10000)

    n = size(A,1)
    BestObj = Inf
    cBest = collect(1:n)

    for k = 1:numtimes
        Cs = LambdaLouvain(A,w,lam,true,maxits)
        c = Cs[:,end]
        obj = LamCCobj(A,c,w,lam)
        if obj < BestObj
            BestObj = obj
            cBest = c
        end
    end

    return cBest, BestObj
end

function LambdaLouvain(A::SparseMatrixCSC{Float64,Int64},w::Vector{Float64},lam::Float64,randflag::Bool=false,maxits::Int64=10000)

    @assert(issymmetric(A))
    n = size(A,1)
    # Run a first step
    c, improved = LambdaLouvain_Step(A,w,lam,randflag,maxits)

    # Keep track of every improvement you find
    if improved
        Cs = c
        c_old = copy(c)
    else
        Cs = c
    end

    while improved
        # Collapse the clustering and run again
        Anew, wnew = collapse_clustering(A,w,c_old)
        cSuper, improved = LambdaLouvain_Step(Anew,wnew,lam,randflag,maxits)
        N = length(wnew)
        if improved
            # Extract what that new clustering is
            c_new = zeros(Int64,n)
            # For each supernode, place all the nodes that make it up into a cluster
            # of the supernodes label
            for i = 1:N

                # What cluster is supernode i in?
                SuperI_Cluster = cSuper[i]

                # What individual nodes are in supernode i?
                SuperI_nodes = findall(x->x==i,c_old)
                c_new[SuperI_nodes] .= SuperI_Cluster
            end
            Cs = [Cs c_new]
            c_old = copy(c_new)
        end
    end

    return Cs
end

function LambdaLouvain_Step(A::SparseMatrixCSC{Float64,Int64},w::Vector{Float64},lam::Float64,randflag::Bool=false,maxits::Int64=Inf)
    @assert(issymmetric(A))
    n = size(A,1)
    # println("Merging $n Communities")

    if randflag
        p = randperm(n)
        A = A[p,p]
        undop = sortperm(p)
        w = w[p]
    end

    c = collect(1:n)
    @assert(size(w,1) == n)
    improving = true
    nextclut = n+1
    its = 0
    Neighbs, degvec = ConstructAdj(A,n)

    # The find function is inefficient for getting nodes in a cluster.
    # We will store a vector of n clusters (by the end, many will be empty clusters)
    Clusters = Vector{Vector{Int64}}()
    for v = 1:n
        push!(Clusters, Vector{Int64}())
        push!(Clusters[v],v)
    end

    # As long as we can keep improving, continue
    while improving && its < maxits
        # println("Iter $its")
        its += 1
        # println("Iteration $its")
        improving = false
        nextclus = 1

        for i = 1:n

            # println("\n\t NODE $i")

            # Cluster index
            Ci_ind = c[i]

            # Get the indices of nodes in i's cluster
            # Ci = findall(x->x == Ci_ind,c)
            Ci = Clusters[Ci_ind]

            # Get the indices of i's neighbors
            # Ni = findall(x->x>0, A[i,:])
            Ni = Neighbs[i]

            # Find the indices of nodes that neighbor i and are in c[i]
            # ci_neighbs = intersect(Ni,Ci)

            # Weight of negative mistakes currently at i
            neg_inner = w[i]*(sum(w[Ci]) - w[i])

            # Weight of positive mistakes if we remove i from Ci
            # pos_inner = sum(A[i,ci_neighbs])
            pos_inner = sum(A[i,Ci])

            # Increase in mistakes if we remove i from Ci
            total_inner = pos_inner - lam*neg_inner

            BestImprove = 0
            BestC = Ci_ind

            # Get the neighboring clusters of i
            NC = unique(c[Ni])
            #NC = Ni

            # @show NC
            # See if it's better to move to a nearby cluster, Cj
            for j = 1:length(NC)
                Cj_ind = NC[j]

                # Check how much it would improve to move i to to cluster j
                if Cj_ind == Ci_ind

                    change = 0

                else

                    # Cj = findall(x->x == Cj_ind,c)
                    Cj = Clusters[Cj_ind]

                    # Find the neighbors of i in Cj
                    # cj_neighbs = intersect(Ni,Cj)

                    # Moving i from Ci to Cj adds negative mistakes
                    neg_outer = w[i]*(sum(w[Cj]))

                    # Moving i from Ci to Cj decreases positive mistakes
                    # pos_outer = sum(A[i,cj_neighbs])
                    pos_outer = sum(A[i,Cj])

                    total_outer = lam*neg_outer - pos_outer

                    # This is the overall change in objective if we move i to Cj
                    change = total_outer + total_inner

                end

                # Check if this would be a more beneficial change than other options
                if change < BestImprove
                    BestImprove = change
                    BestC = Cj_ind
                    improving = true
                end
            end

            # Move i to the best cluster to move it to
            if BestImprove < 0
                ci_old = c[i]
                # if its > 100
                #     println("moved node $i from $ci_old to $BestC")
                #     obj = LamCCobj_slow(A,c,w,lam)
                #     println("obj = $obj")
                # end
                c[i] = BestC

                # Remove i from its old cluster
                Clusters[ci_old] = setdiff(Clusters[ci_old],i)

                # And add it to its new cluster
                push!(Clusters[BestC],i)

                improving = true
            end

        end
    end
    if its == 1
        improved = false
    else
        improved = true
        c, Clusters = renumber(c,Clusters)
    end

    if randflag
        c = c[undop]
    end

    return c, improved

end

# Collapse a clustering into a new network of supernodes and weighted edges
function collapse_clustering(A::SparseMatrixCSC{Float64,Int64},w::Vector{Float64},c::Vector{Int64})

    n = size(A,1)
    # Fill cluster arrays (faster than working with the "find" function)
    Clusters = Vector{Vector{Int64}}()
    for v = 1:maximum(c)
        push!(Clusters, Vector{Int64}())
    end

    for i = 1:n
        push!(Clusters[c[i]],i)
    end

    # Number of supernodes to form
    N = round(Int64,maximum(c))

    I = Vector{Int64}()
    J = Vector{Int64}()
    V = Vector{Float64}()
    wnew = zeros(N)
    Anew = zeros(N,N)
    start = time()
    # Construct a new sparse matrix with new node weights
    for i = 1:N
        #Ci = findall(x->x == i,c)
        Ci = Clusters[i]
        wnew[i] = sum(w[Ci])
        ACi = A[Ci,:]
        for j = i+1:N
            Cj = Clusters[j]
            Eij = sum(ACi[:,Cj])
            Anew[i,j] = Eij
        end
    end
    getedges = time()-start

    start = time()
    # Anew = sparse(I,J,V,N,N)
    Anew = Anew+Anew'
    Anew = sparse(Anew)

    return Anew, wnew
end

function SimpleCliqueExp(H::SparseMatrixCSC{Float64,Int64})

    A = H'*H
    for i = 1:size(A,1)
        A[i,i] = 0.0
    end
    dropzeros!(A)
    return A
end

function WeightedCliqueExp(H::SparseMatrixCSC{Float64,Int64}, order::Vector{Int64})

    m,n = size(H)
    I = Vector{Int64}()
    J = Vector{Int64}()
    vals = Vector{Float64}()
    Hyperedges = incidence2elist(H)
    for e = 1:m
        Edge = Hyperedges[e]
        Ord = order[e]
        for ii = 1:length(Edge)
            for jj = ii+1:length(Edge)
                i = Edge[ii]
                j = Edge[jj]
                push!(I,i); push!(J,j); push!(vals,1/(Ord-1))
            end
        end
    end

    A = sparse(I,J,vals,n,n)
    A = sparse(A+A')
    return A
end

## Sort sizes: Return a set of cluster sizes, arranged in order
function ClusterSizes(c)
    c = renumber(c)

    sizes = Vector{Int64}()
    for i = 1:maximum(c)
        inds = findall(x->x==i,c)
        push!(sizes,length(inds))
    end

    return sort(sizes,rev=true)
end

# return precision, recall, and f1 scores for how well
# the clusters in c "track" the ground truth in labels
function ClusterTracker(c,labels, uLab)

    # uLab = unique(labels)
    nl = length(uLab)

    Clusters = Vector{Vector{Int64}}()
    for j = 1:length(unique(c))
        push!(Clusters,Vector{Int64}())
    end
    for i = 1:length(c)
        push!(Clusters[c[i]],i)
    end

    ## For each ground truth cluster, find the best
    F1s = zeros(nl)
    PRs = zeros(nl)
    REs = zeros(nl)
    for i = 1:nl
        label = uLab[i]

        clus = findall(x->x==label,labels)
        f1 = 0; pr = 0; re = 0
        # @show length(clus)
        # For each cluster, track the best pr, re, and f1
        for j = 1:length(Clusters)
            myclus = Clusters[j]
            p,r,f = PRF(clus,myclus)
            # println("Cluster $i, mycluster $j : $p \t $r \t $f")
            if p > pr
                pr = p
            end
            if r > re
                re = r
            end
            if f > f1
                f1 = f
            end
        end
        F1s[i] = f1
        PRs[i] = pr
        REs[i] = re
    end
    return F1s, PRs, REs
end

# This computes the precision, recall, and F1 score for a set Returned
# compared against a Target set
function PRF(Target,Returned)

    if length(Returned) == 0
        pr = 0; re = 0; F1 = 0
    else
        TruePos = intersect(Returned,Target)
        pr = length(TruePos)/length(Returned)
        re = length(TruePos)/length(Target)
        F1 = 2*(pr*re)/(pr+re)

        if length(TruePos) == 0
            F1 = 0
        end
    end

    return pr, re, F1

end

# ARI scores
using Statistics
using Clustering

function ari(x,y)
    evaluations = randindex(x, y)
    ari = evaluations[1]
    return ari
end


include(pwd()*raw"\GRASPEL.jl")

vars = matread(pwd()*raw"\USPS.mat") #read in matfile with USPS data
fea = vars["fea"] #read in matfile matrix as a Vector
gnd = vec(vars["gnd"].-1) #Read in matfile gnd, labels in gnd are offset by 1?? Also convert to vector
#Graspel = 1 #do graspel iters? unused currently
kNN = 2 #for setting up initial kNNT
samplePerc = 1/100 #percentage of sorted nodes as candidate edges
edge_iter = .01/100 #edge budget per iteration (percentage of nodes)
sig = 1e3 #feature variance
embedTol = 10 #embedding distortion tolerance (should be > 1)
maxIterations = 100 #max number of GRASPEL maxIterations

Subset = [0:1:9;]
num_cluster = length(Subset)
num_eigs = 5*num_cluster
num_test = 10 #number of kmean clustering runs
num_res_test = 1000 #number of resistance TestResistance

#order dataset into subsets
num_node, dim = size(fea)
p = sortperm(gnd)
fea = fea[p,:]
#feaVec = [vec(fea[i,:]) for i = 1:num_node]#fea must be a vector of vectors
gnd = gnd[p]
#preprocessing
##center images by average value
for i = 1:num_node
    fea[i,:] = fea[i,:].-mean(fea[i,:])
end
##normalize
##Julia norm(A,2) returns Frobenius Norm... not the same as Matlab's default 2-norm
##max(svd(A)), in this case svdvals(A)[1] , should eventually test this with larger data sets
fea = (fea/svdvals(fea)[1])*sqrt(dim)

#Do GRASPEL
G, varRatio, costVec, numIter = FiedlerConstruct(fea, samplePerc, edge_iter, sig, maxIterations, kNN, embedTol, num_node, dim)

num_edge=(length(edges(G)));
density=num_edge/num_node
L_J = laplacian_matrix(G)
L_J = L_J + (1/(sig^2))*sparse(LinearAlgebra.I,L_J.m,L_J.n)
A_J = adjacency_matrix(G)
mmwrite(pwd()*raw"\Julia_L.mtx",L_J)
mmwrite(pwd()*raw"\Julia_A.mtx",A_J)
@info "Density: $density"

TestResistance = 1 #resistance for Reff, check against data Distances


export newton

immutable DualFun{F,T}
    f::F
    J::T
end
DualFun(f::Fun)=DualFun(f,
                        SpaceOperator(IdentityOperator(),
                              space(f),
                              space(f)))

domain(df::DualFun)=domain(df.f)

differentiate(d::DualFun)=DualFun(d.f',Derivative(rangespace(d.J))*d.J)
integrate(d::DualFun)=DualFun(integrate(d.f),Integral(rangespace(d.J))*d.J)
function Base.cumsum(d::DualFun)
    Q=Integral(rangespace(d.J))*d.J
    DualFun(cumsum(d.f),(I-Evaluation(rangespace(Q),false))*Q)
end

Base.transpose(d::DualFun)=differentiate(d)

^(d::DualFun,k::Integer)=DualFun(d.f^k,k*d.f^(k-1)*d.J)

for OP in (:+,:-)
    @eval begin
        $OP(a::DualFun,b::Union{Number,Fun})=DualFun($OP(a.f,b),a.J)
        $OP(a::Union{Number,Fun},b::DualFun)=DualFun($OP(a,b.f),b.J)
        $OP(a::DualFun,b::DualFun)=DualFun($OP(a.f,b.f),$OP(a.J,b.J))
    end
end
-(a::DualFun)=DualFun(-a.f,-a.J)

*(a::Union{Number,Fun},b::DualFun)=DualFun(a*b.f,a*b.J)
*(a::DualFun,b::Union{Number,Fun})=DualFun(b*a.f,b*a.J)
*(a::DualFun,b::DualFun)=DualFun(a.f*b.f,a.f*b.J+b.f*a.J)

/(a::Union{Number,Fun},b::DualFun)=DualFun(a/b.f,-a/b.f^2*b.J)
/(a::DualFun,b::Union{Number,Fun})=DualFun(a.f/b,a.J/b)
/(a::DualFun,b::DualFun)=DualFun(a.f/b.f,a.J/b.f-a.f/b.f^2*b.J)

@compat (d::DualFun)(x) = DualFun(d.f(x),Evaluation(rangespace(d.J),x)*d.J)
Base.first(d::DualFun) = DualFun(first(d.f),Evaluation(rangespace(d.J),false)*d.J)
Base.last(d::DualFun) = DualFun(last(d.f),Evaluation(rangespace(d.J),true)*d.J)

jacobian(d::DualFun)=d.J

Base.promote_rule{DF<:DualFun,T<:Number}(::Type{DF},::Type{T})=DualFun
Base.convert(::Type{DualFun},b::Number)=DualFun(b,0)

function Operator(f::Function,ds::Space)
    if (isgeneric(f)&&applicable(f,0)) || (!isgeneric(f)&&arglength(f)==1)
        df=f(DualFun(zeros(ds)))
    elseif (isgeneric(f)&&applicable(f,0,0)) || (!isgeneric(f)&&arglength(f)==2)
        df=f(Fun(ds),DualFun(zeros(ds)))
    else
        error("Not implemented")
    end

    if isa(df,Array)
        map(u->u.J,df)
    else
        df.J
    end
end

Operator(f::Function,d)=Operator(f,Space(d))
Operator(f::Function)=Operator(f,Chebyshev())  #TODO: UnsetSpace

# full operator should be
# N=u->[B*u-bcs;...]
function newton(N,u0::Fun;maxiterations=100,tolerance=1E-15)
    u=u0
    for k=1:maxiterations
        DF=N(DualFun(u))
        J=map(jacobian,DF)
        F=map(d->d.f,DF)
        unew=u-J\F
        if norm(unew-u)≤10tolerance
            return unew
        else
            u=chop(unew,tolerance)
        end
    end
    return u
end

newton(N,d::Domain;opts...) =
    newton(N,zeros(d);opts...)

newton{T<:Number}(N,d::AbstractVector{T};opts...) =
    newton(N,Domain(d);opts...)


using DataFrames
using MixedModels
using MixedModels: dataset
using MixedModelsExtras
using Statistics
using Test

@testset "LMM" begin
    model = fit(MixedModel, @formula(reaction ~ 1 + (1|subj)), dataset(:sleepstudy))
    @test icc(model, :subj) == icc(model, [:subj]) == icc(model)
    @test icc(model, :subj) ≈ 0.37918288298942798287
end

@testset "Binomial" begin
    cbpp = dataset(:cbpp)
    model = fit(MixedModel, @formula((incid/hsz) ~ 1 + (1|herd)),
                cbpp, Binomial(), wts=float(cbpp.hsz))
    @test icc(model, :herd) == icc(model, [:herd]) == icc(model)
    @test icc(model, :herd) ≈ 0.1668 atol=0.0005
end

@testset "Bernoulli" begin
    contra = dataset(:contra)
    modelbern = fit(MixedModel, @formula(use ~ 1 + (1|urban&dist)),
                    contra, Bernoulli(), fast=true)
    # force treating as a Binomial model
    modelbin = fit(MixedModel, @formula(use ~ 1 + (1|urban&dist)),
                   contra, Binomial(), fast=true, wts=ones(length(contra.dist)))
    # Bernoullis are a special case of binomial, so make sure they give the same answer
    @test icc(modelbern, Symbol("urban & dist")) ≈ icc(modelbin, Symbol("urban & dist"))
end

@testset "Poisson" begin
    grouseticks = DataFrame(dataset(:grouseticks))
    grouseticks.ch = grouseticks.height .- mean(grouseticks.height)
    model = fit(MixedModel,
                @formula(ticks ~ 1 + year + ch + (1|index) + (1|brood) + (1|location)),
                grouseticks, Poisson(), fast=true)
    @test icc(model, :index) ≈ 0.13467352262090606652 atol=0.0005
    @test icc(model, [:index, :brood]) ≈ 0.3878770599741494518 atol=0.0005
    @test icc(model, [:index, :brood, :location]) ≈ 0.53293244949745322003 atol=0.0005
    @test icc(model, [:index, :brood, :location]) == icc(model)
end


#Function for angular calculation in Blockade shift
#25/07/2017

function angcalc(lli,jji,mji,ll1,jj1)
  spin = 0.5 #Electron spin
  angvec1 = zeros(3)
  for qq = -1:1
    mj1 = mji - qq
    j1vec = collect(-jj1:jj1)
    if length(find(j1vec -> j1vec==mj1,j1vec))==1
      ang11 = (-1)^float(jji-mji+spin+jj1+1)
      ang12 = sqrt((2*jji+1)*(2*jj1+1)*(2*lli+1)*(2*ll1+1))
      ang13 = Wigner6j(jji,1,jj1,ll1,spin,lli)
      ang14 = Wigner3j(jji,1,jj1,-mji,qq,mj1)
      ang15 = Wigner3j(lli,1,ll1,0,0,0)
      angvec1[qq+2] = ang11*ang12*ang13*ang14*ang15 #Shift for Julia indexing
    end
  end
  return angvec1
end


module StickBreak
include("custom_grads.jl") # cumsum, cumprod

"""
x: real vector of dim K - 1
return: simplex of dim K
"""
function transform(x::T) where T
  size_x = size(x)
  ndim = ndims(x)
  dim_head = size_x[1:end-1]
  K = size_x[end] + 1
  ks = cumsum(one.(x), dims=ndim)
  z = sigmoid.(x .- log.(K .- ks))
  one_minus_z_cumprod = cumprod(1 .- z, dims=ndim)

  if ndim == 1
    ones_pad =  Tracker.collect([one.(layer(x, 1))])
    p = vcat(z, ones_pad) .* vcat(ones_pad, one_minus_z_cumprod)
  else
    ones_pad = one.(layer(x, 1))
    p = cat(z, ones_pad, dims=ndim) .* cat(ones_pad, one_minus_z_cumprod, dims=ndim)
  end

  return p
end

"""
x: real vector of dim K - 1
p: simplex of dim K
return: log abs value of determinant of x
"""
function logabsdetJ(x, p)
  ndim = ndims(x)
  K = size(p)[end]
  ks = cumsum(one.(x), dims=ndim)
  z = sigmoid.(x .- log.(K .- ks))
  detJ = sum(log1p.(-z) .+ log.(head(p)), dims=ndim)

  return detJ
end

end # StickBreak


using JLD
using SharedArrays
using Distributed
include("../examples/lorenz.jl")
function obj_fun(x,y,z)
    return (z-28.0)^2.0
end
function obj_fun_erg_avg(s)
    nSteps = 100000
    u = 2*pi*rand(3)
    J = 0.
    u_trj = step(u,s,nSteps)
    x, y, z = view(u_trj,1,:), view(u_trj,2,:), view(u_trj,3,:)
    J = sum(obj_fun.(x,y,z)/nSteps)
    return J
end
function get_Javg_vs_s()
    n_pts = 10
    n_rep = 16000
    s = [10.0, 28.0, 8/3]
    s_ind = LinRange(27.5,31.5,n_pts)
    J = zeros(n_pts)
    J_proc = SharedArray{Float64}(n_rep)
    for i = 1:n_pts
    	s[2] = s_ind[i]
    	J_proc .= 0.
    	t = @distributed for n=1:n_rep
    		J_proc[n] = obj_fun_erg_avg(s)/n_rep
    	end
    	wait(t)
    	J[i] = sum(J_proc)
    end
    save("../data/obj_erg_avg/lorenz/zm28sq_dJds.jld",
    	 "s2", s_ind,
    	"J", J)
end


function solve_block_tridiag!(
    A::AA, B::BB, C::CC, D::DD
) where {
    AA<:AbstractVector,
    BB<:AbstractVector,
    CC<:AbstractVector,
    DD<:AbstractVector,
}
    dim = length(A)
    @assert length(A) == length(B) == length(C) == length(D)
    B[1] .= inv(B[1])
    C[1] .= B[1] * C[1]
    D[1] .= B[1] * D[1]
    for i in 2:dim
        B[i] -= A[i] * C[i - 1]
        D[i] -= A[i] * D[i - 1]
        B[i] .= inv(B[i])
        C[i] .= B[i] * C[i]
        D[i] .= B[i] * D[i]
    end
    for i in (dim - 1):-1:1
        D[i] -= C[i] * D[i + 1]
    end
end


using UniformScalingMatrixFunctions
using Test

@testset "UniformScalingMatrixFunctions.jl" begin
    @test exp(0I) === 1.0I
    @test exp(-im*π/2*I) ≈ -im*I
    @test exp(-im*π/4*I) ≈ (1-im)/√(2)*I

    @test sin(0I) === 0.0I # This assumes iszero(sin(0)), always true?
    @test sin(π/2*I) ≈ I
    @test cos(0I) === 1.0I # This assumes isone(cos(0)), always true?
    @test cos(π/2*I) ≈ 0I atol=1e-14
    @test tan(π/4*I) ≈ I

    @test asin(I) ≈ π/2*I
    @test acos(I) ≈ 0*I atol=1e-14
    @test atan(I) ≈ π/4*I

    # Test hyperbolics + inverses

    @test log(I) === 0.0I
    @test log(exp(I)) === 1.0I

    @test √(4I) === 2.0I
    @test √(4I) == 2I
    @test √(4.0I) === 2.0I
end


using Base: Float64
const π = 3.1415926535
const μ = 39.4234021
const a = 0.38709893
const α = 1.09778201e-8
const e = 0.20563069

rad_to_arcsec(radians::Float64)::Float64 = radians*180/π*60*60

struct Point
	x::Float64
	y::Float64
	Point(x,y) = new(x,y)
	Point() = new(0.0,0.0)
end

norm(p::Point)::Float64 = hypot(p.x, p.y)
angle(p::Point)::Float64 = mod(atan(p.y, p.x),2π)
dot(p1::Point,p2::Point)::Float64 = p1.x*p2.x + p1.y*p2.y

import Base: +,*,/
(p1::Point + p2::Point)::Point = Point(p1.x+p2.x, p1.y+p2.y)
(p::Point * scalar)::Point = Point(p.x*scalar,p.y*scalar)
(scalar * p::Point)::Point = p*scalar
(p::Point / scalar)::Point = Point(p.x/scalar,p.y/scalar)

function acceleration(r_vec::Point)::Point
	r::Float64 = norm(r_vec)
	β::Float64 = -μ/r^3 * (1+α/r^2)
	return r_vec * β
end

function verlet(tmax::Float64,dt_max::Float64,dt_min::Float64)::Nothing
	r_past = Point(a*(1+e),0.)
	v_past = Point(0.,sqrt((μ/a)*(1-e)/(1+e)))
	a_past = acceleration(r_past)

	r_now = Point()
	v_now = Point()
	a_now = Point()

	event_past = dot(r_past,v_past)
	event_now = 0.0

	dt = dt_max
	t = 0.0
	while t < tmax
		r_now = r_past + v_past*dt + a_past*dt^2/2
		a_now = acceleration(r_now)
		v_now = v_past + dt*(a_now+a_past)/2
		event_now = dot(r_now,v_now)

		if event_past<0 && event_now>0
			if dt<dt_min
				println(t," ",rad2deg(angle((r_now+r_past)/2)-π)*3600)
				dt = dt_max
			else
				# rewind time
				t -= dt
				r_now = r_past
				v_now = v_past
				a_now = a_past
				event_now = event_past
				# adjust dt
				dt /= 2
			end
		end

		r_past = r_now
		v_past = v_now
		a_past = a_now
		event_past = event_now
		t += dt
	end
	return nothing
end

# compile (?)
verlet(1e-1,1e-2,1e-3)

println()

verlet(10.0,1e-6,1e-12)


function excess_gibbs_free_energy(model::ActivityModel,p,T,z)
    γ = activity_coefficient(model,p,T,z)
    return sum(z[i]*R̄*T*log(γ[i]) for i ∈ @comps)
end

function eos(model::ActivityModel,V,T,z)
    Σz = sum(z)
    lnΣz = log(Σz)
    g_E = excess_gibbs_free_energy(model,V,T,z)
    pures = model.puremodel
    p = sum(z[i]*pressure(pures[i],V,T) for i ∈ @comps)/Σz
    g_ideal = sum(z[i]*R̄*T*(log(z[i])-lnΣz) for i ∈ @comps)
    g_pure = sum(z[i]*VT_gibbs_free_energy(pures[i],V,T) for i ∈ @comps)
    return g_E+g_ideal+g_pure-p*V
end

function eos_res(model::ActivityModel,V,T,z)
    Σz = sum(z)
    g_E = excess_gibbs_free_energy(model,V,T,z)
    pures = model.puremodel
    g_pure_res = sum(VT_chemical_potential_res(pures[i],V,T)[1]*z[i] for i ∈ @comps)
    p = sum(z[i]*pressure(pures[i],V,T) for i ∈ @comps)/Σz
    p_res = p - Σz*R̄*T/V
    return g_E+g_pure_res-p_res*V
end

function bubble_pressure(model::ActivityModel,T,x)
    sat = saturation_pressure.(model.puremodel,T)
    p_sat = [tup[1] for tup in sat]
    γ     = activity_coefficient(model,1e-4,T,x)
    p     = sum(x.*γ.*p_sat)
    y     = x.*γ.*p_sat ./ p
    return (p,y)
end

function bubble_temperature(model::ActivityModel,p,x;T0=nothing)
    f(z) = Obj_bubble_temperature(model,z,p,x)
    if T0===nothing
        pure = model.puremodel
        sat = saturation_temperature.(pure,p)
        Ti   = zero(x)
        for i ∈ 1:length(x)
            if isnan(sat[i][1])
                Tc,pc,vc = crit_pure(pure[i])
                g(x) = p-pressure(pure[i],vc,x,[1.])
                Ti[i] = Roots.find_zero(g,(Tc))
            else
                Ti[i] = sat[i][1]
            end
        end
        T = Roots.find_zero(f,(minimum(Ti)*0.9,maximum(Ti)*1.1))
    else
        T = Roots.find_zero(f,T0)
    end
    p,y = bubble_pressure(model,T,x)
    return (T,y)
end

function Obj_bubble_temperature(model::ActivityModel,T,p,x)
    sat = saturation_pressure.(model.puremodel,T)
    p_sat = [tup[1] for tup in sat]
    γ     = activity_coefficient(model,1e-4,T,x)
    y     = x.*γ.*p_sat ./ p
    return sum(y)-1
end


function multi_factorial(a)

    # function to compute multivariate factorials of the form
    # \bm{i!} = i_1!i_2!...i_N!

    return prod([factorial(b) for b in a])

end

function gamma_factorial(a, i)

    # function to compute gamma factors of the form
    # (α)_i = α(α+1)...(α+i-1)

    product = 1.0
    for j in 1:i
        product *= a+i-j
    end

    return product

end

function gamma_closure(sys::MomentEquations, binary_vars::Array{Int,1}=Int[])

    closure = OrderedDict() # leaving symbolic higher order terms
    closure_exp = OrderedDict() # expanding out all higher order terms fully

    N = sys.N
    if sys isa CentralMomentEquations
        M = copy(sys.M)
        μ = central_to_raw_moments(N, sys.m_order)
        μ_symbolic = define_μ(N, sys.q_order)
    else
        # consider covariances explicitly as M (without converting to raw moments)
        # and converting at the end leads to much more tractable expressions
        M = define_M(N, 2)
        μ = copy(sys.μ)
        μ_symbolic = copy(μ)
    end

    # obtain the shape and scale parameters
    # starting with symbolic α initially and substituting the respective
    # central/raw moments after the expressions are simplified

    #iter = Iterators.product(fill(1:N, 2)...)
    #α = define_temp_vars(:α, vcat(iter...))
    #x = define_temp_vars(:x, 1:N)

    #αs = Matrix{typeof(α[1])}(undef, N, N)
    #for (ind, (i,j)) in enumerate(iter)
    #    αs[i, j] = α[ind]
    #end

    @variables αs[1:N, 1:N]
    @variables x[1:N]

    αs = scalarize(αs)
    x = scalarize(x)
    for i in 1:N
        for j in i+1:N
            αs[j, i] = αs[i, j]
        end
    end

    for i in 1:N
        αs[i, i] = x[i]
        for k in 1:N
            if i != k
                αs[i, i] -= αs[i, k]
            end
        end
    end
    # the computations above only work with Num (maybe a bug?), hence `value` only now
    α = value.(αs)
    x = value.(x)

    symbolic_sub = Dict()
    β = Array{Any}(undef, N)
    for i in 1:N
        eᵢ = sys.iter_1[i]
        for j in i+1:N
            eⱼ = sys.iter_1[j]
            #symbolic_sub[αs[i, j]] = M[eᵢ .+ eⱼ] * μ[eᵢ] * μ[eⱼ] / M[2 .* eᵢ] / M[2 .* eⱼ]
            symbolic_sub[αs[i, j]] = M[eᵢ .+ eⱼ] * μ[eᵢ] * μ[eⱼ] * M[2 .* eᵢ]^-1 * M[2 .* eⱼ]^-1
        end
        #symbolic_sub[x[i]] = μ[eᵢ]^2 / M[2 .* eᵢ]
        symbolic_sub[x[i]] = μ[eᵢ]^2 * M[2 .* eᵢ]^-1
        #β[i] = M[2 .* eᵢ] / μ[eᵢ]
        β[i] = M[2 .* eᵢ] * μ[eᵢ]^-1

    end

    # construct the raw moments that follow the multivariate gamma distribution
    iters = Vector(undef, N)

    # TupleTools.sort comes in handy here
    unique_iter_q = unique(sort(i) for i in sys.iter_q)
    sub = Dict()

    iter_2nd_order = filter(x-> sum(x) == 2, sys.iter_m)
    # line below reproduces Lakatos et al. (2015) results - stronger assumptions
    #unique_iter_q = unique(sort(i) for i in vcat(sys.iter_m, sys.iter_q))

    for iter in unique_iter_q

        # construct iterator through the product of sums in the Eq. for raw moments
        for j in 1:N
            iter_j = Iterators.filter(x -> sum(x)==iter[j], sys.iter_all)
            iters[j] = iter_j
        end

        iter_k = Iterators.product(iters...)
        suma = 0.0

        multif_iter = multi_factorial(iter)
        for k in iter_k
            #factor1 = multif_iter / prod([multi_factorial(k_i) for k_i in k])
            factor1 = multif_iter * prod([multi_factorial(k_i) for k_i in k])^-1
            factor2 = 1.0
            for a in 1:N
                factor2 *= gamma_factorial(α[a, a], k[a][a])
                for b in a+1:N
                    factor2 *= gamma_factorial(α[a, b], k[a][b]+k[b][a])
                end
            end
            suma += factor1*factor2
        end

        suma = simplify(suma, expand=true)
        suma = substitute(suma, symbolic_sub)
        suma *= prod([β[j]^iter[j] for j in 1:N])
        μ[iter] = simplify(suma, expand=true)
        closure[μ_symbolic[iter]] = μ[iter]

        # Note that joint moments of multivariate distributions are symmetric in the sense
        # that, for example, M₀₃₀ is equivalent in form to M₃₀₀, only in their respective expressions
        # indices of all lower order moments have to be permuted correspondingly
        # let M₃₀₀ = μ₁₀₀+M₂₁₀, then M₀₃₀ = μ₀₁₀+M₁₂₀ (same form but indices were changed as 1->2 & 2->1)
        # Utilising the symmetry we can greatly reduce the computation time as no redundant
        # expression constructions/simplifications need to be performed anymore
        # TODO: implement for all closures?

        perms = collect(multiset_permutations(iter, length(iter)))[2:end]

        for iter_perm in perms

            iter_perm_ind = sortperm(sortperm(iter_perm))
            for i in sys.iter_1
                sub[μ_symbolic[i]] = μ_symbolic[i[iter_perm_ind]]
            end
            for i in iter_2nd_order
                sub[M[i]] = M[i[iter_perm_ind]]
            end
            #maybe faster to do substitution step by step without dict?
            iter_perm = Tuple(iter_perm)
            μ[iter_perm] = substitute(μ[iter], sub)
            closure[μ_symbolic[iter_perm]] = μ[iter_perm]
        end

    end

    # construct the corresponding truncated expressions of higher order
    # central moments from the obtained gamma raw moment expressions
    if sys isa CentralMomentEquations
        raw_to_central = raw_to_central_moments(N, sys.q_order, μ)
        central_to_raw = central_to_raw_moments(N, sys.q_order)
        closure_M = OrderedDict()
        for i in sys.iter_q
            closure_exp[M[i]] = raw_to_central[i]
            closure_M[M[i]] = simplify(closure[μ_symbolic[i]]-(central_to_raw[i]-M[i]))
        end
        closure = closure_M
    else
        raw_to_central = raw_to_central_moments(N, 2)
        M_to_μ = [M[i] => raw_to_central[i] for i in filter(x -> sum(x)==2, sys.iter_all)]

        # variance of bernoulli variable is expressed in terms of its mean
        #for ind in binary_vars
    #        eᵢ = sys.iter_1[ind]
    #        push!(M_to_μ, M[2 .* eᵢ] => μ[eᵢ] - μ[eᵢ]^2)
    #    end

        for i in sys.iter_q
            closure_exp[sys.μ[i]] = substitute(μ[i], M_to_μ)
            expr = substitute(closure[sys.μ[i]], M_to_μ)
            closure[sys.μ[i]] = simplify(expr)
        end
    end

    if isempty(binary_vars)
        return close_eqs(sys, closure_exp, closure, true)
    else
        # have to perform the bernoulli simplify in the end for gamma closure
        # as otherwise some symbolic terms do not cancel out and break it

        redundant_iter, redundant_eqs, iter_sub = bernoulli_reduce(sys, binary_vars)

        closed_eqs = Equation[]
        for (i, eq) in enumerate(get_eqs(sys.odes))
            if !(i in redundant_eqs)

                closed_rhs = substitute(eq.rhs, closure_exp)
                closed_rhs = expand(closed_rhs)
                closed_rhs = substitute(closed_rhs, iter_sub)
                closed_rhs = simplify(closed_rhs)

                push!(closed_eqs, Equation(eq.lhs, closed_rhs))

            end
        end

        iter_rm = intersect(sys.iter_q, redundant_iter)

        if sys isa CentralMomentEquations
            moments = M
        else
            moments = μ_symbolic
        end

        for i in iter_rm
            delete!(closure, moments[i])
        end

        vars = extract_variables(closed_eqs, N, sys.q_order)
        odename = Symbol(nameof(sys), "_gamma_closure")
        odes = ODESystem(closed_eqs, get_iv(sys.odes), vars, get_ps(sys.odes); name=odename)

        return ClosedMomentEquations(odes, closure, sys)

    end

end


type Softmax <: BanditAlgorithm
  temperature::Float64
  counts::Vector{Int64}
  values::Vector{Float64}
end

function Softmax(n_arms::Int64, temperature::Float64)
  Softmax(temperature, zeros(Int64, n_arms), zeros(n_arms))
end

function initialize(algo::Softmax, n_arms::Int64)
  algo.counts = zeros(Int64, n_arms)
  algo.values = zeros(n_arms)
end

function initialize(algo::Softmax, arms::Vector)
  initialize(algo, length(arms))
end

function select_arm(algo::Softmax)
  z = sum(exp(algo.values ./ algo.temperature))
  probs = exp(algo.values ./ algo.temperature) ./ z
  return rand(Categorical(probs))
end

function update(algo::Softmax, chosen_arm::Int64, reward::Real)
  algo.counts[chosen_arm] += 1
  n = algo.counts[chosen_arm]

  value = algo.values[chosen_arm]
  algo.values[chosen_arm] = ((n - 1) / n) * value + (1 / n) * reward
end


module TestCovariance
    using StreamStats
    using Distributions
    using Base.Test

    # Means of uniform draws
    for n in rand(1:1_000_000, 10)
        xs = randn(n)
        ys = rand() * xs + randn(n)
        stat = StreamStats.Covariance()
        for (x, y) in zip(xs, ys)
            update!(stat, x, y)
        end
        online_cov = cov(stat)
        online_cor = cor(stat)
        online_n = nobs(stat)
        batch_cov = cov(xs, ys)
        batch_cor = cor(xs, ys)

        @test_approx_eq(online_cov, batch_cov)
        @test_approx_eq(online_cor, batch_cor)
        @test online_n == n
    end
end


var1(x) = length(x) > 1 ? var(x) : 0.0
std1(x) = length(x) > 1 ? std(x) : 0.0
sem1(x) = length(x) > 1 ? sem(x) : 0.0
variation1(x) = length(x) > 1 ? variation(x) : 0.0
skewness1(x) = length(x) > 1 ? skewness(x) : 1.0
kurtosis1(x) = length(x) > 1 ? kurtosis(x) : -1.0

for T1 in (:T, :(Union{Missing,T}))
    for (R,F) in ((:runmin, :minimum), (:runmax, :maximum),
                  (:runmean, :mean), (:runmedian, :median),
                  (:runvar, :var1), (:runstd, :std1),
                  (:runskewness, :skewness1), (:runkurtosis, :kurtosis1),
                  (:runsem, :sem1),
                  (:runmad, :mad_not_normalized),
                  (:runmad_normalized, :mad_normalized),
                  (:runvariation, :variation1))
        @eval begin
            $R(data::V, windowspan::Int) where {T, V<:AbstractVector{$T1}} =
                running($F, data, windowspan)
            $R(data::V, windowspan::Int, weighting::AbstractVector{S}) where {T, V<:AbstractVector{$T1}, S} =
                running($F, data, windowspan, weighting)
            $R(data::V, windowspan::Int, weighting::AbstractWeights) where {T, V<:AbstractVector{$T1}} =
                running($F, data, windowspan, weighting.values)
        end
    end
end

runcor(data1::V1, data2::V2, windowspan::Int) where {T, V1<:Union{AbstractVector{T}, AbstractVector{Union{Missing,T}}},
                                                         V2<:Union{AbstractVector{T},AbstractVector{Union{Missing,T}}}} =
    running(cor, data1, data2, windowspan, 1)

runcov(data1::V1, data2::V2, windowspan::Int) where {T, V1<:Union{AbstractVector{T}, AbstractVector{Union{Missing,T}}},
                                                         V2<:Union{AbstractVector{T},AbstractVector{Union{Missing,T}}}} =
    running(cov, data1, data2, windowspan, 0)

runcor(data1::V1, data2::V2, windowspan::Int, weighting::AbstractVector{S}) where {S, T, V1<:Union{AbstractVector{T}, AbstractVector{Union{Missing,T}}},
                                                                                V2<:Union{AbstractVector{T}, AbstractVector{Union{Missing,T}}}} =
    running(cor, data1, data2, windowspan, weighting, 1)

runcov(data1::V1, data2::V2, windowspan::Int, weighting::AbstractVector{S}) where {S, T, V1<:Union{AbstractVector{T}, AbstractVector{Union{Missing,T}}},
                                                                                V2<:Union{AbstractVector{T}, AbstractVector{Union{Missing,T}}}} =
    running(cov, data1, data2, windowspan, weighting, 0)


immutable ThreeBoxSITB <: Cut{3}; groups; end
export ThreeBoxSITB

function getcut(prob::Problem, vars::Variables, ::ThreeBoxSITB, i, j, k)
    cˣ, cʸ = vars.cˣ, vars.cʸ
    ℓˣ, ℓʸ = vars.ℓˣ, vars.ℓʸ
    zˣ, zʸ = vars.zˣ, vars.zʸ

    lbˣ, lbʸ = prob.wlb, prob.hlb
    Lˣ, Lʸ = prob.W, prob.H
    @LinearConstraints(begin
        cˣ[i] ≥ 0.5ℓˣ[i] + lbˣ[k]*zˣ[k,i] + lbˣ[j]*(zˣ[j,k] + zˣ[k,i] - 1)
        cʸ[i] ≥ 0.5ℓʸ[i] + lbʸ[k]*zʸ[k,i] + lbʸ[j]*(zʸ[j,k] + zʸ[k,i] - 1)
        cˣ[i] + 0.5ℓˣ[i] + lbˣ[k]*zˣ[i,k] + lbˣ[j]*(zˣ[i,k] + zˣ[k,j] - 1) ≤ Lˣ
        cʸ[i] + 0.5ℓʸ[i] + lbʸ[k]*zʸ[i,k] + lbʸ[j]*(zʸ[i,k] + zʸ[k,j] - 1) ≤ Lʸ
    end)
end

immutable ThreeBoxNonoverlap <: Cut{3}; groups; end
export ThreeBoxNonoverlap

function getcut(prob::Problem, vars::Variables, ::ThreeBoxNonoverlap, i, j, k)
    cˣ, cʸ = vars.cˣ, vars.cʸ
    ℓˣ, ℓʸ = vars.ℓˣ, vars.ℓʸ
    zˣ, zʸ = vars.zˣ, vars.zʸ

    lbˣ, lbʸ = prob.wlb, prob.hlb
    Lˣ, Lʸ = prob.W, prob.H
    @LinearConstraints(begin
        cˣ[i] - cˣ[j] + 0.5*(ℓˣ[i]+ℓˣ[j]) ≤ Lˣ*(1-zˣ[i,j]) - lbˣ[k]*(zˣ[i,k]+zˣ[k,j]-1)
        cʸ[i] - cʸ[j] + 0.5*(ℓʸ[i]+ℓʸ[j]) ≤ Lʸ*(1-zʸ[i,j]) - lbʸ[k]*(zʸ[i,k]+zʸ[k,j]-1)
    end)
end


## Copyright (c) 2018 Miguel Bazdresch
##
## This file is distributed under the 2-clause BSD License.

# Calculate the path loss using the Okumura-Hata model.
# References:
#   Hata, M, "Empirical Formula for Propagation Loss in Land Mobile Radio
#   Services", IEEE Transactions on Vehicular Technology, August 1980,
#   pp. 317–325.
#
#   Wikipedia: https://en.wikipedia.org/wiki/Hata_model

module OkumuraHata

using NumericIO

export loss_oh

struct LossOH
    large_city :: Float64
    small_city :: Float64
    suburban   :: Float64
    rural      :: Float64
end

function Base.getindex(l::LossOH, i::Symbol)
    i == :large_city && return l.large_city
    i == :small_city && return l.small_city
    i == :suburban && return l.suburban
    i == :rural && return l.rural
end

function Base.show(io::IO, l::LossOH)
    f(x) = formatted(x, :SI, ndigits=5, charset=:ASCII)

    println("Losses according to Okumura-Hata model (all in dB):")
    println("  Loss in large city:    ", f(l.large_city))
    println("  Loss in small city:    ", f(l.small_city))
    println("  Loss in suburban area: ", f(l.suburban))
    println("  Loss in rural area:    ", f(l.rural))
end

"""
loss_oh(f, d, hb, hm)

Calculate the path loss according to the Okumura-Hata model.

  f is the carrier frequency in MHz, in the range 150 <= f <= 1500
  d is the distance between base station and mobile in km; 1 <= d <= 10
  hb is the base station antenna height in m; 30 <= hb <= 200
  hm is the mobile antenna height in m; 1 <= hm <= 10

The loss in decibels is calculated for large cities, small cities,
suburban areas and rural areas.

Returns the four calculated losses in a struct of type LossOH, which can be
accessed by property name (`loss.small_city`) or by indexing
(`loss[:suburban]`).
"""
function loss_oh(f, d, hb, hm)

    (150 <= f <= 1500) || error("Invalid frequency f = $f: must be in range [150, 1500]")
    (1 <= d <= 10) || error("Invalid distance d = $d: must be in rage [1,10]")
    (30 <= hb <= 200) || error("Invalid base station antenna height hb = $hb: must be in range [30,200]")
    (1 <= hm <= 10) || error("Invalid mobile station antenna height hm = $hm: must be in range [1,10]")

    Ch_small = 0.8 + hm*(1.1log10(f)-0.7) - 1.56log10(f)
    if 150 <= f <= 200
        Ch_large = 8.29*(log10(1.54hm))^2 - 1.1
    else
        Ch_large = 3.2*(log10(11.75hm))^2 - 4.97
    end

    Lu = 69.55 + 26.16log10(f) - 13.82log10(hb) + log10(d)*(44.9-6.55log10(hb))

    Lu_large = Lu - Ch_large
    Lu_small = Lu - Ch_small
    Lu_suburban = Lu_small - 2*(log10(f/28)^2) - 5.4
    Lu_rural = Lu_small - 4.78*(log10(f)^2) + 18.33log10(f) - 40.94

    return LossOH(Lu_large, Lu_small, Lu_suburban, Lu_rural)

end

end


using HMMs
using LinearAlgebra
using Distributions

mean = zeros(2)
cov = zeros(2,2) + I
A = [zeros(1,1) ones(1,1)]
b = [0.0]
(x,tol) = HMMs.lenormalcubature( mean, cov, A, b )
@assert( abs(x - 0.5) < tol )

A = zeros(2,2) + I
b = zeros(2)
(x,tol) = HMMs.lenormalcubature( mean, cov, A, b )
@assert( abs(x - 0.25) < tol )

b = [0.0,1.0]
(x,tol) = HMMs.lenormalcubature( mean, cov, A, b )
n1 = cdf( Normal(0,1), 1 )
@assert( abs(x - n1/2) < tol )

b = [1.0,1.0]
(x,tol) = HMMs.lenormalcubature( mean, cov, A, b )
@assert( abs(x - n1^2) < tol )

A = [1.0 0.0;-1.0 -1.0]
b = zeros(2)
@time (x,tol) = HMMs.lenormalcubature( mean, cov, A, b, maxevals=10^7 )
# tol doesn't seem to work for this one
@assert( abs(x - 1/8) < 1e-5 )


using SubsetSelection
using StatsBase, Test

n = 500; p = 1000; k = 10;

indices = sort(sample(1:p, StatsBase.Weights(ones(p)/p), k, replace=false));
w = sample(-1:2:1, k);
X = randn(n,p); Y = X[:,indices]*w;
@time Sparse_Regressor = subsetSelection(OLS(), Constraint(k), Y, X)
for i in 1:k
  @test Sparse_Regressor.indices[i]==indices[i]
end


module Xerxes

using Logging

mutable struct Program
    A::Matrix
    b::Vector
    c::Vector
    nonbasic_vars::Vector
    basic_vars::Vector
    x::Vector # current primal solution
    z::Vector # current dual solution
end

@enum Status Optimal Nonoptimal Unbounded Infeasible

function primal_iterate!(prog::Program)
    B = view(prog.A, :, prog.basic_vars)
    N = view(prog.A, :, prog.nonbasic_vars)
    @debug "Submatrices" B N

    # Check for optimality
    if all(prog.z[prog.nonbasic_vars] .>= 0)
        @info "Program is optimal"
        return Optimal
    end
    @debug "z has negative values, therefore the program is not optimal"

    # Select entering variable
    j = argmin(prog.z)
    if prog.z[j] >= 0
        @debug "No entering variable"
        return Infeasible
    end
    @debug "Entering variable" j

    # Compute primal step direction
    ej = zeros(length(prog.x))
    ej[j] += 1
    primal_step = B \ (N * ej[prog.nonbasic_vars])
    @debug "Primal step direction" primal_step

    # Compute primal step length
    inv_primal_step_length = -Inf
    i = 0
    for k in eachindex(prog.basic_vars)
        if primal_step[k] == 0 && prog.x[prog.basic_vars[k]] == 0
            t = 0
        else
            t = primal_step[k] / prog.x[prog.basic_vars[k]]
        end
        if inv_primal_step_length < t
            inv_primal_step_length = t
            i = prog.basic_vars[k]
        end
    end
    if inv_primal_step_length <= 0
        return Unbounded
    end
    primal_step_length = 1 / inv_primal_step_length
    @debug "Primal step length" primal_step_length
    primal_step *= primal_step_length

    # Select leaving variable (already done, i above)
    @debug "Leaving variable" i

    # Compute dual step direction
    ei = zeros(length(prog.z))
    ei[i] += 1
    dual_step = - N' * (B' \ ei[prog.basic_vars])
    @debug "Dual step direction" dual_step

    # Compute dual step length
    idx = indexin(j, prog.nonbasic_vars)[1]
    dual_step_length = prog.z[j] / dual_step[idx]
    @debug "Dual step length" dual_step_length
    dual_step *= dual_step_length

    # Update current dual and primal solutions
    prog.x[j] = primal_step_length
    prog.x[prog.basic_vars] = prog.x[prog.basic_vars] - primal_step
    prog.z[i] = dual_step_length
    prog.z[prog.nonbasic_vars] = prog.z[prog.nonbasic_vars] - dual_step

    # Update basis
    prog.basic_vars = union(setdiff(prog.basic_vars, [i]), [j])
    prog.nonbasic_vars = union(setdiff(prog.nonbasic_vars, [j]), [i])

    return Nonoptimal
end

end


# This example shows how to use your own field implementation for using them to compute on
# elliptic curves. For making things interesting we shall wrap GaloisFields implementeed binary extension field.
using Test
using CryptoGroups
import CryptoGroups: BinaryField, specialize, PB, bitlength
import GaloisFields: ExtensionField, GaloisField

struct GF₂{F <: ExtensionField} <: BinaryField
    x::F
    GF₂{F}(coeffs::NTuple{N, GaloisField(2)}) where {N, F <: ExtensionField} = new(F(coeffs))

    function GF₂{F}(xrev::BitVector) where F <: ExtensionField
        x = reverse(xrev)
        𝔽₂ = GaloisField(2)
        coeffs = Tuple(𝔽₂(i) for i in x) # May need a reverse
        return GF₂{F}(coeffs)
    end

    GF₂{F}(x::F) where F <: ExtensionField = new{F}(x)
    GF₂(x::F) where F <: ExtensionField = new{F}(x)

    function GF₂(f::BitVector)
        F, x = GaloisField(2, :x=>f)
        return GF₂{F}
    end

    GF₂(poly::Vector{Int}) = GF₂(BitVector(i in poly for i in 0:maximum(poly)))
end

bitlength(::Type{F}) where F <: ExtensionField = length(F.parameters[4]) - 1
bitlength(::Type{GF₂{F}}) where F <: ExtensionField = bitlength(F)

specialize(::Type{GF₂}, basis::PB) = GF₂(basis.f)

Base.convert(::Type{F}, a::BitVector) where F <: GF₂ = F(a)
Base.convert(::Type{BitVector}, a::GF₂) = reverse(BitVector(i.n for i in a.x.coeffs))

#tobits(a::GF₂) = reverse(BitVector(i.n for i in a.x.coeffs)) # For the sake of symmetry

Base.:+(a::F, b::F) where F <: GF₂ = F(a.x + b.x)
Base.:*(a::F, b::F) where F <: GF₂ = F(a.x * b.x)

Base.inv(a::F) where F <: GF₂ = F(inv(a.x))

Base.zero(::Type{GF₂{F}}) where F <: ExtensionField = GF₂(zero(F))
Base.one(::Type{GF₂{F}}) where F <: ExtensionField = GF₂(one(F))

########################## This one we can test easally as follows ##################

import CryptoGroups: BinaryCurve, generator, oncurve, order, @def, AffinePoint, <|, ECPoint, validate
import CryptoGroups: Specs
import .Specs: Curve_B_163_PB, Curve_K_163_PB

let
    B_163v3 = specialize(AffinePoint{BinaryCurve, GF₂}, Curve_B_163_PB)

    g = B_163v3 <| generator(Curve_B_163_PB)
    q = order(Curve_B_163_PB)

    @test oncurve(g)
    @test oncurve(g*3)
    @test g * q == zero(g)
end

@def K_163v3 ECPoint{AffinePoint{BinaryCurve, GF₂}} Curve_K_163_PB

let
    g = K_163v3 <| generator(Curve_K_163_PB)

    @test oncurve(g)
    @test oncurve(g*3)
    @test validate(g)
end


include("../../matplotlib_settings.jl")

set_config(:tab_compensated_summation, false)

using ChargedParticleDynamics.GuidingCenter4d: plot_integral_error, plot_poincare_surface

const coll_title = "Guiding Center 4D Fast 2nd Poincaré Integral Invariant"

XMIN = 1.5
XMAX = 2.0
YMIN = -0.25
YMAX = +0.25

include("../guiding_center_4d_poincare_invariant_2nd.jl")


### A Pluto.jl notebook ###
# v0.16.3

using Markdown
using InteractiveUtils

# ╔═╡ 6446adac-a3e2-4961-bea5-620118149b50
begin
    import Pkg
    Pkg.add(PackageSpec(url="https://github.com/Pocket-titan/DarkMode"))
    import DarkMode
    DarkMode.enable()
    # OR DarkMode.Toolbox(theme="default")
end

# ╔═╡ 81bfcb8c-3345-11ec-21ec-bf37d08f624f
begin
	using DifferentialEquations
	using Plots
	using PlutoUI
	#using DarkMode; DarkMode.Toolbox(theme="default")
end;

# ╔═╡ 369a25d1-6c20-4b22-9f66-a845364b3a0d
function model!(du, u, p, t)
	ω0, ωF, ωM, ϕF, ϕM, μ, ν, αF, αM, ψF, ψM = p
	du[1] = ω0
	du[2] = ωF + μ*sin(u[3]-u[2]-ϕF) + αF*sin(u[1]-u[2]-ψF)
	du[3] = ωM + ν*sin(u[2]-u[3]-ϕM) + αM*sin(u[1]-u[3]-ψM)
end;

# ╔═╡ ecf11f46-b4fa-46c5-8c10-b833adcbd72b
begin
	θ  = rand(3).%(2π) # Initial phases (Sunlight, fibroblast, macrophage)
	tMax = 10.0 	   # Total time
	ω0 = 2π            # Sunlight phase rate of change /day
	ωF = 2π            # Fibroblast phase intrinsic rate of change /day
	ωM = 2π            # Macrophage phase intrinsic rate of change /day
	ϕF = 3π/5          # Fibroblast phase offset in coupling to macrophage
	ϕM = 2π/5 		   # Macrophage phase offset in coupling to fibroblast
	μ  = 1.0           # Amplitude of fibroblast phase coupling to macrophage
	ν  = 0.0           # Amplitude of macrophage phase coupling to fibroblast
	αF = 10.0          # Amplitude of fibroblast phase coupling to sunlight
	αM = 10.0          # Amplitude of macrophage phase coupling to sunlight
	ψF = π/8           # Fibroblast phase offset from sunlight
	ψM = π/4          # Macrophage phase offset from sunlight
	p = [ω0, ωF, ωM, ϕF, ϕM, μ, ν, αF, αM, ψF, ψM]

	prob = ODEProblem(model!,θ,(0.0,tMax),p)
	sol = solve(prob, Tsit5(),saveat=0.01)

	anim = @animate for u in sol.u
		scatter(sin.(u),cos.(u),xlims=(-1.1,1.1),ylims=(-1.1,1.1),aspect_ratio=:equal,marker=[:star8,:circle,:circle],color=[:red,:green,:blue],ms=10,series_annotations=["Light","Fib","Mac"],legend=:none)
	end
	gif(anim,"test.gif",fps=3)

end

# ╔═╡ Cell order:
# ╠═6446adac-a3e2-4961-bea5-620118149b50
# ╠═81bfcb8c-3345-11ec-21ec-bf37d08f624f
# ╠═369a25d1-6c20-4b22-9f66-a845364b3a0d
# ╠═ecf11f46-b4fa-46c5-8c10-b833adcbd72b


@testset "Is integral" begin
  Qx, x = FlintQQ["x"]
  f = x^2 + 1
  K, a = number_field(f, "a")
  Ky, y = K["y"]

  g = y^3 + 3
  L, b = number_field(g, "b")

  @test Hecke.isintegral(b) == true
  @test Hecke.isintegral(fmpq(1, 2)*b) == false

  h = y^4 + 3
  M, c = number_field([g, h], "c")

  @test Hecke.isintegral(c[1]) == true
  @test Hecke.isintegral(fmpq(1, 2)*c[1]) == false
end


using Pkg
Pkg.activate(".")
using Psychometrics
using Distributions
using LinearAlgebra
using Dates

const I = 40
const N = 500

# ITEM PARAMETERS AND LATENTS

items = [Item(i, string("item_",i), ["math"], Parameters2PL()) for i = 1 : I];
examinees = [Examinee(e, string("examinee_",e), Latent1D()) for e = 1 : N];

# Put parameters and latents in matrix
parameters_matrix = get_parameters(items)
latents_matrix = get_latents(examinees)

# RESPONSES

responses = reduce(vcat, [generate_response([e], items[sort!(sample(collect(1:I), 40, replace = false))]) for e in examinees]);
responses_per_item = map(i -> get_responses_by_item_id(i.id, responses), items);  # slow

# Get design matrix
design_matrix = get_design(responses, I, N)

# Get response matrix (zeros are for wrong answers and for missing answers)
response_matrix = get_response_matrix(responses, I, N)

# Go back to a vector of `AbstractResponse`s (starting and ending time of response is Dates.now())
responses2 = get_responses(response_matrix, design_matrix, items, examinees)

# ITEM CHARACTERISTIC FUNCTION (ICF)

## compute probabilities (ICF) by matrix (1PL or 2PL, latent 1D or latend ND) using aθ - b parametrization
p = probability(parameters_matrix, latents_matrix)

## compute probabilities (ICF) on vectors of examinees and items using a(θ - b) parametrization
p_2 = probability(items, examinees)

# ITEM INFORMATION FUNCTION (IIF)

## compute informations (IIF) by matrix (1PL or 2PL, latent 1D) using aθ - b parametrization
information = latent_information(latents_matrix, parameters_matrix)

## compute informations (IIF) on vectors of examinees and items using a(θ - b) parametrization
information_2 = latent_information(examinees, items)

# LIKELIHOOD AND LOG-LIKELIHOOD

# Total Likelihood (only for a(θ - b) parametrization)
L_tot = prod(likelihood(responses, examinees, items));

# Total log-likelihood (only for a(θ - b) parametrization)
l_tot = sum(log_likelihood(responses, examinees, items));
l_tot_per_item = [sum(log_likelihood(responses_per_item[i], examinees, items)) for i in 1:I];

# with gradients
#g_item = zeros(2)
#g_latent = zeros(N)
#l_sum = log_likelihood(responses, examinees, items, g_item, g_latent);

item_expected_information_val = item_expected_information(examinees, items);
# only for 3PL the observed information
#item_observed_information_val = map( r -> item_observed_information(r), responses);

## Estimation

items_est = [Item(string("item_",i), ["math"], Parameters2PL(1.0, [1e-5,5.0], 0.0, [-6.0, 6.0], Product([LogNormal(0, 0.25), Normal(0,1)]),  Product([LogNormal(0, 0.25), Normal(0,1)]), Vector{Vector{Float64}}(undef,2), LinearAlgebra.I(2) )) for i = 1 : I];
examinees_est = [Examinee(string("examinee_",e), Latent1D(0.0, [-6.0, 6.0], Normal(0,1), Normal(0,1), zeros(0), 1.0)) for e = 1 : N];

map responses


####################################################################
function gpmodel(x, flux, σ2, idx; maxiter = 10, maxrandom = 100, seed = 1, show_trace = true)
####################################################################

    #--------------------------------------------------------
    # Auxiliary
    #--------------------------------------------------------

    rg = MersenneTwister(seed)

    JITTER = 1e-6

    N = length(x) # number of observations

    numF  = length(unique(idx)) # number of filters

    @assert(all(sort(unique(idx)) .== collect(1:numF)))

    @assert(N == length(flux) == length(idx) == length(σ2))

    Σdiag = Diagonal(σ2)

    @printf("Running gpmodel with %d observations and %d filters\n", N, numF)

    #--------------------------------------------------------
    # Pre-calculate distances
    #--------------------------------------------------------

    dist2 = pairwise(SqEuclidean(), reshape(x, 1, N), dims=2)

    K     = zeros(N, N)

    #--------------------------------------------------------
    function unpack(param)
    #--------------------------------------------------------

        @assert(length(param) == 2*numF + 1)

        local a     = param[0*numF+1:1*numF]
        local b     = param[1*numF+1:2*numF]
        local logs  = param[2*numF + 1]

        return a, b, logs

    end

    #--------------------------------------------------------
    function objective(param)
    #--------------------------------------------------------

        local a, b, logs = unpack(param)

        -1.0 * logl(a = a, logs = logs, b = b)

    end

    #--------------------------------------------------------
    function blockvector(idx, b)
    #--------------------------------------------------------

        local bv = zeros(eltype(b), N)

        for i in 1:numF
            bv[findall(idx .== i)] .= b[i] # this should be changes, too slow
        end

        return bv

    end

    #--------------------------------------------------------
    function logl(; a = a, b = b, logs = logs)
    #--------------------------------------------------------

        calculatekernelmatrix!(K, dist2, rbf, [0.0; logs])

        local A = Diagonal(blockvector(idx, a))

        local C = (A' * K * A) + Σdiag

        local C = (C.+C')*0.5 + JITTER*I

        local L = try

            cholesky(C).L

        catch err

            if isa(err, PosDefException)

                @warn("Cholesky failed, matrix not posdef. Return -Inf.")

                return -Inf

            else

                throw(err)

            end

        end

        local bvec    = blockvector(idx, b)

        local logdetC = 2*sum(log.(diag(L)))

        local yTCinvy = dot( (flux .- bvec), L'\(L\(flux .- bvec)) ) # this can be improved

        return - 0.5*logdetC - 0.5*yTCinvy - 0.5*N*log(2.0π)

    end

    #--------------------------------------------------------
    # Initialise GP hyperparameters with random search
    # and few iterations
    #--------------------------------------------------------

    initsolutions       = [randn(rg, 2*numF+1)*5 for i=1:maxrandom]
    initialfitness      = @showprogress "\t init random search " map(objective, initsolutions)
    bestinitialsolution = initsolutions[argmin(initialfitness)]

    #--------------------------------------------------------
    # Call optimiser
    #--------------------------------------------------------

    opt    = Optim.Options(iterations = maxiter, show_trace = show_trace, show_every = 100)
    result = Optim.optimize(objective, bestinitialsolution, NelderMead(), opt)

    # opt = NLopt.Opt(:LN_BOBYQA, length(bestinitialsolution))
    #
    # opt.xtol_rel = 1e-8
    #
    # COUNTER, BESTF = 0, Inf
    #
    # function myfunc(x, dummy)
    #
    #     COUNTER +=1
    #
    #     F = objective(x)
    #
    #     BESTF = min(F, BESTF)
    #
    #     mod(COUNTER, 100)==0 ? @printf("Iter %d f=%e\n", COUNTER, BESTF) : nothing
    #
    #     return F
    #
    # end
    #
    # opt.min_objective = myfunc
    #
    # (_minf, minx, _ret) = NLopt.optimize(opt, bestinitialsolution)

    #--------------------------------------------------------
    # Retrieve optimised parameters
    #--------------------------------------------------------

    a, b, logs = unpack(result.minimizer)

    #--------------------------------------------------------
    # Reconstruct curves
    #--------------------------------------------------------
    calculatekernelmatrix!(K, dist2, rbf, [0.0; logs])

    K = K + JITTER*I

    bvec = blockvector(idx, b)
    A    = Diagonal(blockvector(idx, a))
    L    = Σdiag\I
    Λ    = K\I

    Σ = (Λ + A'*L*A)\I

    μ = Σ * (A' * L * (flux .- bvec) )

    figure(101)

    cla()

    clr = ["b","g","r","k","y","m"]

    for (i,l) in enumerate(unique(idx))

        aux = findall(idx .== l)

        plot(t[aux], flux[aux], clr[i]*".", alpha=0.5)

        plot(t, b[i] .+ a[i]*μ, clr[i]*".")

    end

    # xtest = collect(LinRange(minimum(x), maximum(x), 200))
    #
    # # dimensions: N × Ntest
    # k = calculatekernelmatrix(x, xtest, rbf, [0.0;logs])

    # Ntest × 1
    # c = calculatekernelmatrix(xtest, xtest, rbf, uθ)

    # figure(102)
    #
    # cla()
    #
    # μ = k' * ((A'*K*A + σ2*I) \ (flux .- bvec))
    #
    # for (i,l) in enumerate(unique(idx))
    #
    #     plot(xtest, a[i]*μ .+ b[i], clr[i]*"-")
    #
    #     aux = findall(idx .== l)
    #
    #     plot(t[aux], flux[aux], clr[i]*".", alpha=0.5)
    #
    # end

    @show a, b, logs

end


# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using InvertibleNetworks, LinearAlgebra, Test, Statistics

###############################################################################
# Test logdet implementation

# Input
nx = 4
ny = 4
nc = 3
batchsize = 1
X = rand(Float32, nx, ny, nc, batchsize)

# Actnorm and initialize
AN = ActNorm(nc; logdet=true)
AN.forward(X)

# Explicitely compute logdet of Jacobian through probing
# for small number of dimensions
J = zeros(Float32, Int(nx*ny*nc), Int(nx*ny*nc))
for i=1:nc
    count = 1
    for j=1:nx
        for k=1:ny
            E = zeros(Float32, nx, ny, nc, 1)
            E[k, j, i] = 1f0
            local Y = AN.forward(X)[1]
            J[:, (i-1)*nx*ny + count] = vec(AN.backward(E, Y)[1])
            count += 1
        end
    end
end
lgdet1 = log(abs(det(J)))
lgdet2 = AN.forward(X)[2]
@test isapprox((lgdet1 - lgdet2)/lgdet1, 0f0; atol=1f-6)

###############################################################################
# Initialization and invertibility

# Input
nx = 28
ny = 28
nc = 4
batchsize = 1
X = rand(Float32, nx, ny, nc, batchsize)

# Layer and initialization
AN = ActNorm(nc; logdet=false)
Y = AN.forward(X)

# Test initialization
@test isapprox(mean(Y), 0f0; atol=1f-6)
@test isapprox(var(Y), 1f0; atol=1f-3)

# Test invertibility
@test isapprox(norm(X - AN.inverse(AN.forward(X)))/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN.forward(AN.inverse(X)))/norm(X), 0f0, atol=1f-6)

# Reversed layer (all combinations)
AN_rev = reverse(AN)

@test isapprox(norm(X - AN_rev.inverse(AN_rev.forward(X)))/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN_rev.forward(AN_rev.inverse(X)))/norm(X), 0f0, atol=1f-6)

@test isapprox(norm(X - AN_rev.forward(AN.forward(X)))/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN_rev.inverse(AN.inverse(X)))/norm(X), 0f0, atol=1f-6)

@test isapprox(norm(X - AN.forward(AN_rev.forward(X)))/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN.inverse(AN_rev.inverse(X)))/norm(X), 0f0, atol=1f-6)

# Test with logdet enabled
AN = ActNorm(nc; logdet=true)
Y, lgdt = AN.forward(X)

# Test initialization
@test isapprox(mean(Y), 0f0; atol=1f-6)
@test isapprox(var(Y), 1f0; atol=1f-3)

# Test invertibility
@test isapprox(norm(X - AN.inverse(AN.forward(X)[1]))/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN.forward(AN.inverse(X))[1])/norm(X), 0f0, atol=1f-6)

# Reversed layer (all combinations)
AN_rev = reverse(AN)

@test isapprox(norm(X - AN_rev.inverse(AN_rev.forward(X)[1]))/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN_rev.forward(AN_rev.inverse(X))[1])/norm(X), 0f0, atol=1f-6)

@test isapprox(norm(X - AN_rev.forward(AN.forward(X)[1])[1])/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN_rev.inverse(AN.inverse(X)))/norm(X), 0f0, atol=1f-6)

@test isapprox(norm(X - AN.forward(AN_rev.forward(X)[1])[1])/norm(X), 0f0, atol=1f-6)
@test isapprox(norm(X - AN.inverse(AN_rev.inverse(X)))/norm(X), 0f0, atol=1f-6)

###############################################################################
# Gradient Test

AN = ActNorm(nc; logdet=true)
X = randn(Float32, nx, ny, nc, batchsize)
X0 = randn(Float32, nx, ny, nc, batchsize)
dX = X - X0

# Forward pass
Y = AN.forward(X)[1]

function loss(AN, X, Y)

    # Forward pass
    Y_, lgdet = AN.forward(X)

    # Residual and function value
    ΔY = Y_ - Y
    f = .5f0/batchsize*norm(ΔY)^2
    AN.logdet == true && (f -= lgdet)

    # Back propagation
    ΔX, X_ = AN.backward(ΔY./batchsize, Y_)

    # Check invertibility
    isapprox(norm(X - X_)/norm(X), 0f0, atol=1f-6)

    return f, ΔX, get_grads(AN)
end

# Gradient test for X
maxiter = 6
print("\nGradient test actnorm\n")
f0, ΔX = loss(AN, X0, Y)[1:2]
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = loss(AN, X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Gradient test for parameters
AN0 = ActNorm(nc; logdet=true); AN0.forward(randn(Float32, nx, ny, nc, batchsize))
AN_ini = deepcopy(AN0)
θ = get_params(AN_ini)
dθ = get_params(AN)-get_params(AN0)
maxiter = 6
print("\nGradient test actnorm\n")
f0, ΔX, Δθ = loss(AN0, X, Y)
h = 1f0
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)
for j=1:maxiter
    set_params!(AN0, θ+h*dθ)
    f = loss(AN0, X, Y)[1]
    err3[j] = abs(f - f0)
    err4[j] = abs(f - f0 - h*dot(dθ, Δθ))
    print(err3[j], "; ", err4[j], "\n")
    global h = h/2f0
end

@test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=1f1)

###############################################################################
# Gradient Test reversed layer

AN = reverse(ActNorm(nc; logdet=true))
X = randn(Float32, nx, ny, nc, batchsize)
X0 = randn(Float32, nx, ny, nc, batchsize)
dX = X - X0

# Forward pass
Y = AN.forward(X)[1]

# Gradient test for X
maxiter = 6
print("\nGradient test actnorm reverse\n")
f0, ΔX = loss(AN, X0, Y)[1:2]
h = .1f0
err5 = zeros(Float32, maxiter)
err6 = zeros(Float32, maxiter)
for j=1:maxiter
    f = loss(AN, X0 + h*dX, Y)[1]
    err5[j] = abs(f - f0)
    err6[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err5[j], "; ", err6[j], "\n")
    global h = h/2f0
end

@test isapprox(err5[end] / (err5[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err6[end] / (err6[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Gradient test for parameters
AN0 = reverse(ActNorm(nc; logdet=true))
AN0.forward(randn(Float32, nx, ny, nc, batchsize))
AN_ini = deepcopy(AN0)
θ = get_params(AN_ini)
dθ = get_params(AN)-get_params(AN0)
maxiter = 6
print("\nGradient test actnorm reverse\n")
f0, ΔX, Δθ = loss(AN0, X, Y)
h = 1f0
err7 = zeros(Float32, maxiter)
err8 = zeros(Float32, maxiter)
for j=1:maxiter
    set_params!(AN0, θ + h*dθ)
    f = loss(AN0, X, Y)[1]
    err7[j] = abs(f - f0)
    err8[j] = abs(f - f0 - h*dot(dθ, Δθ))
    print(err7[j], "; ", err8[j], "\n")
    global h = h/2f0
end

@test isapprox(err7[end] / (err7[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err8[end] / (err8[1]/4^(maxiter-1)), 1f0; atol=1f1)

###################################################################################################
# Jacobian-related tests

# Gradient test

# Initialization
logdet=true
AN = ActNorm(nc; logdet=logdet); AN.forward(randn(Float32, nx, ny, nc, batchsize))
θ = deepcopy(get_params(AN))
AN0 = ActNorm(nc; logdet=logdet); AN0.forward(randn(Float32, nx, ny, nc, batchsize))
θ0 = deepcopy(get_params(AN0))
X = randn(Float32, nx, ny, nc, batchsize)

# Perturbation (normalized)
dθ = θ-θ0
for i = 1:length(θ)
    dθ[i] = norm(θ0[i])*dθ[i]/(norm(dθ[i]).+1f-10)
end
dX = randn(Float32, nx, ny, nc, batchsize); dX *= norm(X)/norm(dX)

# Jacobian eval
logdet ? ((dY, Y, lgdet, GNdθ) = AN.jacobian(dX, dθ, X)) : ((dY, Y) = AN.jacobian(dX, dθ, X))

# Test
print("\nJacobian test\n")
h = 0.1f0
maxiter = 5
err9 = zeros(Float32, maxiter)
err10 = zeros(Float32, maxiter)
for j=1:maxiter
    set_params!(AN, θ+h*dθ)
    logdet ? ((Y_loc, _) = AN.forward(X+h*dX)) : (Y_loc = AN.forward(X+h*dX))
    err9[j] = norm(Y_loc - Y)
    err10[j] = norm(Y_loc - Y - h*dY)
    print(err9[j], "; ", err10[j], "\n")
    global h = h/2f0
end

@test isapprox(err9[end] / (err9[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err10[end] / (err10[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Adjoint test

set_params!(AN, θ)
logdet ? ((dY, Y, _, _) = AN.jacobian(dX, dθ, X)) : ((dY, Y) = AN.jacobian(dX, dθ, X))
dY_ = randn(Float32, size(dY))
logdet ? ((dX_, dθ_, _, _) = AN.adjointJacobian(dY_, Y)) : ((dX_, dθ_, _) = AN.adjointJacobian(dY_, Y))
a = dot(dY, dY_)
b = dot(dX, dX_)+dot(dθ, dθ_)
@test isapprox(a, b; rtol=1f-3)