# this should be run after running GovDebt.jl
# targeted moments:
varlinc_target  = 0.7  # var(log(annual income))
vardlinc_target = 0.23 # var(one year changes in log(annual income))
mpc_target      = 0.16
pc0_target      = 0.10 # fraction with zero bond holdings
# parameters: pLH, pHL, zlo, sH/sL. We normalize so s and z both have mean 1.

using Distributions
#using PyPlot
using Roots
using JLD
using Optim

function mollifier(z::AbstractFloat,ehi::AbstractFloat,elo::AbstractFloat)
    # mollifier function
    In = 0.443993816237631
    if z<ehi && z>elo
        temp = -1.0 + 2.0*(z-elo)/(ehi-elo)
        out = (2.0/(ehi-elo))*exp(-1.0/(1.0-temp*temp))/In
    else
        out = 0.0
    end

    return out
end

ni = 10000 # number of individuals in the sample we use for simulating skill process
uz = rand(ni,8) # uniform random variable used to get transitory shock z; need 8 quarters to get 1 year change
us = rand(ni,8) # uniform random variable used to get persistent markov chain shock s

# we will create a single z sample which has zlo=0, zhi = 2. We will then rescale it as we change zlo.
nz = 1000
zgrid = collect(linspace(0.,2.,nz))
zprob = zeros(nz)
for i=1:nz
	zprob[i] = 2*mollifier(zgrid[i],2.,0.)/nz
end
zprob = zprob/sum(zprob)
zcdf = cumsum(zprob)

zave = 0.5*zgrid[1:nz-1]+0.5*zgrid[2:nz]

function zsample(uz::Array{Float64,2}, zgrid::AbstractArray, zcdf::AbstractArray, ni::Int, nz::Int)
	zs = zeros(ni,8)
	for i=1:ni
		for t=1:8
			for iz=1:nz-1
				if uz[i,t]>zcdf[iz] && uz[i,t] <= zcdf[iz+1]
					zs[i,t] = zave[iz]
				end
			end
		end
	end
	return zs
end
tic()
zs = zsample(uz,zgrid,zcdf,ni,nz)
toc()

function ssample(us::Array{Float64,2}, P::Array{Float64,2}, πss::AbstractArray, ni::Int)
	shist = ones(Int,ni,8)
	for i=1:ni
		if us[i,1] > πss[1]
			shist[i,1] = 2
		end
		for t=2:8
			if us[i,t] > P[shist[i,t-1],1]
				shist[i,t] = 2
			end
		end
	end
	return shist
end

function ln_annual_inc(zhist::Array{Float64,2}, us::Array{Float64,2}, zlo::AbstractFloat, P::Array{Float64,2}, πss::AbstractArray, sgrid::AbstractArray, ni::Int)
  	s_inds = ssample(us,P,πss,ni)
	linc1 = zeros(ni)
	linc2 = zeros(ni)
	for i=1:ni
		inc1 = 0.
		for t=1:4
#            @show zhist[i, t]-1
#            @show zlo
#            @show sgrid[1]
#            @show sgrid[2]
#            @show s_inds[i,t]-1
			zshock = 1. + (1. - zlo)*(zhist[i,t]-1.)
			sshock = (sgrid[1] + (sgrid[2] - sgrid[1])*(s_inds[i,t] - 1.))
#            @show zshock
#            @show sshock
			inc1 += zshock*sshock
		end
		inc2 = 0.
		for t=5:8
			zshock = 1. + (1. - zlo)*(zhist[i,t]-1.)
			sshock = (sgrid[1] + (sgrid[2] - sgrid[1])*(s_inds[i,t] - 1.))
			inc2 += zshock*sshock
		end
		linc1[i] += log(inc1)
     	linc2[i] += log(inc2)
	end
	return (linc1, linc2)
end

function skill_moments(sH_by_sL::Real, zlo::Real, pLH::AbstractFloat,
                       pHL::AbstractFloat, zhist::Array{Float64,2},
                       us::Array{Float64,2}, ni::Int = 10000)
	πL  = pHL / (pLH+pHL)
	πss = [πL ; 1-πL]
	P   = [[1.0 -pLH pLH] ; [pHL 1.0-pHL]]
	slo = 1.0 / (πL + (1-πL) * sH_by_sL)
	shi = sH_by_sL * slo
	sgrid = [slo; shi]
	(linc1, linc2) = ln_annual_inc(zs, us, zlo, P, πss, sgrid, ni)
	return (var(linc1), var(linc2 - linc1))
end

tic()
(varlinc, vardlinc) = skill_moments(1.2, 0.5, 0.1, 0.1, zs, us)
toc()

function best_fit(pLH::AbstractFloat, pHL::AbstractFloat, varlinc_target::AbstractFloat,
                  vardlinc_target::AbstractFloat, sH_by_sL_grid::AbstractArray,
                  zlo_grid::AbstractArray, n1::Int, n2::Int, zhist::Array{Float64,2},
                  us::Array{Float64,2}, ni::Int = 10000)
	n = n1*n2
	dist = zeros(n)
	varlinc = zeros(n)
	vardlinc = zeros(n)
	for i1=1:n1
		for i2=1:n2
			i = n2*(i1-1)+i2
			(varlinc[i],vardlinc[i]) = skill_moments(sH_by_sL_grid[i1], zlo_grid[i2], pLH, pHL, zs, us)
			dist[i] = abs(varlinc[i] - varlinc_target) + abs(vardlinc[i] - vardlinc_target)
		end
	end
	(mindist,imin) = findmin(dist)
	i2 = mod(imin-1,n2)+1
	i1 = Int((imin-i2)/n2 + 1)
	return (sH_by_sL_grid[i1], zlo_grid[i2], varlinc[imin], vardlinc[imin])
end

np = 100
sH_by_sL_grid = linspace(3.,9.,np)
zlo_grid      = linspace(0.+eps(), 0.8-eps(), np)
# tic()
# varlinc_target  = 0.7
# vardlinc_target = 0.23
# (sH_by_sL, zlo, varlinc, vardlinc) = best_fit(0.02,0.03,varlinc_target,vardlinc_target,sH_by_sL_grid, zlo_grid,np,np,zs,us)
# toc()

# out = [varlinc varlinc_target vardlinc vardlinc_target sH_by_sL zlo]

"""
THIS ONE GETS CALLED
"""
function best_fit(pLH::AbstractFloat, pHL::AbstractFloat, varlinc_target::AbstractFloat, vardlinc_target::AbstractFloat, sH_by_sL_grid::AbstractArray, zlo_grid::AbstractArray,
				  n1::Int, n2::Int, zhist::Array{Float64,2}, us::Array{Float64,2}, T::Int = 8, ni::Int = 10000)
	n = n1*n2
	dist = zeros(n)
	varlinc = zeros(n)
	vardlinc = zeros(n)
	for i1=1:n1
		for i2=1:n2
			i = n2*(i1-1)+i2
			(varlinc[i], vardlinc[i]) = skill_moments(sH_by_sL_grid[i1], zlo_grid[i2], pLH, pHL, zs, us)
			dist[i] = abs(varlinc[i] - varlinc_target) + abs(vardlinc[i] - vardlinc_target)
		end
	end
	(mindist,imin) = findmin(dist)
	i2 = mod(imin-1,n2)+1
	i1 = Int((imin-i2)/n2 + 1)
	return (sH_by_sL_grid[i1], zlo_grid[i2], varlinc[imin], vardlinc[imin])
end

"""

"""
loss(x::Vector{Float64}, target::Vector{Float64}) = sum(abs.(x-target))

function best_fit(pLH::T, pHL::T, target::Vector{T}, lower::Vector{T},
                  upper::Vector{T}, us::Array{Float64,2}, max_iter::Int = 300,
                  initial_guess::Vector{Float64} = [6.3, 0.03]) where T<:AbstractFloat

    skill_moments_f(x) = loss(collect(skill_moments(x[1], x[2], pLH, pHL, zs, us)), target)

    res = optimize(skill_moments_f, lower, upper, initial_guess, Fminbox(NelderMead()),
                   Optim.Options(f_calls_limit = max_iter))

    sH_by_sL_argmin, zlo_argmin = Optim.minimizer(res)
    min_varlinc, min_vardlinc   = skill_moments(sH_by_sL_argmin, zlo_argmin, pLH, pHL, zs, us)

    return (sH_by_sL_argmin, zlo_argmin, min_varlinc, min_vardlinc)
end

function ave_mpc(m::AbstractArray, c::AbstractArray, agrid::AbstractArray, aswts::AbstractArray, na::Int, ns::Int)
	mpc = 0.
	for iss=1:ns
		i = na*(iss-1)+1
		mpc += aswts[i]*m[i]*(c[i+1] - c[i])/(agrid[2] - agrid[1])
		for ia=2:na
			i = na*(iss-1)+ia
			mpc += aswts[i]*m[i]*(c[i] - c[i-1])/(agrid[ia] - agrid[ia-1])
		end
	end
	return mpc
end

function frac_zero(m::AbstractArray, c::AbstractArray, agrid::AbstractArray, aswts::AbstractArray, ns::Int)
	return sum(aswts.*m.*(c.==repeat(agrid,ns)))
end

nLH = 2
nHL = 2

pLH_grid = linspace(0.01,0.1,nLH)
pHL_grid = linspace(0.01,0.1,nHL)
nn = nLH*nHL
sH_by_sL_grid2 = zeros(nn)
zlo_grid2 = zeros(nn)
varlinc_grid = zeros(nn)
vardlinc_grid = zeros(nn)
mpc_grid  = zeros(nn)
pc0_grid  = zeros(nn)

# RECA
sH_by_sL_lo = 3.0
sH_by_sL_hi = 9.0
zlo_lo      = 1e-18
zlo_hi      = 0.8-eps()

lower = [sH_by_sL_lo, zlo_lo]
upper = [sH_by_sL_hi, zlo_hi]
target = [varlinc_target, vardlinc_target]

for iLH = 1:nLH
	for iHL = 1:nHL
		i = nHL*(iLH-1)+iHL

        @time out = best_fit(pLH_grid[iLH], pHL_grid[iHL], target, lower, upper, us)
        @show "Old Method"
		@time (sH_by_sL_grid2[i], zlo_grid2[i], varlinc_grid[i], vardlinc_grid[i]) = best_fit(pLH_grid[iLH],pHL_grid[iHL],varlinc_target,vardlinc_target,sH_by_sL_grid, zlo_grid,np,np,zs,us)
		(f, sgrid, swts) = persistent_skill_process(sH_by_sL_grid2[i], pLH_grid[iLH], pHL_grid[iHL], ns)
		(agrid, awts) = cash_grid(sgrid, ω, H, r, η, γ, T, zlo, na)
		aswts = kron(swts, awts)
		qfunction(x) = mollifier(x, 2. - zlo_grid[i], zlo_grid[i])
		println("parameters:")
		println([sH_by_sL_grid2[i] zlo_grid2[i] pLH_grid[iLH] pHL_grid[iHL]])
		(ell, c, m, β, report) = findss(na, ns, βlo, βhi, R, ω, H, η, T, γ, bg, qfunction, agrid, sgrid,aswts, Win, f)
		mpc_grid[i] = ave_mpc(m,c,agrid,aswts,na,ns)
		pc0_grid[i] = frac_zero(m,c,agrid,aswts,ns)
	end
end


import DataStructures: OrderedDict

@testset "bernoulli" begin

    # random
    x = bernoulli(0.5)

    # logpdf_grad
    f = (x::Bool, prob::Float64) -> logpdf(bernoulli, x, prob)
    args = (false, 0.3,)
    actual = logpdf_grad(bernoulli, args...)
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    args = (true, 0.3,)
    actual = logpdf_grad(bernoulli, args...)
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
end

@testset "beta" begin

    # random
    x = beta(0.5, 0.5)
    @test 0 < x < 1

    # out of support
    @test logpdf(beta, -1, 0.5, 0.5) == -Inf

    # avoid infinities, sanely
    @test logpdf(beta, eps(typeof(0.)), 0.5, 0.5) < logpdf(beta, 0., 0.5, 0.5) < Inf
    @test logpdf(beta, eps(typeof(0.)), 0.5, 1.5) < logpdf(beta, 0., 0.5, 1.5) < Inf
    @test logpdf(beta, 1-eps(typeof(0.)), 0.5, 0.5) < logpdf(beta, 1., 0.5, 0.5) < Inf
    @test logpdf(beta, 1-eps(typeof(0.)), 1.5, 0.5) < logpdf(beta, 1., 1.5, 0.5) < Inf

    # logpdf_grad
    f = (x, alpha, beta_param) -> logpdf(beta, x, alpha, beta_param)
    args = (0.4, 0.2, 0.3)
    actual = logpdf_grad(beta, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "categorical" begin

    # random
    x = categorical([0.2, 0.3, 0.5])
    @test 0 < x < 4

    # out of support
    @test logpdf(categorical, -1, [0.2, 0.3, 0.5]) == -Inf

    # integer-1 probability bug
    N = 20
    @test sum([rand(Gen.Distributions.Categorical([1, 0])) for i in 1:N]) == N

    # logpdf_grad
    f = (x, probs) -> logpdf(categorical, x, probs)
    args = (2, [0.2, 0.3, 0.5])
    actual = logpdf_grad(categorical, args...)
    @test actual[1] == nothing
    @test isapprox(actual[2][1], finite_diff_vec(f, args, 2, 1, dx))
    @test isapprox(actual[2][2], finite_diff_vec(f, args, 2, 2, dx))
    @test isapprox(actual[2][3], finite_diff_vec(f, args, 2, 3, dx))
end

@testset "gamma" begin

    # random
    x = gamma(1, 1)
    @test 0 < x

    # out of support
    @test logpdf(gamma, -1, 1, 1) == -Inf

    # logpdf_grad
    f = (x, shape, scale) -> logpdf(gamma, x, shape, scale)
    args = (0.4, 0.2, 0.3)
    actual = logpdf_grad(gamma, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "inv_gamma" begin

    # random
    x = inv_gamma(1, 1)
    @test 0 < x

    # out of support
    @test logpdf(inv_gamma, -1, 1, 1) == -Inf

    # logpdf_grad
    f = (x, shape, scale) -> logpdf(inv_gamma, x, shape, scale)
    args = (0.4, 0.2, 0.3)
    actual = logpdf_grad(inv_gamma, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "normal" begin

    # random
    x = normal(0, 1)

    # does not overflow
    @test logpdf(normal, 1e13, 0, 1) == -5e25
    @test logpdf_grad(normal, 1e13, 0, 1) == (-1e13, 1e13, 1e26)

    # logpdf_grad
    f = (x, mu, std) -> logpdf(normal, x, mu, std)
    args = (0.4, 0.2, 0.3)
    actual = logpdf_grad(normal, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "zero-dimensional broadcasted normal" begin

    # random
    x = broadcasted_normal(fill(0), fill(1))

    # logpdf_grad
    f = (x, mu, std) -> logpdf(broadcasted_normal, x, mu, std)
    args = (fill(0.4), fill(0.2), fill(0.3))
    actual = logpdf_grad(broadcasted_normal, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx; broadcast=true))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx; broadcast=true))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx; broadcast=true))
