# Main driver routine for WOS fractional Laplacian solver
#
# Type
#
# julia> include("julia_wos.jl"); run_field(0.2,1,true)
#
# to solve with alpha=0.2, with parameters from Example 3.11
# from T. Shardlow "A walk outside spheres for the fractional Laplacian: fields and
# first eigenvalue"
#
# Take advantage of extra cores by initialising with add_procs(n).
#
# You need these packages on your LOAD_PATH (KOS refers to the original
# authors of the algorithm: Kyprianou, Osojnik, Shardlow)
using KOSHELPER
@everywhere using KOS
using KOS_FIELD_SOLVE
#
# other Julia packages
@everywhere using Interpolations
using PyPlot
#
ion()
##########################################
##########################################
# begin: end user routines
function run_field(alpha,egnum,plot_flag)
  # Field solve for built-in examples (egnum)
  # If plot_flag=true, fig_err*.pdf and fig_soln*.pdf files are generated
  # for solution and error relative to exact soln (if available).
  #
  if 1==egnum # Example 3.11 (from the paper)
    p,exact=KOS.example1(alpha) # constant rhs, zero exterior
    myext="_1";l0=5
    R=1.; tol=1e-2
  elseif 2==egnum # Example 3.12
    p,exact=KOS.example2(alpha) # quadratic rhs, zero exterior
    myext="_2"; l0=5
    R=1.; tol=1e-2
  elseif 3==egnum
    p,exact=KOS.example3(alpha) # Green's fn exterior on unit ball, zero rhs
    myext="_3"
    R=1.2; l0=7
    tol=1e-3
  elseif 4==egnum
    p,exact=KOS.example4(alpha) # Green's fn exterior on square, zero rhs
    myext="_3"
    R=1.2; l0=7
    tol=1e-2
  elseif 5==egnum
    p,exact=KOS.example5(alpha) # Example 3.13 (no exact available)
    myext="_5"
    R=1.2; l0=7
    tol=1e-2
  end
  dict=Dict("max_samples"=>Int32(1e6),
  "tol"=>tol,
  "no_samps"=>4000,
  "no_levels"=>7,"l0"=>l0,
  "report"=>true,
  "solve_method"=>"ml")
  G=KOSHELPER.get_tri_mesh([-R,R,-R,R],
                            dict["no_levels"],p.d)
  #
  out,elapsed_time=get_solution(p, dict, G)
  G.get_ext(out,p.fn_ext)
  #
  Z=G.move_to_TriMesh(out)
  #
  exact_Z=G.evalfn(exact, p.fn_ext)
  # two ways of computing L2 error
  err=sqrt(G.sq_int(out-exact_Z,dict["no_levels"]))
  err2=sqrt(G.sq_int2(out,exact,dict["no_levels"]))
  # L2 norm of exact solution
  norm_exact=sqrt(G.sq_int(exact_Z,dict["no_levels"]))
  println("Error norm=",err, " err2=",err2," elapsed time=", elapsed_time,".")
  #
  Ze=G.move_to_TriMesh(exact_Z)
  if plot_flag
    do_plot(G,Z,"",string("fig_soln",myext,".pdf") )
    do_plot(G,abs.(Z-Ze),"",string("fig_err",myext,".pdf") )
  end
  println("Elapsed time=", elapsed_time)
  return elapsed_time, err2, (err2/norm_exact)
end
#################
function run_stats(alpha)
  # Analyse MLMC parameter
  # Find computation times for different initial level
  # from l0 upto no_levels
  #
  #p,exact=KOS.example3(alpha)
  p,exact=KOS.example2(alpha)
  dict=Dict("max_samples"=>Int32(1e5),
  "tol"=>1e-2,
  "no_samps"=>4000,
  "no_levels"=>9,"l0"=>6,
  "report"=>true,
  "solve_method"=>"ml_analyse")
  G=KOSHELPER.get_tri_mesh([-1.,1.,-1.,1.],
                            dict["no_levels"],p.d)
  #
  out,elapsed_time=get_solution(p, dict, G)
  #
end
####################
function run_couple(vec_alpha)
  # Compute stats on the coupling (variance and time per sample)
  #
  #p,exact=KOS.example3(1.0)
  p,exact=KOS.example5(1.0)
  dict=Dict("max_samples"=>Int32(1e5),
  "tol"=>1e-2,
  "no_levels"=>7,"l0"=>3,
  "report"=>true,
  "no_samps"=>40000,
  "solve_method"=>"couple_analyse")
  G=KOSHELPER.get_tri_mesh([-1.,1.,-1.,1.],
                            dict["no_levels"],p.d)
  for alpha in vec_alpha
    p.alpha=alpha
    p,exact=KOS.example5(alpha)
    out,elapsed_time=get_solution(p, dict, G)
  end
end
##################
function run_ev(alpha, lam_tol, no_its,relax_flag,no_restarts,p_restart)
  # Main routine for computing eigenvalues on the unit ball.
  # Inputs: alpha, parameter for fractional Laplacian in (0,2).
  # lam_tol: tolerance for the eigenvalue.
  # no_its: number of Arnoldi iterations.
  # relax_flag: use variable accuracy if true.
  #
  # experimental (set no_restarts=0 for algorithm in paper)
  # no_restarts: no of iteratively restarted its
  # p_restart:Krylov dimension for restart.
  #
  # Ouputs: lam: computed eigenvalue.
  # el_time: elapsed time in seconds.
  # eig_res: eigenvalue residual.
  println("\n==================\nGet smallest eigenvalue.\n")
  #
  # domain is set here via choice of distance function.
  p,exact=KOS.example0(alpha)
  #
  # set the parameters for the method
  dict=Dict("max_samples"=>Int32(1e5), "report"=>true,
    #
    "no_levels"=>7, "l0"=>5, # number of levels
    "solve_method"=>"ml", # don't change this
    #
    "lam_tol"=>lam_tol, # tolerance for eigenvalue solve
    "no_arnoldi_its"=>no_its,
    #
    "relax"=>relax_flag, # for inexact Arnolid
    "B"=>3,#see paper for B (large B smaller WOS tolerance)
    # experimental
    "no_restarts"=>no_restarts, # for implicit restarts
    "p_restart"=>p_restart
   )

  println("Using no_restarts=",dict["no_restarts"],", restart dim=", p_restart,
          ",  variable accuracy=", dict["relax"],
          ", no_levels=",dict["no_levels"], ", lam_tol=",dict["lam_tol"])
  #
  # set up the TriMesh
  G=KOSHELPER.get_tri_mesh([-1.,1.,-1.,1.],
                            dict["no_levels"],p.d)
  # first Arnoldi run
  # initial vector for eigenvector
  Y=ones(G.n_int);
  for i=1:G.n_int
    Y[i]=1-(1-p.d(G.get_pt(i))[1])^2
  end
  arnoldi=Construct_Arnoldi(Y)
  ##########################################
  i=0
  println(">>>>First Arnoldi run")
  lam,evec,delta_lam,arnoldi,el_time,eig_res=wos_ev(p,dict,G,arnoldi,1.)
  dyda_compare(alpha,  lam)
  println("-----------------------------------\n\n")
  #
  delta_lam=1.0
  while (i<dict["no_restarts"])
    i+=1
    KOSHELPER.SmartRestart!(arnoldi, dict["p_restart"])
    println("     >>>>>\n     >>>>Next Arnoldi run number=", i+1)
    #
    lam,evec,delta_lam,arnoldi,el_time,eig_res=wos_ev(p,dict,G,arnoldi,lam)
    #
    dyda_compare(alpha,  lam)
    println("-----------------------------------\n\n")
  end

  println("Using no_restarts=",dict["no_restarts"],", restart dim=",p_restart,
          ", variable accuracy=",dict["relax"],
          ", no_levels=",dict["no_levels"], ", lam_tol=",dict["lam_tol"])
  out=G.reshape(evec,KOS.vec_fn_zero)
  Z=G.move_to_TriMesh(out)
  #
  do_plot(G,Z,"evec", "evec.png")
  return lam, el_time, eig_res
end
#######################################
#######################################
function batch_ev()
  # This script runs the eigenvalue solver for a range of alpha
  #
  lam_tol=1e-2; no_its=5; relax_flag=true; no_restarts=0; p_restart=3;
  valpha=[0.1,0.2,0.5,1.0,1.5,1.8,1.9]
   valpha=[0.5,1.0,1.5]
  n=length(valpha); vl=zeros(n); vel_time=zeros(n); veig_res=zeros(n)
  for i in 1:n
    l,et,er=run_ev(valpha[i],lam_tol,no_its,relax_flag,no_restarts,p_restart)
    vl[i]=l; vel_time[i]=et; veig_res[i]=er;
  end
  println("Vector of computed lambda=", vl)
  println("Vector of elasped times=", vel_time)
  println("Vector of eigenvalue residuals=",veig_res)
end
############################################
#######################################
function batch_ev2()
  # This script runs the eigenvalue solver for a range of alpha
  #
  lam_tol=1e-2; no_its=5; relax_flag=true; no_restarts=0; p_restart=3;
  valpha=[0.1,0.2,0.5,1.0,1.5,1.8,1.9]
   valpha=[0.5,1.0,1.5]
 #  valpha=[0.25, 0.75, 1.25, 1.75]
  n=length(valpha); vl=zeros(n); vel_time=zeros(n); veig_res=zeros(n)
  for i in 1:n
    l,et,er=run_ev(valpha[i],lam_tol,no_its,relax_flag,no_restarts,p_restart)
    vl[i]=l; vel_time[i]=et; veig_res[i]=er;
  end
  println("Vector of computed lambda=", vl)
  println("Vector of elasped times=", vel_time)
  println("Vector of eigenvalue residuals=",veig_res)
  #######################################
  lam_tol=1e-2; no_its=5; relax_flag=false; no_restarts=0; p_restart=3;
  valpha=[0.1,0.2,0.5,1.0,1.5,1.8,1.9]
   valpha=[0.5,1.0,1.5]
  n=length(valpha); vl2=zeros(n); vel_time2=zeros(n); veig_res2=zeros(n)
  for i in 1:n
    l,et,er=run_ev(valpha[i],lam_tol,no_its,relax_flag,no_restarts,p_restart)
    vl2[i]=l; vel_time2[i]=et; veig_res2[i]=er;
  end
  println("alpha=",valpha)
  println("relax_flag=true")
  println("Vector of computed lambda=", vl)
  println("Vector of elasped times=", vel_time)
  println("Vector of eigenvalue residuals=",veig_res)
  println("relax_flag=false")
  println("Vector of computed lambda=", vl2)
  println("Vector of elasped times=", vel_time2)
  println("Vector of eigenvalue residuals=",veig_res2)
end
#######################################
function batch_field()
  # This script runs the field solver for the following set of alpha
  #
  valpha=[0.1,0.2,0.5,1.0,1.5,1.8,1.9]
  #
  n=length(valpha); vrel_err=zeros(n); vel_time=zeros(n); verr=zeros(n)
  for i in 1:n
    et,er,err_rel=run_field(valpha[i],1,false)
    vel_time[i]=et; verr[i]=er; vrel_err[i]=err_rel
  end
  println("Vector of elasped times=", vel_time)
  println("Vector of errors=",verr)
  println("Vector of rel errors=",vrel_err)
end
############################################
############################################
function run_pt(alpha,x0)
  # This script runs WOS for one point inside the domain
  # p,exact=example2(alpha)
  p,exact=KOS.example2(alpha)
  #
  dict=Dict("max_samples"=>Int32(1e5),
  "no_per_batch"=>16,
  "tol"=>1e-3)
  #
  my_step,my_rhs_ball,dist_fn,ext_fn=setup_run_wos(p,dict)
  #
  mean_val, var_sample, err_est, no_samples = fn_run_wos(x0,
    p, dict, my_step, my_rhs_ball, dist_fn, ext_fn)
  #
  println(mean_val," variance=",  var_sample," error=",
   err_est, " no_samples=",no_samples)
  print("exact=",exact(x0), "; error=",exact(x0)-sum(mean_val))
end
####################################
################ end julia_wos.jl
####################################


export len, siz, siz3, sizem, sizen, sizeo, rangem, rangen, rangeo, range

len(a) = length(a)
len(a::AbstractArray{T,N}) where {T,N} = size(a,N)

siz(a) = ones(Int, 2, 1)

function siz(a::Tuple)
    r = ones(Int, 2, 1)
    r[1] = 1
    r[2] = length(a)
    r
end

function siz(a::Union{UnitRange, Char, AbstractString})
    r = ones(Int, 2,1)
    r[2] = length(a)
    r
end

function siz(a::AbstractArray{T,N}) where {T,N}
    r = ones(Int, max(2,N), 1)
    for i = 1:N
        r[i] = size(a,i)
    end
    r
end

function siz3(a::Array{T,N}) where {T,N}
    if N > 3
        error("siz3 can only be used for arrays with ndims <= 3")
    end
    r = ones(Int, 3, 1)
    for i = 1:N
        r[i] = size(a,i)
    end
    r
end
siz3(a) = siz(a)

sizem(a) = size(a,1)
sizen(a) = size(a,2)
sizeo(a) = size(a,3)
rangem(a) = 1:size(a,1)
rangen(a) = 1:size(a,2)
rangeo(a) = 1:size(a,3)

import Base.range
range(a::AbstractArray) = 1:len(a)

import Base.ndims
ndims(a::Tuple) = 1


""" Upper bound computation (prefix is UB for Upper Bound) """

# All functions are public :

"""
Returns an object `parentToChild` : the already known points that won't need to be computed again.

Points we already know are the ones for which the variable `index` is already set to `value`.

Returns parentToChild
"""
function UB_computeParentToChild(prob::BiOKP, UB::UpperBound, index::Int, value::Int)

    CONFIG.debug && DEBUG_UB(prob, UB)

    # known solutions that go from the parent to the child node
	knownSols = nil(Sol)
	knownSolsLastElt = nil(Sol)

    # indicies of the known solutions inside the parent UB.sols
    previousIndicies = nil(Int)
    previousIndiciesLastElt = nil(Int)

    sols = UB.sols

    lenLB = length(sols)

    # attributes of the ParentToChild object : are lexicographic sols known by the child ?
    leftKnown = false
    rightKnown = false

    # right sol to memorize
    rightSol = Sol()

    iter = 1
    atLeastOneSolKnown = false # length(knownSols) >= 1
	while sols != nil(Sol)

		if !atLeastOneSolKnown && sols.head.x[index] == value

            # length(knownSols) >= 1 so atLeastOneSolKnown becomes true
			atLeastOneSolKnown = true

            # append the known solution to knownSols
			knownSols = cons(sols.head, nil(Sol))
			knownSolsLastElt = knownSols

            # append the index of the known solution to previousIndicies
            previousIndicies = cons(iter, nil(Int))
			previousIndiciesLastElt = previousIndicies

            rightSol = sols.head

            # update leftKnown and rightKnown booleans
            if iter == 1 # the left sol is a known solution for the child
                leftKnown = true
            elseif iter == lenLB
                rightKnown = true
            end

		elseif sols.head.x[index] == value

            # append the known solution to knownSols
			knownSolsLastElt.tail = cons(sols.head, nil(Sol))
			knownSolsLastElt = knownSolsLastElt.tail

            # append the index of the known solution to previousIndicies
            previousIndiciesLastElt.tail = cons(iter, nil(Int))
			previousIndiciesLastElt = previousIndiciesLastElt.tail

            rightSol = sols.head

            # update leftKnown and rightKnown booleans
            if iter == 1 # the left sol is a known solution for the child
                leftKnown = true
            elseif iter == lenLB
                rightKnown = true
            end

		end

		sols = sols.tail
        iter += 1
	end

    parentToChild = ParentToChild(knownSols,previousIndicies,rightSol,leftKnown,rightKnown)

    CONFIG.debug && DEBUG_parentToChild(prob, parentToChild, index=index, value=value)

	return parentToChild
