using ApproxFun, OrdinaryDiffEq, Sundials
using DiffEqDevTools
using LinearAlgebra
using Plots; gr()

function cheb(N)
    N==0 && return (0,1)
    x = cos.(pi*(0:N)/N)
    c = [2; ones(N-1,1); 2].*(-1).^(0:N)
    X = hcat([x for i in 1:N+1]...)
    dX = X-X'
    D  = (c*(1 ./c)')./(dX+I)      # off-diagonal entries
    D  = D .- Diagonal(vec(sum(D,dims=2)))                 # diagonal entries
    D,x
end
N = 128
ChebD2,x = cheb(N)
xx = x
x = x[2:N]
w = .53*x + .47*sin.(-1.5*pi*x) - x # use w = u-x to make BCs homogeneous
u = [1;w+x;-1]

ϵ=0.01
D2=ϵ*(ChebD2^2)[2:N, 2:N]
function allen_cahn(du,u,x,t)
    @. du = (u + x) - (u + x)^3
end

prob = SplitODEProblem(DiffEqArrayOperator(D2), allen_cahn, w, (0.0,5.0), x)
sol  = solve(prob, RadauIIA5(autodiff=false); reltol=1e-14,abstol=1e-14)
test_sol = TestSolution(sol)

tslices=[0.0 1.0 2.0 3.0 5.0]
ys=hcat(([1;x.+sol(t);-1] for t in tslices)...)
labels=["t=$t" for t in tslices]
plot(xx,ys,label=labels)

abstols = 0.1 .^ (5:8)
reltols = 0.1 .^ (1:4)
multipliers =  0.5 .^ (0:3)
setups = [Dict(:alg => IMEXEuler(), :dts => 1e-3 * multipliers),
          Dict(:alg => CNAB2(), :dts => 5e-3 * multipliers),
          Dict(:alg => CNLF2(), :dts => 5e-4 * multipliers),
          Dict(:alg => SBDF2(), :dts => 1e-3 * multipliers)]
labels = ["IMEXEuler" "CNAB2" "CNLF2" "SBDF2"]
@time wp1 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true,names=labels,
                            numruns=5,seconds=5,
                            save_everystop=false,appxsol=test_sol,maxiters=Int(1e5));

plot(wp1,label=labels,markershape=:auto,title="IMEX methods, dense linsolve, low order")

abstols = 0.1 .^ (5:8) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (1:4)
multipliers = 0.5 .^ (0:3)
setups = [Dict(:alg => IMEXEuler(linsolve=LinSolveGMRES()), :dts => 1e-3 * multipliers),
          Dict(:alg => CNAB2(linsolve=LinSolveGMRES()), :dts => 5e-3 * multipliers),
          Dict(:alg => CNLF2(linsolve=LinSolveGMRES()), :dts => 5e-4 * multipliers),
          Dict(:alg => SBDF2(linsolve=LinSolveGMRES()), :dts => 1e-3 * multipliers)]
labels = ["IMEXEuler" "CNAB2" "CNLF2" "SBDF2"]
@time wp1 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp1, label=labels, markershape=:auto, title="IMEX methods, Krylov linsolve, low order")

abstols = 0.1 .^ (5:8) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (1:4)
multipliers = 0.5 .^ (0:3)
setups = [Dict(:alg => NorsettEuler(), :dts => 1e-3 * multipliers),
          Dict(:alg => NorsettEuler(krylov=true, m=5), :dts => 1e-3 * multipliers),
          Dict(:alg => NorsettEuler(krylov=true, m=20), :dts => 1e-3 * multipliers),
          Dict(:alg => ETDRK2(), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK2(krylov=true, m=5), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK2(krylov=true, m=20), :dts => 1e-2 * multipliers)]
labels = hcat("NorsettEuler (caching)", "NorsettEuler (m=5)", "NorsettEuler (m=20)",
                            "ETDRK2 (caching)", "ETDRK2 (m=5)", "ETDRK2 (m=20)")
@time wp2 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp2, label=labels, markershape=:auto, title="ExpRK methods, low order")

abstols = 0.1 .^ (5:8) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (1:4)
multipliers = 0.5 .^ (0:3)
setups = [Dict(:alg => CNAB2(), :dts => 5e-3 * multipliers),
          Dict(:alg => CNAB2(linsolve=LinSolveGMRES()), :dts => 5e-3 * multipliers),
          Dict(:alg => ETDRK2(krylov=true, m=5), :dts => 1e-2 * multipliers)]
labels = ["CNAB2 (dense linsolve)" "CNAB2 (Krylov linsolve)" "ETDRK2 (m=5)"]
@time wp3 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp3, label=labels, markershape=:auto, title="Between family, low orders")

abstols = 0.1 .^ (7:13)
reltols = 0.1 .^ (4:10)
setups = [Dict(:alg => KenCarp3()),
          Dict(:alg => KenCarp4()),
          Dict(:alg => KenCarp5()),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=3, linear_solver=:Band, jac_upper=1, jac_lower=1)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=4, linear_solver=:Band, jac_upper=1, jac_lower=1)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:Band, jac_upper=1, jac_lower=1))]
labels = hcat("KenCarp3", "KenCarp4", "KenCarp5",
              "ARKODE3", "ARKODE4", "ARKODE5")
@time wp4 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp4, label=labels, markershape=:auto, title="IMEX methods, dense/band linsolve, medium order")

abstols = 0.1 .^ (7:13)
reltols = 0.1 .^ (4:10)
setups = [Dict(:alg => KenCarp3(linsolve=LinSolveGMRES())),
          Dict(:alg => KenCarp4(linsolve=LinSolveGMRES())),
          Dict(:alg => KenCarp5(linsolve=LinSolveGMRES())),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=3, linear_solver=:GMRES)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=4, linear_solver=:GMRES)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:GMRES))]
labels = ["KenCarp3" "KenCarp4" "KenCarp5" "ARKODE3" "ARKODE4" "ARKODE5"]
@time wp4 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp4, label=labels, markershape=:auto, title="IMEX methods, Krylov linsolve, medium order")

abstols = 0.1 .^ (7:11) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (4:8)
multipliers = 0.5 .^ (0:4)
setups = [Dict(:alg => ETDRK3(), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK3(krylov=true, m=5), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK4(), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK4(krylov=true, m=5), :dts => 1e-2 * multipliers),
          Dict(:alg => HochOst4(), :dts => 1e-2 * multipliers),
          Dict(:alg => HochOst4(krylov=true, m=5), :dts => 1e-2 * multipliers)]
labels = hcat("ETDRK3 (caching)", "ETDRK3 (m=5)", "ETDRK4 (caching)",
              "ETDRK4 (m=5)", "HochOst4 (caching)", "HochOst4 (m=5)")
@time wp5 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp5, label=labels, markershape=:auto, title="ExpRK methods, medium order")

abstols = 0.1 .^ (7:11)
reltols = 0.1 .^ (4:8)
multipliers = 0.5 .^ (0:4)
setups = [Dict(:alg => KenCarp5()),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:Dense)),
          Dict(:alg => KenCarp5(linsolve=LinSolveGMRES())),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:GMRES)),
          Dict(:alg => ETDRK3(krylov=true, m=5), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK4(krylov=true, m=5), :dts => 1e-2 * multipliers)]
labels = hcat("KenCarp5 (dense linsolve)", "ARKODE (dense linsolve)", "KenCarp5 (Krylov linsolve)",
              "ARKODE (Krylov linsolve)", "ETDRK3 (m=5)", "ETDRK4 (m=5)")
@time wp6 = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp6, label=labels, markershape=:auto, title="Between family, medium order")

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


@test Benitor.fibonacci(0) == 0
@test Benitor.fibonacci(1) == 1
@test Benitor.fibonacci(2) == 1
@test Benitor.fibonacci(4) == 3


"""
    unwrap(A[, wrap_around, seed])

Unwraps the values in A modulo 2π, where A is a 1-, 2-, or 3- dimensional array.

# Arguments
- `A::AbstractArray{T, N}`: Array to unwrap
- `wrap_around::Tuple{Bool}`:  When an element of this tuple is `true`, the
    unwrapping process will consider the edges along the corresponding axis
    of the image to be connected
- `seed::Int`: Unwrapping of 2D or 3D images uses a random initialization. This
    sets the seed of the RNG.
"""
function unwrap end

"""
    unwrap!(A[, wrap_around, seed])

In-place version of unwrap.
"""
function unwrap! end

function unwrap(A::AbstractArray{T, N},
                wrap_around::NTuple{N, Bool}=tuple(zeros(Bool, N)...),
                seed::Int=-1) where {T,N}
    A_copy = copy(A)
    unwrap!(A_copy, wrap_around, seed)
    return A_copy
end


# # Baroclinic adjustment
#
# In this example, we simulate the evolution and equilibration of a baroclinically
# unstable front.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add Oceananigans, CairoMakie"
# ```

using Oceananigans
using Oceananigans.Units

# ## Grid
#
# We use a three-dimensional channel that is periodic in the `x` direction:

Lx = 1000kilometers # east-west extent [m]
Ly = 1000kilometers # north-south extent [m]
Lz = 1kilometers    # depth [m]

Nx = 64
Ny = 64
Nz = 40

grid = RectilinearGrid(size = (Nx, Ny, Nz),
                       x = (0, Lx),
                       y = (-Ly/2, Ly/2),
                       z = (-Lz, 0),
                       topology = (Periodic, Bounded, Bounded))

# ## Turbulence closures

# We prescribe the values of vertical viscocity and diffusivity according to the ratio
# of the vertical and lateral grid spacing.

Δx, Δz = Lx/Nx, Lz/Nz
𝒜 = Δz/Δx # Grid cell aspect ratio.

κh = 0.1    # [m² s⁻¹] horizontal diffusivity
νh = 0.1    # [m² s⁻¹] horizontal viscosity
κz = 𝒜 * κh # [m² s⁻¹] vertical diffusivity
νz = 𝒜 * νh # [m² s⁻¹] vertical viscosity

horizontal_diffusive_closure = HorizontalScalarDiffusivity(ν = νh, κ = κh)

vertical_diffusive_closure = VerticalScalarDiffusivity(VerticallyImplicitTimeDiscretization();
                                                       ν = νz, κ = κz)
nothing #hide

# ## Model

# We built a `HydrostaticFreeSurfaceModel` with an `ImplicitFreeSurface` solver.
# Regarding Coriolis, we use a beta-plane centered at 45° South.

model = HydrostaticFreeSurfaceModel(; grid,
                                    coriolis = BetaPlane(latitude = -45),
                                    buoyancy = BuoyancyTracer(),
                                    tracers = :b,
                                    closure = (vertical_diffusive_closure, horizontal_diffusive_closure),
                                    momentum_advection = WENO5(),
                                    tracer_advection = WENO5(),
                                    free_surface = ImplicitFreeSurface())

# We want to initialize our model with a baroclinically unstable front plus some small-amplitude
# noise.

"""
    ramp(y, Δy)

Linear ramp from 0 to 1 between -Δy/2 and +Δy/2.

For example:
```
            y < -Δy/2 => ramp = 0
    -Δy/2 < y < -Δy/2 => ramp = y / Δy
            y >  Δy/2 => ramp = 1
```
"""
ramp(y, Δy) = min(max(0, y/Δy + 1/2), 1)
nothing #hide

# We then use `ramp(y, Δy)` to construct an initial buoyancy configuration of a baroclinically
# unstable front. The front has a buoyancy jump `Δb` over a latitudinal width `Δy`.

N² = 4e-6 # [s⁻²] buoyancy frequency / stratification
M² = 8e-8 # [s⁻²] horizontal buoyancy gradient

Δy = 50kilometers # width of the region of the front
Δb = Δy * M²      # buoyancy jump associated with the front
ϵb = 1e-2 * Δb    # noise amplitude

bᵢ(x, y, z) = N² * z + Δb * ramp(y, Δy) + ϵb * randn()

set!(model, b=bᵢ)

# Let's visualize the initial buoyancy distribution.

using CairoMakie

x, y, z = 1e-3 .* nodes((Center, Center, Center), grid) # convert m -> km

b = model.tracers.b

fig, ax, hm = heatmap(y, z, interior(b)[1, :, :],
                      colormap=:deep,
                      axis = (xlabel = "y [km]",
                              ylabel = "z [km]",
                              title = "b(x=0, y, z, t=0)",
                              titlesize = 24))

Colorbar(fig[1, 2], hm, label = "[m s⁻²]")

save("initial_buoyancy.svg", fig); nothing # hide

# ![](initial_buoyancy.svg)

# Now let's built a `Simulation`.

Δt₀ = 5minutes
stop_time = 40days

simulation = Simulation(model, Δt=Δt₀, stop_time=stop_time)

# We add a `TimeStepWizard` callback to adapt the siulation's time-step,

wizard = TimeStepWizard(cfl=0.2, max_change=1.1, max_Δt=20minutes)

simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(20))