end

@testset "array normal (trivially broadcasted: all args have same shape)" begin

    mu =  [ 0.1   0.2   0.3  ;
            0.4   0.5   0.6  ]
    std = [ 0.01  0.02  0.03 ;
            0.04  0.05  0.06 ]
    x =   [ 1.    2.    3.   ;
            4.    5.    6.   ]

    # random
    broadcasted_normal(mu, std)

    # logpdf_grad
    f = (x, mu, std) -> logpdf(broadcasted_normal, x, mu, std)
    args = (mu, std, x)
    actual = logpdf_grad(broadcasted_normal, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx; broadcast=true))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx; broadcast=true))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx; broadcast=true))
end

@testset "broadcasted normal" begin

    ## Return shape of `broadcasted_normal`
    @test size(broadcasted_normal([0. 0. 0.], 1.)) == (1, 3)
    @test size(broadcasted_normal(zeros(1, 3, 4), ones(2, 1, 4))) == (2, 3, 4)
    @test_throws DimensionMismatch broadcasted_normal([0 0 0], [1 1])

    ## Return shape of `logpdf` and `logpdf_grad`
    @test size(logpdf(broadcasted_normal,
                      ones(2, 4), ones(2, 1), ones(1, 4))) == ()
    @test all(size(g) == ()
              for g in logpdf_grad(
                  broadcasted_normal, ones(2, 4), ones(2, 1), ones(1, 4)))
    # `x` has the wrong shape
    @test_throws DimensionMismatch logpdf(broadcasted_normal,
                                          ones(1, 2), ones(1,3), ones(2,1))
    @test_throws DimensionMismatch logpdf_grad(broadcasted_normal,
                                               ones(1, 2), ones(1,3), ones(2,1))
    # `x` has a shape that is broadcast-compatible with but not equal to the
    # right shape
    @test_throws DimensionMismatch logpdf(broadcasted_normal,
                                          ones(1, 3), ones(1,3), ones(2,1))
    @test_throws DimensionMismatch logpdf_grad(broadcasted_normal,
                                               ones(1, 3), ones(1,3), ones(2,1))
    # `mu` and `std` are broadcast-incompatible
    @test_throws DimensionMismatch logpdf(broadcasted_normal,
                                          ones(2, 1), ones(1,2), ones(1,3))
    @test_throws DimensionMismatch logpdf_grad(broadcasted_normal,
                                               ones(2, 1), ones(1,2), ones(1,3))

    ## Equivalence of broadcast to supplying bigger arrays for `mu` and `std`
    compact = OrderedDict(:x => reshape([ 0.2  0.3  0.4  0.5 ;
                                          0.5  0.4  0.3  0.2 ],
                                        (2, 4)),
                          :mu => reshape([0.7  0.7  0.8  0.6],
                                         (1, 4)),
                          :std => reshape([0.2, 0.1],
                                          (2, 1)))
    expanded = OrderedDict(:x => compact[:x],
                           :mu => repeat(compact[:mu], outer=(2, 1)),
                           :std => repeat(compact[:std], outer=(1, 4)))
    @test (logpdf(broadcasted_normal, values(compact)...) ==
           logpdf(broadcasted_normal, values(expanded)...))
    @test (logpdf_grad(broadcasted_normal, values(compact)...) ==
           logpdf_grad(broadcasted_normal, values(expanded)...))
end

@testset "multivariate normal" begin

    # random
    x = mvnormal([0.0, 0.0], [1.0 0.2; 0.2 1.4])
    @test length(x) == 2

    # logpdf_grad
    f = (x, mu, cov) -> logpdf(mvnormal, x, mu, cov)
    args = ([0.1, 0.2], [0.3, 0.4], [1.0 0.2; 0.2 1.4])
    actual = logpdf_grad(mvnormal, args...)
    @test isapprox(actual[1][1], finite_diff_vec(f, args, 1, 1, dx))
    @test isapprox(actual[1][2], finite_diff_vec(f, args, 1, 2, dx))
    @test isapprox(actual[2][1], finite_diff_vec(f, args, 2, 1, dx))
    @test isapprox(actual[2][2], finite_diff_vec(f, args, 2, 2, dx))
    @test isapprox(actual[3][1, 1], finite_diff_mat_sym(f, args, 3, 1, 1, dx))
    @test isapprox(actual[3][1, 2], finite_diff_mat_sym(f, args, 3, 1, 2, dx))
    @test isapprox(actual[3][2, 1], finite_diff_mat_sym(f, args, 3, 2, 1, dx))
    @test isapprox(actual[3][2, 2], finite_diff_mat_sym(f, args, 3, 2, 2, dx))
end

@testset "uniform" begin

    # random
    x = uniform(-0.5, 0.5)
    @test -0.5 < x < 0.5

    # out of support
    @test logpdf(uniform, -1, -0.5, 0.5) == -Inf

    # logpdf_grad
    f = (x, low, high) -> logpdf(uniform, x, low, high)
    args = (0.0, -0.5, 0.5)
    actual = logpdf_grad(uniform, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "uniform_discrete" begin

    # random
    x = uniform_discrete(10, 20)
    @test 10 <= x <= 20

    # out of support
    @test logpdf(uniform_discrete, -1, 10, 20) == -Inf

    # logpdf_grad
    args = (1, 1, 5)
    actual = logpdf_grad(uniform_discrete, args...)
    @test actual == (nothing, nothing, nothing)
end

@testset "piecewise_uniform" begin

    # random
    x = piecewise_uniform([-0.5, 0.5], [1.0])
    @test -0.5 < x < 0.5

    # out of support
    @test logpdf(piecewise_uniform, -1, [-0.5, 0.5], [1.0]) == -Inf

    # logpdf_grad
    f = (x, bounds, probs) -> logpdf(piecewise_uniform, x, bounds, probs)
    args = (0.5, [-1.0, 0.0, 1.0], [0.4, 0.6])
    actual = logpdf_grad(piecewise_uniform, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2][1], finite_diff_vec(f, args, 2, 1, dx))
    @test isapprox(actual[2][2], finite_diff_vec(f, args, 2, 2, dx))
    @test isapprox(actual[2][3], finite_diff_vec(f, args, 2, 3, dx))
    @test isapprox(actual[3][1], finite_diff_vec(f, args, 3, 1, dx))
    @test isapprox(actual[3][2], finite_diff_vec(f, args, 3, 2, dx))
end

