## Physics:
##
## We simulate 2d flow of ideal gas through a porous medium
## with the help of Darcy's law.
##
## v⃗ = -μ⁻¹ K̂ ⋅∇P,
##
## where K̂, a second-order tensor in the general case, is
## the specific permeability. It depends only on the
## geometry of the mediumm. Here we assume isotropy
## of space, so K is a scalar. μ is the dynamic viscocity.
##
## Continuity equation:
##
## φ ⋅∂ρ / ∂t + div(ρv⃗) = 0
##
## where φ is the porosity.
##
## So, the equations used are the continuity equation and
## Darcy's equation, which replaces the momentum equation,
## used in our last program.
##
## Another change with respect to the ideal gas program
## is that we will set the boundary conditions using ghost cells.
## This way, we can make use of the central difference scheme for
## boundary conditions which consist of partial derivatives,
## which is second-order.
##
## The boundary conditions consist of walls on the top
## and bottom, and constant pressure on either side.
##
## For now, our velocities will be equal to 0 at the beginning.
##
## ---------
## ->     ->
## ->     ->
## ---------
##
## Initial conditions:
## v⃗ = (u, v) = 0⃗ everywhere
## P = P0 everywhere else
##
## Boundary conditions:
## u = 0 at x = 0, 2 (walls)
## P = Pin at y = 0 and x ∈ [0, 1]
## v = 0 at y = 0 and x ∈ [1, 2]
## P = Pout, dv / dy = 0 at y = 2
## ∂P / ∂y = 0 at y = 0, 2

using DelimitedFiles, PyPlot

function update!(; ρnext, unext, vnext, ρwork, uwork, vwork,
        Δt, Δx, Δy, Pin, Pout, φ, K, μ)
    nx, ny = size(ρnext)
    coeff = 8.314 * 298 / 0.029
    for j = 2 : ny - 1
        for i = 2 : nx - 1
            ## Auxiliary variables to avoid typos when
            ## working with partial derivatives
            ##
            ## NW--N--NE
            ##     |
            ## W---C---E
            ##     |
            ## SW--S--SE

            ρc = ρwork[i, j]
            ρe = ρwork[i + 1, j]
            ρw = ρwork[i - 1, j]
            ρn = ρwork[i, j + 1]
            ρs = ρwork[i, j - 1]

            uc = uwork[i, j]
            ue = uwork[i + 1, j]
            uw = uwork[i - 1, j]

            vc = vwork[i, j]
            vn = vwork[i, j + 1]
            vs = vwork[i, j - 1]

            ## Continuity equation to calculate ρ
            ## φ ⋅∂ρ / ∂t + div(ρv⃗) = 0
            ρnext[i, j] = ρc - Δt / φ *
                (uc * (ρe - ρw) / (2 * Δx) +
                 vc * (ρn - ρs) / (2 * Δy) +
                 ρc * (ue - uw) / (2 * Δx) +
                 ρc * (vn - vs) / (2 * Δy))
        end
    end

    ## Boundary conditions on pressure using ghost cells
    ## P = Pin at y = 0 and x ∈ [0, 1]
    @. @views ρnext[1: nx ÷ 2, 1] = 2 * Pin / coeff - ρnext[1: nx ÷ 2, 2]

    ## P = Pout at y = 2
    @. @views ρnext[:, ny] = 2 * Pout / coeff - ρnext[:, ny - 1]

    ## ∂P / ∂x = 0 at x = 0, 2
    @. @views ρnext[nx, :] = ρnext[nx - 1, :]
    @. @views ρnext[1, :] = ρnext[2, :]

    ## ∂P / ∂y = 0 at y = 0 x ∈ [1, 2]
    @. @views ρnext[nx ÷ 2 + 1: nx, 1] = ρnext[nx ÷ 2 + 1: nx, 2]

    ## Here we use a central difference scheme to discretize the second
    ## order derivatives, so we iterate over the 'inner' grid
    for j = 2 : ny - 1
        for i = 2 : nx - 1
            ## Auxiliar variables
            ## можно через (ρwork + ρnext) / 2
            ρe = ρwork[i + 1, j]
            ρw = ρwork[i - 1, j]
            ρn = ρwork[i, j + 1]
            ρs = ρwork[i, j - 1]

            ## Darcy's law
            ## v⃗ = -μ⁻¹ K̂ ⋅∇P
            unext[i, j] = -K / μ * coeff * (ρe - ρw) / (2 * Δx)
            vnext[i, j] = -K / μ * coeff * (ρn - ρs) / (2 * Δy)
        end
    end

    ## Boundary conditions on velocity, again using ghost cells
    ## u = 0 at x = 0, 2 (walls)
    @. @views unext[1, :] = -unext[2, :]
    @. @views unext[nx, :] = -unext[nx - 1, :]

    ## v = 0 at y = 0 and x ∈ [1, 2]
    @. @views vnext[1, nx ÷ 2 + 1: nx] = -vnext[2, nx ÷ 2 + 1: nx]

    ## dv / dy = 0 at y = 2
    @views vnext[:, ny] = vnext[:, ny - 1]

    ## dv / dy = 0 at y = 0 and x ∈ [0, 1]
    @views vnext[1: nx ÷ 2, 1] =  vnext[1: nx ÷ 2, 2]
end

function filtration!(; ρ, u, v, Δx, Δy, Δt, nsteps, K, φ, μ, Pin, Pout)
    ## Helper arrays to store values from previous iteration
    ρwork, uwork, vwork = ρ, u, v
    unext = copy(u)
    vnext = copy(v)
    ρnext = copy(ρ)

    for t = 1 : nsteps
        update!(; ρnext, unext, vnext, ρwork, uwork, vwork,
                Δt, Δx, Δy, Pin, Pout, φ, K, μ)

        ## Swap next and current and continue iteration
        ρnext, ρwork = ρwork, ρnext
        unext, uwork = uwork, unext
        vnext, vwork = vwork, vnext
    end

    ## If in the end ρwork is not the original ρ,
    ## copy the most recent data into ρ, u, v
    if ρwork !== ρ
        ρ .= ρwork
        u .= uwork
        v .= vwork
    end

    return ρ, u, v
end

