module Pade
import SpecialFunctions
## Polynomial type
"""
Polynomial type
"""
struct Polynomial #TODO: make type parametric use eltype(f)
 coeffs
 #constructor
    function Polynomial(coeffs)
        N=length(coeffs)
        n=0
        for i=N:-1:1
         if coeffs[i]==0
             n+=1
         else
             break
         end
        end
        return new(coeffs[1:(N-n)])
    end
end

#make Polynomial callable with Horner-scheme evaluation
function (p::Polynomial)(z,k::Int=0)
    p=derive(p,k)
    f=0.0+0.0im #typing
    for i = length(p.coeffs):-1:1
       f*=z
       f+=p.coeffs[i]
    end
    return f
end

import Base.show
import Base.string

function string(p::Polynomial)
    N=length(p.coeffs)
    if N==0
        txt="0"
    else
        if p.coeffs[1]!=0
            txt="$(p.coeffs[1])"
        else
            txt=""
        end
        for n=2:N
            coeff=string(p.coeffs[n])
            if coeff==0
                continue
            end
            if occursin("+",coeff) || occursin("-",coeff)
                txt*="+($coeff)*z^$(n-1)"
            else
                txt*="+$coeff*z^$(n-1)"
            end
        end
    end
    return txt
end

function show(io::IO,p::Polynomial)
    print(io,string(p))
end

import Base.+
import Base.-
import Base.*
import Base.^

function +(a::Polynomial,b::Polynomial)
    if length(a.coeffs)>length(b.coeffs)
        a,b=b,a
    end
    c=copy(b.coeffs)
    for (idx,val) in enumerate(a.coeffs)
        c[idx]+=val
    end
    return Polynomial(c)
end

function -(a::Polynomial,b::Polynomial)
    return +(a,-1*b)
end

function *(a::Polynomial,b::Polynomial)
    I=length(a.coeffs)
    J=length(b.coeffs)
    K=I+J-1
    c=zeros(ComplexF64,K) #TODO: sanity check for type of b
    idx=1
    for k =1:K
        for i =1:k
            #TODO: figure out when to break the loop
            j=k-i+1
            if i>I ||j>J
                continue
            end
            c[k]+=a.coeffs[i]*b.coeffs[j]
        end
    end
    return Polynomial(c)
end

#scalar multiplication
function *(a::Polynomial,b::Number)
    return a*Polynomial([b])
end
function *(a::Number,b::Polynomial)
    return Polynomial([a])*b
end

function ^(p::Polynomial, k::Int)
    b=Polynomial(copy(p.coeffs))
    for i=1:k-1
        b*=p
    end
    return b
end

function prodrange(start,stop)
    return SpecialFunctions.factorial(stop)/SpecialFunctions.factorial(start)
end

"""
    b=derive(a::Polynomial,k::Int=1)

Compute `k`th derivative of Polynomial `a`.
"""
#TODO: Devlop bigfloat prodrange
# function derive(a::Polynomial,k::Int=1)
#
#     N=length(a.coeffs)
#     if k>=N
#         return Polynomial([]) #TODO: type
#     end
#     d=zeros(typeof(a.coeffs[1]),N-k)
#     for i=1:N-k
#         d[i]=a.coeffs[i+k]*prodrange(i-1,i+k-1)
#     end
#     return Polynomial(d)
# end

"""
    g=shift(f,Δz)

Shift Polynomial with respect to its argument. (Taylor shift)
"""
function shift(f::Polynomial,Δz)
    g=Array{eltype(f)}(undef,length(f.coeffs))
    for n=0:length(f.coeffs)-1
        g[n+1]=f(Δz)/SpecialFunctions.factorial(n*1.0)
        f=Polynomial(f.coeffs[2:end])
        f.coeffs.*=1:length(f.coeffs)
    end
    return Polynomial(g)
end

function scale(p::Polynomial,a::Number)
        p=p.coeffs
        s=1
        for idx =1: length(p)
            p[idx]*=s
            s*=a
        end
        return Polynomial(p)
end

#TODO: write macro nto create !-versions

# tests
# a=Polynomial([0.0,1.0])
# b=a*a
# c=derive(a)
# d=a+b+c
# e=d^3
# f=shift(e,2)
# e(2)-f(0)
# rational Polynomial
## Rational Polynomials
struct Rational_Polynomial
    num::Polynomial
    den::Polynomial

    #default constructor
    function Rational_Polynomial(num::Polynomial,den::Polynomial)
        num=copy(num.coeffs)
        den=copy(den.coeffs)
        if isempty(num) || num==[0]
            den=[1]
        elseif den[1]!=0
            den./=den[1]
            num./=den[1]
        end
        #TODO: else case
        return new(Polynomial(num),Polynomial(den))
    end
end
function (r::Rational_Polynomial)(z,k::Int=0)
    for i=1:k
        r=derive(r)
    end
    return r.num(z)/r.den(z)
end

function derive(a::Rational_Polynomial)
    return Rational_Polynomial(derive(a.num)*a.den-a.num*derive(a.den),a.den^2)
end
##

"""
    derive(a::Any, k:Int)

Take the `k`th derivative of `a` by calling the derive function multiple times.

# Notes
This is a fallback implementation.
"""
function derive(a::Any,k::Int)
    for i=1:k
        a=derive(a)
    end
    return a
end
function derive(p::Polynomial)
    N=length(p.coeffs)
    return Polynomial([i*p.coeffs[i+1] for i=1:N-1] )
end

function pade(p::Polynomial,L,M) #TODO: harmonize symbols for eigenvalues, modes etc
  p=p.coeffs
  A=zeros(ComplexF64,M,M)
  for i=1:M
    for j=1:M
      if L+i-j>=0
        A[i,j]=p[L+i-j+1] #+1 is for 1-based indexing
      end
    end
  end
  b=A\(-p[L+2:L+M+1])
  b=[1;b]
  a=zeros(ComplexF64,L+1)
  for l=0:L
    for m=0:l
      if m<=M
        a[l+1]+=p[l-m+1]*b[m+1] #+1 is for zero-based indexing
      end
    end
  end
  a,b=Polynomial(a),Polynomial(b)
  return Rational_Polynomial(a,b)
end
end# module Pade


### utilities.jl

#
# Strings
#
const rchar = "rml"
const rmchar = "rmean"
const rqchar = "rquant"
const wchar = "waic"
const routchar = "rout"
const runchar = "run"
const chichar = "chi"
const chioutchar = "chiout"
const rlastchar = "rlast"
const onchar = "LCON"
const offchar = "LCOFF"
const fishchar = "FISH"
const txtstr = ".txt"

"""
eig_decompose(M)

Take matrix M and return values and vectors
"""
function eig_decompose(M)
    Meiv = eigen(M)
    return Meiv.values, Meiv.vectors
end

"""
nonzero_rows(T)

Returns an array of row indices that have at least one nonzero element for matrix T
"""
function nonzero_rows(T)
    n = size(T)[2]
    nonzeros = Array{Int}(undef,0)
    for  a in 1:size(T)[1]
        if T[a,:] != zeros(n)
            push!(nonzeros,a)
        end
    end
    return nonzeros
end

"""
normalize_histogram(hist)

Returns normalized histogram hist
"""
normalize_histogram(hist) = hist/sum(hist)

"""
var_update(vartuple, newValue)

vartuple = count, mean, M2
iteration step for online algorithm to compute
mean and variance
returns count, mean, M2
"""
function var_update(vartuple, newValue)
    count, mean, M2 = vartuple
    count += 1
    delta = newValue - mean
    mean .+= delta / count
    delta2 = newValue - mean
    M2 .+= delta .* delta2
    return count, mean, M2
end
"""
cov_update(covtuple,newValue1,newValue2)

update covtuple for online algorithm to compute
    covariance
"""
function cov_update(covtuple,newValue1,newValue2)
    count, meanx, meany, C = covtuple
    count += 1
    dx = newx - meanx
    meanx .+= dx / count
    meany .+= (newy - meany) / count
    C .+= dx * (newy - meany)
    return count, meanx, meany, C
end
"""
finalize(vartuple)

Retrieve the mean, variance and sample variance from an aggregate
"""
function finalize(vartuple)
    count, mean, M2 = vartuple
    if count < 2
        return mean, M2
    else
        return mean, M2 ./ (count - 1)
    end
end

"""
truncatehistogram(x::Array{Int,1},yield::Float64,nhistmax::Int)

only keep columns of histogram x that accounts for yield fraction of total sum
"""
function truncate_histogram(x::Array,yield::Float64,nhistmax::Int)
    counts = sum(x)
    if counts > 0
        nhist = 1
        ratio = 0.
        # Only keep yield of all data to limit support of histogram
        while ratio < yield && nhist <= nhistmax
            ratio = sum(x[1:nhist])/counts
            nhist += 1
        end
        return x[1:min(nhist,length(x))]
    else
        return 0
    end
end

"""
combine_histogram(x1::Array,x2::Array)

add histograms x1 and x2

"""
function combine_histogram(x1::Array,x2::Array)
    n = min(length(x1),length(x2))
    x1[1:n] + x2[1:n]
end

"""
pooled_mean(means,counts)

compute weighted average of means given count totals
"""
pooled_mean(means,counts) = counts'*means/sum(counts)

"""
pooled_variance(vars,counts)

computed weighted average of variances given counts
"""
pooled_variance(vars,counts) = (counts .- 1)'*vars/sum(counts .- 1)

"""
pooled_std(std,counts)

compute weighted average of standard deviations given counts
"""
pooled_std(std,counts) = sqrt(pooled_variance(std.^2,counts))

"""
var_ratio(mua,mub,vara,varb,cov)

Compute ratio of variances of two variables given means, variances, and covariance of variables
"""
var_ratio(mua,mub,vara,varb,cov) = mua^2/mub^2*(vara/mua^2 - 2*cov/(mua*mub) + varb/mub^2)

"""
decimal(x::Vector)

convert digits of base 3 number to base 10
"""
function decimal(x::Vector,base::Int=3)
    nr = length(x)
    x' * base .^ collect(0:nr-1)
end
"""
findbase(l,n,nr)

Find the number of G states for transition rate matrix of size l
"""
function findbase(l,n,nr)
    if l == (n+1)*3^nr
        return 3
    elseif l == (n+1)*2^nr
        return 2
    else
        throw("wrong length")
    end
end

"""
KL(x,y)

Kullback-Leibler distance
"""
KL(x,y) = sum(x .* (log.(max.(x,eps(Float64))) - log.(max.(y,eps(Float64)))))

"""
trim(h::Array,nh::Array)

"""
function trim(h::Array,nh::Array)
    for i in eachindex(h)
        h[i] = h[i][1:nh[i]]
    end
    return h
end

"""
LogNormal_array(param,cv)
Prior distribution arrays
"""
LogNormal_array(param,cv) = distributionarray(param,cv,LogNormal)

"""
Gamma_array(param,cv)

"""
Gamma_array(param,cv) = distributionarray(param,cv,Gamma)

"""
function LogNormalBeta_array(param,cv,ind)

"""
function LogNormalBeta_array(param,cv,ind)
    if ind == 1
        d = [setBeta(param[ind],cv[ind])]
    else
        barind = 1:ind-1
        d = LogNormalarray(param[barind],cv[barind])
        push!(d,setBeta(param[ind],cv[ind]))
    end
    return d
end

"""
distribution_array(param,cv,dist)

fill an array with dist(param,cv)
"""
function distributionBeta_array(param::Vector,cv::Vector,ind::Int,dist=LogNormal)
    if ind == 1
        d = [Beta_meancv(param[ind],cv[ind])]
    else
        barind = 1:ind-1
        d = distribution_array(param[barind],cv[barind],dist)
        push!(d,Beta_meancv(param[ind],cv[ind]))
    end
    return d
end

"""
distribution_array(param::Vector,cv,dist=LogNormal)

"""
function distribution_array(param::Vector,cv,dist=LogNormal)
    d = []
    for i in eachindex(param)
        if dist == LogNormal
            push!(d,LogNormal_meancv(param[i],cv[i]))
        elseif dist == Gamma
            push!(d,Gamma_meancv(param[i],cv[i]))
        else
            push!(d,dist(param[i],cv[i]))
        end

    end
    return d
end

"""
Gamma_meancv(mean,cv)
LogNormal_meancv(mean,cv)
Beta_meancv(mean,cv)

Reparameterized distributions to mean and coefficient of variation arguments
"""
Gamma_meancv(mean,cv) = Gamma(1/cv^2,mean * cv^2)
LogNormal_meancv(mean,cv) = LogNormal(mulognormal(mean,cv), sigmalognormal(cv))
function Beta_meancv(m,cv)
    cv2 = cv^2
    fac = (1-m)/cv2/m - 1
    if fac <= 0.
        fac = 2/m
    end
    alpha = m*fac
    beta = (1-m)*fac
    Beta(alpha,beta)
end
"""
sigmalognormal(cv)
mulognormal(mean,cv)

compute LogNormal mu and sigma parameters
given desired mean and coefficient of variation
"""
sigmalognormal(cv) = sqrt.(log.(1 .+ cv .^ 2))

mulognormal(mean,cv) = log.(mean) - .5*log.(1 .+ cv .^ 2)

"""
Parameter information for GR models
"""
function num_rates(G,R)
    n = G - 1
    2*n + 2*R + 2
end
function Grange(G)
    n = G - 1
    1 : 2*n
end
function initiation(G)
    n = G - 1
    2*n + 1
end
function Rrange(G,R)
    n = G - 1
    2*n + 2 : 2*n + R + 1
end
function Srange(G,R)
    n = G - 1
    2*n + R + 2 : 2*n + 2*R + 1
end

"""
mean_histogram(x)

"""
mean_histogram(x::Vector{Float64}) = collect(0:length(x)-1)' * x/sum(x)

function mean_histogram(x::Vector{Array})
    y = Vector{Float64}(undef,length(x))
    for i in eachindex(x)
        y[i] = mean_histogram(x[i])
    end
    return y
end

"""
    Difference_Zscore(x1,x2,sig1,sig2) = (x1 .- x2)./ sqrt(sig1.^2 + sig2.^2)

"""
Difference_Zscore(x1,x2,sig1,sig2) = (x1 - x2) / sqrt(sig1^2 + sig2^2)

"""
m2_histogram(x)

"""
m2_histogram(x) = (collect(0:length(x)-1).^2)' * x/sum(x)

"""
var_histogram(x)

"""
function var_histogram(x)
    m2_histogram(x) - mean_histogram(x)^2
end

function moment_histogram(x,n)
    y = collect(1:length(x)) .- 1
    v = (y .- mean_histogram(x)).^n
    v' * x/sum(x)
end

function factorial_moment(h::Vector,n)
    m = 0
    for i in n:length(h)
        a = i-1
        for j in 1:n-1
            a *= i-j-1
        end
        m += h[i]*a
    end
    return m / sum(h)
end

