export GLMModel, LogisticModel, LinearModel, QuantileModel
export grad_scratch!

##############################
## Used to specify GLM kind ##
## (link and family pair)   ##
##############################
abstract GLMModel
#Subtypes should define the loss, predict!, and grad_scratch! functions
# if the defaults are not appropriate.

function linpred!(pr::DenseVector{Float64},
                  coefs::Vector{Float64},
                  x::DSMat{Float64};
                  offset::Vector{Float64} = emptyvector(Float64))
        A_mul_B!(pr, x, coefs)
        if !isempty(offset)
            add!(pr, offset)
        end
        pr
end

#
# default predict! for GLMModel uses identity link
#
predict!(m::GLMModel,
        pr::DenseVector{Float64},
        coefs::Vector{Float64},
        x::DSMat{Float64};
        offset::Vector{Float64} = emptyvector(Float64)) = linpred!(pr, coefs, x; offset=offset)

# grad_scratch! calculates the gradient without allocating new memory
# for the residual vector.
# resid is scratch space and should have dims equal to size(y)
# Default is appropriate for GLMs with canonical link
function grad_scratch!(m::GLMModel,
                      gr::Vector{Float64},
                      coefs::Vector{Float64},
                      x::DSMat{Float64},
                      y::Vector{Float64},
                      resid::Vector{Float64};
                      offset::Vector{Float64} = emptyvector(Float64))
    predict!(m, resid, coefs, x; offset=offset)
    subtract!(resid, y) # resid = prediction - y
    alpha = 1.0/size(x, 1)
    At_mul_B!(alpha, x, resid, 0.0, gr)
    #BLAS.gemv!('T', alpha, x, resid, 0.0, gr) #grad \propto x'resid = -x'(y - pred)
    gr
end

grad!(m::GLMModel,
    gr::Vector{Float64},
    coefs::Vector{Float64},
    x::DSMat{Float64},
    y::Vector{Float64};
    offset::Vector{Float64} = emptyvector(Float64)) = grad_scratch!(m, gr, coefs, x, y, similar(y), offset=offset)

#########
## OLS ##
#########
type LinearModel <: GLMModel end

loss(m::LinearModel, pr::DenseVector{Float64}, y::Vector{Float64}) = meansqdiff(pr, y)/2

#########################
## Logistic regression ##
#########################
type LogisticModel <: GLMModel end

function predict!(m::LogisticModel,
                           pr::DenseVector{Float64},
                           coefs::Vector{Float64},
                           x::DSMat{Float64};
                           offset::Vector{Float64} = emptyvector(Float64))
    linpred!(pr, coefs, x; offset=offset)
    map1!(LogisticFun(), pr)
    pr
end

loss(m::LogisticModel, pr::DenseVector{Float64}, y::Vector{Float64}) = mean(NBLL(), pr, y)

#########################
## Quantile regression ##
#########################
type QuantileModel <: GLMModel
    tau :: Float64
end

QuantileModel() = QuantileModel(0.5)

function grad_scratch!(m::QuantileModel,
                               gr::Vector{Float64},
                               coefs::Vector{Float64},
                               x::DSMat{Float64},
                               y::Vector{Float64},
                               scratch::Vector{Float64};
                               offset::Vector{Float64} = emptyvector(Float64))
    predict!(m, scratch, coefs, x; offset=offset)
    subtract!(scratch, y)
    map1!(Qgrad(), scratch, m.tau)
    alpha = 2.0/size(x, 1)
    BLAS.gemv!('T', alpha, x, scratch, 0.0, gr)
    gr
end

loss(m::QuantileModel, pr::DenseVector{Float64}, y::Vector{Float64}) = meanabsdiff(pr, y)


module MCMCUtil

using LinearAlgebra
using Distributions

updateMean(t::Int, Xt::Float64, m0::Float64) = (Xt + m0*(t-1)) / t
updateMean(t::Int, Xt::Array{Float64,1}, m0::Array{Float64,1}) = @. (Xt + m0*(t-1)) / t

function updateVar(t::Int, Xt::Float64, mt::Float64, m0::Float64, V0::Float64)
    return V0*(t-2)/(t-1) + t*(mt-m0)^2
end

function updateVar(t::Int, Xt::Array{Float64,1}, mt::Array{Float64,1},
                   m0::Array{Float64,1}, V0::Array{Float64,2}, eps::Matrix{Float64})
    return V0.*(t-2)/(t-1) + 1/(t-1).*((t-1).*(m0 * transpose(m0)) .- t.* (mt * transpose(mt)) .+ (Xt * transpose(Xt)) + eps)
end

function MH_single(target_log::Function, init, prop)
    r = target_log(prop) - target_log(init)

    if log(rand()) < r
        return prop, 1
    else
        return init, 0
    end
end

function MH_single(target_log::Function, init::T, proposal::Function, init_loglik::T) where T <: Real
    prop = proposal(init)
    tlp = target_log(prop)
    tli = target_log(init, init_loglik)

    if log(rand()) < (tlp[1] - tli[1])
        return prop, 1, tlp[2]
    else
        return init, 0, tlp[2]
    end
end

## single adaptive Metropolis step
function AM_RR_single(target_log::Function, init::Vector{T}, V::Matrix{T},
                      t::Int, t0::Int=2*length(init),
                      s::T=(2.38)^2/length(init), beta::T=0.05) where T <: Real

    d = length(init)
    proposal = zeros(T, d)

    if (t <= t0) || (rand() < beta)
        proposal = rand(MvNormal(init, (0.1)^2*diagm(0 => ones(eltype(init), d))/d))
    else
        proposal = rand(MvNormal(init, s*V))
    end

    if log(rand()) < (target_log(proposal) - target_log(init))
        return proposal, 1
    else
        return init, 0
    end
end

function AM_RR_single(target_log::Function, init::Float64, V::Float64,
                      t::Int, t0::Int=2, s::Float64=2.38, beta::Float64=0.05)

    proposal::Float64 = 0

    if (t <= t0) | (rand(1)[1] < beta)
        proposal = rand(Normal(init, 0.1))
    else
        if(s*V == 0)
            proposal = init
        else
            proposal = rand(Normal(init, s*sqrt(V)))
        end
    end

    r = target_log(proposal) - target_log(init)

    if log(rand()) < r ##min(1.0, r[1])
        return proposal, 1
    else
        return init, 0
    end
end

end


function makejkl(n)

m = (n-2)*(n-1)*n/6;
indj = zeros(Int64,m);
indk = zeros(Int64,m);
indl = zeros(Int64,m);
count = 0;
for j = 1:(n-2)
    for k = (j+1):(n-1)
        for l = (k+1):n
            count = count+1;
            indj[count] = j;
            indk[count] = k;
            indl[count] = l;
        end
    end
end

return indj, indk, indl

end

function neighquadletlist(pairlist, pairnumsum, ilist)

pairnum = pairnumsum[2:end]-pairnumsum[1:end-1];
quadletnum  = (pairnum-2).*(pairnum-1).*pairnum/6;

inum = length(ilist);
quadletlistj = [];
quadletlistk = [];
quadletlistl = [];

for ii = 1:inum
    n1 = pairnumsum[ii];
    n2 = pairnumsum[ii+1];
    g = pairlist[(n1+1):n2]; # a list of neighbors around atom i

    indj, indk, indl = makejkl(n);
    quadletlistj = [quadletlistj; g[indj]];
    quadletlistk = [quadletlistk; g[indk]];
    quadletlistl = [quadletlistl; g[indl]];
end

quadletlist = [quadletlistj[:] quadletlistk[:] quadletlistl[:]];
quadletnum = [0; cumsum(quadletnum)];

quadletlist = Int64.(quadletlist)
quadletnum  = Int64.(quadletnum)

return quadletlist, quadletnum

end


@testset "458.poor-pigs.jl" begin
    @test poor_pigs(1000, 15, 60) == 5
    @test poor_pigs(1000, 15, 15) == 10
end


using SparseArrays
using Profile

"""
    assemble!(dofmap, a_kernel!, L_kernel!, bc_dofs, bc_vals)

Assemble bilinear and linear kernels, `a_kernel`, `L_kernel` and
applies boundary contitions in symmetric way.

Allocates and returns sparse CSC matrix and dense vector
of rhs.

"""
function assemble!(dofmap::DofMap, a_kernel!::Function, L_kernel!::Function,
    bc_dofs::Array{Bool, 1}, bc_vals::Array{Float64, 1})

    mesh = dofmap.mesh
    dofsize = dofmap.dim
    tdim = mesh.reference_cell.get_dimension()::Int64

    dofsize_local = size(dofmap.cell_dofs)[2]
    num_values = dofsize_local ^ 2
    num_cells = num_entities(mesh, tdim)
    cell_vert_conn = mesh.topology[(tdim, 0)]

    vert_per_cell = size(cell_vert_conn)[2]
    gdim = size(mesh.vertices)[2]

    # Vector of non-zero values in the global matrix
    # These values will be scattered according to "sparsity pattern"
    values = Array{Float64, 1}()

    # Storage for coordinates of cell vertices
    cell_coords = zeros(Float64, vert_per_cell, gdim)

    # Storage for insertion index sets
    I = Array{Int64, 1}()
    J = Array{Int64, 1}()

    b = zeros(Float64, dofsize)

    A_local = zeros(Float64, dofsize_local, dofsize_local)
    b_local = zeros(Float64, dofsize_local)

    for cell_id in 1:num_cells

        fill!(A_local, 0.0)
        fill!(b_local, 0.0)

        # Pack vertex coordinates of the cell
        cell_coords[1:vert_per_cell, 1:gdim] = mesh.vertices[cell_vert_conn[cell_id, 1:vert_per_cell], 1:gdim]

        a_kernel!(A_local, cell_coords)
        L_kernel!(b_local, cell_coords)

        # Append values to COO format
        for (i, iglobal) in enumerate(dofmap.cell_dofs[cell_id, 1:dofsize_local])

            if bc_dofs[iglobal]
                A_local[i, 1:dofsize_local] .= 0.0
                b_local[1:dofsize_local] -= A_local[1:dofsize_local, i] * bc_vals[iglobal]
                A_local[1:dofsize_local, i] .= 0.0
                A_local[i, i] = 1.0
                b_local[i] = bc_vals[iglobal]
            end

            b[iglobal] += b_local[i]
            for j in 1:dofsize_local
                # Push only nonzero values to avoid later "eliminate zeros" call
                if A_local[i, j] != 0.0
                    push!(values, A_local[i, j])
                    push!(I, iglobal)
                    push!(J, dofmap.cell_dofs[cell_id, j])
                end
            end
        end
    end

    # Return CSC sparse array
    A = sparse(I, J, values)

    return A, b
end


#=
# 2. Basic flow with a stationary body
In this notebook we will simulate the flow past a stationary body.
=#

using ViscousFlow
#-
using Plots

#=
### The basic steps
From the previous notebook, we add one additional step:
* **Specify the problem**: Set the Reynolds number and free stream
* **Discretize**: Set up a solution domain, grid cell size, time step size
* **Set up bodies**: *Create the body or bodies and specify their motions, if any*
* **Construct the system structure**: Create the operators that will be used to perform the simulation
* **Initialize**: Set the initial flow field and initialize the integrator
* **Solve**: Solve the flow field
* **Examine**: Examine the results
We will go through all of these here. For the examples we will carry out in this notebook,
the first three steps need only be carried out once.
=#

# ### Problem specification
# Set the Reynolds number and free stream
Re = 200 # Reynolds number
U = 1.0 # Free stream velocity
U∞ = (U,0.0);

#=
### Discretize
We set the grid Re to 4.0 here to get a quicker solution, though it is generally
better to make this smaller (it defaults to 2.0).
=#
xlim = (-1.0,3.0)
ylim = (-1.5,1.5)
Δx, Δt = setstepsizes(Re,gridRe=4.0)

#=
### Set up bodies
Here, we will set up a rectangle of half-height 0.5 and half-width 0.25
at 45 degrees angle of attack
=#
body = Rectangle(0.5,0.25,1.5Δx)

#=
We place the body at a desired location and orientation with the `RigidTransform`
function. This function creates an operator `T` that acts in-place on the body:
after the operation is applied, `body` is transformed to the correct location/orientation.
=#
cent = (0.0,0.0) # center of body
α = 45π/180 # angle
T = RigidTransform(cent,α)
T(body) # transform the body to the current configuration

# Let's plot it just to make sure
plot(body,xlim=xlim,ylim=ylim)

#=
### Construct the system structure
This step is like the previous notebook, but now we also provide the body and
the freestream:
=#
sys = NavierStokes(Re,Δx,xlim,ylim,Δt,body,freestream = U∞)

#=
### Initialize
Now, we initialize with zero vorticity. Note that we do this by calling
`newstate` with no argument except for `sys` itself.
=#
u0 = newstate(sys)

#=
and now create the integrator, with a long enough time span to hold the whole
solution history:
=#
tspan = (0.0,20.0)
integrator = init(u0,tspan,sys)

#=
### Solve
Now we are ready to solve the problem. Let's advance the solution to $t = 1$.
=#
@time step!(integrator,1.0)

#=
### Examine
Let's look at the flow field at the end of this interval
=#
plot(
plot(vorticity(integrator),sys,title="Vorticity",clim=(-10,10),levels=range(-10,10,length=30), color = :RdBu,ylim=ylim),
plot(streamfunction(integrator),sys,title="Streamlines",ylim=ylim,color = :Black),
    size=(700,300)
    )

#=
Now let's make a movie, like we did last time.
=#

sol = integrator.sol;
@gif for (u,t) in zip(sol.u,sol.t)
    plot(vorticity(u,sys,t),sys,clim=(-10,10),levels=range(-10,10,length=30), color = :RdBu)
end every 5

#=
#### Compute the force history
To do this, we supply the solution history `sol`, the system `sys`, and the index
of the body (1).
=#
fx, fy = force(sol,sys,1);

#=
Plot the histories. Note that we are actually plotting the drag and lift
coefficient histories here:
$$ C_D = \dfrac{F_x}{\frac{1}{2}\rho U_\infty^2 L}, \quad C_L = \dfrac{F_y}{\frac{1}{2}\rho U_\infty^2 L} $$
Since the quantities in this simulation are already scaled by $\rho$, $U_\infty$, and $L$
(because $\rho$ has been scaled out of the equations, and the free stream speed is
set to 1 and the height of the shape to 1), then we obtain these coefficients by
simply dividing by 1/2, or equivalently, by multiplying by 2:
=#
plot(
plot(sol.t,2*fx,xlim=(0,Inf),ylim=(0,6),xlabel="Convective time",ylabel="\$C_D\$",legend=:false),
plot(sol.t,2*fy,xlim=(0,Inf),ylim=(-6,6),xlabel="Convective time",ylabel="\$C_L\$",legend=:false),
    size=(800,350)
)

# The mean drag and lift coefficients are
meanCD = GridUtilities.mean(2*fx)
#-
meanCL = GridUtilities.mean(2*fy)


# Example 1 : 	Oscillating masses

# Constraints :
# 				Position in box
# 				Velocity in box
#				Acceleration in box
#       		Acceleration rate in box

