export minres_iterable, minres, minres!
using Printf
import LinearAlgebra: BLAS.axpy!, givensAlgorithm
import Base: iterate

mutable struct MINRESIterable{matT, solT, vecT <: DenseVector, smallVecT <: DenseVector, rotT <: Number, realT <: Real}
    A::matT
    skew_hermitian::Bool
    x::solT

    # Krylov basis vectors
    v_prev::vecT
    v_curr::vecT
    v_next::vecT

    # W = R * inv(V) is computed using 3-term recurrence
    w_prev::vecT
    w_curr::vecT
    w_next::vecT

    # Vector of size 4, holding the active column of the Hessenberg matrix
    # rhs is just two active values of the right-hand side.
    H::smallVecT
    rhs::smallVecT

    # Some Givens rotations
    c_prev::rotT
    s_prev::rotT
    c_curr::rotT
    s_curr::rotT

    # Bookkeeping
    mv_products::Int
    maxiter::Int
    tolerance::realT
    resnorm::realT
end

function minres_iterable!(x, A, b;
    initially_zero::Bool = false,
    skew_hermitian::Bool = false,
    tol = sqrt(eps(real(eltype(b)))),
    maxiter = size(A, 2)
)
    T = eltype(x)
    HessenbergT = skew_hermitian ? T : real(T)

    v_prev = similar(x)
    v_curr = similar(x)
    copyto!(v_curr, b)
    v_next = similar(x)
    w_prev = similar(x)
    w_curr = similar(x)
    w_next = similar(x)

    mv_products = 0

    # For nonzero x's, we must do an MV for the initial residual vec
    if !initially_zero
        # Use v_next to store Ax; v_next will soon be overwritten.
        mul!(v_next, A, x)
        axpy!(-one(T), v_next, v_curr)
        mv_products = 1
    end

    resnorm = norm(v_curr)
    reltol = resnorm * tol

    # Last active column of the Hessenberg matrix
    # and last two entries of the right-hand side
    H = zeros(HessenbergT, 4)
    rhs = [resnorm; zero(HessenbergT)]

    # Normalize the first Krylov basis vector
    rmul!(v_curr, inv(resnorm))

    # Givens rotations
    c_prev, s_prev = one(T), zero(T)
    c_curr, s_curr = one(T), zero(T)

    MINRESIterable(
        A, skew_hermitian, x,
        v_prev, v_curr, v_next,
        w_prev, w_curr, w_next,
        H, rhs,
        c_prev, s_prev, c_curr, s_curr,
        mv_products, maxiter, reltol, resnorm
    )
end

converged(m::MINRESIterable) = m.resnorm ≤ m.tolerance

start(::MINRESIterable) = 1

done(m::MINRESIterable, iteration::Int) = iteration > m.maxiter || converged(m)

function iterate(m::MINRESIterable, iteration::Int=start(m))
    if done(m, iteration) return nothing end

    # v_next = A * v_curr - H[2] * v_prev
    mul!(m.v_next, m.A, m.v_curr)

    iteration > 1 && axpy!(-m.H[2], m.v_prev, m.v_next)

    # Orthogonalize w.r.t. v_curr
    proj = dot(m.v_curr, m.v_next)
    m.H[3] = m.skew_hermitian ? proj : real(proj)
    axpy!(-proj, m.v_curr, m.v_next)

    # Normalize
    m.H[4] = norm(m.v_next)
    rmul!(m.v_next, inv(m.H[4]))

    # Rotation on H[1] and H[2]. Note that H[1] = 0 initially
    if iteration > 2
        m.H[1] = m.s_prev * m.H[2]
        m.H[2] = m.c_prev * m.H[2]
    end

    # Rotation on H[2] and H[3]
    if iteration > 1
        tmp = -conj(m.s_curr) * m.H[2] + m.c_curr * m.H[3]
        m.H[2] = m.c_curr * m.H[2] + m.s_curr * m.H[3]
        m.H[3] = tmp
    end

    # Next rotation
    c, s, m.H[3] = givensAlgorithm(m.H[3], m.H[4])

    # Apply as well to the right-hand side
    m.rhs[2] = -conj(s) * m.rhs[1]
    m.rhs[1] = c * m.rhs[1]

    # Update W = V * inv(R). Two axpy's can maybe be one MV.
    copyto!(m.w_next, m.v_curr)
    iteration > 1 && axpy!(-m.H[2], m.w_curr, m.w_next)
    iteration > 2 && axpy!(-m.H[1], m.w_prev, m.w_next)
    rmul!(m.w_next, inv(m.H[3]))

    # Update solution x
    axpy!(m.rhs[1], m.w_next, m.x)

    # Move on: next -> curr, curr -> prev
    m.v_prev, m.v_curr, m.v_next = m.v_curr, m.v_next, m.v_prev
    m.w_prev, m.w_curr, m.w_next = m.w_curr, m.w_next, m.w_prev
    m.c_prev, m.s_prev, m.c_curr, m.s_curr = m.c_curr, m.s_curr, c, s
    m.rhs[1] = m.rhs[2]

    # Due to symmetry of the tri-diagonal matrix
    m.H[2] = m.skew_hermitian ? -m.H[4] : m.H[4]

    # The approximate residual is cheaply available
    m.resnorm = abs(m.rhs[2])

    m.resnorm, iteration + 1
end

"""
    minres!(x, A, b; kwargs...) -> x, [history]

Solve Ax = b for (skew-)Hermitian matrices A using MINRES.

# Arguments

- `x`: initial guess, will be updated in-place;
- `A`: linear operator;
- `b`: right-hand side.

## Keywords

- `initially_zero::Bool = false`: if `true` assumes that `iszero(x)` so that one
  matrix-vector product can be saved when computing the initial
  residual vector;
- `skew_hermitian::Bool = false`: if `true` assumes that `A` is skew-symmetric or skew-Hermitian;
- `tol`: tolerance for stopping condition `|r_k| / |r_0| ≤ tol`. Note that the residual is computed only approximately;
- `maxiter::Int = size(A, 2)`: maximum number of iterations;
- `Pl`: left preconditioner;
- `Pr`: right preconditioner;
- `log::Bool = false`: keep track of the residual norm in each iteration;
- `verbose::Bool = false`: print convergence information during the iterations.

# Return values

**if `log` is `false`**

- `x`: approximate solution.

**if `log` is `true`**

- `x`: approximate solution;
- `history`: convergence history.
"""
function minres!(x, A, b;
    skew_hermitian::Bool = false,
    verbose::Bool = false,
    log::Bool = false,
    tol = sqrt(eps(real(eltype(b)))),
    maxiter::Int = size(A, 2),
    initially_zero::Bool = false
)
    history = ConvergenceHistory(partial = !log)
    history[:tol] = tol
    log && reserve!(history, :resnorm, maxiter)

    iterable = minres_iterable!(x, A, b;
        skew_hermitian = skew_hermitian,
        tol = tol,
        maxiter = maxiter,
        initially_zero = initially_zero
    )

    if log
        history.mvps = iterable.mv_products
    end

    for (iteration, resnorm) = enumerate(iterable)
        if log
            nextiter!(history, mvps = 1)
            push!(history, :resnorm, resnorm)
        end
        verbose && @printf("%3d\t%1.2e\n", iteration, resnorm)
    end

    verbose && println()
    log && setconv(history, converged(iterable))
    log && shrink!(history)

    log ? (iterable.x, history) : iterable.x
end

"""
    minres(A, b; kwargs...) -> x, [history]

Same as [`minres!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
minres(A, b; kwargs...) = minres!(zerox(A, b), A, b; initially_zero = true, kwargs...)


using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots
using DiffEqSensitivity
using Zygote
using ForwardDiff
using LinearAlgebra
using Random
using Statistics
using ProgressBars, Printf
using Flux.Optimise: update!
using Flux.Losses: mae, mse
using BSON: @save, @load

is_restart  = false;
n_epoch     = 50000;
ntotal      = 50
n_plot      = 200;

batch_size  = 50;
opt         = ADAMW(0.005, (0.9, 0.999), 1.f-6);
lb          = 1.e-6;
ode_solver  = AutoTsit5(Rosenbrock23(autodiff=false));

function rober!(du, u, k, t)
    y1, y2, y3  = u
    k1, k2, k3  = k
    du[1]       = -k1 * y1 + k3 * y2 * y3
    du[2]       =  k1 * y1 - k3 * y2 * y3 - k2 * y2^2
    du[3]       =  k2 * y2^2
end

u0          = [1.0, 0, 0];
tspan       = (0.0, 1e5);
t_end       = tspan[2];
k           = [0.04, 3e7, 1e4];
tsteps      = 10 .^ (range(-5, log10(tspan[2]), length=ntotal));
prob_rober  = ODEProblem(rober!, u0, tspan, k);
sol_rober   = solve(prob_rober, ode_solver, saveat=tsteps, atol=1.f-8);
normdata    = Array(sol_rober)

yscale      = maximum(normdata, dims = 2)  #scale for each species

i_slow      = [1, 2, 3];
nslow       = length(i_slow);
node        = 5;
dudt2       = Chain(x -> x,
                  Dense(nslow, node, gelu),
                  Dense(node, node, gelu),
                  Dense(node, node, gelu),
                  Dense(node, node, gelu),
                  Dense(node, node, gelu),
                  Dense(node, node, gelu),
                  Dense(node, nslow))

p, re       = Flux.destructure(dudt2);
rep         = re(p)

yscale_     = yscale[:, 1]
function dudt!(du, u, p, t)
    du      .= rep(u) .* yscale_ /t_end
end

prob        = ODEProblem(dudt!, u0[i_slow], tspan)
sense       = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());
function predict_n_ode(p, sample)
    global rep  = re(p)
    _prob       = remake(prob, p=p, tspan=[0, tsteps[sample]])
    pred        = Array(solve(_prob, ode_solver, saveat=tsteps[1:sample], atol=lb, sensalg=sense))
end
pred        = predict_n_ode(p, ntotal)

function loss_n_ode(p, sample=ntotal)
    pred = predict_n_ode(p, sample)
    loss = mae(pred ./ yscale, normdata[i_slow, 1:size(pred)[2]] ./ yscale)
    return loss
end
loss_n_ode(p, ntotal)

list_loss = []
list_grad = []
iter = 1
cb = function (p, loss_mean, g_norm)
    global list_loss, list_grad, iter
    push!(list_loss, loss_mean)
    push!(list_grad, g_norm)

    if iter % n_plot == 0
        pred = predict_n_ode(p, ntotal)

        list_plt = []
        for i in 1:nslow
            j = i_slow[i]
            plt = scatter(tsteps[2:end], normdata[j,2:end], xscale=:log10, label="data")
            plot!(plt, tsteps[2:end], pred[i,2:end], lw=2, xscale=:log10, label="pred", framestyle=:box)
            plot!(plt, xtickfontsize=11, ytickfontsize=11, xguidefontsize=12, yguidefontsize=12)
            ylabel!(plt, "y$j")
            xlabel!(plt, "Time")
            if i == 1
                plot!(plt, legend=:best)
            else
                plot!(plt, legend=false)
            end
            push!(list_plt, plt)
        end
        plt_all = plot(list_plt..., layouts = (nslow, 1))
        png(plt_all, "figs/pred.png")

        plt_loss = plot(list_loss, xscale=:identity, yscale=:log10, label="loss", legend=:topright)
        plt_grad = plot(list_grad, xscale=:identity, yscale=:log10, label="grad_norm", legend=:bottomright)
        xlabel!(plt_loss, "Epoch")
        xlabel!(plt_grad, "Epoch")
        ylabel!(plt_loss, "Loss")
        ylabel!(plt_grad, "Gradient Norm")
        plt_all = plot([plt_loss, plt_grad]..., framestyle=:box, size = (1000, 400))
        plot!(plt_all, xtickfontsize=10, ytickfontsize=10, xguidefontsize=12, yguidefontsize=12)
        png(plt_all, "figs/loss_grad")

        @save "./checkpoint/mymodel.bson" p opt list_loss list_grad iter
    end
    iter += 1
    return false
end

if is_restart
    @load "./checkpoint/mymodel.bson" p opt list_loss list_grad iter
    iter += 1
end

epochs = ProgressBar(iter:n_epoch);
for epoch in epochs
    global p
    sample = rand(batch_size:ntotal)

    loss = loss_n_ode(p, sample)
    grad = ForwardDiff.gradient(x -> loss_n_ode(x, sample), p)

    grad_norm = norm(grad, 2)
    update!(opt, p, grad)

    set_description(epochs, string(@sprintf("Loss: %.4e grad: %.2e", loss, grad_norm)))
    cb(p, loss, grad_norm)
end


# Figure 5.3a in "Infinite GMRES for parameterized linear systems"

using SparseArrays
using LinearAlgebra
using PyPlot
using MAT
pygui(true)
include("iar.jl");
include("infgmres.jl");

# Problem represented with explicitly available derivatives
struct GeneralProblem
    sigma
    Aders
    Afun
    b
end

function compute_Mlincomb(prob::GeneralProblem,sigma,Y)
    if (prob.sigma != sigma)
        @show sigma
        @show prob.sigma
        error("Wrong sigma");
    end
    z=zeros(eltype(Y),size(Y,1));
    for i=1:size(Y,2)
        z += prob.Aders[i]*Y[:,i];
    end
    return z;
end
function compute_Mder(prob::GeneralProblem,mu)
    return prob.Afun(mu);
end

## Define problem
data=matread("delay_problem_matrices.mat")
A0=data["A0"];
A1=data["A1"];
b=vec(data["b"]);
n=length(b);

# Identity matrix
II=Matrix{Float64}(I,n,n);

Afun= s-> -s*II+A0+A1*exp(-s);
sigma=0;
N=100;
ders=Vector{Matrix{Float64}}(undef,N);

for i=1:N
    ders[i]=((-1)^(i+1))*A1;
end
ders[1] += A0;
ders[2] += -II;

myproblem=GeneralProblem(sigma,ders,Afun,b);

## Run Infinite GMRES
(f,hist1)=infgmres(myproblem,30,method=:iar);
mu_list = range(1.e-5,stop=1.e-1,length=20);
mu_list = mu_list[2:end];
how_many = zeros(length(mu_list),1);
tol = 1.e-12;

for i=1:length(mu_list)
    global mu
    mu = mu_list[i];
    A_of_mu = myproblem.Afun(mu);
    res = zeros(100,1);
    for m=1:200
        global f
        (xx)=f(mu,m)
        res[m] = norm(A_of_mu*xx - myproblem.b)/norm(myproblem.b);
        if res[m] < tol
            how_many[i] = m;
            break;
        end
    end
end

figure(1)
ax = PyPlot.gca()
xlabel("mu")
ylabel("its")
plot(mu_list,how_many,"ro")


# 2D Heat Diffusion
# LBM
#
# Generalize version on complicated mesh, with 1D indexing of lattices.
#
# Hongyang Zhou, hyzhou@umich.edu 05/01/2019

using PyPlot
include("Parameters.jl")
include("lattice.jl")
include("BC.jl")

using .Parameters, .BC, .lattice

function solveHeatDiffusion2D!(param::Param,lat::Lattice,mesh::Mesh)

# Set aliases
ω,Δt,TotalTime = param.ω,param.Δt,param.TotalTime
BCtype = param.BCtype
BCDirichlet,BCRiemann  = Parameters.BCDirichlet, Parameters.BCRiemann
Tleft,Tright,Ttop,Tbot = param.Tleft, param.Tright, param.Ttop, param.Tbot

nw,w,nN  = lat.nw, lat.w, lat.nN
T   = zeros(1,nN)
feq = Array{Float64,2}(undef,nw,nN)
# There will be an issue if I allocate with undef due to values on the boundary.
ftemp = zeros(nw,nN)
t   = 0.0     # Current simulation time

# Advance
while t < TotalTime

   # Calculate equilibrium distribution
   for iw = 1:nw
      @. feq[iw,:] = w[iw]*T[1,:]
   end

   # Collision
   for iw = 1:nw
      @. lat.f[iw,:] = (1.0-ω)*lat.f[iw,:] + ω.*feq[iw,:]
   end

   # Streaming
   for i=1:nN, iw = 1:nw
      iN = lat.IJ[iw,i]
      (iN > nN || iN < 1) && continue
      ftemp[iw,iN] = lat.f[iw,i]
   end
   #copy!(lat.f,ftemp) # This won't work in Julia 1.0!
   lat.f .= ftemp

   # Remaining BCs
   setBCs!(lat,param,mesh)

   # Calculate macroscopic T
   T .= sum(lat.f, dims=1)

   t += Δt
end

end

function init()

param,mesh = setParameters()

Δx,Δt,LT = param.Δx,param.Δt,param.LatticeType

if LT == 1
   lat = LatticeD2Q4(Δx,Δt,mesh)
elseif LT == 2
   lat = LatticeD2Q5(Δx,Δt,mesh)
elseif LT == 3
   lat = LatticeD2Q9(Δx,Δt,mesh)
end

return param,lat,mesh

end

function plotSol(lat::Lattice,mesh::Mesh)
# Visualization
#x = range(0,stop=lat.nI-1,length=lat.nI)
#y = range(0,stop=lat.nJ-1,length=lat.nJ)
T = sum(lat.f, dims=1)

tricontourf(mesh.coord[1,:],mesh.coord[2,:],T[1,:]); colorbar()
title("2D Heat Diffusion in a Plate")
xlabel("x")
ylabel("y")
grid(true)

# 1D cut plot
#plot(0:nx-1,T[1,:,50],linestyle="-",marker="o",label="LBM")
#xlim(0,30)
end

function main()
   param,lat,mesh = init()

   solveHeatDiffusion2D!(param,lat,mesh)

   if param.DoPlot
      plotSol(lat,mesh)
   end

end


@doc raw"""
    gradient_descent(M, F, gradF, x)