end

"""
Computes the segments associated with the given `UB.sols` as a DualSet
"""
function UB_computeSegments!(prob::BiOKP, UB::UpperBound)

    CONFIG.debug && DEBUG_UB(prob, UB)

    sols = UB.sols

	nbConstraints = length(sols) + 1

	A = zeros(Float64, nbConstraints, 2)
	b = zeros(Float64, nbConstraints)

	max1 = sols.head.y[1]
	max2 = sols.head.y[2]

	iter = 1

	while sols.tail != nil(Sol)

        solL = sols.head
        solR = sols.tail.head

		A[iter, 1] = solL.y[2] - solR.y[2]
		A[iter, 2] = solR.y[1] - solL.y[1]

		b[iter] = (solL.y[2] - solR.y[2]) * solR.y[1] + (solR.y[1] - solL.y[1]) * solR.y[2]

		if solL.y[1] > max1
			max1 = solL.y[1]
		end
		if solR.y[1] > max1
			max1 = solR.y[1]
		end
		if solL.y[2] > max2
			max2 = solL.y[2]
		end
		if solR.y[2] > max2
			max2 = solR.y[2]
		end

		sols = sols.tail
		iter += 1
	end

    # add the constraint extreme horizontal and vertical to finish the polytope

    # vertical (on the right)
	A[nbConstraints-1, 1] = 1
	A[nbConstraints-1, 2] = 0

	b[nbConstraints-1] = max1

    # horizontal (on the left)
	A[nbConstraints, 1] = 0
	A[nbConstraints, 2] = 1

	b[nbConstraints] = max2

    UB_setSegments!(UB, DualSet(A, b)) # update segments

    CONFIG.debug && DEBUG_UB(prob, UB)

end


"""
    ℓ = compute_transition_costs(g, t::AbstractArray)
Computes the transition costs `ℓ` given a
function `g` and a time sequence `t`
"""
function compute_transition_costs(g, t::AbstractVector, tol=1e-3)
    T = Float64
    ℓ_lazy = TransitionCostContinuous{T}(g, t)
    N = length(t)
    ℓ = Array{QuadraticForm{T}}(undef, N-1,N)
    for i = 1:N-1
        for ip = i+1:N
            ℓ[i,ip] = ℓ_lazy[i,ip]
        end
    end
    # TODO move/add tests to Lazy, maybe only i to i+1
    for i = 1:N-1, ip = i+1:N
        P, q, r = ℓ[i,ip].P, ℓ[i,ip].q, ℓ[i,ip].r
        minval = -q⋅(P\q)/4+r
        #TODO add epsilon on RHS?
        minval < 0 && error("Transition cost is negative for some values in `compute_transition_costs`, try decreasing rtol.")
    end

    return ℓ
end

"""
    ℓ = compute_transition_costs(g::AbstractArray; t=1:length(g))
Computes the transition costs `ℓ` given a discrete
function `g`.
"""
function compute_discrete_transition_costs(g::AbstractArray, t=1:length(g))
    T = promote_type(eltype(g), Float64)
    ℓ_lazy = TransitionCostDiscrete{T}(g, t)
    N = length(t)
    ℓ = Array{QuadraticForm{T}}(undef, N-1,N)
    for i = 1:N-1
        for ip = i+1:N
            ℓ[i,ip] = ℓ_lazy[i,ip]
        end
    end
    return ℓ
end


#=
Computes quantities related to the Gaussian linear state space model

    x_{t+1} = A x_t + C w_{t+1}

        y_t = G x_t

The shocks {w_t} are iid and N(0, I)

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-07-28

References
----------

TODO: Come back and update to match `LinearStateSpace` type from py side
TODO: Add docstrings

http://quant-econ.net/jl/linear_models.html

=#
import Distributions: MultivariateNormal, rand
import Base: ==

#=
    numpy allows its multivariate_normal function to have a matrix of
    zeros for the covariance matrix; Stats.jl doesn't. This type just
    gives a `rand` method when we pass in a matrix of zeros for Sigma_0
    so the rest of the api can work, unaffected

    The behavior of `rand` is to just pass back the mean vector when
    the covariance matrix is zero.
=#
type FakeMVTNorm{T <: Real}
    mu_0::Array{T}
    Sigma_0::Array{T}
end

==(f1::FakeMVTNorm, f2::FakeMVTNorm) =
    (f1.mu_0 == f2.mu_0) && (f1.Sigma_0 == f2.Sigma_0)

Base.rand{T}(d::FakeMVTNorm{T}) = copy(d.mu_0)
"""
A type that describes the Gaussian Linear State Space Model
of the form:

    x_{t+1} = A x_t + C w_{t+1}

        y_t = G x_t

where {w_t} and {v_t} are independent and standard normal with dimensions
k and l respectively.  The initial conditions are mu_0 and Sigma_0 for x_0
~ N(mu_0, Sigma_0).  When Sigma_0=0, the draw of x_0 is exactly mu_0.

#### Fields

- `A::Matrix` Part of the state transition equation.  It should be `n x n`
- `C::Matrix` Part of the state transition equation.  It should be `n x m`
- `G::Matrix` Part of the observation equation.  It should be `k x n`
- `k::Int` Dimension
- `n::Int` Dimension
- `m::Int` Dimension
- `mu_0::Vector` This is the mean of initial draw and is of length `n`
- `Sigma_0::Matrix` This is the variance of the initial draw and is `n x n` and
                    also should be positive definite and symmetric

"""
type LSS
    A::Matrix
    C::Matrix
    G::Matrix
    k::Int
    n::Int
    m::Int
    mu_0::Vector
    Sigma_0::Matrix
    dist::Union{MultivariateNormal, FakeMVTNorm}
end

function LSS(A::ScalarOrArray, C::ScalarOrArray, G::ScalarOrArray,
             mu_0::ScalarOrArray,
             Sigma_0::Matrix=zeros(size(G, 2), size(G, 2)))
    k = size(G, 1)
    n = size(G, 2)
    m = size(C, 2)

    # coerce shapes
    A = reshape(vcat(A), n, n)
    C = reshape(vcat(C), n, m)
    G = reshape(vcat(G), k, n)

    mu_0 = reshape([mu_0;], n)

    # define distribution
    if all(Sigma_0 .== 0.0)   # no variance -- no distribution
        dist = FakeMVTNorm(mu_0, Sigma_0)
    else
        dist = MultivariateNormal(mu_0, Sigma_0)
    end
    LSS(A, C, G, k, n, m, mu_0, Sigma_0, dist)
end

# make kwarg version
function LSS(A::ScalarOrArray, C::ScalarOrArray, G::ScalarOrArray;
             mu_0::Vector=zeros(size(G, 2)),
             Sigma_0::Matrix=zeros(size(G, 2), size(G, 2)))
    return LSS(A, C, G, mu_0, Sigma_0)
end

function simulate(lss::LSS, ts_length=100)
    x = Array(Float64, lss.n, ts_length)
    x[:, 1] = rand(lss.dist)
    w = randn(lss.m, ts_length - 1)
    for t=1:ts_length-1
        x[:, t+1] = lss.A * x[:, t] .+ lss.C * w[:, t]
    end
    y = lss.G * x

    return x, y
end

"""
Simulate num_reps observations of x_T and y_T given x_0 ~ N(mu_0, Sigma_0).

#### Arguments

- `lss::LSS` An instance of the Gaussian linear state space model.
- `t::Int = 10` The period that we want to replicate values for.
- `num_reps::Int = 100` The number of replications we want

#### Returns

- `x::Matrix` An n x num_reps matrix, where the j-th column is the j_th
              observation of x_T
- `y::Matrix` An k x num_reps matrix, where the j-th column is the j_th
              observation of y_T

"""
function replicate(lss::LSS, t::Integer, num_reps::Integer=100)
    x = Array(Float64, lss.n, num_reps)
    for j=1:num_reps
        x_t, _ = simulate(lss, t+1)
        x[:, j] = x_t[:, end]
    end

    y = lss.G * x
    return x, y
end

replicate(lss::LSS; t::Integer=10, num_reps::Integer=100) =
    replicate(lss, t, num_reps)

"""
Create a generator to calculate the population mean and
variance-convariance matrix for both x_t and y_t, starting at
the initial condition (self.mu_0, self.Sigma_0).  Each iteration
produces a 4-tuple of items (mu_x, mu_y, Sigma_x, Sigma_y) for
the next period.

#### Arguments

- `lss::LSS` An instance of the Gaussian linear state space model

"""
function moment_sequence(lss::LSS)
    A, C, G = lss.A, lss.C, lss.G
    mu_x, Sigma_x = copy(lss.mu_0), copy(lss.Sigma_0)
    while true
        mu_y, Sigma_y = G * mu_x, G * Sigma_x * G'
        produce((mu_x, mu_y, Sigma_x, Sigma_y))

        # Update moments of x
        mu_x = A * mu_x
        Sigma_x = A * Sigma_x * A' + C * C'
    end
end

"""
Compute the moments of the stationary distributions of x_t and
y_t if possible.  Computation is by iteration, starting from the
initial conditions lss.mu_0 and lss.Sigma_0

#### Arguments

- `lss::LSS` An instance of the Guassian linear state space model
- `;max_iter::Int = 200` The maximum number of iterations allowed
- `;tol::Float64 = 1e-5` The tolerance level one wishes to achieve

#### Returns

- `mu_x::Vector` Represents the stationary mean of x_t
- `mu_y::Vector`Represents the stationary mean of y_t
- `Sigma_x::Matrix` Represents the var-cov matrix
- `Sigma_y::Matrix` Represents the var-cov matrix

"""
function stationary_distributions(lss::LSS; max_iter=200, tol=1e-5)
    # Initialize iteration
    m = @task moment_sequence(lss)
    mu_x, mu_y, Sigma_x, Sigma_y = consume(m)

    i = 0
    err = tol + 1.

    while err > tol
        if i > max_iter
            error("Convergence failed after $i iterations")
        else
            i += 1
            mu_x1, mu_y, Sigma_x1, Sigma_y = consume(m)
            err_mu = Base.maxabs(mu_x1 - mu_x)
            err_Sigma = Base.maxabs(Sigma_x1 - Sigma_x)
            err = max(err_Sigma, err_mu)
            mu_x, Sigma_x = mu_x1, Sigma_x1
        end
    end

    return mu_x, mu_y, Sigma_x, Sigma_y
end

function geometric_sums(lss::LSS, bet, x_t)
    I = eye(lss.n)
    S_x = (I - bet .* lss.A) \ x_t
    S_y = lss.G * S_x
    return S_x, S_y
end


## ROC COMPUTATION

"""
    _idx_unique_sorted(v)

Internal function to return the index of unique elements in `v` under the
assumption that the vector `v` is sorted in decreasing order.
"""
function _idx_unique_sorted(v::Vec{<:Real})
    n    = length(v)
    idx  = ones(Int, n)
    p, h = 1, 1
    cur  = v[1]
    @inbounds while h < n
        h     += 1                  # head position
        cand   = v[h]               # candidate value
        cand   < cur || continue    # is it new? otherwise skip
        p     += 1                  # if new store it
        idx[p] = h
        cur    = cand               # and update the last seen value
    end
    p < n && deleteat!(idx, p+1:n)
    return idx
end

"""
    fprs, tprs, ts = roc_curve(ŷ, y) = roc(ŷ, y)

Return the ROC curve for a two-class probabilistic prediction `ŷ` given the
ground  truth `y`. The true positive rates, false positive rates over a range
of thresholds `ts` are returned. Note that if there are `k` unique scores,
there are correspondingly  `k` thresholds and `k+1` "bins" over which the FPR
and TPR are constant:

* `[0.0 - thresh[1]]`
* `[thresh[1] - thresh[2]]`
* ...
* `[thresh[k] - 1]`

consequently, `tprs` and `fprs` are of length `k+1` if `ts` is of length `k`.

To draw the curve using your favorite plotting backend, do `plot(fprs, tprs)`.
"""
function roc_curve(ŷm, ym)
    ŷ, y    = skipinvalid(ŷm, ym)
    n       = length(y)
    lab_pos = levels(y)[2]
    scores  = pdf.(ŷ, lab_pos)
    ranking = sortperm(scores, rev=true)

    scores_sort = scores[ranking]
    y_sort_bin  = (y[ranking] .== lab_pos)

    idx_unique = _idx_unique_sorted(scores_sort)
    thresholds = scores_sort[idx_unique]

    # detailed computations with example:
    # y = [  1   0   0   1   0   0   1]
    # s = [0.5 0.5 0.2 0.2 0.1 0.1 0.1] thresh are 0.5 0.2, 0.1 // idx [1, 3, 5]
    # ŷ = [  0   0   0   0   0   0   0] (0.5 - 1.0] # no pos pred
    # ŷ = [  1   1   0   0   0   0   0] (0.2 - 0.5] # 2 pos pred
    # ŷ = [  1   1   1   1   0   0   0] (0.1 - 0.2] # 4 pos pred
    # ŷ = [  1   1   1   1   1   1   1] [0.0 - 0.1] # all pos pre

    idx_unique_2 = idx_unique[2:end]   # [3, 5]
    n_ŷ_pos      = idx_unique_2 .- 1   # [2, 4] implicit [0, 2, 4, 7]

    cs   = cumsum(y_sort_bin)          # [1, 1, 1, 2, 2, 2, 3]
    n_tp = cs[n_ŷ_pos]                 # [1, 2] implicit [0, 1, 2, 3]
    n_fp = n_ŷ_pos .- n_tp             # [1, 2] implicit [0, 1, 2, 4]

    # add end points
    P = sum(y_sort_bin) # total number of true positives
    N = n - P           # total number of true negatives

    n_tp = [0, n_tp..., P] # [0, 1, 2, 3]
    n_fp = [0, n_fp..., N] # [0, 1, 2, 4]

    tprs = n_tp ./ P  # [0/3, 1/3, 2/3, 1]
    fprs = n_fp ./ N  # [0/4, 1/4, 2/4, 1]

    return fprs, tprs, thresholds
end

const roc = roc_curve


include("misc.jl")

function leastSquares(X,y)

	# Find regression weights minimizing squared error
	w = (X'*X)\(X'*y)

	# Make linear prediction function
	predict(Xhat) = Xhat*w

	# Return model
	return GenericModel(predict)
end


##
using Plots
using FFTW
using DSP

