# Tensordot.contract implements numpy.tensordot for Julia
# numpy.einsum-like interface is available in einengine.jl.

"""
    contract(T1, T2, axesL, axesR)
Contracts tensor 'T1' with 'T2' in axes specified in axesL and axesR.
"""
contract(TL, TR, axesL::Array{Int}, axesR::Array{Int}) = begin
    # Gets shape and axes information from helper.
    shapeInfo = contractprep(size(TL), size(TR), axesL, axesR)

    # Apply transformation
    contractraw(TL, TR, shapeInfo...)
end

# Raw contraction function.
function contractraw(TL, TR, permL::Array{Int}, permR::Array{Int},
        shapeL::Tuple, shapeR::Tuple, extL::Array, extR::Array)

    # Multiply and restore to original shape.
    reshape((reshape(permutedims(TL, permL), shapeL) *
        reshape(permutedims(TR, permR), shapeR)), (extL..., extR...))
end

"Index preparations for contracting."
function contractprep(shapeL::Tuple, shapeR::Tuple, axesL::Array{Int}, axesR::Array{Int})
    shapeL = [shapeL...]
    shapeR = [shapeR...]
    # TODO: Check axes boundary.

    # Dumb index permuting & size extraction
    dumbperm(shape::Array{Int}, pick::Array{Int}) = begin
        sbarrier = sort(pick) .+ 1
        ebarrier = sort(pick) .- 1
        sbarrier = vcat([1], sbarrier)
        ebarrier = vcat(ebarrier, length(shape))
        regular = Int[]
        for i = 1:length(sbarrier)
            append!(regular, sbarrier[i]:ebarrier[i])
        end # for
        return regular
    end # dumbperm

    # External permutation
    permL = dumbperm(shapeL, axesL)
    permR = dumbperm(shapeR, axesR)
    # External shape
    extL = [shapeL[i] for i in permL]
    extR = [shapeR[i] for i in permR]
    # Contractional permutation
    append!(permL, axesL)
    prepend!(permR, axesR)
    outerL = if (length(extL)==0) 1 else reduce(*, extL) end
    outerR = if (length(extR)==0) 1 else reduce(*, extR) end
    innerL = reduce(*, [shapeL[i] for i in axesL])
    innerR = reduce(*, [shapeR[i] for i in axesR])

    permL, permR, (outerL, innerL), (innerR, outerR), extL, extR
end # contract

# Adjoint of contract should refrain from digging into index processing.
@adjoint contractprep(shapeL::Tuple, shapeR::Tuple, axesL::Array{Int}, axesR::Array{Int}) = begin
    contractprep(shapeL, shapeR, axesL, axesR), _ -> nothing
end

#= Further simplification for multiple gradients.
   Disabled by default.

@adjoint permutedims(A, iperm) = begin
    permutedims(A, iperm), dA -> begin
        rperm = invperm(iperm)
        (permutedims(dA, rperm), nothing)
    end
end

@adjoint invperm(iperm) = begin
    invperm(iperm), _ -> nothing
end
=#


import Base: copy, isapprox
using Distributions: Beta, cdf
using LinearAlgebra
using Optim
using Statistics: mean

function wXageF0(ages)
    1.5 .* ages ./ (2 .+ ages)
end

function bandSparse(N, k::Integer)
    m = zeros(N, N)
    if k > 0
        for i in 1:(N-k)
            m[i, i + k] = 1.0
        end
    else
        for i in 1:(N + k)
            m[i - k, i] = 1.0
        end
    end
    m
end

function bandSparse(N, ks::AbstractArray)
    m = zeros(N, N)
    for k in ks
        if k > 0
            for i in 1:(N-k)
                m[i, i + k] = 1.0
            end
        else
            for i in 1:(N + k)
                m[i - k, i] = 1.0
            end
        end
    end
    m
end

function shiftO(N)
  bandSparse(N, 1)
end

function stageO(N, p)
  bandSparse(N, 0) * (1 - p) .+ bandSparse(N, 1) * p
end

function bdiag(A, B)
  calO = zeros(size(A)[1] + size(B)[1], size(A)[2] + size(B)[2])
  calO[1:size(A)[1], 1:size(A)[2]] .= A
  calO[(size(A)[1] + 1):end, (size(A)[2] + 1):end] .= B
  calO
end

function twoBlockO(A, B, p)
  nA = size(A)[1]
  calO = bdiag(A, B)
  calO[nA, nA + 1] = p
  calO
end

function ages_1066(unit = 365)
  x1 = collect(1:18) * 10
  x2 = maximum(x1) .+ collect(1:18) * (30 + 1 / 1.8)
  x3 = maximum(x2) .+ collect(1:18) * 365
  x4 = maximum(x3) .+ collect(1:12) * 365 * 5
  vcat(x1, x2, x3, x4) / unit
end

function calO_1066()
  A = shiftO(18)
  B = stageO(18, 1 / 3)
  C = stageO(18, 1 / 36.5)
  D = stageO(12, 1 / 36.5 / 5)
  calO = twoBlockO(A, B, 1 / 3)
  calO = twoBlockO(calO, C, 1 / 36.5)
  calO = twoBlockO(calO, D, 1 / 36.5 / 5)
  return(calO)
end

function calP_SoI(N, tau, nu, r)
  s = r .* (1 .- nu)
  v = r .* nu
  # Remain in state with probability s, go to next state probability v
  PSoI = bandSparse(N, 0) .* s .+ transpose(bandSparse(N, 1) .* v)
  PSoI = vcat(transpose(1 .- r), PSoI) # rbind
  PSoI = hcat(zeros(10), PSoI) # cbind
  # Uninfected remain uninfected.
  PSoI[1, 1] = 1
  PSoI
end

function SuSblock_Adam(N)
  SuS = bandSparse(N, 0:(N - 1))
  vcat(transpose(zeros(N + 1)), hcat(ones(N), SuS))
end

function ageXimmune(ageInDays)
    ageInDays .> (8 * 365)
end

"""
The original R:
xiF_0 <- function(x, p.xi = 2, N = 9, C = 10, D = 0) {
  diff(pbeta(c(0:N) / N, (1 - exp(-p.xi * x)) * C, exp(-p.xi * x) * C + D))
}
The beta distribution has the same density in both languages. I checked.
"""
function xiF_inner(x; p_xi = 2, N = 9, C = 10, D = 0)
    if (x > 0)
        α = (1 - exp(-p_xi * x)) * C
        β = exp(-p_xi * x) * C + D
        dist = Beta(α, β)
        diff(cdf.(dist, collect(0:N) / N))
    else
        bins = zeros(N)
        bins[1] = 1
        bins
    end
end

# I can't tell whether V is an array or a single float when it gets used.
# If it's an array, use pdf.() above, and it will work out.
function xiF_h(pit, V ; p_xi1 = 5, p_xi2 = 1.5, N = 9, eps = 0.001)
  if V > 0
      v = 1 - exp(-(pit[2] + p_xi2 * pit[3]) / V)
  else
      v = 1
  end
  xi1 = xiF_inner(v, p_xi = p_xi1)
  xi = xi1 .+ eps
  xi / sum(xi)
end

function make_parameters_adam()
    tau = 10
    N = 9
    nuVals = [1/2, 1/6, 1/7, 1/8, 1/9, 1/9, 1/9, 1/12, 0]
    rVals = vcat([0, 1/1000], ones(6) / 250, [1/400])
    calDtrue_SoI =vcat([0], ones(9))
    calDlm_SoI = vcat([0, 0.99, .95], 1 .- collect(3:7).^3 / 8.6^3, [.05, .001])
    calDimm_SoI = vcat([0, 0.95, .9, .6, .3], zeros(5))
    calK_SoI = vcat([0, 0], collect(8:-1:1) / 10)
    calBsevere_SoI = vcat([0, .3], zeros(8))
    calBfever_SoI = vcat([0, 1, .7], exp.(-collect(1:6)), [0])
    rho = 0.2
    etaXstage = vcat([0, 0.99, .95], 1 .- collect(3:7).^3 / 8.6^3, [.05, .001]).^2
    rhoXstage = [0, 1, 1, 1, .8, .8, .6, .6, .4, .4]
    ageForImmune = 8
    delta = 1 / 730
    mu1 = 0.001
    mu2 = 0.01
    mu3 = 0.001
    wF = wXageF0
    agesF = ages_1066
    calOF = calO_1066
    # Probability of taking drugs, given _any_ state. It's a background rate.
    Dd = 1 - exp(-delta / tau)
    r = exp.(-rVals * tau)

    PSoI = calP_SoI(N, tau, nuVals, r)
    SuS = SuSblock_Adam(N)
    ages = agesF()
    aXi = (ages .> 8)
    demog = vcat(diff(ages), [15])
    demog = demog / sum(demog)
    demog29 = (ages .>= 2) .& (ages .< 10)
    demog29 = demog29 / sum(demog29)
    U5 = copy(demog)
    O5 = copy(demog)
    O5[ages .< 5] .= 0
    U5[ages .>= 5] .= 0
    calO = calOF()
    w = wF(ages)
    calBfever = vcat(calBfever_SoI, calBfever_SoI, [0, 0])
    trt = rho .* rhoXstage .* etaXstage
    calBtrtdmal = vcat(trt .* calBfever_SoI, trt .* calBfever_SoI, [0, 0])
    calBuntrmal = vcat((1 .- trt) .* calBfever_SoI, (1 .- trt) .* calBfever_SoI, [0, 0])
    calBsevere = vcat(calBsevere_SoI, calBsevere_SoI, [0, 0])
    calBtrtdsev = vcat(trt .* calBsevere_SoI, trt .* calBsevere_SoI, [0, 0])
    calBuntrsev = vcat((1 .- trt) .* calBsevere_SoI, (1 .- trt) .* calBsevere_SoI, [0, 0])
    calK = vcat(calK_SoI, calK_SoI, [0, 0])
    calK[12] = calK[3]
    Dict{Symbol,Any}(
      :tau => tau,
      :N => N,
      :calDtrue => vcat(calDtrue_SoI, ones(10), [0, 0]),
      :calDlm => vcat(calDlm_SoI, calDlm_SoI, [0, 0]),
      :calDimm => vcat(calDimm_SoI, calDimm_SoI, [0, 0]),
      :calBfever => calBfever,
      :calBtrtdmal => calBtrtdmal,
      :calBuntrmal => calBuntrmal,
      :calBsevere => calBsevere,
      :calBtrtdsev => calBtrtdsev,
      :calBuntrsev => calBuntrsev,
      :calK => calK,
      :nu => nuVals,
      :r => r,
      :Dd => Dd,
      :rho => rho,
      :etaXstage => etaXstage,
      :rhoXstage => rhoXstage,
      :PSoI => PSoI,
      :SuS => SuS,
      :wF => wF,
      :w => w,
      :ages => ages,
      :ageXimmune => ageXimmune,
      :aXi => aXi,
      :demog => demog,
      :demog29 => demog29,
      :U5 => U5,
      :O5 => O5,
      :calO => calO,
      :mu1 => mu1,
      :mu2 => mu2,
      :mu3 => mu3,
      :xiF => xiF_h
    )
end

### Now we make the Adam model from those parameters

"""
A single group. The code may use cohort/population in the opposite sense.
I mean here that this is 22 compartments for one group of people.
"""
mutable struct CohortState
    X::Array{Float64,2}
    pit::Array{Float64,1}
    V::Float64
    ageInDays::Float64
    alpha1::Float64
    alpha2::Float64
end

function emptyX_Adam()
    X = zeros(22, 1)
    X[1] = 1
    CohortState(X, zeros(3), 0, 0, 0, 0)
end

function isapprox(a::CohortState, b::CohortState)
    (a.alpha1 ≈ b.alpha1 && a.alpha2 ≈ b.alpha2 && a.V ≈ b.V) || return(false)
    a.X ≈ b.X || return(false)
    a.pit ≈ b.pit || return(false)
    a.ageInDays ≈ b.ageInDays || return(false)
    true
end

"""
This tells me where they differ.
"""
function differences(a::CohortState, b::CohortState)
    msg = String[]
    a.alpha1 ≈ b.alpha1 || push!(msg, "alpha1")
    a.alpha2 ≈ b.alpha2 || push!(msg, "alpha2")
    a.V ≈ b.V || push!(msg, "V")
    a.X ≈ b.X || push!(msg, "X")
    a.pit ≈ b.pit || push!(msg, "pit")
    a.ageInDays ≈ b.ageInDays || push!(msg, "ageInDays")
    msg
end

"""
Multiple groups. There are 22 compartments for each of the
groups. The code may use cohort to mean this and population to mean
the CohortState, which I get, because this tracks multiple cohorts.
"""
mutable struct PopulationState
    X::Array{Float64,2}
    pit::Array{Float64,2}
    V::Float64
    ageInDays::Array{Float64,2}
    alpha1::Float64
    alpha2::Float64
end

function emptyXX_Adam(; L = 66)
    X = zeros(22, L)
    X[1, :] .= 1
    pit = zeros(3, L)
    ageInDays = zeros(1, L)
    PopulationState(X, pit, 0, ageInDays, 0, 0)
end

function isapprox(a::PopulationState, b::PopulationState)
    (a.alpha1 ≈ b.alpha1 && a.alpha2 ≈ b.alpha2 && a.V ≈ b.V) || return(false)
    a.X ≈ b.X || return(false)
    a.pit ≈ b.pit || return(false)
    a.ageInDays ≈ b.ageInDays || return(false)
    true
end

"""
This tells me where they differ.
"""
function differences(a::PopulationState, b::PopulationState)
    msg = String[]
    a.alpha1 ≈ b.alpha1 || push!(msg, "alpha1")
    a.alpha2 ≈ b.alpha2 || push!(msg, "alpha2")
    a.V ≈ b.V || push!(msg, "V")
    a.X ≈ b.X || push!(msg, "X")
    a.pit ≈ b.pit || push!(msg, "pit")
    a.ageInDays ≈ b.ageInDays || push!(msg, "ageInDays")
    msg
end

