function train_layer_1_rbf_fast(neurons,vars,obs,n_per_part,inds_all,x,y,cc)

    # a_all=zeros(neurons,1000vars)
    a_all=Vector{Vector{Float64}}()
    ooops=0
    for i=1:neurons
        # if i in 1:1:neurons println(Dates.format(now(), "HH:MM:SS")," Calculating weights w for ",i," neuron, of ",neurons," Total.") end
        # if n_per_part[i+1]-n_per_part[i]!=0
            try
                x1=copy(x[inds_all[n_per_part[i]+1:n_per_part[i+1]],:])
                obs1=size(x1,1)
                phi1 = calc_phi(x1,x1,obs1,obs1,cc,vars)
                atmp=phi1\y[inds_all[n_per_part[i]+1:n_per_part[i+1]]]
                # display(plot(y[inds_all[n_per_part[i]+1:n_per_part[i+1]]]))
                # println(mean(y[inds_all[n_per_part[i]+1:n_per_part[i+1]]])," ",n_per_part[i+1]-n_per_part[i])
                # if sum(isnan.(atmp))==0 && sum(isinf.(atmp))==0 a_all[i,1:obs1].=atmp else ooops+=1 end
                # if sum(isnan.(atmp))==0 && sum(isinf.(atmp))==0 push!(a_all,atmp) else push!(a_all,zeros(obs1)); ooops+=1 end
                if sum(isnan.(atmp))==0 && sum(isinf.(atmp))==0 && sum(abs.(atmp[2:end]))>0 push!(a_all,atmp)
                    push!(i_keep,i)
                else
                    ooops+=1
                end
            catch ex
                ooops+=1; println(i," ",ex)
            end
        # end
    end
    println("ooops=",ooops)
    layer1=zeros(obs,neurons)
    for i=1:neurons
        if i in 1:100:neurons println(Dates.format(now(), "HH:MM:SS")," Calculating output of neuron ",i," , of ",neurons," Total.") end
        x1=copy(x[inds_all[n_per_part[i]+1:n_per_part[i+1]],:])
        obs1=size(x1,1)
        # obs1=length(inds_all[n_per_part[i]+1:n_per_part[i+1]])
        # phi1 = calc_phi(x1,x,obs1,obs,cc,vars)
        # # layer1[:,i].=phi1*a_all[i,1:obs1]
        # layer1[:,i].=phi1*a_all[i]

        # phi=zeros(nn,n)
        # obs1=length(inds_all[n_per_part[i]+1:n_per_part[i+1]])
        for ii=1:obs
            phis=0.0
            for j=1:obs1
                dr=0
                for k=1:vars
                    dr = dr + (x[ii,k] - x1[j,k])^2
                    # dr = dr + (x[ii,k] - x[inds_all[n_per_part[i]+j],k])^2
                end
                phis+=exp(-dr/cc)*a_all[i][j]
                # phis+=-(dr^4)/4/cc
            end
            layer1[ii,i]=phis
        end

    end
    println(Dates.format(now(), "HH:MM:SS")," Solving system (",obs," by ",neurons+1,") for weights v in the output Layer")
    a_layer1=[layer1 ones(obs)]\y
    println(Dates.format(now(), "HH:MM:SS")," All weights computed.")

    return a_all,a_layer1,layer1

end


# Hodkin-Huxley spiking layer & net

### HH FUNCTIONS

# helper function: traps for 0 in denominator of rate eqns.
vtrap = (x,y) -> (abs(x/y)<1e-6) ? y*(1 - x/y/2) : x/(exp(x/y) - 1);
# Na channel activation
αM = v-> .1 * vtrap(-(v+40),10)
βM = v-> 4exp(-(v+65)/18)
# Na channel inactivation
αH = v-> .07exp(-(v+65)/20)
βH = v-> 1/(exp(-(v+35)/10)+1)
# K channel activation
αN = v-> .01vtrap(-(v+55),10)
βN = v-> 0.125*exp(-(v+65)/80);
inf = (α,β) -> α/(α+β) # for computing minf, hinf, ninf
Minf = v-> inf(αM(v),βM(v))
Hinf = v-> inf(αH(v),βH(v))
Ninf = v-> inf(αN(v),βN(v))

# Layer parameters
const Cm = 10 # Capacitative density [μF/cm^2]
const gNa = 120 # maximum Na channel conductance
const gK = 36
const gL = 1
const ENa = 115
const EK = -12
const ELeak = 0
const v_r = -65

const e_rev = 0 # reversal potential for AMPA synapses (mV)
const i_rev = -80 # reversal potential for GABA synapses (mV)
const abs_refr=5 # ms
const v_th=-20

function HHNeuron()
    HHNeuron(Minf(v_r), Hinf(v_r), Ninf(v_r), v_r, -Inf, false)
end

function update!(neuron::HHNeuron, Ge_syn::Float64, Gi_syn::Float64, Iinj::Float64, t::Float64)
	m = neuron.m
	h = neuron.h
	n = neuron.n
	v = neuron.v
	Isyn = Ge_syn.*(e_rev-v)-Gi_syn.*(i_rev-v) # total synaptic current
	vdot =  (gNa*m^3*h*(ENa-(v+65)) + gK*n^4*(EK-(v+65)) + gL*(ELeak-(v+65)) + Isyn + Iinj)/Cm

    neuron.v += dt*vdot
    neuron.m += dt*(αM(v)*(1-m) - βM(v)*m);
    neuron.h += dt*(αH(v)*(1-h) - βH(v)*h);
    neuron.n += dt*(αN(v)*(1-n) - βN(v)*n);

    neuron.didSpike = (v > v_th) & (t > neuron.lastSpike+abs_refr)
    if neuron.didSpike
    	neuron.lastSpike = t
    end
end

function reset!(neuron::HHNeuron)
    neuron.m = Minf(v_r)
    neuron.h = Hinf(v_r)
    neuron.n = Ninf(v_r)
    neuron.v = v_r
    neuron.lastSpike = -Inf
    neuron.didSpike = false
end


using FFTW: plan_fft!
using Statistics: mean, cov, var
using Distributed: workers, remotecall, fetch, myid

function linreg(x,y)
  b = cov(x,y)/var(x)
  a = mean(y) - b*mean(x)
  a,b
end

function detrendTS!(outar::AbstractMatrix,xin::AbstractVector{T}) where T
    x=collect(T,eachindex(xin))
    a,b=linreg(x,xin)
    for i in eachindex(xin)
        outar[i,1]=a+b*x[i]
    end
end

function tscale2ind(b::Float64,l::Int)
    i=round(Int,l/b+1)
    return i
end
mirror(i,l)=l-i+2

"""
    filterTSFFT(c)

Filter each time series using a Fourier filter and return the decomposed series
in 4 time windows (Trend, Long-Term Variability, Annual Cycle, Fast Oscillations)

**Input Axis** `Time`axis

**Output Axes** `Time`axis, `Scale`axis

"""
function filterTSFFT(c;kwargs...)
  indims = InDims("time",filter=AnyMissing())
  outdims = OutDims("time",(c,p)->CategoricalAxis("Scale",["Trend", "Long-Term Variability", "Annual Cycle", "Fast Oscillations"]))
  ntime = length(getAxis("Time",c))
  plans = map(workers()) do id
    remotecall(id,ntime,eltype(c)) do nt,et
      testar = zeros(Complex{Base.nonmissingtype(et)},nt)
      fftplan = plan_fft!(testar)
      ifftplan = inv(fftplan)
      fftplan, ifftplan
    end
  end
  plandict = Dict(zip(workers(), plans))
  mapCube(filterTSFFT,c,getNpY(c),plandict;indims=indims,outdims=outdims,kwargs...)
end

function filterTSFFT(outar::AbstractMatrix,y::AbstractVector, annfreq::Number,
  plandict; nharm::Int=3)

  fftplan, ifftplan = fetch(plandict[myid()])

    detrendTS!(outar,y)
    l        = length(y)

    fy       = Complex{Base.nonmissingtype(eltype(y))}[y[i]-outar[i,1] for i=1:l]
    fftplan * fy #Apply fftplan on fy inplace changes fy
    fyout    = similar(fy)
    czero    = zero(eltype(fy))

    #Remove annual cycle
    fill!(fyout,zero(eltype(fyout)))

    for jharm = 1:nharm
        iup=tscale2ind(annfreq*1.1/jharm,l)

        idown=tscale2ind(annfreq*0.9/jharm,l)

        for i=iup:idown
            fyout[i]  = fy[i]
            i2        = mirror(i,l)
            fyout[i2] = fy[i2]
            fy[i]     = czero
            fy[i2]    = czero
        end
    end

    ifftplan * fyout #Apply ifftplan on fyout inplace, changes fyout
    for i=1:l
        outar[i,3]=real(fyout[i])
    end

    iup   = tscale2ind(annfreq*1.1,l)
    idown = tscale2ind(annfreq*0.9,l)
    #Now split the remaining parts
    fill!(fyout,czero)
    for i=2:iup-1
        i2        = mirror(i,l)
        fyout[i]  = fy[i]
        fyout[i2] = fy[i2]
        fy[i]     = czero
        fy[i2]    = czero
    end
    ifftplan * fyout
    ifftplan * fy

    for i=1:l
        outar[i,2]=real(fyout[i])
        outar[i,4]=real(fy[i])
    end
end


type Box
    data::Vector{Int}
    children::Vector{Box}
end

Box() = Box(Int[], Box[])

"""
T: type of the coordinates
P: type of the points stored in the Octree.
"""
type Octree{T,P}
    center::P
    halfsize::T
    rootbox::Box

    points::Vector{P}
    radii::Vector{T}
    expanding_ratio::Float64 # This will add the expanding ratio of boxes

    splitcount::Int
    minhalfsize::T
end

"""
  boundingbox(v)

Compute the bounding cube/square for a Array of Point. The return values
are the center of the bounding box and the half size of the cube.
"""
function boundingbox{P}(v::Vector{P})
  #P are points values (x,y) in 2D or (x,y,z) in 3D
  ll = minimum(v)
  ur = maximum(v)
  #ll => LowerLeft
  #up => UpperRight
  c = (ll + ur)/2
  s = maximum(ur - c)
  #c => centre
  #s => biggest half width, i.e in 2D distance between points result in a rectangular
  #    (so we pick the bigger value here to form a rectangular later)
  return c, s
end

"""
Predicate used for iteration over an Octree. Returns true if two boxes
specified by their centers and halfsizes overlap. More carefull investigation
of the objects within is required to assess collision. False positives are possible.

    boxesoverlap(c1, hs1, c2, hs2)
"""
function boxesoverlap(c1, hs1, c2, hs2)

    tol = sqrt(eps(typeof(hs1)))
    # Checking the type of the problem domain 2D or 3D? and making sure the boxes are from same domain
    dim = length(c1)
    @assert dim == length(c2)
    #Note: I have fixed the condition, now it works for 3D and 2D
    hs = hs1 + hs2
    for i in 1 : dim
        if abs(c1[i] - c2[i]) > hs + tol
            return false
        end
    end

    return true
end

function Octree{T}(points::Vector, radii::Vector{T}, expanding_ratio=1.0, splitcount = 10,  minhalfsize = zero(T))

    n_points = length(points)
    n_dims = length(eltype(points))

    # compute the bounding box taking into account
    # the radius of the objects to be inserted
    radius =  maximum(radii)

	@assert !isempty(points)
	ll = points[1]
	ur = points[1]
	for i in 2:length(points)
		ll = min.(ll, points[i])
		ur = max.(ur, points[i])
	end

	ll = ll .- radius
	ur = ur .+ radius

    #ll = minimum(points) - radius
    #ur = maximum(points) + radius

    center = (ll + ur) / 2
    halfsize = maximum(ur - center)

    # if the minimal box size is not specified,
    # make a reasonable guess
    if minhalfsize == 0
        #TODO generalise
        minhalfsize = 0.1 * halfsize * (splitcount / n_points)^(1/3)
    end

    # Create an empty octree
    rootbox = Box()
    tree = Octree(center, halfsize, rootbox, points, radii, expanding_ratio, splitcount, minhalfsize)

    # populate
    for id in 1:n_points

        point, radius = points[id], radii[id]
        insert!(tree, tree.rootbox, center, halfsize, point, radius, id)

    end

    return tree

end

"""
  Octree(points)

Insert zero radius objects at positions `points` in an Octree
"""

Octree(points) = Octree(points, zeros(eltype(eltype(points)), length(points)))

"""
    childsector(point, center) -> sector

Computes the sector w.r.t. `center` that contains  point. Sector is an Int
that encodes the position of point along each axis in its bit representation
"""
function childsector(point, center)
  # in Case of 3D the Octant they are numbered as follows (+,+,+)->7, (+,+,-)->3, (+,-,+)->5, (+,-,-)->1
  #(-,+,+)->6, (-,+,-)->2, (-,-,+)->4, (-,-,-)->0
  #For 2D (-,-)->0, (-,+)->1, (+,-)->2, (+,+)->3
	sct = 0
	r = point - center
	for (i,x) in enumerate(r)
		if x > zero(x)
			sct |= (1 << (i-1))
		end
	end

	return sct # because Julia has 1-based indexing
end

isleaf(node) = isempty(node.children)

"""
    fitsinbox(pos, radius, center, halfsize) -> true/fasle

Finds out if the object with position (pos) and (raduis) fits inside the box.
It uses the box dimension (centre, and halfsize(w/2)) to do the comparsion
"""
function fitsinbox(pos, radius, center, halfsize)

	for i in 1:length(pos)
		(pos[i] - radius < center[i] - halfsize) && return false
		(pos[i] + radius > center[i] + halfsize) && return false
	end
# the code judege by comapring with the box lower left point and uper right point
	return true
