################################
# Coded by Julien Pascal
# Last modification : 16/04/2016
# Replication of the the following paper:

####################################################################################
# "On the Dynamics of Unemployment and Wage Distributions"
# Paper from Jean Marc Robin (2011)
# Available here: https://sites.google.com/site/jmarcrobin/research/publications
####################################################################################

###############################################################################################
# Module that calculates the wages value functions based on the paramaters estimated by the MSM
module calculate_wages

	using PyPlot, Distributions, StatsFuns, JLD, DataFrames, Copulas

	export Wages

	function Wages()

		####################################################
		#A. Define functions necessary to compute the wages:
		####################################################

	    ##################
	    # Function that returns the inverse of lognormal distribution
	    function logninv(p,mu,sigma)

	        #output: scalar x
	        logx0 = -sqrt(2)*erfcinv(2*p);
	        x = exp(sigma*logx0 + mu);

	        return x
	    end
	    #################

	    ##########################
	    #define match productivity:
	    function yi_m(y_index,x_index, ygrid, xgrid)
	        match = ygrid[y_index,1]*xgrid[x_index,1]
	        return match
	    end
	    ##########################

	    ##########################
	    # define opportunity cost of employment:
	    function zi_m(y_index,x_index, ygrid, xgrid, z0, alpha)

	        op_cost = z0 + alpha*(yi_m(y_index,x_index, ygrid, xgrid) - z0) #call the function defined above

	        return op_cost
	    end

	    #################
	    # Load Parameters
	    #################
		println("loading estimated parameters")
		coeff = readtable(string("tables/estimated_parameters.csv"))

		"""
		z0 = coeff[:Value][1]
		sigma = coeff[:Value][2]
		pho = coeff[:Value][3]
		lambda0 = coeff[:Value][4]
		eta = coeff[:Value][5]
		mu = coeff[:Value][6]
		x_lower_bound = parameters[7]
		delta = coeff[:Value][8] #exogenous layoff rate

		k = 0.12
		lambda1  = k*lambda0
		#x_lower_bound = 0.73
		alpha = 0.64
		tau = 0.5
		#delta = 0.043 #exogenous layoff rate
		r = 0.05/4 #quarterly interest rate
		discount = (1 - delta)/(1 + r)
		epsilon = 0.002
		"""

		"""
		z0 = coeff[:Value][1];
		sigma = coeff[:Value][2];
		lambda0 = coeff[:Value][3];
		eta =  coeff[:Value][4];
		mu =  coeff[:Value][5];
		x_lower_bound = coeff[:Value][6] #lower bound of the ability level
		alpha = coeff[:Value][7]
		delta = coeff[:Value][8]

		pho = 0.94
		k = 0.12;
		lambda1  = k*lambda0;
		#x_lower_bound = 0.73
		#alpha = 0.64
		tau = 0.5
		#delta = 0.043 # exogenous layoff rate

		r = 0.05/4 # interest rate
		discount = (1 - delta)/(1 + r)
		epsilon = 0.002
		"""

		z0 = coeff[:Value][1];
	    sigma = coeff[:Value][2];
	    lambda0 = coeff[:Value][3];
	    eta =  coeff[:Value][4];
	    mu =  coeff[:Value][5];
	    x_lower_bound = coeff[:Value][6] #lower bound of the ability level

	    pho = 0.94
	    k = 0.12;
	    lambda1  = k*lambda0;
	    #x_lower_bound = 0.73
	    alpha = 0.64
	    tau = 0.5
	    delta = 0.042 # exogenous layoff rate

	    r = 0.05/4 # interest rate
	    discount = (1 - delta)/(1 + r)
	    epsilon = 0.002

	    N = 100 #number of states
		M = 500 #number of ability levels

	    # Define the grid for ability x
	    xgrid_space = linspace(x_lower_bound,x_lower_bound+1,M); # column vector with M rows

	    xgrid = Array{Float64}(M,1);
	    for i=1:M
	    xgrid[i,1] = xgrid_space[i];
	    end

	    # Define the grid for match ability yi:
	        # 1. a:
	        agrid = linspace((0+epsilon),(1-epsilon),N);

	        # 2. yi:
	        #intialization:
	        ygrid = Array{Float64}(N,1);

	        for i=1:N
	            ygrid[i] = logninv(agrid[i],0,sigma); #calculate the inverse of the lognormal
	        end

	    #Markov transition matrix:
	    Markov = zeros(N,N);

	    # number of dimension
	    # AR1 parameter
	    cop = Copulas.NormalCopula(2,pho)

	    #Gaussian copula pdf:
	    for j=1:N #move along the column:
	        for i=1:N #move along the lines:
	            Markov[i,j] = Copulas.dnormCopula([agrid[i] agrid[j]], cop)[];
	        end
	    end

	    #normalize so that each row sum to 1:
	    for i = 1:N
	        Markov[i,:] = Markov[i,:]./sum(Markov[i,:]);
	    end

	    # distribution of workers on the grid x
	    lm = Array{Float64}(M,1);

	    # load the beta pdf
	    b = Beta(eta, mu)

	    for i=1:M
	        lm[i,1] = pdf(b, xgrid[i]-x_lower_bound); #evaluation of the beta pdf at a given point
	    end

	    #normalization so that mass of workers sum up to 1:
	    lm[:,1] = lm[:,1]./sum(lm[:,1]);

	    # Pre calculate the opportunity cost along the grid of ability:
	    # N times M matrix
	    # rows = aggregate staes
	    # columns = ability level
	    zi_m_grid = zeros(N,M);
	    for i=1:N
	        for m=1:M
	            zi_m_grid[i,m] = zi_m(i,m, ygrid, xgrid, z0, alpha);
	        end
	    end

	    ##########################################################
	    # Calculate the match surplus by value function iteration:
	    # N times M matrices
	    # N = aggregate level
	    # M = ability level

	    G = zeros(N,M) #intialization
	    p = zeros(N,M)

	    # Pre calculate values:
	    for i = 1:N
	        for m = 1:M
	            G[i,m] = yi_m(i,m, ygrid, xgrid) - zi_m_grid[i,m];
	            p[i,m] = yi_m(i,m, ygrid, xgrid);
	        end
	    end

	    # N times M matrix
	    # column index indicates the workers' type
	    # row index indicates the shock index

	    Si_m = zeros(N,M);

	    #tol = 0.0001;
	    tol = 0.01;
	    maxits = 300;
	    dif = tol+tol;
	    its = 1;

	    #initialization:
	    up = ones(N,M);
	    up_plus1 = zeros(N,M);
	    compare = zeros(N,M);

	    while dif>tol
	        up_plus1 =  G + discount*Markov*max(up,compare)
	        dif = norm(up_plus1 - up)
	        up = up_plus1

	        its = its + 1
	        print(its)

	        if its > maxits
	            break
	        end
	    end

	    Si_m = up;

		##############################################
		# Create the wages by value function iteration
		##############################################
		# 1. Wage low
		W_low = zeros(N,N,M); #Use mutlidimensional arrays
		W_low_star = zeros(N,N,M);

		up_ini = zeros(N,N,M);
		down_ini = zeros(N,N,M);

		###########################
		# Pre caculates some values:
		Z = zeros(N,N,M);

	    for m=1:M
	        for i = 1:N
	            for k = 1:N

	            Z[k,i,m] = zi_m_grid[i,m] - zi_m_grid[k,m]

	            #initialization
	            up_ini[k,i,m] = Z[k,i,m];

	                a = max(up_ini[k,i,m],0)
	                down_ini[k,i,m] = min(a,Si_m[i,m])

	            end
	        end
	    end

		function calculate_W_low(m, N, M, discount, Markov, up_ini, down_ini, Z, Si_m)

		        #########################
		        tol = 0.1;
		        maxits = 100;
		        compare = zeros(N,1);
		        a = 0

		        #move along workers' types

		            its = 1
		            dif = tol+tol

		            up = up_ini #initial value
		            up_plus1 = zeros(N,N)

		            down = down_ini;
		            down_star = down_ini;

		            while dif>tol

		                #move along the columns
		                for i=1:N
		                    for k = 1:N

		                        a = max(up[k,i],0)
		                        down[k,i] = min(a,Si_m[i,m])

		                        sum = 0

		                        for j = 1:N
		                            if Si_m[j,m] > 0
		                                sum = sum + (Markov[k,j] - Markov[i,j])*(lambda1*Si_m[j,m]+(1-lambda1)*(down[j,i]))
		                            end
		                        end

		                        up_plus1[k,i] = (Z[k,i,m])  + discount*sum

		                        a = max(up_plus1[k,i],0)
		                        down_star[k,i] = min(a, Si_m[i,m])
		                    end
		                end

		                dif = vecnorm(up_plus1[:,:]- up[:,:],2)

		                up[:,:] = up_plus1[:,:]

		                its = its + 1 #show the iteration number

		                if its > maxits
		                    break
		                end
		            end

		            #W_low[:,:,m] = up[:,:,m]

		    #println(dif) #show the current worker type
		    return [up; down_star] # concatenate vertically and return the result
		end

		##############
		# 2. Wage high
		##############
		up_high_ini = zeros(N,N,M);
		down2_ini = zeros(N,N,M);

		# Initialization:
		W_high = zeros(N,N,M);
		W_high_star = zeros(N,N,M);

		###########################
		# Pre caculates some values:
		Q = zeros(N,N,M);

		    for m=1:M
		        for i = 1:N
		            for k = 1:N
		                Q[k,i,m] = Z[k,i,m] + Si_m[i,m]

		                #initialization
		                up_high_ini[k,i,m] = Q[k,i,m];

		                    a = max(up_high_ini[k,i,m],0);
		                    down2_ini[k,i,m] = min(a, Si_m[i,m]);
		            end
		        end
		    end

		function calculate_W_high(m, N, M, discount, Markov, up_high_ini, down2_ini, Q, Si_m)

		    #initialization:
		    tol = 0.1;
		    maxits = 100;
		    its = 1;
		    dif = tol+tol;

		    up_high = up_high_ini;
		    up_plus1_high = zeros(N,N);

		    down = down2_ini;
		    down_star = down2_ini;
		    a = 0

		    while dif>tol

		        #move along the columns
		        for i=1:N
		            for k=1:N

		                a = max(up_high[k,i],0);
		                down[k,i] = min(a, Si_m[i,m]);

		                sum = 0;
		                for j = 1:N
		                    if Si_m[j,m] > 0
		                     sum = sum + (Markov[k,j] - Markov[i,j])*(lambda1*Si_m[j,m]+(1-lambda1)*(down[j,i])) ;
		                    end
		                 end

		                up_plus1_high[k,i] =  Q[k,i,m] + discount*sum;

		                a = max(up_plus1_high[k,i], 0);

		                down_star[k,i] = min(a, Si_m[i,m]);
		            end
		        end

		            dif = vecnorm(up_plus1_high[:,:] - up_high[:,:],2)

		            up_high[:,:] = up_plus1_high[:,:];

		            its = its + 1

		            if its > maxits
		                break
		            end
		    end

		    return [up_high; down_star] # concatenate vertically and return the result
		end

		##############
		#B. Iteration:
		##############
		tic()
		transition = zeros(2*N,N);
		for a = 1:M
			println(a)

			transition = calculate_W_low(a, N, M, discount, Markov, up_ini[:,:,a], down_ini[:,:,a], Z, Si_m);

		    W_low[:,:,a] = transition[1:N,:];
			W_low_star[:,:,a] = transition[(N+1):(2*N),:];
		end

		jldopen("surpluses/W_low_star.jld", "w") do file
		    write(file, "W_low_star", W_low_star)
		end

		jldopen("surpluses/W_low.jld", "w") do file
		    write(file, "W_low", W_low)
		end

		transition2 = zeros(2*N,N);
		for a = 1:M
			println(a)

			transition2 = calculate_W_high(a, N, M, discount, Markov, up_high_ini[:,:,a], down2_ini[:,:,a], Q, Si_m)

		    W_high[:,:,a] = transition2[1:N,:];
		    W_high_star[:,:,a] = transition2[(N+1):(2*N),:];

		end
		toc() #Takes 23 mn to run on my computer

		jldopen("surpluses/W_high_star.jld", "w") do file
		        write(file, "W_high_star", W_high_star)
		end

		jldopen("surpluses/W_high.jld", "w") do file
		       write(file, "W_high", W_high)
		end

		xgrid_plot = repmat(ygrid',N,1);
		ygrid_plot = repmat(ygrid,1,N);

		fig = figure("pyplot_surfaceplot",figsize=(10,10))
		ax = fig[:add_subplot](2,1,1)
		cp = ax[:contour](xgrid_plot, ygrid_plot, W_low_star[:,:,100], colors="black", linewidth=2.0)
		ax[:clabel](cp, inline=1, fontsize=10)
		xlabel("X")
		ylabel("Y")
		title("Contour Plot")

		subplot(212)
		ax = fig[:add_subplot](2,1,2)
		cp = ax[:contour](xgrid_plot, ygrid_plot, W_high_star[:,:,100], colors="black", linewidth=2.0)
		ax[:clabel](cp, inline=1, fontsize=10)
		xlabel("X")
		ylabel("Y")
		title("Contour Plot")
	end
end


using ImplicitDomainQuadrature
import Base.==, Base.≈

function allapprox(v1, v2)
    return all(v1 .≈ v2)
end

function allapprox(v1, v2, atol)
    err = maximum(abs.(v1 - v2))
    # if 10err > atol
    #     @warn "Approximation is very tight!"
    # end
    return length(v1) == length(v2) &&
           all([isapprox(v1[i], v2[i], atol = atol) for i = 1:length(v1)])
end

function allequal(v1, v2)
    return all(v1 .== v2)
end

function Base.isequal(c1::CutCellDG.CellMap, c2::CutCellDG.CellMap)
    return allequal(c1.yL, c2.yL) && allequal(c1.yR, c2.yR)
end

function ==(c1::CutCellDG.CellMap, c2::CutCellDG.CellMap)
    return isequal(c1, c2)
end

function plane_distance_function(coords, normal, x0)
    return (coords .- x0)' * normal
end

function closest_point_on_plane(querypoints,normal,x0)
    normalcomp = vec(normal'*(querypoints .- x0))
    disp = hcat([c*normal for c in normalcomp]...)
    return querypoints - disp
end

function circle_distance_function(coords, center, radius)
    difference = (coords .- center) .^ 2
    distance = radius .- sqrt.(vec(mapslices(sum, difference, dims = 1)))
    return distance
end

function closest_point_on_arc(querypoints,center,radius)
    relpos = querypoints .- center
    cqpoints = relpos[1,:] + im*relpos[2,:]
    theta = angle.(cqpoints)
    v = vcat(cos.(theta)',sin.(theta)')
    return center .+ radius*v
end

function corner_distance_function(x::V, xc) where {V<:AbstractVector}
    v = xc - x
    if all(x .<= xc)
        minimum(v)
    elseif all(x .> xc)
        return -sqrt(v' * v)
    elseif x[2] > xc[2]
        return v[2]
    else
        return v[1]
    end
end

function corner_distance_function(points::M, xc) where {M<:AbstractMatrix}
    return vec(
        mapslices(x -> corner_distance_function(x, xc), points, dims = 1),
    )
end

function ≈(q1::QuadratureRule, q2::QuadratureRule)
    flag = allapprox(q1.points, q2.points) && allapprox(q1.weights, q2.weights)
end

function required_quadrature_order(polyorder)
    ceil(Int, 0.5 * (2polyorder + 1))
end

function add_cell_error_squared!(
    err,
    interpolater,
    exactsolution,
    cellmap,
    quad,
)
    detjac = CutCellDG.determinant_jacobian(cellmap)
    for (p, w) in quad
        numsol = interpolater(p)
        exsol = exactsolution(cellmap(p))
        err .+= (numsol - exsol) .^ 2 * detjac * w
    end
end

function mesh_L2_error(nodalsolutions, exactsolution, basis, cellquads, mesh)
    ndofs = size(nodalsolutions)[1]
    err = zeros(ndofs)
    interpolater = InterpolatingPolynomial(ndofs, basis)
    ncells = CutCellDG.number_of_cells(mesh)
    for cellid = 1:ncells
        cellsign = CutCellDG.cell_sign(mesh, cellid)
        CutCellDG.check_cellsign(cellsign)
        if cellsign == +1 || cellsign == 0
            cellmap = CutCellDG.cell_map(mesh, +1, cellid)
            nodeids = CutCellDG.nodal_connectivity(mesh, +1, cellid)
            elementsolution = nodalsolutions[:, nodeids]
            update!(interpolater, elementsolution)
            quad = cellquads[+1, cellid]
            add_cell_error_squared!(
                err,
                interpolater,
                exactsolution,
                cellmap,
                quad,
            )
        end
        if cellsign == -1 || cellsign == 0
            cellmap = CutCellDG.cell_map(mesh, -1, cellid)
            nodeids = CutCellDG.nodal_connectivity(mesh, -1, cellid)
            elementsolution = nodalsolutions[:, nodeids]
            update!(interpolater, elementsolution)
            quad = cellquads[-1, cellid]
            add_cell_error_squared!(
                err,
                interpolater,
                exactsolution,
                cellmap,
                quad,
            )
        end
    end
    return sqrt.(err)
end

function cellwise_L2_error(
    nodalsolutions,
    exactsolution,
    basis,
    cellquads,
    mesh,
)

    ndofs = size(nodalsolutions)[1]
    interpolater = InterpolatingPolynomial(ndofs, basis)
    ncells = CutCellDG.number_of_cells(mesh)
    err = zeros(ndofs, 2, ncells)

    for cellid = 1:ncells
        cellsign = CutCellDG.cell_sign(mesh, cellid)
        CutCellDG.check_cellsign(cellsign)
        if cellsign == +1 || cellsign == 0
            cellmap = CutCellDG.cell_map(mesh, +1, cellid)
            nodeids = CutCellDG.nodal_connectivity(mesh, +1, cellid)
            elementsolution = nodalsolutions[:, nodeids]
            update!(interpolater, elementsolution)
            quad = cellquads[+1, cellid]
            cellerr = zeros(ndofs)
            add_cell_error_squared!(
                cellerr,
                interpolater,
                exactsolution,
                cellmap,
                quad,
            )
            err[:, 1, cellid] = cellerr
        end
        if cellsign == -1 || cellsign == 0
            cellmap = CutCellDG.cell_map(mesh, -1, cellid)
            nodeids = CutCellDG.nodal_connectivity(mesh, -1, cellid)
            elementsolution = nodalsolutions[:, nodeids]
            update!(interpolater, elementsolution)
            quad = cellquads[-1, cellid]
            cellerr = zeros(ndofs)
            add_cell_error_squared!(
                cellerr,
                interpolater,
                exactsolution,
                cellmap,
                quad,
            )
            err[:, 2, cellid] = cellerr
        end
    end
    return sqrt.(err)
end

function average(v)
    return sum(v)/length(v)
end

function convergence_rate(dx, err)
    return diff(log.(err)) ./ diff(log.(dx))
end

function add_cell_norm_squared!(vals, func, cellmap, quad)
    detjac = CutCellDG.determinant_jacobian(cellmap)
    for (p, w) in quad
        v = func(cellmap(p))
        vals .+= v .^ 2 * detjac * w
    end
end

function integral_norm_on_mesh(func, cellquads, mesh, ndofs)
    vals = zeros(ndofs)
    ncells = CutCellDG.number_of_cells(mesh)
    for cellid = 1:ncells
        s = CutCellDG.cell_sign(mesh, cellid)
        @assert s == -1 || s == 0 || s == 1
        if s == 1 || s == 0
            cellmap = CutCellDG.cell_map(mesh, +1, cellid)
            pquad = cellquads[+1, cellid]
            add_cell_norm_squared!(vals, func, cellmap, pquad)
        end
        if s == -1 || s == 0
            cellmap = CutCellDG.cell_map(mesh, -1, cellid)
            nquad = cellquads[-1, cellid]
            add_cell_norm_squared!(vals, func, cellmap, nquad)
        end
    end
    return sqrt.(vals)
end

function add_interface_error_squared!(
    err,
    interpolater,
    exactsolution,
    cellmap,
    quad,
    facescale,
)
    @assert length(quad) == length(facescale)
    for (i, (p, w)) in enumerate(quad)
        numsol = interpolater(p)
        exsol = exactsolution(cellmap(p))
        err .+= (numsol - exsol) .^ 2 * facescale[i] * w
    end
end

function interface_L2_error(
    nodalsolutions,
    exactsolution,
    levelsetsign,
    basis,
    interfacequads,
    mesh,
)
    ndofs = size(nodalsolutions)[1]
    err = zeros(ndofs)
    interpolater = InterpolatingPolynomial(ndofs, basis)

    ncells = CutCellDG.number_of_cells(mesh)

    for cellid = 1:ncells
        if CutCellDG.cell_sign(mesh, cellid) == 0
            quad = interfacequads[levelsetsign, cellid]
            cellmap = CutCellDG.cell_map(mesh, levelsetsign, cellid)
            scalearea = CutCellDG.interface_scale_areas(interfacequads, cellid)

            nodeids = CutCellDG.nodal_connectivity(mesh, levelsetsign, cellid)
            elementsolution = nodalsolutions[:, nodeids]
            update!(interpolater, elementsolution)

            add_interface_error_squared!(
                err,
                interpolater,
                exactsolution,
                cellmap,
                quad,
                scalearea,
            )
        end
    end
    return sqrt.(err)
end

function add_interface_norm_squared!(vals, func, cellmap, quad, facescale)
    @assert length(quad) == length(facescale)
    for (i, (p, w)) in enumerate(quad)
        v = func(cellmap(p))
        vals .+= v .^ 2 * facescale[i] * w
    end
end

function integral_norm_on_interface(
    func,
    interfacequads,
    levelsetsign,
    mesh,
    ndofs,
)
    vals = zeros(ndofs)
    ncells = CutCellDG.number_of_cells(mesh)

    for cellid = 1:ncells
        if CutCellDG.cell_sign(mesh, cellid) == 0
            cellmap = CutCellDG.cell_map(mesh, levelsetsign, cellid)
            quad = interfacequads[levelsetsign, cellid]
            facescale = CutCellDG.interface_scale_areas(interfacequads, cellid)

            add_interface_norm_squared!(vals, func, cellmap, quad, facescale)
        end
    end
    return sqrt.(vals)
end

function update_maxnorm_error!(
    globalerror,
    interpolater,
    exactsolution,
    cellmap,
    quad,
)
    ndofs = length(globalerror)
    for (p, w) in quad
        numsol = interpolater(p)
        exsol = exactsolution(cellmap(p))
        pointerror = abs.(numsol - exsol)
        for i = 1:ndofs
            globalerror[i] = max(globalerror[i], pointerror[i])
        end
    end
end

function interface_maxnorm_error(
    nodalsolutions,
    exactsolution,
    levelsetsign,
    basis,
    interfacequads,
    mesh,
)
    ndofs = size(nodalsolutions)[1]
    err = zeros(ndofs)
    interpolater = InterpolatingPolynomial(ndofs, basis)

    ncells = CutCellDG.number_of_cells(mesh)

    for cellid = 1:ncells
        if CutCellDG.cell_sign(mesh, cellid) == 0
            quad = interfacequads[levelsetsign, cellid]
            cellmap = CutCellDG.cell_map(mesh, levelsetsign, cellid)

            nodeids = CutCellDG.nodal_connectivity(mesh, levelsetsign, cellid)
            elementsolution = nodalsolutions[:, nodeids]
            update!(interpolater, elementsolution)

            update_maxnorm_error!(
                err,
                interpolater,
                exactsolution,
                cellmap,
                quad,
            )
        end
    end
    return err
end

function update_maxnorm!(vals, func, cellmap, quad)
    ndofs = length(vals)
    for (p, w) in quad
        v = abs.(func(cellmap(p)))
        for i = 1:ndofs
            vals[i] = max(vals[i], v[i])
        end
    end
end

function maxnorm_on_interface(func, interfacequads, levelsetsign, mesh, ndofs)
    vals = zeros(ndofs)
    ncells = CutCellDG.number_of_cells(mesh)

    for cellid = 1:ncells
        if CutCellDG.cell_sign(mesh, cellid) == 0
            cellmap = CutCellDG.cell_map(mesh, levelsetsign, cellid)
            quad = interfacequads[levelsetsign, cellid]

            update_maxnorm!(vals, func, cellmap, quad)
        end
    end
    return vals
end

function uniform_mesh_L2_error(nodalsolutions, exactsolution, basis, quad, mesh)
    ndofs, nnodes = size(nodalsolutions)
    err = zeros(ndofs)
    interpolater = InterpolatingPolynomial(ndofs, basis)
    ncells = CutCellDG.number_of_cells(mesh)

    for cellid = 1:ncells
        cellmap = CutCellDG.cell_map(mesh, cellid)
        nodeids = CutCellDG.nodal_connectivity(mesh, cellid)
        elementsolution = nodalsolutions[:, nodeids]
        update!(interpolater, elementsolution)
        add_cell_error_squared!(err, interpolater, exactsolution, cellmap, quad)
    end
    return sqrt.(err)
end

function integral_norm_on_uniform_mesh(func, quad, mesh, ndofs)
    numcells = CutCellDG.number_of_cells(mesh)
    vals = zeros(ndofs)
    for cellid in 1:numcells
        cellmap = CutCellDG.cell_map(mesh,cellid)
        add_cell_norm_squared!(vals, func, cellmap, quad)
    end
    return sqrt.(vals)
end

function lame_lambda(k, m)
    return k - 2m / 3
end


include("serial.jl")
include("../../src/star.jl")

#solves Ax = b in parallel, where A is a symmetric tridiagonal matrix
#upper triangular solve
function usolve(y, z, c, d, l, x)
  #this code performs 5n flops, 4 reads, and 1 write
  n = size(x, 1)
  #y, z are solutions at the previous and next separators
  x[n] = z
  @inbounds for i = n - 1 : -1 : 1
    x[i] = (x[i] - c[i + 1] * x[i + 1] - l[i] * y) / d[i] # 5 flops
  end
end

#lower triangular solve
function lsolve(a, b, c, d, l, x)
  #this code performs 12n flops in total
  #it performs 3 reads, and 3 writes
  n = size(x,1)
  d[1] = a[1]
  x[1] = b[1]
  f = c[1]
  l[1] = f
  r = f / d[1]
  s = f * r
  y = x[1] * r

  @inbounds for i = 2 : n - 1
    l_ = c[i] / d[i - 1]         # 1
    d[i] = a[i] - c[i] * l_      # 2
    x[i] = b[i] - x[i - 1] * l_  # 2
    f *= -l_                     # 2 including negation for fillin
    l[i] = f                     # store fillin
    r = f / d[i]                 # 1
    s += f * r                   # 2
    y += x[i] * r                # 2
  end
  l_ = c[n] / d[n - 1]
  d[n] = a[n] - c[n] * l_
  x[n] = b[n] - x[n - 1] * l_
  f = - r * c[n]  # fillin between previous and next separators

  s, y, f, d[n], x[n]
end

function root(s, y, f, d, x)
  #solve the reduced tridiagonal system
  #d is the diagonal of the reduced system, f the subdiagonal, and
  #x the reduced right hand side
  d[1:end - 1] -= s[2:end]
  x[1:end - 1] -= y[2:end]
  serial_trid(d, x, f[2:end])
  n = length(x)
  x = [0; x]
  return [(x[i], x[i+1]) for i = 1 : n]
end

trid_star(a, b, c, d, l, x) = Star.scan(usolve,
                    root(Star.reduce(lsolve, a, b, c, d, l, x)...), c, d, l, x)

#test routine
function trid_star_test(n)
  P = nworkers() # of processors
  #setup the inputs and outputs
  Da = DArray(I->100 * rand(map(length,I)), (n,), workers(), P, f1)
  Db = DArray(I->rand(map(length,I)), (n,), workers(), P, f1)
  Dc = DArray(I->rand(map(length,I)), (n + 1,), workers(), P, f2)
  Dd = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)
  Dl = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)
  Dx = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)

  t = @elapsed trid_star(Da, Db, Dc, Dd, Dl, Dx)

  # code to test the results
  x = convert(Array{Float64,1}, Dx)
  a = convert(Array{Float64,1}, Da)
  b = convert(Array{Float64,1}, Db)
  c = convert(Array{Float64,1}, Dc)
  d = convert(Array{Float64,1}, Dd)
  c = sub(c, 2:n)
  println("norm(A * x -b)", norm(matrix_vector_product(a, c, c, x) -b))
  println("serial trid - not storing L")
  x[:] = 0 ;
  t1 = @elapsed serial_trid(a, b, c, d, x)
  println("norm(A * x -b)", norm(matrix_vector_product(a, c, c, x) -b))
  println("serial time ", t1, " parallel time ", t)