@testset "beta uniform mixture" begin

    # random
    x = beta_uniform(0.5, 0.5, 0.5)
    @test 0 < x < 1

    # out of support
    @test logpdf(beta_uniform, -1, 0.5, 0.5, 0.5) == -Inf

    # logpdf_grad
    f = (x, theta, alpha, beta) -> logpdf(beta_uniform, x, theta, alpha, beta)
    args = (0.5, 0.4, 10., 2.)
    actual = logpdf_grad(beta_uniform, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
    @test isapprox(actual[4], finite_diff(f, args, 4, dx))
end

@testset "geometric" begin

    # random
    @test geometric(0.5) >= 0

    # out of support
    @test logpdf(geometric, -1, 0.5) == -Inf

    # logpdf_grad
    f = (x, p) -> logpdf(geometric, x, p)
    args = (4, 0.3)
    actual = logpdf_grad(geometric, args...)
    @test actual[1] == nothing
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
end

@testset "binom" begin

    # random
    @test binom(5, 0.3) >= 0

    # out of support
    @test logpdf(binom, -1, 5, 0.3) == -Inf

    # logpdf_grad
    f = (x, n, p) -> logpdf(binom, x, n, p)
    args = (2, 5, 0.3)
    actual = logpdf_grad(binom, args...)
    @test actual[1] == nothing
    @test actual[2] == nothing
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "neg_binom" begin

    # random
    @test neg_binom(5, 0.3) >= 0

    # out of support
    @test logpdf(neg_binom, -1, 5, 0.3) == -Inf

    # logpdf_grad
    f = (x, r, p) -> logpdf(neg_binom, x, r, p)
    args = (2, 5, 0.3)
    actual = logpdf_grad(neg_binom, args...)
    @test actual[1] == nothing
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "exponential" begin

    # random
    @test exponential(0.5) > 0

    # out of support
    @test logpdf(exponential, -1, 0.5) == -Inf

    # logpdf_grad
    f = (x, rate) -> logpdf(exponential, x, rate)
    args = (1.2, 0.5)
    actual = logpdf_grad(exponential, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
end

@testset "poisson" begin

    # random
    @test poisson(1.0) >= 0

    # out of support
    @test logpdf(poisson, -1, 1.0) == -Inf

    # logpdf_grad
    f = (x, lambda) -> logpdf(poisson, x, lambda)
    args = (4, 2.0)
    actual = logpdf_grad(poisson, args...)
    @test actual[1] == nothing
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
end

@testset "laplace" begin

    # random
    x = laplace(1, 1)

    # logpdf_grad
    f = (x, loc, scale) -> logpdf(laplace, x, loc, scale)
    args = (0.4, 0.2, 0.3)
    actual = logpdf_grad(laplace, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end

@testset "cauchy" begin

    # random
    x = cauchy(0, 5)

    # logpdf_grad
    f = (x, x0, gamma) -> logpdf(cauchy, x, x0, gamma)
    args = (0.4, 0.2, 0.3)
    actual = logpdf_grad(cauchy, args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end


"""
    CountTransformer()

Convert a collection of raw documents to matrix representing a bag-of-words structure from
word counts. Essentially, a bag-of-words approach to representing documents in a matrix is
comprised of a count of every word in the document corpus/collection for every document.
This is a simple but often quite powerful way of representing documents as vectors. The
resulting representation is a matrix with rows representing every document in the corpus
and columns representing every word in the corpus. The value for each cell is the raw count
of a particular word in a particular document.

Similarly to the `TfidfTransformer`, the vocabulary considered can be restricted
to words occuring in a maximum or minimum portion of documents.
The parameters `max_doc_freq` and `min_doc_freq` restrict the vocabulary
that the transformer will consider. `max_doc_freq` indicates that terms in only
up to the specified percentage of documents will be considered. For example, if
`max_doc_freq` is set to 0.9, terms that are in more than 90% of documents
will be removed. Similarly, the `min_doc_freq` parameter restricts terms in the
other direction. A value of 0.01 means that only terms that are at least in 1% of
documents will be included.
"""
mutable struct CountTransformer <: AbstractTextTransformer
    max_doc_freq::Float64
    min_doc_freq::Float64
end

function CountTransformer(; max_doc_freq::Float64 = 1.0, min_doc_freq::Float64 = 0.0)
    transformer = CountTransformer(max_doc_freq, min_doc_freq)
    message = MMI.clean!(transformer)
    isempty(message) || @warn message
    return transformer
end

struct CountTransformerResult
    vocab::Vector{String}
end

function _fit(transformer::CountTransformer, verbosity::Int, X::Corpus)
    # process corpus vocab
    update_lexicon!(X)

    # calculate min and max doc freq limits
    if transformer.max_doc_freq < 1 || transformer.min_doc_freq > 0
        # we need to build out the DTM
        dtm_matrix = build_dtm(X)
        n = size(dtm_matrix.dtm, 2) # docs are columns
        high = round(Int, transformer.max_doc_freq * n)
        low = round(Int, transformer.min_doc_freq * n)
        _, vocab = limit_features(dtm_matrix, high, low)
    else
        vocab = sort(collect(keys(lexicon(X))))
    end

    # prepare result
    fitresult = CountTransformerResult(vocab)
    cache = nothing

    return fitresult, cache, NamedTuple()
end

function _transform(::CountTransformer,
                    result::CountTransformerResult,
                    v::Corpus)
    dtm_matrix = build_dtm(v, result.vocab)

    # here we return the `adjoint` of our sparse matrix to conform to
    # the `n x p` dimensions throughout MLJ
    return adjoint(dtm_matrix.dtm)
end

# for returning user-friendly form of the learned parameters:
function MMI.fitted_params(::CountTransformer, fitresult::CountTransformerResult)
    vocab = fitresult.vocab
    return (vocab = vocab,)
end

## META DATA

MMI.metadata_pkg(CountTransformer,
             name="$PKG",
             uuid="7876af07-990d-54b4-ab0e-23690620f79a",
             url="https://github.com/JuliaAI/MLJText.jl",
             is_pure_julia=true,
             license="MIT",
             is_wrapper=false
)

MMI.metadata_model(CountTransformer,
               input_scitype = Union{
                   AbstractVector{<:AbstractVector{STB.Textual}},
                   AbstractVector{<:STB.Multiset{<:ScientificNGram}},
                   AbstractVector{<:STB.Multiset{STB.Textual}}
                   },
               output_scitype = AbstractMatrix{STB.Continuous},
               docstring = "Build Bag-of-Words matrix from word counts for corpus of documents",
               path = "MLJText.CountTransformer"
               )


# SteinProbitRegressionGaussianPrior
#
# We assume we have the following model:
#
# y_i = \Psi (beta_1 x_1 + ... + beta_p x_p)
#
# where \Psi is assumed to be the noraml CDF (and
# a Gaussian prior for beta).

using Distributions: Normal
import Distributions

type SteinProbitRegressionGaussianPrior <: SteinPosterior
    # the X data
    X::Array{Float64,2}
    # the 0/1 y values
    y::Array{Float64,1}
    # priorsigma2 is prior variance
    priorsigma2::Float64
end

defaultsigma2 = 0.1

SteinProbitRegressionGaussianPrior(X, y) =
    SteinProbitRegressionGaussianPrior(X, y, defaultsigma2)

# Returns lower bound of support
function supportlowerbound(d::SteinProbitRegressionGaussianPrior, i::Int)
    -Inf
end

# Returns upper bound of support
function supportupperbound(d::SteinProbitRegressionGaussianPrior, i::Int)
    Inf
end

# Returns the log prior log pi(beta)
function logprior(d::SteinProbitRegressionGaussianPrior,
                  beta::Array{Float64, 1})
    logpriorvalue =
        -0.5 * sum((beta .^ 2) ./ d.priorsigma2) -
        -0.5 * sum(log(2 * pi .* d.priorsigma2))
    logpriorvalue
end

# Returns log pi(d.y[idx] | beta)
function loglikelihood(d::SteinProbitRegressionGaussianPrior,
                       beta::Array{Float64, 1};
                       idx=1:length(d.y))
    y = d.y[idx]
    X = d.X[idx,:]
    ypos = find(y .> 0.0)
    yneg = find(y .<= 0.0)
    u = vec(X * beta)
    gsn = Normal()

    logprobpos = Distributions.logcdf(gsn, u[ypos])
    logprobneg = log(1.0 - Distributions.cdf(gsn, u[yneg]))
    sum(logprobpos) + sum(logprobneg)
end

# Returns the log prior log pi(beta)
function gradlogprior(d::SteinProbitRegressionGaussianPrior,
                      beta::Array{Float64, 1})
    -beta ./ d.priorsigma2
end

# Returns grad_{beta} log pi(d.y[idx] | beta)
function gradloglikelihood(d::SteinProbitRegressionGaussianPrior,
                           beta::Array{Float64,1};
                           idx=1:length(d.y))
    y = d.y[idx]
    X = d.X[idx,:]
    ypos = find(y .> 0.0)
    yneg = find(y .<= 0.0)
    u = vec(X * beta)
    gsn = Normal()

    posweights = Distributions.pdf(gsn, u[ypos]) ./ Distributions.cdf(gsn, u[ypos])
    negweights = -Distributions.pdf(gsn, u[yneg]) ./ (1.0 - Distributions.cdf(gsn, u[yneg]))

    xpos = vec(sum(broadcast(*, X[ypos,:], posweights), 1))
    xneg = vec(sum(broadcast(*, X[yneg,:], negweights), 1))

    xneg + xpos
end

function numdimensions(d::SteinProbitRegressionGaussianPrior)
    size(d.X, 2)
end


export summary_statistics, basic_summary_statistics

function data_size_cutoff(simplices::Vector{Int64}, nverts::Vector{Int64},
                          times::Vector{Int64}, min_simplex_size::Int64,
                          max_simplex_size::Int64)
    new_simplices = Int64[]
    new_nverts = Int64[]
    new_times = Int64[]
    curr_ind = 1
    for (nv, time) in zip(nverts, times)
        if min_simplex_size <= nv <= max_simplex_size
            append!(new_simplices, simplices[curr_ind:(curr_ind + nv - 1)])
            push!(new_nverts, nv)
            push!(new_times, time)
        end
        curr_ind += nv
    end
    return (new_simplices, new_nverts, new_times)
end

"""
summary_statistics
------------------------

Computes several statistics about the dataset.

summary_statistics(dataset::HONData)

Input parameter:
- dataset::HONData: The dataset.
"""
function summary_statistics(dataset::HONData)
    function _summary_statistics(simplices::Vector{Int64}, nverts::Vector{Int64}, times::Vector{Int64})
        if length(nverts) == 0
            return ("0,0,0,0,0,0.0,0.0,0,0,0,0,0.0,0.0,0.0,0.0,0,0")
        end

        A, At, B = basic_matrices(simplices, nverts)
        no, nc = num_open_closed_triangles(A, At, B)
        num_nodes = sum(sum(At, dims=1) .> 0)  # note: includes 1-node simplices
        density = nnz(B) / (num_nodes^2 - num_nodes)
        str1 = @sprintf("%d,%d,%d,%d,%d,%f,%e,%d,%d",
                        num_nodes, length(nverts), nnz(A), mean(nverts),
                        maximum(nverts), mean(nonzeros(B)), density, nc, no)

        # Backbone dataset
        bb_simplices, bb_nverts, bb_times = backbone(simplices, nverts, times)
        (A, At, C) = basic_matrices(bb_simplices, bb_nverts)
        str2 = @sprintf("%d,%d,%f,%f", length(bb_nverts), nnz(A),
                        mean(bb_nverts), mean(nonzeros(C)))

        # Random configuration
        config = configuration_sizes_preserved(bb_simplices, bb_nverts)
        (A, At, B) = basic_matrices(config, bb_nverts)
        num_nodes = sum(sum(At, dims=1) .> 0)
        density = nnz(B) / (num_nodes^2 - num_nodes)
        no, nc = num_open_closed_triangles(A, At, B)
        str3 = @sprintf("%f,%e,%d,%d",
                        mean(nonzeros(B)), density, nc, no)
        return "$str1,$str2,$str3"
    end

    dataset_name = dataset.name
    open("$(dataset_name)-statistics.csv", "w") do f
        write(f, "dataset,")
        write(f, "nnodes,nsimps,nconns,meansimpsize,maxsimpsize,meanprojweight,projdensity,nclosedtri,nopentri,")
        write(f, "nbbnodes,nbbconns,meanbbsimpsize,meanbbprojweight,")
        write(f, "meanbbconfigweight,bbconfigdensity,nbbconfigclosedtri,nbbconfigopentri\n")
        simplices, nverts, times = dataset.simplices, dataset.nverts, dataset.times
        stats = _summary_statistics(simplices, nverts, times)
        write(f, "$(dataset_name),$stats\n")

        # Get the same statistics but for the dataset restricted to 3-node
        # simplices.
        csimplices, cnverts, ctimes = data_size_cutoff(simplices, nverts, times, 3, 3)
        stats = _summary_statistics(csimplices, cnverts, ctimes)
        write(f, "$(dataset_name)-3-3,$stats\n")
    end
end

"""
basic_summary_statistics
------------------------

Prints some basic summary statistics of the dataset.

basic_summary_statistics(dataset::String)

Input parameter:
- dataset::HONData: The dataset.
"""
function basic_summary_statistics(dataset::HONData)
    simplices, nverts, times = dataset.simplices, dataset.nverts, dataset.times
    bb_simplices, bb_nverts, bb_times = backbone(simplices, nverts, times)
    A, At, B = basic_matrices(simplices, nverts)
    num_nodes = sum(sum(At, dims=1) .> 0)
    num_edges = nnz(B) / 2
    num_simps = length(nverts)
    num_bb_simps = length(bb_nverts)
    println("dataset & # nodes & # edges in proj. graph & # simplices & # unique simplices")
    @printf("%s & %d & %d & %d & %d\n", dataset.name, num_nodes, num_edges, num_simps, num_bb_simps)
end


using NetCDF, Plots, Printf

sub_directory = "/Geostrophic_Start_20"
filename = pwd() * sub_directory * "_middepth.nc"
ncinfo(filename)
x = Array(NetCDF.open(filename, "xC"))
y = Array(NetCDF.open(filename, "yC"))
z = Array(NetCDF.open(filename, "zC"))
t = Array(NetCDF.open(filename, "time"))
sim_day = t ./ 86400
z = (z[2:end] + z[1:end-1]) / 2
b = NetCDF.open(filename, "b")
u = NetCDF.open(filename, "u")
v = NetCDF.open(filename, "v")
w = NetCDF.open(filename, "w")
##
sti = 150 # starting time index
cmax = maximum(Array(b[2:end-1,2:end-1, 1, sti:end]))
cmin = minimum(Array(b[:,:, 1, sti:end]))
clims = (cmin, cmax)
anim = @animate for i in sti:1:length(b[1, 1, 1, :])
    b_array = Array(b[:, :, 1, i])
    # p1 = contourf(y, z, b_array, fill = true, linewidth = 0, color = :ocean, clim = (-0.0013871555898403098, -3.3776441214941526e-6))
    day_label = @sprintf("%.2f ", sim_day[i])
    p1 = contourf(x, y, b_array',
    color = :ocean, title = "Buoyancy at Middepth at day " * day_label ,
     xlabel = "Meridional [m]", ylabel = "Depth [m]"
     , clims = clims)
end
gif(anim, filename *".gif", fps = 15)
##
filename = pwd() * sub_directory * "_zonal_average.nc"
ncinfo(filename)
y = Array(NetCDF.open(filename, "yC"))
z = Array(NetCDF.open(filename, "zC"))
b = NetCDF.open(filename, "Bz")
u = NetCDF.open(filename, "Uz")
t = Array(NetCDF.open(filename, "time"))
sim_day = t ./ 86400
cmax = maximum(Array(b[2:end-1, 2:end-1, 1:1:end]))
cmin = minimum(Array(b[2:end-1, 2:end-1, 1:1:end]))
clims = (cmin, cmax)
anim = @animate for i in 1:1:length(b[ 1, 1, :])
    b_array = Array(b[:, :, i])
    day_label = @sprintf("%.2f ", sim_day[i])
    p1 = contourf(y, z, b_array',
        color = :ocean, title = "Zonal Average Buoyancy at day " * day_label ,
        xlabel = "Meridional [m]", ylabel = "Depth [m]"
        , clims = clims)
end
gif(anim, filename * ".gif", fps = 15)
##
i = length(b[ 1, 1, :])
u_array = Array(u[:, :, i])
day_label = @sprintf("%.2f ", sim_day[i])
p1 = contourf(y, z, u_array',
color = :ocean, title = "Zonal Average Zonal Velocity at day " * day_label ,
        xlabel = "Meridional [m]", ylabel = "Depth [m]")

##
i = length(b[1, 1, :])
b_array = Array(b[:, :, i])
cmax = maximum(b_array)
cmin = minimum(b_array)
clims = (cmin, cmax)
day_label = @sprintf("%.2f ", sim_day[i])
p1 = contourf(y, z, b_array',
        color = :ocean, title = "Zonal Average Buoyancy at day " * day_label ,
        xlabel = "Meridional [m]", ylabel = "Depth [m]"
        , clims = clims)
plot(p1)
##
filename = pwd() * sub_directory * "_surface.nc"
ncinfo(filename)
y = Array(NetCDF.open(filename, "yC"))
x = Array(NetCDF.open(filename, "xC"))
b = NetCDF.open(filename, "b")
t = Array(NetCDF.open(filename, "time"))
sim_day = t ./ 86400
start_time = 1
skip = 1
end_time = 180
cmax = maximum(Array(b[2:end-1, 2:end-1, 1, start_time:skip:end_time]))
cmin = minimum(Array(b[2:end-1, 2:end-1, 1, start_time:skip:end_time]))
clims = (cmin, cmax)
anim = @animate for i in start_time:end_time
    b_array = Array(b[:, :,1, i])
    day_label = @sprintf("%.2f ", sim_day[i])
    p1 = contourf(x, y, b_array',
        color = :thermometer, title = "Buoyancy at 50 [m] on day " * day_label ,
        xlabel = "Zonal [m]", ylabel = "Meridional [m]"
        , clims = clims, linewidth = 0)
end
gif(anim, filename * ".gif", fps = 15)
##
filename = pwd() *  "/Windstress_Convection_4_meridional.nc"
ncinfo(filename)
x = Array(NetCDF.open(filename, "yC"))
y = Array(NetCDF.open(filename, "yC"))
z = Array(NetCDF.open(filename, "zC"))
t = Array(NetCDF.open(filename, "time"))
sim_day = t ./ 86400
# z = (z[2:end] + z[1:end-1]) / 2
b = NetCDF.open(filename, "b")
u = NetCDF.open(filename, "u")
v = NetCDF.open(filename, "v")
w = NetCDF.open(filename, "w")
Δz = z[2:end] - z[1:end-1]
Δx = x[2:end] - x[1:end-1]
Δy = y[2:end] - y[1:end-1]
a = abs.(w[1, 1, 2:end-1, 1] ./ (Δz))
a2 = abs.(w[1, 1:end-1, 16, 1] ./ (Δx))
b = abs.(u[1, 1:end-1, 1, 1] ./ (Δy))
c = abs.(v[1, 2:end-1, 1, 1] ./ (Δy))

##
filename = pwd() *  "/Windstress_Convection_4_checkpoint_iteration4851612.jld2"
ncinfo(filename)


# routine tests for wrappers
# This section is just to check existence and sanity of methods.
# Serious accuracy verification is elsewhere.

using LinearAlgebra
using LinearAlgebra: sorteig!
using Test
using GenericSchur

# allow for different phase factors and perhaps ordering
# v1 is the reference value against which we are checking v2
function _chkeigvecs(v1::AbstractMatrix{Tref}, v2::AbstractMatrix{T},
                     isnrml, sameorder=true
                     ) where {Tref,T}
    m,n = size(v2)
    if isnrml
        @test norm(v2' * v2 - I) < 50 * n * eps(real(T))
    end
    tmat = v1' * v2
    zT = zero(real(T))
    nvecs_ok = 0
    Tcheck = (precision(real(Tref)) < precision(real(T))) ? real(Tref) : real(T)
    myeps = eps(Tcheck)
    for j in 1:m
        # largest and next-largest projections
        if sameorder
            p1 = abs(tmat[j,j])
            tmat[j,j] = 0
            p2 = maximum(abs.(tmat[:,j]))
        else
            p1,p2 = zT,zT
            for i in 1:m
                ta = abs(tmat[i,j])
                if ta > p1
                    p1 = ta
                elseif ta > p2
                    p2 = ta
                end
            end
        end
        δ = abs(p1-1)
        ok = (δ < sqrt(myeps))
        if (verbosity[] > 0) && (δ > 100*myeps)
            @warn "vector projection error is $δ angle $(acos(clamp(p1,0,1)))"
        end
        if isnrml
            ok &= (abs(p2) < sqrt(myeps))
        end
        nvecs_ok += ok
    end
    @test nvecs_ok == m
end

let T = BigFloat
    n = 10
    # FIXME: should protect against accidental poor condition
    A = rand(T,n,n)
    for (w,f) in zip([:bare, :hermitian, :symmetric],[identity, Hermitian, Symmetric])
        @testset "wrappers $w $T" begin
            Awrk = f(A)
            E = eigen(Awrk)
            @test norm(Awrk*E.vectors - E.vectors * Diagonal(E.values)) < sqrt(eps(T))
            v = eigvecs(Awrk)
            _chkeigvecs(E.vectors, v, w != :bare)
            λ = eigvals(Awrk)
            if (verbosity[] > 0) && (! (λ ≈ E.values))
                @warn "eigval order comparison (eigvals, eigen, diff): "
                    display(hcat(λ, E.values, λ .- E.values))
                    println()
            end
            @test λ ≈ E.values
        end
    end
end

let T = Complex{BigFloat}
    n = 10
    # FIXME: should protect against accidental poor condition
    A = rand(T,n,n)
    for (w,f) in zip([:bare, :hermitian],[identity, Hermitian])
        @testset "wrappers $w $T" begin
            Awrk = f(A)
            E = eigen(Awrk)
            @test norm(Awrk*E.vectors - E.vectors * Diagonal(E.values)) < sqrt(eps(real(T)))
            v = eigvecs(Awrk)
            _chkeigvecs(E.vectors, v, w != :bare)
            λ = eigvals(Awrk)
            @test λ ≈ E.values
            if (verbosity[] > 0) && (! (λ ≈ E.values))
                @warn "eigval order comparison (eigvals, eigen, diff): "
                    display(hcat(λ, E.values, λ .- E.values))
                    println()
            end
        end
    end
end

function _chkrcond(x,y,rtol,atol=1e3*eps(eltype(x)))
    ok = true
    for (xi,yi) in zip(x,y)
        x1 = clamp(xi,0,1)
        y1 = clamp(yi,0,1)
        if abs(x1 - y1) > rtol * abs(y1) + atol
            ok = false
        end
    end
    if (verbosity[] > 0 && !ok) || (verbosity[] > 1)
        if !ok
            @warn "rcond mismatch (trial, ref, ratio): "
        else
            @info "rcond comparison (trial, ref, ratio): "
        end
        display(hcat(x, y, x ./ y))
        println()
    end
    return ok
end

# This section tests compatibility with Julia PR #38483 which was reverted.
# It should be rewritten but is kept to facilitate study of questionable cases.
if ((VERSION > v"1.7.0-DEV.976") && (VERSION < v"1.7.0-beta3.37"))
    for (T,Tref) in ((Complex{BigFloat},ComplexF64),(BigFloat, Float64))
        @testset "extended eigen() $T" begin
            n = 10
            Aref = rand(Tref,n,n)
            A = T.(Aref)
            old = precision(real(T))
            @assert old >= 53
            if T <: Real
                Eref1 = eigen(Aref, jvl=true)
                # we convert to complex for condition nrs, so this is a fairer check
                Eref2 = eigen(Aref .+ 0im, jvl=true, jce=true, jcv=true)
            else
                Eref1 = eigen(Aref, jvl=true, jce=true, jcv=true)
                Eref2 = Eref1
            end
            setprecision(real(T), 53) do
                E = eigen(A, jvl=true, jce=true, jcv=true)
                @test norm(E.vectorsl' * A - Diagonal(E.values) * E.vectorsl') < sqrt(eps(real(T)))
                if verbosity[] > 0
                    @info "eigval comparison (trial, ref, diff): "
                    display(hcat(E.values, Eref1.values, E.values .- Eref1.values))
                    println()
                end
                @test sorteig!(E.values) ≈ sorteig!(Eref1.values)
                # stdlib inverts condition nrs for real matrices (WTF?)
                # if we switch to the real forms, we would need this:
                # rconde = (Tref <: Real) ? (1.0 ./ Eref1.rconde) : Eref1.rconde
                # rcondv = (Tref <: Real) ? (1.0 ./ Eref1.rcondv) : Eref1.rcondv
                rconde = Eref2.rconde
                rcondv = Eref2.rcondv
                @test _chkrcond(E.rconde, rconde, 0.1)
                # CHECKME: we should not need to be so lenient.
                # The nasty cases are not common (empirically 1 vector in < 10% of 10x10),
                # but nightlys on Github are good at finding them.
                # This accuracy is sufficient for typical applications.
                @test _chkrcond(E.rcondv, rcondv, 5.0)
                _chkeigvecs(Eref1.vectors, E.vectors, false)
                _chkeigvecs(Eref1.vectorsl, E.vectorsl, false)
            end
        end
    end
end


# Smith Normal Form

function divisable(y::R, x::R ) where {R}
  x == zero(R) && return y == zero(R)
  return div(y,x)*x == y
end

function divide(y::R, x::R) where {R}
    if x != -one(R)
        return div(y,x)
    else
        return y * x
    end
end

function rcountnz(X, i)
    n = size(X, 1)
    c = 0
    for j in 1:n
        if X[j, i] != 0
           c += 1
        end
    end
    return c
end

function ccountnz(X, j)
    m = size(X, 2)
    c = 0
    for i in 1:m
        if X[j, i] != 0
           c += 1
        end
    end
    return c
end

function rswap!(M::AbstractArray{R,2}, r1::Int, r2::Int) where {R}
    r1 == r2 && return M
    m = size(M, 2)
    @inbounds for i in 1:m
        tmp = M[r1, i]
        M[r1, i] = M[r2, i]
        M[r2, i] = tmp
    end
    return M
end

function cswap!(M::AbstractArray{R,2}, c1::Int, c2::Int) where {R}
    c1 == c2 && return M
    n = size(M, 1)
    @inbounds for j in 1:n
        tmp = M[j, c1]
        M[j, c1] = M[j, c2]
        M[j, c2] = tmp
    end
    return M
end

function rowelimination(D::AbstractArray{R,2}, a::R, b::R, c::R, d::R, i::Int, j::Int) where {R}
    m = size(D, 2)
    @inbounds for k in 1:m
        t = D[i,k]
        s = D[j,k]
        D[i,k] = a*t + b*s
        D[j,k] = c*t + d*s
    end
    return D
end

function colelimination(D::AbstractArray{R,2}, a::R, b::R, c::R, d::R, i::Int, j::Int) where {R}
    n = size(D, 1)
    @inbounds for k in 1:n
        t = D[k,i]
        s = D[k,j]
        D[k,i] = a*t + b*s
        D[k,j] = c*t + d*s
    end
    return D
end

function rowpivot(U::AbstractArray{R,2},
                  Uinv::AbstractArray{R,2},
                  D::AbstractArray{R,2},
                  i, j; inverse=true) where {R}
    for k in reverse!(findall(!iszero, view(D, :, j)))
        a = D[i,j]
        b = D[k,j]

        i == k && continue

        s, t, g = bezout(a, b)
        x = divide(a, g)
        y = divide(b, g)

        rowelimination(D, s, t, -y, x, i, k)
        inverse && rowelimination(Uinv, s, t, -y, x, i, k)
        colelimination(U, x, y, -t, s, i, k)
    end
end

function colpivot(V::AbstractArray{R,2},
                  Vinv::AbstractArray{R,2},
                  D::AbstractArray{R,2},
                  i, j; inverse=true) where {R}
    for k in reverse!(findall(!iszero, view(D, i, :)))
        a = D[i,j]
        b = D[i,k]

        j == k && continue

        s, t, g = bezout(a, b)
        x = divide(a, g)
        y = divide(b, g)

        colelimination(D, s, t, -y, x, j, k)
        inverse && colelimination(Vinv, s, t, -y, x, j, k)
        rowelimination(V, x, y, -t, s, j, k)
    end
end

function smithpivot(U::AbstractArray{R,2},
                    Uinv::AbstractArray{R,2},
                    V::AbstractArray{R,2},
                    Vinv::AbstractArray{R,2},
                    D::AbstractArray{R,2},
                    i, j; inverse=true) where {R}

    pivot = D[i,j]
    @assert pivot != zero(R) "Pivot cannot be zero"
    while ccountnz(D,i) > 1 || rcountnz(D,j) > 1
        colpivot(V, Vinv, D, i, j, inverse=inverse)
        rowpivot(U, Uinv, D, i, j, inverse=inverse)
    end
end

function init(M::AbstractSparseMatrix{R,Ti}; inverse=true) where {R, Ti}
    D = copy(M)
    rows, cols = size(M)

    U = spzeros(R, rows, rows)
    for i in 1:rows
        U[i,i] = one(R)
    end
    Uinv = inverse ? copy(U) : spzeros(R, 0, 0)

    V = spzeros(R, cols, cols)
    for i in 1:cols
        V[i,i] = one(R)
    end
    Vinv = inverse ? copy(V) : spzeros(R, 0, 0)

    return U, V, D, Uinv, Vinv
end

function init(M::AbstractMatrix{R}; inverse=true) where {R}
    D = copy(M)
    rows, cols = size(M)

    U = zeros(R, rows, rows)
    for i in 1:rows
        U[i,i] = one(R)
    end
    Uinv = inverse ? copy(U) : zeros(R, 0, 0)

    V = zeros(R, cols, cols)
    for i in 1:cols
        V[i,i] = one(R)
    end
    Vinv = inverse ? copy(V) : zeros(R, 0, 0)

    return U, V, D, Uinv, Vinv
end

function snf(M::AbstractMatrix{R}; inverse=true, debug=false) where {R}
    rows, cols = size(M)
    U, V, D, Uinv, Vinv = init(M, inverse=inverse)

    t = 1
    for j in 1:cols
        debug && println("Working on column $j out of ", size(D,2))
        debug && display(collect(D))

        rcountnz(D,j) == 0 && continue

        prow = 1
        if D[t,t] != zero(R)
            prow = t
        else
            # Good pivot row for j-th column is the one
            # that have a smallest number of elements
            rsize = typemax(R)
            for i in 1:rows
                if D[i,j] != zero(R)
                    c = count(!iszero, view(D, i, :))
                    if c < rsize
                        rsize = c
                        prow = i
                    end
                end
            end
        end

        debug && "Pivot Row selected: t = $t, pivot_row = $prow" |> println
        debug && display(collect(D))
        debug && "Swapping rows. ( $t, $prow )" |> println

        rswap!(D, t, prow)
        inverse && rswap!(Uinv, t, prow)
        cswap!(U, t, prow)

        debug && println("Performing the pivot step at (t=$t, j=$j)")
        debug && display(collect(D))

        smithpivot(U, Uinv, V, Vinv, D, t, j, inverse=inverse)

        cswap!(D, t, j)
        inverse && cswap!(Vinv, t, j)
        rswap!(V, t, j)

        t += 1

        debug && println("D: $(size(D))")
        debug && display(collect(D))
        debug && println("U: $(size(U))")
        debug && display(collect(U))
        debug && println("V: $(size(V))")
        debug && display(collect(V))
        debug && println("Vinv:")
        debug && display(collect(Vinv))
    end

    # Make sure that d_i is divisible be d_{i+1}.
    r = minimum(size(D))
    pass = true
    while pass
        pass = false
        for i in 1:r-1
            divisable(D[i+1,i+1], D[i,i]) && continue
            pass = true
            D[i+1,i] = D[i+1,i+1]

            colelimination(Vinv, one(R), one(R), zero(R), one(R), i, i+1)
            rowelimination(V, one(R), zero(R), -one(R), one(R), i, i+1)

            smithpivot(U, Uinv, V, Vinv, D, i, i, inverse=inverse)
        end
    end

    if issparse(D)
        return dropzeros!(U), dropzeros!(V), dropzeros!(D), dropzeros!(Uinv), dropzeros!(Vinv)
    else
        return U, V, D, Uinv, Vinv
    end
end


"""
    temporal_connected(g, v1, t1, v2, t2)

Return `true` if there is path from `v1` at timestamp `t1` to `v2` at
timestamp `t2` and `false` otherwise.
"""
temporal_connected{V, T}(g::AbstractEvolvingGraph{V, T}, v1, t1, v2, t2) =
     (v2, t2) in breadth_first_impl(g, v1, t1) ? true : false
function temporal_connected{V, T}(g::EvolvingGraph{V, T}, v1, t1, v2, t2)
    nv1 = find_node(g, v1)
    nv2 = find_node(g, v2)

    # check if v1 and v2 are nodes of the graph
    if nv1 == false || nv2 == false
        return false
    end

    if (nv2, T(t2)) in breadth_first_impl(g, nv1, T(t1))
        return true
    else
        return false
    end
end
"""
    weak_connected(g, v1, v2)

Return `true` if there is a temporal path from `v1` to `v2` at any timestamps.
"""
function weak_connected(g::AbstractEvolvingGraph, v1, v2)
    sum = false
    ts = timestamps(g)
    n = length(ts)
    for i in 1:n
        for j in i:n
            sum += temporal_connected(g, v1, ts[i], v2, ts[j])
            if sum > 0
                return true
            end
        end
    end
    return false
end

"""
    weak_connected_components(g, valuesonly = true)

Find the weakly connected components of an evolving graph `g`, i.e,
each node in the set is weakly connected to all the other nodes.
If valuesonly = false, returns a dictionary with the starting
the search as dictionary key.
"""
function weak_connected_components{V}(g::AbstractEvolvingGraph{V},
                                                                           valuesonly::Bool = true)
    components = Dict{Tuple, Vector{Tuple}}()
    nodelist = V[]
    reachable = Tuple[]

    t = timestamps(g)
    n = length(t)

    for n1 in nodes(g)
        i = 1
        if ! (n1 in nodelist)

            reachable = breadth_first_impl(g, n1, t[i])
             while length(reachable) == 1
                 if i < n
                     i += 1
                 else
                     break
                 end
                 reachable = breadth_first_impl(g, n1, t[i])
            end

            append!(nodelist, map(x -> x[1], reachable))

            components[(n1, t[i])] = reachable
            ks = keys(components)

            for n2 in ks
                if !(n2 == (n1, t[i])) &&
                    (temporal_connected(g, n1, t[i], n2[1], n2[2]))
                    components[n2] = reachable
                    delete!(components, (n1, t[i]))
                end
            end
        end
    end
    if valuesonly
        return collect(values(components))
    else
        return components
    end
end


using ConditionalMoments
using Test

## Observations

const X = [1.,1.,4.,5.,6.,7,1.,2.,4.,5.,6.,7,3.,1.,4.,5.,6.,7]
const dt = 0.1
const ob = Observation(X,dt)

@testset "Observations" begin
    @test ob[3] == 4.0
    @test ob[ob .> 3] == X[X .> 3]
    @test ob[(ob .> 3)'] == X[X .> 3]
    @test ob[:] == X
    @test ob.dt == dt
    @test ob.npoints == 18

    @test_throws MethodError Observation([1.,1.,4.,5.,6.,7+2im],dt)
end

## Moments

const τ_indices = [1,3,4]
const bin_edges = LinRange(1.5,6.5,5)
const histogramSettings = HistogramSettings(τ_indices,bin_edges)

@testset "Moment settings" begin
    @test histogramSettings.n_time_points == 3
    @test histogramSettings.n_evaluation_points == 5

    @test_throws MethodError HistogramSettings([1,2,3.1],bin_edges)
    @test_throws MethodError HistogramSettings([1,2,3+1im],bin_edges)
end

const moment_sol = MomentSolution(ob, histogramSettings)
@testset "Moment calculations" begin
    @test moment_sol.observation.npoints == 18
    @test size(moment_sol.moments.M1) == (3,4)
    @test size(moment_sol.moments.M2) == (3,4)
    #@test size(moment_sol.errors) == (3,4)
    @test moment_sol.moments.M1[2,4] == -4.5
    @test moment_sol.moments.M2[2,4] == 0.25
end

## Drift and noise functions
const model1 = WhiteNoiseProcess(moment_sol)
const model2 = WhiteNoiseProcess(moment_sol,i=2)

@testset "Drift and noise calculations" begin
    @test size(model1.drift) == (4,)
    @test size(model1.noise) == (4,)
    #@test size(model1.errors) == (4,)
    @test size(model2.drift) == (4,)
    @test size(model2.noise) == (4,)
    #@test size(model2.errors) == (4,)
    @test ~isnan(model1.drift[4])
    @test ~isnan(model2.drift[4])
    @test ~isnan(model1.noise[4])
    @test ~isnan(model2.noise[4])
end

## Ensemble observations

const X_ensemble = repeat(X,1,2)

const ensemble_ob = EnsembleObservation(X_ensemble,dt)

@testset "Ensemble observations" begin
    @test ensemble_ob[3] == 4.0
    @test ensemble_ob[:,1] == X
    @test ensemble_ob[:,2] == X
    @test ensemble_ob[3,:] == [4.0,4.0]
    @test ensemble_ob[:] == X_ensemble
    @test ensemble_ob.dt == dt
    @test ensemble_ob.npoints == 18
    @test ensemble_ob.nsample == 2

    @test_throws MethodError Observation([1.,1.,4.,5.,6.,7+2im],dt)
end

## Ensemble moments

const ensemble_sol = MomentSolution(ensemble_ob, histogramSettings)
@testset "Mapping moments" begin
    @test ensemble_sol.observation.npoints == 18
    @test ensemble_sol.observation.nsample == 2
    @test size(ensemble_sol.moments.M1) == (3,4)
    @test size(ensemble_sol.moments.M2) == (3,4)
    #@test size(ensemble_moments.errors) == (3,4)
    @test ensemble_sol.moments.M1[2,4] == -4.5
    @test ensemble_sol.moments.M2[2,4] == 0.25
end

## Larger test

const t = 0:0.001:200
const dt_large = t.step.hi
const X_large = exp.(-t) .*
    cumsum(exp.(t) .* (sqrt(dt_large)*pushfirst!(randn(length(t)-1),0)))
const ob_large = Observation(X_large,dt_large)

@testset "Larger observation" begin
    @test ob_large[ob_large .> 3] == X_large[X_large .> 3]
    @test ob_large[(ob_large .> 3)'] == X_large[X_large .> 3]
    @test ob_large[:] == X_large
    @test ob_large.dt == dt_large
    @test ob_large.npoints == length(t)
end

const τ_indices_large = 1:20
const τ_vec = dt_large*τ_indices_large
const bin_edges_large = LinRange(-1,1,10)
const histogramSettings_large =
        HistogramSettings(τ_indices_large,bin_edges_large)
const moments_large = MomentSolution(ob_large, histogramSettings_large)

const M1_τ = moments_large.moments.M1 ./ τ_vec
const M2_τ = moments_large.moments.M2 ./ (2*τ_vec)

@testset "Larger moments" begin
    @test moments_large.observation.npoints == length(t)
    expected_size = (length(τ_indices_large),length(bin_edges_large)-1)
    @test size(moments_large.moments.M1) == expected_size
    @test size(moments_large.moments.M2) == expected_size
    #@test size(moments_large.errors) == expected_size
end

#NOTE: Not type stable
const model1_large = WhiteNoiseProcess(moments_large)
const model2_large = WhiteNoiseProcess(moments_large,i=2)

@testset "Larger drift and noise functions" begin
    expected_length = length(bin_edges_large)-1
    @test size(model1_large.drift) == (expected_length,)
    @test size(model1_large.noise) == (expected_length,)
    #@test size(model1_large.errors) == (expected_length,)
    @test size(model2_large.drift) == (expected_length,)
    @test size(model2_large.noise) == (expected_length,)
    #@test size(model2_large.errors) == (expected_length,)
    @test ~isnan(model1_large.drift[4])
    @test ~isnan(model2_large.drift[4])
    @test ~isnan(model1_large.noise[4])
    @test ~isnan(model2_large.noise[4])
end


# This file is part of the Taylor1Series.jl Julia package, MIT license
#
# Luis Benet & David P. Sanders
# UNAM
#
# MIT Expat license
#

"""
    AbstractSeries{T<:Number} <: Number

Parameterized abstract type for [`Taylor1`](@ref),
[`HomogeneousPolynomial`](@ref) and [`TaylorN`](@ref).
"""
abstract type AbstractSeries{T<:Number} <: Number end

## Constructors ##

######################### Taylor1
"""
    Taylor1{T<:Number} <: AbstractSeries{T}

DataType for polynomial expansions in one independent variable.

**Fields:**

- `coeffs :: Array{T,1}` Expansion coefficients; the ``i``-th
    component is the coefficient of degree ``i-1`` of the expansion.
- `order  :: Int` Maximum order (degree) of the polynomial.

Note that `Taylor1` variables are callable. For more information, see
[`evaluate`](@ref).
"""
struct Taylor1{T<:Number} <: AbstractSeries{T}
    coeffs :: Array{T,1}
    order :: Int

    ## Inner constructor ##
    function Taylor1{T}(coeffs::Array{T,1}, order::Int) where T<:Number
        resize_coeffs1!(coeffs, order)
        return new{T}(coeffs, order)
    end
end

## Outer constructors ##
Taylor1(x::Taylor1{T}) where {T<:Number} = x
Taylor1(coeffs::Array{T,1}, order::Int) where {T<:Number} = Taylor1{T}(coeffs, order)
Taylor1(coeffs::Array{T,1}) where {T<:Number} = Taylor1(coeffs, length(coeffs)-1)
function Taylor1(x::T, order::Int) where {T<:Number}
    v = fill(zero(x), order+1)
    v[1] = x
    return Taylor1(v, order)
end

# Methods using 1-d views to create Taylor1's
Taylor1(a::SubArray{T,1}, order::Int) where {T<:Number} = Taylor1(a.parent[a.indices...], order)
Taylor1(a::SubArray{T,1}) where {T<:Number} = Taylor1(a.parent[a.indices...])

# Shortcut to define Taylor1 independent variables
"""
    Taylor1([T::Type=Float64], order::Int)

Shortcut to define the independent variable of a `Taylor1{T}` polynomial of
given `order`. The default type for `T` is `Float64`.

```julia
julia> Taylor1(16)
 1.0 t + 𝒪(t¹⁷)

julia> Taylor1(Rational{Int}, 4)
 1//1 t + 𝒪(t⁵)
```
"""
Taylor1(::Type{T}, order::Int) where {T<:Number} = Taylor1( [zero(T), one(T)], order)
Taylor1(order::Int) = Taylor1(Float64, order)

######################### HomogeneousPolynomial
"""
    HomogeneousPolynomial{T<:Number} <: AbstractSeries{T}

DataType for homogenous polynomials in many (>1) independent variables.

**Fields:**

- `coeffs  :: Array{T,1}` Expansion coefficients of the homogeneous
polynomial; the ``i``-th component is related to a monomial, where the degrees
of the independent variables are specified by `coeff_table[order+1][i]`.
- `order   :: Int` order (degree) of the homogenous polynomial.

Note that `HomogeneousPolynomial` variables are callable. For more information,
see [`evaluate`](@ref).
"""
struct HomogeneousPolynomial{T<:Number} <: AbstractSeries{T}
    coeffs  :: Array{T,1}
    order   :: Int

    function HomogeneousPolynomial{T}(coeffs::Array{T,1}, order::Int) where T<:Number
        resize_coeffsHP!(coeffs, order)
        return new{T}(coeffs, order)
    end
end

HomogeneousPolynomial(x::HomogeneousPolynomial{T}) where {T<:Number} = x
HomogeneousPolynomial(coeffs::Array{T,1}, order::Int) where {T<:Number} =
    HomogeneousPolynomial{T}(coeffs, order)
HomogeneousPolynomial(coeffs::Array{T,1}) where {T<:Number} =
    HomogeneousPolynomial(coeffs, orderH(coeffs))
HomogeneousPolynomial(x::T, order::Int) where {T<:Number} =
    HomogeneousPolynomial([x], order)

# Shortcut to define HomogeneousPolynomial independent variable
"""
    HomogeneousPolynomial([T::Type=Float64], nv::Int])

Shortcut to define the `nv`-th independent `HomogeneousPolynomial{T}`.
The default type for `T` is `Float64`.

```julia
julia> HomogeneousPolynomial(1)
 1.0 x₁

julia> HomogeneousPolynomial(Rational{Int}, 2)
 1//1 x₂
```
"""
function HomogeneousPolynomial(::Type{T}, nv::Int) where {T<:Number}
    @assert 0 < nv ≤ get_numvars()
    v = zeros(T, get_numvars())
    @inbounds v[nv] = one(T)
    return HomogeneousPolynomial(v, 1)
end
HomogeneousPolynomial(nv::Int) = HomogeneousPolynomial(Float64, nv)

######################### TaylorN
"""
    TaylorN{T<:Number} <: AbstractSeries{T}

DataType for polynomial expansions in many (>1) independent variables.

**Fields:**

- `coeffs  :: Array{HomogeneousPolynomial{T},1}` Vector containing the
`HomogeneousPolynomial` entries. The ``i``-th component corresponds to the
homogeneous polynomial of degree ``i-1``.
- `order   :: Int`  maximum order of the polynomial expansion.

Note that `TaylorN` variables are callable. For more information, see
[`evaluate`](@ref).
"""
struct TaylorN{T<:Number} <: AbstractSeries{T}
    coeffs  :: Array{HomogeneousPolynomial{T},1}
    order   :: Int

    function TaylorN{T}(v::Array{HomogeneousPolynomial{T},1}, order::Int) where T<:Number
        coeffs = zeros(HomogeneousPolynomial{T}, order)
        @inbounds for i in eachindex(v)
            ord = v[i].order
            if ord ≤ order
                coeffs[ord+1] += v[i]
            end
        end
        new{T}(coeffs, order)
    end
end

TaylorN(x::TaylorN{T}) where {T<:Number} = x
function TaylorN(x::Array{HomogeneousPolynomial{T},1}, order::Int) where {T<:Number}
    if order == 0
        order = maxorderH(x)
    end
    return TaylorN{T}(x, order)
end
TaylorN(x::Array{HomogeneousPolynomial{T},1}) where {T<:Number} =
    TaylorN(x, maxorderH(x))
TaylorN(x::HomogeneousPolynomial{T}, order::Int) where {T<:Number} =
    TaylorN( [x], order )
TaylorN(x::HomogeneousPolynomial{T}) where {T<:Number} = TaylorN(x, x.order)
TaylorN(x::T, order::Int) where {T<:Number} =
    TaylorN(HomogeneousPolynomial([x], 0), order)

# Shortcut to define TaylorN independent variables
"""
    TaylorN([T::Type=Float64], nv::Int; [order::Int=get_order()])

Shortcut to define the `nv`-th independent `TaylorN{T}` variable as a
polynomial. The order is defined through the keyword parameter `order`,
whose default corresponds to `get_order()`. The default of type for
`T` is `Float64`.

```julia
julia> TaylorN(1)
 1.0 x₁ + 𝒪(‖x‖⁷)

julia> TaylorN(Rational{Int},2)
 1//1 x₂ + 𝒪(‖x‖⁷)
```
"""
TaylorN(::Type{T}, nv::Int; order::Int=get_order()) where {T<:Number} =
    TaylorN( HomogeneousPolynomial(T, nv), order )
TaylorN(nv::Int; order::Int=get_order()) = TaylorN(Float64, nv, order=order)

# A `Number` which is not an `AbstractSeries`
const NumberNotSeries = Union{Real,Complex}

# A `Number` which is not `TaylorN` nor a `HomogeneousPolynomial`
const NumberNotSeriesN = Union{Real,Complex,Taylor1}

## Additional Taylor1 outer constructor ##
Taylor1{T}(x::S) where {T<:Number,S<:NumberNotSeries} = Taylor1([convert(T,x)], 0)


using Luxor, Test, Random
Random.seed!(42)
function circlecircletangent_test(fname)
    Drawing(800, 800, fname)
    origin()
    background("white")
    setopacity(0.7)
    sethue("black")
    # overlapping circles will return four Os
    c1center = c2center = O
    c1radius = 50
    c2radius = 100
    p1, p2, p3, p4 = circlecircleoutertangents(c1center, c1radius, c2center, c2radius)
    @test all(pt -> isequal(pt, O), (p1, p2, p3, p4)) == true
    for i in 1:30
        randomhue()
        c1center = rand(BoundingBox() * 0.5)
        c2center = rand(BoundingBox() * 0.75)
        c1radius = rand(5:30)
        c2radius = rand(30:50)
        circle(c1center, c1radius, :fill)
        circle(c2center, c2radius, :fill)
        pt1, pt2, pt3, pt4 = circlecircleoutertangents(c1center, c1radius, c2center, c2radius)
        poly([pt1, pt2, pt4, pt3], :fill)
        sethue("black")
        circle(c2center, c2radius, :stroke)
        circle(c1center, c1radius, :stroke)
    end
    @test finish() == true
    println("...finished test: output in $(fname)")
    preview()
end

function circlecircletangent_test2(fname)
    Drawing(800, 800, fname)
    origin()
    background("white")
    setline(0.5)
    sethue("black")
    # overlapping circles will return four Os
    c1center = c2center = O
    c1radius = 50
    c2radius = 100
    p1, p2, p3, p4 = circlecircleinnertangents(c1center, c1radius, c2center, c2radius)
    @test all(pt -> isequal(pt, O), (p1, p2, p3, p4)) == true

    for θ in range(0, step=2π/12, length=12)
        @layer begin
            randomhue()
            c1center = polar(100, θ)
            c2center = polar(300, θ)
            c1radius = 20
            c2radius = rand(50:70)
            circle(c1center, c1radius, :fill)
            circle(c2center, c2radius, :fill)

            pt1, pt2, pt3, pt4 = circlecircleinnertangents(c1center, c1radius, c2center, c2radius)

            if !all(pt -> isequal(pt, O), (pt1, pt2, pt3, pt4))
                sethue("green")
                circle(c2center, c2radius, :stroke)
                circle(c1center, c1radius, :stroke)
                circle.((pt1, pt2, pt3, pt4), 4, :fill)
                sethue("black")
                arrow(pt1, pt2, linewidth=2)
                arrow(pt3, pt4, linewidth=2)
            end

            pt1, pt2, pt3, pt4 = circlecircleoutertangents(c1center, c1radius, c2center, c2radius)

            if !all(pt -> isequal(pt, O), (pt1, pt2, pt3, pt4))
                sethue("red")
                arrow(pt1, pt2, linewidth=2)
                arrow(pt3, pt4, linewidth=2)
                circle.((pt1, pt2, pt3, pt4), 3, :fill)
            end
        end
    end
    @test finish() == true
    println("...finished test: output in $(fname)")
end

circlecircletangent_test("circlecircletangent-outer.png")
circlecircletangent_test2("circlecircletangent-inner.png")


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

using LinearAlgebra, InvertibleNetworks, Test, Flux
import Flux.Optimise.update!

# Input
nx = 28
ny = 28
n_in = 4
n_hidden = 8
batchsize = 2

# Flux networks
model = Chain(
    Conv((3,3), n_in => n_hidden; pad=1),
    Conv((3,3), n_hidden => n_in; pad=1)
)

model0 = Chain(
    Conv((3,3), n_in => n_hidden; pad=1),
    Conv((3,3), n_hidden => n_in; pad=1)
)

# Input
X = glorot_uniform(nx, ny, n_in, batchsize)
X0 = glorot_uniform(nx, ny, n_in, batchsize)
dX = X - X0

# Flux blocks
RB = FluxBlock(model)
RB0 = FluxBlock(model0)

# Observed data
Y = RB.forward(X)

function loss(RB, X, Y)
    Y_ = RB.forward(X)
    ΔY = Y_ - Y
    f = .5f0*norm(ΔY)^2
    ΔX = RB.backward(ΔY, X)
    return f, ΔX, RB.params[1].grad, RB.params[2].grad
end

###################################################################################################
# Gradient tests

# Gradient test w.r.t. input
f0, ΔX = loss(RB, X0, Y)[1:2]
h = 0.1f0
maxiter = 5
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test convolutions\n")
for j=1:maxiter
    f = loss(RB, X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

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

# Gradient test for weights
RB_ini = deepcopy(RB0)
dW1 = RB.params[1].data - RB0.params[1].data
dW2 = RB.params[2].data - RB0.params[2].data

f0, ΔX, ΔW1, ΔW2 = loss(RB0, X, Y)
h = 0.1f0
maxiter = 5
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)

print("\nGradient test convolutions\n")
for j=1:maxiter
    RB0.params[1].data[:] = RB_ini.params[1].data + h*dW1
    RB0.params[2].data[:] = RB_ini.params[2].data + h*dW2
    f = loss(RB0, X, Y)[1]
    err3[j] = abs(f - f0)
    err4[j] = abs(f - f0 - h*dot(dW1, ΔW1) - h*dot(dW2, ΔW2))
    print(err3[j], "; ", err4[j], "\n")
    global h = h/2f0
end

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


## Signal creation

# With no elements
sig = Signal(1:100)
@test length(sig) == 0
@test isempty(collect(sig))
@test sig.time == 1s:100s
@test samplingfreq(sig) == 1/s
@test samplingrate(sig) == 1s
@test_throws BoundsError sig[0]
@test_throws BoundsError sig[1]

# A matrix of 8 channels
sig = Signal(1:100, reshape(1:100*8, 100, 8))
@test length(sig) == 8
@test sig[1] == [1:100]
@test all([sig[i] == [(1:100)+100*(i-1)] for i=1:length(sig)])
@test length(collect(sig)) == 8
@test_throws BoundsError sig[0]
@test_throws BoundsError sig[9]
sig2 = reshape(sig, 4, 2)
@test collect(sig) == collect(sig2)

sig = Signal(1:100, [1:100 [1.:100.]*2 [1:100]*4])
@test length(sig) == 3
@test sig[1]*4 == sig[2]*2 == sig[3] == [1:100]*4
@test_throws BoundsError sig[0]
@test_throws BoundsError sig[4]

# A Signal with a group of nested Signals (like spikes with snippets)
t1 = (-.2:.025:.8)./1e3
snips = Signal(t1, randn(length(t1), 100))
@test length(snips) == 100
@test length(snips[1]) == length(snips.time) == length(t1)
@test isregular(snips)
@test samplingfreq(snips) == 40_000.0/s
@test samplingrate(snips) == (2.5e-5)s
spikes = Signal(cumsum(rand(100)), snips)
@test length(spikes) == 1
@test spikes[1] == snips
@test !isregular(spikes)

# Ensure we cannot build a Signal with bad time vectors
@test_throws ArgumentError Signal([1,2,3,2])
@test_throws ArgumentError Signal([1.,2.,2.])

## Regularization:
sig = Signal([0:1/100:2pi], [sin, cos])
@test !isregular(sig)
@test isregular(regularize(sig))

# Indexing by vectors returns Signals
sig = Signal(1:100, reshape(1:100*8, 100, 8))
subsig = sig[2:3]
@test length(subsig) == 2
@test subsig[1] == sig[2]
@test subsig[2] == sig[3]


function rsch_integ_rk4!(
    E::Float64,
    Z::Int64, l::Int64, rmesh, V, P, Q;
    MAXVAL_STOP=1e6
)
    Nr = size(rmesh,1)

    # Calculate Vmid
    Vmid = zeros(Float64, Nr-1)
    for i in 1:Nr-1
        rmid = 0.5*( rmesh[i] + rmesh[i+1] )
        Vmid[i] = radial_interp(rmesh, V, rmid, i+1)
    end

    #@printf("Some Vmid\n")
    #for i in 1:4
    #    @printf("%8d %20.10e\n", i, Vmid[i])
    #end
    #@printf("...\n")
    #for i in Nr-4:Nr-1
    #    @printf("%8d %20.10e\n", i, Vmid[i])
    #end

    # BC at r -> 0
    y0 = zeros(Float64,2)
    if l == 0
        y0[1] = 1.0 - Z*rmesh[1]
        y0[2] = -Z
    else
        y0[1] = rmesh[1]^l
        y0[2] = l*rmesh[1]^(l-1)
    end

    C1 = zeros(Float64,Nr)
    C2 = zeros(Float64,Nr)
    for i in 1:Nr
        C1[i] = 2*(V[i] - E) + l*(l + 1)/rmesh[i]
        C2[i] = -2/rmesh[i]
    end

    C1mid = zeros(Float64,Nr-1)
    C2mid = zeros(Float64,Nr-1)
    for i in 1:Nr-1
        rmid = 0.5*(rmesh[i] + rmesh[i+1])
        C1mid[i] = 2*(Vmid[i] - E) + l*(l + 1)/rmid^2
        C2mid[i] = -2/rmid
    end

    y1 = zeros(Float64,Nr)
    y2 = zeros(Float64,Nr)
    imax = sch_rk4_step!( rmesh, y0, C1, C2, C1mid, C2mid, y1, y2, MAXVAL_STOP )

    for i in 1:imax
        P[i] = y1[i]*rmesh[i] # P(r) = r*R(r)
        Q[i] = y2[i]*rmesh[i] + y1[i]
    end

    return imax
end


### estimators for π0 (pi0) ###

## Storey estimator ##

"""
Storey π0 estimator

**Parameters**

- λ : tuning parameter, FloatingPoint, default: 0.1

**Examples**

```julia
Storey()
Storey(0.1)
```

**References**

Storey, JD (2002). "A Direct Approach to False Discovery Rates." Journal of the
Royal Statistical Society, doi:10.1111/1467-9868.00346

"""
type Storey <: Pi0Estimator
    λ::AbstractFloat

    Storey(λ) = isin(λ, 0., 1.) ? new(λ) : throw(DomainError())
end

Storey() = Storey(0.1)

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::Storey)
    storey_pi0(pValues, pi0estimator.λ)
end

function storey_pi0{T<:AbstractFloat}(pValues::Vector{T}, lambda::T)
    validPValues(pValues)
    pi0 = (sum(pValues .>= lambda) / length(pValues)) / (1.-lambda)
    pi0 = min(pi0, 1.)
    return pi0
end

## Storey bootstrap estimator ##
"""
Storey closed-form bootstrap π0 estimator

StoreyBootstrap(λseq, q)

Reference: David Robinson, 2012
"""
type StoreyBootstrap <: Pi0Estimator
    ## check range of arguments
    λseq::Vector{AbstractFloat}
    q   ::AbstractFloat

    StoreyBootstrap(λseq, q) =
        isin(λseq, 0., 1.) && isin(q, 0., 1.) ? new(λseq, q) : throw(DomainError())
end

StoreyBootstrap() = StoreyBootstrap(collect(0.05:0.05:0.95), 0.1)

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::StoreyBootstrap)
    bootstrap_pi0(pValues, pi0estimator.λseq, pi0estimator.q)
end

function bootstrap_pi0{T<:AbstractFloat,S<:AbstractFloat}(pValues::Vector{T}, lambda::Vector{S} = [0.05:0.05:0.95;], q::S = 0.1)
    #validPValues(pValues)
    n = length(pValues)
    w = Int[sum(pValues .>= l) for l in lambda] ## TODO: check >= or >
    pi0 = w ./ n ./ (1. - lambda)
    min_pi0 = quantile(pi0, q)
    mse = (w ./ (n.^2 .* (1. - lambda).^2 )) .* (1. - w/n) + (pi0 - min_pi0).^2
    pi0 = min(pi0[indmin(mse)], 1.)
    pi0
end

## Least SLope (LSL) estimator ##

"""
Least SLope (LSL) π0 estimator

LeastSlope()
"""
type LeastSlope <: Pi0Estimator
end

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::LeastSlope)
    lsl_pi0(pValues)
end

function lsl_pi0{T<:AbstractFloat}(pValues::Vector{T})
    n = length(pValues)
    if !issorted(pValues)
        pValues = sort(pValues)
    end
    s0 = lsl_slope(1, n, pValues)
    sx = 0.
    for i in 2:n
        s1 = lsl_slope(i, n, pValues)
        if (s1 - s0) < 0.
            sx = s1
            break
        end
        s0 = s1
    end
    pi0 = min( 1/sx + 1, n ) / n
    return(pi0)
end

function lsl_slope{T<:AbstractFloat}(i::Int, n::Int, pval::Vector{T})
    s = (1 - pval[i]) / (n - i + 1)
    return s
end

## alternative, vectorized version
## used for comparison and compactness
function lsl_pi0_vec{T<:AbstractFloat}(pValues::Vector{T})
    n = length(pValues)
    if !issorted(pValues)
        pValues = sort(pValues)
    end
    s = (1 - pValues) ./ (n - collect(1:n) + 1)
    d = diff(s) .< 0
    idx = findfirst(d) + 1
    pi0 = min( 1/s[idx] + 1, n ) / n
    return(pi0)
end

## Oracle
"""
Oracle π0

Oracle(π0)
"""
type Oracle <: Pi0Estimator
    π0::AbstractFloat

    Oracle(π0) = isin(π0, 0., 1.) ? new(π0) : throw(DomainError())
end

Oracle() = Oracle(1.0)

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::Oracle)
    pi0estimator.π0
end

## Two-Step estimator: Benjamini, Krieger and Yekutieli (2006) ##

"""
Two step π0 estimator

TwoStep(α)

Reference: Benjamini, Krieger and Yekutieli, 2006
"""
type TwoStep <: Pi0Estimator
    α::AbstractFloat
    ## method::PValueAdjustmentMethod

    TwoStep(α) = isin(α, 0., 1.) ? new(α) : throw(DomainError())
end

TwoStep() = TwoStep(0.05)

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::TwoStep)
    twostep_pi0(pValues, pi0estimator.α)
end

function twostep_pi0{T<:AbstractFloat}(pValues::Vector{T}, alpha::T, method::PValueAdjustmentMethod = BenjaminiHochberg())
    padj = adjust(pValues, method)
    pi0 = sum(padj .>= (alpha/(1+alpha))) / length(padj)
    return(pi0)
end

# RightBoundary procedure as defined in Definition 2 of Liang and Nettleton 2012
# "Adaptive and dynamic adaptive procedures for false discovery rate control and estimation"
"""
Right boundary π0 estimator

RightBoundary(λseq)
"""
type RightBoundary <: Pi0Estimator
    ## check range of arguments
    λseq::Vector{Float64}

    RightBoundary(λseq) =
        isin(λseq, 0., 1.) ? new(λseq) : throw(DomainError())
end

# λseq used in Liang, Nettleton 2012
RightBoundary() = RightBoundary([collect(0.02:0.02:0.1); collect(0.15:0.05:0.95)])

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::RightBoundary)
    rightboundary_pi0(pValues, pi0estimator.λseq)
end

function rightboundary_pi0{T<:AbstractFloat}(pValues::Vector{T}, λseq)
    validPValues(pValues)
    n = length(pValues)
    if !issorted(λseq)
        λseq = sort(λseq)
    end
    # make sure we catch p-values equal to 1 despite left closure
    # use closed=:left because we have been using >= convention in this package
    # note that original paper uses > convention.
    h = fit(Histogram, pValues, [λseq; Inf], closed=:left)
    pi0_estimates = reverse(cumsum(reverse(h.weights)))./(1.-λseq)./n
    pi0_decrease = diff(pi0_estimates) .>= 0
    pi0_decrease[end] = true
    pi0 = pi0_estimates[findfirst(pi0_decrease, true) + 1]
    return(min(pi0,1))
end

## Censored BUM
"""
Censored BUM π0 estimator

CensoredBUM(γ0, λ, xtol, maxiter)
"""
type CensoredBUM <: Pi0Estimator
    γ0::Float64
    λ::Float64
    xtol::Float64
    maxiter::Int64

    function CensoredBUM(γ0, λ, xtol, maxiter)
        if isin(γ0, 0., 1.) && isin(λ, 0., 1.) && isin(xtol, 0., 1.) && maxiter > 0
            new(γ0, λ, xtol, maxiter)
        else
            throw(DomainError())
        end
    end
end

CensoredBUM() = CensoredBUM(0.5, 0.05, 1e-6, 10000)
CensoredBUM(γ0, λ) = CensoredBUM(γ0, λ, 1e-6, 10000)

type CensoredBUMFit <: Pi0Fit
    π0::Float64
    param::Vector{Float64}
    is_converged::Bool
end

function fit(pi0estimator::CensoredBUM, pValues;
             kw...)
    π0, param, is_converged = cbum_pi0(pValues, pi0estimator.γ0, pi0estimator.λ,
                                       pi0estimator.xtol, pi0estimator.maxiter;
                                       kw...)
    return CensoredBUMFit(π0, param, is_converged)
end

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::CensoredBUM)
    estimate_pi0(fit(pi0estimator, pValues))