function moment_param_estimates(h)
    e1 = factorial_moment(h,1)
    e2 = factorial_moment(h,2)
    e3 = factorial_moment(h,3)
    r1 = e1
    r2 = e2/e1
    r3 = e3/e2
    println(r1,", ",r2,", ",r3)

    lambda = 2*r1*(r3-r2)/(r1*r2-2*r1*r3+r2*r3)
    mu = 2*(r2-r1)*(r1-r3)*(r3-r2)/(r1*r2-2*r1*r3+r2*r3)/(r1-2*r2+r3)
    nu = (-r1*r2 + 2*r1*r3-r2*r3)/(r1-2*r2+r3)

    return lambda,mu,nu
end
"""
tstat_2sample(x1,x2)

Compute t statistics of histograms x1 and x2
"""
tstat_2sample(x1,x2) = (mean_histogram(x1) - mean_histogram(x2))/(sqrt(var_histogram(x1)/length(x1) + var_histogram(x2)/length(x2)))

"""
log_2sample(x1,x2)

compute the log of the ratio of means of histograms x1 and x2
"""
log_2sample(x1,x2) = log(mean_histogram(x1)) - log(mean_histogram(x2))

"""
delta_2sample(x1,x2)

compute the difference of means of x1 and x2
"""
delta_2sample(x1,x2) = mean_histogram(x1)-mean_histogram(x2)

"""
delta_2frac(x1,x2)

compute E(x1) - E(x2) of means normalize dot mean of x1
"""
delta_2frac(x1,x2) = delta_2sample(x1,x2)/mean_histogram(x1)

"""
mediansmooth(xin,window)

median smooth histogram xin over a window
"""
function mediansmooth(xin,window)
    x = copy(xin)
    halfwin = div(window-1,2)
    for i in 1:length(x)-window+1
        x[i+halfwin] = median(x[i:i+window-1])
    end
    x
end

"""
residenceprob_G(file,G,header)
residenceprob_G(r,n)

Residence probability of G states
given by exact steady state solution
of master equation

"""
function residenceprob_G(file::String,G,header=false)
    r = get_all_rates(file,header)
    m = size(r)[1]
    p = Array{Any,2}(undef,m,G+1)
    n = G-1
    for i in 1:m
        p[i,1] = r[i,1]
        p[i,2:end] = residenceprob_G(r[i,2:2*n+1],n)
    end
    return p
end

function residenceprob_G(r::Vector,n::Int)
    Gss = Array{Float64,2}(undef,1,n+1)
    Gss[1,1] = 1.
    for k in 1:n
        Gss[1,k+1] = Gss[1,k]*r[2*k-1]/r[2*k]
    end
    Gss ./= sum(Gss)
end

"""
splicesiteusage()

splice site usage probability is the probabilty of ejection times
the probability of not ejection prior to that point
"""
splicesiteusage(model::GRSMmodel) = splicesiteusage(model.rates,model.G-1,model.R)
function splicesiteusage(r::Vector,n::Int,nr::Int)
    nu = get_nu(r,n,nr)
    eta = get_eta(r,n,nr)
	ssf = zeros(nr)
	survival = 1
	for j in 1:nr
		ssf[j] = eta[j]/(nu[j+1]+eta[j])*survival
		survival *= nu[j+1]/(nu[j+1]+eta[j])
	end
	return ssf
end

"""
burstoccupancy(n,nr,r)
Burst size distribution of GRS  model
for total pre-mRNA occupancy and
unspliced (visible) occupancy
obtained by marginalizing over conditional steady state distribution
"""
burstoccupancy(model::GRSMmodel) = burstsize(model.G-1,model.R,model.rates)

function burstoccupancy(n::Int,nr::Int,r::Vector)
    T =  mat_GSR_T(r,n,nr)
    pss = normalized_nullspace(T)
	Rss = zeros(nr)
	Rssvisible = zeros(nr)
	ssf = zeros(nr)
	asum = 0
	for w=1:nr
		for i = 1:n+1, z = 1:3^nr
			zdigits = digits(z-1,base=3,pad=nr)
			a = i + (n+1)*(z-1)
			if sum(zdigits .== 2) == w
				Rssvisible[w] += pss[a]
			end
			if sum(zdigits .> 0) == w
				Rss[w] += pss[a]
			end
		end
	end
	Rss ./= sum(Rss), Rssvisible ./= sum(Rssvisible)
end

model2_mean(ron,roff,eject,decay,nalleles) = 2*model2_mean(ron,roff,eject,decay)

model2_mean(ron,roff,eject,decay) = ron/(ron+roff)*eject/decay

model2_variance(ron,roff,eject,decay,nalleles) = 2*model2_variance(ron,roff,eject,decay)

model2_variance(ron,roff,eject,decay) = ron/(ron+roff)*eject/decay + ron*roff/(ron+roff)^2 * eject^2/decay/(ron + roff + decay)

"""
test_steadystatemodel(model,nhist)

Compare chemical master solution to Gillespie simulation for steadystate mRNA distribution

"""
function test_steadystatemodel(model::AbstractGMmodel,nhist)
    G = model.G
    r = model.rates
    g1 = steady_state(r[1:2*G],G-1,nhist,model.nalleles)
    g2 = simulatorGM(r[1:2*G],G-1,nhist,model.nalleles)
    return g1,g2
end
#
# function fit_rna_test(root)
#     gene = "CENPL"
#     cell = "HCT116"
#     fish = false
#     data = data_rna(gene,"MOCK","data/HCT116_testdata",fish,"scRNA_test",root)
#     model = model_rna(gene,cell,2,fish,.01,[1,2,3],(),"scRNA_test","scRNA_test",1,".",data,.05,1.0)
#     options = MHOptions(10000,2000,0,120.,1.,100.)
#     fit,stats,waic = run_mh(data,model,options,1);
#     return stats.meanparam, fit.llml
# end
#

# function online_covariance(data1, data2)
#     meanx = meany = C = n = 0
#     for x in data1, y in data2
#         n += 1
#         dx = x - meanx
#         meanx += dx / n
#         meany += (y - meany) / n
#         C += dx * (y - meany)
#         population_covar = C / n
#         # Bessel's correction for sample variance
#         sample_covar = C / (n - 1)
#     end
# end


using StaticArrays

typealias Vec2d SVector{2,Float64};

type CartesianMesh
    nx
    ny
    Δx
    Δy
end

CartesianMesh(nx, ny) = CartesianMesh(nx, ny, 1.0/nx, 1.0/ny);

immutable Cell
    id  :: Int
    x   :: Vec2d
    vol :: Float64
end

immutable Face
    x     :: Vec2d
    s     :: Vec2d
    owner :: Int
    neigh :: Int
end;

immutable Patch
    name
    faces
end

function cell(msh::CartesianMesh, id)
    i,j = rem(id-1, msh.nx)+1, div(id-1, msh.nx)+1
    return Cell(id, Vec2d( (i-1/2)*msh.Δx, (j-1/2)*msh.Δy ), msh.Δx*msh.Δy )
end;

cells(msh::CartesianMesh) = @task begin
    id = 1
    vol = msh.Δx * msh.Δy
    for y in range(msh.Δy/2, msh.Δy, msh.ny)
        for x in range(msh.Δx/2, msh.Δx, msh.nx)
            produce( Cell(id, Vec2d(x,y), vol) )
            id += 1
        end
    end
end;

internalfaces(msh::CartesianMesh) = @task begin
    id(i,j) = i + (j-1)*msh.nx
    Δx, Δy = msh.Δx, msh.Δy

    s = Vec2d(Δy, 0.0)
    for j=1:msh.ny
        for i=1:msh.nx-1
            owner, neigh = id(i,j), id(i+1,j)
            x = Vec2d( i*Δx, (j-1/2)*Δy )
            produce( Face(x, s, owner, neigh) )
        end
    end

    s = Vec2d(0.0, Δx)
    for j=1:msh.ny-1
        for i=1:msh.nx
            owner, neigh = id(i,j), id(i,j+1)
            x = Vec2d( (i-1/2)*Δx, j*Δy )
            produce( Face(x, s, owner, neigh) )
        end
    end

end;

function patches(msh::CartesianMesh)
    id(i,j) = i + (j-1)*msh.nx
    Δx, Δy = msh.Δx, msh.Δy

    leftfaces() = @task begin
        s = Vec2d(-Δy, 0.0)
        for j=1:msh.ny
            produce( Face(Vec2d(0.0, (j-1/2)*Δy), s, id(1,j), 0) )
        end
    end

    rightfaces() = @task begin
        s = Vec2d(Δy, 0.0)
        for j=1:msh.ny
            produce( Face(Vec2d(msh.nx*Δx, (j-1/2)*Δy), s, id(msh.nx,j), 0) )
        end
    end

    bottomfaces() = @task begin
        s = Vec2d(0.0,-Δx)
        for i=1:msh.nx
            produce( Face(Vec2d((i-1/2)*Δx, 0.0), s, id(i,1), 0) )
        end
    end

    topfaces() = @task begin
        s = Vec2d(0.0,Δx)
        for i=1:msh.nx
            produce( Face(Vec2d((i-1/2)*Δx, msh.ny*Δy), s, id(i,msh.ny), 0) )
        end
    end

    return [ Patch("left", leftfaces), Patch("right", rightfaces),
        Patch("bottom", bottomfaces), Patch("top", topfaces) ]
end;

name(p::Patch) = p.name;
boundaryfaces(p::Patch) = p.faces();


# An implementation of MINRES for the solution of the
# linear system Ax = b, or the linear least-squares problem
#
#  minimize ‖Ax - b‖
#
# where A is square and symmetric.
#
# MINRES is formally equivalent to applying the conjugate residuals method
# to Ax = b when A is positive definite, but is more general and also applies
# to the case where A is indefinite.
#
# This implementation follows the original implementation by
# Michael Saunders described in
#
# C. C. Paige and M. A. Saunders, Solution of Sparse Indefinite Systems
# of Linear Equations, SIAM Journal on Numerical Analysis, 12(4), 617-629, 1975.
#
# Dominique Orban, <dominique.orban@gerad.ca>
# Brussels, Belgium, June 2015.
# Montreal, August 2015.

export minres

"""Solve the shifted linear least-squares problem

  minimize ‖b - (A + λ I)x‖₂²

or the shifted linear system

  (A + λ I) x = b

using the MINRES method, where λ ≥ 0 is a shift parameter,
where A is square and symmetric.

MINRES is formally equivalent to applying CR to Ax=b when A is positive
definite, but is typically more stable and also applies to the case where
A is indefinite.

MINRES produces monotonic residuals ‖r‖₂ and optimality residuals ‖A'r‖₂.

A preconditioner M may be provided in the form of a linear operator and is
assumed to be symmetric and positive definite.
"""
function minres(A :: AbstractLinearOperator, b :: AbstractVector{T};
                M :: AbstractLinearOperator=opEye(size(A,1)), λ :: Float64=0.0,
                atol :: Float64=1.0e-12, rtol :: Float64=1.0e-12, etol :: Float64=1.0e-8,
                window :: Int=5, itmax :: Int=0, conlim :: Float64=1.0e+8,
                verbose :: Bool=false) where T <: Number

  m, n = size(A)
  m == n || error("System must be square")
  size(b, 1) == m || error("Inconsistent problem size")
  verbose && @printf("MINRES: system of size %d\n", n)

  ϵM = eps(T)
  x = zeros(T, n)
  ctol = conlim > 0.0 ? 1 ./ conlim : 0.0;

  # Initialize Lanczos process.
  # β₁ M v₁ = b.
  r1 = copy(b)
  y = M * r1
  β₁ = @kdot(m, r1, y)
  β₁ < 0.0 && error("Preconditioner is not positive definite")
  β₁ == 0.0 && return (x, SimpleStats(true, true, [0.0], [0.0], "x = 0 is a zero-residual solution"))
  β₁ = sqrt(β₁)
  β = β₁

  oldβ = 0.0
  δbar = 0.0
  ϵ = 0.0
  rNorm = β₁
  rNorms = [β₁]
  ϕbar = β₁
  rhs1 = β₁
  rhs2 = 0.0
  γmax = 0.0
  γmin = Inf
  cs = -1.0
  sn = 0.0
  v = zeros(T, n)
  w = zeros(T, n)
  w1 = zeros(T, n)
  w2 = zeros(T, n)
  r2 = copy(r1)

  ANorm² = 0.0
  ANorm = 0.0
  Acond = 0.0
  ArNorm = 0.0
  ArNorms = [ArNorm]
  xNorm = 0.0

  xENorm² = 0.0
  err_lbnd = 0.0
  err_vec = zeros(T, window)

  verbose && @printf("%5s  %7s  %7s  %7s  %8s  %8s  %7s  %7s\n",
                     "Aprod", "‖r‖", "‖A'r‖", "β", "cos", "sin", "‖A‖", "κ(A)")
  verbose && @printf("%5d  %7.1e  %7.1e  %7.1e  %8.1e  %8.1e  %7.1e  %7.1e\n",
                     0, rNorm, ArNorm, β, cs, sn, ANorm, Acond)

  iter = 0
  itmax == 0 && (itmax = n)

  tol = atol + rtol * β₁
  status = "unknown"
  solved = solved_mach = solved_lim = (rNorm ≤ rtol)
  tired  = iter ≥ itmax
  ill_cond = ill_cond_mach = ill_cond_lim = false
  zero_resid = zero_resid_mach = zero_resid_lim = (rNorm ≤ tol)
  fwd_err = false

  while ! (solved || tired || ill_cond)
    iter = iter + 1

    # Generate next Lanczos vector.
    v = copy(y)
    @kscal!(n, 1 ./ β, v)
    y = A * v
    λ != 0.0 && @kaxpy!(n, -λ, v, y)          # (y = y - λ * v)
    iter ≥ 2 && @kaxpy!(n, -β / oldβ, r1, y)  # (y = y - β / oldβ * r1)

    α = dot(v, y)
    @kaxpy!(n, -α / β, r2, y)  # y = y - α / β * r2
    r1 = copy(r2)
    r2 = copy(y)
    y = M * r2
    oldβ = β
    β = @kdot(n, r2, y)
    β < 0.0 && error("Preconditioner is not positive definite")
    β = sqrt(β)
    ANorm² = ANorm² + α * α + oldβ * oldβ + β * β

    # Apply rotation to obtain
    #  [ δₖ    ϵₖ₊₁    ] = [ cs  sn ] [ δbarₖ  0    ]
    #  [ γbar  δbarₖ₊₁ ]   [ sn -cs ] [ αₖ     βₖ₊₁ ]
    oldϵ = ϵ
    δ = cs * δbar + sn * α
    γbar = sn * δbar - cs * α
    ϵ = sn * β
    δbar = -cs * β
    root = sqrt(γbar * γbar + δbar * δbar)
    ArNorm = ϕbar * root  # = ‖A'rₖ₋₁‖
    push!(ArNorms, ArNorm)

    # Compute the next plane rotation.
    γ = sqrt(γbar * γbar + β * β)
    γ = max(γ, ϵM)
    cs = γbar / γ
    sn = β / γ
    ϕ = cs * ϕbar
    ϕbar = sn * ϕbar

    # Update x.
    w1 = copy(w2)
    w2 = copy(w)
    w = (v - oldϵ * w1 - δ * w2) / γ
    @kaxpy!(n, ϕ, w, x)  # x = x + ϕ * w
    xENorm² = xENorm² + ϕ * ϕ

    # Compute lower bound on forward error.
    err_vec[mod(iter, window) + 1] = ϕ
    iter ≥ window && (err_lbnd = norm(err_vec))

    γmax = max(γmax, γ)
    γmin = min(γmin, γ)
    ζ = rhs1 / γ
    rhs1 = rhs2 - δ * ζ
    rhs2 = -ϵ * ζ

    # Estimate various norms.
    ANorm = sqrt(ANorm²)
    xNorm = @knrm2(n, x)
    ϵA = ANorm * ϵM
    ϵx = ANorm * xNorm * ϵM
    ϵr = ANorm * xNorm * rtol
    d = γbar
    d == 0.0 && (d = ϵA)

    rNorm = ϕbar

    test1 = rNorm / (ANorm * xNorm)
    test2 = root / ANorm
    push!(rNorms, rNorm)

    Acond = γmax / γmin

    verbose && @printf("%5d  %7.1e  %7.1e  %7.1e  %8.1e  %8.1e  %7.1e  %7.1e\n",
                       iter, test1, test2, β, cs, sn, ANorm, Acond)

    if iter == 1
      # A'b = 0 so x = 0 is a minimum least-squares solution
      β / β₁ ≤ 10 * ϵM && return (x, SimpleStats(true, true, [β₁], [0.0], "x is a minimum least-squares solution"))
    end

    # Stopping conditions that do not depend on user input.
    # This is to guard against tolerances that are unreasonably small.
    ill_cond_mach = (1.0 + 1.0 / Acond ≤ 1.0)
    solved_mach = (1.0 + test2 ≤ 1.0)
    zero_resid_mach = (1.0 + test1 ≤ 1.0)
    # solved_mach = (ϵx ≥ β₁)

    # Stopping conditions based on user-provided tolerances.
    tired = iter ≥ itmax
    ill_cond_lim = (1.0 / Acond ≤ ctol)
    solved_lim = (test2 ≤ tol)
    zero_resid_lim = (test1 ≤ tol)
    iter ≥ window && (fwd_err = err_lbnd ≤ etol * sqrt(xENorm²))

    zero_resid = zero_resid_mach | zero_resid_lim
    ill_cond = ill_cond_mach | ill_cond_lim
    solved = solved_mach | solved_lim | zero_resid | fwd_err
  end

  tired         && (status = "maximum number of iterations exceeded")
  ill_cond_mach && (status = "condition number seems too large for this machine")
  ill_cond_lim  && (status = "condition number exceeds tolerance")
  solved        && (status = "found approximate minimum least-squares solution")
  zero_resid    && (status = "found approximate zero-residual solution")
  fwd_err       && (status = "truncated forward error small enough")

  stats = SimpleStats(solved, !zero_resid, rNorms, ArNorms, status)
  return (x, stats)