end

function trid_star_run(n, num_runs = 1)
  P = nworkers() # of processors
  #setup the inputs and outputs
  Da = DArray(I->100 * rand(map(length,I)), (n,), workers(), P, f1)
  Db = DArray(I->rand(map(length,I)), (n,), workers(), P, f1)
  Dc = DArray(I->rand(map(length,I)), (n + 1,), workers(), P, f2)
  Dd = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)
  Dl = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)
  Dx = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)
  T = zeros(num_runs)

  for nruns = 1 : num_runs
    t = @elapsed trid_star(Da, Db, Dc, Dd, Dl, Dx)
    T[nruns] = t
  end
  println(" parallel time ", T)
  println(" min ", minimum(T), " max ", maximum(T), " mean ",  mean(T))
end

function test_serial(n, num_runs = 1)
  P = nworkers() # of processors
  #setup the inputs and outputs
  a = 100 * rand(n)
  b = rand(n)
  c = rand(n)
  d = zeros(n)
  x = zeros(n)
  T = zeros(num_runs)
  for nruns = 1 : num_runs
    t = @elapsed serial_trid(a, b, c, d, x)
    T[nruns] = t
  end
  println("serial time ", T)
  println(" min ", minimum(T), " max ", maximum(T), " mean ",  mean(T))