end
"""
  itsinbox(Opj_c, Obj_rad, box_c, box_hf, ratio) -> true/fasle

  Finds out if the object with center position (Opj_c) and (Obj_rad) fits inside the box.
  It uses the box dimension (box_c, and box_hf(w/2)) to do the comparsion
  ratio is relative to box half width, so if you want to be just within the box
  then ratio =1, if the boundaries of the object is slightly bigger and you still
  want to include them make ratio bigger ex:1.2
"""
function fitsinbox(Opj_c, Obj_rad, box_c, box_hf, ratio)
  for i in 1:length(Opj_c)
		(Opj_c[i] - Obj_rad < box_c[i] - (box_hf*ratio)) && return false
		(Opj_c[i] + Obj_rad > box_c[i] + (box_hf*ratio)) && return false
	end
# the code judege by comapring with the box lower left point and uper right point
	return true

end

"""
  childcentersize(center, halfsize, sector) -> center, halfsize

Computes the center and halfsize of the child of the input box
that resides in octant `sector`
"""
@generated function childcentersize(center, halfsize, sector)
  D = length(center)
  xp1 = :(halfsize = halfsize / 2)
  xp2 = Expr(:call, center)
  for d in 1:D
    push!(xp2.args, :(center[$d] + (sector & (1 << $(d-1)) == 0 ? -halfsize : +halfsize)))
  end
  xp = quote
    $xp1
    return $xp2, halfsize
  end
  #@show xp
  #xp
end

"""
insert!(tree, box, center, halfsize, point, radius, id)

tree:     the tree in which to insert
box:      the box in which to try insertion
center:   center of the box
halfsize: 0.5 times the length of the box side
point:    the point at which to insert
radius:   the radius of the item to insert
id:       a unique id that identifies the inserted item uniquely
"""
function insert!{P,T}(tree, box, center::P, halfsize::T, point::P, radius::T, id)

    # if not saturated: insert here
    # if saturated and not internal : create children and redistribute
    # if saturated and internal and not fat: insert!(childbox,...)
    # if saturated and internal and fat: insert here
    # also if not saturated but there are non empty internal boxes try to insert there not here

    # or shorter:

    # sat & not internal: create children and redistibute
    # sat & internal & not fat: insert in childbox
    # all other cases: insert here
    # Will find out first if we are solveing octree or quadtree 3D/2D
    dim = length(P)
    nch = 2^dim
    expanding_ratio=tree.expanding_ratio

    saturated = (length(box.data) + 1) > tree.splitcount
    fat       = !fitsinbox(point, radius, center, halfsize,expanding_ratio) # we test if the point is within the box and the expanswion if there is any
    internal  = !isleaf(box)

    if (!internal && !saturated) || (saturated && internal && fat)
        # this will only insert in top level in two cases
        # 1) the object is fat anyway
        # 2) there is no child boxes and still there is a place in this box ( it is not saturated)
        push!(box.data, id)

    elseif internal && !fat
        # now if the box has a space or not but it has a childern try to insert in the child not here
        sct = childsector(point, center)
        chdbox = box.children[sct+1]
        chdcenter, chdhalfsize = childcentersize(center, halfsize, sct)
        if fitsinbox(point, radius, chdcenter, chdhalfsize,expanding_ratio) # check if it is fat and please consider the expansion
          insert!(tree, chdbox, chdcenter, chdhalfsize, point, radius, id)
        else
          push!(box.data, id)
        end

    else # saturated && not internal

        # if their was a previous attempt to subdivide this box,

        # insert the new element in this box for now as we will sort them after we create childrens
        push!(box.data, id)

        # if we are not allowed to subdivide any further stop. This will
        # avoid the contruction of a tree with N levels when N equal points
        # are inserted.
        if halfsize/2 < tree.minhalfsize
            return
        end

        # Create an array of childboxes
        box.children = Array{Box}(nch)
        for i in 1:nch
            box.children[i] = Box(Int[], Box[])
        end

        # subdivide:
        # for every id in this box
        #   find the correspdoning child sector
        #   if it fits in the child box, insert
        #   if not, add to the list of unmovables
        # replace the current box data with the list of unmovables

        unmovables = Int[]
        for id in box.data

            point = tree.points[id]
            radius = tree.radii[id]

            sct = childsector(point, center)
            chdbox = box.children[sct+1]
            chdcenter, chdhalfsize = childcentersize(center, halfsize, sct)
            if fitsinbox(point, radius, chdcenter, chdhalfsize,expanding_ratio)# # check if it is fat for the child and please consider the expansion
                push!(chdbox.data, id)
            else
                push!(unmovables, id)
            end

        end

        box.data = unmovables

    end

end

import Base.length
function length(tree::Octree)

    # Traversal order:
    #   data in the box
    #   data in all children
    #   data in the sibilings
    level = 0

    box = tree.rootbox
    sz = length(box.data)

    box_stack = [tree.rootbox]
    sct_stack = [1]

    sz = -0

    box = tree.rootbox
    sct = 0

    box_stack = Box[]
    sct_stack = Int[]

    while true

        # if this is the first time the box is visited, count the data
        if sct == 0
            sz += length(box.data)
            if length(box.data) != 0
                println("Adding ", length(box.data), " contributions at level: ", level)
            end
        end

        # if this box has unprocessed children
        # push this box on the stack and process the children
        if sct < length(box.children)
            push!(box_stack, box)
            push!(sct_stack, sct+1)
            level += 1

            box = box.children[sct+1]
            sct = 0

            continue
        end

        # if box and its children are processed,
        # and their is no parent above this box:
        # end the traversal:
        if isempty(box_stack)
            break
        end

        # if either no children or all children processed:
        # move up one level
        box = pop!(box_stack)
        sct = pop!(sct_stack)
        level -= 1

    end

    return sz

end

type BoxIterator{T,P,F}
    predicate::F
    tree::Octree{T,P}
end

type BoxIteratorStage{T,P}
    box::Box
    sct::Int
    center::P
    halfsize::T
end

boxes(tree::Octree, pred = (ctr,hsz)->true) = BoxIterator(pred, tree)

import Base: start, next, done
function start(bi::BoxIterator)

    pred     = bi.predicate
    center   = bi.tree.center
    halfsize = bi.tree.halfsize

    state = [ BoxIteratorStage(
        bi.tree.rootbox, 0, center, halfsize
    ) ]

    # If the rootbox does not satisfy the predicate,
    # fast forward to the next eligible box
    # take care of this deeply annoying corner case:
    if !pred(center, halfsize)
        box, state = next(bi, state)
    end

    return state

end

function next(bi::BoxIterator, state)

    item = last(state).box.data

    box = last(state).box   # current box
    sct = last(state).sct   # next child to visit
    hsz = last(state).halfsize
    ctr = last(state).center

    while true

        # scan for a next child box that meets the criterium
        childbox_found = false
        while sct < length(box.children)
            chd_ctr, chd_hsz = childcentersize(ctr, hsz, sct)
            if bi.predicate(chd_ctr, chd_hsz)
                childbox_found = true
                break
            end
            sct += 1
        end

        if childbox_found

            # if this box has unvisited children, increment
            # the next child sct counter and move down the tree
            last(state).sct = sct + 1
            ctr, hsz = childcentersize(ctr, hsz, sct)
            stage = BoxIteratorStage(box.children[sct+1], 0, ctr, hsz)
            push!(state, stage)

        else

            pop!(state)

        end

        # if we popped the root, we're finished
        if isempty(state)
            break
        end

        box = last(state).box
        sct = last(state).sct
        hsz = last(state).halfsize
        ctr = last(state).center

        # only stop the iteration when a new box is found
        # and if that box is non-empty
        # (sct == 0) implies that this is the first visit
        if sct == 0 && !isempty(box.data)
            break
        end

    end

    return item, state
end

done(bi::BoxIterator, state) = isempty(state)

import Base.find

"""
    find(octree, pos, tolerance=sqrt(eps(eltype(pos))))

Return an array containing the indices of values at `pos` (up to a tolerance)
"""
function find(tr::Octree, v; tol = sqrt(eps(eltype(v))))

    pred = (c,s) -> fitsinbox(v, 0.0, c, s)# i didn't change it because find will get the point anyway, it only chck for its existance
    I = Int[]
    for b in boxes(tr, pred)
      for i in b
        if norm(tr.points[i] - v) < tol
          push!(I, i)
        end
      end
    end

    return I

end


@testset "affine" begin

    let rng = MersenneTwister(123456)

        @test params(DiagAffine(1.0, 1.0)) == [[1.0], [1.0]]
        @test params(DiagAffine(1.0, 1.0)) == params(DiagAffine([1.0], [1.0]))

        # Test that the application yields sane results.
        x1, x2 = randn(rng, 2, 5), randn(rng, 1, 5)
        @test apply(DiagAffine(0.0, 0.0), x2) == x2
        @test apply(DiagAffine(1.0, 0.0), x2) == x2 + 1
        @test apply(DiagAffine(ones(2), zeros(2)), x1) == x1 + ones(x1)
        @test apply(DiagAffine(1.0, 0.0), [1.0]) ≈ [2.0]

        # Test that the logdetJ yields some sane results.
        @test logdetJ(DiagAffine(0.0, 0.0), [5.0]) == zero(5.0)
        @test logdetJ(DiagAffine(0.0, 1.0), [5.0]) == one(5.0)
        @test logdetJ(DiagAffine(0.0, 2.0), [5.0]) ≈ 2 * one(5.0)
        @test logdetJ(DiagAffine(0, 0), randn(1, 5)) == zeros(5)
        @test logdetJ(DiagAffine(0, 1), randn(1, 7)) == ones(7)
    end
end


g = IntAdjacencyList(4,3)
add_edge!(g, 1, 2, 1)
add_edge!(g, 2, 3, 1)
add_edge!(g, 1, 4, 2)

@test num_edges(g) == 3

@test forward_neighbors(g, 1, 1) == [(2,1), (1,2)]


# Multithreaded matrix-vector multiplication for 'Transpose{<:Any,<:SparseMatrixCSC}'
# using Threads.@threads
# Why the unused α and β? See: https://github.com/JuliaLang/julia/blob/381693d3dfc9b7072707f6d544f82f6637fc5e7c/stdlib/LinearAlgebra/src/matmul.jl#L202
function mul!(y::StridedVector, transA::Transpose{<:Any, <:SparseMatrixCSC}, x::StridedVector, α::Number, β::Number)
    size(transA)[2] == length(x) || throw(DimensionMismatch("transA.n != length(x)"))
    A = transA.parent
    nzval = A.nzval
    colval = A.rowval
    rowptr = A.colptr
    m = A.n
    @inbounds Threads.@threads for i = 1:m
        y[i] = 0
        @inbounds for j = rowptr[i]:(rowptr[i + 1] - 1)
            y[i] += nzval[j] * x[colval[j]]
        end
    end
    y
end
*(transA::Transpose{<:Any, <:SparseMatrixCSC}, x::StridedVector{Tx}) where {Tx} =
    (T = promote_op(matprod, eltype(transA), Tx); mul!(similar(x, T, size(transA, 1)), transA, x, true, false))

# Multithreaded matrix-vector multiplication for 'Adjoint{<:Any,<:SparseMatrixCSC}'
# using Threads.@threads
# Why the unused α and β? See: https://github.com/JuliaLang/julia/blob/381693d3dfc9b7072707f6d544f82f6637fc5e7c/stdlib/LinearAlgebra/src/matmul.jl#L202
function mul!(y::StridedVector, adjA::Adjoint{<:Any, <:SparseMatrixCSC}, x::StridedVector, α::Number, β::Number)
    size(adjA)[2] == length(x) || throw(DimensionMismatch("adjA.n != length(x)"))
    A = adjA.parent
    nzval = A.nzval
    colval = A.rowval
    rowptr = A.colptr
    m = A.n
    @inbounds Threads.@threads for i = 1:m
        y[i] = 0
        @inbounds for j = rowptr[i]:(rowptr[i + 1] - 1)
            y[i] += nzval[j] * x[colval[j]]
        end
    end
    y
end
*(transA::Adjoint{<:Any, <:SparseMatrixCSC}, x::StridedVector{Tx}) where {Tx} =
    (T = promote_op(matprod, eltype(transA), Tx); mul!(similar(x, T, size(transA, 1)), adjA, x, true, false))


module Digits

export
  reversedigits,
  reversedigits!,
  undigit,
  digithist,
  isanagram,
  ispalindrome,
  contains,
  startswith,
  endswith,
  crop,
  combine,
  crosssum,
  select,
  replace,
  replace!,
  digitroot

function undigit(l::Vector{Int}, base::Int = 10)
  acc = 0
  for (i, digit) in enumerate(l)
    acc += digit * base ^ (i-1)
  end
  acc::Integer
end

undigit(n::Integer) = n

reversedigits(n::Integer) = undigit(reverse!(digits(n)))

reversedigits(l::Array{Int,1}) = reverse(l)

reversedigits!(l::Array{Int,1}) = reverse!(l)

function digithist(l::Array{Int,1})
  return hist(abs(l),-0.5:10)[2]
end

digithist(n::Integer) = digithist(digits(n))

function isanagram(a::Array{Int,1}, b::Array{Int,1})
  if sign(a) != sign(b)
    return false
  end
  return digithist(a) == digithist(b)
end

isanagram(a::Integer, b::Integer) = isanagram(digits(a),digits(b))
isanagram(a::Integer, b::Vector{Int}) = isanagram(digits(a),b)
isanagram(a::Vector{Int}, b::Integer) = isanagram(a,digits(b))

function ispalindrome(l::Array{Int,1})
  m = div(length(l),2)
  for i = 1:m
    if l[i] != l[end-i+1]
      return false
    end
  end
  return true
end

ispalindrome(n::Integer) = ispalindrome(digits(n))

function contains(a::Array{Int,1},b::Array{Int,1})
  la = length(a)
  lb = length(b)
  if la > lb
    for i = 1:la-lb+1
      if a[i:i+lb-1] == b
        return true
      end
    end
  elseif lb > la
    for i = 1:lb-la+1
      if b[i:i+la-1] == a
        return true
      end
    end
  elseif a == b
    return true
  end
  return false