end

function estimate_pi0(pi0fit::CensoredBUMFit)
    π0 = pi0fit.is_converged ? π0 = pi0fit.π0 : NaN
    return π0
end

function cbum_pi0{T<:AbstractFloat}(pValues::Vector{T}, γ0::T = 0.5, λ::T = 0.05,
                                    xtol::Float64 = 1e-6, maxiter::Int = 10000)
    n = length(pValues)
    idx_right = pValues .>= λ
    n2 = sum(idx_right)
    n1 = n - n2
    sz = (1-γ0)*n
    szr = (1-γ0)*n2
    szl = sz - szr
    ## compute constant values only once
    pr = pValues[idx_right]
    lpr = log(pr)
    ll = log(λ)
    zr = fill(1-γ0, n2)
    pi0_old = γ0 = α = γ = Inf
    for i in 1:maxiter
        γ = 1 - sz/n
        α = -szr / ( ll * szl + sum(zr .* lpr) )
        γ = max(min(γ, 1.), 0.)
        α = max(min(α, 1.), 0.)
        xl = (1-γ) * (λ^α)
        szl = (xl ./ (γ*λ + xl)) * n1
        xr = (1-γ) * α * pr.^(α-1)
        zr = xr ./ (γ + xr)
        szr = sum(zr)
        sz = szl + szr
        pi0_new = γ + (1-γ)*α
        if abs(pi0_new - pi0_old) <= xtol
            return pi0_new, [γ, α], true
        end
        γ0 = γ
        pi0_old = pi0_new
    end
    return NaN, [γ, α], false