# problem definition
module prb

	using LinearAlgebra, StaticArrays

	const nx = 8											# state dim
	const nu = 4											# imput dim
	const N = 30											# horizon length
	const Δ = 0.25 											# sampling time
	const tvec = [(j-1)*Δ for j in 1:N]   					# time vector (t=0:N-1)

	# solution method settings
		# common
		const kmax = 500										# max iterations for main algorithms (and outer loop)
		const scale_Hg = true

		# JuMP
		const ϵ_pd = 1e-9										# solver precision
		const ϵ_gap = 1e-9										# solver precision

		# PIPG
		const ϵ_pipg = 1e-4										# termination precision
		const restrt_idx = 40									# index for restarting

		# ADMM
		const α_admm = 1.0										# step size
		const ϵ_nest_admm = 1e-10								# Nesterov's method termination precision
		const kmax_nest = 100									# max iterations for Nesterov's method

		# Chambolle and Pock (PDHG)
		const β⁰ = 2.0											# initial step size
		const γ⁰ = 1.0											# initial step size

	# system parameters
	const nxb2 = Int64(nx/2)									# position veloity dim
	const pmax = 2.0											# position
	const vmax = 2.0											# velocity
	const umax = 2.0					 						# input (forcing)
	const dumax = 0.5											# input rate

	# reference trajectory
	const xref = [SVector{nx}(cat(ones(nxb2)...,zeros(nxb2)...,dims=1)) for _ in 1:N-1]	# state reference (t=1:N-1)
	const uref = [SVector{nu}(zeros(nu)) for _ in 1:N-1]							# input reference (t=0:N-2)
	const x0 = 1 .* ones(nx)								    					# initial condition (t=0)

	# dynamics (t=0:N-2)
	const Ac = Array([zeros(nxb2,nxb2) Diagonal(ones(nxb2))
						       Tridiagonal(ones(nxb2-1),-2 .* ones(nxb2),ones(nxb2-1)) zeros(nxb2,nxb2)])
	const Ad = SMatrix{nx,nx}(exp(Ac*Δ))
	const Bc = [zeros(nxb2,nxb2)
				Diagonal(ones(nxb2,nxb2))]
	const Bd = SMatrix{nx,nu}(inv(Ac)*(Ad-Diagonal(ones(nx)))*Bc)

	const A = [Ad for _ in 1:N-1]
	const B = [Bd for _ in 1:N-1]
	const e = [SVector{nx}(zeros(nx)) for _ in 1:N-1]

	# state and control penalty (LQR weights)

	const q_tmp = [1.0,1.0]
	const r_tmp = 1.0

	# const q_tmp = [0.5,1.0]
	# const r_tmp = 1.0

	const Q = [SMatrix{nx,nx}(Array(Diagonal(cat(q_tmp[1] .* ones(nxb2), q_tmp[2] .* ones(nxb2), dims=1)))) for _ in 1:N-1] # t=1:N-1
	const R = [SMatrix{nu,nu}(r_tmp .* Array(Diagonal(ones(nu)))) for _ in 1:N-1] 											# t=0:N-2

end

include("proj_funcs.jl") 		# projection library

# assemble quantities for vectorized PIPG+ implementation
module asm

	using ..prb
	using ..proj 								# projection onto simple sets
	using LinearAlgebra, StaticArrays, JuMP

	const scale_Hg = prb.scale_Hg
	const nH = 2*prb.nu*(prb.N-2)				# no. of inequality constraints

	function construct_H(nx,nu,N,A,B,Inx,Inu)
		nxnu = nx+nu
		H = zeros(nx*(N-1) + nH, nxnu*(N-1))
		H[1:nx,1:nxnu] .= [-B[1] Inx]
		for j=1:N-2
			H[j*nx+1:(j+1)*nx, nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+2*nx+nu] .= [-A[j] -B[j] Inx]

			H[nx*(N-1)+(j-1)*nu+1:nx*(N-1)+j*nu, (j-1)*nxnu+1:(j-1)*nxnu+nu+nxnu] .= [Inu zeros(nu,nx) -Inu]

			H[nx*(N-1)+nu*(N-2)+(j-1)*nu+1:nx*(N-1)+nu*(N-2)+j*nu, (j-1)*nxnu+1:(j-1)*nxnu+nu+nxnu] .= [-Inu zeros(nu,nx) Inu]
		end

		norm_row_H = zeros(nx*(N-1)+nH)
		for j=1:nx*(N-1)+nH
			norm_row_H[j] = norm(H[j,1:end],Inf)
		end

		if scale_Hg
			# print("Scaling rows of H\n")
			for j=1:length(norm_row_H)
				H[j,1:end] .= H[j,1:end] ./ norm_row_H[j]
			end
		end

		return H,norm_row_H
	end

	function construct_G(nx,nu,N,Q,R)
		nxnu = nx+nu
		G = zeros(nxnu*(N-1),nxnu*(N-1))
		for j=1:N-1
			G[(j-1)*nxnu+1:j*nxnu,(j-1)*nxnu+1:j*nxnu] .= [R[j] 		zeros(nu,nx)
														   zeros(nx,nu) Q[j]]
		end
		return G
	end

	H_scl,norm_row_H = construct_H(prb.nx,prb.nu,prb.N,prb.A,prb.B,Diagonal(ones(prb.nx)),Diagonal(ones(prb.nu)))
	const H = deepcopy(H_scl)
	const HT = transpose(H)

	const zref = cat([cat(prb.uref[t],prb.xref[t],dims=1) for t=1:prb.N-1]...,dims=1)
	const h = cat([cat(-prb.R[t]*prb.uref[t], -prb.Q[t]*prb.xref[t], dims=1) for t=1:prb.N-1]...,dims=1)

	if scale_Hg
		# print("Scaling components of g\n")
		const scl_val_Hg = norm_row_H
	else
		const scl_val_Hg = 1.0
	end

	const g = cat(prb.A[1]*prb.x0+prb.e[1],cat([prb.e[t] for t=2:prb.N-1]...,dims=1), -prb.dumax .* ones(2*prb.nu*(prb.N-2)),dims=1)

	@assert length(h) == (prb.nx+prb.nu)*(prb.N-1) "Length of h is incorrect."
	@assert length(g) == 2*prb.nu*(prb.N-2)+prb.nx*(prb.N-1) "Length of g is incorrect."

	const Heq = [H -1 .* Array(Diagonal(ones(prb.nx*(prb.N-1) + nH)))]
	const HeqT = transpose(Heq)

	const G = construct_G(prb.nx,prb.nu,prb.N,prb.Q,prb.R)

	const svd_HTH = svd(HT*H)
	const svd_HeqTHeq = svd(HeqT*Heq)

	const svd_G = svd(G)

	const σ = max(svd_HTH.S...)
	const σ_eq = max(svd_HeqTHeq.S...)
	const λ = max(svd_G.S...)
	const μ = min(svd_G.S...)

	function construct_xu(z::Array{Float64,1})
		return [z[nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+nx] for j=1:N-1], [z[(j-1)*nxnu+1:(j-1)*nxnu+nu] for j=1:N-1]
	end

	# useful constants
	const nx = copy(prb.nx)
	const nu = copy(prb.nu)
	const nxnu = nx+nu
	const nxb2 = Int64(nx/2)
	const N = copy(prb.N)

	const κ_x = MVector{nx}(randn(nx))
	const κ_u = MVector{nu}(randn(nu))

	# projection
	function project_U!(u)
		κ_u .= u
		proj.box!(u,κ_u,-prb.umax,prb.umax,nu)
	end
	function project_X!(x)
		κ_x .= x
		proj.box!(x[1:nxb2],κ_x[1:nxb2],-prb.pmax,prb.pmax,nxb2)
		proj.box!(x[nxb2+1:nx],κ_x[nxb2+1:nx],-prb.vmax,prb.vmax,nxb2)
	end

	# set state and input constraints in JuMP (vectorized version)
	function set_constraint_jump_vectorized!(model,zz)
		# acceleration, position and velocity in box
		@constraint(model,[j=1:N-1],-prb.umax .≤ zz[(j-1)*nxnu+1:(j-1)*nxnu+nu] .≤  prb.umax)
		@constraint(model,[j=1:N-1],-prb.pmax .≤ zz[nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+nxb2] .≤  prb.pmax)
		@constraint(model,[j=1:N-1],-prb.vmax .≤ zz[nu+(j-1)*nxnu+nxb2+1:nu+(j-1)*nxnu+nx] .≤  prb.vmax)
	end

	# set state and input constraints (everything other than dynamics and initial condition) in JuMP (decomposed version)
	function set_constraint_jump_decomposed!(model,x,u)
		@constraint(model,[t=1:N-1], -prb.pmax .≤ x[1:nxb2,t+1] .≤ prb.pmax)
		@constraint(model,[t=1:N-1], -prb.vmax .≤ x[nxb2+1:nx,t+1] .≤ prb.vmax)
		@constraint(model,[t=1:N-1], -prb.umax .≤ u[1:nu,t] .≤ prb.umax)
		@constraint(model,[t=1:N-2],-prb.dumax .≤ u[:,t+1] .- u[:,t] .≤ prb.dumax)
	end

end


print("\tmodel.jl")

m = 8
n = 8
A = sprand(m, n, 0.5)
c = 1.0 - 2.0*rand(n)
b = 1.0 + rand(m)
u = 10.0 + sprand(n, 0.8)
p = nnz(u)

# Tests for PrimalDualPoint
v1 = Tulip.PrimalDualPoint(rand(n), rand(p), rand(m), rand(n), rand(p))
v2 = copy(v1)
@test n == sum(v1.x .== v2.x)
@test p == sum(v1.w .== v2.w)
@test m == sum(v1.y .== v2.y)
@test n == sum(v1.s .== v2.s)
@test p == sum(v1.z .== v2.z)

v3 = v1 + v2
@test n == sum(v3.x .== (v1.x + v2.x))
@test p == sum(v3.w .== (v1.w + v2.w))
@test m == sum(v3.y .== (v1.y + v2.y))
@test n == sum(v3.s .== (v1.s + v2.s))
@test p == sum(v3.z .== (v1.z + v2.z))

# Create random instance and run checks
model = Tulip.Model(A, b, c, u.nzind, u.nzval)

@test m == model.nconstr
@test n == model.nvars
@test model.status == :Built

sol = model.sol
@test typeof(sol) <: Tulip.PrimalDualPoint
@test n == size(sol.x, 1)
@test p == size(sol.w, 1)
@test m == size(sol.y, 1)
@test n == size(sol.s, 1)
@test p == size(sol.z, 1)

println("\tPassed.")


using PtFEM

data = Dict(
  # Frame(nels, nn, ndim, nst, nip, finite_element(nod, nodof))
  :struc_el => Frame(7, 8, 2, 1, 1, Line(2, 3)),
  :properties => [
    1.0e10 1.0e6 20.0;
    1.0e10 1.0e6 50.0;
    1.0e10 1.0e6 80.0
     ],
  :etype => [1, 2, 2, 1, 3, 3, 1],
  :x_coords => [0.0,  0.0, 10.0, 20.0, 20.0, 35.0, 50.0, 50.0],
  :y_coords => [0.0, 15.0, 15.0, 15.0,  0.0, 15.0, 15.0,  0.0],
  :g_num => [
    1 2 3 5 4 6 8;
    2 3 4 4 6 7 7],
  :support => [
    (1, [0 0 0]),
    (5, [0 0 0]),
    (8, [0 0 0])
    ],
  :loaded_nodes => [
    (2, [4.0   0.0 0.0]),
    (3, [0.0  -6.0 0.0]),
    (6, [0.0 -12.0 0.0]),
    ],
  :limit => 250,
  :tol => 0.0001,
  :incs => 8,
  :dload => [0.5; 0.3; 0.2; 0.2; 0.1; 0.05; 0.02; 0.01]
)

data |> display
println()

@time m = p45(data)
println()


#=Copyright (c) 2017 Gabriel Goh

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.=#

# α = 0.076
# β = 0.513

α = 0.2
β = 0.99
x = 2
y = 0
X = zeros(100)
for i = 1:100
    y  = β*y - x
    x  = x - α*y
    X[i] = x
end
figure(); plot(X,"k"); plot(X,"k."); xlabel(L"k"); ylabel(L"f(x)")


isdefined(Base, :__precompile__) && __precompile__()

module WoodburyMatrices

using Compat

import Base: *, \, A_ldiv_B!, convert, copy, det, full, show, similar, size

export Woodbury, SymWoodbury, liftFactor

#### Woodbury matrices ####
"""
`W = Woodbury(A, U, C, V)` creates a matrix `W` identical to `A + U*C*V` whose inverse will be calculated using
the Woodbury matrix identity.
"""
type Woodbury{T,AType,UType,VType,CType,CpType} <: AbstractMatrix{T}
    A::AType
    U::UType
    C::CType
    Cp::CpType
    V::VType
    tmpN1::Vector{T}
    tmpN2::Vector{T}
    tmpk1::Vector{T}
    tmpk2::Vector{T}
end

function Woodbury{T}(A, U::AbstractMatrix{T}, C, V::AbstractMatrix{T})
    N = size(A, 1)
    k = size(U, 2)
    if size(A, 2) != N || size(U, 1) != N || size(V, 1) != k || size(V, 2) != N
        throw(DimensionMismatch("Sizes of U ($(size(U))) and/or V ($(size(V))) are inconsistent with A ($(size(A)))"))
    end
    if k > 1
        if size(C, 1) != k || size(C, 2) != k
            throw(DimensionMismatch("C should be $(k)x$(k)"))
        end
    end
    Cp = inv(inv(C) + V*(A\U))
    # temporary space for allocation-free solver
    tmpN1 = Array(T, N)
    tmpN2 = Array(T, N)
    tmpk1 = Array(T, k)
    tmpk2 = Array(T, k)
    # don't copy A, it could be huge
    Woodbury{T,typeof(A),typeof(U),typeof(V),typeof(C),typeof(Cp)}(A, copy(U), copy(C), Cp, copy(V), tmpN1, tmpN2, tmpk1, tmpk2)
end

Woodbury{T}(A, U::Vector{T}, C, V::Matrix{T}) = Woodbury(A, reshape(U, length(U), 1), C, V)

size(W::Woodbury) = size(W.A)

function show(io::IO, W::Woodbury)
    println(io, summary(W), ":")
    print(io, "A:\n", W.A)
    print(io, "\nU:\n")
    Base.print_matrix(io, W.U)
    if isa(W.C, Matrix)
        print(io, "\nC:\n")
        Base.print_matrix(io, W.C)
    else
        print(io, "\nC: ", W.C)
    end
    print(io, "\nV:\n")
    Base.print_matrix(io, W.V)
end

full{T}(W::Woodbury{T}) = convert(Matrix{T}, W)
convert{T}(::Type{Matrix{T}}, W::Woodbury{T}) = full(W.A) + W.U*W.C*W.V

copy(W::Woodbury) = Woodbury(W.A, W.U, W.C, W.V)

## Woodbury matrix routines ##

*(W::Woodbury, B::StridedVecOrMat)=W.A*B + W.U*(W.C*(W.V*B))

function \(W::Woodbury, R::StridedVecOrMat)
    AinvR = W.A\R
    AinvR - W.A\(W.U*(W.Cp*(W.V*AinvR)))