"""
Acts on members of the cohort state.
"""
function calP_Adam(alpha, pit, V, params)
    SuS = params[:SuS]
    PSoI = params[:PSoI]
    alpha = min(alpha, 1)
    xi = params[:xiF](pit, V)
    E1 = SuS .* vcat(0, xi)  # multiply each column
    e1 = sum(E1; dims=(1,))
    e1[isapprox.(e1, 0)] .= 1
    hatxi = E1 * ((1 ./ e1) .* I(10))
    Deta = min.(params[:rho] * params[:etaXstage] .* params[:rhoXstage], 1)
    vecD = Deta .+ (1 .- Deta) .* params[:Dd]
    notTreat = Diagonal(1 .- vecD)
    treat = vcat(vecD, vecD)
    B1 = PSoI * notTreat
    @assert size(B1) == (10, 10)
    B2 = hatxi * notTreat
    @assert size(B2) == (10, 10)
    chi = vcat([1], cumsum(xi))
    calP = vcat(
        hcat((1 .- alpha .* chi) .* B1, (1 -alpha) .* B2),
        hcat(alpha * chi .* B1, alpha * B2),
        transpose(treat),
        transpose(0 * treat)
    )
    calP = hcat(calP, zeros(size(calP)[1], 2))
    calP[22, 21] = 1
    calP[1, 22] = 1
    @assert size(calP) == (22, 22)
    calP
end

function PxX_Adam!(vecX::CohortState, alpha, params)
    vecX.ageInDays += params[:tau]
    w = params[:wF](vecX.ageInDays / 365)
    Xt = calP_Adam(w * alpha, vecX.pit, vecX.V, params) * vecX.X

    vecX.V = alpha + vecX.V * (1 - params[:mu1])
    vecX.pit[1] = vecX.pit[1] * (1 - params[:mu1]) + vecX.alpha2
    vecX.pit[2] = (vecX.pit[2] * (1 - params[:mu2]) +
            dot(params[:calDimm],  vecX.X))
    vecX.pit[3] = (vecX.pit[3] * (1 - params[:mu3]) +
            vecX.alpha2 * params[:ageXimmune](vecX.ageInDays))
    vecX.alpha2 = vecX.alpha1
    vecX.alpha1 = w * alpha
    vecX.X = Xt
end

"""
Simulates a cohort from birth over mx time steps and returns it
as a PopulationState.
"""
function cohortXX_Adam(alpha, params; mx = 2920)
    # Grow a cohort so you can assign its stages to a population.
    vecX = emptyX_Adam()
    XX = emptyXX_Adam(L = mx)
    vecX.V = alpha / params[:mu1]
    for i in 1:mx
        # An array column is size (22,) but an array with one column is (22,1)
        XX.X[:, i] .= vecX.X[:, 1]
        XX.pit[:, i] = vecX.pit[:, 1]
        XX.ageInDays[1, i] = vecX.ageInDays
        PxX_Adam!(vecX, alpha, params)
    end
    XX.alpha2 = vecX.alpha2
    XX.alpha1 = vecX.alpha1
    XX.V = vecX.V
    XX
end

function cohort2ages_Adam(cXX::PopulationState, params)
    ages = params[:ages]
    # ageInDays is a row vector. Easier to work with column here.
    ageInYears = vec(cXX.ageInDays / 365)
    XX = emptyXX_Adam()
    XX.V = cXX.V
    XX.ageInDays = zeros(1, length(ages))
    no_ages = []
    lower = 0.0
    for i in 1:length(ages)
        upper = ages[i]
        ix = (ageInYears .>= lower) .& (ageInYears .< upper)
        age_cnt = sum(ix)
        if age_cnt > 0
            XX.X[:, i] = sum(cXX.X[:, ix], dims = 2) / age_cnt
            XX.pit[:, i] = sum(cXX.pit[:, ix], dims = 2) / age_cnt
            XX.ageInDays[1, i] = mean(cXX.ageInDays[ix])
        else
            push!(no_ages, ages[i])
        end
        lower = upper
    end
    if length(no_ages) > 0
        println("no ages for $no_ages")
        @assert length(no_ages) == 0
    end
    XX
end

function test_cohort2ages_Adam()
    params = make_parameters_adam()
    XX = cohortXX_Adam(0.3, params)
    cohort2ages_Adam(XX, params)
end

function PtimesX_Adam!(XX::PopulationState, alpha, params)
    ageInDays = XX.ageInDays .+ params[:tau] * params[:calO]
    Xt = copy(XX.X)
    for a in 1:66
        Xt[:, a] .= calP_Adam(
            alpha * params[:w][a],
            XX.pit[:, a],
            XX.V,
            params
            ) * XX.X[:, a]
    end
    Xt *= params[:calO]
    Xt[1, 1] = 1

    XX.V = alpha + XX.V * (1 - params[:mu1])
    XX.pit[1, :] .= XX.pit[1, :] * (1 - params[:mu1]) .+ alpha
    XX.pit[2, :] .= (XX.pit[2, :] * (1 - params[:mu2]) .+
            transpose(XX.X) * params[:calDimm])
    XX.pit[3, :] .= XX.pit[3, :] * (1 - params[:mu3]) .+ XX.alpha2 * params[:aXi]

    XX.pit *= params[:calO]
    XX.ageInDays = XX.ageInDays * params[:calO]
    XX.alpha2 = XX.alpha1
    XX.alpha1 = alpha
    XX.X = Xt
end

function XX2pr29_Adam(XX, params)
    transpose(params[:calDlm]) * XX.X * params[:demog29]
end

function ar2stableWave_Adam(alpha, params; tol = 0.01)
    cXX = cohortXX_Adam(alpha, params)
    XX = cohort2ages_Adam(cXX, params)
    for i in 1:100
        PtimesX_Adam!(XX, alpha, params)
    end
    x = XX2pr29_Adam(XX, params)
    df = 1
    while df > tol
        original = copy(XX.X)
        PtimesX_Adam!(XX, alpha, params)
        x = XX2pr29_Adam(XX, params)
        df = sum(abs.(original .- XX.X))
    end
    XX
end

function ar2pr_Adam(alpha, params)
    XX = ar2stableWave_Adam(alpha, params)
    XX2pr29_Adam(XX, params)
end

"""
    pr2arSS_Adam(x, params)

Given a PfPR value and model parameters, return the attack rate
that would produce this PfPR for a population. This is an optimization
step.
"""
function pr2arSS_Adam(x, params)
    objective = x -> (x - ar2pr_Adam(alpha, params))^2
    lower = 0.0
    upper = 1.0
    res = optimize(objective, lower, upper)
    Optim.minimizer(res), Optim.iterations(res), Optim.converged(res)
end

function pr2ar_Adam(xx, params; alpha0 = nothing, XX0 = nothing, tol = 0.01)
    if alpha0 === nothing
        alpha = pr2arSS_Adam(xx[1], params)
    else
        alpha = alpha0
    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Geometry{Dim,T}

A geometry embedded in a `Dim`-dimensional space with coordinates of type `T`.
"""
abstract type Geometry{Dim,T} end

"""
    embeddim(geometry)

Return the number of dimensions of the space where the `geometry` is embedded.
"""
embeddim(::Type{<:Geometry{Dim,T}}) where {Dim,T} = Dim
embeddim(g::Geometry) = embeddim(typeof(g))

"""
    paramdim(geometry)

Return the number of parametric dimensions of the `geometry`. For example, a
sphere embedded in 3D has 2 parametric dimension (polar and azimuthal angles).
"""
paramdim(g::Geometry) = paramdim(typeof(g))

"""
    coordtype(geometry)

Return the machine type of each coordinate used to describe the `geometry`.
"""
coordtype(::Type{<:Geometry{Dim,T}}) where {Dim,T} = T
coordtype(g::Geometry) = coordtype(typeof(g))

"""
    centroid(geometry)

Return the centroid of the `geometry`.
"""
centroid(g::Geometry) = center(g)

"""
    measure(geometry)

Return the measure of the `geometry`, i.e. the length, area, or volume.
"""
function measure end

"""
    extrema(geometry)

Return the top left and bottom right corners of the
bounding box of the `geometry`.
"""
Base.extrema(g::Geometry) = extrema(boundingbox(g))

"""
    boundary(geometry)

Return the boundary of the `geometry`.
"""
function boundary end

"""
    g₁ ⊆ g₂

Tells whether or not the geometry `g₁` is a subset of geometry `g₂`.
"""
Base.issubset(g₁::Geometry, g₂::Geometry) = all(p ∈ g₂ for p in vertices(g₁))

"""
    p ∈ g

Tells whether or not the point `p` is in the geometry `g`.
"""
Base.in(p::Point, g::Geometry)

"""
    isconvex(geometry)

Tells whether or not the `geometry` is convex.
"""
isconvex(g::Geometry) = isconvex(typeof(g))

"""
    issimplex(geometry)

Tells whether or not the `geometry` is simplex.
"""
issimplex(::Type{<:Geometry}) = false
issimplex(g::Geometry) = issimplex(typeof(g))

# ----------------
# IMPLEMENTATIONS
# ----------------

include("polytopes.jl")
include("primitives.jl")

"""
    PointOrGeometry{Dim,T}

A union type that can represent either a point or a geometry.
"""
const PointOrGeometry{Dim,T} = Union{Point{Dim,T},Geometry{Dim,T}}

"""
    Multi(entities)

A collection of points or geometries seen as a single entity.

In geographic information systems (GIS) it is common to represent
multiple polygons as a single entity. In this case the polygons
refer to the same object in the real world (e.g. country with islands).
"""
struct Multi{Dim,T,I<:PointOrGeometry{Dim,T}} <: Geometry{Dim,T}
  items::Vector{I}
end

Base.getindex(multi::Multi, ind) = getindex(multi.items, ind)
Base.length(multi::Multi) = length(multi.items)
Base.eltype(multi::Multi) = eltype(multi.items)
Base.firstindex(multi::Multi) = firstindex(multi.items)
Base.lastindex(multi::Multi) = lastindex(multi.items)
Base.iterate(multi::Multi, state=1) =
  state > length(multi) ? nothing : (multi[state], state+1)

paramdim(multi::Multi) = maximum(paramdim, multi.items)

==(multi₁::Multi, multi₂::Multi) =
  length(multi₁) == length(multi₂) &&
  all(g -> g[1] == g[2], zip(multi₁, multi₂))

vertices(multi::Multi) =
  [vertex for geom in multi for vertex in vertices(geom)]

function centroid(multi::Multi)
  cs = coordinates.(centroid.(multi.items))
  Point(sum(cs) / length(cs))
end

measure(multi::Multi) = sum(measure, multi.items)

area(multi::Multi{Dim,T,<:Polygon}) where{Dim,T} = measure(multi)

boundary(multi::Multi) = Multi([boundary(geom) for geom in multi])

chains(multi::Multi{Dim,T,<:Polygon}) where {Dim,T} =
  [chain for geom in multi for chain in chains(geom)]

Base.in(point::Point, multi::Multi) = any(geom -> point ∈ geom, multi)

function Base.show(io::IO, multi::Multi{Dim,T}) where {Dim,T}
  n = length(multi.items)
  G = eltype(multi.items)
  print(io, "$n Multi-$(nameof(G)){$Dim,$T}")
end

function Base.show(io::IO, ::MIME"text/plain", multi::Multi)
  println(io, multi)
  print(io, io_lines(multi, "  "))
end


#!/usr/bin/env julia

export pe7

# need Primes.jl
# https://github.com/JuliaMath/Primes.jl
using Primes

pe7()=prime(10001)

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe7()|>print
end


using Test, SparseArrays
import LuxurySparse: SparseMatrixCOO, allocated_coo, isdense

coo1 = SparseMatrixCOO([1,4,2,3,3,3], [1,1,2,4,3,4], [0.1, 0.2, 0.4im, 0.5, 0.3, 0.5im], 4, 4)

@testset "basic" begin
    @test eltype(coo1) == ComplexF64
    @test size(coo1, 1) == size(coo1, 2) == 4

    coo2 = allocated_coo(ComplexF64, 4, 4, 6)
    @test eltype(coo2) == ComplexF64
    @test nnz(coo2) == 6
    @test size(coo2) == (4, 4)
    @test size(coo2, 1) == size(coo2, 2) == 4

    @test coo1 == copy(coo1)
    @test coo1 == copyto!(coo2, coo1)

    @test Matrix(coo1) == [0.1 0 0 0; 0 0.4im 0 0; 0 0 0.3 0.5+0.5im; 0.2 0 0 0]
    @test sparse(coo1) == sparse(findnz(coo1)...)
    @test isdense(coo1) == false

    A = sprand(50,50, 0.2) |> Matrix
    @test SparseMatrixCOO(A) == A
end

@testset "sparse" begin
    p1 = copy(coo1)
    @test nonzeros(p1) == p1.vs
    p1.vs[2] = 0
    @test dropzeros!(p1) == [0.1 0 0 0; 0 0.4im 0 0; 0 0 0.3 0.5+0.5im; 0 0 0 0]
end

@testset "basicmath" begin
    @test coo1*2 == Matrix(coo1)*2
    @test 2*coo1 == 2*Matrix(coo1)
    @test coo1/2 == Matrix(coo1)/2
    @test coo1*2 isa SparseMatrixCOO
    @test coo1/2 isa SparseMatrixCOO
    @test 2*coo1 isa SparseMatrixCOO
end


##################################################
# Check if the obtained transfer matrices are
# markov.
##################################################

"""
    is_markov(M::AbstractArray{T, 2}) -> Bool

Is the matrix Markov? """
function is_markov(M::AbstractArray{T, 2}) where T
    all(sum(M, dims = 2) .≈ 1)
end

"""
    is_markov(TO::AbstractTransferOperator) -> Bool

Is the transfer operator Markov? """
function is_markov(TO::AbstractTransferOperator)
    is_markov(TO.transfermatrix)
end

"""
    is_almostmarkov(M::AbstractArray{T, 2}; tol = 0.01) -> Bool

Is the matrix almost Markov?
"""
function is_almost_markov(M::AbstractArray{T, 2}; tol = 1e-3) where T
    all(sum(M, dims = 2) .> (1 - tol))
end

"""
    is_almostmarkov(TO::AbstractTransferOperator; tol = 0.01) -> Bool