end

function cbum_pi0_naive(pValues, γ0 = 0.5, λ = 0.05,
                        xtol = 1e-6, maxiter = 10000)
    n = length(pValues)
    z = fill(1-γ0, n)
    idx_left = pValues .< λ
    idx_right = !idx_left
    pi0_old = γ0 = α = γ = Inf
    ## compute constant values only once
    lpr = log(pValues[idx_right])
    ll = log(λ)
    for i in 1:maxiter
        γ = sum(1-z) / n
        α = -sum(z[idx_right])
        α = α / ( ll * sum(z[idx_left]) + sum(z[idx_right] .* lpr) )
        xl = (1-γ) * (λ^α)
        z[idx_left] = xl ./ (γ*λ + xl)
        xr = (1-γ) * α * pValues[idx_right].^(α-1)
        z[idx_right] = xr ./ (γ + xr)
        pi0_new = γ + (1-γ)*α
        if abs(pi0_new - pi0_old) <= xtol
            return pi0_new, [γ, α], true
        end
        γ0 = γ
        pi0_old = pi0_new
    end
    return NaN, [γ, α], false
end

## BUM
"""
BUM π0 estimator

BUM(γ0, xtol, maxiter)
"""
type BUM <: Pi0Estimator
    γ0::Float64
    xtol::Float64
    maxiter::Int64

    function BUM(γ0, xtol, maxiter)
        if isin(γ0, 0., 1.) && isin(xtol, 0., 1.)
            new(γ0, xtol, maxiter)
        else
            throw(DomainError())
        end
    end