end

contains(a::Array{Int,1},b::Integer) = contains(a,digits(b))
contains(a::Integer,b::Array{Int,1}) = contains(digits(a),b)
contains(a::Integer,b::Integer) = contains(digits(a),digits(b))

function endswith(a::Array{Int,1},b::Array{Int,1})
  if a == b
    return true
  end
  la = length(a)
  lb = length(b)
  if la > lb
    if a[1:lb] == b
      return true
    end
  else
    if b[1:la] == a
      return true
    end
  end
  return false
end

endswith(a::Array{Int,1},b::Integer) = endswith(a,digits(b))
endswith(a::Integer,b::Array{Int,1}) = endswith(digits(a),b)
endswith(a::Integer,b::Integer) = endswith(digits(a),digits(b))

function startswith(a::Array{Int,1},b::Array{Int,1})
  if a == b
    return true
  end
  la = length(a)
  lb = length(b)
  if la > lb
    if a[end-lb+1:end] == b
      return true
    end
  else
    if b[end-la+1:end] == a
      return true
    end
  end
  return false
end

startswith(a::Array{Int,1},b::Integer) = startswith(a,digits(b))
startswith(a::Integer,b::Array{Int,1}) = startswith(digits(a),b)
startswith(a::Integer,b::Integer) = startswith(digits(a),digits(b))

function crop(l::Array{Int,1},i::Int)
  if abs(i) > length(l)
    error("Crop value too big!")
  elseif abs(i) == length(l)
    return [0]
  elseif i == 0
    return l
  elseif i > 0
    return l[1:end-i]
  else i < 0
    return l[1+abs(i):end]
  end
end

crop(n::Integer,i::Int) = undigit(crop(digits(n),i))

function replace!(l::Array{Int,1},idx::AbstractArray{Int,1},d::Array{Int,1})
  l[idx] = d
  return l
end

function replace!(l::Array{Int,1},old::Int,new::Int)
  l[l.==old] = new
  return l
end

replace(l::Array{Int,1},idx::AbstractArray{Int,1},d::Array{Int,1}) = replace!(deepcopy(l),idx,d)
replace(l::Array{Int,1},old::Int,new::Int) = replace!(deepcopy(l),old,new)

function replace(n::Integer,idx::AbstractArray{Int,1},d::Array{Int,1})
  return undigit(reverse(replace(reverse(digits(n)),idx,d)))
end

function replace(n::Integer,old::Int,new::Int)
  return undigit(replace(digits(n),old,new));
end

function select(l::Array{Int,1},idx::AbstractArray{Int,1})
  return l[idx]
end

function select(n::Integer,idx::AbstractArray{Int,1})
  return undigit(reverse(reverse(digits(n))[idx]))
end

function combine(a::Integer,b::Integer)
  if a == 0
    return b
  end
  f = 10^ndigits(b)
  return a*f+b
end

combine(a::Array{Int,1},b::Array{Int,1}) = [b;a]
combine(a::Integer,b::Vector{Int}) = combine(a,undigit(b))
combine(a::Vector{Int},b::Integer) = combine(a,digits(b))

crosssum(l::Array{Int,1}) = sum(l)
crosssum(n::Integer) = sum(digits(n))

function digitroot(n::Integer)
  x = abs(n)
  r = x-9*div(x-1,9)
  return copysign(r,n)
end

digitroot(l::Array{Int,1}) = digitroot(undigit(l))

end # module


include("./utils.jl")
include("../src/unionfind.jl")
using .UnionFindAlgo

PROJECT_PATH = joinpath(@__DIR__, "..")
tinyuf = loadufdata(joinpath(PROJECT_PATH, "data/tinyUF.txt"))
mediumuf = loadufdata(joinpath(PROJECT_PATH, "data/mediumUF.txt"))

@testset "Union find" begin
    uf = UnionFind(tinyuf)
    @test uf.ncomp == 2
    @test isconnected(uf, 3, 4) == false
    @test isconnected(uf, 4, 5) == true
    @test isconnected(uf, 9, 10) == true

    uf = UnionFind(mediumuf)
    @test uf.ncomp == 3
end


mutable struct KeyValuePair{T, S}
    key::T
    priority::S
end

mutable struct PriorityQueue{T<:KeyValuePair} #<: AbstractQueue{T}
    heap::BinaryHeap{T}
    function PriorityQueue{T}(data) where T
        heap = BinaryHeap{T}(data, (x,y)->isless(y.priority, x.priority))
        new(heap)
    end
end

PriorityQueue(data::Vector) = PriorityQueue{eltype(data)}(data)

function enqueue!(q::PriorityQueue, item)
    insert!(q.heap, item)
end

function dequeue!(q::PriorityQueue)
    remove!(q.heap)
end

length(p::PriorityQueue) = length(p.heap)
isempty(p::PriorityQueue) = isempty(p.heap)


#using Lint

function readSMAT(filename::AbstractString)
    (rows,header) = readdlm(filename;header=true)
    A = sparse(
               convert(Array{Int64,1},rows[1:parse(Int,header[3]),1])+1,
               convert(Array{Int64,1},rows[1:parse(Int,header[3]),2])+1,
               convert(Array{Int64,1},rows[1:parse(Int,header[3]),3]),
               parse(Int,header[1]),
               parse(Int,header[2])
               )
    return A
end

type MatrixNetworkMetadata
    A::SparseMatrixCSC{Int64,Int64}
    labels::Vector{AbstractString}
    xy::Array{Float64,2}
    source::AbstractString
end

function load_matrix_network_all(name::AbstractString)
    A = load_matrix_network(name)
    pathname = joinpath(Pkg.dir("MatrixNetworks"),"data")

    meta_source = joinpath(pathname,"$(name).source")
    if isfile(meta_source)
        source = open(readall, meta_source)
    else
        source = "(None given)"
    end

    meta_xy = joinpath(pathname,"$(name).xy")
    if isfile(meta_xy)
        xy = readdlm(meta_xy)
    else
        xy = zeros(0,2)
    end

    meta_labels = joinpath(pathname,"$(name).labels")
    if isfile(meta_labels)
        labels = open(readlines, meta_labels)
    else
        labels = map(x -> @sprintf("%i",x), 1:size(A,1))
    end

    return MatrixNetworkMetadata(A,labels,xy,source)
end

function load_matrix_network(name::AbstractString)
    pathname = joinpath(Pkg.dir("MatrixNetworks"),"data")
    smatfile = joinpath(pathname,"$(name).smat")
    if isfile(smatfile)
        return readSMAT(smatfile)
    else
        error(@sprintf "The example datafile '%s' does not seem to exist where it should\n" name)
    end
end

function load_matrix_network_metadata(name::AbstractString)
    pathname = joinpath(Pkg.dir("MatrixNetworks"),"data")
    smatfile = joinpath(pathname,"$(name).smat")
    meta_xy = joinpath(pathname,"$(name).xy")
    meta_labels = joinpath(pathname,"$(name).labels")
    if isfile(smatfile)
        if isfile(meta_xy) && isfile(meta_labels)
            xy = readdlm(meta_xy)
            labels = readdlm(meta_labels)
            return (readSMAT(smatfile),xy,labels)
        end
    else
        error(@sprintf "The example datafile '%s' does not seem to exist where it should\n" name)
    end
end

function matrix_network_datasets()
    datasets_location = joinpath(Pkg.dir("MatrixNetworks"),"data")
    content = readdir(datasets_location)
    smat_files = filter(x->contains(x,".smat"),content)
    for i = 1:length(smat_files)
        # no need for this lintpragma anymore
        # @lintpragma( "Ignore use of undeclared variable end")
        smat_files[i] = smat_files[i][1:end-5]
    end
    return smat_files
end


"""
	MCprobability!(state::QState; kwargs...)
Get probability distribution of qubit configurations by Monte Carlo sampling.

# Argument
- `state`: Any QState with **measure!** defined
- `samplesize`: Keyword argument to specify sample size. By default `samplesize = 1024`.
- `percentage`: Keyword argument. By default `percentage = false`. Return counts of configurations if set to `false`; return percentage frequency if set to `true`.

# Return Value
Returns a dinctionary data structure. The key of each pair is the qubit configuration; the value of the pair is frequency of appearance in sampling.

# Example
``` julia
julia> state = MPSState(2)
2-qubit MPSState

julia> apply!(state, H, ActPosition(1))
2-qubit MPSState

julia> d =MCprobability!(state; samplesize=1000)
Dict{Int64,Real} with 2 entries:
  0 => 509
  1 => 491

julia> d =MCprobability!(state; samplesize=1000, percentage = true)
Dict{Int64,Real} with 2 entries:
  0 => 0.47168
  1 => 0.504883
```
"""
function MCprobability!(state::QState; kwargs...)
	samplesize = get(kwargs, :samplesize, 1024)
	result = measure!(state, [1:length(state);], samplesize; binary = false)
	return frequency(result; kwargs...)
end

function frequency(v::Vector; kwargs...)
	percentage = get(kwargs, :percentage, false)
 	samplesize = get(kwargs, :samplesize, length(v))
 	cutoff = get(kwargs, :cutoff, 0)
 	abscutoff = get(kwargs, :abscutoff, 0)
 	if abscutoff !=0 && cutoff !=0
 		error("cannot specify both cutoff and abscutoff")
 	end
 	abscutoff ==0 && (abscutoff = cutoff*samplesize)
	dict = Dict{Int, Real}()
	for key ∈ v
		push!(dict, key => get!(dict, key, 0)+1)
	end
	for (k,s) ∈ dict
		if s <= abscutoff
			delete!(dict, k)
		elseif percentage
			push!(dict, k => s/samplesize)
		end
	end

	return dict
end


function Jdist(x::Vector{Float64},Vs::Float64,vs::Float64,ts::Float64,C::Float64,
    W1end::Float64,c0::Float64,rho::Float64,beta::Float64,h::Float64)
    J = zeros(2,2);

    J[1,1] = C.*rho.*vs.^2./(2.*Vs).*((W1end./vs .- x[1])./8 .+ c0/vs).^2;
    J[1,2] = 1;
    J[2,1] = 2.*ts.*vs.^5.*rho.^2./(Vs.*beta.^2).*(((W1end./vs .- x[1])./8 .+ c0./vs).^3.*0.5.*
        (W1end./vs .+ x[1]) .- ((W1end./vs .- x[1])./8 .+ c0./vs).^4);
    J[2,2] = ts./h;

    return J
end


# Generate scenarios of the demands for use within the J-SAA and J+-SAA frameworks