# Also, we add a callback to print a message about how the simulation is going,

using Printf

wall_clock = [time_ns()]

function print_progress(sim)
    @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): (%6.3e, %6.3e, %6.3e) m/s, next Δt: %s\n",
            100 * (sim.model.clock.time / sim.stop_time),
            sim.model.clock.iteration,
            prettytime(sim.model.clock.time),
            prettytime(1e-9 * (time_ns() - wall_clock[1])),
            maximum(abs, sim.model.velocities.u),
            maximum(abs, sim.model.velocities.v),
            maximum(abs, sim.model.velocities.w),
            prettytime(sim.Δt))

    wall_clock[1] = time_ns()

    return nothing
end

simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100))

# ## Diagnostics/Output

# Add some diagnostics. Here, we save the buoyancy, ``b``, at the edges of our domain as well as
# the zonal (``x``) averages of buoyancy and zonal velocity ``u``.

u, v, w = model.velocities

B = Field(Average(b, dims=1))
U = Field(Average(u, dims=1))

filename = "baroclinic_adjustment"
save_fields_interval = 0.5day

slicers = (west = (1, :, :),
           east = (grid.Nx, :, :),
           south = (:, 1, :),
           north = (:, grid.Ny, :),
           bottom = (:, :, 1),
           top = (:, :, grid.Nz))

for side in keys(slicers)
    indices = slicers[side]

    simulation.output_writers[side] = JLD2OutputWriter(model, (; b);
                                                       filename = filename * "_$(side)_slice",
                                                       schedule = TimeInterval(save_fields_interval),
                                                       indices)
end

simulation.output_writers[:zonal] = JLD2OutputWriter(model, (b=B,);
                                                     schedule = TimeInterval(save_fields_interval),
                                                     filename = filename * "_zonal_average")

# Now let's run!

@info "Running the simulation..."

run!(simulation)

@info "Simulation completed in " * prettytime(simulation.run_wall_time)

# ## Visualization

# Now we are ready to visualize our resutls! We use `CairoMakie` in this example.
# On a system with OpenGL `using GLMakie` is more convenient as figures will be
# displayed on the screen.

using CairoMakie

# We load the saved buoyancy output on the top, bottom, and east surface as `FieldTimeSeries`es.

filename = "baroclinic_adjustment"

sides = keys(slicers)

slice_filenames = NamedTuple(side => filename * "_$(side)_slice.jld2" for side in sides)

b_timeserieses = (east   = FieldTimeSeries(slice_filenames.east, "b"),
                  north  = FieldTimeSeries(slice_filenames.north, "b"),
                  bottom = FieldTimeSeries(slice_filenames.bottom, "b"),
                  top    = FieldTimeSeries(slice_filenames.top, "b"))

avg_b_timeseries = FieldTimeSeries(filename * "_zonal_average.jld2", "b")

nothing #hide

# We build the coordinates. We rescale horizontal coordinates so that they correspond to kilometers.

x, y, z = nodes(b_timeserieses.east)

x = x .* 1e-3 # convert m -> km
y = y .* 1e-3 # convert m -> km

x_xz = repeat(x, 1, Nz)
y_xz_north = y[end] * ones(Nx, Nz)
z_xz = repeat(reshape(z, 1, Nz), Nx, 1)

x_yz_east = x[end] * ones(Ny, Nz)
y_yz = repeat(y, 1, Nz)
z_yz = repeat(reshape(z, 1, Nz), grid.Ny, 1)

x_xy = x
y_xy = y
z_xy_top = z[end] * ones(grid.Nx, grid.Ny)
z_xy_bottom = z[1] * ones(grid.Nx, grid.Ny)
nothing #hide

# Then we create a 3D axis. We use `zonal_slice_displacement` to control where the plot of the instantaneous
# zonal average flow is located.

fig = Figure(resolution = (900, 520))

zonal_slice_displacement = 1.2

ax = Axis3(fig[2, 1], aspect=(1, 1, 1/5),
           xlabel="x (km)", ylabel="y (km)", zlabel="z (m)",
           limits = ((x[1], zonal_slice_displacement * x[end]), (y[1], y[end]), (z[1], z[end])),
           elevation = 0.45, azimuth = 6.8,
           xspinesvisible = false, zgridvisible=false,
           protrusions=40,
           perspectiveness=0.7)

nothing #hide

# We use Makie's `Observable` to animate the data. To dive into how `Observable`s work we
# refer to [Makie.jl's Documentation](https://makie.juliaplots.org/stable/documentation/nodes/index.html).

n = Observable(1)

# Now let's make a 3D plot of the buoyancy and in front of it we'll use the zonally-averaged output
# to plot the instantaneous zonal-average of the buoyancy.

b_slices = (east   = @lift(interior(b_timeserieses.east[$n], 1, :, :)),
            north  = @lift(interior(b_timeserieses.north[$n], :, 1, :)),
            bottom = @lift(interior(b_timeserieses.bottom[$n], :, :, 1)),
            top    = @lift(interior(b_timeserieses.top[$n], :, :, 1)))

avg_b = @lift interior(avg_b_timeseries[$n], 1, :, :)

clims = @lift 1.1 .* extrema(b_timeserieses.top[$n][:])

kwargs = (colorrange = clims, colormap = :deep)

surface!(ax, x_yz_east, y_yz, z_yz;    color = b_slices.east, kwargs...)
surface!(ax, x_xz, y_xz_north, z_xz;   color = b_slices.north, kwargs...)
surface!(ax, x_xy, y_xy, z_xy_bottom ; color = b_slices.bottom, kwargs...)
surface!(ax, x_xy, y_xy, z_xy_top;     color = b_slices.top, kwargs...)

sf = surface!(ax, zonal_slice_displacement .* x_yz_east, y_yz, z_yz; color = avg_b, kwargs...)

contour!(ax, y, z, avg_b; transformation = (:yz, zonal_slice_displacement * x[end]),
         levels = 15, linewidth = 2, color = :black)

Colorbar(fig[2, 2], sf, label = "m s⁻²", height = 200, tellheight=false)

# Finally, we add a figure title with the time of the snapshot and then record a movie.

times = avg_b_timeseries.times

title = @lift "Buoyancy at t = " * string(round(times[$n] / day, digits=1)) * " days"

fig[1, 1:2] = Label(fig, title; textsize = 24, tellwidth = false, padding = (0, 0, -120, 0))

frames = 1:length(times)

record(fig, filename * ".mp4", frames, framerate=8) do i
    msg = string("Plotting frame ", i, " of ", frames[end])
    print(msg * " \r")
    n[] = i
end
nothing #hide

# ![](baroclinic_adjustment.mp4)


module Plasmons

export dielectric, polarizability, dispersion
export julia_main

using CUDA
using Adapt
using LinearAlgebra
using ArgParse
using HDF5
using Dates
using Logging
using LoggingExtras

include("utilities.jl")

"""
    fermidirac(E; mu, kT) -> f

Return Fermi-Dirac distribution ``f`` at energy `E`, chemical potential `mu`, and temperature
`kT`. Note that `kT` is assumed to be temperature multiplied by the Boltzmann constant, i.e.
physical dimension of `kT` is the same as `E` (e.g. electron-volts).
"""
@inline fermidirac(E; mu, kT) = 1 / (1 + exp((E - mu) / kT))

function _build_g(ħω::Complex{ℝ}, E::AbstractVector{ℝ}; mu::ℝ, kT::ℝ) where {ℝ <: Real}
    f = map(x -> fermidirac(x, mu = mu, kT = kT), E)
    n = length(E)
    G = similar(E, complex(ℝ), n, n)
    @inbounds for j in 1:n
        @. G[:, j] = (f - f[j]) / (E - E[j] - ħω)
    end
    return G
end
function _postprocess_g(
    ::Type{T},
    G::AbstractMatrix,
    blocks::Bool,
) where {T <: Union{Real, Complex}}
    split_if(x) = T <: Real ? (real(x), imag(x)) : x
    compress_if(x) = blocks ? ThreeBlockMatrix(x) : x
    return _map_tuple(compress_if, split_if(G))
end
"""
    _map_tuple(f::Function, x)

Apply `f` to every element of `x`, but only if `x` is a `Tuple`. If not, just apply `f` to
the whole `x`.
"""
_map_tuple(f::Function, x::Tuple) = f.(x)
_map_tuple(f::Function, x::Any) = f(x)

@doc raw"""
    g_matrix(::Type{ℝ}, ħω::ℂ, E; mu::ℝ, kT::ℝ, blocks::Bool) -> (Gᵣ, Gᵢ)
    g_matrix(::Type{ℂ}, ħω::ℂ, E; mu::ℝ, kT::ℝ, blocks::Bool) -> G

Compute matrix ``G``
```math
    G_{ij} = \frac{f(E_i) - f(E_j)}{E_i - E_j - \hbar\omega}
```
where ``f`` is [Fermi-Dirac
distribution](https://en.wikipedia.org/wiki/Fermi%E2%80%93Dirac_statistics) at chemical
potential `mu` (``\mu``) and temperature `kT` (``k_B T``). `E` is a vector of eigenenergies.
`ħω` is a complex frequency including Landau damping (i.e.  ``\hbar\omega + i\eta``). All
arguments are assumed to be in energy units.

Matrix ``G`` is always complex, but the first argument specifies which type should be used
to represent it. For example, if you call `g_matrix(Float32, ...)` then a tuple of two
values will be returned -- representing real and imaginary parts of ``G``. If, however, you
call `g_matrix(ComplexF64, ...)` then ``G`` will not be split into real and imaginary parts
and a single object is returned.

Sometimes one can further exploit the structure of ``G``. For ``E \ll \mu`` or ``E \gg \mu``
Fermi-Dirac distribution is just a constant and ``G`` goes to 0 for all ``\omega``. We can
use this fact to construct a block-sparse version of ``G`` to speed up matrix-matrix
products containing it. To return ``G`` in block-sparse form, set `blocks = true`.

This function works with both CPU and GPU arrays. If vector `E` is stored on the GPU, then
the returned matrix (or all submatrices for block-sparse version) will also reside on GPU.
"""
g_matrix(::Type{T}, ħω::Complex{ℝ}, E::Vector{ℝ}; mu::ℝ, kT::ℝ, blocks::Bool) where {T, ℝ <: Real} =
    _postprocess_g(T, _build_g(ħω, E; mu = mu, kT = kT), blocks)
g_matrix(
    ::Type{T},
    ħω::Complex{ℝ},
    E::CuVector{ℝ};
    mu::ℝ,
    kT::ℝ,
    blocks::Bool,
) where {T, ℝ <: Real} = _map_tuple(
    x -> adapt(CuArray, x),
    _postprocess_g(T, _build_g(ħω, Vector(E); mu = mu, kT = kT), blocks),
)

@doc raw"""
    polarizability(ħω::ℂ, E, ψ; mu::ℝ, kT::ℝ, method = :batched, blocks = true) -> χ

Compute polarizability matrix ``\chi`` using method `method` (either `:simple` or
`:batched`). `ħω` specifies complex frequency including Landau damping (i.e. ``\hbar\omega +
i\eta``). `mu` and `kT` specify chemical potential and temperature respectively. `E` and `ψ`
are eigenvalues and eigenvectors of the system Hamiltonian. They can be either CPU or GPU
arrays.
"""
function polarizability(
    ħω::Complex,
    E::AbstractVector{<:Real},
    ψ::AbstractMatrix{<:Union{Real, Complex}};
    mu::Real,
    kT::Real,
    method::Symbol = :batched,
    blocks::Bool = true,
)
    if size(ψ, 1) != size(ψ, 2)
        throw(DimensionMismatch("'ψ' has wrong shape: $(size(ψ)); expected a square matrix"))
    end
    if size(E, 1) != size(ψ, 1)
        throw(DimensionMismatch(
            "dimensions of 'E' and 'ψ' do not match: $(size(E)) & $(size(ψ)); " *
            "expected 'ψ' to be of the same dimension as 'E'",
        ))
    end
    if method == :simple
        polarizability_simple(ħω, E, ψ; mu = mu, kT = kT)
    elseif method == :batched
        polarizability_batched(ħω, E, ψ; mu = mu, kT = kT, blocks = blocks)
    else
        throw(ArgumentError("invalid 'method': $method; expected either :simple or :batched"))
    end
end

struct Workspace{M <: AbstractMatrix}
    A::M
    temp::M
end

@doc raw"""
    polarizability_batched(ħω::ℂ, E, ψ; mu::ℝ, kT::ℝ, blocks::Bool) -> χ

Compute polarizability matrix ``\chi`` using "batched" method.
"""
function polarizability_batched(
    ħω::Complex,
    E::AbstractVector{ℝ},
    ψ::AbstractMatrix{ℂ};
    mu::ℝ,
    kT::ℝ,
    blocks::Bool,
) where {ℝ <: Real, ℂ <: Union{ℝ, Complex{ℝ}}}
    G = g_matrix(ℂ, ħω, E; mu = mu, kT = kT, blocks = blocks)
    ℂ <: Real ? polarizability_batched(G[1], G[2], ψ) : polarizability_batched(G, ψ)