##
# Number of points
N = 2^14 - 1
# Sample period
Ts = 1 / (1.1 * N)
# Start time
t0 = 0
tmax = t0 + N * Ts
# time coordinate
t = t0:Ts:tmax

# signal
signal = sin.(2π * 60 .* t) # sin (2π f t)

# Fourier Transform of it
F = fft(signal) |> fftshift
w = fftfreq(length(t), 1.0/Ts) |> fftshift

# plots
time_domain = plot(t, signal, title = "Signal")
freq_domain = plot(w, abs.(F), title = "Spectrum", xlim=(-1000, +1000))
plot(time_domain, freq_domain, layout = 2)

##
f(x,y) = x + y

function g(x,y)
    x * y
end


# test FourierOperators
function testFT(N=16)
  # random image
  x = zeros(ComplexF64,N,N)
  for i=1:N,j=1:N
    x[i,j] = rand()
  end

  # FourierMatrix
  idx = CartesianIndices((N,N))[collect(1:N^2)]
  F = [ exp(-2*pi*im*((idx[j][1]-1)*(idx[k][1]-1)+(idx[j][2]-1)*(idx[k][2]-1))/N) for j=1:N^2, k=1:N^2 ]
  F_adj = F'

  # Operators
  tr = CartesianTrajectory(N,N)
  F_nfft = NFFTOp((N,N),tr,symmetrize=false)
  F_exp = ExplicitOp((N,N),tr,zeros(ComplexF64,N,N),symmetrize=false)

  # test agains FourierOperators
  y = vec( ifftshift(reshape(F*vec(fftshift(x)),N,N)) )
  y_adj = vec( ifftshift(reshape(F_adj*vec(fftshift(x)),N,N)) )

  y_nfft = F_nfft*vec(x)
  y_adj_nfft = adjoint(F_nfft) * vec(x)

  y_exp = F_exp*vec(x)
  y_adj_exp = adjoint(F_exp) * vec(x)

  @test (norm(y-y_nfft)/norm(y)) < 1e-2
  @test (norm(y_adj-y_adj_nfft)/norm(y_adj)) < 1e-2
  @test (norm(y-y_exp)/norm(y)) < 1e-2
  @test (norm(y_adj-y_adj_exp)/norm(y_adj)) < 1e-2
end

function testFT3d(N=12)
  # random image
  x = zeros(ComplexF64,N,N,N)
  for i=1:N,j=1:N,k=1:N
    x[i,j,k] = rand()
  end

  # FourierMatrix
  idx = CartesianIndices((N,N,N))[collect(1:N^3)]
  F = [ exp(-2*pi*im*((idx[j][1]-1)*(idx[k][1]-1)+(idx[j][2]-1)*(idx[k][2]-1)+(idx[j][3]-1)*(idx[k][3]-1))/N) for j=1:N^3, k=1:N^3 ]
  F_adj = F'

  # Operators
  tr = CartesianTrajectory3D(N,N,numSlices=N)
  F_nfft = NFFTOp((N,N,N),tr,symmetrize=false)
  F_exp = ExplicitOp((N,N,N),tr,zeros(ComplexF64,N,N,N),symmetrize=false)

  # test agains FourierOperators
  y = vec( ifftshift(reshape(F*vec(fftshift(x)),N,N,N)) )
  y_adj = vec( ifftshift(reshape(F_adj*vec(fftshift(x)),N,N,N)) )

  y_nfft = F_nfft*vec(x)
  y_adj_nfft = adjoint(F_nfft) * vec(x)

  y_exp = F_exp*vec(x)
  y_adj_exp = adjoint(F_exp) * vec(x)

  relError = (norm(y-y_exp)/norm(y))
  println("Relative error ExplicitOp: ", relError)
  @test  relError< 1e-2
  relError = (norm(y_adj-y_adj_exp)/norm(y_adj))
  println("Relative error adj. ExplicitOp: ", relError)
  @test  relError< 1e-2
  relError = (norm(y-y_nfft)/norm(y))
  println("Relative error NFFT: ", relError)
  @test relError < 1e-2
  relError = (norm(y_adj-y_adj_nfft)/norm(y_adj))
  println("Relative error adj. NFFT: ", relError)
  @test  relError< 1e-2
end

# test FieldmapNFFTOp
function testFieldmapFT(N=16)
  # random image
  x = zeros(ComplexF64,N,N)
  for i=1:N,j=1:N
    x[i,j] = rand()
  end

  tr = CartesianTrajectory(N,N;TE=0.0,AQ=0.01)
  times = readoutTimes(tr)
  nodes = kspaceNodes(tr)
  cmap = im*quadraticFieldmap(N,N)[:,:,1]

  # FourierMatrix
  idx = CartesianIndices((N,N))[collect(1:N^2)]
  F = [exp(-2*1im*pi*(nodes[1,k]*(idx[l][1]-size(x,1)/2-1)+nodes[2,k]*(idx[l][2]-size(x,2)/2-1))-cmap[idx[l][1],idx[l][2]]*times[k]) for k=1:size(nodes,2), l=1:length(x)]
  F_adj = F'

  # Operators
  F_nfft = FieldmapNFFTOp((N,N),tr,cmap,symmetrize=false)

  # test agains FourierOperators
  y = F*vec(x)
  y_adj = F_adj*vec(x)

  y_nfft = F_nfft*vec(x)
  y_adj_nfft = adjoint(F_nfft) * vec(x)

  @test (norm(y-y_nfft)/norm(y)) < 1e-2
  @test (norm(y_adj-y_adj_nfft)/norm(y_adj)) < 1e-2
end

function testOperators()
  @testset "Linear Operator" begin
    testFT()
    testFT3d()
    testFieldmapFT()
  end
end

testOperators()


#=
Threshold methods

Running thresholds

=#

export ThresholdMeanP, ThresholdMeanN

#running median

#running mean

mutable struct ThresholdMeanP <: Threshold
    m_k::Float64
    k::Int64
    m_l::Float64
    s_k::Float64
    s_l::Float64
    mean::Float64
    stds::Float64
end

ThresholdMeanP()=ThresholdMeanP(0.0,1,0.0,0.0,0.0,0.0,3.0)

ThresholdMeanP(stds::Float64)=ThresholdMeanP(0.0,1,0.0,0.0,0.0,0.0,stds)

function threshold(t::ThresholdMeanP,sort::Sorting,v,i)

    sort.t.k+=1

    sort.t.m_k=sort.t.m_l+(v[i,sort.id]-sort.t.m_l)/sort.t.k
    sort.t.s_k=sort.t.s_l+(v[i,sort.id]-sort.t.m_l)*(v[i,sort.id]-sort.t.m_k)

    myvar=sort.t.s_k/(sort.t.k-1)

    sort.t.mean+=v[i,sort.id]
    mymean=sort.t.mean/sort.t.k

    sort.thres=mymean+sort.t.stds*sqrt(myvar)

    sort.t.m_l=sort.t.m_k
    sort.t.s_l=sort.t.s_k

    nothing
end

mutable struct ThresholdMeanN <: Threshold
    m_k::Float64
    k::Int64
    m_l::Float64
    s_k::Float64
    s_l::Float64
    mean::Float64
    stds::Float64
end

ThresholdMeanN()=ThresholdMeanN(0.0,1,0.0,0.0,0.0,0.0,3.0)

ThresholdMeanN(stds::Float64)=ThresholdMeanN(0.0,1,0.0,0.0,0.0,0.0,stds)

function threshold(t::ThresholdMeanN,sort::Sorting,v,i)

    sort.t.k+=1

    sort.t.m_k=sort.t.m_l+(v[i,sort.id]-sort.t.m_l)/sort.t.k
    sort.t.s_k=sort.t.s_l+(v[i,sort.id]-sort.t.m_l)*(v[i,sort.id]-sort.t.m_k)

    myvar=sort.t.s_k/(sort.t.k-1)

    sort.t.mean+=v[i,sort.id]
    mymean=sort.t.mean/sort.t.k

    sort.thres=mymean-sort.t.stds*sqrt(myvar)

    sort.t.m_l=sort.t.m_k
    sort.t.s_l=sort.t.s_k

    nothing
end


#####
##### Truncated exponential distribition
#####

function mean(d::Truncated{<:Exponential,Continuous})
    θ = d.untruncated.θ
    l, r = extrema(d)           # l is always finite
    if isfinite(r)
        if abs(l-r) ≤ √eps()*θ
            # linear is a good approximation, just take the midpoint
            middle(l, r)
        else
            Δ = r - l
            R = -Δ/θ
            θ + middle(l, r) + Δ/2*(exp(R)+1)/expm1(R)
        end
    else
        θ + l
    end
end


# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools

function main()
    # Define beam
    λ = LinRange(200,1000,600) # wavelength range [nm]
    θ = [0.] # angle of incidence [degrees]
    beam = PlaneWave(λ, θ)
    # Define layers
    l0 = LayerTMMO(RIdb.air(beam.λ))
    l1 = LayerTMMO(RIdb.dummy(beam.λ,1.45,0.0); type=:OT, d=1/4.)
    l2 = LayerTMMO(RIdb.dummy(beam.λ,3.45,0.0); type=:OT, d=1/4.)
    layers = [l0, l1, l2, l1, l2, l1, l2, l1, l2, l0]
    # Reference wavelenth [nm]
    λ0 = 700.
    # Call main script
    sol = tmm_optics(beam, layers; λ0=λ0, emfflag=true, h=10, pbgflag=true)
    return sol
end

sol = main()

# Plot spectra
plot(Spectrum1D(),
    sol.beam.λ, [sol.Spectra.Rp, sol.Spectra.Tp, 1.0.-(sol.Spectra.Rp.+sol.Spectra.Tp)],
    label=["Reflectance" "Transmittance" "Absorbance"],
    line=([:solid :dash :dashdot]),
    ylims=(0.0,1.0),
    xlims=(sol.beam.λ[1], sol.beam.λ[end]),
);
gui()

# Plot the EMF pattern for normal incidence
plot(EMF2D(),
    sol.beam.λ, sol.Misc.ℓ, log10.(sol.Field.emfp[:,1,:]),
    title=("Log of EMF intensity"),
);
gui()

# plot the refractive index profile
plot(RIprofile(), sol; plotemf=true)
gui()

# plot the photonic dispersion with custom function
plot(PBGDispersion1D(), sol.Bloch)
gui()

plot(PBGDispersion1Dalt(), sol.Bloch)
gui()

plot(PBGDispersion1Dimre(), sol.Bloch)
gui()


using LinearAlgebraicRepresentation
LAR = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL

function stresstest2d()
	function input_collection(data::Array)::LAR.LAR
		assembly = LAR.Struct(data)
		return LAR.struct2lar(assembly)
	end

	V,(_,EV,FV) = LAR.cuboidGrid([4,1],true);
	W,(_,EW,FW) = LAR.cuboidGrid([1,5],true);
	mycircle(r,n) = LAR.circle(r)([n])
	data2d1 = (V,EV)
	data2d2 = LAR.Struct([ LAR.t(2,2), LAR.r(pi/3), LAR.t(-1.5,-2.5), (W,EW) ])
	data2d3 = LAR.Struct([ LAR.t(2,2), mycircle(2.5,160) ])
	data2d4 = LAR.Struct([ LAR.t(3.5,3.5), mycircle(.25,160) ])
	data2d5a = LAR.Struct([ LAR.t(5,3.5), mycircle(.75,160) ])
	data2d5 = LAR.Struct([ LAR.t(5,3.5), mycircle(.5,160) ])
	data2d6 = LAR.Struct([ LAR.t(5,3.5), mycircle(.25,160) ])
	data = [ data2d1, data2d2, data2d3, data2d4, data2d5,  data2d5a, data2d6 ]
	model2d = input_collection( [ LAR.Struct(data), LAR.t(-pi/6,pi/3), LAR.Struct(data) ] )
	V,EV = model2d
	GL.VIEW([ GL.GLLines(V,EV, GL.COLORS[1]) ]);

	model = Lar.Model(V);
	Lar.addModelCells!(model, 1, EV)

	model = Lar.Arrangement.planar_arrangement(model)

	EVs = Lar.FV2EVs(model)
	FVs = Lar.triangulate2D(model)
	GL.VIEW(GL.GLExplode(model.G, FVs, 1.2,1.2,1.2,99,1));
	GL.VIEW(GL.GLExplode(model.G, EVs, 1.2,1.2,1.2,99,1));
	GL.VIEW(GL.GLExplode(model.G, EVs, 1.2,1.2,1.2, 1,1));
end


using FIRLSFilterDesign
using Test

function generate_bands_vec(N; n_duplicates = 0, fs = 1.)
    bands_vec = rand(N)
    bands_vec = vcat(bands_vec, bands_vec[1:n_duplicates])
    sort!(bands_vec)
    bands_vec .-= bands_vec[1]
    bands_vec ./= bands_vec[end]
    bands_vec .*= fs/2
    return bands_vec
end

vec2mat(vec) = @views hcat(vec[1:end-1], vec[2:end])

function all_perms_q(f,W)
    M = 5
    q_odd = FIRLSFilterDesign.get_q(M, f, W, FIRLSFilterDesign.FIR_I())
    q_even = FIRLSFilterDesign.get_q(M, f, W, FIRLSFilterDesign.FIR_II())
    return q_odd, q_even
end

function all_perms_b(f,W,D)
    M = 5
    b_FIR_I     = FIRLSFilterDesign.get_b(M, f, D, W, FIRLSFilterDesign.FIR_I())
    b_FIR_II    = FIRLSFilterDesign.get_b(M, f, D, W, FIRLSFilterDesign.FIR_II())
    b_FIR_III   = FIRLSFilterDesign.get_b(M, f, D, W, FIRLSFilterDesign.FIR_III())
    b_FIR_IV    = FIRLSFilterDesign.get_b(M, f, D, W, FIRLSFilterDesign.FIR_IV())
    return b_FIR_I, b_FIR_II, b_FIR_III, b_FIR_IV
end

function all_perms_firlsdesign(args...; kwargs...)
    order_odd = 10
    order_even = order_odd + 1
    h_FIR_I     = firls_design(order_odd,  args..., false; kwargs...)
    h_FIR_II    = firls_design(order_even, args..., false; kwargs...)
    h_FIR_III   = firls_design(order_odd,  args..., true;  kwargs...)
    h_FIR_IV    = firls_design(order_even, args..., true;  kwargs...)
    return h_FIR_I, h_FIR_II, h_FIR_III, h_FIR_IV
end

function all_perms_toimpulseresponse(a)
    h_FIR_I     = FIRLSFilterDesign._to_impulse_response(a, FIRLSFilterDesign.FIR_I())
    h_FIR_II    = FIRLSFilterDesign._to_impulse_response(a, FIRLSFilterDesign.FIR_II())
    h_FIR_III   = FIRLSFilterDesign._to_impulse_response(a, FIRLSFilterDesign.FIR_III())
    h_FIR_IV    = FIRLSFilterDesign._to_impulse_response(a, FIRLSFilterDesign.FIR_IV())
    return h_FIR_I, h_FIR_II, h_FIR_III, h_FIR_IV
end