perform a gradient descent

```math
x_{k+1} = \operatorname{retr}_{x_k}\bigl( s_k\operatorname{grad}f(x_k) \bigr)
```

with different choices of ``s_k`` available (see `stepsize` option below).

# Input
* `M` – a manifold ``\mathcal M``
* `F` – a cost function ``F: \mathcal M→ℝ`` to minimize
* `gradF` – the gradient ``\operatorname{grad}F: \mathcal M → T\mathcal M`` of F
* `x` – an initial value ``x ∈ \mathcal M``

# Optional
* `evaluation` – ([`AllocatingEvaluation`](@ref)) specify whether the gradient works by allocation (default) form `gradF(M, x)`
  or [`MutatingEvaluation`](@ref) in place, i.e. is of the form `gradF!(M, X, x)`.
* `retraction_method` – (`default_retraction_method(M)`) a `retraction(M,x,ξ)` to use.
* `return_options` – (`false`) – if activated, the extended result, i.e. the
    complete [`Options`](@ref) are returned. This can be used to access recorded values.
    If set to false (default) just the optimal value `x_opt` if returned
* `stepsize` – ([`ConstantStepsize`](@ref)`(1.)`) specify a [`Stepsize`](@ref)
  functor.
* `stopping_criterion` – ([`StopWhenAny`](@ref)`(`[`StopAfterIteration`](@ref)`(200), `[`StopWhenGradientNormLess`](@ref)`(10.0^-8))`)
  a functor inheriting from [`StoppingCriterion`](@ref) indicating when to stop.
...
and the ones that are passed to [`decorate_options`](@ref) for decorators.

# Output
* `x_opt` – the resulting (approximately critical) point of gradientDescent
OR
* `options` - the options returned by the solver (see `return_options`)
"""
function gradient_descent(
    M::AbstractManifold, F::TF, gradF::TDF, x; kwargs...
) where {TF,TDF}
    x_res = allocate(x)
    copyto!(M, x_res, x)
    return gradient_descent!(M, F, gradF, x_res; kwargs...)
end
@doc raw"""
    gradient_descent!(M, F, gradF, x)

perform a gradient_descent ``x_{k+1} = \mathrm{retr}_{x_k} s_k\operatorname{grad}f(x_k)`` in place of `x`
with different choices of ``s_k`` available.

# Input
* `M` – a manifold ``\mathcal M``
* `F` – a cost function ``F:\mathcal M→ℝ`` to minimize
* `gradF` – the gradient ``\operatorname{grad}F:\mathcal M→ T\mathcal M`` of F
* `x` – an initial value ``x ∈ \mathcal M``

For more options, especially [`Stepsize`](@ref)s for ``s_k``, see [`gradient_descent`](@ref)
"""
function gradient_descent!(
    M::AbstractManifold,
    F::TF,
    gradF::TDF,
    x;
    stepsize::Stepsize=ConstantStepsize(1.0),
    retraction_method::AbstractRetractionMethod=default_retraction_method(M),
    stopping_criterion::StoppingCriterion=StopAfterIteration(200) |
                                          StopWhenGradientNormLess(10.0^-8),
    direction=IdentityUpdateRule(),
    evaluation::AbstractEvaluationType=AllocatingEvaluation(),
    return_options=false,
    kwargs..., #collect rest
) where {TF,TDF}
    p = GradientProblem(M, F, gradF; evaluation=evaluation)
    o = GradientDescentOptions(
        M,
        x;
        stopping_criterion=stopping_criterion,
        stepsize=stepsize,
        direction=direction,
        retraction_method=retraction_method,
    )
    o = decorate_options(o; kwargs...)
    resultO = solve(p, o)
    if return_options
        return resultO
    else
        return get_solver_result(resultO)
    end
end
#
# Solver functions
#
function initialize_solver!(p::GradientProblem, o::GradientDescentOptions)
    o.gradient = get_gradient(p, o.x)
    return o
end
function step_solver!(p::GradientProblem, o::GradientDescentOptions, iter)
    s, o.gradient = o.direction(p, o, iter)
    retract!(p.M, o.x, o.x, -s * o.gradient, o.retraction_method)
    return o
end
get_solver_result(o::GradientDescentOptions) = o.x


using DistributedArrays, Krylov, LinearOperators, MatrixMarket

function residuals(A, b, shifts, x)
  nshifts = size(shifts, 1);
  r = [ (b - A * x[:,i] - shifts[i] * x[:,i]) for i = 1 : nshifts ];
  return r;
end

# Parallel reduce.
preduce(func, darray) = reduce(func,
                               map(fetch,
                                   [ (@spawnat p reduce(func, localpart(darray))) for p = procs(darray) ] ));

"""The Lanczos version of the conjugate gradient method to solve a family
of shifted systems

  (A + αI) x = b  (α = α₁, α₂, ...)

The method does _not_ abort if A + αI is not definite.