Is the transfer operator almost Markov?
"""
function is_almost_markov(TO::AbstractTransferOperator; tol = 1e-3)
    is_almost_markov(TO.transfermatrix)
end

function zerocols(M::AbstractArray{T, 2}) where T
    findall(sum(M, dims = 1) .== 0)
end

function zerorows(M::AbstractArray{T, 2}) where T
    findall(sum(M, dims = 2) .== 0)
end

function zerocols(TO::AbstractTransferOperator)
    findall(sum(TO.transfermatrix, dims = 1) .== 0)
end

function zerorows(TO::AbstractTransferOperator)
    findall(sum(TO.transfermatrix, dims = 2) .== 0)
end

export
zerocols,
zerorows,
is_markov,
is_almost_markov

# #####################################################
# # How much do the exact and approximate triangulation
# # estimators differ?
# #####################################################
# """
#     max_discrep(to1::ExactSimplexTransferOperator,
#                 to2::ApproxSimplexTransferOperator)
#
# Measures the maximum discrepancy between an exact and an approximate estimate
# of the transfer operator.
#
# For this function to work, the transfer operators must have been generated from
# the same triangulation (otherwise, the number of simplices don't match up).
# """
# function max_discrep(exact::ExactSimplexTransferOperator,
#                      approx::ApproxSimplexTransferOperator)
#
# end


#---------------------------------------------------------
# # [ImageGeom](@id ImageGeom)
#---------------------------------------------------------

# This page explains
# [`ImageGeom`](https://github.com/JeffFessler/MIRT.jl/blob/master/src/fbp/image_geom.jl)
# in
# [MIRT.jl](http://github.com/JeffFessler/MIRT.jl).

# ### Setup

# Packages needed here.

using MIRT
using MIRTjim: jim, prompt
using Plots; default(markerstrokecolor=:auto)

# The following line is helpful when running this example.jl file as a script;
# this way it will prompt user to hit a key after each figure is displayed.

isinteractive() ? jim(:prompt, true) : prompt(:draw);

# ### Overview

# When performing tomographic image reconstruction,
# one must specify the geometry of the grid of image pixels.
# (In contrast, in image denoising and image deblurring problems,
# one works with the given discrete image
# and no physical coordinates are needed.)

# The key parameters of a grid of image pixels are
# * the size (dimensions) of the grid, e.g., `128 × 128`,
# * the spacing of the pixels, e.g., `1mm × 1mm`,
# * the offset of the pixels relative to the origin, e.g., `(0,0)`

# The data type `ImageGeom` describes such a geometry,
# for arbitrary dimensions (2D, 3D, etc.).

# There are several ways to construct this structure.
# The default is a `128 × 128` grid
# with pixel size ``\Delta_X = \Delta_Y = 1`` (unitless) and zero offset:

ig = ImageGeom()

# Here is a 3D example with non-cubic voxel size:

ig = ImageGeom( (512,512,128), (1,1,2), (0,0,0) )

# To avoid remembering the order of the arguments,
# named keyword pairs are also supported:

ig = ImageGeom( dims=(512,512,128), deltas=(1,1,2), offsets=(0,0,0) )

# ### Units

# The pixel dimensions `deltas` can (and should!) be values with units.

# Here is an example for a video (2D+time) with 12 frames per second:
using UnitfulRecipes
using Unitful: mm, s

ig = ImageGeom( dims=(640,480,1000), deltas=(1mm,1mm,(1//12)s) )

# ### Properties

# An ImageGeom object has *many* useful properties;
# there are too many to remember, so there is built-in help:

ig.help

# This small 2D example illustrates how some properties are used:

ig = ImageGeom( dims=(12,10), deltas=(1mm,1mm), offsets=(0.5,0.5) )

showgrid = (ig) -> # x,y grid locations of pixel centers
    scatter(ig.xg, ig.yg, label="", xlabel="x", ylabel="y",
        xlims = maximum(abs, ig.x) * 1.2 .* (-1,1),
        xticks=[ig.x[1], zero(eltype(ig.x)), ig.x[end]],
        ylims = maximum(abs, ig.y) * 1.2 .* (-1,1),
        yticks=[ig.y[1], zero(eltype(ig.y)), ig.y[end]],
        aspect_ratio=1, title="offsets $(ig.offsets)")
showgrid(ig)

#-
prompt();

# ### Offsets (unitless translation of grid)

# The default `offsets` are zeros,
# corresponding to symmetric sampling around origin:

ig = ImageGeom( dims=(12,10), deltas=(1mm,1mm) )
p = showgrid(ig)

#-
prompt();

# That default offset is natural for tomography
# when considering finite pixel size:

square = (x,y,Δ) -> plot!(p, label="", color=:black,
    x .+ Δ[1] * ([0,1,1,0,0] .- 0.5),
    y .+ Δ[2] * ([0,0,1,1,0] .- 0.5),
)
showgrid(ig)
square2 = (x,y) -> square(x, y, ig.deltas)
square2.(ig.xg,ig.yg)
plot!(p)

#-
prompt();

# In that default geometry, the center `(0,0)` of the image
# is at a corner of the middle 4 pixels (for even image sizes).
# That default is typical for tomographic imaging (e.g., CT, PET, SPECT).
# One must be careful when using operations like `imrotate` or `fft`.

# ### Mask

# In tomographic image reconstruction, patients are usually more "round"
# than "square" so often we only want to estimate the pixels inside some
# support `mask`: a Bool array indicating which pixels are to be estimated.
# (The rest are constrained to be zero.)
# The `ImageGeom` struct has an entry to store this `mask`.
# The default is `Trues(dims)` which is a "lazy" Bool `AbstractArray`
# from the `FillArrays` package that is conceptually similar to `trues(dims)`
# but requires `O(1)` storage.  So there is essentially no memory penalty
# to storing this entry in the `ImageGeom` for users who do not want
# to think about a `mask`.
# For users who do want a `mask`, fortunately Julia uses a special `BitArray`
# type to store Bool arrays, so the storage is 8× less than using bytes
# in most other languages.

# Often we use a "circle inscribed in the square" as a generic support mask,
# and one of the built-in properties can generate such a circular mask:

mask = ig.circ()
ig = ImageGeom(ig.dims, ig.deltas, ig.offsets, mask)
ig.plot(jim)

# Note that `jim` displays the axes with the units naturally;
# see [MIRTjim.jl](http://github.com/JeffFessler/MIRTjim.jl).

# There is an older interface `image_geom` similar to the function
# of the same name in
# [the Matlab version of MIRT](https://github.com/JeffFessler/mirt)
# provided for backward compatibility,
# but using `ImageGeom` is recommended for Julia work.

# ### AxisArrays

# There is a natural connection between `ImageGeom` and `AxisArrays`.
# Note the automatic labeling by
# [MIRTjim.jim](https://github.com/JeffFessler/MIRTjim.jl).

using AxisArrays
using Unitful: mm
ig = ImageGeom( dims=(60,48), deltas=(1.5mm,1mm) )
za = AxisArray( ig.circ() * 10/mm ; x=ig.x, y=ig.y)
jim(za, "AxisArray example")


struct Boundary{T}
    points::Vector{NTuple{2,T}}

    function Boundary(points::Vector{NTuple{2,T}}) where T
        if first(points) != last(points)
            push!(points,first(points))
        end
        return new{T}(points)
    end
end

const PlasmaBoundary = Boundary
const FluxSurface = Boundary
const Wall = Boundary

function Base.show(io::IO, b::Boundary)
    print(io,"Boundary: npoints = $(length(b.points))")
end

Base.broadcastable(b::Boundary) = (b,)

function Boundary(x::Vector,y::Vector)
    Boundary(collect(zip(x,y)))
end

function Base.getproperty(b::Boundary,s::Symbol)
    if s == :points
        return getfield(b,s)
    end
    if s in (:r,:x)
        v = getfield(b,:points)
        return getindex.(v,1)
    end
    if s in (:z,:y)
        v = getfield(b,:points)
        return getindex.(v,2)
    end
    error("type $(typeof(b)) has no field $s")
end

in_boundary(b::Boundary, p) = PolygonOps.inpolygon(p,b.points,in=true,on=true,out=false)
in_boundary(b::Boundary,x,y) = in_boundary(b,(x,y))

const in_vessel = in_boundary
const in_plasma = in_boundary

function boundary(M::AbstractEquilibrium, psi, dx=0.01,dy=0.01)
    xlims, ylims = limits(M)
    x = range(xlims...,step=dx)
    y = range(ylims...,step=dy)
    Psi = [M(xx,yy) for xx in x, yy in y]
    cl = Contour.contour(x,y,Psi,psi)
    l = Contour.lines(cl)[1]
    xc, yc = Contour.coordinates(l)
    return Boundary(xc,yc)
end

const flux_surface = boundary

function circumference(b::Boundary)
    p = b.points
    return sum(@inbounds norm(p[i+1] .- p[i]) for i=1:(length(p)-1))
end

function average(b::Boundary, F)
    p = b.points
    return sum(@inbounds norm(p[i+1] .- p[i])*(F(p[i][1],p[i][2])) for i=1:(length(p)-1))/circumference(b)
end

area(s::Boundary) = PolygonOps.area(s.points)

function surface_area(b::Boundary)
    p = b.points
    return 2pi*sum(@inbounds norm(p[i+1] .- p[i])*(p[i+1][1] + p[i][1])/2 for i=1:(length(p)-1))
end

function surface_area_average(b::Boundary, F)
    A = surface_area(b)
    p = b.points
    return (2pi/A)*sum(@inbounds norm(p[i+1] .- p[i])*((p[i+1][1] + p[i][1])/2)*
                   (F(p[i+1][1],p[i+1][2]) + F(p[i][1],p[i][2]))/2 for i=1:(length(p)-1))
end

function surface_area_average(b::Boundary, F::Vector)
    A = surface_area(b)
    p = b.points
    return (2pi/A)*sum(@inbounds norm(p[i+1] .- p[i])*((p[i+1][1] + p[i][1])/2)*
                       (F[i+1] + F[i])/2 for i=1:(length(p)-1))
end

function area_average(b::Boundary, F; dx=0.01, dy=0.01)
    x = range(extrema(b.r)...,step=dx) .+ dx/2
    y = range(extrema(b.z)...,step=dy) .+ dy/2

    A = zero(dx)
    for xx in x, yy in y
        if in_boundary(b,(xx,yy))
            A = A + F(xx,yy)*dx*dy
        end
    end
    return A/area(b)
end

function area_average(b::Boundary, F::AbstractMatrix, x::AbstractRange, y::AbstractRange)
    dx = step(x)
    dy = step(y)
    A = zero(eltype(x))
    @inbounds for ix in eachindex(x), iy in eachindex(y)
        if in_boundary(b, (x[ix],y[iy]))
            A = A + F[ix,iy]*dx*dy
        end
    end
    return A/area(b)
end

function volume(b::Boundary; dx=0.01,dy=0.01)
    x = range(extrema(b.r)...,step=dx) .+ dx/2
    y = range(extrema(b.z)...,step=dy) .+ dy/2

    V = zero(eltype(x))
    for xx in x, yy in y
        if in_boundary(b, (xx,yy))
            V = V + xx*dx*dy
        end
    end
    return 2pi*V
end

function volume_average(b::Boundary, F; dx=0.01, dy=0.01)
    x = range(extrema(b.r)...,step=dx) .+ dx/2
    y = range(extrema(b.z)...,step=dy) .+ dy/2

    A = zero(dx)
    for xx in x, yy in y
        if in_boundary(b, (xx,yy))
            A = A + F(xx,yy)*xx*dx*dy
        end
    end
    return 2pi*A/volume(b)
end

function volume_average(b::Boundary, F::AbstractMatrix, x::AbstractRange, y::AbstractRange)
    dx = step(x)
    dy = step(y)
    A = zero(eltype(x))
    @inbounds for ix in eachindex(x), iy in eachindex(y)
        if in_boundary(b, (x[ix],y[iy]))
            A = A + F[ix,iy]*x[ix]*dx*dy
        end
    end
    return 2pi*A/volume(b)
end


module thick_pipe_examples
using FinEtools
using PyCall

struct MyIData
    c::FInt
    r::FFltVec
    s::FFltVec
end

function thick_pipe_axi()
    println("Thick pipe with internal pressure: axially symmetric model")
    #=
    This is a simple modification of the full three-dimensional simulation of
    the tutorial pub_thick_pipe that implements the axially-symmetric model
    reduction procedure.

    An infinitely long thick walled cylindrical pipe
    with inner boundary radius of 3 mm and outer boundary radius of 9 mm is
    subjected to an internal pressure of 1.0 MPa. A wedge   with thickness of
    2 mm and a 90-degree angle sector is considered for the finite element
    analysis. The material properties are taken as  isotropic linear elastic
    with $E=1000$ MPa and $\nu=0.4999$ to represent nearly incompressible
    behavior. This problem has been proposed to by MacNeal and Harder as a
    test of an element's ability to represent the  response of a nearly
    incompressible material. The plane-strain condition is assumed in the
    axial direction of the pipe which together with the radial symmetry
    confines the material in all but the radial direction and therefore
    amplifies the numerical difficulties associated with the confinement of
    the nearly incompressible material.

    There is an analytical solution to this problem. Timoshenko and Goodier
    presented the original solution of Lame in their textbook. We are going
    to compare with  both the stress distribution (radial and hoop stresses)
    and the displacement of the inner  cylindrical surface.

    References:
    - Macneal RH, Harder RL (1985) A proposed standard set of problems to test
    finite element accuracy. Finite Elements in Analysis and Design 1: 3-20.
    - Timoshenko S. and Goodier J. N., Theory of Elasticity, McGraw-Hill, 2nd ed., 1951.

    =#

    # Internal radius of the pipe.
    a = 3*phun("MM");
    ##
    # External radius of the pipe.
    b = 9*phun("MM");
    ##
    # Thickness of the slice.
    t = 2*phun("MM");

    ##
    # Geometrical tolerance.
    tolerance   = a/10000.;
    ##
    # Young's modulus and Poisson's ratio.
    E = 1000*phun("MEGA*PA");
    nu = 0.499;
    ##
    # Applied pressure on the internal surface.
    press =   1.0*phun("MEGA*PA");

    ##
    # Analytical solutions.   Radial stress:
    radial_stress(r)  = press*a.^2/(b^2-a^2).*(1-b^2./r.^2);
    ##
    # Circumferential (hoop) stress:
    hoop_stress(r) = press*a.^2/(b^2-a^2).*(1+b^2./r.^2);

    ##
    # Radial displacement:
    radial_displacement(r) = press*a^2*(1+nu)*(b^2+r.^2*(1-2*nu))/(E*(b^2-a^2).*r);;

    ##
    # Therefore the radial displacement of the loaded surface will be:
    urex  =  radial_displacement(a);

    ##
    # The mesh parameters: The numbers of element edges axially,
    # and through the thickness of the pipe wall (radially).

    na = 1; nt = 10;

    ##
    # Note that the material object needs to be created with the proper
    # model-dimension reduction in effect.  In this case that is the axial symmetry
    # assumption.
    MR = DeforModelRed2DAxisymm
    axisymmetric = true

    # Create the mesh and initialize the geometry.  First we are going
    # to construct the block of elements with the first coordinate
    # corresponding to the thickness in the radial direction, and the second
    # coordinate is the thickness in the axial direction.
    fens,fes  =   Q8block(b-a, t, nt, na);

    # Extract the boundary  and mark the finite elements on the
    # interior surface.
    bdryfes = meshboundary(fes);

    bcl = selectelem(fens, bdryfes, box=[0.,0.,-Inf,Inf], inflate=tolerance);
    internal_fenids= connectednodes(subset(bdryfes,bcl));
    # Now  shape the block  into  the actual wedge piece of the pipe.
    for i=1:count(fens)
        fens.xyz[i,:] = fens.xyz[i,:] + [a; 0.0];
    end

    # now we create the geometry and displacement fields
    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz,1),2)) # displacement field

    # The plane-strain condition in the axial direction  is specified by selecting nodes
    # on the plane y=0 and y=t.
    l1 = selectnode(fens; box=[-Inf Inf 0.0 0.0], inflate = tolerance)
    setebc!(u, l1, true, 2, 0.0)
    l1 = selectnode(fens; box=[-Inf Inf t t], inflate = tolerance)
    setebc!(u, l1, true, 2, 0.0)

    applyebc!(u)
    numberdofs!(u)

    # The traction boundary condition is applied in the radial
    # direction.

    el1femm =  FEMMBase(IntegData(subset(bdryfes,bcl), GaussRule(1, 3), axisymmetric))
    fi = ForceIntensity([press; 0.0]);
    F2= distribloads(el1femm, geom, u, fi, 2);

    # Property and material
    material = MatDeforElastIso(MR,  E, nu)

    femm = FEMMDeforLinear(MR, IntegData(fes, GaussRule(2, 2), axisymmetric), material)

    K = stiffness(femm, geom, u)
    #K=cholesky(K)
    U =  K\(F2)
    scattersysvec!(u,U[:])

    # Transfer the solution of the displacement to the nodes on the
    # internal cylindrical surface and convert to
    # cylindrical-coordinate displacements there.
    uv=u.values[internal_fenids,:]
    # Report the  relative displacement on the internal surface:
    println("(Approximate/true displacement) at the internal surface: $( mean(uv[:,1])/urex*100  ) %")

    # Produce a plot of the radial stress component in the cylindrical
    # coordinate system. Note that this is the usual representation of
    # stress using nodal stress field.

    fld = fieldfromintegpoints(femm, geom, u, :Cauchy, 1)

    File =  "thick_pipe_sigmax.vtk"
    vtkexportmesh(File, fens, fes; scalars=[("sigmax", fld.values)])

    # Produce a plot of the solution components in the cylindrical
    # coordinate system.

    function inspector(idat::MyIData, elnum, conn, xe,  out,  xq)
        push!(idat.r, xq[1])
        push!(idat.s, out[idat.c])
        return idat
    end

    idat = MyIData(1, FFltVec[], FFltVec[])
    idat = inspectintegpoints(femm, geom, u, collect(1:count(fes)), inspector, idat, :Cauchy)

    @pyimport matplotlib.pyplot as plt
    plt.style[:use]("seaborn-whitegrid")
    fig = plt.figure()
    ax = plt.axes()
    # Plot the analytical solution.
    r = linspace(a,b,100);
    ax[:plot](r, radial_stress(r), linestyle="solid", marker=:o, label="Analytical")
    # Plot the computed  integration-point data
    ax[:plot](idat.r, idat.s, linestyle="none", marker=:o, label="FEA")
    # plt.xlim(0.0, 1.0)
    # plt.axis("equal")
    plt.legend()
    ax[:set_xlabel]("Radial distance")
    ax[:set_ylabel]("Radial stress")
    ax[:grid](linestyle="--", linewidth=0.5, color=".25", zorder=-10)
    plt.show()

    # ##
    # # *Regular quadratic triangle*
    # ##
    # # We start with the workhorse of most commonly used finite element
    # # packages, the quadratic triangle. Similarly to the quadratic
    # # tetrahedron in the 3-D version of this tutorial (pub_thick_pipe), the
    # # stress is polluted with oscillations. Definitely not as bad as in the
    # # plane-strain simulations when the triangles were distorted into
    # # shapes with curved edges, but the disturbances are there.
    # description ='T6';# tetrahedron
    # mf =@T6_block;
    # femmf =@(fes)femm_deformation_linear(struct('fes',fes,...
    #     'material',mater,'integration_rule',tri_rule(struct('npts',3))));
    # surface_integration_rule=gauss_rule(struct('dim',1, 'order', 3));
    # execute_simulation (description, mf, femmf, surface_integration_rule);

    # ##
    # # The same remedy of selective reduced integration as in full 3-D models
    # # will also work here.  Is demonstrated by the simulation with the
    # # selective reduced integration quadratic triangle.
    # ##
    # # *Selective reduced
    # # integration quadratic triangle*
    # description ='T6-SRI';
    # mf =@T6_block;
    # femmf =@(fes)femm_deformation_linear_sri(struct('fes',fes,...
    #     'material',mater,...
    #     'integration_rule_volumetric',tri_rule(struct('npts',1)),...
    #     'integration_rule_deviatoric',tri_rule(struct('npts',3))));
    # surface_integration_rule=gauss_rule(struct('dim',1, 'order', 3));
    # execute_simulation (description, mf, femmf, surface_integration_rule);

    # ##
    # # The selective reduced integration works very well with the T6 triangle.
    # ##
    # # An element that is often used in these situations is the uniformly
    # # under integrated serendipity (8-node) quadrilateral.
    # ##
    # # *Reduced integration serendipity quadrilateral*
    # ##
    # #  The same finite
    # # element model machine as above is used, and the integration is the 2
    # # x 2 Gauss rule (one order lower than that required for full
    # # integration which would be 3 x 3).
    # description ='Q8R';
    # mf =@Q8_block;
    # femmf =@(fes)femm_deformation_linear(struct('fes',fes,...
    #     'material',mater,...
    #     'integration_rule',gauss_rule(struct('dim',2, 'order',2))));
    # surface_integration_rule=gauss_rule(struct('dim',1, 'order',3));
    # execute_simulation (description, mf, femmf, surface_integration_rule);

    # ##
    # # *Full integration serendipity quadrilateral*
    # ##
    # # Using the full-integration Gauss rule of 3 x 3 points clearly leads to disaster.
    # description ='Q8';
    # mf =@Q8_block;
    # femmf =@(fes)femm_deformation_linear(struct('fes',fes,...
    #     'material',mater,...
    #     'integration_rule',gauss_rule(struct('dim',2, 'order',3))));
    # surface_integration_rule=gauss_rule(struct('dim',1, 'order',3));
    # execute_simulation (description, mf, femmf, surface_integration_rule);

    ##
    # The stress is now totally unacceptable.

    ## Discussion
    #
    ##
    # The axially symmetric model is clearly very effective
    # computationally, as the size is much reduced compared to the 3-D
    # model.  In conjunction with  uniform or selective reduced integration
    # it can be very accurate as well.

    #pub_thick_pipe_axi()
    # end