end


using HaploSoup
using Test

@testset "HaploSoup.jl" begin
    H =  Array{Int8, 2}([0 1 0 1 0 1; 1 1 0 0 0 1; 1 1 1 1 1 1; 0 1 1 1 1 0; 0 0 0 0 0 0; 1 0 0 0 1 0; 1 1 0 0 0 1; 0 1 0 1 1 0])
    ppa = Array{Int32, 2}([1 1 5 5 5 5 5; 2 4 6 6 6 2 6; 3 5 1 1 2 7 8; 4 8 4 8 7 1 4; 5 2 8 2 1 6 2; 6 3 2 7 8 8 7; 7 6 3 4 4 4 1; 8 7 7 3 3 3 3])
    div = Array{Int32, 2}([1 2 3 4 5 6 7; 1 1 2 2 2 3 6; 1 1 3 3 3 1 5; 1 1 1 1 1 5 4; 1 2 1 2 5 6 7; 1 1 2 1 1 5 1; 1 1 1 4 4 4 5; 1 1 1 2 2 2 6])
    @test build_prefix_and_divergence_arrays(H) == (ppa, div)
end


using MathPhysicalConstants, Measurements, SpecialFunctions, QuadGK, Calculus, Unitful
import MathPhysicalConstants.SI: α, O_M, c, e, ε0, h, µ0, ℯ, PlanckConstantHbar, ƛe
using Test

import Base: isapprox
import Measurements: value, uncertainty

if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

@testset "Comparisons and Tests" begin
    # isapprox is overloaded only for tests.  Make sure it works as expected.
    @test measurement(big"0.75", 0.01) == 3//4 ± 1//100
    @test measurement(big(π)) ≠ π
    @test ℯ ≠ measurement(Float32(ℯ))
    @test 3//4 == measurement(Float32(0.75), Float32(0)) ≠ 4//3
    @test MathPhysicalConstants.SI.PlanckConstantHbar*MathPhysicalConstants.SI.c == 3.1615262057012976e-26
    @test MathPhysicalConstants.SI.PlanckConstantHbar/(MathPhysicalConstants.SI.MassElectron*MathPhysicalConstants.SI.c) == 3.861592046068738e-13
    @test MathPhysicalConstants.SI.GravitationalConstant^2 == 4.45433438464e-21
    @test MathPhysicalConstants.SI.MassElectron*MathPhysicalConstants.SI.MassProton*MathPhysicalConstants.SI.MassHydrogen == 2.550178071499845e-84
    @test (2*MathPhysicalConstants.SI.PlanckConstantHbar^2)/(MathPhysicalConstants.SI.SanchezGravitationalConstant*MathPhysicalConstants.SI.MassElectron*MathPhysicalConstants.SI.MassProton*MathPhysicalConstants.SI.MassNeutron) == 1.3056254088279212e26
    @test MathPhysicalConstants.SI.KotovCosmicPeriodicity == 9600.061
    @test MathPhysicalConstants.SI.SanchezGravitationalConstant == 6.675453818e-11
    @test MathPhysicalConstants.SI.O_M == 80801742479451287588645990496171075700575436800000000
    @test isinteger(-3 ± 0)
    @test !isinteger(3 ± 1)
    @test !isinteger(pi ± 0)
    @test iszero(0 ± 0)
    @test iszero(zero(Measurement))
    @test !iszero(0 ± 1)
    @test !iszero(1 ± 0)
    @test isone(1 ± 0)
    @test isone(one(Measurement))
    @test !isone(1 ± 1)
    @test !isone(0 ± 0)
end

@testset "Base" begin
#    @test ustrip(big(h)) == big"6.626070040e-34"
    @test setprecision(BigFloat, 768) do; precision(ustrip(big(c))) end == 768
    @test measurement(h) === measurement(h)
    @test iszero(measurement(α) - measurement(α))
#    @test isone(measurement(BigFloat, ƛe) / measurement(BigFloat, ƛe))
#    @test iszero(measurement(BigFloat, PlanckConstantHbar) - (measurement(BigFloat, h) / 2big(pi)))
#    @test isone(measurement(BigFloat, PlanckConstantHbar) / (measurement(BigFloat, h) / 2big(pi)))
end

@testset "Utils" begin
    @test c === ustrip(float(c))
    @test_throws ErrorException h.foo
end

#@testset "Promotion" begin
#    x = @inferred(inv(μ0 * c ^ 2))
#    T = promote_type(typeof(ε0), typeof(x))
#    u = u"A^2 * kg^-1 * m^-3 * s^4"
#    @test promote_type(typeof(α), BigInt) === BigFloat
#    @test promote_type(typeof(α), typeof(1u"m/cm")) === Float64
#    @test T === Quantity{Float64, dimension(x), typeof(u)}
#    @test convert(T, ε0) === (ε0 / unit(ε0)) * u
#    @test convert(Float32, α) === float(Float32, α)
#    @test uconvert(u"cm/s", c) === uconvert(u"cm/s", float(c))
#end
@testset "Maths" begin
#    @test α ≈ @inferred(ℯ^2/(4 * pi * ε0 * PlanckConstantHbar * c))
    @test @inferred(α + 2) ≈ 2 + float(α)
    @test @inferred(5 + α) ≈ float(α) + 5
    @test @inferred(α + 2.718) ≈ 2.718 + float(α)
    @test @inferred(-3.14 + α) ≈ float(α) - 3.14
    @test ε0 ≈ @inferred(1 / (μ0 * c ^ 2))
    @test @inferred(big(0) + α) == big(α)
    @test @inferred(α * 1.0) == float(α)
end

@testset "Show" begin
    @test repr(c) ==
        "2.99792458e8"

end
# using Base.Test
#include("../src/RunTests.jl")


using Test, Documenter, Quaternionic

DocMeta.setdocmeta!(Quaternionic, :DocTestSetup, :(using Quaternionic); recursive=true)
doctest(Quaternionic)


"""
$(SIGNATURES)

Compute the convex hull of `points` and store the result in `hull` using the
[Jarvis march (gift wrapping) algorithm](https://en.wikipedia.org/wiki/Gift_wrapping_algorithm).
This algorithm has ``O(nh)`` complexity, where ``n`` is the number of points and ``h`` is
the number of vertices of the convex hull.
"""
function jarvis_march!(hull::ConvexHull, points::AbstractVector{<:PointLike}; atol=eps(eltype(hull)))
    # Adapted from https://www.algorithm-archive.org/contents/jarvis_march/jarvis_march.html.
    op = orientation_comparator(hull)
    n = length(points)
    vertices = hull.vertices
    @inbounds begin
        # Preallocate
        resize!(vertices, n)
        if n <= 2
            vertices .= points
        else
            # Find an initial hull vertex using lexicographic ordering.
            current = first(points)
            for i in 2 : n
                p = points[i]
                if Tuple(p) < Tuple(current)
                    current = p
                end
            end

            i = 1
            while true
                # Add point
                vertices[i] = current

                # Next point is the one with extremal internal angle.
                next = first(points)
                δnext = next - current
                for j in 2 : n
                    p = points[j]
                    δ = p - current
                    c = cross2(δnext, δ)

                    # Note the last clause here, which ensures strong convexity in the presence of
                    # collinear points by accepting `p` if it's farther away from `current` than
                    # `next`.
                    if next == current || (op(0, c) && abs(c) > atol) || (abs(c) <= atol && δ ⋅ δ > δnext ⋅ δnext)
                        next = p
                        δnext = δ
                    end
                end
                current = next
                current == first(vertices) && break
                i += 1
                if i > n
                    # println(IOContext(stdout, :compact=>false), points)
                    # println(IOContext(stdout, :compact=>false), vertices)
                    error("should never get here.")
                end
            end

            # Shrink to computed number of vertices.
            resize!(vertices, i)
        end
    end
    return hull
end


# # 2D decaying turbulence
#
#md # This example can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/literated/twodnavierstokes_decaying.ipynb).
#
# A simulation of decaying two-dimensional turbulence.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add GeophysicalFlows, Printf, Random, Plots"
# ```

# ## Let's begin
# Let's load `GeophysicalFlows.jl` and some other needed packages.
#
using GeophysicalFlows, Printf, Random, Plots

using Random: seed!
using GeophysicalFlows: peakedisotropicspectrum

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)
nothing # hide

# ## Numerical, domain, and simulation parameters
#
# First, we pick some numerical and physical parameters for our model.

n, L  = 128, 2π             # grid resolution and domain length
nothing # hide

# Then we pick the time-stepper parameters
    dt = 1e-2  # timestep
nsteps = 4000  # total number of steps
 nsubs = 20    # number of steps between each plot
nothing # hide

# ## Problem setup
# We initialize a `Problem` by providing a set of keyword arguments. The
# `stepper` keyword defines the time-stepper to be used.
prob = TwoDNavierStokes.Problem(dev; nx=n, Lx=L, ny=n, Ly=L, dt=dt, stepper="FilteredRK4")
nothing # hide

# Next we define some shortcuts for convenience.
sol, clock, vars, grid = prob.sol, prob.clock, prob.vars, prob.grid
x, y = grid.x, grid.y
nothing # hide

# ## Setting initial conditions

# Our initial condition tries to reproduce the initial condition used by McWilliams (_JFM_, 1984).
seed!(1234)
k₀, E₀ = 6, 0.5
ζ₀ = peakedisotropicspectrum(grid, k₀, E₀, mask=prob.timestepper.filter)
TwoDNavierStokes.set_ζ!(prob, ζ₀)
nothing # hide

# Let's plot the initial vorticity field. Note that when plotting, we decorate the variable
# to be plotted with `Array()` to make sure it is brought back on the CPU when `vars` live on
# the GPU.
heatmap(x, y, Array(vars.ζ'),
    aspectratio = 1,
              c = :balance,
           clim = (-40, 40),
          xlims = (-L/2, L/2),
          ylims = (-L/2, L/2),
         xticks = -3:3,
         yticks = -3:3,
         xlabel = "x",
         ylabel = "y",
          title = "initial vorticity",
     framestyle = :box)

# ## Diagnostics

# Create Diagnostics -- `energy` and `enstrophy` functions are imported at the top.
E = Diagnostic(TwoDNavierStokes.energy, prob; nsteps=nsteps)
Z = Diagnostic(TwoDNavierStokes.enstrophy, prob; nsteps=nsteps)
diags = [E, Z] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.
nothing # hide

# ## Output

# We choose folder for outputing `.jld2` files and snapshots (`.png` files).
filepath = "."
plotpath = "./plots_decayingTwoDNavierStokes"
plotname = "snapshots"
filename = joinpath(filepath, "decayingTwoDNavierStokes.jld2")
nothing # hide

# Do some basic file management
if isfile(filename); rm(filename); end
if !isdir(plotpath); mkdir(plotpath); end
nothing # hide

# And then create Output
get_sol(prob) = prob.sol # extracts the Fourier-transformed solution
get_u(prob) = irfft(im * grid.l .* grid.invKrsq .* sol, grid.nx)
out = Output(prob, filename, (:sol, get_sol), (:u, get_u))
saveproblem(out)
nothing # hide

# ## Visualizing the simulation

# We initialize a plot with the vorticity field and the time-series of
# energy and enstrophy diagnostics.

p1 = heatmap(x, y, Array(vars.ζ'),
         aspectratio = 1,
                   c = :balance,
                clim = (-40, 40),
               xlims = (-L/2, L/2),
               ylims = (-L/2, L/2),
              xticks = -3:3,
              yticks = -3:3,
              xlabel = "x",
              ylabel = "y",
               title = "vorticity, t=" * @sprintf("%.2f", clock.t),
          framestyle = :box)

p2 = plot(2, # this means "a plot with two series"
               label = ["energy E(t)/E(0)" "enstrophy Z(t)/Z(0)"],
              legend = :right,
           linewidth = 2,
               alpha = 0.7,
              xlabel = "t",
               xlims = (0, 41),
               ylims = (0, 1.1))

l = @layout Plots.grid(1, 2)
p = plot(p1, p2, layout = l, size = (800, 360))

# ## Time-stepping the `Problem` forward

# We time-step the `Problem` forward in time.

startwalltime = time()

anim = @animate for j = 0:Int(nsteps/nsubs)
  if j % (1000 / nsubs) == 0
    cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy])

    log = @sprintf("step: %04d, t: %d, cfl: %.2f, ΔE: %.4f, ΔZ: %.4f, walltime: %.2f min",
        clock.step, clock.t, cfl, E.data[E.i]/E.data[1], Z.data[Z.i]/Z.data[1], (time()-startwalltime)/60)

    println(log)
  end

  p[1][1][:z] = Array(vars.ζ)
  p[1][:title] = "vorticity, t=" * @sprintf("%.2f", clock.t)
  push!(p[2][1], E.t[E.i], E.data[E.i]/E.data[1])
  push!(p[2][2], Z.t[Z.i], Z.data[Z.i]/Z.data[1])

  stepforward!(prob, diags, nsubs)
  TwoDNavierStokes.updatevars!(prob)