function plot(u, v, ρ)
    fig = PyPlot.figure(figsize=(10, 10))
    ax = PyPlot.axes()
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    #legend(loc="upper right",fancybox="true")

    nx, ny= size(u)

    rgx = range(0, 2, length=nx)
    rgy = range(0, 2, length=ny)
    x = [rgx;]'
    y = [rgy;]'
    X = repeat([rgx;]', length(x))
    Y = repeat([rgy;],1, length(y))

    ## Velocity vector field
    ## scale = 100.0
    quiver(x, y, u', v')

    ## Streamplot of velocity vector field
    #streamplot(x, y, u', v')

    ## Density contour map
    pos = ax.contourf(X, Y, ρ', alpha=0.5, cmap=matplotlib.cm.viridis)
    fig.colorbar(pos, ax=ax)
    cp = contour(X, Y, ρ', cmap=matplotlib.cm.viridis)

    PyPlot.title("2d filtration of ideal gas through a porous medium")

    # PyPlot.svg(true)

    savefig("img/porous-medium-darcy.png")
    savefig("img/porous-medium-darcy.svg")
end

let
    duration = 2000
    Δt = 0.001
    nsteps = round(Int64, duration / Δt)
    # nsteps = 1000

    ## Space grid [0, 2] × [0, 2]
    Δx = 0.05
    Δy = Δx
    nx = convert(Int64, 2 / Δx)
    ny = convert(Int64, 2 / Δy)

    ## Porosity
    φ = 0.7

    ## Specific permeability
    K = 1e-12

    ## Dynamic viscocity (for air)
    μ = 18e-6

    ## P = ρRT / μ, (μ - molar mass)
    coeff = 8.314 * 298 / 0.029

    ## Pressures at top and bottom, and initial pressure
    Pin = 1e6
    Pout = 1e5
    P0 = Pout

    u = zeros(nx, ny)
    v = zeros(nx, ny)
    ρ = zeros(nx, ny)

    ρ .= P0 / coeff

    problem = (
               ρ = ρ,
               u = u,
               v = v,
               Δx = Δx,
               Δy = Δy,
               Δt = Δt,
               nsteps = nsteps,
               K = K,
               μ = μ,
               φ = φ,
               Pin = Pin,
               Pout = Pout
              )

    ## Run our solver
    ρ, u, v = filtration!(; problem...)

    # writedlm("density.txt", ρ, ' ')

    ## Plotting
    plot(u, v, ρ)
end


# Utility function which returns bitstring mask for one output of the
# packed representation.
# Examples:  output_mask(2) returns  0b1111 = 0XF
#        and output_mask(3) returns  0b11111111 = 0XFF
function output_mask(num_inputs)
    mask = convert(BitString, 0x1)
    for i in 0:(num_inputs-1)
        mask |= (mask << 2^i)
    end
    return mask
end

# Returns an array of BitStrings which are the standard inputs to a chromosome
#    when it is executed.
# num_inputs  is the number of inputs of the chromosome.
function std_input_context(num_inputs)
   if 2^num_inputs > 8*sizeof(BitString)
      error("number of inputs is too large for size of BitString in std_input_context")
   end
   std_in = Array(BitString,num_inputs)
   std_in[1] = 0b10
   for i in 2:num_inputs
      for j in i:-1:2
         std_in[j] = std_in[j-1] $ std_in[j-1] << 2^(i-1)
      end
      std_in[1] = output_mask(i-1) << 2^(i-1)
   end
   return std_in
end


function make_landmarks(nlandmarks::Int, boundaries, margin)
    xmin, xmax, ymin, ymax = boundaries
    bx = margin*(xmax - xmin)
    by = margin*(ymax - ymin)
    rand([xmin + bx:xmax - bx; ymin + by:ymax - by], 2, nlandmarks)
end

"""
Return indices of landmarks within sensor acceptance (r < rmax, phi < +/)
"""
function nearby_landmark_indices(pose, lm, rmax)
    dx = lm[1,:] - pose[1]
    dy = lm[2,:] - pose[2]
    phi = pose[3]

    inearby = Array{Int}(0)
    for i in 1:size(lm, 2)
        # Select landmarks within +/- 90 degrees of current heading
        if (dx[i]*cos(phi) + dy[i]*sin(phi)) > 0
            # Select landmarks inside sensor range
            if ((dx[i]^2 + dy[i]^2) < rmax^2)
                push!(inearby, i)
            end
        end
    end
    return inearby
end

"""
Add Gaussian noise to nominal control values
Noise is taken to be uncorrelated (assume diagonal control_cov)
"""
function add_control_noise!(vehicle::Vehicle, control_cov::AbstractMatrix)
    vehicle.measured_speed = vehicle.target_speed + randn(1)[1]*sqrt(control_cov[1,1])
    vehicle.measured_gamma = vehicle.target_gamma + randn(1)[1]*sqrt(control_cov[2,2])
end

function add_control_noise(speed, gamma, control_cov)
    # Add Gaussian noise to nominal control values
    # Noise is taken to be uncorrelated (assume diagonal control_cov)
    speed + randn(1)[1]*sqrt(control_cov[1,1]),
    gamma + randn(1)[1]*sqrt(control_cov[2,2])
end

"""
Return array with columns z = [range; angle] for all nearby landmark
observations as well as the corresponding list of landmark IDs.
"""
function get_observations(vehicle::Vehicle, scene::Scene, R)
    lm = scene.landmarks
    x,y,phi = vehicle.pose

    inearby = nearby_landmark_indices(vehicle.pose, lm, vehicle.sensor_range)

    dx  = lm[1, inearby] - x
    dy  = lm[2, inearby] - y

    # 2 by nz observation matrix - columns are range and bearing
    z = [sqrt(dx.^2 + dy.^2) atan2(dy, dx) - phi]'

    # Add Gaussian noise to z to simulate measurement resolution
    ncols = size(z, 2)
    if ncols > 0
        z = z + [randn(1, ncols)*sqrt(R[1,1]); randn(1, ncols)*sqrt(R[2,2])]
    end

    # Unique identifier for each landmark
    landmark_tags = 1:size(lm, 2)

    return z, landmark_tags[inearby]
end


module LassoTest

const  NUMBER_REPEAT = 1

using CoordinateDescent
using Test
using ProximalBase
using Random
using LinearAlgebra
using SparseArrays

Random.seed!(1)

##############################################
#
#  Lasso
#
##############################################

@testset "lasso" begin

  @testset "zero" begin
    n = 100
    p = 10

    X = randn(n, p)
    Y = X * ones(p) + 0.1 * randn(n)
    Xy = X' * Y / n

    lambda = maximum(abs.(Xy)) + 0.1
    out = lasso(X, Y, lambda)
    @test out.x == SparseIterate(p)
  end

  @testset "non-zero" begin
    for i=1:NUMBER_REPEAT
      n = 100
      p = 10
      s = 5

      X = randn(n, p)
      Y = X[:,1:s] * ones(s) + 0.1 * randn(n)

      λ = fill(0.3, p)
      beta = lasso(X, Y, 1., λ, CDOptions(;optTol=1e-12))

      f = CDQuadraticLoss(X'X/n, -X'Y/n)
      g = ProximalBase.ProxL1(1., λ)
      x1 = SparseIterate( p )
      coordinateDescent!(x1, f, g, CDOptions(;optTol=1e-12))
      @test beta.x ≈ x1 atol=1e-5

      @test (maximum(abs.(X'*(Y - X*beta.x) / n)) - 0.3) / 0.3 ≈ 0. atol=1e-5
    end
  end

  @testset "different interfaces" begin
    n = 500
    p = 500
    s = 50

    X = randn(n, p)
    β = randn(s)
    Y = X[:,1:s] * β + randn(n)

    λ = 0.1
    x1 = lasso(X, Y, λ)
    x2 = lasso(X, Y, λ, ones(p))

    @test Vector(x1.x) ≈ Vector(x2.x) atol=1e-5
  end

end

@testset "cd lasso" begin

  for i=1:NUMBER_REPEAT
    n = 200
    p = 50
    s = 10

    X = randn(n, p)
    β = randn(s)
    Y = X[:,1:s] * β + 0.1 * randn(n)

    g = ProximalBase.ProxL1(0.2)
    f1 = CDQuadraticLoss(X'X/n, -X'Y/n)
    f2 = CDLeastSquaresLoss(Y, X)

    x1 = SparseIterate(p)
    x2 = SparseIterate(p)
    coordinateDescent!(x1, f1, g, CDOptions(;optTol=1e-12))
    coordinateDescent!(x2, f2, g, CDOptions(;optTol=1e-12))

    @test maximum(abs.(x1 - x2)) ≈ 0. atol=1e-5
    @test (maximum(abs.(X'*(Y - X*x1) / n)) - 0.2) / 0.2 ≈ 0. atol=1e-5
    @test (maximum(abs.(X'*(Y - X*x2) / n)) - 0.2) / 0.2 ≈ 0. atol=1e-5
  end

end

@testset "cd sqrt-lasso" begin

  @testset "kkt" begin
    for i=1:NUMBER_REPEAT
      n = 100
      p = 50
      s = 5

      X = randn(n, p)
      β = randn(s)
      Y = X[:,1:s] * β + randn(n)

      λ = 2.8
      g = ProximalBase.ProxL1(λ)
      f = CDSqrtLassoLoss(Y, X)

      x1 = SparseIterate(p)
      coordinateDescent!(x1, f, g, CDOptions(;maxIter=5000, optTol=1e-8))

      @test max(0, maximum(abs.(X'*(Y - X*x1) / norm(Y - X*x1))) - λ) / λ  ≈ 0. atol=1e-3
    end
  end

  @testset "interfaces" begin
    for i=1:NUMBER_REPEAT
      n = 500
      p = 500
      s = 50

      X = randn(n, p)
      β = randn(s)
      Y = X[:,1:s] * β + randn(n)

      opt1 = CDOptions(;maxIter=5000, optTol=1e-10, warmStart=true, randomize=false)
      opt2 = CDOptions(;maxIter=5000, optTol=1e-10, warmStart=true, randomize=true)
      opt3 = CDOptions(;maxIter=5000, optTol=1e-10, warmStart=false, randomize=false)
      opt4 = CDOptions(;maxIter=5000, optTol=1e-10, warmStart=false, randomize=true)

      x1 = SparseIterate(sprand(p, 0.6))
      x2 = SparseIterate(sprand(p, 0.6))
      x3 = SparseIterate(sprand(p, 0.6))
      x4 = SparseIterate(sprand(p, 0.6))

      λ = 1.5
      g = ProximalBase.ProxL1(λ)
      f = CDSqrtLassoLoss(Y, X)

      coordinateDescent!(x1, f, g, opt1)
      coordinateDescent!(x2, f, g, opt2)
      coordinateDescent!(x3, f, g, opt3)
      coordinateDescent!(x4, f, g, opt4)

      @test Vector(x1) ≈ Vector(x2) atol=1e-4
      @test Vector(x3) ≈ Vector(x2) atol=1e-4
      @test Vector(x4) ≈ Vector(x2) atol=1e-4

      y1 = sqrtLasso(X, Y, λ, opt1, standardizeX=false)
      y2 = sqrtLasso(X, Y, λ, opt2, standardizeX=false)
      y3 = sqrtLasso(X, Y, λ, opt3, standardizeX=false)
      y4 = sqrtLasso(X, Y, λ, opt4, standardizeX=false)

      @test Vector(y1.x) ≈ Vector(x2) atol=1e-4
      @test Vector(y2.x) ≈ Vector(x2) atol=1e-4
      @test Vector(y3.x) ≈ Vector(x2) atol=1e-4
      @test Vector(y4.x) ≈ Vector(x2) atol=1e-4

      z1 = sqrtLasso(X, Y, λ, ones(p), opt1)
      z2 = sqrtLasso(X, Y, λ, ones(p), opt2)
      z3 = sqrtLasso(X, Y, λ, ones(p), opt3)
      z4 = sqrtLasso(X, Y, λ, ones(p), opt4)

      @test Vector(z1.x) ≈ Vector(x2) atol=1e-4
      @test Vector(z2.x) ≈ Vector(x2) atol=1e-4
      @test Vector(z3.x) ≈ Vector(x2) atol=1e-4
      @test Vector(z4.x) ≈ Vector(x2) atol=1e-4

    end
  end

end

@testset "scaled lasso" begin

  for i=1:NUMBER_REPEAT
    n = 1000
    p = 500
    s = 50

    X = randn(n, p)
    β = randn(s)
    Y = X[:,1:s] * β + randn(n)

    λ = rand() / 5.

    opt1 = IterLassoOptions(;maxIter=100, optTol=1e-8, optionsCD=CDOptions(;maxIter=5000, optTol=1e-8))
    opt2 = IterLassoOptions(;maxIter=100, optTol=1e-8, initProcedure=:InitStd, σinit=2., optionsCD=CDOptions(;maxIter=5000, optTol=1e-8))

    x1 = SparseIterate(p)
    x2 = SparseIterate(p)
    sol1 = scaledLasso!(x1, X, Y, λ, ones(p), opt1)
    sol2 = scaledLasso!(x2, X, Y, λ, ones(p), opt2)

    σ1 = sol1.σ
    σ2 = sol2.σ

    @test max.((maximum(abs.(X'*(Y - X*x1) / n)) - λ*σ1), 0.) / (σ1*λ) ≈ 0. atol=1e-4
    @test max.((maximum(abs.(X'*(Y - X*x2) / n)) - λ*σ2), 0.) / (σ2*λ) ≈ 0. atol=1e-4
    @test Vector(x1) ≈ Vector(x2) atol=1e-4
  end

end

@testset "lasso path" begin

  @testset "standardizeX = false" begin
    n = 1000
    p = 500
    s = 50

    X = randn(n, p)
    β = randn(s)
    Y = X[:,1:s] * β + randn(n)

    λ1 = 0.3
    λ2 = 0.1
    opt = CDOptions(;maxIter=5000, optTol=1e-8)

    x1 = lasso(X, Y, λ1)
    x2 = lasso(X, Y, λ2)

    λpath = [λ1, λ2]
    path = LassoPath(X, Y, λpath, opt; standardizeX=false)

    @test typeof(path) == LassoPath{Float64}
    @test Vector(path.βpath[1]) ≈ Vector(x1.x) atol=1e-5
    @test Vector(path.βpath[2]) ≈ Vector(x2.x) atol=1e-5

    S1 = findall(!iszero, x1.x)
    S2 = findall(!iszero, x2.x)
    rf = refitLassoPath(path, X, Y)

    @test rf[S1] ≈ X[:,S1] \ Y atol=1e-5
    @test rf[S2] ≈ X[:,S2] \ Y atol=1e-5
  end

  @testset "standardizeX = true" begin
    n = 1000
    p = 500
    s = 50

    X = randn(n, p)
    β = randn(s)
    Y = X[:,1:s] * β + randn(n)

    loadingX = Array{Float64}(undef, p)
    CoordinateDescent._stdX!(loadingX, X)

    λ1 = 0.3
    λ2 = 0.1
    opt = CDOptions(;maxIter=5000, optTol=1e-8)

    x1 = lasso(X, Y, λ1, loadingX, opt)
    x2 = lasso(X, Y, λ2, loadingX, opt)

    λpath = [λ1, λ2]
    path = LassoPath(X, Y, λpath, opt)

    @test typeof(path) == LassoPath{Float64}
    @test Vector(path.βpath[1]) ≈ Vector(x1.x) atol=1e-5
    @test Vector(path.βpath[2]) ≈ Vector(x2.x) atol=1e-5

    S1 = findall(!iszero, x1.x)
    S2 = findall(!iszero, x2.x)
    rf = refitLassoPath(path, X, Y)

    @test rf[S1] ≈ X[:,S1] \ Y atol=1e-5
    @test rf[S2] ≈ X[:,S2] \ Y atol=1e-5
  end

end

end


using TrajectoryOptimization
using StaticArrays
using LinearAlgebra
using TrajOptPlots
using GeometryTypes
using MeshCat
const TO = TrajectoryOptimization
import TrajectoryOptimization.Controllers: RBState

# Model
Rot = UnitQuaternion{Float64,CayleyMap}
model = Dynamics.FreeBody{Rot,Float64}()
n,m = size(model)

if !isdefined(Main,:vis)
    vis = Visualizer()
    open(vis)
    set_mesh!(vis, model)
end

# Discretization
N = 101
tf = 5.0

# Initial and final conditions
X0 = zero(RBState)
Xf = RBState([1,1,1.], expm(deg2rad(70) * @SVector rand(3)), zeros(3), zeros(3))
x0 = Dynamics.build_state(model, X0)
xf = Dynamics.build_state(model, Xf)

# Objective
Qd = Dynamics.fill_state(model, 0e-0, 0e-1, 1e-2, 1e-2)
Rd = @SVector fill(1e-2, 6)
costfun = QuatLQRCost(Diagonal(Qd), Diagonal(Rd), xf, w=0.)
obj = Objective(costfun, N)

# Constraint
struct PointingConstraint <: TO.AbstractConstraint{Equality,State,1}
    p::SVector{3,Float64}
    b::SVector{3,Float64}
    PointingConstraint(p::SVector{3}, b=@SVector [0,0,1.]) = new(p,b)
end
TO.state_dim(::PointingConstraint) = 13

function TO.evaluate(con::PointingConstraint, x::SVector)
    q = UnitQuaternion(x[4],x[5],x[6],x[7])
    x3 = q*con.b
    r = @SVector [x[1],x[2],x[3]]
    dx = con.p - r
    return @SVector [x3'normalize(dx) - 1]
end

point = @SVector [1,-.5,1]
add_point!(vis, point)
con = PointingConstraint(point)
point2 = @SVector [-1,1,0.]
add_point!(vis, point2, name="point2")
con2 = PointingConstraint(point2, @SVector [1,0,0])

conSet = ConstraintSet(n,m,N)
add_constraint!(conSet, con, N:N)
add_constraint!(conSet, con2, N:N)

# Initialization
U0 = [@SVector zeros(6) for k = 1:N-1]

# Problem
prob = Problem(model, obj, xf, tf, x0=x0, constraints=conSet)

# Solve
solver = iLQRSolver(prob)
rollout!(solver)
cost(solver)
solver.opts.verbose = true
solve!(solver)

solver = AugmentedLagrangianSolver(prob)
initial_controls!(solver, U0)
rollout!(solver)
solver.opts.verbose = true
TO.update_constraints!(solver)
TO.constraint_jacobian!(solver)
max_violation(solver)
get_constraints(solver)
cost(solver)
evaluate(con, states(solver)[end])

solve!(solver)
err = Xf ⊖ RBState(model, states(solver)[end])
norm(err)

visualize!(vis, solver)

qf = Xf.q
e3 = @SVector [0,0,1.]
qf*e3

rotmat(qf)


# Draw a series of graphs as svg files

using LightGraphs
using GraphPlot
using PETScBinaryIO
using Compose
using Colors
using HDF5
using MatrixMarket
using ArgParse
using PerceptualColourMaps

s = ArgParseSettings()
@add_arg_table s begin
    "--eig"
        action = :store_true
        help = "color nodes by eigenvalue"
    "INPUT.petsc"
        help = "input graph"
        required = true
    "OUTPUT.pdf"
        help = "output filename"
        required = true
end
parsed_args = parse_args(ARGS, s)

function write_chaco(filename, g)
    open(filename, "w") do f
        write(f, "$(nv(g)) $(ne(g))\n")
        for i = 1:nv(g)
            for n in neighbors(g, i)
                write(f, "$n ")
            end
            write(f, "\n")
        end
    end
end

function map_colormap(vals, clm)
    min = minimum(vals)
    max = maximum(vals)
    map(x -> clm[Int(floor(1+(x - min)/(max-min)*(length(clm)-1)))], vals)
end

function compute_eig(A)
    mmwrite("tmp.mtx", spdiagm(vec(sum(A,1))) - A)
    run(`build/bin/ligmg_eigensolve tmp.mtx tmp.h5 -eps_smallest_magnitude -st_pc_type ligmg -st_ksp_type preonly -eps_monitor -st_type sinvert -eps_target 0 -eps_nev 1`)
    d = h5read("tmp.h5", "/")
    rm("tmp.h5")
    rm("tmp.mtx")
    d["eigen_vector_0"]
end

function layout_graph(A)
    write_chaco("tmp.graph", A)
    x = readdlm(read(`../ogdf_layout/build/ogdf_layout tmp.graph`))
    rm("tmp.graph")
    x[:,1], x[:,2]
end

println("Reading")
infile = parsed_args["INPUT.petsc"]
A = readPETSc(infile)
A = (A + A.')/2
if all(x->x==0, sum(A,1))
    A = spdiagm(diag(A)) - A
end
g = Graph(A)

plot = if parsed_args["eig"]
    println("Computing smallest eigenvalue")
    eig = compute_eig(A)
    nodefill = map_colormap(eig, cmap("D1"))
    gplot(g; layout=layout_graph, nodefillc=nodefill, NODESIZE=1/sqrt(nv(g)), edgestrokec=RGBA(0,0,0,0.01))
else
    gplot(g; layout=layout_graph, nodefillc=colorant"purple", EDGELINEWIDTH=15/sqrt(nv(g)), NODESIZE=0.5/sqrt(nv(g)), edgestrokec=RGBA(0,0,0,0.3))
end

outfile = parsed_args["OUTPUT.pdf"]
println("Saving to $outfile")
draw(PNG(outfile, 2400px, 2400px, ), plot)


using Test
using XPORTA: IEQ

@testset "./src/game_conversions.jl" begin

using BellScenario

@testset "Conversions to BellGames" begin
    @testset  "Vector{Int64} -> BellGame, LocalSignaling" begin
        norm_facet = [2,1,0,0,1,2,-2,-1,-2,2]
        gen_facet = [2,1,0,0,0,1,2,0,0,1,0,2,4]
        scenario = LocalSignaling(3,4,2)

        bell_game = convert(BellGame, norm_facet, scenario, rep="normalized")
        bell_game2 = convert(BellGame, gen_facet, scenario, rep="generalized")

        @test bell_game isa BellGame
        @test bell_game2 isa BellGame
        @test bell_game == [2 0 0;1 1 1;0 2 0;0 0 2]
        @test bell_game.β == 4
        @test bell_game2 == [2 0 0;1 1 1;0 2 0;0 0 2]
        @test bell_game2.β == 4

        @test_throws DomainError convert(BellGame, norm_facet, scenario, rep="non-signaling")
    end

    @testset "IEQ -> BellGames, BlackBox" begin
        gen_ieq = IEQ(inequalities = [1 1 0 0 0 1 2;1 1 0 0 0 0 1] )
        scenario = BlackBox(3,2)

        bell_games = convert(Vector{BellGame}, gen_ieq, scenario, rep="generalized")

        @test bell_games isa Vector{BellGame}
        @test bell_games == BellGame.([[1 0;1 0;0 1],[1 0;1 0;0 0]],[2,1])

        norm_ieq = IEQ(inequalities = [1 1 -1 -1 1;1 1 0 0 1])

        bell_games = convert(Vector{BellGame}, norm_ieq, scenario, rep="normalized")
        bell_games2 = convert(Vector{BellGame}, norm_ieq, scenario)

        @test bell_games isa Vector{BellGame}
        @test bell_games2 isa Vector{BellGame}
        @test bell_games == BellGame.([[1 0;1 0;0 1],[1 0;1 0;0 0]],[2,1])
        @test bell_games2 == BellGame.([[1 0;1 0;0 1],[1 0;1 0;0 0]],[2,1])
    end

    @testset "Vector{Int64} -> BellGame, BipartiteNonSignaling" begin
        scenario = BipartiteNonSignaling(2,2,2,2)

        ns_facet = [1, 0, 1, 0, -1, -1, -1, 1, 1]
        bell_game = convert(BellGame, ns_facet, scenario)

        @test bell_game isa BellGame
        @test bell_game == [1 0 0 1;1 1 1 0;1 1 1 0;0 1 1 0]
        @test bell_game.β == 3

        norm_facet = [0, 0, 1, 0, 1, -1, 0, 0, 0, 1, 0, 1, 1]
        bell_game = convert(BellGame, norm_facet, scenario, rep="normalized")

        @test bell_game isa BellGame
        @test bell_game.β == 2
        @test bell_game == [0 1 0 1;0 2 0 0;1 0 0 1;0 1 0 0]
    end
end

@testset "Conversions from BellGames" begin

    @testset "BellGame -> Vector{Int64}" begin
        BG = BellGame([2 0 0;1 1 1;0 2 0;0 0 2], 4)

        gen_facet = convert(Vector{Int64}, BG, rep="generalized")

        @test gen_facet isa Vector{Int64}
        @test gen_facet == [2,1,0,0,0,1,2,0,0,1,0,2,4]

        norm_facet = convert(Vector{Int64}, BG, rep="normalized")

        @test norm_facet isa Vector{Int64}
        @test norm_facet == [2,1,0,0,1,2,-2,-1,-2,2]

        @test_throws DomainError convert(Vector{Int64}, BG, rep="non-signaling")
    end

    @testset "BellGames -> IEQ" begin
        bell_games = BellGame.([[1 0;0 1],[1 0;0 0]],[2,1])
        gen_ieq = convert(IEQ, bell_games, rep="generalized")

        @test gen_ieq isa IEQ
        @test gen_ieq.inequalities == [1 0 0 1 2;1 0 0 0 1]

        norm_ieq = convert(IEQ, bell_games, rep="normalized")

        @test norm_ieq isa IEQ
        @test norm_ieq.inequalities == [1 -1 1;1 0 1]

        norm_ieq2 = convert(IEQ, bell_games)

        @test norm_ieq2 isa IEQ
        @test norm_ieq2.inequalities == [1 -1 1;1 0 1]
    end

    @testset "BellGame -> Facet BipartiteNonSignaling" begin
        @testset "chsh spot check" begin
            BG = BellGame([1 0 0 1;1 1 1 0;1 1 1 0;0 1 1 0], 3)
            scenario = BipartiteNonSignaling(2,2,2,2)

            facet = convert(Vector{Int64}, BG, scenario)

            @test facet isa Vector{Int64}
            @test length(facet) == LocalPolytope.vertex_dims(scenario, "non-signaling") + 1
            @test facet == [1,0,1,0,-1,-1,-1,1,1]
        end
    end
end

@testset "identity conversions" begin
    @testset "BellGames -> IEQ -> BellGames" begin
        bell_games = BellGame.([[2 0 0 0;0 2 0 0;0 0 2 0;1 1 1 0],[1 0 0 0;1 0 0 0;1 0 0 0;0 0 0 0]],[4,1])
        scenario = LocalSignaling(4,4,2)

        ieq = convert(IEQ, bell_games)
        converted_bell_games = convert(Vector{BellGame}, ieq, scenario)

        @test converted_bell_games == bell_games
    end

    @testset "BipartiteNonSignaling facet -> BellGame -> facet" begin
        @testset "simple scenarios full check" for scenario_params in [
            (2,2,2,2),(3,2,2,2),(2,3,2,2),(2,2,3,2),(2,2,2,3),
            (3,3,2,2),(3,2,3,2),(3,2,2,3),(2,3,3,2),(2,3,2,3),
            (2,2,3,3)
        ]
            scenario = BipartiteNonSignaling(scenario_params...)

            vertices = LocalPolytope.vertices(scenario)
            facet_dict = LocalPolytope.facets(vertices)

            for facet in facet_dict["facets"]
                BG = convert(BellGame, facet, scenario)

                @test BG isa BellGame

                facet2 = convert(Vector{Int64}, BG, scenario)

                @test facet2 isa Vector{Int64}
                @test facet2 == facet
            end
        end
    end
end

end


@symbol_func function phi_CS(cur_reactor::AbstractReactor)
  cur_phi = cur_reactor.k

  cur_phi *= cur_reactor.B_0

  cur_phi *= sqrt( h_CS(cur_reactor) )

  cur_phi *= sqrt( G_LO(cur_reactor) )

  cur_phi *= sqrt( K_LP(cur_reactor) )

  cur_phi *= sqrt( cur_reactor.R_0 )

  cur_phi /= sqrt( 1e-7 )

  cur_phi
end


module MathTools

using StaticArrays
using Base.Iterators
using Base: product
using Interpolations

struct Point2D{T <: Real} <: FieldVector{2, T}
    x :: T
    y :: T
end

x(p :: Point2D) = p.x
y(p :: Point2D) = p.y

struct Point3D{T <: Real} <: FieldVector{3, T}
    x :: T
    y :: T
    z :: T
end

x(p :: Point3D) = p.x
y(p :: Point3D) = p.y
z(p :: Point3D) = p.z

# Copying NumPy's linspace function
linspace(min, max, step) = min:(max - min)/step:max
columns(M) = tuple((view(M, :, i) for i in 1:size(M, 2))...)

## Haskell Master Race
#===========================================================================#

span(pred, iter) = (takewhile(pred, iter), dropwhile(pred, iter))
splitat(n, xs) = @views (xs[1:n,:], xs[n+1:end,:])

lisa(pred, iter) = span(!pred, iter)

Base.Iterators.partition(pred, xs, f, g) = f.(filter(pred, xs)), g.(filter(!pred, xs))
Base.Iterators.partition(pred, xs, f = identity) = partition(pred, xs, f, f)

## Renaming math operations
#===========================================================================#

"""
"Lenses" to access subfields on lists of objects.
"""
|>(obj, fields :: Array{Symbol}) = foldl(getproperty, fields, init = obj)
|>(list_objs :: Array{T}, fields :: Array{Symbol}) where T <: Any = list_objs .|> [fields]

field << obj = getfield(obj, field)

# Convert homogeneous struct entries to lists
structtolist(x) = [ name << x for name ∈ (fieldnames ∘ typeof)(x) ]

## Renaming math operations
#===========================================================================#

⊗(A, B)    = kron(A, B)

×(xs, ys)   = product(xs, ys)
dot(V₁, V₂) = sum(V₁ .* V₂)
# ×(xs, ys) = (collect ∘ zip)(xs' ⊗ (ones ∘ length)(ys), (ones ∘ length)(xs)' ⊗ ys)

## Basic transformations/geometry
#===========================================================================#

# Transforms (x, y) to the coordinate system with (x_s, y_s) as origin oriented at α_s.
affine_2D(x, y, x_s, y_s, α_s) = rotation(x - x_s, y - y_s, α_s)
inverse_rotation(x, y, angle) = SVector(x * cos(angle) - y * sin(angle), x * sin(angle) + y * cos(angle))
rotation(x, y, angle) = SVector(x * cos(angle) + y * sin(angle), -x * sin(angle) + y * cos(angle))

# Matrix versions
rotation(θ)         = [ cos(θ) sin(θ) ;
                       -sin(θ) cos(θ) ]
inverse_rotation(θ) = rotation(-θ)

# Cartesian-polar coordinates
magnitude(U, W) = sqrt(U^2 + W^2)
angle(U, W) = atan(W, U)

slope(x1, y1, x2, y2) = (y2 - y1) / (x2 - x1)

## Array conversions
#===========================================================================#

tuparray(xs)  = tuple.(eachcol(xs)...)
vectarray(xs) = SVector.(eachcol(xs)...)

extend_yz(coords) = @views [ coords[:,1] zero(coords[:,1]) coords[:,2] ]

reflect_mapper(f, xs) = @views [ f(xs[:,end:-1:1]) xs ]

## Difference operations
#===========================================================================#

forward_difference_matrix(n) = [ I zeros(n) ] - [ zeros(n) I ]

forward_sum(xs) = @views @. xs[2:end] + xs[1:end-1]
forward_difference(xs) = @views @. xs[2:end] - xs[1:end-1]
forward_division(xs) = @views @. xs[2:end] / xs[1:end-1]
ord2diff(xs) = @views @. xs[3:end] - 2 * xs[2:end-1] + xs[1:end-2]

adj3(xs) = @views zip(xs[1:end-2,:], xs[2:end-1,:], xs[3:end,:])

# Central differencing schema for pairs except at endpoints
midpair_map(f :: H, xs) where {H} =
    @views  [        f.(xs[1,:], xs[2,:])'        ;
               f.(xs[1:end-2,:], xs[3:end,:]) / 2 ;
                 f.(xs[end-1,:], xs[end,:])'      ]

# stencil(xs, n) = [ xs[n+1:end] xs[1:length(xs) - n] ]
# parts(xs) = let adj = stencil(xs, 1); adj[1,:], adj[end,:] end

# Lazy? versions
# stencil(xs, n) = zip(xs[n+1:end], xs[1:length(xs) - n])
# parts(xs) = (first ∘ stencil)(xs, 1), (last ∘ stencil)(xs, 1)

# function midgrad(xs)
#     first_two_pairs, last_two_pairs = permutedims.(parts(xs))
#     central_diff_pairs = stencil(xs, 2)

#     [first_two_pairs; central_diff_pairs; last_two_pairs]
# end

## Spacing formulas
#===========================================================================#

uniform_spacing(x1, x2, n) = range(x1, x2, length = n)
linear_spacing(x_center, len, n :: Integer) = @. x_center + len * 0:1/(n-1):1
cosine_spacing(x_center, diameter, n :: Integer = 40) = @. x_center + (diameter / 2) * cos(-π:π/(n-1):0)

function sine_spacing(x1, x2, n :: Integer = 40)
    d = x2 - x1
    if n < 0
       @. x2 - d * sin(π/2 * (1. - (1:1/(n+1):0)))[end:-1:1]
    else
       @. x1 + d * sin(π/2 * (0:1/(n-1):1))
    end
end

function cosine_interp(xs, n :: Integer = 40)
    d = maximum(xs) - minimum(xs)
    x_center = (maximum(xs) + minimum(xs)) / 2
    x_circ = cosine_spacing(x_center, d, n)
end

## Iterator methods
#===========================================================================#

# Need to improve this via recursion; the for loop seems really unnecessary
function accumap(f, n, xs)
    data = [ xs ]
    for i = 1:n
        ys = map(f, xs)
        data = [ reduce(vcat, data); ys ]
        xs = ys
    end
    return reduce(hcat, data)
end

## Helper functions
#===========================================================================#

"""
    weighted(x1, x2, μ)

Compute the weighted value between two values ``x_1`` and ``x_2`` with weight ``\\mu \\in [0,1]``.
"""
weighted(x1, x2, μ) = (1 - μ) * x1 + μ * x2

"""
Compute the weighted average (μ) of two vectors.
"""
weighted_vector(x1, x2, μ) = weighted.(x1, x2, μ)

## Macros
#===========================================================================#

# macro getter(obj)
#     return :((function name(x)
#                 getfield(name, x)
#             end))
# end

reshape_array(arr, inds, sizes) = @views [ reshape(arr[i1+1:i2], size) for (i1, i2, size) in zip(inds[1:end-1], inds[2:end], sizes) ]

end


using EffectiveWaves
using LinearAlgebra
using Plots
pyplot()

medium = Medium(ρ=1.0, c=1.0)
species = [Specie(ρ=0.1,r=1.0, c=0.01, volfrac=0.3)]

ωs = [0.5]
ω = ωs[1]
k = ω/medium.c
tol = 1e-8

kp = 1.0+0.4im
kps = collect(0.2:0.2:1.0) .+ collect(0.3:0.2:1.1) .* im
hankel_order = ho = 2

wavesystemPlane = wavematrix3DPlane(ωs[1], medium, species; tol = 1e-5, hankel_order=hankel_order)
detP(kp) = det(wavesystemPlane(kp))

# dispersion = dispersion_function(ω, medium, species; tol = low_tol, dim=dim, symmetry = :plane, hankel_order = ho)

# dim = 3
# k_effs = wavenumbers_path(ω, medium, species;
#     num_wavenumbers=7, dim=dim, symmetry = :plane,
#     hankel_order = ho, tol=tol)

# detsplanes = det.(wavesystemPlane.(kps))
# MPs = wavesystemPlane.(kps)

wavesystem3D = wavematrix3D(ωs[1], medium, species; tol = 1e-5, hankel_order=hankel_order);
detR(kp) = sqrt(det(wavesystem3D(kp)))

wavesystem3D = wavematrix3D(ωs[1], medium, species; tol = 1e-5, hankel_order=1);
detR(kp) = sqrt(det(wavesystem3D(kp)))

# wavesystem3D_2 = wavematrix3D_allocate(ωs[1], medium, species; tol = 1e-4, hankel_order=hankel_order);
# detR2(kp) = det(wavesystem3D_2(kp))

# @time detR.(kps);
#   1.996288 seconds (469.46 k allocations: 52.257 MiB, 1.06% gc time)
#
# @time detR.(kps);
#   2.103966 seconds (327.72 k allocations: 45.012 MiB, 0.39% gc time)
#
# @time detR.(kps);
# 0.084732 seconds (32.86 k allocations: 3.510 MiB)

# Effective medium
    dim = 3; T= Float64
    β = medium.ρ*medium.c^2 # medium bulk modulus
    φ = sum(volume_fraction.(species; dim = dim))
    β_eff = one(T)/((1-φ)/β + sum(volume_fraction(s; dim = dim)/(s.ρ*s.c^2) for s in species))
    ρ_frac = sum(volume_fraction(s; dim = dim)*(medium.ρ - s.ρ)/(medium.ρ + 2*s.ρ) for s in species)

    ρ_eff = medium.ρ*(1 - ρ_frac)/(1 + T(dim - 1) * ρ_frac)
    eff_medium = Medium(ρ=ρ_eff, c=sqrt(β_eff/ρ_eff))
    k_low1 = ω/eff_medium.c

    ρ_eff = medium.ρ*(1 - ρ_frac)/(1 + (ho/(2ho+3))*ρ_frac)
    eff_medium = Medium(ρ=ρ_eff, c=sqrt(β/ρ_eff))
    k_low2 = ω/eff_medium.c

    ρ_eff = medium.ρ*(1 - ρ_frac)/(1 + ((ho-1)/(2ho+1))*ρ_frac)
    eff_medium = Medium(ρ=ρ_eff, c=sqrt(β_eff/ρ_eff))
    k_low3 = ω/eff_medium.c

kp0s = wavenumbers(ω, medium, species; dim = dim, symmetry = :plane, hankel_order = ho,
    num_wavenumbers=20, tol = tol, mesh_points = 14, mesh_size = 3.0)

kp_vecs = reduce_kvecs([[real(kp),imag(kp)] for kp in kp0s], 1e-5)
kps = [kp[1]+im*kp[2] for kp in kp_vecs]
kps = kps[1:30]

# inds = findall(abs.(detP.(kps)) .< 1e-6)
# kps = kp0s[inds]

# For low frequency
    # k_lows = [k_low1,k_low2,k_low3]
    # k1 = minimum(real.(k_lows[[1,3]]))
    # k2 = maximum(real.(k_lows))
    # x = LinRange(0.9*k1,1.1*k2,110)
    # y = k2 .* LinRange(-0.01,0.2,90)

# Given some kps
    # dk1 = real(kps[1]) / 4
    # dk2 = imag(kps[1]) / 4
    k1 = maximum(abs.(real.(kps)))
    k2 = imag(kps[end])
    x = LinRange(-k1,k1,6600)
    y = LinRange(-0.01,k2,160)

X = repeat(x',length(y),1)
Y = repeat(y,1,length(x))

ZPs = map( (x,y) -> abs(detP(x+y*im)),X,Y)
# ZRs = map( (x,y) -> abs(detR(x+y*im)),X,Y)

minR = 1e-18
minP = 0.8

# ZRs2 = [ (abs(ZRs[ind]) < minR ? ZRs[ind] : NaN ) for ind in LinearIndices(ZRs)];
ZPs2 = [ (abs(ZPs[ind]) < minP ? ZPs[ind] : NaN ) for ind in LinearIndices(ZPs)];

# h1 = heatmap(x,y,ZPs, xlab = "Re k*", ylab = "Im k*", title="Planewave det(I + G), ka~1", clims = (0.,1.0))

# h2 = heatmap(x,y,ZRs2, xlab = "Re k*", ylab = "Im k*", title="General det(I + G), ka~$(ω*species[1].r)", clims = (0.,minR))
h1 = heatmap(x,y,ZPs2, xlab = "Re k*", ylab = "Im k*", title="Planewave det(I + G), ka~$(ω*species[1].r)", clims = (0.,minP))
scatter!(kps)

# Z = map( (x,y) -> abs(((x^2-1)^2*(y*im - 1)*(y*im - 2))),X,Y)
# Z = [ (abs(Z[ind]) < 1.0 ? Z[ind] : NaN ) for ind in LinearIndices(Z)]
# heatmap(x,y,Z, xlab = "Re k*", ylab = "Im k*", clims = (0.,1.0))

abs.(detsplanes)
abs.(dets3D)

h1 = heatmap(x,y,ZPs, xlab = "Re k*", ylab = "Im k*", title="Planewave det(I + G), ka~$(ω*species[1].r)", clims = (0.,1.0))
scatter!(kps)


# -----------------------------------------------------------
#
# Implementing a nodal basis for DG sparse grids
#
# -----------------------------------------------------------

#=

We begin by defining analogues of our DG functions
to satisfy the following conditions:

    1) There are a set of collocation points respecting the hierarchical structure
    2) The nodal basis functions at a given level are those that vanish
        on all nodal points at the above levels and vanish on all but one
        of the cells on the level below
    3) Third constraint?

To do this, we _restrict_ to the continuous part of the DG basis, and focus
on piecewise polynomials that vanish on their boundaries. We then pick the
collocation points to be exactly at these boundaries.

For this, it is useful to define lagrange polynomials, and moreover
use k=2, 3, 5.
The k=2 case has our basis is exactly the hat function basis constructed
previously.

For k=2, the coarsest level has nodal points at {0, 1},
    the next level (the first hat) has nodal point at 0.5
    the next level has two cells with points at 0.25, 0.75 etc.

For k=3 the coarsest level has nodal points at {0, 0.5, 1},
    the next level has two nodal functions with points at 0.25, 0.75,
    the next level has two cells and two nodes per cell at all odd multiples of 0.125 etc.

For k=5 the coarsest level has nodal points at {0, 0.25, 0.5, 0.75, 1}
    the next level has four nodal functions with points at all multiples of 0.125
    the next level has two cells and four nodes per cell at all odd multiples of 0.0625 etc.

Notice that k=2,3,5 work because at each cell of the multiresolution basis the space of
    picewise continuous polynomials is k-1 = 2^i for i=0,1,2 in this case.
    A construction for other k would be harder, if not impossible.

=#

# Lagrange polynomials interpolating at the coarsest level
function lag_nodal(k::Int, mode::Int, x::Real)
    (k < 2 || k > 5) && throw(ArgumentError("k is not between 2 and 5"))
    (mode > k || mode < 1) && throw(ArgumentError("mode is not in [1, k]"))

    (x < 0 || x > 1) && return zero(x)
    if k == 2
        mode == 1 && return x
        mode == 2 && return 1 - x
    elseif k == 3
        mode == 1 && return 2 * (x - 1//2) * (x - 1)
        mode == 2 && return -4 * x * (x - 1)
        mode == 3 && return 2 * x * (x - 1//2)
    elseif k == 4
        throw(ArgumentError("k = 4 not implemented"))
    elseif k == 5
        mode == 1 && return 32//3 * (x - 1//4) * (x - 1//2) * (x - 3//4) * (x - 1)
        mode == 2 && return -128//3 * x * (x - 1//2) * (x - 3//4) * (x - 1)
        mode == 3 && return 64 * x * (x - 1//4) * (x - 3//4) * (x - 1)
        mode == 4 && return -128//3 * x * (x - 1//4) * (x - 1//2) * (x - 1)
        mode == 5 && return 32//3 * x * (x - 1//4) * (x - 1//2) * (x - 3//4)
    end
end

# Multiresolution functions for the levels below
function h_nodal(k::Int, mode::Int, x::Real)
    (k < 2 || k > 5) && throw(ArgumentError("k is not between 2 and 5"))
    (mode > k || mode < 1) && throw(ArgumentError("mode is not in [1, k]"))

    if k == 2
        mode == 1 && return max(zero(x), 1-abs(x))
        if mode == 2
            (abs(x) > 1 || abs(x) == 0) && return zero(x)
            x < 0 && return (x + 1)
            x > 0 && return zero(x)# (x - 1)
        end
    elseif k == 3
        mode == 1 && return max(zero(x), -4 * x * (x + 1))
        mode == 2 && return max(zero(x), -4 * x * (x - 1))
        if mode == 3
            (abs(x) > 1 || abs(x) == 0) && return zero(x)
            x < 0 && return 2 * (x + 1) * (x + 1//2)
            x > 0 && return zero(x) # -2 * (x - 1) * (x - 1//2)
        end
    elseif k == 4
        throw(MethodError("k = 4 not implemented"))
    elseif k == 5
        mode == 1 && return (x < -1 || x > 0) ? zero(x) : -128//3 * x * (x + 1) * (x + 1//2) * (x + 1//4)
        mode == 2 && return (x < -1 || x > 0) ? zero(x) : -128//3 * x * (x + 1) * (x + 1//2) * (x + 3//4)
        mode == 3 && return (x > 1  || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 3//4)
        mode == 4 && return (x > 1  || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 1//4)
        if mode == 5
            (abs(x) > 1 || abs(x) == 0) && return zero(x)
            x < 0 && return 32//3 * (x + 1) * (x + 3//4) * (x + 1//2) * (x + 1//4)
            x > 0 && return zero(x) #32//3 * (x - 1) * (x - 3//4) * (x - 1//2) * (x - 1//4)
        end
    end
end

# This yields our nodal basis:
function v_nodal(k::Int, level::Int, cell::Int, mode::Int, x::Real)
    if level==0
        return lag_nodal(k, mode, x)
    else
        return h_nodal(k, mode, (1<<level)*x - (2*cell-1))
    end
end

function v_nodal(k::Int, level::Int, cell::Int, mode::Int)
    return x->v_nodal(k, level, cell, mode, x)
end

# -----------------------------------------------------------
# Made for evaluating the dicontinuous elements in the nodal
# basis at their points of discontinuity
# -----------------------------------------------------------

# Evaluates the modal basis at a given point, averaging over jumps
function eval_v_periodic(k::Int, level::Int, cell::Int, mode::Int, x::T) where {T<:Real}
    xl = max(x - 1e-16, 0)
    xr = min(x + 1e-16, 1)
    (xl == 0) && (xl = 1 - 1e-16)
    (xr == 1) && (xr = 1e-16)
    return 1//2 * (v(k, level, cell, mode, xl) + v(k, level, cell, mode, xr))
end

# Find the size of the discontinuity of the modal basis at a given point
function gap_v(k::Int, level::Int, cell::Int, mode::Int, x::T) where {T<:Real}
    xl = max(x - 1e-16, 0)
    xr = min(x + 1e-16, 1)
    return 1//2 * (v(k, level, cell, mode, xr) - v(k, level, cell, mode, xl))
end

# Evaluates the modal basis at a given point, averaging over jumps
function eval_v_nodal_periodic(k::Int, level::Int, cell::Int, mode::Int, x::T) where {T<:Real}
    xl = max(x - 1e-30, 0)
    xr = min(x + 1e-30, 1)
    (xl == 0) && (xl = 1 - 1e-30)
    (xr == 1) && (xr = 1e-30)
    return 1//2 * (v_nodal(k, level, cell, mode, xl) + v_nodal(k, level, cell, mode, xr))
end

# Find the size of the discontinuity of the modal basis at a given point
function avg_v_nodal(k::Int, level::Int, cell::Int, mode::Int, x::T) where {T<:Real}
    xl = x - 1e-16
    xr = x + 1e-16
    (xl < 0) && (xl = xr)
    (xr > 1) && (xr = xl)
    return 1//2 * (v_nodal(k, level, cell, mode, xr) + v_nodal(k, level, cell, mode, xl))
end

# Find the size of the discontinuity of the modal basis at a given point
function gap_v_nodal(k::Int, level::Int, cell::Int, mode::Int, x::T) where {T<:Real}
    xl = x - 1e-16
    xr = x + 1e-16
    (xl < 0) && (xl = xr)
    (xr > 1) && (xr = xl)
    return 1//2 * (v_nodal(k, level, cell, mode, xr) - v_nodal(k, level, cell, mode, xl))
end

function eval_v_nodal_right(k::Int, level::Int, cell::Int, mode::Int, x::T) where {T<:Real}
    xr = x + 1e-16
    return v_nodal(k, level, cell, mode, xr)
end

function eval_v_nodal_left(k::Int, level::Int, cell::Int, mode::Int, x::T) where {T<:Real}
    xl = x - 1e-16
    return v_nodal(k, level, cell, mode, xl)
end

# -----------------------------------------------------------
# Next we define the sparse matrices transforming between
# the heirarchical, nodal, and point bases
# -----------------------------------------------------------

# Huge TODO:
#     Replace all of these for loops and multiple function calls
#     by cleverly using the iterable interface in julia

# Evaluation of a given nodal function along all relevant gridpoints
function eval_points_1D(k::Int, max_level::Int, level::Int, cell::Int, mode::Int)
    i = 1
    I = Int[]
    V = Float64[]
    for l in 0:max_level
        l == 0 ? (node_range = 0:(1//(k-1)):(1 - 1//(k-1))) :
                 (node_range = 1//(2*(k-1)):1//(k-1):(1-1//(2*(k-1))))

        for c in 1:1<<max(0, l-1)
            for node in node_range
                x = (c - 1 + node)//(1<<max(0, l-1))
                val = eval_v_nodal_right(k, level, cell, mode, x)
                (val != zero(x)) && (push!(I, i); push!(V, val))
                i += 1
            end

            if l == 0
                x = 1
                val = eval_v_nodal_left(k, level, cell, mode, x)
            else
                x = (c - 1//2)/(1<<max(0, l-1))
                val = eval_v_nodal_left(k, level, cell, mode, x)
            end
            (val != zero(x)) && (push!(I, i); push!(V, val))
            i += 1
        end
    end
    return sparsevec(I, V)
end

# Sparse matrix converting nodal basis element to its value
# on the collocation points
function nodal2points_1D(k::Int, max_level::Int; atol=1e-15)
    I = Int[]
    J = Int[]
    V = Float64[]
    j = 1

    for level in 0:max_level
        mode_range = 1:k

        for cell in 1:1<<max(0, level-1)
            for mode in mode_range
                coeffs = eval_points_1D(k, max_level, level, cell, mode)
                for (i, val) in zip(findnz(coeffs)...)
                    push!(I, i)
                    push!(J, j)
                    push!(V, val)
                end
                j += 1
            end
        end
    end
    return threshold(sparse(I, J, V), atol)
end

# Sparse matrix inverting the above construction
# Converting a set of values on the collocation points to the DG space
# in terms of the nodal basis
function points2nodal_1D(k::Int, max_level::Int; atol=1e-15)
    n2p = nodal2points_1D(k, max_level)
    return threshold(inv(Matrix(n2p)), atol)
end

# Converts a given nodal basis element to a sum in the standard
# hieararchical DG scheme
# TODO: Remove hquadrature and do this directly
function nodal2modal_1D(k::Int, level::Int, cell::Int, mode::Int; rtol=1e-10, atol=1e-15, max_evals=1500)
    I = Int[]
    V = Float64[]
    i = 1

    for hier_level in 0:level
        hier_cells = 1<<max(0, hier_level-1)
        for hier_cell in 1:hier_cells
            for hier_mode in 1:k
                x_min = (hier_cell-1)//hier_cells
                x_max = (hier_cell)//hier_cells
                val = hquadrature(x->v_nodal(k, level, cell, mode, x)*
                                    v(k, hier_level, hier_cell, hier_mode, x),
                                    x_min, x_max,
                                    reltol=rtol, abstol=atol, maxevals=max_evals)[1]
                (abs(val) > eps(atol)) && (push!(I, i); push!(V, val))
                i += 1
            end
        end
    end
    return dropzeros!(sparsevec(I, V))
end

function nodal2modal_1D(k::Int, max_level::Int; rtol=1e-10, atol=1e-15, max_evals=1500)
    I = Int[]
    J = Int[]
    V = Float64[]
    j = 1

    for nodal_level in 0:max_level
        num_nodes = k
        nodal_cells = 1<<max(0, nodal_level-1)
        for nodal_cell in 1:nodal_cells
            for nodal_mode in 1:num_nodes
                coeffs = nodal2modal_1D(k, nodal_level, nodal_cell, nodal_mode;
                                        rtol=rtol,
                                        atol=atol,
                                        max_evals=max_evals)
                for (i, val) in zip(findnz(coeffs)...)
                    push!(I, i)
                    push!(J, j)
                    push!(V, val)
                end
                j += 1
            end
        end
    end
    return threshold(sparse(I, J, V), atol)
end

# Accounting for ambiguity in the definition of the hierarchical
# basis across regions
function eval_v(k, level, cell, mode, x)
    xl = max(x - 1e-16, 0)
    xr = min(x + 1e-16, 1)
    (xl == 0) && (xl = 1 - 1e-16)
    (xr == 1) && (xr = 1e-16)
    return 0.5 * (v(k, level, cell, mode, xl) + v(k, level, cell, mode, xr))
end

function nodal2pos_1D(k, max_level, nodal_level, nodal_cell, nodal_mode;
                   rtol=1e-10, atol=1e-15, max_evals=1500)
    I = Int[]
    V = Float64[]
    i = 1

    nodal_min = (nodal_cell-1)/(1<<max(0, nodal_level-1))
    nodal_med = (nodal_cell-1//2)/(1<<max(0, nodal_level-1))
    nodal_max = (nodal_cell)/(1<<max(0, nodal_level-1))
    for c in 1:1<<max_level
        pos_min = (c-1)/(1<<max_level)
        pos_med = (c-1//2)/(1<<max_level)
        pos_max = c/(1<<max_level)

        sm = 1e-20

        if (pos_min > nodal_max) || (pos_max < nodal_min)
            i += k
            continue
        end
        for m in 1:k
            val1  = hquadrature(x->v_nodal(k, nodal_level, nodal_cell, nodal_mode, x)*basis(max_level, c, m, x),
                              pos_min + sm, pos_med - sm,
                              rtol=rtol, atol=atol, maxevals=max_evals)[1]

            val2 = hquadrature(x->v_nodal(k, nodal_level, nodal_cell, nodal_mode, x)*basis(max_level, c, m, x),
                              pos_med + sm, pos_max - sm,
                              rtol=rtol, atol=atol, maxevals=max_evals)[1]

            val = val1 + val2
            (abs(val) > eps(atol)) && (push!(I, i); push!(V, val))
            i += 1
        end
    end
    return sparsevec(I, V)
end

function nodal2pos_1D(k, max_level; rtol=1e-10, atol=1e-15, max_evals=1500)
    I = Int[]
    J = Int[]
    V = Float64[]
    j = 1

    for nodal_level in 0:max_level
        num_nodes = k
        nodal_cells = 1<<max(0, nodal_level-1)
        for nodal_cell in 1:nodal_cells
            for nodal_mode in 1:num_nodes
                coeffs = nodal2pos_1D(k, max_level, nodal_level, nodal_cell, nodal_mode;
                                        rtol=rtol,
                                        atol=atol,
                                        max_evals=max_evals)
                for (i, val) in zip(findnz(coeffs)...)
                    push!(I, i)
                    push!(J, j)
                    push!(V, val)
                end
                j += 1
            end
        end
    end
    return threshold(sparse(I, J, V), atol)
end

function pos2nodal_1D(k, max_level; rtol=1e-10, atol=1e-12, max_evals=1500)
    n2pos = nodal2pos_1D(k, max_level; rtol=rtol, atol=atol, max_evals=max_evals)
    return threshold(inv(Matrix(n2pos)), atol)
end

# Final function, combining all of the above:
function transform_1D(k::Int, n::Int, from::String, to::String; atol=1e-12)
    if from == "nodal" && to == "pos"
        return nodal2pos_1D(k, n; atol=atol)
    elseif from == "pos" && to == "nodal"
        return pos2nodal_1D(k, n; atol=atol)
    elseif from == "nodal" && to == "points"
        return nodal2points_1D(k, n; atol=atol)
    elseif from == "points" && to == "nodal"
        return points2nodal_1D(k, n; atol=atol)
    elseif from == "nodal" && to == "modal"
        return threshold(pos2hier(k, n) * nodal2pos_1D(k, n), atol)
    elseif from == "modal" && to == "nodal"
        return threshold(pos2nodal_1D(k, n) * hier2pos(k, n), atol)
    elseif from == "modal" && to == "pos"
        return hier2pos(k, n; atol=atol)
    elseif from == "pos" && to == "modal"
        return pos2hier(k, n; atol=atol)
    elseif from == "points" && to == "pos"
        return threshold(nodal2pos_1D(k, n) * points2nodal_1D(k, n), atol)
    elseif from == "pos" && to == "points"
        return threshold(nodal2points_1D(k, n) * pos2nodal_1D(k, n), atol)
    elseif from == "modal" && to == "points"
        m2n = transform_1D(k, n, "modal", "nodal"; atol=atol)
        n2p = transform_1D(k, n, "nodal", "points"; atol=atol)
        return threshold(n2p * m2n, atol)
    elseif from == "points" && to == "modal"
        p2n = transform_1D(k, n, "points", "nodal"; atol=atol)
        n2m = transform_1D(k, n, "nodal", "modal"; atol=atol)
        return threshold(n2m * p2n, atol)
    else throw(ArgumentError("A basis was undefined"))
    end
end


"""
Return the minimum distance between
two array of arrays.

Arguments
-----------
:Type{MinDistance}
    must be data type MinDistance

array_1:Array
    input 1

array_2:Array
    input 2
"""
function gage(::Type{MinDistance},array_1::Array, array_2::Array)
    if length(array_1) == 0
        @debug print_with_color(:red, "WARNING! input array 1 is empty for gage(MinDistance,...)\n\n")
        return 0.
    end

    if length(array_2) == 0
        @debug print_with_color(:red, "WARNING! input array 2 is empty for gage(MinDistance,...)\n\n")
        return 0.
    end

    if issubtype(typeof(array_1[1]),Number)
        array_1 = Array[array_1]
    end

    if issubtype(typeof(array_2[1]),Number)
        array_2 = Array[array_2]
    end

    return minimum(Distances.pairwise(Distances.Euclidean(), hcat(array_1...), hcat(array_2...)))
end


using DelayEmbeddings, DynamicalSystemsBase, PyPlot, Test
pygui(true)

function ar_process(u0::T, α::T, p::T, N::Int) where {T<:Real}
    x = zeros(T, N+10)
    x[1] = u0
    for i = 2:N+10
        x[i] = α*x[i-1] + p*randn()
    end
    return x[11:end]
end

trajectories = Dict(
    "White noise" => Dataset(randn(5000)),
    "AR(1)" => Dataset(ar_process(.2,.9,.1,5000)),
    "Hénon (chaotic)" => trajectory(Systems.henon(a=1.4, b=0.3), 199, Ttr=1000),
    "Hénon (periodic)" => trajectory(Systems.henon(a=1.054, b=0.3), 199, Ttr=1000),
    "Lorenz" => trajectory(Systems.lorenz(), 1000.0; Ttr=1000, Δt = 0.05),
    "Lorenz96" => trajectory(Systems.lorenz96(5; F = 8.0), 1000.0; Ttr=1000, Δt = 0.05),
    "Roessler" => trajectory(Systems.roessler(), 1000.0; Ttr=1000, Δt = 0.05),
    "Sine wave" => Dataset(map(x->[sin.(x) cos.(x)], StepRangeLen(0.0,0.2,200)))
)
dict_keys = ["White noise","AR(1)","Hénon (chaotic)","Hénon (periodic)","Lorenz","Lorenz96","Roessler","Sine wave"]

# set noise level
σ = 0.01
# set maximum encountered dimension
dmax = 15
γs = 0:dmax-1
dims = γs .+ 1

slope_thres = .1

for names in dict_keys
    tr = trajectories[names]
    tr = standardize(tr)
    s = tr[:, 1].+ σ.*randn(length(tr))
    τ = estimate_delay(s, "mi_min")

    # do a plot of the differently found methods

    fs = (afnn, fnn, f1nn, ifnn)
    methods = ["afnn", "fnn", "f1nn", "ifnn"]
    fig, axs  = subplots(2, 1; sharex = true)
    for (i, f) in enumerate(fs)
        x = f(s, τ, γs)

        if i == 1
            x .= 1 .- x
            E2 = stochastic_indicator(s, τ, γs)
        end

        axs[1].plot(dims, x, label = string(f); color = "C$(i)")
        axs[2].plot(dims[1:end-1], diff(x); color = "C$(i)")
        # lrs, tans = linear_regions(dims, x, tol = 0.5)
        # axs[1].scatter(dims[lrs], x[lrs], color = "C$(i)")
        z = slope_thres # threshold of slopes
        axs[2].axhspan(-z, z; color = "0.5", alpha = 0.2, label = "±$(z)")

        # Estimations of condition for convergence:
        # earliest point where x is both less than a thresohold, 0.05, and local slope is
        # between ± 0.05
        r = 0.05 # threshold of convergence near 0
        d_max = dmax
        Y, _, rat =  optimal_traditional_de(s, "mi_min", methods[i]; dmax = d_max, slope_thres = slope_thres)
        println("$(methods[i]) embedding dimension for $names: $(size(Y,2)) ")
        axs[1].scatter(dims[size(Y,2)], x[size(Y,2)]; color = "C$(i)", edgecolors  = "k")
    end
    axs[1].grid(); axs[2].grid()
    axs[1].legend(loc = "upper right", ncol = 2)
    axs[1].set_ylabel("estimator")
    axs[2].set_xlabel("dimension")
    axs[2].set_ylabel("local slope")
    axs[1].set_title(names)
end


@testset "sm" begin

    D_in = 5
    v1 = rand(D_in)
    v2 = rand(D_in)

    αs₁ = rand(3)
    αs₂ = rand(D_in, 3)
    γs = rand(D_in, 3)
    ωs = rand(D_in, 3)

    k1 = spectral_mixture_kernel(αs₁, γs, ωs)
    k2 = spectral_mixture_product_kernel(αs₂, γs, ωs)

    t = v1 - v2

    @test k1(v1, v2) ≈ sum(αs₁ .* exp.(-(t' * γs)'.^2 ./ 2) .* cospi.((t' * ωs)')) atol=1e-5

    @test isapprox(
        k2(v1, v2),
        prod(
            [sum(αs₂[i,:]' .* exp.(-(γs[i,:]' * t[i]).^2 ./ 2) .*
            cospi.(ωs[i,:]' * t[i])) for i in 1:length(t)],
        );
        atol=1e-5,
    )

    @test_throws DimensionMismatch spectral_mixture_kernel(rand(5) ,rand(4,3), rand(4,3))
    @test_throws DimensionMismatch spectral_mixture_kernel(rand(3) ,rand(4,3), rand(5,3))
    @test_throws DimensionMismatch spectral_mixture_product_kernel(rand(5,3) ,rand(4,3), rand(5,3))

    # Standardised tests. Choose input dims carefully.
    @testset "ColVecs" begin
        x0 = ColVecs(randn(D_in, 3))
        x1 = ColVecs(randn(D_in, 3))
        x2 = ColVecs(randn(D_in, 2))
        TestUtils.test_interface(k1, x0, x1, x2)
        TestUtils.test_interface(k2, x0, x1, x2)
    end
    @testset "RowVecs" begin
        x0 = RowVecs(randn(3, D_in))
        x1 = RowVecs(randn(3, D_in))
        x2 = RowVecs(randn(2, D_in))
        TestUtils.test_interface(k1, x0, x1, x2)
        TestUtils.test_interface(k2, x0, x1, x2)
    end
    # test_ADs(x->spectral_mixture_kernel(exp.(x[1:3]), reshape(x[4:18], 5, 3), reshape(x[19:end], 5, 3)), vcat(log.(αs₁), γs[:], ωs[:]), dims = [5,5])
    @test_broken "No tests passing (BaseKernel)"
end


export getDiamondModel,getBallModel ;

function getDiamondModel(n::Array{Int64})
n_tup = tuple(n...);
n = collect(n_tup);
u = zeros(Float32,n_tup);
X,Y,Z = meshgrid(-div(n[1],2):div(n[1],2),-div(n[2],2):div(n[2],2),-div(n[3],2):div(n[3],2))
I = (abs.(X)+abs.(Y)+abs.(Z)) .<= div(n[1],4);
u[I] .= 1.0;
return u;
end

function getBallModel(n::Array{Int64})
n_tup = tuple(n...);
n = collect(n_tup);
u = zeros(Float32,n_tup);
X,Y,Z = meshgrid(-div(n[1],2):div(n[1],2),-div(n[2],2):div(n[2],2),-div(n[3],2):div(n[3],2))
I = (sqrt.(X.^2+Y.^2 + Z.^2)) .<= div(n[1],4);
u[I] .= 1.0;
return u;
end


using Mimi
using MimiDICE2010
using Distributions

"""
Case: We want to do an SCC calculation with MimiDICE2010, which consists of running both a base and modified model
(the latter being a model including an additional emissions pulse,
see the create_marginal_model function or create your own two models).
We then take the difference between the consumption level in these two models and
obtain the discounted net present value to get the SCC.
"""

## STEP 1 Defining typical variables for a simulation, number of trials and simulation definition

# Define the number of trials
N = 100

# define your simulation (defaults to Monte Carlo sampling)
sd = @defsim begin
    t2xco2 = Truncated(Gamma(6.47815626,0.547629469), 1.0, Inf) # a dummy distribution
end

## STEP 2 Payload object

"""
Simulation definitions can hold a user-defined payload object which is not used or modified by Mimi.
In this example, we will use the payload to hold an array of pre-computed
discount factors that we will use in the SCC calculation, as well as a storage array for saving the SCC values.
"""
# Choose what year to calculate the SCC for
scc_year = 2015
year_idx = findfirst(isequal(scc_year), MimiDICE2010.model_years)

# Pre-coimpute the discount factors for each discount rate
discount_rates = [0.03, 0.05, 0.07]
nyears = length(MimiDICE2010.model_years)
discount_factors = [[zeros(year_idx - 1)... [(1/(1 + r))^((t-year_idx)*10) for t in year_idx:nyears]...] for r in discount_rates]

# Create an array to store the compute SCC in each trial for each discount rate
scc_results = zeros(N, length(discount_rates))

# Set the payload object in the simulation definition
my_payload_object = (discount_factors, scc_results) # In this case, the payload object is a tuple which holds both both arrays
Mimi.set_payload!(sd, my_payload_object)

## STEP 3 Post-trial function
"""
User may ned to perform other calculations before or after each trial is run
For example, the SCC is calculated using two models, so this calculation needs to happen in a post-trial function
"""
function my_scc_calculation(sim_inst::SimulationInstance, trialnum::Int, ntimesteps::Int, tup::Nothing)
    mm = sim_inst.models[1]
    discount_factors, scc_results = Mimi.payload(sim_inst)  # Unpack the payload object

    marginal_damages = mm[:neteconomy, :C] * -1 * 10^12 * 12/44 # convert from trillion $/ton C to $/ton CO2; multiply by -1 to get positive value for damages
    for (i, df) in enumerate(discount_factors)
        scc_results[trialnum, i] = sum(df .* marginal_damages .* 10)
    end
end

# STEP 4: RUN THE SIMULATION
# Build the marginal model
mm = MimiDICE2010.get_marginal_model(year = scc_year)   # The additional emissions pulse will be added in the specified year

# Run
si = run(sd, mm, N; trials_output_filename = "ecs_sample.csv", post_trial_func = my_scc_calculation)

# View the scc_results by retrieving them from the payload object
scc_results = Mimi.payload(si)[2]   # Recall that the SCC array was the second of two arrays we stored in the payload tuple


using Multibreak

function part1(input)
    heights = reduce(hcat, (parse.(Int, split(line, "")) for line in eachline(input)))
    risk = 0
    neighbors = CartesianIndex.(((1, 0), (-1, 0), (0, 1), (0, -1)))
    C = CartesianIndices(heights)
    @multibreak for I in C
        for Δ in neighbors
            if I + Δ in C && heights[I + Δ] <= heights[I]
                break; continue
            end
        end
        risk += 1 + heights[I]
    end
    return risk
end

function part2(input)
    heights = reduce(hcat, (parse.(Int, split(line, "")) for line in eachline(input)))
    mark = heights .== 9
    basin_sizes = Int[]
    neighbors = CartesianIndex.(((1, 0), (-1, 0), (0, 1), (0, -1)))
    C = CartesianIndices(heights)
    queue = CartesianIndex[]
    for I in C
        if !mark[I]
            basin_size = 1
            push!(queue, I)
            mark[I] = true
            while !isempty(queue)
                x = pop!(queue)
                for Δ in neighbors
                    if x + Δ in C && !mark[x + Δ]
                        push!(queue, x + Δ)
                        basin_size += 1
                        mark[x + Δ] = true
                    end
                end
            end
            push!(basin_sizes, basin_size)
        end
    end
    return prod(sort(basin_sizes, rev = true)[1:3])
end


include("../../src/evaluation/roc_auc.jl")
using Test

data_actual = [1,1,0,0,1,0,0,0,1,1,1]
data_pred =   [1,1,1,0,1,1,1,0,0,0,1]
data_proba = [0.9, 0.8, 0.52, 0.2, 0.55, 0.52, 0.6, 0.1, 0.45, 0.3, 0.8]
roc_data_output = roc_data(data_actual, data_proba)

@testset "ROC and AUC" begin
    @test round(sum(roc_data_output[1]), digits=4) == 63.3333
    @test sum(roc_data_output[2]) == 38.8

    @test calc_auc([1,2,3], [1,2,3]) == 4.0
    @test calc_auc([1,3,2], [1,3,2]) == 4.0
    @test round(calc_auc(roc_data_output[2], roc_data_output[1]), digits=4) == 0.7667

    @test !isa(try roc_curve(data_actual, data_proba) catch ex ex end, Exception)
end


using Combinatorics

export expval_as_corrs, corr_as_expvals

function term2corr(A::QuTerm,C::Tuple)
    corrs = noaliascopy(A.corrs)
    sizehint!(corrs,length(corrs)+length(C))
    for grp in C
        cc = Corr(BaseOpProduct(collect(map(i -> A.bares.v[i], grp))))
        push!(corrs,cc)
    end
    sort!(corrs)
    QuTerm(A.nsuminds,A.δs,A.params,A.expvals,corrs,BaseOpProduct())
end

"""
    expval_as_corrs(expr::QuExpr)

Take an expression `expr=A B C + D E...` and write its expectation value in
terms of correlations ``⟨A⟩_c, ⟨B⟩_c, ⟨AB⟩_c, ⟨ABC⟩_c, \\ldots``. Note that
``⟨A⟩_c = ⟨A⟩``.

E.g., `expval_as_corrs(adag(:n)*a(:n))` returns ``⟨a^\\dagger_n a_n⟩_c +
⟨a^\\dagger_n⟩_c ⟨a_n⟩_c`` (which is equal to ``⟨a^\\dagger_n a_n⟩``), while
`expval_as_corrs(adag(:n)*a(:m)*a(:n))` returns ``\\langle a_{n}^\\dagger a_{m}
a_{n} \\rangle_{c} + \\langle a_{n}^\\dagger \\rangle_{c} \\langle a_{m}
\\rangle_{c} \\langle a_{n} \\rangle_{c} + \\langle a_{n}^\\dagger \\rangle_{c}
\\langle a_{m} a_{n} \\rangle_{c} + \\langle a_{m} \\rangle_{c} \\langle
a_{n}^\\dagger a_{n} \\rangle_{c} + \\langle a_{n} \\rangle_{c} \\langle
a_{n}^\\dagger a_{m} \\rangle_{c}``.

See also: [`expval`](@ref), [`corr`](@ref)"""
function expval_as_corrs(A::QuExpr)
    newA = QuExpr()
    corrterms = QuExpr()
    for (t,s) in A.terms
        # first calculate the correlation for the term in the sum with the "bare" indices, which means that the sum index
        # is assumed to be distinct from the indices of the expressions
        # then, calculate the correction expression for the term where the sum index is the same as any of the non-summed indices,
        # and add the correction between that and the "incorrect" term where index identity is not taken into account
        _add_corrs!(newA,t,s)

        if t.nsuminds>0 && !isempty(t.bares)
            extinds = unique(extindices(t.bares))
            for ii in 1:t.nsuminds
                sumind = sumindex(ii)
                for ind in extinds
                    f = replace_inds(sumind=>ind,(sumindex(jj)=>sumindex(jj-1) for jj=ii+1:t.nsuminds)...)
                    tb = f(t.bares)
                    tbc = noaliascopy(tb)
                    normal_order!(tb,corrterms)
                    @assert isempty(corrterms)
                    if tb != tbc
                        tright = _normalize_without_commutation(QuTerm(t.nsuminds-1,f.(t.δs),f.(t.params),f.(t.expvals),f.(t.corrs),tb))
                        _add_corrs!(newA,tright,s)
                        twrong = QuTerm(tright.nsuminds,tright.δs,tright.params,tright.expvals,tright.corrs,tbc)
                        _add_corrs!(newA,twrong,-s,_add_with_normal_order!)
                    end
                end
            end
        end
    end
    newA
end

function _add_corrs!(A,t,s,addfun! = _add_sum_term!)
    if isempty(t.bares)
        addfun!(A,t,s)
    else
        for C in expval2corrs_inds(length(t.bares))
            addfun!(A,term2corr(t,C),s)
        end
    end
end

function _add_expvs!(A,t,s,addfun! = _add_sum_term!)
    if isempty(t.bares)
        addfun!(A,t,s)
    else
        for (C,sC) in corr2expvals_inds(length(t.bares))
            addfun!(A,term2expvals(t,C),s*sC)
        end
    end
end

"""
    CorrTup_isless(a,b)

isless for Tuples of integers that represent Corr of sorted operators (with n representing An such that n<m == An<Am).
(n,m,...) ≡ Corr(An*Am*...). Defined in such a way that the same order is obtained as with `BaseOperator` objects
"""
CorrTup_isless(a::NTuple{N,Int},b::NTuple{M,Int}) where {N,M} = N==M ? a<b : N<M

"""
    CorrPerm_isless(a,b)

isless for Tuples of Tuples representing products of Corr (see above) of a permutation of operators.
E.g., ((1,3),(2,)) represents <A1 A3>_C <A2>_C.
It is assumed that the total number of operators in a and b is equal, i.e., that `sum(length.(a)) == sum(length.(b))`."""
function CorrPerm_isless(a::Tuple,b::Tuple)
    for (ca,cb) in zip(a,b)
        ca == cb || return CorrTup_isless(ca,cb)
    end
    # if we reach here, they are equal
    return false
end

"""
    expval2corrs_inds(N::Int)

for N operators, create an array of tuples of tuples that represents the terms
in a sum of products of correlators. Each tuple corresponds to a sum term, see
[`CorrTup_isless`](@ref) and [`CorrPerm_isless`](@ref) for details of the
format. The returned array and terms are sorted such that if the N operators are
sorted, the represented expression is also sorted with the conventions of the
QuantumAlgebra package. This allows to directly return a normal-ordered form.
"""
function expval2corrs_inds(N::Int)
    # get return value for N from cache if present, otherwise calculate and cache it
    get!(_EXPVAL2CORRS_CACHE,N) do
        terms = Any[Tuple(tuple.(1:N))]
        for n = 2:N-1
            append!(terms,ncomb_inds(n,1:N))
        end
        [(tuple(1:N...),),sort!(terms,lt=CorrPerm_isless)...]
    end
end
# preload the cache for N=0 and N=1 so we do not have to special-case those above
const _EXPVAL2CORRS_CACHE = Dict{Int,Vector{Any}}(0 => [((),)], 1 => [((1,),)])

function ncomb_inds(n,inds,used_combs=Set())
    terms = []
    for c in combinations(inds,n)
        C = Tuple(c)
        C in used_combs && continue
        push!(used_combs,C)
        notCinds = setdiff(inds,c)
        S = tuple.(notCinds)
        push!(terms,Tuple(sort!([C,S...],lt=CorrTup_isless)))
        if n<length(inds)
            # pass a copy of used_combs here so getting, e.g., <jk>*C here does not prevent a term <jk> <a><b> later
            for CC in ncomb_inds(n,notCinds,copy(used_combs))
                push!(terms,Tuple(sort!([C,CC...],lt=CorrTup_isless)))
            end
        end
    end
    terms
end

function term2expvals(A::QuTerm,C::Tuple)
    expvals = noaliascopy(A.expvals)
    sizehint!(expvals,length(expvals)+length(C))
    for grp in C
        cc = ExpVal(BaseOpProduct(collect(map(i -> A.bares.v[i], grp))))
        push!(expvals,cc)
    end
    sort!(expvals)
    QuTerm(A.nsuminds,A.δs,A.params,expvals,A.corrs,BaseOpProduct())
end

_canon_op_to_num(A::BaseOperator) = _canon_ind_to_num(only(A.inds))
_expval2tuple(A::ExpVal) = Tuple(_canon_op_to_num.(A.ops.v))
const _CORR2EXPVALS_CACHE = Dict{Int,Vector{Any}}()
function corr2expvals_inds(N)
    get!(_CORR2EXPVALS_CACHE,N) do
        A = QuExpr(QuTerm(BaseOpProduct(BosonDestroy.(:a,_canon_ind.(1:N)))))
        # expval_as_corrs(As...) = <As...> in terms of correlators <A1 A2...>c
        # <As...>c = <As...> - <As...> + <As...>c, where RHS expression only contains <As...> and lower-order correlators
        cA = expval(A) - expval_as_corrs(A) + corr(A)
        # rewrite the lower-order correlators in terms of expectation values
        ncA = QuExpr()
        for (t,s) in cA.terms
            nt = noaliascopy(t)
            empty!(nt.corrs)
            nts = QuExpr(nt)
            for co in t.corrs
                nts *= QuExpr((term2expvals(QuTerm(co.ops),C)=>sC for (C,sC) in corr2expvals_inds(length(co))))
            end
            for (ntt,ns) in nts.terms
                _add_with_normal_order!(ncA,ntt,ns*s)
            end
        end
        [Tuple(_expval2tuple.(t.expvals))=>s for (t,s) in sort!(collect(ncA.terms),by=x->x[1])]
    end
end

function corr_as_expvals(A::QuExpr)
    newA = QuExpr()
    corrterms = QuExpr()
    for (t,s) in A.terms
        # first calculate the correlation for the term in the sum with the "bare" indices, which means that the sum index
        # is assumed to be distinct from the indices of the expressions
        # then, calculate the correction expression for the term where the sum index is the same as any of the non-summed indices,
        # and add the correction between that and the "incorrect" term where index identity is not taken into account
        _add_expvs!(newA,t,s)

        if t.nsuminds>0 && !isempty(t.bares)
            extinds = unique(extindices(t.bares))
            for ii in 1:t.nsuminds
                sumind = sumindex(ii)
                for ind in extinds
                    f = replace_inds(sumind=>ind,(sumindex(jj)=>sumindex(jj-1) for jj=ii+1:t.nsuminds)...)
                    tb = f(t.bares)
                    tbc = noaliascopy(tb)
                    normal_order!(tb,corrterms)
                    @assert isempty(corrterms)
                    if tb != tbc
                        tright = _normalize_without_commutation(QuTerm(t.nsuminds-1,f.(t.δs),f.(t.params),f.(t.expvals),f.(t.corrs),tb))
                        _add_expvs!(newA,tright,s)
                        twrong = QuTerm(tright.nsuminds,tright.δs,tright.params,tright.expvals,tright.corrs,tbc)
                        _add_expvs!(newA,twrong,-s,_add_with_normal_order!)
                    end
                end
            end
        end
    end
    newA
end


using Libdl
using MPI
using MSTM
using OffsetArrays
using Test

@testset "MSTM" begin
    MPI.Init()

    mstm_lib = ("CI" => "true") in ENV ? "/home/runner/work/MSTM.jl/MSTM.jl/shared/mstm" : "../shared/mstm"
    mstm = Libdl.dlopen(mstm_lib)

    @testset "Constants" begin
        @testset "init!($notd)" for notd in [5, 10, 20, 50, 100]
            @test begin
                ctx = MSTM.Constants.init(notd)
                bcof_julia, _, _, vwh_coef_julia = MSTM.Constants.get_offset_constants(ctx)

                MSTM.Wrapper.init!(mstm, notd)
                bcof_fortran = MSTM.Wrapper.bcof(mstm, notd)
                vwh_coef_fortran = MSTM.Wrapper.vwh_coef(mstm, notd)

                isapprox(bcof_julia, bcof_fortran) && isapprox(vwh_coef_julia, vwh_coef_fortran)
            end
        end
    end

    @testset "SpecialFunctions" begin
        ϵ = 1e-8
        @testset "ricbessel($n, $ds)" for (n, ds) in [(50, 10.0), (100, 100.0), (150, 200.0)]
            @test begin
                Ψ_julia = MSTM.SpecialFunctions.ricbessel(n, ds, ϵ)
                Ψ_fortran = MSTM.Wrapper.ricbessel(mstm, n, ds, ϵ)
                isapprox(Ψ_julia, Ψ_fortran)
            end
        end

        @testset "richankel($n, $ds)" for (n, ds) in [(50, 10.0), (100, 100.0), (150, 200.0)]
            @test begin
                Ξ_julia = MSTM.SpecialFunctions.richankel(n, ds)
                Ξ_fortran = MSTM.Wrapper.richankel(mstm, n, ds)
                isapprox(Ξ_julia, Ξ_fortran)
            end
        end

        @testset "cricebessel($n, $ds)" for (n, ds) in [
            (10, ComplexF64(0, 0)),
            (20, ComplexF64(10, 10)),
            (500, ComplexF64(0.5, 5.1)),
        ]
            @test begin
                Ψ_julia = MSTM.SpecialFunctions.cricbessel(n, ds)
                Ψ_fortran = MSTM.Wrapper.cricbessel(mstm, n, ds)
                isapprox(Ψ_julia, Ψ_fortran)
            end
        end

        @testset "crichankel($n, $ds)" for (n, ds) in [
            (10, ComplexF64(0, 0)),
            (20, ComplexF64(10, 10)),
            (500, ComplexF64(0.5, 5.1)),
        ]
            @test begin
                Ξ_julia = MSTM.SpecialFunctions.crichankel(n, ds)
                Ξ_fortran = MSTM.Wrapper.crichankel(mstm, n, ds)
                isapprox(Ξ_julia, Ξ_fortran)
            end
        end

        @testset "cspherebessel($n, $z)" for (n, z) in [
            (10, ComplexF64(0, 0)),
            (20, ComplexF64(10, 10)),
            (500, ComplexF64(0.5, 5.1)),
        ]
            @test begin
                csj_julia, csy_julia = MSTM.SpecialFunctions.cspherebessel(n, z)
                csj_fortran, csy_fortran = MSTM.Wrapper.cspherebessel(mstm, n, z)
                isapprox(csj_julia, csj_fortran) && isapprox(csy_julia, csy_fortran)
            end
        end

        @testset "vcfunc($m, $n, $k, $l)" for (m, n, k, l) in [(3, 4, 3, 10), (30, 40, 30, 100), (-4, 10, 4, 5)]
            @test begin
                ctx = MSTM.Constants.init()
                vcn_julia = MSTM.SpecialFunctions.vcfunc(ctx, m, n, k, l)
                vcn_fortran = MSTM.Wrapper.vcfunc(mstm, m, n, k, l)
                isapprox(vcn_julia, vcn_fortran)
            end
        end

        @testset "normalizedlegendre($cbe, $mmax, $nmax)" for (cbe, mmax, nmax) in
                                                              [(0.5, 10, 10), (-0.21, 15, 20), (0.98, 30, 34)]
            @test begin
                ctx = MSTM.Constants.init()
                dc_julia = MSTM.SpecialFunctions.normalizedlegendre(ctx, cbe, mmax, nmax)
                dc_fortran = MSTM.Wrapper.normalizedlegendre(mstm, cbe, mmax, nmax)
                isapprox(dc_julia, dc_fortran)
            end
        end

        @testset "rotcoef($cbe, $kmax, $nmax)" for (cbe, kmax, nmax) in [(0.5, 10, 10), (-0.21, 15, 20), (0.98, 30, 34)]
            @test begin
                ctx = MSTM.Constants.init()
                dc_julia = MSTM.SpecialFunctions.rotcoef(ctx, cbe, kmax, nmax)
                dc_fortran = MSTM.Wrapper.rotcoef(mstm, cbe, kmax, nmax)
                isapprox(dc_julia, dc_fortran)
            end
        end

        @testset "taufunc($cb, $nmax)" for (cb, nmax) in [(0.5, 10), (-0.21, 20), (0.98, 34)]
            @test begin
                ctx = MSTM.Constants.init()
                τ_julia = MSTM.SpecialFunctions.taufunc(ctx, cb, nmax)
                τ_fortran = MSTM.Wrapper.taufunc(mstm, cb, nmax)
                isapprox(τ_julia, τ_fortran)
            end
        end

        @testset "pifunc($cb, $E_ϕ, $nmax, $ndim)" for (cb, E_ϕ, nmax, ndim) in [
            (0.5, 1 + 1.0im, 10, 10),
            (-0.21, 0.2 - 0.3im, 14, 20),
            (0.98, -1.2 + 2.5im, 15, 15),
        ]
            @test begin
                ctx = MSTM.Constants.init()
                π_julia = MSTM.SpecialFunctions.pifunc(ctx, cb, E_ϕ, nmax, ndim)
                π_fortran = MSTM.Wrapper.pifunc(mstm, cb, E_ϕ, nmax, ndim)
                isapprox(π_julia, π_fortran)
            end
        end

        @testset "planewavecoef($α, $β, $nodr)" for (α, β, nodr) in
                                                    [(0.5, 0.6, 10), (-0.21, -0.45, 14), (0.98, -0.45, 20)]
            @test begin
                ctx = MSTM.Constants.init()
                pmnp0_julia = MSTM.SpecialFunctions.planewavecoef(ctx, α, β, nodr)
                pmnp0_fortran = MSTM.Wrapper.planewavecoef(mstm, α, β, nodr)
                isapprox(pmnp0_julia, pmnp0_fortran)
            end
        end

        @testset "gaussianbeamcoef($α, $β, $cbeam, $nodr)" for (α, β, cbeam, nodr) in [
            (0.5, 0.6, 15.0, 10),
            (-0.21, -0.45, 0.5, 14),
            (0.98, -0.45, 100.0, 20),
        ]
            @test begin
                ctx = MSTM.Constants.init()
                pmnp0_julia = MSTM.SpecialFunctions.gaussianbeamcoef(ctx, α, β, cbeam, nodr)
                pmnp0_fortran = MSTM.Wrapper.gaussianbeamcoef(mstm, α, β, cbeam, nodr)
                isapprox(pmnp0_julia, pmnp0_fortran)
            end
        end

        @testset "sphereplanewavecoef" begin
            @test begin
                ctx = MSTM.Constants.init()

                nsphere = 3
                nodr = [3, 3, 3]
                α = 0.5
                β = 0.5
                rpos = [[0.0, 0.0, 0.0] [0.0, 0.0, -0.5] [0.0, 0.0, -1.0]]
                hostsphere = [0, 1, 2]
                numberfieldexp = [3, 3, 3]
                rimedium = [1.1 + 1.3im, 1.0 - 0.5im]

                pmnp_julia = MSTM.SpecialFunctions.sphereplanewavecoef(
                    ctx,
                    nodr,
                    α,
                    β,
                    rpos,
                    hostsphere,
                    numberfieldexp,
                    rimedium,
                )

                pmnp_fortran =
                    MSTM.Wrapper.sphereplanewavecoef(mstm, nodr, α, β, rpos, hostsphere, numberfieldexp, rimedium)

                isapprox(pmnp_julia, pmnp_fortran)
            end
        end

        @testset "axialtrancoefrecurrence($itype, $r, $nmax, $lmax)" for (itype, r, nmax, lmax) in [
            (1, 0.0, 10, 10),
            (1, 1.3, 15, 20),
            (3, 0.0, 4, 5),
            (3, 4.3, 20, 15),
        ]
            @test begin
                ri = [1.0 + 0.5im, 1.0 - 0.5im]
                ctx = MSTM.Constants.init()
                ac_julia = MSTM.SpecialFunctions.axialtrancoefrecurrence(ctx, itype, r, ri, nmax, lmax)
                ac_fortran = MSTM.Wrapper.axialtrancoefrecurrence(mstm, itype, r, ri, nmax, lmax)
                isapprox(ac_julia, ac_fortran)
            end
        end

        @testset "tranordertest($r, $ri, $lmax, $ϵ)" for (r, ri, lmax, ϵ) in [
            (0.01, 1.1 + 0.3im, 10, 1e-4),
            (0.1, 1.1 + 0.03im, 20, 1e-6),
            (0.35, 1.1 + 0.003im, 50, 1e-6),
        ]
            @test begin
                ctx = MSTM.Constants.init()
                nmax_julia = MSTM.SpecialFunctions.tranordertest(ctx, r, ri, lmax, ϵ)
                nmax_fortran = MSTM.Wrapper.tranordertest(mstm, r, ri, lmax, ϵ)
                nmax_julia == nmax_fortran
            end
        end

        @testset "gentrancoef($itype, $xptran, $nrow0, $nrow1, $ncol0, $ncol1, $iaddrow0, $iaddcol0)" for (
            itype,
            xptran,
            nrow0,
            nrow1,
            ncol0,
            ncol1,
            iaddrow0,
            iaddcol0,
        ) in [
            (1, [0.0, 0.0, 0.0], 1, 5, 1, 5, 0, 0),
            (3, [0.0, 0.0, 0.0], 1, 5, 1, 5, 0, 0),
            (1, [1.0, -0.05, 0.04], 1, 5, 1, 5, 0, 0),
            (3, [1.0, -0.05, 0.04], 1, 5, 1, 5, 0, 0),
            (1, [1.0, -0.05, 0.04], 1, 5, 1, 5, 5, 5),
            (3, [1.0, -0.05, 0.04], 1, 5, 1, 5, 2, 3),
        ]
            @test begin
                ri = [1.0 + 0.5im, 1.0 - 0.5im]
                ctx = MSTM.Constants.init()
                ac_julia = MSTM.SpecialFunctions.gentrancoef(
                    ctx,
                    itype,
                    xptran,
                    ri,
                    nrow0,
                    nrow1,
                    ncol0,
                    ncol1,
                    iaddrow0,
                    iaddcol0,
                )
                ac_fortran =
                    MSTM.Wrapper.gentrancoef(mstm, itype, xptran, ri, nrow0, nrow1, ncol0, ncol1, iaddrow0, iaddcol0)
                isapprox(ac_julia, ac_fortran)
            end
        end

        @testset "eulerrotation($xp, $eulerangf, $dir)" for (xp, eulerangf, dir) in [
            ([1.02, 2.03, 3.04], [0.95, 2.33, 1.34], 1),
            ([2.43, -23.2, -19.4], [10.2, 12.3, 11.3], -1),
            ([-0.42, -0.2, -0.001], [-23.2, -23.4, 12.2], 1),
        ]
            @test begin
                xprot_julia = MSTM.SpecialFunctions.eulerrotation(xp, eulerangf, dir)
                xprot_fortran = MSTM.Wrapper.eulerrotation(mstm, xp, eulerangf, dir)
                isapprox(xprot_julia, xprot_fortran)
            end
        end

        @testset "planewavetruncationorder($r, $rimedium, $ϵ)" for (r, rimedium, ϵ) in [
            (0.01, [1.1 + 0.3im, 0.8 - 0.2im], 1e-4),
            (0.1, [1.1 + 0.03im, 1.1 - 0.05im], 1e-6),
            (0.35, [1.1 + 0.003im, 1.2 + 1.1im], 1e-8),
        ]
            @test begin
                nodr_julia = MSTM.SpecialFunctions.planewavetruncationorder(r, rimedium, ϵ)
                nodr_fortran = MSTM.Wrapper.planewavetruncationorder(mstm, r, rimedium, ϵ)
                nodr_julia == nodr_fortran
            end
        end

        @testset "vwhcalc($rpos, $nodr, $itype)" for (rpos, nodr, itype) in [
            ([0.0, 0.0, 0.0], 5, 1),
            ([0.0, 0.0, 0.0], 5, 3),
            ([1.0, -0.05, 0.04], 10, 1),
            ([1.0, -0.05, 0.04], 10, 3),
            ([1.0, -0.05, 0.04], 20, 1),
            ([1.0, -0.05, 0.04], 20, 3),
        ]
            @test begin
                ri = [1.0 + 0.5im, 1.0 - 0.5im]
                ctx = MSTM.Constants.init()
                vwh_julia = MSTM.SpecialFunctions.vwhcalc(ctx, rpos, ri, nodr, itype)
                vwh_fortran = MSTM.Wrapper.vwhcalc(mstm, rpos, ri, nodr, itype)
                isapprox(vwh_julia, vwh_fortran)
            end
        end

        @testset "vwhaxialcalc($rpos, $nodr, $itype)" for (rpos, nodr, itype) in [
            ([0.0, 0.0, 0.0], 5, 1),
            ([0.0, 0.0, 0.0], 5, 3),
            ([1.0, -0.05, 0.04], 10, 1),
            ([1.0, -0.05, 0.04], 10, 3),
            ([1.0, -0.05, 0.04], 20, 1),
            ([1.0, -0.05, 0.04], 20, 3),
        ]
            @test begin
                ri = [1.0 + 0.5im, 1.0 - 0.5im]
                ctx = MSTM.Constants.init()
                vwh_julia = MSTM.SpecialFunctions.vwhaxialcalc(ctx, rpos, ri, nodr, itype)
                vwh_fortran = MSTM.Wrapper.vwhaxialcalc(mstm, rpos, ri, nodr, itype)
                isapprox(vwh_julia, vwh_fortran)
            end
        end

        @testset "twobytwoinverse($mat)" for mat in [
            [1.1+1.0im 2.2-1.5im; 3.3-0.2im -2.5+0.4im],
            [1.1+1.0im 1.2+1.5im; 0.3-0.2im -1.5+0.3im],
        ]
            @test begin
                imat_julia = MSTM.SpecialFunctions.twobytwoinverse(mat)
                imat_fortran = MSTM.Wrapper.twobytwoinverse(mstm, mat)
                isapprox(imat_julia, imat_fortran)
            end
        end
    end

    @testset "Data" begin
        @testset "read MSTM config from .toml file" begin
            config = MSTM.Data.MSTMConfig("fixtures/default.toml")
            @test length(config.spheres) == 2
        end
    end

    @testset "Mie" begin
        ri = OffsetArray(
            [[1.1 + 1.3im, 0.1 + 0.0im] [1.2 + 0.2im, 0.1 + 0.0im] [1.5 - 0.4im, 0.1 + 0.0im] [
                1.5 + 1.5im,
                0.1 + 0.0im,
            ]],
            1:2,
            0:3,
        )

        @testset "mieoa($x, $ri, $nodr0, $qeps, $ri_medium)" for (x, ri, nodr0, qeps, ri_medium) in [
            (0.35, [1.1 + 1.0im, 2.2 - 1.5im], 3, 1e-6, nothing),
            (0.15, [1.1 + 1.0im, 2.2 - 1.5im], 3, 1e-6, [1.0 + 0.1im, 1.0 + 0.1im]),
            (0.25, [1.1 + 1.0im, 2.2 - 1.5im], 4, 0.0, [1.0 + 0.1im, 1.0 + 0.1im]),
            (0.45, [1.1 + 1.0im, 2.2 - 1.5im], 4, -6.0, [1.0 + 0.1im, 1.0 + 0.1im]),
        ]
            @test begin
                qext_julia, qsca_julia, qabs_julia, nodr_julia, _, _, _, _, _, _ =
                    MSTM.Mie.mieoa(x, ri, nodr0, qeps, ri_medium = ri_medium)
                qext_fortran, qsca_fortran, qabs_fortran, nodr_fortran =
                    MSTM.Wrapper.mieoa(mstm, x, ri, nodr0, qeps, ri_medium = ri_medium)
                isapprox(qext_julia, qext_fortran) &&
                    isapprox(qsca_julia, qsca_fortran) &&
                    isapprox(qabs_julia, qabs_fortran) &&
                    nodr_julia == nodr_fortran
            end
        end

        @testset "miecoefcalc($xsp, $hostsphere, $numberfieldexp, $qeps) (No T-Matrix files)" for (
            xsp,
            hostsphere,
            numberfieldexp,
            qeps,
        ) in [
            ([1.2, 0.4, 0.1], [0, 1, 2], [2, 2, 2], 1e-4),
            ([1.2, 0.4, 0.1], [0, 1, 2], [3, 3, 3], 1e-6),
            ([1.2, 0.4, 0.1], [0, 1, 2], [3, 3, 3], 1e-8),
            ([1.2, 0.4, 0.1], [0, 1, 2], [4, 4, 4], 1e-6),
        ]
            @test begin
                params = MSTM.Data.MSTMParameters()
                mie_julia = MSTM.Mie.miecoefcalc(params, xsp, ri, hostsphere, numberfieldexp, qeps)
                mie_fortran = MSTM.Wrapper.miecoefcalc(mstm, xsp, ri, hostsphere, numberfieldexp, qeps)
                isapprox(mie_julia.an, mie_fortran.an) &&
                    isapprox(mie_julia.cn, mie_fortran.cn) &&
                    isapprox(mie_julia.qext, mie_fortran.qext) &&
                    isapprox(mie_julia.qabs, mie_fortran.qabs) &&
                    mie_julia.order == mie_fortran.order &&
                    mie_julia.offset == mie_fortran.offset &&
                    mie_julia.block == mie_fortran.block &&
                    mie_julia.block_offset == mie_fortran.block_offset &&
                    mie_julia.number_field_expansions == mie_fortran.number_field_expansions &&
                    mie_julia.is_optically_active == mie_fortran.is_optically_active &&
                    mie_julia.number_eqns == mie_fortran.number_eqns
            end
        end

        @testset "onemiecoeffmult($i, $mie_coefficent)" for (i, mie_coefficent) in
                                                            [(1, 'a'), (2, 'c'), (3, 'd'), (2, 'u'), (3, 'v')]
            @test begin
                xsp = [1.2, 0.4, 0.1]
                hostsphere = [0, 1, 2]
                numberfieldexp = [3, 3, 3]
                qeps = 1e-6
                cb = 0.5
                E_ϕ = 1.0 + 1.0im
                params = MSTM.Data.MSTMParameters()
                mie = MSTM.Mie.miecoefcalc(params, xsp, ri, hostsphere, numberfieldexp, qeps)
                nodr = mie.order[i]
                ctx = MSTM.Constants.init()
                MSTM.Wrapper.miecoefcalc(mstm, xsp, ri, hostsphere, numberfieldexp, qeps)
                cx = MSTM.SpecialFunctions.pifunc(ctx, cb, E_ϕ, nodr, nodr)
                cy_julia = MSTM.Mie.onemiecoeffmult(mie, i, cx, mie_coefficent)
                cy_fortran = MSTM.Wrapper.onemiecoeffmult(mstm, i, nodr, cx, mie_coefficent)
                isapprox(cy_julia, cy_fortran)
            end
        end

        @testset "multmiecoeffmult" begin
            xsp = [1.2, 0.4, 0.1]
            hostsphere = [0, 1, 2]
            numberfieldexp = [3, 2, 1]
            qeps = 1e-6
            cb = 0.5
            E_ϕ = 1.0 + 1.0im
            params = MSTM.Data.MSTMParameters()
            mie = MSTM.Mie.miecoefcalc(params, xsp, ri, hostsphere, numberfieldexp, qeps)
            MSTM.Wrapper.miecoefcalc(mstm, xsp, ri, hostsphere, numberfieldexp, qeps)
            for i in 1:4
                @test begin
                    nrhs = i + 3
                    ain = 1.0im * rand(mie.number_eqns * nrhs) + rand(mie.number_eqns * nrhs)
                    aout_julia = MSTM.Mie.multmiecoeffmult(mie, nrhs, i % 2, ain)
                    aout_fortran = MSTM.Wrapper.multmiecoeffmult(mstm, mie.number_eqns, nrhs, i % 2, ain)
                    isapprox(aout_julia, aout_fortran)
                end
            end
        end
    end

    Libdl.dlclose(mstm)
end


# Find the smallest order of magnitude that is larger than xspan This is a
# little opaque because I want to avoid assuming the log function is defined
# over typeof(xspan)
function bounding_order_of_magnitude(xspan)
    DT = typeof(xspan)
    one_dt = one(DT)

    a = 1
    step = 1
    while xspan < 10.0^a * one_dt
        a -= step
    end

    b = 1
    step = 1
    while xspan > 10.0^b * one_dt
        b += step
    end

    while a + 1 < b
        c = div(a + b, 2)
        if xspan < 10.0^c * one_dt
            b = c
        else
            a = c
        end
    end

    return b
end

# Empty catchall
optimize_ticks() = {}

# Find some reasonable values for tick marks.
#
# This is basically Wilkinson's ad-hoc scoring method that tries to balance
# tight fit around the data, optimal number of ticks, and simple numbers.
#
# Args:
#   x_min: minimum value occuring in the data.
#   x_max: maximum value occuring in the data.
#
# Returns:
#   A Float64 vector containing tick marks.
#
function optimize_ticks{T}(x_min::T, x_max::T; extend_ticks::Bool=false)
    if x_min == x_max
        return [x_min], x_min - one(T), x_min + one(T)
    end

    # tick intervals and scores
    # TODO: these should perhaps be part of the theme
    const Q = {(1,1), (5, 0.9), (2, 0.7), (25, 0.5), (3, 0.2)}
    const n = length(Q)

    # number of ticks
    const k_min   = 2
    const k_max   = 10
    const k_ideal = 5

    # generalizing "order of magnitude"
    xspan = x_max - x_min
    z = bounding_order_of_magnitude(xspan)
    one_t = one(T)

    high_score = -Inf
    z_best = 0.0
    k_best = 0.0
    r_best = 0.0
    q_best = 0.0

    while k_max * 10.0^(z+1) * one_t > xspan
        for k in k_min:k_max
            for (q, qscore) in Q
                span = (k - 1) * q * 10.0^z * one_t
                if span < xspan
                    continue
                end

                r = ceil((x_max - span) / (q*10.0^z * one_t))
                while r*q*10.0^z * one_t <= x_min
                    has_zero = r <= 0 && abs(r) < k

                    # simplicity
                    s = has_zero ? 1.0 : 0.0

                    # granularity
                    g = 0 < k < 2k_ideal ? 1 - abs(k - k_ideal) / k_ideal : 0.0

                    # coverage
                    c = 1.5 * xspan/span

                    score = (1/4)g + (1/6)s + (1/3)c + (1/4)qscore

                    # strict limits on coverage
                    if span >= 2.0*xspan || span < xspan
                        score -= 1000
                    end

                    if score > high_score
                        (q_best, r_best, k_best, z_best) = (q, r, k, z)
                        high_score = score
                    end
                    r += 1
                end
            end
        end
        z -= 1
    end

    span = q_best * 10.0^z_best * one_t
    if extend_ticks
        S = Array(typeof(1.0 * one_t), int(3 * k_best))
        for i in 0:(3*k_best - 1)
            S[i+1] = (r_best + i - k_best) * span
        end
        viewmin, viewmax = S[k_best + 1], S[2 * k_best]
    else
        S = Array(typeof(1.0 * one_t), int(k_best))
        for i in 0:(k_best - 1)
            S[i+1] = (r_best + i) * span
        end
        viewmin, viewmax = S[1], S[end]
    end

    S, viewmin, viewmax
end

function optimize_ticks(x_min::Date, x_max::Date; extend_ticks::Bool=false)
    # This can be pretty simple. We are choosing ticks on one of three
    # scales: years, months, days.
    if year(x_max) - year(x_min) <= 1
        if year(x_max) == year(x_min) && month(x_max) - month(x_min) <= 1
            ticks = Date[]
            if x_max - x_min > days(7)
                # This will probably need to be smarter
                push!(ticks, x_min)
                while true
                    next_month = date(year(ticks[end]), month(ticks[end])) + month(1)
                    while ticks[end] + week(1) < next_month - days(2)
                        push!(ticks, ticks[end] + week(1))
                    end
                    push!(ticks, next_month)
                    if next_month >= x_max
                        break
                    end
                end
            else
                push!(ticks, x_min)
                while ticks[end] < x_max
                    push!(ticks, ticks[end] + day(1))
                end
            end

            viewmin, viewmax = ticks[1], ticks[end]
            ticks, viewmin, viewmax
        else
            ticks = Date[]
            push!(ticks, date(year(x_min), month(x_min)))
            while ticks[end] < x_max
                push!(ticks, ticks[end] + month(1))
            end
            viewmin, viewmax = ticks[1], ticks[end]

            ticks, x_min, x_max
        end
    else
        ticks, viewmin, viewmax =
            optimize_ticks(year(x_min), year(x_max), extend_ticks=extend_ticks)
        Date[date(y) for y in ticks], date(viewmin), date(viewmax)
    end
end


# This file contains functions for the Revenue Efficiency Russell DEA model
"""
    RevenueRussellDEAModel
An data structure representing a revenue Russell DEA model.
"""
struct RevenueRussellDEAModel <: AbstractRevenueDEAModel
    n::Int64
    m::Int64
    s::Int64
    rts::Symbol
    monetary::Bool
    dmunames::Union{Vector{String},Nothing}
    eff::Vector
    lambda::SparseMatrixCSC{Float64, Int64}
    techeff::Vector
    alloceff::Vector
    normalization::Vector
    Xtarget::Matrix
    Ytarget::Matrix
end

"""
    dearevenuerussell(X, Y, P)
Compute revenue efficiency using Russell data envelopment analysis for
inputs `X`, outputs `Y` and price of outputs `P`.

# Optional Arguments
- `rts=:VRS`: chooses variable returns to scale. For constant returns to scale choose `:CRS`.
- `names`: a vector of strings with the names of the decision making units.

# Examples
```jldoctest
julia> X = [1; 1; 1; 1; 1; 1; 1; 1];

julia> Y = [7 7; 4 8; 8 4; 3 5; 3 3; 8 2; 6 4; 1.5 5];

julia> P = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1];

julia> dearevenuerussell(X, Y, P)
Russell Revenue DEA Model
DMUs = 8; Inputs = 1; Outputs = 2
Orientation = Output; Returns to Scale = VRS
──────────────────────────────────
   Revenue  Technical   Allocative
──────────────────────────────────
1  0.0       0.0       0.0
2  0.25      0.0       0.25
3  0.25      0.0       0.25
4  1.0       0.866667  0.133333
5  1.33333   1.33333   2.22045e-16
6  1.0       0.5       0.5
7  0.5       0.458333  0.0416667
8  2.5       2.05556   0.444444
──────────────────────────────────
```
"""
function dearevenuerussell(X::Union{Matrix,Vector}, Y::Union{Matrix,Vector},
    P::Union{Matrix,Vector}; rts::Symbol = :VRS,
    monetary::Bool = false,
    names::Union{Vector{String},Nothing} = nothing,
    optimizer::Union{DEAOptimizer,Nothing} = nothing)::RevenueRussellDEAModel

    # Check parameters
    nx, m = size(X, 1), size(X, 2)
    ny, s = size(Y, 1), size(Y, 2)

    np, sp = size(P, 1), size(P, 2)

    if nx != ny
        throw(DimensionMismatch("number of rows in X and Y ($nx, $ny) are not equal"));
    end
    if np != ny
        throw(DimensionMismatch("number of rows in P and Y ($np, $ny) are not equal"));
    end
    if sp != s
        throw(DimensionMismatch("number of columns in P and Y ($sp, $s) are not equal"));
    end

    # Default optimizer
    if optimizer === nothing
        optimizer = DEAOptimizer(:LP)
    end

    # Get maximum revenue targets and lambdas
    n = nx

    Xtarget = X[:,:]
    Ytarget, rlambdaeff = deamaxrevenue(X, Y, P, rts = rts, dispos = :Strong, optimizer = optimizer)

    # Revenue, technical and allocative efficiency
    PY = P .* Y
    PY[PY .== 0] .= PY[PY .== 0] .+ Inf
    sno0 = sum(Y .!= 0, dims = 2)

    refficiency  = vec( sum(P .* Ytarget, dims = 2)  .- sum(P .* Y, dims = 2) )
    normalization = vec(sno0 .* minimum(PY, dims = 2) )
    techefficiency = efficiency(dearussell(X, Y, orient = :Output, rts = rts, slack = false, optimizer = optimizer)) .- 1

    if monetary
        techefficiency = techefficiency .* normalization
    else
        refficiency = refficiency ./ normalization
    end

    allocefficiency = refficiency .- techefficiency

    return RevenueRussellDEAModel(n, m, s, rts, monetary, names, refficiency, rlambdaeff, techefficiency, allocefficiency, normalization, Xtarget, Ytarget)
end

function Base.show(io::IO, x::RevenueRussellDEAModel)
    compact = get(io, :compact, false)

    n = nobs(x)
    m = ninputs(x)
    s = noutputs(x)
    dmunames = names(x)

    eff = efficiency(x)
    techeff = efficiency(x, :Technical)
    alloceff = efficiency(x, :Allocative)

    if !compact
        print(io, "Russell Revenue DEA Model \n")
        print(io, "DMUs = ", n)
        print(io, "; Inputs = ", m)
        print(io, "; Outputs = ", s)
        print(io, "\n")
        print(io, "Orientation = Output")
        print(io, "; Returns to Scale = ", string(x.rts))
        print(io, "\n")

        show(io, CoefTable(hcat(eff, techeff, alloceff), ["Revenue", "Technical", "Allocative"], dmunames))
    end

end


# Commutative semirings have commutative multiplication (addition is commutative by definition).
@traitdef IsCommutative{S<:AbstractSemiringElement}
@traitimpl IsCommutative{S} <- iscommutative(S)
is_commutative(x) = false

# Idempotent semirings are semirings with idempotent addition: `a+a = a`
@traitdef IsIdempotent{S<:AbstractSemiringElement}
@traitimpl IsIdempotent{S} <- isidempotent(S)
isidempotent(::Type{<:AbstractSemiringElement}) = false

# Star semirings have an additional unary operator `*`: `a^* = \sum_{i = 0}^{\infty} a^i`.
@traitdef IsStar{S<:AbstractSemiringElement}
@traitimpl IsStar{S} <- isstar(S)
isstar(::Type{<:AbstractSemiringElement}) = false