function generateJSAAScenarios(outputs::Array{Float64,2},inputs::Array{Float64,2},covariates_obs::Array{Float64,1},regressionMethod::String)

	numDependentVar::Int64 = size(outputs,2)
	numIndependentVar::Int64 = size(inputs,2)
	numSamples::Int64 = size(outputs,1)
	demand_scen_JSAA = zeros(Float64,numSamples,numDependentVar)
	demand_scen_JpSAA = zeros(Float64,numSamples,numDependentVar)

	if(size(outputs,1) != size(inputs,1))
		throw(ErrorException("Number of output samples != number of input samples"))
	end

	if(regressionMethod == "ols")

		if(numSamples < numIndependentVar)
			throw(ErrorException("Expected number of samples > number of covariates in OLS!"))
		end

		XtX_fact = lufact((inputs')*inputs)
		XtY = (inputs')*outputs

		# first, compute the OLS coefficients using all data points
		coeff_full = zeros(Float64,numIndependentVar,numDependentVar)
		for d = 1:numDependentVar
			coeff_full[:,d] = XtX_fact\(XtY[:,d])
		end

		# next, compute the leave-one-out coefficients
		for s = 1:numSamples
			coeff_loo = zeros(Float64,numIndependentVar,numDependentVar)

			if(numSamples == numIndependentVar)
				range_loo = union([1:s-1,s+1:numSamples]...)
				for d = 1:numDependentVar
					coeff_loo[:,d] = (inputs[range_loo,:])\(outputs[range_loo,d])
				end
			else
				rhs_pert = -inputs[s,:]*((outputs[s,:])')
				term1 = XtX_fact\(inputs[s,:])
				for d = 1:numDependentVar
					term2 = XtX_fact\(rhs_pert[:,d])
					term3 = coeff_full[:,d] + term2

					coeff_loo[:,d] = term3 + ((((inputs[s,:])')*term3)/(1.0 - ((inputs[s,:])')*term1))*term1
				end
			end

			demand_scen_JSAA[s,:] = ((covariates_obs')*coeff_full)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
			demand_scen_JpSAA[s,:] = ((covariates_obs')*coeff_loo)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
		end

	elseif(regressionMethod == "lasso")

		# first, compute the OLS coefficients using all data points
		coeff_full = getLassoCoeff(outputs,inputs)

		# next, compute the leave-one-out coefficients
		for s = 1:numSamples
			range_loo = union([1:s-1,s+1:numSamples]...)
			coeff_loo = getLassoCoeff(outputs[range_loo,:],inputs[range_loo,:])

			demand_scen_JSAA[s,:] = ((covariates_obs')*coeff_full)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
			demand_scen_JpSAA[s,:] = ((covariates_obs')*coeff_loo)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
		end

	end

	return demand_scen_JSAA, demand_scen_JpSAA
end


mutable struct REINFORCE{T,TP,TV} <: AbstractPolicyAgent where{T<:Real,TP<:AbstractPolicy,TV<:AbstractFuncApprox}
    π::TP
    vf::TV
    αₚ::T
    αᵥ::T
    γ::T
    avgGrad::Bool
    discountGrad::Bool
    states::Vector
    actions::Vector
    rewards::Array{T, 1}
    function REINFORCE(π::AbstractPolicy, vf::AbstractFuncApprox, αₚ::T, αᵥ::T, γ::T, avgGrad::Bool=false, discountGrad::Bool=true) where {T<:Real}
        new{T,typeof(π),typeof(vf)}(clone(π), clone(vf), αₚ, αᵥ, γ, avgGrad, discountGrad, [], [], T[])
    end
end

function act!(agent::REINFORCE, env::AbstractEnvironment, rng::AbstractRNG)
    push!(agent.states, env.state)

    get_action!(agent.π, env.state, rng)
    push!(agent.actions, agent.π.action)

    reward = step!(env, agent.π.action, rng)
    push!(agent.rewards, reward)

    if is_terminal(env)
        Gs = compute_returns(agent.rewards, agent.γ)
        update_withBaseline!(agent, Gs)
        new_episode!(agent, rng)
    end
end

function compute_returns(rewards::Array{T, 1}, γ::T)::Array{T, 1} where {T<:Real}
    G = 0.
    Gs = zeros(T, length(rewards))
    for t in length(rewards):-1:1
        G = rewards[t] + γ * G
        Gs[t] = G
    end
    return Gs
end

function update_withBaseline!(agent::REINFORCE, Gs::Array{T, 1}) where {T<:Real}
    γₜ = 1.0
    gₚ = zeros(get_num_params(agent.π))
    gᵥ = zeros(get_num_params(agent.vf))
    ∇ = zeros(get_num_params(agent.π))
    for t in 1:length(agent.states)
        v = call_gradient!(gᵥ, agent.vf, agent.states[t], 1)
        δ = Gs[t] - v
        @. gᵥ *= δ * agent.αᵥ
        add_to_params!(agent.vf, gᵥ)

        gradient_logp!(gₚ, agent.π, agent.states[t], agent.actions[t])
        @. ∇ += δ * agent.αₚ * γₜ * gₚ

        if agent.discountGrad
            γₜ *= agent.γ
        end
    end
    if agent.avgGrad
        ∇ ./= length(agent.states)
    end
    add_to_params!(agent.π, ∇)
end

function new_episode!(agent::REINFORCE, rng::AbstractRNG)
    empty!(agent.states)
    empty!(agent.actions)
    empty!(agent.rewards)
end

function clone(agent::REINFORCE)::REINFORCE
    a = REINFORCE(agent.π, agent.vf, agent.αₚ, agent.αᵥ, agent.γ, agent.avgGrad, agent.discountGrad)
    a.states = deepcopy(agent.states)
    a.actions = deepcopy(agent.actions)
    a.rewards = deepcopy(agent.rewards)
    return a
end


import sbp
using MAT
using LinearAlgebra
include("../common/wave_equation.jl")
include("../common/bound.jl")
include("../common/discretization.jl")
disc = Discretization
pde=WaveEquation

datafile = ARGS[1]
nx = 100
nsamples = 100
alpha = range(0.0, stop=10.0, length=nsamples)
order = 4

Mnorm = zeros(nsamples)
sigma_norm_upper = zeros(nsamples)
sigma_norm_lower = zeros(nsamples)

ops = disc.OperatorsType(nx, order)
h = ops.h

M0 = disc.discretization(pde.A(), ops)

sigma_d = LinearAlgebra.svd(Matrix(ops.Dx) * h)
rd = sigma_d.S[2]
sd = sigma_d.S[1]
rl = ops.Hxi[1,1] * h
sl = rl

# Left penalty term
Pl = disc.projection_matrix(pde.A(pde.LEFT), zeros(1,1), zeros(1,1))
Ml, sat_term_data = pde.sat(pde.A(pde.LEFT), pde.LEFT, ops, Pl)

Sl =  pde.A() * Pl

println("Estimated constants: $rd <= h\\|D\\|_2 <= $sd")
println("Estimated constants: $rl <= h\\|L\\|_2 <= $sl")

i = 1
for alphai in alpha

        # Right penalty term
        Sr, Mr, sat_term_data = pde.naive_leakage(alphai, ops)

        M = M0 + Ml + Mr

        sigma_norm_upper[i] = LinearAlgebra.svd(Matrix(Sr)).S[1] + LinearAlgebra.svd(Matrix(Sl)).S[1]
        sigma_norm_lower[i] = LinearAlgebra.svd(Matrix(Sr)).S[1] - LinearAlgebra.svd(Matrix(Sl)).S[1]

        eig_max = maximum(disc.check_stability(M))
        if (eig_max > 0)
                println("Warning: Unstable. eig_max = $eig_max")
        end
        Mnorm[i] = h * LinearAlgebra.opnorm(M)
        println("\\|M\\|_2 = ", Mnorm[i], " ", sigma_norm_lower[i], " <= \\|Sigma\\|_2 <= ", sigma_norm_upper[i])
        global i += 1
end
matwrite(datafile, Dict("alpha" => collect(alpha),
                                   "Mnorm" => Mnorm,
                                   "rd" => rd,
                                   "sd" => sd,
                                   "rl" => rl,
                                   "sl" => sl,
                                   "sigma_norm_upper" => sigma_norm_upper,
                                   "sigma_norm_lower" => sigma_norm_lower)
         )


# This file includes an example script to plot the basis functions of PSK modulator.

using DigitalCommunications
using Plots; plotlyjs()

# Simulation parameters
k = 2
M = 2^k
nsymbols = 2
nbits = k * nsymbols
ebno = collect(0 : 10)          # Snr per bit
esno = ebno .+ 10 * log10(k)    # Snr ber symbol

# Communcation system components
gen = BitGenerator(nbits)
coding = GrayCoding(M)
modulator = Modulator(PSK(M))
channel = WaveformAWGNChannel(esno[1], basis(modulator), modulator.tsample, modulator.pulse.duration)

# Simulate the system
@run gen.bits |> coding |> modulator
tx = gen.bits |> coding |> modulator
rx = tx |> channel

# Plots
sps = Int(modulator.pulse.duration / modulator.tsample)
t = collect(Iterators.partition(0 : modulator.tsample : modulator.pulse.duration * nsymbols, sps))
plt = plot()
for (ti, txi, rxi) in zip(t, tx, rx)
    plot!(ti, txi)
    plot!(ti, rxi)
end
display(plt)


using ValueOrientedRiskManagementInsurance
using DataFrames
using Base.Test

### TESTS
println("start test ChainLadder_Test.jl...")
# We use the example from Mack's 1993-paper
# Cumulated claim amounts
cum_triangle =
  convert(Array{Real,2},
          readtable("test/ChainLadderMack_Input.csv", header = false))

triangle =cum2claims(cum_triangle)
@test_approx_eq(cum_triangle, claims2cum(triangle))

mack = Mack(cum_triangle, cum=true)

# Test: development factor
@test_approx_eq(
  [3.491, 1.747, 1.457, 1.174, 1.104, 1.086, 1.054, 1.077, 1.018],  round(mack.f,3))
# Test: future claims
@test_approx_eq(sum(mack.futureclaims), mack.tot_res)
# Test: standard error
@test_approx_eq(
  [0.80, 0.26, 0.19, 0.27, 0.29, 0.26, 0.22, 0.23, 0.29],
  round(√(mack.mse) ./ mack.res,2)[2:end])
# Test: total standard error
@test_approx_eq(0.13, round(√(mack.tot_mse)/mack.tot_res, 2))
println("...end test ChainLadder_Test.jl")


module UnscentedKalmanFilter
export UKFState, predict!, update!

using LinearAlgebra: I, dot, diagm

include("Helpers.jl")
using ..SigmaPoints
using .Helpers

"""
# State of the Unscented Kalman Filter

- `dim_x::Int64`: State dimension
- `dim_z::Int64`: Measurement dimension
- `x::AbstractArray{Float64, 1}`: State of the filter
- `P::AbstractArray{Float64, 2}`: Covariance matrix of the filter.
- `K::AbstractArray{Float64, 2}`: Kalman gain of the filter.
- `S::AbstractArray{Float64, 2}`: System uncertainty.
- `S_inv::AbstractArray{Float64, 2}`: Inverse of system uncertainty.
- `σ_parameters::MerweScaled`: Parameters for Van der Merwe sigma points.
- `σ_fx::AbstractArray{Float64, 2}`: Sigma points after `fx`.
- `σ_hx::AbstractArray{Float64, 2}`: Sigma points after `hx`.
- `add_x::Function`: Function to calculate addition between two states.
- `residual_x::Function`: Function to calculate residual between two states.
- `residual_z::Function`: Function to calculate residual between state and measurement.
- `mean_x::Function`: Function to calculate mean between Σ-points and weights
    after passing them through `fx`.
- `mean_z::Function`: Function to calculate mean between Σ-points and weights
    after passing them through `hx`.
"""
struct UKFState
    dim_x::Int64
    dim_z::Int64

    x::AbstractArray{Float64, 1}
    P::AbstractArray{Float64, 2}

    K::AbstractArray{Float64, 2}
    y::AbstractArray{Float64, 1}

    S::AbstractArray{Float64, 2}
    S_inv::AbstractArray{Float64, 2}

    σ_parameters::MerweScaled
    σ_fx::AbstractArray{Float64, 2}
    σ_hx::AbstractArray{Float64, 2}

    add_x::Function
    residual_x::Function
    residual_z::Function
    mean_x::Function
    mean_z::Function
end

"""
```julia
UKFState(
    ;dim_x::Int64, dim_z::Int64, σ_parameters::MerweScaled,
    add_x::Function = default_add_x,
    residual_x::Function = default_residual,
    residual_z::Function = default_residual,
    mean_x::Function = default_mean_f,
    mean_z::Function = default_mean_f,
)
```

# Arguments
- `dim_x::Int64`: State dimension
- `dim_z::Int64`: Measurement dimension
- `σ_parameters::MerweScaled`: Parameters for Van der Merwe sigma points.
- `add_x::Function`:
    Function to calculate addition between two states.
    Should accept two positional arguments `y` and `x`
    both of type `AbstractArray{Float64, 1}` and return `AbstractArray{Float64, 1}` as well.
- `residual_x::Function`: Function to calculate residual between two states.
    Should accept two positional arguments `y` and `x`,
    but should have methods both for `AbstractArray{Float64, 1}`, `AbstractArray{Float64, 1}`
    and for `AbstractArray{Float64, 2}`, `AbstractArray{Float64, 1}` types of arguments.
    In second case, residual will be calculated between `Σ`-points
    and state `x` and return matrix of the same shape as `Σ`.
- `residual_z::Function`: Function to calculate residual between state and measurement.
    Same as for `residual_x`, you should define methods for both kinds
    of arguments.
- `mean_x::Function`: Function to calculate mean between `Σ`-points
    and weights after passing them through `fx`.
    Arguments are positional.
- `mean_z::Function`: Function to calculate mean between `Σ`-points
    and weights after passing them through `hx`.
    Arguments are positional.

!!! note
    Do not forget to implement methods for both kinds of arguments for
    `residual_x` and/or `residual_z` as described above.
"""
function UKFState(
    ;dim_x::Int64, dim_z::Int64, σ_parameters::MerweScaled,
    add_x::Function = default_add_x,
    residual_x::Function = default_residual,
    residual_z::Function = default_residual,
    mean_x::Function = default_mean_f,
    mean_z::Function = default_mean_f,
)
    x = zeros(Float64, dim_x)
    P = Matrix{Float64}(I, dim_x, dim_x)
    K = zeros(Float64, dim_x, dim_z)
    y = zeros(Float64, dim_z)
    σ_fx = zeros(Float64, num_σ_points(σ_parameters), dim_x)
    σ_hx = zeros(Float64, num_σ_points(σ_parameters), dim_z)
    S = zeros(Float64, dim_z, dim_z)
    S_inv = zeros(Float64, dim_z, dim_z)

    UKFState(
        dim_x, dim_z,
        x, P, K, y, S, S_inv,
        σ_parameters, σ_fx, σ_hx,
        add_x, residual_x, residual_z, mean_x, mean_z
    )
end

function Base.show(io::IO, ukf::UKFState)
    print(
        io, "UKF State:\n",
        "dim_x: ", repr(MIME("text/plain"), ukf.dim_x, context=io), "\n",
        "dim_z: ", repr(MIME("text/plain"), ukf.dim_z, context=io), "\n",
        "x: ", repr(MIME("text/plain"), ukf.x, context=io), "\n",
        "P: ", repr(MIME("text/plain"), ukf.P, context=io), "\n",
        "K: ", repr(MIME("text/plain"), ukf.K, context=io), "\n",
        "y: ", repr(MIME("text/plain"), ukf.y, context=io), "\n",
        "S: ", repr(MIME("text/plain"), ukf.S, context=io), "\n",
        "S_inv: ", repr(MIME("text/plain"), ukf.S_inv, context=io), "\n",
        "σ_parameters: ", repr(MIME("text/plain"), ukf.σ_parameters, context=io), "\n",
        "σ_fx: ", repr(MIME("text/plain"), ukf.σ_fx, context=io), "\n",
        "σ_hx: ", repr(MIME("text/plain"), ukf.σ_hx, context=io), "\n",
        "add_x: ", repr(MIME("text/plain"), ukf.add_x, context=io), "\n",
        "residual_x: ", repr(MIME("text/plain"), ukf.residual_x, context=io), "\n",
        "residual_z: ", repr(MIME("text/plain"), ukf.residual_z, context=io), "\n",
        "mean_x: ", repr(MIME("text/plain"), ukf.mean_x, context=io), "\n",
        "mean_z: ", repr(MIME("text/plain"), ukf.mean_z, context=io), "\n",
    )
end

"""
Pass σ-points for the current UKF state `(x, P)` through `fx` function.
"""
function calculate_process_σ!(
    ;ukf_state::UKFState, δt::Float64, fx::Function, fx_args...
)
    σ_points = calculate_σ_points(
        σ_parameters=ukf_state.σ_parameters, x=ukf_state.x, P=ukf_state.P
    )
    @inbounds for i = 1:size(σ_points, 1)
        ukf_state.σ_fx[i, :] = fx(x=σ_points[i, :], δt=δt; fx_args...)
    end
end

"""
Calculate cross-variance between state and measurement.
"""
function cross_variance(
    ;ukf::UKFState, x::AbstractArray{Float64, 1}, z::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 2}
    Δx = ukf.residual_x(ukf.σ_fx, x)
    Δz = ukf.residual_z(ukf.σ_hx, z)
    Δx' * diagm(ukf.σ_parameters.Σ_w) * Δz
end

"""
```julia
unscented_transform(
    ;ukf=UKFState, Σ::AbstractArray{Float64, 2}, residual_f::Function,
    noise_cov::Union{Nothing, AbstractArray{Float64, 2}} = nothing,
)::Tuple{AbstractArray{Float64, 1}, AbstractArray{Float64, 2}}
```

Unscented transformation of Σ-points and their weights.
This is used to calculate new mean `x` and covariance `P`
for the state of Unscented Kalman Filter.

# Arguments
- `ukf::UKFState`: State of the UKF filter which contains weights for `Σ`.
- `Σ::AbstractArray{Float64, 2}`: `[2x+1, n]` matrix of Σ-points
- `residual_f::Function`: Function to calculate residual between `y` and `x`.
- `mean_f::Function`: Function to calculate mean between Σ points and weights.
- `noise_cov::Union{Nothing, AbstractArray{Float64, 2}} = nothing`:
    Optional noise, that will be added to final covariance matrix.
    Its size should match that of a covariance matrix.
"""
function unscented_transform(
    ;ukf=UKFState, Σ::AbstractArray{Float64, 2},
    residual_f::Function, mean_f::Function,
    noise_cov::Union{Nothing, AbstractArray{Float64, 2}} = nothing,
)::Tuple{AbstractArray{Float64, 1}, AbstractArray{Float64, 2}}
    x = mean_f(Σ, ukf.σ_parameters.m_w)
    Y = residual_f(Σ, x)
    P = Y' * diagm(ukf.σ_parameters.Σ_w) * Y
    if isa(noise_cov, AbstractArray{Float64, 2})
        P .+= noise_cov
    end
    x, P
end

"""
```julia
predict!(
    ;ukf::UKFState, δt::Float64, Q::Union{Nothing, AbstractArray{Float64, 2}},
    fx::Function, fx_args...
)
```

Perform prediction of the Unscented Kalman Filter (UKF).
This updates `x` and `P` of the `ukf`
to contain predicted state and covariance.

# Arguments
- `ukf::UKFState`: State of the UKF on which to perform prediction.
- `δt::Float64`: Time delta.
- `Q::Union{Nothing, AbstractArray{Float64, 2}}`: Process noise.
- `fx::Function(;x::AbstractArray{Float64, 1}, δt::Float64, ...)`:
    State transition function.
    Should accept keyword-only arguments, with mandatory `x` and `δt`.
    All other arguments may serve as control input or provide additional info.
- `fx_args...`: Arguments to pass to `fx` function.
    Must contain at least `x` and `δt` arguments.

!!! note
    Just to emphasize, `fx` should accept keyword-only arguments,
    with mandatory `x` and `δt`. All other arguments may serve
    as control input or provide additional info.
"""
function predict!(
    ;ukf::UKFState, δt::Float64, Q::Union{Nothing, AbstractArray{Float64, 2}},
    fx::Function, fx_args...
)
    calculate_process_σ!(ukf_state=ukf, δt=δt, fx=fx; fx_args...)
    ukf.x[:], ukf.P[:] = unscented_transform(
        ukf=ukf, Σ=ukf.σ_fx, noise_cov=Q,
        residual_f=ukf.residual_x, mean_f=ukf.mean_x
    )
    ukf.σ_fx .= calculate_σ_points(
        σ_parameters=ukf.σ_parameters, x=ukf.x, P=ukf.P
    )
end

"""
```julia
update!(
    ;ukf::UKFState, z::AbstractArray{Float64, 1}, R::Union{Nothing, AbstractArray{Float64, 2}},
    hx::Function, hx_args...
)
```

Perform update step of the Unscented Kalman Filter (UKF)
given measurements.
This updates `x` and `P` of the `ukf_state`
to contain updated state and covariance.

# Arguments
- `ukf::UKFState`: State of the UKF on which to perform update step.
- `z::AbstractArray{Float64, 1}`: Measurements.
- `R::Union{Nothing, AbstractArray{Float64, 2}}`: Measurement noise.
- `hx::Function`:
    Measurement function.
    Should accept keyword-only arguments with mandatory `x` argument for state.
    Given state `x` of UKF it must transform it to the respective measurement.
- `hx_args...`: Arguments that will be passed to `hx` function.
    Should contain at least `x` argument.

!!! note
    Just to emphasize, `hx` should accept keyword-only arguments
    with mandatory `x` argument for state.
    Given state `x` of UKF it must transform it to the respective measurement.
"""
function update!(
    ;ukf::UKFState, z::AbstractArray{Float64, 1},
    R::Union{Nothing, AbstractArray{Float64, 2}},
    hx::Function, hx_args...
)
    @inbounds for i = 1:size(ukf.σ_fx, 1)
        ukf.σ_hx[i, :] = hx(x=ukf.σ_fx[i, :]; hx_args...)
    end

    zp, ukf.S[:] = unscented_transform(
        ukf=ukf, Σ=ukf.σ_hx, noise_cov=R,
        residual_f=ukf.residual_z, mean_f=ukf.mean_z
    )
    ukf.S_inv .= inv(ukf.S)

    V = cross_variance(ukf=ukf, x=ukf.x, z=zp)
    ukf.K .= V * ukf.S_inv
    ukf.y .= ukf.residual_z(z, zp)

    ukf.x .= ukf.add_x(ukf.x, ukf.K * ukf.y)
    ukf.P .-= ukf.K * ukf.S * ukf.K'
end

end


function eftAddGTE{T<:StdFloat}(a::T,b::T,c::T)
    s,t = eftAddGTE(b, c)
    x,u = eftAddGTE(a, s)
    y,z = eftAddGTE(u, t)
    x,y = eftAddGTE(x, y)
    x,y,z
end

function eftAddGTEas2{T<:StdFloat}(a::T,b::T,c::T)
    s,t = eftAddGTE(b, c)
    x,u = eftAddGTE(a, s)
    y   = u+t
    x,y = eftAddGTE(x, y)
    x,y
end

function eftAdd{T<:StdFloat}(a::T,b::T,c::T)
    s,t = eftAdd(b, c)
    x,u = eftAdd(a, s)
    y,z = eftAdd(u, t)
    x,y = eftAddGTE(x, y)
    x,y,z
end

function eftAddAs2{T<:StdFloat}(a::T,b::T,c::T)
    s,t = eftAdd(b, c)
    x,u = eftAdd(a, s)
    y   = u+t
    x,y = eftAddGTE(x, y)
    x,y
end

function eftMul{T<:StdFloat}(a::T, b::T, c::T)
    p,e = eftMul(a,b)
    x,y = eftMul(p,c)
    z,t = eftMul(e,c)
    y,z = eftAdd(y,z)
    x,y = eftAdd(x,y)
    eftAddGTE(x,y,z,t)
end

function eftMulas3{T<:StdFloat}(a::T, b::T, c::T)
    p,e = eftMul(a,b)
    x,y = eftMul(p,c)
    z,t = eftMul(e,c)
    y,z = eftAdd(y,z)
    x,y = eftAdd(x,y)
    eftAddGTEas3(x,y,z,t)
end

function eftMulAs2{T<:StdFloat}(a::T, b::T, c::T)
    p,e = eftMul(a,b)
    x,y = eftMul(p,c)
    z   = e*c
    y,z = eftAdd(y,z)
    x,y = eftAdd(x,y)
    x,y
end

function eftFMA{T<:StdFloat}(a::T, b::T, c::T)
    x = fma(a,b,c)
    u1,u2 = eftMul(a,b)
    a1,a2 = eftAdd(u2,c)
    b1,b2 = eftAdd(u1,a1)
    g = (b1-x)+b2
    y,z = eftAddGTE(g,a2)
    x,y,z
end

function eftFMAas2{T<:StdFloat}(a::T, b::T, c::T)
    x = fma(a,b,c)
    u1,u2 = eftMul(a,b)
    a1,a2 = eftAdd(u2,c)
    b1,b2 = eftAdd(u1,a1)
    g = (b1-x)+b2
    y = g+a2
    x,y
end

function eftFMS{T<:StdFloat}(a::T, b::T, c::T)
    x = fma(a,b,-c)
    u1,u2 = eftMul(a,b)
    a1,a2 = eftSub(u2,c)
    b1,b2 = eftAdd(u1,a1)
    g = (b1-x)+b2
    y,z = eftAddGTE(g,a2)
    x,y,z
end

function eftFMSas2{T<:StdFloat}(a::T, b::T, c::T)
    x = fma(a,b,-c)
    u1,u2 = eftMul(a,b)
    a1,a2 = eftSub(u2,c)
    b1,b2 = eftAdd(u1,a1)
    g = (b1-x)+b2
    y = g+a2
    x,y
end

for fn in (:eftAdd, :eftAddGTE, :eftAddAs2, :eftMul, :eftMulAs2, :eftMulAs4, :eftFMA, :eftFMS, :eftFMAas2, :eftFMSas2)
  @eval begin
    ($fn){T1<:Real,T2<:StdFloat,T3<:StdFloat}(a::T1, b::T2, c::T3) = ($fn)(convert(Float64,a), b, c)
    ($fn){T1<:StdFloat,T2<:Real,T3<:StdFloat}(a::T1, b::T2, c::T3) = ($fn)(a, convert(Float64,b), c)
    ($fn){T1<:StdFloat,T2<:StdFloat,T3<:Real}(a::T1, b::T2, c::T3) = ($fn)(a, b, convert(Float64,c))
    ($fn){T1<:StdFloat,T2<:Real,T3<:Real}(a::T1, b::T2, c::T3) = ($fn)(a, convert(Float64,b), convert(Float64,c))
    ($fn){T1<:Real,T2<:StdFloat,T3<:Real}(a::T1, b::T2, c::T3) = ($fn)(convert(Float64,a), b, convert(Float64,c))
    ($fn){T1<:Real,T2<:Real,T3<:StdFloat}(a::T1, b::T2, c::T3) = ($fn)(convert(Float64,a), convert(Float64,b), c)
    ($fn){T1<:Real,T2<:Real,T3<:Real}(a::T1, b::T2, c::T3) = ($fn)(convert(Float64,a), convert(Float64,b), convert(Float64,c))
  end
end


function sbt(x)
  K = length(x)
  offset = (K + 1) .- cumsum(ones(K), dims=1)
  z = sigmoid.(x .- log.(offset))
  z_cumprod = cumprod(1 .- z, dims=1)
  # y = cat(z, pad, dims=dims) .* cat(pad, z_cumprod, dims=1)
  y = vcat(z, 1) .* vcat(1, z_cumprod)

  return y
end

function sbt_inv(p)
  # see: https://pytorch.org/docs/stable/_modules/torch/distributions/transforms.html
  error("unimplemented")
end

function sbt_logabsdetJ(x, p)
  K = length(x)

  offset = (K + 1) .- cumsum(ones(K), dims=1)
  z = sigmoid.(x .- log.(offset))
  detJ = sum(log.(1 .- z) .+ log.(p[1:end-1]))

  return detJ
end

#= TEST
x = param(randn(3))
p = sbt(x)
sbt_logabsdetJ(x, p)
logpdf(Dirichlet(ones(4)), p) + sbt_logabsdetJ(x, p)
=#


module GaussianWeightedMeanPrecisionTest

using Test
using ForneyLab
using ForneyLab: outboundType, isApplicable, isProper, unsafeMean, unsafeMode, unsafeVar, unsafeCov, unsafeMeanCov, unsafePrecision, unsafeMeanPrecision, unsafeWeightedMean, unsafeWeightedMeanPrecision
using ForneyLab: SPGaussianWeightedMeanPrecisionOutNPP, VBGaussianWeightedMeanPrecisionOut

@testset "dims" begin
    @test dims(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=0.0, w=1.0)) == ()
    @test dims(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=ones(2), w=diageye(2))) == (2,)
end

@testset "vague" begin
    @test vague(GaussianWeightedMeanPrecision) == ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=0.0, w=tiny)
    @test vague(GaussianWeightedMeanPrecision, (2,)) == ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=zeros(2), w=tiny*eye(2))
end

@testset "isProper" begin
    # Univariate
    @test isProper(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=0.0, w=1.0))
    @test !isProper(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=0.0, w=-1.0))

    # Multivariate
    @test isProper(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[0.0], w=mat(1.0)))
    @test isProper(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=ones(2), w=diageye(2)))
    @test !isProper(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[0.0], w=mat(-1.0)))
end

@testset "==" begin
    # Univariate
    @test ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=0.0, w=1.0) == ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=0.0, w=1.0)
    @test ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=0.0, w=1.0) == ProbabilityDistribution(Univariate, GaussianMeanVariance, m=0.0, v=1.0)

    # Multivariate
    @test ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[0.0], w=mat(1.0)) == ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[0.0], w=mat(1.0))
    @test ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[0.0], w=mat(1.0)) == ProbabilityDistribution(Multivariate, GaussianMeanVariance, m=[0.0], v=mat(1.0))