end
function polarizability_batched(
    Gᵣ::Union{AbstractMatrix{ℝ}, ThreeBlockMatrix{ℝ}},
    Gᵢ::Union{AbstractMatrix{ℝ}, ThreeBlockMatrix{ℝ}},
    ψ::AbstractMatrix{ℝ},
) where {ℝ <: Real}
    N = size(ψ, 1)
    χ = similar(ψ, complex(ℝ))
    fill!(χ, zero(eltype(χ)))
    ws = Workspace(similar(ψ), similar(ψ))
    for b in 1:N
        a = b:N
        _batched_mat_el!(view(χ, a, b), ws, a, b, Gᵣ, Gᵢ, ψ)
    end
    χ .= 2 .* (χ .+ transpose(χ))
    χ[diagind(χ)] ./= 2
    χ
end
function polarizability_batched(
    G::Union{AbstractMatrix{ℂ}, ThreeBlockMatrix{ℂ}},
    ψ::AbstractMatrix{ℂ},
) where {ℂ <: Complex}
    N = size(ψ, 1)
    χ = similar(ψ)
    fill!(χ, zero(eltype(χ)))
    ws = Workspace(similar(ψ), similar(ψ))
    for b in 1:N
        _batched_mat_el!(view(χ, :, b), ws, b, G, ψ)
    end
    χ .*= 2
    χ
end
function _batched_mat_el!(
    out::AbstractVector{Complex{ℝ}},
    ws::Workspace,
    as::UnitRange{Int},
    b::Int,
    Gᵣ::Union{AbstractMatrix{ℝ}, ThreeBlockMatrix{ℝ}},
    Gᵢ::Union{AbstractMatrix{ℝ}, ThreeBlockMatrix{ℝ}},
    ψ::AbstractMatrix{ℝ},
) where {ℝ <: Real}
    batch_size = length(as)
    A = view(ws.A, 1:batch_size, :)
    temp = view(ws.temp, 1:batch_size, :)
    A .= view(ψ, as, :) .* transpose(view(ψ, b, :))
    _mul_with_strides!(temp, A, Gᵣ, one(ℝ), zero(ℝ))
    outᵣ = view(reinterpret(ℝ, out), 1:2:(2 * batch_size))
    dot_batched!(outᵣ, A, temp)
    _mul_with_strides!(temp, A, Gᵢ, one(ℝ), zero(ℝ))
    outᵢ = view(reinterpret(ℝ, out), 2:2:(2 * batch_size))
    dot_batched!(outᵢ, A, temp)
    out
end
function _batched_mat_el!(
    out::AbstractVector{ℂ},
    ws::Workspace,
    b::Int,
    G::Union{AbstractMatrix{ℂ}, ThreeBlockMatrix{ℂ}},
    ψ::AbstractMatrix{ℂ},
) where {ℂ <: Complex}
    A = ws.A
    temp = ws.temp
    A .= ψ .* conj.(transpose(view(ψ, b, :)))
    _mul_with_strides!(temp, A, G, one(ℂ), zero(ℂ))
    dot_batched!(out, A, temp)
    out
end

@doc raw"""
    polarizability_simple(ħω, E, ψ; mu, kT) -> χ

Calculates the polarizability matrix ``\chi`` by direct evaluation of eq. (3). This approach
is single-threaded and doesn't utilize the CPU well. It is provided here for testing
purposes only.
"""
function polarizability_simple(
    ħω::Complex{ℝ},
    E::AbstractVector{ℝ},
    ψ::AbstractMatrix{<:Union{ℝ, Complex{ℝ}}};
    mu::ℝ,
    kT::ℝ,
) where {ℝ <: Real}
    G = _build_g(ħω, E; mu = mu, kT = kT)
    χ = similar(G)
    @inbounds for a in 1:size(χ, 2)
        for b in 1:size(χ, 1)
            acc = zero(eltype(χ))
            for i in 1:length(E)
                for j in 1:length(E)
                    acc += G[i, j] * conj(ψ[a, j]) * ψ[a, i] * conj(ψ[b, i]) * ψ[b, j]
                end
            end
            χ[a, b] = 2 * acc
        end
    end
    return χ
end

@doc raw"""
    dielectric(χ, V) -> ε

Compute dielectric function ``\varepsilon`` given polarizability matrix ``\chi`` and Coulomb
interaction potential ``V``. Dielectric function is simply ``\varepsilon = 1 - \chi V``.
"""
function dielectric(χ::AbstractMatrix{Complex{ℝ}}, V::AbstractMatrix{ℝ}) where {ℝ <: Real}
    ℂ = complex(ℝ)
    if size(χ, 1) != size(χ, 2) || size(χ) != size(V)
        throw(DimensionMismatch(
            "dimensions of χ and V do not match: $(size(χ)) != $(size(V)); " *
            "expected two square matrices of the same size",
        ))
    end
    ε = similar(χ)
    fill!(ε, zero(ℂ))
    @inbounds ε[diagind(ε)] .= one(ℂ)
    mul!(ε, V, χ, -one(ℂ), one(ℂ))
    return ε
end

function _momentum_eigenvectors(
    qs::AbstractVector{NTuple{3, ℝ}},
    rs::AbstractVector{NTuple{3, ℝ}}
) where {ℝ <: Real}
    out = similar(rs, complex(ℝ), length(rs), length(qs))
    @inbounds for j in 1:size(out, 2)
        (qˣ, qʸ, qᶻ) = qs[j]
        @simd for i in 1:size(out, 1)
            (xᵢ, yᵢ, zᵢ) = rs[i]
            out[i, j] = (1 / sqrt(length(rs))) * exp(1im * (qˣ * xᵢ + qʸ * yᵢ + qᶻ * zᵢ))
        end
    end
    out
end
function dispersion_function(
    qs::AbstractVector{NTuple{3, ℝ}},
    rs::AbstractVector{NTuple{3, ℝ}},
) where {ℝ <: Real}
    return let r_q = _momentum_eigenvectors(qs, rs),
        temp = similar(r_q, length(rs), size(r_q, 2))

        A -> begin
            mul!(temp, A, r_q)
            out = similar(r_q, length(qs))
            @inbounds for i in 1:length(out)
                out[i] = dot(view(r_q, :, i), view(temp, :, i))
            end
            out
        end
    end
end
"""
    dispersion(A(ω, r, r'), q, x, y, z; n = 100) -> A(ω, q)

Calculate diagonal elements of the Fourier transform of matrix A.
"""
# function dispersion(As, qs, x, y, z)
#     fn! = _dispersion_function(qs, x, y, z)
#     matrix = Any[]
#     for (i, A) in enumerate(As)
#         out = similar(x, complex(eltype(x)), length(qs))
#         fn!(out, A)
#         push!(matrix, out)
#     end
#     hcat(matrix...)
# end
dispersion(A::AbstractMatrix, qs::AbstractVector, rs::AbstractVector) = dispersion_function(qs, rs)(A)

function main(
    E::AbstractVector{ℝ},
    ψ::AbstractMatrix{ℂ};
    kT::Real,
    μ::Real,
    η::Real,
    ωs::Vector{<:Real},
    out::Union{HDF5.File, HDF5.Group},
    V::Union{AbstractMatrix{ℝ}, Nothing} = nothing,
) where {ℝ <: Real, ℂ <: Union{ℝ, Complex{ℝ}}}
    if kT <= 0
        throw(ArgumentError("invalid 'kT': $kT; expected a positive real number"))
    end
    if η <= 0
        throw(ArgumentError("invalid 'η': $η; expected a positive real number"))
    end
    HDF5.attributes(out)["kT"] = kT
    HDF5.attributes(out)["μ"] = μ
    HDF5.attributes(out)["η"] = η
    group_χ = create_group(out, "χ")
    group_ε::Union{HDF5.Group, Nothing} = isnothing(V) ? nothing : create_group(out, "ε")
    @info "Polarizability matrices χ(ω) will be saved to group 'χ'"
    if !isnothing(V)
        @info "Dielectric functions ε(ω) will be saved to group 'ε'"
    else
        @warn "Coulomb interaction matrix not provided: dielectric function ε(ω) will not be computed"
    end

    for (i, ω) in enumerate(map(x -> x + 1im * η, ωs))
        @info "Calculating χ(ω = $ω) ..."
        name = string(i, pad = 4)
        t₀ = time_ns()
        χ = Array(polarizability(
            convert(complex(ℝ), ω),
            E,
            ψ;
            mu = convert(ℝ, μ),
            kT = convert(ℝ, kT),
        ))
        t₁ = time_ns()
        group_χ[name] = χ
        HDF5.attributes(group_χ[name])["ħω"] = ω
        HDF5.attributes(group_χ[name])["time"] = (t₁ - t₀) / 1e9
        flush(group_χ)
        if !isnothing(V)
            t₀ = time_ns()
            ε = Array(dielectric(χ, V))
            t₁ = time_ns()
            group_ε[name] = ε
            HDF5.attributes(group_ε[name])["ħω"] = ω
            HDF5.attributes(group_ε[name])["time"] = (t₁ - t₀) / 1e9
            flush(group_ε)
        end
    end
end
function main!(H::AbstractMatrix{ℂ}; kwargs...) where {ℂ}
    @warn "Eigenvalues or eigenvectors not provided: diagonalizing the Hamiltonian..."
    E, ψ = _eigen!(Hermitian(H))
    @info "Diagonalization completed"
    main(E, ψ; kwargs...)
end

ArgParse.parse_item(::Type{Vector{Float64}}, x::AbstractString) =
    map(s -> parse(Float64, s), split(x, ","))

function _parse_commandline()
    s = ArgParseSettings()
    #! format: off
    @add_arg_table! s begin
        "--kT"
            help = "Temperature kʙ·T in eletron-volts"
            arg_type = Float64
            required = true
        "--mu"
            help = "Chemical potential μ in eletron-volts"
            arg_type = Float64
            required = true
        "--damping"
            help = "Landau damping η in eletron-volts"
            arg_type = Float64
            required = true
        "--frequency"
            help = "Comma-separated list of ħω values in eletron-volts for which to compute χ and ε"
            arg_type = Vector{Float64}
            required = true
        "--hamiltonian"
            help = "Path to Hamiltonian matrix in the input HDF5 file"
            arg_type = String
        "--coulomb"
            help = "Path to unscreened Coulomb interaction matrix in the input HDF5 file"
            arg_type = String
        "--eigenvalues"
            help = "Path to eigenvalues vector in the input HDF5 file"
            arg_type = String
        "--eigenvectors"
            help = "Path to eigenvectors matrix in the input HDF5 file"
            arg_type = String
        "--cuda"
            help = "Index of CUDA device to use"
            arg_type = Int
        "input_file"
            help = "Input HDF5 file"
            required = true
        "output_file"
            help = "Output HDF5 file"
            required = true
    end
    #! format: on
    return parse_args(s, as_symbols = true)
end

function tryread(io::HDF5.File, path::Union{<:AbstractString, Nothing} = nothing)
    isnothing(path) && return nothing
    if !haskey(io, path)
        @warn "Path '$path' not found in input HDF5 file"
        return nothing
    end
    read(io, path)
end

const date_format = "yyyy-mm-dd HH:MM:SS"

timestamp_logger(logger) =
    TransformerLogger(logger) do log
        merge(log, (; message = "$(Dates.format(now(), date_format))  $(log.message)"))
    end

function julia_main()::Cint
    CUDA.allowscalar(false)
    args = _parse_commandline()
    H, E, ψ, V = h5open(args[:input_file], "r") do io
        tryread(io, get(args, :hamiltonian, nothing)),
        tryread(io, get(args, :eigenvalues, nothing)),
        tryread(io, get(args, :eigenvectors, nothing)),
        tryread(io, get(args, :coulomb, nothing))
    end
    if (isnothing(E) || isnothing(ψ)) && isnothing(H)
        @error "if Hamiltonian 'H' is not specified, both eigenvalues 'E'" *
               " and eigenvectors 'ψ' must be provided"
        return 1
    end
    cuda_device = args[:cuda]
    if !isnothing(cuda_device)
        @info "Setting CUDA device to $cuda_device..."
        device!(cuda_device)
        if !isnothing(H)
            H = adapt(CuArray, H)
        end
        if !isnothing(E)
            E = adapt(CuArray, E)
        end
        if !isnothing(ψ)
            ψ = adapt(CuArray, ψ)
        end
        if !isnothing(V)
            V = adapt(CuArray, V)
        end
    end
    h5open(args[:output_file], "w") do io
        kwargs = (
            kT = args[:kT],
            μ = args[:mu],
            η = args[:damping],
            ωs = args[:frequency],
            out = io,
            V = V,
        )
        with_logger(ConsoleLogger(stdout, Logging.Info) |> timestamp_logger) do
            if (isnothing(E) || isnothing(ψ))
                main!(H; kwargs...)
            else
                main(E, ψ; kwargs...)
            end
        end
    end
    return 0