# Randomized inputs
fs_rand = 1 + rand() * 10
bands_D_mat_rand = vec2mat(generate_bands_vec(100, n_duplicates = 10, fs = fs_rand))
D_mat_rand = rand(size(bands_D_mat_rand)...)
bands_W_mat_rand = vec2mat(generate_bands_vec(100, n_duplicates = 10, fs = fs_rand))
W_mat_rand = rand(size(bands_W_mat_rand)...)
# Fixed inputs
fs = 1.
bands_D_mat = [0. .25; 0.25 0.25; .25 .5]
D_mat = [2. 1.; 10. 10.; 0. 0.]
bands_W_mat = [0. .25; 0.25 0.25; .25 .5]
W_mat = [1. 2.; 11. 11.; 2. 1.]

hs_rand1 = all_perms_firlsdesign(bands_D_mat_rand, D_mat_rand; fs = fs_rand)
hs_rand2 = all_perms_firlsdesign(bands_D_mat_rand, D_mat_rand, fill(rand(), size(D_mat_rand)); fs = fs_rand)
# hs_rand3 = all_perms_firlsdesign(bands_D_mat_rand, D_mat_rand, bands_D_mat_rand, ones(size(D_mat_rand)); fs = fs_rand)

qs_fixed = all_perms_q(bands_D_mat, W_mat)
qs_fixed_data = [
    [
        1.5
        0.0
        -0.2026423672846756
        0.0
        0.0
        0.0
        -0.02251581858718621
        0.0
        0.0
        0.0
        -0.008105694691387062
   ],
   [
        1.5
        0.0
        -0.2026423672846756
        0.0
        0.0
        0.0
        -0.02251581858718621
        0.0
        0.0
        0.0
        -0.008105694691387062
        0.0
   ]
]

bs_fixed = all_perms_b(bands_D_mat, D_mat, W_mat)
bs_fixed_data = [
    [
        1.0833333333333333
        0.6919896805485019
        0.0
        -0.24427841957880944
        -0.025330295910584485
        0.12128235798585402
    ],
    [
        0.9761255156564024
        0.32778347331995833
       -0.2012359779620491
       -0.16111118889304135
        0.08495418615790015
        0.080981441609415
    ],
    [
        NaN
        0.6919896805485019
        0.7011228412339804
        0.2442784195788095
        0.0
        0.12128235798585399
    ],
    [
        0.40432442716331324
        0.7913393068108028
        0.48582662723339226
        0.06673443948955127
        0.03518917608697089
        0.19550649463397424
    ]
]

hs_fixed = all_perms_firlsdesign(bands_D_mat, D_mat, W_mat)
hs_fixed_data = [
    [
        0.07788241325853798
        -0.0017268410299129148
        -0.08309128924173821
        0.1009958144655533
        0.5217570192015153
        0.7495102634610803
        0.5217570192015153
        0.1009958144655533
        -0.08309128924173821
        -0.0017268410299129148
        0.07788241325853798
    ],
    [
        0.056631088208928274
        0.05686190577259725
        -0.05915886754220385
        -0.034409300873168014
        0.30433805233189765
        0.6880664081058022
        0.6880664081058022
        0.30433805233189765
        -0.034409300873168014
        -0.05915886754220385
        0.05686190577259725
        0.056631088208928274
    ],
    [
        0.10513740773549088
        0.056950501912317146
        0.2320211851531876
        0.4742540923933572
        0.432645024065513
        0.0
       -0.432645024065513
       -0.4742540923933572
       -0.2320211851531876
       -0.056950501912317146
       -0.10513740773549088
    ],
    [
        0.14347815111515985
        0.06421192516629874
        0.12735250873084986
        0.3641022891441058
        0.5103455604394119
        0.24763921096456235
        -0.24763921096456235
        -0.5103455604394119
        -0.3641022891441058
        -0.12735250873084986
        -0.06421192516629874
        -0.14347815111515985
    ]
]

function get_half(h)
    N = length(h)
    M = N÷2
    if iseven(N)
        @views h[1:M], h[end:-1:M+1]
    else
        @views h[1:M+1], h[end:-1:M+1]
    end
end

results = all_perms_toimpulseresponse(rand(5))
halfs = get_half.(results)

all_approx(x,y) = all(@. abs(x - y) < 1e16)

@testset "FIRLSFilterDesign.jl" begin
    # q vector test
    @test all_approx(qs_fixed[1], qs_fixed[2][1:end-1])
    @test all_approx(qs_fixed[1], qs_fixed_data[1]) # odd filter length
    @test all_approx(qs_fixed[2], qs_fixed_data[2]) # even filter length
    # b vector test
    @test all_approx(bs_fixed[1],        bs_fixed_data[1])        # odd  filter length, symmetric
    @test all_approx(bs_fixed[2],        bs_fixed_data[2])        # even filter length, symmetric
    @test all_approx(bs_fixed[3][2:end], bs_fixed_data[3][2:end]) # odd  filter length, antisymmetric
    @test all_approx(bs_fixed[4],        bs_fixed_data[4])        # even filter length, antisymmetric
    # Symmetry of filter coefficients
    @test all(halfs[1][1] .==  halfs[1][2]) # odd  filter length, symmetric
    @test all(halfs[2][1] .==  halfs[2][2]) # even filter length, symmetric
    @test all(halfs[3][1] .== -halfs[3][2]) # odd  filter length, antisymmetric
    @test all(halfs[4][1] .== -halfs[4][2]) # even filter length, antisymmetric
    # Tests that show equivalence between passing no weights and passing all equal weights
    @test all_approx(hs_rand1[1], hs_rand2[1]) # odd  filter length, symmetric
    @test all_approx(hs_rand1[2], hs_rand2[2]) # even filter length, symmetric
    @test all_approx(hs_rand1[3], hs_rand2[3]) # odd  filter length, antisymmetric
    @test all_approx(hs_rand1[4], hs_rand2[4]) # even filter length, antisymmetric
    # Fixed inputs
    @test all_approx(hs_fixed[1], hs_fixed_data[1]) # odd  filter length, symmetric
    @test all_approx(hs_fixed[2], hs_fixed_data[2]) # even filter length, symmetric
    @test all_approx(hs_fixed[3], hs_fixed_data[3]) # odd  filter length, antisymmetric
    @test all_approx(hs_fixed[4], hs_fixed_data[4]) # even filter length, antisymmetric
end


module GeometryReferenceGenericSolver

using LinearAlgebra: norm
using LinearAlgebra: kron, I
using SparseArrays: sparse

using Parametron
using Parametron:setdirty!
using OSQP
using OSQP.MathOptInterfaceOSQP: OSQPSettings

import MathOptInterface
const MOI = MathOptInterface

function _constrain_convex_combination_coefficients(model, coeffs)
    M = length(coeffs)
    for i = 1:M
        @constraint(model, coeffs[i] >= 0)
    end
    s = sum(coeffs)
    @constraint(model, s == 1)
end

function _convex_combination(coeffs, points, m)
    s = @expression coeffs[1] * points[:, 1]
    for i = 2:m
        s = @expression s + coeffs[i] * points[:, i]
    end
    return s
end

function squared_distance(p1, p2, n)
    s = @expression (p1[1] - p2[1])^2
    for i = 2:n
        s = @expression s + (p1[i] - p2[i])^2
    end
    return s
end

struct SolverError{T} <: Exception
    status::T
end

function find_minimum_distance_parametron(convex1::NTuple{M1, PointT}, convex2::NTuple{M2, PointT}) where {M1, M2, PointT}
    model = Model(OSQP.Optimizer(
        max_iter=4000,
        warm_start=false,   # for deterministic results
        polish=true,
        eps_prim_inf=1e-8,
        eps_dual_inf=1e-8,
        eps_abs=1e-8,
        eps_rel=1e-6,
        scaling=0)
    )

    # from https://github.com/tkoolen/Parametron.jl/blob/d824c9f9ff4ba7bc38a9407129504bb8ecca31ee/test/model.jl#L22
    #MOI.set(model.optimizer, OSQPSettings.AdaptiveRhoInterval(), 25) # required for deterministic behavior
    #MOI.set(model.optimizer, OSQPSettings.Verbose(), false)
    #MOI.set(model.optimizer, OSQPSettings.Scaling(), false)
    N = length(PointT)      # TODO works for GeometryType.Point, general trait?
    # points
    c1 = Parameter(model, val=zeros(N, M1))
    c2 = Parameter(model, val=zeros(N, M2))

    # TODO use broadcast machinery to set ]up problem
    # coefficients of convex combination
    a1 = [Variable(model) for _ in 1:M1]
    a2 = [Variable(model) for _ in 1:M2]

    _constrain_convex_combination_coefficients(model, a1)
    _constrain_convex_combination_coefficients(model, a2)

    p1 = _convex_combination(a1, c1, M1)
    p2 = _convex_combination(a2, c2, M2)
    d_sq = squared_distance(p1, p2, N)

    @objective(model, Minimize, d_sq)

    # set parameters
    for i = 1:M1
        c1.val[][:,i] = convex1[i]
    end
    for i = 1:M2
        c2.val[][:,i] = convex2[i]
    end

    setdirty!(c1)
    setdirty!(c2)
    #=data_before_solve = deepcopy(
     unsafe_load_data(
        unsafe_load(
            unsafe_load(model.optimizer.inner.workspace).data))
    )
    =#
    data_before_solve = "syke"
    solve!(model)
    status =  MOI.get(model.optimizer, MOI.TerminationStatus())

    if status != MOI.OPTIMAL
        #throw(SolverError(status))
    end

    # TODO remove after done debugging
    g = value(model, a1[1])

    # extract decision varibles and re-compute expressions
    # TODO avoid expressing convex combinations and norm twice
    v_a1 = value.(model, a1)
    v_a2 = value.(model, a2)

    v_p1 = sum(v_a1 .* convex1)
    v_p2 = sum(v_a2 .* convex2)
    v_d_sq = sum((v_p1 - v_p2).^2)
    return (d_sq = v_d_sq, p1=v_p1, p2=v_p2, model=model,
        model_parts = (
            a1=a1, a2=a2, c1=c1, c2=c2,
        ),
        data_before_solve=data_before_solve
    )
end

"""
produce U from A such that `x' * U * x == x' * A * x` ∀ x
overwrite A in-place
"""
function upper_triangularize_quadratic!(A)
  s = size(A)
  n = first(s)
  @assert s == (n, n)

  # iterate below diagonal
  for i=2:n, j=1:(i-1)
    A[j,i] += A[i,j]
    A[i,j] = 0
  end
  return nothing
end