end

mp4(anim, "twodturb.mp4", fps=18)

# Last we can save the output by calling
# ```julia
# saveoutput(out)`
# ```

# ## Radial energy spectrum

# After the simulation is done we plot the instantaneous radial energy spectrum to illustrate
# how `FourierFlows.radialspectrum` can be used,

E  = @. 0.5 * (vars.u^2 + vars.v^2) # energy density
Eh = rfft(E)                  # Fourier transform of energy density
kr, Ehr = FourierFlows.radialspectrum(Eh, grid, refinement=1) # compute radial specturm of `Eh`
nothing # hide

# and we plot it.
plot(kr, abs.(Ehr),
    linewidth = 2,
        alpha = 0.7,
       xlabel = "kᵣ", ylabel = "∫ |Ê| kᵣ dk_θ",
        xlims = (5e-1, grid.nx),
       xscale = :log10, yscale = :log10,
        title = "Radial energy spectrum",
       legend = false)


@inline front(t::Tuple, m::Integer) = ntuple(i -> t[i], m)
@inline front(t::Tuple, ::StaticOrVal{M}) where {M} = ntuple(i -> t[i], Val(M))
@inline front(t::NTuple{L,Any}) where {L} = front(t, Val(L - 1))

@inline function tail(t::Tuple, n::Integer)
    m = length(t) - n
    ntuple(i -> t[i+m], n)
end
@inline function tail(t::NTuple{L,Any}, ::StaticOrVal{N}) where {L,N}
    M = L - N
    ntuple(i -> t[i+M], Val(N))
end
@inline tail(t::NTuple{L,Any}) where {L} = tail(t, Val(L - 1))

@inline function tuple_split(t::NTuple{L,Any}, ::StaticOrVal{M}) where {L,M}
    front(t, Val(M)), tail(t, Val(L - M))
end
@inline function tuple_split(t::NTuple{L,Any}, m::Integer) where {L}
    front(t, Val(m)), tail(t, L - m)
end

@pure tuple_length(T::Type{<:Tuple}) = length(T.parameters)
@pure tuple_length(t::Tuple) = length(t)

@pure tuple_prod(T::Type{<:Tuple}) = length(T.parameters) == 0 ? 1 : *(T.parameters...)
@pure tuple_prod(t::Tuple) = prod(t)

@pure function tuple_minimum(T::Type{<:Tuple})
    length(T.parameters) == 0 ? 0 : minimum(tuple(T.parameters...))
end
@pure tuple_minimum(t::Tuple) = minimum(t)


using Gnuplot
function ribbon(x, y, yerr; scolor = "black", linecolor = "black", leg = "y")
    @gp x y "w l lw 1 lc '$(linecolor)' t '$(leg)'"
    @gp :- x y y .+ yerr "w filledcu lc '$(scolor)' t ''"
    @gp :- "set style fill transparent solid 0.2 noborder"
    @gp :- x y y .- yerr "w filledcu lc '$(scolor)' t ''"
end
# example
x = LinRange(-π,3π,100)
y = sin.(x)
yerr =0.15
ribbon(x, y, yerr; scolor = "orange", linecolor = "red")
save(term ="pngcairo size 600,400", output ="ribbon.png")


module Examples

# Utilities
include("reparametrization.jl")

# Dynamical systems
include("pitchfork.jl")
include("transcritical.jl")
include("bautin.jl")
include("cusp.jl")
include("calcium.jl")
include("predator_prey.jl")
include("bazykin_85.jl")
include("bazykin_khibnik_81.jl")
include("duffing_van_der_pol.jl")
include("morris_lecar.jl")

example_modules() = [
    Pitchfork,
    Transcritical,
    Bautin,
    Cusp,
    Calcium,
    PredatorPrey,
    Bazykin85,
    BazykinKhibnik81,
    DuffingVanDerPol,
    MorrisLecar,
]

examples() = [nameof(ex) => ex for ex in example_modules()]

end  # module

using .Examples: examples
using .Examples.Reparametrization: reparametrize


using TeukEvolution
using Test

include("TestRadial.jl")
include("TestSphere.jl")

import .TestRadial, .TestSphere

#============================================================================
## Testing Radial functions
============================================================================#
# Test 4th order convergence of derivatives

n11, n12 = TestRadial.norms_diff(128,48)
n21, n22 = TestRadial.norms_diff(256,60)

@test n11/n21 >= 15.5
@test n12/n22 >= 15.5

# Test filter is TVD

n1, n2 = TestRadial.norms_filter(128,48)
n3, n4 = TestRadial.norms_filter(180,90)

@test n2 < n1
@test n4 < n3

#============================================================================
## Testing Spherical functions
============================================================================#

println("Testing spherical spin-weighted spherical harmonic inner product")
TestSphere.test_swal_inner_product(48,-2,-3)
TestSphere.test_swal_inner_product(48, 0,-3)
TestSphere.test_swal_inner_product(48, 1,-3)

TestSphere.test_swal_inner_product(48,-1,-2)
TestSphere.test_swal_inner_product(48, 0,-2)
TestSphere.test_swal_inner_product(48, 2,-2)

TestSphere.test_swal_inner_product(48,-2,-1)
TestSphere.test_swal_inner_product(48, 0,-1)
TestSphere.test_swal_inner_product(48, 1,-1)

TestSphere.test_swal_inner_product(48,-1,0)
TestSphere.test_swal_inner_product(48, 0,0)
TestSphere.test_swal_inner_product(48, 2,0)

TestSphere.test_swal_inner_product(48,-2,1)
TestSphere.test_swal_inner_product(48, 0,1)
TestSphere.test_swal_inner_product(48, 1,1)

TestSphere.test_swal_inner_product(48,-1,2)
TestSphere.test_swal_inner_product(48, 0,2)
TestSphere.test_swal_inner_product(48, 2,2)

TestSphere.test_swal_inner_product(48,-2,3)
TestSphere.test_swal_inner_product(48, 0,3)
TestSphere.test_swal_inner_product(48, 1,3)

println("Testing spherical spin-weighted laplacian operator")
TestSphere.test_norm_swal_lap(32, 0,0,2)

TestSphere.test_norm_swal_lap(32, 0,1,2)
TestSphere.test_norm_swal_lap(32, 1,1,2)

TestSphere.test_norm_swal_lap(32,-1,-2,2)
TestSphere.test_norm_swal_lap(32, 0,-2,2)
TestSphere.test_norm_swal_lap(32, 2,-2,2)

TestSphere.test_norm_swal_lap(32,-1,2,2)
TestSphere.test_norm_swal_lap(32, 0,2,2)
TestSphere.test_norm_swal_lap(32, 2,2,2)

TestSphere.test_norm_swal_lap(48,-1,5,23)
TestSphere.test_norm_swal_lap(48, 0,5,23)
TestSphere.test_norm_swal_lap(48, 2,5,23)

println("Testing spherical spin-weighted raising operator")
TestSphere.test_norm_swal_raising(32, 0,0,2)

TestSphere.test_norm_swal_raising(32, 0,1,2)
TestSphere.test_norm_swal_raising(32, 1,1,2)

TestSphere.test_norm_swal_raising(32,-1,-2,2)
TestSphere.test_norm_swal_raising(32, 0,-2,2)
TestSphere.test_norm_swal_raising(32, 2,-2,2)

TestSphere.test_norm_swal_raising(32,-1,2,2)
TestSphere.test_norm_swal_raising(32, 0,2,2)
TestSphere.test_norm_swal_raising(32, 2,2,2)

TestSphere.test_norm_swal_raising(48,-1,5,23)
TestSphere.test_norm_swal_raising(48, 0,5,23)
TestSphere.test_norm_swal_raising(48, 2,5,23)

println("Testing spherical spin-weighted lowering operator")
TestSphere.test_norm_swal_lowering(32, 0,0,2)

TestSphere.test_norm_swal_lowering(32, 0,1,2)
TestSphere.test_norm_swal_lowering(32, 1,1,2)

TestSphere.test_norm_swal_lowering(32,-1,-2,2)
TestSphere.test_norm_swal_lowering(32, 0,-2,2)
TestSphere.test_norm_swal_lowering(32, 2,-2,2)

TestSphere.test_norm_swal_lowering(32,-1,2,2)
TestSphere.test_norm_swal_lowering(32, 0,2,2)
TestSphere.test_norm_swal_lowering(32, 2,2,2)

TestSphere.test_norm_swal_lowering(48,-1,5,23)
TestSphere.test_norm_swal_lowering(48, 0,5,23)
TestSphere.test_norm_swal_lowering(48, 2,5,23)

println("Testing spherical spin-weighted filter operator")
TestSphere.test_norm_filter(32,-2,2,2)
TestSphere.test_norm_filter(32,-1,3,4)
TestSphere.test_norm_filter(32, 0,1,5)
TestSphere.test_norm_filter(32, 1,0,2)
TestSphere.test_norm_filter(32, 2,-1,2)


using MAT

@testset "simulateDSM" begin
    # MATLAB output
    test = matread("resources/simulateDSM.mat")

    # ntf input test
    u1 = 0.5 * sin.(2*pi*85/8192 * (0:8191))
    u2 = 0.5 * sin.(2*pi*1045/8192 * (0:8191))
    u3 = 8 * sin.(2*pi*146/8192 * (0:8191))
    u4 = 16 * undbv(-3) * sin.(2*pi*164/16384 * (0:16383))
    u5 = 16 * undbv(-30) * sin.(2*pi*164/16384 * (0:16383))
    ntf1 = synthesizeNTF(5, 32, 1)
    ntf2 = synthesizeNTF(8, 64, 1, 1.5, 0.125)
    ntf3 = synthesizeNTF(7, 8, 1, 2)
    ntf4 = synthesizeNTF(7, 8, 1, 8)
    ntf5 = synthesizeNTF(6, 25, 1, 4)
    u    = [  u1,   u2,   u3,   u3,   u4,   u5]
    ntf  = [ntf1, ntf2, ntf3, ntf4, ntf5, ntf5]
    nlev = [   2,    2,   17,   17,   17,   17]
    for i = 1:length(u)
        v, _, _, y = simulateDSM(u[i], ntf[i], nlev[i])
        @test vec(test["v_$i"]) == v
        # Allow higher error tolerance because of differences in MATLAB and
        # Julia state space models. This also why xn and xmax are ignored;
        # the x vectors aren't the same. Default tolerance is sqrt(eps).
        @test isapprox(vec(test["y_$i"]), y, rtol=10*sqrt(eps(Float64)))
    end

    # ABCD input test
    u = range(0, 0.6; length=30)
    ntf = synthesizeNTF(5, 42, 1)
    a, g, b, c = realizeNTF(ntf)
    ABCD = stuffABCD(a, g, b, c)
    for i = 1:length(u)
        v, xn, xmax, y = simulateDSM(u[i]*ones(10000), ABCD)
        @test vec(test["v"][i:i,:]) == v
        @test isapprox(test["xn"][i,:,:], xn)
        @test isapprox(test["xmax"][i,:,:], xmax)
        @test isapprox(test["y"][i,:,:], y)
    end
end


function bottom_drag_quad!(Bu,Bv,KEu,KEv,sqrtKE,sqrtKE_u,sqrtKE_v,u,v,η,H,h,u²,v²,h_u,h_v)
    # quadratic bottom drag Bu,Bv = c_D/h * | u⃗ | * u⃗

    thickness!(h,η,H)
    Ix!(h_u,h)
    Iy!(h_v,h)

    speed!(u²,v²,u,v)
    Ix!(KEu,u²)
    Iy!(KEv,v²)

    # sqrt of KE, which is actually the kinetic energy without the 0.5 factor
    m,n = size(sqrtKE)
    @boundscheck (m+ep,n+2) == size(KEu) || throw(BoundsError())
    @boundscheck (m+2,n) == size(KEv) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            sqrtKE[i,j] = sqrt(KEu[i+ep,j+1] + KEv[i+1,j])
        end
    end

    Ix!(sqrtKE_u,sqrtKE)
    Iy!(sqrtKE_v,sqrtKE)

    m,n = size(Bu)
    @boundscheck (m,n) == size(sqrtKE_u) || throw(BoundsError())
    @boundscheck (m,n) == size(h_u) || throw(BoundsError())
    @boundscheck (m+2+ep,n+2) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bu[i,j] = c_D*sqrtKE_u[i,j] * u[i+1+ep,j+1] / h_u[i,j]
        end
    end

    m,n = size(Bv)
    @boundscheck (m,n) == size(sqrtKE_v) || throw(BoundsError())
    @boundscheck (m,n) == size(h_v) || throw(BoundsError())
    @boundscheck (m+2,n+2) == size(v) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bv[i,j] = c_D*sqrtKE_v[i,j] * v[i+1,j+1] / h_v[i,j]
        end
    end
end

"""Linear bottom drag computed as r_B*(u,v). r_B is negative and contains the
grid spacing Δ as gradient operators are dimensionless."""
function bottom_drag_lin!(Bu,Bv,KEu,KEv,sqrtKE,sqrtKE_u,sqrtKE_v,u,v,η,H,h,u²,v²,h_u,h_v)
    m,n = size(Bu)
    @boundscheck (m+2+ep,n+2) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bu[i,j] = r_B * u[i+1+ep,j+1]
        end
    end

    m,n = size(Bv)
    @boundscheck (m+2,n+2) == size(v) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bv[i,j] = r_B * v[i+1,j+1]
        end
    end
end

function no_bottom_drag!(Bu,Bv,KEu,KEv,sqrtKE,sqrtKE_u,sqrtKE_v,u,v,η,H,h,u²,v²,h_u,h_v)
    Bu .= zeero
    Bv .= zeero
end

if bottom_friction == "linear"
    bottom_drag! = bottom_drag_lin!
elseif bottom_friction == "quadratic"
    bottom_drag! = bottom_drag_quad!
elseif bottom_friction == "none"
    bottom_drag! = no_bottom_drag!
else
    throw(error("Bottom friction not correctly specified. Only linear, quadratic or none allowed."))
end


module Inflation

using SymPy
using DifferentialEquations
using ODEInterfaceDiffEq
#using DiffEqParamEstim
using LinearAlgebra
using CurveFit
using Serialization
using JLD2
using Parameters
using RuntimeGeneratedFunctions

export inflation_setup
export background_evolve
export transport_perturbations
export SetupOptions, BackgroundOptions, PerturbationOptions
export output_data