end

@testset "unsafe statistics" begin
    # Univariate
    @test unsafeMean(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == 0.5
    @test unsafeMode(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == 0.5
    @test unsafeVar(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == 0.25
    @test unsafeCov(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == 0.25
    @test unsafeMeanCov(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == (0.5, 0.25)
    @test unsafePrecision(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == 4.0
    @test unsafeMeanPrecision(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == (0.5, 4.0)
    @test unsafeWeightedMean(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == 2.0
    @test unsafeWeightedMeanPrecision(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)) == (2.0, 4.0)

    # Multivariate
    @test unsafeMean(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == [0.5]
    @test unsafeMode(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == [0.5]
    @test unsafeVar(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == [0.25]
    @test unsafeCov(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == mat(0.25)
    @test unsafeMeanCov(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == ([0.5], mat(0.25))
    @test unsafePrecision(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == mat(4.0)
    @test unsafeMeanPrecision(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == ([0.5], mat(4.0))
    @test unsafeWeightedMean(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == [2.0]
    @test unsafeWeightedMeanPrecision(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))) == ([2.0], mat(4.0))
end

@testset "log pdf" begin
    @test isapprox(logPdf(ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=2.0), 1.0), -0.5723649429247)
    @test isapprox(logPdf(ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0, 2.0], w=[2.0 0.0; 0.0 2.0]), [1.0, 0.0]), -2.1447298858494)
end

@testset "convert" begin
    @test convert(ProbabilityDistribution{Univariate, GaussianWeightedMeanPrecision}, ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=0.5, w=4.0)) == ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)
    @test convert(ProbabilityDistribution{Univariate, GaussianWeightedMeanPrecision}, ProbabilityDistribution(Univariate, GaussianMeanVariance, m=0.5, v=0.25)) == ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=2.0, w=4.0)
    @test convert(ProbabilityDistribution{Multivariate, GaussianWeightedMeanPrecision}, ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=[0.5], w=mat(4.0))) == ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))
    @test convert(ProbabilityDistribution{Multivariate, GaussianWeightedMeanPrecision}, ProbabilityDistribution(Multivariate, GaussianMeanVariance, m=[0.5], v=mat(0.25))) == ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[2.0], w=mat(4.0))
    @test convert(ProbabilityDistribution{Multivariate, GaussianWeightedMeanPrecision}, ProbabilityDistribution(Univariate, GaussianWeightedMeanPrecision, xi=1.0, w=2.0)) == ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[1.0], w=mat(2.0))
end

#-------------
# Update rules
#-------------

@testset "SPGaussianWeightedMeanPrecisionOutNPP" begin
    @test SPGaussianWeightedMeanPrecisionOutNPP <: SumProductRule{GaussianWeightedMeanPrecision}
    @test outboundType(SPGaussianWeightedMeanPrecisionOutNPP) == Message{GaussianWeightedMeanPrecision}
    @test isApplicable(SPGaussianWeightedMeanPrecisionOutNPP, [Nothing, Message{PointMass}, Message{PointMass}])
    @test !isApplicable(SPGaussianWeightedMeanPrecisionOutNPP, [Message{PointMass}, Nothing, Message{PointMass}])

    @test ruleSPGaussianWeightedMeanPrecisionOutNPP(nothing, Message(Univariate, PointMass, m=1.0), Message(Univariate, PointMass, m=2.0)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=1.0, w=2.0)
    @test ruleSPGaussianWeightedMeanPrecisionOutNPP(nothing, Message(Multivariate, PointMass, m=[1.0]), Message(MatrixVariate, PointMass, m=mat(2.0))) == Message(Multivariate, GaussianWeightedMeanPrecision, xi=[1.0], w=mat(2.0))
end

@testset "VBGaussianWeightedMeanPrecisionOut" begin
    @test VBGaussianWeightedMeanPrecisionOut <: NaiveVariationalRule{GaussianWeightedMeanPrecision}
    @test outboundType(VBGaussianWeightedMeanPrecisionOut) == Message{GaussianWeightedMeanPrecision}
    @test isApplicable(VBGaussianWeightedMeanPrecisionOut, [Nothing, ProbabilityDistribution, ProbabilityDistribution])
    @test !isApplicable(VBGaussianWeightedMeanPrecisionOut, [ProbabilityDistribution, ProbabilityDistribution, Nothing])

    @test ruleVBGaussianWeightedMeanPrecisionOut(nothing, ProbabilityDistribution(Univariate, GaussianMeanVariance, m=1.0, v=2.0), ProbabilityDistribution(Univariate, PointMass, m=3.0)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=1.0, w=3.0)
    @test ruleVBGaussianWeightedMeanPrecisionOut(nothing, ProbabilityDistribution(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0)), ProbabilityDistribution(MatrixVariate, PointMass, m=mat(3.0))) == Message(Multivariate, GaussianWeightedMeanPrecision, xi=[1.0], w=mat(3.0))
end

end # module


ishidden = s -> startswith(basename(s),'.')
isimg = s -> !ishidden(s) && any(endswith.(lowercase(s),[".tif",".tiff",".png",".jpg",".jpeg"]))

"""
	interpimage (img)
Returns a callable function that performs interpolation/extrapolation of a filtered image.
"""
function interpimage(img)
	X = interpolate(img, BSpline(Linear()))
	Z = extrapolate(X,0)  # extend by zero
	return (i,j) -> float(Z(i,j))
end

"Return data relevant for cornea detection to be used on the given image."
function detectiondata(img,m,n;options=get_defaults())
	# select the angles for optimization of the circle residual
	θ = 2π*(-180:180)/360
	select = falses(size(θ))
	for rng in options.angle_ranges
		select[ rng[1] .<= θ .<= rng[2] ] .= true
	end
	θ = θ[select]

	X = imresize(img,m,n)
	purk = findpurkinje(X)

	# optimization initializations
	u_init = [ (m/size(X,1)).*i for i in initvals(X,purk) ]  # "smart"
	append!(u_init,options.initializer(m,n))            # dumb

	G = green.(X)

	# try to screen out the eyelid lines, using bright blue values as the indicator.
	C = blue.(X)
	Cmax = mapwindow(maximum,C,(15,15))  # windowed max
	thresh = 0.75*maximum(C)
	G[@. (Cmax > thresh) & (G > 0.25) ] .= 0.33

	# Screen out purkinje, using a generous overestimate
	if !isempty(purk)
		iran,jran = [ extrema(i[k] for i in purk) for k in 1:2 ]
		h,w = iran[end]-iran[1],jran[end]-jran[1]
		cen = round.(Int,(median(iran),median(jran)))
		rows = clamp.(cen[1]-2h:cen[1]+2h,1,m)
		cols = clamp.(cen[2]-2w:cen[2]+2w,1,n)
		G[rows,cols] .= median(G[rows,cols])
		end

	# smooth and interpolate green channel
	k = m/options.blur_width
	ker = KernelFactors.gaussian((k,k))
	Z = interpimage(imfilter(G,ker))

	return Z,θ,u_init
end

function detect(sz,Z,θ,u_init;options=get_defaults())
	# Given data and multiple intitializations and options, try them all and find the best.
	u,fmin,best = [],Inf,[]
	for ui in u_init, method in options.optimizers
		unew,fnew = fitcircle(Z,sz...,ui...,method,θ;options)
		@debug "latest: $unew,$fnew,$ui,$method"
		if fnew < fmin
			u,fmin = unew,fnew
			best = (ui,method)
		end
	end
	return u,fmin,best
end

function detect(img::AbstractMatrix{T} where T <: Colorant,sz=size(img);options=get_defaults())
	detect(sz,detectiondata(img,sz...)...;options)
end

"""
	detect(files[,size])
	detect(img[,size])
Perform cornea detection on all the image files named in the vector `files`, or on the image `img`. If `size` is given, all images to be resized to match it.

Returns a named tuple with the following fields:
* `cenrow` row location of the cornea center, normalized by the image height
* `cencol` column location of the cornea center, normalized by the image height
* `radius` radius of the cornea center, normalized by the image height
* `fmin` optimized value of the optimization objective function
* `init` initial condition that led to the optimal result
* `method` info about the optimizer that got the optimal result
"""
function detect(folder::AbstractVector{String},sz=[];options=get_defaults())
	result = (cenrow = Float32[], cencol = Float32[], radius = Float32[], fmin = Float32[], init = [], method = [] )
	if isempty(sz)
		sz = size(load(folder[1]))
	end
	@showprogress for fname in folder
		img = load(fname)
		Z,θ,u_init = detectiondata(img,sz...;options)
		if !isempty(result.radius)
			new_ui = sz[1].*(result.cenrow[end],result.cencol[end],result.radius[end])
			push!(u_init,new_ui)
		end
		u,fmin,best = detect(sz,Z,θ,u_init;options)

		push!(result.cenrow,u[1]/sz[1])
		push!(result.cencol,u[2]/sz[1])
		push!(result.radius,u[3]/sz[1])
		push!(result.fmin,fmin)
		push!(result.init,best[1])
		push!(result.method,best[2])
	end
	return result
end


using Test, FinancialMonteCarlo, DualNumbers;
@show "Black Scholes Model"
S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

underlying_name = "ENJ"
underlying_name_2 = "ABPL"

underlying_ = underlying_name * "_" * underlying_name_2 * "_TESL"

Nsim = 10000;
Nstep = 30;
sigma = 0.2;
mc = MonteCarloConfiguration(Nsim, Nstep);
mc1 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.AntitheticMC());
toll = 0.8