In this version, the shifted systems are dispatched on the available processors,
and operations specific to each shift is carried out on the processor hosting it.
"""
function cg_lanczos_shift_par{Tb <: Real, Ts <: Real}(A :: AbstractLinearOperator, b :: Array{Tb,1}, shifts :: Array{Ts,1};
                                                      atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, itmax :: Int=0,
                                                      check_curvature :: Bool=false, verbose :: Bool=false)

  n = size(b, 1);
  (size(A, 1) == n & size(A, 2) == n) || error("Inconsistent problem size");

  dshifts = distribute(shifts);
  nshifts = size(shifts, 1);
  nchunks = length(dshifts.indexes)
  verbose && @printf("CG Lanczos: system of %d equations in %d variables with %d shifts\n", n, n, nshifts);

  # Initial state.
  ## Distribute x similarly to shifts.
  dx = dzeros((n, nshifts), workers(), [1, nchunks]);
  β = norm(b);
  β == 0 && return convert(Array, dx),
            Krylov.LanczosStats(true, zeros(nshifts), false, 0.0, 0.0, "x = 0
                                is a zero-residual solution")
  v = b / β;
  v_prev = v;

  # Distribute p similarly to shifts.
  # Initialize each p to b.
  dp = dzeros((n, nshifts), workers(), [1, nchunks]);
  @sync for proc in procs(dp)
    @spawnat proc begin
      p_loc = localpart(dp);
      for i = 1 : size(p_loc, 2)
        p_loc[:,i] = b;
      end
    end
  end

  # Keep track of shifted systems that have converged.
  dsolved = dfill(false, (nshifts,), workers(), [nchunks]);
  dconverged = dfill(false, (nshifts,), workers(), [nchunks]);
  dindefinite= dfill(false, (nshifts,), workers(), [nchunks]);
  iter = 0;
  itmax == 0 && (itmax = 2 * n);

  # Initialize some constants used in recursions below.
  dσ = dfill(β, (nshifts,), workers(), [nchunks]);
  dδ = dzeros((nshifts,), workers(), [nchunks]);
  dω = dzeros((nshifts,), workers(), [nchunks]);
  dγ = dones((nshifts,), workers(), [nchunks]);

  # Define stopping tolerance.
  drNorm = dfill(β, (nshifts,), workers(), [nchunks]);
  ε = atol + rtol * β;

  # Build format strings for printing.
  if verbose
    fmt = "%5d" * repeat("  %8.1e", nshifts) * "\n";
    c_printf(fmt, iter, drNorm...);
  end

  solved = false;
  tired = iter >= itmax;
  status = "unknown";

  # Main loop.
  while ! (solved || tired)
    # Form next Lanczos vector.
    v_next = A * v;
    δ = dot(v, v_next);
    BLAS.axpy!(n, -δ, v, 1, v_next, 1);  # Faster than v_next = v_next - δ * v;
    if iter > 0
      BLAS.axpy!(n, -β, v_prev, 1, v_next, 1);  # Faster than v_next = v_next - β * v_prev;
      v_prev = v;
    end
    β = norm(v_next);
    v = v_next / β;

    # Compute next CG iterate for each shift.
    @sync for proc in procs(dp)
      @spawnat proc begin
        solved_loc = localpart(dsolved);
        converged_loc = localpart(dconverged);
        indefinite_loc = localpart(dindefinite);
        shifts_loc = localpart(dshifts);
        δ_loc = localpart(dδ);
        γ_loc = localpart(dγ);
        ω_loc = localpart(dω);

        rNorm_loc = localpart(drNorm);

        # Check curvature: v'(A + sᵢI)v = v'Av + sᵢ ‖v‖² = δ + sᵢ because ‖v‖ = 1.
        # It is possible to show that σⱼ² (δⱼ + sᵢ - ωⱼ₋₁ / γⱼ₋₁) = pⱼᵀ (A + sᵢ I) pⱼ.
        # Stop iterating on indefinite problems if requested.
        δ_loc[:] = δ + shifts_loc
        γ_loc[:] = 1 ./ (δ_loc - ω_loc ./ γ_loc)
        indefinite_loc[:] |= (γ_loc .<= 0.0);
        not_cv = check_curvature ? find(! (converged_loc | indefinite_loc)) : find(! converged_loc);

        if length(not_cv) > 0
          # Fetch parts of relevant arrays for which the residual has not yet converged.
          σ_loc = localpart(dσ);
          x_loc = localpart(dx);
          p_loc = localpart(dp);
          shifts_loc = shifts_loc[not_cv];

          ω_loc[not_cv] = β * γ_loc[not_cv];
          σ_loc[not_cv] .*= -ω_loc[not_cv];
          ω_loc[not_cv] .*= ω_loc[not_cv];
          rNorm_loc[not_cv] = abs(σ_loc[not_cv]);

          # Can't seem to use BLAS3 GEMM calls here.
          x_loc[:, not_cv] += (p_loc[:, not_cv] * diagm(γ_loc[not_cv]));  # diagm?!
          p_loc[:, not_cv] = (v * σ_loc[not_cv]' + p_loc[:, not_cv] * diagm(ω_loc[not_cv]));

          # Update list of systems that have converged.
          converged_loc[not_cv] = rNorm_loc[not_cv] .<= ε;
        end

        solved_loc[:] = converged_loc | indefinite_loc;
      end
    end

    iter = iter + 1;
    verbose && c_printf(fmt, iter, drNorm...);

    solved = preduce(&, dsolved);
    tired = iter >= itmax;
  end

  status = tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"
  stats = Krylov.LanczosStats(solved, convert(Array, drNorm), convert(Array, dindefinite), 0.0, 0.0, status);
  return (dx, stats);
end

cg_lanczos_shift_par{TA <: Number, Tb <: Number, Ts <: Number}(A :: Array{TA,2}, b :: Array{Tb,1}, shifts :: Array{Ts,1};
                                                               atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, itmax :: Int=0,
                                                               check_curvature :: Bool=false, verbose :: Bool=false) =
  cg_lanczos_shift_par(LinearOperator(A), b, shifts, atol=atol, rtol=rtol, itmax=itmax, check_curvature=check_curvature, verbose=verbose);

cg_lanczos_shift_par{TA <: Number, Tb <: Number, Ts <: Number, IA <: Integer}(A :: SparseMatrixCSC{TA,IA}, b :: Array{Tb,1}, shifts :: Array{Ts,1};
                                                                             atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, itmax :: Int=0,
                                                                             check_curvature :: Bool=false, verbose :: Bool=false) =
  cg_lanczos_shift_par(LinearOperator(A), b, shifts, atol=atol, rtol=rtol, itmax=itmax, check_curvature=check_curvature, verbose=verbose);

# mtx = "data/1138bus.mtx";
mtx = "data/bcsstk09.mtx";

A = MatrixMarket.mmread(mtx);
n = size(A, 1);
b = ones(n); b_norm = norm(b);

# Define a linear operator with preallocation.
Ap = zeros(n);
op = LinearOperator(n, n, true, true, p -> A_mul_B!(1.0,  A, p, 0.0, Ap))

# Solve (A+αI)x = b in parallel.
shifts = [1, 2, 3, 4];
(x, stats) = cg_lanczos_shift_par(op, b, shifts, verbose=false);
# @profile (x, stats) = cg_lanczos_shift_par(op, b, shifts);
@time (x, stats) = cg_lanczos_shift_par(op, b, shifts, verbose=false);
r = residuals(A, b, shifts, convert(Array, x));
resids = map(norm, r) / b_norm;
@printf("Relative residuals with shifts:\n");
for resid in resids
  @printf(" %8.1e", resid);
end
@printf("\n");


#### Specialized matrix types ####

## Hermitian tridiagonal matrices
immutable SymTridiagonal{T} <: AbstractMatrix{T}
    dv::Vector{T}                        # diagonal
    ev::Vector{T}                        # subdiagonal
    function SymTridiagonal(dv::Vector{T}, ev::Vector{T})
        length(ev)==length(dv)-1 || throw(DimensionMismatch(""))
        new(dv,ev)
    end
end

SymTridiagonal{T}(dv::Vector{T}, ev::Vector{T}) = SymTridiagonal{T}(dv, ev)

function SymTridiagonal{Td,Te}(dv::Vector{Td}, ev::Vector{Te})
    T = promote_type(Td,Te)
    SymTridiagonal(convert(Vector{T}, dv), convert(Vector{T}, ev))
end

SymTridiagonal(A::AbstractMatrix) = diag(A,1)==diag(A,-1)?SymTridiagonal(diag(A), diag(A,1)):throw(DimensionMismatch("matrix is not symmetric; cannot convert to SymTridiagonal"))
full{T}(M::SymTridiagonal{T}) = convert(Matrix{T}, M)
convert{T}(::Type{SymTridiagonal{T}}, S::SymTridiagonal) = SymTridiagonal(convert(Vector{T}, S.dv), convert(Vector{T}, S.ev))
convert{T}(::Type{AbstractMatrix{T}}, S::SymTridiagonal) = SymTridiagonal(convert(Vector{T}, S.dv), convert(Vector{T}, S.ev))
convert{T}(::Type{Matrix{T}}, M::SymTridiagonal{T})=diagm(M.dv)+diagm(M.ev,-1)+conj(diagm(M.ev,1))

size(m::SymTridiagonal) = (length(m.dv), length(m.dv))
size(m::SymTridiagonal, d::Integer) = d<1 ? error("dimension out of range") : (d<=2 ? length(m.dv) : 1)

#Elementary operations
for func in (:copy, :round, :iround, :conj)
    @eval begin
        ($func)(M::SymTridiagonal) = SymTridiagonal(($func)(M.dv), ($func)(M.ev))
    end
end

transpose(M::SymTridiagonal) = M #Identity operation
ctranspose(M::SymTridiagonal) = conj(M)

function diag{T}(M::SymTridiagonal{T}, n::Integer=0)
    absn = abs(n)
    absn==0 ? M.dv : absn==1 ? M.ev : absn<size(M,1) ? zeros(T,size(M,1)-absn) : throw(BoundsError())
end

+(A::SymTridiagonal, B::SymTridiagonal) = SymTridiagonal(A.dv+B.dv, A.ev+B.ev)
-(A::SymTridiagonal, B::SymTridiagonal) = SymTridiagonal(A.dv-B.dv, A.ev-B.ev)
*(A::SymTridiagonal, B::Number) = SymTridiagonal(A.dv*B, A.ev*B)
*(B::Number, A::SymTridiagonal) = A*B
/(A::SymTridiagonal, B::Number) = SymTridiagonal(A.dv/B, A.ev/B)
==(A::SymTridiagonal, B::SymTridiagonal) = (A.dv==B.dv) && (A.ev==B.ev)

## Solver
function \{T<:BlasFloat}(M::SymTridiagonal{T}, rhs::StridedVecOrMat{T})
    if stride(rhs, 1) == 1
        return LAPACK.gtsv!(copy(M.ev), copy(M.dv), copy(M.ev), copy(rhs))
    end
    solve(Tridiagonal(M), rhs)  # use the Julia "fallback"
end

#Wrap LAPACK DSTE{GR,BZ} to compute eigenvalues
eigfact!{T<:BlasFloat}(m::SymTridiagonal{T}) = Eigen(LAPACK.stegr!('V', m.dv, m.ev)...)
eigfact!{T<:BlasFloat}(m::SymTridiagonal{T}, irange::UnitRange) = Eigen(LAPACK.stegr!('V', 'I', m.dv, m.ev, 0.0, 0.0, irange.start, irange.stop)...)
eigfact!{T<:BlasFloat}(m::SymTridiagonal{T}, vl::Real, vu::Real) = Eigen(LAPACK.stegr!('V', m.dv, m.ev, convert(T, vl), convert(T, vu), 0, 0)...)
eigvals!{T<:BlasFloat}(m::SymTridiagonal{T}) = LAPACK.stegr!('N', 'A', m.dv, m.ev, 0.0, 0.0, 0, 0)[1]
eigvals!{T<:BlasFloat}(m::SymTridiagonal{T}, irange::UnitRange) = LAPACK.stegr!('N', 'I', m.dv, m.ev, 0.0, 0.0, irange.start, irange.stop)[1]
eigvals!{T<:BlasFloat}(m::SymTridiagonal{T}, vl::Real, vu::Real) = LAPACK.stegr!('N', 'V', m.dv, m.ev, vl, vu, 0, 0)[1]

#Computes largest and smallest eigenvalue
eigmax(m::SymTridiagonal) = eigvals(m, size(m, 1), size(m, 1))[1]
eigmin(m::SymTridiagonal) = eigvals(m, 1, 1)[1]

#Compute selected eigenvectors only corresponding to particular eigenvalues
eigvecs(m::SymTridiagonal) = eigfact(m)[:vectors]
eigvecs{T<:BlasFloat,Eigenvalue<:Real}(m::SymTridiagonal{T}, eigvals::Vector{Eigenvalue}) = LAPACK.stein!(m.dv, m.ev, eigvals)

###################
# Generic methods #
###################

#Needed for inv_usmani()
type ZeroOffsetVector
    data::Vector
end
getindex (a::ZeroOffsetVector, i) = a.data[i+1]
setindex!(a::ZeroOffsetVector, x, i) = a.data[i+1]=x

#Implements the inverse using the recurrence relation between principal minors
# a, b, c are assumed to be the subdiagonal, diagonal, and superdiagonal of
# a tridiagonal matrix.
#Ref:
#    R. Usmani, "Inversion of a tridiagonal Jacobi matrix",
#    Linear Algebra and its Applications 212-213 (1994), pp.413-414
#    doi:10.1016/0024-3795(94)90414-6
function inv_usmani{T}(a::Vector{T}, b::Vector{T}, c::Vector{T})
    n = length(b)
    θ = ZeroOffsetVector(zeros(T, n+1)) #principal minors of A
    θ[0] = 1
    n>=1 && (θ[1] = b[1])
    for i=2:n
        θ[i] = b[i]*θ[i-1]-a[i-1]*c[i-1]*θ[i-2]
    end
    φ = zeros(T, n+1)
    φ[n+1] = 1
    n>=1 && (φ[n] = b[n])
    for i=n-1:-1:1
        φ[i] = b[i]*φ[i+1]-a[i]*c[i]*φ[i+2]
    end
    α = Array(T, n, n)
    for i=1:n, j=1:n
        sign = (i+j)%2==0 ? (+) : (-)
        if i<j
            α[i,j]=(sign)(prod(c[i:j-1]))*θ[i-1]*φ[j+1]/θ[n]
        elseif i==j
            α[i,i]=                       θ[i-1]*φ[i+1]/θ[n]
        else #i>j
            α[i,j]=(sign)(prod(a[j:i-1]))*θ[j-1]*φ[i+1]/θ[n]
        end
    end
    α
end

#Implements the determinant using principal minors
#Inputs and reference are as above for inv_usmani()
function det_usmani{T}(a::Vector{T}, b::Vector{T}, c::Vector{T})
    n = length(b)
    θa = one(T)
    n==0 && return θa
    θb = b[1]
    for i=2:n
        θb, θa = b[i]*θb-a[i-1]*c[i-1]*θa, θb
    end
    return θb
end

inv(A::SymTridiagonal) = inv_usmani(A.ev, A.dv, A.ev)
det(A::SymTridiagonal) = det_usmani(A.ev, A.dv, A.ev)

function getindex{T}(A::SymTridiagonal{T}, i::Integer, j::Integer)
    (1<=i<=size(A,2) && 1<=j<=size(A,2)) || throw(BoundsError())
    i==j ? A.dv[i] : i==j+1 ? A.ev[j] : i+1==j ? A.ev[i] : zero(T)
end

## Tridiagonal matrices ##
immutable Tridiagonal{T} <: AbstractMatrix{T}
    dl::Vector{T}    # sub-diagonal
    d::Vector{T}     # diagonal
    du::Vector{T}    # sup-diagonal
    du2::Vector{T}   # supsup-diagonal for pivoting
end
function Tridiagonal{T}(dl::Vector{T}, d::Vector{T}, du::Vector{T})
    n = length(d)
    if (length(dl) != n-1 || length(du) != n-1)
        error(string("Cannot make Tridiagonal from incompatible lengths of subdiagonal, diagonal and superdiagonal: (", length(dl), ", ", length(d), ", ", length(du),")"))
    end
    Tridiagonal(dl, d, du, zeros(T,n-2))
end
function Tridiagonal{Tl, Td, Tu}(dl::Vector{Tl}, d::Vector{Td}, du::Vector{Tu})
    Tridiagonal(map(v->convert(Vector{promote_type(Tl,Td,Tu)}, v), (dl, d, du))...)
end

size(M::Tridiagonal) = (length(M.d), length(M.d))
size(M::Tridiagonal, d::Integer) = d<1 ? error("dimension out of range") : (d<=2 ? length(M.d) : 1)

full{T}(M::Tridiagonal{T}) = convert(Matrix{T}, M)
function convert{T}(::Type{Matrix{T}}, M::Tridiagonal{T})
    A = zeros(T, size(M))
    for i = 1:length(M.d)
        A[i,i] = M.d[i]
    end
    for i = 1:length(M.d)-1
        A[i+1,i] = M.dl[i]
        A[i,i+1] = M.du[i]
    end
    A
end
function similar(M::Tridiagonal, T, dims::Dims)
    if length(dims) != 2 || dims[1] != dims[2]
        throw(DimensionMismatch("Tridiagonal matrices must be square"))
    end
    Tridiagonal{T}(similar(M.dl), similar(M.d), similar(M.du), similar(M.du2))
end

# Operations on Tridiagonal matrices
copy!(dest::Tridiagonal, src::Tridiagonal) = Tridiagonal(copy!(dest.dl, src.dl), copy!(dest.d, src.d), copy!(dest.du, src.du), copy!(dest.du2, src.du2))

#Elementary operations
for func in (:copy, :round, :iround, :conj)
    @eval begin
        ($func)(M::Tridiagonal) = Tridiagonal(map(($func), (M.dl, M.d, M.du, M.du2))...)
    end
end

transpose(M::Tridiagonal) = Tridiagonal(M.du, M.d, M.dl)
ctranspose(M::Tridiagonal) = conj(transpose(M))

diag{T}(M::Tridiagonal{T}, n::Integer=0) = n==0 ? M.d : n==-1 ? M.dl : n==1 ? M.du : abs(n)<size(M,1) ? zeros(T,size(M,1)-abs(n)) : throw(BoundsError())
function getindex{T}(A::Tridiagonal{T}, i::Integer, j::Integer)
    (1<=i<=size(A,2) && 1<=j<=size(A,2)) || throw(BoundsError())
    i==j ? A.d[i] : i==j+1 ? A.dl[j] : i+1==j ? A.du[i] : zero(T)
end

###################
# Generic methods #
###################

+(A::Tridiagonal, B::Tridiagonal) = Tridiagonal(A.dl+B.dl, A.d+B.d, A.du+B.du)
-(A::Tridiagonal, B::Tridiagonal) = Tridiagonal(A.dl-B.dl, A.d-B.d, A.du+B.du)
*(A::Tridiagonal, B::Number) = Tridiagonal(A.dl*B, A.d*B, A.du*B)
*(B::Number, A::Tridiagonal) = A*B
/(A::Tridiagonal, B::Number) = Tridiagonal(A.dl/B, A.d/B, A.du/B)

==(A::Tridiagonal, B::Tridiagonal) = (A.dl==B.dl) && (A.d==B.d) && (A.du==B.du)
==(A::Tridiagonal, B::SymTridiagonal) = (A.dl==A.du==B.ev) && (A.d==B.dv)
==(A::SymTridiagonal, B::Tridiagonal) = (B.dl==B.du==A.ev) && (B.d==A.dv)

inv(A::Tridiagonal) = inv_usmani(A.dl, A.d, A.du)
det(A::Tridiagonal) = det_usmani(A.dl, A.d, A.du)

# Elementary operations that mix Tridiagonal and SymTridiagonal matrices
convert(::Type{Tridiagonal}, A::SymTridiagonal) = Tridiagonal(A.ev, A.dv, A.ev)
+(A::Tridiagonal, B::SymTridiagonal) = Tridiagonal(A.dl+B.ev, A.d+B.dv, A.du+B.ev)
+(A::SymTridiagonal, B::Tridiagonal) = Tridiagonal(A.ev+B.dl, A.dv+B.d, A.ev+B.du)
-(A::Tridiagonal, B::SymTridiagonal) = Tridiagonal(A.dl-B.ev, A.d-B.dv, A.du-B.ev)
-(A::SymTridiagonal, B::Tridiagonal) = Tridiagonal(A.ev-B.dl, A.dv-B.d, A.ev-B.du)

convert{T}(::Type{Tridiagonal{T}},M::Tridiagonal) = Tridiagonal(convert(Vector{T}, M.dl), convert(Vector{T}, M.d), convert(Vector{T}, M.du), convert(Vector{T}, M.du2))
convert{T}(::Type{AbstractMatrix{T}},M::Tridiagonal) = convert(Tridiagonal{T}, M)
convert{T}(::Type{Tridiagonal{T}}, M::SymTridiagonal{T}) = Tridiagonal(M)
convert{T}(::Type{SymTridiagonal{T}}, M::Tridiagonal) = M.dl==M.du ? (SymTridiagonal(M.dl, M.d)) :
    error("Tridiagonal is not symmetric, cannot convert to SymTridiagonal")
convert{T}(::Type{SymTridiagonal{T}},M::SymTridiagonal) = SymTridiagonal(convert(Vector{T}, M.dv), convert(Vector{T}, M.ev))

function A_mul_B!(C::AbstractVecOrMat, A::Tridiagonal, B::AbstractVecOrMat)
    size(C,1) == size(B,1) == (nA = size(A,1)) || throw(DimensionMismatch(""))
    size(C,2) == (nB = size(B,2)) || throw(DimensionMismatch(""))
    l = A.dl
    d = A.d
    u = A.du
    @inbounds begin
        for j = 1:nB
            C[1,j] = d[1]*B[1,j] + u[1]*B[2,j]
            for i = 2:nA-1
                C[i,j] = l[i-1]*B[i-1,j] + d[i]*B[i,j] + u[i]*B[i+1,j]
            end
            C[nA,j] = l[nA-1]*B[nA-1,j] + d[nA]*B[nA,j]
        end
    end
    C
end
*(A::Tridiagonal, B::AbstractVecOrMat) = A_mul_B!(similar(B), A, B)

A_ldiv_B!(A::Tridiagonal,B::AbstractVecOrMat) = A_ldiv_B!(lufact!(A), B)
At_ldiv_B!(A::Tridiagonal,B::AbstractVecOrMat) = At_ldiv_B!(lufact!(A), B)
Ac_ldiv_B!(A::Tridiagonal,B::AbstractVecOrMat) = Ac_ldiv_B!(lufact!(A), B)


using StructuralCausalModels, Test

ProjDir = @__DIR__

d_str = "DAG(u ~ x + w, x ~ v + z, y ~ v, s1 ~ u, w ~ y, k ~ y, s2 ~ w + s3 + k, s3 ~ s1)"

dag = DAG("test_desc", d_str)

to_dagitty(dag)

#=
fname = joinpath(ProjDir, "test_descendents_01.dot")
to_graphviz(dag, fname)
Sys.isapple() && run(`open -a GraphViz.app $(fname)`)
=#

bs = basis_set(dag)
as = adjustment_sets(dag, :u, :s2)

@testset "Test_descendents_01" begin

  @test dag.a[:x, :u] == 0
  @test dag.e[:x, :u] == 1
  @test length(bs) == 83
  @test bs[70] == [:y, :s1, :u]
  @test bs[71] == [:y, :s1, :v, :u]
  @test bs[81] == [:z, :k, :y]
  @test bs[82] == [:z, :s2, :k, :w, :s3]
  @test bs[83] == [:z, :s3, :s1]

end

@testset "Adjustment_sets" begin

  @test as[1] == [:x, :w]
  @test as[2] == [:v, :w]
  @test as[3] == [:y, :w]
  @test as[4] == [:w, :k]

end


# ======================================
# Construction of the linearized system
# ======================================

"""
    build_matrix(F₁, F₂, N)