"""
Make QP matrices for the problem of finding the minimum distance between two convex hulls.
"""
function setup_qp_data(c1, c2)
  # order all a1 and then all a2
  m1 = length(c1)
  m2 = length(c2)
  n = length(c1[1]) # assume all same dimension

  M = [m1, m2]
  M_offset = [0, m1]

  _C = [c1, c2]

  C = zeros(2*n, m1 + m2)
  for point_set_i = 1:2
    for _j = 1:M[point_set_i]
      j = _j + M_offset[point_set_i]
      n_offset = (point_set_i-1) * n
      C[(1:n) .+ n_offset, j] .= _C[point_set_i][_j]
    end
  end

  p1_minus_p2_op = kron([1, -1]', I(n))

  # maps convex combination coefficients to the difference
  p1_minus_p2__Ta = p1_minus_p2_op * C

  QP_P = 2 * (p1_minus_p2__Ta' * p1_minus_p2__Ta)

  n_constraints = (m1 + m2 + length(M))
  n_variables = m1 + m2

  A = zeros(n_constraints, n_variables)
  l = zeros(n_constraints)
  u = zeros(n_constraints)

  A[1:n_variables, 1:n_variables] = I(n_variables)
  l[1:n_variables] .= 0
  u[1:n_variables] .= Inf

  A[n_variables+1, (1:m1) ] .= ones(m1)
  A[n_variables+2, (1:m2) .+ m1 ] .= ones(m2)
  l[(n_variables+1):end] .= 1
  u[(n_variables+1):end] .= 1
  #upper_triangularize_quadratic!(QP_P)
  return (
    A=A, l=l, u=u, P=QP_P, q=zeros(n_variables),
    parts=(
        m1=m1,
        m2=m2,
        p1_minus_p2__Ta=p1_minus_p2__Ta,
        C = C,
    )
  )
end

quad(x) = x' * x

function setup_model(data)
  model = OSQP.Optimizer().inner

  #model = OSQP.Model()
  #MOI.set(model.optimizer, OSQPSettings.Scaling(), false)
  #OSQP.update_settings!(model; scaling=0)

  OSQP.setup!(model;
          P=sparse(data.P),
          q=data.q,
          A=sparse(data.A),
          l=data.l,
          u=data.u,

          verbose=false,
          max_iter=4000,
          warm_start=false,   # for deterministic results
          polish=true,
          eps_prim_inf=1e-8,
          eps_dual_inf=1e-8,
          eps_abs=1e-8,
          eps_rel=1e-6,
          scaling=10,
          )
  model
end

function find_minimum_distance(c1, c2)
  qp_problem_data = setup_qp_data(c1, c2)

  model = setup_model(qp_problem_data)
  result = OSQP.solve!(model)
  coeffs = result.x
  a1 = coeffs[1:qp_problem_data.parts.m1 .+ 0]
  a2 = coeffs[(1:qp_problem_data.parts.m2) .+ qp_problem_data.parts.m1]
  p1 = sum(c1 .* a1)
  p2 = sum(c2 .* a2)

  return (
    d_sq=result.info.obj_val, p1=p1, p2=p2,
    model=model, result=result)
end

end # module


module SSpline

    using LinearAlgebra

    import Base: ==, +, -, *, /
    import Base: convert

    export  spline0, stepspline, nearestneighbor, nspline,
            spline1,linearspline, lspline,
            spline2, quadraticspline, qspline,
            spline3, cubicspline, cspline, cub,
            naturalspline, clampedspline, periodicspline,
            hermitespline, hspline,
            nakspline,
            interpolate, interp, extrapolate, extrap,
            deriv1, grad, slope,
            deriv2, curvature,
            deriv3, curvrate,
            CubicSpline, HermiteSpline, QuadraticSpline, LinearSpline, NearestNeighborSpline

    include("types.jl")
    include("splinefunctions.jl")
    include("interpolation.jl")
    include("extrapolation.jl")
    include("derivatives.jl")
end #module


# Polar heatmaps

Using Plots

x = range(0, 2π, length = 9)
y = 0:4
z = (1:4) .+ (1:8)'
heatmap(x, y, z, projection = :polar)


#=
Here I show the exact correspondence
between my spiking implementation and code by
Litwin-Kumar / Auguste Schulz
For one large random network
and inputs in currents (high freq input still missing)
=#

push!(LOAD_PATH, abspath(@__DIR__,".."))

using LinearAlgebra,Statistics,StatsBase,Distributions
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs
using BenchmarkTools
using ProgressMeter

function onesparsemat(w::Real)
  return sparse(fill(w,(1,1)))
end

##

using Random; Random.seed!(0)

function initweights(Ne, Ni, p, Jee0, Jei0, Jie, Jii)
	""" initialise the connectivity weight matrix based on initial
	E-to-E Jee0 initial weight e to e plastic
	I-to-E Jei0 initial weight i to e plastic
	E-to-I Jie0 constant weight e to i not plastic
	I-to-I Jii constant weight i to i not plastic"""
	Ncells = Ne+Ni

	# initialise weight matrix
	# w[i,j] is the weight from pre- i to postsynaptic j neuron
	weights = zeros(Float64,Ncells,Ncells)
	weights[1:Ne,1:Ne] .= Jee0
	weights[1:Ne,(1+Ne):Ncells] .= Jie
	weights[(1+Ne):Ncells,1:Ne] .= Jei0
	weights[(1+Ne):Ncells,(1+Ne):Ncells] .= Jii
	# set diagonal elements to 0
	# for cc = 1:Ncells
	# 	weights[cc,cc] = 0
	# end
	weights[diagind(weights)] .= 0.0
	# ensure that the connection probability is only p
	weights = weights.*(rand(Ncells,Ncells) .< p)
	return weights
end
function weightpars(;Ne = 4000, Ni = 1000, p = 0.2 )
	"""Ne, Ni number of excitatory, inhibitory neurons
	p initial connection probability"""
	Jee0 = 2.86 #initial weight e to e plastic
	Jei0 = 48.7 #initial weight i to e plastic
	Jie = 1.27 #constant weight e to i not plastic
	Jii = 16.2 #constant weight i to i not plastic
	return Ne,Ni,p,Jee0,Jei0,Jie,Jii
end
initweights() = initweights(weightpars()...)

# this is a bare-bone version of the Schulz code
function runsimulation_static_simplified(Ne::Integer,Ni::Integer,
  weights::Array{Float64,2},spiketimes::Array{<:Integer,2},
  v_start::Vector{Float64},input_constant::Vector{Float64}; dt = 0.1, T = 2000)

		# Naming convention
		# e corresponds to excitatory
		# i corresponds to inhibitory
		# x corresponds to external

		#membrane dynamics
		taue = 30 #e membrane time constant ms
		taui = 20 #i membrane time constant ms
		vreste = -60 #e resting potential mV
		vresti = -70 #i resting potential mV
		vpeak = 20 #cutoff for voltage.  when crossed, record a spike and reset mV
		eifslope = 2 #eif slope parameter mV
		C = 300 #capacitance pF
		erev = 0 #e synapse reversal potential mV
		irev = -75 #i synapse reversal potntial mV
		vth0 = -52 #initial spike voltage threshold mV
		thrchange = false # can be switched off to have vth constant at vth0
		ath = 10 #increase in threshold post spike mV
		tauth = 30 #threshold decay timescale ms
		vresete = vreste #reset potential mV
    vreseti = vresti
		taurefrac = 15 #absolute refractory period ms

		# total number of neurons
		Ncells = Ne+Ni

    # constant input
    in_const = input_constant

		# synaptic kernel
		tauerise = 1 #e synapse rise time
		tauedecay = 6 #e synapse decay time
		tauirise = .5 #i synapse rise time
		tauidecay = 2 #i synapse decay time

		# Initialisation of spike arrays ------------------------------
		# spike count and spike times
		# Scope problems - changing arrays are therefore defined here directly

		totalspikes = zeros(Int,Ncells)
		totsp::Int64 = 0; # total numberof spikes
		spmax::Int64 = size(spiketimes,1); # maximum recorded spikes

		# further arrays for storing inputs and synaptic integration
		forwardInputsE = zeros(Float64,Ncells) #sum of all incoming weights from excitatpry inputs both external and EE and IE
		forwardInputsI = zeros(Float64,Ncells) #sum of all incoming weights from inhibitory inputs both II and EI
		forwardInputsEPrev = zeros(Float64,Ncells) #as above, for previous timestep
		forwardInputsIPrev = zeros(Float64,Ncells)

		xerise = zeros(Float64,Ncells) #auxiliary variables for E/I currents (difference of exponentials)
		xedecay = zeros(Float64,Ncells)
		xirise = zeros(Float64,Ncells)
		xidecay = zeros(Float64,Ncells)

		v = zeros(Float64,Ncells) #membrane voltage

		# initialisation of membrane potentials and poisson inputs
    copy!(v,v_start)
		# for cc = 1:Ncells
		# 	v[cc] = vreset # + (vth0-vreset)*rand()
		# end

		vth = vth0*ones(Float64,Ncells) #adaptive threshold
		lastSpike = -100*ones(Ncells) #last time the neuron spiked

		# ---------------------------------- set up storing avg weights -----------------

		# total number of simulation steps, steps when to normalise, and save
		Nsteps = round(Int,T/dt)

		# true time
		t::Float64 = 0.0

		# bool counter if a neuron has just had a spike
		spiked = zeros(Bool,Ncells)
  v_out = Matrix{Float64}(undef,Ncells,Nsteps)
	@time	for tt = 1:Nsteps

				if mod(tt,Nsteps/100) == 1  #print percent complete
					print("\r",round(Int,100*tt/Nsteps))
				end

				forwardInputsE[:] .= 0.
				forwardInputsI[:] .= 0.
				t = dt*tt

				fill!(spiked,zero(Bool)) # reset spike bool without new memory allocation

				for cc = 1:Ncells
					xerise[cc] += -dt*xerise[cc]/tauerise + forwardInputsEPrev[cc]
					xedecay[cc] += -dt*xedecay[cc]/tauedecay + forwardInputsEPrev[cc]
					xirise[cc] += -dt*xirise[cc]/tauirise + forwardInputsIPrev[cc]
					xidecay[cc] += -dt*xidecay[cc]/tauidecay + forwardInputsIPrev[cc]

					if cc <= Ne # excitatory
						if thrchange
						vth[cc] += dt*(vth0 - vth[cc])/tauth;
						end
					end

					if t > (lastSpike[cc] + taurefrac) #not in refractory period
						# update membrane voltage

						ge = (xedecay[cc] - xerise[cc])/(tauedecay - tauerise);
						gi = (xidecay[cc] - xirise[cc])/(tauidecay - tauirise);

						if cc <= Ne #excitatory neuron (eif), has adaptation
							dv = (vreste - v[cc] + eifslope*exp((v[cc]-vth[cc])/eifslope))/taue +
                  ge*(erev-v[cc])/C + gi*(irev-v[cc])/C +
                  in_const[cc]/taue
							v[cc] += dt*dv
							if v[cc] > vpeak
								spiked[cc] = true
							end
						else
							dv = (vresti - v[cc])/taui +
                 ge*(erev-v[cc])/C + gi*(irev-v[cc])/C +
                  in_const[cc]/taui
							v[cc] += dt*dv
							if v[cc] > vth0
								spiked[cc] = true
							end
						end

						if spiked[cc] #spike occurred
							spiked[cc] = true;
							v[cc] = cc<=Ne ? vresete : vreseti;
							lastSpike[cc] = t;
							totalspikes[cc] += 1;
							totsp += 1;
							if totsp < spmax
								spiketimes[totsp,1] = tt; # time index as a sparse way to save spiketimes
								spiketimes[totsp,2] = cc; # cell id
							elseif totsp == spmax
								spiketimes[totsp,1] = tt; # time index
								spiketimes[totsp,2] = cc; # cell id

								totsp = 0 # reset counter total number of spikes
							end

							if cc <= Ne && thrchange # only change for excitatory cells and when thrchange == true
								vth[cc] = vth0 + ath;
							end

							#loop over synaptic projections
							for dd = 1:Ncells # postsynaptic cells dd  - cc presynaptic cells
								if cc <= Ne #excitatory synapse
									forwardInputsE[dd] += weights[cc,dd];
								else #inhibitory synapse
									forwardInputsI[dd] += weights[cc,dd];
								end
							end

						end #end if(spiked)
					end #end if(not refractory)
        v_out[cc,tt] = v[cc]
			end #end loop over cells

			forwardInputsEPrev = copy(forwardInputsE)
			forwardInputsIPrev = copy(forwardInputsI)

		end # tt loop over time
	print("\r")

	println("simulation finished")

	return totalspikes,v_out

end

##

const Ne = 4000
const Ni = 1000
const Ttot = 60.0
const all_in_e = (60.0-52.5)
const all_in_i = (70.0-53.0)

weights = permutedims(initweights(weightpars(;Ne=Ne,Ni=Ni)...))

v_start = let n=Ne+Ni,
  (vmin,vmax) = (-60.0,-51.0)
  rand(Uniform(vmin,vmax),n)
end
const n_spikes_count  = 100_000
test_spiketimes = fill(-1,(n_spikes_count,2))

totspikes,vtest = let dt=0.1,
  _weights = permutedims(weights),
  input = vcat(fill(all_in_e,Ne),fill(all_in_i,Ni)),
  T=round(Integer,Ttot*1E3)
  runsimulation_static_simplified(Ne,Ni,_weights,test_spiketimes,v_start,
    input;dt=dt, T=T)
end
rates_test = totspikes ./ Ttot

_ = let plt = plot(),nneu=22,
  ts = (1:size(vtest,2)).*0.1E-3
  plot!(plt,ts,vtest[nneu,:];linewidth=2,xlims=(0,1))
end

histogram(rates_test[1:Ne])

##
#=

		#membrane dynamics
		taue = 30 #e membrane time constant ms
		taui = 20 #i membrane time constant ms
		vreste = -60 #e resting potential mV
		vresti = -70 #i resting potential mV
		vpeak = 20 #cutoff for voltage.  when crossed, record a spike and reset mV
		eifslope = 2 #eif slope parameter mV
		C = 300 #capacitance pF
		erev = 0 #e synapse reversal potential mV
		irev = -75 #i synapse reversal potntial mV
		vth0 = -52 #initial spike voltage threshold mV
		thrchange = false # can be switched off to have vth constant at vth0
		ath = 10 #increase in threshold post spike mV
		tauth = 30 #threshold decay timescale ms
		vreset = -70 #reset potential mV
		taurefrac = 50 #absolute refractory period ms
		aw_adapt = 4 #adaptation parameter a nS conductance
		bwfactor = 100
		bw_adapt = bwfactor*0.805 #adaptation parameter b pA current

    # constant input
    in_const = fill(0.0,Ncells)
    in_const[2] = -51.0 - vreste

		# synaptic kernel
		tauerise = 1 #e synapse rise time
		tauedecay = 6 #e synapse decay time
		tauirise = .5 #i synapse rise time
		tauidecay = 2 #i synapse decay time

=#

##
# time in seconds, voltage in mV
const dt = 0.1E-3
myτe = 30E-3 # seconds
myτi = 20E-3 # seconds
τrefr= 15E-3 # refractoriness
vth_e = 20.   # mV
vthexp = -52.0 # actual threshold for spike-generation
vth_i = vthexp
eifslope = 2.0
v_rest_e = -60.0
v_rest_i = -70.0
v_rev_e = 0.0
v_rev_i = -75.0
v_leak_e = v_rest_e
v_leak_i = v_rest_i
v_reset_e = v_rest_e
v_reset_i = v_rest_i
Cap = 300.0 #capacitance mF
in_const_e = all_in_e*Cap/myτe
in_const_i = all_in_i*Cap/myτi

# synaptic kernel
tauerise = 1E-3 #e synapse rise time
tauedecay = 6E-3 #e synapse decay time
taueplus,taueminus  = tauedecay, tauerise
tauirise = 0.5E-3 #i synapse rise time
tauidecay = 2E-3 #i synapse decay time
tauiplus,tauiminus = tauidecay,tauirise

##

nt_e = let sker = S.SKExpDiff(taueplus,taueminus)
  sgen = S.SpikeGenEIF(vthexp,eifslope)
  S.NTLIFConductance(sker,sgen,myτe,Cap,
   vth_e,v_reset_e,v_rest_e,τrefr,v_rev_e)
end
ps_e = S.PSLIFConductance(nt_e,Ne)

nt_i = let sker = S.SKExpDiff(tauiplus,tauiminus)
  sgen = S.SpikeGenNone()
  S.NTLIFConductance(sker,sgen,myτi,Cap,
   vth_i,v_reset_i,v_rest_i,τrefr,v_rev_i)
end
ps_i = S.PSLIFConductance(nt_i,Ni)

# static inputs
ps_input_e = S.PSSimpleInput(S.InputSimpleOffset(in_const_e))
ps_input_i = S.PSSimpleInput(S.InputSimpleOffset(in_const_i))

## connections
conn_ee = let w_ee = weights[1:Ne,1:Ne]
  S.ConnGeneralIF2(sparse(w_ee))
end
conn_ii = let w_ii = weights[Ne+1:end,Ne+1:end]
  S.ConnGeneralIF2(sparse(w_ii))
end
conn_ei = let w_ei = weights[1:Ne,Ne+1:end]
  S.ConnGeneralIF2(sparse(w_ei))
end
conn_ie = let w_ie = weights[Ne+1:end,1:Ne]
  S.ConnGeneralIF2(sparse(w_ie))
end

# input populations
pop_e = S.Population(ps_e,(S.FakeConnection(),ps_input_e),
   (conn_ee,ps_e),(conn_ei,ps_i) )
pop_i = S.Population(ps_i,(S.FakeConnection(),ps_input_i),
   (conn_ii,ps_i),(conn_ie,ps_e) )

##
# that's it, let's make the network
myntw = S.RecurrentNetwork(dt,pop_e,pop_i)

# record spiketimes and internal potential
krec = 1
n_e_rec = 1000
n_i_rec = 1000
t_wup = 0.0
rec_state_e = S.RecStateNow(ps_e,krec,dt,Ttot;idx_save=collect(1:n_e_rec),t_warmup=t_wup)
rec_state_i = S.RecStateNow(ps_i,krec,dt,Ttot;idx_save=collect(1:n_i_rec),t_warmup=t_wup)
rec_spikes_e = S.RecSpikes(ps_e,5.0,Ttot;idx_save=collect(1:n_e_rec),t_warmup=t_wup)
rec_spikes_i = S.RecSpikes(ps_i,5.0,Ttot;idx_save=collect(1:n_i_rec),t_warmup=t_wup)

## Run

times = (0:myntw.dt:Ttot)
nt = length(times)
# clean up
S.reset!.([rec_state_e,rec_spikes_e])
S.reset!.([rec_state_i,rec_spikes_i])
S.reset!.([ps_e,ps_i])
S.reset!(conn_ei)
# initial conditions
ps_e.state_now .= v_start[1:Ne]
ps_i.state_now .= v_start[Ne+1:end]

@time begin
  @showprogress 5.0 "network simulation " for (k,t) in enumerate(times)
    rec_state_e(t,k,myntw)
    rec_state_i(t,k,myntw)
    rec_spikes_e(t,k,myntw)
    rec_spikes_i(t,k,myntw)
    S.dynamics_step!(t,myntw)
  end
end

S.add_fake_spikes!(1.0vth_e,rec_state_e,rec_spikes_e)
S.add_fake_spikes!(0.0,rec_state_i,rec_spikes_i)
##

rates_e = let rdic=S.get_mean_rates(rec_spikes_e,dt,Ttot)
  ret = fill(0.0,n_i_rec)
  for (k,v) in pairs(rdic)
    ret[k] = v
  end
  ret
end
rates_i = let rdic=S.get_mean_rates(rec_spikes_i,dt,Ttot)
  ret = fill(0.0,n_i_rec)
  for (k,v) in pairs(rdic)
    ret[k] = v
  end
  ret
end

_ = let plt=plot(;leg=false),
  netest = n_e_rec
  scatter!(rates_e,rates_test[1:netest];ratio=1,
    xlabel="Dylan's model",ylabel="Suchulz et al",title="E rates")
  plot!(plt,identity; linewidth=2)
end

_ = let plt=plot(;leg=false),
  ntest = n_i_rec
  scatter!(rates_i,rates_test[Ne+1:Ne+ntest];ratio=1,
    xlabel="Dylan's model",ylabel="Suchulz et al",title="I rates")
  plot!(plt,identity;linewidth=2)
end

##

_ = let neu = 202,
  plt=plot()
  plot!(plt,rec_state_e.times,rec_state_e.state_now[neu,:];linewidth=2,leg=false,
    xlims=(0,1),ylims=(-70,-40))
  ts = (1:size(vtest,2)).*0.1E-3
  plot!(plt,ts,vtest[neu,:];linewidth=2,linestyle=:dash)
end
_ = let plt=plot(),neu=16
  plot!(plt,rec_state_i.times,rec_state_i.state_now[neu,:];linewidth=2,leg=false,
    xlims=(0,1))
  ts = (1:size(vtest,2)).*0.1E-3
  plot!(plt,ts,vtest[Ne+neu,:];linewidth=2,linestyle=:dash)
end


##
# Abstract ontology.
abstract type AbstractOntology end

##
# Abstract propositional ontology.
abstract type AbstractPropositionalOntology <: AbstractOntology end
# Abstract modal ontology.
abstract type AbstractModalOntology         <: AbstractOntology end

##
# Abstract temporal ontology.
abstract type AbstractTemporalOntology          <: AbstractModalOntology end
# Abstract spatial ontology.
abstract type AbstractSpatialOntology           <: AbstractModalOntology end
# Abstract spatio-temporal ontology.
abstract type AbstractSpatioTemporalOntology    <: AbstractModalOntology end

##
# Propositional state.
struct PropositionalState <: AbstractPropositionalOntology end

##
# (Temporal) Point ontology.
struct Point{T<:Int} <: AbstractTemporalOntology
    x::T
end

const point(x) = Point(x)

# (Temporal) Interval ontology.
struct Interval{T<:Int} <: AbstractTemporalOntology
    x::Point{T}     # start point
    y::Point{T}     # end point
end

const interval(x,y) = Interval(point(x), point(y))

# (Spatial) Rectangle ontology.
struct Rectangle{T<:Int} <: AbstractSpatialOntology
    x::Interval{T}  # first spatial dimension
    y::Interval{T}  # second spatial dimension
end

const rectangle((xx,xy),(yx,yy)) = Rectangle(interval(xx,xy), interval(yx,yy))

# (Spatio-temporal) Cube ontology.
struct Cube{S<:Int,T<:Int} <: AbstractSpatioTemporalOntology
    x::Interval{S}  # first spatial dimension
    y::Interval{S}  # second spatial dimension
    z::Interval{T}  # temporal dimension
end

# const cube((xx,xy),(yx,yy),(x,y)) = Cube(rectangle((xx,xy),(yx,yy)),interval(x,y))


"""
    emd(μ, ν, C; kwargs...)

Compute the optimal transport plan for the Monge-Kantorovich problem with source and target
marginals `μ` and `ν` and cost matrix `C` of size `(length(μ), length(ν))`.

The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle.
```

This function is a wrapper of the function
[`emd`](https://pythonot.github.io/all.html#ot.emd) in the Python Optimal Transport
package. Keyword arguments are listed in the documentation of the Python function.

# Examples

```jldoctest
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> emd(μ, ν, C)
3×2 Matrix{Float64}:
 0.0  0.5
 0.0  0.2
 0.0  0.3
```

See also: [`emd2`](@ref)
"""
function emd(μ, ν, C; kwargs...)
    return pot.lp.emd(μ, ν, PyCall.PyReverseDims(permutedims(C)); kwargs...)
end

"""
    emd2(μ, ν, C; kwargs...)

Compute the optimal transport cost for the Monge-Kantorovich problem with source and target
marginals `μ` and `ν` and cost matrix `C` of size `(length(μ), length(ν))`.

The optimal transport cost is the scalar value
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle.
```

This function is a wrapper of the function
[`emd2`](https://pythonot.github.io/all.html#ot.emd2) in the Python Optimal Transport
package. Keyword arguments are listed in the documentation of the Python function.

# Examples

```jldoctest
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> emd2(μ, ν, C)
0.95
```

See also: [`emd`](@ref)
"""
function emd2(μ, ν, C; kwargs...)
    return pot.lp.emd2(μ, ν, PyCall.PyReverseDims(permutedims(C)); kwargs...)
end

"""
    emd_1d(xsource, xtarget; kwargs...)

Compute the optimal transport plan for the Monge-Kantorovich problem with univariate
discrete measures with support `xsource` and `xtarget` as source and target marginals.

This function is a wrapper of the function
[`emd_1d`](https://pythonot.github.io/all.html#ot.emd_1d) in the Python Optimal Transport
package. Keyword arguments are listed in the documentation of the Python function.

# Examples

```jldoctest
julia> xsource = [0.2, 0.5];

julia> xtarget = [0.8, 0.3];

julia> emd_1d(xsource, xtarget)
2×2 Matrix{Float64}:
 0.0  0.5
 0.5  0.0

julia> histogram_source = [0.8, 0.2];

julia> histogram_target = [0.7, 0.3];

julia> emd_1d(xsource, xtarget; a=histogram_source, b=histogram_target)
2×2 Matrix{Float64}:
 0.5  0.3
 0.2  0.0
```

See also: [`emd`](@ref), [`emd2_1d`](@ref)
"""
function emd_1d(xsource, xtarget; kwargs...)
    return pot.lp.emd_1d(xsource, xtarget; kwargs...)
end

"""
    emd2_1d(xsource, xtarget; kwargs...)

Compute the optimal transport cost for the Monge-Kantorovich problem with univariate
discrete measures with support `xsource` and `xtarget` as source and target marginals.

This function is a wrapper of the function
[`emd2_1d`](https://pythonot.github.io/all.html#ot.emd2_1d) in the Python Optimal Transport
package. Keyword arguments are listed in the documentation of the Python function.

# Examples

```jldoctest
julia> xsource = [0.2, 0.5];

julia> xtarget = [0.8, 0.3];

julia> round(emd2_1d(xsource, xtarget); sigdigits=6)
0.05

julia> histogram_source = [0.8, 0.2];

julia> histogram_target = [0.7, 0.3];

julia> round(emd2_1d(xsource, xtarget; a=histogram_source, b=histogram_target); sigdigits=6)
0.201
```

See also: [`emd2`](@ref), [`emd2_1d`](@ref)
"""
function emd2_1d(xsource, xtarget; kwargs...)
    return pot.lp.emd2_1d(xsource, xtarget; kwargs...)
end

"""
    sinkhorn(μ, ν, C, ε; kwargs...)

Compute the optimal transport plan for the entropic regularization optimal transport problem
with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, and entropic regularization parameter `ε`.

The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma),
```
where ``\\Omega(\\gamma) = \\sum_{i,j} \\gamma_{i,j} \\log \\gamma_{i,j}`` is the entropic
regularization term.

This function is a wrapper of the function
[`sinkhorn`](https://pythonot.github.io/all.html#ot.sinkhorn) in the Python Optimal
Transport package. Keyword arguments are listed in the documentation of the Python function.

# Examples

```jldoctest sinkhorn
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> sinkhorn(μ, ν, C, 0.01)
3×2 Matrix{Float64}:
 0.0  0.5
 0.0  0.2
 0.0  0.3
```

It is possible to provide multiple target marginals as columns of a matrix. In this case the
optimal transport costs are returned:

```jldoctest sinkhorn
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn(μ, ν, C, 0.01); sigdigits=6)
2-element Vector{Float64}:
 0.95
 0.45
```

See also: [`sinkhorn2`](@ref)
"""
function sinkhorn(μ, ν, C, ε; kwargs...)
    return pot.sinkhorn(μ, ν, PyCall.PyReverseDims(permutedims(C)), ε; kwargs...)
end

"""
    sinkhorn2(μ, ν, C, ε; kwargs...)

Compute the optimal transport cost for the entropic regularization optimal transport problem
with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, and entropic regularization parameter `ε`.

The optimal transport cost is the scalar value
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma),
```
where ``\\Omega(\\gamma) = \\sum_{i,j} \\gamma_{i,j} \\log \\gamma_{i,j}`` is the entropic
regularization term.

This function is a wrapper of the function
[`sinkhorn2`](https://pythonot.github.io/all.html#ot.sinkhorn2) in the Python Optimal
Transport package. Keyword arguments are listed in the documentation of the Python function.

# Examples

```jldoctest sinkhorn2
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> round.(sinkhorn2(μ, ν, C, 0.01); sigdigits=6)
1-element Vector{Float64}:
 0.95
```

It is possible to provide multiple target marginals as columns of a matrix.

```jldoctest sinkhorn2
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn2(μ, ν, C, 0.01); sigdigits=6)
2-element Vector{Float64}:
 0.95
 0.45
```

See also: [`sinkhorn`](@ref)
"""
function sinkhorn2(μ, ν, C, ε; kwargs...)
    return pot.sinkhorn2(μ, ν, PyCall.PyReverseDims(permutedims(C)), ε; kwargs...)
end

"""
    empirical_sinkhorn_divergence(xsource, xtarget, ε; kwargs...)

Compute the Sinkhorn divergence from empirical data, where `xsource` and `xtarget` are
arrays representing samples in the source domain and target domain, respectively, and  `ε`
is the regularization term.

This function is a wrapper of the function
[`ot.bregman.empirical_sinkhorn_divergence`](https://pythonot.github.io/gen_modules/ot.bregman.html#ot.bregman.empirical_sinkhorn_divergence)
in the Python Optimal Transport package. Keyword arguments are listed in the documentation of the Python function.

# Examples

```jldoctest
julia> xsource = [1];

julia> xtarget = [2, 3];

julia> ε = 0.01;

julia> empirical_sinkhorn_divergence(xsource, xtarget, ε) ≈
       sinkhorn2([1], [0.5, 0.5], [1 4], ε) -
       (
           sinkhorn2([1], [1], zeros(1, 1), ε) +
           sinkhorn2([0.5, 0.5], [0.5, 0.5], [0 1; 1 0], ε)
       ) / 2
true
```

See also: [`sinkhorn2`](@ref)
"""
function empirical_sinkhorn_divergence(xsource, xtarget, ε; kwargs...)
    return pot.bregman.empirical_sinkhorn_divergence(
        reshape(xsource, Val(2)), reshape(xtarget, Val(2)), ε; kwargs...
    )
end

"""
    sinkhorn_unbalanced(μ, ν, C, ε, λ; kwargs...)

Compute the optimal transport plan for the unbalanced entropic regularization optimal
transport problem with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, entropic regularization parameter `ε`, and marginal relaxation
term `λ`.

The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma)
+ \\lambda \\mathrm{KL}(\\gamma 1, \\mu)
+ \\lambda \\mathrm{KL}(\\gamma^{\\mathsf{T}} 1, \\nu),
```
where ``\\Omega(\\gamma) = \\sum_{i,j} \\gamma_{i,j} \\log \\gamma_{i,j}`` is the entropic
regularization term and ``\\mathrm{KL}`` is the Kullback-Leibler divergence.

This function is a wrapper of the function
[`sinkhorn_unbalanced`](https://pythonot.github.io/all.html#ot.sinkhorn_unbalanced) in the
Python Optimal Transport package. Keyword arguments are listed in the documentation of the
Python function.

# Examples

```jldoctest sinkhorn_unbalanced
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> sinkhorn_unbalanced(μ, ν, C, 0.01, 1_000)
3×2 Matrix{Float64}:
 0.0  0.499964
 0.0  0.200188
 0.0  0.29983
```

It is possible to provide multiple target marginals as columns of a matrix. In this case the
optimal transport costs are returned:

```jldoctest sinkhorn_unbalanced
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn_unbalanced(μ, ν, C, 0.01, 1_000); sigdigits=6)
2-element Vector{Float64}:
 0.949709
 0.449411
```

See also: [`sinkhorn_unbalanced2`](@ref)
"""
function sinkhorn_unbalanced(μ, ν, C, ε, λ; kwargs...)
    return pot.sinkhorn_unbalanced(
        μ, ν, PyCall.PyReverseDims(permutedims(C)), ε, λ; kwargs...
    )
end

"""
    sinkhorn_unbalanced2(μ, ν, C, ε, λ; kwargs...)

Compute the optimal transport cost for the unbalanced entropic regularization optimal
transport problem with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, entropic regularization parameter `ε`, and marginal relaxation
term `λ`.

The optimal transport cost is the scalar value
```math
\\inf_{\\gamma} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma)
+ \\lambda \\mathrm{KL}(\\gamma 1, \\mu)
+ \\lambda \\mathrm{KL}(\\gamma^{\\mathsf{T}} 1, \\nu),
```
where ``\\Omega(\\gamma) = \\sum_{i,j} \\gamma_{i,j} \\log \\gamma_{i,j}`` is the entropic
regularization term and ``\\mathrm{KL}`` is the Kullback-Leibler divergence.

This function is a wrapper of the function
[`sinkhorn_unbalanced2`](https://pythonot.github.io/all.html#ot.sinkhorn_unbalanced2) in the
Python Optimal Transport package. Keyword arguments are listed in the documentation of the
Python function.

# Examples

```jldoctest sinkhorn_unbalanced2
julia> μ = [0.5, 0.2, 0.3];

julia> ν = [0.0, 1.0];

julia> C = [0.0 1.0; 2.0 0.0; 0.5 1.5];

julia> round.(sinkhorn_unbalanced2(μ, ν, C, 0.01, 1_000); sigdigits=6)
1-element Vector{Float64}:
 0.949709
```

It is possible to provide multiple target marginals as columns of a matrix:

```jldoctest sinkhorn_unbalanced2
julia> ν = [0.0 0.5; 1.0 0.5];

julia> round.(sinkhorn_unbalanced2(μ, ν, C, 0.01, 1_000); sigdigits=6)
2-element Vector{Float64}:
 0.949709
 0.449411
```

See also: [`sinkhorn_unbalanced`](@ref)
"""
function sinkhorn_unbalanced2(μ, ν, C, ε, λ; kwargs...)
    return pot.sinkhorn_unbalanced2(
        μ, ν, PyCall.PyReverseDims(permutedims(C)), ε, λ; kwargs...
    )
end

"""
    barycenter(A, C, ε; kwargs...)

Compute the entropically regularized Wasserstein barycenter with histograms `A`, cost matrix
`C`, and entropic regularization parameter `ε`.

The Wasserstein barycenter is a histogram and solves
```math
\\inf_{a} \\sum_{i} W_{\\varepsilon,C}(a, a_i),
```
where the histograms ``a_i`` are columns of matrix `A` and ``W_{\\varepsilon,C}(a, a_i)``
is the optimal transport cost for the entropically regularized optimal transport problem
with marginals ``a`` and ``a_i``, cost matrix ``C``, and entropic regularization parameter
``\\varepsilon``. Optionally, weights of the histograms ``a_i`` can be provided with the
keyword argument `weights`.

This function is a wrapper of the function
[`barycenter`](https://pythonot.github.io/all.html#ot.barycenter) in the
Python Optimal Transport package. Keyword arguments are listed in the documentation of the
Python function.

# Examples

```jldoctest
julia> A = rand(10, 3);

julia> A ./= sum(A; dims=1);

julia> C = rand(10, 10);

julia> isapprox(sum(barycenter(A, C, 0.01; method="sinkhorn_stabilized")), 1; atol=1e-4)
true
```
"""
barycenter(A, C, ε; kwargs...) = pot.barycenter(A, C, ε; kwargs...)

"""
    barycenter_unbalanced(A, C, ε, λ; kwargs...)

Compute the entropically regularized unbalanced Wasserstein barycenter with histograms `A`, cost matrix
`C`, entropic regularization parameter `ε` and marginal relaxation parameter `λ`.

The Wasserstein barycenter is a histogram and solves
```math
\\inf_{a} \\sum_{i} W_{\\varepsilon,C,\\lambda}(a, a_i),
```
where the histograms ``a_i`` are columns of matrix `A` and ``W_{\\varepsilon,C,\\lambda}(a, a_i)}``
is the optimal transport cost for the entropically regularized optimal transport problem
with marginals ``a`` and ``a_i``, cost matrix ``C``, entropic regularization parameter
``\\varepsilon`` and marginal relaxation parameter ``\\lambda``. Optionally, weights of the histograms ``a_i`` can be provided with the
keyword argument `weights`.

This function is a wrapper of the function
[`barycenter_unbalanced`](https://pythonot.github.io/gen_modules/ot.unbalanced.html#ot.unbalanced.barycenter_unbalanced) in the
Python Optimal Transport package. Keyword arguments are listed in the documentation of the
Python function.

# Examples

```jldoctest
julia> A = rand(10, 3);

julia> A ./= sum(A; dims=1);

julia> C = rand(10, 10);

julia> isapprox(sum(barycenter_unbalanced(A, C, 0.01, 1; method="sinkhorn_stabilized")), 1; atol=1e-4)
false

julia> isapprox(sum(barycenter_unbalanced(
           A, C, 0.01, 10_000; method="sinkhorn_stabilized", numItermax=5_000
       )), 1; atol=1e-4)
true
```

See also: [`barycenter`](@ref)
"""
function barycenter_unbalanced(A, C, ε, λ; kwargs...)
    return pot.barycenter_unbalanced(A, C, ε, λ; kwargs...)
end


"""
    slerp(p1, p2, t)

The slerp interpolation at `t` between two vectors.
"""
function slerp(p1,p2,t)
    0 <= t <= 1 || throw(DomainError("Interpolation parameter should be between 0 and 1"))
    m = (1-t)*norm(p1) + t*norm(p2)
    p1 = normalize(p1)
    p2 = normalize(p2)
    d = p1'p2
    if d > 0.99
        return m*((1-t)*p1 + t*p2)
    end
    Ω  = acos(d)
    m*((sin((1-t)*Ω)/sin(Ω))*p1 + (sin(t*Ω)/sin(Ω))*p2)
end

# function sinkhorn_interpolate(a1,c1,a2,c2,ai,ci,α)
#     error("Update to new interface")
#     α == 0 && return a1,c1,0
#     α == 1 && return a2,c2,0
#     # ai = (1-α)*a1 + α*a2
#     # ci  = (1-α)*c1 + α*c2
#     distmat = loss_spectral_ot(a1,c1,a2,c2)[2]
#     function loss(r)
#         ar,cr = r[1:length(ai)], r[length(ai)+1:end]
#         ar[1] = 1
#         (1-α)*loss_spectral_ot!(distmat,a1,c1,ar,cr) +  α*loss_spectral_ot!(distmat,a2,c2,ar,cr)
#     end
#     res = Optim.optimize(loss, [ai;ci], BFGS(), Optim.Options(store_trace=false, show_trace=true, iterations=500, allow_f_increases=false, time_limit=20, x_tol=1e-6, f_tol=1e-6), inplace=false, autodiff=:forward)
#
#     ai,ci = res.minimizer[1:length(ai)], res.minimizer[length(ai)+1:end]
#     ai[1] = 1
#     ai, ci, res
# end

function centraldiff(v::AbstractVector)
    dv = diff(v)./2
    a1 = [dv[1];dv]
    a2 = [dv;dv[end]]
    a = a1+a2
end

"""
    $(SIGNATURES)

Perform displacement interpolation between two models.
"""
function interpolator(d::RationalOptimalTransportDistance,A1,A2)
    @assert d.p == 2 "Interpolation only supported for p=2, you have p=$(d.p)"
    interval   = (0., d.interval[2])
    @assert (interval[1] == 0 || interval[1] == -interval[2])
    e1   = sqrt(2)/sqrt(spectralenergy(domain(d), A1))
    e2   = sqrt(2)/sqrt(spectralenergy(domain(d), A2))
    f1    = w -> evalfr(domain(d), magnitude(d), w, A1, e1)
    f2    = w -> evalfr(domain(d), magnitude(d), w, A2, e2)
    sol1  = c∫(f1,interval...)
    sol2  = c∫(f2,interval...)
    σ1    = sol1(interval[2]) # The total energy in the spectrum
    σ2    = sol2(interval[2]) # The total energy in the spectrum
    F1(w) = sol1(w)/σ1
    F2(w) = sol2(w)/σ2
    iF1   = inv(F1, interval)
    iF2   = inv(F2, interval)
    function (w,t)
        t == 0 && return map(f1,w)
        t == 1 && return map(f2,w)
        r = tmap(w) do w
            finv(e->(1-t)*iF1(e) + t*iF2(e), w, (0,max(σ1, σ2)))
        end
        centraldiff(r) ./ centraldiff(w)
    end
end

function interpolator(d::EuclideanRootDistance,A1,A2; normalize=false)
    p = d.p
    @assert p == 2 "Interpolation only supported for p=2, you have p=$p"
    RT = domain(d) isa Continuous ? ContinuousRoots : DiscreteRoots
    e1,e2 = roots(domain(d), A1), roots(domain(d), A2)
    I1,I2 = d.assignment(e1, e2)
    e1,e2 = RT(e1[I1]),RT(e2[I2])
    w1,w2 = d.weight(e1), d.weight(e2)
    function (w,t)
        A = AR(RT(((1-t)*w1.*e1 + t*w2.*e2)./((1-t).*w1.+t.*w2)))
        e = normalize ? 1/sqrt(spectralenergy(domain(d), A)) : 1
        tmap(w) do w
            evalfr(domain(d), magnitude(d), w, A, e)
        end
    end
end


#!/usr/bin/env julia

NUM_TRIALS = 10000
PRIOR_LOWER_MAX = 100

# Runs a single trial where an envelope is chosen.  If the chosen envelope has
# a value < cutoff, the function will switch envelopes, otherwise it will keep
# the envelope it has chosen. Returns the value of the envelope it ultimately
# selects.
function single_trial(cutoff)
  lower_value = rand() * PRIOR_LOWER_MAX
  higher_value = 2 * lower_value
  if rand(0:1) == 0
    return lower_value >= cutoff ? lower_value : higher_value
  else
    return higher_value >= cutoff ? higher_value : lower_value
  end
end

# Runs many trials at a given cutoff to approximate the expected value.
function multi_trial(cutoff)
  total_result = 0
  for _ = 1:NUM_TRIALS
    total_result = total_result + single_trial(cutoff)
  end
  return total_result / NUM_TRIALS
end

for cutoff = 1:2*PRIOR_LOWER_MAX
  println("cutoff=", cutoff, " expected_value=", multi_trial(cutoff))
end


using Distributions, StatsBase
export Hollywood

struct Hollywood
    α::Real
    θ::Real
    ν::DiscreteUnivariateDistribution
    K::Real
    function Hollywood(α::Real, θ::Real, ν::DiscreteUnivariateDistribution, K::Real)
        # @show θ, -K*α
        @assert ((0 < α < 1.0) & (θ > -α)) | ((α < 0) & (θ == (- K * α))) "Check parameters satisfy required constraints"
        new(α, θ, ν, K)
    end
end

Hollywood(α::Real, θ::Real, ν::DiscreteUnivariateDistribution) = Hollywood(α, θ, ν, Inf)
Hollywood(α::Real, ν::DiscreteUnivariateDistribution, K::Int) = Hollywood(α, (-K)*α, ν, K)

function StatsBase.sample!(
    out::InteractionSequence,
    model::Hollywood
)
    # Init
    t = 0
    V = 0
    counts = Int[]
    prob_new(t, V, θ, α) = (θ + α * V) / (t - 1 + θ)
    for i in 1:length(out)
        m = rand(model.ν)
        I_tmp = zeros(Int, m)
        for j in 1:m
            t += 1
            if t==1
                I_tmp[j] = 1
                V += 1
                push!(counts, 1)
            elseif rand() < prob_new(t, V, model.θ, model.α)
                I_tmp[j] = V+1
                V+=1
                push!(counts, 1)
            else
                μ = cumsum(counts) / sum(counts)
                pushfirst!(μ, 0.0)
                # @show μ
                val, p = rand_multivariate_bernoulli(μ)
                I_tmp[j] = val
                counts[val] += 1
            end
            out[i] = I_tmp
        end
    end
end

function StatsBase.sample(
    model::Hollywood,
    n::Int
    )
    out = fill(Int[], n)
    sample!(out, model)
    return out
end


mutable struct Impact{T} <: Contact{T}
    Nx::Adjoint{T,SVector{3,T}}
    p::SVector{3,T}
    offset::SVector{6,T}

    function Impact(body::Body{T}, normal::AbstractVector; p::AbstractVector = zeros(3),  offset::AbstractVector = zeros(3)) where T
        # Derived from plane equation a*v1 + b*v2 + distance*v3 = p - offset
        V1, V2, V3 = orthogonalcols(normal) # gives two plane vectors and the original normal axis
        A = [V1 V2 V3]
        Ainv = inv(A)
        ainv3 = Ainv[3,SA[1; 2; 3]]
        Nx = ainv3'
        offset = [offset;0.0;0.0;0.0]

        new{T}(Nx, p, offset), body.id
    end
end

function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, constraint::Impact{T}) where {T}
    summary(io, constraint)
    println(io,"")
    println(io, " Nx:     "*string(constraint.Nx))
    println(io, " offset: "*string(constraint.offset))
end


export NALU

"""
    NALU(in::Int, out::Int; initNAC=glorot_uniform, initG=glorot_uniform, initb=glorot_uniform)

Neural Arithmetic Logic Unit. Layer that is capable of learing multiplication,
division, power functions, addition, and subtraction.
As proposed in: https://arxiv.org/abs/1808.00508
"""
struct NALU{Tg<:AbstractMatrix,Tb<:AbstractVector,Tϵ<:Real}
    nac::NAC
    G::Tg
    b::Tb
    ϵ::Tϵ
end

NALU(nac::NAC, G::AbstractMatrix, b::AbstractVector) = NALU(nac, G, b, 1e-8)

function NALU(in::Int, out::Int;
              initNAC=glorot_uniform, initG=glorot_uniform, initb=glorot_uniform)
    nac = NAC(in, out, initW=initNAC, initM=initNAC)
    G = initG(out, in)
    b = initb(out)
    NALU(nac, G, b)
end

add(nalu::NALU, x) = nalu.nac(x)
mult(nalu::NALU, x) = exp.(nalu.nac(log.(abs.(x) .+ nalu.ϵ)))
gate(nalu::NALU, x) = σ.(nalu.G*x .+ nalu.b)

function (nalu::NALU)(x)
    a = add(nalu, x)
    m = mult(nalu, x)
    g = gate(nalu, x)
    g .* a .+ (1 .- g) .* m
end

function Base.show(io::IO, l::NALU)
    in = size(l.G, 2)
    out = size(l.G, 1)
    print(io, "NALU(in=$in, out=$out)")
end

Flux.@functor NALU


export leastSquares

# if Φ is a row matrix and is passed as a 1 dimensional array
# function leastSquares(Φ::AbstractVector{<:Real}, Y::AbstractVector{<:Real})
#     Φ = reshape(Φ, length(Φ), 1)
#     leastSquares(Φ, Y)
# end

# Ordinary least-squares regression
function leastSquares(Φ::AbstractMatrix{<:Real}, Y::AbstractVector{<:Real})
    @assert size(Φ, 1) == length(Y) "Dimension mismatch in matrix Φ and observation vector Y."

    # FUTURE: need to account for correlated data, if present

    coefficients = zeros(size(Y))

    condition = 1/cond(Φ) # Matrix condition number

    # Fastest, but least accurate (squares condition number)
    # if condition < 1e6
    #     ΦTΦ = Φ' * Φ
    #     coefficients = ΦTΦ \ (Φ' * Y)

    # (Much) more precise but ~ twice as slow
    if condition > eps()
        coefficients = Φ \ Y

    # Slowest, but best precision for very ill-conditioned matrices. For better conditioned matrices it is not as good as other methods.
    else
        @warn("Matrix condition number is really high. Results can become inaccurate.")
        coefficients = pinv(Φ) * Y # TODO: calibrate tolerances of pinv
    end

    return coefficients
end


function rel_energy_err_sde(sol)

    NQ = ndims(sol.q)

    if NQ==2

        en_ref  = 0.5 * ( sol.q.d[1,1]^2 + sol.q.d[2,1]^2 )
        en_last = 0.5 * ( sol.q.d[1,end]^2 + sol.q.d[2,end]^2 )

        return abs((en_last-en_ref)/en_ref)

    elseif NQ==3

        en_ref  = 0.5 * ( sol.q.d[1,1,:].^2 + sol.q.d[2,1,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:].^2 + sol.q.d[2,end,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    elseif NQ==4

        en_ref  = 0.5 * ( sol.q.d[1,1,:,:].^2 + sol.q.d[2,1,:,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:,:].^2 + sol.q.d[2,end,:,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    end
end

function rel_energy_err_psde(sol)

    NQ = ndims(sol.q)

    if NQ==2

        en_ref  = 0.5 * ( sol.q.d[1,1]^2 + sol.p.d[1,1]^2 )
        en_last = 0.5 * ( sol.q.d[1,end]^2 + sol.p.d[1,end]^2 )

        return abs((en_last-en_ref)/en_ref)

    elseif NQ==3

        en_ref  = 0.5 * ( sol.q.d[1,1,:].^2 + sol.p.d[1,1,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:].^2 + sol.p.d[1,end,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    elseif NQ==4

        en_ref  = 0.5 * ( sol.q.d[1,1,:,:].^2 + sol.p.d[1,1,:,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:,:].^2 + sol.p.d[1,end,:,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    end
end


using Test, LinearAlgebra, SparseArrays, Random
using QDLDL
rng = Random.MersenneTwister(2701)

@testset "inertia checks" begin

  m = 20
  n = 30

  A = random_psd(n) + Diagonal(rand(n))
  B = sprandn(m,n,0.2)
  C = -Diagonal(rand(m))
  M = [A B' ; B C]

  b = randn(m+n)
  F = qdldl(M)

  @test positive_inertia(F) == n

end

nothing


using SpectralDistances, Distances, DSP

@testset "Sinkhorn convolutional" begin
    @info "Testing Sinkhorn convolutional"
    m,n,T = 8, 10, Float64
    for m = 8:2:12, n = 8:2:12, T = (Float64, )
        # @show m,n,T

        a = zeros(T, m, n)
        a[2,2] = 1
        b = zeros(T, m, n)
        b[3,3] = 1
        d1 = sinkhorn_convolutional(a,b,τ=1e2)[1]
        b = zeros(T, m, n)
        b[4,4] = 1
        d2 = sinkhorn_convolutional(a,b,τ=1e2)[1]
        @test 2sqrt(d1) ≈ sqrt(d2) rtol = 0.1
        b = zeros(T, m, n)
        b[6,6] = 1
        d3 = sinkhorn_convolutional(a,b,τ=1e2)[1]
        @test 2sqrt(d2) ≈ sqrt(d3) rtol = 0.1

        # Do the same as above but with different log stabilization
        a = zeros(T, m, n)
        a[2,2] = 1
        b = zeros(T, m, n)
        b[3,3] = 1
        d4 = sinkhorn_convolutional(a,b,τ=1e5)[1]
        b = zeros(T, m, n)
        b[4,4] = 1
        d5 = sinkhorn_convolutional(a,b,τ=1e5)[1]
        @test 2sqrt(d4) ≈ sqrt(d5) rtol = 0.1
        b = zeros(T, m, n)
        b[6,6] = 1
        d6 = sinkhorn_convolutional(a,b,τ=1e5)[1]
        @test 2sqrt(d5) ≈ sqrt(d6) rtol = 0.1

        @test d1 ≈ d4 # Test that the result is the same with and without log stabilization.
        @test d2 ≈ d5
        @test d3 ≈ d6
    end
end

@testset "Invariant axis" begin
    @info "Testing Invariant axis"
    w,h = 5,5
    for w = [5, 6, 20], h = [5,6,20]
        @show w,h
        β = 0.05
        D = [abs2((i-j)/(h-1)) for i = 1:h, j = 1:h]
        B = [zeros(h,w) for _ in 1:2]

        B[1][1,1] = 1
        B[1][1,5] = 1
        B[2][2,1] = 1
        B[2][2,5] = 1

        B = s1.(B)
        c,V,U = sinkhorn_convolutional(B[1], B[2], β=β, τ=1e20)
        c1,V1,U1 = sinkhorn_convolutional(B[1], B[1], β=β, τ=1e20)
        c2,V2,U2 = sinkhorn_convolutional(B[2], B[2], β=β, τ=1e20)
        c3 = sqrt((c - 0.5(c1+c2))) # The size of the summed over dimension needs to multiply here
        v,u = sum(B[1], dims=2)[:], sum(B[2], dims=2)[:] # Integrate over the dimension to be sensitive to (because we need to be invariant to this dimension in this step)
        @assert sum(u) ≈ 1
        @assert sum(v) ≈ 1
        Γ = discrete_grid_transportplan(v, u)    # Solve 1D OT between integrated measures
        invariant_cost = sqrt(dot(Γ, D))

        @test c3 ≈ invariant_cost rtol=0.15

        if isinteractive()
            Plots.heatmap(Γ, layout=5, title="Transport plan")
            Plots.heatmap!(V, sp=2, title="V")
            Plots.heatmap!(U, sp=3, title="U")
            Plots.plot!(v, sp=4, lab=invariant_cost)
            Plots.plot!(u, sp=5, lab=c3)
        end

        invariant_cost = dot(Γ, D)
        di = ConvOptimalTransportDistance(β=β, invariant_axis=1)
        d = ConvOptimalTransportDistance(β=β, invariant_axis=0)
        cdist = d(B[1], B[2]) - 0.5*(d(B[1], B[1]) + d(B[2], B[2]))
        @test cdist ≈ c3^2
        ci = (di(B[1], B[2]) - 0.5*(di(B[1], B[1]) + di(B[2], B[2])))

        @test ci ≈ cdist-invariant_cost rtol=0.15 atol=sqrt(eps())

        @test di(B[1], B[2]) ≈ d(B[1], B[2])-invariant_cost rtol=0.15

    end
end

@testset "Misc convolutional" begin
    @info "Testing Misc convolutional"

    @test -0.72 < SpectralDistances.default_dynamic_floor(rand(100,100)) < -0.6
    a1 = zeros(Float64, 10,10)
    a1[2,2] = 1
    a2 = zeros(Float64, 10,10)
    a2[6,6] = 1
    A = [a1,a2]
    β = 0.001
    λ = [0.5, 0.5]
    b = barycenter_convolutional(A,λ,β=β)
    @test maximum(b) == b[4,4]
    @test sum(b) ≈ 1

    d = ConvOptimalTransportDistance(β=0.01)
    l = barycentric_coordinates(d, A, a1)[1]
    @test l ≈ [1, 0]
    @test sum(l) ≈ 1
    l = barycentric_coordinates(d, A, a2)[1]
    @test l ≈ [0, 1]

    l = barycentric_coordinates(d, A, b)[1]
    @test l ≈ [0.5, 0.5]

    A1 = spectrogram(sin.(1:10000) .+ 0.1randn(10000), 256, window=hanning)
    A2 = spectrogram(sin.(1:10000) .+ 0.1randn(10000), 256, window=hanning)
    A = [A1, A2]

    @test -5 < SpectralDistances.default_dynamic_floor(A1) < -4
    @test -5 < SpectralDistances.default_dynamic_floor(A) < -4

    o = similar(A1.power)
    @test normalize_spectrogram!(o,A1) == normalize_spectrogram(A1)
    A1n = normalize_spectrogram(A1)
    A1f = mask_filter(A1n)
    @test sum(A1f) < sum(A1n)

    B = barycenter_convolutional(A,λ,β=β, dynamic_floor=-3)
    @test B isa typeof(A1)

    d = ConvOptimalTransportDistance(β=β, dynamic_floor=-3.0)
    B2 = barycenter(d, A)
    @test power(B) == power(B2)
    @test time(B) == time(B2)

    @test d(A1,A2) > d(A1,A1)
    @test d(A1,A2) > d(A2,A2)
    @test d(A1,A2) - 0.5*(d(A1,A1) + d(A2,A2)) > 0

    A3 = spectrogram(sin.(LinRange(0.8,1.2,300_000) .* (1:300_000)) .+ 0.01randn(300_000), 256, window=hanning)

    d = ConvOptimalTransportDistance(β=0.05, dynamic_floor=-3.0)
    @time D = distance_profile(d, A1, A3, tol=1e-6, stride=15)
    isinteractive() && plot(D)
    @test 40 <= argmin(D) <= 70

    dslide = SlidingConvOptimalTransportDistance(d)
    @test dslide(A1, A3; tol=1e-6, stride=15) ≈ minimum(D)
    @test SlidingConvOptimalTransportDistance(β=0.05, dynamic_floor=-3.0).d.β == SlidingConvOptimalTransportDistance(ConvOptimalTransportDistance(β=0.05, dynamic_floor=-3.0)).d.β

end

# @test B.power == barycenter_convolutional([A1.power, A2.power],λ,β=β)

# w = SpectralDistances.BCWorkspace(A, β)
# @btime SpectralDistances.barycenter_convolutional($w,$A,$λ)
# 1.198 ms (46 allocations: 1.26 MiB)
# 585.176 μs (46 allocations: 1.26 MiB) # add @avx
# 469.132 μs (16 allocations: 313.13 KiB) # operate inplace
# 387.881 μs (0 allocations: 0 bytes) # views and generator in sum

##
# m,n = 1280,1280
# x = zeros(m,n)
# x[10,10] = 1
# β = 0.0003 # The threshold for imfilter to be faster is around 0.0003 on size 128
#
# t = LinRange(0, 1, m)
# Y, X = meshgrid(t, t)
# xi1 = @. exp(-(X - Y)^2 / β)
#
# t = LinRange(0, 1, n)
# Y, X = meshgrid(t, t)
# xi2 = @. exp(-(X - Y)^2 / β)
#
# K1(x) = xi1 * x * xi2
#
#
# using ImageFiltering
#
# t = LinRange(-0.5, 0.5, m)
# k1 = @. exp(-t^2 / β)
# t = LinRange(-0.5, 0.5, n)
# k2 = @. exp(-t^2 / β)
#
# function kfg(x)
#     l = round(Int,5x)
#     l = isodd(l) ? l : l+1
#     KernelFactors.gaussian(x,l)
# end
#
# kk1 = kfg(sqrt(0.5β)*m)
# kk2 = kfg(sqrt(0.5β)*n)
# kk1,kk2 = kk1 ./ maximum(kk1), kk2 ./ maximum(kk2)
# # kernf = kernelfactors((centered(k1), centered(k2)))
#
# # kernf = KernelFactors.gaussian((sqrt(β)*m, sqrt(β)*m))
# kernf = kernelfactors((kk1, kk2))
# kernp = broadcast(*, kernf...)
#
# CPU1(Algorithm.FIR())
# K2(x) = imfilter(x, kernf)
#
# plot(heatmap(K1(x)), heatmap(K2(x)), heatmap(K2(x)-K1(x)))
#
# using FFTW
# FFTW.set_num_threads(2)
#
#
# using ImageFiltering.ComputationalResources
# @btime $xi1 * $x * $xi2
# @btime imfilter!($(similar(x)), $x, $kernf)
#
# using LoopVectorization
# function filter2davx!(out::AbstractMatrix, A::AbstractMatrix, kerns)
#     for kern in kerns
#         @avx for J in CartesianIndices(out)
#             tmp = zero(eltype(out))
#             for I ∈ CartesianIndices(kern)
#                 tmp += A[I + J] * kern[I]
#             end
#             out[J] = tmp
#         end
#     end
#     out
# end
#
# out = similar(x)
#
# out = ImageFiltering.OffsetArray(similar(x, size(x).-2), 1, 1)
#
# ka = convert(AbstractArray, kernf[1])
# kb = convert(AbstractArray, kernf[2])
#
# @btime filter2davx!($out, $x, $((ka, kb)));
#
# plot(heatmap(K1(x)), heatmap(K2(x)), heatmap(K2(x)-K1(x)), heatmap(out-K1(x)))


# test MMD distance

# testing only
# using Revise
# using Gadfly
# Gadfly.set_default_plot_size(35cm,25cm)

using Test
using ApproxManifoldProducts
using DataFrames

const AMP = ApproxManifoldProducts

@testset "Test mmd distance between Euclidean beliefs" begin

# when not plottin
offsets = 1:1 # -30:0.25:30
mc      = 1   # 1:100

checkGrid = zeros(length(offsets), length(mc))

df = DataFrame(offsets=[], mmd=[])

for i in 1:length(offsets), j in mc

P = [[randn();]  for k in 1:100]
Q = [[randn()+offsets[i];]  for k in 1:100]

res = zeros(1)

AMP.mmd!(res, P, Q, AMP.Euclid, bw=[0.001])

checkGrid[i,j] = res[1]
push!(df, (offsets[i],res[1]))

end

# when not plotting
@test 0 < df[!,:mmd][1] < 1.0

# plot(df, x=:offsets, y=:mmd, Geom.boxplot,
# Theme(default_color="MidnightBlue"))

# (manikde!(P, (:Euclid,)) |> getBW)[1,1]

end

@testset "Test mmd distance between 2D Euclidean beliefs" begin

# when not plottin
offsets = 1:1 # -30:0.25:30
mc      = 1   # 1:100

checkGrid = zeros(length(offsets), length(mc))

df = DataFrame(offsets=[], mmd=[])

for i in 1:length(offsets), j in mc

P = [randn(2) for k in 1:100]
Q = [randn(2) .+ offsets[i] for k in 1:100]

res = zeros(1)

AMP.mmd!(res, P, Q, AMP.Euclid2, bw=[0.001])

checkGrid[i,j] = res[1]
push!(df, (offsets[i],res[1]))

end

# when not plotting
@test 0 < df[!,:mmd][1] < 1.0

# plot(df, x=:offsets, y=:mmd, Geom.boxplot,
# Theme(default_color="MidnightBlue"))

# (manikde!(P, (:Euclid,)) |> getBW)[1,1]

end


using ScatteringLayers
using CSVFiles
using DataFrames
using Dates
using DiffEqOperators
using Images
using StaticArrays

println("Testing echogram types...")
sp = SpacePoint(SVector(randn(3)...))
sp = SpacePoint(randn(3))
stp = SpaceTimePoint(SVector(randn(3)...), Dates.now())
stp = SpaceTimePoint(randn(3), Dates.now())
sps = [SpacePoint(SVector(randn(3)...)) for i in 1:10]
z = collect(1:5.0)
eg = Echogram(randn(5, 10), z, sps)
depths(eg)
spacetime(eg)
png = getping(eg, 1)

println("Loading test data...")
echo_df = DataFrame(load("data\\test_data.csv"))
fmt = DateFormat("yyyymmddHH:MM:SS.s")
echo_df[:datetime] = DateTime.(string.(echo_df[:Date_M]) .* echo_df[:Time_M], fmt)

echo = unstack(echo_df, :Layer_depth_max, :datetime, :Sv_mean)#[1:end, 2:end]
z = float.(echo[:Layer_depth_max])
t = DateTime.(string.(names(echo)[2:end]))
echo = disallowmissing(Matrix(echo[:, 2:end]))
echo = echo[end:-1:1, :]
echo[echo .< -90] .= -90
echo = 10 .^(echo/10)
echo = imfilter(echo, Kernel.gaussian(0.5))
stp = [SpaceTimePoint([36.712110, -122.187028], ti) for ti in t]
echo = Echogram(echo, z, stp)
m, n = size(echo)
dz = 1.0
zz = 880.0:-dz:8.0
order = 1
D2 = DerivativeOperator{Float64}(2, order, dz, m, :Dirichlet0, :Dirichlet0)

println("Testing...")
ping = getping(echo, 5)

THRESH = 10^(-80/10)
μ, p, c = getpeaks(ping, thresh=THRESH)
mix = guesslayers(ping, thresh=THRESH)

basis = getbasis(mix, zz)
N1 = refine_heights!(mix, basis, ping)
ping_nnls = mixpdf(mix, zz)

fit1 = refine_μσ!(mix, ping)
ping_refined1 = mixpdf(mix, zz)
mix = guesslayers(ping, thresh=THRESH)
fit!(mix, ping, trace=true)
ping_refined2 = mixpdf(mix, zz)

mix1 = fit(mix, ping, trace=true)
mix1 = fit(mix, ping, refine=:σ)
mix1 = fit(mix, ping, refine=:σμ)
mix1 = fit(mix, ping, refine=:widths)
mix1 = fit(mix, ping, refine=:widths_locs)

x = SpacePoint([36.712110, -122.187028])
echogram = Echogram(imfilter(echo, Kernel.gaussian(0.5)), zz, repeat([x], n))
layers = fitlayers(echogram, thresh=THRESH)