end

BUM() = BUM(0.5, 1e-6, 10000)

BUM(y0::Float64) = BUM(y0, 1e-6, 10000)

type BUMFit <: Pi0Fit
    π0::Float64
    param::Vector{Float64}
    is_converged::Bool
end

function fit(pi0estimator::BUM, pValues;
             kw...)
    π0, param, is_converged = cbum_pi0(pValues, pi0estimator.γ0, eps(),
                                       pi0estimator.xtol, pi0estimator.maxiter;
                                       kw...)
    return BUMFit(π0, param, is_converged)
end

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::BUM)
    estimate_pi0(fit(pi0estimator, pValues))
end

function estimate_pi0(pi0fit::BUMFit)
    π0 = pi0fit.is_converged ? π0 = pi0fit.π0 : NaN
    return π0
end

## Longest constant interval in the Grenander estimator: Langaas et al., 2005
# Leave unexported until thoroughly tested

"""
Flat Grenander π0 estimator

FlatGrenander()

Estimates π0 by the longest constant interval in the Grenander estimator

Reference: Langaas et al., 2005: section 4.3
"""
type FlatGrenander <: Pi0Estimator
end

function estimate_pi0{T<:AbstractFloat}(pValues::Vector{T}, pi0estimator::FlatGrenander)
    flat_grenander_pi0(pValues)