Compute the Carleman linearization matrix associated to the quadratic
system ``x' = F₁x + F₂(x⊗x)``, truncated at order ``N``.

### Input

- `F₁` -- sparse matrix of size ``n × n``
- `F₂` -- sparse matrix of size ``n × n^2``
- `N`  -- integer representing the truncation order, should be at least two
- `compress` -- if `true` produces the matrix corresponding to the commutative monomials only

### Output

Sparse matrix `A`.

### Algorithm

See references [1] and [2] of CARLIN.md.
"""
function build_matrix(F₁, F₂, N; compress=false)
    F₁ = sparse(F₁)
    F₂ = sparse(F₂)
    if compress
        n = size(F₁)[1]
        monoms = generate_monomials(n, N)
        nonzero_monoms = firstrest(monoms)[2]
        monom_to_ind = Dict(m => i for (i, m) in enumerate(nonzero_monoms))
        result = spzeros(length(monoms) - 1, length(monoms) - 1)
        for (ind, m) in enumerate(nonzero_monoms)
            for (i, j, c) in zip(findnz(F₁)...)
                if m[i] > 0
                    deriv = m
                    if i != j
                        deriv = m .+ Tuple((k == i) ? -1 : ((k == j) ? 1 : 0) for k in 1:n)
                    end
                    result[ind, monom_to_ind[deriv]] = m[i] * c
                end
            end

            if sum(m) < N
                for (i, j, c) in zip(findnz(F₂)...)
                    j0 = ((j - 1) % n) + 1
                    j1 = ((j - 1) ÷ n) + 1
                    if m[i] > 0
                        deriv = [m...,]
                        deriv[i] -= 1
                        deriv[j0] += 1
                        deriv[j1] += 1
                        deriv = (deriv...,)
                        result[ind, monom_to_ind[deriv]] += m[i] * c
                    end
                end
            end
        end
        return result
    end

    # No compression
    if N < 1
        throw(ArgumentError("expected the truncation order to be at least 1, got N=$N"))
    elseif N == 1
        _build_matrix_N1(F₁, F₂)
    elseif N == 2
        _build_matrix_N2(F₁, F₂)
    elseif N == 3
        _build_matrix_N3(F₁, F₂)
    elseif N == 4
        _build_matrix_N4(F₁, F₂)
    else
        _build_matrix_N(F₁, F₂, N) # general case
    end
end

"""
    lift_vector(x0, N)

returns a vector of monomials in x0 (hyperrectangle) of degree at most N
"""
function lift_vector(x0, N)
    monoms = generate_monomials(dim(x0), N)
    nonzero_monoms = firstrest(monoms)[2]
    result = []
    intervals = [interval(low(x0, i), high(x0, i)) for i in 1:dim(x0)]
    for m in nonzero_monoms
        push!(result, prod(intervals .^ m))
    end
    return Hyperrectangle(low=[i.lo for i in result], high=[i.hi for i in result])
end

"""
    generate_monomials(n, N)

returns a list of n-tuples of nonegative integers with the sum at most N
ordered by the total degree
"""
function generate_monomials(n, N)
    if n == 1
        return [(i,) for i in 0:N]
    end
    result = []
    prev = generate_monomials(n - 1, N)
    for ord in 0:N
        pind = 1
        while pind <= length(prev) && sum(prev[pind]) <= ord
            push!(result, (ord - sum(prev[pind]), prev[pind]...))
            pind += 1
        end
    end
    return result
end

function build_matrix_compressed(F1, F2, N)

end

function _build_matrix_N1(F₁, F₂)
    return F₁
end

function _build_matrix_N2(F₁, F₂)
    n = size(F₁, 1)
    a = hcat(kron_sum(F₁, 1), kron_sum(F₂, 1))
    b = hcat(spzeros(n^2, n), kron_sum(F₁, 2))
    return vcat(a, b)
end

function _build_matrix_N3(F₁, F₂)
    n = size(F₁, 1)
    a = hcat(kron_sum(F₁, 1), kron_sum(F₂, 1), spzeros(n, n^3))
    b = hcat(spzeros(n^2, n), kron_sum(F₁, 2), kron_sum(F₂, 2))
    c = hcat(spzeros(n^3, n), spzeros(n^3, n^2), kron_sum(F₁, 3))
    return vcat(a, b, c)
end

function _build_matrix_N4(F₁, F₂)
    n = size(F₁, 1)
    a = hcat(kron_sum(F₁, 1), kron_sum(F₂, 1), spzeros(n, n^3), spzeros(n, n^4))
    b = hcat(spzeros(n^2, n), kron_sum(F₁, 2), kron_sum(F₂, 2), spzeros(n^2, n^4))
    c = hcat(spzeros(n^3, n), spzeros(n^3, n^2), kron_sum(F₁, 3), kron_sum(F₂, 3))
    d = hcat(spzeros(n^4, n), spzeros(n^4, n^2), spzeros(n^4, n^3), kron_sum(F₁, 4))
    return vcat(a, b, c, d)
end

function _build_matrix_N(F₁, F₂, N)
    @assert N >= 3 "expected N to be at least 3, got N=$N"
    n = size(F₁, 1)

    out = Vector{typeof(F₁)}()
    for j in 1:N
        if j == N
            F12_j = kron_sum(F₁, j)
            Zleft_j = spzeros(n^j, sum(n^i for i in 1:j-1))
            push!(out, hcat(Zleft_j, F12_j))
        else
            F12_j = hcat(kron_sum(F₁, j), kron_sum(F₂, j))
            if j == 1
                Zright_j = spzeros(n^j, sum(n^i for i in j+2:N))
                push!(out, hcat(F12_j, Zright_j))
            elseif j == N-1
                Zleft_j = spzeros(n^j, sum(n^i for i in 1:j-1))
                push!(out, hcat(Zleft_j, F12_j))
            else # general case
                Zleft_j = spzeros(n^j, sum(n^i for i in 1:j-1))
                Zright_j = spzeros(n^j, sum(n^i for i in j+2:N))
                push!(out, hcat(Zleft_j, F12_j, Zright_j))
            end
        end
    end
    return reduce(vcat, out)
end


@inline function two_sum_tail(a::T, b::T, x::T) where {T}
    v = x - a
    return (a - (x - v)) + (b - v)
end

@inline function fast_two_sum_tail(a::T, b::T, x::T) where {T}
    v = x - a
    return b - v
end

@inline function two_sum(a::T, b::T) where {T}
    x = a + b
    return two_sum_tail(a, b, x), x
end

@inline function fast_two_sum(a::T, b::T) where {T}
    x = a + b
    return fast_two_sum_tail(a, b, x), x
end

@inline function two_diff_tail(a::T, b::T, x::T) where {T}
    b1 = a - x
    a1 = x + b1
    return (a - a1) + (b1 - b)
end

@inline function two_diff(a::T, b::T) where {T}
    x = a - b
    return two_diff_tail(a, b, x), x
end

@inline two_prod_tail(a::T, b::T, x::T) where {T} = fma(a, b, -x)

@inline function two_prod(a::T, b::T) where {T}
    x = a * b
    return two_prod_tail(a, b, x), x
end

@inline grow_expansion(a::T, b::T) where {T} = two_sum(a, b)
@inline grow_expansion(a) = a

@generated function grow_expansion(e::NTuple{N,T}, b::T) where {N,T}
    h = [Symbol(:h, i) for i in 1:N]
    body = [:(($(h[i]), b) = two_sum(b, e[$i])) for i in 1:N]
    return quote
        $(body...)
        return ($(h...), b)
    end
end

@generated function expansion_sum(e::NTuple{N,T}, f::NTuple{M,T}) where {N,M,T}
    h = [Symbol(:h, i) for i in 1:(N+M)]
    body = [
        :(($(h[i:i+N]...),) = grow_expansion(($(h[i:i+N-1]...),),f[$i])) for i in 1:M
    ]
    return quote
        ($(h[1:N]...),) = e
        $(body...)
        return ($(h...),)
    end
end

@inline function expansion_diff(e::NTuple{2,T}, b::T) where {T}
    x1, x2 = two_diff(e[1], b)
    x2, x3 = two_sum(e[2], x2)
    return x1, x2, x3
end

@inline function expansion_diff(e::NTuple{2,T}, f::NTuple{2,T}) where {T}
    x1, x2, x3 = expansion_diff((e[2], e[1]), f[1])
    x2, x3, x4 = expansion_diff((x3, x2), f[2])
    return x1, x2, x3, x4
end

@generated function expansion_scale(e::NTuple{N,T}, b::T) where {N,T}
    h = [Symbol(:h, i) for i in 1:(2N-1)]
    Q = [Symbol(:Q, i) for i in 2:2N]
    t = [Symbol(:t, i) for i in 2:N]
    s = [Symbol(:s, i) for i in 2:N]
    body = [
        quote
            ($(t[i-1]), $(s[i-1])) = two_prod(e[$i], b)
            ($(h[2i-2]), $(Q[2i-2])) = two_sum($(t[i-1]), $(Q[2i-3]))
            ($(h[2i-1]), $(Q[2i-1])) = fast_two_sum($(s[i-1]), $(Q[2i-2]))
        end for i in 2:N
    ]
    return quote
        (h1, Q2) = two_prod(e[1], b)
        $(body...)
        return ($(h...), $(Q[end]))
    end
end

approximate(e) = sum(e)


module ResidueClass

import Base: print, println,  *, /, -, +, ==, divrem, div, rem, copy, one, zero, sign, convert

using Polynomial

export Residue

### TO DO

##
# residue class ring with value a mod p^i
immutable Residue{T}

	val::T
	## generator of ideal
	p::T
	# power of ideal
	i::Int64
	# Ideal m = p^i, saved so that its computational more efficient
	m::T

	function Residue(val::T, p::T, i::Int, m::T)
		sign(val) >= 0 ? new(val,p,i,m) : new(val+m,p,i,m)
	end
end

## Factory ##
Residue{T}(a::T, p::T) = Residue{T}(rem(a,p), p, 1, p)
Residue{T}(a::T, p::T, i::Int) = (m = p^i; Residue{T}(a % m, p, i, m))
Residue{T}(a::T, p::T, i::Int, m::T) = Residue{T}(a % m, p, i, m)
Residue{T}(newVal::T, a::Residue{T}) = Residue{T}(newVal % a.m, a.p, a.i, a.m)
ResidueInRange{T}(newVal::T, a::Residue{T}) = Residue{T}(newVal, a.p, a.i, a.m)
# erzeugt aus einer Liste von Polynomen die entsprechenden Polynome über dem Restklassenring
function Residue{T}(a::Vector{Poly{T}}, p::T, i::Int)
	m = p^i
	array = Array{Poly{Residue{T}}}(size(a))
	for k=1:length(a)
		array[k] = Poly(Residue{T}(a[k].a, p, i, m), a[k].var)
	end
	return array
end

function Residue{T}(a::Vector{T}, p::T, i::Int, m::T)
	array = Array{Residue{T}}(size(a))
	for k=1:length(a)
		array[k] = Residue{T}(a[k] % m, p, i, m)
	end
	return array
end

Residue{T}(a::Vector{T}, p::T, i::Int) = Residue(a, p, i, p^i)
Residue{T}(a::Vector{T}, p::T) = Residue(a, p, 1, p)

convert{T}(t::Type{Residue{T}}, a::Residue{T}) = a

copy(a::Residue) = Residue(a.val, a)
one{T}(a::Residue{T}) = Residue(one(T), a.p, a.i, a.m)
zero{T}(a::Residue{T}) = Residue(zero(T), a)

# REVIEW
one{T}(t::Type{Residue{T}}) = one(T)
zero{T}(t::Type{Residue{T}}) = zero(T)

# ## Printing ##

print(io::IO, a::Residue) = (a.i == 1) ? print(io, "$(a.val) mod $(a.p)") :  print(io, "$(a.val) mod $(a.p)^$(a.i)=$(a.m)")

# ## Equates ##
==(a::Residue, b::Residue) = ((a.p != b.p) || (a.i != b.i)) ? false : a.val == b.val
==(a::WeakRef,b::Residue{WeakRef}) = (Residue(a % b.m, b) == b)
==(a::Residue{WeakRef},b::WeakRef) = ==(b,a)

## Arithmetics ##

sign{T}(a::Residue{T}) = sign(a.val)

# plus #
+{T}(a::Residue{T}, b::Residue{T}) = (a.m != b.m) ? error("Can't add residues!") : Residue(a.val+b.val, a)
+{T}(a::Residue{T}, b::T) = Residue(a.val+b, a)
+{T}(a::T, b::Residue{T}) = +(b,a)

# minus #
-{T}(a::Residue{T}, b::Residue{T}) = (a.m != b.m) ? error("Can't add residues!") : Residue(a.val - b.val, a)
-{T}(a::Residue{T}, b::T) = Residue(rem(a.val - b,a.m), a)
-{T}(a::T, b::Residue{T}) = -(b,a)

# unary minus #
-{T}(a::Residue{T}) = Residue(-a.val, a)

# # multiply #
*{T}(a::T, b::Residue{T}) = Residue(rem(a * b.val, b.m), b)
*{T}(a::Residue{T}, b::T) = *(b,a)
*{T}(a::Residue{T}, b::Residue{T}) = (a.m != b.m) ? error("Can't multiply! Quotients ($(a.m) and $(b.m))") : Residue(a.val*b.val, b)
function *{T}(a::Vector{Residue{T}}, b::Residue{T})
	c = Array{Residue{T}}(size(a))
	for i=1:length(a)
		c[i] = a[i] * b
	end
	return c
end
 *{T}(a::Residue{T}, b::Vector{Residue{T}}) = *(b,a)

function ^{T}(a::Residue{T}, n::Int)

    if n == 0
    	return ResidueInRange(one(T), a)
    end
    # a^-n = (a^-1)^n
    a_new = n > 0 ? a : inverse(a)
    b = copy(a_new)
    for k in reverse(digits(abs(n),2))[2:end]
        b = (k == 1) ? (b * b) * a_new : b * b
    end
    return b
end

# a is invertibale iff gcd(a, a.m) == 1, bezout coefficient s is then the inverse
function inverse{T}(a::Residue{T})
	(gcd, s, t) = gcdx(a.val, a.m)
    if gcd != one(T)
    	error("Value has no Inverse")
    end
    return Residue(s, a)
end

# division #

function divrem{T}(a::Residue{T}, b::Residue{T})
	if (a.m != b.m)
		error("Quotients ($(a.m) and $(b.m))")
	end
	(div, rem) = divrem(a.val, b.val)
	return (Residue(div, a), Residue(rem, a))
end

div{T}(a::Residue{T}, b::Residue{T}) = divrem(a,b)[1]
rem{T}(a::Residue{T}, b::Residue{T}) = divrem(a,b)[2]
/{T}(a::Residue{T}, b::Residue{T}) = a * inverse(b)

# function gcdx{T}(a::Residue{T}, b::Residue{T})
# 	(gcd, s, t) = gcdx(a.val, b.val)
# 	return (Residue(gcd, a), Residue(s, a), Residue(t, a))
# end

# gcd{T}(a::Residue{T}, b::Residue{T}) = Residue(gcd(a.val, b.val), a)

end


export auc, RocCurve, PrecisionRecallCurve, ratio_at_threshold, fpr_at_threshold, tpr_at_threshold, auc_at_threshold, precision_at_threshold, recall_at_threshold

using LinearAlgebra
using Statistics

function auc(x,y)
    direction = 1
    dx = [x[i] - x[i-1] for i in 2:length(x)]
    if any(x-> x < 0, dx)
        if all(x-> x < 0, dx)
            direction = -1
        else
            throw(DomainError("Function is neither increasing nor decreasing"))
        end
    end
    return direction * _trapz(y,x)
end

auc(c::confusion_matrix) = auc(false_positive_rate(c; average="binary"), true_positive_rate(c; average="binary"))

struct RocCurve
    thresholds::Vector{Number}
    ratios::Vector{Tuple{Number, Number}}
    options::NamedTuple
end

function RocCurve(y_true, y_pred, thresholding_func, thresholds; ith_class = nothing, class_name = nothing, average = "macro", normalize = false, weights = nothing)
    return _calculate_for_curves(y_true, y_pred, thresholding_func, thresholds; ith_class = ith_class, class_name = class_name,
        average = average, normalize = normalize, weights = weights)
end

function RocCurve(c::Vector{confusion_matrix} , thresholds, first_func = false_positive_rate, second_func = true_positive_rate; ith_class = nothing,
        class_name = nothing, average = "macro", normalize = false, weights = nothing)
        return _calculate_for_curves_with_matrices(c, thresholding_func, thresholds; ith_class = ith_class, class_name = class_name,
            average = average, normalize = normalize, weights = weights)
end

ratio_at_threshold(r::RocCurve, threshold) = r.ratios[findfirst(x -> x == threshold, r.thresholds)]
fpr_at_threshold(r::RocCurve, threshold) = ratio_at_threshold(r, threshold)[1]
tpr_at_threshold(r::RocCurve, threshold) = ratio_at_threshold(r, threshold)[2]

function Base.show(io::IO, ::MIME"text/plain", r::RocCurve)
    println(io,"Roc Curve:")
    println(io, repeat("=", 60))
    for i in 1:length(r.thresholds)
        println(io, lpad("Threshold : ", 30), r.thresholds[i],  lpad("False Positive Rate : ", 30), r.ratios[i][1], lpad("True Positive Rate : ", 30), r.ratios[i][2])
    end
    println("These values have been calculated with the following parameters : ", r.options)
end

##

struct PrecisionRecallCurve
    thresholds::Vector{Number}
    ratios::Vector{Tuple{Number, Number}}
    options::NamedTuple
end

function PrecisionRecallCurve(y_true, y_pred, thresholding_func, thresholds; ith_class = nothing, class_name = nothing, average = "macro", normalize = false, weights = nothing)
    return _calculate_for_curves(y_true, y_pred, thresholding_func, thresholds, precision_score, recall_score; ith_class = ith_class, class_name = class_name,
        average = average, normalize = normalize, weights = weights)
end

function PrecisionRecallCurve(c::Vector{confusion_matrix} , thresholds, first_func = precision_score, second_func = recall_score; ith_class = nothing,
        class_name = nothing, average = "macro", normalize = false, weights = nothing)
        return _calculate_for_curves_with_matrices(c, thresholding_func, thresholds, precision_score, recall_score; ith_class = ith_class, class_name = class_name,
            average = average, normalize = normalize, weights = weights)
end

ratio_at_threshold(r::PrecisionRecallCurve, threshold) =  r.ratios[findfirst(x -> x == threshold, r.thresholds)]
precision_at_threshold(r::PrecisionRecallCurve, threshold) = ratio_at_threshold(r, threshold)[1]
recall_at_threshold(r::PrecisionRecallCurve, threshold) = ratio_at_threshold(r, threshold)[2]

function Base.show(io::IO, ::MIME"text/plain", r::PrecisionRecallCurve)
    println(io,"Precision Recall Curve:")
    println(io, repeat("=", 60))
    for i in 1:length(r.thresholds)
        println(io, lpad("Threshold : ", 30), r.thresholds[i],  lpad("Precision Score : ", 30), r.ratios[i][1], lpad("Recall Score : ", 30), r.ratios[i][2])
    end
    println("These values have been calculated with the following parameters : ", r.options)
end

##

#TODO
# roc_auc_score and other metrics


##
"""
IDEAL OBSERVER MODEL: Probability of switch compared to no switch at poke k after consecutive failures
k = consecutive failure
C = estimated state (L or R, turned in High and Low as in being already in the rewarded side)
Pr = probability