SymPy.diff(x::T,::SymPy.Sym) where {T <: Real} = 0
SymPy.diff(x::T,::SymPy.Sym,::SymPy.Sym) where {T <: Real} = 0

RuntimeGeneratedFunctions.init(@__MODULE__)

# needed as of SymPy 1.0.38
import SymPy.fn_map
fn_map["Mul"] = :__prod__

# change this whenever there's a breaking change in the cache functions
cache_version = 1.3
# unit conversion
Mpc_to_Lpl = 2.6245e-57

include("options.jl")
include("quoting.jl")
include("symbolic_setup.jl")
include("solvers.jl")
include("output.jl")

end # module


module TestSolubility

using Base.Test
using SoilTracers.PhysChem: solub_co2, solub_cos

@time @testset "Solubility functions" begin
    temp_0C = 273.15
    temp_25C = 298.15

    # test solub_co2

    k_co2_0C = solub_co2(temp_0C)
    @test typeof(k_co2_0C) <: Real
    @test isapprox(k_co2_0C, 1.739, rtol=1e-3)

    k_co2_25C = solub_co2(temp_25C)
    @test typeof(k_co2_25C) <: Real
    @test isapprox(k_co2_25C, 0.831, rtol=1e-3)

    # test solub_cos

    k_cos_0C = solub_cos(temp_0C)
    @test typeof(k_cos_0C) <: Real
    @test isapprox(k_cos_0C, 1.549, rtol=1e-3)

    k_cos_25C = solub_cos(temp_25C)
    @test typeof(k_cos_25C) <: Real
    @test isapprox(k_cos_25C, 0.488, rtol=1e-3)
end

end


function infidelity!(u,δt,Ut,Hd,Hc,H,U,X,D,V,u_last)
    # Calculate forward propagators
    calc_fprops!(U,X,D,V,u,δt,Hd,Hc,H,u_last)
    Uf = X[end]; N² = size(Ut,1)^2
    # Calculate infidelity:
    # fₑ = 1 - Φ where Φ = |⟨Ut,Uf⟩|²/N²
    return 1 - abs2(inner(Ut,Uf))/N²
end

function infidelityprime!(fp,u,δt,Ut,Hd,Hc,H,A,U,X,D,V,P,u_last)
    # Calculate forward and backward propagators
    calc_fprops!(U,X,D,V,u,δt,Hd,Hc,H,u_last)
    calc_bprops!(P,U,Ut)
    n = length(U); m = length(Hc); N² = size(Ut,1)^2
    # Calculate exact derivative of the infidelity function:
    # ∂Φ/∂uₖⱼ  = ⟨Pⱼ,∂Uⱼ/∂uₖⱼ*Xⱼ₋₁⟩⟨Xⱼ,Pⱼ⟩/N² + c.c.
    # ∂fₑ/∂uₖⱼ = -∂Φ/∂uₖⱼ
    #          = -2*Re(⟨Pⱼ,Jₖⱼ*Xⱼ₋₁⟩⟨Uf,Ut⟩)/N² where Jₖⱼ = ∂Uⱼ/∂uₖⱼ
    Jkj = similar(A)
    cisDj = Vector{ComplexF64}(length(D[1]))
    a = inner(X[end],Ut)
    for j = 1:n
        cisDj .= cis.(D[j])
        for k = 1:m
            Jmat!(Jkj,Hc[k],cisDj,D[j],V[j],δt,A)
            j==1 ? copy!(A,Jkj) : mul!(A,Jkj,X[j-1])
            fp[(k-1)*n+j] = -2*real(inner(P[j],A)*a)/N²
        end
    end
    return fp
end

function infidelityprimeapprox!(fp,u,δt,Ut,Hd,Hc,H,A,U,X,D,V,P,u_last)
    # Calculate forward and backward propagators
    calc_fprops!(U,X,D,V,u,δt,Hd,Hc,H,u_last)
    calc_bprops!(P,U,Ut)
    n = length(U); m = length(Hc); N² = size(Ut,1)^2
    # Calculate approximate derivative of fidelity error function:
    # ∂Φ/∂uₖⱼ  = ⟨Pⱼ,∂Uⱼ/∂uₖⱼ*Xⱼ₋₁⟩⟨Xⱼ,Pⱼ⟩/N² + c.c.
    # ∂fₑ/∂uₖⱼ = -∂Φ/∂uₖⱼ
    #          ≈ 2*Re(⟨Pⱼ,iδt*Hₖ*Xⱼ⟩⟨Uf,Ut⟩)/N²
    a = inner(X[end],Ut)
    for j = 1:n, k = 1:m
        mul!(A,Hc[k],X[j])
        fp[(k-1)*n+j] = -2δt*imag(inner(P[j],A)*a)/N²
    end
    return fp
end

function gen_opt_fun(Ut::Operator,Hd::Operator,Hc::Vector{<:Operator},t::Real,n::Integer)
    N = prod(dims(Hd))
    m = length(Hc)
    # Make sure we pass dense operators
    Ut_d = full(Ut)
    Hd_d = full(Hd)
    Hc_d = full.(Hc)
    # Generate cache for various objects
    H = promote_type(typeof(Hd_d),eltype(Hc_d))(N,N)
    A = Matrix{ComplexF64}(N,N)
    U = [similar(A) for i=1:n]
    X = deepcopy(U)
    P = deepcopy(U)
    # For the exact derivative we need to store eigenvectors and eigenvalues
    V = [similar(H) for i=1:n]
    D = [Vector{Float64}(N) for i=1:n] # eigenvalues are always real
    # Storage for last control ampitudes, NaN for first run
    u_last = fill(NaN64,m*n)
    # Create optimization function object
    f = (u) -> infidelity!(u,t/n,Ut_d,Hd_d,Hc_d,H,U,X,D,V,u_last)
    g! = (fp,u) -> infidelityprime!(fp,u,t/n,Ut_d,Hd_d,Hc_d,H,A,U,X,D,V,P,u_last)
    #g! = (fp,u) -> infidelityprimeapprox!(fp,u,t/n,Ut_d,Hd_d,Hc_d,H,A,U,X,D,V,P,u_last)
    # Return a OnceDifferentiable object with appropriate seed
    return OnceDifferentiable(f,g!,zeros(m*n)),X[end]
end

function grape(Ut::Operator,Hd::Operator,Hc::Vector{<:Operator},u_init,t::Real,n::Integer)
    m = length(Hc)
    n == size(u_init,1) || throw(ArgumentError("control amplitude matrix not consistent with number of timesteps"))
    m == size(u_init,2) || throw(ArgumentError("control amplitude matrix not consistent with number of control hamiltonians"))
    # Generate OnceDifferentiable object, and get reference to final U
    od, Uf = gen_opt_fun(Ut,Hd,Hc,t,n)
    # Optimization options
    #opt = Optim.Options(g_tol = 1E-9)
    # Run optimization
    res = optimize(od,vec(transpose(u_init)),ConjugateGradient())
    #res = optimize(od,vec(transpose(u_init)))
    # Reshape final control amplitude matrix
    uf = transpose(reshape(Optim.minimizer(res),m,n))
    return uf, Uf, res
end

function opt_2Q_QFT1(n)
    Hd = 0.5 * (σx⊗σx + σy⊗σy + σz⊗σz)
    Hc = [0.5*σx⊗σ0, 0.5*σy⊗σ0, 0.5*σ0⊗σx, 0.5*σ0⊗σy]
    t = 6.0
    u_init = zeros(n,4)
    #u_init = rand(n,4) .- 0.5
    Ut = Operator([0.5  0.5    0.5  0.5
                   0.5  0.5im -0.5 -0.5im
                   0.5 -0.5    0.5 -0.5
                   0.5 -0.5im -0.5  0.5im],
                  (2,2))
    grape(Ut,Hd,Hc,u_init,t,n)
end


using Test
using GeophysicalModelGenerator

function main_gravity()
############# Input #############
# survey
x        = [-20.,   100.]
y        = [100., 200.]
z        = [-100., 10.]
nx       = 111
ny       = 101
nz       = 91

# spheres
centerX  = [  50.,   10.,   80.]
centerY  = [ 150.,  150.,  180.]
centerZ  = [ -30.,  -10.,  -50.]
radius   = [  10.,    3.,   10.]
rho      = [2650., 2600., 2750.]

# background density model
background = [0. 2700.]
#################################

######### Set things up #########
# constants
G        = 6.67408e-11
mGal     = 1e5

# coordinates
x_vec    = LinRange(x[1],x[2],nx)
y_vec    = LinRange(y[1],y[2],ny)
z_vec    = LinRange(z[1],z[2],nz)
Y,X,Z    = meshgrid(y_vec,x_vec,z_vec)

# reference model
RefMod   = zeros(nz)
RefMod[findall(x->x>=0, z_vec)] .= background[1]
RefMod[findall(x->x<0,  z_vec)] .= background[2]

# densities
RHO      = zeros(nx,ny,nz)
RHO[findall(x->x>=0, Z)] .= background[1]
RHO[findall(x->x<0,  Z)] .= background[2]

# put in spheres
if !(length(centerX) == length(centerY) && length(centerX) == length(centerZ) && length(centerX) == length(radius))
    error("Sphere definition is wrong")
end

numSpheres = length(rho)

for i = 1 : numSpheres
    d = (X .- centerX[i]).^2 .+ (Y .- centerY[i]).^2 .+ (Z .- centerZ[i]).^2
    RHO[findall(x->x<radius[i]^2, d)] .= rho[i]
end
#################################

############ Compute ############
# Analytical Solution
ana = zeros(nx,ny)
for iS = 1 : numSpheres
    for iX = 1 : nx
        for iY = 1 : ny
            d          = ((x_vec[iX] - centerX[iS])^2 + (y_vec[iY] - centerY[iS])^2 + (centerZ[iS])^2)^0.5
            depth      = -centerZ[iS]
            ana[iX,iY] = ana[iX,iY] + (4 * π * G * (radius[iS]^3) * (rho[iS] - background[2]) * depth) / (3 * (d^3)) * mGal
        end
    end
end

# Voxel Code
dg1, gradX1, gradY1 = voxGrav(X, Y, Z, RHO, refMod = RefMod, outName = "Benchmark1", printing = false)

# Other options (same result)
dg2, gradX2, gradY2 = voxGrav(X, Y, Z, RHO, rhoTol = 1, refMod = "SW", outName = "Benchmark2", printing = false)

# Other options (different result)
dg3, gradX3, gradY3 = voxGrav(X, Y, Z, RHO, rhoTol = 70, refMod = "NW", outName = "Benchmark3", printing = false)
#################################

############# Test ##############
maxVal = maximum(broadcast(abs, ana))
maxErr = maxVal / 20

# check differences
diff1 = broadcast(abs, ana-dg1)
diff2 = broadcast(abs, ana-dg2)
diff3 = broadcast(abs, ana-dg3)

# check
@test maximum(diff1) < maxErr
@test maximum(diff2) < maxErr
@test maximum(diff3) > maxErr
#################################
end

main_gravity()


"""
    Catenary(trans::AffineMap, xmin, xmax, a)

Construct a `Catenary` object. The `AffineMap` brings the catenary into a frame
frame in which the `x` axis is aligned with vertically hanging catenary. In this frame,
the `x` values of the catenary lie between `xmin` and `xmax`, the `y` component is zero,
and `z` is parameterized by `x` as:

z = a * (cosh(x / a) - 1)
"""
struct Catenary{Tfm <: AffineTransformation, T <: Real}
    transform::Tfm # transformation from global frame to quadratic's frame
    lmin::T  # "Longitudinal" coordinate at end of catenary
    lmax::T  # "Longitudinal" coordinate at start of catenary
    a::T # Catenary paramter equivalent to 1/2a (such that ax^2+bx+c)
end

"""
    Catenary(Θ, x, y, z, lmin, lmax, a)

Create a catenary from the parameters stored in the Roames database.

### Arguments
* `Θ` : Angle of catenary in XY plane
* `x`: x position of lowest point on catenary
* `y`: y position of lowest point on catenary
* `z`: z position of lowest point on catenary
* `lmin`: Longitudinal coordinate at start of catenary
* `lmax`: Longitudinal coordinate at end of catenary
* `a`: Catenary paramter equivalent to 1/2a (such that ax^2+bx+c)
"""
function Catenary(Θ::T, x::T, y::T, z::T, lmin::T, lmax::T, a::T) where {T}
    trans = inv(Translation(SVector(x, y, z)) ∘ LinearMap(RotZ(Θ)))
    return Catenary(trans, lmin, lmax, a)
end
@inline Catenary(args...) = Catenary(promote(args...)...)

eltype(::Catenary{T}) where {T} = SVector{3,T}
eltype(::Type{<:Catenary{T}}) where {T} = SVector{3,T}

# Use getindex to get points along the line
function getindex(c::Catenary, x)
    inv(c.transform)(SVector(x, 0, c.a * (cosh(x / c.a) - 1)))
end
length(c::Catenary) = c.lmax - c.lmin
#lastindex(c::Catenary) = c.lmax # seems unsafe with `begin`

getindex(c::Catenary, v::AbstractVector) = map(x -> c[x], v)

# Transformations (which compose into AffineMap's)
# Note: Julia's dispatch system makes it difficult to support all `Transformations`.
(trans::Translation{V})(c::Catenary) where {V} = Catenary(c.transform ∘ inv(trans), c.lmin, c.lmax, c.a)
(trans::LinearMap{M})(c::Catenary) where {M} = Catenary(c.transform ∘ inv(trans), c.lmin, c.lmax, c.a)
(trans::AffineMap{M,V})(c::Catenary) where {M,V} = Catenary(c.transform ∘ inv(trans), c.lmin, c.lmax, c.a)

"""
    Quadratic(c::Catenary)

Create a quadratic which approximates `c`. This is an approximate operation, however
the fitted quadratic is tuned for the relevant segment of the catenary and
`Catenary(Quadratic(c))` should be equal to `c` up to numerical error.
"""
function Quadratic(c::Catenary)
    Quadratic(c[c.lmin], c[0.5*(c.lmin + c.lmax)], c[c.lmax])
end

"""
    Catenary(q::Quadratic)

Create a quadratic which approximates `q`. This is an approximate operation, however
the fitted catenary is tuned for the relevant segment of the quadratic and
`Quadratic(Catenary(q))` should be equal to `q` up to numerical error.
"""
function Catenary(q::Quadratic)
    Catenary(q[0], q[0.5*end], q[end])
end

"""
    Catenary(p_start, p_mid, p_end)

Takes three coordinates for the start, somewhere in the middle, and the end of a catenary
and constructs the appropriate `Catenary` object. A numerical fitting method is used to
fit the catenary parameters.
"""
function Catenary(p_start::StaticVector{3,T}, p_mid::StaticVector{3,T}, p_end::StaticVector{3,T}) where {T}
    dir = SVector(p_end[1] - p_start[1], p_end[2] - p_start[2])
    len = norm(dir) # horizontal length
    @static if VERSION < v"0.7"
        θ = atan2(dir[2], dir[1]) # rotation in x-y plane
    else
        θ = atan(dir[2], dir[1]) # rotation in x-y plane
    end
    tform = inv(AffineMap(RotZ(θ), p_start))

    # Now fit a catenary in the x-z plane. Ignore y value of midpoint on the basis that
    # `Catenary` does not swing (unlike `Quadratic`)
    p_mid_trans = tform(p_mid)
    x_mid = p_mid_trans[1]
    z_mid = p_mid_trans[3]

    x_end = len
    z_end = p_end[3] - p_start[3]

    (a, x0, z0) = fit_catenary_origin_2d(x_mid, z_mid, x_end, z_end)

    # Bring transformation to midpoint
    tform2 = Translation(SVector(-x0, zero(T), -z0)) ∘ tform

    return Catenary(tform2, -x0, -x0+len, a)