end

function flat_grenander_pi0(pv::Vector{Float64})
    p, f, F = grenander(pv)
    pi0 = longest_constant_interval(p, f)
    return pi0
end

function longest_constant_interval(p::Vector{Float64}, f::Vector{Float64})
    p = [0.0; p]
    f = [Inf; f]
    i2 = length(f)
    Δp_max = Δp = 0.0
    pi0 = 1.0
    for i1 in (length(f)-1):-1:1
        if f[i2] ≈ f[i1] ## within constant interval
            Δp = p[i2] - p[i1]
        else
            if Δp >= Δp_max
                Δp_max = Δp
                pi0 = f[i2]
            end
            i2 = i1
            Δp = 0.0
        end
        if f[i1] > 1.0
            break
        end
    end
    return pi0
end


"""
    convergence_rates(problem, solution, method, dts; kwargs...)

Compute the errors rates of `method` on `problem` by comparing to `solution`
on the set of `dt` values in `dts`. Extra `kwargs` are passed to `solve`

`solution` should be a function with a method `solution(u0, p, t)`.
"""
function convergence_errors(prob, sol, method, dts; kwargs...)
  errs = map(dts) do dt
       # copy the problem so we don't mutate u0
      prob_copy = deepcopy(prob)
      u = solve(prob_copy, method; dt=dt, kwargs...)
      norm(u .- sol(prob.u0, prob.p, prob.tspan[end]))
  end
  return errs
end

"""
  convergence_order(problem, solution, method, dts; kwargs...)

Estimate the order of the rate of convergence of `method` on `problem` by comparing to
`solution` the set of `dt` values in `dts`.

`solution` should be a function with a method `solution(u0, p, t)`.
"""
function convergence_order(prob, sol, method, dts; kwargs...)
  errs = convergence_errors(prob, sol, method, dts; kwargs...)
  # find slope coefficient in log scale
  _,order_est = hcat(ones(length(dts)), log2.(dts)) \ log2.(errs)
  return order_est
end

"""
    DirectSolver

A linear solver which forms the full matrix of a linear operator and its LU factorization.
"""
struct DirectSolver end

DirectSolver(args...) = DirectSolver()

function (::DirectSolver)(x,A,b,matrix_updated; kwargs...)
  n = length(x)
  M = mapslices(y -> mul!(similar(y), A, y), Matrix{eltype(x)}(I,n,n), dims=1)
  x .= M \ b
end


"""
    NormInfinityBridge{T}

The `NormInfinityCone` is representable with LP constraints, since
``t \\ge \\max_i \\lvert x_i \\rvert`` if and only if
``t \\ge x_i`` and ``t \\ge -x_i`` for all ``i``.
"""
struct NormInfinityBridge{T,F,G} <:
       SetMapBridge{T,MOI.Nonnegatives,MOI.NormInfinityCone,F,G}
    constraint::CI{F,MOI.Nonnegatives}
end
function concrete_bridge_type(
    ::Type{<:NormInfinityBridge{T}},
    G::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.NormInfinityCone},
) where {T}
    F = MOIU.promote_operation(+, T, G, G)
    return NormInfinityBridge{T,F,G}
end

function map_set(::Type{<:NormInfinityBridge}, set::MOI.NormInfinityCone)
    return MOI.Nonnegatives(2 * (MOI.dimension(set) - 1))
end
function inverse_map_set(::Type{<:NormInfinityBridge}, set::MOI.Nonnegatives)
    return MOI.NormInfinityCone(div(MOI.dimension(set), 2) + 1)
end

function map_function(::Type{<:NormInfinityBridge{T}}, func) where {T}
    scalars = MOIU.eachscalar(func)
    t = scalars[1]
    lb = scalars[2:end]
    ub = MOIU.operate(-, T, lb)
    f_new = MOIU.operate(vcat, T, ub, lb)
    for i in 1:(2*(length(scalars)-1))
        MOIU.operate_output_index!(+, T, i, f_new, t)
    end
    return f_new
end
function inverse_map_function(::Type{<:NormInfinityBridge{T}}, func) where {T}
    scalars = MOIU.eachscalar(func)
    t = MOIU.operate!(/, T, sum(scalars), T(length(scalars)))
    d = div(length(scalars), 2)
    x = MOIU.operate!(
        /,
        T,
        MOIU.operate!(-, T, scalars[(d+1):end], scalars[1:d]),
        T(2),
    )
    return MOIU.operate(vcat, T, t, x)
end
# Given a_i is dual on t - x_i >= 0 and b_i is dual on t + x_i >= 0,
# the dual on (t, x) in NormInfinityCone is (u, v) in NormOneCone, where
# v_i = -a_i + b_i and u = sum(a) + sum(b).
function adjoint_map_function(::Type{<:NormInfinityBridge}, func)
    scalars = MOIU.eachscalar(func)
    t = sum(scalars)
    d = div(length(scalars), 2)
    x = (scalars[(d+1):end] - scalars[1:d])
    return vcat(t, x)
end
function inverse_adjoint_map_function(
    ::Type{<:NormInfinityBridge{T}},
    func::AbstractVector{T},
) where {T}
    # This is used by `MOI.ConstraintDualStart`.
    # The result should belong to `MOI.Nonnegatives` and the sum of the elements should
    # be `t`.
    # Only one of dual of the upper and lower bound should be active so one of the two
    # duals is zero. We know which one since they should be nonnegative.
    # Then if `t = sum abs(x_i)`, we will indeed have only one of them being zero.
    t = func[1]
    y = func[2:end]
    lb = [x > 0 ? x : zero(x) for x in y]
    ub = [x < 0 ? -x : zero(x) for x in y]
    x = [ub; lb]
    return x .+ (t - sum(x)) / length(x)
end

"""
    NormOneBridge{T}

The `NormOneCone` is representable with LP constraints, since
``t \\ge \\sum_i \\lvert x_i \\rvert`` if and only if there exists a vector y such that
``t \\ge \\sum_i y_i`` and ``y_i \\ge x_i``, ``y_i \\ge -x_i`` for all ``i``.
"""
struct NormOneBridge{T,F,G} <: AbstractBridge
    y::Vector{MOI.VariableIndex}
    nn_index::CI{F,MOI.Nonnegatives}
end
function bridge_constraint(
    ::Type{NormOneBridge{T,F,G}},
    model::MOI.ModelLike,
    f::G,
    s::MOI.NormOneCone,
) where {T,F,G}
    f_scalars = MOIU.eachscalar(f)
    d = MOI.dimension(s)
    y = MOI.add_variables(model, d - 1)
    ge = MOIU.operate(-, T, f_scalars[1], MOIU.operate(sum, T, y))
    lb = f_scalars[2:d]
    ub = MOIU.operate(-, T, lb)
    lb = MOIU.operate!(+, T, lb, MOI.VectorOfVariables(y))
    ub = MOIU.operate!(+, T, ub, MOI.VectorOfVariables(y))
    f_new = MOIU.operate(vcat, T, ge, ub, lb)
    nn_index = MOI.add_constraint(model, f_new, MOI.Nonnegatives(2d - 1))
    return NormOneBridge{T,F,G}(y, nn_index)
end

function MOI.supports_constraint(
    ::Type{NormOneBridge{T}},
    ::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.NormOneCone},
) where {T}
    return true
end
function MOIB.added_constrained_variable_types(::Type{<:NormOneBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:NormOneBridge{T,F}}) where {T,F}
    return [(F, MOI.Nonnegatives)]
end
function concrete_bridge_type(
    ::Type{<:NormOneBridge{T}},
    G::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.NormOneCone},
) where {T}
    S = MOIU.scalar_type(G)
    F = MOIU.promote_operation(
        vcat,
        T,
        MOIU.promote_operation(+, T, S, S),
        MOIU.promote_operation(-, T, S, S),
    )
    return NormOneBridge{T,F,G}
end

# Attributes, Bridge acting as a model
MOI.get(b::NormOneBridge, ::MOI.NumberOfVariables) = length(b.y)
MOI.get(b::NormOneBridge, ::MOI.ListOfVariableIndices) = b.y
function MOI.get(
    b::NormOneBridge{T,F},
    ::MOI.NumberOfConstraints{F,MOI.Nonnegatives},
) where {T,F}
    return 1
end
function MOI.get(
    b::NormOneBridge{T,F},
    ::MOI.ListOfConstraintIndices{F,MOI.Nonnegatives},
) where {T,F}
    return [b.nn_index]
end

# References
function MOI.delete(model::MOI.ModelLike, c::NormOneBridge)
    MOI.delete(model, c.nn_index)
    return MOI.delete(model, c.y)