end

end # module


using OptimalTransport

using Distances
using ForwardDiff
using ReverseDiff
using LogExpFunctions
using PythonOT: PythonOT

using LinearAlgebra
using Random
using Test

const POT = PythonOT

Random.seed!(100)

@testset "sinkhorn_gibbs.jl" begin
    # size of source and target
    M = 250
    N = 200

    # create two random histograms
    μ = normalize!(rand(M), 1)
    ν = normalize!(rand(N), 1)

    # create random cost matrix
    C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)

    # regularization parameter
    ε = 0.01

    @testset "example" begin
        # compute optimal transport plan and optimal transport cost
        γ = sinkhorn(μ, ν, C, ε, SinkhornGibbs(); maxiter=5_000, rtol=1e-9)
        c = sinkhorn2(μ, ν, C, ε, SinkhornGibbs(); maxiter=5_000, rtol=1e-9)

        # check that plan and cost are consistent
        @test c ≈ dot(γ, C)

        # compare with default algorithm
        γ_default = sinkhorn(μ, ν, C, ε; maxiter=5_000, rtol=1e-9)
        c_default = sinkhorn2(μ, ν, C, ε; maxiter=5_000, rtol=1e-9)
        @test γ_default == γ
        @test c_default == c

        # compare with POT
        γ_pot = POT.sinkhorn(μ, ν, C, ε; numItermax=5_000, stopThr=1e-9)
        c_pot = POT.sinkhorn2(μ, ν, C, ε; numItermax=5_000, stopThr=1e-9)[1]
        @test γ_pot ≈ γ rtol = 1e-6
        @test c_pot ≈ c rtol = 1e-7

        # compute optimal transport cost with regularization term
        c_w_regularization = sinkhorn2(
            μ, ν, C, ε, SinkhornGibbs(); maxiter=5_000, regularization=true
        )
        @test c_w_regularization ≈ c + ε * sum(x -> iszero(x) ? x : x * log(x), γ)
        @test c_w_regularization ==
              sinkhorn2(μ, ν, C, ε; maxiter=5_000, regularization=true)

        # ensure that provided plan is used and correct
        c2 = sinkhorn2(similar(μ), similar(ν), C, rand(), SinkhornGibbs(); plan=γ)
        @test c2 ≈ c
        @test c2 == sinkhorn2(similar(μ), similar(ν), C, rand(); plan=γ)
        c2_w_regularization = sinkhorn2(
            similar(μ), similar(ν), C, ε, SinkhornGibbs(); plan=γ, regularization=true
        )
        @test c2_w_regularization ≈ c_w_regularization
        @test c2_w_regularization ==
              sinkhorn2(similar(μ), similar(ν), C, ε; plan=γ, regularization=true)

        # batches of histograms
        d = 10
        for (size2_μ, size2_ν) in
            (((), (d,)), ((1,), (d,)), ((d,), ()), ((d,), (1,)), ((d,), (d,)))
            # generate batches of histograms
            μ_batch = repeat(μ, 1, size2_μ...)
            ν_batch = repeat(ν, 1, size2_ν...)

            # compute optimal transport plan and check that it is consistent with the
            # plan for individual histograms
            γ_all = sinkhorn(
                μ_batch, ν_batch, C, ε, SinkhornGibbs(); maxiter=5_000, rtol=1e-9
            )
            @test size(γ_all) == (M, N, d)
            @test all(view(γ_all, :, :, i) ≈ γ for i in axes(γ_all, 3))
            @test γ_all == sinkhorn(μ_batch, ν_batch, C, ε; maxiter=5_000, rtol=1e-9)

            # compute optimal transport cost and check that it is consistent with the
            # cost for individual histograms
            c_all = sinkhorn2(
                μ_batch, ν_batch, C, ε, SinkhornGibbs(); maxiter=5_000, rtol=1e-9
            )
            @test size(c_all) == (d,)
            @test all(x ≈ c for x in c_all)
            @test c_all == sinkhorn2(μ_batch, ν_batch, C, ε; maxiter=5_000, rtol=1e-9)
        end
    end

    # different element type
    @testset "Float32" begin
        # create histograms and cost matrix with element type `Float32`
        μ32 = map(Float32, μ)
        ν32 = map(Float32, ν)
        C32 = map(Float32, C)
        ε32 = Float32(ε)

        # compute optimal transport plan and optimal transport cost
        γ = sinkhorn(μ32, ν32, C32, ε32, SinkhornGibbs(); maxiter=5_000, rtol=1e-6)
        c = sinkhorn2(μ32, ν32, C32, ε32, SinkhornGibbs(); maxiter=5_000, rtol=1e-6)
        @test eltype(γ) === Float32
        @test typeof(c) === Float32

        # check that plan and cost are consistent
        @test c ≈ dot(γ, C32)

        # compare with default algorithm
        γ_default = sinkhorn(μ32, ν32, C32, ε32; maxiter=5_000, rtol=1e-6)
        c_default = sinkhorn2(μ32, ν32, C32, ε32; maxiter=5_000, rtol=1e-6)
        @test γ_default == γ
        @test c_default == c

        # compare with POT
        γ_pot = POT.sinkhorn(μ32, ν32, C32, ε32; numItermax=5_000, stopThr=1e-6)
        c_pot = POT.sinkhorn2(μ32, ν32, C32, ε32; numItermax=5_000, stopThr=1e-6)[1]
        @test map(Float32, γ_pot) ≈ γ rtol = 1e-3
        @test Float32(c_pot) ≈ c rtol = 1e-3

        # batches of histograms
        d = 10
        for (size2_μ, size2_ν) in
            (((), (d,)), ((1,), (d,)), ((d,), ()), ((d,), (1,)), ((d,), (d,)))
            # generate batches of histograms
            μ32_batch = repeat(μ32, 1, size2_μ...)
            ν32_batch = repeat(ν32, 1, size2_ν...)

            # compute optimal transport plan and check that it is consistent with the
            # plan for individual histograms
            γ_all = sinkhorn(
                μ32_batch, ν32_batch, C32, ε32, SinkhornGibbs(); maxiter=5_000, rtol=1e-6
            )
            @test size(γ_all) == (M, N, d)
            @test all(view(γ_all, :, :, i) ≈ γ for i in axes(γ_all, 3))
            @test γ_all ==
                  sinkhorn(μ32_batch, ν32_batch, C32, ε32; maxiter=5_000, rtol=1e-6)

            # compute optimal transport cost and check that it is consistent with the
            # cost for individual histograms
            c_all = sinkhorn2(
                μ32_batch, ν32_batch, C32, ε32, SinkhornGibbs(); maxiter=5_000, rtol=1e-6
            )
            @test size(c_all) == (d,)
            @test all(x ≈ c for x in c_all)
            @test c_all ==
                  sinkhorn2(μ32_batch, ν32_batch, C32, ε32; maxiter=5_000, rtol=1e-6)
        end
    end

    # https://github.com/JuliaOptimalTransport/OptimalTransport.jl/issues/86
    @testset "AD" begin
        # compute gradients with respect to source and target marginals separately and
        # together. test against gradient computed using analytic formula of Proposition 2.3 of
        # Cuturi, Marco, and Gabriel Peyré. "A smoothed dual approach for variational Wasserstein problems." SIAM Journal on Imaging Sciences 9.1 (2016): 320-343.
        #
        ε = 0.05 # use a larger ε to avoid having to do many iterations
        # target marginal
        for Diff in [ReverseDiff, ForwardDiff]
            ∇ = Diff.gradient(log.(ν)) do xs
                sinkhorn2(μ, softmax(xs), C, ε, SinkhornGibbs(); regularization=true)
            end
            ∇default = Diff.gradient(log.(ν)) do xs
                sinkhorn2(μ, softmax(xs), C, ε; regularization=true)
            end
            @test ∇ == ∇default

            solver = OptimalTransport.build_solver(μ, ν, C, ε, SinkhornGibbs())
            OptimalTransport.solve!(solver)
            # helper function
            function dualvar_to_grad(x, ε)
                x = -ε * log.(x)
                x .-= sum(x) / size(x, 1)
                return -x
            end
            ∇_ot = dualvar_to_grad(solver.cache.v, ε)
            # chain rule because target measure parameterised by softmax
            J_softmax = ForwardDiff.jacobian(log.(ν)) do xs
                softmax(xs)
            end
            ∇analytic_target = J_softmax * ∇_ot
            # check that gradient obtained by AD matches the analytic formula
            @test ∇ ≈ ∇analytic_target rtol = 1e-6

            # source marginal
            ∇ = Diff.gradient(log.(μ)) do xs
                sinkhorn2(softmax(xs), ν, C, ε, SinkhornGibbs(); regularization=true)
            end
            ∇default = Diff.gradient(log.(μ)) do xs
                sinkhorn2(softmax(xs), ν, C, ε; regularization=true)
            end
            @test ∇ == ∇default

            # check that gradient obtained by AD matches the analytic formula
            solver = OptimalTransport.build_solver(μ, ν, C, ε, SinkhornGibbs())
            OptimalTransport.solve!(solver)
            J_softmax = ForwardDiff.jacobian(log.(μ)) do xs
                softmax(xs)
            end
            ∇_ot = dualvar_to_grad(solver.cache.u, ε)
            ∇analytic_source = J_softmax * ∇_ot
            @test ∇ ≈ ∇analytic_source rtol = 1e-6

            # both marginals
            ∇ = Diff.gradient(log.(vcat(μ, ν))) do xs
                sinkhorn2(
                    softmax(xs[1:M]),
                    softmax(xs[(M + 1):end]),
                    C,
                    ε,
                    SinkhornGibbs();
                    regularization=true,
                )
            end
            ∇default = Diff.gradient(log.(vcat(μ, ν))) do xs
                sinkhorn2(softmax(xs[1:M]), softmax(xs[(M + 1):end]), C, ε; regularization=true)
            end
            @test ∇ == ∇default
            ∇analytic = vcat(∇analytic_source, ∇analytic_target)
            @test ∇ ≈ ∇analytic rtol = 1e-6
        end
    end

    @testset "deprecations" begin
        # check `sinkhorn`
        γ = sinkhorn(μ, ν, C, ε; atol=1e-6)
        @test (@test_deprecated sinkhorn(μ, ν, C, ε; tol=1e-6)) == γ
        γ = sinkhorn(μ, ν, C, ε; check_convergence=5)
        @test (@test_deprecated sinkhorn(μ, ν, C, ε; check_marginal_step=5)) == γ

        # check `sinkhorn2`
        c = sinkhorn2(μ, ν, C, ε; atol=1e-6)
        @test (@test_deprecated sinkhorn2(μ, ν, C, ε; tol=1e-6)) == c
        c = sinkhorn2(μ, ν, C, ε; check_convergence=5)
        @test (@test_deprecated sinkhorn2(μ, ν, C, ε; check_marginal_step=5)) == c
    end
end


# BFS algorithm using Julia stdlib primitives

using SparseArrays