end

function fit_catenary_origin_2d(x1::T, z1::T, x2::T, z2::T) where {T}
    # Note: one point is going through the origin

    # First guess - make a quadratic approximation: z = q_a*x^2 + q_b*x + q_c
    q_a = (z2/x2 - z1/x1)/(x2 - x1)
    q_b = (x1*z2/x2 - x2*z1/x1)/(x1 - x2)
    # q_c = 0

    # Find the catenary which matches best at the minimum

    # defend against straight catenaries and degeneracies
    if isnan(q_a) || abs(q_a*max(x1,x2)) < 1e-6 # Second term selects cases where quadratic term changes z by less than 1e-6th of it's length
        # It's either flat or degenerate

        # Choose a bottom which is one million times further than points
        q_b = z2/x2 # slope

        # allow x0 to be close when it's (almost?) flat, but force it to be far when slope is large
        x0 = -1e6*z2 # = -1e6*x2*sign(z2)*abs(q_b)
        a = (q_b == zero(T)) ? inv(2*eps(T)) : -x0/q_b

        # Iteratively improve (as below, but optimizing only `a` w.r.t. `z2`)
        for i = 1:10
            Δz₂ = a*(cosh((x2-x0)/a) - cosh(-x0/a)) - z2
            dΔz₂_da = (cosh((x2-x0)/a) - cosh(-x0/a)) - (((x2-x0)/a)*sinh((x2-x0)/a) - (-x0)/a*sinh(-x0/a))
            a = a - dΔz₂_da \ Δz₂
        end

        z0 = a*(1 - cosh(-x0/a))

        if !isfinite(z0)||!isfinite(x0)||!isfinite(a) # On rare occassions this path proves unstable when the standard one succeeds
            # Continue with standard path
            q_b = (x1*z2/x2 - x2*z1/x1)/(x1 - x2)
        else
            return (a, x0, z0)
        end
    end

    a = inv(2*q_a)
    x0 = -q_b * a # -q_b/(2*q_a)
    # z0 = q_a*x0*x0 + q_b*x0 - a

    # Now iteratively improve upon this guess, fixing z0 = a(1-cosh(-x0/a))
    for i = 1:10
        Δz₁ = a*(cosh((x1-x0)/a) - cosh(-x0/a)) - z1
        Δz₂ = a*(cosh((x2-x0)/a) - cosh(-x0/a)) - z2
        dΔz₁_dx0 = -(sinh((x1-x0)/a) - sinh(-x0/a))
        dΔz₂_dx0 = -(sinh((x2-x0)/a) - sinh(-x0/a))
        dΔz₁_da = (cosh((x1-x0)/a) - cosh(-x0/a)) - (((x1-x0)/a)*sinh((x1-x0)/a) - (-x0)/a*sinh(-x0/a))
        dΔz₂_da = (cosh((x2-x0)/a) - cosh(-x0/a)) - (((x2-x0)/a)*sinh((x2-x0)/a) - (-x0)/a*sinh(-x0/a))

        J = @SMatrix [dΔz₁_dx0 dΔz₁_da; dΔz₂_dx0 dΔz₂_da]
        x = @SVector [Δz₁, Δz₂]
        sol = J \ x
        x0 = x0 - sol[1]
        a = a - sol[2]

        if Δz₁*Δz₁ + Δz₂*Δz₂ < eps(T)
            break
        end
    end

    z0 = a*(1 - cosh(-x0/a))

    return (a, x0, z0)
end


using Base.Test
using MIPVerify
using MIPVerify: BlurringPerturbationFamily
isdefined(:TestHelpers) || include("../../../TestHelpers.jl")
using TestHelpers: batch_test_adversarial_example

@testset "FC + Softmax" begin

### Parameters for neural net
batch = 1
in1_height = 8
in1_width = 8
in1_channels = 2

A_height = 5
A_width = batch*in1_height*in1_width*in1_channels

B_height = 3
B_width = A_height

srand(5)
input_size = (batch, in1_height, in1_width, in1_channels)
x0 = rand(input_size)

kernelf1 = rand(-10:10, A_width, A_height)
biasf1 = rand(-10:10, A_height)

kernelf2 = rand(-10:10, B_width, B_height)
biasf2 = rand(-10:10, B_height)

nn = Sequential([Flatten(4),
        Linear(kernelf1, biasf1), ReLU(),
        Linear(kernelf2, biasf2)
    ],
    "tests.integration.generated_weights.fc+softmax"
)

@testset "Blurring perturbations accept multiple input channels." begin
    pp_blur = BlurringPerturbationFamily((5, 5))

    expected_objective_values = Dict(
        (1, pp_blur, Inf, 0) => 0,
        (2, pp_blur, Inf, 0) => 0.1936444023676,
        (3, pp_blur, Inf, 0) => 0.4151904104813
    )

    batch_test_adversarial_example(nn, x0, expected_objective_values)
end

end


# OP1
"""
`faceswith(mesh::Mesh, idx)`

Return a vector `I` of the indices of faces containing vertex with given `idx`.
"""
faceswith(mesh::Mesh, idx) =
  findall(face -> idx in face, mesh.faces)

function faceswith(mesh::HalfEdgeMesh, idx)
  faces = UInt32[]
  itersameorigin(mesh, idx) do edge
    push!(faces, edge.face)
  end
  faces
end

# OP2
"""
`vneighbors(mesh::Mesh, idx; depth=1)`

Return a vector `I` containing the indices of vertices surrounding vertex with given `idx`.
"""
function vneighbors(mesh::Mesh, idx; depth=1)
  neighbors = Set{UInt32}()
  distance = zerobased(fill(depth + 1, length(mesh.vertices)))
  usedfaces = fill(false, length(mesh.faces))
  distance[idx] = 0

  for depth in 1:depth, (i, face) in enumerate(mesh.faces)
    usedfaces[i] && continue
    dist = minimum(v -> distance[v], face)
    if dist < depth
      usedfaces[i] = true
      for v in face
        distance[v] = min(distance[v], dist + 1) # works only for triangles
        push!(neighbors, v)
      end
    end
  end
  delete!(neighbors, idx)

  collect(neighbors)
end

function vneighbors(mesh::HalfEdgeMesh, idx; depth=1)
  neighbors = Set{UInt32}()
  distance = zerobased(fill(depth + 1, length(mesh.vertices)))
  distance[idx] = 0

  queue = Queue{UInt32}()
  enqueue!(queue, idx)
  while !isempty(queue)
    idxv = dequeue!(queue)
    itersameorigin(mesh, idxv) do edge
      for vert in (mesh.halfedges[edge.next].origin, mesh.halfedges[edge.prev].origin)
        if distance[idxv] + 1 < distance[vert]
          distance[vert] = distance[idxv] + 1
          push!(neighbors, vert)
          if distance[vert] != depth
            enqueue!(queue, vert)
          end
        end
      end
    end
  end
  delete!(neighbors, idx)

  collect(neighbors)
end

# OP3
"""
`fneighbors(mesh::Mesh, idx; depth=1)`

Return a vector `I` containing the indices of faces surrounding face with given `idx`.
"""
function fneighbors(mesh::TriangleMesh, idx; depth=1)
  neighbors = Set{UInt32}()
  distance = zerobased(fill(depth + 1, length(mesh.vertices)))
  usedfaces = fill(false, length(mesh.faces))

  distance[face(mesh, idx)] .= 0
  usedfaces[idx] = true

  for depth in 1:depth, (i, face) in enumerate(mesh.faces)
    usedfaces[i] && continue
    dist = sort(map(v -> distance[v], face))
    if dist[2] < depth
      usedfaces[i] = true
      push!(neighbors, i)
      for v in face
        distance[v] = min(distance[v], dist[2] + 1) # works only for triangles
      end
    end
  end
  delete!(neighbors, idx)

  collect(neighbors)
end

function fneighbors(mesh::HalfEdgeMesh, idx; depth=1)
  neighbors = Set{UInt32}()
  distance = fill(depth + 1, length(mesh.faces))
  distance[idx] = 0

  queue = Queue{UInt32}()
  enqueue!(queue, idx)
  while !isempty(queue)
    facev = dequeue!(queue)
    itersameface(mesh, facev) do edge
      edge.twin != 0 || return # next edge
      twin = mesh.halfedges[edge.twin].face
      if distance[facev] + 1 < distance[twin]
        distance[twin] = distance[facev] + 1
        push!(neighbors, twin)
        if distance[twin] != depth
          enqueue!(queue, twin)
        end
      end
    end
  end
  delete!(neighbors, idx)

  collect(neighbors)
end

# OP4
"""
`traverse(mesh::Mesh{3,Point})`
"""
traverse(mesh::Mesh{3,Point}, faceidx, target::Point) where Point =
Channel{Int32}() do channel
  # compute surface norms
  usedfaces = fill(false, length(mesh.faces))

  @inbounds while true
    @label start
    usedfaces[faceidx] && break
    usedfaces[faceidx] = true
    facev = face(mesh, faceidx)
    put!(channel, faceidx) # yield

    # find all edges
    edgepoints = [(vertex(mesh, S), vertex(mesh, T)) for (S, T) in edges(facev)]

    # compute edge norms
    edgenorms = [normalize((T .- S) × facenorm(mesh, faceidx)) for (S, T) in edgepoints]
    # compute directional vecotrs
    dirs = [normalize(target .- midpoint(S, T)) for (S, T) in edgepoints]

    # compute dot products
    dotprods = [e ⋅ d for (e, d) in zip(edgenorms, dirs)] # e,d normalized
    # find smallest angle
    ord = sortperm(dotprods; rev=true) # cosθ₁ < cosθ₂ => θ₁ > θ₂

    dotprods[ord[1]] <= 0 && break # point is inside the current face
    for v in ord
      S, T = facev[[v, v % 3 + 1]]
      # check if edge has neighbouring face
      for newfaceidx in eachindex(mesh.faces)
        (newfaceidx == faceidx || usedfaces[newfaceidx]) && continue
        newface = mesh.faces[newfaceidx]
        if S in newface && T in newface
          faceidx = newfaceidx
          @goto start
        end
      end
    end
  end
end

# OP5
"""
`flipdiag!(mesh::AbstractMesh{N,V,Triangle}, faceidx₁, faceidx₂)`
"""
function flipdiag!(mesh::TriangleMesh, Aidx, Bidx)
  A = face(mesh, Aidx)
  B = face(mesh, Bidx)

  olddiag = A ∩ B
  @assert length(olddiag) == 2 "Face $Aidx should have exacly one edge in common with face $Bidx."

  newdiag = (setdiff(A, olddiag)[1], setdiff(B, olddiag)[1])

  mesh.faces[Aidx] = Triangle(newdiag[1], newdiag[2], olddiag[1])
  mesh.faces[Bidx] = Triangle(olddiag[2], newdiag[2], newdiag[1])
  newdiag
end

# OP4
traverse(mesh::HalfEdgeMesh{3,Point}, faceidx, target::Point) where Point =
Channel{Int32}() do channel
  usedfaces = fill(false, length(mesh.faces))

  @inbounds while true
    usedfaces[faceidx] && break
    usedfaces[faceidx] = true
    halfedges, facev = face(mesh, faceidx; keephalfedges=true)
    put!(channel, faceidx) # yield

    # find all edges
    edgepoints = [(vertex(mesh, S), vertex(mesh, T)) for (S, T) in edges(facev)]

    # compute edge norms
    edgenorms = [normalize((T .- S) × facenorm(mesh, faceidx)) for (S, T) in edgepoints]
    # compute directional vecotrs
    dirs = [normalize(target .- midpoint(S, T)) for (S, T) in edgepoints]

    # compute dot products
    dotprods = [e ⋅ d for (e, d) in zip(edgenorms, dirs)] # e,d normalized
    # find smallest angle
    ord = sortperm(dotprods; rev=true) # cosθ₁ < cosθ₂ => θ₁ > θ₂

    dotprods[ord[1]] <= 0 && break # point is inside the current face
    for v in ord
      edge = halfedges[v]
      # check if edge has neighbouring face
      twin = mesh.halfedges[edge].twin
      twin == 0 && continue
      newfaceidx = mesh.halfedges[twin].face
      (newfaceidx == faceidx || usedfaces[newfaceidx]) && continue
      faceidx = newfaceidx
      break
    end
  end
end

# OP5
function flipdiag!(mesh::HalfEdgeMesh{N,V,Triangle}, Aidx, Bidx) where {N,V}
  # find common edge
  start = ait = mesh.faces[Aidx]
  found = false
  while ait != start
    edge = mesh.halfedges[ait]
    if edge.twin != 0 && mesh.halfedges[edge.twin].face == Bidx
      found = true
      break
    end
    ait = edge.next
  end

  @assert found "Face $Aidx should have exacly one edge in common with face $Bidx."

  # swap
  @inbounds begin
    bit = mesh.halfedges[ait].twin

    Aprev = mesh.halfedges[ait].prev
    Bprev = mesh.halfedges[bit].prev

    Anext = mesh.halfedges[ait].next
    Bnext = mesh.halfedges[bit].next

    mesh.vertexhalfedge[meshes.halfedges[ait].origin] = Bnext
    mesh.vertexhalfedge[meshes.halfedges[bit].origin] = Anext

    mesh.halfedges[ait].origin = mesh.halfedges[Aprev].origin
    mesh.halfedges[bit].origin = mesh.halfedges[Bprev].origin

    A = (Anext, ait, Bprev, Anext, ait)
    B = (Bnext, bit, Aprev, Bnext, bit)

    for i in 2:4
      mesh.halfedges[A[i]].next = A[i+1]
      mesh.halfedges[A[i]].prev = A[i-1]
      mesh.halfedges[B[i]].next = B[i+1]
      mesh.halfedges[B[i]].prev = B[i-1]
    end

    mesh.halfedges[Bprev].face = Aidx
    mesh.halfedges[Aprev].face = Bidx

    mesh.faces[Aidx] = ait
    mesh.faces[Bidx] = bit

    (mesh.halfedges[ait].origin, mesh.halfedges[bit].origin)
  end
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

module TestSVD

using Test, LinearAlgebra, Random
using LinearAlgebra: BlasComplex, BlasFloat, BlasReal, QRPivoted

