# Convert a prolongation matrix P to a tentative operator S.
# P is piecewise constant over some number of aggregates.
# S has a n x n for each aggregate where n is the size of the aggregate.

using PETScBinaryIO
using SparseArrays
using LinearAlgebra

function p_to_s(P)
    rows = rowvals(P)
    vals = nonzeros(P)
    m, n = size(P)
    is = Vector{Int}()
    js = Vector{Int}()
    ks = Vector{Float64}()
    j = 0
    bs = 16
    for i = 1:bs:n
        # get row indices of the aggregate
        inds = sort(unique(vcat(map(x -> rows[nzrange(P, i+x)], 0:(bs-1))...)))
        agg_size = length(inds)
        if agg_size > 10000
            println("Aggregate is too large ($agg_size)")
            exit(-1)
        end
        if agg_size == 1
            @warn "Ignoring aggregate of size 1"
            continue
        end
        B = Matrix(P[inds, i:(i+bs-1)])

        # I - B(BᵀB)⁻¹Bᵀ
        block = I - B * inv(B' * B) * B'
        # block = I - fill(1/agg_size, (agg_size, agg_size)) # B(BᵀB)⁻¹Bᵀ for piecewise constant

        U, S, _ = svd(block)
        l = findlast(x -> x > 1e-14, S)
        for x = 1:agg_size
            for y = 1:l
                push!(is, inds[x])
                push!(js, j + y)
                push!(ks, U[x,y])
            end
        end
        j += l

        # verify that S projects out the coarse space
        @assert norm(U[:, 1:l]' * B, Inf) < 1e14
    end
    sparse(is,js,ks,m,j)
end

if joinpath(pwd(), PROGRAM_FILE) == @__FILE__
    P = readpetsc(ARGS[1])[1]
    writepetsc(ARGS[2], p_to_s(P))
end


using AssemblyCalculus
using Graphs
using Plots

function time_sim(assem_size::Int; mean_degree=100)
    n = assem_size^2
    p = mean_degree / n
    β = 0.1
    ba = Brain(num_areas=1, n=n, k=assem_size, p=p, β=β)
    stims = [rand_stim(ba[1])]
    assems = Assembly{Float64}[]
    timesteps = 50
    td = @timed as, sp = simulate!(stims, assems, timesteps)
    return td.time
end

# Precompile
time_sim(5)

assem_sizes = collect(100:50:350)
times = time_sim.(assem_sizes)
neurons = 2 .* assem_sizes .^ 2

plot(neurons, times, linewidth=3, color=:teal, legend=:bottomright)
xticks!(neurons, string.(neurons))
xlabel!("Number of neurons")
ylabel!("Seconds to run 50 iterations (PartialNeuralArea)")
title!("Time Complexity of Generative Synapses")


# Model: Anand, McFarland, Kihm, Wong (Optimization of Aerosol Penetration through Transport Lines, 1992)
# Based on: straightTubePenetrationAnand1992.m

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

module tubeStraight

import physicalConstants: k, g, p, MM

struct Parameters
 # CarrierFluidProfile
    T :: Float64        # Carrier fluid temperature (T) SI units: K
    Ro :: Float64       # Carrier fluid density (ro) SI units: kg m-3
    Q :: Float64        # Carrier fluid flow rate (Q) SI units: m3 s-1
    mu :: Float64       # Carrier fluid dynamic viscosity (mu) SI units: N s m-2
end

## Input
print("Parameter details \n")
print("Enter T, Ro, Q, mu: ")
params = split(readline(), ",")

T = params[1]
Ro = params[2]
Q = params[3]
mu = params[4]

# print(T)
# print(Ro)
# print(Q)
# print(mu)
# print("\n")

mutable struct Properties
 # Droplet
    ro_d :: Float64                         # Droplet density (ro_d) SI units: kg m-3
    mu_d :: Float64                         # Droplet dynamic viscosity (mu_d) SI units: N s m-2
    d_d :: Float64                          # Droplet diameter (dd) SI units: m
    md :: Float64                           # Droplet mass (md) SI units: kg

 # Tube
    d :: Float64                            # Tube inner diameter (d) SI units: m
    phi :: Float64                          # Inclination angle in radians (phi) SI units: rad
    L :: Float64                            # Distance from the tube inlet (L) SI units: m
end

## Input
print("Droplet Properties \n")
print("Enter ro_d, mu_d, d_d: ")
propsDroplet = split(readline(), ",")

ro_d = propsDroplet[1]
mu_d = propsDroplet[2]
d_d = propsDroplet[3]

# print(ro_d)
# print(mu_d)
# print(d_d)
# print("\n")

print("Tube Properties \n")
print("Enter d, phi, L: ")
propsTube = split(readline(), ",")

function meanFlow(Q,  d)
        U=4*Q/(pi*d^2)                       # Fluid flow mean velocity (U) SI units: m s-1
end

function reynold(U,  d, ro, mu)
        Re=(U*d*ro)/mu                       # Reynolds number (Re)
end

function fff(Re)
        f=0.316/(4*Re^0.25)                  # Fanning friction factor, Blasius equation for turbulent pipe flow (f) dimensionless
end

# The Blasius equation is not applicable for the following range of Reynolds number:
function blasius(Re)
    if (Re .< 3000 || Re .> 100000)
        extrapolation=extrapolation+1
end

function frictionV(U, f)
        Vf = U*sqrt(f/2)                      # Friction velocity (Vf) SI units: m s-1
end

function CunninghamSlipCorrection(mu, T, d_d)
         lambda = (mu/p)*sqrt((pi*k*T)/(2*MM))                   # Model for the mean free path of the carrier fluid (lambda)
         C = 1 + (lambda/dd)*(2.34+1.05*exp(-.39*d_d/lambda))    #Cunningham slip correction (C)
end

function relaxationTime(C, ro_d, d_d,  mu)
        tau = (C*ro_d*d_d^2)/(18*mu)          # Droplet relaxation time (tau)
end

function Rplus(d_d, Vf, ro, mu)
        Rplus = (d_d*Vf*ro)/(2*mu)            # Parameter Rplus (Rplus)
end

function paramS(tau, Vf, d_d)
        ParamS = 0.9*tau*Vf + d_d/2            # Parameter S (ParamS)
end

function Splus(ParamS, Vf, ro, mu)
        Splus = (ParamS*Vf*ro)/(mu)            # Parameter Splus (Splus)
end

# Dimensionless depositional velocity (Vs)
if (Splus >= 0 && Splus <= 10)
    Vs = 0.05*Splus
else
    Vs = 0.5 + 0.0125*(Splus - 10)
    extrapolation=extrapolation+1
end

# Dimensionless depositional velocity (Vr)
if (Rplus >= 0 && Rplus <= 10)
    Vr = 0.05*Rplus
else
    Vr = 0.5 + 0.0125*(Splus - 10)
    extrapolation=extrapolation+1
end

function diffVelTur(Vf, Vs, Vr)
        Vt = (Vf*(Vs+Vr))/(4)                  # Turbulent diffusional velocity (Vt)
end

function gSettlingVel(tau, g, phi)
        Vg = tau*g*cos(phi)                    # Gravitational settling velocity (Vg)
end

function depositionV(Vt, Vb, K)
        Vd = Vt + Vb*K                         # Sum of Vt and Vb is defined as Vd (Vd)
end

function critAngle(Vd, Vg)
        TETAc = asin(Vd/Vg)                    # Critical angle (TETAc)
end

function effectiveDepositionV(Vd, TETAc, Vg)
        Ve=Vd*TETAc/pi + Vd/2 + Vg*cos(TETAc)/pi        # Ve is the effective depositional velocity, given by the vector sum of velocities due to turbulent diffusion (Vt),
end

# If Vd>Vg, there is no need to use TETAc, since the integral that originated the Ve equation can be evaluated from 0 to 2pi. Also, arcsine(x) for x>1 is not defined in the real domain, therefore, TETAc would be imaginary.
if (Vd>Vg)
    Ve=Vd
end

function penetration(d, Ve, L, Q)
        P = -exp(-(pi*d*Ve*L)/(Q))             # P = penetration = C/C0 = ratio between concentration at distance x and initial concentration
end

    d = parse(Float64, propsTube[1])
phi = parse(Float64, propsTube[2])
L = parse(Float64, propsTube[3])

print(typeof(d))

# print(d)
# print(phi)
# print(L)
# print("\n")

# function dropletMass(ro_d, d_d)
#     md = ro_d*(4/3)*(pi*(d_d/2)^3)     # Droplet mass (md) SI units: kg
# end

# md = ro_d*(4/3)*(pi*(d_d/2)^3)     # Droplet mass (md) SI units: kg

# dropletMass(ro_d, d_d)

# print(md)

# function velocityBr(k, T, md)
#     Vb = sqrt((k*T)/(2*pi*md))          # Brownian diffusion velocity (Vb) SI units: m s-1
# end

# test = velocityBr(k, T, md)

# print(test)

function velocityTu(Q, d)
end
#Debug test
# println(k)

end

end


#prec = 15 #precision

##Single Particle
eigensp(s::State, prec=15) = sort([round(eigvals(Matrix(rhosp(s)))[i], digits = prec) for i in 1:dim(ope(s))], rev=true)
eigensp(s::State_fixed, prec=15) = sort([round(eigvals(Matrix(rhosp(s)))[i], digits = prec) for i in 1:dim(ope(s))], rev=true)
# I can not compute eigenvalues
#directly from sparse matrices

function ssp(sta::State, prec=15)
    eigen = eigensp(sta, prec)
    lene = length(eigen)
    s = 0
    for i in 1:lene
        if eigen[i] != 0 && eigen[i] != 1
            s = s - (eigen[i]*log(2,eigen[i]) + (1 - eigen[i])*log(2,1-eigen[i]))
        end
    end
    return s/lene
end

function ssp(sta::State_fixed, prec=15)
    eigen = eigensp(sta, prec)
    lene = length(eigen)
    s = 0
    for i in 1:lene
        if eigen[i] != 0 && eigen[i] != 1
            s = s - (eigen[i]*log(2,eigen[i]) + (1 - eigen[i])*log(2,1-eigen[i]))
        end
    end
    return s/lene
end
## Quasi Particles
eigenqsp(s::State, prec=15) = sort([round(eigvals(Matrix(rhoqsp(s)))[i], digits = prec) for i in 1:(2*dim(ope(s)))], rev=true)
# I can not compute eigenvalues
#directly from sparse matrices

function sqsp(sta::State, prec=15)
    eigen = eigenqsp(sta, prec)
    lene = length(eigen)
    s = 0
    for i in 1:lene
        if eigen[i] != 0 && eigen[i] != 1
            s = s - (eigen[i]*log(2,eigen[i]) + (1 - eigen[i])*log(2,1-eigen[i]))
        end
    end
    return s/lene
end

#This function outputs 1 if s1 majorizes s2,
#-1 if s2 majorizes s1,
#0 if there is no majorization relation
function majorization_sp(s1, s2)
    e1 = eigensp(s1)
    e2 = eigensp(s2)
    count1 = 0
    count2 = 0
    d = length(e1)
    for j in 1:(d-1) #el último valor a veces hace problemas
        if sum([round(e1[i]-e2[i], digits = 10) for i in 1:j]) < 0
            count1 = count1 + 1
        elseif sum([round(e2[i]-e1[i], digits = 10) for i in 1:j]) < 0
            count2 = count2 + 1
        end
    end
    if count1 == (d-1)
        #println(@Name(s2)," sp majorizes ", @Name(s1))
        return -1
    elseif count2 == (d-1)
        #println(@Name(s1)," sp majorizes ", @Name(s2))
        return 1
    else
        #println("There is no sp majorization relation")
        return 0
    end
end

function majorization_qsp(s1, s2)
    e1 = eigenqsp(s1)
    e2 = eigenqsp(s2)
    count1 = 0
    count2 = 0
    d = length(e1)
    for j in 1:(d-1) #el último valor a veces hace problemas
        if sum([round(e1[i]-e2[i], digits = 10) for i in 1:j]) < 0
            count1 = count1 + 1
        elseif sum([round(e2[i]-e1[i], digits = 10) for i in 1:j]) < 0
            count2 = count2 + 1
        end
    end
    if count1 == (d-1)
        #println(@Name(s2)," qsp majorizes ", @Name(s1))
        return -1
    elseif count2 == (d-1)
        #println(@Name(s1)," qsp majorizes ", @Name(s2))
        return 1
    else
        #println("There is no qsp majorization relation")
        return 0
    end
end

function n_avg(s::State)
    navg = Real(tr(rhosp(s)))
    return navg
end

####################################################
# Following are the matrices of coefficients in
# the m-(N-m) schmidt like decomposition
function coef(s::State, m::Int64, prec=15)
    d = dim(s)
    num = Int(round(n_avg(s),digits = 8))
    bas, _ = basis_m(d, m)
    estat = st(s)
    ty = typ(s)
    cof = cof_mat(d, num, m, bas, estat, ty, prec)
    return round.(cof,digits=prec)
end

function coef(s::State_fixed, m::Int64, prec=15)
    d = dim(s)
    num = nume(s)
    bas_tot = basis(s)
    bas, _ = basis_m(d, m)
    estat = st(s)
    ty = typ(s)
    cof = cof_mat_fixed(d, num, m, bas, bas_tot, estat, ty, prec)
    return round.(cof,digits=prec)
end

#The following are the rhom matrices without diagonalization
# i.e. in the original basis
# For general m with fixed_state
function rhom(s::Union{State,State_fixed}, m::Int64, prec=15)
    cof = coef(s, m, prec)
    rhom = cof*cof'
    return round.(rhom,digits=prec)
end

#The following are the rhom matrices with diagonalization
function rhomd(s::Union{State,State_fixed}, m::Int64, prec=15)
        r = rhom(s,m,prec)
        dim = size(r)[1]
        d = non_zero(eigvals(Matrix(r)))[2]
        row = [i for i in 1:length(d)]
        col = row
        rhomd = sparse(row, col, d, dim, dim)
        return round.(rhomd,digits=prec)
end

#A more efficient 2 body DM in fixed space
function rhom2(s::State_fixed, prec=15)
    m = 2
    o = ope(s)
    num = nume(s)
    d = dim(o)
    rhomnd = spzeros(typ(s),binomial(d, m), binomial(d, m))
    bas,_ = basis_m(d,m)
    estat = st(s)
    for i in 1:binomial(d, m)
        indi = indx(bas[i,:])
        for j in 1:binomial(d,m)
            indj = indx(bas[j,:])
            #rhomd[i,j] = round(estat'*fixed(and[j]*and[i]',num)*estat, digits = 14)
            if indi[2] != indj[1] && indi[1] != indj[2]
                rhomnd[i,j] = round(estat'*ada(o,indj[1],indi[1])*ada(o,indj[2],indi[2])*estat, digits = prec)
            else
                rhomnd[i,j] = round(-estat'*ada(o,indj[1],indi[2])*ada(o,indj[2],indi[1])*estat, digits = prec)
            end
        end
    end
    return rhomnd
end

# Following function is used for constructing coef and rhom
#for example, for a vector [0,1,0,0,1]->[2,5]
function indx(arr)
    l = length(arr)
    ind = spzeros(Int,0)
    for i=1:l
        if arr[i] != 0
            ind = sparse([ind; i])
        end
    end
    return ind
end

# Core functons of the coeficient matrices
function cof_mat(d, num, m, bas, estat, ty, prec=15)
    cmat = spzeros(ty, binomial(d,m),binomial(d,num-m))
    am = sparse([bas[i,:] for i in 1:binomial(d, m)])
    bas_nume_m, _ = basis_m(d, num - m)
    amn = sparse([bas_nume_m[i,:] for i in 1:binomial(d, num-m)])
    for i in 1:binomial(d, m)
        ai = am[i]
        indi = indx(ai)
        for j in 1:binomial(d, num-m)
            aj = amn[j]
            if ai'*aj != 0
                cmat[i,j] = 0
            else
                permut = 0
                for k in 1:length(indi)
                    permut = permut + sum(aj[1:floor(Int,indi[k]-1)])
                end
                indij = indx(ai+aj)
                elem = Int(sum([2^(d-i) for i in indij]) + 1)
                cmat[i,j] = round(estat[elem] * (-1)^permut, digits=prec)
            end
        end
    end
    return cmat
end

function cof_mat_fixed(d, num, m, bas, bas_tot, estat, ty, prec=15)
    cmat = spzeros(ty,binomial(d,m),binomial(d,num-m))
    am = sparse([bas[i,:] for i in 1:binomial(d, m)])
    bas_nume_m, _ = basis_m(d, num - m)
    amn = sparse([bas_nume_m[i,:] for i in 1:binomial(d, num-m)])
    for i in 1:binomial(d, m)
        ai = am[i]
        indi = indx(ai)
        for j in 1:binomial(d, num-m)
            aj = amn[j]
            if ai'*aj != 0
                cmat[i,j] = 0
            else
                permut = 0
                for k in 1:length(indi)
                    permut = permut + sum(aj[1:floor(Int,indi[k]-1)])
                end
                indij = indx(ai+aj)
                elem = 1
                while indx(bas_tot[elem,:]) != indij
                    elem = elem + 1
                end
                cmat[i,j] = round(estat[elem] * (-1)^permut, digits=prec)
            end
        end
    end
    return cmat
end

#Partial trace in a given basis of modes
function trp(state::State, modos::Array{Int64,1})
    d = dim(ope(state))
    bas = basis(ope(state))
    sta = st(state)
    lm = length(modos)
    zvr = zeros(typ(state),2^lm,2^(d-lm))
    full = [i for i in 1:d]
    lista = sort(modos)
    listb = filter(x->x ∉ lista,full)
    for k in 1:2^d
        indicea = parse(Int,join([Int(bas[k,i]) for i in lista]), base=2)+1
        indiceb = parse(Int,join([Int(bas[k,i]) for i in listb]), base=2)+1
        signi = 0
        for i in lista
            listac = filter(x->x <= i-1,listb)
            signl = 0
            for l in listac
                signl = signl+bas[k,l]
            end
            signi = signi + signl*bas[k,i]
        end
        sign = (-1)^(signi)
        zvr[indicea,indiceb] = sta[k]*sign
    end
    rhoa = zvr*zvr'
    #rhob=zvr'*zvr;
    return rhoa
end

#= This was used for finding the diagonalizing
# operators in m
function non_diag_ops(o, d, nume, m, bas, i)
    vec = indx(bas[i, :])
    #le es nume-1 me parece
    mel = 1
    for j in 1:m
        mel = mel*ad(o, Int(vec[j]))
    end
    return mel
end

function diag_ops(o, d, nume, m, bas, u, and, i)
    adi = spzeros(2^d, 2^d)
    for j in 1:binomial(d,m)
        if round(u[j,i], digits=14) != 0
            adi = adi + u[j, i]*and[j]
        end
    end
    return adi
end
=#

#= Work in Progress: non diagonal operators with fixed states
function rhomnd(s::Union{State_fixed,State_sparse_fixed}, m::Int64)
    o = ope(s)
    num = nume(s)
    d = dim(o) #tengo que pasar a matrix lamentablemente
    rhomd = spzeros(binomial(d, m), binomial(d, m))
    bas, _ = basis_m(d, m)
    estat = st(s)
    and = sparse([non_diag_ops(o, d, num, m, bas, j) for j in 1:binomial(d, m)])
    for i in 1:binomial(d, m)
        for j in 1:binomial(d,m)
            rhomd[i,j] = round(estat'*and[j]*and[i]'*estat, digits = 14)
        end
    end
    return rhomd
end

#main function
function rhomnd(s::Union{State_complex_fixed,State_sparse_complex_fixed}, m::Int64)
    o = ope(s)
    num = nume(s)
    d = dim(o) #tengo que pasar a matrix lamentablemente
    rhomd = spzeros(Complex{Float64},binomial(d, m), binomial(d, m))
    bas, _ = basis_m(d, m)
    estat = st(s)
    and = sparse([non_diag_ops(o, d, num, m, bas, j) for j in 1:binomial(d, m)])
    for i in 1:binomial(d, m)
        for j in 1:binomial(d,m)
            rhomd[i,j] = round(estat'*and[j]*and[i]'*estat, digits = 14)
        end
    end
    return rhomd
end
=#

#= Work in Progress: partial trace for fixed particle number
function trp(state::Union{State_fixed,State_sparse_fixed},modos::Array{Int64,1})
    d = dim(ope(state))
    bas = basis(ope(state))
    sta = st(state)
    lm = length(modos)
    zvr = zeros(Complex{Float64},2^lm,2^(d-lm))
    full = [i for i in 1:d]
    lista = sort(modos)
    listb = filter(x->x ∉ lista,full)
    for k in 1:2^d
        indicea=parse(Int,join([Int(bas[k,i]) for i in lista]), base=2)+1
        indiceb=parse(Int,join([Int(bas[k,i]) for i in listb]), base=2)+1
        signi = 0
        for i in lista
            listac = filter(x->x <= i-1,listb)
            signl = 0
            for l in listac
                signl = signl+bas[k,l]
            end
            signi = signi + signl*bas[k,i]
        end
        sign = (-1)^(signi)
        zvr[indicea,indiceb]=sta[k]*sign
    end
    rhoa=zvr*conj(zvr')
    #rhob=zvr'*zvr;
    return rhoa
end

function trp(state::Union{State_complex_fixed,State_sparse_complex_fixed},modos::Array{Int64,1})
    d = dim(ope(state))
    bas = basis(ope(state))
    sta = st(state)
    lm = length(modos)
    zvr = zeros(Complex{Float64},2^lm,2^(d-lm))
    full = [i for i in 1:d]
    lista = sort(modos)
    listb = filter(x->x ∉ lista,full)
    for k in 1:2^d
        indicea=parse(Int,join([Int(bas[k,i]) for i in lista]), base=2)+1
        indiceb=parse(Int,join([Int(bas[k,i]) for i in listb]), base=2)+1
        signi = 0
        for i in lista
            listac = filter(x->x <= i-1,listb)
            signl = 0
            for l in listac
                signl = signl+bas[k,l]
            end
            signi = signi + signl*bas[k,i]
        end
        sign = (-1)^(signi)
        zvr[indicea,indiceb]=sta[k]*sign
    end
    rhoa=zvr*zvr'
    #rhob=zvr'*zvr;
    return rhoa
end
=#

#= This code is faster for many iterations (majorization)
count = 0
@time for i in 1:repeticiones
    #definimos estado premedida
    state_ran = spzeros(Complex{Float64},l)
    for i in 1:l
        if sum(basis(o)[i,:]) == nume
            state_ran[i] = 2*rand(Complex{Float64},1)[1]-1-im
        end
    end
    state_ran = state_ran/sqrt(state_ran'*state_ran)
    state_ran = State_sparse_complex(state_ran, o);
    #e = sort(eigvals(Matrix(rhoqsp(state_ran))), rev=true)

    #ahora medimos
    state_ran_medido = operator*st(state_ran);
    state_ran_medido = state_ran_medido/sqrt(state_ran_medido'*state_ran_medido)
    state_ran_medido = State_sparse_complex(state_ran_medido, o)
    #e_medido = sort(eigvals(Matrix(rhoqsp(state_ran_medido))), rev=true)

    #y chequeamos mayorización
    global count = count + majorization_qsp(state_ran_medido, state_ran)
end

if count == repeticiones
    println("Measured state qsp majorizes unmeasured state")
elseif count == -repeticiones
    println("Unmeasured state qsp majorizes measured state")
else
    println("There is no qsp majorization")
end

=#

#This macro serves for printing as a string the variable's name
#=
macro Name(arg)
   string(arg)
end
=#

#= function rhom(s::State, m::Int64, prec=15)
#     o = ope(s)
#     num = Int(round(n_avg(s),digits = 8))
#     d = dim(o) #tengo que pasar a matrix lamentablemente
#     rhomd = spzeros(typ(s),binomial(d, m), binomial(d, m))
#     bas,_ = basis_m(d, m)
#     estat = st(s)
#     and = sparse([non_diag_ops(o, d, num, m, bas, j) for j in 1:binomial(d, m)])
#     for i in 1:binomial(d, m)
#         for j in 1:binomial(d,m)
#             rhomd[i,j] = round(estat'*and[j]*and[i]'*estat, digits = prec)
#         end
#     end
#     return rhomd
# # end
#
# #main function
# function rhomd(s::State, m::Int64, prec=15)
#     o = ope(s)
#     num = Int(round(n_avg(s),digits = 8))
#     d = dim(o) #tengo que pasar a matrix lamentablemente
#     rhomd = spzeros(typ(s),binomial(d, m), binomial(d, m))
#     bas, _ = basis_m(d, m)
#     estat = st(s)
#     ty = typ(s)
#     u, _ = svd(cof_mat(d, num, m, bas, estat, ty))
#     and = sparse([non_diag_ops(o, d, num, m, bas, j) for j in 1:binomial(d, m)])
#     adi = sparse([diag_ops(o, d, num, m, bas, u, and, j) for j in 1:binomial(d, m)])
#     #adi[1] es evaluar la funcion diagonal en i=1
#     for i in 1:binomial(d, m)
#         rhomd[i,i] = round(estat'*adi[i]*adi[i]'*estat, digits = prec)
#     end
#     return rhomd
# end

método directo (muy lento)
function rhom2(s::State)
    d = dim(ope(s))
    o = ope(s)
    bas,_ = basis_m(d,2)
    l = Int(binomial(d,2))
    rhom = spzeros(l,l);
    stat = st(s)
    for i in 1:l
        indi = indx(bas[i,:])
        for j in 1:l
            indj = indx(bas[j,:])
                rhom[i,j]=(stat'*ad(o,Int(indi[2]))*ad(o,Int(indi[1]))*a(o,Int(indj[1]))*a(o,Int(indj[2]))*stat)
        end
    end
    return rhom
end

rhom(State(1/sqrt(2)*(ad(o,1)*ad(o,2)*ad(o,3)*ad(o,4)+ad(o,3)*ad(o,4)*ad(o,5)*ad(o,6))*vacuum(o),o),2)

# It works, but it is not efficient,
# we need to optimieze ir for state_fixed
function rhom(s::State_fixed, o::Op, m::Int64, prec=15)
    d = dim(o)
    #o = Op(d)
    num = nume(s)
    rhomd = spzeros(typ(s),binomial(d, m), binomial(d, m))
    bas,_ = basis_m(d, m)
    estat = st(s)
    and = sparse([non_diag_ops(o, d, num, m, bas, j) for j in 1:binomial(d, m)])
    for i in 1:binomial(d, m)
        for j in 1:binomial(d,m)
            rhomd[i,j] = round(estat'*fixed(and[j]*and[i]',num)*estat, digits = prec)
        end
    end
    return rhomd
end

#
# # It works, but it is not efficient,
# # we need to optimieze ir for state_fixed
# function rhomd(s::State_fixed, o::Op, m::Int64, prec=15)
#     d = dim(o)
#     #o = Op(d)
#     num = nume(s)
#     rhomd = spzeros(typ(s),binomial(d, m), binomial(d, m))
#     bas_tot, _ = basis_m(d, num)
#     bas, _ = basis_m(d, m)
#     estat = st(s)
#     ty = typ(s)
#     u, _ = svd(cof_mat_fixed(d, num, m, bas, bas_tot, estat, ty))
#     and = sparse([non_diag_ops(o, d, num, m, bas, i) for i in 1:binomial(d, m)])
#     adi = sparse([diag_ops(o, d, num, m, bas, u, and, i) for i in 1:binomial(d, m)])
#     #ad[1] es evaluar la funcion diagonal en i=1
#     for i in 1:binomial(d, m)
#         rhomd[i,i] = round(estat'*fixed(adi[i]*adi[i]',num)*estat, digits = prec)
#     end
#     return rhomd
# end

=#


### CONCRETE TYPE: ITERATIVE PROX EVALUATION

using LinearAlgebra
using IterativeSolvers

struct QuadraticIterative{R <: Real, M, V <: AbstractVector{R}} <: Quadratic
  Q::M
  q::V
  temp::V
end

is_prox_accurate(f::QuadraticIterative) = false

function QuadraticIterative(Q::M, q::V) where {R <: Real, M, V <: AbstractVector{R}}
  if size(Q, 1) != size(Q, 2) || length(q) != size(Q, 2)
    error("Q must be squared and q must be compatible with Q")
  end
  QuadraticIterative{R, M, V}(Q, q, similar(q))
end

function (f::QuadraticIterative{R, M, V})(x::AbstractArray{R}) where {R, M, V}
  mul!(f.temp, f.Q, x)
  return 0.5*dot(x, f.temp) + dot(x, f.q)
end

function prox!(y::AbstractArray{R}, f::QuadraticIterative{R, M, V}, x::AbstractArray{R}, gamma::R=one(R)) where {R, M, V}
  f.temp .= x./gamma .- f.q
  op = Shift(f.Q, one(R)/gamma)
  IterativeSolvers.cg!(y, op, f.temp)
  mul!(f.temp, f.Q, y)
  fy = 0.5*dot(y, f.temp) + dot(y, f.q)
  return fy
end

function gradient!(y::AbstractArray{R}, f::QuadraticIterative{R, M, V}, x::AbstractArray{R}) where {R, M, V}
  mul!(y, f.Q, x)
  y .+= f.q
  return 0.5*(dot(x, y) + dot(x, f.q))
end

function prox_naive(f::QuadraticIterative, x, gamma=1.0)
  y = IterativeSolvers.cg(gamma*f.Q + I, x - gamma*f.q)
  fy = 0.5*dot(y, f.Q*y) + dot(y, f.q)
  return y, fy
end


using LinearAlgebra

@testset "PointForce" begin

  Re = 200
  xlim, ylim = (-2,2), (-2,4)
  Δx,Δt = setstepsizes(Re,gridRe=4)

  sys = NavierStokes(Re,Δx,xlim,ylim,Δt)
  w₀ = Nodes(Dual,size(sys))

  x0 = [1.5,0.0]
  f0 = 10.0
  t0 = 4
  σt = 1
  pulse = PointForce(w₀,x0,f0,t0,σt,sys)

  @test isapprox(abs(sum(pulse(t0))-f0),0,atol=1e-2)

  x0 = (1.5,0.0)
  pulse = PointForce(w₀,x0,f0,t0,σt,sys)

  @test isapprox(abs(sum(pulse(t0))-f0),0,atol=1e-2)

end

@testset "Spatial pulse" begin

σx = 0.5
σy = 0.5
x0 = 0
y0 = 0
A = 1
dgaussx = SpatialGaussian(σx,σy,x0,y0,A,deriv=1)
dgaussy = SpatialGaussian(σx,σy,x0,y0,A,deriv=2)
x = 0.1
y = 0.2
@test dgaussx(0.1,0.2) == dgaussy(0.2,0.1) ≈ -2*A*x/π/σx^3/σy*exp(-x^2/σx^2)*exp(-y^2/σy^2)

end


using LinearAlgebra: norm

"""
    DiatomicHarmonic(r₀=1.0)

Harmonic interaction between two particles.
"""
Parameters.@with_kw struct DiatomicHarmonic <: AdiabaticModel
    r₀::Float64 = 1.0
end

NQCModels.ndofs(::DiatomicHarmonic) = 3

function NQCModels.potential(model::DiatomicHarmonic, R::AbstractMatrix)
    return (norm(R[:,1] .- R[:,2]) - model.r₀)^2 / 2
end

function NQCModels.derivative!(model::DiatomicHarmonic, D::AbstractMatrix, R::AbstractMatrix)
    diff = R[:,1] .- R[:,2]
    leng = norm(diff)
    D .= (leng - model.r₀) / leng .* diff
    D[:,2] .*= -1
end


@doc raw"""
    GeneralizedBurresWassertseinMetric{T<:AbstractMatrix} <: AbstractMetric

The generalized Bures Wasserstein metric for symmetric positive definite matrices, see[^HanMishraJawanpuriaGao2021].

This metric internally stores the symmetric positive definite matrix ``M`` to generalise the metric,
where the name also follows the mentioned preprint.

[^HanMishraJawanpuriaGao2021]:
    > Han, A., Mishra, B., Jawanpuria, P., Gao, J.:
    > _Generalized Bures-Wasserstein geometry for positive definite matrices_.
    > arXiv: [2110.10464](https://arxiv.org/abs/2110.10464).
"""
struct GeneralizedBuresWassersteinMetric{T<:AbstractMatrix} <: RiemannianMetric
    M::T
    GeneralizedBuresWassersteinMetric(MM::TT) where {TT<:AbstractMatrix} = new{TT}(MM)
end

@doc raw"""
    change_representer(M::MetricManifold{ℝ,SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric}, E::EuclideanMetric, p, X)

Given a tangent vector ``X ∈ T_p\mathcal M`` representing a linear function on the tangent
space at `p` with respect to the [`EuclideanMetric`](@ref) `g_E`,
this is turned into the representer with respect to the (default) metric,
the [`GeneralizedBuresWassersteinMetric`](@ref) on the [`SymmetricPositiveDefinite`](@ref) `M`.

To be precise we are looking for ``Z∈T_p\mathcal P(n)`` such that for all ``Y∈T_p\mathcal P(n)``
it holds

```math
⟨X,Y⟩ = \operatorname{tr}(XY) = ⟨Z,Y⟩_{\mathrm{BW}}
```
for all ``Y`` and hence we get ``Z = 2pXM + 2MXp``.
"""
change_representer(
    ::MetricManifold{ℝ,SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric},
    ::EuclideanMetric,
    p,
    X,
)

function change_representer!(
    M::MetricManifold{ℝ,<:SymmetricPositiveDefinite,<:GeneralizedBuresWassersteinMetric},
    Y,
    ::EuclideanMetric,
    p,
    X,
)
    Y .= 2 .* (p * X * M.metric.M + M.metric.M * X * p)
    return Y
end

@doc raw"""
    distance(::MatricManifold{SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric}, p, q)

Compute the distance with respect to the [`BuresWassersteinMetric`](@ref) on [`SymmetricPositiveDefinite`](@ref) matrices, i.e.

```math
d(p,q) = \operatorname{tr}(M^{-1}p) + \operatorname{tr}(M^{-1}q)
       - 2\operatorname{tr}\bigl( (p^{\frac{1}{2}}M^{-1}qM^{-1}p^{\frac{1}{2}} \bigr)^{\frac{1}{2}},
```
"""
function distance(
    M::MetricManifold{ℝ,<:SymmetricPositiveDefinite,<:GeneralizedBuresWassersteinMetric},
    p,
    q,
)
    luM = lu(M.metric.M)
    luMp = luM \ p
    luMq = luM \ q
    return sqrt(tr(luMp) + tr(luMq) - 2 * tr(sqrt(luMq * luMp)))
end

@doc raw"""
    exp(::MatricManifold{ℝ,SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric}, p, X)

Compute the exponential map on [`SymmetricPositiveDefinite`](@ref) with respect to
the [`GeneralizedBuresWassersteinMetric`](@ref) given by

```math
    \exp_p(X) = p+X+\mathcal ML_{p,M}(X)pML_{p,M}(X)
```

where ``q=L_{M,p}(X)`` denotes the generalized Lyapunov operator, i.e. it solves ``pqM + Mqp = X``.
"""
exp(::MetricManifold{ℝ,SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric}, p, X)

function exp!(
    M::MetricManifold{ℝ,<:SymmetricPositiveDefinite,<:GeneralizedBuresWassersteinMetric},
    q,
    p,
    X,
)
    m = M.metric.M
    Y = lyapc(p, m, -X) #lyap solves qpM + Mpq - X =0
    q .= p .+ X .+ m * Y * p * Y * m
    return q
end

@doc raw"""
    inner(::MetricManifold{ℝ,SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric}, p, X, Y)

Compute the inner product [`SymmetricPositiveDefinite`](@ref) with respect to
the [`GeneralizedBuresWassersteinMetric`](@ref) given by

```math
    ⟨X,Y⟩ = \frac{1}{2}\operatorname{tr}(L_{p,M}(X)Y)
```

where ``q=L_{M,p}(X)`` denotes the generalized Lyapunov operator, i.e. it solves ``pqM + Mqp = X``.
"""
function inner(
    M::MetricManifold{ℝ,<:SymmetricPositiveDefinite,<:GeneralizedBuresWassersteinMetric},
    p,
    X,
    Y,
)
    return dot(lyapc(p, M.metric.M, -X), Y) / 2
end

@doc raw"""
    log(::MatricManifold{SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric}, p, q)

Compute the logarithmic map on [`SymmetricPositiveDefinite`](@ref) with respect to
the [`BuresWassersteinMetric`](@ref) given by

```math
    \log_p(q) = M(M^{-1}pM^{-1}q)^{\frac{1}{2}} + (qM^{-1}pM^{-1})^{\frac{1}{2}}M - 2 p.
```
"""
log(::MetricManifold{ℝ,SymmetricPositiveDefinite,GeneralizedBuresWassersteinMetric}, p, q)

function log!(
    M::MetricManifold{ℝ,<:SymmetricPositiveDefinite,<:GeneralizedBuresWassersteinMetric},
    X,
    p,
    q,
)
    m = M.metric.M
    lum = lu(m)
    lum_p_lum = lum \ p / lum
    X .= real.(Symmetric(m * sqrt(lum_p_lum * q) + sqrt(q * lum_p_lum) * m) - 2 * p)
    return X
end


using GeodesicTracer
using ComputedGeodesicEquations
using GeodesicBase
using CarterBoyerLindquist

using StaticArrays
using Optim

# include("/Users/lucyackland-snow/Desktop/DevEnv1/dev/AccretionFormulae/src/johannsen.jl")
include("hunting-circular-orbits.jl")
include("alt-emissivity-deviations.jl")

using Plots

# import GeodesicBase, GeodesicTracer
# function GeodesicBase.metric(m::GeodesicTracer.AbstractAutoDiffStaticAxisSymmetricParams{T}, u) where {T}
#     rθ = @SVector [u[2], u[3]]
#     comps = GeodesicTracer.metric_components(m, rθ)
#     @SMatrix [
#     comps[1] 0 0 comps[5]
#     0 comps[2] 0 0
#     0 0 comps[3] 0
#     comps[5] 0 0 comps[4]
#     ]
# end

a=0.0
M=1.0
# R_isco = CarterBoyerLindquist.rms(M, a)
r_range = 4.0:0.1:8.0
# m = BoyerLindquist(M=M, a=a)

e_lz = find_orbit_range(;m=m, r_range = r_range, a = a, upper = 0.1)

E_vals = []

for i in 1:length(r_range)
        E = e_lz[i, 1]
        push!(E_vals, E)
end

plt = plot(r_range, E_vals, xlabel="r (M)", ylabel="E")
display(plt)

r_isco_num = r_range[findmin(E_vals)[2]]
r_isco_true = AccretionFormulae.r_isco(a, M)

@show(r_isco_num)
@show(r_isco_true)


using SpecialFunctions, StatsFuns, Base.Iterators

"""
    log_int_prod_2_Gammas(α1, β1, α2, β2)

Compute the log of ``\\int_{\\mathbb{R}^+} Gamma(α1, β1)Gammas(α2, β2)`` exactly. This formula is only valid for ``α1 + α2 > 1``.

# Examples
```julia-repl
julia> log_int_prod_2_Gammas(1, 1, 1, 1) #== -log(2)
-0.6931471805599453
julia> log_int_prod_2_Gammas(2, 0.5, 1, 0.5) #==  log(0.5^2*0.5 )
-0.6931471805599453
```
"""
function log_int_prod_2_Gammas(α1, β1, α2, β2)
    @assert α1 + α2 > 1 "The formula for the L2 distance is only valid for α1 + α2 > 1"
    return α1 * log(β1) + α2 * log(β2) + lgamma_local(α1 + α2 - 1) - lgamma_local(α1) - lgamma_local(α2) - (α1 + α2 - 1) * log(β1 + β2)
end

"""
    int_prod_2_Gammas(α1, β1, α2, β2)

Compute ``\\int_{\\mathbb{R}^+} Gamma(α1, β1)Gammas(α2, β2)`` exactly. This formula is only valid for ``α1 + α2 > 1``.

# Examples
```julia-repl
julia> int_prod_2_Gammas(1, 1, 1, 1)
0.5
julia> int_prod_2_Gammas(2, 0.5, 1, 0.5) #==  0.5^2*0.5
-0.6931471805599453
```
"""
int_prod_2_Gammas(α1, β1, α2, β2) = log_int_prod_2_Gammas(α1, β1, α2, β2) |> exp

function log_L2_dist_Gamma_mixtures(logw1, αlist_1, βlist_1, logw2, αlist_2, βlist_2)
    I = length(logw1)
    J = length(logw2)
    @assert I == length(αlist_1)
    @assert I == length(βlist_1)
    @assert J == length(αlist_2)
    @assert J == length(βlist_2)
    # println("I=$I, J=$J")

    first_logterm = logsumexp(logw1[i] + logw1[j] + log_int_prod_2_Gammas(αlist_1[i], βlist_1[i], αlist_1[j], βlist_1[j]) for (i,j) in Base.Iterators.product(1:I, 1:I))

    second_logterm =  logsumexp(logw2[i] + logw2[j] + log_int_prod_2_Gammas(αlist_2[i], βlist_2[i], αlist_2[j], βlist_2[j]) for (i,j) in Base.Iterators.product(1:J, 1:J))

    third_log_term = logsumexp(log(2) + logw1[i] + logw2[j] + log_int_prod_2_Gammas(αlist_1[i], βlist_1[i], αlist_2[j], βlist_2[j]) for (i,j) in Base.Iterators.product(1:I, 1:J))

    # println([first_logterm, second_logterm, third_log_term])
    # println(exp.([first_logterm, second_logterm, third_log_term]))
    # println(sum(exp.([first_logterm, second_logterm, third_log_term]) .* [1,1,-1]))
    #
    # lx = [first_logterm, second_logterm, third_log_term]
    # signs = [1,1,-1]
    # m = maximum(lx)
    #
    # println("m= $m")
    # scaled_sum = sum(signs .* exp.(lx .- m))
    # println("scaled_sum= $scaled_sum")
    #
    # println("abs(scaled_sum) <= eps(Float64): $(abs(scaled_sum) <= eps(Float64))")

    log_squared_L2_dist = ExactWrightFisher.signed_logsumexp([first_logterm, second_logterm, third_log_term], [1,1,-1])[2] #The result will always be positive
    return 0.5*log_squared_L2_dist
end

"""
    log_int_prod_2_Dir(α1, α2)

Compute the log of ``\\int_{\\nabla_K} Dir(α1)Dir(α2)`` exactly. This formula is only valid for ``\\forall j \\in \\{1,\\ldots, K\\}, \\alpha_{1,j} + \\alpha_{2,j} > 1``.

# Examples
```julia-repl
julia> DualOptimalFiltering.log_int_prod_2_Dir([1,1], [1,1])
0.0
julia> log_int_prod_2_Gammas(2, 0.5, 1, 0.5) #==  log(0.5^2*0.5 )
-0.6931471805599453
```
"""
function log_int_prod_2_Dir(α1, α2)
    @assert all(α1 .+ α2 .> 1) "The formula for the L2 distance is only valid for α1 + α2 > 1 for each component of α"
    return lmvbeta(α1 .+ α2 .- 1) - lmvbeta(α1) - lmvbeta(α2)
end

"""
    lmvbeta(α::AbstractArray{T,1}) for T <: Real

Compute the log of the [multivariate beta function](https://en.wikipedia.org/wiki/Beta_function#Multivariate_beta_function): ``\\Beta(\\alpha_1,\\alpha_2,\\ldots\\alpha_n) = \\frac{\\Gamma(\\alpha_1)\\Gamma(\\alpha_2) \\cdots \\Gamma(\\alpha_n)}{\\Gamma(\\alpha_1 + \\alpha_2 + \\cdots + \\alpha_n)}``.

# Examples
```julia-repl
julia> lmvbeta([1,1])
0.0
```
"""
function lmvbeta(α::AbstractArray{T,1}) where T <: Real
    return sum(lgamma_local.(α)) - lgamma_local(sum(α))
end

# function log_L2_dist_Dirichlet_mixtures(logw1::RealVector, αlist_1::Array{RealVector, 1}, logw2::RealVector, αlist_2::Array{RealVector, 1})
function log_L2_dist_Dirichlet_mixtures(logw1::RealVector, αlist_1::Array{T, 1}, logw2::RealVector, αlist_2::Array{T, 1}) where T <: RealVector
    I = length(logw1)
    J = length(logw2)
    first_logterm = logsumexp(logw1[i] + logw1[j] + log_int_prod_2_Dir(αlist_1[i], αlist_1[j]) for (i,j) in Base.Iterators.product(1:I, 1:I))
    second_logterm =  logsumexp(logw2[i] + logw2[j] + log_int_prod_2_Dir(αlist_2[i], αlist_2[j]) for (i,j) in Base.Iterators.product(1:J, 1:J))
    third_log_term = logsumexp(log(2) + logw1[i] + logw2[j] + log_int_prod_2_Dir(αlist_1[i], αlist_2[j]) for (i,j) in Base.Iterators.product(1:I, 1:J))

    # println([first_logterm, second_logterm, third_log_term])
    # println(exp.([first_logterm, second_logterm, third_log_term]))
    # println(sum(exp.([first_logterm, second_logterm, third_log_term]) .* [1,1,-1]))
    #
    # lx = [first_logterm, second_logterm, third_log_term]
    # signs = [1,1,-1]
    # m = maximum(lx)
    #
    # println("m= $m")
    # scaled_sum = sum(signs .* exp.(lx .- m))
    # println("scaled_sum= $scaled_sum")
    #
    # println("abs(scaled_sum) <= eps(Float64): $(abs(scaled_sum) <= eps(Float64))")

    log_squared_L2_dist = ExactWrightFisher.signed_logsumexp([first_logterm, second_logterm, third_log_term], [1,1,-1])[2] #The result will always be positive
    return 0.5*log_squared_L2_dist
end

# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.6,0.4]), [[1,2,1],[4,2,3]])
# #
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]])
#
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.200001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.2000001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.20000001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# DualOptimalFiltering.log_L2_dist_Dirichlet_mixtures(log.([0.2,0.3,0.5]), [[1,2,3],[1,2,4],[5,3,2]], log.([0.2000000001,0.3,0.5] |> DualOptimalFiltering.normalise), [[1,2,3],[1,2,4],[5,3,2]])
#
# import DualOptimalFiltering.RR
#
# ExactWrightFisher.signed_logsumexp([log(RR(3.00000000000001)), log(RR(3)), log(RR(6))],[1,1,-1])
#
# ExactWrightFisher.signed_logsumexp([log(3), log(3), log(6)],[1,1,-1])


using OptiMo
using Random
using CUTEst

# problem
problems = CUTEst.select( min_var=1, max_var=10, min_con=1, max_con=10 );
problem = problems[1]
nlp = CUTEstModel( problem )

prob = NLPOptiModel( nlp );
x = copy( prob.meta.x0 )
nvar = prob.meta.nvar
ncon = prob.meta.ncon
R = eltype( x )

y = randn(R,ncon)
mu = rand(R,ncon)

alprob = AugLagOptiModel( prob, mu, y )

mu = rand(R,ncon)
update!( alprob, mu, y )

y = randn(R,ncon)
update!( alprob, mu, y )

lx = obj( alprob, x )
lx, dlx = objgrad( alprob, x )
gz, z = objprox( alprob, x, 0.1 )

finalize( nlp )


module optim

using LinearAlgebra, Statistics
export SGD, ADAM, Newton, optim_step, done

mutable struct SGD
    α
    x
    max_iter
    grad_tol
    log_file
    iter
    done

    function SGD(α :: Float64, x0 :: Array{Float64}, max_iter :: Int, grad_tol :: Float64, log_name :: String = nothing)
        if log_name !== nothing
            log_file = open(log_name, "w+")
        else
            log_file = nothing
        end

        new(α, x0, max_iter, grad_tol, log_file, 0, false)
    end
end

function log(o :: SGD, s :: String)
    if o.log_file !== nothing
        write(o.log_file, s)
        flush(o.log_file)
    end
end

function optim_step(o :: SGD, ss :: Matrix{Float64}, target_Es :: Vector{Float64})
    g = (mean(ss, dims=2) - target_Es)[:, 1]
    optim_step(o, g)
end

function optim_step(o :: SGD, g :: Vector{Float64})
    gn = norm(g)

    if o.iter > o.max_iter
        log(o, "Maximum iterations ($(o.max_iter)) reached. Terminating.\n")
        finish(o)
    elseif gn ≤ o.grad_tol
        log(o, "Gradient tolerance (|g| = $gn ≤ $(o.grad_tol)) reached. Terminating.\n")
        finish(o)
    else
        log(o, "Iteration $(o.iter): |g| = $gn.\n")
        o.x -= o.α * g
        o.iter += 1
    end

    o.x
end

function finish(o :: SGD)
    o.done = true
    close(o.log_file)
end

function done(o :: SGD)
    o.done
end

mutable struct ADAM
    α
    β1
    β2
    ε
    x
    m
    v
    max_iter
    grad_tol
    log_file
    iter
    done

    function ADAM(
        α :: Float64, β1 :: Float64, β2 :: Float64, ε :: Float64,
        x0 :: Array{Float64}, max_iter :: Int, grad_tol :: Float64,
        log_name :: String = nothing
    )
        if log_name !== nothing
            log_file = open(log_name, "w+")
        else
            log_file = nothing
        end

        m0 = zeros(size(x0))
        v0 = zeros(size(x0))

        new(
            α, β1, β2, ε,
            x0, m0, v0,
            max_iter, grad_tol, log_file,
            0, false
        )
    end
end

function log(o :: ADAM, s :: String)
    if o.log_file !== nothing
        write(o.log_file, s)
        flush(o.log_file)
    end
end

function optim_step(o :: ADAM, ss :: Matrix{Float64}, target_Es :: Vector{Float64})
    G = ss .- target_Es
    g = mean(G, dims=2)
    gn = norm(g)

    if o.iter > o.max_iter
        log(o, "Maximum iterations ($(o.max_iter)) reached. Terminating.\n")
        finish(o)
    elseif gn ≤ o.grad_tol
        log(o, "Gradient tolerance (|g| = $gn ≤ $(o.grad_tol)) reached. Terminating.\n")
        finish(o)
    else
        log(o, "Iteration $(o.iter): |g| = $gn.\n")
        o.iter += 1
        o.m = o.β1 * o.m + (1 - o.β1) * g
        o.v = o.β2 * o.v + (1 - o.β2) * mean(G .^ 2, dims=2)
        mc = o.m / (1 - o.β1 ^ o.iter)
        vc = o.v / (1 - o.β2 ^ o.iter)
        o.x -= o.α * mc ./ (sqrt.(vc) .+ o.ε)
    end

    o.x
end

function finish(o :: ADAM)
    o.done = true
    close(o.log_file)
end

function done(o :: ADAM)
    o.done
end

mutable struct Newton
    x
    max_iter
    grad_tol
    log_file
    iter
    done

    function Newton(x0 :: Vector{Float64}, max_iter :: Int, grad_tol :: Float64, log_name :: String = nothing)
        if log_name !== nothing
            log_file = open(log_name, "w+")
        else
            log_file = nothing
        end

        new(x0, max_iter, grad_tol, log_file, 0, false)
    end
end

function log(o :: Newton, s :: String)
    if o.log_file !== nothing
        write(o.log_file, s)
        flush(o.log_file)
    end
end

function optim_step(o :: Newton, ss :: Matrix{Float64}, target_Es :: Vector{Float64})
    p, n = size(ss)
    g = mean(ss, dims=2) - target_Es
    gn = norm(g)

    if o.iter > o.max_iter
        log(o, "Maximum iterations ($(o.max_iter)) reached. Terminating.\n")
        finish(o)
    elseif gn ≤ o.grad_tol
        log(o, "Gradient tolerance (|g| = $gn ≤ $(o.grad_tol)) reached. Terminating.\n")
        finish(o)
    else
        log(o, "Iteration $(o.iter): |g| = $gn.\n")
        H = 1 / (n - 1) * ss * ss'
        o.x -= H^-1 * g
        o.iter += 1
    end

    o.x
end

function finish(o :: Newton)
    o.done = true
    close(o.log_file)
end

function done(o :: Newton)
    o.done
end

end


# variational inference sample code

module VariationalInference
  using DataFrames, CSV, StatsBase, Plots, Distributions
  using SpecialFunctions
  pyplot()

  function main(;test=false)
    # read original data
    data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_600.txt")
    all_data = CSV.read(data_path, DataFrame,
                        header=["date", "time", "ir", "z"],
                        delim=' ')
    data = sort(sample(all_data.z, 3000))
    data = DataFrame(z=data)

    # initialize responsibility
    K = 3 # num of class
    n = Int(ceil(size(data)[1]/K))
    for k in 1:K
      data_k = []
      for i in 1:size(data)[1]
        if k == Int(ceil(i/n))
          append!(data_k, 1.0)
        else
          append!(data_k, 0.0)
        end
      end

      if k == 1
        data = hcat(data, DataFrame(one=data_k))
      elseif k == 2
        data = hcat(data, DataFrame(two=data_k))
      else
        data = hcat(data, DataFrame(three=data_k))
      end
    end

    # iteration
    params_history = Dict()
    params = DataFrame()
    for t in 1:10000
      data, params = one_step(data, K)
      # record per 10 times
      if t % 10 == 0
        params_history[t] = params
        println("$(t) $(params)")

        if length(params_history) < 2
          continue
        end

        if all([(abs(params_history[t-10][k, "tau"] - params_history[t][k, "tau"]) < 10e-5) for k in 1:K])
          break
        end
      end
    end
    draw_sensor_dist(data, params, K,
                     "src/variational_inference/draw_sensor_dist_K3_t340.png",
                     test)
  end

  function update_parameters(ds, k, mu_avg=600, zeta=1, alpha=1, beta=1, tau=1)
    if k == 1
      key_class = "one"
    elseif k == 2
      key_class = "two"
    else
      key_class = "three"
    end

    # calculate R, S, T
    R = sum([ds[i, key_class] for i in 1:size(ds)[1]])
    S = sum([ds[i, key_class]*ds[i, :z] for i in 1:size(ds)[1]])
    T = sum([ds[i, key_class]*(ds[i, :z]^2) for i in 1:size(ds)[1]])

    # store calculated parameters temporarily
    hat = Dict()
    hat["tau"] = R + tau
    hat["zeta"] = R + zeta
    hat["mu_avg"] = (S + zeta*mu_avg)/hat["zeta"]
    hat["alpha"] = R/2 + alpha
    hat["beta"] = (T + zeta*(mu_avg^2) - hat["zeta"]*(hat["mu_avg"]^2))/2 + beta
    hat["z_std"] = sqrt(hat["beta"]/hat["alpha"])

    return DataFrame(hat)
  end

  function draw_sensor_dist(data, params, K, save_name, test)
    pi_ = rand(Dirichlet([params[k, "tau"] for k in 1:K]))
    pdfs = [Normal(params[k, "mu_avg"], params[k, "z_std"]) for k in 1:K]

    xs = range(600, 650, length=Int64(floor(50/0.5)))

    # histogram of observation
    histogram(data.z, bins=(maximum(data.z)-minimum(data.z)), color=:gray, label="histogram")
    # draw gaussian mixture distribution
    ys = [sum([pdf(pdfs[k], x)*pi_[k] for k in 1:K])*size(data)[1] for x in xs]
    plot!(xs, ys, color=:red, label="Mixture")
    # gaussian distribution of each cluster
    for k in 1:K
      ys = [pdf(pdfs[k], x)*pi_[k]*size(data)[1] for x in xs]
      plot!(xs, ys, color=:blue, label="Cluster$(k)")
    end
    if test == false
      save_path = joinpath(split(@__FILE__, "src")[1], save_name)
      savefig(save_path)
    end
  end

  function responsibility(z, K, params)
    tau_sum = sum([params[k, "tau"] for k in 1:K])

    r = Dict()
    for k in 1:K
      log_rho = (digamma(params[k, "alpha"]) - log(params[k, "beta"]))/2 - (1/params[k, "zeta"] + ((params[k, "mu_avg"]-z)^2)*params[k, "alpha"]/params[k, "beta"])/2 + digamma(params[k, "tau"]) - digamma(tau_sum)
      if k == 1
        r["one"] = exp(log_rho)
      elseif k == 2
        r["two"] = exp(log_rho)
      else
        r["three"] = exp(log_rho)
      end

    end

    # sum of r_k
    rks = []
    for k in 1:K
      if k == 1
        append!(rks, r["one"])
      elseif k == 2
        append!(rks, r["two"])
      else
        append!(rks, r["three"])
      end
    end
    w = sum(rks)

    # normalization
    for k in 1:K
      if k == 1
        r["one"] /= w
      elseif k == 2
        r["two"] /= w
      else
        r["three"] /= w
      end
    end

    return r
  end

  function one_step(data, K)
    # update parameters
    params = DataFrame()
    for k in 1:K
      if k == 1
        params = update_parameters(data, k)
      else
        params = vcat(params, update_parameters(data, k))
      end
    end

    # update responsibility
    rs = [responsibility(data[i, :z], K, params) for i in 1:size(data)[1]]
    for k in 1:K
      if k == 1
        data.one = [rs[i]["one"] for i in 1:size(data)[1]]
      elseif k == 2
        data.two = [rs[i]["two"] for i in 1:size(data)[1]]
      else
        data.three = [rs[i]["three"] for i in 1:size(data)[1]]
      end
    end

    return data, params
  end

end


using CairoMakie, AbstractPlotting, GeoMakie

lons = LinRange(-179.5, 179.5, 360)
lats = LinRange(-89.5, 89.5, 180)

field = [exp(cosd(l)) + 3(y/90) for l in lons, y in lats]

source = LonLat()
dest = WinkelTripel()

points, faces = GeoMakie.triangulated_grid(lons, lats)

tpoints = Proj4.transform.(source, dest, points)

scene = mesh(
    tpoints, faces;
    color = GeoMakie.img2colorvec(field),
    shading = false,
    show_axis = false, scale_plot = false, # don't show the normal axis and
                                           # ensure the aspect ratio is correct
    resolution = (1200, 600)               # just a nice resolution for this projection
)

geoaxis!(scene, extrema(lons), extrema(lats); crs = (src = source, dest = dest,))

coastlines!(scene; crs = (src = source, dest = dest))

save("geomakie-readme.png", scene) # works
save("geomakie-readme.svg", scene) # works
save("geomakie-readme.pdf", scene) # works but is really large


## Two electron integrals

function coulomb(aexpn::Float64,ax::Float64,ay::Float64,az::Float64,
                 aI::Int64,aJ::Int64,aK::Int64,
                 bexpn::Float64,bx::Float64,by::Float64,bz::Float64,
                 bI::Int64,bJ::Int64,bK::Int64,
                 cexpn::Float64,cx::Float64,cy::Float64,cz::Float64,
                 cI::Int64,cJ::Int64,cK::Int64,
                 dexpn::Float64,dx::Float64,dy::Float64,dz::Float64,
                 dI::Int64,dJ::Int64,dK::Int64)
    # This is the slow method of computing integrals from Huzinaga et al.
    # Use the HRR/VRR scheme from Head-Gordon & Pople instead

    rab2 = dist2(ax-bx,ay-by,az-bz)
    rcd2 = dist2(cx-dx,cy-dy,cz-dz)

    px,py,pz = gaussian_product_center(aexpn,ax,ay,az,bexpn,bx,by,bz)
    qx,qy,qz = gaussian_product_center(cexpn,cx,cy,cz,dexpn,dx,dy,dz)
    rpq2 = dist2(px-qx,py-qy,pz-qz)

    g1 = aexpn+bexpn
    g2 = cexpn+dexpn
    delta = 0.25*(1/g1+1/g2)

    Bx = Barray(aI,bI,cI,dI,px,ax,bx,qx,cx,dx,g1,g2,delta)
    By = Barray(aJ,bJ,cJ,dJ,py,ay,by,qy,cy,dy,g1,g2,delta)
    Bz = Barray(aK,bK,cK,dK,pz,az,bz,qz,cz,dz,g1,g2,delta)

    s = 0
    #println("$(aI+bI+cI+dI),$(aJ+bJ+cJ+dJ),$(aK+bK+cK+dK)")
    for I in 0:(aI+bI+cI+dI)
        for J in 0:(aJ+bJ+cJ+dJ)
            for K in 0:(aK+bK+cK+dK)
                #println("coul: $I,$J,$K,$(Bx[I+1]),$(By[J+1]),$(Bz[K+1])")
                s += Bx[I+1]*By[J+1]*Bz[K+1]*Fgamma(I+J+K,0.25*rpq2/delta)
            end
        end
    end
    return 2pi^(2.5)/(g1*g2*sqrt(g1+g2))*exp(-aexpn*bexpn*rab2/g1)*exp(-cexpn*dexpn*rcd2/g2)*s
end

function coulomb(a::PGBF,b::PGBF,c::PGBF,d::PGBF)
    return a.norm*b.norm*c.norm*d.norm*coulomb(a.expn,a.x,a.y,a.z,a.I,a.J,a.K,
        b.expn,b.x,b.y,b.z,b.I,b.J,b.K,
        c.expn,c.x,c.y,c.z,c.I,c.J,c.K,
        d.expn,d.x,d.y,d.z,d.I,d.J,d.K)
end

fB(i::Int64,l1::Int64,l2::Int64,p::Float64,a::Float64,b::Float64,r::Int64,g::Float64) = binomial_prefactor(i,l1,l2,p-a,p-b)*B0(i,r,g)
B0(i::Int64,r::Int64,g::Float64) = fact_ratio2(i,r)*(4g)^(r-i)
fact_ratio2(a::Int64,b::Int64) = factorial(a)factorial(b)/factorial(a-2b)

function Bterm(i1::Int64,i2::Int64,r1::Int64,r2::Int64,u::Int64,
               l1::Int64,l2::Int64,l3::Int64,l4::Int64,
               Px::Float64,Ax::Float64,Bx::Float64,Qx::Float64,Cx::Float64,Dx::Float64,
               gamma1::Float64,gamma2::Float64,delta::Float64)
    # THO eq. 2.22
    #print("Bterm($i1,$i2,$r1,$r2,$u,$l1,$l2,$l3,$l4,$Px,$Ax,$Bx,$Qx,$Cx,$Dx,$gamma1,$gamma2,$delta)=")
    val = (-1)^(i2+u)*fB(i1,l1,l2,Px,Ax,Bx,r1,gamma1)*fB(i2,l3,l4,Qx,Cx,Dx,r2,gamma2)*(
            fact_ratio2(i1+i2-2*(r1+r2),u)*(Qx-Px)^(i1+i2-2*(r1+r2)-2*u)/delta^(i1+i2-2*(r1+r2)-u))
    #println("$val")
    return val
end

function Barray(l1::Int64,l2::Int64,l3::Int64,l4::Int64,p::Float64,a::Float64,b::Float64,
                q::Float64,c::Float64,d::Float64,g1::Float64,g2::Float64,delta::Float64)
    Imax = l1+l2+l3+l4+1
    B = zeros(Float64,Imax)
    for i1 in 0:(l1+l2)
        for i2 in 0:(l3+l4)
            for r1 in 0:div(i1,2)
                for r2 in 0:div(i2,2)
                    for u in 0:(div(i1+i2,2)-r1-r2)
                        I = i1+i2-2*(r1+r2)-u
                        B[I+1] += Bterm(i1,i2,r1,r2,u,l1,l2,l3,l4,p,a,b,q,c,d,g1,g2,delta)
                    end
                end
            end
        end
    end
    return B
end

coulomb(a::CGBF,b::CGBF,c::CGBF,d::CGBF) = contract(coulomb,a,b,c,d)

function test_two_terms()
    @testset "test_two_terms" begin
    @test fB(0,0,0,0.0,0.0,0.0,0,2.0) == 1
    @test fB(0,0,0,1.0,1.0,1.0,0,2.0) == 1
    @test fB(0,0,0,0.0,0.0,0.0,0,2.0 ) == 1
    @test fB(1,0,1,0.0,0.0,0.0,0,2.0 ) == 0.125
    @test B0(0,0,2.0) == 1
    @test fact_ratio2(0,0) == 1
    @test Bterm(0,0,0,0,0,0,0,0,0,0.0,0.0,0.0,0.0,0.0,0.0,2.0,2.0,0.25)==1
    @test Bterm(0,1,0,0,0,0,0,0,1,0.0,0.0,0.0,0.0,0.0,0.0,2.0,2.0,0.25)==0
    end #testset
end

function test_coul1()
    @testset "test_coul1" begin
    s = pgbf(1.0)
    px = pgbf(1.0,0,0,0,1,0,0)
    @test coulomb(s,s,s,px)==0 # 0
    @test coulomb(s,s,px,px) ≈ 0.9403159725793305
    end #testset
end

function coulomb_hgp(a::PGBF,b::PGBF,c::PGBF,d::PGBF)
    return a.norm*b.norm*c.norm*d.norm*hrr(a.expn,a.x,a.y,a.z,a.I,a.J,a.K,
        b.expn,b.x,b.y,b.z,b.I,b.J,b.K,
        c.expn,c.x,c.y,c.z,c.I,c.J,c.K,
        d.expn,d.x,d.y,d.z,d.I,d.J,d.K)
end
coulomb_hgp(a::CGBF,b::CGBF,c::CGBF,d::CGBF) = contract(coulomb_hgp,a,b,c,d)


module TestSynthesis
using CustomTest
using Base.Test
using ControlSystems

P = tf(1,[1,1])
C = tf([1,1],[1,0])
L = P*C
Lsys = ss(L)

B = [1]
A = [1,1]
R = [1,1]
S = [1]
T = [1]

@test isapprox(minreal(feedback(P,C),1e-5), minreal(feedback(L),1e-5))
@test isapprox(numpoly(minreal(feedback(L),1e-5))[1].a, numpoly(tf(1,[1,1]))[1].a)# This test is ugly, but numerical stability is poor for minreal
@test feedback2dof(B,A,R,S,T) == tf(B.*T, conv(A,R) + [0;0;conv(B,S)])
@test feedback2dof(P,R,S,T) == tf(B.*T, conv(A,R) + [0;0;conv(B,S)])
@test isapprox(pole(minreal(tf(feedback(Lsys)),1e-5)) , pole(minreal(feedback(L),1e-5)), atol=1e-5)

@test_throws ErrorException feedback(ss(1),ss(1))
@test_throws ErrorException feedback(ss(eye(2), ones(2,2), ones(1,2),0))

end


using AmbientForcing
using OrdinaryDiffEq, Distributions

function rober(du,u,p,t)
  y₁,y₂,y₃ = u
  du[1] = -0.04 * y₁ + 1e4 * y₂ * y₃
  du[2] =  0.04 * y₁ - 1e4 * y₂ * y₃ - 3e7 * y₂^2
  du[3] =  y₁ + y₂ + y₃ - 1
  nothing
end

M = [1. 0  0
     0  1. 0
     0  0  0];

ode_rober = ODEFunction(rober, mass_matrix = M)

u0 = [1.0,0.0,0.0];

g_rober = constraint_equations(ode_rober)
g_rober(u0)

Frand = random_force(ode_rober, [0.0,1], Uniform)
z_new = ambient_forcing(ode_rober, u0, 2.0, Frand)

g_rober(z_new)

h = [0, 1 ,0]
z_new = ambient_forcing(ode_rober, u0, 2.0, h)

g_rober(z_new)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


fourx(x::T) where {T<:IEEEFloat} = four(T) * x
twox(x::T) where {T<:IEEEFloat} = two(T) * x
halfx(x::T) where {T<:IEEEFloat} = onehalf(T) * x
quarterx(x::T) where {T<:IEEEFloat} = onequarter(T) * x

# this is "TwoSum"
@inline function add_2(a::T, b::T) where {T<:AbstractFloat}
    s = a + b
    v = s - a
    e = (a - (s - v)) + (b - v)
    return s, e
end

@inline add_(a::T, b::T) where {T<:AbstractFloat} = add_2(a, b)

@inline function add_1(a::T, b::T) where {T<:AbstractFloat}
    s = a + b
    return s
end

# ThreeSum
function add_3(a::T,b::T,c::T) where {T<:AbstractFloat}
    s, t = add_(b, c)
    x, u = add_(a, s)
    y, z = add_(u, t)
    x, y = add_(x, y)
    return x, y, z
end

@inline add_(a::T, b::T, c::T) where {T<:AbstractFloat} = add_3(a, b, c)

function add_2(a::T,b::T,c::T) where {T<:AbstractFloat}
    s, t = add_(b, c)
    x, u = add_(a, s)
    y    = u + t
    x, y = add_(x, y)
    return x, y
end

function add_1(a::T,b::T,c::T) where {T<:AbstractFloat}
    s, t = add_(b, c)
    x, u = add_(a, s)
    y    = u + t
    x   += y
    return x
end

# FourSum
function add_4(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = add_(a ,  b)
    t0, t2 = add_(t0,  c)
    a,  t3 = add_(t0,  d)
    t0, t1 = add_(t1, t2)
    b,  t2 = add_(t0, t3)
    c,  d  = add_(t1, t2)
    return a, b, c, d
end

function add_4(ac::Tuple{T, T}, bd::Tuple{T, T}) where {T<: AbstractFloat}
    a, c = ac
    b, d = bd
    t0, t1 = add_(a ,  b)
    t0, t2 = add_(t0,  c)
    a,  t3 = add_(t0,  d)
    t0, t1 = add_(t1, t2)
    b,  t2 = add_(t0, t3)
    c,  d  = add_(t1, t2)
    return a, b, c, d
end

@inline add_(a::T, b::T, c::T, d::T) where {T<:AbstractFloat} = add_4(a, b, c, d)

function add_3(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = add_(a ,  b)
    t0, t2 = add_(t0,  c)
    a,  t3 = add_(t0,  d)
    t0, t1 = add_(t1, t2)
    b,  t2 = add_(t0, t3)
    c  = t1 + t2
    return a, b, c
end

function add_2(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = add_(a ,  b)
    t0, t2 = add_(t0,  c)
    a,  t3 = add_(t0,  d)
    t0  = t1 + t2
    b   = t0 + t3
    return a, b
end

function add_1(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = add_(a ,  b)
    t0, t2 = add_(t0,  c)
    a = t0 + d
    return a
end

# FiveSum
function add_5(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_(y, z)
    t0, t3 = add_(x, t0)
    t0, t2 = add_(w, t0)
    a, t1  = add_(v, t0)
    t0, t3 = add_(t3, t4)
    t0, t2 = add_(t2, t0)
    b, t1  = add_(t1, t0)
    t0, t2 = add_(t2, t3)
    c, t1  = add_(t1, t0)
    d, e   = add_(t1, t2)
    return a, b, c, d, e
end

@inline add_(a::T, b::T, c::T, d::T, e::T) where {T<:AbstractFloat} =
    add_5(a, b, c, d, e)

function add_4(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_(y, z)
    t0, t3 = add_(x, t0)
    t0, t2 = add_(w, t0)
    a, t1  = add_(v, t0)
    t0, t3 = add_(t3, t4)
    t0, t2 = add_(t2, t0)
    b, t1  = add_(t1, t0)
    t0, t2 = add_(t2, t3)
    c, t1  = add_(t1, t0)
    d      = t1 + t2
    return a, b, c, d
end

function add_3(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_(y, z)
    t0, t3 = add_(x, t0)
    t0, t2 = add_(w, t0)
    a, t1  = add_(v, t0)
    t0, t3 = add_(t3, t4)
    t0, t2 = add_(t2, t0)
    b, t1  = add_(t1, t0)
    t0, t2 = add_(t2, t3)
    c      = t1 + t0
    return a, b, c
end

function add_2(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_(y, z)
    t0, t3 = add_(x, t0)
    t0, t2 = add_(w, t0)
    a, t1  = add_(v, t0)
    t0, t3 = add_(t3, t4)
    t0, t2 = add_(t2, t0)
    b      = t1 + t0
    return a, b
end

function add_1(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_(y, z)
    t0, t3 = add_(x, t0)
    t0, t2 = add_(w, t0)
    a  = v + t0
    return a
end

# this is TwoDiff
@inline function sub_2(a::T, b::T) where {T<:AbstractFloat}
    s = a - b
    v = s - a
    e = (a - (s - v)) - (b + v)
    return s, e
end

@inline sub_(a::T, b::T) where {T<:AbstractFloat} = sub_2(a, b)

# a - b - c
function sub_3(a::T, b::T, c::T) where {T<:AbstractFloat}
    s, t = sub_(-b, c)
    x, u = add_(a, s)
    y, z = add_(u, t)
    x, y = add_(x, y)
    return x, y, z
end
# a - b - c
function sub_2(a::T, b::T, c::T) where {T<:AbstractFloat}
    s, t = sub_(-b, c)
    x, u = add_(a, s)
       y = u + t
    x, y = add_(x, y)
    return x, y
end

# a - b - c - d
function sub_4(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = sub_(a ,  b)
    t0, t2 = sub_(t0,  c)
    a,  t3 = sub_(t0,  d)
    t0, t1 = add_(t1, t2)
    b,  t2 = add_(t0, t3)
    c,  d  = add_(t1, t2)
    return a, b, c, d
end

# a - b - c - d
function sub_3(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = sub_(a ,  b)
    t0, t2 = sub_(t0,  c)
    a,  t3 = sub_(t0,  d)
    t0, t1 = add_(t1, t2)
    b,  t2 = add_(t0, t3)
    c  = t1 + t2
    return a, b, c
end

# a - b - c - d
function sub_2(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = sub_(a ,  b)
    t0, t2 = sub_(t0,  c)
    a,  t3 = sub_(t0,  d)
    t0  = t1 + t2
    b   = t0 + t3
    return a, b
end

@inline sub_(a::T, b::T, c::T) where {T<:AbstractFloat} = sub_3(a, b, c)

function sub_(ac::Tuple{T, T}, bd::Tuple{T, T}) where {T<: AbstractFloat}
    a, c = ac
    b, d = bd
    t0, t1 = sub_(a ,  b)
    t0, t2 = add_(t0,  c)
    a,  t3 = sub_(t0,  d)
    t0, t1 = add_(t1, t2)
    b,  t2 = add_(t0, t3)
    c,  d  = add_(t1, t2)
    return a, b, c, d
end

# this is QuickTwoSum, presumes abs(a) >= abs(b)
@inline function add_hilo_2(a::T, b::T) where {T<:AbstractFloat}
    s = a + b
    e = b - (s - a)
    return s, e
end

@inline add_hilo_(a::T, b::T) where {T<:AbstractFloat} =
    add_hilo_2(a, b)

# QuickThreeSum, presumes abs(a) >= abs(b) >= abs(c)
function add_hilo_3(a::T,b::T,c::T) where {T<:AbstractFloat}
    s, t = add_hilo_(b, c)
    x, u = add_hilo_(a, s)
    y, z = add_hilo_(u, t)
    x, y = add_hilo_(x, y)
    return x, y, z
end

@inline add_hilo_(a::T, b::T, c::T) where {T<:AbstractFloat} =
    add_hilo_3(a, b, c)

function add_hilo_2(a::T,b::T,c::T) where {T<:AbstractFloat}
    s, t = add_hilo_(b, c)
    x, u = add_hilo_(a, s)
    y    = u + t
    x, y = add_hilo_(x, y)
    return x, y
end

# QuickFourSum, presumes abs(a) >= abs(b) >= abs(c) >= abs(d)
function add_hilo_4(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = add_hilo_(a ,  b)
    t0, t2 = add_hilo_(t0,  c)
    a,  t3 = add_hilo_(t0,  d)
    t0, t1 = add_hilo_(t1, t2)
    b,  t2 = add_hilo_(t0, t3)
    c,  d  = add_hilo_(t1, t2)
    return a, b, c, d
end

@inline add_hilo_(a::T, b::T, c::T, d::T) where {T<:AbstractFloat} =
    add_hilo_4(a, b, c, d)

function add_hilo_3(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = add_hilo_(a ,  b)
    t0, t2 = add_hilo_(t0,  c)
    a,  t3 = add_hilo_(t0,  d)
    t0, t1 = add_hilo_(t1, t2)
    b,  t2 = add_hilo_(t0, t3)
    c      = t1 + t2
    return a, b, c
end

function add_hilo_2(a::T,b::T,c::T,d::T) where {T<: AbstractFloat}
    t0, t1 = add_hilo_(a ,  b)
    t0, t2 = add_hilo_(t0,  c)
    a,  t3 = add_hilo_(t0,  d)
    t0, t1 = add_hilo_(t1, t2)
    b      = t0 + t3
    return a, b
end

# QuickFiveSum presumes abs(v) >= abs(w) >= abs(x) >= abs(y) >= abs(z)
function add_hilo_5(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_hilo_(y, z)
    t0, t3 = add_hilo_(x, t0)
    t0, t2 = add_hilo_(w, t0)
    a, t1  = add_hilo_(v, t0)
    t0, t3 = add_hilo_(t3, t4)
    t0, t2 = add_hilo_(t2, t0)
    b, t1  = add_hilo_(t1, t0)
    t0, t2 = add_hilo_(t2, t3)
    c, t1  = add_hilo_(t1, t0)
    d, e   = add_hilo_(t1, t2)
    return a, b, c, d, e
end

@inline add_hilo_(a::T, b::T, c::T, d::T, e::T) where {T<:AbstractFloat} =
    add_hilo_5(a, b, c, d, e)

function add_hilo_4(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_hilo_(y, z)
    t0, t3 = add_hilo_(x, t0)
    t0, t2 = add_hilo_(w, t0)
    a, t1  = add_hilo_(v, t0)
    t0, t3 = add_hilo_(t3, t4)
    t0, t2 = add_hilo_(t2, t0)
    b, t1  = add_hilo_(t1, t0)
    t0, t2 = add_hilo_(t2, t3)
    c, t1  = add_hilo_(t1, t0)
    d      = t1 + t2
    return a, b, c, d
end

function add_hilo_3(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_hilo_(y, z)
    t0, t3 = add_hilo_(x, t0)
    t0, t2 = add_hilo_(w, t0)
    a, t1  = add_hilo_(v, t0)
    t0, t3 = add_hilo_(t3, t4)
    t0, t2 = add_hilo_(t2, t0)
    b, t1  = add_hilo_(t1, t0)
    t0, t2 = add_hilo_(t2, t3)
    c      = t1 + t0
    return a, b, c
end

function add_hilo_2(v::T, w::T, x::T, y::T, z::T) where {T<:AbstractFloat}
    t0, t4 = add_hilo_(y, z)
    t0, t3 = add_hilo_(x, t0)
    t0, t2 = add_hilo_(w, t0)
    a, t1  = add_hilo_(v, t0)
    t0, t3 = add_hilo_(t3, t4)
    t0, t2 = add_hilo_(t2, t0)
    b      = t1 + t0
    return a, b
end

# this is QuickTwoDiff, requires abs(a) >= abs(b)
@inline function sub_hilo_2(a::T, b::T) where {T<:AbstractFloat}
    s = a - b
    e = (a - s) - b
    s, e
end

@inline sub_hilo_(a::T, b::T) where {T<:AbstractFloat} =
    sub_hilo_2(a, b)

function sub_hilo_3(a::T,b::T,c::T) where {T<:AbstractFloat}
    s, t = sub_hilo_(-b, c)
    x, u = add_hilo_(a, s)
    y, z = add_hilo_(u, t)
    x, y = add_hilo_(x, y)
    return x, y, z
end

@inline sub_hilo_(a::T, b::T, c::T) where {T<:AbstractFloat} =
    sub_hilo_3(a, b, c)

# this is TwoProdFMA
@inline function mul_2(a::T, b::T) where {T<:AbstractFloat}
    p = a * b
    e = fma(a, b, -p)
    return p, e
end

@inline mul_(a::T, b::T) where {T<:AbstractFloat} = mul_2(a, b)

function mul_3(a::T, b::T, c::T) where {T<:AbstractFloat}
    abhi, ablo = two_prod(a, b)
    hi, abhiclo = two_prod(abhi, c)
    ablochi, abloclo = two_prod(ablo, c)
    md, lo, tmp  = three_sum(ablochi, abhiclo, abloclo)
    return hi, md, lo
end

@inline function mul_2(a::T, b::T, c::T) where {T<:AbstractFloat}
    abhi, ablo = two_prod(a, b)
    hi, abhiclo = two_prod(abhi, c)
    ablochi, abloclo = two_prod(ablo, c)
    md = ablochi + (abhiclo + abloclo)
    return hi, md
end

@inline function mul_(a::T, b::T, c::T) where {T<:AbstractFloat}
    return mul_3(a, b, c)
end

# a squared
@inline function powr2_2(a::T) where {T<:AbstractFloat}
    p = a * a
    e = fma(a, a, -p)
    return p, e
end

@inline powr2_(a::T) where {T<:AbstractFloat} = powr2_2(a)

#= a cubed
function powr3_2(a::T) where {T<:AbstractFloat}
    y = a*a; z = fma(a, a, -y)
    x = y*a; y = fma(y, a, -x)
    z = fma(z,a,y)
    return x, z
end

@inline powr3_(a::T) where {T<:AbstractFloat} = powr3_2(a)

function powr3_3(a::T) where {T<:AbstractFloat}
    y, z = mul_(a, a)
    x, y = mul_(y, a)
    z, t = mul_(z, a)
    y, z = add_hilo_(y, z)
    z += t
    return x, y, z
end
=#

#=
   xfma algorithm from
   Sylvie Boldo and Jean-Michel Muller
   Some Functions Computable with a Fused-mac
=#

"""
    fma_(a, b, c) => (x, y, z)

Computes `x = fl(fma_(a, b, c))` and `y, z = fl(err(fma_(a, b, c)))`.
"""
function fma_(a::T, b::T, c::T) where {T<:AbstractFloat}
     x = fma(a, b, c)

     y, z = mul_(a, b)
     t, z = add_(c, z)
     t, u = add_(y, t)
     y = ((t - x) + u)

     y, z = add_hilo_(y, z)
     return T<:IEEEFloat ? (x, y, z) : add_(x, y, z)
end

fma_3(a::T, b::T, c::T) where {T<:AbstractFloat} = fma_(a, b, c)

function fma_2(a::T, b::T, c::T) where {T<:AbstractFloat}
     x = fma(a, b, c)

     y, z = mul_2(a, b)
     t, z = add_2(c, z)
     t, u = add_2(y, t)
     y = ((t - x) + u)

     y += z

     return T<:IEEEFloat ? (x, y) : add_(x, y)
end

"""
    fms_(a, b, c) => (x, y, z)

Computes `x = fl(fms_(a, b, c))` and `y, z = fl(err(fms_(a, b, c)))`.
"""
@inline function fms_(a::T, b::T, c::T) where {T<:IEEEFloat}
     return fma_(a, b, -c)
end

fms_(a::T, b::T, c::T) where {T<:AbstractFloat} =
    add_(fms(a, b, c)...,)

fms_3(a::T, b::T, c::T) where {T<:AbstractFloat} = fms_(a, b, c)
fms_2(a::T, b::T, c::T) where {T<:AbstractFloat} = fma_2(a, b, -c)


using VortexDistributions, Plots

Nx=1000; Ny=1000;
θx = LinRange(0,4*pi,Nx); θy = LinRange(0,4*pi,Ny);
psix = exp.(im*θx) + 0.1*(randn(Nx)+im*randn(Nx)); psiy = exp.(im*θy) + 0.1*(randn(Ny)+im*randn(Ny));
ϕx = angle.(psix); ϕy = angle.(psiy);
ϕ = [ϕx ϕy]
ϕu = unwrap(ϕ)
ϕui = similar(ϕu)
unwrap!(ϕui,ϕ)
@test ϕui == ϕu

ϕ = zeros(2, Nx)
ϕ[1, :] .= ϕx
ϕ[2, :] .=ϕy
ϕu = unwrap(ϕ,2)
ϕui = similar(ϕu)
unwrap!(ϕui,ϕ,2)
@test ϕui == ϕu


module ConvexHull

using JuMP, GenericLinearAlgebra, RowEchelon, Combinatorics

using Polyhedra
using Printf: @printf

export homogeneous_system, read_ine, write_ine, read_ext, write_ext,
       double_description, canonicalize!,
       get_extrema, is_approx_included

const ε = 10eps()

include("double_description.jl")
include("jump_interface.jl")
include("readers_writers.jl")
include("polyhedron.jl")

@static if VERSION < v"1.1"
    @eval isnothing(x) = x === nothing
end

end


module sphere_examples
using FinEtools
using FinEtoolsAcoustics
using FinEtools.MeshExportModule
using LinearAlgebra

function sphere_dipole_1()
    println("The interior sphere accelerates in the positive x-direction, generating
    positive pressure ahead of it, negative pressure behind
    ")
    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    a_amplitude=1.0*phun("mm/s^2");# amplitude of the  acceleration of the sphere
    # omega = 2000*phun("rev/s");      # frequency of the incident wave
    R = 5.0*phun("mm"); # radius of the interior sphere
    Ro = 4*R # radius of the external sphere
    P_amplitude = R*rho*a_amplitude; # pressure amplitude
    nref=2;
    nlayers=40;
    tolerance = Ro/(nlayers)/100
    frequency=2000; # Hz
    omega=2*pi*frequency;
    k = omega/c*[1.0 0.0 0.0];

    println("""
    Spherical domain filled with acoustic medium, no scatterer.
    Hexahedral mesh.
    """)

    t0  =  time()

    # Hexahedral mesh
    fens,fes  =  H8sphere(R,nref);
    fens1,fes1  =  mirrormesh(fens, fes, [-1.0, 0.0, 0.0], [0.0, 0.0, 0.0],
    renumb =  r(c) = c[[1, 4, 3, 2, 5, 8, 7, 6]]);
    fens,newfes1,fes2 =  mergemeshes(fens1, fes1, fens, fes, tolerance)
    fes = cat(newfes1,fes2)

    # Derive the finite element sets for the boundary
    bfes  =  meshboundary(fes)
    # Outer spherical boundary
    function dout(xyz)
        return xyz/norm(xyz)
    end

    louter = selectelem(fens, bfes, facing = true, direction = dout)
    outer_fes = subset(bfes, louter);

    fens,fes  = H8extrudeQ4(fens, outer_fes,
    nlayers, (xyz, layer)->(R+layer/nlayers*(Ro-R))*xyz/norm(xyz));
    connected = findunconnnodes(fens, fes);
    fens, new_numbering = compactnodes(fens, connected);
    fess = renumberconn!(fes, new_numbering);

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(FCplxFlt,size(fens.xyz,1),1))

    bfes  =  meshboundary(fes)

    # File  =   "Sphere.vtk"
    # vtkexportmesh(File, bfes.conn, geom.values, FinEtools.MeshExportModule.Q4)
    # @async run(`"paraview.exe" $File`)

    # File  =   "Sphere.vtk"
    # vtkexportmesh (File, fes.conn, fens.xyz, MeshExportModule.H8)
    # @async run(`"C:/Program Files (x86)/ParaView 4.2.0/bin/paraview.exe" $File`)

    linner = selectelem(fens, bfes, distance = R, from = [0.0 0.0 0.0],
    inflate = tolerance)
    louter = selectelem(fens, bfes, facing = true, direction = dout)

    println("Pre-processing time elapsed  =  ",time() - t0,"s")

    t1  =  time()

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, GaussRule(3, 2)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    abcfemm  =  FEMMAcoustSurf(IntegDomain(subset(bfes, louter), GaussRule(2, 2)), material)
    D  =  acousticABC(abcfemm, geom, P);

    # Inner sphere pressure loading
    function dipole(dpdn, xyz, J, label)
        n = cross(J[:,1],J[:,2]);
        n = vec(n/norm(n));
        #println("$( (-1.0im)*dot(vec(k),n)*exp(-1.0im*dot(vec(k),vec(xyz))) )")
        dpdn[1] = -rho*a_amplitude*n[1]
    end

    fi  =  ForceIntensity(FCplxFlt, 1, dipole);
    dipfemm  =  FEMMAcoustSurf(IntegDomain(subset(bfes, linner), GaussRule(2, 2)), material)
    F  = distribloads(dipfemm, geom, P, fi, 2);

    K = lu((1.0+0.0im)*(-omega^2*S + omega*1.0im*D + C)) # We fake a complex matrix here
    p = K\F  #

    scattersysvec!(P, p[:])

    println(" Minimum/maximum pressure= $(minimum(real(p)))/$(maximum(real(p)))")

    println("Computing time elapsed  =  ",time() - t1,"s")
    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "sphere_dipole_1.vtk"
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.H8;
    scalars = [( "realP", real(P.values))])
    @async run(`"paraview.exe" $File`)

    true

end # sphere_dipole_1

function sphere_inc_example()
    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega = 2000*phun("rev/s");      # frequency of the incident wave
    Ro = 300.0*phun("mm"); # radius of the enclosure
    nPerradius = 18;#number of elements on the radius of the sphere
    tolerance = Ro/(2^nPerradius)/100
    pincampl = 1.0*phun("kilo*Pa");
    k = omega/c*[1.0 0.0 0.0];

    println("""
    Spherical domain filled with acoustic medium, no scatterer.
    Incident wave.
    Hexahedral mesh.
    """)

    t0  =  time()

    # Hexahedral mesh
    fens,fes  =  H8spheren(Ro, nPerradius);
    fens1,fes1 = mirrormesh(fens, fes, [-1.0, 0.0, 0.0], [0.0, 0.0, 0.0]; renumb = (c ->  (c[[1, 4, 3, 2, 5, 8, 7, 6]])))
    fens,newfes1,fes2 =  mergemeshes(fens1, fes1, fens, fes, tolerance)
    fes = cat(newfes1,fes2)

    # Derive the finite element sets for the boundary
    bfes  =  meshboundary(fes)
    # Outer spherical boundary
    function dout(xyz)
        return xyz/norm(xyz)
    end

    louter = selectelem(fens, bfes, facing = true, direction = dout)
    outer_fes = subset(bfes, louter);

    # File  =   "Sphere.vtk"
    # vtkexportmesh (File, fes.conn, fens.xyz, MeshExportModule.H8)
    # @async run(`"C:/Program Files (x86)/ParaView 4.2.0/bin/paraview.exe" $File`)

    println("Pre-processing time elapsed  =  ",time() - t0,"s")

    t1  =  time()

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(FCplxFlt,size(fens.xyz,1),1))

    numberdofs!(P)

    pinc = deepcopy(P);

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, GaussRule(3, 2)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    abcfemm  =  FEMMAcoustSurf(IntegDomain(outer_fes, GaussRule(2, 2)), material)
    D  =  acousticABC(abcfemm, geom, P);

    pincf(loc)  =  pincampl*exp((-1.0im*(vec(k)'*vec(loc)))[1])

    # Incident pressure loading
    for j = 1:size(geom.values,1)
        pinc.values[j] = pincampl[1,1]*exp((-1.0im*(vec(k)'*vec(geom.values[j,:])))[1])
    end
    vpinc = gathersysvec(pinc)

    F  =  0.0 * vpinc;                # Start with a zero vector

    #pfemm = FEMMBase(inner_fes, GaussRule(order = 2,dim = 2))
    function dpincdn(dpdn, xyz, J, label)
        n = cross(J[:,1],J[:,2]);
        n = vec(n/norm(n));
        #println("$( (-1.0im)*dot(vec(k),n)*exp(-1.0im*dot(vec(k),vec(xyz))) )")
        dpdn[1] =  pincampl*(-1.0im)*dot(vec(k),n)*exp(-1.0im*dot(vec(k),vec(xyz)))
    end

    fi  =  ForceIntensity(FCplxFlt, 1, dpincdn);
    #F  =  F - distribloads(pfemm, nothing, geom, P, fi, 2);
    F  =  F + distribloads(abcfemm, geom, P, fi, 2);

    K = lu((1.0+0.0im)*(-omega^2*S + C)) # We fake a complex matrix here
    p = K\F  #+omega*1.0im*D

    scattersysvec!(P,p[:])

    println("Computing time elapsed  =  ",time() - t1,"s")
    println("Total time elapsed  =  ",time() - t0,"s")

    ptot = deepcopy(P)
    ptot.values = P.values+pinc.values

    # File  =   "Sphereptot.vtk"
    # vtkexportmesh (File, fes.conn, geom.values, MeshExportModule.H8; scalars = abs(ptot.values), scalars_name  = "ptot")
    # @async run(`"C:/Program Files (x86)/ParaView 4.2.0/bin/paraview.exe" $File`)

    File  =   "Spherepinc.vtk"
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.H8;
    scalars = [("realpinc", real(pinc.values))])
    @async run(`"paraview.exe" $File`)

    File  =   "SphereP.vtk"
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.H8;
    scalars = [( "realP", real(P.values))])
    @async run(`"paraview.exe" $File`)

    true

end # sphere_inc_example

function sphere_scatterer_example()
    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    R = 50.0*phun("mm");# radius of the piston
    Ro = 300.0*phun("mm"); # radius of the enclosure
    nPerradius=24;# number of elements along the radius of the scatterer
    nlayers=78;                     # number of layers of elements surrounding the piston
    tolerance = R/(nPerradius)/100
    kR =  1.0;
    k = kR/R*[1.0 0.0 0.0]; k = vec(k);
    omega = norm(k)*c;    # angular frequency of the incident wave
    pincampl = 1.0*phun("Pa");
    h = maximum([(Ro-R)/nlayers pi*Ro/(nPerradius)]);

    println("""

    Rigid sphere scattering.

    Frequency: $(omega/2/pi). Wave vector: $k. kR = $(norm(k)*R).
    Number of elements per wavelength: $(2*pi/(norm(k)*h)).

    Hexahedral mesh.
    """)

    t0  =  time()

    # Hexahedral mesh
    fens, fes  =  H8spheren(R, nPerradius);
    bfes  =  meshboundary(fes)
    l = selectelem(fens, bfes, facing = true, direction = [1.0 1.0  1.0], dotmin= 0.001)
    ex(xyz, layer) = (R+layer/nlayers*(Ro-R))*xyz/norm(xyz)
    fens,fes  =  H8extrudeQ4(fens, subset(bfes,l), nlayers, ex);
    fens1,fes1 = mirrormesh(fens, fes, [-1.0, 0.0, 0.0], [0.0, 0.0, 0.0]; renumb =  c ->  c[[1, 4, 3, 2, 5, 8, 7, 6]])
    fens,newfes1,fes2 =  mergemeshes(fens1, fes1, fens, fes, tolerance)
    fes = cat(newfes1,fes2)

    # Derive the finite element sets for the boundary
    bfes  =  meshboundary(fes)
    # Outer spherical boundary
    function dout(xyz)
        return xyz/norm(xyz)
    end

    louter = selectelem(fens, bfes, facing = true, direction = dout)
    outer_fes = subset(bfes,louter);

    # Inner spherical boundary
    function din(xyz)
        return -xyz/norm(xyz)
    end
    linner = selectelem(fens, bfes, facing = true, direction = din)
    inner_fes = subset(bfes,linner);

    # File  =   "Sphere.vtk"
    # vtkexportmesh(File, inner_fes.conn, fens.xyz, MeshExportModule.Q4)
    # @async run(`"paraview.exe" $File`)
    println("Pre-processing time elapsed  =  ",time() - t0,"s")

    t1  =  time()

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(FCplxFlt,size(fens.xyz,1),1))

    numberdofs!(P)

    pinc = deepcopy(P);

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, GaussRule(3, 3)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    abcfemm  =  FEMMAcoustSurf(IntegDomain(outer_fes, GaussRule(2, 3)), material)
    D  =  acousticABC(abcfemm, geom, P);

    # Incident pressure loading
    for j = 1:size(geom.values,1)
        xyz = vec(geom.values[j,:]);
        pinc.values[j] = pincampl*exp((-1.0im*(k'*xyz))[1])
    end
    vpinc = gathersysvec(pinc)

    F  =  - (-omega^2*S + C) * vpinc;

    #pfemm = FEMMBase(inner_fes, GaussRule(order = 2,dim = 2))
    function dpincdn(dpdn, xyz, J, label)
        xyz = vec(xyz);
        n = cross(J[:,1],J[:,2]);
        n = vec(n/norm(n));
        dpdn[1] = pincampl*(-1.0im)*dot(vec(k),n)*exp(-1.0im*dot(vec(k),vec(xyz)))
    end

    fi  =  ForceIntensity(FCplxFlt, 1, dpincdn);
    F  =  F + distribloads(abcfemm, geom, P, fi, 2);

    K = ((-omega^2*S + omega*1.0im*D + C))
    p = K\F

    scattersysvec!(P,p[:])

    println("Computing time elapsed  =  ",time() - t1,"s")
    println("Total time elapsed  =  ",time() - t0,"s")

    ptot = deepcopy(P)
    ptot.values = P.values+pinc.values

    # File  =   "Sphereptot.vtk"
    # vtkexportmesh (File, fes.conn, geom.values, MeshExportModule.H8; scalars = abs(ptot.values), scalars_name  = "ptot")
    # @async run(`"C:/Program Files (x86)/ParaView 4.2.0/bin/paraview.exe" $File`)

    # File  =   "Spherepinc.vtk"
    # vtkexportmesh (File, fes.conn, geom.values, MeshExportModule.H8; scalars = real(pinc.values), scalars_name  = "pinc")
    # @async run(`"C:/Program Files (x86)/ParaView 4.2.0/bin/paraview.exe" $File`)

    File  =   "SphereP.vtk"
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.H8; scalars = [("P", abs.(P.values))])
    @async run(`"paraview.exe" $File`)

    # File  =   "SphereP.vtk"
    # vtkexportmesh (File, cat(outer_fes,inner_fes).conn, geom.values, MeshExportModule.Q4; scalars = abs(P.values), scalars_name  = "P")
    # @async run(`"C:/Program Files (x86)/ParaView 4.2.0/bin/paraview.exe" $File`)

    # using Winston
    # pl  =  FramedPlot(title = "Matrix",xlabel = "x",ylabel = "Re P, Im P")
    # setattr(pl.frame, draw_grid = true)
    # add(pl, Curve([1:length(C[:])],vec(C[:]), color = "blue"))

    # # pl = plot(geom.values[nLx,1][ix],scalars[nLx][ix])
    # # xlabel("x")
    # # ylabel("Pressure")
    # display(pl)

    true

end # sphere_scatterer_example

function allrun()
    println("#####################################################")
    println("# sphere_dipole_1 ")
    sphere_dipole_1()
    println("#####################################################")
    println("# sphere_inc_example ")
    sphere_inc_example()
    println("#####################################################")
    println("# sphere_scatterer_example ")
    sphere_scatterer_example()
end # function allrun

end # module sphere_examples


function intersectingvertices(S1, S2; tol::Float64 = 1/10^10)
    # Dimension
    n = size(S1, 1)
    # Centroid and radii
	 c1, c2 = Circumsphere(S1)[2:n+1], Circumsphere(S2)[2:n+1]
  	 r1, r2 = Circumsphere(S1)[1], Circumsphere(S2)[1]

  	 # Orientation of simplices
  	 orientation_S1 = det([ones(1, n + 1); S1])
  	 orientation_S2 = det([ones(1, n + 1); S2])

  	 if abs(orientation_S1) < tol || abs(orientation_S2) < tol
        return Array{Float64, 2}(undef, 0, 0)
  	 end

  	# Set volume to zero initially. Change only if there is intersection
  	 IntVol = 0.0

    # -------------------------------------
    # Simplices intersect in some way
    # -------------------------------------

    # If the (distance between centroids)^2-(sum of radii)^2 < 0,
    # then the simplices intersect in some way.
    dist_difference::Float64 = (transpose(c1 - c2) * (c1 - c2) - (r1 + r2)^2)[1]

    if dist_difference < 0
        #@show "The simplices intersect in some way (possibly only along a boundary, with zero volume intersection)"
        # Find the number of points of each simplex contained within the
        # circumsphere of the other simplex
        vertices1InCircum2 = SomeVertexInCircumsphere(S1, r2, c2)
        vertices2InCircum1 = SomeVertexInCircumsphere(S2, r1, c1)

        if vertices1InCircum2 + vertices2InCircum1 >= 1
            βs1in2, βs2in1, ordered_vertices1, ordered_vertices2, numof1in2, numof2in1 =
            BarycentricCoordinates(S1,S2,orientation_S1,orientation_S2,tol)
            # Trivial intersections
            TriviallyContained = heaviside0([numof1in2 numof2in1] .- (n+1))
            IsSomeContained = sum(TriviallyContained, dims=2)[1]

            if IsSomeContained == 2.0 # The simplices coincide
                #@show "The simplices coincide"
                return copy(transpose(S1))
            elseif IsSomeContained == 1.0 # One simplex is contained in the other
                if TriviallyContained[1] == 1.0 # Simplex1 is contained in Simplex2
                    #@show "Simplex 1 is contained in simplex 2"
                    return copy(transpose(S1))
                else # Simplex2 is contained in Simplex1
                    #@show "Simplex 2 is contained in simplex 1"
                    return copy(transpose(S2))
                end
            else # No simplex contains the other
                #@show "No simplex contains the other"
                Ncomm, ordered_vertices1, ordered_vertices2 = SharedVertices(βs1in2,ordered_vertices1,ordered_vertices2,numof1in2,numof2in1)

                # Is there any shared face?
                if Ncomm == n
                    #@show "The simplices share a face"
                    IntVol = SharedFaceVolume(S2, βs1in2, ordered_vertices1, ordered_vertices2)
                    return transpose(SharedFaceVertices(S2, βs1in2, ordered_vertices1, ordered_vertices2))
                else # The simplices do not share a face.
                    #@show "The simplices do not share a face"
                    IntVert, ConvexExpIntVert = IntersectionOfBoundaries_NoStorage(S1,S2,βs1in2,βs2in1, ordered_vertices1, ordered_vertices2, numof1in2, numof2in1, Ncomm, tol)

                    if !isempty(IntVert)
                        IntVert,ConvexExpIntVert = PolytopeGeneratingVertices(S1,S2,IntVert,ConvexExpIntVert,βs1in2,βs2in1,ordered_vertices1,ordered_vertices2,numof1in2,numof2in1,Ncomm);
                        IntVol = VolumeComputation(IntVert, ConvexExpIntVert)

                        return IntVert

                    else
                        return Array{Float64, 2}(undef, 0, 0)
                    end
                end
            end
        else
            IntVert = Array{Float64, 2}(undef, 0, 0)
        end
    else
        IntVert = Array{Float64, 2}(undef, 0, 0)
    end

    return IntVert
end


#compute the output of the neural net for input x0
#requires that the functions nnAct() and nnDAct() be defined,
#representing the activation function and its derivative, respectively
function nnForward(nn::neuralNet,x0::Array{Float64},nnAct::Function=nnAct,nnDAct::Function=nnDAct)
    #define inputs
    nn.x[nn.xIdx[1]] = x0;
    #loop over levels
    for l=1:nn.nLevels-1
        Wt = reshape(nn.w[nn.wIdx[l]],nn.nNeurons[l+1],nn.nNeurons[l]); #assumes that we've stored W' (row-major)
        z  = Wt*nn.x[nn.xIdx[l]] + nn.b[nn.bIdx[l]];
        nn.x[nn.xIdx[l+1]] = nnAct(z);
        nn.da[nn.bIdx[l]] = nnDAct(z); #save derivatives for later
    end

    return nn.x[nn.xIdx[nn.nLevels]]; #return NN output in case that's helpful
end


################################
###    Summary statistics    ###
################################
"""
	poissonFixation(observedValues,λds, λdn)

Divergence sampling from Poisson distribution. The expected neutral and selected fixations are subset through their relative expected rates ([`fixNeut`](@ref), [`fixNegB`](@ref), [`fixPosSim`](@ref)). Empirical values are used are used to simulate the locus *L* along a branch of time *T* from which the expected *Ds* and *Dn* raw count estimated given the mutation rate (``\\mu``). Random number generation is used to subset samples arbitrarily given the success rate ``\\lambda`` in the distribution.

```math
\\mathbb{E}[D_N] = X \\in Poisson\\left(\\lambda = D \\times \\left[\\frac{\\mathbb{E}[D_+] + \\mathbb{E}[D_-]}{\\mathbb{E}[D_+] + \\mathbb{E}[D_-] + \\mathbb{E}[D_0]}\\right]\\right)
```
```math
\\mathbb{E}[D_S] = X \\in Poisson\\left(\\lambda = D \\times \\left[\\frac{\\mathbb{E}[D_0]}{\\mathbb{E}[D_+] + \\mathbb{E}[D_-] + \\mathbb{E}[D_0]}\\right]\\right)
```
# Arguments
 - `observedValues::Array`: Array containing the total observed divergence.
 - ` λds::Float64`: expected neutral fixations rate.
 - ` λdn::Float64`: expected selected fixations rate.
# Returns
 - `Array{Int64,1}` containing the expected count of neutral and selected fixations.

"""
function poissonFixation(;observedValues::Array, λds::Array, λdn::Array,λweak::Array,λstrong::Array)

	ds = @. λds / (λds + λdn) * observedValues
	dn = @. λdn / (λds + λdn) * observedValues
	dweak = @. λweak / (λds + λdn) * observedValues
	dstrong = @. λstrong / (λds + λdn) * observedValues

	sampledDs     = pois_rand.(ds)
	sampledDn     = pois_rand.(dn)
	sampledWeak   = pois_rand.(dweak)
	sampledStrong = pois_rand.(dstrong)

	alphas = @. [sampledWeak/sampledDn sampledStrong/sampledDn (sampledWeak+sampledStrong)/sampledDn]

	out = alphas,sampledDn, sampledDs
	return out
end

"""
	poissonPolymorphism(observedValues,λps,λpn)

Polymorphism sampling from Poisson distributions. The total expected neutral and selected polimorphism are subset through the relative expected rates at the frequency spectrum ([`fixNeut`](@ref), [`DiscSFSNeutDown`](@ref),). Empirical sfs are used to simulate the locus *L* along a branch of time *T* from which the expected *Ps* and *Pn* raw count are estimated given the mutation rate (``\\mu``). Random number generation is used to subset samples arbitrarily from the whole sfs given each frequency success rate ``\\lambda`` in the distribution.

The success rate managing the Poisson distribution by the observed count each frequency.  We considered both sampling variance and process variance is affecting the number of variable alleles we sample from SFS. This variance arises from the random mutation-fixation process along the branch. To incorporate this variance we do one sample per frequency-bin and use the total sampled variation and the SFS for the summary statistics.

```math
\\mathbb{E}[P_N] = \\sum_{x=0}^{x=1} X \\in Poisson\\left(\\lambda = SFS_{(x)} \\times \\left[\\frac{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}]}{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}] + \\mathbb{E}[P_{0(x)}]}\\right]\\right)
```

```math
\\mathbb{E}[P_S] = \\sum_{x=0}^{x=1} X \\in Poisson\\left(\\lambda = SFS_{(x)} \\times \\left[\\frac{\\mathbb{E}[P_{0(x)}]}{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}] + \\mathbb{E}[P_{0(x)}]}\\right]\\right)
```

# Arguments
 - `observedValues::Array{Int64,1}`: Array containing the total observed divergence.
 - ` λps::Array{Float64,1} `: expected neutral site frequency spectrum rate.
 - ` λpn::Array{Float64,1} `: expected selected site frequency spectrum rate.
# Returns
 - `Array{Int64,2}` containing the expected total count of neutral and selected polymorphism.

"""
function poissonPolymorphism(;observedValues::Array, λps::Array, λpn::Array)

	# Neutral λ;
	λ1 = @. λps / (λps + λpn) * observedValues
	# Selected λ;
	λ2 = @. λpn / (λps + λpn) * observedValues

	# Replace negative and 0 rates values. Strange behaviour depending on θ and Γ values
	# Relative rates output NaN values due to 0 divisons.

	replace!(λ1,NaN=>1)
	replace!(λ2,NaN=>1)

	sampledPs = pois_rand.(λ1)
	sampledPn = pois_rand.(λ2)

	return (sampledPn, sampledPs)
end

"""
	sampledAlpha(observedValues,λds, λdn)

Ouput the expected values from the Poisson sampling process. Please check [`poissonFixation`](@ref) and [`poissonPolymorphism`](@ref) to understand the samplingn process. α(x) is estimated through the expected values of Dn, Ds, Pn and Ps.

# Arguments
 - `param::parameters`: Array containing the total observed divergence.
 - `d::Array`: observed divergence.
 - `afs::Array`: observed polymorphism.
 - ` λdiv::Array{Float64,2}`: expected fixations rate.
 - ` λdiv::Array{Float64,2}`: expected site frequency spectrum rates.
# Returns
αS,expDn,expDs,expPn,expPs,ssAlpha
 - `Array{Int64,2}` containing α(x) values.
 - `Array{Int64,1}` expected non-synonymous divergence.
 - `Array{Int64,1}` expected synonymous divergence.
 - `Array{Int64,1}` expected non-synonymous polymorphism.
 - `Array{Int64,1}` expected synonymous polymorphism.
 - `Array{Int64,1}` expected synonymous polymorphism.
 - `Array{Int64,1}` expected synonymous polymorphism.
 - `Array{Int64,2}` containing α(x) binned values.

"""
function sampledAlpha(;d::Array,afs::Array,observed::Matrix,λdiv::Array,λpol::Array,dac::Array)

	## Outputs
	alphas, expDn, expDs = poissonFixation(observedValues=d,λds=λdiv[1],λdn=λdiv[2],λweak=λdiv[3],λstrong=λdiv[4])
	expPn, expPs         = poissonPolymorphism(observedValues=afs,λps=λpol[1],λpn=λpol[2])

	θ_pn = @. observed[1,1] / (expPn[1,:]);
	θ_ps = @. observed[1,2] / (expPs[1,:]);

	r_pn = ones(size(expPn));
	r_ps = ones(size(expPs));
	for j = 2:size(dac,1)
		r_pn[j,:] .= @. (observed[j,1]) / (expPn[j,:]);
		r_pn[j,:] = @. r_pn[j,:] / θ_pn;

		r_ps[j,:] .= @. (observed[j,2]) / (expPs[j,:]);
		r_ps[j,:] .= @. r_ps[j,:] / θ_ps;

	end

	expPn_rj = zeros(size(expPn))
	expPs_rj = zeros(size(expPs))

	expPn_rj[1,:] = @. expPn[1,:] * θ_pn;
	expPs_rj[1,:] = @. expPs[1,:] * θ_ps;
	for j = 2:size(dac,1)
		expPn_rj[j,:] = @. expPn[j,:] * θ_pn * r_pn[j,:];
		expPs_rj[j,:] = @. expPs[j,:] * θ_ps * r_ps[j,:];
	end

	## Alpha from expected values. Used as summary statistics
	αS = @. round(1 - ((expDs/expDn) * (expPn/expPs)'),digits=5)
	αS_v2 = @. round(1 - ((expDs/expDn) * (expPn_rj/expPs_rj)'),digits=5)

	return αS,alphas,expDn,expDs,expPn,expPs
end

"""
	samplingFromRates(gammaL,gammaH,pposL,pposH,observedData,nopos)
"""
function samplingFromRates(m::Array,s::Array,d::Array,nt::Array,sl::Array,x::Array,observed::Matrix,dac::Array)

	ds             = x[:,1]
	dn             = x[:,2]
	dweak          = x[:,3]
	dstrong        = x[:,4]
	gn             = abs.(m[:,4])
	sh             = round.(m[:,end-1],digits=5)

	alxSummStat, alphasDiv, expectedDn, expectedDs, expectedPn, expectedPs = sampledAlpha(d=d,afs=s,observed=observed,λdiv=[ds,dn,dweak,dstrong],λpol=[permutedims(nt),permutedims(sl)],dac=dac)
	expectedValues = hcat(round.(alphasDiv,digits=5),gn,sh,alxSummStat)

	return(expectedValues)
end

"""
	summaryStatsFromRates(param::parameters,rates::JLD2.JLDFile,analysisFolder::String,summstatSize::Int64,replicas::Int64,bootstrap::Bool)

Estimate summary statistics using observed data and analytical rates. *analysisFolder* will check for the SFS and divergence file and will be used to output summary statistics

# Arguments
 - `param::parameters` : Mutable structure containing the models
 - `rates::JLD2.JLDFile` : HDF5 containing solved models, fixation and polymorphic rates
 - `analysisFolder::String` : Folder containing the SFS and divergence files. It will be used to output the observed data and summary estatistics.
 - `summstatSize::Int64` : Number of summary statistics
 - `replicas::Int64` : Number of bootstrap replicas
 - `bootstrap::Bool` : Boolean to perform or not bootstrapping
# Output
 - Obserded data and summary statistics to ABC inference
"""
function summaryStatsFromRates(;param::parameters,h5file::JLD2.JLDFile,analysisFolder::String,summstatSize::Int64,replicas::Int64,bootstrap::Bool)

	#Opening files
	sFile   = filter(x -> occursin("sfs",x), readdir(analysisFolder,join=true));
	dFile   = filter(x -> occursin("div",x), readdir(analysisFolder,join=true));

	sfs,d,α,observed = openSfsDiv(sFile,dFile,param.dac,replicas,bootstrap);

	#Open rates
	tmp    = h5file[string(param.N) * "/" * string(param.n)]

	#Subset index
	idx    = sample(1:size(tmp["models"],1),summstatSize,replace=false)
	models = Array(view(tmp["models"],idx,:));
	dsdn   = Array(view(tmp["dsdn"],idx,:));

	# Filtering polymorphic rate by dac
	n    = hcat(map(x -> view(tmp["neut"][x],:),param.dac)...);
	s    = hcat(map(x -> view(tmp["sel"][x],:),param.dac)...);
	neut = Array(view(n,idx,:));
	sel  = Array(view(s,idx,:));

	#Making summaries
	f(x,y,o,m=models,nt=neut,sl=sel,d=dsdn,dac=param.dac) = samplingFromRates(m,x,y,nt,sl,d,o,dac);
	expectedValues = f.(sfs,d,observed);

	w(x,name) = CSV.write(name,DataFrame(x,:auto),delim='\t',header=false);

	# Controling outlier cases
	fltInf(e) = replace!(e, -Inf=>NaN)
	expectedValues = fltInf.(expectedValues)
	fltNan(e) = e[vec(.!any(isnan.(e),dims=2)),:]
	expectedValues = fltInf.(expectedValues)
	expectedValues = fltNan.(expectedValues)
	expectedValues = map(x -> x[(x[:,3] .> 0 ) .& (x[:,3] .<1 ),:],expectedValues)

	# Writting ABCreg input
	#=w(vcat(α...), analysisFolder * "/alphas.txt");
	w(expectedValues, analysisFolder * "/summstat.txt");=#
	progress_map(w,α, analysisFolder * "/alphas_" .* string.(1:size(sFile,1)) .* ".txt";progress= Progress(size(sFile,1),desc="Writting α "));
	progress_pmap(w,expectedValues,  analysisFolder * "/summstat_" .* string.(1:size(sFile,1)) .* ".txt";progress= Progress(size(sFile,1),desc="Writting summaries "));

	return(expectedValues)
end


# Load packages

using Polynomials: roots, Polynomial # For finding roots
using StatsBase: fit, Histogram # For binning the roots
using ColorSchemes: viridis
using Images: save # For saving the histogram to an image

# Create every combination of -1 and 1

polynomials = Iterators.product(repeat([[1, -1]], 2)...) |> collect |> vec .|> Polynomial;

# Solve the polynomials

complex_roots =
    polynomials .|>
    roots |>
    Iterators.flatten |>
    collect |>
    (Z -> filter(z -> imag(z) != 0, Z));

# Bin the roots

heatmap = fit(Histogram, (imag.(complex_roots), real.(complex_roots)), nbins = 2000).weights;

# Normalise

heatmap_normalised = log.(1 .+ heatmap);

# Apply a colour scheme

heatmap_coloured = get(viridis, heatmap_normalised, :extrema);

# Save the histogram to a file

save(joinpath("images", "littlewood.png"), heatmap_coloured);

# ![The Littlewood fractal](images/littlewood.png)


@testset "dffunctions_test.jl" begin
    # test reading out a csv file into a DataFrame, then interpolate it to produce a function
    spG = LaserLab.CsvG(';',',')  # spanish: decimals represented with ',' delimited with ';'
    enG = LaserLab.CsvG(',','.')  # english
    df = LaserLab.load_df_from_csv("/Users/jj/JuliaProjects/LaserLab/data/objectives", "LMU-40X-NUV.csv", enG)
    wl=200.0:20:800.0
    fdf =LaserLab.dftof(wl, df, "Transmission(%)")
    @test all([fdf(df[1,"Wavelength(µm)"]) == df[1,"Transmission(%)"] for i in 1:31])

    green_area = LaserLab.qpdf(fdf, 400.0, 600.0)
    red_area = LaserLab.qpdf(fdf, 600.0, 800.0)
    @test isapprox(green_area/red_area, 1.88, rtol=0.01)

    total_area = LaserLab.qpdf(fdf, 200.0, 800.0)
    N, dfpdf = LaserLab.ftopdf(wl, fdf)
    @test total_area ≈ N
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

"""
    cirrange(number[, max]) -> restricted_number

### Purpose ###

Force a number into a given range \$[0, \\text{max})\$.

### Argument ###

* `number`: the number to modify.  Can be a scalar or an array.
* `max` (optional numerical argument): specify the extremum of the range \$[0, \\text{max})\$
 into which the number should be restricted.  If omitted, defaults to `360.0`.

### Output ###

The converted number or array of numbers, as `AbstractFloat`.

### Example ###

Restrict an array of numbers in the range \$[0, 2\\pi)\$ as if they are angles
expressed in radians:

``` julia
cirrange([4pi, 10, -5.23], 2.0*pi)
# => 3-element Array{Float64,1}:
#     0.0
#     3.71681
#     1.05319
```

### Notes ###

This function does not support the `radians` keyword like IDL implementation.
Use `2.0*pi` as second argument to restrict a number to the same interval.
"""
cirrange(number::Real, max::Real=360.0) =
    mod(promote(float(number), float(max))...)

function cirrange{N<:Real}(numbers::AbstractArray{N}, max::Real=360.0)
    result = similar(numbers, typeof(float(one(N))))
    for i in eachindex(numbers)
        result[i] = cirrange(numbers[i], max)
    end
    return result
end


@doc doc"""
compute_ps(Z::Array{Int, 2})

Estimate of Posterior Similarity Matrix

Assuming Z to be [Trails x Observations]

""" ->
function compute_psm(Z::Array{Int, 2})

	(N, M) = size(Z)

	valuerange = collect(1:M)
	inRange = true

	for n in 1:N
		for m in 1:M
				if !(Z[n, m] in valuerange)
					inRange = false
				end
		end
	end

	@assert inRange "Values are not in range: All values must be larger equal 1 and smaller equal "

	psm = zeros(N, N)

	for m in 1:M
		for i in 1:N
			for j in 1:N
				if Z[i, m] == Z[j, m]
					psm[i, j] += 1
				end
			end
		end
	end

	return psm ./ M

end