"""
    bfs_primitive(A::SparseMatrixCSC{Bool, U}, source::U) where {U <: Integer}

Compute BFS on the graph represented by `A` starting from `source` using Julia primitives.

# Examples
```julia-repl
julia> dim = 100;

julia> A = sparse_square_matrix(dim, Bool);

julia> source = 1;

julia> bfs_primitive(A,source);
```
"""
function bfs_primitive(
    A::SparseMatrixCSC{Bool, U},
    source::U
    ) where {U <: Integer}

    source > 0 || throw(ArgumentError("source must be positive"))

    dim = size(A, 2)
    # Input matrix must be square
    size(A, 1) == dim || throw(DimensionMismatch())

    q = BitArray(undef, dim)
    q[source] = true

    # Create v, make it dense
    v = zeros(Int64, dim)
    v[source] = 1

    aux = similar(q)
    toupdate = spzeros(Bool, dim)

    @inbounds for level in 1:dim

        aux = broadcast(!iszero, A'*q)
        copy!(q, aux)

        # toupdate[i] = true iff vertex i is
        # in the frontier and was not visited before
        @inbounds for i in eachindex(toupdate)
            toupdate[i] = iszero(v[i]) & aux[i]
        end

        v .+= toupdate .* (level+1)

        if(!reduce(max,q))
            break
        end
    end

    return v
end


#using Debugger
using StructuralCausalModels, Test

ProjDir = @__DIR__
cd(ProjDir) #do

fname = scm_path("..", "examples", "Shipley", "fig2.6", "fig2.6.dot")
#Sys.isapple() && run(`open -a GraphViz.app $(fname)`)

# Read `=>` as `~` in regression models, or `<-` in causal models.
d = OrderedDict(
  :u => [:x, :v],
  :s1 => [:u],
  :w => [:v, :y],
  :s2 => [:w]
);

dag = DAG("fig2.6", d);

@testset "d_separation" begin

  @test to_ggm(dag) == "DAG(u ~ x + v, s1 ~ u, w ~ v + y, s2 ~ w)"
  @test d_separation(dag, :x, :v; debug=false) == true
  @test d_separation(dag, :x, [:v]; c=:u, debug=false) == false
  @test d_separation(dag, [:x], :v; c=:s1, debug=false) == false
  @test d_separation(dag, [:u], [:w]; debug=false) == false
  @test d_separation(dag, :u, [:w]; c=:v, debug=false) == true
  @test d_separation(dag, [:x], :y; debug=false) == true
  @test d_separation(dag, :x, :y; c=[:u, :w], debug=false) == false
  @test d_separation(dag, :x, :y; c=[:s1, :s2], debug=false) == false
  @test d_separation(dag, :x, :y; c=[:u, :v, :w], debug=false) == true
  @test d_separation(dag, :x, [:v]; c=:u, debug=false) == false
  @test d_separation(dag, [:x], :v; c=:s1, debug=false) == false

  include("test_descendents_02.jl")
  @test d_separation(dag, :k1, :v, c=[:x1]) == true
  @test d_separation(dag, :k1, :v, c=[:w]) == true
  @test d_separation(dag, :k1, :x1, c=[:w]) == true
  @test d_separation(dag, :k1, :x2, c=[:v]) == true
  @test d_separation(dag, :k1, :x2, c=[:x1]) == true
  @test d_separation(dag, :k1, :x2, c=[:w]) == true
  @test d_separation(dag, :k1, :x3, c=[:w]) == true
  @test d_separation(dag, :k1, :y, c=[:k2, :x2, :x3]) == true
  @test d_separation(dag, :k1, :y, c=[:k2, :v, :x3]) == true
  @test d_separation(dag, :k1, :y, c=[:k2, :x1, :x3]) == true
  @test d_separation(dag, :k1, :y, c=[:k2, :w]) == true
  @test d_separation(dag, :k2, :v, c=[:x1]) == true
  @test d_separation(dag, :k2, :v, c=[:w]) == true
  @test d_separation(dag, :k2, :v, c=[:k1]) == true
  @test d_separation(dag, :k2, :w, c=[:k1]) == true
  @test d_separation(dag, :k2, :x1, c=[:w]) == true
  @test d_separation(dag, :k2, :x1, c=[:k1]) == true
  @test d_separation(dag, :k2, :x2, c=[:v]) == true
  @test d_separation(dag, :k2, :x2, c=[:x1]) == true
  @test d_separation(dag, :k2, :x2, c=[:w]) == true
  @test d_separation(dag, :k2, :x2, c=[:k1]) == true
  @test d_separation(dag, :k2, :x3, c=[:w]) == true
  @test d_separation(dag, :k2, :x3, c=[:k1]) == true
  @test d_separation(dag, :v, :w; c=:x1) == true
  @test d_separation(dag, :v, :x3) == true
  @test d_separation(dag, :v, :y; c=[:k2, :x2, :x3]) == true
  @test d_separation(dag, :v, :y; c=[:k1, :x2, :x3]) == true
  @test d_separation(dag, :v, :y; c=[:w, :x2, :x3]) == true
  @test d_separation(dag, :v, :y; c=[:x2, :x1]) == true
  @test d_separation(dag, :w, :x2; c=[:v]) == true
  @test d_separation(dag, :w, :x2; c=[:x1]) == true
  @test d_separation(dag, :w, :y; c=[:k2, :x2, :x3]) == true
  @test d_separation(dag, :w, :y; c=[:k2, :v, :x3]) == true
  @test d_separation(dag, :w, :y; c=[:k1, :x2, :x3]) == true
  @test d_separation(dag, :w, :y; c=[:k1, :v, :x3]) == true
  @test d_separation(dag, :w, :y; c=[:k1, :x1, :x3]) == true
  @test d_separation(dag, :w, :y; c=[:k2, :x2, :x3]) == true
  @test d_separation(dag, :x1, :x2; c=[:v]) == true
  @test d_separation(dag, :x1, :x3) == true
  @test d_separation(dag, :x1, :y; c=[:k2, :x2, :x3]) == true
  @test d_separation(dag, :x1, :y; c=[:k1, :x2, :x3]) == true
  @test d_separation(dag, :x1, :y; c=[:w, :x2, :x3]) == true
  @test d_separation(dag, :x1, :y; c=[:k2, :v, :x3]) == true
  @test d_separation(dag, :x1, :y; c=[:k1, :v, :x3]) == true
  @test d_separation(dag, :x1, :y; c=[:v, :w, :x3]) == true
  @test d_separation(dag, :x2, :x3) == true

end


## Julia 1.1.0

z = 1 + 2im
for T in [real, imag, conj, abs, abs2, angle, sqrt, sin, cos, sinh, cosh, exp]
    println("$(T)($(z)) -> $(T(z))")
end


#------------------------------
# functions to simulate data
#------------------------------

function generate_xs(n, p, true_u0, sol_group1, sol_group2; t_start=1.5, t_end=10, maxntps = 10, dt=0.1, σ_var=0.1, σ_ind=0.5)

    # (1) generate artifical group labels
    groups = zeros(n)
    groups[randperm(n)[1:Int(floor(n/2))]] .= 1 # for groups of the same size
    group1 = findall(x -> x==1, groups)
    group2 = findall(x -> x==0, groups)

    # (2) generate artificial time stamps
    ntps = rand(1:maxntps, n)
    tvals = [sort(rand(t_start:dt:t_end,ntps[i])) for i in 1:n]

    # (3) obtain true values as solutions of the ODEs at the initial time point and the drawn second time point
    # check for equal number of variables:
    if p%2 != 0
        error("Please select an even number of variables")
    end
    # true starting point
    z_t0_p1 = true_u0[1] # for variables 1-(p/2)
    z_t0_p2 = true_u0[2] # for variables (p/2+1)-p
    z_t0 = repeat([z_t0_p1, z_t0_p2], inner=Int(p/2))

    # now use ODE solutions to obtain true temporal development value
    # for all individuals in both variables u1 and u2
    z_later_ts = collect((i ∈ group1) ? (Array(sol_group1)[:,Int.(round.(tvals[i].*(1 ./dt)).+1)]) : (Array(sol_group2)[:,Int.(round.(tvals[i].*(1 ./dt)).+1)]) for i in 1:n)

    # (4) sample variable- specific and individual-specific errors at both time points
    # variable specific random effect (general difficulty measuring that specific variable)
    us = rand(Normal(0,σ_var),p)

    xs = []
    for i in 1:n
        # make time series structure, should have shape (p x ntps[i])
        cur_timeseries = zeros(p, ntps[i]+1)
        for j in 1:p
            cur_timeseries[j,1] = z_t0[j] + us[j] + randn() .* σ_ind
            for tp in 1:ntps[i]
                if j <= Int(p/2)
                    cur_timeseries[j,tp+1] = z_later_ts[i][1,tp] + us[j] + randn() .* σ_ind
                else
                    cur_timeseries[j,tp+1] = z_later_ts[i][2,tp] + us[j] + randn() .* σ_ind
                end
            end
        end
        push!(xs, cur_timeseries)
    end

    return xs, tvals, group1, group2
end

function generate_baseline(n, q, q_info, group1; σ_info=1, σ_noise=1)
    zs = fill(1,(n,1))
    zs[group1].= -1
    means = fill(0,n,q)
    means[:,1:q_info] .= zs

    vars=fill(σ_noise,q)
    vars[1:q_info] .= σ_info
    x_params = [cat([rand(Normal(means[i,j],vars[j])) for j in 1:q]..., dims=1) for i in 1:n]
    return x_params
end

function generate_baseline(n, q, q_info, group1, true_odeparams_group1, true_odeparams_group2; σ_info=0.1, σ_noise=0.1)
    signs = fill(1,(n,1))
    signs[group1] .= -1
    z1s = zeros(n,1)
    z1s[signs .== -1] .= true_odeparams_group1[1]
    z1s[signs .== 1] .= true_odeparams_group2[1]
    z2s = zeros(n,1)
    z2s[signs .== -1] .= true_odeparams_group1[4]
    z2s[signs .== 1] .= true_odeparams_group2[4]

    means = zeros(n,q)
    means[:,1:Int(floor(q_info/2))] .= z1s
    means[:,Int(floor(q_info/2))+1:q_info] .=z2s

    vars=fill(σ_noise,q)
    vars[1:q_info] .= σ_info
    x_params = [cat([rand(Normal(means[i,j],vars[j])) for j in 1:q]..., dims=1) for i in 1:n]
    return x_params
end

function generate_baseline_4params(n, q, q_info, group1, true_odeparams_group1, true_odeparams_group2; σ_info=0.1, σ_noise=0.1)
    signs = fill(1,n)
    signs[group1] .= -1
    zs = cat(repeat([true_odeparams_group2], n)..., dims=2)
    zs[:,findall(signs .==-1)] .= true_odeparams_group1
    zs = zs'

    means = zeros(n,q)
    means[:,1:Int(floor(q_info/4))] .= zs[:,1]
    means[:,Int(floor(q_info/4))+1:Int(floor(q_info/2))] .= zs[:,2]
    means[:,Int(floor(q_info/2))+1:Int(floor(q_info*3/4))] .= zs[:,3]
    means[:,Int(floor(q_info*3/4))+1:q_info] .= zs[:,4]

    vars=fill(σ_noise, q)
    vars[1:q_info] .= σ_info
    x_params = [[rand(Normal(means[i,j],vars[j])) for j in 1:q] for i in 1:n]
    return x_params
end


# Themelis, Stella, Patrinos, "Douglas-Rachford splitting and ADMM
# for nonconvex optimization: Accelerated and Newton-type algorithms",
# arXiv:2005.10230, 2020.
#
# https://arxiv.org/abs/2005.10230

using Base.Iterators
using ProximalAlgorithms.IterationTools
using ProximalOperators: Zero
using LinearAlgebra
using Printf

function drls_default_gamma(f, muf, Lf, alpha, lambda)
    if muf !== nothing && muf > 0
        return 1 / (alpha * muf)
    end
    if ProximalOperators.is_convex(f)
        return alpha / Lf
    else
        return alpha * (2 - lambda) / (2 * Lf)
    end
end

function drls_C(f, muf, Lf, gamma, lambda, beta)
    a = muf === nothing || muf <= 0 ? gamma * Lf : 1 / (gamma * muf)
    m = ProximalOperators.is_convex(f) ? max(a - lambda / 2, 0) : 1
    return (lambda / ((1 + a)^2) * ((2 - lambda) / 2 - a * m))
end

Base.@kwdef struct DRLSIteration{R,C<:Union{R,Complex{R}},Tx<:AbstractArray{C},Tf,Tg,TH}
    f::Tf = Zero()
    g::Tg = Zero()
    x0::Tx
    alpha::R = real(eltype(x0))(0.95)
    beta::R = real(eltype(x0))(0.5)
    lambda::R = real(eltype(x0))(1)
    muf::Maybe{R} = nothing
    Lf::Maybe{R} = nothing
    gamma::R = drls_default_gamma(f, muf, Lf, alpha, lambda)
    c::R = beta * drls_C(f, muf, Lf, gamma, lambda, beta)
    dre_sign::Int = muf === nothing || muf <= 0 ? 1 : -1
    max_backtracks::Int = 20
    H::TH = LBFGS(x0, 5)
end

Base.IteratorSize(::Type{<:DRLSIteration}) = Base.IsInfinite()

Base.@kwdef mutable struct DRLSState{R,Tx,TH}
    x::Tx
    u::Tx
    v::Tx
    w::Tx
    res::Tx
    res_prev::Tx = zero(x)
    xbar::Tx
    xbar_prev::Tx = zero(x)
    d::Tx = zero(x)
    x_d::Tx = zero(x)
    gamma::R
    f_u::R
    g_v::R
    H::TH
    tau::Maybe{R} = nothing
end

DRE(f_u::Number, g_v::Number, x, u, res, gamma) = f_u + g_v - real(dot(x - u, res)) / gamma + 1 / (2 * gamma) * norm(res)^2

DRE(state::DRLSState) = DRE(state.f_u, state.g_v, state.x, state.u, state.res, state.gamma)

function Base.iterate(iter::DRLSIteration)
    x = copy(iter.x0)
    u, f_u = prox(iter.f, x, iter.gamma)
    w = 2 * u - x
    v, g_v = prox(iter.g, w, iter.gamma)
    res = u - v
    xbar = x - iter.lambda * res
    state = DRLSState(
        x=x, u=u, v=v, w=w, res=res, xbar=xbar, gamma=iter.gamma, f_u=f_u,
        g_v=g_v, H=iter.H,
    )
    return state, state
end

function Base.iterate(iter::DRLSIteration{R}, state::DRLSState) where {R}
    DRE_curr = DRE(state)

    mul!(state.d, iter.H, -state.res)
    state.x_d .= state.x .+ state.d
    state.xbar_prev, state.xbar = state.xbar, state.xbar_prev
    state.res_prev, state.res = state.res, state.res_prev
    state.tau = R(1)
    state.x .= state.x_d

    for k = 1:iter.max_backtracks
        state.x .= state.tau .* state.x_d .+ (1 - state.tau) .* state.xbar_prev

        state.f_u = prox!(state.u, iter.f, state.x, iter.gamma)
        state.w .= 2 .* state.u .- state.x
        state.g_v = prox!(state.v, iter.g, state.w, iter.gamma)
        state.res .= state.u .- state.v

        if k == 1
            update!(iter.H, state.d, state.res - state.res_prev)
        end

        state.xbar .= state.x .- iter.lambda * state.res
        DRE_candidate = DRE(state)

        if iter.dre_sign * DRE_candidate <= iter.dre_sign * DRE_curr - iter.c / iter.gamma * norm(state.res_prev)^2
            break
        end

        state.tau = k < iter.max_backtracks - 1 ? state.tau / 2 : R(0)
    end

    return state, state
end

# Solver

struct DRLS{R, K}
    maxit::Int
    tol::R
    verbose::Bool
    freq::Int
    kwargs::K
end

function (solver::DRLS)(x0; kwargs...)
    stop(state::DRLSState) = norm(state.res, Inf) / state.gamma <= solver.tol
    disp((it, state)) = @printf(
        "%5d | %.3e | %.3e | %.3e\n",
        it,
        state.gamma / state.gamma,
        norm(state.res, Inf),
        (state.tau === nothing ? 0.0 : state.tau)
    )
    iter = DRLSIteration(; x0=x0, solver.kwargs..., kwargs...)
    iter = take(halt(iter, stop), solver.maxit)
    iter = enumerate(iter)
    if solver.verbose
        iter = tee(sample(iter, solver.freq), disp)
    end
    num_iters, state_final = loop(iter)
    return state_final.u, state_final.v, num_iters
end

DRLS(; maxit=1_000, tol=1e-8, verbose=false, freq=10, kwargs...) =
    DRLS(maxit, tol, verbose, freq, kwargs)


struct RMSPropIterable{T, F, S, R}
    w0::T
    f::F
    stepsize::S
    alpha::R
    epsilon::R
end

RMSPropIterable(w0, f; stepsize=1e-3, alpha=0.99, epsilon=1e-8) = RMSPropIterable(w0, f, to_iterator(stepsize), alpha, epsilon)

Base.IteratorSize(::Type{<:RMSPropIterable}) = Base.IsInfinite()

mutable struct RMSPropState{T, G, S}
    w::T
    f_w
    grad_f_w::G
    stepsize_iterator::S
    mean_squared_grad::T
end

function Base.iterate(iter::RMSPropIterable)
    w = copy(iter.w0)
    grad_f_w, f_w = gradient(iter.f, w)
    state = RMSPropState(w, f_w, grad_f_w, Iterators.Stateful(iter.stepsize), zero(w))
    return IterationOutput(state.w, state.f_w, state.grad_f_w), state
end

function Base.iterate(iter::RMSPropIterable, state::RMSPropState)
    state.mean_squared_grad .= (1 - iter.alpha) .* (state.grad_f_w .^ 2) .+ iter.alpha .* state.mean_squared_grad
    stepsize = popfirst!(state.stepsize_iterator)
    state.w .-= stepsize .* state.grad_f_w ./ (sqrt.(state.mean_squared_grad) .+ iter.epsilon)
    state.grad_f_w, state.f_w = gradient(iter.f, state.w)
    return IterationOutput(state.w, state.f_w, state.grad_f_w), state
end

struct RMSProp
    kwargs
    RMSProp(; kwargs...) = new(kwargs)
end

(alg::RMSProp)(args...) = RMSPropIterable(args...; alg.kwargs...)


module day17

include("inputs.jl")

ex = inputs.ex_d17()
input = inputs.input_d17()

function parse_input(input, dim)
    input = split(input, "\n") .|> String
    a = [char=='#' for line in input for char in line]
    @assert rem(length(a), sqrt(length(a)))==0.0 # square-able dim
    side = Int(sqrt(length(a)))
    input2d = reshape(a, (side, side)) |> permutedims
    if dim == 2
        return input2d
    elseif dim == 3
        return reshape(input2d, size(input2d)..., 1)
    elseif dim == 4
        return reshape(input2d, size(input2d)..., 1, 1)
    end
end

function count_actives(input, dim)
    cycles = 6
    parsed_input = parse_input(input, dim)

    grid_size = size(parsed_input) .+ (2*cycles)
    grid = zeros(Bool, grid_size)

    seed_loc = UnitRange.(cycles+1, cycles .+ size(parsed_input))
    grid[seed_loc...] = parsed_input

    adj = CartesianIndices((repeat([-1:1], dim)...,))

    next_grid = copy(grid)
    for cycle in 1:cycles
        for coord ∈ CartesianIndices(grid)
            n_adj = sum(get(grid, coord + a, 0) for a in adj) - grid[coord]
            if grid[coord] && n_adj ∈ [2, 3]
                # println("keep ", coord, "true")
                next_grid[coord] = true
            elseif grid[coord] && n_adj ∉ [2, 3]
                 next_grid[coord] = false
            elseif !grid[coord] && n_adj == 3
                # println("set ", coord, "from false to true")
                next_grid[coord] = true
            elseif !grid[coord] && n_adj != 3
                 next_grid[coord] = false
            end
        end
        grid .= next_grid
    end
    return sum(grid)
end

part1(input) = count_actives(input, 3)

@assert part1(ex) == 112
@info "Day 17, Part 1 answer: $(part1(input))"

part2(input) = count_actives(input, 4)
@assert part2(ex) == 848
@info "Day 17, Part 2 answer: $(part2(input))"

end


using Test
import DataFrames: DataFrame, DataFrameRow

using JMcDM

const testUtilityFunctions = true
const testMCDMFunctions = true
const testSCDMFunctions = true
const testLPBasedFunctions = true

if testUtilityFunctions
    include("./testutility.jl")
end

if testSCDMFunctions
    include("./testscdm.jl")
end # Test SCDM Tools

if testMCDMFunctions
    include("./testmcdm.jl")
end  # Test MCDM Tools

if testLPBasedFunctions
    include("./testlp.jl")
end # Test LP Based Tools


# REVERSE OPERATORS

function MC_1_is_equal(y, x, tol)
    bool1 = isapprox(y.cc,x.cc,atol=tol)
    bool2 = isapprox(y.cv,x.cv,atol=tol)
    bool3 = isapprox(y.cv_grad[1], x.cv_grad[1], atol=tol)
    bool4 = isapprox(y.cc_grad[1], x.cc_grad[1], atol=tol)
    bool5 = isapprox(y.Intv.lo, x.Intv.lo, atol=tol)
    bool6 = isapprox(y.Intv.hi, x.Intv.hi, atol=tol)
    return (bool1 && bool2 && bool3 && bool4 && bool5 && bool6)
end

@testset "Reverse Extrema" begin
    a = MC{1,NS}(2.0, Interval{Float64}(0.4,3.0), 1)
    b = MC{1,NS}(Interval{Float64}(-10.0,-1.0))
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))

    aout1, bout1, cout1 = max_rev(a, b, c)
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0
    @test cout1.Intv.lo == 1.1
    @test cout1.Intv.hi == 3.0

    aout1, cout1, bout1 = max_rev(a, c, b)
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0
    @test cout1.Intv.lo == 1.1
    @test cout1.Intv.hi == 3.0

    empt = empty(a)
    aout1, bout1, cout1 = max_rev(empt, b, c)
    @test isempty(aout1)
    @test isempty(bout1)
    @test isempty(cout1)

    a = MC{1,NS}(2.0, Interval{Float64}(0.4,3.0), 1)
    b = MC{1,NS}(Interval{Float64}(10.0, 20.0))
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))

    aout1, bout1, cout1 = min_rev(a, b, c)
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0
    @test cout1.Intv.lo == 1.1
    @test cout1.Intv.hi == 3.0

    aout1, cout1, bout1 = min_rev(a, c, b)
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0
    @test cout1.Intv.lo == 1.1
    @test cout1.Intv.hi == 3.0

    empt = empty(a)
    aout1, bout1, cout1 = min_rev(empt, b, c)
    @test isempty(aout1)
    @test isempty(bout1)
    @test isempty(cout1)

    a = MC{1,NS}(2.0, Interval{Float64}(0.4,3.0), 1)
    b = -1.1
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))
    aout1, bout1, cout1 = max_rev(a, b, c)
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0

    aout1, bout1, cout1 = max_rev(a, c, b)
    @test bout1.cv == 2.0
    @test bout1.cc == 2.0

    b = 5.1
    aout1, bout1, cout1 = min_rev(a, b, c)
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0

    aout1, bout1, cout1 = min_rev(a, c, b)
    @test bout1.cv == 2.0
    @test bout1.cc == 2.0