rfCurve = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOptionND(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model_enj = BlackScholesProcess(sigma, Underlying(S0, d));
Model_abpl = BlackScholesProcess(sigma, Underlying(S0, d));
Model_tesl = BlackScholesProcess(sigma, Underlying(S0, d));
rho_1 = [1.0 0.0 0.1; 0.0 1.0 0.0; 0.1 0.0 1.0];
Model = GaussianCopulaNVariateProcess(rho_1, Model_enj, Model_abpl, Model_tesl)

display(Model)

mktdataset = underlying_ → Model
mktdataset_1 = underlying_name → Model_enj

portfolio_ = [EUData];
portfolio = underlying_ → EUData
portfolio_1 = underlying_name → AMData
price_mkt = FinancialMonteCarlo.delta(mktdataset, rfCurve, mc, portfolio, underlying_name, 1e-7)
price_mkt_1 = FinancialMonteCarlo.delta(mktdataset_1, rfCurve, mc, portfolio_1, underlying_name, 1e-7)

Model_enj_dual = BlackScholesProcess(sigma, Underlying(dual(S0, 1.0), d));
Model_dual = GaussianCopulaNVariateProcess(rho_1, Model_enj_dual, Model_abpl, Model_tesl)

mktdataset_dual = underlying_ → Model_dual
mktdataset_1_dual = underlying_name → Model_enj_dual
price_mkt_dual = pricer(mktdataset_dual, rfCurve, mc, portfolio)
price_mkt_1_dual = pricer(mktdataset_1_dual, rfCurve, mc, portfolio_1)

@test(abs(price_mkt_dual.epsilon - price_mkt) < 1e-4)
@test(abs(price_mkt_1_dual.epsilon - price_mkt_1) < 1e-4)
@test(abs(price_mkt_1_dual.epsilon - FinancialMonteCarlo.delta(Model_enj, rfCurve, mc, AMData, 1e-7)) < 1e-4)
@test(abs(price_mkt_1_dual.epsilon - FinancialMonteCarlo.delta(Model_enj, rfCurve, mc, AMData * 1.0, 1e-7)) < 1e-4)
@test(abs(price_mkt_1_dual.epsilon - FinancialMonteCarlo.delta(Model_enj, rfCurve, mc, [AMData], 1e-7)[1]) < 1e-4)

@test_throws(ErrorException, FinancialMonteCarlo.delta(mktdataset_1, rfCurve, mc, portfolio_1, underlying_, 1e-7))
@test(FinancialMonteCarlo.delta(mktdataset_1, rfCurve, mc, portfolio_1, "ciao", 1e-7) == 0.0)


## Example for using DSANet

# Make sure all the required packages are available
cd(@__DIR__)
using Pkg; Pkg.activate(".")
Pkg.instantiate()

@info "Loading packages"
using FluxArchitectures
using Random
using Plots

# Load some sample data
@info "Loading data"
poollength = 10
horizon = 15
datalength = 4000
input, target = get_data(:exchange_rate, poollength, datalength, horizon) |> gpu

# Define the network architecture
inputsize = size(input, 1)
local_length = 3
n_kernels = 3
d_model = 4
hiddensize = 1
n_layers = 3
n_head = 2

# Define the neural net
@info "Creating model and loss"
Random.seed!(123)
model = DSANet(inputsize, poollength, local_length, n_kernels, d_model,
               hiddensize, n_layers, n_head) |> gpu

# MSE loss
function loss(x, y)
    Flux.reset!(model)
    return Flux.mse(model(x), permutedims(y))
end

# Callback for plotting the training
cb = function ()
    Flux.reset!(model)
    pred = model(input) |> permutedims |> cpu
    Flux.reset!(model)
    p1 = plot(pred, label="Predict")
    p1 = plot!(cpu(target), label="Data", title="Loss $(loss(input, target))")
    display(plot(p1))
end

# Training loop
@info "Starting training"
@info "Start loss" loss = loss(input, target)
Flux.train!(loss, Flux.params(model),Iterators.repeated((input, target), 50),
             ADAM(0.005), cb=cb)

@info "Finished"
@info "Final loss" loss = loss(input, target)


    a = wted_chimera(102,2)
    spectral_drawing(a)


#aproximates liquid volume at a known pressure and t,
#by using isothermal compressibility

"""
    volume_compress(model,p,T,z=SA[1.0];V0=x0_volume(model,p,T,z,phase=:liquid),max_iters=100)

Main routine to calculate a volume, given a pressure, temperature, composition and intitial volume guess. each step is taken by locally aproximating the EoS as an isothermal compressibility process.
The new volume is calculated by the following recurrence formula:

```julia
v[i+1] = v[i]*exp(β[i]*(p-p(v[i])))
```

In the liquid root region, the iterations follow `v0 < v[i] < v[i+1] < v(p)`, allowing the calculation of the liquid root without entering the metastable region.
"""
function volume_compress(model,p,T,z=SA[1.0];V0=x0_volume(model,p,T,z,phase=:liquid),max_iters=100)
    _0 = zero(p+T+first(z))
    _nan = _0/_0
    logV0 = log(V0)
    lb_v = lb_volume(model,z)
    function logstep(_V)
        _V < log(lb_v) && return zero(_V)/zero(_V)
        _V = exp(_V)
        _p,dpdV = p∂p∂V(model,_V,T,z)
        _Δ = (p-_p)/(_V*dpdV)
        return _Δ
    end

    function f_fixpoint(_V)
        Δ = logstep(_V)
        _V + sign(Δ)*abs(Δ)^(one(Δ)+Δ)
    end

    Δ0 = logstep(logV0)
    Δ00 = Δ0
    Δ1 = logstep(logV0+Δ0)
    logV1 = logV0
    iter = 0

    #convergence hill. we do Successive substitution until we land on a
    #good initial point for aitken acceleration
    for i in 1:20
        if abs(Δ00) > abs(Δ1)
            break
        end
        if !isfinite(Δ1)
            return _nan
        end
        Δ0 = Δ1
        Δ1 = logstep(logV1+Δ0)
        logV1 = logV1 + Δ1
        iter +=1
    end
    if iter == 20
        return _nan
    end
    res = @nan(Solvers.fixpoint(f_fixpoint,logV1,Solvers.SSFixPoint(),rtol = 1e-12,max_iters=max_iters),_nan)

    return exp(res)
end

"""
    volume_virial(model::EoSModel,p,T,z=SA[1.0])
    volume_virial(B::Real,p,T,z=SA[1.0])

Calculates an aproximation to the gas volume at specified pressure, volume and composition, by aproximating:

```julia

Z(v) ≈ 1 + B(T)/v
```
where `Z` is the compressibility factor and `B` is the second virial coefficient.
If `B>0`, (over the inversion temperature) returns `NaN`. If the solution to the problem is complex (`Z = 1 + B/v` implies solving a quadratic polynomial), returns `-2*B`.

If you pass an `EoSModel` as the first argument, `B` will be calculated from the EoS at the input `T`. You can provide your own second virial coefficient instead of a model.

"""
function volume_virial end

function volume_virial(model::EoSModel,p,T,z=SA[1.0])
    B = second_virial_coefficient(model,T,z)
    return volume_virial(B,p,T,z)
end

function volume_virial(B::Real,p,T,z=SA[1.0])
    _0 = zero(B)
    B > _0 && return _0/_0
    a = p/(R̄*T)
    b = -1
    c = -B
    Δ = b*b-4*a*c
    n = sum(z)
    if Δ <= 0
        #virial approximation could not be calculated
        #return value at spinodal
        return -2*B
    end
    #only the left root has physical meaning
    return (-b + sqrt(b*b-4*a*c))/(2*a)
end

#(z = pV/RT)
#(RT/p = V/z)
"""
    volume(model::EoSModel,p,T,z=SA[1.0];phase=:unknown,threaded=true)

calculates the volume (m³) of the compound modelled by `model` at a certain pressure,temperature and moles.

`phase` is a Symbol that determines the initial volume root to look for:

- If `phase =:liquid` it will return the volume of the phase using a liquid initial point.

- If `phase =:vapor` it will return the volume of the phase using a gas initial point.

The default is `phase =:unknown`. with this, both liquid and volume roots will be calculated, and
the phase with the least amount of energy is returned.

The calculation of both volume roots can be calculated in serial (`threaded=false`) or in parallel (`threaded=true`)

"""
function volume(model::EoSModel,p,T,z=SA[1.0];phase=:unknown,threaded=true)

    fp(_V) = log(pressure(model,_V,T,z)/p)

#Threaded version
    phase = Symbol(phase)
    if phase != :unknown
        V0 = x0_volume(model,p,T,z,phase=phase)
        #return Solvers.ad_newton(fp,Vg0)
        return volume_compress(model,p,T,z,V0=V0)
    end

    if threaded
            Vg0 = x0_volume(model,p,T,z,phase=:v)
            Vl0 = x0_volume(model,p,T,z,phase=:l)
            _Vg = Threads.@spawn volume_compress(model,$p,$T,$z;V0=Vg0)
            _Vl = Threads.@spawn volume_compress(model,$p,$T,$z;V0=$Vl0)
            Vg = fetch(_Vg)
            Vl = fetch(_Vl)
    else
            Vg0 = x0_volume(model,p,T,z,phase=:v)
            Vl0 = x0_volume(model,p,T,z,phase=:l)
            Vg =  volume_compress(model,p,T,z,V0=Vg0)
            Vl =  volume_compress(model,p,T,z,V0=Vl0)
    end

# Serial version
#=
    Vg0 = volume_virial(model,p,T,z)
    Vg =Solvers.ad_newton(fp,Vg0,rtol=1e-08)

    Vl =  volume_compress(model,p,T,z)
=#
    function gibbs(V)
        _df,f =  ∂f(model,V,T,z)
        dV,dt = _df
        if abs((p+dV)/p) > 0.03
            return Inf
        else
            return f  +p*V
        end
    end
    #this catches the supercritical phase as well

    isnan(Vl) && return Vg
    isnan(Vg) && return Vl
    (isnan(Vl) & isnan(Vg)) && error("Failed to converge to a root")

    gg = gibbs(Vg)
    gl = gibbs(Vl)
    #@show Vg,Vl
    #@show gg,gl
    return ifelse(gg<gl,Vg,Vl)

end

function volume(model::IdealModel,p,T,z=SA[1.0];phase=:unknown,threaded=false)
    return sum(z)*R̄*T/p
end

export volume


abstract type AbstractCivecm end

eigvals(obj::AbstractCivecm) = eigvals(convert(VAR, obj))

function loglikelihood(obj::AbstractCivecm)
    nobs = size(obj.endogenous, 1) - obj.lags
    -nobs*(logdet(cholesky!(residualvariance(obj))) - size(obj.endogenous, 2)*(log(nobs) - 1 - log(2π)))/2
end
# loglikelihood(A::StridedMatrix) = -0.5*size(A, 1)*(2sum(log(abs(diag(qr(A).factors)/sqrt(size(A,1))))) - size(A, 2)*(log(size(A, 1)) - 1 - log(2π)))
loglikelihood(A::StridedMatrix) =
    -size(A, 1)*(logdet(Statistics.covzm(A, 1, corrected=false)) - size(A, 2)*(log(size(A, 1)) - 1 - log(2π)))/2

# aic(obj::Civecm) = 2*(npars(obj) - loglikelihood(obj))

function residualvariance(obj::AbstractCivecm)
    mresiduals = residuals(obj)
    mOmega = mresiduals'mresiduals / size(mresiduals, 1)
    return mOmega
end

normalitytest(obj::AbstractCivecm) = normalitytest(residuals(obj))

## LR test
mutable struct LRTest{T<:AbstractCivecm}
    H0::T
    HA::T
    value::Float64
    df::Int64
end

lrtest(obj0::AbstractCivecm, objA::AbstractCivecm, df::Integer) = LRTest(obj0, objA, 2*(loglikelihood(objA) - loglikelihood(obj0)), df)

function bootstrap(obj::LRTest, reps::Integer, simH0 = true)
    lrvals = Vector{Float64}(undef, reps)
    bootH0 = copy(obj.H0)
    bootHA = copy(obj.HA)
    bootHA.endogenous .= bootH0.endogenous
    bootVAR = simH0 ? convert(VAR, bootH0) : convert(VAR, bootHA)
    iT, p = size(obj.H0.Z0)
    # simResiduals = residuals(obj.H0)
    simResiduals = residuals(obj.HA)
    mbr = mean(simResiduals, dims=1)
    bi = Vector{Float64}(undef, iT)
    bootResiduals = similar(simResiduals)
    for i = 1:reps
        bi[:] = randn(iT)
        bootResiduals[:] = copy(simResiduals)
        for k = 1:p
            for j = 1:iT
                 bootResiduals[j,k] -= mbr[k]
                 bootResiduals[j,k] *= bi[j]
                 bootResiduals[j,k] += mbr[k]
            end
        end
        bootH0.endogenous[:] = simulate(bootVAR, bootResiduals)
        auxilliaryMatrices!(bootH0)
        auxilliaryMatrices!(bootHA)
        estimate!(bootH0)
        estimate!(bootHA)
        lrvals[i] = -2*(loglikelihood(bootH0) - loglikelihood(bootHA))
    end
    return lrvals
end


module TestNoEmptyModules

   using Base.Test
   using Brim

   A = [1 0 0; 0 1 0; 0 0 1]
   S = [1 0 0; 0 1 0; 0 1 0]

   M = no_empty_modules!(Modular(A, S))

   @test size(M.S)[1] == 3
   @test size(M.S)[2] == 2
end


"""
    surf(data::Array{<:Real,2}, target::Array{<:Integer,1}, m::Signed=-1,
         dist_func::Any=(e1, e2) -> sum(abs.(e1 .- e2), dims=2);
         f_type::String="continuous")::Array{Float64,1}

Compute feature weights using SURF algorithm. The f_type argument specifies whether the features are continuous or discrete
and can either have the value of "continuous" or "discrete".

---
# Reference:
- Casey S. Greene, Nadia M. Penrod, Jeff Kiralis, and Jason H. Moore.
Spatially uniform ReliefF (SURF) for computationally-efficient filtering
of gene-gene interactions. BioData mining, 2(1):5–5, Sep 2009.
"""
function surf(data::Array{<:Real,2}, target::Array{<:Integer,1}, m::Signed=-1,
              dist_func::Any=(e1, e2) -> sum(abs.(e1 .- e2), dims=2);
              f_type::String="continuous")::Array{Float64,1}

    # Initialize feature weights vector.
    weights = zeros(Float64, 1, size(data, 2))

    # Compute vectors of maximum and minimum feature values.
    max_f_vals = maximum(data, dims=1)
    min_f_vals = minimum(data, dims=1)

    # Sample m samples without replacement. If m has signal value -1, use all samples.
    sample_idxs = StatsBase.sample(1:size(data, 1), if (m==-1) size(data,1) else m end, replace=false)
    if (m == -1) m = size(data, 1) end # If m has signal value -1, set m to total number of examples.

    # Compute pairwise distances between samples (vector form).
    dists = Array{Float64}(undef, Int64((size(data, 1)^2 - size(data, 1))/2 + 1))

    dists[1] = 0  # Set first value of distances vector to 0 - accessed when i == j in square form indices.

    # Construct pairwise distances vector using vectorized distance function.
    top_ptr = 2
    @inbounds for idx = 1:size(data,1)-1
        upper_lim = top_ptr + size(data, 1) - idx - 1
        dists[top_ptr:upper_lim] = dist_func(data[idx:idx, :], data[idx+1:end, :])
        top_ptr = upper_lim + 1
    end

    # Get mean distance between all samples.
    mean_dist = Statistics.mean(dists[2:end])

    # Go over sampled indices.
    @inbounds for idx = sample_idxs

        # Row and column indices for querying pairwise distance vector.
        row_idxs = repeat([idx - 1], size(data, 1))
        col_idxs = collect(0:size(data, 1)-1)

        # Get indices in distance vector (from square form indices).
        neigh_idx = square_to_vec(row_idxs, col_idxs, size(data, 1)) .+ 2

        # Query distances to neighbors and get mask for neighbors that fall within
        # hypersphere with radius mean_dist.
        neigh_mask = dists[neigh_idx[neigh_idx .!= 1]] .<= mean_dist
        insert!(neigh_mask, idx, 0)

        # Get masks for nearest hits and nearest misses.
        hit_neigh_mask = neigh_mask .& (target .== target[idx])
        miss_neigh_mask = neigh_mask .& (target .!= target[idx])

        # Ignore samples that have no hits or misses in radius.
        if sum(hit_neigh_mask) == 0 || sum(miss_neigh_mask) == 0
            continue
        end

        # Compute probability weights for misses in considered region.
        miss_classes = target[miss_neigh_mask]

        # Compute weights for nearest misses and compute weighting vector.
        weights_mult = Array{Float64}(undef, length(miss_classes))  # Allocate weights multiplier vector.
        cm = countmap(miss_classes)                                 # Count unique values.
        u = collect(keys(cm))
        c = collect(values(cm))
        neighbor_weights = c ./ length(miss_classes)       # Compute misses' weights.
        @inbounds for (w, val) = zip(neighbor_weights, u)  # Build multiplier vector.
            find_res = findall(miss_classes .== val)
            weights_mult[find_res] .= w
        end

        ### Weights Update ###

        if f_type == "continuous"
            # If features continuous.

            # Penalty term
            penalty = sum(abs.(data[idx:idx, :] .- data[hit_neigh_mask, :])./(max_f_vals .- min_f_vals .+ eps(Float64)), dims=1)

            # Reward term
            reward = sum(weights_mult .* abs.(data[idx:idx, :] .- data[miss_neigh_mask, :])./(max_f_vals .- min_f_vals .+ eps(Float64)), dims=1)

            # Weights update
            weights = weights .- penalty ./ (size(data, 1)*sum(hit_neigh_mask) + eps(Float64)) .+
            reward ./ (size(data, 1)*sum(miss_neigh_mask) + eps(Float64))

        elseif f_type == "discrete"
            # If features discrete.

            # Penalty term
            penalty = sum(Int64.(data[idx:idx, :] .!= data[hit_neigh_mask, :]), dims=1)

            # Reward term
            reward = sum(weights_mult .* Int64.(data[idx:idx, :] .!= data[miss_neigh_mask, :]), dims=1)

            # Weights update
            weights = weights .- penalty ./ (size(data, 1)*sum(hit_neigh_mask) + eps(Float64)) .+
            reward ./ (size(data, 1)*sum(miss_neigh_mask) + eps(Float64))

        else
            throw(DomainError(f_type, "f_type can only be equal to \"continuous\" or \"discrete\"."))
        end

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

    end

    # Return computed feature weights.
    return vec(weights)
end


struct SRDirectCache{Tc,T,G,E,TT}
    S_c::Tc
    S::T
    F::G
    ∇out::E

    time_evolving::TT
end

function _sr_direct_cache(algo::SR, prob, net, par_cache)
    T = eltype(trainable_first(net))
    g = grad_cache(T, net)
    gv = vec_data(g)

    S = [similar(gv, T, length(gv), length(gv)) for gv in vec_data(g)]
    if T <: Real
        Sc = [similar(gv, Complex{T}, length(gv), length(gv)) for gv in vec_data(g)]
    else # is real
        Sc = S
    end

    return SRDirectCache(Sc, S, g, grad_cache(T, net),
                         algo.time_evo)
end

function setup_algorithm!(g::SRDirectCache, ∇C, Ô, par_cache)
    α = g.time_evolving isa Val{true} ? true * im : true

    O = Ô
    for (Sc, S, F) in zip(g.S_c, g.S, vec_data(g.F))
        T = eltype(S)
        N = size(Ô, 2)

        mul!(Sc, O, O')

        if T <: Real
            S .= real.(Sc ./ N)
            F .= real.(α .* ∇C)
        else
            # We take the conjugate because O is actually Oconj matrix respect
            # to the standard SR implementation..
            S .= conj.(Sc) ./ N
            F .= α .* ∇C
        end

        workers_sum!(S, par_cache)
    end
end

function precondition!(data::SRDirectCache, params::SR, iter_n)
    ϵ = params.sr_diag_shift

    success = true

    for (Δw, S, F) in zip(vec_data(data.∇out), data.S, vec_data(data.F))
        # new matrix
        if params.precondition_type == sr_none
            Sprecond = S
        elseif params.precondition_type == sr_shift
            #Sprecond = S + convert(eltype(S), ϵ)*I
            @inbounds @simd for i = 1:size(S, 1)
                S[i, i] += convert(eltype(S), ϵ)
            end
            Sprecond = S
        elseif params.precondition_type == sr_multiplicative
            λ0 = params.λ0
            b = params.b
            λmin = params.λmin
            λ = convert(eltype(S), max(λ0 * b^iter_n, λmin))
            Sprecond = S + λ * Diagonal(diag(S))
        end

        @assert is_iterative(params) == false
        if params.algorithm == sr_diag
            Sinv = pinv(Sprecond)
            mul!(Δw, Sinv, F)
        elseif params.algorithm == sr_cholesky
            C = cholesky!(Hermitian(S), check = true)
            copyto!(Δw, F)
            ldiv!(C, Δw)
        elseif params.algorithm == sr_div
            copyto!(Δw, F)
            ldiv!(Hermitian(S), Δw)
        else
            throw("Alg not known")
        end
    end
    return data.∇out
end
###########


using CUDA, Plots, Printf
# pyplot()
do_viz = false

function compute_flux!(qx, qy, T, λ, dx, dy, nx, ny)
    ix = (blockIdx().x-1) * blockDim().x + threadIdx().x # thread ID, dimension x
    iy = (blockIdx().y-1) * blockDim().y + threadIdx().y # thread ID, dimension y

    if (ix<=nx-1 && 2<=iy<=ny-1)  qx[ix,iy-1] = -λ*(T[ix+1,iy]-T[ix,iy])/dx  end
    if (2<=ix<=nx-1 && iy<=ny-1)  qy[ix-1,iy] = -λ*(T[ix,iy+1]-T[ix,iy])/dy  end
    return
end

function update_T!(T, qx, qy, dt, ρCp, dx, dy, nx, ny)
    ix = (blockIdx().x-1) * blockDim().x + threadIdx().x # thread ID, dimension x
    iy = (blockIdx().y-1) * blockDim().y + threadIdx().y # thread ID, dimension y

    if (2<=ix<=nx-1 && 2<=iy<=ny-1)  T[ix,iy] = T[ix,iy] - dt/ρCp*((qx[ix,iy-1]-qx[ix-1,iy-1])/dx + (qy[ix-1,iy]-qy[ix-1,iy-1])/dy)  end
    return
end

@views function heat_2D_gpu()
    # Physics
    lx, ly = 10.0, 10.0
    λ      = 1.0
    ρCp    = 1.0
    nt     = 200
    # Numerics
    BLOCKX = 16
    BLOCKY = 16
    GRIDX  = 8
    GRIDY  = 8
    nx, ny = BLOCKX*GRIDX, BLOCKY*GRIDY
    nout   = 10
    # Derived numerics
    dx, dy = lx/nx, ly/ny
    dt     = min(dx^2,dy^2)/ρCp/λ/4.1
    xc     = LinRange(dx/2, lx-dx/2, nx)
    yc     = LinRange(dy/2, ly-dy/2, ny)
    # Array allocation
    qx     = CUDA.zeros(Float64, nx-1,ny-2)
    qy     = CUDA.zeros(Float64, nx-2,ny-1)
    # Initial condition
    T      = CuArray( exp.(.-(xc.-lx./2.0).^2 .-(yc.-ly./2.0)'.^2) )
    cuthreads = (BLOCKX, BLOCKY, 1)
    cublocks  = (GRIDX,  GRIDY,  1)
    # Time loop
    for it = 1:nt
        if (it==11) global t0 = Base.time() end
        @cuda blocks=cublocks threads=cuthreads compute_flux!(qx, qy, T, λ, dx, dy, nx, ny)
        synchronize()
        @cuda blocks=cublocks threads=cuthreads update_T!(T, qx, qy, dt, ρCp, dx, dy, nx, ny)
        synchronize()
        # Visualise
        if mod(it,nout)==0 && do_viz
            display(heatmap(xc, yc, Array(T)', xlabel="lx", ylabel="ly", title="heat diffusion, it=$it", clims=(0.,1.)))
            # sleep(.01)
        end
    end
    time_s = (Base.time()-t0)
    @printf("Time = %1.4e s, T_eff = %1.2f GB/s \n", time_s, round((2/1e9*nx*ny*sizeof(lx))/(time_s/(nt-10)), sigdigits=2))
    return
end

@time heat_2D_gpu()


function cpairsum(a::TwoFloat{T}, b::TwoFloat{T}) where {T} =
    TwoFloat{T}(cpairsum(a[1],a[2],b[1],b[2]))

function cpairdiff(a::TwoFloat{T}, b::TwoFloat{T}) where {T} =
    TwoFloat{T}(cpairdiff(a[1],a[2],b[1],b[2]))

function cpairprod(a::TwoFloat{T}, b::TwoFloat{T}) where {T} =
    TwoFloat{T}(cpairprod(a[1],a[2],b[1],b[2]))

function cpairdiv(a::TwoFloat{T}, b::TwoFloat{T}) where {T} =
    TwoFloat{T}(cpairdiv(a[1],a[2],b[1],b[2]))

function cpairinv(a::TwoFloat{T}) where {T} =
    TwoFloat{T}(cpairinv(a[1],a[2]))

function cpairinv(a::TwoFloat{T}) where {T} =
    TwoFloat{T}(cpairinv(a[1],a[2]))

@inline function cpairsum(ahi::T, alo::T, bhi::T, blo::T) where {T}
    hihi, hilo = two_sum_ordered(ahi, bhi)
    t  = ahi + bhi - hihi
    lo = (t + hilo) + (alo + blo)
    return hi, lo
end

@inline function cpairsum(ahi::T, alo::T, bhi::T, blo::T) where {T}
    hi = ahi + bhi
    t  = ahi + bhi - hi
    lo = t + (alo + blo)
    return hi, lo
end

@inline function cpairdiff(ahi::T, alo::T, bhi::T, blo::T) where {T}
    hi = ahi - bhi
    t  = ahi - bhi - hi
    lo = t + (alo - blo)
    return hi, lo
end

function cpairprod(ahi::T, alo::T, bhi::T, blo::T) where {T}
    hi = a * b
    t  = fma(ahi, bhi, -hi)
    lo = t + (ahi*blo + bhi*alo)
    return hi, lo
end

function cpairinv(bhi::T, blo::T) where {T}
    hi = inv(bhi)
    t  = fma(-hi, bhi, one(T))
    lo = fma(-hi, blo, t) / (bhi + blo)
    return hi, lo
end

function cpairdiv(ahi::T, alo::T, bhi::T, blo::T) where {T}
    hi = ahi / bhi
    t  = fma(-hi, bhi, ahi) + alo
    lo = fma(-hi, blo, t) / (bhi + blo)
    return hi, lo
end

function cpairsqrt(ahi::T, alo::T) where {T}
    hi = sqrt(ahi)
    t  = -fma(hi, hi, -ahi)
    lo = (t + alo) / (2hi)
    return hi, lo
end

Base.:(-)(a::TwoFloat{T}) where {T} = (-a[1], -a[2])
Base.signbit(a::TwoFloat{T}) where {T} = signbit(a[1])
Base.abs(a::TwoFloat{T}) where {T} = signbit(a[1]) ? (abs(a[1]), -a[2]) : a
Base.copysign(a::TwoFloat, b) = signbit(b) ? (signbit(a[1]) ? a : -a) : a
Base.flipsign(a::TwoFloat, b) = signbit(b) ? -a : a

function Base.:(+)(a::TwoFloat{T}, b::TwoFloat{T}) where {T}
    if abs(b[1]) <= abs(a[1])
         ahi, alo = a
         bhi, blo = b
    else
         ahi, alo = b
         bhi, blo = a
    end
    return Base.:(+)(ahi, alo, bhi, blo)
end

function Base.:(+)(ahi::T, alo::T, bhi::T, blo::T)  where {T<:FastFloat}
    t0, t1 = two_hilo_sum(ahi, alo)
    t2, t3 = two_hilo_sum(bhi, blo)
    hi, t4 = two_hilo_sum(t0, t2)
    lo = t4 + (t1 + t3)
    hi, lo = two_hilo_sum(hi, lo)

    return TwoFloat{T}((hi, lo))
end

function Base.:(+)(a::TwoFloat{T}, b::TwoFloat{T}) where {T}
    if abs(b[1]) <= abs(a[1])
         ahi, alo = a
         bhi, blo = b
    else
         ahi, alo = b
         bhi, blo = a
    end
    t0, t1 = two_hilo_sum(ahi, alo)
    t2, t3 = two_hilo_sum(bhi, blo)
    hi, t4 = two_hilo_sum(t0, t2)
    lo = t4 + (t1+t3)
    hi, lo = two_hilo_sum(hi, lo)
    return TwoFloat{T}((hi, lo))
end

function Base.:(-)(a::TwoFloat{T}, b::TwoFloat{T}) where {T}
    ahi, alo = a
    bhi, blo = b
    ahi, alo, bhi, blo = mxxmnn_abs(ahi, alo, -bhi, -blo)
    t0, t1 = two_hilo_sum(ahi, alo)
    t2, t3 = two_hilo_sum(bhi, blo)
    hi, t4 = two_hilo_sum(t0, t2)
    lo = t4 + (t1+t3)
    hi, lo = two_hilo_sum(hi, lo)
    return TwoFloat{T}((hi, lo))
end

function Base.:(*)(a::TwoFloat{T}, b::TwoFloat{T}) where {T}
    ahi, alo = a
    bhi, blo = b
    if abs(alo) > abs(bhi)
        alo, bhi = bhi, alo
        alo, blo = maxmin_abs(alo,blo)
    end
    hi = ahi * bhi
    t = fma(ahi, bhi, -hi)
    lo = t + (ahi*blo + bhi*alo)
    return TwoFloat{T}((hi, lo))
end

function Base.:(/)(a::TwoFloat{T}, b::TwoFloat{T}) where {T}
    ahi, alo = a
    bhi, blo = b
    hi = ahi / bhi
    t = fma(-bhi, hi, ahi)
    lo = ((t + alo) - hi*blo) / (bhi + blo)
    return TwoFloat{T}((hi, lo))
end

function Base.inv(b::TwoFloat{T}) where {T}
    bhi, blo = b
    hi = inv(bhi)
    t = fma(-bhi, hi, one(T))
    lo = -(hi * blo) / (bhi + blo)
    return TwoFloat{T}((hi, lo))
end

function Base.sqrt(a::TwoFloat{T}) where {T}
    ahi, alo = a
    hi = sqrt(ahi)
    t = fma(hi, -hi, ahi)
    lo = (t + alo) / (2 * hi)
    return TwoFloat{T}((hi, lo))
end


# export jlmie_Qsca, jlmie_Qsca_a_n, jlmie_Qsca_b_n

function jlmie_Qsca(m,x,nmax::Int=-1)
# Calculate scattering efficiency of a sphere
# Input:
# m: relative refractive index (n_material / n_environment)
# x: size parameter (wavenumber * radius)
# namx: maximun order of resonances being considered
#       (nmax = 1: dipole, 2: dipole+quadrupole, ...,
#        -1: sufficiently large value determined by x)
# Output:
# Qsca: Scattering efficiency (total up to n-th order resonances)
    if nmax == -1
        nmax = largenmax(x)
    end

    Qsca = zeros(length(x))
    for i = 1:nmax
        Qsca = Qsca .+ jlmie_Qsca_n(m,x,i)
    end
    return Qsca
end

function jlmie_Qsca_n(m,x,n::Int)
    an,bn,_,_ = jlmie_abcd(m,x,n)
    nth_Qsca = 2 ./ x.^2 .* (2*n+1).*(abs.(an).^2 .+ abs.(bn).^2)
    return nth_Qsca
end

function jlmie_Qsca_a_n(m,x,n::Int)
# calculate scattering efficiency due to the n-th order electric multipole
# (n = 1: dipole, 2: quadrupole, ...)
    an,_,_,_ = jlmie_abcd(m,x,n)
    nth_Qsca_a = 2 ./ x.^2 .* (2*n+1).*(abs.(an).^2)
    return nth_Qsca_a
end

function jlmie_Qsca_b_n(m,x,n::Int)
# calculate scattering efficiency due to the n-th order magnetic multipole
# (n = 1: dipole, 2: quadrupole, ...)
    _,bn,_,_ = jlmie_abcd(m,x,n)
    nth_Qsca_b = 2 ./ x.^2 .* (2*n+1).*(abs.(bn).^2)
    return nth_Qsca_b
end