Qk💠 = Pr(Ck+1 ≠ Ck|t[k0:k],C[k0:k]) / Pr(Ck+1 = Ck|t[k0:k],C[k0:k]) {P(switch)/P(no switch)}

i = trial
t = confidence, depends on
    yi =  p of correct discrimination(ABSENT IN THE Flipping)  equal to Prew; they used the subjective psychometric function
    λi =  p of transition at trial i given the hazard of Psw

P of not reward (in absence of switch) = λk0 +  Σ[i = k0:k](λi(Π[j = k0:i-1])) / Π[i = k0:k](1 - yi)(1 - λi)

CONFIDENCE-BASED MODEL
χΣ = trial by trial evidence for a switch, increment with a magnitude dipendent on the expected accuracy (confidence)
    it has a gaussian distribution at each trial with μχ and σχ  such that CBM was equal to IOM for 1B-Er. σχ grows linearly
    with the number of error trial (failures) (σ²χ for N failure = N σχ of first error)
Θ = threshold

Making CBM close to IOM
    Qk🌸 = P(χΣ > = Θ) \ P(χΣ < Θ) = Integral(Θ,inf)(N(μχ,σχ)dχΣ) / 1 - Integral(Θ,inf)(N(μχ,σχ)dχΣ) = Qk💠

    REARRANGED IN
    Integral(-inf, Θ) (N(μχ,σχ)dχΣ) = 1 - (Qk💠 / 1 - Qk💠)

α = perseverance factor

being μχΣ fixed by the IOM and its relationship with σχΣ the model can be reduced to 2 parameter σχΣ and α Ψ = {σχΣ, α}
probability of switch become

    p(Xy/n = SW|ψ) Integral(θ,inf) (N(α(μχ),σχ)dχΣ)
"""
##
using QuadGK
using Distributions
using Optim
using FileIO
using DataFrames
using StatsBase
using StatsFuns
using StructArrays
using Plots
import Flipping.nextcount
using LinearAlgebra
##
quadgk(f, a,b,c...; rtol=sqrt(eps), atol=0, maxevals=10^7, order=7, norm=norm)
fun(x) = exp(x)
quadgk(fun,-1.0,1.0)

plot(pdf.(Normal(), -1.0:0.2:1.0))

rand(Normal(1,0.5))
#########

function χ(p, leave::Bool, fail)
    failure_index = Int64(fail) + 1
    μ,σ = p
    if σ < 0
        return 10000000
    else
        store = Vector{Float64}(undef,failure_index)
        for i in 1:length(failure_index)
            store[i] = rand(Normal(μ,i*σ))
        end
        accumulation = Base.accumulate(+,store)
        result = accumulation[end]
        return - log(exp(-result)+1) - !leave * result
    end
end

function smart_data_evidence(p, leaves, failure_indices)
    soa = StructArray((leaves, failure_indices))
    cm = countmap(soa)
    smart_data_evidence(p, cm)
end

function smart_data_evidence(p, cm)
    sum(n*χ(p, value...) for (value, n) in cm)
end
###############
function failure_idx(count::T, rewarded) where {T <: Number}
    if rewarded
        count = 0.0
    else
        count + 1.0
    end
end
###############

filename = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/Stimulations/DRN_Opto_Flipping/pokesDRN_Opto_Flipping.csv"
data = FileIO.load(filename) |> DataFrame
data[!,:LastPoke] .= false
data[!,:FailuresIdx] .= 0.0
by(data,[:Session,:Streak]) do dd
    dd.LastPoke[end] = true
    dd[:,:FailuresIdx] = Base.accumulate(failure_idx, occursin.("true",dd.Reward);init=0.0)
end
checkpoint = by(data,:Session) do d
    d.PokeIn[1] == d.PokeIn[2]
end
any(checkpoint[:,2])
println(describe(data))
filtered = filter(t -> t[:Protocol].!="45/15", data)
#################
any(isnan.(filtered[:,:LastPoke]))
any(isnan.(filtered[:,:FailuresIdx]))
#################
coll =by(filtered,[:Protocol,:Wall,:MouseID]) do dd
    leaves = dd[:,:LastPoke]
    failures_indices = dd[:,:FailuresIdx]
    ϵ = 10^(-2)
    res = optimize(p -> -smart_data_evidence(p, leaves, failures_indices) + ϵ*norm(p)^2, [1.0,1.0])
    mins = Optim.minimizer(res)
    DataFrame(Mu_G = mins[1], Sigma_N = mins[2])
    #DataFrame(Inverse_Temp = mins[1], Integration = mins[2])
end
println(describe(coll))
f = plot(; legend = :bottomright)
param_plot(coll,:Wall,[:Mu_G,:Sigma_N])
###############
function param_plot(df::AbstractDataFrame,split_var::Symbol,params::Union{Tuple,AbstractVector};groups = :MouseID)
    z = []
    sort!(df, [groups,split_var])
    dof = union(df[:,split_var])
    idxs_x = df[:,split_var].== dof[1]
    idxs_y = df[:,split_var].== dof[end]
    for par in params
        o = scatter(df[idxs_x,par],df[idxs_y,par],
        group = df[idxs_x,groups],
        legend = false,
        xlabel = dof[1],
        ylabel = dof[end],
        title = string(par),
        markersize = 1)
        Plots.abline!(1,0)
        push!(z,o)
    end
    plot(z...)
end


using Base.Test
using MPFR

# constructors
x = MPFRFloat{53}()
x = MPFRFloat(12)
y = MPFRFloat(x)
@test x == y
y = MPFRFloat(0xc)
@test x == y
y = MPFRFloat(12.)
@test x == y
y = MPFRFloat(BigInt(12))
@test x == y
y = MPFRFloat(BigFloat(12))
@test x == y
y = MPFRFloat("12")
@test x == y
y = MPFRFloat(float32(12.))
@test x == y
y = MPFRFloat(12//1)
@test x == y

# +
x = MPFRFloat(12)
y = MPFRFloat(30)
@test x + y == MPFRFloat(42)

# -
x = MPFRFloat(12)
y = MPFRFloat(-30)
@test x - y == MPFRFloat(42)

# *
x = MPFRFloat(6)
y = MPFRFloat(9)
@test x * y != MPFRFloat(42)
@test x * y == MPFRFloat(54)

# /
x = MPFRFloat(9)
y = MPFRFloat(6)
@test x / y == MPFRFloat(9/6)

# < / > / <= / >=
x = MPFRFloat(12)
y = MPFRFloat(42)
z = MPFRFloat(30)
@test y > x
@test y >= x
@test y > z
@test y >= z
@test x < y
@test x <= y
@test z < y
@test z <= y
@test y - x >= z
@test y - x <= z
@test !(x >= z)
@test !(y <= z)

# ^
x = MPFRFloat(12)
y = MPFRFloat(4)
@test x^y == MPFRFloat(20736)

# ceil
x = MPFRFloat(12.042)
@test MPFRFloat(13) == ceil(x)

# copysign
x = MPFRFloat(1)
y = MPFRFloat(-1)
@test copysign(x, y) == y
@test copysign(y, x) == x

# isfinite / isinf
x = MPFRFloat(Inf)
y = MPFRFloat(1)
@test isinf(x) == true
@test isinf(y) == false
@test isfinite(x) == false
@test isinf(x) == true

# isnan
x = MPFRFloat(NaN)
y = MPFRFloat(1)
@test isnan(x) == true
@test isnan(y) == false

# convert to
@test convert(MPFRFloat{53}, 1//2) == MPFRFloat("0.5")
@test convert(MPFRFloat{53}, 0.5) == MPFRFloat("0.5")
@test convert(MPFRFloat{53}, 40) == MPFRFloat("40")
@test convert(MPFRFloat{53}, float32(0.5)) == MPFRFloat("0.5")

# convert from
@test convert(Float64, MPFRFloat(0.5)) == 0.5
@test convert(Float32, MPFRFloat(0.5)) == float32(0.5)

# exponent
x = MPFRFloat(0)
@test_fails exponent(x)
x = MPFRFloat(Inf)
@test_fails exponent(x)
x = MPFRFloat(15.674)
@test exponent(x) == exponent(15.674)

# sqrt DomainError
@test_fails sqrt(MPFRFloat(-1))

# precision
old_precision = get_default_precision()
x = MPFRFloat(0)
@test prec(x) == old_precision
set_default_precision(256)
x = MPFRFloat(0)
@test prec(x) == 256
set_default_precision(old_precision)
z = with_precision(240) do
    z = x + 20
    return z
end
@test float(z) == 20.
@test prec(z) == 240
x = MPFRFloat(12)
@test prec(x) == old_precision
@test_fails set_default_precision(1)

# integer_valued
@test integer_valued(MPFRFloat(12))
@test !integer_valued(MPFRFloat(12.12))

# nextfloat / prevfloat
with_precision(53) do
    x = MPFRFloat(12.12)
    @test MPFRFloat(nextfloat(12.12)) == nextfloat(x)
    @test MPFRFloat(prevfloat(12.12)) == prevfloat(x)
end
@test isnan(nextfloat(MPFRFloat(NaN)))
@test isnan(prevfloat(MPFRFloat(NaN)))

# comparisons
x = MPFRFloat(1)
y = MPFRFloat(-1)
z = MPFRFloat(NaN)
ipl = MPFRFloat(Inf)
imi = MPFRFloat(-Inf)
@test x > y
@test x >= y
@test x >= x
@test y < x
@test y <= x
@test y <= y
@test x < ipl
@test x <= ipl
@test x > imi
@test x >= imi
@test imi == imi
@test ipl == ipl
@test imi < ipl
@test z != z
@test !(z == z)
@test !(z <= z)
@test !(z < z)
@test !(z >= z)
@test !(z > z)

# modf
x = MPFRFloat(12)
y = MPFRFloat(0.5)
@test modf(x+y) == (y, x)
x = MPFRFloat(NaN)
@test map(isnan, modf(x)) == (true, true)
x = MPFRFloat(Inf)
y = modf(x)
@test (isnan(y[1]), isinf(y[2])) == (true, true)

# rem
with_precision(53) do
    x = MPFRFloat(2)
    y = MPFRFloat(1.67)
    @test rem(x,y) == rem(2, 1.67)
    y = MPFRFloat(NaN)
    @test isnan(rem(x,y))
    @test isnan(rem(y,x))
    y = MPFRFloat(Inf)
    @test rem(x,y) == x
    @test isnan(rem(y,x))
end

# min/max
x = MPFRFloat(4)
y = MPFRFloat(2)
@test max(x,y) == x
@test min(x,y) == y
y = MPFRFloat(NaN)
@test max(x,y) == x
@test min(x,y) == x
@test isnan(max(y,y))
@test isnan(min(y,y))

# sum
x = MPFRFloat(1)
y = MPFRFloat(2)
z = MPFRFloat(3)
w = MPFRFloat(4)
@test sum([x,y,z,w]) == MPFRFloat(10)
big_array = ones(MPFRFloat, 100)
@test sum(big_array) == MPFRFloat(100)

# promotion
# the array converts everyone to the DEFAULT_PRECISION!
x = MPFRFloat(12)
y = with_precision(60) do
    MPFRFloat(42)
end
@test [x,y] == [MPFRFloat(12), MPFRFloat(42)]

# log / log2 / log10
x = MPFRFloat(42)
@test log(x) == log(42)
@test isinf(log(MPFRFloat(0)))
@test_fails log(MPFRFloat(-1))
@test log2(x) == log2(42)
@test isinf(log2(MPFRFloat(0)))
@test_fails log2(MPFRFloat(-1))
@test log10(x) == log10(42)
@test isinf(log10(MPFRFloat(0)))
@test_fails log10(MPFRFloat(-1))

# exp / exp2 / exp10
x = MPFRFloat(10)
@test exp(x) == exp(10)
@test exp2(x) == 1024
@test exp10(x) == 10000000000

# convert to integer types
x = MPFRFloat(12.1)
y = MPFRFloat(42)
@test_fails convert(Int32, x)
@test_fails convert(Int64, x)
@test_fails convert(BigInt, x)
@test_fails convert(Uint32, x)
@test_fails convert(Uint32, x)
@test convert(Int32, y) == 42
@test convert(Int64, y) == 42
@test convert(BigInt, y) == 42
@test convert(Uint32, y) == 42
@test convert(Uint32, y) == 42

# iround
x = MPFRFloat(42.42)
y = with_precision(256) do
    MPFRFloat("9223372036854775809.2324")
end
z = BigInt("9223372036854775809")
@test iround(x) == 42
@test iround(y) == z
@test typeof(iround(Uint8, x)) == Uint8 && iround(Uint8, x) == 0x2a
@test typeof(iround(Uint16, x)) == Uint16 && iround(Uint16, x) == 0x2a
@test typeof(iround(Uint32, x)) == Uint32 && iround(Uint32, x) == 0x2a
@test typeof(iround(Uint64, x)) == Uint64 && iround(Uint64, x) == 0x2a
@test typeof(iround(Int64, x)) == Int64 && iround(Int64, x) == 42
@test typeof(iround(Int, x)) == Int && iround(Int, x) == 42
@test typeof(iround(Uint, x)) == Uint && iround(Uint, x) == 0x2a

# factorial
with_precision(256) do
    x = MPFRFloat(42)
    @test factorial(x) == factorial(BigInt(42))
    x = MPFRFloat(10)
    @test factorial(x) == factorial(10)
    @test_fails factorial(MPFRFloat(-1))
    @test_fails factorial(MPFRFloat(331.3))
end

# bessel functions
@test_approx_eq besselj(4, MPFRFloat(2)) besselj(4, 2.)
@test_approx_eq besselj0(MPFRFloat(2))  besselj0(2.)
@test_approx_eq besselj1(MPFRFloat(2))  besselj1(2.)
@test_approx_eq bessely(4, MPFRFloat(2))  bessely(4, 2.)
@test_approx_eq bessely0(MPFRFloat(2))  bessely0(2.)
@test_approx_eq bessely1(MPFRFloat(2))  bessely1(2.)

# trigonometric functions
for f in (:sin,:cos,:tan,:sec,:csc,:cot,:acos,:asin,:atan,
        :cosh,:sinh,:tanh,:sech,:csch,:coth,:asinh),
    j in (-1., -0.5, -0.25, .25, .5, 1.)
    @eval begin
        @test_approx_eq ($f)(MPFRFloat($j)) ($f)($j)
    end
end
for f in (:acos,:asin,:acosh,:atanh),
    j in (-2, -1.5)
    @eval begin
        @test_fails ($f)(MPFRFloat($j))
    end
end
for f in (:sin,:cos,:tan,:sec,:csc,:cot,:cosh,:sinh,:tanh,
        :sech,:csch,:coth,:acosh,:asinh),
    j in (1., 1.5, 1.9)
    @eval begin
        @test_approx_eq ($f)(MPFRFloat($j)) ($f)($j)
    end
end
for j in (.25, .5)
    @test_approx_eq atanh(MPFRFloat(j)) atanh(j)
end


export CartesianTrajectory, cartesianTrajectory2dNodes, cartesian2dDensity

"""
    CartesianTrajectory(numProfiles, numSamplingPerProfile
                  ; TE::Float64=0.0
                  , AQ::Float64=1.e-3
                  , kmin=(-0.5,-0.5)
                  , kmax=(0.5,0.5)
                  , kargs...)