end

@testset "Reverse Multiplication" begin

    # THE BINARY OPERATOR
    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    b = MC{1,NS}(Interval{Float64}(-10.0,-1.0))
    c = MC{1,NS}(2.0, Interval{Float64}(1.1,4.5), 1)

    aout1, bout1, cout1 = mult_rev(a,b,c)

    @test bout1.Intv.lo == Inf
    @test bout1.Intv.hi == -Inf
    @test cout1.Intv.lo == Inf
    @test cout1.Intv.hi == -Inf

    bout1 = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    cout1 = MC{1,NS}(Interval{Float64}(-10.0,-1.0))
    aout1 = bout1*cout1

    aout1_a, bout1_a, cout1_a = mult_rev(aout1, bout1, cout1)

    MC_1_is_equal(aout1_a, aout1, 0.00001)
    MC_1_is_equal(bout1_a, bout1, 0.00001)
    MC_1_is_equal(cout1_a, cout1, 0.00001)

    bout2 = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    cout2 = MC{1,NS}(Interval{Float64}(-10.0,-1.0))
    aout2 = 0.3*bout1*cout1+1.0

    aout2_a, bout2_a, cout2_a = mult_rev(aout2, bout2, cout2)

    MC_1_is_equal(aout2_a, aout2, 0.00001)
    MC_1_is_equal(bout2_a, bout2, 0.00001)
    @test cout2_a.Intv.lo == -10.0
    @test cout2_a.Intv.hi == -1.0
    @test cout2_a.cv == -6.32
    @test cout2_a.cc == -1.0
    @test cout2_a.cv_grad[1] == -8.38
    @test cout2_a.cc_grad[1] == 0.0

    # THE BINARY OPERATOR
    a = MC{1,NS}(Interval{Float64}(0.4,3.0))
    b = 0.5
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))

    aout2, bout2, cout2 = mult_rev(a, b, c)
    @test cout2.cv == 1.1
    @test cout2.cc == 4.5

    aout2, bout2, cout2 = mult_rev(a, c, b)
    @test bout2.cv == 1.1
    @test bout2.cc == 4.5

    aout2, bout2, cout2 = mult_rev(a, 0.0, c)
    @test isempty(cout2)

    aout2, bout2, cout2 = mult_rev(a, c, 0.0)
    @test isempty(bout2)
end

@testset "Reverse Addition" begin
    # THE ADDITION OPERATOR
    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    b = MC{1,NS}(Interval{Float64}(-10.0,-1.0))
    c = MC{1,NS}(2.0, Interval{Float64}(1.1,4.5), 1)

    aout1, bout1, cout1 = plus_rev(a,b,c)

    @test isapprox(bout1.Intv.lo, -4.1000000000000005)
    @test bout1.Intv.hi == -1.0
    @test cout1.Intv.lo == 1.4
    @test cout1.Intv.hi == 4.5

    # THE ADDITION OPERATOR
    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    b = MC{1,NS}(Interval{Float64}(-10.0,-1.0))
    c = 3.0

    aout1, bout1, cout1 = plus_rev(a, b, c)
    @test bout1.cv == -2.0
    @test cout1.cv == 3.0

    a = MC{1,NS}(45.0, Interval{Float64}(40.0, 50.0), 1)
    aout1, bout1, cout1 = plus_rev(a, b, c)
    @test isempty(aout1)
    @test isempty(bout1)
    @test isempty(cout1)

    a = empty(MC{1,NS}(45.0, Interval{Float64}(40.0, 50.0), 1))
    aout1, bout1, cout1 = plus_rev(a, b, c)
    @test isempty(aout1)
    @test isempty(bout1)
    @test isempty(cout1)

    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    aout1, cout1, bout1 = plus_rev(a, c, b)
    @test bout1.cv == -2.0
    @test cout1.cv == 3.0

    c = MC{1,NS}(2.0, Interval{Float64}(1.1,4.5), 1)
    aout1, bout1 = plus_rev(a, c)                    # should be NaN
    @test isnan(bout1.cv)
    @test isnan(bout1.cc)

    c = MC{1,NS}(2.0, Interval{Float64}(1.1,4.5), 1)
    b = MC{1,NS}(Interval{Float64}(-10.0, 10.0))
    aout1, bout1 = plus_rev(c, b)
    @test bout1.cv == 2.0
    @test bout1.cc == 2.0