end

function trid_no_star(n, num_runs = 1)
  P = nworkers() # of processors
  #setup the inputs and outputs
  Da = DArray(I->100 * rand(map(length,I)), (n,), workers(), P, f1)
  Db = DArray(I->rand(map(length,I)), (n,), workers(), P, f1)
  Dc = DArray(I->rand(map(length,I)), (n + 1,), workers(), P, f2)
  Dd = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)
  Dl = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)
  Dx = DArray(I->zeros(map(length,I)), (n,), workers(), P, f1)

  ref = Array(Any, P)
  T = zeros(num_runs)
  for nruns = 1 : num_runs
    t = @elapsed begin
      for i = 1 : P
        ref[i] = @spawnat workers()[i] lsolve(localpart(Da), localpart(Db),
                    localpart(Dc), localpart(Dd), localpart(Dl), localpart(Dx))
      end
      s = Array(Float64, P)
      y = Array(Float64, P)
      f = Array(Float64, P)
      d = Array(Float64, P)
      x = Array(Float64, P)
      results = [fetch(ref[i]) for i = 1 : P]
      for i = 1 : P
        (s[i], y[i], f[i], d[i], x[i]) = results[i]
      end
      X = root(s, y, f, d, x)
      @sync begin
        for i = 1 : P
          @spawnat workers()[i] usolve(X[i][1], X[i][2], localpart(Dc),
                                  localpart(Dd), localpart(Dl), localpart(Dx))
        end
      end
    end #@elapsed
    T[nruns] = t
  end #for
  println(" parallel time ", T)
  println(" min ", minimum(T), " max ", maximum(T), " mean ",  mean(T))
#=
  # code to test the results
  x = convert(Array{Float64,1}, Dx)
  a = convert(Array{Float64,1}, Da)
  b = convert(Array{Float64,1}, Db)
  c = convert(Array{Float64,1}, Dc)
  d = convert(Array{Float64,1}, Dd)
  c = sub(c, 2:n)
  println("norm(A * x -b)", norm(matrix_vector_product(a, c, c, x) -b))
  println("serial trid - not storing L")
  x[:] = 0 ;
  t1 = @elapsed serial_trid(a, b, c, d, x)
  println("norm(A * x -b)", norm(matrix_vector_product(a, c, c, x) -b))
  println("serial time ", t1, " parallel time ", t)
=#
end

#auxiliary functions to partition the darrays among processors.
function f1(dims, procs) # auxiliary function to partition all darrays but the
                         # subdiagonal
  n = dims[1]
  P = size(procs, 1)
  leaf_size::Int64 = max(floor(n / P), 2) #size of a leaf
  cuts = ones(Int, P + 1) * leaf_size
  cuts[1] = 1
  cuts = cumsum(cuts)
  cuts[end] = n + 1
  indexes = Array(Any, P)
  for i = 1:P
    indexes[i] = (cuts[i] : cuts[i+1] - 1,)
  end
  indexes, Any[cuts]
end

function f2(dims, procs) # auxiliary function to partition the subdiagonal
  n = dims[1]
  P = size(procs, 1)
  leaf_size::Int64 = max(floor((n - 1) / P), 2) #size of a leaf
  cuts = ones(Int, P + 1) * leaf_size
  cuts[1] = 1
  cuts = cumsum(cuts)
  cuts[end] = n + 1
  indexes = Array(Any, P)
  for i = 1:P
    indexes[i] = (cuts[i] : cuts[i+1] - 1,)
  end
  indexes, Any[cuts]
end

#overload the DArray constructor to provide user defined partitions.
#=
function DArray(init, dims, procs, dist, distfunc::Function)
    np = prod(dist)
    procs = procs[1:np]
    idxs, cuts = distfunc([dims...], procs)

    chunks = Array(RemoteRef, dist...)
    for i = 1:np
        chunks[i] = remotecall(procs[i], init, idxs[i])
    end
    return DistributedArrays.construct_darray(dims, chunks, procs, idxs, cuts)
end
=#

#function DArray(init, dims, procs, dist)
function DArray(init, dims, procs, dist, distfunc::Function)
    np = prod(dist)
    procs = reshape(procs[1:np], ntuple(i->dist[i], length(dist)))
    idxs, cuts = distfunc([dims...], procs)
    #idxs, cuts = chunk_idxs([dims...], dist)
    identity = DistributedArrays.next_did()

    return DistributedArrays.construct_darray(identity, init, dims, procs, idxs, cuts)
end


# Simulation example 1

using CSV
using DataFrames
using Cairo
using Revise
using BayesNegativeBinomial
using BNPRegressionGGA2021
using Distributions
using Gadfly
using Random
using Statistics
using Test
using LinearAlgebra

function simulate_sample(rng, N0, N1, D)
    dϵ = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    x0 = [-0.5, 0.5]
    X0 = [ones(N0) rand(rng, x0, N0, D)]
    y0 = X0 * [ones(2); zeros(D - 1)] + rand(rng, dϵ, N0)
    y1 = LinRange(minimum(y0), maximum(y0), N1) |> collect |> x -> repeat(x, 4)
    x1 = kron(x0[[1, 1, 2, 2]], ones(N1))
    x2 = kron(x0[[1, 2, 1, 2]], ones(N1))
    X1 = [ones(4 * N1) x1 x2 zeros(4 * N1, D - 2)]
    return y0, X0, y1, X1
end