end # thick_pipe_axi

function thick_pipe_ps()

    ## Thick pipe with internal pressure: plane strain
    #

    ##
    # Link to the  <matlab:edit('pub_thick_pipe_ps') m-file>.

    ## Description
    ##
    # This is a simple modification of the full three-dimensional simulation of
    # the tutorial pub_thick_pipe takes advantage of the plane-strain model
    # reduction procedure.
    ##
    # An infinitely long thick walled cylindrical pipe
    # with inner boundary radius of 3 mm and outer boundary radius of 9 mm is
    # subjected to an internal pressure of 1.0 MPa. A wedge   with thickness of
    # 2 mm and a 90-degree angle sector is considered for the finite element
    # analysis. The material properties are taken as  isotropic linear elastic
    # with $E=1000$ MPa and $\nu=0.4999$ to represent nearly incompressible
    # behavior. This problem has been proposed to by MacNeal and Harder as a
    # test of an element's ability to represent the  response of a nearly
    # incompressible material. The plane-strain condition is assumed in the
    # axial direction of the pipe which together with the radial symmetry
    # confines the material in all but the radial direction and therefore
    # amplifies the numerical difficulties associated with the confinement of
    # the nearly incompressible material.
    ##
    # There is an analytical solution to this problem. Timoshenko and Goodier
    # presented the original solution of Lame in their textbook. We are going
    # to compare with  both the stress distribution (radial and hoop stresses)
    # and the displacement of the inner  cylindrical surface.

    ##
    #
    #
    #

    #
    #

    #