end

@testset "Reverse Division" begin
    a = MC{1,NS}(1.14, Interval{Float64}(0.1,1.2), 1)
    b = inv(a)
    out = inv_rev(b,a)
    @test out[2].cv == a.cv
    @test out[2].cc == a.cc

    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    b = MC{1,NS}(Interval(0.5, 12.0))
    c = MC{1,NS}(2.0, Interval{Float64}(1.1,4.5), 1)

    aout1, bout1, cout1 = div_rev(a, b, c)
    @test bout1.cv == 1.46
    @test bout1.cc == 3.5
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0

    aout1, bout1, cout1 = div_rev(a, 2.0, c)
    @test cout1.cv == 2.0
    @test cout1.cc == 2.0

    a = MC{1,NS}(Interval{Float64}(0.4,3.0))
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))
    aout1, bout1, cout1 = div_rev(a, c, 3.0)
    @test isapprox(bout1.cv, 1.2000000000000002, atol=1E-6)
    @test bout1.cc == 4.5

    a = empty(MC{1,NS}(Interval{Float64}(0.4,3.0)))
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))
    aout1, bout1, cout1 = div_rev(a, c, 3.0)
    @test isempty(bout1)

    a = empty(MC{1,NS}(Interval{Float64}(0.4,3.0)))
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))
    aout1, bout1, cout1 = div_rev(a, 3.0, c)
    @test isempty(cout1)

end

@testset "Reverse Subtraction" begin
    a = MC{1,NS}(Interval{Float64}(0.4,3.0))
    b = MC{1,NS}(Interval(0.5, 12.0))
    c = MC{1,NS}(Interval{Float64}(1.1,4.5))

    aout1, bout1, cout1 = minus_rev(a, b, c)
    @test bout1.cv == 1.5
    @test bout1.cc == 7.5
    @test cout1.cv == 1.1
    @test cout1.cc == 4.5

    aout1, bout1, cout1 = minus_rev(a, b, 3.0)
    @test bout1.cv == 3.4
    @test bout1.cc == 6.0

    aout1, bout1, cout1 = minus_rev(a, 3.0, b)
    @test cout1.cv == 0.5
    @test cout1.cc == 2.6

    aout1, bout1 = minus_rev(-a, c)
    @test bout1.cv == 1.1
    @test bout1.cc == 3.0
end

@testset "Reverse Exponential" begin
    a = MC{1,NS}(1.0,Interval{Float64}(0.4,3.0),1)
    expa = exp(a)*1.1
    y,x = exp_rev(expa,a)

    @test MC_1_is_equal(expa, y, 0.00001)
    @test x.cc == 1.0
    @test x.cv == 1.0
    @test isapprox(x.Intv.lo, 0.49531, atol=1E-5)
    @test x.Intv.hi == 3.0
    @test x.cc_grad[1] == 1.0
    @test x.cv_grad[1] == 1.0

    exp2a = exp2(a)*1.1
    y,x = exp2_rev(exp2a,a)

    @test MC_1_is_equal(exp2a, y, 0.00001)
    @test x.cc == 1.0
    @test x.cv == 1.0
    @test isapprox(x.Intv.lo, 0.537503, atol=1E-5)
    @test x.Intv.hi == 3.0
    @test x.cc_grad[1] == 1.0
    @test x.cv_grad[1] == 1.0

    exp10a = exp10(a)*1.1
    y,x = exp10_rev(exp10a,a)

    @test MC_1_is_equal(exp10a, y, 0.00001)
    @test x.cc == 1.0
    @test x.cv == 1.0
    @test isapprox(x.Intv.lo, 0.441392, atol=1E-5)
    @test x.Intv.hi == 3.0
    @test x.cc_grad[1] == 1.0
    @test x.cv_grad[1] == 1.0

    expm1a = expm1(a)*1.1
    y,x = expm1_rev(expm1a,a)

    @test MC_1_is_equal(expm1a, y, 0.00001)
    @test x.cc == 1.0
    @test x.cv == 1.0
    @test isapprox(x.Intv.lo, 0.432436, atol=1E-5)
    @test x.Intv.hi == 3.0
    @test x.cc_grad[1] == 1.0
    @test x.cv_grad[1] == 1.0
end