end

det(W::Woodbury)=det(W.A)*det(W.C)/det(W.Cp)

function A_ldiv_B!(W::Woodbury, B::AbstractVector)
    length(B) == size(W, 1) || throw(DimensionMismatch("Vector length $(length(B)) must match matrix size $(size(W,1))"))
    copy!(W.tmpN1, B)
    Alu = lufact(W.A) # Note. This makes an allocation (unless A::LU). Alternative is to destroy W.A.
    A_ldiv_B!(Alu, W.tmpN1)
    A_mul_B!(W.tmpk1, W.V, W.tmpN1)
    A_mul_B!(W.tmpk2, W.Cp, W.tmpk1)
    A_mul_B!(W.tmpN2, W.U, W.tmpk2)
    A_ldiv_B!(Alu, W.tmpN2)
    for i = 1:length(W.tmpN2)
        @inbounds B[i] = W.tmpN1[i] - W.tmpN2[i]
    end
    B
end

"""
Given a type `T`, an integer `n` and `m` tuples of the form (i, j, v), build
sparse matrices `rows`, `vals`, `cols` such that the product
`out = rows * vals * cols` is equivalent to:

```julia
out = zeros(T, n, n)

for (i, j, v) in args
    out[i, j] = v
end
```

The first two components (`i` and `j`) of each tuple should be integers
whereas the third component should be of type `T`

Example:

```
julia> r, v, c = WoodburyMatrices.sparse_factors(Float64, 3,
                                                 (1, 1, 2.0),
                                                 (2, 2, 3.0),
                                                 (3, 3, 4.0));

julia> r*c*v - Diagonal([2.0, 3.0, 4.0])
3x3 sparse matrix with 0 Float64 entries:
```

"""
function sparse_factors{T}(::Type{T}, n::Int, args::@compat(Tuple{Int, Int, Any})...)
    m = length(args)
    rows = spzeros(T, n, m)
    cols = spzeros(T, m, n)
    vals = zeros(T, m, m)

    ix = 1
    for (i, (row, col, val)) in enumerate(args)
        rows[row, ix] = 1
        cols[ix, col] = 1
        vals[ix, ix] = val
        ix += 1
    end

    rows, vals, cols
end

include("SymWoodburyMatrices.jl")

end


using NumericalRationalFunctions
using LinearAlgebra
using Test

@testset "AAA" begin
    Z = range(-1, stop=1, length=20)
    F = exp.(Z)
    r, pol, res, zer, z, f, w, errvec = aaa(F, Z)
    @test length(z) == 7
    @test length(pol) == length(res) == length(zer) == 6
    @test r(Z) ≈ F atol=1e-14
    x = 0.12345
    @test r(x) ≈ exp(x) atol=1e-14
    @test norm(r(zer), Inf) < 1e-9
    @test errvec[end] < 1e-13
end

@testset "Laplace" begin
    u, err, f, Z, Zplot, A, inpoly = laplace(:iso, tol=1e-10, boundary=z->log(abs(z)))
    @test err < 1e-10
    u, err, f, Z, Zplot, A, inpoly = laplace(:circleL)
    @test err < 1e-6
    u, err, f, Z, Zplot, A, inpoly = laplace(complex.([[1, 0.5], 1+1im, -1+1im, -1]), boundary=[1, 0, 0, 0])
    @test err < 1e-6

    #=
    sx = sy = range(-1.5, stop=1.5, length=200)
    p = contour(sx, sy, (x, y)->begin
                    z = x + im*y
                    inpoly(z) == 1 ? u(z) : NaN
                end, levels=30, c=:viridis, aspect_ratio=1)
    nw = length(Zplot)
    for k in 1:nw-1
        kn = k+1
        plot!(p, [Zplot[k], Zplot[kn]], color=:black, lab=false, linewidth=2)
    end
    p
    =#
end


"""
timederiv(x) Take time derivative of symbolic expression x.

timederiv(x,subdiffout,subin) also substitutes the derivative of subdiffout with
the symbols in subin.  IE the symbols in subin are the time derivatives of
the symbols in subdiffout.

timederiv(x,body) uses the properties of body to determine the relationship between
variables (that the velocity states are the time derivative of position states, etc.).

timederiv(x,model) same as timederiv(x,model) but with the info of the whole model
"""
function timederiv(x::Array{SymPy.Sym})
  dxdt = diff(x,t)
end
function timederiv(x::SymPy.Sym)
  dxdt = diff(x,t)
end

"""
Take time derivative of x, where subdiffout specifies what symbols in x depend on time.
Also replace the time derivatives of subdiffout with the symbols in subin.

For example if you have a variable x that depends on time, this looks for x,
replaces it with x(t), take the derivative which produces expressions like:
Derivative(x(t),t), which then are replaced with the symbol defined in subin,
for example vx.

julia> @syms x vx
julia> pos = x^2
julia> vel = timederiv(pos,x,vx)
2*x*vx
"""
function timederiv(x,subdiffout::Array{SymPy.Sym,1},subin::Array{SymPy.Sym,1})
  if length(subdiffout) != length(subin)
    error("subdiffout and subin must have same dimensions")
  end

  # Make subdiffout depend on time
  subdiffoutT = AddTimeDependence(subdiffout)
  x = subs(x,subdiffout,subdiffoutT)

  # Take derivative
  dxdt = timederiv(x)

  # Substitute deriviative of subdiffout for diffin
  diffout = diff(subdiffoutT,t)
  for i in eachindex(diffout)
    dxdt = subs(dxdt,diffout[i],subin[i])
  end

  # remove time dependence
  return dxdt = subs(dxdt,subdiffoutT,subdiffout)
end

# function timederiv(x::SymPy.Sym,v...)
#   return dxdt = timederiv([x],v...)
# end

function timederiv(x::Array{SymPy.Sym},body::Body)
  # the derivative of the position and velocity states are the velocity and accelerations
  # respectively
  return timederiv(x,body.connection.q,body.connection.u)
end


import StringUtils

function char_to_digit(c::Char)
    if isdigit(c)
        return Int(c) - Int('0')
    else
        throw(ArgumentError(@sprintf("%c is not a digit", c)))
    end
end

function main()
    if length(ARGS) > 1
        filename = ARGS[1]
        n = parse(Int64, ARGS[2])
        a = readlines(filename)
        s = StringUtils.conc(a)
        max_product = -1
        for i in 1:(length(s)-n+1)
            product = reduce(*,
                             map(char_to_digit,
                                 collect(s[i:(i+n-1)])))
            if product > max_product
                max_product = product
            end
        end
        @printf("The maximum %d length product has value %d\n", n, max_product)
    else
        error("this program needs a single numeric argument")
    end
end

main()


@testset "Structured Matrices" begin
    @testset "/ and \\ on Square Matrixes" begin
        @testset "//, $T on the RHS" for T in (Diagonal, UpperTriangular, LowerTriangular)
            RHS = T(randn(T == Diagonal ? 10 : (10, 10)))
            test_rrule(/, randn(5, 10), RHS)
        end

        @testset "\\ $T on LHS" for T in (Diagonal, UpperTriangular, LowerTriangular)
            LHS = T(randn(T == Diagonal ? 10 : (10, 10)))
            test_rrule(\, LHS, randn(10))
            test_rrule(\, LHS, randn(10, 10))
        end
    end

    @testset "Diagonal" begin
        N = 3
        test_rrule(Diagonal, randn(N); output_tangent=randn(N, N))
        D = Diagonal(randn(N))
        test_rrule(Diagonal, randn(N); output_tangent=D)
        # Concrete type instead of UnionAll
        test_rrule(typeof(D), randn(N); output_tangent=D)

        # TODO: replace this with a `rrule_test` once we have that working
        # see https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/24
        res, pb = rrule(Diagonal, [1, 4])
        @test pb(10*res) == (NO_FIELDS, [10, 40])
        comp = Composite{typeof(res)}(; diag=10*res.diag)  # this is the structure of Diagonal
        @test pb(comp) == (NO_FIELDS, [10, 40])
    end
    @testset "dot(x, ::Diagonal, y)" begin
        N = 4
        test_rrule(dot, randn(ComplexF64, N), Diagonal(randn(ComplexF64, N)), randn(ComplexF64, N))
    end
    @testset "::Diagonal * ::AbstractVector" begin
        N = 3
        test_rrule(*, Diagonal(randn(N)), randn(N))
    end
    @testset "diag" begin
        N = 7
        test_rrule(diag, randn(N, N))
        test_rrule(diag, Diagonal(randn(N)))
        test_rrule(diag, randn(N, N) ⊢ Diagonal(randn(N)))
        test_rrule(diag, Diagonal(randn(N)) ⊢ Diagonal(randn(N)))
        VERSION ≥ v"1.3" && @testset "k=$k" for k in (-1, 0, 2)
            test_rrule(diag, randn(N, N), k ⊢ nothing)
        end
    end
    @testset "diagm" begin
        @testset "without size" begin
            M, N = 7, 9
            s = (8, 8)
            a, ā = randn(M), randn(M)
            b, b̄ = randn(M), randn(M)
            c, c̄ = randn(M - 1), randn(M - 1)
            ȳ = randn(s)
            ps = (0 => a, 1 => b, 0 => c)
            y, back = rrule(diagm, ps...)
            @test y == diagm(ps...)
            ∂self, ∂pa, ∂pb, ∂pc = back(ȳ)
            @test ∂self === NO_FIELDS
            ∂a_fd, ∂b_fd, ∂c_fd = j′vp(_fdm, (a, b, c) -> diagm(0 => a, 1 => b, 0 => c), ȳ, a, b, c)
            for (p, ∂px, ∂x_fd) in zip(ps, (∂pa, ∂pb, ∂pc), (∂a_fd, ∂b_fd, ∂c_fd))
                ∂px = unthunk(∂px)
                @test ∂px isa Composite{typeof(p)}
                @test ∂px.first isa AbstractZero
                @test ∂px.second ≈ ∂x_fd
            end
        end
        VERSION ≥ v"1.3" && @testset "with size" begin
            M, N = 7, 9
            a, ā = randn(M), randn(M)
            b, b̄ = randn(M), randn(M)
            c, c̄ = randn(M - 1), randn(M - 1)
            ȳ = randn(M, N)
            ps = (0 => a, 1 => b, 0 => c)
            y, back = rrule(diagm, M, N, ps...)
            @test y == diagm(M, N, ps...)
            ∂self, ∂M, ∂N, ∂pa, ∂pb, ∂pc = back(ȳ)
            @test ∂self === NO_FIELDS
            @test ∂M === DoesNotExist()
            @test ∂N === DoesNotExist()
            ∂a_fd, ∂b_fd, ∂c_fd = j′vp(_fdm, (a, b, c) -> diagm(M, N, 0 => a, 1 => b, 0 => c), ȳ, a, b, c)
            for (p, ∂px, ∂x_fd) in zip(ps, (∂pa, ∂pb, ∂pc), (∂a_fd, ∂b_fd, ∂c_fd))
                ∂px = unthunk(∂px)
                @test ∂px isa Composite{typeof(p)}
                @test ∂px.first isa AbstractZero
                @test ∂px.second ≈ ∂x_fd
            end
        end
    end
    @testset "$f, $T" for
        f in (Adjoint, adjoint, Transpose, transpose),
        T in (Float64, ComplexF64)

        n = 5
        m = 3
        @testset "$f(::Matrix{$T})" begin
            A = randn(T, n, m)
            Y = f(A)
            Ȳ_mat = randn(T, m, n)
            Ȳ_composite = Composite{typeof(Y)}(parent=collect(f(Ȳ_mat)))

            test_rrule(f, A; output_tangent=Ȳ_mat)

            _, pb = rrule(f, A)
            @test pb(Ȳ_mat) == pb(Ȳ_composite)
        end

        @testset "$f(::Vector{$T})" begin
            a = randn(T, n)
            y = f(a)
            ȳ_mat = randn(T, 1, n)
            ȳ_composite = Composite{typeof(y)}(parent=collect(f(ȳ_mat)))

            test_rrule(f, a; output_tangent=ȳ_mat)

            _, pb = rrule(f, a)
            @test pb(ȳ_mat) == pb(ȳ_composite)
        end

        @testset "$f(::Adjoint{$T, Vector{$T})" begin
            a = randn(T, n)'
            ā = randn(T, n)'
            test_rrule(f, a ⊢ ā; output_tangent=randn(T, n))
        end

        @testset "$f(::Transpose{$T, Vector{$T})" begin
            a = transpose(randn(T, n))
            ā = transpose(randn(T, n))
            test_rrule(f, a ⊢ ā; output_tangent=randn(T, n))
        end
    end
    @testset "$T" for T in (UpperTriangular, LowerTriangular)
        n = 5
        test_rrule(T, randn(n, n); output_tangent=T(randn(n, n)))
    end
    @testset "$Op" for Op in (triu, tril)
        n = 7
        test_rrule(Op, randn(n, n))
        @testset "k=$k" for k in -2:2
            test_rrule(Op, randn(n, n), k ⊢ nothing)
        end
    end

    @testset "det and logdet $S" for S in (Diagonal, UpperTriangular, LowerTriangular)
        @testset "$op" for op in (det, logdet)
            @testset "$T" for T in (Float64, ComplexF64)
                n = 5
                # rand (not randn) so det will be postive, so logdet will be defined
                X = S(3*rand(T, (n, n)) .+ 1)
                X̄_acc = Diagonal(rand(T, (n, n)))  # sensitivity is always a diagonal for these types
                test_rrule(op, X ⊢ X̄_acc)
            end
            @testset "return type" begin
                X = S(3*rand(6, 6) .+ 1)
                _, op_pullback = rrule(op, X)
                X̄ = op_pullback(2.7)[2]
                @test X̄ isa Diagonal
            end
        end
    end
end


module ClimateModels

using Zarr, AWSCore, DataFrames, CSV, CFTime, Dates, Statistics

export cmip