returns a 2d cartesian trajectory.

# Arguments
* `numProfiles::Int64`            - number of profiles
* `numSamplingPerProfile::Int64`  - number of sampling points per profile
* (`TE::Float64=0.0`)             - echo time in s
* (`AQ::Float64=1.e-3`)           - readout duration in s (per profile)
* (`kmin=(-0.5,-0.5)`)            - minimum values of the covered k-space (for partial Fourier imaging)
* (`kmax=(-0.5,-0.5)`)            - maximum values of the covered k-space (for partial Fourier imaging)
"""
function CartesianTrajectory(numProfiles, numSamplingPerProfile
                  ; TE::Float64=0.0
                  , AQ::Float64=1.e-3
                  , kmin=(-0.5,-0.5)
                  , kmax=(0.5,0.5)
                  , kargs...)
  nodes = cartesian2dNodes(numProfiles, numSamplingPerProfile; kmin=kmin,kmax=kmax)
  times = readoutTimes(numProfiles, numSamplingPerProfile; TE=TE, AQ=AQ)
  return  Trajectory("Cartesian", nodes, times, TE, AQ, numProfiles, numSamplingPerProfile, 1, true, false)
end

function cartesian2dNodes(numProfiles, numSamplingPerProfile
                  ; kmin=(-0.5,-0.5)
                  , kmax=(0.5,0.5)
                  , kargs...)

  nodes = zeros(2,numSamplingPerProfile, numProfiles)
  posX = collect( -ceil(Int64, (numSamplingPerProfile-1)/2.):floor(Int64, (numSamplingPerProfile-1)/2.) ) / numSamplingPerProfile
  posY = collect( -ceil(Int64, (numProfiles-1)/2.):floor(Int64, (numProfiles-1)/2.) ) / numProfiles

  for l = 1:numProfiles
    for k = 1:numSamplingPerProfile
      nodes[1,k,l] = posX[k]
      nodes[2,k,l] = posY[l]
    end
  end

  # rescale nodes according to the region covered by kmin & kmax
  nodes[1,:,:] .*= (kmax[1]-kmin[1])
  nodes[2,:,:] .*= (kmax[2]-kmin[2])
  # shift nodes to the region defined by kmin & kmax
  nodes[1,:,:] .+= 0.5*(kmin[1]+kmax[1])
  nodes[2,:,:] .+= 0.5*(kmin[2]+kmax[2])

  return reshape(nodes, 2, numSamplingPerProfile*numProfiles)
end

function cartesian2dDensity(numProfiles::Int64, numSamplingPerProfile::Int64)
  density = zeros(numSamplingPerProfile, numProfiles)
  for l = 1:numProfiles
    for k = 1:numSamplingPerProfile
      density[k,l] = 1
    end
  end
  return vec(density)
end


using JSON

mutable struct Body
    name::String
    mass::Float64
    x::Float64
    y::Float64
    z::Float64
    vx::Float64
    vy::Float64
    vz::Float64
end

function read_bodies(filename)
    bodies = Body[]
    open(filename, "r") do handle
        json = JSON.parse(handle)
        unitmass = if haskey(json, "unitmass")
            json["unitmass"]
        else
            1.0
        end
        for (key, value) in json
            if key != "unitmass"
                body = Body(key, unitmass * value["mass"],
                            value["position"]..., value["velocity"]...)
                push!(bodies, body)
            end
        end
    end
    bodies
end

function energy(bodies)
    E = 0.0
    for i in 1:length(bodies)
        E += kinetic_energy(bodies[i])
        for j in (i+1):length(bodies)
            E += potential_energy(bodies[i], bodies[j])
        end
    end
    E
end

kinetic_energy(body) = 0.5 * body.mass * (body.vx^2 + body.vy^2 + body.vz^2)

function potential_energy(a, b)
    dx = a.x - b.x
    dy = a.y - b.y
    dz = a.z - b.z
    -a.mass * b.mass / sqrt(dx^2 + dy^2 + dz^2)
end

function advance!(bodies, dt)
    for i in 1:length(bodies)
        a = bodies[i]
        for j in (i+1):length(bodies)
            b = bodies[j]

            dx = a.x - b.x
            dy = a.y - b.y
            dz = a.z - b.z
            d2 = dx^2 + dy^2 + dz^2
            ds = sqrt(d2)
            magnitude = dt / (d2*d2)
            magnitude *= ds

            a.vx -= b.mass * magnitude * dx
            a.vy -= b.mass * magnitude * dy
            a.vz -= b.mass * magnitude * dz

            b.vx += a.mass * magnitude * dx
            b.vy += a.mass * magnitude * dy
            b.vz += a.mass * magnitude * dz
        end
    end

    for b in bodies
        b.x += dt * b.vx
        b.y += dt * b.vy
        b.z += dt * b.vz
    end
end

function main(filename, time_steps)
    bodies = read_bodies(filename)

    for body in bodies
        println(body)
    end

    @printf "Initial Energy: %.9f\n" energy(bodies)
    duration = @elapsed for _ in 1:time_steps
        advance!(bodies, 0.01)
    end
    @printf "Final Energy:   %.9f\n" energy(bodies)
    @printf "Elapsed: %fs\n\n" duration
end

if length(ARGS) < 2
    error("usage: julia nbody.jl <json-file> <time-steps>")
end

const filename = ARGS[1]
const time_steps = parse(Int, ARGS[2])

main(filename, time_steps)


"""
    univariatecalibrationrecipe(rmodel::UnivariateCalibration, X, Y;
                                            decimals = 3, text_location = :topleft)

Plots the calibration trend line of a UnivariateCalibration object given an input `X` and ground-truth `Y`.
"""
@recipe function univariatecalibrationrecipe(rmodel::UnivariateCalibration, X, Y;
                                            decimals = 3, text_location = :topleft)
    Yhat = rmodel.(X)
    seriestype := :scatter
    title := "Univariate Regression Plot"
    xlabel := Symbol("Fitted Values")
    ylabel := Symbol("Residuals")
    label := "Actual Values"
    @series y := (X, Y)
    seriestype := :straightline
    label := "Trend-Line"
    slope, bias, Rsq = string.( round.( [ rmodel.Slope, rmodel.Offset, rmodel.Rsq ], digits = decimals ) )
    @series begin
        annotation := (text_location, "y = " * slope * " x + " * bias * "\n R2 = " * Rsq )
        y := ( [0, 1], [rmodel.Offset, rmodel.Offset + rmodel.Slope])
    end
end

"""
    standard_addition_recipe(rmodel::StandardAddition, X, Y;
                                            decimals = 3, text_location = :topleft)

Plots a trend line of a StandardAddition object given an input `X` and ground-truth `Y`.
"""
@recipe function standard_addition_recipe(rmodel::StandardAddition, X, Y;
                                            decimals = 3, text_location = :topleft)
    seriestype := :scatter
    legend := :bottomright
    title := "Standard Addition Plot"
    xlabel := Symbol("Spike")
    ylabel := Symbol("Response")
    label := "Experiment"
    @series y := (X, Y)
    seriestype := :straightline
    label := "Trend-Line"
    slope, bias, Rsq = string.( round.( [ rmodel.Slope, rmodel.Offset, rmodel.Rsq ], digits = decimals ) )
    @series begin
        y := ( [0, 1], [rmodel.Offset, rmodel.Offset + rmodel.Slope])
    end
    seriestype := :scatter
    label := "Unknown"
    @series y := ( [ rmodel.Unknown ], [ 0 ] )
    seriestype := :hline
    label := ""
    color := :black
    @series y := [ 0 ]
    seriestype := :vline
    label := ""
    color := :black
    @series begin
        annotation := (text_location, "y = " * slope * " x + " * bias * "\n R2 = " * Rsq )
        x := [ X[1] ]
    end
end

"""
    residualsplotrecipe(rmodel::ChemometricsTools.RegressionModels, X, Y)

Plots the residuals of a RegressionModel object given an input `X` and ground-truth `Y`.
"""
@recipe function residualsplotrecipe(rmodel::ChemometricsTools.RegressionModels, X, Y)
    Yhat = rmodel(X)
    seriestype := :scatter
    title := "Residual vs Fitted Values Plot"
    xlabel := Symbol("Fitted Values")
    ylabel := Symbol("Residuals")
    @series y := (Yhat .- Y, Yhat)
end

struct QQ
    StoredTuple::Tuple
end

"""
    QQ( Y1, Y2; Quantiles = collect( 1 : 99 ) ./ 100 )

Returns a plotable object of a Quantile-Quantile plot between vectors `Y1` and `Y2` at the desired `Quantiles`.
"""
function QQ( Y1, Y2; Quantiles = collect( 1 : 99 ) ./ 100 )
    return QQ( (Statistics.quantile!(Y1, Quantiles), Statistics.quantile!(Y2, Quantiles) ))
end

"""
    quantilequantile(qq::QQ)

Plots a Quantile-Quantile plot.
"""
@recipe function quantilequantile(qq::QQ)
    m = sum( diff( qq.StoredTuple[ 1 ] ) ) / sum( diff( qq.StoredTuple[ 2 ] ) )
    int = qq.StoredTuple[2][2] - (m * qq.StoredTuple[1][2])

    seriestype := :scatter
    title := "Quantile-Quantile Plot"
    xlabel := Symbol("Quantile 1")
    ylabel := Symbol("Quantile 2")
    @series y := qq.StoredTuple

    seriestype := :straightline
    label := "Trend-Line"
    @series y := ([0, 1], [int, int + m])
end

struct BlandAltman
    means::Array{Float64, 1}
    differences::Array{Float64, 1}
    UpperLimit::Float64
    Center::Float64
    LowerLimit::Float64
    Outliers::Array{Float64}
end

"""
    BlandAltman(Y1, Y2; Confidence = 1.96)

Returns a Plottable object of a Bland-Altman plot between vectors `Y1` and `Y2` with a confidence limit of `Confidence`.
"""
function BlandAltman(Y1, Y2; Confidence = 1.96)
    means = (Y1 .+ Y2) ./ 2.0
    diffs = Y2 .- Y1
    MeanofDiffs = StatsBase.mean( diffs )
    StdofDiffs = StatsBase.std( diffs )

    UpperLimit = MeanofDiffs + Confidence * StdofDiffs
    Center = MeanofDiffs
    LowerLimit = MeanofDiffs - Confidence * StdofDiffs
    #To:Do Add trend-line....
    Outliers = findall( (diffs .> MeanofDiffs + Confidence*StdofDiffs) )
    Outliers = vcat(Outliers, findall( diffs .< MeanofDiffs - Confidence*StdofDiffs ) )
    return BlandAltman( means, diffs, UpperLimit, Center, LowerLimit, Outliers )
end

"""
    blandaltman(BA::BlandAltman)

Plots a Bland-Altman plot.
"""
@recipe function blandaltman(BA::BlandAltman)
    seriestype := :scatter
    title := "Bland Altman"
    xlabel := Symbol("Means")
    ylabel := Symbol("Differences")
    @series y := ( BA.means, BA.differences )
    seriestype := :straightline
    color := :red
    @series y := ([0, 1], [BA.UpperLimit, BA.UpperLimit])
    seriestype := :straightline
    color := :red
    @series y := ([0, 1], [BA.Center, BA.Center])
    seriestype := :straightline
    color := :red
    @series y := ([0, 1], [BA.LowerLimit, BA.LowerLimit])
end

#rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])

"""
    IntervalOverlay(Spectra, Intervals, Err)

Creates a Plottable object to display the relative error(`Err`) of each `interval` ontop of a `Spectra`.
"""
struct IntervalOverlay
    Spectra
    Intervals
    Err
end

"""
    intervaloverlay(IO::IntervalOverlay)