end

# Attributes, Bridge acting as a constraint
function MOI.get(
    model::MOI.ModelLike,
    ::MOI.ConstraintFunction,
    c::NormOneBridge{T,F,G},
) where {T,F,G}
    nn_func =
        MOIU.eachscalar(MOI.get(model, MOI.ConstraintFunction(), c.nn_index))
    t = MOIU.operate!(/, T, nn_func[1] + sum(nn_func), T(2))
    d = div(length(nn_func) - 1, 2)
    x = MOIU.operate!(
        /,
        T,
        MOIU.operate!(-, T, nn_func[(d+2):end], nn_func[2:(d+1)]),
        T(2),
    )
    return MOIU.convert_approx(
        G,
        MOIU.remove_variable(MOIU.operate(vcat, T, t, x), c.y),
    )
end
function MOI.get(model::MOI.ModelLike, ::MOI.ConstraintSet, c::NormOneBridge)
    dim = div(
        MOI.dimension(MOI.get(model, MOI.ConstraintSet(), c.nn_index)) + 1,
        2,
    )
    return MOI.NormOneCone(dim)
end
function MOI.supports(
    ::MOI.ModelLike,
    ::Union{MOI.ConstraintPrimalStart,MOI.ConstraintDualStart},
    ::Type{<:NormOneBridge},
)
    return true
end
function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimalStart,
    bridge::NormOneBridge{T},
    value,
) where {T}
    x_value = value[1 .+ (1:length(bridge.y))]
    y_value = abs.(x_value)
    for i in eachindex(bridge.y)
        MOI.set(model, MOI.VariablePrimalStart(), bridge.y[i], y_value[i])
    end
    nn_value = vcat(
        value[1] - reduce(+, y_value, init = zero(T)),
        y_value - x_value,
        y_value + x_value,
    )
    MOI.set(model, attr, bridge.nn_index, nn_value)
    return
end
function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimal,MOI.ConstraintPrimalStart},
    bridge::NormOneBridge,
)
    nn_primal = MOI.get(model, attr, bridge.nn_index)
    t = (nn_primal[1] + sum(nn_primal)) / 2
    d = length(bridge.y)
    x = (nn_primal[(d+2):end] - nn_primal[2:(d+1)]) / 2
    return vcat(t, x)
end
# Given a_i is dual on y_i - x_i >= 0 and b_i is dual on y_i + x_i >= 0 and c is dual on t - sum(y) >= 0,
# the dual on (t, x) in NormOneCone is (u, v) in NormInfinityCone, where
# v_i = -a_i + b_i and u = c.
function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintDual,MOI.ConstraintDualStart},
    bridge::NormOneBridge,
)
    nn_dual = MOI.get(model, attr, bridge.nn_index)
    d = length(bridge.y)
    x = nn_dual[(d+2):end] - nn_dual[2:(d+1)]
    return vcat(nn_dual[1], x)
end
# value[1 + i] = nn_dual[1 + d + i] - nn_dual[1 + i]
# and `nn_dual` is nonnegative. By complementarity slackness, only one of each
# `nn_dual` can be nonzero (except if `x = 0`) so we can set
# depending on the sense of `value[1 + i]`.
function MOI.set(
    model::MOI.ModelLike,
    ::MOI.ConstraintDualStart,
    bridge::NormOneBridge,
    value,
)
    d = length(bridge.y)
    nn_dual = zeros(eltype(value), 2d + 1)
    nn_dual[1] = value[1]
    for i in eachindex(bridge.y)
        v = value[1+i]
        if v < 0
            nn_dual[1+i] = -v
        else
            nn_dual[1+d+i] = v
        end
    end
    MOI.set(model, MOI.ConstraintDualStart(), bridge.nn_index, nn_dual)
    return
end


# Per Lyndon. Needs adaptation, and/or needs redefinition of map to use functions rather
# than AbstractOp.
#function rrule(map, f, xs)
#    # Rather than 3 maps really want 1 multimap
#    ys_and_pullbacks = map(x->rrule(f, x), xs) #Take this to ys = map(f, x)
#    ys = map(first, ys_and_pullbacks)
#    pullbacks = map(last, ys_and_pullbacks)
#    function map_pullback(dys)
#        _call(f, x) = f(x)
#        dfs_and_dxs = map(_call, pullbacks, dys)
#        # but in your case you know it will be NoTangent() so can  skip
#        df = sum(first, dfs_and_dxs)
#        dxs = map(last, dfs_and_dxs)
#        return NoTangent(), df, dxs
#    end
#    return ys, map_pullback
#end
macro scalarmaprule(func, derivative)
    return ChainRulesCore.@strip_linenos quote
        function ChainRulesCore.frule(
            (_, _, $(esc(:ΔA))),
            ::typeof(Base.map),
            ::typeof($(func)),
            $(esc(:A))::GBArray
        )
            $(esc(:Ω)) = map($(esc(func)), $(esc(:A)))
            return $(esc(:Ω)), $(esc(derivative)) .* unthunk($(esc(:ΔA)))
        end
        function ChainRulesCore.rrule(
            ::typeof(Base.map),
            ::typeof($(func)),
            $(esc(:A))::GBArray
        )
            $(esc(:Ω)) = map($(esc(func)), $(esc(:A)))
            function mapback($(esc(:ΔA)))
                NoTangent(), NoTangent(), $(esc(derivative)) .* $(esc(:ΔA))
            end
            return $(esc(:Ω)), mapback
        end
    end
end

function ChainRulesCore.frule(
    (_,_,ΔA),
    ::typeof(map),
    ::typeof(sqrt),
    A::Array
)
    Ω = map(sqrt, A)
    return Ω, inv.(2 .* Ω)
end

#Trig
@scalarmaprule UnaryOps.SIN cos.(A)
@scalarmaprule UnaryOps.COS -sin.(A)
@scalarmaprule UnaryOps.TAN @. 1 + (Ω ^ 2)

#Hyperbolic Trig
@scalarmaprule UnaryOps.SINH cosh.(A)
@scalarmaprule UnaryOps.COSH sinh.(A)
@scalarmaprule UnaryOps.TANH @. 1 - (Ω ^ 2)

@scalarmaprule UnaryOps.MINV -(Ω .^ 2)
@scalarmaprule UnaryOps.EXP Ω

@scalarmaprule UnaryOps.ABS sign.(A)
#Anything that uses MINV fails the isapprox tests :().
# Since in the immortal words of Miha - "FiniteDiff is smarter than you", these shouldn't be enabled.
#@scalarmaprule UnaryOps.ASIN @. inv(sqrt.(1 - A ^ 2))
#@scalarmaprule UnaryOps.ACOS @. inv(sqrt.(1 - A ^ 2))
#@scalarmaprule UnaryOps.ATAN @. inv(1 + A ^ 2)
#@scalarmaprule UnaryOps.SQRT inv.(2 .* Ω)

function frule(
    (_, _, ΔA),
    ::typeof(map),
    ::typeof(UnaryOps.IDENTITY),
    A::GBArray
)
    return (A, ΔA)
end
function rrule(::typeof(map), ::typeof(UnaryOps.IDENTITY), A::GBArray)
    return A, (ΔΩ) -> (NoTangent(), NoTangent(), ΔΩ)
end


@testset "Determinant by minor expansion" begin
    for d = 1:5
        for testcase = 1:10
            mat_space = Nemo.MatrixSpace(Nemo.QQ, d, d)
            rnd_matrix = mat_space([mod(rand(Int), 1000) for i = 1:d, j = 1:d])
            @test det(rnd_matrix) == det_minor_expansion(rnd_matrix)
        end
    end
end


using Pkg
Pkg.activate(joinpath(Pkg.devdir(), "MLCourse"))
using MLCourse, CSV, MLJ, DataFrames, MLJLinearModels, Random, Flux, MLJFlux, MLJXGBoostInterface, Plots
include("./Data_Processing.jl")

#Load the data
drop = generate(option = "drop", std = "fale", valid = "false", test = "true");
med = generate(option = "med", std = "false", valid = "false", test = "true");
drop_std = generate(option = "drop", std = "true", valid = "false", test = "true");
med_std = generate(option = "med", std = "true", valid = "false", test = "true");

#First scheme of possible neural network on med data set and drop data set, standardized or not.
Random.seed!(2711);
NN_1 = machine(NeuralNetworkClassifier(builder = MLJFlux.Short(n_hidden = 200, dropout = 0.1, σ = relu),
                                                batch_size = 1000, epochs = 100),
                                                select(drop_std.train[:,:], Not(:precipitation_nextday)),
                                                drop_std.train.precipitation_nextday) |> fit!;

MLJ.auc(predict(NN_1, select(drop_std.test[:,:], Not(:precipitation_nextday))), drop_std.test.precipitation_nextday)
# AUC_DROP_STD = 0.912
# Relu seems to perform better than sigmoid function.
# Drop seems to perform better than med.
# No particular differences between standardized and normal data.

#Tuned model
#Use this section to try multiple model with changes in some hyper-parametrs in order to see if there is some increase in the AUC
Random.seed!(2711)
model = NeuralNetworkClassifier( builder = MLJFlux.Short(n_hidden = 100, dropout = 0.1, σ = relu),
    epochs = 100,
    batch_size = 100);

result=[];
for _ in 1:10 #Change the second number if you want to try some changes
    NN_2 = machine(model, select(drop_std.train, Not(:precipitation_nextday)), drop_std.train.precipitation_nextday)
    fit!(NN_2)
    AUC_drop = MLJ.auc(predict(NN_2, select(drop_std.test[:,:], Not(:precipitation_nextday))), drop_std.test.precipitation_nextday)
    push!(result, AUC_drop)
    model.batch_size = model.batch_size + 100
end
scatter(1:10, result, ylims = [0.9, 0.95])

# Best model found in the previous part
Random.seed!(2711)
best_model = NeuralNetworkClassifier( builder = MLJFlux.Short(n_hidden = 200, dropout = 0.1, σ = relu),
    epochs = 300,
    batch_size = 1000)
best_NN = machine(best_model, select(drop_std.train, Not(:precipitation_nextday)), drop_std.train.precipitation_nextday)
fit!(best_NN, verbosity = 2)
MLJ.auc(predict(best_NN, select(drop_std.test[:,:], Not(:precipitation_nextday))), drop_std.test.precipitation_nextday)

#Write in the submission file with a machine trained on all the data
train, test = generate(option = "drop", std = "true", valid = "false", test = "false");
best_mach = machine(best_model, select(train, Not(:precipitation_nextday)), train.precipitation_nextday);
fit!(best_mach)
sample = CSV.read(joinpath(@__DIR__, "data", "sample_submission.csv"), DataFrame);
pred = pdf.(predict(best_mach, test), true);
sample.precipitation_nextday = pred;
CSV.write(joinpath(@__DIR__, "results", "MLP_submission.csv"), sample);


using ConvexTests, SCS

@info "Starting SCS tests"

do_tests("SCS", () -> SCS.Optimizer(verbose=0, eps=1e-6); exclude = [r"mip"], description = "Tests run with `eps=1e-6`.")


function drawline!(data::IBuffer, p1::Vector{Float64}, p2::Vector{Float64},
                   color::Vector{Float64})
    all(isfinite.(p1)) && all(isfinite.(p2)) || return    # refuse to draw NaN and Inf

    x = p1[1]
    y = p1[2]
    z = p1[3]
    dx = p2[1] - p1[1]
    dy = p2[2] - p1[2]
    dz = p2[3] - p1[3]

    if abs(dy) < abs(dx)
        if p1[1] > p2[1]
            drawline!(data, p2, p1, color)
        else    # drawline: low
            mx = p2[1]
            yi = 1
            if dy < 0
                yi = -1
                dy = -dy
            end

            d = 2dy - dx
            dist = abs(mx - x)
            dz = dist > 0 ? dz / dist : zero(dz)
            for i in x:mx
                plot!(data, RInt[z, y, i], color)
                d += 2dy
                z += dz

                if d > 0
                    y += yi
                    d -= 2dx
                end
            end
        end
    else
        if p1[2] > p2[2]
            drawline!(data, p2, p1, color)
        else    # drawline: high
            my = p2[2]
            xi = 1
            if dx < 0
                xi = -1
                dx = -dx
            end

            d = 2dx -dy
            dist = abs(my - y)
            dz = dist > 0 ? dz / dist : zero(dz)

            for i in y:my
                plot!(data, RInt[z, i, x], color)
                d += 2dx
                z += dz

                if d > 0
                    x += xi
                    d -= 2dy
                end
            end
        end
    end
end


module GeometryModule
export Geometry, readGeometryXYZ, angstrom2bohr, angstrom2bohr!, computeEnergyInteratomicRepulsion

using ..BaseModule
using ..AtomModule

type Geometry
	atoms::Array{Atom,1}
end

function readGeometryXYZ(filename::AbstractString; unit::Symbol=:Bohr)
	geo = Geometry([])
	lines = open(readlines,filename)
	for line in lines[3:end]	# line 1 is n. of atoms, line 2 is comment
		columns = split(line)
		elem = Element(columns[1])
		pos = Position(map(parse,columns[2:4])...)	# float is not recommended for use on tuples
		append!(geo.atoms,[Atom(elem,pos)])
	end
	if (unit == :Bohr) angstrom2bohr!(geo) end	# Should be default as Bohr is an atomic unit
	return geo
end

angstrom2bohr(x::Real) = (return x/0.529177210)
angstrom2bohr(pos::Position) = (return Position(angstrom2bohr(pos.x),angstrom2bohr(pos.y),angstrom2bohr(pos.z)))
function angstrom2bohr!(geo::Geometry)
  for (atom in geo.atoms)
    atom.position = angstrom2bohr(atom.position)
  end
end

function computeEnergyInteratomicRepulsion(
  atom1::Atom,
  atom2::Atom)
  q1 = atom1.element.atomicNumber
  q2 = atom2.element.atomicNumber
  A = atom1.position
  B = atom2.position
  return q1*q2/(distance(A,B))
end

function computeEnergyInteratomicRepulsion(geo::Geometry)
  result = 0.::Float64
  for (A in 1:length(geo.atoms))
    for (B in 1:A-1)
      result += computeEnergyInteratomicRepulsion(geo.atoms[A],geo.atoms[B])
    end
  end
  return result
end

end #module