@testset "Simple svdvals / svd tests" begin
    ≊(x,y) = isapprox(x,y,rtol=1e-15)

    m = [2, 0]
    @test @inferred(svdvals(m)) ≊ [2]
    @test @inferred(svdvals!(float(m))) ≊ [2]
    for sf in (@inferred(svd(m)), @inferred(svd!(float(m))))
        @test sf.S ≊ [2]
        @test sf.U'sf.U ≊ [1]
        @test sf.Vt'sf.Vt ≊ [1]
        @test sf.U*Diagonal(sf.S)*sf.Vt' ≊ m
    end
    F = @inferred svd(m, full=true)
    @test size(F.U) == (2, 2)
    @test F.S ≊ [2]
    @test F.U'F.U ≊ Matrix(I, 2, 2)
    @test F.Vt'*F.Vt ≊ [1]
    @test @inferred(svdvals(3:4)) ≊ [5]
    A = Matrix(1.0I, 2, 2)
    Z = svd(Hermitian(A); full=true)
    @test Z.S ≈ ones(2)
    @test Z.U'Z.U ≈ I(2)

    m1 = [2 0; 0 0]
    m2 = [2 -2; 1 1]/sqrt(2)
    m2c = Complex.([2 -2; 1 1]/sqrt(2))
    @test @inferred(svdvals(m1))  ≊ [2, 0]
    @test @inferred(svdvals(m2))  ≊ [2, 1]
    @test @inferred(svdvals(m2c)) ≊ [2, 1]

    sf1 = @inferred svd(m1)
    sf2 = @inferred svd(m2)
    @test sf1.S ≊ [2, 0]
    @test sf2.S ≊ [2, 1]
    # U & Vt are unitary
    I22 = Matrix(I, 2, 2)
    @test sf1.U*sf1.U'   ≊ I22
    @test sf1.Vt*sf1.Vt' ≊ I22
    @test sf2.U*sf2.U'   ≊ I22
    @test sf2.Vt*sf2.Vt' ≊ I22
    # SVD not uniquely determined, so just test we can reconstruct the
    # matrices from the factorization as expected.
    @test sf1.U*Diagonal(sf1.S)*sf1.Vt' ≊ m1
    @test sf2.U*Diagonal(sf2.S)*sf2.Vt' ≊ m2

    @test ldiv!([0., 0.], svd(Matrix(I, 2, 2)), [1., 1.]) ≊ [1., 1.]
    @test inv(svd(Matrix(I, 2, 2))) ≈ I
    @test inv(svd([1 2; 3 4])) ≈ [-2.0 1.0; 1.5 -0.5]
    @test inv(svd([1 0 1; 0 1 0])) ≈ [0.5 0.0; 0.0 1.0; 0.5 0.0]
    @test_throws SingularException inv(svd([0 0; 0 0]))
    @test inv(svd([1+2im 3+4im; 5+6im 7+8im])) ≈ [-0.5 + 0.4375im 0.25 - 0.1875im; 0.375 - 0.3125im -0.125 + 0.0625im]
end

n = 10

Random.seed!(1234321)

areal = randn(n,n)/2
aimg  = randn(n,n)/2

@testset for eltya in (Float32, Float64, ComplexF32, ComplexF64, Int)
    aa = eltya == Int ? rand(1:7, n, n) : convert(Matrix{eltya}, eltya <: Complex ? complex.(areal, aimg) : areal)
    asym = aa' + aa                 # symmetric indefinite
    for a in (aa, view(aa, 1:n, 1:n))
        usv = svd(a)
        @testset "singular value decomposition" begin
            @test usv.S === svdvals(usv)
            @test usv.U * (Diagonal(usv.S) * usv.Vt) ≈ a
            @test convert(Array, usv) ≈ a
            @test usv.Vt' ≈ usv.V
            @test_throws ErrorException usv.Z
            b = rand(eltya,n)
            @test usv\b ≈ a\b
            @test Base.propertynames(usv) == (:U, :S, :V, :Vt)
            @test size(usv) == size(a)
            if eltya <: BlasFloat
                svdz = svd!(Matrix{eltya}(undef,0,0))
                @test svdz.U ≈ Matrix{eltya}(I, 0, 0)
                @test svdz.S ≈ real(zeros(eltya,0))
                @test svdz.Vt ≈ Matrix{eltya}(I, 0, 0)
            end
        end
        @testset "singular value decomposition of adjoint/transpose" begin
            for transform in (adjoint, transpose)
                usv = svd(transform(a))
                @test usv.S === svdvals(usv)
                @test usv.U * (Diagonal(usv.S) * usv.Vt) ≈ transform(a)
                @test convert(Array, usv) ≈ transform(a)
                @test usv.Vt' ≈ usv.V
                @test_throws ErrorException usv.Z
                b = rand(eltya,n)
                @test usv\b ≈ transform(a)\b
            end
        end
        @testset "Generalized svd" begin
            a_svd = a[1:div(n, 2), :]
            gsvd = svd(a,a_svd)
            @test Base.propertynames(gsvd) == (:alpha, :beta, :vals, :S, :D1, :D2, :R0, :U, :V, :Q, :a, :b, :k, :l, :R)
            @test gsvd.U*gsvd.D1*gsvd.R*gsvd.Q' ≈ a
            @test gsvd.V*gsvd.D2*gsvd.R*gsvd.Q' ≈ a_svd
            @test usv.Vt' ≈ usv.V
            @test_throws ErrorException usv.Z
            @test_throws ErrorException gsvd.Z
            @test gsvd.vals ≈ svdvals(a,a_svd)
            α = eltya == Int ? -1 : rand(eltya)
            β = svd(α)
            @test β.S == [abs(α)]
            @test svdvals(α) == abs(α)
            u,v,q,d1,d2,r0 = svd(a,a_svd)
            @test u ≈ gsvd.U
            @test v ≈ gsvd.V
            @test d1 ≈ gsvd.D1
            @test d2 ≈ gsvd.D2
            @test q ≈ gsvd.Q
            @test gsvd.a.^2 + gsvd.b.^2 ≈ fill(1, length(gsvd.a))
            @test gsvd.alpha.^2 + gsvd.beta.^2 ≈ ones(eltya, length(gsvd.a))
            #testing the other layout for D1 & D2
            b = rand(eltya,n,2*n)
            c = rand(eltya,n,2*n)
            gsvd = svd(b,c)
            @test gsvd.U*gsvd.D1*gsvd.R*gsvd.Q' ≈ b
            @test gsvd.V*gsvd.D2*gsvd.R*gsvd.Q' ≈ c
        end
    end
    @testset "singular value decomposition of Hermitian/real-Symmetric" begin
        for T in (eltya <: Real ? (Symmetric, Hermitian) : (Hermitian,))
            usv = svd(T(asym))
            @test usv.S === svdvals(usv)
            @test usv.U * (Diagonal(usv.S) * usv.Vt) ≈ T(asym)
            @test convert(Array, usv) ≈ T(asym)
            @test usv.Vt' ≈ usv.V
            @test_throws ErrorException usv.Z
            b = rand(eltya,n)
            @test usv\b ≈ T(asym)\b
        end
    end
    if eltya <: LinearAlgebra.BlasReal
        @testset "Number input" begin
            x, y = randn(eltya, 2)
            @test svd(x)    == svd(fill(x, 1, 1))
            @test svdvals(x)    == first(svdvals(fill(x, 1, 1)))
            @test svd(x, y) == svd(fill(x, 1, 1), fill(y, 1, 1))
            @test svdvals(x, y) ≈  first(svdvals(fill(x, 1, 1), fill(y, 1, 1)))
        end
    end
    if eltya != Int
        @testset "isequal, ==, and hash" begin
            x, y   = rand(eltya), convert(eltya, NaN)
            Fx, Fy = svd(x), svd(y)
            @test   Fx == Fx
            @test !(Fy == Fy)
            @test isequal(Fy, Fy)
            @test hash(Fx)          == hash(Fx)
            @test hash(Fx, UInt(1)) == hash(Fx, UInt(1))
            @test hash(Fy)          == hash(Fy)
            @test hash(Fy, UInt(1)) == hash(Fy, UInt(1))
        end
    end
end

@testset "SVD Algorithms" begin
    ≊(x,y) = isapprox(x,y,rtol=1e-15)

    x = [0.1 0.2; 0.3 0.4]

    for alg in [LinearAlgebra.QRIteration(), LinearAlgebra.DivideAndConquer()]
        sx1 = svd(x, alg = alg)
        @test sx1.U * Diagonal(sx1.S) * sx1.Vt ≊ x
        @test sx1.V * sx1.Vt ≊ I
        @test sx1.U * sx1.U' ≊ I
        @test all(sx1.S .≥ 0)

        sx2 = svd!(copy(x), alg = alg)
        @test sx2.U * Diagonal(sx2.S) * sx2.Vt ≊ x
        @test sx2.V * sx2.Vt ≊ I
        @test sx2.U * sx2.U' ≊ I
        @test all(sx2.S .≥ 0)
    end
end

@testset "REPL printing of SVD" begin
    svdd = svd(randn(3, 3))
    svdstring = sprint((t, s) -> show(t, "text/plain", s), svdd)
    ustring = sprint((t, s) -> show(t, "text/plain", s), svdd.U)
    sstring = sprint((t, s) -> show(t, "text/plain", s), svdd.S)
    vtstring = sprint((t, s) -> show(t, "text/plain", s), svdd.Vt)
    @test svdstring == "$(summary(svdd))\nU factor:\n$ustring\nsingular values:\n$sstring\nVt factor:\n$vtstring"
end

@testset "REPL printing of Generalized SVD" begin
    a = randn(3, 3)
    b = randn(3, 3)
    svdd = svd(a, b)
    svdstring = sprint((t, s) -> show(t, "text/plain", s), svdd)
    ustring = sprint((t, s) -> show(t, "text/plain", s), svdd.U)
    qstring = sprint((t, s) -> show(t, "text/plain", s), svdd.Q)
    vstring = sprint((t, s) -> show(t, "text/plain", s), svdd.V)
    d1string = sprint((t, s) -> show(t, "text/plain", s), svdd.D1)
    d2string = sprint((t, s) -> show(t, "text/plain", s), svdd.D2)
    r0string = sprint((t, s) -> show(t, "text/plain", s), svdd.R0)
    @test svdstring == "$(summary(svdd))\nU factor:\n$ustring\nV factor:\n$vstring\nQ factor:\n$qstring\nD1 factor:\n$d1string\nD2 factor:\n$d2string\nR0 factor:\n$r0string"
end

@testset "c-tor with varying input eltypes" begin
    A = randn(Float64, 10, 10)
    U, S, V = svd(A)
    Ut = convert.(Float16, U)
    Vt = convert.(Float32, V)
    svdc = SVD{ComplexF32}(Ut, S, Vt)
    @test svdc isa SVD{ComplexF32}
    Uc, Sc, Vc = svdc
    @test Uc * diagm(0=>Sc) * transpose(V) ≈ complex.(A) rtol=1e-3
end

@testset "adjoint of SVD" begin
    n = 5
    B = randn(5, 2)

    @testset "size(b)=$(size(b))" for b in (B[:, 1], B)
        @testset "size(A)=$(size(A))" for A in (
            randn(n, n),
            # Wide problems become minimum norm (in x) problems similarly to LQ
            randn(n + 2, n),
            randn(n - 2, n),
            complex.(randn(n, n), randn(n, n)))

            F = svd(A)
            x = F'\b
            @test x ≈ A'\b
            @test length(size(x)) == length(size(b))
        end
    end
end

end # module TestSVD


#=
Provides a type called RBLQ for solving robust linear quadratic control
problems.

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

@date : 2014-08-19

References
----------

https://lectures.quantecon.org/jl/robustness.html
=#

@doc doc"""
Represents infinite horizon robust LQ control problems of the form

```math
    \min_{u_t} \sum_t \beta^t {x_t' R x_t + u_t' Q u_t }
```

subject to

```math
    x_{t+1} = A x_t + B u_t + C w_{t+1}
```

and with model misspecification parameter ``\theta``.

##### Fields

- `Q::Matrix{Float64}` :  The cost(payoff) matrix for the controls. See above for more. ``Q`` should be `k x k` and symmetric and positive definite
- `R::Matrix{Float64}` :  The cost(payoff) matrix for the state. See above for more. ``R`` should be `n x n` and symmetric and non-negative definite
- `A::Matrix{Float64}` :  The matrix that corresponds with the state in the state space system. ``A`` should be `n x n`
- `B::Matrix{Float64}` :  The matrix that corresponds with the control in the state space system.  ``B`` should be `n x k`
- `C::Matrix{Float64}` :  The matrix that corresponds with the random process in the state space system. ``C`` should be `n x j`
- `beta::Real` : The discount factor in the robust control problem
- `theta::Real` The robustness factor in the robust control problem
- `k, n, j::Int` : Dimensions of input matrices

"""
mutable struct RBLQ
    A::Matrix
    B::Matrix
    C::Matrix
    Q::Matrix
    R::Matrix
    k::Int
    n::Int
    j::Int
    bet::Real
    theta::Real
end

function RBLQ(Q::ScalarOrArray, R::ScalarOrArray, A::ScalarOrArray,
              B::ScalarOrArray, C::ScalarOrArray, bet::Real, theta::Real)
    k = size(Q, 1)
    n = size(R, 1)
    j = size(C, 2)

    # coerce sizes
    A = reshape([A;], n, n)
    B = reshape([B;], n, k)
    C = reshape([C;], n, j)
    R = reshape([R;], n, n)
    Q = reshape([Q;], k, k)
    RBLQ(A, B, C, Q, R, k, n, j, bet, theta)
end

@doc doc"""
The ``D`` operator, mapping ``P`` into

```math
    D(P) := P + PC(\theta I - C'PC)^{-1} C'P
```

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type
- `P::Matrix{Float64}` : size is `n x n`

##### Returns

- `dP::Matrix{Float64}` : The matrix ``P`` after applying the ``D`` operator

"""
function d_operator(rlq::RBLQ, P::Matrix)
    C, theta, Im = rlq.C, rlq.theta, Matrix(I, rlq.j, rlq.j)
    S1 = P*C
    dP = P + S1*((theta.*Im - C'*S1) \ (S1'))

    return dP
end

@doc doc"""
The ``D`` operator, mapping ``P`` into

```math
    B(P) := R - \beta^2 A'PB(Q + \beta B'PB)^{-1}B'PA + \beta A'PA
```

and also returning

```math
    F := (Q + \beta B'PB)^{-1} \beta B'PA
```

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type
- `P::Matrix{Float64}` : size is `n x n`

##### Returns

- `F::Matrix{Float64}` : The ``F`` matrix as defined above
- `new_p::Matrix{Float64}` : The matrix ``P`` after applying the ``B`` operator

"""
function b_operator(rlq::RBLQ, P::Matrix)
    A, B, Q, R, bet = rlq.A, rlq.B, rlq.Q, rlq.R, rlq.bet

    S1 = Q + bet.*B'*P*B
    S2 = bet.*B'*P*A
    S3 = bet.*A'*P*A

    F = S1 \ S2
    new_P = R - S2'*F + S3

    return F, new_P
end