Plots a barchart overlay over a spectra according to an `IntervalOverlay` object.
"""
@recipe function intervaloverlay(IO::IntervalOverlay)
    RelativeErr = IO.Err ./ sum(IO.Err);

    seriestype := :path
    title := "Spectral Error Overlay"
    xlabel := Symbol("Bins")
    ylabel := Symbol("Absorbance / Relative CV Error")
    legend := false
    @series y := IO.Spectra'

    seriestype := :shape
    opacity := 0.5
    for (i, line) in enumerate( IO.Intervals )
        scalethis = reduce(max, IO.Spectra)  / reduce(max, RelativeErr)
        w = line[end] - line[1]
        h = RelativeErr[i] * scalethis
        xloc = line[1]
        #@series y := rectangle(w,h,xloc,0)
        #seriestype := :shape
        @series ( xloc .+ [ 0, w, w, 0 ], [ 0, 0, h, h ] )
    end
end

"""
    PCA/LDA(::Union{PCA, LDA}; Axis = [1,2])

Plots scores of `PCA`/`LDA` object using the defined `Axis'`.
"""
@recipe function f(DA::Union{PCA, LDA}; Axis = [1,2])
    @assert(length(Axis) == 2, "Axis for PCA/LDA plot must be of length 2." )
    exvar = round.(ExplainedVariance(DA) .* 10000) /100
    seriestype := :scatter
    if isa(DA, PCA)
        title := "PCA Scores Plot"
        xlabel := "PCA $(Axis[1]) [$(exvar[Axis[1]]) %]"
        ylabel := "PCA $(Axis[2]) [$(exvar[Axis[2]]) %]"
    else
        title := "LDA Scores Plot"
        xlabel := "DA $(Axis[1]) [$(exvar[Axis[1]]) %]"
        ylabel := "DA $(Axis[2]) [$(exvar[Axis[2]]) %]"
    end
    @series y := ( DA.Scores[:,Axis[2]], DA.Scores[:,Axis[1]] )
end

"""
    DiscriminantAnalysisPlot(DA, GD, YHot, LblEncoding, UnlabeledData,
                                Axis = [1,2], Confidence = 0.90)

An object to hold all the information needed to plot a discriminant analysis.
`DA` = `LDA`/`PCA` object
`GD` = `GuassianDiscriminant` object
`YHot` = One hot encoded y vectors
`LblEncoding` = a `ClassificationLabel` object
`UnlabeledData` = data that does not have a known label (see where it falls in the plot).
`Axis` = principal or discriminant axis' to display
`Confidence` = Significance value from 0-1 for the confidence ellipses.
"""
struct DAPlot
    DA
    GD
    YHot
    LblEncoding
    UnlabeledData
    Axis
    Confidence
end

DiscriminantAnalysisPlot(DA, GD, YHot, LblEncoding, UnlabeledData;
                    Axis = [1,2], Confidence = 0.90) = DAPlot( DA, GD, YHot,
                    LblEncoding, UnlabeledData, Axis, Confidence)

"""
    discrimanalyplot(dap::DAPlot)

Plots two DA axis, with ellipsoids according to a `DiscriminantAnalysisPlot` object.
"""
@recipe function discrimanalyplot(dap::DAPlot)
    exvar = round.(ExplainedVariance(dap.DA) .* 10000) /100
    @assert all(dap.Axis .<= size(dap.UnlabeledData)[2])
    A = []
    Classes = size(dap.YHot)[2]
    colors = get_color_palette(:auto, plot_color(:white), Classes)
    for c in 1:Classes
        #Labelled data
        whichrows = dap.YHot[:,c] .== 1.0
        seriestype := :scatter
        color := colors[c]
        label := dap.LblEncoding.ToCold[c]
        @series y := ( dap.DA.Scores[whichrows,dap.Axis[1]], dap.DA.Scores[whichrows,dap.Axis[2]] )
        #Make confidence ellipse
        elipse = ConfidenceEllipse(dap.GD.ProjectedClassCovariances[c], dap.GD.ProjectedClassMeans[c,:],
                    dap.Confidence, dap.Axis; pointestimate = 180 );

        seriestype := :path
        color := :black
        linestyle := :dash
        linewidth := 2
        label := ""
        @series y := (elipse[:,1], elipse[:,2])
    end
    #New Data
    seriestype := :scatter
    color := :black
    markershape = :star
    label := ""
    legend := :topleft
    xlabel := "DA $(dap.Axis[1]) [$(exvar[dap.Axis[1]]) %]"
    ylabel := "DA $(dap.Axis[2]) [$(exvar[dap.Axis[2]]) %]"
    @series y := ( dap.UnlabeledData[:,dap.Axis[1]] , dap.UnlabeledData[:,dap.Axis[2]] )
end


mutable struct StructuredGrid{nD} <: AbstractArray{Real,nD}
    nodes::NTuple{nD,Vector{Float64}}
end

function StructuredGrid(nodes...)
    StructuredGrid{length(nodes)}(nodes)
end

Base.size(grid::StructuredGrid) = length.(grid.nodes)

nodes(sgrid::StructuredGrid) = sgrid.nodes
xnodes(grid::StructuredGrid) = grid.nodes[1]
ynodes(grid::StructuredGrid) = grid.nodes[2]
znodes(grid::StructuredGrid) = grid.nodes[3]

xnodes(grid::StructuredGrid, i::Int) = getindex(xnodes(grid), i)
ynodes(grid::StructuredGrid, j::Int) = getindex(ynodes(grid), j)
znodes(grid::StructuredGrid, k::Int) = getindex(znodes(grid), k)

Base.getindex(grid::StructuredGrid{1}, i::Int) = (xnodes(grid, i))
Base.getindex(grid::StructuredGrid{2}, i::Int, j::Int) = (xnodes(grid, i), ynodes(grid, j))
Base.getindex(grid::StructuredGrid{3}, i::Int, j::Int, k::Int) = (xnodes(grid, i), ynodes(grid, j), znodes(grid, k))

cell_sizes(sgrid::StructuredGrid) = StructuredGrid(diff.(nodes(sgrid)))
cell_volumes(sgrid::StructuredGrid) = map(prod, cell_sizes(sgrid))
_avg(nodes::AbstractVector) = 0.5(nodes[2:end] + nodes[1:end-1])
cell_centroids(sgrid::StructuredGrid) = StructuredGrid(_avg.(nodes(sgrid)))

abstract type AbstractNodeType end
struct Primal <: AbstractNodeType end
struct Dual <: AbstractNodeType end

function primal_nodes(ax::FieldComponent, reference_grid::StructuredGrid{Dim}) where {Dim}
    StructuredGrid(ntuple(i -> i == Int(ax) ? _avg(nodes(reference_grid)[i]) : nodes(reference_grid)[i], Dim))
end
function dual_nodes(ax::FieldComponent, reference_grid::StructuredGrid{Dim}) where {Dim}
    StructuredGrid(ntuple(i -> i == Int(ax) ? nodes(reference_grid)[i] : _avg(nodes(reference_grid)[i]), Dim))
end

function nodes(::Type{Primal}, ax::FieldComponent, reference_grid::StructuredGrid)
    primal_nodes(ax, reference_grid)
end
function nodes(::Type{Dual}, ax::FieldComponent, reference_grid::StructuredGrid)
    dual_nodes(ax, reference_grid)
end


__precompile__(false)
module ANNBN
    # __precompile__()
    using Plots, Statistics, Dates, LinearAlgebra, Random, SpecialFunctions, Clustering, IterativeSolvers
    using Base.Threads
    # using CUDAdrv, CUDAnative, CuArrays
    # __precompile__()

    function sigm1(x) 1.0/(1.0+exp(-x)) end
    function isigm1(x) -log((1.0-x)/x) end
    # function sigm1(x) atan(x) end
    # function isigm1(x) tan(x) end
    # function sigm1(x) 2atan(pi*x/2)/pi end
    # function isigm1(x) 2tan(pi*x/2)/pi end
    # function sigm1(x) x^2 end
    # function isigm1(x) sqrt(x) end
    # function sigm1(x) max(0,x) end
    # function isigm1(x) if x>0.1 return x else return -1.0 end end
    # function sigm1(x) log(1.0+exp(x)) end
    # function isigm1(x) log(exp(x)-1.0) end
    # function sigm1(x) if x<-1.0 -1.0 elseif x<1.0 x else 1.0 end end
    # function isigm1(x) if x<-1.0 -10.0 elseif x<1.0 x else 10.0 end end
    # function sigm1(x) 1.0/(1.0+exp(-1000.0x)) end
    # function isigm1(x) -log((1.0-1000.0x)/1000.0x) end
    # function sigm1(x) erf(x) end
    # function isigm1(x) erfi(x) end
    # function sigm1(x) exp(x) end
    # function isigm1(x) log(x) end
    # function sigm1(x) x^3 end
    # function isigm1(x) x^(1/3) end

    include("___clustering.jl")
    # include("new_clustering.jl")
    # include("___augment_data.jl")

    include("calc_phi.jl")
    include("calc_phi_deriv.jl")
    include("calc_phi_deriv_deriv.jl")
    include("predict_new.jl")
    include("predict_new_rbf.jl")
    include("predict_new_rbf_deriv.jl")
    include("predict_new_rbf_deriv_deriv.jl")
    include("train_layer_1_rbf.jl")
    include("train_layer_1_rbf_pde.jl")
    include("train_layer_1_sigmoid.jl")
    include("train_layer_1_sigmoid_fast.jl")
    include("train_layer_1_sigmoid_fast2.jl")
    include("train_layer_1_rbf_fast.jl")
    include("train_layer_1_rbf_deriv.jl")

    include("fit_nfolds!.jl")
    include("predict_nfolds.jl")
    include("fit_nfolds_sigmoid.jl")
    include("predict_nfolds_sigmoid.jl")

    include("deep_nnm.jl")
    include("initialize_deep_weights_per_layer.jl")
    # include("run_iit.jl")

    # include("train_layer_1_rbf_fast.jl")
    include("train_layer_1_rbf_new.jl")

    export ___clustering,new_clustering,
    calc_phi,calc_phi_deriv,calc_phi_deriv_deriv,
    predict_new,predict_new_rbf,predict_new_rbf_deriv,predict_new_rbf_deriv_deriv,
    train_layer_1_rbf,train_layer_1_rbf_pde,train_layer_1_sigmoid,
    fit_nfolds!,predict_nfolds,fit_nfolds_sigmoid,predict_nfolds_sigmoid,train_layer_1_rbf_fast,
    deep_nnm,initialize_deep_weights_per_layer,run_iit,train_layer_1_rbf_new,train_layer_1_rbf_deriv,
	train_layer_1_sigmoid_fast,train_layer_1_sigmoid_fast2

    # precompile(sigm1, (Float64,))
    # precompile(isigm1, (Float64,))
    # precompile(deep_nnm, (Array{Float64,1},Array{Float64,2},Array{Float64,2},Int64,Int64,Int64))
    # precompile(calc_phi, (Array{Float64,2},Array{Float64,2},Int64,Int64,Float64,Int64))

end


# TODO: add impulse response iterators

# features of the input output arrays to plan(w::𝕎)
# ====================================

# features of the grid from w.period
# =================================

function Δpix(w::𝕎{Tf}) where {Tf}
    rTf = real(Tf)
    rTf.(w.period ./ w.sz)
end

function Δfreq(w::𝕎{Tf}) where {Tf}
    rTf = real(Tf)
    rTf.((2π) ./ w.period)
end

function nyq(w::𝕎{Tf}) where {Tf}
    rTf = real(Tf)
    rTf.(π ./ Δpix(w))
end

# Note: this gives the area element of
# only the fourier tranformed coordinates
Ωpix(w::𝕎) = prod(Δr[1] for Δr in zip(Δpix(w), w.region) if Δr[2])

Ωfreq(w::𝕎) = prod(Δr[1] for Δr in zip(Δfreq(w), w.region) if Δr[2])

# 𝕎 scalings
# =================================

function unscale(w::𝕎{Tf,d}) where {Tf,d}
    𝕎{Tf,d}(w.sz, w.region, true, w.period)
end

function Base.real(w::𝕎{Tf,d}) where {Tf,d}
    𝕎{real(Tf),d}(w.sz, w.region, w.scale, w.period)
end

function Base.complex(w::𝕎{Tf,d}) where {Tf,d}
    𝕎{Complex{real(Tf)},d}(w.sz, w.region, w.scale, w.period)
end

"`nv_scale(w::𝕎)->Number` returns the multiplicative scale of the inverse of w::𝕎"
function inv_scale(w::𝕎{Tf}) where {Tf}
    rTf = real(Tf)
    ifft_normalization = FFTW.normalization(
                rTf,
                w.sz,
                tuple(findall(w.region)...)
            )
    return rTf(ifft_normalization / w.scale)
end

function unitary_scale(w::𝕎{Tf}) where {Tf}
    rTf = real(Tf)
    return rTf(prod(1/√i[1] for i in zip(w.sz, w.region) if i[2]))
end

function ordinary_scale(w::𝕎{Tf}) where {Tf}
    rTf = real(Tf)
    rTf(Ωpix(w) / ((2π)^(sum(w.region)/2)))
end

# Test that inv_scale(ordinary_scale(w))== Ωfreq(w) / ((2π)^(sum(w.region)/2))

# pix grid
# =================================

# low level
# -----------

function pix(n::Int, p::Real)
    Δx = p/n
    x  = (0:n-1) * Δx |> collect
    return x
end

# main API
# -----------

function pix(w::𝕎{Tf,d}) where {d,Tf}
    rTf = real(Tf)
    return map(w.sz, w.period) do nᵢ, pᵢ
        rTf.(pix(nᵢ, pᵢ))
    end::NTuple{d,Array{rTf,1}}
end

function fullpix(i::Int, w::𝕎{Tf,d,Tsf,Tp}) where {Tf,d,Tsf,Tp}
    xi     = pix(w)[i]
    rTf    = real(Tf)
    xifull = zeros(rTf, w.sz)
    for Ic ∈ CartesianIndices(xifull)
        xifull[Ic] = getindex(xi, Tuple(Ic)[i])
    end
    return xifull::Array{rTf,d}
end

function fullpix(w::𝕎{Tf,d,Tsi,Tp}) where {Tf,d,Tsi,Tp}
    rTf = real(Tf)
    return map(tuple(1:d...)) do i
        fullpix(i, w)
    end::NTuple{d,Array{rTf,d}}
end

# frequency grid
# =================================

# low level
# -----------

function freq(n::Int, p::Real)
    k    = _fft_output_index_2_freq.(1:n, n, p)
    return k
end

function rfreq(n::Int, p::Tp) where Tp
    freq(n, p)[1:(n÷2+1)]
end

# main API
# -----------

function freq(w::𝕎{Tf,d}) where {d, Tf<:FFTC}
    rTf = real(Tf)
    return map(w.sz, w.period, w.region) do nᵢ, pᵢ, rᵢ
        kᵢ = rᵢ ? freq(nᵢ, pᵢ) : pix(nᵢ, pᵢ)
        rTf.(kᵢ)
    end::NTuple{d,Array{rTf,1}}
end

function freq(w::𝕎{Tf,d}) where {d, Tf<:FFTR}
    ir = findfirst(w.region)
    return map(w.sz, w.period, w.region, tuple(1:d...)) do nᵢ, pᵢ, rᵢ, i
        kᵢ = (i==ir) ? rfreq(nᵢ, pᵢ) :
                  rᵢ ?  freq(nᵢ, pᵢ) : pix(nᵢ, pᵢ)
        Tf.(kᵢ)
    end::NTuple{d,Array{Tf,1}}
end

function fullfreq(i::Int, w::𝕎{Tf,d,Tsf,Tp}) where {Tf,d,Tsf,Tp}
    ki  = freq(w)[i]
    rTf = real(Tf)
    kifull = zeros(rTf, size_out(w))
    for Ic ∈ CartesianIndices(kifull)
        kifull[Ic] = getindex(ki, Tuple(Ic)[i])
    end
    return kifull::Array{rTf,d}
end

function fullfreq(w::𝕎{Tf,d,Tsf,Tp}) where {Tf,d,Tsf,Tp}
    return map(tuple(1:d...)) do i
        fullfreq(i,w)
    end::NTuple{d,Array{real(Tf),d}}
end

function wavenum(w::𝕎{Tf,d,Tsf,Tp}) where {Tf,d,Tsf,Tp}
    k   = freq(w)
    rTf = real(Tf)
    λ   = zeros(rTf, size_out(w))
    for Ic ∈ CartesianIndices(λ)
        λ[Ic] = sqrt(sum(abs2, getindex.(k,Tuple(Ic))))
    end
    return λ::Array{rTf,d}
end

# Internal function
# -----------------------------------
#`k_pre = _fft_output_index_2_freq.(1:n, n, p)` computes the 1-d freq
function _fft_output_index_2_freq(ind, nside, period)
    kpre = (2π / period) * (ind - 1)
    nyq  = (2π / period) * (nside/2)

    # • Both of the following options are valid
    # • Both options return the same value when nside is odd
    # • Using (kpre <= nyq) sets the convention  that the
    #   frequencies fall in (-nyq, nyq] when nside is even
    # • Using (kpre < nyq) sets the convention  that the
    #   frequencies fall in [-nyq, nyq) when nside is even
    #----------------
    # • Here are the two options:
    return  ifelse(kpre <= nyq, kpre, kpre - 2nyq) # option 1
    # return ifelse(kpre < nyq, kpre, kpre - 2nyq)  # option 2
end

# function _get_npd(;nᵢ, pᵢ=nothing, Δxᵢ=nothing)
#     @assert !(isnothing(pᵢ) & isnothing(Δxᵢ)) "either pᵢ or Δxᵢ needs to be specified (note: pᵢ = Δxᵢ .* nᵢ)"
#     d = length(nᵢ)
#     if isnothing(pᵢ)
#         @assert d == length(Δxᵢ) "Δxᵢ and nᵢ need to be tuples of the same length"
#         pᵢ = tuple((prod(xn) for xn in zip(Δxᵢ,nᵢ))...)
#     end
#     @assert d == length(pᵢ) "pᵢ and nᵢ need to be tuples of the same length"
#     nᵢ, pᵢ, d
# end

# #%% Used for constructing the covariance matrix of a subset of frequencies
# function get_rFFTimpulses(::Type{F}) where {T,nᵢ,pᵢ,dnᵢ,F<:rFFTgeneric{T,nᵢ,pᵢ,dnᵢ}}
#     g  = Grid(F)
#     CI = CartesianIndices(Base.OneTo.(g.nki))
#     LI = LinearIndices(Base.OneTo.(g.nki))

#     function _get_dual_k(k,n)
#         dk = n-k+2
#         mod1(dk,n)
#     end

#     function get_dual_ci(ci::CartesianIndex{dnᵢ})
#         return CartesianIndex(map(_get_dual_k, ci.I, g.nxi))
#     end

#     function rFFTimpulses(ci::CartesianIndex{dnᵢ})
#         rimpls = zeros(Complex{T}, g.nki...)
#         cimpls = zeros(Complex{T}, g.nki...)
#         dual_ci = get_dual_ci(ci)
#         if (ci==first(CI)) || (ci==dual_ci)
#             rimpls[ci]  = 1
#         elseif dual_ci ∈ CI
#             rimpls[ci]  = 1/2
#             cimpls[ci]  = im/2
#             rimpls[dual_ci]  =  1/2
#             cimpls[dual_ci]  = -im/2
#         else
#             rimpls[ci]  = 1/2
#             cimpls[ci]  = im/2
#         end
#         return rimpls, cimpls
#     end

#     return rFFTimpulses, CI, LI, get_dual_ci
# end


struct Scatterer
    m::ComplexF64
    x::Float64
end

function Scatterer(m::ComplexF64, radius::Float64, wavelength::Float64, n_medium::Float64 = 1.00)
    m /= n_medium
    wavelength /= n_medium  # effective wavelength in the environment

    k = 2π / wavelength  # free space wavenumber in the environment
    x = k .* radius
    return Scatterer(m, x)
end

scattering(s::Scatterer, m) = scattering(s.m, s.x, m)
absorption(s::Scatterer, m) = absorption(s.m, s.x, m)
extinction(s::Scatterer, m) = extinction(s.m, s.x, m)
backscattering(s::Scatterer, m) = backscattering(s.m, s.x, m)

@doc raw"""
mx(m, radius, wavelength, n_medium=1.00)