Figure 1. Definition of the geometry of the internally pressurized thick pipe

    #

    #

    ##
    # References:
    #
    # # Macneal RH, Harder RL (1985) A proposed standard set of problems to test
    # finite element accuracy. Finite Elements in Analysis and Design 1: 3-20.
    #
    # # Timoshenko S. and Goodier J. N., Theory of Elasticity, McGraw-Hill, 2nd ed., 1951.

    ## Solution
    #

    ##
    # Internal radius of the pipe.
    a = 3*phun("MM");
    ##
    # External radius of the pipe.
    b = 9*phun("MM");
    ##
    # Thickness of the slice.
    t = 2*phun("MM");

    ##
    # Geometrical tolerance.
    tolerance  =a/10000.;
    ##
    # Young's modulus and Poisson's ratio.
    E = 1000*phun("MEGA*PA");
    nu = 0.499;
    ##
    # Applied pressure on the internal surface.
    press = 1.0*phun("MEGA*PA");

    ##
    # Analytical solutions.   Radial stress:
    radial_stress(r) =press*a.^2/(b^2-a^2).*(1-(b^2)./r.^2);
    ##
    # Circumferential (hoop) stress:
    hoop_stress(r)=press*a.^2/(b^2-a^2).*(1+(b^2)./r.^2);

    ##
    # Radial displacement:
    radial_displacement(r)=press*a^2*(1+nu)*(b^2+r.^2*(1-2*nu))/(E*(b^2-a^2).*r);;

    ##
    # Therefore the radial displacement of the loaded surface will be:
    urex = radial_displacement(a);

    ##
    # The mesh parameters: The numbers of element edges axially,
    # and through the thickness of the pipe wall (radially).

    nc=3; nt=3;

    ##
    # Note that the material object needs to be created with the proper
    # model-dimension reduction in mind.  In this case that is the axial symmetry
    # assumption.
    MR = DeforModelRed2DStrain

    # Create the mesh and initialize the geometry.  First we are going
    # to construct the block of elements with the first coordinate
    # corresponding to the angle, and the second
    # coordinate is the thickness in the radial direction.
    anglrrange = 90.0/180*pi;
    fens,fes =  Q8block(anglrrange, b-a, nc, nt);

    # Extract the boundary  and mark the finite elements on the
    # interior surface.
    bdryfes = meshboundary(fes);
    bcl = selectelem(fens, bdryfes, box=[-Inf,Inf,0.,0.], inflate=tolerance);
    internal_fenids = connectednodes(subset(bdryfes,bcl));
    # Now  shape the block  into  the actual wedge piece of the pipe.
    ayr=fens.xyz;
    for i=1:count(fens)
        angl=ayr[i,1]; r=a+ayr[i,2];
        fens.xyz[i,:] = [r*sin(angl),(r*cos(angl))];
    end

    # now we create the geometry and displacement fields
    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz,1),2)) # displacement field

    # The symmetry boundary condition  is specified by selecting nodes
    # on the plane x=0.
    l1 = selectnode(fens; box=[0.0 0.0 -Inf Inf], inflate = tolerance)
    setebc!(u, l1, true, 1, 0.0)
    # The second symmetry boundary condition is specified by selecting
    # nodes on the plane y=0.
    l1 = selectnode(fens; box=[-Inf Inf 0.0 0.0], inflate = tolerance)
    setebc!(u, l1, true, 2, 0.0)

    applyebc!(u)
    numberdofs!(u)

    # The traction boundary condition is applied in the radial
    # direction.

    el1femm =  FEMMBase(IntegData(subset(bdryfes,bcl), GaussRule(1, 3)))
    function pressureloading!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        copy!(forceout, XYZ/norm(XYZ)*press)
        return forceout
    end
    fi = ForceIntensity(FFlt, 2, pressureloading!); # pressure normal to the internal cylindrical surface
    F2 = distribloads(el1femm, geom, u, fi, 2);

    # Property and material
    material = MatDeforElastIso(MR, E, nu)

    femm = FEMMDeforLinear(MR, IntegData(fes, GaussRule(2, 2)), material)

    K =stiffness(femm, geom, u)
    #K=cholesky(K)
    U=  K\(F2)
    scattersysvec!(u, U[:])

    # Transfer the solution of the displacement to the nodes on the
    # internal cylindrical surface and convert to
    # cylindrical-coordinate displacements there.
    uv = u.values[internal_fenids,:];
    ur = zeros(FFlt,length(internal_fenids));

    for j = 1:length(internal_fenids)
        n = fens.xyz[internal_fenids[j],:];
        n = n'/norm(n);# normal to the cylindrical internal surface
        ur[j] = dot(vec(uv[j,:]),vec(n));
    end

    # Report the  relative displacement on the internal surface:
    println("(Approximate/true displacement) at the internal surface: $( mean(ur)/urex*100  ) %")

    # Produce a plot of the radial stress component in the cylindrical
    # coordinate system. Note that this is the usual representation of
    # stress using nodal stress field.

    fld = fieldfromintegpoints(femm, geom, u, :Cauchy, 1)

    File =  "thick_pipe_sigmax.vtk"
    vtkexportmesh(File, fens, fes; scalars=[("sigmax", fld.values)])

    # Produce a plot of the solution components in the cylindrical
    # coordinate system.
    # Plot the analytical solution.

    function inspector(idat::MyIData, elnum, conn, xe,  out,  xq)
        function outputRm(c)
            theNormal=c;
            r=norm(theNormal);# distance from the axis of symmetry
            theNormal =theNormal/r;# compute the unit normal vector
            e1p=[theNormal';0.];# local cylind. coordinate  basis vectors
            e3p=[0.,0.,1.]';# this one points along the axis of the cylinder
            e2p=cross(vec(e3p),vec(e1p));# this one is along the hoop direction
            R= [vec(e1p) vec(e2p) vec(e3p)];# transformation matrix for the stress
            return R
        end
        Rm=outputRm(xq)
        tm=zeros(FFlt,3,3)
        stress4vto3x3t!(tm, out);# stress in global XYZ
        tpm = Rm'*tm*Rm;#  stress matrix in cylindrical coordinates
        sp=zeros(FFlt,6)
        stress3x3tto6v!(sp, tpm);# stress vector in cylindr. coord.
        push!(idat.r,norm(xq))
        push!(idat.s,sp[idat.c])
        return idat
    end

    idat = MyIData(1, FFltVec[], FFltVec[])
    idat = inspectintegpoints(femm, geom, u, collect(1:count(fes)),
    inspector, idat, :Cauchy)
    # show(idat)

    @pyimport matplotlib.pyplot as plt
    plt.style[:use]("seaborn-whitegrid")
    fig = plt.figure()
    ax = plt.axes()
    # Plot the analytical solution.
    r = linspace(a,b,100);
    ax[:plot](r, radial_stress(r), linestyle="solid", marker=:o, label="Analytical")
    # Plot the computed  integration-point data
    ax[:plot](idat.r, idat.s, linestyle="none", marker=:o, label="FEA")
    # plt.xlim(0.0, 1.0)
    # plt.axis("equal")
    plt.legend()
    ax[:set_xlabel]("Radial distance")
    ax[:set_ylabel]("Radial stress")
    ax[:grid](linestyle="--", linewidth=0.5, color=".25", zorder=-10)
    plt.show()

end # thick_pipe_ps

function thick_pipe_ps_T6()

    ## Thick pipe with internal pressure: plane strain
    #

    ##
    # Link to the  <matlab:edit('pub_thick_pipe_ps') m-file>.

    ## Description
    ##
    # This is a simple modification of the full three-dimensional simulation of
    # the tutorial pub_thick_pipe takes advantage of the plane-strain model
    # reduction procedure.
    ##
    # An infinitely long thick walled cylindrical pipe
    # with inner boundary radius of 3 mm and outer boundary radius of 9 mm is
    # subjected to an internal pressure of 1.0 MPa. A wedge   with thickness of
    # 2 mm and a 90-degree angle sector is considered for the finite element
    # analysis. The material properties are taken as  isotropic linear elastic
    # with $E=1000$ MPa and $\nu=0.4999$ to represent nearly incompressible
    # behavior. This problem has been proposed to by MacNeal and Harder as a
    # test of an element's ability to represent the  response of a nearly
    # incompressible material. The plane-strain condition is assumed in the
    # axial direction of the pipe which together with the radial symmetry
    # confines the material in all but the radial direction and therefore
    # amplifies the numerical difficulties associated with the confinement of
    # the nearly incompressible material.
    ##
    # There is an analytical solution to this problem. Timoshenko and Goodier
    # presented the original solution of Lame in their textbook. We are going
    # to compare with  both the stress distribution (radial and hoop stresses)
    # and the displacement of the inner  cylindrical surface.

    ##
    #
    #
    #

    #
    #

    #
Figure 1. Definition of the geometry of the internally pressurized thick pipe

    #

    #

    ##
    # References:
    #
    # # Macneal RH, Harder RL (1985) A proposed standard set of problems to test
    # finite element accuracy. Finite Elements in Analysis and Design 1: 3-20.
    #
    # # Timoshenko S. and Goodier J. N., Theory of Elasticity, McGraw-Hill, 2nd ed., 1951.

    ## Solution
    #

    ##
    # Internal radius of the pipe.
    a = 3*phun("MM");
    ##
    # External radius of the pipe.
    b = 9*phun("MM");
    ##
    # Thickness of the slice.
    t = 2*phun("MM");

    ##
    # Geometrical tolerance.
    tolerance  =a/10000.;
    ##
    # Young's modulus and Poisson's ratio.
    E = 1000*phun("MEGA*PA");
    nu = 0.499;
    ##
    # Applied pressure on the internal surface.
    press = 1.0*phun("MEGA*PA");

    ##
    # Analytical solutions.   Radial stress:
    radial_stress(r) =press*a.^2/(b^2-a^2).*(1-(b^2)./r.^2);
    ##
    # Circumferential (hoop) stress:
    hoop_stress(r)=press*a.^2/(b^2-a^2).*(1+(b^2)./r.^2);

    ##
    # Radial displacement:
    radial_displacement(r)=press*a^2*(1+nu)*(b^2+r.^2*(1-2*nu))/(E*(b^2-a^2).*r);;

    ##
    # Therefore the radial displacement of the loaded surface will be:
    urex = radial_displacement(a);

    ##
    # The mesh parameters: The numbers of element edges axially,
    # and through the thickness of the pipe wall (radially).

    nc=3; nt=3;

    ##
    # Note that the material object needs to be created with the proper
    # model-dimension reduction in mind.  In this case that is the axial symmetry
    # assumption.
    MR = DeforModelRed2DStrain

    # Create the mesh and initialize the geometry.  First we are going
    # to construct the block of elements with the first coordinate
    # corresponding to the angle, and the second
    # coordinate is the thickness in the radial direction.
    anglrrange = 90.0/180*pi;
    fens,fes =  T6block(anglrrange, b-a, nc, nt);

    # Extract the boundary  and mark the finite elements on the
    # interior surface.
    bdryfes = meshboundary(fes);
    bcl = selectelem(fens, bdryfes, box=[-Inf,Inf,0.,0.], inflate=tolerance);
    internal_fenids = connectednodes(subset(bdryfes,bcl));
    # Now  shape the block  into  the actual wedge piece of the pipe.
    ayr=fens.xyz;
    for i=1:count(fens)
        angl=ayr[i,1]; r=a+ayr[i,2];
        fens.xyz[i,:] = [r*sin(angl),(r*cos(angl))];
    end

    # now we create the geometry and displacement fields
    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz,1),2)) # displacement field

    # The symmetry boundary condition  is specified by selecting nodes
    # on the plane x=0.
    l1 = selectnode(fens; box=[0.0 0.0 -Inf Inf], inflate = tolerance)
    setebc!(u, l1, true, 1, 0.0)
    # The second symmetry boundary condition is specified by selecting
    # nodes on the plane y=0.
    l1 = selectnode(fens; box=[-Inf Inf 0.0 0.0], inflate = tolerance)
    setebc!(u, l1, true, 2, 0.0)

    applyebc!(u)
    numberdofs!(u)

    # The traction boundary condition is applied in the radial
    # direction.

    el1femm =  FEMMBase(IntegData(subset(bdryfes,bcl), GaussRule(1, 3)))
    function pressureloading!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        copy!(forceout, XYZ/norm(XYZ)*press)
        return forceout
    end
    fi = ForceIntensity(FFlt, 2, pressureloading!); # pressure normal to the internal cylindrical surface
    F2 = distribloads(el1femm, geom, u, fi, 2);

    # Property and material
    material = MatDeforElastIso(MR, E, nu)

    femm = FEMMDeforLinear(MR, IntegData(fes, TriRule(3)), material)

    K =stiffness(femm, geom, u)
    #K=cholesky(K)
    U=  K\(F2)
    scattersysvec!(u, U[:])

    # Transfer the solution of the displacement to the nodes on the
    # internal cylindrical surface and convert to
    # cylindrical-coordinate displacements there.
    uv = u.values[internal_fenids,:];
    ur = zeros(FFlt,length(internal_fenids));

    for j = 1:length(internal_fenids)
        n = fens.xyz[internal_fenids[j],:];
        n = n'/norm(n);# normal to the cylindrical internal surface
        ur[j] = dot(vec(uv[j,:]),vec(n));
    end

    # Report the  relative displacement on the internal surface:
    println("(Approximate/true displacement) at the internal surface: $( mean(ur)/urex*100  ) %")

    # Produce a plot of the radial stress component in the cylindrical
    # coordinate system. Note that this is the usual representation of
    # stress using nodal stress field.

    fld = fieldfromintegpoints(femm, geom, u, :Cauchy, 1)

    File =  "thick_pipe_sigmax.vtk"
    vtkexportmesh(File, fens, fes; scalars=[("sigmax", fld.values)])

    # Produce a plot of the solution components in the cylindrical
    # coordinate system.
    # Plot the analytical solution.

    function inspector(idat::MyIData, elnum, conn, xe,  out,  xq)
        function outputRm(c)
            theNormal=c;
            r=norm(theNormal);# distance from the axis of symmetry
            theNormal =theNormal/r;# compute the unit normal vector
            e1p=[theNormal';0.];# local cylind. coordinate  basis vectors
            e3p=[0.,0.,1.]';# this one points along the axis of the cylinder
            e2p=cross(vec(e3p),vec(e1p));# this one is along the hoop direction
            R= [vec(e1p) vec(e2p) vec(e3p)];# transformation matrix for the stress
            return R
        end
        Rm=outputRm(xq)
        tm=zeros(FFlt,3,3)
        stress4vto3x3t!(tm, out);# stress in global XYZ
        tpm = Rm'*tm*Rm;#  stress matrix in cylindrical coordinates
        sp=zeros(FFlt,6)
        stress3x3tto6v!(sp, tpm);# stress vector in cylindr. coord.
        push!(idat.r,norm(xq))
        push!(idat.s,sp[idat.c])
        return idat
    end

    idat = MyIData(1, FFltVec[], FFltVec[])
    idat = inspectintegpoints(femm, geom, u, collect(1:count(fes)),
    inspector, idat, :Cauchy)
    # show(idat)

    @pyimport matplotlib.pyplot as plt
    plt.style[:use]("seaborn-whitegrid")
    fig = plt.figure()
    ax = plt.axes()
    # Plot the analytical solution.
    r = linspace(a,b,100);
    ax[:plot](r, radial_stress(r), linestyle="solid", marker=:o, label="Analytical")
    # Plot the computed  integration-point data
    ax[:plot](idat.r, idat.s, linestyle="none", marker=:o, label="FEA")
    # plt.xlim(0.0, 1.0)
    # plt.axis("equal")
    plt.legend()
    ax[:set_xlabel]("Radial distance")
    ax[:set_ylabel]("Radial stress")
    ax[:grid](linestyle="--", linewidth=0.5, color=".25", zorder=-10)
    plt.show()

    #pub_thick_pipe_ps()

end # thick_pipe_ps_T6

function allrun()
    println("#####################################################")
    println("# thick_pipe_axi ")
    thick_pipe_axi()
    println("#####################################################")
    println("# thick_pipe_ps ")
    thick_pipe_ps()
    println("#####################################################")
    println("# thick_pipe_ps_T6 ")
    thick_pipe_ps_T6()
    return true
end # function allrun

end # module thick_pipe_examples


import Base: show, ==, push!, append!

## nice-valued ranges for histograms
function histrange{T}(v::AbstractArray{T}, n::Integer, closed::Symbol=:right)
    F = float(T)
    nv = length(v)
    if nv == 0 && n < 0
        throw(ArgumentError("number of bins must be ≥ 0 for an empty array, got $n"))
    elseif nv > 0 && n < 1
        throw(ArgumentError("number of bins must be ≥ 1 for a non-empty array, got $n"))
    elseif nv == 0
        return zero(F):zero(F)
    end

    lo, hi = extrema(v)
    histrange(F(lo), F(hi), n, closed)
end

function histrange{F}(lo::F, hi::F, n::Integer, closed::Symbol=:right)
    if hi == lo
        start = F(hi)
        step = one(F)
        divisor = one(F)
        len = one(F)
    else
        bw = (F(hi) - F(lo)) / n
        lbw = log10(bw)
        if lbw >= 0
            step = exp10(floor(lbw))
            r = bw / step
            if r <= 1.1
                nothing
            elseif r <= 2.2
                step *= 2
            elseif r <= 5.5
                step *= 5
            else
                step *= 10
            end
            divisor = one(F)
            start = step*floor(lo/step)
            len = ceil((hi - start)/step)
        else
            divisor = exp10(-floor(lbw))
            r = bw * divisor
            if r <= 1.1
                nothing
            elseif r <= 2.2
                divisor /= 2
            elseif r <= 5.5
                divisor /= 5
            else
                divisor /= 10
            end
            step = one(F)
            start = floor(lo*divisor)
            len = ceil(hi*divisor - start)
        end
    end
    # fix up endpoints
    if closed == :right #(,]
        while lo <= start/divisor
            start -= step
        end
        while (start + (len-1)*step)/divisor < hi
            len += one(F)
        end
    else
        while lo < start/divisor
            start -= step
        end
        while (start + (len-1)*step)/divisor <= hi
            len += one(F)
        end
    end
    FloatRange(start,step,len,divisor)
end

histrange{N}(vs::NTuple{N,AbstractVector},nbins::NTuple{N,Integer},closed::Symbol) =
    map((v,n) -> histrange(v,n,closed),vs,nbins)
histrange{N}(vs::NTuple{N,AbstractVector},nbins::Integer,closed::Symbol) =
    map(v -> histrange(v,nbins,closed),vs)

## histograms ##
function sturges(n)  # Sturges' formula
    n==0 && return one(n)
    ceil(Integer, log2(n))+1
end

abstract AbstractHistogram{T<:Real,N,E}

# N-dimensional histogram object
type Histogram{T<:Real,N,E} <: AbstractHistogram{T,N,E}
    edges::E
    weights::Array{T,N}
    closed::Symbol
    function Histogram(edges::NTuple{N,AbstractArray},weights::Array{T,N},closed::Symbol)
        closed == :right || closed == :left || error("closed must :left or :right")
        map(x -> length(x)-1,edges) == size(weights) || error("Histogram edge vectors must be 1 longer than corresponding weight dimensions")
        new(edges,weights,closed)
    end
end

Histogram{T,N}(edges::NTuple{N,AbstractVector},weights::AbstractArray{T,N},closed::Symbol=:right) =
    Histogram{T,N,typeof(edges)}(edges,weights,closed)

Histogram{T,N}(edges::NTuple{N,AbstractVector},::Type{T},closed::Symbol=:right) =
    Histogram(edges,zeros(T,map(x -> length(x)-1,edges)...),closed)

Histogram{N}(edges::NTuple{N,AbstractVector},closed::Symbol=:right) =
    Histogram(edges,Int,closed)

function show(io::IO, h::AbstractHistogram)
    println(io, typeof(h))
    println(io,"edges:")
    for e in h.edges
        println(io,"  ",e)
    end
    println(io,"weights: ",h.weights)
    print(io,"closed: ",h.closed)
end

(==)(h1::Histogram,h2::Histogram) = (==)(h1.edges,h2.edges) && (==)(h1.weights,h2.weights) && (==)(h1.closed,h2.closed)

# 1-dimensional
Histogram{T}(edge::AbstractVector, weights::AbstractVector{T}, closed::Symbol=:right) =
    Histogram((edge,), weights, closed)

Histogram{T}(edge::AbstractVector, ::Type{T}, closed::Symbol=:right) =
    Histogram(edge, zeros(T,length(edge)-1), closed)

Histogram(edge::AbstractVector,closed::Symbol=:right) =
    Histogram(edge, Int, closed)

function push!{T,E}(h::Histogram{T,1,E}, x::Real,w::Real)
    i = if h.closed == :right
        searchsortedfirst(h.edges[1], x) - 1
    else
        searchsortedlast(h.edges[1], x)
    end
    if 1 <= i <= length(h.weights)
        @inbounds h.weights[i] += w
    end
    h
end
push!{T,E}(h::AbstractHistogram{T,1,E}, x::Real) = push!(h,x,one(T))

function append!{T}(h::AbstractHistogram{T,1}, v::AbstractVector)
    for x in v
        push!(h,x)
    end
    h
end
function append!{T}(h::AbstractHistogram{T,1}, v::AbstractVector,wv::WeightVec)
    for (x,w) in zip(v,wv.values)
        push!(h,x,w)
    end
    h
end

fit(::Type{Histogram},v::AbstractVector, edg::AbstractVector; closed::Symbol=:right) =
    append!(Histogram(edg,closed), v)
fit(::Type{Histogram},v::AbstractVector; closed::Symbol=:right, nbins=sturges(length(v))) =
    fit(Histogram, v, histrange(v,nbins,closed); closed=closed)

fit{W}(::Type{Histogram},v::AbstractVector, wv::WeightVec{W}, edg::AbstractVector; closed::Symbol=:right) =
    append!(Histogram(edg,W,closed), v, wv)
fit(::Type{Histogram},v::AbstractVector, wv::WeightVec; closed::Symbol=:right, nbins=sturges(length(v))) =
    fit(Histogram, v, wv, histrange(v,nbins,closed); closed=closed)

# N-dimensional
function push!{T,N}(h::Histogram{T,N},xs::NTuple{N,Real},w::Real)
    is = if h.closed == :right
        map((edge, x) -> searchsortedfirst(edge,x) - 1, h.edges, xs)
    else
        map(searchsortedlast, h.edges, xs)
    end
    try
        h.weights[is...] += w
    catch e
        isa(e,BoundsError) || rethrow(e)
    end
    h
end
push!{T,N}(h::AbstractHistogram{T,N},xs::NTuple{N,Real}) = push!(h,xs,one(T))

function append!{T,N}(h::AbstractHistogram{T,N}, vs::NTuple{N,AbstractVector})
    for xs in zip(vs...)
        push!(h,xs)
    end
    h
end
function append!{T,N}(h::AbstractHistogram{T,N}, vs::NTuple{N,AbstractVector},wv::WeightVec)
    for (xs,w) in zip(zip(vs...),wv.values)
        push!(h,xs,w)
    end
    h
end

fit{N}(::Type{Histogram}, vs::NTuple{N,AbstractVector}, edges::NTuple{N,AbstractVector}; closed::Symbol=:right) =
    append!(Histogram(edges,closed), vs)
fit{N}(::Type{Histogram}, vs::NTuple{N,AbstractVector}; closed::Symbol=:right, nbins=sturges(length(vs[1]))) =
    fit(Histogram, vs, histrange(vs,nbins,closed); closed=closed)

fit{N,W}(::Type{Histogram}, vs::NTuple{N,AbstractVector}, wv::WeightVec{W}, edges::NTuple{N,AbstractVector}; closed::Symbol=:right) =
    append!(Histogram(edges,W,closed), vs, wv)
fit{N}(::Type{Histogram},vs::NTuple{N,AbstractVector}, wv::WeightVec; closed::Symbol=:right, nbins=sturges(length(vs[1]))) =
    fit(Histogram, vs, wv, histrange(vs,nbins,closed); closed=closed)


"Convert the indices in a graph to 32-bit ints.  This takes less storage, but does not speed up much"
shortIntGraph(a::SparseMatrixCSC) = SparseMatrixCSC{Float64,Int32}(convert(Int32,a.m), convert(Int32,a.n), convert(Array{Int32,1},a.colptr), convert(Array{Int32,1},a.rowval), a.nzval)

"Convert the nonzero entries in a graph to Float64"
floatGraph(a::SparseMatrixCSC) = SparseMatrixCSC{Float64,Int64}(a.m, a.n, a.colptr, a.rowval, convert(Array{Float64,1},a.nzval))

"""Create a Laplacian matrix from an adjacency matrix.
We might want to do this differently, say by enforcing symmetry"""
lap(a) = spdiagm(a*ones(size(a)[1])) - a

"""Create an adjacency matrix and a diagonal vector from a Laplacian with added diagonal weights"""
function adj{Tv,Ti}(la::SparseMatrixCSC{Tv,Ti})
  a = -(2 * la - triu(la) - tril(la))
  d = diag(la) - diag(lap(a))
  return a,d
end

"""Create a new graph in that is the same as the original, but with all edge weights 1"""
function unweight{Tval,Tind}(ain::SparseMatrixCSC{Tval,Tind})
    a = copy(ain)
    m = length(a.nzval)
    for i in 1:m
        a.nzval[i] = 1
    end
    return a
end # unweight

"""Change the weight of every edge in a to 1"""
function unweight!{Tval,Tind}(a::SparseMatrixCSC{Tval,Tind})
    m = length(a.nzval)
    for i in 1:m
        a.nzval[i] = 1
    end
end # unweight

"""Create a new graph that is the same as the original, but with f applied to each nonzero entry of a. For example, to make the weight of every edge uniform in [0,1], we could write

~~~julia
b = mapweight(a, x->rand(1)[1])
~~~
"""
function mapweight{Tval,Tind}(a::SparseMatrixCSC{Tval,Tind},f)
    b = triu(a,1)
    b.nzval = map(f, b.nzval)
    b = b + b'

    return b

end # mapweight

"""Put a uniform [0,1] weight on every edge.  This is an example of how to use mapweight."""
uniformWeight{Tval,Tind}(a::SparseMatrixCSC{Tval,Tind}) = mapweight(a,x->rand(1)[1])

"""Set the weight of every edge to 1"""
function uniformWeight!(mat::SparseMatrixCSC)
    mat.nzval = ones(length(mat.nzval))
end

"""The Cartesian product of two graphs.  When applied to two paths, it gives a grid."""
function productGraph(a0::SparseMatrixCSC, a1::SparseMatrixCSC)
  n0 = size(a0)[1]
  n1 = size(a1)[1]
  a = kron(speye(n0),a1) + kron(a0,speye(n1));

end # productGraph

"""The signed edge-vertex adjacency matrix"""
function edgeVertexMat(mat::SparseMatrixCSC)
    (ai,aj) = findnz(triu(mat,1))
    m = length(ai)
    n = size(mat)[1]
    return sparse(collect(1:m),ai,1.0,m,n) - sparse(collect(1:m),aj,1.0,m,n)
end

"""Create a new graph from the old, but keeping edge edge with probability `p`"""
function subsampleEdges(a::SparseMatrixCSC{Float64,Int64}, p::Float64)
  (ai, aj, av) = findnz(triu(a))
  n = size(a)[1]
  mask = map(x -> convert(Float64,x < p), rand(length(ai)))
  as = sparse(ai,aj,mask.*av,n,n)
  as = as + as'
  return as

end # subsampleEdges

"""Creats a 2-lift of a.  `flip` is a boolean indicating which edges cross"""
function twoLift(a, flip::AbstractArray{Bool,1})
    (ai,aj,av) = findnz(triu(a))
    m = length(ai)
    #flip = rand(false:true,m)
    n = size(a)[1]
    a0 = sparse(ai[flip],aj[flip],1,n,n)
    a1 = sparse(ai[!flip],aj[!flip],1,n,n)
    a00 = a0 + a0'
    a11 = a1 + a1'
    return [a00 a11; a11 a00]
end

twoLift(a) = twoLift(a,rand(false:true,div(nnz(a),2)))

twoLift(a, k::Integer) = twoLift(a,randperm(div(nnz(a),2)) .> k)

"""
 Create a disjoint union of graphs a and b,
 and then put k random edges between them
"""
function joinGraphs{Tval,Tind}(a::SparseMatrixCSC{Tval,Tind}, b::SparseMatrixCSC{Tval,Tind}, k::Integer)
    na = size(a)[1]
    nb = size(b)[1]

    (ai,aj,av) = findnz(a)
    (bi,bj,bv) = findnz(b)
    bi = bi + na
    bj = bj + na

    ji = rand(1:na,k)
    jj = rand(1:nb,k)+na

    ab = sparse([ai;bi;ji;jj],[aj;bj;jj;ji],[av;bv;ones(Tval,2*k)],na+nb,na+nb)
end

"""
 Create a disjoint union of graphs a and b,
  with no edges between them.
"""
disjoin(a,b) = joinGraphs(a,b,0)

"""Plots graph gr with coordinates (x,y)"""
function plotGraph(gr,x,y,color=[0,0,1];dots=true,setaxis=true,number=false)
  (ai,aj,av) = findnz(triu(gr))
  arx = [x[ai]';x[aj]';NaN*ones(length(ai))']
  ary = [y[ai]';y[aj]';NaN*ones(length(ai))']
  p = plot(arx[:],ary[:],color=color)

  if dots
    plot(x,y,color=color,marker="o",linestyle="none")
  end #if

  if number
    for i in 1:length(x)
      annotate(i, xy = [x[i]; y[i]])
    end
  end

  axis("off")

  if setaxis
  ax = axes()

  minx = minimum(x)
  maxx = maximum(x)
  miny = minimum(y)
  maxy = maximum(y)
  delx = maxx - minx
  dely = maxy - miny
  ax[:set_ylim]([miny - dely/20, maxy + dely/20])
  ax[:set_xlim]([minx - delx/20, maxx + delx/20])
  end

  return p
end # plotGraph

"""Computes spectral coordinates, and then uses plotGraph to draw"""
function spectralDrawing(a)

    x, y = spectralCoords(a)
    plotGraph(a,x,y)

end # spectralDrawing

"""Computes the spectral coordinates of a graph"""
function spectralCoords(a)

    E = eigs(lap(a), nev = 3, which=:SR)
    V = E[2]
    return V[:,2], V[:,3]

end # spectralCoords

"""Creates a unit vector of length n from a given set of integers, with weights based on the number of occurences"""
function toUnitVector(a::Array{Int64,1}, n)

  v = zeros(n)

  for i in a
    v[i] = v[i] + 1
  end

  sum = 0
  for i in 1:length(v)
    sum = sum + v[i] * v[i]
  end

  for i in 1:length(v)
    v[i] = v[i] / sqrt(sum)
  end

  return v

end # toUnitVector

"Returns the diagonal matrix(as a sparse matrix) of a graph"
function diagmat{Tv, Ti}(G::SparseMatrixCSC{Tv, Ti})

  dw = zeros(Tv, G.n)

  u,v,w = findnz(G)

  for i in 1:length(u)
    dw[v[i]] = dw[v[i]] + w[i]
  end

  return sparse(collect(1:G.n), collect(1:G.n), dw)

end # diagmat

"""
 Constructs a generalized necklace graph starting with two graphs A and H. The
resulting new graph will be constructed by expanding each vertex in H to an
instance of A. k random edges will be generated between components. Thus, the
resulting graph may have weighted edges.
"""
function generalizedNecklace{Tv, Ti}(A::SparseMatrixCSC{Tv, Ti}, H::SparseMatrixCSC, k::Int64)
  a = findnz(A)
  h = findnz(H)

  # these are square matrices
  n = A.n
  m = H.n

  newI = Ti[]
  newJ = Ti[]
  newW = Tv[]

  # duplicate the vertices in A so that each vertex in H corresponds to a copy of A
  for i in 1:m
    newI = append!(newI, a[1] + n * (i - 1))
    newJ = append!(newJ, a[2] + n * (i - 1))
    newW = append!(newW, a[3])
  end

  # for each edge in H, add k random edges between two corresponding components
  # multiedges will be concatenated to a single edge with higher cost
  for i in 1:length(h[1])
    u = h[1][i]
    v = h[2][i]

    if (u < v)
      #component x is from 1 + (x - 1) * n to n + (x - 1) * n
      for edgeToAdd in 1:k
        newU = rand(1:n) + n * (u - 1)
        newV = rand(1:n) + n * (v - 1)
        append!(newI, [newU, newV])
        append!(newJ, [newV, newU])
        append!(newW, [1, 1])
      end
    end
  end

  return sparse(newI, newJ, newW)
end # generalizedNecklace


#using python interface for Julia
using PyCall

#to test math functions and data strucrts from Python to Julia
@pyimport math

#using @test for unit testing

#test using Python Lists in Julia
let x = PyObject(test[1,2])

	@test pop!(x) == 2
	@test collect(x) == 1
	@time collect(x)
	@test collect(push!(x, 2)) == [1,2]
	@test collect(append!(x, 3,4,5)) == [1,2,3,4,5]
	@time append!(x, 3,4,5)
	@test collect(prepend!(x,0)) == [0,1,2,3,4,5]
	@test collect(pop!(x)) == [0,1,2,3,4]
	@time pop!(x)
	@test collect(append!(x,x)) == [0,1,2,3,4,0,1,2,3,4]

#quick test for strings
let y = PyObject(test2[one,two])

	@test collect(y) == [one,two]
	@test collect(push!(y,three)) == [one,two,three]
	@time push!(y,three)
	@test collect(pop!(y)) == [one,two]

	#this test above fails for some reason, probably wrong syntax

#basic matrix tests

	@test [10,20] + [10,20] == [30,30]
	@test [10,20] - [10,20] == [0,0]

#math sin functions

@test math.sin(math.pi / 4 - sin(pi / 4)) == 0
@test math.pi - pi == 0
@test math.pi/2 - pi/2 == 0

@test math.pow(2,2) == 4
@time math.pow(2,2)
#using Python syntax to run basic function

let x = 5
py"""
def testFunction(x):
	return 2 * x
"""
	@test py"testFunction"(5) == 10
	@time py"testFunction"(5)
	@test py"testFunction"(0) == 0

#passing Julia variable in python function

let x = 2
py"""
def testFunction2(x):
	return x * $x
"""

	@test py"testFunction2"(5) == 10
	# 2 * x
	@test py"2 * $x" == 4
	@time py"2 * $x"


using SparseArrays
using DelimitedFiles: readdlm
using ProgressBars
using PyPlot

function fast_kcore(nodes::Vector{Int64}, edges::Matrix{Int64})
    n = length(nodes)
    adjacency_list = Dict{Int64, Vector{Int64}}()
    for i in 1:n
        adjacency_list[i] = Vector{Int64}()
    end
    for edge in eachrow(edges)
         neighbors = get!(Vector{Int64}, adjacency_list, edge[1])
         push!(neighbors, edge[2])
         neighbors = get!(Vector{Int64}, adjacency_list, edge[2])
         push!(neighbors, edge[1])
    end

    degrees = zeros(Int64, n)

    for edge in eachrow(edges)
        degrees[edge[1]] += 1
        degrees[edge[2]] += 1
    end

    max_degree = maximum(degrees)
    eta = Dict{Int64, Int64}()
    c = 0
    for i in ProgressBar(1:n)
        v = argmin(degrees)
        c = max(c, degrees[v])
        eta[v] = c
        degrees[v] = max_degree + 1
        for neighbor in adjacency_list[v]
            degrees[neighbor] -= 1
        end
    end

    return eta
end

filename = "com-amazon.ungraph.txt"
edges = readdlm(filename, '\t', Int, comments=true)
nodes = Vector(1:maximum(edges))
@time core_numbers = fast_kcore(nodes, edges)
println("Core value for ", filename, " = ", maximum(values(core_numbers)))

filename = "com-lj.ungraph.txt"
edges = readdlm(filename, '\t', Int, comments=true)
edges = edges .+ 1
nodes = Vector(1:maximum(edges))
@time core_numbers = fast_kcore(nodes, edges)
println("Core value for ", filename, " = ", maximum(values(core_numbers)))

filename = "scholar/net.txt"
edges = readdlm(filename, '\t', Int, comments=true)
edges = edges .+ 1
nodes = Vector(1:maximum(edges))
@time core_numbers = fast_kcore(nodes, edges)

n = length(nodes)
degrees = zeros(Int64, n)
for edge in eachrow(edges)
    degrees[edge[1]] += 1
    degrees[edge[2]] += 1
end

coreness = zeros(Int64, n)
for key in keys(core_numbers)
    coreness[key] = core_numbers[key]
end

density = zeros(Int64, maximum(degrees), maximum(coreness))
for i in 1:n
    density_matrix[degrees[i], coreness[i]] += 1
end

a = Vector()
b = Vector()
c = Vector()

for element in findall(!iszero, density)
    push!(a, element[1])
    push!(b, element[2])
    push!(c, density[element])
end

figure()
title("Google scholar dataset")
scatter(x=a, y=b, c=c, s = c/maximum(c) * 2000)
xscale("log")
xlabel("Degree")
ylabel("Coreness")
colorbar()
plot(1:14, 1:14, c="black")


global rs = nothing
@info "Precompiling extra Reduce methods (set `ENV[\"REDPRE\"]=\"0\"` to disable)"
Reduce.Load(); atexit(() -> kill(rs))

rcall(:((1+pi)^2)) == convert(Expr,RExpr(rcall("(1+pi)**2")))
try; "1/0" |> rcall; false; catch; true; end
RExpr("(x+i)^3")
Reduce._syme(Reduce.r_to_jl)
R"x+2" == R"2+x-1+1"
:((x+1+π)^2; int(1/(1+x^3),x)) |> RExpr |> Reduce.parse
string(R"x+1")
RExpr(:x) == R"x"
Algebra.:*(RExpr(:x),R"x") == R"x^2"
convert(RExpr,R"x").str == convert(Array{String,1},R"x")
load_package([:rlfi]) == load_package(:rlfi,:rlfi)
rcall(:(x^2+2x+1),off=[:factor])
rcall("x + 1","factor")
Expr(:function,:(fun(z)),:(return begin; x = 700; y = x; end)) |> RExpr |> Reduce.parse
Expr(:for,:(i=2:34),:(product(i))) |> rcall
try; Expr(:type,false,:x) |> RExpr; false; catch; true; end
try; :(@time f(x)) |> RExpr; false; catch; true; end
Expr(:function,:(fun(a,b)),:(return y=a^3+3*a^2*b+3*a*b^2+b^3)) |> factor |> expand
try; Expr(:for,:(i=2:34),:(product(i))) |> RExpr |> parse; false; catch; true; end
R"begin; 1:2; end" |> Reduce.parse |> RExpr |> string
latex(:(x+1))
#Algebra.length(:(x+y))
Algebra.log(:(ℯ^x))
#Algebra.nextprime(100)
#Algebra.ceiling(1.2)
Algebra.impart(:(1+2*im))
#Algebra.impart(2+1.7im)
#Algebra.bernoulli(2)
Sys.islinux() && Reduce.RSymReplace("!#03a9; *x**2 + !#03a9;")
Algebra.int(:(x^2+y),:x) |> RExpr == Algebra.int("x^2+y","x") |> RExpr
R"/(2,begin 2; +(7,4); return +(4,*(2,7))+9 end)" |> Reduce.parse
Algebra.df(Expr(:function,:(fun(z,c)),:(return begin; zn = z^2+c; nz = z^3-1; end))|>RExpr,:z)
:([1 2; 3 4]) |> RExpr |> Reduce.parse |> RExpr == [1 2; 3 4] |> RExpr

Algebra.nextprime("3")
expand("(x-2)^2") |> RExpr == R"(x-2)^2"
nat("x+1")
Algebra.:^(:x,2) == :(x^2)
Algebra.://(NaN,NaN)
join(split(R"x+1;x+2"))
Algebra.sub(:x=>7,Algebra.:+(:x,7)) == Algebra.sub([:x=>7,:z=>21],Algebra.:-(:z,:x))
#squash(Expr(:function,:(fun(x)),:(z=3;z+=:x))).args[2] == squash(:(y=:x;y+=3))
squash(:(sqrt(x)^2))
Expr(:block,:(x+1)) |> RExpr == R"1+x"
Algebra.limit(Algebra.:^(Algebra.:-(1,Algebra.:/(1,:n)),Algebra.:-(:n)),:n,Inf)
Algebra.log(Algebra.exp(:pi))
Algebra.://(2,Inf)
Algebra.://(Inf,2)

rcall("x")
Algebra.operator(:x); Algebra.clear(:x);
#try Algebra.det([:x :y]) catch; true end
join([R"1",R"1"]) == R"1;1"
list([R"1",R"x"]) == list((1,:x))
Reduce.lister(:x) == R"x"
!latex(false)
@rounded @factor x^2-2x+1
@rounded @off_factor @rcall x^2
Algebra.det([:a :b; :c :d])
Algebra.tp([:a;:b]) == [:a :b;]
transpose(:x)
adjoint(:x)
Algebra.operator(:cbrt)
Algebra.rlet(:(cbrt(~x))=>:(x^(1/3)))
Algebra.rlet([:(cbrt(~x))=>:(x^(1/3))])
Algebra.rlet(Dict(:(cbrt(~x))=>:(x^(1/3))))
Algebra.clear(:cbrt)
Algebra.:+(1,R"x")
Algebra.inv([:a :b; :c :d])
Algebra.inv(1)
Algebra.:\([1],[2])
#Algebra.://(1.0,1.0)
Algebra.:/([:a :b; :c :d],2)
Algebra.:+([:a :b; :c :d],1) == Algebra.:+(1,[:a :b; :c :d])
Algebra.:+([:x],1) == Algebra.:+(1,[:x])
Algebra.:+([:x,:y]',1) == Algebra.:+(1,[:x,:y]')
Algebra.solve(:(x-1),:x) == Algebra.solve((:(x-1),),:x)
Algebra.order(nothing); Algebra.korder(nothing)


using SparseArrays, LinearAlgebra
using JLD, CSV, DataFrames, DataFramesMeta
#using DataDeps
using ArtifactUtils, Pkg.Artifacts

artifacts_path = joinpath(pkgdir(VarianceComponentsHDFE), "benchmark\\Artifacts.toml")

# add in code to generate .jld, .csv, etc. in the benchmark/data directory
# This should only generate if the file doesn't exist, or if this `force_generate = true`
force_generate = false

pkg_dir = pkgdir(VarianceComponentsHDFE)

# Obtain environment variables
run_large_benchmark = get(ENV, "VCHDFE_LARGE_BENCHMARK", "false")  == "true" ? true : false
run_huge_benchmark = get(ENV, "VCHDFE_HUGE_BENCHMARK", "false")  == "true" ? true : false

## Medium-Sized Network Generator
function rademacher!(R; demean = false)
    R .= R - (R .== 0)

    if demean == true
        R .= R .- mean(R, dims = 2)
    end
    return nothing
end

function compute_X_No_Controls(data)
    id = data.id
    firmid = data.firmid
    y = data.y

    NT = size(y,1);
    J = maximum(firmid);
    N = maximum(id);
    K = 0
    nparameters = N + J + K

    #Worker Dummies
    D = sparse(collect(1:NT),id,1);

    #Firm Dummies
    F = sparse(collect(1:NT),firmid,1);

    # N+J x N+J-1 restriction matrix
    S= sparse(1.0I, J-1, J-1);
    S=vcat(S,sparse(-zeros(1,J-1)));
    X_Laplacian = hcat(D, -F)
    X_GroundedLaplacian = hcat(D, -F*S)
    S_xx = Symmetric(X_GroundedLaplacian'*X_GroundedLaplacian)
    S_xx_lap = Symmetric(SparseMatrixCSC{Float64,Int64}(X_Laplacian'*X_Laplacian))

    A, diag = adj(SparseMatrixCSC{Float64,Int64}(X_Laplacian'*X_Laplacian))

    return X_Laplacian, X_GroundedLaplacian, S_xx, S_xx_lap, A
end

function compute_X_Controls(data)
    id = data.id
    firmid = data.firmid
    y = data.y

    NT=size(y,1);
    J=maximum(firmid);
    N=maximum(id);
    K = 2
    nparameters = N + J + K

    #Worker Dummies
    D=sparse(collect(1:NT),id,1);

    #Firm Dummies
    F=sparse(collect(1:NT),firmid,1);

    # N+J x N+J-1 restriction matrix
    S= sparse(1.0I, J-1, J-1);
    S=vcat(S,sparse(-zeros(1,J-1)));

    #Assuming columns 5 and 6 are the controls
    controls = hcat(data.control1[id], data.control2[id])

    Xcontrols = hcat(D,F*S,controls)
    S_xx = Symmetric(Xcontrols'*Xcontrols)

    return Xcontrols, S_xx
end

# We don't use the symmetric flag here, it is making the cluster kill the process
# Also flagged X_laplacian as a sparse matrix of Floats
function compute_X_No_Controls_Huge(data)
    id = data.id
    firmid = data.firmid
    y = data.y

    NT = size(y,1);
    J = maximum(firmid);
    N = maximum(id);
    K = 0
    nparameters = N + J + K

    #Worker Dummies
    D = sparse(collect(1:NT),id,1);

    #Firm Dummies
    F = sparse(collect(1:NT),firmid,1);

    # N+J x N+J-1 restriction matrix
    S= sparse(1.0I, J-1, J-1);
    S=vcat(S,sparse(-zeros(1,J-1)));
    X_Laplacian = SparseMatrixCSC{Float64,Int64}(hcat(D, -F))
    X_GroundedLaplacian = hcat(D, -F*S)
    S_xx = X_GroundedLaplacian'*X_GroundedLaplacian
    S_xx_lap = X_Laplacian'*X_Laplacian
    A, diag = adj(S_xx_lap)

    return X_Laplacian, X_GroundedLaplacian, S_xx, S_xx_lap, A
end

if ~isfile(pkg_dir*"/benchmark/data/medium_main.jld") || force_generate
    #data = CSV.read(datadep"VarianceComponentsHDFE/medium_nocontrols_pruned.csv"; header=true)
    data_hash = artifact_hash("medium_nocontrols_pruned",artifacts_path)
    data = CSV.read(joinpath(artifact_path(data_hash), "medium_nocontrols_pruned.csv"))

    X_Laplacian, X_GroundedLaplacian, S_xx , S_xx_lap, A = compute_X_No_Controls(data)
    save(pkg_dir*"/benchmark/data/medium_main.jld", "X_Laplacian", X_Laplacian, "X_GroundedLaplacian", X_GroundedLaplacian, "S_xx", S_xx, "S_xx_lap", S_xx_lap, "A", A)
end

if ~isfile(pkg_dir*"/benchmark/data/medium_controls_main.jld") || force_generate
    #data = CSV.read(datadep"VarianceComponentsHDFE/medium_controls_pruned.csv"; header=true)
    data_hash = artifact_hash("medium_controls_pruned",artifacts_path)
    data = CSV.read(joinpath(artifact_path(data_hash), "medium_controls_pruned.csv"))

    Xcontrols, S_xx = compute_X_Controls(data)
    save(pkg_dir*"/benchmark/data/medium_controls_main.jld", "Xcontrols", Xcontrols, "S_xx", S_xx)
end

if run_large_benchmark && (~isfile(pkg_dir*"/benchmark/data/large_main.jld") || force_generate)
    #data = CSV.read(datadep"VarianceComponentsHDFE/large_nocontrols_pruned.csv"; header=true)
    data_hash = artifact_hash("large_nocontrols_pruned",artifacts_path)
    data = CSV.read(joinpath(artifact_path(data_hash), "large_nocontrols_pruned.csv"))

    X_Laplacian, X_GroundedLaplacian, S_xx , S_xx_lap, A = compute_X_No_Controls(data)
    save(pkg_dir*"/benchmark/data/large_main.jld", "X_Laplacian", X_Laplacian, "X_GroundedLaplacian", X_GroundedLaplacian, "S_xx", S_xx, "S_xx_lap", S_xx_lap, "A", A)
 end

 if run_large_benchmark && (~isfile(pkg_dir*"/benchmark/data/large_controls_main.jld") || force_generate)
    #lar_csv = CSV.read(joinpath(artifact"large_nocontrols_pruned", "large_nocontrols_pruned.csv"))
    data_hash = artifact_hash("large_controls_pruned",artifacts_path)
    data = CSV.read(joinpath(artifact_path(data_hash), "large_controls_pruned.csv"))

    Xcontrols, S_xx = compute_X_Controls(data)
    save(pkg_dir*"/benchmark/data/large_controls_main.jld", "Xcontrols", Xcontrols, "S_xx", S_xx)
 end

 if run_huge_benchmark && (~isfile(pkg_dir*"/benchmark/data/huge_main.jld") || force_generate)
     #data = CSV.read(datadep"VarianceComponentsHDFE/huge_pruned_main.csv"; header=true)
     data_hash = artifact_hash("huge_nocontrols_pruned",artifacts_path)
     data = CSV.read(joinpath(artifact_path(data_hash), "huge_nocontrols_pruned.csv"))

     X_Laplacian, X_GroundedLaplacian, S_xx , S_xx_lap, A = compute_X_No_Controls_Huge(data)
     save(pkg_dir*"/benchmark/data/huge_main.jld", "X_Laplacian", X_Laplacian, "X_GroundedLaplacian", X_GroundedLaplacian, "S_xx", S_xx, "S_xx_lap", S_xx_lap, "A", A)
 end


module ApproximationGrids

### Imports
using FFTW

### Exports
export ChebyshevGrid, RegularGrid,
       interpolate, jacobian, nodes, weights

### Implemetation
abstract type ApproximationGrid end

# Objects' string representations
function Base.show(io::IO, g::T) where {T <: ApproximationGrid}
    print(io, "$(nameof(T))(($(g.a),$(g.b)), $(g.n))")
end

include("regular.jl")
include("chebyshev.jl")

end # module


import GR
using Triplot

include("common.jl")

function plot_contours(fname, n, f; levels=32)
    x,y,t = cartmesh(n)
    z = f.(x,y)

    GR.beginprint(fname)
    GR.setcolormap(1)
    GR.setlinewidth(1.0)
    tricontourf(x,y,z,t,levels)
    GR.setcolormapfromrgb([0.0,0.0],[0.0,0.0],[0.0,0.0])
    tricontour(x,y,z,t,levels)
    GR.endprint()
end

plot_contours(joinpath(@__DIR__, "diag.png"), 5, (x,y) -> x + y)
plot_contours(joinpath(@__DIR__, "peaks.png"), 100, peaks_unit)
plot_contours(joinpath(@__DIR__, "sine.png"), 100, (x,y) -> sin(2*pi*x)*cos(2*pi*y))


function extend_to_face(x,faceid,cell::IntervalBox{2})
    if faceid == 1
        return extend(x,2,cell[2].lo)
    elseif faceid == 2
        return extend(x,1,cell[1].hi)
    elseif faceid == 3
        return extend(x,2,cell[2].hi)
    elseif faceid == 4
        return extend(x,1,cell[1].lo)
    else
        throw(ArgumentError("Expected faceid ∈ {1,2,3,4}, got faceid = $faceid"))
    end
end

function extend_face_roots(roots,faceids,cell)
    extended_roots = hcat([extend_to_face(r,f,cell) for (r,f) in zip(roots,faceids)]...)
    return extended_roots
end

function roots_without_end(f,xL,xR,atol)
    r = sort!(find_zeros(f,xL,xR))
    if length(r) > 0
        isapprox(r[end],xR,atol=atol) ? pop!(r) : nothing
    end
    return r
end

function roots_without_start(f,xL,xR,atol)
    r = sort!(find_zeros(f,xL,xR))
    if length(r) > 0
        isapprox(r[1],xL,atol=atol) ? popfirst!(r) : nothing
    end
    return r
end

function roots_on_edges(funcs,xL,xR,yL,yR;atol=0.0)
    r1 = roots_without_end(funcs[1],xL,xR,atol)
    f1 = repeat([1],length(r1))

    r2 = roots_without_end(funcs[2],yL,yR,atol)
    f2 = repeat([2],length(r2))

    r3 = roots_without_start(funcs[3],xL,xR,atol)
    f3 = repeat([3],length(r3))

    r4 = roots_without_start(funcs[4],yL,yR,atol)
    f4 = repeat([4],length(r4))

    r = vcat(r1,r2,r3,r4)
    f = vcat(f1,f2,f3,f4)

    return r,f
end

function element_face_intersections(poly,cell::IntervalBox{2})

    xL,xR = (cell[1].lo,cell[1].hi)
    yL,yR = (cell[2].lo,cell[2].hi)

    funcs = restrict_on_faces(poly,cell)
    roots,faceids = roots_on_edges(funcs,xL,xR,yL,yR)

    @assert length(roots) == 2

    x1 = extend_to_face(roots[1],faceids[1],cell)
    x2 = extend_to_face(roots[2],faceids[2],cell)

    return x1,x2
end

function gradient_descent_to_zero!(x,P,∇P,atol,maxiter)
    @assert atol > 0.0
    @assert maxiter > 0

    val = P(x)
    iter = 1
    while abs(val) > atol && iter < maxiter
        dval = vec(∇P(x))
        alpha = val/dot(dval,dval)
        x .-= alpha*dval
        val = P(x)
        iter += 1
    end
    if iter == maxiter
        throw(ErrorException("Newton iteration failed to converge in $maxiter iterations"))
    end
end

function gradient_descent_to_zero!(x,P::InterpolatingPolynomial,atol,maxiter)
    gradient_descent_to_zero!(x,P,x->gradient(P,x),atol,maxiter)
end

function resolve_zero_levelset(poly,refpoints,xL,xR;atol=1e-12,maxiter=50)
    dim = dimension(poly)
    @assert dim == 2
    np = length(refpoints)

    points = zeros(2,np)
    dx = xR - xL

    for (idx,p) in enumerate(refpoints)
        x = xL + 0.5*(1.0+p)*dx
        gradient_descent_to_zero!(x,poly,atol,maxiter)
        points[:,idx] = x
    end
    return points
end

function fit_zero_levelset(poly,basis1d,quad1d,mass,cell;atol=1e-12,maxiter=50)
    xL,xR = element_face_intersections(poly,cell)
    intp = resolve_zero_levelset(poly,quad1d.points,xL,xR,atol=atol,maxiter=maxiter)
    rhs = linear_form(intp,basis1d,quad1d)
    return mass\rhs
end


const IT = Interpolations

function pairs(curve::Contour.Curve2)
    collect(zip(Contour.coordinates(curve)...))
end

function LinearInterpolation(xs,ys,A)
    itp = IT.interpolate(A, IT.BSpline(IT.Linear(IT.OnGrid())))
    IT.scale(itp, xs, ys)
end

function CubicInterpolation(xs,ys,A)
    itp = IT.interpolate(A, IT.BSpline(IT.Cubic(IT.Line(IT.OnGrid()))))
    IT.scale(itp, xs, ys)
end

const _DEFAULT_ISOLINE_KWARGS =
    OrderedDict(:colors => ["MidnightBlue",
                            "LightSeaGreen",
                            "Gold",
                            "Tomato"],
                :lift => false,
                :interpolation => :cubic,
		:levels => 10)

"""
    isolines(xs, ys, zs; lift = false, interpolation = :cubic)

Plot the contour lines of the function whose values are represented
by the array (or function) `zs`. If `lift` is true, plot in 3D.

# Examples
```julia-repl
julia> isolines(0:10, 0:10, (x,y) -> (100 - x^2 + y^2)/10, lift = true)
"""
function isolines(args...;kwargs...)
    D = deepcopy(_DEFAULT_ISOLINE_KWARGS)
    merge!(D,Dict(kwargs))
    isolinekwargs, penkwargs = splitkwargs(kwargs,_DEFAULT_ISOLINE_KWARGS)
    grlist = GraphicElement[]
    if length(args) ≥ 3 && args[3] isa Function
        args = (args[1],
                args[2],
                float([args[3](x,y) for x in args[1], y in args[2]]))
    end
    isolinedict = Dict(isolinekwargs)
    C = Contour.contours(args...,isolinedict[:levels])
    m, M = extrema([Contour.level(cl) for cl in Contour.levels(C)])
    if D[:lift]
        if D[:interpolation] == :linear
            interp = LinearInterpolation(args[1:3]...)
        elseif D[:interpolation] == :cubic
            interp = CubicInterpolation(args[1:3]...)
        else
            error("`interpolation` should be either :cubic or :linear")
        end
    end
    for cl in Contour.levels(C)
        L = Contour.level(cl)
        for line in Contour.lines(cl)
            pairs(line)
            pts = (D[:lift] ?
                   [(x,y,interp(x,y)) for (x,y) in pairs(line)] :
                   pairs(line))
            push!(grlist,
                  Path(pts; pen =
                       Pen(color=AsyPlots.cmap(NamedColor.(D[:colors]),
                                               (L-m)/(M-m));penkwargs...)))
        end
    end
    Plot(grlist)
end


using BenchmarkTools, ConstrainedSystems

Ns, Nc = 1000, 100
A,B2,B1t,C,rhs1v,rhs2v,solex = ConstrainedSystems.basic_linalg_problem(Ns=Ns,Nc=Nc);

Δt = 1.0e-2
prob1, xexact, yexact = ConstrainedSystems.basic_constrained_problem(tmax=3Δt)
prob2, xexact, yexact = ConstrainedSystems.cartesian_pendulum_problem(tmax=3Δt)

BenchmarkTools.DEFAULT_PARAMETERS.gcsample = true

SUITE = BenchmarkGroup()

SUITE["saddle setup"] = BenchmarkGroup()
SUITE["saddle solve"] = BenchmarkGroup()

rhs = SaddleVector(rhs1v,rhs2v)
sol = similar(rhs)

SUITE["saddle setup"]["basic problem"] = @benchmarkable ConstrainedSystems.SaddleSystem(A,B2,B1t,C,rhs)

As = ConstrainedSystems.SaddleSystem(A,B2,B1t,C,rhs)

SUITE["saddle solve"]["basic problem"] = @benchmarkable sol .= As\rhs

SUITE["timemarch setup"] = BenchmarkGroup()
SUITE["timemarch basic problem"] = BenchmarkGroup()
SUITE["timemarch pendulum problem"] = BenchmarkGroup()

SUITE["timemarch setup"]["IFHEEuler"] = @benchmarkable ConstrainedSystems.init(prob2, IFHEEuler(),dt=Δt)
SUITE["timemarch setup"]["LiskaIFHERK"] = @benchmarkable ConstrainedSystems.init(prob2, LiskaIFHERK(),dt=Δt)

integrator11 = ConstrainedSystems.init(prob1, IFHEEuler(),dt=Δt)
integrator12 = ConstrainedSystems.init(prob1, LiskaIFHERK(),dt=Δt)
integrator21 = ConstrainedSystems.init(prob2, IFHEEuler(),dt=Δt)
integrator22 = ConstrainedSystems.init(prob2, LiskaIFHERK(),dt=Δt)

#SUITE["basic problem"]["LiskaIFHERK"] = @benchmarkable solve(prob1, LiskaIFHERK(),dt=Δt)
#SUITE["basic problem"]["IFHEEuler"] = @benchmarkable solve(prob1, IFHEEuler(),dt=Δt)
#SUITE["pendulum problem"]["LiskaIFHERK"] = @benchmarkable solve(prob2, LiskaIFHERK(),dt=Δt)
#SUITE["pendulum problem"]["IFHEEuler"] = @benchmarkable solve(prob2, IFHEEuler(),dt=Δt)

SUITE["timemarch basic problem"]["IFHEEuler"] = @benchmarkable step!(integrator11,Δt)
SUITE["timemarch basic problem"]["LiskaIFHERK"] = @benchmarkable step!(integrator12,Δt)
SUITE["timemarch pendulum problem"]["IFHEEuler"] = @benchmarkable step!(integrator21,Δt)
SUITE["timemarch pendulum problem"]["LiskaIFHERK"] = @benchmarkable step!(integrator22,Δt)

run(SUITE)