@doc doc"""
Solves the robust control problem.

The algorithm here tricks the problem into a stacked LQ problem, as described in
chapter 2 of Hansen- Sargent's text "Robustness".  The optimal control with
observed state is

```math
    u_t = - F x_t
```

And the value function is ``-x'Px``

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type

##### Returns

- `F::Matrix{Float64}` : The optimal control matrix from above
- `P::Matrix{Float64}` : The positive semi-definite matrix defining the value function
- `K::Matrix{Float64}` : the worst-case shock matrix ``K``, where ``w_{t+1} = K x_t`` is the worst case shock

"""
function robust_rule(rlq::RBLQ)
    A, B, C, Q, R = rlq.A, rlq.B, rlq.C, rlq.Q, rlq.R
    bet, theta, k, j = rlq.bet, rlq.theta, rlq.k, rlq.j

    # Set up LQ version
    # I = eye(j)
    Z = zeros(k, j)
    Ba = [B C]
    Qa = [Q Z
          Z' -bet.*I.*theta]
    lq = QuantEcon.LQ(Qa, R, A, Ba, bet=bet)

    # Solve and convert back to robust problem
    P, f, d = stationary_values(lq)
    F = f[1:k, :]
    K = -f[k+1:end, :]

    return F, K, P
end

@doc doc"""
Solve the robust LQ problem

A simple algorithm for computing the robust policy ``F`` and the
corresponding value function ``P``, based around straightforward
iteration with the robust Bellman operator.  This function is
easier to understand but one or two orders of magnitude slower
than `self.robust_rule()`.  For more information see the docstring
of that method.

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type
- `P_init::Matrix{Float64}(zeros(rlq.n, rlq.n))` : The initial guess for the
value function matrix
- `;max_iter::Int(80)`: Maximum number of iterations that are allowed
- `;tol::Real(1e-8)` The tolerance for convergence

##### Returns

- `F::Matrix{Float64}` : The optimal control matrix from above
- `P::Matrix{Float64}` : The positive semi-definite matrix defining the value function
- `K::Matrix{Float64}` : the worst-case shock matrix ``K``, where ``w_{t+1} = K x_t`` is the worst case shock

"""
function robust_rule_simple(rlq::RBLQ,
                            P::Matrix=zeros(Float64, rlq.n, rlq.n);
                            max_iter=80,
                            tol=1e-8)
    # Simplify notation
    A, B, C, Q, R = rlq.A, rlq.B, rlq.C, rlq.Q, rlq.R
    bet, theta, k, j = rlq.bet, rlq.theta, rlq.k, rlq.j
    iterate, e = 0, tol + 1.0

    F = similar(P)  # instantiate so available after loop

    while iterate <= max_iter && e > tol
        F, new_P = b_operator(rlq, d_operator(rlq, P))
        e = sqrt(sum((new_P - P).^2))
        iterate += 1
        copyto!(P, new_P)
    end

    if iterate >= max_iter
        @warn("Maximum iterations in robust_rul_simple")
    end

    # I = eye(j)
    K = (theta.*I - C'*P*C)\(C'*P)*(A - B*F)

    return F, K, P
end

@doc doc"""
Compute agent 2's best cost-minimizing response ``K``, given ``F``.

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type
- `F::Matrix{Float64}`: A `k x n` array representing agent 1's policy

##### Returns

- `K::Matrix{Float64}` : Agent's best cost minimizing response corresponding to ``F``
- `P::Matrix{Float64}` : The value function corresponding to ``F``

"""
function F_to_K(rlq::RBLQ, F::Matrix)
    # simplify notation
    R, Q, A, B, C = rlq.R, rlq.Q, rlq.A, rlq.B, rlq.C
    bet, theta = rlq.bet, rlq.theta

    # set up lq
    Q2 = bet * theta
    R2 = - R - F'*Q*F
    A2 = A - B*F
    B2 = C
    lq = QuantEcon.LQ(Q2, R2, A2, B2, bet=bet)

    neg_P, neg_K, d = stationary_values(lq)

    return -neg_K, -neg_P
end

@doc doc"""
Compute agent 1's best cost-minimizing response ``K``, given ``F``.

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type
- `K::Matrix{Float64}`: A `k x n` array representing the worst case matrix

##### Returns

- `F::Matrix{Float64}` : Agent's best cost minimizing response corresponding to ``K``
- `P::Matrix{Float64}` : The value function corresponding to ``K``

"""
function K_to_F(rlq::RBLQ, K::Matrix)
    R, Q, A, B, C = rlq.R, rlq.Q, rlq.A, rlq.B, rlq.C
    bet, theta = rlq.bet, rlq.theta

    A1, B1, Q1, R1 = A+C*K, B, Q, R-bet*theta.*K'*K
    lq = QuantEcon.LQ(Q1, R1, A1, B1, bet=bet)

    P, F, d = stationary_values(lq)

    return F, P
end

@doc doc"""
Given ``K`` and ``F``, compute the value of deterministic entropy, which is
``\sum_t \beta^t x_t' K'K x_t`` with ``x_{t+1} = (A - BF + CK) x_t``.

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type
- `F::Matrix{Float64}` The policy function, a `k x n` array
- `K::Matrix{Float64}` The worst case matrix, a `j x n` array
- `x0::Vector{Float64}` : The initial condition for state

##### Returns

- `e::Float64` The deterministic entropy

"""
function compute_deterministic_entropy(rlq::RBLQ, F, K, x0)
    B, C, bet = rlq.B, rlq.C, rlq.bet
    H0 = K'*K
    C0 = zeros(Float64, rlq.n, 1)
    A0 = A - B*F + C*K
    return var_quadratic_sum(A0, C0, H0, bet, x0)
end

@doc doc"""
Given a fixed policy ``F``, with the interpretation ``u = -F x``, this function
computes the matrix ``P_F`` and constant ``d_F`` associated with discounted cost
``J_F(x) = x' P_F x + d_F``.

##### Arguments

- `rlq::RBLQ`: Instance of `RBLQ` type
- `F::Matrix{Float64}` :  The policy function, a `k x n` array

##### Returns

- `P_F::Matrix{Float64}` : Matrix for discounted cost
- `d_F::Float64` : Constant for discounted cost
- `K_F::Matrix{Float64}` : Worst case policy
- `O_F::Matrix{Float64}` : Matrix for discounted entropy
- `o_F::Float64` : Constant for discounted entropy

"""
function evaluate_F(rlq::RBLQ, F::Matrix)
    R, Q, A, B, C = rlq.R, rlq.Q, rlq.A, rlq.B, rlq.C
    bet, theta, j = rlq.bet, rlq.theta, rlq.j

    # Solve for policies and costs using agent 2's problem
    K_F, P_F = F_to_K(rlq, F)
    # I = eye(j)
    H = inv(I - C'*P_F*C./theta)
    d_F = log(det(H))

    # compute O_F and o_F
    sig = -1.0 / theta
    AO = sqrt(bet) .* (A - B*F + C*K_F)
    O_F = solve_discrete_lyapunov(AO', bet*K_F'*K_F)
    ho = (tr(H .- 1) - d_F) / 2.0
    trace = tr(O_F*C*H*C')
    o_F = (ho + bet*trace) / (1 - bet)

    return K_F, P_F, d_F, O_F, o_F
end


module TestKernelOptimization
using Test
using LinearAlgebra
using GaussianDistributions
using GaussianDistributions: nlml, optimize, optimize!

@testset "kernel optimization" begin
    function k(θ)
        a, l = exp.(θ)
        (x, y) -> a*exp(-(x-y)^2/(2l^2))
    end
    f(x) = sin(x)
    n = 8
    x = randn(n)
    y = f.(x)
    θ = fill(-2., 2) # WARNING: without noise covariance, this has to lead to full rank system
    k0 = k(θ)
    optimize!(k, θ, x, y)
    @test nlml(k(θ), x, y) < nlml(k0, x, y)

    # using Plots
    # G = Gaussian(k(θ))
    # C = G | (x, y)
    # plot(C)
    # scatter!(x, y)
    # gui()

    # with noisy observations
    σ = 5e-2
    @. y += σ * randn()

    θ = fill(-2., 2) # WARNING: without noise covariance, this has to lead to full rank system
    k0 = k(θ)
    θ = optimize(k, θ, x, y, σ^2)
    @test nlml(k(θ), x, y, σ^2) < nlml(k0, x, y, σ^2)

    # G = Gaussian(k(θ))
    # C = G | (x, y, σ^2)
    # plot(C)
    # scatter!(x, y)
    # gui()
end

@testset "noise variance optimization" begin
    l = .5
    k(x, y) = exp(-(x-y)^2/(2l^2))
    f(x) = sin(2π*x)
    n = 128
    x = randn(n)
    y = f.(x)
    σ_truth = 5e-2     # with noisy observations
    @. y += σ_truth * randn()

    σ²_new = .1
    σ²_new = optimize(σ²_new, k, x, y) # optimizing noise variance
    @test nlml(k, x, y, σ²_new) < nlml(k, x, y, σ_truth^2)
    @test isapprox(σ²_new, σ_truth^2, rtol = .5)

    # using Plots
    # G = Gaussian(k)
    # C = G | (x, y, σ²_new)
    # plot(C)
    # scatter!(x, y)
    # gui()
end

end # TestKernelOptimization


"""Generate documents with LDA and HLDA
"""

module LDA
using Distributions
using StatsBase
using DirichletProcess

vocabulary = ["see", "spot", "run"]
n_terms = length(vocabulary)
mean_document_length = 5
term_dirichlet_parameter = 1
topic_dirichlet_parameter = 1

mutable struct Document
    topics::Array{Vector{Float64}}
    words::Array{String}
    len::Int
    Document() = new([], [], 0)
end

function add_word!(doc::Document, word, topic)
    push!(doc.topics, topic)
    push!(doc.words, word)
    doc.len += 1
    return nothing
end

function print_topic_counts(documents)
    for (i, doc) in enumerate(documents)
        println("Doc: $(i)")
        for (topic, count) in countmap(doc.topics)
            println("      count: $(count) | topic: $(round.(topic, digits=2))")
        end
    end
end

function finite(n_documents)
    n_topics = 2

    term_dirichlet_vector = ones(n_terms) * term_dirichlet_parameter
    term_distributions = [rand(Dirichlet(term_dirichlet_vector)) for _ in 1:n_topics]
    base_distribution = () -> rand(term_distributions)

    documents = [Document() for _ in 1:n_documents]

    for doc in documents
        topic_distribution = DP(base_distribution, topic_dirichlet_parameter)
        document_length = rand(Poisson(mean_document_length))
        for _ in 1:document_length
            p = dp_sample!(topic_distribution)
            word = vocabulary[rand(Categorical(p))]
            add_word!(doc, word, p)
        end
        println(doc.words)
    end

    print_topic_counts(documents)

    return nothing
end

function infinite(n_documents)
    base_dp_parameter = 10
    base_distribution = () -> rand(Dirichlet(term_dirichlet_parameter * ones(n_terms)))
    base_dp = DP(base_distribution, base_dp_parameter)

    documents = [Document() for _ in 1:n_documents]

    for doc in documents
        topic_distribution = DP(() -> dp_sample!(base_dp), topic_dirichlet_parameter)
        document_length = rand(Poisson(mean_document_length))
        for _ in 1:document_length
            p = dp_sample!(topic_distribution)
            word = vocabulary[rand(Categorical(p))]
            add_word!(doc, word, p)
        end
        println(doc.words)
    end

    print_topic_counts(documents)

    return nothing
end

end


using RecursiveSparseBlocks

function run(T)
    out = open("benchmark.csv", "w")
    println(out, "method,n,p,t")

    N = 100
    for (maxn, p) in [(1000000, 1e-4)]
        for _ in 1:N
            n = rand(1:maxn)
            A = sprand(T, n, n, p)
            Arsb = SparseMatrixRSB(A)
            x = rand(T, n)

            y = A*x
            tic()
            y = A*x
            t = toc()
            println(out, "SparseMatrixCSC,", n, ",", p, ",", t)

            y = Arsb*x
            tic()
            y = Arsb*x
            t = toc()
            println(out, "SparseMatrixRSB,", n, ",", p, ",", t)
        end
    end

    close(out)
end

run(Float64)

#@time bench_csc()
#@time bench_csc()
#@profile bench_csc()
#Profile.print(C=true)

#@time bench_rsb()
#@time bench_rsb()
#@profile bench_rsb()
#Profile.print(C=true)


struct Knnr
    X::Array{Float64}
    Y::Array{Float64}
    fm
    nlvdis::Int
    metric::String
    h::Real
    k::Int
    tol::Real
end

"""
    knnr(X, Y; nlvdis = 0, metric = "eucl", h = Inf, k = 1, tol = 1e-4)
k-Nearest-Neighbours regression (KNNR).
* `X` : X-data.
* `Y` : Y-data.
* `nlvdis` : Number of latent variables (LVs) to consider in the global PLS used for the dimension reduction before
    calculating the dissimilarities. If `nlvdis = 0`, there is no dimension reduction.
* `metric` : Type of dissimilarity used to select the neighbors. Possible values are "eucl" (default; Euclidean distance)
    and "mahal" (Mahalanobis distance).
* `h` : A scalar defining the shape of the weight function. Lower is h, sharper is the function. See function `wdist`.
* `k` : The number of nearest neighbors to select for each observation to predict.
* `tol` : For stabilization when very close neighbors.

For each new observation to predict, the prediction is the weighted mean
over the selected (training) neighborhood. The weights are defined from
the dissimilarities between the new observation and the selected neighborhood,
and are computed from function 'wdist'.

The dissimilarities can be calculated from the original X-data
or after a dimension reduction, depending on argument `nlvdis`.
In the last case, global PLS scores (LVs) are computed from (`X`, `Y`)
and the dissimilarities are computed over these scores.
In general, for high dimensional X-data, using the Mahalanobis distance requires
psuch a reliminary dimensionality reduction.
"""
function knnr(X, Y; nlvdis = 0, metric = "eucl", h = Inf, k = 1, tol = 1e-4)
    X = ensure_mat(X)
    Y = ensure_mat(Y)
    if nlvdis == 0
        fm = nothing
    else
        fm = plskern(X, Y; nlv = nlvdis)
    end
    return Knnr(X, Y, fm, nlvdis, metric, h, k, tol)
end

"""
    predict(object::Knnr, X)
Compute the Y-predictions from the fitted model.
* `object` : The fitted model.
* `X` : X-data for which predictions are computed.
"""
function predict(object::Knnr, X)
    X = ensure_mat(X)
    m = size(X, 1)
    q = size(object.Y, 2)
    # Getknn
    if isnothing(object.fm)
        res = getknn(object.X, X; k = object.k, metric = object.metric)
    else
        res = getknn(object.fm.T, transform(object.fm, X); k = object.k, metric = object.metric)
    end
    listw = copy(res.d)
    for i = 1:m
        w = wdist(res.d[i]; h = object.h)
        w[w .< object.tol] .= object.tol
        listw[i] = w
    end
    # End
    pred = zeros(m, q)
    @inbounds for i = 1:m
        s = res.ind[i]
        w = mweights(listw[i])
        pred[i, :] .= colmeans(vrow(object.Y, s), w)
    end
    (pred = pred, listnn = res.ind, listd = res.d, listw = listw)
end