"""
    cmip(institution_id,source_id,variable_id)

Access CMIP6 climate model archive (https://bit.ly/2WiWmoh) via
`AWSCore.jl` and `Zarr.jl` and compute (1) time mean global map and
(2) time evolving global mean.

This was partly inspired by @rabernat 's https://bit.ly/2VRMgvl notebook

```
using ClimateModels
(mm,gm,meta)=cmip()
nm=meta["long_name"]*" in "*meta["units"]

using Plots
heatmap(mm["lon"], mm["lat"], transpose(mm["m"]),
        title=nm*" (time mean)")
plot(gm["t"][1:12:end],gm["y"][1:12:end],xlabel="time",ylabel=nm,
     title=meta["institution_id"]*" (global mean, month by month)")
display.([plot!(gm["t"][i:12:end],gm["y"][i:12:end], leg = false) for i in 2:12])

```
"""
function cmip(institution_id="IPSL",source_id="IPSL-CM6A-LR",
    variable_id="tas")

    #choose model and variable
    S=[institution_id, source_id, variable_id]

    #initiate cloud connection
    ⅁ = AWSCore.aws_config(creds=nothing, region="",
    service_host="googleapis.com", service_name="storage")

    #get list of contents for cloud storage unit
    β = S3Store("cmip6","", aws=⅁, listversion=1)
    ξ = CSV.read(IOBuffer(β["cmip6-zarr-consolidated-stores.csv"]))

    # get model grid cell areas
    ii=findall( (ξ[!,:source_id].==S[2]).&(ξ[!,:variable_id].=="areacella") )
    μ=ξ[ii,:]
    i1=findfirst("cmip6",μ.zstore[end])[end]+2
    P = μ.zstore[end][i1:end]
    ζ = zopen(S3Store("cmip6", P, aws=⅁, listversion=1))
    Å = ζ["areacella"][:, :];

    # get model solution ensemble list
    i=findall( (ξ[!,:activity_id].=="CMIP").&(ξ[!,:table_id].=="Amon").&
    (ξ[!,:variable_id].==S[3]).&(ξ[!,:experiment_id].=="historical").&
    (ξ[!,:institution_id].==S[1]) )
    μ=ξ[i,:]

    # access one model ensemble member
    i1=findfirst("cmip6",μ.zstore[end])[end]+2
    P = μ.zstore[end][i1:end]
    ζ = zopen(S3Store("cmip6", P, aws=⅁, listversion=1))

    meta=Dict("institution_id" => institution_id,"source_id" => source_id,
        "variable_id" => variable_id, "units" => ζ[S[3]].attrs["units"],
        "long_name" => ζ[S[3]].attrs["long_name"])

    # time mean global map
    m = convert(Array{Union{Missing, Float32},3},ζ[S[3]][:,:,:])
    m = dropdims(mean(m,dims=3),dims=3)

    mm=Dict("lon" => ζ["lon"], "lat" => ζ["lat"], "m" => m)

    # time evolving global mean
    t = ζ["time"]
    t = timedecode(t[:], t.attrs["units"], t.attrs["calendar"])

    y = ζ[S[3]][:,:,:]
    y=[sum(y[:, :, i].*Å) for i in 1:length(t)]./sum(Å)

    gm=Dict("t" => t, "y" => y)

    return mm,gm,meta
end

end # module


# ======= ALL ROTATIONS ALLOWED START =======
#lx + ly + lz = 1
@constraints(model, begin
    [i=1:length(boxes)], (lpx_length_of_box_parallel_to_x[i] +
    lpy_length_of_box_parallel_to_y[i] +
    lpz_length_of_box_parallel_to_z[i] == 1)
end)

#wx + wy + wz = 1
@constraints(model, begin
    [i=1:length(boxes)], (wpx_width_of_box_parallel_to_x[i] +
    wpy_width_of_box_parallel_to_y[i] +
    wpz_width_of_box_parallel_to_z[i] == 1)
end)

#hx + hy + hz = 1
@constraints(model, begin
    [i=1:length(boxes)], (hpx_height_of_box_parallel_to_x[i] +
    hpy_height_of_box_parallel_to_y[i] +
    hpz_height_of_box_parallel_to_z[i] == 1)
end)

#lx + wx + hx = 1
@constraints(model, begin
    [i=1:length(boxes)], (lpx_length_of_box_parallel_to_x[i] +
    wpx_width_of_box_parallel_to_x[i] +
    hpx_height_of_box_parallel_to_x[i] == 1)
end)

#ly + wy + hy = 1
@constraints(model, begin
    [i=1:length(boxes)], (lpy_length_of_box_parallel_to_y[i] +
    wpy_width_of_box_parallel_to_y[i] +
    hpy_height_of_box_parallel_to_y[i] == 1)
end)

#lz + wz + hz = 1
@constraints(model, begin
    [i=1:length(boxes)], (lpz_length_of_box_parallel_to_z[i] +
    wpz_width_of_box_parallel_to_z[i] +
    hpz_height_of_box_parallel_to_z[i] == 1)
end)
# ======= ALL ROTATIONS ALLOWED END =======


# Code for Video 2.3:
# The "iris" dataset contains 150 measurements on 3 species
# of Iris flowers, namely I. setosa, I. versicolor and I. virginica.
# The columns 1 to 4 respectively contain the:
# sepal length - sepal width -petal length - petal width.
# The species is specified in the last column.
# The dataset can be found in the file iris.txt

module IrisTools

### exports:
export Iris
export species
export read_file, counting_species, length_matrix, calculations_setosa, to_string

### module variables:
###################################################
fname = "iris.txt"
species = ["I. setosa", "I. versicolor", "I. virginica"]

irism = [] # will contain the measurements of the dataset, each array element is a line containing 4 measurements
iris = [] # temporary array  to contain the names of iris specimens
lengths = [] # matrix with the length measurements
species_freq = Dict{ASCIIString, Int64}()

### types:
###################################################
type Iris
    sepal_length::Float64
    sepal_width::Float64
    petal_length::Float64
    petal_width::Float64
    species::ASCIIString
end

type Setosa
    # some fields
end

type Versicolor
    # some fields
end

type Virginica
    # some fields
end

IrisU = Union{Setosa, Versicolor, Virginica}

### functions:
###################################################
"executed when the module is loaded"
function __init__()
    read_file()
end

"reads in the file iris.txt with iris measurements"
function read_file()
    open(fname) do file
        global irism = readlines(file)
    end
    # println(irism)
end

"gives a count of measurements for each species"
function counting_species()
    # gather names of species in iris array:
    for line in irism
        # println(line)
        # the species name starts from column 17:
        species = line[17:end]
        # removes the trailing \r\n:
        species = strip(species)
        push!(iris, species)
    end
    # count species in dictionary species-freq:
    for species in iris
        haskey(species_freq, species) ?
            species_freq[species] += 1 :
            species_freq[species] = 1
    end
    # sort the the keys and print out the frequencies:
    println("Species : frequency \n")
    species = sort!(collect(keys(species_freq)))
    for sp in species
        println("$sp : $(species_freq[sp])")
    end
end