@testset "Reverse Logarithm" begin

   a = MC{2,NS}(9.5,9.6,Interval{Float64}(8.0,11.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   b = MC{2,NS}(9.5,9.6,Interval{Float64}(9.4,9.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   log1a = log(b)
   y1,x1 = log_rev(log1a,a)

   @test isapprox(9.5, x1.cv, atol=1E-5)
   @test isapprox(9.6, x1.cc, atol=1E-5)
   @test isapprox(9.399999999999997, x1.Intv.lo, atol=1E-4)
   @test isapprox(9.800000000000002, x1.Intv.hi, atol=1E-4)
   @test isapprox(1.0, x1.cv_grad[1], atol=1E-4)
   @test isapprox(1.0, x1.cc_grad[1], atol=1E-4)

   @test isapprox(2.2511278633761, y1.cv, atol=1E-5)
   @test isapprox(2.2617630, y1.cc, atol=1E-5)
   @test isapprox(2.240709689275958, y1.Intv.lo, atol=1E-4)
   @test isapprox(2.2823823856765264, y1.Intv.hi, atol=1E-4)
   @test isapprox( 0.1041817410014211, y1.cv_grad[1], atol=1E-4)
   @test isapprox(0.10416666666666667, y1.cc_grad[1], atol=1E-4)

   a = MC{2,NS}(9.5,9.6,Interval{Float64}(8.0,11.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   b = MC{2,NS}(9.5,9.6,Interval{Float64}(9.4,9.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   log2a = log2(b)
   y2,x2 = log2_rev(log2a,a)

   @test isapprox(3.247691004899668, y2.cv, atol=1E-5)
   @test isapprox(3.2630344, y2.cc, atol=1E-5)
   @test isapprox(3.232660756790275, y2.Intv.lo, atol=1E-4)
   @test isapprox(3.292781749227846, y2.Intv.hi, atol=1E-4)
   @test isapprox(0.15030248109392783, y2.cv_grad[1], atol=1E-4)
   @test isapprox(0.1502807334259337, y2.cc_grad[1], atol=1E-4)

   a = MC{2,NS}(9.5,9.6,Interval{Float64}(8.0,11.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   b = MC{2,NS}(9.5,9.6,Interval{Float64}(9.4,9.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   log10a = log10(b)
   y3,x3 = log10_rev(log10a,a)

   @test isapprox(0.9776524091228977, y3.cv, atol=1E-5)
   @test isapprox(0.9822712, y3.cc, atol=1E-5)
   @test isapprox(0.9731278535996987, y3.Intv.lo, atol=1E-4)
   @test isapprox(0.991226075692495, y3.Intv.hi, atol=1E-4)
   @test isapprox(0.04524555523199077, y3.cv_grad[1], atol=1E-4)
   @test isapprox(0.04523900853158873, y3.cc_grad[1], atol=1E-4)

   a = MC{2,NS}(9.5,9.6,Interval{Float64}(8.0,11.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   b = MC{2,NS}(9.5,9.6,Interval{Float64}(9.4,9.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   log1pa = log1p(b)
   y4,x4 = log1p_rev(log1pa,a)

   @test isapprox(2.3512408881430384, y4.cv, atol=1E-5)
   @test isapprox(2.3608540, y4.cc, atol=1E-5)
   @test isapprox(2.3418058061473266, y4.Intv.lo, atol=1E-4)
   @test isapprox(2.3795461341301745, y4.Intv.hi, atol=1E-4)
   @test isapprox(0.09435081995711961, y4.cv_grad[1], atol=1E-4)
   @test isapprox(0.09433962264150944, y4.cc_grad[1], atol=1E-4)
end

@testset "Reverse Trignometric" begin

    a = MC{2,NS}(9.5,9.6,Interval{Float64}(8.0,11.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    b = MC{2,NS}(9.5,9.6,Interval{Float64}(9.4,9.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)

    sin1a = sin(b)
    msin = sin_rev(sin1a, b)
    @test msin[2].cv == 9.5
    @test msin[2].cc == 9.6

    cos1a = cos(b)
    mcos = cos_rev(cos1a, b)
    @test mcos[2].cv == 9.5
    @test mcos[2].cc == 9.6

    tan1a = tan(b)
    mtan = tan_rev(tan1a, b)
    @test mtan[2].cv == 9.5
    @test mtan[2].cc == 9.6

    b = MC{2,NS}(0.5,0.6,Interval{Float64}(0.4,0.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    sec1a = sec(b)
    msec = sec_rev(sec1a, b)
    @test msec[2].cv == 0.5
    @test msec[2].cc == 0.6

    csc1a = csc(b)
    mcsc = csc_rev(csc1a, b)
    @test mcsc[2].cv == 0.5
    @test mcsc[2].cc == 0.6

    cot1a = cot(b)
    mcot = cot_rev(cot1a, b)
    @test mcot[2].cv == 0.5
    @test mcot[2].cc == 0.6

    asin1a = asin(b)
    masin = asin_rev(asin1a, b)
    @test isapprox(masin[2].cv, 0.48692255094800774, atol=1E-6)
    @test isapprox(masin[2].cc, 0.6205203125624899, atol=1E-6)

    acos1a = acos(b)
    macos = acos_rev(acos1a, b)
    @test isapprox(macos[2].cv, 0.5000000000000001, atol=1E-6)
    @test isapprox(macos[2].cc, 0.5999999999999999, atol=1E-6)

    atan1a = atan(b)
    matan = atan_rev(atan1a, b)
    @test matan[2].cv == 0.5
    @test matan[2].cc == 0.6

    b = MC{2,NS}(1.1,2.2,Interval{Float64}(1.1,2.2), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    asec1a = asec(b)
    masec = asec_rev(asec1a, b)
    @test masec[2].cv == 1.1
    @test masec[2].cc == 2.2

    acsc1a = acsc(b)
    macsc = acsc_rev(acsc1a, b)
    @test macsc[2].cv == 1.1
    @test macsc[2].cc == 2.2

    b = MC{2,NS}(0.5,0.6,Interval{Float64}(0.4,0.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    acot1a = acot(b)
    macot = acot_rev(acot1a, b)
    @test macot[2].cv == 0.5
    @test macot[2].cc == 0.6

    sind1a = sind(b)
    msind = sind_rev(sind1a, b)
    @test msind[2].cv == 0.5
    @test msind[2].cc == 0.6

    cosd1a = cosd(b)
    mcosd = cosd_rev(cosd1a, b)
    @test mcosd[2].cv == 0.5
    @test mcosd[2].cc == 0.6

    tand1a = tand(b)
    mtand = tand_rev(tand1a, b)
    @test mtand[2].cv == 0.5
    @test mtand[2].cc == 0.6

    secd1a = secd(b)
    msecd = secd_rev(secd1a, b)
    @test msecd[2].cv == 0.5
    @test msecd[2].cc == 0.6

    cscd1a = cscd(b)
    mcscd = cscd_rev(cscd1a, b)
    @test mcscd[2].cv == 0.5
    @test mcscd[2].cc == 0.6

    cotd1a = cotd(b)
    mcotd = cotd_rev(cotd1a, b)
    @test mcotd[2].cv == 0.5
    @test mcotd[2].cc == 0.6

    b = MC{1,NS}(Interval{Float64}(0.85, 0.95))
    b1 = MC{1,NS}(Interval{Float64}(0.1, 0.9))*60.0
    asind1a = asind(b)
    masind = asind_rev(asind1a, b1)
    @test masind[2].cv == Inf
    @test masind[2].cc == -Inf

    acosd1a = acosd(b)
    macosd = acosd_rev(acosd1a, b1)
    @test macosd[2].cv == Inf
    @test macosd[2].cc == -Inf

    #atand1a = atand(b)
    #matand = atand_rev(atand1a, b1)
    #@test matand[2].cv == Inf
    #@test matand[2].cc == -Inf

    asecd1a = asecd(b)
    masecd = asecd_rev(asecd1a, b1)
    @test isnan(masecd[2].cv)
    @test isnan(masecd[2].cc)

    b = MC{1,NS}(Interval{Float64}(1.1,2.2))*60
    acscd1a = acscd(b)
    macscd = acscd_rev(acscd1a, b)
    @test macscd[2].cv == 66.0
    @test macscd[2].cc == 132.0

    acotd1a = acotd(b)
    macotd = acotd_rev(acotd1a, b)
    @test macotd[2].cv == 66.0
    @test macotd[2].cc == 132.0
end

@testset "Reverse Hyperbolic" begin

   a = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   b = MC{2,NS}(2.4,2.6,Interval{Float64}(2.2,2.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   sinh1 = sinh(b)
   y1,x1 = sinh_rev(sinh1,a)
   @test isapprox(y1.cv, 5.466229213676094, atol=1E-7)
   @test isapprox(y1.cc, 6.946980626335908, atol=1E-7)
   @test isapprox(y1.Intv.lo, 4.457105170535894, atol=1E-7)
   @test isapprox(y1.Intv.hi, 8.191918354235915, atol=1E-7)
   @test isapprox(y1.cv_grad[1], 5.556947166965507, atol=1E-7)
   @test isapprox(y1.cc_grad[1], 6.224688639500037, atol=1E-7)

   coshb = cosh(b)
   y2,x2 = cosh_rev(coshb,a)
   @test isapprox(y2.cv, 5.556947166965507, atol=1E-7)
   @test isapprox(y2.cc, 7.024455054206832, atol=1E-7)
   @test isapprox(y2.Intv.lo, 4.567908328898228, atol=1E-7)
   @test isapprox(y2.Intv.hi, 8.252728416861133, atol=1E-7)
   @test isapprox(y2.cv_grad[1], 5.466229213676094, atol=1E-7)
   @test isapprox(y2.cc_grad[1], 6.141366813271513, atol=1E-7)

   tanh1 = tanh(b)
   y3,x3 = tanh_rev(tanh1,a)
   @test isapprox(y3.cv, 0.9813725934213438, atol=1E-7)
   @test isapprox(y3.cc, 0.9890274022010992, atol=1E-7)
   @test isapprox(y3.Intv.lo, 0.9757431300314515, atol=1E-7)
   @test isapprox(y3.Intv.hi, 0.9926315202011281, atol=1E-7)
   @test isapprox(y3.cv_grad[1], 0.0281473169494608, atol=1E-7)
   @test isapprox(y3.cc_grad[1], 0.021824797695345177, atol=1E-7)

   @test isapprox(x1.cv, 2.5, atol=1E-7)
   @test isapprox(x1.cc, 2.6, atol=1E-7)
   @test isapprox(x1.Intv.lo, 2.19999, atol=1E-5)
   @test isapprox(x1.Intv.hi, 2.80001, atol=1E-5)

   a = MC{2,NS}(2.5,2.6,Interval{Float64}(1.3,4.5), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   b = MC{2,NS}(2.4,2.6,Interval{Float64}(2.3,3.4), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)

   asinhb = asinh(b)
   y1,x1 = asinh_rev(asinhb,a)
   @test isapprox(y1.cv, 1.6036967920467529, atol=1E-7)
   @test isapprox(y1.cc, 1.6837431439977444, atol=1E-7)
   @test isapprox(y1.Intv.lo, 1.570278543484978, atol=1E-7)
   @test isapprox(y1.Intv.hi, 1.9378792776645006, atol=1E-7)
   @test isapprox(y1.cv_grad[1], 0.3341824856177478, atol=1E-7)
   @test isapprox(y1.cc_grad[1], 0.358979079308869, atol=1E-7)

   acoshb = acosh(b)
   y2,x2 = acosh_rev(acoshb,a)
   @test isapprox(y2.cv, 1.5131824386442316, atol=1E-7)
   @test isapprox(y2.cc, 1.6094379124341003, atol=1E-7)
   @test isapprox(y2.Intv.lo, 1.47504478124142, atol=1E-7)
   @test isapprox(y2.Intv.hi, 1.894559012672298, atol=1E-7)
   @test isapprox(y2.cv_grad[1], 0.3813765740280664, atol=1E-7)
   @test isapprox(y2.cc_grad[1], 0.45834924851410563, atol=1E-7)

   a = MC{2,NS}(0.31,0.34,Interval{Float64}(-0.2,0.5), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   b = MC{2,NS}(0.31,0.34,Interval{Float64}(0.3,0.4), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)

   atanhb = atanh(b)
   y3,x3 = atanh_rev(atanhb,a)
   @test isapprox(y3.cv, 0.3205454093019461, atol=1E-7)
   @test isapprox(y3.cc, 0.35517133459930783, atol=1E-7)
   @test isapprox(y3.Intv.lo, 0.3095196042031117, atol=1E-7)
   @test isapprox(y3.Intv.hi, 0.42364893019360184, atol=1E-7)
   @test isapprox(y3.cv_grad[1], 1.1063170704723972, atol=1E-7)
   @test isapprox(y3.cc_grad[1], 1.1412932599049004, atol=1E-7)

   @test isapprox(x3.cv, 0.31, atol=1E-7)
   @test isapprox(x3.cc, 0.34, atol=1E-7)
   @test isapprox(x3.Intv.lo, 0.29999999999999993, atol=1E-7)
   @test isapprox(x3.Intv.hi, 0.4000000000000001, atol=1E-7)
   @test isapprox(x3.cv_grad[1], 1.0, atol=1E-7)
   @test isapprox(x3.cc_grad[1], 1.0, atol=1E-7)

   asechv = asech(b)
   asecha, asechb = asech_rev(asechv, a)
   @test asechb.cv == 0.31
   @test asechb.cc == 0.34

   acschv = acsch(b)
   acscha, acschb = acsch_rev(acschv, a)
   @test acschb.cv == 0.31
   @test acschb.cc == 0.34

   sechv = sech(b)
   secha, sechb = sech_rev(sechv, a)
   @test sechb.cv == 0.31
   @test sechb.cc == 0.34

   cschv = csch(b)
   cscha, cschb = csch_rev(cschv, a)
   @test cschb.cv == 0.31
   @test cschb.cc == 0.34

   a0 = MC{2,NS}(0.1, 0.2, Interval{Float64}(0.1,0.2), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   a1 = MC{2,NS}(Interval{Float64}(0.01,2.1))
   cothv = coth(a0)
   cotha, cothb = coth_rev(cothv, a1)
   @test isapprox(cothb.cv, 1.0000000038566186, atol=1E-6)
   @test isapprox(cothb.cc, 1.0000794969264033, atol=1E-6)

   b1 = MC{2,NS}(2.4,2.6,Interval{Float64}(2.2,2.8), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
   a2 = MC{2,NS}(Interval{Float64}(0.01,3.1))
   acothv = acoth(b1)
   acotha, acothb = acoth_rev(acothv, a2)
   @test isapprox(acothb.cv, 2.362216567294689, atol=1E-6)
   @test isapprox(acothb.cc, 2.6376185113149493, atol=1E-6)
end

@testset "Reverse Power" begin

    a = MC{1,NS}(1.14, Interval{Float64}(0.1,1.2), 1)
    b = sqrt(a)
    out = sqrt_rev(b,a)
    @test out[2].cv == a.cv
    @test out[2].cc == a.cc

    a = MC{1,NS}(1.14, Interval{Float64}(0.1,1.2), 1)
    b = a^2
    out = sqr_rev(b,a)
    @test out[2].cv == a.cv
    @test out[2].cc == a.cc

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^-3
    aout1, bout1, cout1 = power_rev(a, b, -3)
    @test bout1.cv == 2.5
    @test bout1.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^-2
    aout1, bout, cout1 = power_rev(a, b, -2)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^-1
    aout1, bout, cout1 = power_rev(a, b, -1)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^0
    aout1, bout, cout1 = power_rev(a, b, 0)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^1
    aout1, bout, cout1 = power_rev(a, b, 1)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^2
    aout1, bout, cout11 = power_rev(a, b, 2)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^3
    aout1, bout, cout1 = power_rev(a, b, 3)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^4
    aout1, bout, cout1 = power_rev(a, b, 4)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^2.5
    aout1, bout, cout1 = power_rev(a, b, 2.5)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    c = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = b^c
    aout1, bout, cout1 = power_rev(a, b, c)
    @test bout.cv == 2.5
    @test bout.cc == 2.6

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    c = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = empty(b)
    aout1, bout, cout1 = power_rev(a, b, c)
    @test isempty(aout1)
    @test isempty(bout)
    @test isempty(cout1)

    b = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    c = empty(MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false))
    a = b^c
    aout1, bout, cout1 = power_rev(a, b, c)
    @test isempty(aout1)
    @test isempty(bout)
    @test isempty(cout1)

    c = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    a = 3.0^c
    aout1, bout, cout1 = power_rev(a, 3.0, c)
    @test bout.cv == 3.0
    @test bout.cc == 3.0

    c = empty(MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false))
    a = 3.0^c
    aout1, bout, cout1 = power_rev(a, 3.0, c)
    @test isempty(aout1)
    @test isempty(cout1)

    c = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    c1 = empty(MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false))
    a = 3.0^c
    aout1, bout, cout1 = power_rev(a, 3.0, c1)
    @test isempty(aout1)
    @test isempty(cout1)

    c = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    c1 = MC{2,NS}(2.5,2.6,Interval{Float64}(2.0,3.0), seed_gradient(1,Val(2)), seed_gradient(1,Val(2)), false)
    aout1, bout, cout1 = power_rev(c, c1, 0)
    @test isempty(aout1)
    @test isempty(cout1)
end

@testset "Reverse Other Operators" begin
    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    b = MC{1,NS}(Interval(0.2, 12.0))

    aout1, bout1 = deg2rad_rev(a, b)
    @test isempty(bout1)

    c = MC{1,NS}(Interval{Float64}(0.4,3.0))*0.5*pi/180.0
    aout1, bout1 = deg2rad_rev(c, a)
    @test bout1.cv == 1.0
    @test bout1.cc == 1.0
    @test bout1.Intv.lo == 0.4
    @test isapprox(bout1.Intv.hi, 1.5000000000000009, atol=1E-6)

    aout1, bout1 = rad2deg_rev(a, b)
    @test isempty(bout1)

    c = MC{1,NS}(Interval{Float64}(0.4,3.0))*0.5*180.0/pi
    aout1, bout1 = rad2deg_rev(c, a)
    @test bout1.cv == 1.0
    @test bout1.cc == 1.0
    @test bout1.Intv.lo == 0.4
    @test isapprox(bout1.Intv.hi, 1.5000000000000007, atol=1E-6)

    aout1, bout1 = step_rev(a, b)
    @test b.cv == bout1.cv
    @test b.cc == bout1.cc

    aout1, bout1 = sign_rev(a, b)
    @test b.cv == bout1.cv
    @test b.cc == bout1.cc

    a = MC{1,NS}(0.5, Interval{Float64}(-0.2,0.8), 1)
    aout1, bout1 = step_rev(a, b)
    @test isnan(bout1)

    a = MC{1,NS}(0.5, Interval{Float64}(-1.2,0.8), 1)
    aout1, bout1 = sign_rev(a, b)
    @test isnan(bout1)

    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    aout1, bout1 = abs_rev(a, b)
    @test b.cv == bout1.cv
    @test b.cc == bout1.cc

    aout1, bout1 = zero_rev(a, b)
    @test isempty(bout1)

    c = MC{1,NS}(Interval(-2.2, 12.0))
    aout1, bout1 = zero_rev(c, a)
    @test bout1.cv == a.cv
    @test bout1.cc == a.cc

    aout1, bout1 = one_rev(a, b)
    @test b.cv == bout1.cv
    @test b.cc == bout1.cc

    a = MC{1,NS}(2.6, Interval{Float64}(2.4,3.0), 1)
    aout1, bout1 = one_rev(a, b)
    @test isempty(bout1)

    a = MC{1,NS}(1.0, Interval{Float64}(0.4,3.0), 1)
    aout1, bout1 = real_rev(a, b)
    @test aout1.cv == bout1.cv
    @test aout1.cc == bout1.cc
end

# reverse of an empty is an empty (should override NaN)
@testset "Empty Propagation" begin
    for T in (NS, MV, Diff)
        a = MC{1,T}(3.1, Interval{Float64}(2.5, 6.0), 1)
        c = a + 2.0
        b = empty(a)
        for f in (exp_rev, exp2_rev, exp10_rev, expm1_rev, log_rev, log2_rev,
                  log10_rev, log1p_rev, sin_rev, cos_rev, tan_rev, asin_rev,
                  acos_rev, atan_rev, sinh_rev, cosh_rev, tanh_rev, asinh_rev,
                  acosh_rev, atanh_rev, abs_rev, sqrt_rev, minus_rev, plus_rev,
                  zero_rev, real_rev, one_rev, step_rev, sign_rev, deg2rad_rev,
                  rad2deg_rev, sind_rev, cosd_rev, tand_rev, sec_rev, csc_rev,
                  cot_rev, secd_rev, cscd_rev, cotd_rev, asind_rev, acosd_rev,
                  atand_rev, asec_rev, acsc_rev, acot_rev, asecd_rev, acscd_rev,
                  acot_rev, inv_rev)
            bout, aout = f(b, a)
            @test isempty(aout)
        end
        for f in (plus_rev, minus_rev, mult_rev, div_rev)
            bout, cout, aout = f(b, c, a)
            ~isempty(aout) && (println("f = $(f)"))
            ~isempty(cout) && (println("f = $(f)"))
            @test isempty(aout)
            @test isempty(cout)
        end
        # special cases, ^, max, min
    end
end