Conversion of parameters for calculation of light scattering by a small sphere based on Mie theory

# Arguments
- `m`: ``n_{\mathrm{material}}``; refractive index of the material of the sphere
- `radius`: radius of the sphere
- `wavelength`: target vacuum wavelength (range)
- `n_medium`: n_environement; refractive index of the environment
  around the sphere (default: air)

# Return values
- `m`: relative refractive index (``n_\mathrm{material} / n_\mathrm{environment}``)
- `x`: size parameter (wavenumber * radius)
"""
function mx(m, radius, wavelength, n_medium = 1.00)
    m /= n_medium
    wavelength /= n_medium  # effective wavelength in the environment

    k = 2π / wavelength  # free space wavenumber in the environment
    x = k .* radius
    return m, x
end


using Random

const Series = AbstractVector{Int}

abstract type InfoDist end

mutable struct MIDist <: InfoDist
    joint::Matrix{Int}
    m1::Vector{Int}
    m2::Vector{Int}
    N::Int
    MIDist() = new(zeros(Int, 2, 2), zeros(Int, 2), zeros(Int, 2), 0)
end

MIDist(xs::Series, ys::Series) = accumulate!(MIDist(), xs, ys)

function entropy(dist::MIDist)
    mi = 0.0
    for i in 1:2, j in 1:2
        p = dist.joint[i,j]
        if !iszero(p)
            mi += p * log2(p / (dist.m1[i] * dist.m2[j]))
        end
    end
    log2(dist.N) + mi/dist.N
end

function accumulate!(dist::MIDist, xs::Series, ys::Series)
    dist.N += length(xs)
    @inbounds for i in eachindex(xs)
        x, y = xs[i], ys[i]
        dist.m1[x] += 1
        dist.m2[y] += 1
        dist.joint[x, y] += 1
    end
    dist
end

@inline function clear!(dist::MIDist)
    dist.joint[:] .= 0
    dist.m1[:] .= 0
    dist.m2[:] .= 0
    dist.N = 0
    dist
end

function mutualinfo!(dist::MIDist, xs::Series, ys::Series)
    entropy(accumulate!(dist, xs, ys))
end

mutualinfo(xs::Series, ys::Series) = mutualinfo!(MIDist(), xs, ys)

mutable struct TEDist <: InfoDist
    k::Int
    states::Vector{Int}
    histories::Vector{Int}
    sources::Vector{Int}
    predicates::Vector{Int}
    N::Int

    function TEDist(k::Int)
        q = 2^k
        states = zeros(Int, 4q)
        histories = zeros(Int, q)
        sources = zeros(Int, 2q)
        predicates = zeros(Int, 2q)
        new(k, states, histories, sources, predicates, 0)
    end
end

@inline function clear!(dist::TEDist)
    dist.states[:] .= 0
    dist.histories[:] .= 0
    dist.sources[:] .= 0
    dist.predicates[:] .= 0
    dist.N = 0
    dist
end

function accumulate!(dist::TEDist, xs::Series, ys::Series)
    rng = dist.k:(length(ys)-1)
    dist.N = length(rng)
    history, q = 0, 1
    @inbounds for i in 1:dist.k
        q *= 2
        history = 2history + ys[i] - 1;
    end
    @inbounds for i in rng
        src = xs[i] - 1
        future = ys[i + 1] - 1
        source = 2history + src
        predicate = 2history + future
        state = 2predicate + src

        dist.states[state + 1] += 1
        dist.histories[history + 1] += 1
        dist.sources[source + 1] += 1
        dist.predicates[predicate + 1] += 1

        history = predicate - q*(ys[i - dist.k + 1] - 1)
    end
    dist
end

function entropy(xs::AbstractArray{Int}, N::Int)
    h = N * log2(N)
    for i in eachindex(xs)
        n = xs[i]
        if !iszero(n)
            h -= n * log2(n)
        end
    end
    h / N
end

function entropy(dist::TEDist)
    entropy(dist.sources, dist.N) +
    entropy(dist.predicates, dist.N) -
    entropy(dist.states, dist.N) -
    entropy(dist.histories, dist.N)
end

function transferentropy!(dist::TEDist, xs::Series, ys::Series)
    entropy(accumulate!(dist, xs, ys))
end

function transferentropy(xs::Series, ys::Series, k::Int)
    transferentropy!(TEDist(k), xs, ys)
end

function significance(rng::AbstractRNG, dist::InfoDist, measure::Function,
                      xs::Series, ys::Series;
                      nperms=1000, pvalue=0.05, spotcheck=1000)
    gt = measure(dist, xs, ys)
    clear!(dist)

    count = 0
    xsperm = xs[:]

    if nperms <= spotcheck
        for _ in 1:nperms
            count += (measure(dist, shuffle!(rng, xsperm), ys) ≥ gt)
            clear!(dist)
        end
        p = count / (nperms + 1)
        se = sqrt((p * (1 - p)) / (nperms + 1))

        return gt, p, se
    else
        for _ in 1:spotcheck
            count += (measure(dist, shuffle!(rng, xsperm), ys) ≥ gt)
            clear!(dist)
        end
        p = count / (spotcheck + 1)
        se = sqrt((p * (1 - p)) / (spotcheck + 1))

        if abs(p - pvalue) < 2pvalue
            for _ in 1:(nperms - spotcheck)
                count += (measure(dist, shuffle!(rng, xsperm), ys) ≥ gt)
                clear!(dist)
            end
            p = count / (nperms + 1)
            se = sqrt((p * (1 - p)) / (nperms + 1))
        end

        return gt, p, se
    end
end

function issig(datum::NTuple{3, Float64}; p::Float64=0.05)
    if p < zero(p) || p > one(p)
        throw(DomainError(p, "must be between 0.0 and 1.0"))
    end
    datum[2] < p
end

function signify(data::AbstractArray{NTuple{3,Float64}}; p::Float64=0.05)
    sigdata = first.(data)
    sigdata[.!issig.(data; p=p)] .= 0.0
    sigdata
end


@test SpecificV(0.0035, 300.0) ≈ 0.394_913_866E2
@test SpecificV(0.0035u"MPa", 300.0u"K") ≈ 0.394_913_866E2u"m^3/kg"

@test SpecificV(0.0035, 700.0) ≈ 0.923_015_898E2
@test SpecificV(0.0035u"MPa", 700.0u"K") ≈ 0.923_015_898E2u"m^3/kg"

@test SpecificV(30.0, 700.0) ≈ 0.542_946_619E-2
@test SpecificV(30.0u"MPa", 700.0u"K") ≈ 0.542_946_619E-2u"m^3/kg"

@test SpecificH(0.0035, 300.0) ≈ 0.254_991_145E4
@test SpecificH(0.0035u"MPa", 300.0u"K") ≈ 0.254_991_145E4u"kJ/kg"

@test SpecificH(0.0035, 700.0) ≈ 0.333_568_375E4
@test SpecificH(0.0035u"MPa", 700.0u"K") ≈ 0.333_568_375E4u"kJ/kg"

@test SpecificH(30.0, 700.0) ≈ 0.263_149_474E4
@test SpecificH(30.0u"MPa", 700.0u"K") ≈ 0.263_149_474E4u"kJ/kg"

@test SpecificU(0.0035, 300.0) ≈ 0.241_169_160E4
@test SpecificU(0.0035u"MPa", 300.0u"K") ≈ 0.241_169_160E4u"kJ/kg"

@test SpecificU(0.0035, 700.0) ≈ 0.301_262_819E4
@test SpecificU(0.0035u"MPa", 700.0u"K") ≈ 0.301_262_819E4u"kJ/kg"

@test SpecificU(30.0, 700.0) ≈ 0.246_861_076E4
@test SpecificU(30.0u"MPa", 700.0u"K") ≈ 0.246_861_076E4u"kJ/kg"

@test SpecificS(0.0035, 300.0) ≈ 0.852_238_967E1
@test SpecificS(0.0035u"MPa", 300.0u"K") ≈ 0.852_238_967E1u"kJ/kg/K"

@test SpecificS(0.0035, 700.0) ≈ 0.101_749_996E2
@test SpecificS(0.0035u"MPa", 700.0u"K") ≈ 0.101_749_996E2u"kJ/kg/K"

@test SpecificS(30.0, 700.0) ≈ 0.517_540_298E1
@test SpecificS(30.0u"MPa", 700.0u"K") ≈ 0.517_540_298E1u"kJ/kg/K"

@test SpecificCP(0.0035, 300.0) ≈ 0.191_300_162E1
@test SpecificCP(0.0035u"MPa", 300.0u"K") ≈ 0.191_300_162E1u"kJ/kg/K"

@test SpecificCP(0.0035, 700.0) ≈ 0.208_141_274E1
@test SpecificCP(0.0035u"MPa", 700.0u"K") ≈ 0.208_141_274E1u"kJ/kg/K"

@test SpecificCP(30.0, 700.0) ≈ 0.103_505_092E2
@test SpecificCP(30.0u"MPa", 700.0u"K") ≈ 0.103_505_092E2u"kJ/kg/K"

@test SpeedOfSound(0.0035, 300.0) ≈ 0.427_920_172E3
@test SpeedOfSound(0.0035u"MPa", 300.0u"K") ≈ 0.427_920_172E3u"m/s"

@test SpeedOfSound(0.0035, 700.0) ≈ 0.644_289_068E3
@test SpeedOfSound(0.0035u"MPa", 700.0u"K") ≈ 0.644_289_068E3u"m/s"

@test SpeedOfSound(30.0, 700.0) ≈ 0.480_386_523E3
@test SpeedOfSound(30.0u"MPa", 700.0u"K") ≈ 0.480_386_523E3u"m/s"


function simplify_constants(O::Operation)
    O_last = nothing
    _O = O
    while _O != O_last
        O_last = _O
        _O = _simplify_constants(_O)
        if typeof(_O) <: Operation
            _O = Operation(_O.op,simplify_constants.(_O.args))
        end
    end
    _O
end

function _simplify_constants(O)
    if Symbol(O.op) == :*
        # If any variable is `Constant(0)`, zero the whole thing
        # If any variable is `Constant(1)`, remove that `Constant(1)` unless
        # they are all `Constant(1)`, in which case simplify to a single variable
        if any(x->typeof(x)<:Variable && (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
            return Constant(0)
        elseif any(x->typeof(x)<:Variable && isequal(x,Constant(1)),O.args)
            idxs = find(x->typeof(x)<:Variable && isequal(x,Constant(1)),O.args)
            _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
            if isempty(_O.args)
                return Constant(1)
            elseif length(_O.args) == 1
                return _O.args[1]
            else
                return _O
            end
        else
            return O
        end
    elseif Symbol(O.op) == :+ && any(x->typeof(x)<:Variable && (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
        # If there are Constant(0)s in a big `+` expression, get rid of them
        idxs = find(x->typeof(x)<:Variable && (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
        _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
        if isempty(_O.args)
            return Constant(0)
        elseif length(_O.args) == 1
            return _O.args[1]
        else
            return O
        end
    elseif O.op == identity
        return O.args[1]
    elseif Symbol(O.op) == :- && length(O.args) == 1
        return Operation(*,Expression[-1,O.args[1]])
    else
        return O
    end
end
simplify_constants(x::Variable) = x
_simplify_constants(x::Variable) = x

export simplify_constants