" extracts a matrix with the length measurements"
function length_matrix()
    # extract the measurements as numbers in a matrix lengths:
    lengths_temp = Array(Float64, 1, 4)
    for line in irism
        # The columns 1 to 16 respectively contain the lengths of 1 measurement:
        mlengths = line[1:16] # measurement lengths
        arr_temp = split(mlengths) # split on spaces to an array of strings
        # convert to an array of Float64 values (a column vector 4 x 1)
        arr_mlength = [parse(Float64, str) for str in arr_temp]
        # transpose arr_mlength to obtain a 1 x 4 array
        # and append it vertically to lengths
        lengths_temp = vcat(lengths_temp, arr_mlength')
        # we must indicate that we mean the variable lengths
        # defined in global scope by prepending it with global
    end
    # cut off the first row which contained arbitrary values
    global lengths = lengths_temp[2:end, :]
end

" performs some calculations on I. setosa measurements"
function calculations_setosa()
    # extract I. setosa measurements:
    lengths_setosa = lengths[1:50, :]
    # do some calculations:
    println("sum of each column")
    println(sum(lengths_setosa, 1))
    println("mean of the features:")
    println(mean(lengths_setosa, 1))
    println("smallest values of the features:")
    println(sort(lengths_setosa, 1)[1, :])
    println("largest values of the features:")
    println(sort(lengths_setosa, 1)[end, :])
end

function show(i::Iris)
    println("This $(i.species) has measurements: ")
    println("Sepal length: $(i.sepal_length)")
    println("Sepal width: $(i.sepal_width)")
    println("Petal length: $(i.petal_length)")
    println("Petal width: $(i.petal_width)")
end

function to_string(i::Iris)
    "This $(i.species) has measurements: \n" *
    "\tSepal length: $(i.sepal_length)\n" *
    "\tSepal width: $(i.sepal_width)\n" *
    "\tPetal length: $(i.petal_length)\n" *
    "\tPetal width: $(i.petal_width)\n"
end

end  # module end


# abstract type AbstractBoxPartition{B <: Box} end

"""
A generalized box with
`center`:   vector where the box's center is located
`radius`:   vector of radii, length of the box in each dimension

"""
struct Box{N,T <: AbstractFloat}
    center::SVector{N,T}
    radius::SVector{N,T}

    function Box(center, radius)
        N = length(center)

        if length(radius) != N
            throw(DimensionMismatch("Center vector and radius vector must have same length ($N)"))
        end

        if any(x -> x <= 0, radius)
            throw(DomainError(radius, "radius must be positive in every component"))
        end

        T = promote_type(eltype(center), eltype(radius))
        if !(T <: AbstractFloat)
            T = Float64
        end

        return new{N,T}(SVector{N,T}(center), SVector{N,T}(radius))
    end
end

Base.in(point, box::Box) = all(box.center - box.radius  .<=  point  .<  box.center + box.radius)

volume(box::Box) = prod(2 * box.radius)

function Base.show(io::IO, box::Box)
    print(io, "Box: center = $(box.center), radii = $(box.radius)")
end


function xpow(x)
return [x x^2 x^3 x^4]
end
function xpow_loop(n)
s = 0
for i = 1:n
s = s + xpow(i)[2]
end
return s
end


#Parameter extraction of a parametric sinusoidal "simulation"
#-------------------------------------------------------------------------------
module CMDimData_SampleUsage

using CMDimData
using CMDimData.MDDatasets
using CMDimData.EasyPlot
CMDimData.@includepkg EasyPlotInspect

#==Constants
===============================================================================#
WIDTH_LEGEND = 250
WIDTH_DATA = 600 #Approx
WIDTH_PLOT = WIDTH_DATA+WIDTH_LEGEND
HEIGHT_PLOT = round(Int, WIDTH_DATA/3) #Total is roughly square
HEIGHT_CANVAS = round(Int, 3*HEIGHT_PLOT*1.2) #Add extra for overhead (titles, etc)

#==Read in results of parametric "simulation"
===============================================================================#
signal = evalfile("parametric_sin_data.jl")

#==Parameter extraction
===============================================================================#
sdim1 = 4e3 #Sample dimension 1 @ 𝑓=4kHz
sdim2 = .002 #Sample dimension 2 @ A=.002

println()
@info("Information on `signal`:")
@show ndims(signal), paramlist(signal)

#Convenience: x-value resulting from paramter reduction operations:
x_paramreduce = reverse(paramlist(signal))

#Generate DataRS with value of amplitude (A) in each parameter combination:
ampvalue = parameter(signal, "A")

#Normalize signal amplitudes for ALL signals SIMULTANEOUSLY:
signal_norm = signal / ampvalue

#Compute continuous-time signal rate for ALL signals SIMULTANEOUSLY:
#   signal(𝜙, A, 𝑓, t) -> rate(𝜙, A, 𝑓):
rate = deriv(signal)

println()
@info("Reduction 0: Compute `fallx`: First fall-crossing point for ALL signals SIMULTANEOUSLY:")
#   signal(𝜙, A, 𝑓, t) -> fallx(𝜙, A, 𝑓):
fallx = xcross1(signal, xstart=0, allow=CrossType(:fall))
@show ndims(fallx), paramlist(fallx)
@show xred1 = x_paramreduce[1] #x-value for 1st fallx reduction

println()
@info("Reduction 1: Evaluate `fallx` @ $xred1=$sdim1 (=>`fallx_red1`):")
#   fallx(𝜙, A, 𝑓) -> fallx_red1(𝜙, A):
fallx_red1 = value(fallx, x=sdim1)
@show ndims(fallx_red1), paramlist(fallx_red1)
@show xred2 = x_paramreduce[2] #x-value for 2nd fallx reduction

println()
@info("Reduction 2: Evaluate `fallx_red1` @ $xred2=$sdim2 (=>`fallx_red2`):")
#   fallx_red1(𝜙, A) -> fallx_red2(𝜙):
fallx_red2 = value(fallx_red1, x=sdim2)
@show ndims(fallx_red2), paramlist(fallx_red2)
@show xred3 = x_paramreduce[3] #x-value for 3rd fallx reduction

data = (
	signal = signal, signal_norm = signal_norm,
	rate = rate, fallx = fallx,
	fallx_red1 = fallx_red1, fallx_red2 = fallx_red2,
	xred1 = xred1, xred2 = xred2,	xred3 = xred3,
	sdim1 = sdim1,	sdim2 = sdim2,
)

#==Generate plots
===============================================================================#
PB = EasyPlot.load_plotbuilders(@__DIR__,
	initial = "bld_parametric_sin_initial.jl",
	explore = "bld_parametric_sin_explore.jl",
)

function adjust_legend(gtkplot)
	mplot = gtkplot.src
	for sp in mplot.subplots
		sp.layout[:halloc_legend] = WIDTH_LEGEND
	end
end
plotdisplay = EasyPlot.GUIDisplay(:InspectDR, postproc=adjust_legend)
	opt = EasyPlot.ShowOptions(dim=set(w=WIDTH_PLOT, h=HEIGHT_CANVAS))

plotset1 = EasyPlot.build(PB[:initial], data)
	plotgui1 = display(plotdisplay, plotset1)
	EasyPlot._write(:png, "parametric_sin_initial.png", opt, plotgui1)
plotset2 = EasyPlot.build(PB[:explore], data)
	plotgui2 = display(plotdisplay, plotset2)
	EasyPlot._write(:png, "parametric_sin_explore.png", opt, plotgui2)

end #module
:SampleCode_Executed


## PART 2 -- Note: you have to look at p1.jl for this to make sense

# Find integrated autocorrelation times

include("../q1/autocorrelation.jl");

println("\nPart 2");
println("\n\tFirst, plot autocorrelation to estimate τint");

# Fuse md trials
md = [md_1069 md_1304]

# Number of quantities to look at
N = 6

# Same (almost) as q1p3:
lengthmd = size(md)[1];

# Create x and y arrays for plotting
xcorr = [0:1:100];
ycorr = zeros( length( xcorr ), N );
for a in [1:N]
  # Calculate autocorrelations
  ycorr[:,a] = autocorrelation(lengthmd, md, a, xcorr);
  # Normalize by dividing through by Ĉ[1,a]
  ycorr[:,a] /= ycorr[1,a];
end

# Make plots for each alpha
println("\n\tMaking autocorrelation plots and placing in 'mdautocorr' vector");
using Gadfly
Gadfly.set_default_plot_size(20Gadfly.cm, 12Gadfly.cm)
mdautocorr = Plot[];
quantity = ["Temperature at T = 1.069", "Potential Energy at T = 1.069",
            "Virial at T = 1.069", "Temperature at T = 1.304",
            "Potential Energy at T = 1.304", "Virial at T = 1.304"];

for i in [1:N]
  # Plot
  push!(mdautocorr, Gadfly.plot ( x = 10xcorr, y = ycorr[:,i], Geom.line, Geom.point,
      Guide.xlabel("Step Separation"), Guide.ylabel("Autocorrelation"),
      Guide.title(join(["Autocorrelation for ",quantity[i]]," ")),
      Guide.xticks(ticks=[0:100:1000])))
end

# Find integrated autocorrelation times τint (similar to q1p4)

# Pick ncut for temperature, potential energy, and the virial
#   based on observed graph of autocorrelation
ncutt_1069 = 250;
ncutp_1069 = 250;
ncutv_1069 = 250;
ncutt_1304 = 100;
ncutp_1304 = 100;
ncutv_1304 = 100;

τtemp_1069 = τ(lengthmd, temps_1069, ncutt_1069);
τpe_1069 = τ(lengthmd, pe_1069, ncutp_1069);
τvir_1069 = τ(lengthmd, vir_1069, ncutv_1069);
τtemp_1304 = τ(lengthmd, temps_1304, ncutt_1304);
τpe_1304 = τ(lengthmd, pe_1304, ncutp_1304);
τvir_1304 = τ(lengthmd, vir_1304, ncutv_1304);

τmd = [τtemp_1069, τpe_1069, τvir_1069, τtemp_1304, τpe_1304, τvir_1304]

println("\n\tAutocorrelation times:\n");
for i in [1:N]
  println("\t\tτ̂v",quantity[i]":\t",τmd[i]);
end


begin
	f500kHz = 500.0kHz
	w100mus = 100.0μs
	wmus100mus=w100mus/μs # window in μs
	tlaser500kHz = uconvert(μs, 1.0/f500kHz)
	tlus2 = tlaser500kHz/μs
end

md" - Parameters:
 - laser frequency= $f500kHz
 - DAQ window= $w100mus
 - Laser pulses every $tlaser500kHz
"

md"# BOLD062\_A2\_500kHz\_730nm\_100us"

md"## Read waveform"

fn8n5k1m = 1  # test file

begin
	fls8n5k1m = glob("../labdata/BOLD062_A2_500kHz_730nm_100us/C1*.csv")
	n8n5k1m = length(fls8n5k1m)

	ffl8n5k1m = 200.0e+6
	wvfm8n5k1m = read_waveform(fls8n5k1m,fn8n5k1m; wmus=wmus100mus, wsym=true)
	first(wvfm8n5k1m,3)
end

md"
- Files in directory = $n8n5k1m
- sampling rate= $(round(sr8n5k1m, sigdigits=2)) GHz
- sampling period = $(round(sp8n5k1m, sigdigits=2) )ns
"

begin
	sr8n5k1m = sampling_rate(wvfm8n5k1m,wmus100mus)/GHz
	sp8n5k1m = sampling_period(wvfm8n5k1m,wmus100mus)/ns
end

md"
Unfiltered wvfm:
- mean = $(round(mean8n5k1m, sigdigits=4))
- std = $(round(std8n5k1m, sigdigits=4))
"
begin
	mean8n5k1m, std8n5k1m, thrp8n5k1m, thrn8n5k1m = thrx(wvfm8n5k1m; nsigma=3.0, waveform="raw");
	pfw8n5k1m = plot_filtered_waveform(wvfm8n5k1m; window=false, waveform="raw", thrp=thrp8n5k1m, thrn=thrn8n5k1m, scatter=false)
	pfft8n5k1m =plot_fft(wvfm8n5k1m,wmus100mus)
	plot(pfw8n5k1m, pfft8n5k1m)
end

md"### Filter signal
- Filter frequency (in Hz) = $ffl8n5k1m

Filtered wvfm:
- mean = $fmean8n5k1m
- std = $fstd8n5k1m
"
begin
    filter_signal_lp!(wvfm8n5k1m, wmus100mus; filtertype="Butterworth", flhz=ffl8n5k1m, n=4)
    first(wvfm8n5k1m,3)
end

begin
    fmean8n5k1m, fstd8n5k1m, fthrp8n5k1m, fthr8n5k1m = thrx(wvfm8n5k1m; nsigma=2.0, waveform="filtered")
    plot_filtered_waveform(wvfm8n5k1m; window=false, thrp=fthrp8n5k1m, thrn=fthr8n5k1m)
end

md"### Select peaks"

sprom8n5k1m = 2.0  #prominence cut

md"
- Peak Prominence cut = $sprom8n5k1m
"

begin
	r8n5k1m = select_filtered_peaks(wvfm8n5k1m, fthrp8n5k1m; promsel = sprom8n5k1m, wsel=0.0);
	if r8n5k1m != Nothing
		wvflt8n5k1m, spk8n5k1m = r8n5k1m
		plot_filtered_peaks(wvflt8n5k1m, spk8n5k1m)
	end
end

md"### Optimizing prominence cut"

begin
	P8n5k1m, W8n5k1m = proms(fls8n5k1m, 1, n8n5k1m;
	                 wmus=wmus100mus, wsym=true, flhz=ffl8n5k1m, promsel=0.5)
	hP8n5k1m = lfi.LaserLab.hist1d(P8n5k1m, 20, 0.0, 10.0)
	hW8n5k1m = lfi.LaserLab.hist1d(W8n5k1m, 20, 5.0, 15.0)
end

lfi.LaserLab.hist1d(hP8n5k1m, "Prominence (promsel=0.5)")

lfi.LaserLab.hist1d(hW8n5k1m, "Width")

md"### Selection"

begin
    spks8n5k1m = select_peaks(fls8n5k1m, 1, n8n5k1m; wmus=wmus100mus, wsym=true, flhz=ffl8n5k1m, promsel=sprom8n5k1m)
    t8n5k1m = reduce(vcat,[subtract_laser(spks8n5k1m[i].xs, tlus2) for i in 1:length(spks8n5k1m)])
    hp8n5k1m = lfi.LaserLab.hist1d(t8n5k1m, 20, 0.0, 2.0)
end
lfi.LaserLab.hist1d(hp8n5k1m, "times")

md"### Fits"

begin
    cofe8n5k1m, stder8n5k1m, tft8n5k1m = fit_peaks(hp8n5k1m; pa0=[100.0, 0.5], i0=1)
    pp8n5k1m = plot_fit(hp8n5k1m, cofe8n5k1m, tft8n5k1m; savefig=true,
	              fn="b62A2_500kHz_730nm_100mus_fit.png")
	plot(pp8n5k1m)
end


using Test
using LinearAlgebra
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using HDGElasticity

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

sbasis = TensorProductBasis(1,1)
squad = tensor_product_quadrature(1,2)
facequads = repeat([squad],4)

testmatrix = [2/3 0.0 1/3 0.0
              0.0 2/3 0.0 1/3
              1/3 0.0 2/3 0.0
              0.0 1/3 0.0 2/3]

HH = HDGElasticity.HHop(sbasis,squad,1.)
@test allapprox(HH,testmatrix)

cellmap = HDGElasticity.CellMap([0.,0.],[1.,0.5])
facescale = HDGElasticity.face_determinant_jacobian(cellmap)
HH = [HDGElasticity.HHop(sbasis,facequads[i],facescale[i]) for i = 1:4]

testmatrix = [2/3 0.0 1/3 0.0
              0.0 2/3 0.0 1/3
              1/3 0.0 2/3 0.0
              0.0 1/3 0.0 2/3]

@test allapprox(HH[1],0.5*testmatrix)
@test allapprox(HH[2],0.25*testmatrix)
@test allapprox(HH[3],0.5*testmatrix)
@test allapprox(HH[4],0.25*testmatrix)

HH = reshape([0.0],1,1)
sfunc(x) = [0.5*x[1]*(x[1]+1)]
squad = tensor_product_quadrature(1,3)
components = [1.]

HDGElasticity.HHop!(HH,sfunc,squad,components,1.,1)
@test allapprox(HH,[4/15])

fill!(HH,0.0)
HDGElasticity.HHop!(HH,sfunc,squad,components,2.,1)
@test allapprox(HH,2*[4/15])

fill!(HH,0.0)
components = [1.0]
HDGElasticity.HHop!(HH,sfunc,squad,components,3.,1)
@test allapprox(HH,3*[4/15])

n = [1.,1.]/sqrt(2.)
normals = repeat(n,inner=(1,3))
imap = InterpolatingPolynomial(2,sbasis)
coeffs = [-0.5  -1.0
          -1.0  +0.0]
update!(imap,coeffs)
cellmap = HDGElasticity.CellMap([0.,0.],[2.,1.])

HH = zeros(4,4)
facescale = HDGElasticity.scale_area(cellmap,normals)
HDGElasticity.HHop!(HH,sbasis,squad,imap,facescale,2)
@test allapprox(HH,0.5/sqrt(2.0)*testmatrix)

fill!(HH,0.0)
HDGElasticity.HHop!(HH,sbasis,squad,imap,2facescale,2)
@test allapprox(HH,1.0/sqrt(2.0)*testmatrix)

HH = HDGElasticity.HHop(sbasis,squad,imap,normals,cellmap)
@test allapprox(HH,0.5/sqrt(2.0)*testmatrix)

cellmap = HDGElasticity.CellMap([0.,0.],[1.,0.5])
coeffs = [1.0,-1.0,1.,1.]
update!(imap,coeffs)
normals = repeat([1.0,0.0],inner=(1,length(squad)))
HH = zeros(4,4)

facescale = HDGElasticity.scale_area(cellmap,normals)
HDGElasticity.HHop!(HH,sbasis,squad,imap,normals,facescale,2)
testHH = 0.25*testmatrix
testHH[2,:] .= 0.0
testHH[4,:] .= 0.0
@test allapprox(HH,testHH)

fill!(HH,0.0)
HDGElasticity.HHop!(HH,sbasis,squad,imap,normals,2facescale,2)
@test allapprox(HH,2testHH)

HH = HDGElasticity.HHop(sbasis,squad,imap,normals,normals,cellmap)
@test allapprox(HH,testHH)


using Photometry.Aperture: circular_overlap,
                           circular_overlap_core,
                           circular_overlap_single_exact,
                           circular_overlap_single_subpixel,
                           area_arc,
                           area_triangle

@testset "circular overlap" for grid_size in [50, 500, 1000], circ_size in (0.2, 0.4, 0.8), method in [:exact, :center, (:subpixel, 2), (:subpixel, 5), (:subpixel, 10)]

    g = circular_overlap(-1, -1, 1, 1, grid_size, grid_size, circ_size, method = method)
    any(g .> 0) &&  @test maximum(g) ≈ 1.0
    any(g .< 1) &&  @test minimum(g) ≈ 0.0
end

@testset "area triangle" begin
    @test area_triangle(0, 0, 0, 0, 0, 0) ≈ 0

    @test area_triangle(0, 0, 1, 0, 0, 1) ≈ 1 / 2

    @test area_triangle(1, 1, 1, 4, 6, 1) ≈ 7.5

end

@testset "area arc (r = $r)" for r in 1:10
    @test area_arc(0, 0, 0, 0, r) ≈ 0
    @test area_arc(0, r, r, 0, r) ≈ r^2 / 2 * (π / 2 - 1)
end

@testset "overlap core (r = $r)" for r in 1:10
    @test circular_overlap_core(0, 0, 0, 0, r) ≈ 0
    @test circular_overlap_single_exact(0, 0, r, r, r) ≈ r^2 * π / 4
    @test circular_overlap_core(0, 0, r, r, r) ≈ r^2 * π / 4
end

@testset "overlap single exact" begin
    @test circular_overlap_single_exact(0, 0, 0, 0, 1) ≈ 0
    @test circular_overlap_single_exact(0, 0, 20, 20, 15) ≈ 176.714586764429
end

@testset "overlap subpixel" begin
    @test circular_overlap_single_subpixel(0, 0, 0, 0, 1, 1) ≈ 1

    @test circular_overlap_single_subpixel(0, 0, 20, 20, 15, 2) ≈ 0.25
    @test circular_overlap_single_subpixel(0, 0, 20, 20, 15, 5) ≈ 0.44
    @test circular_overlap_single_subpixel(0, 0, 20, 20, 15, 10) ≈ 0.43
    @test circular_overlap_single_subpixel(0, 0, 20, 20, 15, 100) ≈ 0.4423
end

@testset "overlap grid" begin
    @test circular_overlap(0, 0, 0, 0, 1, 1, 1) == ones(1, 1)

    @test circular_overlap(0, 20, 0, 20, 5, 5, 15) ≈ [1 1 1 0.70506901 0;
    1 1 1 0.42240835 0;
    1 1 0.74542733 0.02213982 0;
    0.70506901 0.42240835 0.02213982 0 0;
    0 0 0 0 0]

    @test circular_overlap(-10, 10, -10, 10, 5, 5, 4) ≈ [0 0         0        0        0;
    0 0.07878669 0.45661148 0.07878669 0;
    0 0.45661148 1         0.45661148 0;
    0 0.07878669 0.45661148 0.07878669 0;
    0 0         0         0        0]

end


#========================================================================

	Steady state solution for an Aiyagari model with a diffusion process

	Translated julia code based on code from Ben Moll's Website:
		http://www.princeton.edu/%7Emoll/HACTproject.htm

========================================================================#
using LinearAlgebra, SparseArrays, PyPlot

# Model parameters
γ = 2 # CRRA utility
α = 0.35 # parameter for cobb-douglas production function
δ = 0.1 # captal depreciation
z_mean = 1.0 # mean of levels of diffusion process
σ_sq = (.10)^2 # also for diffusion process
corr = exp(-0.3)
ρ = 0.05

K = 3.8 #inital guess of aggregate capital, need this to be close for algorithm to converge
relax = 0.99 # a relaxation parameter
J=40

z_min = 0.5
z_max = 1.5
a_min = -1.0
a_max = 30.0
H =100

# Paramters for the simulation
maxit = 100
maxitK = 100
crit = 10^(-6)
critK = 1e-5
Δ = 1000 #step size for HJB

# Diffusion process in levels
θ = -log(corr)
Var = σ_sq/(2*θ)

# Grid spaces for variables
a = LinRange(a_min, a_max, H)
a = convert(Array,a)
da = (a_max-a_min)/(H-1)

z = LinRange(z_min, z_max, J)
z = convert(Array,z)
dz = (z_max-z_min)/(J-1)
dz_sq = dz^2

aa = a*ones(1,J)
zz = ones(H,1)*z' #check later

# Other variables
μ = θ*(z_mean.-z) #drift for HJB via Ito's lemma
s_sq = σ_sq.*ones(1,J)

# Spaces for the finite difference terms
Vaf, Vab, Vzf, Vzb, Vzz, c = [zeros(H,J) for i in 1:6]

# Construct the matrix to summarize the evolution of z over time
	# this matrix will not change over the finite differences

yy = -(s_sq/dz_sq)' - μ/dz
χ = s_sq/(2*dz_sq)
ζ = μ/dz + (s_sq/(2*dz_sq))'

# Define the Diagonals of this matrix
updiag_z = zeros(H,1)
   for j = 1:J
	   global updiag_z =[updiag_z; repeat([ζ[j]], H, 1)]
   end
updiag_z =(updiag_z[:])

centerdiag_z=repeat([χ[1]+yy[1]],H,1)
   for j = 2:J-1
	   global centerdiag_z = [centerdiag_z; repeat([yy[j]], H, 1)]
   end
centerdiag_z=[centerdiag_z; repeat([yy[J]+ζ[J]], H, 1)]
centerdiag_z = centerdiag_z[:]

lowdiag_z = repeat([χ[2]], H, 1)
   for j=3:J
	   global lowdiag_z = [lowdiag_z; repeat([χ[j]],H,1)]
   end
lowdiag_z=lowdiag_z[:]

Aswitch = sparse(Diagonal(centerdiag_z))+ [zeros(H,H*J);  sparse(Diagonal(lowdiag_z)) zeros(H*(J-1), H)]+ sparse(Diagonal(updiag_z))[(H+1):end,1:(H*J)]

# Inital Guess for the prices
r = α * K^(α-1)-δ
w = (1-α) * K^(α)
v0 = (w*zz + r.* aa).^(1-γ)/(1-γ)/ρ
global v = v0
dist = zeros(1,maxit)

# The main loop, this makes sure markets clear

for iter in 1:maxitK
	println("Main Loop iteration $(iter)")

	for n in 1:maxit
		V=v

		# Forward difference
		Vaf[1:H-1,:] = (V[2:H,:]-V[1:H-1,:])/da
		Vaf[H,:] = (w*z .+ r.*a_max).^(-γ)

		# Backward difference
		Vab[2:H,:] = (V[2:H,:]-V[1:H-1,:])/da
		Vab[1,:] = (w*z .+ r.*a_min).^(-γ)

		# caculate consumption and savings
		# First the forward difference case
		cf = Vaf.^(-1/γ)
		sf = w*zz + r.*aa -cf

		# consumption and savings for the backward difference
		cb = Vab.^(-1/γ)
		sb = w*zz + r.*aa -cb

		#Consumption and V' at the steady state
		c0 = w*zz + r.*aa
		Va0 = c0.^(-γ)

		# Setup to the Upwind scheme
		If = sf .> 0
		Ib = sb .< 0
		I0 = (1 .- If - Ib)

		#Now implement the upwind scheme
		Va_Upwind = Vaf.*If + Vab.*Ib + Va0.*I0

		c= Va_Upwind.^(-1/γ)
		u = c.^(1-γ)/(1-γ)

		X = -min.(sb,0)/da
		Y = -max.(sf,0)/da + min.(sb,0)/da
		Z = max.(sf,0)/da

		updiag = 0
		   for j = 1:J
			   updiag =[updiag; Z[1:H-1,j];0]
		   end
		updiag =(updiag[:])

		centerdiag=reshape(Y,H*J,1)

		lowdiag = X[2:H,1]
		   for j=2:J
			   lowdiag = [lowdiag; 0; X[2:H,j]]
		   end
		lowdiag=lowdiag[:]

		AA = sparse(Diagonal(centerdiag[:]))+ [zeros(1, H*J); sparse(Diagonal(lowdiag)) zeros(H*J-1,1)] + sparse(Diagonal(updiag))[2:end, 1:(H*J)]

		global A= AA + Aswitch
		B = (1/Δ + ρ)*sparse(I,H*J,H*J)-A

		u_stacked = reshape(u, H*J, 1)
		V_stacked = reshape(V,H*J, 1)

		b = u_stacked + V_stacked/Δ

		V_stacked = B\b

		V = reshape(V_stacked, H,J)
		V_change = V-v
		global v = V

		dist[n] = findmax(V_change)[1]

		if dist[n] < crit
			println("Value Function Converge Iteration = $(n)")
			break
		end
	end
	# Now solve the Fokker-Planck equation
	AT = A'
	b = zeros(H*J,1)

	# Dirty fix to avoid matrix singularities
	i_fix=1
	b[i_fix] =.1
	row = [zeros(1,i_fix-1) 1 zeros(1,H*J-i_fix)]
	AT[i_fix,:] = row

	# Now solve the linear system
	gg = AT\b
	g_sum = gg'*ones(H*J,1)*da*dz
	gg = gg./g_sum

	global g = reshape(gg,H,J)

	# Update aggregate capital
	S = sum(g'*a*da*dz)
	println(S)

	if abs(K-S)<critK
		break
	end

	#update prices
	global K =relax*K + (1-relax)*S
	global r = α * K^(α-1) - δ
	global w = (1-α) * K^(α)
end

# Graphs

ss = w*zz + r.*aa - c
icut = 50
acut=a[1:icut]
sscut=ss[1:icut,:]
gcut = g[1:icut,:]

surface(acut,z,sscut', α=0.7,camera=(45,25),
        colorbar=false, zrotation=90,
        xlabel="Wealth", ylabel="Density")
png("OptimalSavings")

surface(acut,z,gcut', α=0.7,camera=(45,25))
png("KFE_pdf")


using OmegaCore
using OmegaTest
using Distributions
using Test

@testset "multivariate" begin
  μ = 5
  xs = 1 ~ Mv(Normal(5, 1), (1000,))
  samples = randsample(xs)
  @test mean(samples) ≈ μ atol = 0.3
  @test isinferred(randsample, xs)
end

@testset "Multivariate 2" begin
  x = 1 ~ Normal(0, 1)
  function f(id, ω)
    x(ω) + Uniform(0, 1)(id, ω)
  end
  xs = 2 ~ Mv(f, (3, 3))
  randsample((x, xs))
end


export p_speed
export far_field_pattern, pair_field_pattern, diff_far_field_pattern
export wavenumber_challis, one_species_low_wavenumber,
        two_species_approx_wavenumber
export wienerhopf_reflection_coefficient, wienerhopf_mode_amplitudes

include("acoustic_mediums.jl")

include("low_frequency.jl")

include("effective_wave/asymptotics.jl")

include("effective_wave/planewaves/eigensystem.jl")
include("effective_wave/planewaves/material_scattering_coefficients.jl")
include("effective_wave/planewaves/boundary_condition.jl")
# include("effective_wave/planewaves/wavemode-wienerhopf.jl")

# regular means a smooth field assumption and includes cases such as a sphere filled with particles.
include("effective_wave/regular/eigensystem.jl")
include("effective_wave/regular/boundary_condition.jl")
include("effective_wave/regular/material_scattering_coefficients.jl")

include("effective_wave/far_fields.jl")
include("effective_wave/low_volumefraction.jl")
include("effective_wave/alternative_wavenumbers.jl")
include("effective_wave/two_species_approximate.jl")

include("discrete_wave/export.jl")


using Polyhedra, CDDLib
P = polyhedron(vrep(eye(2)), CDDLibrary())
P * P


module Day6

using SparseArrays

M = spdiagm(1 => ones(Int32, 8))
M[7, 1] = M[9, 1] = 1

test = "3,4,3,1,2"

sim(input::String, day::Int) = begin
    pops = foldl(split(input, ","), init=zeros(Int, 9)) do acc, age
        acc[parse(Int, age)+1] += 1
        acc
    end
    sum(M^day * pops)
end

@assert sim(test, 80) == 5934
@assert sim(test, 256) == 26984457539

input = read("resources/day6_input.txt", String)
sim(input, 80) |> display
sim(input, 256) |> display

@time M^256

function do_it_really_really_fast()
    M2 = M^2
    M4 = M2^2
    M8 = M4^2
    M16 = M8^2
    M32 = Array(M16^2)
    M64 = M32^2
    M128 = M64^2
    M256 = M128^2
end

end


# Model
include_model("hopper")
model_ft = free_time_model(no_slip_model(model))

# Horizon]
T = 51

# Time step
tf = 0.5
h = tf / (T - 1)

# Bounds
_uu = Inf * ones(model_ft.m)
_uu[model_ft.idx_u] = model_ft.uU
_uu[end] = 2.0 * h
_ul = zeros(model_ft.m)
_ul[model_ft.idx_u] = model_ft.uL
_ul[end] = 0.5 * h
ul, uu = control_bounds(model_ft, T, _ul, _uu)

# Initial and final states
z_h = 0.25
q1 = [0.0, 0.5 + z_h, 0.0, 0.25]

xl, xu = state_bounds(model_ft, T,
		[model_ft.qL; model_ft.qL],
		[model_ft.qU; model_ft.qU],
        x1 = [q1; Inf * ones(model.nq)])

# Objective
include_objective("velocity")
obj_tracking = quadratic_time_tracking_objective(
    [Diagonal(zeros(model_ft.n)) for t = 1:T],
    [Diagonal([1.0, 1.0, zeros(model_ft.m - model_ft.nu - 1)..., 0.0]) for t = 1:T-1],
    [zeros(model_ft.n) for t = 1:T],
    [zeros(model_ft.m) for t = 1:T],
    1.0)
obj_contact_penalty = PenaltyObjective(1.0e5, model_ft.m - 1)
obj_velocity = velocity_objective(
    [Diagonal(1.0 * ones(model_ft.nq)) for t = 1:T-1],
    model_ft.nq,
    h = h,
    idx_angle = collect([3]))
obj = MultiObjective([obj_tracking, obj_contact_penalty, obj_velocity])

# Constraints
include_constraints(["free_time", "contact_no_slip", "loop"])
con_free_time = free_time_constraints(T)
con_contact = contact_no_slip_constraints(model_ft, T)
con_loop = loop_constraints(model, (1:model.n), 1, T)
con = multiple_constraints([con_free_time, con_contact, con_loop])

# Problem
prob = trajectory_optimization_problem(model_ft,
               obj,
               T,
               xl = xl,
               xu = xu,
               ul = ul,
               uu = uu,
               con = con)

# Trajectory initialization
x0 = [[q1; q1] for t = 1:T] # linear interpolation on state
u0 = [[1.0e-5 * rand(model_ft.m-1); h] for t = 1:T-1] # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

#NOTE: may need to run examples multiple times to get good trajectories
# Solve nominal problem

optimize = true
if optimize
	include_snopt()
	@time z̄ , info = solve(prob, copy(z0),
		nlp = :ipopt,
		tol = 1.0e-5, c_tol = 1.0e-5, mapl = 5,
		time_limit = 60)
	@show check_slack(z̄, prob)
	x̄, ū = unpack(z̄, prob)
	tf, t, h̄ = get_time(ū)

	# projection
	Q = [Diagonal(ones(model_ft.n)) for t = 1:T]
	R = [Diagonal(0.1 * ones(model_ft.m)) for t = 1:T-1]
	x_proj, u_proj = lqr_projection(model_ft, x̄, ū, h̄[1], Q, R)

	@show tf
	@show h̄[1]
	@save joinpath(pwd(), "examples/trajectories/hopper_vertical_gait_no_slip.jld2") x̄ ū h̄ x_proj u_proj
else
	@load joinpath(pwd(), "examples/trajectories/hopper_vertical_gait_no_slip.jld2") x̄ ū h̄ x_proj u_proj
end

using Plots
plot(hcat(ū...)[1:2, :]',
    linetype = :steppost,
    label = "",
    color = :red,
    width = 2.0)

plot!(hcat(u_proj...)[1:2, :]',
    linetype = :steppost,
    label = "", color = :black)

plot(hcat(state_to_configuration(x̄)...)',
    color = :red,
    width = 2.0,
	label = "")
plot!(hcat(state_to_configuration(x_proj)...)',
    color = :black,
	label = "")

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model_ft, state_to_configuration(x_proj), Δt = h̄[1])


function calc_cstar!(c,cm1,Δt,Δt2)
    @mthreads for j in TRANGE
        r = (Δt/Δt2)
        @inbounds @simd for i in REAL_RANGES[j]
            cstar = muladd(r,(c[i]-cm1[i]),c[i])
            cm1[i] = c[i]
            #c[i] = max(0.0,cstar)
            c[i] = cstar
        end
    end
end

function average_c!(c,Δ2)
    fullfourier!(c)
    ind = CartesianIndices((Base.OneTo(NY),Base.OneTo(NZ)))
    @mthreads for jk in ind
        j,k = Tuple(jk)
            @inbounds @msimd for i in Base.OneTo(NX)
                G = Gaussfilter(Δ2,i,j,k)
                c[i,j,k] *= G
            end
    end

    real!(c)
    return nothing
end


using Underscores
using Utils

struct BusChoice
    arrival::Int
    choices::Array{Union{Int64,Missing}}
end

function parse_file(f)
    lines = eachline(f)
    line, next = iterate(lines)
    time = parse(Int, line)
    line, next = iterate(lines, next)
    choices = @_ line |> split(__, ',') |> map(_ == "x" ? missing : parse(Int64, _), __)
    BusChoice(time, choices)
end

function load()
    open(joinpath(@__DIR__, "13_bus-ids.txt"), "r") do f
        parse_file(f)
    end
end

function ex1(plan)
    nexts = @_ map(cld(plan.arrival, _) * _, plan.choices)
    minutes = @_ map(_ - plan.arrival, nexts)
    perm = sortperm(nexts)
    minutes[perm[1]] * plan.choices[perm[1]]
end

function combine_moduli((a1, n1), (a2, n2))
    d, m1, m2 = gcdx(n1, n2)
    @assert d == 1
    x = a1 * n2 * m2 + a2 * n1 * m1
    p = n1 * n2
    (p + x) % p, p
end

function ex2_crt(busses)
    only_nums = skipmissing(busses)

    @assert all(((a, b),) -> a == b || gcd(a, b) == 1, Iterators.product(only_nums, only_nums))

    busses = @_ only_nums |> eachindex |> map(let id = only_nums[_]; (convert(Int128, id - _ + 1), id) end, __)

    a, n = reduce(combine_moduli, busses)
    while a < 0
        a += n
    end
    a % n
end

plan = load()

println("First bus times wait time: ", ex1(plan))
println("Perfect alignment: ", ex2_crt(plan.choices))

using Test

@testset "Shuttle Search" begin
    input = """
    939
    7,13,x,x,59,x,31,19
    """

    example = parse_file(IOBuffer(input))
    @test example.arrival == 939
    @test length(example.choices) == 8

    @testset "example 1" begin
        @test ex1(example) == 295
    end

    @testset "example 2" begin
        for f in [ex2_crt]
            @test f([67,7,59,61]) == 754018
            @test f([17,missing,13,19]) == 3417
            @test f([67,missing,7,59,61]) == 779210
            @test f([67,7,missing,59,61]) == 1261476
            @test f([1789,37,47,1889]) == 1202161486
            @test f(example.choices) == 1068781
        end
    end

    @testset "results" begin
        @test ex1(plan) == 8063
        @test ex2_crt(plan.choices) == 775230782877242
    end
end


@testset "DMRG" begin
    @testset "DAS" begin
        #Hubbard
        l = 3
        pchs = U1Charges(-3:3)
        vchs = U1Charges(-8:8)
        pds = [1 for i in pchs]
        vds = [5 for i in vchs]

        mpo = hubbardMPO(1, 1, pchs, l, T = Float64)
        mps = randmps(l, Float64, (U1(), (pchs, vchs), (pds, vds)))
        gs, es = dmrg(mps, mpo)

        left, blk, right = mpo[[1,2,l]]
        @tensor h[1,2,3,4,5,6] := left[1,-1,4] * blk[2,-1,-2,5] * right[3,-2,6];
        hmattens = fuselegs(h, ((1,2,3),(4,5,6)), InOut(1,-1))[1]
        hmat = hmattens[DASSector(U1Charge(0),U1Charge(0))]
        ref = minimum(eigvals(hmat))
        res = es[end][end]
        @test  ref ≈ res

        #tfIsing
        for l in 2:10 #edge cases
            mps = randmps(l, ComplexF64, (Z2(), (Z2Charges(), Z2Charges()), ([1,1], [5,5])));
            @test dmrg(mps, tfisingMPO(0,l))[2][end] ≈ -(l-1)
            @test dmrg(mps, tfisingMPO(1e6,l))[2][end]/1e6 ≈ -l
        end
        l = 22
        mps = randmps(l, ComplexF64, (Z2(), (Z2Charges(), Z2Charges()), ([1,1], [5,5])));
        gs = dmrg(mps, tfisingMPO(0,l))[1]
        @test isapprox(onesiteexp(gs, sigmaZop(), 11), 0; atol = 1e-9)
        @test isapprox(onesitevar(gs, sigmaZop(), 11), 1; atol = 1e-9)

        gs = dmrg(mps, tfisingMPO(1e6,l))[1]
        @test isapprox(onesiteexp(gs, sigmaZop(), 11), 1; atol = 1e-9)
        @test isapprox(onesitevar(gs, sigmaZop(), 11), 0; atol = 1e-9)

        gs = dmrg(mps, tfisingMPO(-1e6,l))[1]
        @test isapprox(onesiteexp(gs, sigmaZop(), div(l,2)), -1; atol = 1e-9)
        @test isapprox(onesitevar(gs, sigmaZop(), div(l,2)), 0; atol = 1e-9)
    end

    @testset "Dense" begin
        #Hubbard
        l = 3
        ansatz = randmps(l, Float64, (7,85))
        mpo = hubbardMPO(1, 1, U1Charges(-3:3), l, T = Float64)
        gsa, eas = dmrg(ansatz, todense(mpo))

        left, blk, right = todense(mpo)[[1,2,l]]
        @tensor h[1,2,3,4,5,6] := left[1,-1,4] * blk[2,-1,-2,5] * right[3,-2,6];
        hmat = fuselegs(h, ((1,2,3),(4,5,6)))[1]
        ref = min(eigvals(toarray(hmat))...)
        res = eas[end][end]
        @test ref ≈ res

        #tfIsing
        l = 22
        mps = randmps(l, ComplexF64, (Z2(), (Z2Charges(), Z2Charges()), ([1,1], [5,5])));
        dmps = todense(mps)
        gs = dmrg(dmps, todense(tfisingMPO(0,l)))[1]
        @test isapprox(onesiteexp(gs, todense(sigmaZop()), div(l,2)), 0; atol = 1e-9)
        @test isapprox(onesitevar(gs, todense(sigmaZop()), div(l,2)), 1; atol = 1e-9)

        gs = dmrg(dmps, todense(tfisingMPO(1e6,l)))[1]
        @test isapprox(onesiteexp(gs, todense(sigmaZop()), div(l,2)), 1; atol = 1e-9)
        @test isapprox(onesitevar(gs, todense(sigmaZop()), div(l,2)), 0; atol = 1e-9)

        gs = dmrg(dmps, todense(tfisingMPO(-1e6,l)))[1]
        @test isapprox(onesiteexp(gs, todense(sigmaZop()), div(l,2)), -1; atol = 1e-9)
        @test isapprox(onesitevar(gs, todense(sigmaZop()), div(l,2)), 0; atol = 1e-9)
    end
end


function getRealVs(V,evals)

    # get real vectors made from the eigenvectors

    # by paying attention to this order, things will always rotate CCW
    if abs(evals[1])>0  # if the eigenvalue with negative imaginary component comes first
        Vr = [V[:,1] + V[:,2], (V[:,1] - V[:,2])*im];
    else
        Vr = [V[:,2] + V[:,1], (V[:,2] - V[:,1])*im]
    end
    Vr = Vr / sqrt(2)
    return Vr
    #=
    # now get axes aligned so that plan is spread mostly along the horizontal axis()
    testProj = (Vr"*Ared[1:numAnalyzedTimes:end,:]')"; # just picks out the plan times
    rotV = princomp[testProj]
    crossProd = cross([rotV[:,1];0], [rotV[:,2];0])
    if crossProd[3] < 0, rotV[:,2] = -rotV[:,2]; end   # make sure the second vector is 90 degrees clockwise from the first
    Vr = Vr*rotV;

    # flip both axes if necessary so that the maximum move excursion is in the positive direction
    testProj = (Vr"*Ared')";  # all the times
    if max(abs(testProj[:,2])) > max(testProj[:,2])  # 2nd column is the putative "muscle potent" direction.
        Vr = -Vr
    end
    =#
end


using EmpiricalBayes
using Test
using Random: seed!

marginal_grid = collect(range(-6,stop=6,length=1001));
prior_grid = collect(range(-3,stop=3,length=121));
ds = MixingNormalConvolutionProblem(Normal, 0.2, prior_grid, marginal_grid);

true_dist = MixtureModel([ Normal(-2,.2), Normal(+2,.2)])
d_center = NormalConvolutionProblem(true_dist, marginal_grid);
f_center = BinnedMarginalDensity(d_center)

m = 1_000
ma = MinimaxCalibrator(ds, f_center, m, MarginalDensityTarget(0.0); C=Inf,
    ε=0.0, max_smoother=false);

seed!(1)
Xs = rand(true_dist, m)


# Basic way of constructing Neighborhoods.

struct KDECalibrator{K} <: LinearEstimator where K<:ContinuousUnivariateDistribution
    kernel::K
    Q::BinnedCalibrator
    m::Int64
end

function KDECalibrator(::Type{T}, m::Int64, marginal_grid;
    h= default_bandwidth(T, m)) where T<:ContinuousUnivariateDistribution
    kernel = T(h)
    Qc = BinnedCalibrator(marginal_grid, pdf.(kernel, marginal_grid))
    KDECalibrator(kernel, Qc, m)
end

# TODO: Make more general maybe
inference_target(a::KDECalibrator) = MarginalDensityTarget(0.0)

struct SincKernel <: ContinuousUnivariateDistribution
    h::Float64 #Bandwidth
end

pretty_label(kd::KDECalibrator{SincKernel})  = "Sinc Kernel"

default_bandwidth(a::Type{SincKernel}, m) = 1/sqrt(log(m))

cf(a::SincKernel, t) = one(Float64)*(-1/a.h <= t <= 1/a.h)

function pdf(a::SincKernel, t)
    if t==zero(Float64)
        return(one(Float64)/pi/a.h)
    else
        return(sin(t/a.h)/pi/t)
    end
end

struct DeLaValleePoussinKernel <: ContinuousUnivariateDistribution
    h::Float64 #Bandwidth
end

pretty_label(kd::KDECalibrator{DeLaValleePoussinKernel})  = "DLVP Kernel"

default_bandwidth(a::Type{DeLaValleePoussinKernel}, m) = 1/sqrt(log(m))

function cf(a::DeLaValleePoussinKernel, t)
    if abs(t * a.h) <= 1
        return(one(Float64))
    elseif abs(t * a.h) <= 2
        return(2*one(Float64) - abs(t * a.h))
    else
        return(zero(Float64))
    end
end

function pdf(a::DeLaValleePoussinKernel, t)
    if t==zero(Float64)
        return(3*one(Float64)/2/pi/a.h)
    else
        return(a.h*(cos(t/a.h)-cos(2*t/a.h))/pi/t^2)
    end
end

function sinc_kde(Xs, marginal_grid, ::Type{T};
     ws=KernelDensity.UniformWeights(length(Xs))) where T<:ContinuousUnivariateDistribution
    m = length(Xs)
    h = default_bandwidth(T,m)
    ker = T(h)
    (grid_min, grid_max) = extrema(marginal_grid)
    # hack for now to avoid weird floating point tricks wherein the conversion from
    # vector to range fails.. Should switch to ranges everywhere though
    marginal_h = marginal_grid[2] - marginal_grid[1] - 1e-10
    f_marginal = kde(Xs, ws,
         grid_min:marginal_h:grid_max, ker);
    f_marginal
end

function sinc_kde(Xs, marginal_grid; kwargs...)
    sinc_kde(Xs, marginal_grid, DeLaValleePoussinKernel; kwargs... )
end

mutable struct BinnedMarginalDensityNeighborhood
    f::BinnedMarginalDensity
    f_kde::UnivariateKDE
    C_bias::Float64
    C_std::Float64
    η_infl::Float64
end

function BinnedMarginalDensityNeighborhood(f, f_kde, C_bias,C_std)
    BinnedMarginalDensityNeighborhood(f, f_kde, C_bias, C_std, 0.01)
end
# should be a fit function
function fit(::Type{BinnedMarginalDensityNeighborhood}, Xs,
                marginal_grid::Vector{Float64}, ::Type{T};
                nboot=101) where T<:ContinuousUnivariateDistribution

    marginal_h = marginal_grid[2] - marginal_grid[1]
    f_kde = sinc_kde(Xs, marginal_grid, T)

    m = length(Xs)
    C_stds = Vector{Float64}(undef, nboot)

    for k =1:nboot
        # Poisson bootstrap to estimate certainty band
        Z_pois = rand(Poisson(1), m)
        ws =  Weights(Z_pois/sum(Z_pois))
        f_kde_pois =  sinc_kde(Xs, marginal_grid, T; ws=ws)
        C_stds[k] = maximum(abs.(f_kde.density .- f_kde_pois.density))
    end

    C_std = median(C_stds)

    f = BinnedMarginalDensity(f_kde.density * marginal_h, marginal_grid, marginal_h)

    BinnedMarginalDensityNeighborhood(f, f_kde, 0.0, C_std)
end

function fit(::Type{BinnedMarginalDensityNeighborhood}, Xs,
       marginal_grid::Vector{Float64};
       kwargs...)

       fit(BinnedMarginalDensityNeighborhood, Xs,
              marginal_grid, DeLaValleePoussinKernel;
              kwargs...)
end

function fit(::Type{BinnedMarginalDensityNeighborhood}, Xs,
       ds::MixingNormalConvolutionProblem, ::Type{T};
       kwargs...) where T<:ContinuousUnivariateDistribution

    m = length(Xs)
    marginal_grid = ds.marginal_grid

    f = fit(BinnedMarginalDensityNeighborhood, Xs, marginal_grid, T; kwargs...)

    #cb = ComteButucea(MarginalDensityTarget(0.0), m, marginal_grid)
    cb = KDECalibrator(T, m, marginal_grid)

    max_bias = check_bias(cb, ds; maximization=true)
    f.C_bias = max_bias
    f
end

function fit(::Type{BinnedMarginalDensityNeighborhood}, Xs,
       ds::MixingNormalConvolutionProblem;
       kwargs...)

       fit(BinnedMarginalDensityNeighborhood, Xs,
              ds::MixingNormalConvolutionProblem, DeLaValleePoussinKernel;
              kwargs...)
end

function check_bias(cb::Union{ComteButucea, KDECalibrator},
                  ds::MixingNormalConvolutionProblem;
                  maximization=true)

        f = BinnedMarginalDensity([],[],0)

        check_bias(cb.Q, ds, f, cb.m, inference_target(cb);
                         C=Inf, maximization=maximization)
end


push!(LOAD_PATH, "./src")

# "Packages"
using Revise # reduce need for recompile
using Plots
using Documenter
using LinearAlgebra
using SparseArrays

# "User defined modules"
using CommonUtils
using Basis1D
using Basis3DHex
using UniformHexMesh

using SetupDG
using UnPack

N = 3
K1D = 8
CFL = .75
T = .25 # endtime

VX,VY,VZ,EToV = uniform_hex_mesh(K1D,K1D,K1D)

rd = init_reference_hex(N,gauss_quad(0,0,N))
md = init_mesh((VX,VY,VZ),EToV,rd)

# make domain periodic
@unpack Nfaces = rd
@unpack xf,yf,zf,K,mapM,mapP,mapB = md
LX = 2; LY = 2; LZ = 2
mapPB = build_periodic_boundary_maps(xf,yf,zf,LX,LY,LZ,
                                    Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB
@pack! md = mapP

"Time integration"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (N+1)*(N+2)*3/2  # estimated trace constant
dt = CFL * 2 / (CN*K1D)
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

function rhs(u,rd::RefElemData,md::MeshData)

    @unpack Dr,Ds,Dt,LIFT,Vf = rd
    @unpack rxJ,sxJ,txJ,ryJ,syJ,tyJ,rzJ,szJ,tzJ,J = md
    @unpack nxJ,nyJ,nzJ,sJ = md
    @unpack mapP,mapB = md

    uf = Vf*u
    du = uf[mapP]-uf
    uflux = @. .5*(du*nxJ - 0*abs(nxJ)*du)

    ur,us,ut = (A->A*u).((Dr,Ds,Dt))
    uxJ      = @. rxJ*ur + sxJ*us + txJ*ut
    rhsu     = uxJ + LIFT*uflux

    return @. -rhsu/J
end

# set initial conditions
@unpack x,y,z = md
u = @. exp(-25*(x^2+y^2))

resQ = zeros(size(x))
for i = 1:Nsteps
    for INTRK = 1:5
        rhsQ    = rhs(u,rd,md)
        @. resQ = rk4a[INTRK]*resQ + dt*rhsQ
        @. u   += rk4b[INTRK]*resQ
    end

    if i%10==0 || i==Nsteps
        println("Number of time steps: $i out of $Nsteps")
    end
end

# "plotting nodes"
gr(aspect_ratio=1,legend=false,
   markerstrokewidth=0,markersize=2)

@unpack Vp = rd
xp,yp,zp = (x->Vp*x).((x,y,z))
vv = Vp*u
ids = map(x->x[1],findall(@. abs(zp[:]-1)<1e-10))
(xp,yp,zp,vv) = (x->x[ids]).((xp,yp,zp,vv))
scatter(xp,yp,vv,zcolor=vv,camera=(0,90))


function naive_index_from_time(sample_rate, sample_time)
    # This stepping computation is prone to roundoff error, so we'll work in high precision
    sample_time_in_seconds = big(Dates.value(Nanosecond(sample_time))) //
                             big(TimeSpans.NS_IN_SEC)
    # At time 0, we are at index 1
    t = Rational{BigInt}(0 // 1)
    index = 1
    while true
        # Now step forward in time; one index, and time 1/sample_rate
        t += 1 // sample_rate
        index += 1
        if t > sample_time_in_seconds
            # we just passed it, so previous index is the last one before the time of interest
            return index - 1
        end
    end
end

# Modified from
# https://github.com/beacon-biosignals/TimeSpans.jl/blob/e3c999021336e51a08d118e6defb792e38ac1cc7/test/runtests.jl#L116-L126
@testset "index_and_error_from_time" begin
    for rate in (101 // 2, 1001 // 10, 200, 256, 1, 10)
        for sample_time in
            (Nanosecond(12345), Minute(5), Nanosecond(Minute(5)) + Nanosecond(1),
             Nanosecond(1), Nanosecond(10^6), Nanosecond(6970297031))
            # compute with a very simple algorithm
            index = naive_index_from_time(rate, sample_time)
            # Check against our `TimeSpans.index_from_time`:
            @test index ==
                  AlignedSpans.index_and_error_from_time(rate, sample_time, RoundDown)[1]
            # Works even if `rate` is in Float64 precision:
            @test index ==
                  AlignedSpans.index_and_error_from_time(Float64(rate), sample_time,
                                                         RoundDown)[1]
        end
    end
end

# https://github.com/beacon-biosignals/TimeSpans.jl/blob/e3c999021336e51a08d118e6defb792e38ac1cc7/test/runtests.jl#L93-L96
@testset "time_from_index" begin
    @test_throws ArgumentError AlignedSpans.time_from_index(200, 0)
    @test AlignedSpans.time_from_index(100, 1) == Nanosecond(0)
end

@testset "`n_samples`" begin
    times = [Nanosecond(12345), Minute(5), Nanosecond(Minute(5)) + Nanosecond(1),
             Nanosecond(1), Nanosecond(10^6), Nanosecond(6970297031)]
    append!(times, Nanosecond.(Second(1):Second(1):Second(20)))
    sort!(times)

    for rate in (101 // 2, 1001 // 10, 200, 256, 1, 10)
        for i in 1:length(times), j in (i + 1):length(times)
            span = TimeSpan(times[i], times[j])

            if duration(span) < Nanosecond(floor(Int, 10^9 * inv(rate)))
                @test n_samples(rate, TimeSpans.duration(span)) == 0
                continue
            end

            aligned_span = AlignedSpan(rate, span, RoundInward)
            n = n_samples(aligned_span)
            @test n == length(indices(aligned_span))

            # `n` is the *actual* number of samples, given the whole span
            # `n_samples(rate, TimeSpans.duration(span))` gives us the minimum
            # number of samples of any span of that duration.
            # Thus, we should have:
            @test n >= n_samples(rate, TimeSpans.duration(span)) >= (n - 1)
        end
    end
end


function uncertainty_set_with_index(i, data, ϵ)
    set = uncertainty_set(data[i, :], ϵ)
    return collect(zip(fill(i, length(set)), set))
end

function combine(a::NTuple{K, Tuple{CartesianIndex{M}, Vector{Float64}}}, s::NTuple{N, Int})::Array{Float64, N} where {M, N, K}
    @assert M == N - 1
    p = zeros(s...)
    for (i, v) in a
        p[i, :] = v
    end
    return p
end

function uncertainty_sets(data::Array{Float64, N}, ϵ::Float64)::Vector{Array{Float64, N}} where N
    # TODO: ϵs::Vector{Float64}, size: |S_j|, j::Node
    s = size(data)
    sets = [uncertainty_set_with_index(i, data, ϵ) for i in CartesianIndices(s[1:end-1])];
    prods = vec(map(a -> combine(a, s), product(sets...)));
    return prods
end

function worst_case(model, objs)
    @variable(model, x)
    for obj in objs
        @constraint(model, x ≤ obj)
    end
    return x
end


using MacroTools: @q
import PyCall
using MLStyle
using Lazy

import SymPy
using SymPy: Sym, sympy, symbols

# stats = PyCall.pyimport_conda("sympy.stats", "sympy")
# import_from(stats)

export sym
sym(s::Symbol) = SymPy.symbols(s)
sym(s) = Base.convert(Sym, s)
function sym(expr::Expr)
    @match expr begin
        Expr(:call, f, args...) => :($f($(map(sym,args)...)))
        :($x[$j]) => begin
            j = symbols(:j, cls=sympy.Idx)
            x = sympy.IndexedBase(x)
            return x[j]
        end
        _ => begin
                 @show expr
                 error("sym: Argument type not implemented")
             end
    end
end

export symlogpdf
function symlogpdf(m::Model)
    result = @q begin
        ctx = Dict()

        ℓ = zero(Soss.Sym)
    end

    exprs = @as x m begin
        canonical(x)
        dropLines(x)
        x.body
        symlogpdf.(x)
    end

    append!(result.args, exprs)

    # result

    push!(result.args, :(ctx,ℓ))

    result

    # expandSums(ℓ)
end

function symlogpdf(st::Soss.Follows)
    d = st.rhs
    x = st.x
    :(ℓ += $(symlogpdf(d,x)))
end

function symlogpdf(st::Soss.Let)
    val = st.rhs
    x = st.x
    :(ctx[$(QuoteNode(x))] = Expr(:$, $val))
    # :($x = $val)
end

function symlogpdf(d::Expr, x::Symbol)
    @match d begin
        :(iid($n,$dist)) => begin
            j = symbols(:j, cls=sympy.Idx)
            dist = sym(dist)
            x = sympy.IndexedBase(x)
            n = sym(n)
            :(Soss.sympy.Sum(logpdf($dist,$x[$j]), ($j,1,$n)))
        end

        :(For($f, 1:$n)) => begin
            n = sym(n)
            @match f begin
                :(($j,) -> begin $lineno; $dist end) => begin
                            j = symbols(j) # , cls=sympy.Idx)
                            # @show j
                            dist = sym(dist)
                            # @show dist
                            x = sympy.IndexedBase(x)
                            return :(Soss.sympy.Sum(logpdf($dist,$x[$j]), ($j,1,$n)))
                end

                f => begin
                    @show f
                    error("symlogpdf: bad argument")
                end
            end

        end

        _ => :(logpdf($(sym(d)), $(sym(x))))
    end
end

function expandSums(s::Sym)
    s.args==() && return s
    if s.func==sympy.Sum
        s2 = expandSum(s)
        # Maybe it didn't simplify and we're done
        s2 == s && return s2
        # Or not, and we recurse
        s = s2
    end
    newargs = [expandSums(t) for t in s.args]
    s.func(newargs...)
end

function expandSum(s::Sym)
    # println("expandSum")
    # @show s
    # println()
    @assert s.func == sympy.Sum
    sfunc = s.args[1].func
    sargs = s.args[1].args
    limits = s.args[2]
    ix = limits.args[1]
    if sfunc == sympy.Add
        return sfunc([maybesum(t, limits) for t in sargs]...)
    elseif sfunc == sympy.Mul
        factors = sargs
        constants = [fac for fac in factors if !(ix in fac)]
        newconst = foldl(*,constants)
        newfacs = foldl(*,setdiff(factors, constants))
        newsum = sympy.Sum(newfacs, limits)
        return newconst * newsum
    else
        return s
    end

end

import Base.in
function Base.in(j::Sym, s::Sym)
    for t in s.args
        if j==t || in(j,t)
            return true
        end
    end
    return false
end

function maybesum(t::Sym, limits::Sym)
    # println("maybeSum")
    # @show t,limits
    # println()
    j = limits.args[1]
    thesum = sympy.Sum(t, limits)
    ifelse(j in t, thesum, thesum.doit())
end

# integrate(exp(ℓ), (sym(:μ), -oo, oo), (sym(:logσ),-oo,oo))
export marginal
function marginal(ℓ,v)
    f = ℓ.func
    f == sympy.Add || return ℓ
    newargs = filter(t -> sym(v) in t, collect(ℓ.args))
    foldl(+,newargs)
end

marginal(m::Model, v) = marginal(m |> symlogpdf, v)

# We should be able to reason about a marginal from its derivative
export dmarginal
function dmarginal(ℓ, v)
    @as x ℓ begin
        marginal(x,sym(v))
        diff(x, sym(v))
        expand(x)
        sympy.collect(x, sym(v))
    end
end

dmarginal(m::Model, v) = dmarginal(m |> symlogpdf, v)

# https://discourse.julialang.org/t/pyobjects-as-keys/26521/2
const symfuncs = Dict()
function __init__()
    merge!(symfuncs, Dict(
          sympy.log => log
        , sympy.Pow => :^
        , sympy.Abs => abs
        , sympy.Indexed => getindex
    ))
end

export codegen
function codegen(s::Sym)
    s.func == sympy.Add && begin
        @gensym add
        ex = @q begin
            $add = 0.0
        end
        for arg in s.args
            t = codegen(arg)
            push!(ex.args, :($add += $t))
        end
        push!(ex.args, add)
        # @show ex
        return ex
    end

    s.func == sympy.Mul && begin
        @gensym mul
        ex = @q begin
            $mul = 1.0
        end
        for arg in s.args
            t = codegen(arg)
            push!(ex.args, :($mul *= $t))
        end
        push!(ex.args, mul)
        return ex
    end

    # s.func == sympy.Sum && begin

    # end

    s.func ∈ keys(symfuncs) && begin
        # @show s
        @gensym symfunc
        argnames = gensym.("arg" .* string.(1:length(s.args)))
        argvals = codegen.(s.args)
        ex = @q begin end
        for (k,v) in zip(argnames, argvals)
            push!(ex.args, :($k = $v))
        end
        f = symfuncs[s.func]
        push!(ex.args, :($symfunc = $f($(argnames...))))
        push!(ex.args, symfunc)
        return ex
    end

    s.func == sympy.Sum && begin
        @gensym sum
        @gensym Δsum
        @gensym lo
        @gensym hi

        summand = codegen(s.args[1])
        (ix, ixlo, ixhi) = s.args[2].args

        ex = @q begin
            $sum = 0.0
            $lo = $(codegen(ixlo))
            $hi = $(codegen(ixhi))
            @inbounds @simd for $(codegen(ix)) = $lo:$hi
                $Δsum = $summand
                $sum += $Δsum
            end
            $sum
        end

        return ex
    end

    s.func == sympy.Symbol && return Symbol(string(s))
    s.func == sympy.Idx && return Symbol(string(s))
    s.func == sympy.IndexedBase && return Symbol(string(s))

    # @show s
    SymPy.is_real(s) && begin
        return SymPy.N(s)
    end

    @show s.func
    error("codegen")
end

codegen(s::AbstractFloat) = s

# export codegen
function codegen(m::Model)
    code = codegen(symlogpdf(m))
    unknowns = parameters(m) ∪ arguments(m)
    unkExpr = Expr(:tuple,unknowns...)
    @gensym logdensity
    result = @q begin
        function $logdensity(pars)
            @unpack $(unkExpr) = pars
            $code
        end
    end

    flatten(result)
end

# s = symlogpdf(normalModel).args[7].args[3]

# export fexpr
# fexpr = quote
#     f = function(μ,σ,x)
#         a = $(codegen(s))
#         return a
#     end
# end


#=
   My faster rewrites of `min`, `max`, `minmax`, `maxmin`.

   These are processing isomorphs for all normal uses.
   Julia's NaN, NaN32 are handled properly.
   Negatively signed NaNs, which cannot appear unless
   intentionally generated, are not IEEE754 compliant.

   For use with `Float32`, `Float64`.
=#

const FastFloat = Union{Float32, Float64}

function fastermin(a::T, b::T) where {T<:FastFloat}
    a_b = a - b
    (signbit(a_b) || isnan(a)) ? a : b
end

function fastermax(a::T, b::T) where {T<:FastFloat}
    b_a = b - a
    (signbit(b_a) || isnan(a)) ? a : b
end

fasterminmax(a::T, b::T) where {T<:FastFloat} =
    (min(a, b), max(a, b))

fastermaxmin(a::T, b::T) where {T<:FastFloat} =
    (max(a, b), min(a, b))

#=
   their utilzation
=#

min(x::Robust32, y::Robust32) = Rob32(fastermin(value64(x), value64(y)))
max(x::Robust32, y::Robust32) = Rob32(fastermax(value64(x), value64(y)))
minmax(x::Robust32, y::Robust32) = Rob32(fasterminmax(value64(x), value64(y)))
maxmin(x::Robust32, y::Robust32) = Rob32(fastmaxmin(value64(x), value64(y)))

min(x::Robust32, y::Float64) = Rob32(fastermin(value64(x), y))
max(x::Robust32, y::Float64) = Rob32(fastermax(value64(x), y))
minmax(x::Robust32, y::Float64) = Rob32(fasterminmax(value64(x), y))
maxmin(x::Robust32, y::Float64) = Rob32(fasterminmax(value64(x), y))

min(x::Float64, y::Robust32) = Rob32(fastermin(x, value64(y)))
max(x::Float64, y::Robust32) = Rob32(fastermax(x, value64(y)))
minmax(x::Float64, y::Robust32) = Rob32(fasterminmax(x, value64(y)))
maxmin(x::Float64, y::Robust32) = Rob32(fasterminmax(x, value64(y)))

min(x::Robust32, y::Real) = Rob32(fastermin(value64(x), Float64(y)))
max(x::Robust32, y::Real) = Rob32(fastermax(value64(x), Float64(y)))
minmax(x::Robust32, y::Real) = Rob32(fasterminmax(value64(x), Float64(y)))
maxmin(x::Robust32, y::Real) = Rob32(fasterminmax(value64(x), Float64(y)))

min(x::Real, y::Robust32) = Rob32(fastermin(Float64(x), value64(y)))
max(x::Real, y::Robust32) = Rob32(fastermax(Float64(x), value64(y)))
minmax(x::Real, y::Robust32) = Rob32(fasterminmax(Float64(x), value64(y)))
maxmin(x::Real, y::Robust32) = Rob32(fasterminmax(Float64(x), value64(y)))


@withkw @trait immutable IFParameter
    τm::Float = 20ms
    τe::Float = 5ms
    τi::Float = 10ms
    Vt::Float = -50mV
    Vr::Float = -60mV
    El::Float = Vr
end

@withkw @trait type IF
    param::IFParameter = IFParameter()
    N::Int = 100
    v::Vector{Float} = param.Vr + rand(N) * (param.Vt - param.Vr)
    ge::Vector{Float} = zeros(N)
    gi::Vector{Float} = zeros(N)
    fire::Vector{Bool} = zeros(Bool, N)
    I::Vector{Float} = zeros(N)
    records::Dict = Dict()
end

@replace function integrate!(p::IF, param::IFParameter, dt::Float)
    @inbounds for i = 1:N
        v[i] += dt * (ge[i] + gi[i] - (v[i] - El) + I[i]) / τm
        ge[i] += dt * -ge[i] / τe
        gi[i] += dt * -gi[i] / τi
    end
    @inbounds for i = 1:N
        fire[i] = v[i] > Vt
        v[i] = ifelse(fire[i], Vr, v[i])
    end
end


using Turing
using Turkie
using OnlineStats
using GLMakie
@model function demo(x)
    v ~ InverseGamma(3, 2)
    s ~ InverseGamma(2, v)
    m ~ Normal(0, √s)
    for i in eachindex(x)
        x[i] ~ Normal(m, √s)
    end
end

xs = randn(100) .+ 1;
m = demo(xs);

cb = TurkieCallback(m);
# chain = sample(m,  HMC(0.5, 10), 40; callback = cb);
chain = sample(m,  NUTS(0.65), 500; callback = cb);

# using Makie
# record(cb.scene, joinpath(@__DIR__, "video.gif")) do io
#     addIO!(cb, io)
#     sample(m,  NUTS(0.65), 300; callback = cb)
# end


using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra
using SparseArrays

@testset "MSLP" begin

    target=0;
    nev=3;
    nep=nep_gallery("beam",5);

    # Full matrix test
    λ1,v1=mslp(nep)
    @test norm(compute_Mlincomb(nep,λ1,v1))<1e-10

    Av=[sparse(nep.A[1]),sparse(nep.A[2])];
    nep2 = DEP(nep.n, Av, nep.tauv)
    # Sparse matrix test
    λ2,v2 = mslp(nep2)
    @test norm(compute_Mlincomb(nep2,λ2,v2)) < 1e-10

    @test λ1≈λ2 # They should be exactly the same

    @info "mslp noconv + double"
    @bench @testset  "mslp + double" begin
        nep3 = nep_gallery("dep_double");
        @test_throws NoConvergenceException mslp(nep3, λ=9im, maxit=10)
        λ,v = mslp(nep3, λ=9im, maxit=100,errmeasure=ResidualErrmeasure(nep3))
        @test norm(compute_Mlincomb(nep3, λ, v)) < eps()*100
    end
end