# Estimate the density for one realization
p = begin
    N0, N1, D = 1000, 50, 10
    rng = MersenneTwister(1);
    y0, X0, y1, X1 = simulate_sample(rng, N0, N1, D);
    smpl = BNPRegressionGGA2021.NormalSampler(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample(rng, smpl; mcmcsize = 10000, burnin = 5000);

    fb = mean(chainf);
    plot(x = y0, color = X0[:, 2], Geom.histogram(density = true, bincount = 50))
    p = plot(x = y1, y = fb, color = string.(X1[:, 2]), xgroup = string.(X1[:, 3]),  Geom.subplot_grid(Geom.line))
    img = SVG("figures/fig01.svg", 14cm, 8cm)
    draw(img, p)
    p
end

# Estimate p(\beta[j] | y) for each j, averaged over 100 samples
function get_γchains(R = 1)
    N0, N1, D = 1000, 50, 10
    γchains = zeros(D + 1, 100)
    rng = MersenneTwister(1);
    for i in 1:R
        y0, X0, y1, X1 = simulate_sample(rng, N0, N1, D);
        smpl = BNPRegressionGGA2021.NormalSampler(; y0, X0, y1, X1);
        chainf, chainβ = BNPRegressionGGA2021.sample(rng, smpl; mcmcsize = 5000, burnin = 2500);
        γchains[:, i] .= 1 .- mean([iszero.(chainβ[it]) for it in 1:length(chainβ)])
        println(i)
    end
    return γchains
end
γchains = get_γchains(100)
df = DataFrame(γchains', :auto)
CSV.write("data/data.csv", df)


include("../src/ode_tool.jl")
using ODE

tspan=[0,2*pi]
y0=Float64[1 0; 0 0]
function F(t, yy)
    y=yy[1]
    d0=[0 1; -1 0]*y[:,1]
    d1=[1,1]
    {[d0 d1],1}
end
#(T,Y)=ode23_abstract(F, tspan, {y0,0}) ## get samples from the whole ode path
(t1,y1)=ode23_abstract_end(F, tspan, {y0,0}) ## just get the ending value at the terminal time


###########################################################################################
#Define the marginal likelihood function
###########################################################################################
@everywhere function MargLik(x, y, g)
    N = size(y, 1)
    k = size(x, 2)
    P = eye(N) - x * inv(x'x) * x'
    yty = (y - mean(y))' * (y - mean(y))
    ymy = y'P * y
    g1 = g / (g + 1)
    g2 = 1 / (g + 1)
    ln_fy = .5 * k * log.(g1) - .5 * (N - 1) * log.(g2 * ymy + g1 * yty)
    return ln_fy[1]
end

###########################################################################################
#Use MCMCMC method for bayesian model averaging
###########################################################################################
@everywhere function bma_forc(Y, Y_lag, X_lag, vec_î, h, θ, pre_conv, post_conv)
    LHS = Y[1:vec_î - h]
    N = length(LHS)
    x = [Y_lag[1:N] X_lag[1:N, :]]
    x = x .- mean(x, 1)
    K = size(x, 2)
    τ_vec = zeros(Bool, K)
    k = sum(τ_vec)
    RHS = [ones(N) x[:, find(τ_vec)]]
    ln_mod_prior = log(θ ^ k * (1 - θ) ^ (K - k))
    # Compute marginal likelihood for initial model
    g = min(1 / N, 1 / K ^ 2)
    ln_fy = MargLik(RHS, LHS, g)
    y_forc = 0.
    # Begin MC3 simulations
    for s = 1:(pre_conv + post_conv)
        # Generate candidate Neighborhood Model
        neighbor_indx = rand(1:(1 + K))
        # Specify candidate model
        τ_vec_new = deepcopy(τ_vec)
        if neighbor_indx >= 2
            τ_vec_new[neighbor_indx - 1] = !τ_vec[neighbor_indx - 1]
        end
        k = sum(τ_vec_new)
        println("s = $(s), k = $(k): $(find(τ_vec_new))")
        RHS_new = [ones(N) x[:, find(τ_vec_new)]]
        ln_mod_prior_new = log(θ ^ k * (1 - θ) ^ (K - k))
        # Compute marginal likelihood for candidate model
        ln_fy_new = -Inf
        if k < K
            writedlm("RSM_new.csv", RHS_new, ',')
            writedlm("LHS.csv", LHS, ',')
            ln_fy_new = MargLik(RHS_new, LHS, g)
        end
        # MH Step
        prob_acc = min(exp(ln_fy_new + ln_mod_prior_new - (ln_fy + ln_mod_prior)), 1)
        if rand() <= prob_acc
            ln_fy = ln_fy_new
            τ_vec = deepcopy(τ_vec_new)
        end
        # If post-convergence, form forecasts
        if s > pre_conv
            x_selected = x[:, find(τ_vec)]
            post_β̄ = [mean(LHS); inv((1 + g) * (x_selected' * x_selected)) * (x_selected' * LHS)]
            x_forc = [Y_lag[1:vec_î] X_lag[1:vec_î, :]]
            x_forc = x_forc .- mean(x_forc, 1)
            x_forc = [ones(size(x_forc, 1)) x_forc[:, find(τ_vec)]]
            x_crnt = x_forc[end, :]
            y_forc = y_forc + x_crnt' * post_β̄
        end
    end
    return y_forc / (post_conv)
end

###########################################################################################
#The inner function
###########################################################################################
@everywhere function bma_inner_RMSE(J, h, h_lag, J_cv, Y, Y_lag, X_lag, θ, pre_conv, post_conv, oos_i)
    ## Vector Index for oos forecast goes from J_cv - (h_lag + h) to J - (1 + h_lag + h)
    #println("θ = $θ")
    ŷ_cv = zeros(Float64, size(J_cv - (h_lag + h):J + (oos_i - 1) - (1 + h_lag + h)))
    for (n, vec_î) in enumerate(J_cv - (h_lag + h):J + (oos_i - 1) - (1 + h_lag + h))
        println("n = $n")
        ŷ_cv[n] = bma_forc(Y, Y_lag, X_lag, vec_î, h, θ, pre_conv, post_conv)
    end
    return [sqrt(mean((Y[J_cv - (h_lag + h):J - (1 + h_lag + h)] - ŷ_cv) .^ 2)), θ]
end

###########################################################################################
#The outer function
###########################################################################################
function bma_oos_forc(J, h, h_lag, J_cv, Y, Y_lag, X_lag, θ_vec, pre_conv, post_conv, oos_i)
    # Step1 : pick the best θ
    RMSE_cv = zeros(length(θ_vec), 2)
    result = pmap(θ -> bma_inner_RMSE(J, h, h_lag, J_cv, Y, Y_lag, X_lag, θ, pre_conv, post_conv, oos_i),
                  θ_vec)::Array{Array{Float64,1},1}
    for i in 1:length(θ_vec)
        RMSE_cv[i, :] = result[i]
    end
    index = findmin(RMSE_cv[:, 1])[2]
    θ = θ_vec[index]
    # Step 2: oos forecasting
    ## Vector Index for oos forecast goes from J - (h_lag + h) to length(Y)
    vec_î = (J - (h_lag + h):length(Y))[oos_i]
    ŷ = bma_forc(Y, Y_lag, X_lag, vec_î, h, θ, pre_conv, post_conv)
    return θ, ŷ, RMSE_cv
end


using Test
using Random
using Statistics
using BenchmarkTools
using PolynomialRoots

using QuarticSolver

#==
These tests simply compare PolynomialRoots using BigFloats to the QuarticSolver
using Float64 (and the Complex versions of both).
==#

# Generate random coefficients (PolynomialRoots requires coeffs in reverse order)
rng = MersenneTwister(235)

rcoeffs = rand(rng, Float64, 5, 1000)
rncoeffs = randn(rng, Float64, 5, 1000)
ccoeffs = rand(rng, ComplexF64, 5, 1000)
cncoeffs = randn(rng, ComplexF64, 5, 1000)
largeccoeffs = rand(rng, ComplexF64, 5, 1000)*100

revrcoeffs = reverse(rcoeffs, dims=1)
revrncoeffs = reverse(rncoeffs, dims=1)
revccoeffs = reverse(ccoeffs, dims=1)
revcncoeffs = reverse(cncoeffs, dims=1)
largerevccoeffs = reverse(largeccoeffs, dims=1)

@testset "real, positive coefficient samples" begin
    pr_rresults = [sort(roots(collect(big.(coeffs))),
                        by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(revrcoeffs)]
    qs_rresults = [sort([QuarticSolver.solvequartic(coeffs...)...],
                        by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(rcoeffs)]

    rmatches = isapprox.(pr_rresults, qs_rresults, atol=1e-8)
    rfraction = sum(rmatches)/length(rmatches)
    rdiff = pr_rresults - qs_rresults
    avgdiff = mean(rdiff)
    avgdiffreal = mean(real.(rdiff))
    avgdiffimag = mean(imag.(rdiff))
    mindiffreal = minimum(real.(rdiff))
    mindiffimag = minimum(imag.(rdiff))
    maxdiffreal = maximum(real.(rdiff))
    maxdiffimag = maximum(imag.(rdiff))
    maxabsdiffreal = maximum(maximum(abs.(x) for x in real.(rdiff)))
    maxabsdiffimag = maximum(maximum(abs.(x) for x in imag.(rdiff)))

    argmax(real.(rdiff))
    argmax(imag.(rdiff))
    argmin(real.(rdiff))
    argmin(imag.(rdiff))

    @test rfraction == 1
    @test all(avgdiffreal .< 1e-10)
    @test all(avgdiffimag .< 1e-10)
    @test maxabsdiffreal < 1e-8
    @test maxabsdiffimag < 1e-8
end

@testset "real, positive and negative coefficient samples" begin
    pr_rnresults = [sort(roots(collect(big.(coeffs))),
                         by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(revrncoeffs)]
    qs_rnresults = [sort([QuarticSolver.solvequartic(coeffs...)...],
                         by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(rncoeffs)]

    rnmatches = isapprox.(pr_rnresults, qs_rnresults, atol=1e-8)
    rnfraction = sum(rnmatches)/length(rnmatches)
    rndiff = pr_rnresults - qs_rnresults
    avgdiff = mean(rndiff)
    avgdiffreal = mean(real.(rndiff))
    avgdiffimag = mean(imag.(rndiff))
    mindiffreal = minimum(real.(rndiff))
    mindiffimag = minimum(imag.(rndiff))
    maxdiffreal = maximum(real.(rndiff))
    maxdiffimag = maximum(imag.(rndiff))
    maxabsdiffreal = maximum(maximum(abs.(x) for x in real.(rndiff)))
    maxabsdiffimag = maximum(maximum(abs.(x) for x in imag.(rndiff)))

    argmax(imag.(rndiff))
    argmin(imag.(rndiff))

    @test rnfraction == 1
    @test all(avgdiffreal .< 1e-10)
    @test all(avgdiffimag .< 1e-10)
    @test maxabsdiffreal < 1e-8
    @test maxabsdiffimag < 1e-8
end

@testset "complex, positive coefficient samples" begin
    pr_cresults = [sort(roots(collect(big.(coeffs))),
                        by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(revccoeffs)]
    qs_cresults = [sort([QuarticSolver.solvequartic(coeffs...)...],
                        by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(ccoeffs)]

    cmatches = isapprox.(pr_cresults, qs_cresults, atol=1e-8)
    cfraction = sum(cmatches)/length(cmatches)
    cdiff = pr_cresults - qs_cresults
    avgdiff = mean(cdiff)
    avgdiffreal = mean(real.(cdiff))
    avgdiffimag = mean(imag.(cdiff))
    mindiffreal = minimum(real.(cdiff))
    mindiffimag = minimum(imag.(cdiff))
    maxdiffreal = maximum(real.(cdiff))
    maxdiffimag = maximum(imag.(cdiff))
    maxabsdiffreal = maximum(maximum(abs.(x) for x in real.(cdiff)))
    maxabsdiffimag = maximum(maximum(abs.(x) for x in imag.(cdiff)))

    argmax(imag.(cdiff))
    argmin(imag.(cdiff))

    @test cfraction == 1
    @test all(avgdiffreal .< 1e-10)
    @test all(avgdiffimag .< 1e-10)
    @test maxabsdiffreal < 1e-8
    @test maxabsdiffimag < 1e-8
end

@testset "complex, positive and negative coefficient samples" begin
    pr_cnresults = [sort(roots(collect(big.(coeffs))),
                         by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(revcncoeffs)]
    qs_cnresults = [sort([QuarticSolver.solvequartic(coeffs...)...],
                         by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(cncoeffs)]

    cnmatches = isapprox.(pr_cnresults, qs_cnresults, atol=1e-8)
    cnfraction = sum(cnmatches)/length(cnmatches)
    cndiff = pr_cnresults - qs_cnresults
    avgdiff = mean(cndiff)
    avgdiffreal = mean(real.(cndiff))
    avgdiffimag = mean(imag.(cndiff))
    mindiffreal = minimum(real.(cndiff))
    mindiffimag = minimum(imag.(cndiff))
    maxdiffreal = maximum(real.(cndiff))
    maxdiffimag = maximum(imag.(cndiff))
    maxabsdiffreal = maximum(maximum(abs.(x) for x in real.(cndiff)))
    maxabsdiffimag = maximum(maximum(abs.(x) for x in imag.(cndiff)))

    argmax(imag.(cndiff))
    argmin(imag.(cndiff))

    @test cnfraction == 1
    @test all(avgdiffreal .< 1e-10)
    @test all(avgdiffimag .< 1e-10)
    @test maxabsdiffreal < 1e-8
    @test maxabsdiffimag < 1e-8
end

@testset "complex, positive coefficients larger than 1" begin
    pr_lcresults = [sort(roots(collect(big.(coeffs))),
                         by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(largerevccoeffs)]
    qs_lcresults = [sort([QuarticSolver.solvequartic(coeffs...)...],
                         by=x->(round(real(x), digits=6), round(imag(x), digits=6))) for coeffs in eachcol(largeccoeffs)]

    lcmatches = isapprox.(pr_lcresults, qs_lcresults, atol=1e-8)
    lcfraction = sum(lcmatches)/length(lcmatches)
    lcdiff = pr_lcresults - qs_lcresults
    avgdiff = mean(lcdiff)
    avgdiffreal = mean(real.(lcdiff))
    avgdiffimag = mean(imag.(lcdiff))
    mindiffreal = minimum(real.(lcdiff))
    mindiffimag = minimum(imag.(lcdiff))
    maxdiffreal = maximum(real.(lcdiff))
    maxdiffimag = maximum(imag.(lcdiff))
    maxabsdiffreal = maximum(maximum(abs.(x) for x in real.(lcdiff)))
    maxabsdiffimag = maximum(maximum(abs.(x) for x in imag.(lcdiff)))

    argmax(imag.(lcdiff))
    argmin(imag.(lcdiff))

    @test lcfraction == 1
    @test all(avgdiffreal .< 1e-10)
    @test all(avgdiffimag .< 1e-10)
    @test maxabsdiffreal < 1e-8
    @test maxabsdiffimag < 1e-8
end

singleccoeffs = ccoeffs[:, 13]
singlerevccoeffs = reverse(singleccoeffs)
@benchmark roots($singlerevccoeffs)
@benchmark QuarticSolver.solvequartic($singleccoeffs[1],
                                      $singleccoeffs[2],
                                      $singleccoeffs[3],
                                      $singleccoeffs[4],
                                      $singleccoeffs[5])


function scheduling(processing_times, due_dates, ::Val{:raw})
    m = model(; kind = :scheduling)
    n = length(processing_times) # number of jobs
    max_time = sum(processing_times)

    d = domain(0:max_time)
    foreach(_ -> variable!(m, d), 1:n) # C (1:n)
    foreach(_ -> variable!(m, d), 1:n) # S (n+1:2n)

    minus_eq_param(param) = x -> abs(x[1] - x[2] - param)
    less_than_param(param) = x -> abs(x[1] - param)
    sequential_tasks(x) = x[1] ≤ x[2] || x[3] ≤ x[4]

    for i in 1:n
        constraint!(m, minus_eq_param(processing_times[i]), [i, i + n])
        constraint!(m, less_than_param(due_dates[i]), [i])
    end
    for i in 1:n, j in 1:n
        i == j && continue
        constraint!(m, sequential_tasks, [i, j + n, j, i + n])
    end

    return m
end

function scheduling(processing_times, due_dates, ::Val{:JuMP})
    model = JuMP.Model(CBLS.Optimizer)
    n = length(processing_times) # number of jobs
    max_time = sum(processing_times)

    @variable(model, 0 ≤ C[1:n] ≤ max_time, Int) # completion
    @variable(model, 0 ≤ S[1:n] ≤ max_time, Int) # start

    for i in 1:n
        @constraint(model, [C[i], S[i]] in MinusEqualParam(processing_times[i]))
        @constraint(model, [C[i]] in LessThanParam(due_dates[i]))
    end
    for i in 1:n, j in 1:n
        i == j && continue
        @constraint(model, [C[i], S[j], C[j], S[i]] in SequentialTasks())
    end

    return model, C, S
end

"""
    scheduling(processing_time, due_date; modeler=:JuMP)

Create a model for the n-queens problem with `n` queens. The `modeler` argument accepts :JuMP (default), which refer to the JuMP model.

!!! warning

    The model seems to have a flaw. Needs to be investigated.

"""
scheduling(processing_times, due_dates; modeler=:JuMP) = scheduling(processing_times, due_dates, Val(modeler))


function getsum(a::Int, b::Int)::Int
    start,stop = minmax(a,b)
    N = stop-start+1
    (N*(N-1))÷2 + N*start
end


using Distributions, LinearAlgebra
using Random: seed!
using AugmentedGaussianProcesses
using CSV, DataFrames
seed!(42)
N = 3000
D = 20
K = 4
k = RBFKernel(1.0)
X = rand(N,D)
y = rand(MvNormal(kernelmatrix(X,k)+1e-3I))
df = DataFrame(X)
df.y_reg = y
df.y_class = sign.(y)
width = maximum(y)-minimum(y);normy = (y.-minimum(y))/width*K
df.y_multi = floor.(Int64,normy)
CSV.write(joinpath(@__DIR__,"benchmarkdata.csv"),df)


# Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors    #src
# This Source Code Form is subject to the terms of the Mozilla Public License   #src
# v.2.0. If a copy of the MPL was not distributed with this file, You can       #src
# obtain one at https://mozilla.org/MPL/2.0/.                                   #src

# # Robust uncertainty sets

# Computes the Value at Risk for a data-driven uncertainty set; see "Data-Driven
# Robust Optimization" (Bertsimas 2013), section 6.1 for details. Closed-form
# expressions for the optimal value are available.

using JuMP
import SCS
import LinearAlgebra
import Test

function example_robust_uncertainty()
    R = 1
    d = 3
    𝛿 = 0.05
    ɛ = 0.05
    N = ceil((2 + 2 * log(2 / 𝛿))^2) + 1
    c = randn(d)
    μhat = rand(d)
    M = rand(d, d)
    Σhat = 1 / (d - 1) * (M - ones(d) * μhat')' * (M - ones(d) * μhat')
    Γ1(𝛿, N) = R / sqrt(N) * (2 + sqrt(2 * log(1 / 𝛿)))
    Γ2(𝛿, N) = 2 * R^2 / sqrt(N) * (2 + sqrt(2 * log(2 / 𝛿)))
    model = Model(SCS.Optimizer)
    set_silent(model)
    @variable(model, Σ[1:d, 1:d], PSD)
    @variable(model, u[1:d])
    @variable(model, μ[1:d])
    @constraint(model, [Γ1(𝛿 / 2, N); μ - μhat] in SecondOrderCone())
    @constraint(model, [Γ2(𝛿 / 2, N); vec(Σ - Σhat)] in SecondOrderCone())
    @SDconstraint(model, [((1 - ɛ) / ɛ) (u - μ)'; (u - μ) Σ] >= 0)
    @objective(model, Max, LinearAlgebra.dot(c, u))
    optimize!(model)
    I = Matrix(1.0 * LinearAlgebra.I, d, d)
    exact =
        LinearAlgebra.dot(μhat, c) +
        Γ1(𝛿 / 2, N) * LinearAlgebra.norm(c) +
        sqrt((1 - ɛ) / ɛ) * sqrt(LinearAlgebra.dot(c, (Σhat + Γ2(𝛿 / 2, N) * I) * c))
    Test.@test objective_value(model) ≈ exact atol = 1e-3
    return
end

example_robust_uncertainty()


using FFTW

function recenter_block(x::Block)
    return Block(x.data .- 128)
end

function decenter_block(x::Block)
    return Block(x.data .+ 128)
end

function dct(g::Block)::Matrix{Float32}
    rg = recenter_block(g).data
    return FFTW.dct(rg)
end

function invdct(g::Matrix{Float32})::Block
    rg = round.(FFTW.idct(g))
    b = decenter_block(Block(rg))
    return b
end


## kepler.jl
##     Functions for orbital elements and Kepler motion

###################################################################
#                       Orbital elements
###################################################################

"""
# Orbital elements

- `a` : Semi-major axis
- `e` : Eccentricity
- `I` : Inclination
- `ω` : Argument of periapsis
- `Ω` : Longitude of ascending node
- `Φ` : Mean anomaly : tₚ = - Φ / n
- `tₚ` : Periapsis passage time

- `μ` : Standard gravitational parameter of the system (GM)
- `n` : Mean motion    : n = √(μ / a^3)
- `T` : Orbital period : T = 2π / n
"""
struct OrbitalElements{T}
    a::T
    e::T
    I::T
    ω::T
    Ω::T
    Φ::T
    tₚ::T

    μ::T
    n::T
    T::T
end

function OrbitalElements(params)

    a = params[:a] * AU      # semi-mojor axis [m]
    e = params[:e]           # eccentricity
    I = deg2rad(params[:I])  # inclination [rad]
    Ω = deg2rad(params[:Ω])  # longitude of the ascending node [rad]
    ω = deg2rad(params[:ω])  # argument of periapsis [rad]

    μ = params[:μ]  # Gravitational parameter
    n = √(μ / a^3)  # mean motion  [rad/sec]
    T = 2π / n      # orbital period [sec]

    if haskey(params, :Φ)
        Φ  = deg2rad(params[:Φ])  # mean anomaly at the epoch [rad]
        tₚ = - Φ / n              # periapsis passage time [sec]
    elseif haskey(params, :tₚ)
        tₚ = params[:tₚ]
        Φ  = - n * tp
    else
        println("Give [:Φ] or [:tp] in Dict.")
    end

    OrbitalElements(a, e, I, ω, Ω, Φ, tₚ, μ, n, T)
end

function OrbitalElements(a, e, I, μ)
    a *= AU
    I = deg2rad(I)
    ω = 2π * rand()
    Ω = 2π * rand()
    Φ = 2π * rand()

    n = √(μ / a^3)
    T = 2π / n
    tₚ = - Φ / n

    OrbitalElements(a, e, I, ω, Ω, Φ, tₚ, μ, n, T)
end

function Base.show(io::IO, elms::OrbitalElements)
    println(io, "Orbital elements")
    println(io, "----------------")

    println("Semi-mojor axis        : a  = ", elms.a / AU,     " [AU]")
    println("Eccentricity           : e  = ", elms.e,          " [-]")
    println("Lon. of ascending node : Ω  = ", rad2deg(elms.Ω), " [deg]")
    println("Argument of periapsis  : ω  = ", rad2deg(elms.ω), " [deg]")
    println("Inclination            : I  = ", rad2deg(elms.I), " [deg]")
    println("Periapsis passage time : tₚ = ", elms.tₚ,         " [sec]")
    println("Mean anomaly           : Φ  = ", rad2deg(elms.Φ), " [deg]")

    println()
    println("Other parameters")
    println("----------------")

    println("Gravitational parameter : μ = ", elms.μ,                      " [m^3/s^2]")
    println("Mean motion             : n = ", rad2deg(elms.n) * (3600*24), " [deg/day]")
    println("Orbital period          : T = ", elms.T / (3600*24),          " [day]")
end

###################################################################
#                   Coordinate transformation
#                   R (X, Y, Z) -> r (x, y, z)
###################################################################

ref_to_orb!(r, elms) = ref_to_orb!(r, elms.ω, elms.I, elms.Ω)
ref_to_orb(r, elms) = ref_to_orb(r, elms.ω, elms.I, elms.Ω)

"""
Transform from the reference coordinate system to the orbital plane system

# Parameters
- r : vector to be transformed
- ω : argument of periapsis
- I : inclination
- Ω : longitude of the ascneding node
"""
function ref_to_orb!(r, ω, I, Ω)
    rotateZ!(r, Ω)
    rotateX!(r, I)
    rotateZ!(r, ω)
end

function ref_to_orb(r, ω, I, Ω)
    r = rotateZ(r, Ω)
    r = rotateX(r, I)
    r = rotateZ(r, ω)
end

###################################################################
#                   Coordinate transformation
#                    r(x, y, z) -> R(X, Y, Z)
###################################################################

orb_to_ref!(r, elms) = orb_to_ref!(r, elms.ω, elms.I, elms.Ω)
orb_to_ref(r, elms) = orb_to_ref(r, elms.ω, elms.I, elms.Ω)

"""
Transform from the orbital plane system to the reference coordinate system

# Parameters
- r : vector to be transformed
- ω : argument of periapsis
- I : inclination
- Ω : longitude of the ascneding node
"""
function orb_to_ref!(r, ω, I, Ω)
    rotateZ!(r, -ω)
    rotateX!(r, -I)
    rotateZ!(r, -Ω)
end

function orb_to_ref(r, ω, I, Ω)
    r = rotateZ(r, -ω)
    r = rotateX(r, -I)
    r = rotateZ(r, -Ω)
end

###################################################################
#           R(X, Y, Z), V(Ẋ, Ẏ, Ż) -> Orbital elements
#             @Reference frame
###################################################################

OrbitalElements(R, V, M, m, t) = OrbitalElements(R, V, G*(M+m), t)

function OrbitalElements(R::AbstractVector, V::AbstractVector, μ, t)

    h = R × V

    a = get_a(R, V, μ)
    e = get_e(h, μ, a)
    I = get_I(h)
    Ω = get_Ω(h, I)

    ω = get_ω(R, V, a, e, I, Ω)
    f = get_f(R, V, a, e)

    E = f2E(f, e)   # eccentric anomaly (E or u)
    n = √(μ / a^3)  # mean motion
    T = 2π / n      # orbital period

    tₚ = get_tₚ(t%T, E, e, n)  # time of periapsis passage
    M = - n * tₚ              # mean anomaly (M or Φ)

    OrbitalElements(a, e, I, ω, Ω, M, tₚ, μ, n, T)
end

function OrbitalElements(p1, p2, t)
    ps = setParticles(p1, p2)
    r_G, v_G, a_G = getBaryCenter(ps)
    OrbitalElements(p2.r - r_G, p2.v - v_G, p1.m, p2.m, t)
end

get_a(R, V, M, m) = get_a(R, V, G*(M+m))
get_a(R, V, μ) = 1 / (2/norm(R) - norm(V)^2/μ)

get_e(R, V, μ, a) = get_e(R × V, μ, a)
get_e(h, μ, a) = √(1 - norm(h)^2/(μ*a))

get_I(R, V) = get_I(R × V)
get_I(h) = acos(h[3]/norm(h))

get_Ω(R, V, I) = get_Ω(R × V, I)

function get_Ω(h, I)
    # sinΩ =  sign(h[3]) * h[1] / (norm(h)*sin(I))
    # cosΩ = -sign(h[3]) * h[2] / (norm(h)*sin(I))
    # Ω = atan(sinΩ, cosΩ)

    Ω = atan(sign(h[3])*h[1], -sign(h[3])*h[2])
    Ω ≥ 0 ? Ω : Ω+2π
end

function get_ω(R, V, a, e, I, Ω)
    ω₊f = get_ω₊f(R, V, I, Ω)
    f = get_f(R, V, a, e)

    ω = ω₊f - f
    # ω ≥ 0 ? ω : ω+2π
end

function get_ω₊f(R, V, I, Ω)
    I == 0 && return asin((R[2]*cos(Ω) - R[1]*sin(Ω))/norm(R))

    sin_ω₊f = R[3] / (norm(R)*sin(I))
    cos_ω₊f = sec(Ω) * (R[1]/norm(R) + sin(Ω)*sin_ω₊f*cos(I))

    ω₊f = atan(sin_ω₊f, cos_ω₊f)
    # ω₊f ≥ 0 ? ω₊f : ω₊f + 2π
end

function get_f(R, V, a, e)
    h = R × V
    Ṙ = sign(R ⋅ V) * √(norm(V)^2 - norm(h)^2/norm(R)^2)

    sinf = a*(1-e*e) / (norm(h)*e) * Ṙ
    cosf = (a*(1-e*e) / norm(R) - 1) / e

    f = atan(sinf, cosf)
    # f ≥ 0 ? f : f+2π
end

function get_E(R, a, e)
    E = acos((a - norm(R)) / (a*e))  # eccentric anomaly (E or u)
    # E ≥ 0 ? E : E+2π
end

E2f(E, e) = 2 * atan(√((1+e)/(1-e)) * sin(E/2), cos(E/2))
f2E(f, e) = 2 * atan(√((1-e)/(1+e)) * sin(f/2), cos(f/2))

get_tₚ(t, E, e, M, m, a) = get_tₚ(t, E, G*(M+m), a)
get_tₚ(t, E, e, μ, a) = get_tₚ(t, E, e, √(μ/a^3))
get_tₚ(t, E, e, n) = t - (E - e*sin(E)) / n

# ****************************************************************
# ****************************************************************
# ****************************************************************

function getOrbitalElementsError(elms, elms_ref)
    Δa = elms.a - elms_ref.a
    Δe = elms.e - elms_ref.e
    Δω = elms.ω - elms_ref.ω
    ΔT = elms.T - elms_ref.T

    return Δa, Δe, Δω, ΔT
end

# ****************************************************************
#           Orbital elements -> r(x, y, z), v(ẋ, ẏ, ż)
#                                @Orbital plane frame
# ****************************************************************

solveKeplerEquation1(elms, ts::T) where T<:AbstractArray = [solveKeplerEquation1(elms, t) for t in ts]
solveKeplerEquation2(elms, ts::T) where T<:AbstractArray = [solveKeplerEquation2(elms, t) for t in ts]

solveKeplerEquation1(elms, t) = solveKeplerEquation1(elms.e, elms.n, elms.tₚ, t)
solveKeplerEquation2(elms, t) = solveKeplerEquation2(elms.e, elms.n, elms.tₚ, t)

get_rv(elms, u) = get_rv(elms.a, elms.e, elms.n, u)
get_r(elms, u) = get_r(elms.a, elms.e, u)
get_v(elms, u) = get_v(elms.a, elms.e, elms.n, u)

heliocentric_distance(elms, u) = norm(get_r(elms, u))
heliocentric_distance(elms, us::T) where T<:AbstractArray = [heliocentric_distance(elms, u) for u in us]

"""
    solveKeplerEquation1(e, n, tₚ, t) -> u_i

Solve Kepler's equation with the 1st order method

# Parameters
- `e`  : eccentricity
- `n`  : mean motion
- `tₚ` : periapsis passage time [sec]
- `t`  : time

# Return
- `u_i` : eccentric anomaly (i-th order solution)
"""
function solveKeplerEquation1(e, n, tₚ, t)

    Φ = n * (t - tₚ)  # mean anomaly

    u_i = Φ  # eccentirc anomaly at order 0

    for i in 1:20
        u_pri = u_i
        u_i = Φ + e * sin(u_i)
        err = abs(1 - u_pri / u_i)
        # println(i, " : ", u_i, " : ", err)
        err < 1e-10 && return u_i
    end
    u_i
end

"""
    solveKeplerEquation2(e, n, tₚ, t) -> u_i

Solve Kepler's equation with the 2nd order method (Newton's method)

# Parameters
- `e`  : eccentricity
- `n`  : mean motion
- `tₚ` : periapsis passage time [sec]
- `t`  : time

# Return
- `u_i` : eccentric anomaly (i-th order solution)
"""
function solveKeplerEquation2(e, n, tₚ, t)

    Φ = n * (t - tₚ)  # mean anomaly

    u_i = Φ  # eccentirc anomaly at order 0

    for i in 1:20
        u_pri = u_i
        u_i -= (u_i - e * sin(u_i) - Φ) / (1 - e * cos(u_i))
        err = abs(1 - u_pri / u_i)
        # println(i, " : ", u_i, " : ", err)
        err < 1e-10 && return u_i
    end
    u_i
end

"""
    u2ν(u, e) -> ν

Calculate true anomaly ν from eccentric anomaly u and eccentricity e

# Parameters
- `u` : eccentric anomaly
- `e` : eccentricity

# Return
- `ν` : true anomaly

"""
u2ν(u, e) = 2 * atan(√((1+e)/(1-e)) * sin(u*0.5), cos(u*0.5))
u2ν(u, elms::OrbitalElements) = u2ν(u, elms.e)

u2ν(us::T, e) where T<:AbstractArray = [u2ν(u, e) for u in us]
u2ν(us::T, elms::OrbitalElements) where T<:AbstractArray = u2ν(us, elms.e)

"""
    getrv(a, e, n, u) -> r, v

Get a body's positon and velocity based on Kepler's motion

# Parameters
- `a` : semi-major axis
- `e` : eccentricity
- `n` : mean motion
- `u` : eccentric anomaly

# Returns
- `r` : position
- `v` : velocity
(the coordinate system is in the orbital plane and its origin is the eclipse focus)
"""
function get_rv(a, e, n, u)
    sin_u = sin(u)
    cos_u = cos(u)

    x = a * (cos_u - e)
    y = a * √(1 - e*e) * sin_u
    z = 0.

    r = SA_F64[x, y, z]

    x = - a * n * sin_u / (1 - e * cos_u)
    y = a * n * √(1 - e*e) * cos_u / (1 - e * cos_u)
    z = 0.

    v = SA_F64[x, y, z]

    return r, v
end

"""
    get_r(a, e, u) -> r

Get a body's positon based on Kepler's motion

# Parameters
- `a` : semi-major axis
- `e` : eccentricity
- `u` : eccentric anomaly

# Return
- `r` : position
(the coordinate system is in the orbital plane and its origin is the eclipse focus)
"""
function get_r(a, e, u)
    x = a * (cos(u) - e)
    y = a * √(1 - e*e) * sin(u)
    z = 0.

    r = SA_F64[x, y, z]
end

"""
    get_v(a, e, n, u) -> v

Get a body's velocity based on Kepler's motion

# Parameters
- `a` : semi-major axis
- `e` : eccentricity
- `n` : mean motion
- `u` : eccentric anomaly

# Return
- `v` : velocity
(the coordinate system is in the orbital plane and its origin is the eclipse focus)
"""
function get_v(a, e, n, u)
    sin_u = sin(u)
    cos_u = cos(u)

    x = - a * n * sin_u / (1 - e * cos_u)
    y = a * n * √(1 - e*e) * cos_u / (1 - e * cos_u)
    z = 0.

    v = SA_F64[x, y, z]
end

################################################################
#                     Rotation matrix
################################################################

# Rx(θ) = [
#     1   0      0
#     0  cos(θ) sin(θ)
#     0 -sin(θ) cos(θ)
# ]

# Ry(θ) = [
#     cos(θ) 0 -sin(θ)
#      0     1   0
#     sin(θ) 0  cos(θ)
# ]

# Rz(θ) = [
#      cos(θ) sin(θ) 0
#     -sin(θ) cos(θ) 0
#       0      0     1
# ]

# orb2ref(r, elms) = orb2ref(r, elms.ω, elms.I, elms.Ω)
# ref2orb(r, elms) = ref2orb(r, elms.ω, elms.I, elms.Ω)

# orb2ref(r, ω, I, Ω) = Rz(-Ω) * Rx(-I) * Rz(-ω) * r
# ref2orb(r, ω, I, Ω) = Rz( ω) * Rx( I) * Rz( Ω) * r


######### Stan program example  ###########

using StanSample, MCMCChains

ProjDir = @__DIR__

bernoullimodel = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
"

datatheta = joinpath(ProjDir, "bernoulli.data.json")
inittheta = joinpath(ProjDir , "bernoulli.init.json")

sm = SampleModel("bernoulli", bernoullimodel);

rc = stan_sample(sm, data=datatheta, init=inittheta)
#rc = stan_sample(sm, data=datatheta)

if success(rc)
  chn = read_samples(sm, :mcmcchains)
  chn |> display
end


using Mimi

@defcomp grosseconomy begin
    YGROSS  = Variable(index=[time])    # Gross output
    K       = Variable(index=[time])    # Capital
    l       = Parameter(index=[time])   # Labor
    tfp     = Parameter(index=[time])   # Total factor productivity
    s       = Parameter(index=[time])   # Savings rate
    depk    = Parameter()               # Depreciation rate on capital - Note that it has no time index
    k0      = Parameter()               # Initial level of capital
    share   = Parameter()               # Capital share

    function run_timestep(p, v, d, t)
        # Define an equation for K
        if is_first(t)
            v.K[t]  = p.k0  # Note the use of v. and p. to distinguish between variables and parameters
        else
            v.K[t]  = (1 - p.depk)^5 * v.K[t-1] + v.YGROSS[t-1] * p.s[t-1] * 5
        end

        # Define an equation for YGROSS
        v.YGROSS[t] = p.tfp[t] * v.K[t]^p.share * p.l[t]^(1-p.share)
    end
end

@defcomp emissions begin
    E       = Variable(index=[time])    # Total greenhouse gas emissions
    sigma   = Parameter(index=[time])   # Emissions output ratio
    YGROSS  = Parameter(index=[time])   # Gross output - Note that YGROSS is now a parameter

    function run_timestep(p, v, d, t)
        # Define an eqation for E
        v.E[t] = p.YGROSS[t] * p.sigma[t]   # Note the p. in front of YGROSS
    end
end

@Mimi.defmodel model begin

    index[time] = 2015:5:2110

    # Order matters here. If the emissions component were defined first, the model would not run.
    component(grosseconomy)
    component(emissions)

    # Set parameters for the grosseconomy component
    grosseconomy.l = [(1. + 0.015)^t * 6404 for t in 1:20]
    grosseconomy.tfp = [(1 + 0.065)^t * 3.57 for t in 1:20]
    grosseconomy.s = ones(20) * 0.22
    grosseconomy.depk = 0.1
    grosseconomy.k0 = 130.
    grosseconomy.share = 0.3

    # Set parameters for the emissions component
    emissions.sigma = [(1. - 0.05)^t * 0.58 for t in 1:20]

    # Connect parameters
    grosseconomy.YGROSS => emissions.YGROSS
end

# Above macro yields this:
# quote
#     tutorial = (Mimi.Model)()
#     (Mimi.set_dimension!)(tutorial, :time, 2015:5:2110)
#     (Mimi.add_comp!)(tutorial, Main.grosseconomy, :grosseconomy)
#     (Mimi.add_comp!)(tutorial, Main.emissions, :emissions)
#     (Mimi.set_param!)(tutorial, :grosseconomy, :l, [(1.0 + 0.015) ^ t * 6404 for t = 1:20])
#     (Mimi.set_param!)(tutorial, :grosseconomy, :tfp, [(1 + 0.065) ^ t * 3.57 for t = 1:20])
#     (Mimi.set_param!)(tutorial, :grosseconomy, :s, ones(20) * 0.22)
#     (Mimi.set_param!)(tutorial, :grosseconomy, :depk, 0.1)
#     (Mimi.set_param!)(tutorial, :grosseconomy, :k0, 130.0)
#     (Mimi.set_param!)(tutorial, :grosseconomy, :share, 0.3)
#     (Mimi.set_param!)(tutorial, :emissions, :sigma, [(1.0 - 0.05) ^ t * 0.58 for t = 1:20])
#     (Mimi.connect_param!)(tutorial, :emissions, :YGROSS, :grosseconomy, :YGROSS)
# end

run(model)

# Show model results
model[:emissions, :E]

# Or, as a DataFrame
getdataframe(model, :emissions, :E)


using PyPlot

d = 0:0.01:1

function kp(d, A...)
    s = d * (1-d)
    r = 0.0
    for (i, a) in enumerate(A)
        r += a * (1-d)^(i-1)
    end
    return s * r
end

fig, axs = subplots()
plot(d, kp.(d, 0.0, 1.0), "-")
plot(d, kp.(d, 1.0, 0.0), "--")
plot(d, kp.(d, 0.1, 1.0), "-.")
plot(d, kp.(d, 1.0, -1.0), ":")
gcf()


# -------------------------------------------------------------------------------
# Projection PatMat hinge
# -------------------------------------------------------------------------------
# general solver
function projection_exact(α0::AbstractVector,
                          β0::AbstractVector,
                          δ0::Real,
                          C1::Real,
                          C2::Real;
                          solver::Any = Optimizer(verbose = false))

    α = Convex.Variable(length(α0))
    β = Convex.Variable(length(β0))
    δ = Convex.Variable()

    objective   = Convex.sumsquares(α - α0)/2 + Convex.sumsquares(β - β0)/2 +
                  Convex.sumsquares(δ - δ0)/2
    constraints = [Convex.sum(α) == Convex.sum(β),
                   α <= C1,
                   α >= 0,
                   β <= C2*δ,
                   β >= 0]

    problem = Convex.minimize(objective, constraints)
    Convex.solve!(problem, solver)

    return vec(α.value), vec(β.value), δ.value
end

# our solution
function f(μ::Real, α0::AbstractVector, β0::AbstractVector, δ0::Real, C1::Real, C2::Real)
    return C2*δ0 + C2^2*sum(max.(β0 .- μ, 0)) - μ
end

function h(μ::Real, α0::AbstractVector, β0::AbstractVector, δ0::Real, C1::Real, C2::Real)
    λ = f(μ, α0, β0, δ0, C1, C2)
    return sum(min.(max.(α0 .- λ, 0), C1)) - sum(min.(max.(β0 .+ λ, 0), λ + μ))
end

function projection(α0::AbstractVector, β0::AbstractVector, δ0::Real, C1::Real, C2::Real)

    if δ0 <= - C2*sum(max.(β0 .+ maximum(α0), 0))
        return zero(α0), zero(β0), zero(δ0)
    end

    μ_ub1 = minimum(β0)
    μ_ub2 = minimum(C2*δ0 + C2^2*sum(β0) .- α0)/(C2^2*length(β0) + 1)
    μ_ub  = min(μ_ub1, μ_ub2) - 1e8
    μ_lb  = maximum(β0) + C2*max(δ0, 0) + 1e8

    μ = Roots.find_zero(μ -> h(μ, α0, β0, δ0, C1, C2), (μ_lb, μ_ub))
    λ = f(μ, α0, β0, δ0, C1, C2)

    α = @. min(max(α0 - λ, 0), C1)
    β = @. min(max(β0 + λ, 0), λ + μ)
    δ = (λ + μ)/C2

    return α, β, δ
end

# -------------------------------------------------------------------------------
# Projection PatMat quadratic
# -------------------------------------------------------------------------------
# general solver
function projection_exact(α0::AbstractVector,
                          β0::AbstractVector,
                          δ0::Real;
                          solver::Any = Optimizer(verbose = false))

    α = Convex.Variable(length(α0))
    β = Convex.Variable(length(β0))
    δ = Convex.Variable()

    objective   = Convex.sumsquares(α - α0)/2 + Convex.sumsquares(β - β0)/2 +
                  Convex.sumsquares(δ - δ0)/2
    constraints = [Convex.sum(α) == Convex.sum(β),
                   α >= 0,
                   β >= 0,
                   δ >= 0]

    problem = Convex.minimize(objective, constraints)
    Convex.solve!(problem, solver)

    return vec(α.value), vec(β.value), δ.value
end

# our solution
function f(λ::Real, α0::AbstractVector, β0::AbstractVector, δ0::Real)
    return sum(max.(α0 .- λ, 0)) - sum(max.(β0 .+ λ, 0))
end

function projection(α0::AbstractVector, β0::AbstractVector, δ0::Real)

    if - maximum(β0) > maximum(α0)
        return zero(α0), zero(β0), zero(δ0)
    end

    λ = Roots.find_zero(λ -> f(λ, α0, β0, δ0), (- maximum(β0), maximum(α0)))

    α = @. max(α0 - λ, 0)
    β = @. max(β0 + λ, 0)
    δ = max(δ0, 1e-4)

    return α, β, δ
end

# -------------------------------------------------------------------------------
# Projection TopPushK hinge
# -------------------------------------------------------------------------------
# general solver
function projection_exact(α0::AbstractVector,
                          β0::AbstractVector,
                          C::Real,
                          K::Integer;
                          solver::Any = Optimizer(verbose = false))

    α = Convex.Variable(length(α0))
    β = Convex.Variable(length(β0))

    objective   = Convex.sumsquares(α - α0)/2 + Convex.sumsquares(β - β0)/2
    constraints = [Convex.sum(α) == Convex.sum(β),
                   α <= C,
                   α >= 0,
                   β <= Convex.sum(α)/K,
                   β >= 0]

    problem = Convex.minimize(objective, constraints)
    Convex.solve!(problem, solver)

    return vec(α.value), vec(β.value)
end

# our solution
function f(μ::Real, s::AbstractVector, K::Integer)
    i, j, d  = 2, 1, 1
    λ, λ_old = s[1], 0
    g, g_old = -K*μ, -K*μ

    while g < 0
        g_old = g
        λ_old = λ

        if s[i] <= s[j] + μ
            g += d*(s[i] - λ)
            λ  = s[i]
            d += 1
            i += 1
        else
            g += d*(s[j] + μ - λ)
            λ  = s[j] + μ
            d -= 1
            j += 1
        end
    end
    return -(λ - λ_old)/(g - g_old)*g_old + λ_old
end

function h(μ::Real, s::AbstractVector, α0::AbstractVector, β0::AbstractVector, C::Real, K::Integer)
    λ = f(μ, s, K)
    return sum(min.(max.(α0 .- λ .+ sum(max.(β0 .+ λ .- μ, 0))/K, 0), C)) - K*μ
end

function projection(α0::AbstractVector, β0::AbstractVector, C::Real, K::Real)

    if K >= length(β0)
        @error "No feasible solution: K < length(β0) needed."
        return α0, β0
    end

    if mean(partialsort(β0, 1:K; rev = true)) + maximum(α0) <= 0
        return zero(α0), zero(β0)
    end

    s    = vcat(.- sort(β0; rev = true), Inf)
    μ_lb = 1e-10
    μ_ub = length(α0)*C/K + 1e-6

    μ = Roots.find_zero(μ -> h(μ, s, α0, β0, C, K), (μ_lb, μ_ub))
    λ = f(μ, s, K)
    δ = sum(max.(β0 .+ λ .- μ, 0))/K

    α = @. max(min(α0 - λ + δ, C), 0)
    β = @. max(min(β0 + λ, μ), 0)

    return α, β
end

# ---------------------------------------------------------------------------------------
# Projection TopPushK quadratic
# ---------------------------------------------------------------------------------------
# general solver
function projection_exact(α0::AbstractVector,
                          β0::AbstractVector,
                          K::Integer;
                          solver::Any = Optimizer(verbose = false))

    α = Convex.Variable(length(α0))
    β = Convex.Variable(length(β0))

    objective   = Convex.sumsquares(α - α0)/2 + Convex.sumsquares(β - β0)/2
    constraints = [Convex.sum(α) == Convex.sum(β),
                   α >= 0,
                   β <= Convex.sum(α)/K,
                   β >= 0]

    problem = Convex.minimize(objective, constraints)
    Convex.solve!(problem, solver)

    return vec(α.value), vec(β.value)
end

# our solution
function h(μ::Real, s::AbstractVector, α0::AbstractVector, β0::AbstractVector, K::Integer)
    λ = f(μ, s, K)
    return sum(max.(α0 .- λ .+ sum(max.(β0 .+ λ .- μ, 0))/K, 0)) - K*μ
end

function projection(α0::AbstractVector, β0::AbstractVector, K::Integer)

    if K >= length(β0)
        @error "No feasible solution: K < length(β0) needed."
        return α0, β0
    end

    if mean(partialsort(β0, 1:K; rev = true)) + maximum(α0) <= 0
        return zero(α0), zero(β0)
    end

    s    = vcat(.- sort(β0; rev = true), Inf)
    μ_lb = 1e-10
    μ_ub = length(α0)*(maximum(α0) + maximum(β0))/K

    μ = Roots.find_zero(μ -> h(μ, s, α0, β0, K), (μ_lb, μ_ub))
    λ = f(μ, s, K)
    δ = sum(max.(β0 .+ λ .- μ, 0))/K

    α = @. max(α0 - λ + δ, 0)
    β = @. max(min(β0 + λ, μ), 0)

    return α, β
end


module TestGeneralised
using Test

using Diversity
using Diversity.ShortNames

metacom = Metacommunity([1 2; 1 2]./6)

# Basic checks for the diversity() function
@testset "diversity()" begin
    @test diversity([metacommunityDiversity], [ρ̄], metacom, 1)[!,:diversity] ≈ [1.0]
    @test diversity([subcommunityDiversity], [ρ̄], metacom, 1)[!,:diversity] ≈ [1.0, 1.0]
end

pops = rand(5, 4)
pops = pops / sum(pops)
meta = Metacommunity(pops)
@testset "all-in-one norm/raw_alpha/beta/rho/gamma_div()" begin
    ra = α(meta)
    na = ᾱ(meta)
    rb = β(meta)
    nb = β̄(meta)
    rr = ρ(meta)
    nr = ρ̄(meta)
    g  = Γ(meta)
    # Alpha diversities
    @test raw_sub_alpha(meta, 0)[!,:diversity] ≈ subdiv(ra, 0)[!,:diversity]
    @test raw_meta_alpha(meta, 1)[!,:diversity] ≈ metadiv(ra, 1)[!,:diversity]
    @test norm_sub_alpha(meta, 2)[!,:diversity] ≈ subdiv(na, 2)[!,:diversity]
    @test norm_meta_alpha(meta, 3)[!,:diversity] ≈ metadiv(na, 3)[!,:diversity]
    # Beta diversities
    @test raw_sub_beta(meta, 0.0)[!,:diversity] ≈ subdiv(rb, 0.0)[!,:diversity]
    @test raw_meta_beta(meta, 1.0)[!,:diversity] ≈ metadiv(rb, 1.0)[!,:diversity]
    @test norm_sub_beta(meta, 2.0)[!,:diversity] ≈ subdiv(nb, 2.0)[!,:diversity]
    @test norm_meta_beta(meta, 3.0)[!,:diversity] ≈ metadiv(nb, 3.0)[!,:diversity]

    @test raw_sub_rho(meta, Inf)[!,:diversity] ≈ subdiv(rr, Inf)[!,:diversity]
    @test raw_meta_rho(meta, Inf)[!,:diversity] ≈ metadiv(rr, Inf)[!,:diversity]
    @test norm_sub_rho(meta, Inf)[!,:diversity] ≈ subdiv(nr, Inf)[!,:diversity]
    @test norm_meta_rho(meta, Inf)[!,:diversity] ≈ metadiv(nr, Inf)[!,:diversity]

    # Gamma diversities
    @test sub_gamma(meta, 0.5)[!,:diversity] ≈ subdiv(g, 0.5)[!,:diversity]
    @test meta_gamma(meta, 1.5)[!,:diversity] ≈ metadiv(g, 1.5)[!,:diversity]

end

end


using StatsModels
import Statistics.quantile
#import Splines2.bs
import Base.show
using Random
using BSplines
const bsPLINE_CONTEXT = Any
#x = Random.rand(20)
#x = collect(range(0,stop=10,step=0.1))
#y = log.(x)
#scatter(x,y)

#using Splines2
#f = Splines2.bs_(x,df=5,intercept=true)
#f(x)

function offsetArrayToZeros!(oneRow, spl)
    oneRow[spl.offsets[1]+1:spl.offsets[1]+length(spl)] = parent(spl)
end

function genSplFunction(x, df)
    p = range(0.0, length = df - 1, stop = 1.0)
    breakpoints = quantile(x, p)
    basis = BSplineBasis(4, breakpoints) # 4 = cubic
    return x -> splFunction(x, basis)
end
function splFunction(x, basis)
    df = length(basis)

    large = zeros(Union{Missing,Float64},length(x), df)

    bs_eval = bsplines.(Ref(basis), x)

    for k = 1:length(bs_eval)
        if isnothing(bs_eval[k])
            @warn("spline prediction outside of possible range, putting to missing")
            large[k,:] .= missing
        else
            offsetArrayToZeros!(view(large, k, :), bs_eval[k])
        end
    end

    return large
end

#spl(x,df) = Splines2.bs(x,df=df,intercept=true) # assumes intercept
spl(x, df) = 1

# make a nice call if the function is called via REPL
spl(t::Symbol, d::Int) = uf_bsplineTerm(term(t), term(d))

mutable struct uf_bsplineTerm{T,D} <: AbstractTerm
    term::T
    df::D
    fun::Any # function handle
end
function uf_bsplineTerm(term, df)
    uf_bsplineTerm(term, df, nothing)
end

Base.show(io::IO, p::uf_bsplineTerm) = print(io, "spl($(p.term), $(p.df))")

function StatsModels.apply_schema(
    t::FunctionTerm{typeof(spl)},
    sch::StatsModels.Schema,
    Mod::Type{<:bsPLINE_CONTEXT},
)
    apply_schema(uf_bsplineTerm(t.args_parsed...), sch, Mod)
end
function StatsModels.apply_schema(
    t::uf_bsplineTerm,
    sch::StatsModels.Schema,
    Mod::Type{<:bsPLINE_CONTEXT},
)
    term = apply_schema(t.term, sch, Mod)
    isa(term, ContinuousTerm) ||
        throw(ArgumentError("uf_bsplineTerm only works with continuous terms (got $term)"))

    isa(t.df, ConstantTerm) ||
        throw(ArgumentError("uf_bsplineTerm df must be a number (got $(t.df))"))
    uf_bsplineTerm(term, t.df.n)
end
function StatsModels.modelcols(p::uf_bsplineTerm, d::NamedTuple)

    col = modelcols(p.term, d)

    if isnothing(p.fun)
        p.fun = genSplFunction(col, p.df)#Splines2.bs_(col,df=p.df+1,intercept=true)
    end
    #X = Splines2.bs(col, df=p.df+1,intercept=true)
    X = p.fun(col)

    # remove middle X to negate intercept = true, generating a pseudo effect code
    X[:, Not(Int(ceil(end / 2)))]
end
#StatsModels.terms(p::uf_bsplineTerm) = terms(p.term)
StatsModels.termvars(p::uf_bsplineTerm) = StatsModels.termvars(p.term)
StatsModels.width(p::uf_bsplineTerm) = p.df
StatsModels.coefnames(p::uf_bsplineTerm) =
    "spl(" .* coefnames(p.term) .* "," .* string.(1:p.df) .* ")"


using Plots, MAT

nprocs = (2, 2) # nprocs (x, y) dim

@views function vizme2D_mpi(nprocs)
    H  = []
    ip = 1
    for ipx = 1:nprocs[1]
        for ipy = 1:nprocs[2]
            file = matopen("H_$(ip-1).mat"); H_loc = read(file, "H"); close(file)
            nx_i, ny_i = size(H_loc,1)-2, size(H_loc,2)-2
            ix1, iy1   = 1+(ipx-1)*nx_i, 1+(ipy-1)*ny_i
            if (ip==1)  H = zeros(nprocs[1]*nx_i, nprocs[2]*ny_i)  end
            H[ix1:ix1+nx_i-1,iy1:iy1+ny_i-1] .= H_loc[2:end-1,2:end-1]
            ip += 1
        end
    end
    fontsize = 12
    opts = (aspect_ratio=1, yaxis=font(fontsize, "Courier"), xaxis=font(fontsize, "Courier"),
        ticks=nothing, framestyle=:box, titlefontsize=fontsize, titlefont="Courier",
        xlabel="Lx", ylabel="Ly", xlims=(1, size(H,1)), ylims=(1, size(H,2)) )
    display(heatmap(H'; c=:davos, title="diffusion 2D MPI", opts...))
    return
end

vizme2D_mpi(nprocs)


##########################
### Variance estimator ###
##########################

abstract type VarEstimator{T} end

# NOTE: this naive variance estimator is used only in testing
mutable struct NaiveVar{T} <: VarEstimator{T}
    n :: Int
    S :: Vector{Vector{T}}
end

NaiveVar() = NaiveVar(0, Vector{Vector{Float64}}())
NaiveVar(::Int) = NaiveVar()

function add_sample!(nc::NaiveVar, s::AbstractVector)
    nc.n += 1
    push!(nc.S, s)
end

function reset!(nc::NaiveVar{T}) where {T<:AbstractFloat}
    nc.n = 0
    nc.S = Vector{Vector{T}}()
end

function get_var(nc::NaiveVar{T})::Vector{T} where {T<:AbstractFloat}
    @assert nc.n >= 2 "Cannot get variance with only one sample"
    return Statistics.var(nc.S)
end

# Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
mutable struct WelfordVar{T<:AbstractFloat,AT<:AbstractVector{T}} <: VarEstimator{T}
    n :: Int
    μ :: AT
    M :: AT
end

WelfordVar(d::Int) = WelfordVar(0, zeros(d), zeros(d))

function reset!(wv::WelfordVar{T,AT}) where {T<:AbstractFloat,AT<:AbstractVector{T}}
    wv.n = 0
    wv.μ .= zero(T)
    wv.M .= zero(T)
end

function add_sample!(wv::WelfordVar, s::AbstractVector)
    wv.n += 1
    δ = s .- wv.μ
    wv.μ .+= δ ./ wv.n
    wv.M .+= δ .* (s .- wv.μ)
end

# https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/var_adaptation.hpp
function get_var(wv::WelfordVar{T,AT})::AT where {T<:AbstractFloat,AT<:AbstractVector{T}}
    n, M = wv.n, wv.M
    @assert n >= 2 "Cannot get covariance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M .+ 1e-3 * (5 / (n + 5))
end

abstract type CovEstimator{T} end

# NOTE: this naive covariance estimator is used only in testing
mutable struct NaiveCov{T} <: CovEstimator{T}
    n :: Int
    S :: Vector{Vector{T}}
end

NaiveCov() = NaiveCov(0, Vector{Vector{Float64}}())
NaiveCov(::Int) = NaiveCov()

function add_sample!(nc::NaiveCov, s::AbstractVector)
    nc.n += 1
    push!(nc.S, s)
end

function reset!(nc::NaiveCov{T}) where {T<:AbstractFloat}
    nc.n = 0
    nc.S = Vector{Vector{T}}()
end

function get_cov(nc::NaiveCov{T})::Matrix{T} where {T<:AbstractFloat}
    @assert nc.n >= 2 "Cannot get covariance with only one sample"
    return Statistics.cov(nc.S)
end

# Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
mutable struct WelfordCov{T<:AbstractFloat} <: CovEstimator{T}
    n :: Int
    μ :: Vector{T}
    M :: Matrix{T}
end

WelfordCov(d::Int) = WelfordCov(0, zeros(d), zeros(d,d))

function reset!(wc::WelfordCov{T}) where {T<:AbstractFloat}
    wc.n = 0
    wc.μ .= zero(T)
    wc.M .= zero(T)
end

function add_sample!(wc::WelfordCov, s::AbstractVector)
    wc.n += 1
    δ = s .- wc.μ
    wc.μ .+= δ ./ wc.n
    wc.M .+= (s .- wc.μ) * δ'
end
# Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/covar_adaptation.hpp
function get_cov(wc::WelfordCov{T})::Matrix{T} where {T<:AbstractFloat}
    n, M = wc.n, wc.M
    @assert n >= 2 "Cannot get variance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M + 1e-3 * (5 / (n + 5)) * LinearAlgebra.I
end

################
### Adaptors ###
################
# TODO: integrate metric into preconditioner
abstract type AbstractPreConditioner <: AbstractAdaptor end
struct UnitPreConditioner <: AbstractPreConditioner end

string(::UnitPreConditioner) = "I"
adapt!(::UnitPreConditioner, ::AbstractVector{<:Real}, ::AbstractFloat, is_update::Bool=true) = nothing
reset!(::UnitPreConditioner) = nothing
getM⁻¹(dpc::UnitPreConditioner) = nothing

mutable struct DiagPreConditioner{T<:Real,AT<:AbstractVector{T}} <: AbstractPreConditioner
    n_min   :: Int
    ve  :: VarEstimator{T}
    var :: AT
end

function DiagPreConditioner(d::Int, n_min::Int=10)
    ve = WelfordVar(0, zeros(d), zeros(d))
    return DiagPreConditioner(n_min, ve, Vector(ones(d)))
end

function string(dpc::DiagPreConditioner)
    return string(dpc.var)
end

function adapt!(dpc::DiagPreConditioner, θ::AbstractVector{<:Real}, α::AbstractFloat, is_update::Bool=true)
    add_sample!(dpc.ve, θ)
    if dpc.ve.n >= dpc.n_min && is_update
        dpc.var .= get_var(dpc.ve)
    end
end

reset!(dpc::DiagPreConditioner) = reset!(dpc.ve)

function getM⁻¹(dpc::DiagPreConditioner)
    return dpc.var
end

mutable struct DensePreConditioner{T<:AbstractFloat} <: AbstractPreConditioner
    n_min :: Int
    ce    :: CovEstimator{T}
    covar :: Matrix{T}
end

function DensePreConditioner(d::Integer, n_min::Int=10)
    ce = WelfordCov(d)
    # TODO: take use of the line below when we have an interface to set which pre-conditioner to use
    # ce = NaiveCov()
    return DensePreConditioner(n_min, ce, LinearAlgebra.diagm(0 => ones(d)))
end

function string(dpc::DensePreConditioner)
    return string(LinearAlgebra.diag(dpc.covar))
end

function adapt!(dpc::DensePreConditioner, θ::AbstractVector{<:AbstractFloat}, α::AbstractFloat, is_update::Bool=true)
    add_sample!(dpc.ce, θ)
    if dpc.ce.n >= dpc.n_min && is_update
        dpc.covar .= get_cov(dpc.ce)
    end
end

reset!(dpc::DensePreConditioner) = reset!(dpc.ce)

function getM⁻¹(dpc::DensePreConditioner)
    return dpc.covar
end


module TestStepwise
using Test
using CompressedSensing: gaussian_data, perturb, rmp, foba

@testset "Stepwise Regression Algorithms" begin
    n, m, k = 32, 64, 3
    A, x, b = gaussian_data(n, m, k)
    δ = 1e-2
    y = perturb(b, δ)

    @testset "Relevance Matching Pursuit" begin
        xrmp = rmp(A, b, k)
        # noiseless
        @test xrmp.nzind == x.nzind
        @test xrmp.nzval ≈ x.nzval

        # noisy
        xrmp = rmp(A, y, δ)
        @test xrmp.nzind == x.nzind
        @test isapprox(xrmp.nzval, x.nzval, atol = 2δ)

        # multiple outer loop iterations
        xrmp = rmp(A, y, δ, 3)
        @test xrmp.nzind == x.nzind
        @test isapprox(xrmp.nzval, x.nzval, atol = 2δ)
    end

    @testset "FoBa" begin
        # noiseless
        xfoba = foba(A, b, δ)
        @test xfoba.nzind == x.nzind
        @test xfoba.nzval ≈ x.nzval

        # noisy
        xfoba = foba(A, y, δ)
        @test xfoba.nzind == x.nzind
        @test isapprox(xfoba.nzval, x.nzval, atol = 2δ)
    end
end

end


immutable Gamma <: ContinuousUnivariateDistribution
    shape::Float64
    scale::Float64
    function Gamma(sh::Real, sc::Real)
        if sh > 0.0 && sc > 0.0
            new(float64(sh), float64(sc))
        else
            error("Both shape and scale must be positive")
        end
    end
end

Gamma(sh::Real) = Gamma(sh, 1.0)
Gamma() = Gamma(1.0, 1.0) # Standard exponential distribution

@_jl_dist_2p Gamma gamma

function entropy(d::Gamma)
    x = (1.0 - d.shape) * digamma(d.shape)
    x += lgamma(d.shape) + log(d.scale) + d.shape
    return x
end

insupport(d::Gamma, x::Number) = isreal(x) && isfinite(x) && 0.0 <= x

kurtosis(d::Gamma) = 6.0 / d.shape

mean(d::Gamma) = d.shape * d.scale

function mgf(d::Gamma, t::Real)
    k, theta = d.shape, d.scale
    return (1.0 - t * theta)^(-k)
end

function cf(d::Gamma, t::Real)
    k, theta = d.shape, d.scale
    return (1.0 - im * t * theta)^(-k)
end

function modes(d::Gamma)
    if d.shape >= 1.0
        d.scale * (d.shape - 1.0)
    end
end

# rand()
#
#  A simple method for generating gamma variables - Marsaglia and Tsang (2000)
#  http://www.cparity.com/projects/AcmClassification/samples/358414.pdf
#  Page 369
#  basic simulation loop for pre-computed d and c
function randg2(d::Float64, c::Float64)
    while true
        x = v = 0.0
        while v <= 0.0
            x = randn()
            v = 1.0 + c * x
        end
        v = v^3
        U = rand()
        x2 = x^2
        if U < 1.0 - 0.331 * x2^2 ||
           log(U) < 0.5 * x2 + d * (1.0 - v + log(v))
            return d * v
        end
    end
end

function rand(d::Gamma)
    dpar = (d.shape <= 1.0 ? d.shape + 1.0 : d.shape) - 1.0 / 3.0
    return d.scale *
           randg2(dpar, 1.0 / sqrt(9.0 * dpar)) *
           (d.shape > 1.0 ? 1.0 : rand()^(1.0 / d.shape))
end

function rand!(d::Gamma, A::Array{Float64})
    dpar = (d.shape <= 1.0 ? d.shape + 1.0 : d.shape) - 1.0 / 3.0
    cpar = 1.0 / sqrt(9.0 * dpar)
    for i in 1:length(A)
        A[i] = randg2(dpar, cpar)
    end
    if d.shape <= 1.0
        ainv = 1.0 / d.shape
        for i in 1:length(A)
            A[i] *= rand()^ainv
        end
    end
    return d.scale * A
end

skewness(d::Gamma) = 2.0 / sqrt(d.shape)

var(d::Gamma) = d.shape * d.scale * d.scale


using DSP, Test

# Naive rational resampler
function naivefilt(h::Vector, x::Vector, resamplerate::Union{Integer, Rational}=1)

    upfactor     = numerator(resamplerate)
    downfactor   = denominator(resamplerate)
    xLen         = length(x)
    xZeroStuffed = zeros(eltype(x), length(x) * upfactor)

    for n in 0:length(x)-1
        xZeroStuffed[n*upfactor+1] = x[n+1]
    end

    y = filt(h, one(eltype(x)), xZeroStuffed)
    y = [y[n] for n = 1:downfactor:length(y)]
end

# Naive arbitrary resampler
function naivefilt(h::Vector, x::Vector, resamplerate::AbstractFloat, numfilters::Integer=32)
    xLen          = length(x)
    xInterpolated = naivefilt(h, x, numfilters//1)
    xLen          = length(xInterpolated)
    yLen          = ceil(Int, xLen * resamplerate)
    y             = similar(x, yLen)
    yIdx          = 1
    xIdx          = 1
    α             = 0.0
    (δ, ϕStride)  = modf(numfilters/resamplerate)
    ϕStride       = convert(Int, ϕStride)

    while xIdx < xLen
        yLower  = xInterpolated[xIdx]
        yUpper  = xInterpolated[xIdx+1]
        y[yIdx] = yLower + α*(yUpper - yLower)
        yIdx   += 1
        α      += δ
        xIdx   += floor(Int, α) + ϕStride
        α       = mod(α, 1.0)
    end

    resize!(y, yIdx-1)

    return y
end

# Disable time and printf macros when not running interactivly ( for travis )
if isinteractive()
    macro timeifinteractive(ex)
        :(@time $(esc(ex)))
    end

    macro printfifinteractive(args...)
        :(@printf($(map(esc, args)...)))
    end
else
    macro timeifinteractive(ex)
        esc(ex)
    end

    macro printfifinteractive(args...)
    end
end

#
# Single rate filtering
#

function test_singlerate(h, x)
    xLen       = length(x)
    hLen       = length(h)
    pivotPoint = min(rand(50:150), div(xLen, 4))
    x1         = x[1:pivotPoint]
    x2         = x[pivotPoint+1:end]

    @printfifinteractive( "\n\n" )
    @printfifinteractive( "____ _ _  _ ____ _    ____    ____ ____ ___ ____\n" )
    @printfifinteractive( "[__  | |\\ | | __ |    |___    |__/ |__|  |  |___\n" )
    @printfifinteractive( "___] | | \\| |__] |___ |___    |  \\ |  |  |  |___\n" )
    @printfifinteractive( "\nTesting single-rate fitering, h is %s, x is %s. xLen = %d, hLen = %d", string(eltype(h)), string(eltype(x)), xLen, hLen )

    @printfifinteractive( "\n\tBase.filt\n\t\t")
    @timeifinteractive naiveResult = filt(h, 1.0, x)

    @printfifinteractive( "\n\tDSP.filt( h, x, 1//1 )\n\t\t" )
    @timeifinteractive statelesResult = DSP.filt( h, x )
    @test naiveResult ≈ statelesResult

    @printfifinteractive( "\n\tDSP.filt. length(x1) = %d, length(x2) = %d\n\t\t", length(x1), length(x2) )
    myfilt = DSP.FIRFilter(h, 1//1)
    @timeifinteractive begin
        y1 = DSP.filt(myfilt, x1)
        y2 = DSP.filt(myfilt, x2)
    end
    statefulResult = [y1; y2]
    @test naiveResult ≈ statefulResult

    @printfifinteractive( "\n\tDSP.filt filt. Piecewise for first %d inputs\n\t\t", length(x1) )
    DSP.reset!(myfilt)
    @timeifinteractive begin
        for i in 1:length(x1)
            y1[i] = DSP.filt(myfilt, x1[i:i])[1]
        end
        y2 = DSP.filt(myfilt, x2)
    end
    piecewiseResult = [y1; y2]
    @test naiveResult ≈ piecewiseResult

    DSP.reset!(myfilt)
    @test inputlength(myfilt, length(piecewiseResult)) == xLen
end

#
# Decimation
#

function test_decimation(h, x, decimation)
    xLen       = length(x)
    hLen       = length(h)
    pivotPoint = min(rand(50:150), div(xLen, 4))
    x1         = x[1:pivotPoint]
    x2         = x[pivotPoint+1:end]

    @printfifinteractive( "\n\n" )
    @printfifinteractive( "___  ____ ____ _ _  _ ____ ___ _ ____ _  _ \n" )
    @printfifinteractive( "|  \\ |___ |    | |\\/| |__|  |  | |  | |\\ | \n" )
    @printfifinteractive( "|__/ |___ |___ | |  | |  |  |  | |__| | \\| \n" )
    @printfifinteractive( "\nTesting decimation. h::%s, x::%s. xLen = %d, hLen = %d, decimation = %d", string(typeof(h)), string(typeof(h)), xLen, hLen, decimation )

    @printfifinteractive( "\n\tNaive decimation\n\t\t")
    @timeifinteractive naiveResult = naivefilt(h, x, 1//decimation)

    @printfifinteractive( "\n\tDSP.filt( h, x, 1//%d)\n\t\t", decimation )
    @timeifinteractive statelesResult = DSP.filt(h, x, 1//decimation)
    @test naiveResult ≈ statelesResult

    @printfifinteractive( "\n\tDSP.filt decimation. length(x1) = %d, length(x2) = %d\n\t\t", length(x1), length(x2) )
    myfilt = DSP.FIRFilter(h, 1//decimation)
    @timeifinteractive begin
        y1 = DSP.filt(myfilt, x1)
        y2 = DSP.filt(myfilt, x2)
    end
    statefulResult = [y1; y2]
    @test naiveResult ≈ statefulResult

    @printfifinteractive( "\n\tDSP.filt decimation. Piecewise for first %d inputs.\n\t\t", length(x1) )
    DSP.reset!(myfilt)
    y1 = similar( x, 0 )
    @timeifinteractive begin
        for i in 1:length(x1)
            append!(y1, DSP.filt(myfilt, x1[i:i]))
        end
        y2 = DSP.filt(myfilt, x2)
    end
    piecewiseResult = [y1; y2]
    @test ≈(naiveResult, piecewiseResult, atol=sqrt(eps(real(one(eltype(x))))))

    DSP.reset!(myfilt)
    @test inputlength(myfilt, length(piecewiseResult)) == xLen
end

#
# Interpolation
#

function test_interpolation(h, x, interpolation)
    xLen       = length(x)
    hLen       = length(h)
    pivotPoint = min(rand(50:150), div(xLen, 4))
    x1         = x[1:pivotPoint]
    x2         = x[pivotPoint+1:end]

    @printfifinteractive( "\n\n" )
    @printfifinteractive( "_ _  _ ___ ____ ____ ___  _    ____ ____ ___ _ ____ _  _ \n" )
    @printfifinteractive( "| |\\ |  |  |___ |__/ |__] |    |  | |__|  |  | |  | |\\ | \n" )
    @printfifinteractive( "| | \\|  |  |___ |  \\ |    |___ |__| |  |  |  | |__| | \\| \n" )
    @printfifinteractive( "\nTesting interpolation, h::%s, x::%s. xLen = %d, hLen = %d, interpolation = %d", typeof(h), typeof(x), xLen, hLen, interpolation )

    @printfifinteractive( "\n\tNaive interpolation with Base.filt\n\t\t")
    @timeifinteractive begin
        xZeroStuffed = zeros(eltype(x), xLen * interpolation)
        for n = 0:xLen-1;
            xZeroStuffed[n*interpolation+1] = x[n+1]
        end
        naiveResult = filt(h, one(eltype(h)), xZeroStuffed)
    end

    @printfifinteractive( "\n\tDSP.filt( h, x, %d//1 )\n\t\t", interpolation )
    @timeifinteractive statelesResult = DSP.filt( h, x, interpolation//1 )
    @test naiveResult ≈ statelesResult

    @printfifinteractive( "\n\tDSP.filt interpolation. length(x1) = %d, length(x2) = %d\n\t\t", length(x1), length(x2) )
    myfilt = DSP.FIRFilter( h, interpolation//1 )
    @timeifinteractive begin
        y1 = DSP.filt(myfilt, x1)
        y2 = DSP.filt(myfilt, x2)
    end
    statefulResult = [y1; y2]
    @test naiveResult ≈ statefulResult

    DSP.reset!(myfilt)
    @test inputlength(myfilt, length(statefulResult)) == xLen

    @printfifinteractive( "\n\tDSP.filt interpolation. Piecewise for first %d inputs\n\t\t", length(x1) )
    DSP.reset!(myfilt)
    y1 = similar(x, 0)
    @timeifinteractive begin
        for i in 1:length(x1)
            append!(y1, DSP.filt(myfilt, x1[i:i]))
        end
        y2 = DSP.filt(myfilt, x2)
    end
    piecewiseResult = [y1; y2]
    @test ≈(naiveResult, piecewiseResult, atol=sqrt(eps(real(one(eltype(x))))))
end

#
# Rational resampling
#

function test_rational(h, x, ratio)
    xLen       = length(x)
    hLen       = length(h)
    pivotPoint = min(rand(50:150), div(xLen, 4))
    x1         = x[1:pivotPoint]
    x2         = x[pivotPoint+1:end]
    upfactor   = numerator(ratio)
    downfactor = denominator(ratio)
    resultType = promote_type(eltype(h), eltype(x))

    @printfifinteractive( "\n\n" )
    @printfifinteractive( "      ____ ____ ___ _ ____ _  _ ____ _    \n" )
    @printfifinteractive( "      |__/ |__|  |  | |  | |\\ | |__| |    \n" )
    @printfifinteractive( "      |  \\ |  |  |  | |__| | \\| |  | |___ \n" )
    @printfifinteractive( "                                          \n" )
    @printfifinteractive( "____ ____ ____ ____ _  _ ___  _    _ _  _ ____\n" )
    @printfifinteractive( "|__/ |___ [__  |__| |\\/| |__] |    | |\\ | | __\n" )
    @printfifinteractive( "|  \\ |___ ___] |  | |  | |    |___ | | \\| |__]\n" )
    @printfifinteractive( "\n\nTesting rational resampling, h::%s, x::%s. xLen = %d, hLen = %d, ratio = %d//%d", string(typeof(h)), string(typeof(x)), xLen, hLen, upfactor, downfactor )

    @printfifinteractive( "\n\tNaive rational resampling\n\t\t")
    @timeifinteractive naiveResult = naivefilt(h, x, ratio)

    @printfifinteractive( "\n\tDSP.filt( h, x, %d//%d )\n\t\t", upfactor, downfactor )
    @timeifinteractive statelesResult = DSP.filt(h, x, ratio)
    @test naiveResult ≈ statelesResult

    @printfifinteractive( "\n\tDSP.filt rational resampling. length(x1) = %d, length(x2) = %d\n\t\t", length(x1), length(x2) )
    myfilt = DSP.FIRFilter(h, ratio)
    @timeifinteractive begin
        s1 = DSP.filt(myfilt, x1)
        s2 = DSP.filt(myfilt, x2)
    end
    statefulResult = [s1; s2]
    @test naiveResult ≈ statefulResult

    @printfifinteractive( "\n\tDSP.filt rational. Piecewise for all %d inputs\n\t\t", length( x ) )
    reset!(myfilt)
    y1 = similar(x, 0)
    @timeifinteractive begin
        for i in 1:length(x)
            append!(y1, DSP.filt(myfilt, x[i:i]))
        end
    end
    piecewiseResult = y1
    @test naiveResult ≈ piecewiseResult

    DSP.reset!(myfilt)
    @test inputlength(myfilt, length(piecewiseResult)) == xLen
end

#
# Arbitrary resampling
#

function test_arbitrary(Th, x, resampleRate, numFilters)
    cutoffFreq      = 0.45
    transitionWidth = 0.05
    h               = digitalfilter(Lowpass(cutoffFreq, fs=numFilters), FIRWindow(transitionwidth=transitionWidth/numFilters)) .* numFilters
    h               = convert(Vector{Th}, h)
    myfilt          = DSP.FIRFilter(h, resampleRate, numFilters)
    xLen            = length(x)

    @printfifinteractive( "____ ____ ___      ____ ____ ____ ____ _  _ ___  _    _ _  _ ____\n" )
    @printfifinteractive( "|__| |__/ |__]     |__/ |___ [__  |__| |\\/| |__] |    | |\\ | | __\n" )
    @printfifinteractive( "|  | |  \\ |__] .   |  \\ |___ ___] |  | |  | |    |___ | | \\| |__]\n" )
    @printfifinteractive( "\n\nh::%s, x::%s, rate = %f, Nϕ = %d, xLen = %d, ", string(typeof(h)), string(typeof(x)), resampleRate, numFilters, length(x) )

    @printfifinteractive( "\n\tNaive arbitrary resampling\n\t\t" )
    @timeifinteractive naiveResult = naivefilt(h, x, resampleRate, numFilters)

    @printfifinteractive( "\n\tStateless arbitrary resampling\n\t\t" )
    @timeifinteractive statelessResult = DSP.filt(h, x, resampleRate, numFilters)

    @printfifinteractive( "\n\tStateful arbitrary resampling\n\t\t" )
    @timeifinteractive statefulResult = DSP.filt(myfilt, x)

    # DSP.reset!(myfilt)
    # TODO: figure out why this fails
    # @test inputlength(myfilt, length(statefulResult)) == xLen

    @printfifinteractive( "\n\tPiecewise arbitrary resampling\n\t\t" )
    reset!(myfilt)
    piecwiseResult = eltype(x)[]
    sizehint!(piecwiseResult, ceil(Int, length(x)*resampleRate))
    @timeifinteractive for i in 1:length(x)
        thisY = filt(myfilt, x[i:i])
        append!(piecwiseResult, thisY)
    end

    commonLen = min(length(naiveResult), length(statelessResult), length(statefulResult), length(piecwiseResult))
    resize!(naiveResult, commonLen)
    resize!(statelessResult, commonLen)
    resize!(statefulResult, commonLen)
    resize!(piecwiseResult, commonLen)

    @test naiveResult ≈ statelessResult
    @test naiveResult ≈ statefulResult
    @test naiveResult ≈ piecwiseResult
end

#
# Run the tests
#

@testset "interp.=$interpolation, dec.=$decimation, Th=$Th, Tx=$Tx" for interpolation in [1, 5, 14, 23],
            decimation in [1, 9, 17, 21],
                Th in [Float32, Float64],
                    Tx in [Float32, Float64, ComplexF32, ComplexF64]

    h     = rand(Th, rand(16:128))
    xLen  = rand(200:300)
    xLen  = xLen-mod(xLen, decimation)
    x     = rand(Tx, xLen)
    ratio = interpolation//decimation
    if ratio == 1
        test_singlerate(h, x)
    end
    if decimation != 1
        test_decimation(h, x, decimation)
    end
    if interpolation != 1
        test_interpolation(h, x, interpolation)
    end
    if numerator(ratio) == interpolation && denominator(ratio) == decimation && numerator(ratio) != 1 && denominator(ratio) != 1
        test_rational(h, x, ratio)
        if Tx in [Float32, ComplexF32]
            test_arbitrary(Th, x, convert(Float64, ratio)+rand(), 32)
        end
    end
    if decimation == 1
        test_rational(h, x, interpolation)
    end
end