text:
using Test, AdmittanceModels
using ProgressMeter: @showprogress
using LinearAlgebra: norm, normalize

# set to false to create the figures in the paper, though it will take longer
run_fast = true
# to save data so figures in the paper can be created, set to true
write_data = false
# to make plots and print information set to true
show_results = false

if write_data
    using DataFrames, CSV, NPZ
    output_folder = joinpath(@__DIR__, "hybridization_results")
    try
        mkdir(output_folder)
    catch SystemError
    end
end

if show_results
    using PlotlyJS
end

#######################################################
# Parameters
#######################################################

begin
    ν = 1.2e8 # propagation_speed
    Z0 = 50.0 # characteristic_impedance
    δ = 50e-6 # discretization length
    pfilter_num = 202
    pfilter_len = pfilter_num * δ # purcell filter length
    pfilter_res_0_coupler_num = floor(Int, pfilter_num/2) - 5
    pfilter_res_coupler_0_len = pfilter_res_0_coupler_num * δ # position of res 0 coupler on filter
    pfilter_res_1_coupler_num = floor(Int, pfilter_num/2) + 5
    pfilter_res_coupler_1_len = pfilter_res_1_coupler_num * δ # position of res 1 coupler on filter
    res_0_num = 102
    res_0_len = res_0_num * δ
    res_1_num = 101
    res_1_len = res_1_num * δ
    res_coupler_num = 16
    res_coupler_len = res_coupler_num * δ # res coupler length from open
    capacitance_0 = 12e-15 # res 0 - purcell filter capacitance
    capacitance_1 = 12e-15 # res 1 - purcell filter capacitance
end

#######################################################
# Two resonators and Purcell filter
# Circuit version
#######################################################

begin
    pfilter_circ = Circuit(TransmissionLine(["p/short_0", "p/short_1"], ν, Z0, pfilter_len, δ=δ),
                           "p/", [pfilter_num])
    res_0_circ = Circuit(TransmissionLine(["r0/open", "r0/short"], ν, Z0, res_0_len, δ=δ),
                         "r0/", [res_0_num])
    res_1_circ = Circuit(TransmissionLine(["r1/open", "r1/short"], ν, Z0, res_1_len, δ=δ),
                         "r1/", [res_1_num])
    pfilter_res_0_coupler_name = "p/$(1+pfilter_res_0_coupler_num)"
    res_0_coupler_name = "r0/$(1+res_coupler_num)"
    pfilter_res_1_coupler_name = "p/$(1+pfilter_res_1_coupler_num)"
    res_1_coupler_name = "r1/$(1+res_coupler_num)"
    capacitor_0 = Circuit(SeriesComponent(pfilter_res_0_coupler_name, res_0_coupler_name, 0, 0, capacitance_0))
    capacitor_1 = Circuit(SeriesComponent(pfilter_res_1_coupler_name, res_1_coupler_name, 0, 0, capacitance_1))
    circ = connect(pfilter_circ, res_0_circ, res_1_circ, capacitor_0, capacitor_1)
    ground = AdmittanceModels.ground
    circ = unite_vertices(circ, ground, "p/short_0", "p/short_1", "r0/short", "r1/short")

    input_output(stub_num) = ("p/$(1+stub_num)", "p/$(pfilter_num-stub_num+1)")
end

function bbox_from_circ(stub_num::Int, ω::Vector{<:Real})
    input, output = input_output(stub_num)
    pso = PSOModel(circ, [(input, ground), (output, ground)], ["in", "out"])
    return Blackbox(ω, pso) |> canonical_gauge
end

#######################################################
# Two resonators and Purcell filter
# Admittance model version
#######################################################

function components_from_model(stub_num::Int, discretization::Real=δ)
    pfilter = TransmissionLine(["p/short_0", "in", "p/res_0_coupler", "p/middle",
                                "p/res_1_coupler", "out", "p/short_1"],
                               ν, Z0, pfilter_len, locations=[stub_num * δ, pfilter_res_coupler_0_len,
                                                              .5 * pfilter_len, pfilter_res_coupler_1_len,
                                                              pfilter_len - stub_num * δ], δ=discretization)
    res_0 = TransmissionLine(["r0/open", "r0/pfilter_coupler", "r0/short"],
                             ν, Z0, res_0_len, locations=[res_coupler_len], δ=discretization)
    res_1 = TransmissionLine(["r1/open", "r1/pfilter_coupler", "r1/short"],
                         ν, Z0, res_1_len, locations=[res_coupler_len], δ=discretization)
    capacitor_0 = SeriesComponent("p/res_0_coupler", "r0/pfilter_coupler", 0, 0, capacitance_0)
    capacitor_1 = SeriesComponent("p/res_1_coupler", "r1/pfilter_coupler", 0, 0, capacitance_1)
    return [pfilter, res_0, res_1, capacitor_0, capacitor_1]
end

function bbox_from_pso(stub_num::Int, ω::Vector{<:Real}, discretization::Real=δ)
    model = connect(PSOModel.(components_from_model(stub_num, discretization)))
    model = short_ports(model, ["p/short_0", "p/short_1", "r0/short", "r1/short"])
    model = open_ports_except(model, ["in", "out"])
    return canonical_gauge(Blackbox(ω, model))
end

function bbox_from_bbox(stub_num::Int, ω::Vector{<:Real})
    model = connect(Blackbox.(Ref(ω), components_from_model(stub_num)))
    model = short_ports(model, ["p/short_0", "p/short_1", "r0/short", "r1/short"])
    model = open_ports_except(model, ["in", "out"])
    return canonical_gauge(model)
end

#######################################################
# Compare them
#######################################################

begin
    ω_circ = collect(range(5.4, stop=6.4, length=run_fast ? 100 : 2000)) * 2π * 1e9

    bbox_circ_10 = bbox_from_circ(10, ω_circ)
    bbox_circ_30 = bbox_from_circ(30, ω_circ)

    bbox_pso_10 = bbox_from_pso(10, ω_circ)
    bbox_pso_30 = bbox_from_pso(30, ω_circ)

    bbox_bbox_10 = bbox_from_bbox(10, ω_circ)
    bbox_bbox_30 = bbox_from_bbox(30, ω_circ)
end

s21(bbox) = [x[1,2] for x in scattering_matrices(bbox, [Z0, Z0])]
freqs = ω_circ/(2π)

if write_data
    df = DataFrame(freqs = freqs,
        circ_S_p5 = s21(bbox_circ_10),
        circ_S_1p5 = s21(bbox_circ_30),
        pso_S_p5 = s21(bbox_pso_10),
        pso_S_1p5 = s21(bbox_pso_30),
        bbox_S_p5 = s21(bbox_bbox_10),
        bbox_S_1p5 = s21(bbox_bbox_30))
    CSV.write(joinpath(output_folder, "S12.csv"), df)
end

if show_results
    plot([scatter(x=freqs/1e6, y=abs.(s21(bbox_circ_10)), name="0.5mm circ"),
          scatter(x=freqs/1e6, y=abs.(s21(bbox_circ_30)), name="1.5mm circ"),
          scatter(x=freqs/1e6, y=abs.(s21(bbox_pso_10)), name="0.5mm pso"),
          scatter(x=freqs/1e6, y=abs.(s21(bbox_pso_30)), name="1.5mm pso"),
          scatter(x=freqs/1e6, y=abs.(s21(bbox_bbox_10)), name="0.5mm bbox", line_dash="dash"),
          scatter(x=freqs/1e6, y=abs.(s21(bbox_bbox_30)), name="1.5mm bbox", line_dash="dash")],
          Layout(xaxis_title="frequency", yaxis_title="|S21|"))
end

if show_results
    abs_diff_10 = abs.(s21(bbox_circ_10) .- s21(bbox_bbox_10))
    abs_diff_30 = abs.(s21(bbox_circ_30) .- s21(bbox_bbox_30))
    plot([scatter(x=ω_circ/(2π*1e6), y=abs_diff_10, name="0.5mm error"),
        scatter(x=ω_circ/(2π*1e6), y=abs_diff_30, name="1.5mm error")],
        Layout(xaxis_title="frequency", yaxis_title="|S21_circ - S21_bbox|"))
end

if show_results
    density(m) = count(!iszero, m)/(size(m,1) * size(m,2))
    model = connect(PSOModel.(components_from_model(10, δ)))
    model = short_ports(model, ["p/short_0", "p/short_1", "r0/short", "r1/short"])
    println("pso matrix densities: $(map(density, get_Y(model)))")
end

#######################################################
# All three models now match so we can use the circuit
# one for making mode plots
#######################################################

function resistor_circ(stub_num::Int)
    input, output = input_output(stub_num)
    resistor_0 = Circuit(ParallelComponent(input, 0, 1/Z0, 0))
    resistor_1 = Circuit(ParallelComponent(output, 0, 1/Z0, 0))
    return connect(circ, resistor_0, resistor_1)
end

begin
    pfilter_vertices = pfilter_circ.vertices[3:end-1]
    res_0_vertices = res_0_circ.vertices[2:end-1]
    res_1_vertices = res_1_circ.vertices[2:end-1]
    vertices = [pfilter_vertices; res_0_vertices; res_1_vertices]
    tree = SpanningTree(ground, [(v, ground) for v in vertices])
    port_edges = [("p/$(floor(Int, (2+pfilter_num)/2))", ground),
                  ("r0/open", ground),
                  ("r1/open", ground)]
    ports = ["pfilter_max", "res_0_max", "res_1_max"]
    mode_names = ["purcell filter", "resonator 0", "resonator 1"]
end

function resistor_pso(stub_num::Int)
    return PSOModel(resistor_circ(stub_num), port_edges, ports, tree)
end

function eigenmodes(stub_num::Int)
    pso = resistor_pso(stub_num)
    eigenvalues, eigenvectors = lossy_modes_dense(pso, min_freq=4e9, max_freq=8e9)
    port_inds = ports_to_indices(pso, ports)
    mode_inds = AdmittanceModels.match_vectors(get_P(pso)[:, port_inds], eigenvectors)
    return eigenvalues[mode_inds], eigenvectors[:,mode_inds]
end

begin
    stub_nums = run_fast ? (1:10:60) : (1:60)
    eigenmodes_arr = @showprogress [eigenmodes(i) for i in stub_nums]
    complex_freqs_arr = hcat([p[1] for p in eigenmodes_arr]...)
    eigenvectors = [p[2] for p in eigenmodes_arr]
end

if write_data
    df = DataFrame(Dict(Symbol(mode_names[i]) => complex_freqs_arr[i,:] for i in 1:3))
    df[:stub_len] = stub_nums * δ
    CSV.write(joinpath(output_folder, "complex_frequencies.csv"), df)

    df = DataFrame(transpose(hcat(eigenvectors...)))
    labels = vcat([[(i, m) for m in mode_names] for i in stub_nums]...)
    df[:stub_len] = [l[1] * δ for l in labels]
    df[:mode_name] = [l[2] for l in labels]
    CSV.write(joinpath(output_folder, "eigenvectors.csv"), df)
end

freq_func(s) = imag(s)/(2π)
decay_func(s) = -2 * real(s)/(2π)
T1_func(s) = 1/(-2 * real(s))

if show_results
    frequency_plots = [scatter(x=stub_nums * δ * 1e3, y=freq_func.(complex_freqs_arr[i, :])/1e9,
        name="$(mode_names[i])", mode="lines+markers") for i in 1:3]
    plot(frequency_plots, Layout(xaxis_title="stub length [mm]", yaxis_title="frequency [GHz]",
        title="Frequencies"))
end

if show_results
    decay_plots = [scatter(x=stub_nums * δ * 1e3, y=decay_func.(complex_freqs_arr[i, :]),
        name="$(mode_names[i])", mode="lines+markers") for i in 1:3]
    plot(decay_plots, Layout(xaxis_title="stub length [mm]", yaxis_title="decay rate [Hz]",
        yaxis_type="log", title="Decay rates"))
end

function spatial_profile(v)
    w = [0; v[1:length(pfilter_vertices)]; 0;
        v[length(pfilter_vertices)+1:length(pfilter_vertices)+length(res_0_vertices)]; 0;
        v[length(pfilter_vertices)+length(res_0_vertices)+1:end]; 0]
    return normalize(abs.(w))
end

mode_1 = run_fast ? 2 : 10
mode_2 = run_fast ? 4 : 30

function spatial_plot(mode_num::Int)
    plot([scatter(y=spatial_profile(eigenvectors[mode_1][:, mode_num]), name= "mode 1"),
          scatter(y=spatial_profile(eigenvectors[mode_2][:, mode_num]), name= "mode 2")],
         Layout(title=mode_names[mode_num]))
end

if show_results
    spatial_plot(1)
end

if show_results
    spatial_plot(2)
end

if show_results
    spatial_plot(3)
end

if write_data
    df = DataFrame(Dict(Symbol(mode_names[mode_num]) => spatial_profile(eigenvectors[mode_1][:, mode_num])
        for mode_num in 1:3))
    CSV.write(joinpath(output_folder, "mode_profile_p5.csv"), df)

    df = DataFrame(Dict(Symbol(mode_names[mode_num]) => spatial_profile(eigenvectors[mode_2][:, mode_num])
        for mode_num in 1:3))
    CSV.write(joinpath(output_folder, "mode_profile_1p5.csv"), df)
end

function regional_support(v)
    x = norm(v[1:length(pfilter_vertices)])
    y = norm(v[length(pfilter_vertices)+1:length(pfilter_vertices)+length(res_0_vertices)])
    z = norm(v[length(pfilter_vertices)+length(res_0_vertices)+1:end])
    return normalize([x, y, z]).^2
end

begin
    support_vects = [[regional_support(vects[:, i]) for vects in eigenvectors] for i in 1:3]
    standard_basis = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    hybridizations = [[norm(v - standard_basis[i],2) for v in support_vects[i]] for i in 1:3]
    get_x(v) = v[2] + v[3]/2
    get_y(v) = v[3] * sqrt(3)/2
    xs = [get_x.(support_vects[i]) for i in 1:3]
    ys = [get_y.(support_vects[i]) for i in 1:3]
end

if show_results
    plot([scatter(x=stub_nums * δ * 1e3, y=hybridizations[i], name=mode_names[i], mode="lines+markers") for i in 1:3],
        Layout(xaxis_title="stub length [mm]", yaxis_title="Distance"))
end

if show_results
    plot([scatter(x=xs[i], y=ys[i], name=mode_names[i]) for i in 1:3],
    Layout(xaxis_range=[0,1], yaxis_range=[0,sqrt(3)/2]))
end

if write_data
    df = DataFrame(Dict(Symbol(mode_names[i]) => hybridizations[i] for i in 1:3))
    df[:stub_len] = stub_nums * δ
    CSV.write(joinpath(output_folder, "hybridization_curves.csv"), df)

    df = DataFrame(Dict(Symbol(mode_names[i]) => xs[i] for i in 1:3))
    CSV.write(joinpath(output_folder, "hybridization_scatter_x.csv"), df)

    df = DataFrame(Dict(Symbol(mode_names[i]) => ys[i] for i in 1:3))
    CSV.write(joinpath(output_folder, "hybridization_scatter_y.csv"), df)
end

#######################################################
# Convergence
#######################################################

function eigenvals(stub_num::Int, discretization::Real)
    comp = components_from_model(stub_num, discretization)
    model = connect(PSOModel.(comp))
    model = short_ports(model, ["p/short_0", "p/short_1", "r0/short", "r1/short"])
    terminations = PSOModel.([ParallelComponent("in",  0, 1/Z0, 0),
                              ParallelComponent("out", 0, 1/Z0, 0)])
    pso = connect([model; terminations])
    eigenvalues, eigenvectors = lossy_modes_dense(pso, min_freq=4e9, max_freq=8e9)
    port_inds = ports_to_indices(pso, "p/middle", "r0/open", "r1/open")
    mode_inds = AdmittanceModels.match_vectors(get_P(pso)[:, port_inds], eigenvectors)
    return eigenvalues[mode_inds]
end

if !run_fast
    discretizations = [60e-6, 50e-6, 40e-6]
    eigenvals_arr = [@showprogress [eigenvals(i, d) for i in stub_nums] for d in discretizations]

    if write_data
        npzwrite(joinpath(output_folder, "convergence_eigenvalues.npy"),
            cat([hcat(eigenvals_arr[i]...) for i in 1:length(eigenvals_arr)]..., dims=3))
    end
end

#######################################################
# Declare victory
#######################################################

@testset "hybridization" begin
    @test true
end


text:
compose_demo_v2.jl

# plots a demo quantum circuit in julia with compose julia framework
# Greenberger-Horne-Zeilinger (GHZ) state with Quantum Circuit

# import Pkg; Pkg.add("Yao")
# import Pkg; Pkg.add("Compose")
# using Compose
# using Yao, Compose, Colors, Measures

module YaoPlots

export plot
using Yao, Compose, Colors, Measures

# Size of quantum circuit
set_default_graphic_size(14cm,4cm)

# square(ctx, name) where name = name of your quantum circuit
function square(ctx, name)
    compose(ctx,
        (context(), text(0.5, 0.5, name, hcenter, vcenter), font("Helvetica-Bold")),
        (context(), rectangle(), fill("transparent"), stroke("black")),
    )
end

line_position(n, k) = grid_height(n) / 2 + (k - 1) * (interval_v(n) + grid_height(n))

# connector() plots the dash lines connecting the boxes for
# :: A return type can also be specified in the function declaration using the :: operator. This converts the return value to the specified type.

function connector(ctx::Context, x0, x1, xs, ys)
    @assert length(xs) == length(ys) "number of qubits mismatch"
    out = compose(ctx); n = length(xs)
    # xs,ys = list for size of circuit
    for (x, y) in zip(xs, ys)
        compose!(out, (context(), line([(x0, line_position(n, x)), (x1, line_position(n, y))]), stroke("black")))
    end
    return out
end

# AbstractBlock = data type declaration the block representing the quantum circuit
# Currently only takes 1 qubit
plot(x::AbstractBlock; figsize=(10cm, 1cm), fontsize=10) = compose(context(0, 0, figsize...),
    connector(context(0.0, 0), 0, 0.1, 1:nqubits(x), 1:nqubits(x)),
    compose(context(0.1, 0, 0.8, 1), x),
    connector(context(0.9, 0), 0, 0.1, 1:nqubits(x), 1:nqubits(x))
)

plot(x) = compose()

# Add gates - TODO: right now only have X and H
Compose.compose(ctx::Context, x::AbstractBlock) = compose!(copy(ctx), x)
Compose.compose(ctx::Context, x::AbstractBlock) = compose!(copy(ctx), x)
Compose.compose!(ctx::Context, x::XGate) = square(ctx, "X") # Pauli X
Compose.compose!(ctx::Context, x::YGate) = square(ctx, "Y") # Pauli Y
Compose.compose!(ctx::Context, x::ZGate) = square(ctx, "Z") # Pauli Z
Compose.compose!(ctx::Context, x::HGate) = square(ctx, "H") # Hadamard H
Compose.compose!(ctx::Context, x::MGate) = square(ctx, "M") # Measurement Gate [M symbol]

# Do not work yet
# Compose.compose!(ctx::Context, x::SGate) = square(ctx, "S") # Phase(S,P)
# Compose.compose!(ctx::Context, x::TGate) = square(ctx, "T") # pi/8 (T)
# Compose.compose!(ctx::Context, x::CNOTGate) = square(ctx, "CNOT") # Controlled Not
# Compose.compose!(ctx::Context, x::CZGate) = square(ctx, "CZ") # Controlled Z
# Compose.compose!(ctx::Context, x::SWAPGate) = square(ctx, "SWAP") # SWAP
# Compose.compose!(ctx::Context, x::ToffoliGate) = square(ctx, "Toffoli") # Toffoli(CCNOT, CCX, TOFF)

grid_width(l::Int; α=0.4) = 1/(l * exp(α * (l - 1)))
interval_h(l::Int; α=0.4) = l == 1 ? 0 : (exp(α * (l - 1)) - 1) / ((l - 1) * exp(α * (l - 1)))

grid_height(n::Int; α=0.4) = 1/(n * exp(α * (n - 1)))
interval_v(n::Int; α=0.4) = n == 1 ? 0 : (exp(α * (n - 1)) - 1) / ((n - 1) * exp(α * (n - 1)))

#TODO: print a circle/dot - for CNOT gates and have it connect to different wires given two inputs

function filled_black_Circle_Node()
filled_black_circle = [compose(context(minwidth=figsize + 2mm, minheight=figsize),
                  circle(0.5, 0.5, figsize/2), fill(LCHab(0.0, 0.0, 0.0)))]

compose(context(), filled_black_circle)

function unfilled_black_Circle_Node()
unfilled_black_circle = [compose(context(minwidth=figsize + 2mm, minheight=figsize),
                    circle(0.5, 0.5, figsize/2), fill(LCHab(0.0, 0.0, 0.0)))]

compose(context(), unfilled_black_circle)

# TODO: print Ket notation for the input qubits for Q_0 to Q_n where n = number of qubits in circuit
# start at 0 or start at 1?
"""
q0_0
q0_1
.
.
.
q0_n
"""

"""
|0> q0
|0> q1
|0> q2
.
.
.
|0> qn
"""

# Chainblock is one of compose's 15 methods
function Compose.compose!(ctx::Context, blk::ChainBlock)
    x = 0.0; n = nqubits(blk); l = length(blk)
    out = compose(ctx)

    for (k, each) in enumerate(subblocks(blk))
        out = compose!(out, (context(x, 0, grid_width(l), 1), each))
        prev_x = x # store x to plot lines later
        x = x + (interval_h(l) + grid_width(l))

        if k != lastindex(blk)
            out = compose!(out, connector(context(), prev_x + grid_width(l), x, 1:n, 1:n))
        end
    end
    return out
end

# KronBlock is one of compose's 15 methods
function Compose.compose!(ctx::Context, blk::KronBlock)
    out = compose(ctx)
    y = 0; n = nqubits(blk)
    for k in 1:n
        if k in occupied_locs(blk)
            out = compose!(out, (context(0, y, 1, grid_height(n)), blk[k]))
        else
            out = compose!(out,
                (context(0, y + grid_height(n)/2), line([(0, 0), (1, 0)]), stroke("black"))
            )
        end
        y += interval_v(n) + grid_height(n)
    end
    return out
end

end # module

"""
TODO: qubit = 2 right now it only works for 1 (don't know how to start this)
TODO: try writing your own software that does this

"""

ext: jl


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

"""
    PointFolding()

A method for creating folds from a spatial object that
are single points.
"""
struct PointFolding <: FoldingMethod end

function folds(object, ::PointFolding)
  n = nelms(object)

  function pair(i)
    # source and target indices
    sinds  = setdiff(1:n, i)
    tinds  = [i]

    sinds, tinds
  end

  (pair(i) for i in 1:n)
end


text:
using Oceananigans.Fields: AbstractField

import Oceananigans.AbstractOperations: ConditionalOperation, get_condition, truefunc
import Oceananigans.Fields: condition_operand, conditional_length

#####
##### Reduction operations involving immersed boundary grids exclude the immersed periphery,
##### which includes both external nodes and nodes on the immersed interface.
#####

struct NotImmersed{F} <: Function
    func :: F
end

# ImmersedField
const IF = AbstractField{<:Any, <:Any, <:Any, <:ImmersedBoundaryGrid}

@inline condition_operand(func::Function,         op::IF, cond,      mask) = ConditionalOperation(op; func, condition=NotImmersed(cond),     mask)
@inline condition_operand(func::Function,         op::IF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersed(truefunc), mask)
@inline condition_operand(func::typeof(identity), op::IF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersed(truefunc), mask)

@inline conditional_length(c::IF)       = conditional_length(condition_operand(identity, c, nothing, 0))
@inline conditional_length(c::IF, dims) = conditional_length(condition_operand(identity, c, nothing, 0), dims)

@inline function get_condition(condition::NotImmersed, i, j, k, ibg, co::ConditionalOperation, args...)
    LX, LY, LZ = location(co)
    return get_condition(condition.func, i, j, k, ibg, args...) & !(immersed_peripheral_node(LX(), LY(), LZ(), i, j, k, ibg))
end


text:
using SemiclassicalOrthogonalPolynomials
using ClassicalOrthogonalPolynomials, ContinuumArrays, BandedMatrices, QuasiArrays, Test, LazyArrays, LinearAlgebra
import BandedMatrices: _BandedMatrix
import SemiclassicalOrthogonalPolynomials: op_lowering, RaisedOP
import ClassicalOrthogonalPolynomials: recurrencecoefficients, orthogonalityweight, symtridiagonalize, Expansion

@testset "Jacobi" begin
    P = Normalized(Legendre())
    Q = RaisedOP(P, 1)
    @test symtridiagonalize(jacobimatrix(Q))[1:10,1:10] ≈ jacobimatrix(Normalized(Jacobi(1,0)))[1:10,1:10]
    Q = RaisedOP(P, -1)
    @test symtridiagonalize(jacobimatrix(Q))[1:10,1:10] ≈ jacobimatrix(Normalized(Jacobi(0,1)))[1:10,1:10]

    L = op_lowering(P,1)
    L̃ = P \ WeightedJacobi(1,0)
    # off by scaling
    @test (L ./ L̃)[5,5] ≈ (L ./ L̃)[6,5]
    L = op_lowering(P,-1)
    L̃ = P \ WeightedJacobi(0,1)
    @test (L ./ L̃)[5,5] ≈ (L ./ L̃)[6,5]

    t = 2
    P̃ = SemiclassicalJacobi(t, 0, 0, 0)
    @test P̃ isa SemiclassicalJacobi{Float64}
    @test P̃ == SemiclassicalJacobi{Float64}(t,0,0,0)

    @test P̃[0.1,1:10] ≈ P[2*0.1-1,1:10]/P[0.1,1]
end

@testset "SemiclassicalJacobiWeight" begin
    a,b,c = 0.2,0.1,0.3
    w = SemiclassicalJacobiWeight(2,a,b,c)
    @test w[0.1] ≈ 0.1^a * (1-0.1)^b * (2-0.1)^c
    @test sum(w) ≈ 0.8387185832077594 #Mathematica
    @test Expansion(w)[0.1] ≈ w[0.1]
    @test copy(w) == w
    @test Expansion(w) == w
    @test w == Expansion(w)
end

@testset "Half-range Chebyshev" begin
    @testset "T and W" begin
        T = SemiclassicalJacobi(2, -1/2, 0, -1/2)
        W = SemiclassicalJacobi(2,  1/2, 0, -1/2, T)
        w_T = orthogonalityweight(T)
        w_W = orthogonalityweight(W)

        @test copy(T) == T
        @test W ≠ T

        X = jacobimatrix(T)
        A, B, C = recurrencecoefficients(T)
        L = T \ (SemiclassicalJacobiWeight(2,1,0,0) .* W)
        R = W \ T
        @test bandwidths(L) == (1,0)

        @time @testset "Relationship with Lanczos" begin
            P₋ = jacobi(0,-1/2,0..1)
            x = axes(P₋,1)
            y = @.(sqrt(x)*sqrt(2-x))
            T̃ = LanczosPolynomial(1 ./ y, P₋)
            @test T[0.1,1:10] ≈ T̃[0.1,1:10]/T̃[0.1,1]

            @test Normalized(T).scaling[1:10] ≈ fill(1/sqrt(sum(w_T)), 10)

            P₊ = jacobi(0,1/2,0..1)
            W̃ = LanczosPolynomial(@.(sqrt(x)/sqrt(2-x)), P₊)
            A_W̃, B_W̃, C_W̃ = recurrencecoefficients(W̃)

            @test Normalized(W)[0.1,1:10] ≈ W̃[0.1,1]W[0.1,1:10] ≈ W̃[0.1,1:10]

            # this expresses W in terms of W̃
            kᵀ = cumprod(A)
            k_W̃ = cumprod(A_W̃)
            W̄ = (x,n) -> n == 1 ? one(x) : kᵀ[n]*L[n+1,n]/(W̃[0.1,1]k_W̃[n-1]) *W̃[x,n]
            @test W̄.(0.1,1:5) ≈ W[0.1,1:5]

            @testset "x*W == T*L tells us k_n for W" begin
                @test T[0.1,1] == 1
                @test T[0.1,2] ≈ A[1]*0.1 + B[1] ≈ kᵀ[1]*0.1 + B[1]
                @test T[0.1,3] ≈ (A[2]*0.1 + B[2])*(A[1]*0.1 + B[1]) - C[2] ≈ kᵀ[2]*0.1^2 + (A[2]B[1]+B[2]A[1])*0.1 + B[2]B[1]-C[2]
                @test A[1]*L[2,1] ≈ 1
                @test 0.1 ≈ dot(T[0.1,1:2],L[1:2,1])
                @test 0.1*W̄(0.1,2) ≈ dot(T[0.1,2:3],L[2:3,2])
                @test 0.1*W̄(0.1,3) ≈ dot(T[0.1,3:4],L[3:4,3])
            end

            @testset "Jacobi operator" begin
                X_W_N = (L[1:12,1:12] \ X[1:12,1:11] * L[1:11,1:10])[1:11,:]
                @test 0.1*W̄.(0.1,1:10)' ≈ W̄.(0.1,1:11)' * X_W_N

                @test L[1:2,1:2] \ X[1:2,1:2]*L[1:2,1] ≈ X_W_N[1:2,1]
                @test L[1:3,1:3] \ X[1:3,2:3]*L[2:3,2] ≈ X_W_N[1:3,2]
                @test L[2:4,2:4] \ X[2:4,3:4]*L[3:4,3] ≈ X_W_N[2:4,3]

                x = axes(W,1)
                X_W = W \ (x .* W);
                @test X_W isa ClassicalOrthogonalPolynomials.SymTridiagonal
                @test X_W[1:11,1:10] ≈ X_W_N
            end
        end

        @testset "Evaluation" begin
            x = axes(W,1)
            A_W,B_W,C_W = recurrencecoefficients(W)
            X_W = W \ (x .* W);

            @test W[0.1,2] ≈ A_W[1]*0.1 + B_W[1]
            @test W[0.1,3] ≈ (A_W[2]*0.1 + B_W[2])*(A_W[1]*0.1 + B_W[1]) - C_W[2]

            @test W[0.1,1:11]'*X_W[1:11,1:10] ≈ 0.1 * W[0.1,1:10]'
            @test 0.1*W[0.1,1:10]' ≈ T[0.1,1:11]' * L[1:11,1:10]

            @test  W[0.1,1:3] ≈ [1,-2.029703176007828,2.2634440711109365]
        end

        @testset "Mass matrix" begin
            @test (T'*(w_T .* T))[1:10,1:10] ≈ sum(w_T)I
            M = W'*(w_W .* W)
            @test (sum(w_T)*inv(R)'L)[1:10,1:10] ≈ M[1:10,1:10]
            @test T[0.1,1:10]' ≈ W[0.1,1:10]' * R[1:10,1:10]
        end
    end

    @time @testset "T and V" begin
        T = SemiclassicalJacobi(2, -1/2, 0, -1/2)
        V = SemiclassicalJacobi(2, -1/2, 0, 1/2, T)
        w_T = orthogonalityweight(T)
        w_V = orthogonalityweight(V)
        X = jacobimatrix(T)
        A, B, C = recurrencecoefficients(T)
        L = T \ (SemiclassicalJacobiWeight(2,0,0,1) .* V);
        @test eltype(L) == Float64
        @test bandwidths(L) == (1,0)

        @testset "Relationship with Lanczos" begin
            P₋ = jacobi(0,-1/2,0..1)
            x = axes(T,1)
            Ṽ = LanczosPolynomial(@.(sqrt(2-x)/sqrt(x)), P₋)
            A_Ṽ, B_Ṽ, C_Ṽ = recurrencecoefficients(Ṽ)

            @test Normalized(V)[0.1,1:10] ≈ Ṽ[0.1,1:10]

            # this expresses W in terms of W̃
            kᵀ = cumprod(A)
            k_Ṽ = cumprod(A_Ṽ)
            V̄ = (x,n) -> n == 1 ? one(x) : -kᵀ[n]*L[n+1,n]/(Ṽ[0.1,1]k_Ṽ[n-1]) *Ṽ[x,n]

            @test V̄.(0.1,1:5) ≈ V[0.1,1:5]

            @testset "x*W == T*L tells us k_n for W" begin
                @test (2-0.1)*V̄(0.1,2) ≈ dot(T[0.1,2:3],L[2:3,2])
                @test (2-0.1)*V̄(0.1,3) ≈ dot(T[0.1,3:4],L[3:4,3])
            end

            @testset "Jacobi operator" begin
                X_V_N = (L[1:12,1:12] \ X[1:12,1:11] * L[1:11,1:10])[1:11,:]
                @test 0.1*V̄.(0.1,1:10)' ≈ V̄.(0.1,1:11)' * X_V_N

                @test L[1:2,1:2] \ X[1:2,1:2]*L[1:2,1] ≈ X_V_N[1:2,1]
                @test L[1:3,1:3] \ X[1:3,2:3]*L[2:3,2] ≈ X_V_N[1:3,2]
                @test L[2:4,2:4] \ X[2:4,3:4]*L[3:4,3] ≈ X_V_N[2:4,3]

                x = axes(V,1)
                X_V = V \ (x .* V)
                @test X_V isa ClassicalOrthogonalPolynomials.SymTridiagonal
                @test X_V[1:11,1:10] ≈ X_V_N
            end
        end

        @testset "Evaluation" begin
            x = axes(V,1)
            A_V,B_V,C_V = recurrencecoefficients(V)
            X_V = V \ (x .* V)

            @test V[0.1,2] ≈ A_V[1]*0.1 + B_V[1]
            @test V[0.1,3] ≈ (A_V[2]*0.1 + B_V[2])*(A_V[1]*0.1 + B_V[1]) - C_V[2]

            @test V[0.1,1:11]'*X_V[1:11,1:10] ≈ 0.1 * V[0.1,1:10]'
            @test (2-0.1)*V[0.1,1:10]' ≈ T[0.1,1:11]' * L[1:11,1:10]
        end

        @testset "Mass matrix" begin
            R = V \ T
            @test T[0.1,1:10]' ≈ V[0.1,1:10]' * R[1:10,1:10]
        end
    end

    @time @testset "U" begin
        T = SemiclassicalJacobi(2, -1/2, 0, -1/2)
        W = SemiclassicalJacobi(2, 1/2, 0, -1/2, T)
        V = SemiclassicalJacobi(2, -1/2, 0, 1/2, T)
        U = SemiclassicalJacobi(2, 1/2, 0, 1/2, W)

        @testset "Jacobi matrix derivation" begin
            L_1 = T \ (SemiclassicalJacobiWeight(2,0,0,1) .* V);
            L_2 = V \ (SemiclassicalJacobiWeight(2,1,0,0) .* U);

            X_V = jacobimatrix(V)
            X_U = jacobimatrix(U)
            @test (inv(L_2[1:12,1:12]) * X_V[1:12,1:11] * L_2[1:11,1:10])[1:10,1:10] ≈ X_U[1:10,1:10]

            @test 0.1*U[0.1,1:10]' ≈ V[0.1,1:11]' * L_2[1:11,1:10]
            @test 0.1 * V[0.1,1:10]'V[0.0,1:10] ≈ V[0.1,1:11]'*X_V[1:11,1:10]*V[0.0,1:10]

            L = T \ (SemiclassicalJacobiWeight(2,1,0,1) .* U)
            @test L[1:10,1:10] ≈ L_1[1:10,1:10] * L_2[1:10,1:10]
            @test (2-0.1)*0.1 * U[0.1,1:10]' ≈ T[0.1,1:12]' * L[1:12,1:10]

            L̃_1 = T \ (SemiclassicalJacobiWeight(2,1,0,0) .* W)
            L̃_3 = inv(L̃_1[1:11,1:11])*L[1:11,1:10]
            @test (2-0.1) * U[0.1,1:10]' ≈ W[0.1,1:11]' * L̃_3[1:11,1:10]
            L_3 = W \ (SemiclassicalJacobiWeight(2,0,0,1) .* U);
            @test L_3 isa ClassicalOrthogonalPolynomials.Bidiagonal
            @test L̃_3 ≈ L_3[1:11,1:10]
            @test (2-0.1) * U[0.1,1:10]' ≈ W[0.1,1:11]' * L_3[1:11,1:10]

            R = U \ T
            @test T[0.1,1:10]' ≈ U[0.1,1:10]' * R[1:10,1:10]
        end

        @testset "Evaluation" begin
            @test U[0.1,1:4] ≈ [1.1283791670955114, -2.03015300715061, 2.131688648563451, -1.3816046198529939] * sqrt(sum(SemiclassicalJacobiWeight(2, 1/2,0,1/2)))
        end
    end

    @testset "Expansions" begin
        T = SemiclassicalJacobi(2, -1/2, 0, -1/2)
        U = SemiclassicalJacobi(2, 1/2, 0, 1/2, T)
        x = axes(T,1)

        u = T * (T \ exp.(x))
        @test u[0.1] ≈ exp(0.1)

        @test T[:,1:20] \ exp.(x) ≈ u.args[2][1:20]

        u = U * (U \ exp.(x))
        @test u[0.1] ≈ exp(0.1)
        @test U[:,1:20] \ exp.(x) ≈ u.args[2][1:20]
    end

    @time @testset "Derivation" begin
        P₋ = jacobi(0,-1/2,0..1)
        P₊ = jacobi(0,1/2,0..1)
        x = axes(P₋,1)
        y = @.(sqrt(x)*sqrt(2-x))
        T = LanczosPolynomial(1 ./ y, P₋)
        W = LanczosPolynomial(@.(sqrt(x)/sqrt(2-x)), P₊)
        X = T \ (x .* T)

        @testset "Christoffel–Darboux" begin
            x,y = 0.1,0.2
            n = 10
            β = X[n,n+1]
            @test (x-y) * T[x,1:n]'*T[y,1:n] ≈ T[x,n:n+1]' * [0 -β; β 0] * T[y,n:n+1]

            # y = 0.0
            @test x * T[x,1:n]'T[0,1:n] ≈ -X[n,n+1]*(T[x,n]*T[0,n+1] - T[x,n+1]*T[0,n])

            # y = 2.0
            @test (2-x) * T[x,1:n]'*Base.unsafe_getindex(T,2,1:n) ≈ T[x,n:n+1]' * [0 β; -β 0] * Base.unsafe_getindex(T,2,n:n+1) ≈
                        β*(T[x,n]*Base.unsafe_getindex(T,2,n+1) - T[x,n+1]*Base.unsafe_getindex(T,2,n))

            @testset "T and W" begin
                W̃ = (x,n) -> -X[n,n+1]*(T[x,n]*T[0,n+1] - T[x,n+1]*T[0,n])/x
                @test norm(diff(W̃.([0.1,0.2,0.3],5) ./ W[[0.1,0.2,0.3],5])) ≤ 5E-13

                L = _BandedMatrix(Vcat((-X.ev .* T[0,2:end])', (X.ev .* T[0,1:end])'), ∞, 1, 0)
                x = 0.1
                @test x*W̃(x,1) ≈ T[x,1:2]' * L[1:2,1]
                @test (x*W̃.(x,1:10)')[:,1:9] ≈ W̃.(x,1:10)' * (L[1:10,1:10] \ X[1:10,1:10] * L[1:10,1:9])

                @test (L[1:10,1:10] \ X[1:10,1:10] * L[1:10,1:9])[2:4,3] ≈ L[2:4,2:4] \ (X*L)[2:4,3]
            end
        end
    end

    @testset "P" begin
        P = SemiclassicalJacobi(2.0,0,0,0)
        P̃ = Normalized(legendre(0..1))
        @test P̃[0.1,1:10] ≈ P[0.1,1:10]
        Q = SemiclassicalJacobi(2.0,0,0,1)
        R = Q \ P
        @test Q[0.1,1:10]' * R[1:10,1:10] ≈ P[0.1,1:10]'
        x = axes(Q,1)
        X = Q \ (x .* Q)
        @time X[1:1000,1:1000];
    end

    @testset "BigFloat" begin
        # T = SemiclassicalJacobi(2, -BigFloat(1)/2, 0, -BigFloat(1)/2)
    end
end

@testset "Legendre" begin
    t = 2
    P =   SemiclassicalJacobi(t, 0, 0, 0)
    P¹¹ = SemiclassicalJacobi(t, 1, 1, 0, P)
    Q = SemiclassicalJacobi(t, 0, 0, 1, P)
    R = SemiclassicalJacobi(t, 1, 1, 1, P)

    @time @test P[0.1,1:10] ≈ Normalized(legendre(0..1))[0.1,1:10]
    @test P¹¹[0.1,1:10] ≈ Normalized(jacobi(1,1,0..1))[0.1,1:10]/Normalized(jacobi(1,1,0..1))[0.1,1]
    x = axes(P,1)
    @time @test LanczosPolynomial(t .- x, legendre(0..1))[0.1,1:10] ≈ Q[0.1,1:10] / sqrt(sum(orthogonalityweight(Q)))

    @time @testset "expansion" begin
        @test (P * (P \ exp.(x)))[0.1] ≈ exp(0.1)
        @test (Q * (Q \ exp.(x)))[0.1] ≈ exp(0.1)
        @test (R * (R \ exp.(x)))[0.1] ≈ exp(0.1)
    end
end

@testset "Hierarchy" begin
    Ps = SemiclassicalJacobi.(2, 0:100,0,0)
    for m = 0:100
        @test jacobimatrix(Ps[m+1])[1:10,1:10] ≈ jacobimatrix(Normalized(jacobi(0,m,0..1)))[1:10,1:10]
    end

    Ps = SemiclassicalJacobi.(2, 0,0:100,0)
    for m = 0:100
        @test jacobimatrix(Ps[m+1])[1:10,1:10] ≈ jacobimatrix(Normalized(jacobi(m,0,0..1)))[1:10,1:10]
    end

    Ps = SemiclassicalJacobi.(2, 0,0,0:100)
    m = 5
    x = Inclusion(0..1)
    @test jacobimatrix(LanczosPolynomial(@. (2-x)^m))[1:10,1:10] ≈ jacobimatrix(Ps[m+1])[1:10,1:10]

    @test SemiclassicalJacobi(2,0,0,2)[0.1,1:5] ≈ SemiclassicalJacobi(2,0,0,2,Ps[1])[0.1,1:5] ≈ Ps[3][0.1,1:5]
end

@testset "Semiclassical operator asymptotics" begin
    t = 2.2
    P = SemiclassicalJacobi(t, 0, 0, 0)
    # ratio asymptotics
    φ = z -> (z + sqrt(z-1)sqrt(z+1))/2
    U = ChebyshevU()

    @testset "ratio asymptotics" begin
        n = 200;
        @test 2φ(t)*Base.unsafe_getindex(U,t,n)/(Base.unsafe_getindex(U,t,n+1)) ≈ 1
        @test 2φ(2t-1)*Base.unsafe_getindex(P,t,n)/(Base.unsafe_getindex(P,t,n+1)) ≈ 1 atol=1E-3

        L1 = P \ WeightedSemiclassicalJacobi(t,0,0,1,P)
        @test L1[n+1,n]/L1[n,n] ≈ -1/(2φ(2t-1)) atol=1E-3
    end

    @testset "single raising" begin
        R_0 = SemiclassicalJacobi(t, 1, 0, 0, P) \ P;
        R_1 = SemiclassicalJacobi(t, 0, 1, 0, P) \ P;
        R_t = SemiclassicalJacobi(t, 0, 0, 1, P) \ P;

        sqrt(0.5)
        @test R_0[999,999:1000] ≈ sqrt.([0.5,0.5]) atol=1e-2
        @test R_1[999,999:1000] ≈ [sqrt(0.5),-sqrt(0.5)] atol=1e-2
        @test R_t[200,201]/R_t[200,200] ≈ -1/(2*φ(2t-1)) atol=1e-2
    end

    @testset "T,V,W,U" begin
        T = SemiclassicalJacobi(t, -1/2, 0, -1/2)
        V = SemiclassicalJacobi(t, -1/2, 0, 1/2, T)
        U = SemiclassicalJacobi(t,  1/2, 0, 1/2, T)

        R_t = V \ T;
        n = 1000
        c = -1/(2φ(2*t-1))
        @test R_t[n,n+1]/R_t[n,n] ≈ c atol=1E-3
        R = U \ T;
        @test R[n,n+1]/R[n,n] ≈ 1+c atol=1E-3
        @test R[n,n+2]/R[n,n]  ≈ c atol=1E-3
        L =T \ (SemiclassicalJacobiWeight(t, 1, 0, 1) .* U)
        @test L[n+1,n]/L[n,n] ≈ 1+c atol=1E-3
        @test L[n+2,n]/L[n,n] ≈ c atol=1E-3
    end

    @testset "P,Q" begin
        Q = SemiclassicalJacobi(t, 1, 1, 1, P)
        R = Q \ P
        c = -1/(2*φ(2t-1))
        # (1 + c*z)*(1-z^2) == 1 + c*z - z^2 - c*z^2
        @test R[200,201]/R[200,200] ≈ c atol=1e-2
        @test R[200,202]/R[200,200] ≈ -1 atol=1e-2
        @test R[200,203]/R[200,200] ≈ -c atol=1e-2
    end
end

@testset "equilibrium measure" begin
    ρ = 0.5
    t = inv(1-ρ^2)
    P = SemiclassicalJacobi(t, -1/2, -1/2, 0)
    Q = SemiclassicalJacobi(t, -1/2, -1/2, 1/2)

    P̃ = SemiclassicalJacobi(t, 1/2, 1/2, -1);
    Q̃ = SemiclassicalJacobi(t, 1/2, 1/2, 0);

    # finite rank perturbation
    @test jacobimatrix(P)[2:5,2:5] ≈ jacobimatrix(P̃)[2:5,2:5] ≈ SymTridiagonal(fill(0.5,4), fill(0.25,3))
end

@testset "Normalized" begin
    # here we derive the formula for the Jacobi operators
    for (a,b,c) in ((-0.5,-0.4,-0.3),(0,0,0))
        P = SemiclassicalJacobi(2,a,b,c)
        Q = SemiclassicalJacobi(2,a,b,c+1,P)

        X_P = jacobimatrix(P)
        L = (P \ (SemiclassicalJacobiWeight(2,0,0,1) .* Q))
        X = jacobimatrix(Normalized(Q))
        ℓ = OrthogonalPolynomialRatio(P,2)
        a,b = X_P.dv,X_P.ev
        @test X[1,1] ≈ a[1]-b[1]ℓ[1]
        @test X[1,2] ≈ X[2,1] ≈ sqrt(b[1] * (a[1]ℓ[1]+b[1]*(1-ℓ[1]^2)-a[2]ℓ[1]))
        for n = 2:5
            @test X[n,n] ≈ ℓ[n-1]b[n-1]+a[n]-b[n]ℓ[n]
            @test X[n,n+1] ≈ X[n+1,n] ≈ sqrt(b[n] * (b[n-1]ℓ[n-1]ℓ[n]+a[n]ℓ[n]+b[n]*(1-ℓ[n]^2)-a[n+1]ℓ[n]))
        end

        T = Float64
        bl = b .* ℓ
        X̃ = ClassicalOrthogonalPolynomials.SymTridiagonal(a .- bl .+ Vcat(zero(T),bl),
                           sqrt.(bl .* a .+ b .^2 .- bl .^ 2 .- b .* a[2:∞] .* ℓ .+ Vcat(zero(T),bl .* bl[2:∞])));
        @test X̃[1:10,1:10] ≈ X[1:10,1:10]
    end
end

include("test_derivative.jl")
include("test_twoband.jl")
include("test_lowering.jl")


text:
include("load.jl")
ThreadedScansTests.runtests()


text:
using Printf
using SpecialFunctions

include("constants.jl")
include("alias.jl")
include("Atoms.jl")
include("utils.jl")
include("PGBF.jl")
include("CGBF.jl")
include("overlap.jl")
include("kinetic.jl")
include("nuclear.jl")
include("coulomb.jl")
include("sto3g.jl")
include("BasisSet.jl")

function test_two_terms()
    @printf("Calling test_two_terms\n")

    @assert fB(0,0,0,0.0,0.0,0.0,0,2.0) == 1
    @assert fB(0,0,0,1.0,1.0,1.0,0,2.0) == 1
    @assert fB(0,0,0,0.0,0.0,0.0,0,2.0 ) == 1
    @assert fB(1,0,1,0.0,0.0,0.0,0,2.0 ) == 0.125
    @assert B0(0,0,2.0) == 1
    @assert fact_ratio2(0,0) == 1
    @assert Bterm(0,0,0,0,0,0,0,0,0,0.0,0.0,0.0,0.0,0.0,0.0,2.0,2.0,0.25)==1
    @assert Bterm(0,1,0,0,0,0,0,0,1,0.0,0.0,0.0,0.0,0.0,0.0,2.0,2.0,0.25)==0

    @printf("Pass test_two_terms\n")
end

function test_coul1()
    @printf("Calling test_coul1\n")
    s = PGBF(1.0)
    px = PGBF( 1.0, (0.0,0.0,0.0),  (1,0,0) )
    @assert coulomb(s,s,s,px)==0 # 0
    @assert isapprox(coulomb(s,s,px,px), 0.9403159725793305 )
    @printf("Pass test_coul1\n")
end

test_two_terms()

test_coul1()


text:
export FFTOp,
LinearOperator,
MatrixMultiplyOp,
WeightingOp
include("FFTOp.jl")
include("LinearOperator.jl")
include("MatrixMultiplyOp.jl")
include("WeightingOp.jl")


text:
__precompile__()

module ZStd

if isfile(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))
    include(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))
else
    error("ZStd.jl is not properly installed. Please run Pkg.build(\"ZStd\") " *
          "and restart Julia.")
end

immutable ZStdError <: Exception
    msg::String
end

Base.showerror(io::IO, ex::ZStdError) = print(io, "ZStd: " * ex.msg)

# Determine whether the input represents a zstd error, yes => throw it, no => return it
function check_zstd_error(code::Csize_t)
    iserr = Bool(ccall((:ZSTD_isError, libzstd), Cuint, (Csize_t, ), code))
    if iserr
        msg = unsafe_string(ccall((:ZSTD_getErrorName, libzstd), Ptr{Cchar}, (Csize_t, ), code))
        throw(ZStdError(msg))
    end
    return code # input is not an error
end

"""
    ZStd.MAX_COMPRESSION

An integer representing the maximum compression level available.
"""
const MAX_COMPRESSION = Int(ccall((:ZSTD_maxCLevel, libzstd), Cint, ()))

"""
    maxcompressedsize(srcsize)

Get the maximum compressed size in the worst-case scenario for a given input size.
"""
function maxcompressedsize(srcsize::Csize_t)
    return ccall((:ZSTD_compressBound, libzstd), Csize_t, (Csize_t, ), srcsize)
end

maxcompressedsize(srcsize::Int) = Int(maxcompressedsize(Csize_t(srcsize)))

"""
    ZStd.ZSTD_VERSION

The version of Zstandard in use.
"""
const ZSTD_VERSION = let
    ver = Int(ccall((:ZSTD_versionNumber, libzstd), Cuint, ()))
    str = join(match(r"(\d+)(\d{2})(\d{2})$", string(ver)).captures, ".")
    VersionNumber(str)
end

end # module


text:
using MAT
using Statistics

# Labels for Amazon datasets
labels = [1; 2; 3; 12; 18; 17; 25; 15; 24]
names = ["Amazon Fashion", "All Beauty", "Appliances",
"Gift Cards", "Magazine Subscriptions", "Luxury Beauty", "Software",
"Industrial and Scientific","Prime Pantry" ]
sizes = [31; 85; 48; 148; 157; 1581; 802; 5334; 4970]
lnum = length(labels)

# See outer parameters
epsilon = 1.0
s1 = 10
s2 = 50
s3 = 200
g1 = 200
g2 = 2000
g3 = 10000
outputmat = "Output/Smallest_$s1"*"_$g1"*"_$epsilon.mat"

println("\nSeeds = ($s1, $s2, $s3), Grow by ($g1, $g2, $g3), Epsilon = $epsilon")
mat = matread(outputmat)
hl_time = round.(mean(mat["hl_time"],dims = 2),digits = 1)
hl_size= round.(mean(mat["hl_size"],dims = 2),digits = 1)
r = round.(mean(mat["r_f1"],dims = 2),digits = 2)
hl = round.(mean(mat["hl_f1"],dims = 2),digits = 2)
b1 = round.(mean(mat["b1_f1"],dims = 2),digits = 2)
b2 = round.(mean(mat["b2_f1"],dims = 2),digits = 2)
nS = round.(mean(mat["newS"],dims = 2),digits = 2)

for i = 1:size(hl_time,1)
    println(names[i]*" & $(sizes[i]) & $(hl_time[i])& $(hl[i]) & $(b1[i]) & $(b2[i])  & $(r[i]) \\\\")
end

## Get medium clusters
outputmat = "Output/Medium_$s2"*"_$g2"*"_$epsilon.mat"
mat = matread(outputmat)
hl_time = round.(mean(mat["hl_time"],dims = 2),digits = 1)
hl_size= round.(mean(mat["hl_size"],dims = 2),digits = 1)
nS = round.(mean(mat["newS"],dims = 2),digits = 2)
r = round.(mean(mat["r_f1"],dims = 2),digits = 2)
hl = round.(mean(mat["hl_f1"],dims = 2),digits = 2)
b1 = round.(mean(mat["b1_f1"],dims = 2),digits = 2)
b2 = round.(mean(mat["b2_f1"],dims = 2),digits = 2)

for i = 1:size(hl_time,1)
    println(names[i+5]*" & $(sizes[i+5]) & $(hl_time[i])& $(hl[i]) & $(b1[i]) & $(b2[i])  & $(r[i]) \\\\")
end

## Get large clusters
outputmat = "Output/Large_$s3"*"_$g3"*"_$epsilon.mat"
mat = matread(outputmat)
hl_time = round.(mean(mat["hl_time"],dims = 2),digits = 1)
hl_size= round.(Int64,mean(mat["hl_size"],dims = 1))
r = round.(mean(mat["r_f1"],dims = 2),digits = 2)
hl = round.(mean(mat["hl_f1"],dims = 2),digits = 2)
b1 = round.(mean(mat["b1_f1"],dims = 2),digits = 2)
b2 = round.(mean(mat["b2_f1"],dims = 2),digits = 2)
nS = round.(mean(mat["newS"],dims = 2),digits = 2)

for i = 1:size(hl_time,1)
    println(names[i+7]*" & $(sizes[i+7]) & $(hl_time[i]) & $(hl[i]) & $(b1[i]) & $(b2[i])  & $(r[i]) \\\\")
end

max_stars_repo_name: nveldt/HypergraphFlowClustering

max_issues_repo_name: nveldt/HypergraphFlowClustering

max_forks_repo_name: nveldt/HypergraphFlowClustering


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -19.736879506320445 <= q <= 22.217442927624916)

                     add_NL_constraint(m, :(swish(-0.3136061325511519 + 0.2821341474705279*swish(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + 0.4396178028948987*swish(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.24849525966922226*swish(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + 0.5972646324288218*swish(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + -0.8620507270647968*swish(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + swish(-0.9705953045797062 + -0.6259685253929317*swish(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + 0.41918630672941015*swish(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.9388772793667424*swish(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + -0.18849282727638839*swish(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + -0.7819718693068158*swish(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + swish(0.8285261526161638 + 0.6863822914423507*swish(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + -0.7060930690780838*swish(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.2956343589188948*swish(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + 0.32353098534638214*swish(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + 0.6963022628505415*swish(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + swish(-0.2246585184074279 + 0.8591320330102441*swish(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + -0.5037522661683527*swish(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + 0.9621050832972355*swish(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + -0.1995700954953512*swish(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + 0.442251462608958*swish(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + swish(0.5761249271362421 + -0.042724357144423575*swish(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + -0.4537006876708345*swish(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.12519521024691826*swish(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + -0.7426009177864223*swish(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + -0.8758186284143128*swish(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
using Main.SweepContractor

LTN = LabelledTensorNetwork{Char}()
LTN['A'] = Tensor(['D','B'], [i^2-2j for i=0:2, j=0:2], 0, 1)
LTN['B'] = Tensor(['A','D','C'], [-3^i*j+k for i=0:2, j=0:2, k=0:2], 0, 0)
LTN['C'] = Tensor(['B','D'], [j for i=0:2, j=0:2], 1, 0)
LTN['D'] = Tensor(['A','B','C'], [i*j*k for i=0:2, j=0:2, k=0:2], 1, 1)

brute = 0.0
for ab=1:3, ad=1:3, bc=1:3, bd=1:3, cd=1:3
    global brute += LTN['A'].arr[ad,ab] * LTN['B'].arr[ab,bd,bc] * LTN['C'].arr[bc,cd] *
        LTN['D'].arr[ad,bd,cd]
end

sweep = sweep_contract(LTN, 2, 4; fast=true)

(mps,dangexp) = sweep_contract_dangling(LTN, 2, 4; fast=true)
dang=0.0
for i1=1:2, i2=1:2, da=1:3, db=1:3, dc=1:3
    global dang+=mps[1][1,da,i1]*mps[2][i1,db,i2]*mps[3][i2,dc,1]*LTN['D'].arr[da,db,dc]
end

println("example_ABCD")
println("brute = $brute")
println("sweep = $(ldexp(sweep...))")
println("sweep_dangling = $(ldexp(dang,dangexp))")


text:
# The original license on the Mann-Whitney tests from HypothesisTests.jl
# is as follows (the filename is wrong, though). We're including it here
# because we're mostly copying the documentation of the functions.
#
# Wilcoxon.jl
# Wilcoxon rank sum (Mann-Whitney U) tests
#
# Copyright (C) 2012   Simon Kornblith
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import HypothesisTests.MannWhitneyUTest

"""
    MannWhitneyUTest(d1::AbstractUncertainValue, d2::AbstractUncertainValue,
        n::Int = 1000) -> MannWhitneyUTest

Let `s1` and `s2` be samples of `n` realisations from the distributions
furnishing the uncertain values `d1` and `d2`.

Perform a Mann-Whitney U test of the null hypothesis that the probability that
an observation drawn from the same population as `s1` is greater than an
observation drawn from the same population as `s2` is equal to the probability
that an observation drawn from the same population as `s2` is greater than an
observation drawn from the same population as `s1` against the alternative
hypothesis that these probabilities are not equal.

The Mann-Whitney U test is sometimes known as the Wilcoxon rank-sum test.
When there are no tied ranks and ≤50 samples, or tied ranks and ≤10 samples,
`MannWhitneyUTest` performs an exact Mann-Whitney U test. In all other cases,
`MannWhitneyUTest` performs an approximate Mann-Whitney U test.
"""
function MannWhitneyUTest(d1::AbstractUncertainValue, d2::AbstractUncertainValue,
        n::Int = 1000)

    x = resample(d1, n)
    y = resample(d2, n)
    MannWhitneyUTest(x, y)
end

"""
    MannWhitneyUTest(d1::UncertainDataset, d2::UncertainDataset,
        n::Int = 1000) -> MannWhitneyUTest

Let ``s_{1_{i}}`` be a sample of `n` realisations of the distribution
furnishing the uncertain value `d1[i]`, where ``i \\in [1, 2, \\ldots, N]``
and ``N`` is the number of uncertain values in `d1`.  Next, gather the
samples for all ``s_{1_{i}}`` in a pooled sample ``S_1``.  Do the same
for the second uncertain dataset `d2`, yielding the pooled sample  ``S_2``.

Perform a Mann-Whitney U test of the null hypothesis that the probability that
an observation drawn from the same population as ``S_1`` is greater than an
observation drawn from the same population as ``S_2`` is equal to the probability
that an observation drawn from the same population as ``S_2`` is greater than an
observation drawn from the same population as ``S_1`` against the alternative
hypothesis that these probabilities are not equal.

The Mann-Whitney U test is sometimes known as the Wilcoxon rank-sum test.
When there are no tied ranks and ≤50 samples, or tied ranks and ≤10 samples,
`MannWhitneyUTest` performs an exact Mann-Whitney U test. In all other cases,
`MannWhitneyUTest` performs an approximate Mann-Whitney U test.
"""
function MannWhitneyUTestPooled(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000)

    x = vcat(resample(d1, n)...,)
    y = vcat(resample(d2, n)...,)

    MannWhitneyUTest(x, y)
end

"""
    MannWhitneyUTest(d1::UncertainDataset, d2::UncertainDataset,
        n::Int = 1000) -> Vector{MannWhitneyUTest}

Assume `d1` and `d2` consist of the same number of uncertain values.
Let ``s_{1_{i}}`` be a sample of `n` realisations of the distribution
furnishing the uncertain value `d1[i]`, where ``i \\in [1, 2, \\ldots, N]``
and ``N`` is the number of uncertain values in `d1`. Let ``s_{2_{i}}`` be
the corresponding sample for `d2[i]`. This function

Perform an element-wise Mann-Whitney U test of the null hypothesis that the
probability that an observation drawn from the same population as ``s_{1_{i}}``
is greater than an observation drawn from the same population as ``s_{2_{i}}``
is equal to the probability that an observation drawn from the same population
as ``s_{2_{i}}`` is greater than an observation drawn from the same population
as ``s_{1_{i}}`` against the alternative hypothesis that these probabilities
are not equal.

The Mann-Whitney U test is sometimes known as the Wilcoxon rank-sum test.
When there are no tied ranks and ≤50 samples, or tied ranks and ≤10 samples,
`MannWhitneyUTest` performs an exact Mann-Whitney U test. In all other cases,
`MannWhitneyUTest` performs an approximate Mann-Whitney U test.
"""
function MannWhitneyUTestElementWise(d1::UncertainDataset,
        d2::UncertainDataset,
        n::Int = 1000)

    N = length(d1)
    [MannWhitneyUTest(resample(d1[i], n), resample(d2[i], n)) for i = 1:N]
end

export
MannWhitneyUTest,
MannWhitneyUTestPooled,
MannWhitneyUTestElementWise


text:
struct Grouping
  groups::Vector{Vector{Int}}
  sources::Vector{Int}
end
components(x) = vcat(x.groups...)
Base.sum(fn::Function,x::Grouping) = sum(fn,zip(x.groups,x.sources))
iterable(x::Grouping) = zip(x.groups,x.sources)

struct TrackedSources{S,F}
  sources::Vector{S}
  freqs::Vector{F}
  params::PriorTracking
end
function TrackedSources(d,params::PriorTracking)
  TrackedSources([zero(params.source_prior,d) for i in 1:params.max_sources],
                 [Beta(0.0,0.0) for i in 1:params.max_sources],params)
end

function logpdf(track::TrackedSources,C::AbstractArray,sumcomponents,
                grouping::Grouping)
  # observed, modeled sources
  logsum = sum(grouping) do (kk,i)
    observed = sumcomponents(C,kk)
    logpdf(track.params.source_prior+track.sources[i], observed) +
      logpdf(track.params.freq_prior + track.freqs[i],true)
  end

  # unobserved (but modeled) sources
  unobserved = setdiff(1:track.params.max_sources,grouping.sources)
  if !isempty(unobserved)
    logsum += sum(unobserved) do i
      logpdf(track.params.freq_prior + track.freqs[i],false)
    end
  end

  logsum
end

function mult!(track::TrackedSources,x)
  for i in 1:track.params.max_sources
    mult!(track.sources[i],x)
    track.freqs[i] *= x
  end
end

function update!(track::TrackedSources,Csource,veccomponent,grouping,w=1.0)
  for i in grouping.sources
    update!(track.sources[i],veccomponent(Csource,i),w)
    track.freqs[i] = update(track.freqs[i],true,w)
  end
  for i in setdiff(1:track.params.max_sources,grouping.sources)
    track.freqs[i] = update(track.freqs[i],false,w)
  end
end

function downdate!(track::TrackedSources,Csource,veccomponent,grouping,w=1.0)
  for i in grouping.sources
    downdate!(track.sources[i],veccomponent(Csource,i),w)
    track.freqs[i] = downdate(track.freqs[i],true,w)
  end
  for i in setdiff(1:track.params.max_sources,grouping.sources)
    track.freqs[i] = downdate(track.freqs[i],false,w)
  end
end

max_stars_repo_name: JuliaTagBot/AuditoryBistabilityLE

max_issues_repo_name: JuliaTagBot/AuditoryBistabilityLE

max_forks_repo_name: JuliaTagBot/AuditoryBistabilityLE


text:
@doc doc"""
    CholeskySpace{N} <: Manifold

The manifold of lower triangular matrices with positive diagonal and
a metric based on the cholesky decomposition. The formulae for this manifold
are for example summarized in Table 1 of [^Lin2019].

# Constructor

    CholeskySpace(n)

Generate the manifold of $n\times n$ lower triangular matrices with positive diagonal.

[^Lin2019]:
    > Lin, Zenhua: "Riemannian Geometry of Symmetric Positive Definite Matrices via
    > Cholesky Decomposition", arXiv: [1908.09326](https://arxiv.org/abs/1908.09326).
"""
struct CholeskySpace{N} <: Manifold end

CholeskySpace(n::Int) = CholeskySpace{n}()

@doc doc"""
    check_manifold_point(M::CholeskySpace, x; kwargs...)

Check whether the matrix `x` lies on the [`CholeskySpace`](@ref) `M`, i.e.
it's size fits the manifold, it is a lower triangular matrix and has positive
entries on the diagonal.
The tolerance for the tests can be set using the `kwargs...`.
"""
function check_manifold_point(M::CholeskySpace, x; kwargs...)
    if size(x) != representation_size(M)
        return DomainError(
            size(x),
            "The point $(x) does not lie on $(M), since its size is not $(representation_size(M)).",
        )
    end
    if !isapprox(norm(strictlyUpperTriangular(x)), 0.0; kwargs...)
        return DomainError(
            norm(UpperTriangular(x) - Diagonal(x)),
            "The point $(x) does not lie on $(M), since it strictly upper triangular nonzero entries",
        )
    end
    if any(diag(x) .<= 0)
        return DomainError(
            min(diag(x)...),
            "The point $(x) does not lie on $(M), since it hast nonpositive entries on the diagonal",
        )
    end
    return nothing
end

"""
    check_tangent_vector(M::CholeskySpace, x, v; kwargs... )

Checks whether `v` is a tangent vector to `x` on the [`CholeskySpace`](@ref) `M`, i.e.
atfer [`check_manifold_point`](@ref)`(M,x)`, `v` has to be of same dimension as `x`
and a symmetric matrix.
The tolerance for the tests can be set using the `kwargs...`.
"""
function check_tangent_vector(M::CholeskySpace, x, v; kwargs...)
    mpe = check_manifold_point(M, x)
    mpe !== nothing && return mpe
    if size(v) != representation_size(M)
        return DomainError(
            size(v),
            "The vector $(v) is not a tangent to a point on $(M) since its size does not match $(representation_size(M)).",
        )
    end
    if !isapprox(norm(strictlyUpperTriangular(v)), 0.0; kwargs...)
        return DomainError(
            norm(UpperTriangular(v) - Diagonal(v)),
            "The matrix $(v) is not a tangent vector at $(x) (represented as an element of the Lie algebra) since it is not lower triangular.",
        )
    end
    return nothing
end

@doc doc"""
    distance(M::CholeskySpace, x, y)

Compute the Riemannian distance on the [`CholeskySpace`](@ref) `M` between two
matrices `x`, `y` that are lower triangular with positive diagonal. The formula
reads

````math
d_{\mathcal M}(x,y) = \sqrt{\sum_{i>j} (x_{ij}-y_{ij})^2 +
\sum_{j=1}^m (\log x_{jj} - \log y_{jj})^2
}
````
"""
function distance(::CholeskySpace, x, y)
    return sqrt(
        sum((strictlyLowerTriangular(x) - strictlyLowerTriangular(y)) .^ 2) +
        sum((log.(diag(x)) - log.(diag(y))) .^ 2),
    )
end

@doc doc"""
    exp(M::CholeskySpace, x, v)

Compute the exponential map on the [`CholeskySpace`](@ref) `M` eminating from the lower
triangular matrix with positive diagonal `x` towards the lower triangular matrix `v`
The formula reads

````math
\exp_x v = \lfloor x \rfloor + \lfloor v \rfloor + \operatorname{diag}(x)
\operatorname{diag}(x)\exp\bigl( \operatorname{diag}(v)\operatorname{diag}(x)^{-1}\bigr),
````

where $\lfloor x\rfloor$ denotes the strictly lower triangular matrix of $x$ and
$\operatorname{diag}(x)$ the diagonal matrix of $x$
"""
exp(::CholeskySpace, ::Any...)

function exp!(::CholeskySpace, y, x, v)
    y .= (
        strictlyLowerTriangular(x) +
        strictlyLowerTriangular(v) +
        Diagonal(diag(x)) * Diagonal(exp.(diag(v) ./ diag(x)))
    )
    return y
end

@doc doc"""
    inner(M::CholeskySpace, x, v, w)

Compute the inner product on the [`CholeskySpace`](@ref) `M` at the
lower triangular matric with positive diagonal `x` and the two tangent vectors
`v`,`w`, i.e they are both lower triangular matrices with arbitrary diagonal.
The formula reads

````math
    g_{x}(v,w) = \sum_{i>j} v_{ij}w_{ij} + \sum_{j=1}^m v_{ii}w_{ii}x_{ii}^{-2}
````
"""
function inner(::CholeskySpace, x, v, w)
    return (
        sum(strictlyLowerTriangular(v) .* strictlyLowerTriangular(w)) +
        sum(diag(v) .* diag(w) ./ (diag(x) .^ 2))
    )
end

@doc doc"""
    log(M::CholeskySpace, v, x, y)

Compute the logarithmic map on the [`CholeskySpace`](@ref) `M` for the geodesic eminating
from the lower triangular matrix with positive diagonal `x` towards `y`.
The formula reads

````math
\log_x v = \lfloor x \rfloor - \lfloor y \rfloor
+\operatorname{diag}(x)\log\bigl(\operatorname{diag}(y)\operatorname{diag}(x)^{-1}\bigr),
````

where $\lfloor x\rfloor$ denotes the strictly lower triangular matrix of $x$ and
$\operatorname{diag}(x)$ the diagonal matrix of $x$
"""
log(::Cholesky, ::Any...)

function log!(::CholeskySpace, v, x, y)
    return copyto!(
        v,
        strictlyLowerTriangular(y) - strictlyLowerTriangular(x) +
        Diagonal(diag(x) .* log.(diag(y) ./ diag(x))),
    )
end

@doc doc"""
    manifold_dimension(M::CholeskySpace)

Return the manifold dimension for the [`CholeskySpace`](@ref) `M`, i.e. $\frac{N(N+1)}{2}$.
"""
@generated manifold_dimension(::CholeskySpace{N}) where {N} = div(N * (N + 1), 2)

@doc doc"""
    reporesentation_size(M::CholeskySpace)

Return the representation size for the [`CholeskySpace`](@ref)`{N}` `M`, i.e. `(N,N)`.
"""
@generated representation_size(::CholeskySpace{N}) where {N} = (N, N)

# two small helpers for strictly lower and upper triangulars
strictlyLowerTriangular(x) = LowerTriangular(x) - Diagonal(diag(x))

strictlyUpperTriangular(x) = UpperTriangular(x) - Diagonal(diag(x))

@doc doc"""
    vector_transport_to(M::CholeskySpace, x, v, y, ::ParallelTransport)

Parallely transport the tangent vector `v` at `x` along the geodesic to `y`
on to the [`CholeskySpace`](@ref) manifold `M`. The formula reads

````math
\mathcal P_{y\gets x}(v) = \lfloor v \rfloor
+ \operatorname{diag}(y)\operatorname{diag}(x)^{-1}\operatorname{diag}(v),
````

where $\lfloor\cdot\rfloor$ denotes the strictly lower triangular matrix,
and $\operatorname{diag}$ extracts the diagonal matrix.
"""
vector_transport_to(::CholeskySpace, ::Any, ::Any, ::Any, ::ParallelTransport)

function vector_transport_to!(::CholeskySpace, vto, x, v, y, ::ParallelTransport)
    return copyto!(
        vto,
        strictlyLowerTriangular(x) + Diagonal(diag(y) .* diag(v) ./ diag(x)),
    )
end

@doc doc"""
    zero_tangent_vector(M::CholeskySpace, x)

Return the zero tangent vector on the [`CholeskySpace`](@ref) `M` at `x`.
"""
zero_tangent_vector(::CholeskySpace, ::Any...)

zero_tangent_vector!(M::CholeskySpace, v, x) = fill!(v, 0)


text:
# Mason Protter  2 days ago
# Here's my version with updated semantics to match Base.structdiff more closely, and vararg support:
Keys(::Type{<:NamedTuple{names}}) where {names} = names
Keys(::Type{Type{T}}) where {T} = Keys(T)

@generated function structintersect(a::NamedTuple{an}, rest::Union{NamedTuple, Type{<:NamedTuple}}...) where {an}
    names = Tuple(intersect(an, Keys.(rest)...))
    data  = Expr(:tuple, (:(getproperty(a, $(QuoteNode(name)))) for name in names)...)
    :(NamedTuple{$names}($data))
end

# julia> structintersect((;a=1, b=2, c=3, d=4), (;b=1, c=3, d=4), NamedTuple{(:b, :d)})
# (b = 2, d = 4)


text:
using TropicalYao.Reversible.LogLikeNumbers

@testset "Tropical algebra" begin
	Random.seed!(4)
	@test check_inv(i_muladd, (Tropical(2.0), Tropical(3.0), Tropical(4.0), false))
	@test check_inv(i_muladd, (Tropical(2.0), Tropical(3.0), Tropical(-5.0), false))
	x = Tropical.(randn(100, 100))
	y = Tropical.(randn(100, 100))
	v = Tropical.(randn(100))
	out1 = x * y
	out2 = one.(out1)
	i_gemm!(out2, x, y)
	out4 = Tropical.(zeros(100))
	out3 = x * v
	bk = zeros(Bool, 100)
	i_gemv!(out4, x, v, bk)
	@test out1 ≈ out2
	@test out3 ≈ out4
end

@testset "i_sum" begin
	i_sum(TropicalF32(0), TropicalF32.([3.0,10.0,2.0]))[1].n == 10
end


text:
"""
    align(x::Union{Int,Ptr}, [n])

Return aligned memory address with minimum increment. `align` assumes `n` is a
power of 2.
"""
function align end
@inline align(x::Integer) = (x + Int(register_size()-One())) & Int(-register_size())
@inline align(x::Ptr{T}, arg) where {T} = reinterpret(Ptr{T}, align(reinterpret(UInt, x), arg))
@inline align(x::Ptr{T}) where {T} = reinterpret(Ptr{T}, align(reinterpret(UInt, x)))
@inline align(x::Integer, n) = (nm1 = n - One(); (x + nm1) & -n)
@inline align(x::Integer, ::StaticInt{N}) where {N} = (nm1 = N - 1; (x + nm1) & -N)
@inline align(x::Integer, ::Type{T}) where {T} = align(x, register_size() ÷ static_sizeof(T))

# @generated align(::Val{L}, ::Type{T}) where {L,T} = align(L, T)
aligntrunc(x::Integer, n) = x & -n
aligntrunc(x::Integer) = aligntrunc(x, register_size())
aligntrunc(x::Integer, ::Type{T}) where {T} = aligntrunc(x, register_size() ÷ sizeof(T))
alignment(x::Integer, N = 64) = reinterpret(Int, x) % N

function valloc(N::Integer, ::Type{T} = Float64, a = max(register_size(), cache_linesize())) where {T}
    # We want alignment to both vector and cacheline-sized boundaries
  size_T = max(1, sizeof(T))
  reinterpret(Ptr{T}, align(reinterpret(UInt,Libc.malloc(size_T*N + a - 1)), a))
end


text:
"""
    BoolMask(n::Int)

Iterate over Boolean representaion of 1:2^n - 1
"""
struct BoolMask
    N::Int
end

Base.length(BM::BoolMask) = 2^BM.N
Base.eltype(::Type{BoolMask}) = Array{Bool,1}

function Base.iterate(BM::BoolMask, state::Int=0)
    n = BM.N
    if state == 2^n
        return
    end

    mask = digits(Bool, state, base=2, pad=n)

    return (mask, state + 1)
end


text:
module Scans
import ArgParse: ArgParseSettings, parse_args, parse_item, @add_arg_table!
import Logging: @info, @warn
import Printf: @sprintf
import Base: length
import Luna: @hlock, Utils
import Pidfile: mkpidlock
import HDF5
import Distributed: @spawnat, addprocs, rmprocs, fetch, Future, @everywhere
import Dates

"""
    AbstractExec

Abstract supertype for scan execution modes.
"""
abstract type AbstractExec end

"""
    LocalExec

Execution mode to simply run the whole scan in the current Julia session in a `for` loop.
"""
struct LocalExec <: AbstractExec end

"""
    RangeExec(range)

Execution mode to run a subsection of the scan, given by a `UnitRange`, in the current Julia session.
"""
struct RangeExec <: AbstractExec
    r::UnitRange{Int}
end

"""
    BatchExec(Nbatches, batch)

Execution mode to divide the scan into `Nbatches` chunks and run only the given `batch`.
"""
struct BatchExec <: AbstractExec
    Nbatches::Int
    batch::Int
end

"""
    QueueExec(nproc=0, queuefile="")

Execution mode to run a scan using a file-based queueing system. Can be run in multiple separate
Julia sessions, or can spawn `nproc` subprocesses which then take items from the queue to run.

Possible values for `nproc` are:
- `0`: run only in the current Julia process
- `n > 0`: spawn `n` subprocesses and run on these
- `-1`: spawn as many subprocesses as the number of logical cores on the CPU
    (`Base.Sys.CPU_THREADS`)

If `queuefile` is given, the queuefile is stored at that path. If omitted, the queuefile is
stored in `Utils.cachedir()`. Note that the queuefile is deleted at the end of the scan.
"""
struct QueueExec <: AbstractExec
    nproc::Int
    queuefile::String
end

QueueExec(nproc=0) = QueueExec(nproc, "")

"""
    CondorExec(scriptfile, ncores)

Execution mode which submits a scan to an HTCondor queue system claiming `ncores` cores.

!!! note
    `scriptfile` must **always** be `@__FILE__`
"""
struct CondorExec <: AbstractExec
    scriptfile::String
    ncores::Int
end

"""
    SSHExec(localexec, scriptfile, hostname, subdir)

Execution mode which transfers the `scriptfile` file to the host given by `hostname` via SSH
and executes the scan on that host with a mode defined by `localexec`. `subdir` gives the
subdirectory (relative to the home directory) where scans are stored on the remote host. A
subfolder with automatically chosen name will be created in `subdir` to store this scan.

!!! note
    `scriptfile` must **always** be `@__FILE__`
"""
struct SSHExec{eT} <: AbstractExec
    localexec::eT
    scriptfile::String
    hostname::String
    subdir::String
end

function SSHExec(le::CondorExec, hostname, subdir)
    SSHExec(le, le.scriptfile, hostname, subdir)
end

struct Scan{eT}
    name::String
    variables::Vector{Symbol}
    arrays::Vector
    exec::eT
end

"""
    Scan(name; kwargs...)
    Scan(name, ex::AbstractExec; kwargs...)

Create a new `Scan` with name `name` and variables given as keyword arguments. The execution
mode `ex` can be given directly or via command-line arguments to the script. **If given,
command-line arguments overwrite any explicitly passed execution mode.**

If neither an explicit execution mode nor command-line arguments are given,
`ex` defaults to `LocalExec`, i.e. running the whole scan locally in the current Julia process.
"""
function Scan(name, cmdlineargs::Vector{String}=ARGS; kwargs...)
    Scan(name, makeexec(cmdlineargs); kwargs...)
end

function Scan(name, ex::AbstractExec; kwargs...)
    if !isempty(ARGS)
        cmdlineargs = copy(ARGS)
        # remove command-line arguments to avoid infinite recursion:
        [pop!(ARGS) for _ in eachindex(ARGS)]
        return Scan(name, cmdlineargs; kwargs...)
    end
    variables = Symbol[]
    arrays = Vector[]
    for (var, arr) in kwargs
        push!(variables, var)
        push!(arrays, arr)
    end
    Scan(name, variables, arrays, ex)
end

length(s::Scan) = (length(s.arrays) == 0) ? 0 : prod(length, s.arrays)

"""
    addvariable!(scan, variable::Symbol, array)
    addvariable!(scan; kwargs...)

Add scan variable(s) to the `scan`, either as a single pair of `Symbol` and array, or as a
sequence of keyword arguments.
"""
function addvariable!(scan, variable::Symbol, array)
    push!(scan.variables, variable)
    push!(scan.arrays, array)
end

function addvariable!(scan; kwargs...)
    for (var, arr) in kwargs
        push!(scan.variables, var)
        push!(scan.arrays, arr)
    end
end

"""
    makefilename(scan, scanidx)

Make an appropriate file name for an `HDF5Output` or `prop_capillary` output for the
`scan` at the current `scanidx`.

# Examples
```
scan = Scan("scan_example"; energy=collect(range(5e-6, 200e-6; length=64)))
runscan(scan) do scanidx, energyi
    prop_capillary(125e-6, 3, :HeJ, 0.8; λ0=800e-9, τfwhm=10e-15, energy=energyi
                   filepath=makefilename(scan, scanidx))
end
```
"""
makefilename(scan::Scan, scanidx) = makefilename(scan.name, scanidx)
makefilename(name::AbstractString, scanidx) = @sprintf("%s_%05d.h5", name, scanidx)

function makeexec(args::Vector{String})
    isempty(args) && return LocalExec()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--local", "-l"
            help = "Execute the whole scan locally."
            action = :store_true
        "--range", "-r"
            help = "Linear range of scan indices to execute."
            arg_type = UnitRange{Int}
        "--batch", "-b"
            help = "Number of batches and batch index to execute"
            arg_type = Tuple{Int, Int}
        "--queue", "-q"
            help = """Use a file-based queue to execute the scan. Can be run in parallel
                      using the --procs/-p option."""
            action = :store_true
        "--procs", "-p"
            help = """Number of processes to use with queued execution. If 0, use only the
                      main Julia instance. If -1, use as many processes as the machine
                      has logical cores."""
            arg_type = Int
            default = 0
    end
    args = parse_args(args, s)
    for k in keys(args)
        isnothing(args[k]) && delete!(args, k)
    end
    args["local"] && return LocalExec()
    haskey(args, "range") && return RangeExec(args["range"])
    haskey(args, "r") && return RangeExec(args["r"])
    haskey(args, "batch") && return BatchExec(args["batch"]...)
    haskey(args, "b") && return BatchExec(args["b"]...)
    args["queue"] && return QueueExec(args["procs"])
    error("Command-line arguments do not define a valid execution mode.")
end

# Enable parsing of command-line arguments of the form "1:5" to a UnitRange
parse_item(::Type{UnitRange{Int}}, x::AbstractString) = eval(Meta.parse(x))

# Enable parsing of command-line arguments of the form "1,5" to a Tuple of integers
parse_item(::Type{Tuple{Int, Int}}, x::AbstractString) = Tuple(parse(Int, xi) for xi in split(x, ","))

function logiter(scan, scanidx, args)
    logmsg = @sprintf("Running scan: %s (%d points)\nIndex: %05d\nVariables:\n",
                      scan.name, length(scan), scanidx)
    for (variable, value) in zip(scan.variables, args)
        logmsg *= @sprintf("\t%s: %g\n", variable, value)
    end
    @info logmsg
end

function getvalue(scan, variable, scanidx)
    values = vec(collect(Iterators.product(scan.arrays...)))[scanidx]
    idx = findfirst(scan.variables .== variable)
    values[idx]
end

"""
    runscan(f, scan)

Run the function `f` in a scan with arguments defined by the `scan::Scan`.
The function `f` must have the signature `f(scanidx, args...)` where the length of `args`
is the number of variables to be scanned over. Can be used with the `do` block syntax.

The exact subset and order of scan points which is run depends on `scan.exec`, see
[`Scan`](@ref).

# Examples
```
scan = Scan("scan_example"; energy=collect(range(5e-6, 200e-6; length=64)))
runscan(scan) do scanidx, energyi
    prop_capillary(125e-6, 3, :HeJ, 0.8; λ0=800e-9, τfwhm=10e-15, energy=energyi)
end
```
"""
function runscan(f, scan::Scan{LocalExec})
    for (scanidx, args) in enumerate(Iterators.product(scan.arrays...))
        logiter(scan, scanidx, args)
        try
            f(scanidx, args...)
        catch e
            bt = catch_backtrace()
            msg = "Error at scanidx $scanidx:\n"*sprint(showerror, e, bt)
            @warn msg
        end
    end
end

function runscan(f, scan::Scan{RangeExec})
    combos = vec(collect(Iterators.product(scan.arrays...)))
    for (scanidx, args) in enumerate(combos[scan.exec.r])
        logiter(scan, scanidx, args)
        try
            f(scanidx, args...)
        catch e
            bt = catch_backtrace()
            msg = "Error at scanidx $scanidx:\n"*sprint(showerror, e, bt)
            @warn msg
        end
    end
end

"""
    chunks(a::AbstractArray, n::Int)

Split array a into n chunks, spreading the entries of a evenly.

# Examples
```jldoctest
julia> a = collect(range(1, length=10));
julia> Scans.chunks(a, 3)
3-element Array{Array{Int64,1},1}:
 [1, 4, 7, 10]
 [2, 5, 8]
 [3, 6, 9]
```
"""
function chunks(a::AbstractArray, n::Int)
    N = length(a)
    done = 0
    out = [Array{eltype(a), 1}() for ii=1:n]
    while done < N
        push!(out[mod(done, n)+1], a[done+1])
        done += 1
    end
    return out
end

function runscan(f, scan::Scan{BatchExec})
    combos = vec(collect(Iterators.product(scan.arrays...)))
    linidx = collect(1:length(scan))
    chs = chunks(linidx, scan.exec.Nbatches)
    idcs = chs[scan.exec.batch]
    scanidcs_this = linidx[idcs]
    combos_this = combos[idcs]
    for (scanidx, args) in zip(scanidcs_this, combos_this)
        logiter(scan, scanidx, args)
        try
            f(scanidx, args...)
        catch e
            bt = catch_backtrace()
            msg = "Error at scanidx $scanidx:\n"*sprint(showerror, e, bt)
            @warn msg
        end
    end
end

function runscan(f, scan::Scan{QueueExec})
    if scan.exec.nproc == 0
        _runscan(f, scan)
    else
        nproc = (scan.exec.nproc == -1) ? Base.Sys.CPU_THREADS : scan.exec.nproc
        procs = addprocs(nproc)
        @everywhere eval(:(using Luna))
        futures = Future[]
        for p in procs
            fut = @spawnat p _runscan(f, scan)
            push!(futures, fut)
        end
        for fut in futures
            fetch(fut)
        end
        rmprocs(procs)
    end
end

function _runscan(f, scan::Scan{QueueExec})
    if isempty(scan.exec.queuefile)
        h = string(hash(scan.name); base=16)
        qfile = joinpath(Utils.cachedir(), "qfile_$h.h5")
    else
        qfile = scan.exec.queuefile
    end
    lockpath = joinpath(Utils.cachedir(), basename(qfile)*"_lock")

    combos = vec(collect(Iterators.product(scan.arrays...)))
    while true
        mkpidlock(lockpath) do
            # first process to catch the pidlock creates the queue file
            if ~isfile(qfile)
                @hlock HDF5.h5open(qfile, "cw") do file
                    file["qdata"] = zeros(Int, length(scan))
                end
            end
            # read the queue data
            global qdata = @hlock HDF5.h5open(qfile) do file
                read(file["qdata"])
            end
            # find the first index which is neither done nor in progress
            global scanidx = findfirst(qdata) do qi
                qi == 0
            end
            if ~isnothing(scanidx)
                # mark the index as in progress
                @hlock HDF5.h5open(qfile, "r+") do file
                    file["qdata"][scanidx] = 1
                end
            end
        end # release pidlock
        if isnothing(scanidx) # no scan points left to start
            if all(qdata .> 1) # completely done--either all done or failed
                # this point is only reached by one process
                rm(qfile) # remove the queue file
            end
            break # break out of the loop
        end
        logiter(scan, scanidx, combos[scanidx])
        code = 2 # code for finished successfully
        try
            f(scanidx, combos[scanidx]...) # run scan function
        catch e
            code = 3 # code for failed
            bt = catch_backtrace()
            msg = "Error at scanidx $scanidx:\n"*sprint(showerror, e, bt)
            @warn msg
        end
        mkpidlock(lockpath) do # acquire lock on qfile again
            @hlock HDF5.h5open(qfile, "r+") do file
                file["qdata"][scanidx] = code # mark as done/failed
            end
        end
    end
end

function runscan(f, scan::Scan{CondorExec})
    # make submission file for HTCondor
    cmd = split(string(Base.julia_cmd()))[1]
    julia = strip(cmd, ['`', '\''])
    script = scan.exec.scriptfile
    cores = scan.exec.ncores
    name = scan.name
    @info "Submitting Condor job for $script running on $cores cores."
    # Adding the --queue command-line argument below means that when running the Condor job,
    # the CondorExec is ignored even if explicitly defined inside the script.
    lines = [
        "executable = $julia",
        """arguments = "$(basename(script)) --queue" """,
        "log = $name.log.\$(Process)",
        "output = $name.out.\$(Process)",
        "error = $name.err.\$(Process)",
        "stream_error = True",
        "initialdir = $(dirname(script))",
        "request_cpus = 1",
        "queue $cores"
    ]
    subfile = joinpath(dirname(script), "doit.sub")
    @info "Writing job file to $subfile..."
    open(subfile, "w") do file
        for l in lines
            write(file, l*"\n")
        end
    end
    @info "Submitting job..."
    out = read(`condor_submit $subfile`, String)
    @info "Condor submission output:\n$out"
end

function changexec(scan, newexec)
    newscan = Scan(scan.name, newexec)
    for (var, arr) in zip(scan.variables, scan.arrays)
        addvariable!(newscan, var, arr)
    end
    newscan
end

function runscan(f, scan::Scan{<:SSHExec})
    if gethostname() == scan.exec.hostname
        # running on the machine defined in the SSH Exec? just run the scan
        runscan(f, changexec(scan, scan.exec.localexec))
    else
        # running somewhere else? submit the job via SSH
        host = scan.exec.hostname
        subdir = scan.exec.subdir
        script = scan.exec.scriptfile
        scriptfile = basename(script)
        name = scan.name
        folder = Dates.format(Dates.now(), "yyyymmdd_HHMMSS") * "_$name"
        @info "Making directory \$HOME/$subdir/$folder"
        read(`ssh $host "mkdir -p \$HOME/$subdir/$folder"`)
        @info "Transferring file..."
        read(`scp $script $host:\$HOME/$subdir/$folder`)
        @info "Running Luna script on remote host $host"
        read(`ssh $host julia \$HOME/$subdir/$folder/$scriptfile`, String)
    end
end

end


text:
"""
Compute Kummer's confluent hypergeometric function `M(a, b, z) = ₁F₁(a; b; z)`.
"""
function _₁F₁(a, b, z)
    if real(z) ≥ 0
        return _₁F₁maclaurin(a, b, z)
    else
        return exp(z)*_₁F₁(b-a, b, -z)
    end
end

"""
Compute Kummer's confluent hypergeometric function `M(a, b, z) = ₁F₁(a; b; z)`.
"""
const M = _₁F₁

"""
Compute Tricomi's confluent hypergeometric function `U(a, b, z) ∼ z⁻ᵃ ₂F₀([a, a-b+1]; []; -z⁻¹)`.
"""
function U(a, b, z)
    return z^-a*pFq([a, a-b+1], typeof(z)[], -inv(z))
end


text:
module MPIMapReduce
using MPI

include("concat.jl")

export pmapreduce
export pmapgatherv
export Cat

abstract type Operation end
struct Reduce <: Operation end
struct Concat <: Operation end

isroot(comm::MPI.Comm, root::Integer) = MPI.Comm_rank(comm) == root

function nelementsdroptake(len, np, rank)
    0 <= rank <= np - 1 || throw(ArgumentError("rank = $rank does not satisfy 0 <= rank <= $(np - 1)"))
    d, r = divrem(len, np)
    drop = d*rank + min(r, rank)
    lastind = d*(rank + 1) + min(r, rank + 1)
    take = lastind - drop
    drop, take
end

_reduce(x, op, root, comm) = MPI.Reduce(x, op, root, comm)

function _vbuffer(x, counts, sizes, c::Cat)
    rcvbuf = _similar_cat(c, x, sizes, counts)
    MPI.VBuffer(rcvbuf, counts)
end

# Return the sizes of the arrays on each worker
# These may differ, so we do this in two steps
# At the first step fetch the ndims from all workers to estimate the number of values to read
# Given this, read the correct number of values from each worker

# Treat numbers like 1-elements arrays in Gather
_sizevec(x::AbstractArray) = collect(size(x))
_sizevec(x::Number) = 1
_ndims(x::AbstractArray) = ndims(x)
_ndims(x::Number) = 1

function _sizes(x, root, comm)
    np = MPI.Comm_size(comm)

    dims = MPI.Gather(_ndims(x), root, comm)
    size_local = _sizevec(x)

    sz = if isroot(comm, root)
        vbuf =  _vbuffer(size_local, dims, nothing, Cat(Val(1)))
        MPI.Gatherv!(size_local, vbuf, root, comm)
    else
        MPI.Gatherv!(size_local, nothing, root, comm)
    end

    if isroot(comm, root)
        sizes = Vector{Vector{Int}}(undef, np)
        offset = 0
        for (ind, nd) in enumerate(dims)
            sizes[ind] = sz[offset .+ (1:nd)]
            offset += nd
        end
        return map(Tuple, sizes)
    end
    return nothing
end

function __reshapedviews(xg, offset, sz)
    inds = offset .+ (1:prod(sz))
    v = @view xg[inds]
    reshape(v, sz...)
end
function _reshapedviews(xg, sizes, counts)
    offsets = cumsum(reduce(vcat, [[zero(eltype(counts))], counts[1:end-1]]), dims = 1)
    [__reshapedviews(xg, offset, sz) for (offset, sz) in zip(offsets, sizes)]
end

# _checkcontiguity(sizes, ::Val{N}) where {N} = _checkcontiguity(sizes, N)
_checkcontiguity(sizes, dims::Integer) = all(x -> length(x) <= dims, sizes)
_checkcontiguity(sizes, dims) = false

_gatherv(x, op, root, comm) = __gatherv(x, op, root, comm, dims(op))

function __gatherv(x, op, root, comm, dims)
    sizes = _sizes(x, root, comm)
    np = MPI.Comm_size(comm)

    if isroot(comm, root)
        counts = map(prod, sizes)
        vbuf =  _vbuffer(x, counts, sizes, Cat(dims))
        xg = MPI.Gatherv!(x, vbuf, root, comm)

        contiguous = _checkcontiguity(sizes, dims)
        if !contiguous
            X = _reshapedviews(xg, sizes, counts)
            xg = reduce(op, X)
        end
        xg
    else
        MPI.Gatherv!(x, nothing, root, comm)
    end
end

function _split_iterators(itzip, Niter, comm)
    rank = MPI.Comm_rank(comm)
    np = MPI.Comm_size(comm)

    drop, take = nelementsdroptake(Niter, np, rank)
    # split the iterators into parts
    it_local = Iterators.take(Iterators.drop(itzip, drop), take)

    np_mapreduce = min(Niter, np)
    emptyiter = isempty(it_local)

    it_local, np_mapreduce, np == np_mapreduce, emptyiter
end

struct Elementwise{F}
    f :: F
end

(o::Elementwise)(x, y) = o.f.(x, y)
(o::Elementwise)(x::Array, y) = x .= o.f.(x, y)
(o::Elementwise)(x, y::Array) = y .= o.f.(x, y)
(o::Elementwise)(x::Array, y::Array) = x .= o.f.(x, y)

_mapreduce_zipsection(f, op, it) = mapreduce(f, Elementwise(op), it)
function _mapreduce_zipsection(f, op::MPI.Op, it)
    m = map(f, it)
    length(m) == 1 || throw(ArgumentError("more than one value returned"))
    _reduce(first(m), op, 0, MPI.COMM_SELF)
end

function _mapreduce_local_zipsection(::Reduce, f, op, it)
    isempty(it) && return nothing
    _mapreduce_zipsection(x -> f(x...), op, it)
end
function _mapreduce_local_zipsection(::Concat, f, op, it)
    isempty(it) && return nothing
    mapreduce(x -> f(x...), op, it)
end

function _collectroot(reducefn, m, op, comm, root, emptyiter, allactive)
    # If there are more processes than elements in the iterator, then some processes might
    # not be involved in the mapreduce. Split the communicator to exclude these processes
    # from the final reduction.
    comm_reduce = allactive ? comm : MPI.Comm_split(comm, Int(emptyiter), MPI.Comm_rank(comm))

    np = MPI.Comm_size(comm_reduce)
    if 0 <= root < np
        # Carry out the reduction only on processes that contain some elements
        if !emptyiter
            return reducefn(m, op, root, comm_reduce)
        else
            return nothing
        end
    else
        # In this case the iterator on root is empty, so no reduction happens there
        # We carry out the reduction on comm_reduce and transfer the result to root
        s = !emptyiter ? reducefn(m, op, 0, comm_reduce) : nothing
        tag = 32 # arbitrary
        if MPI.Comm_rank(comm) == 0
            MPI.send(s, root, tag, comm)
            return nothing
        elseif isroot(comm, root)
            r, status = MPI.recv(0, tag, comm)
            if status.source == 0 && status.tag == tag
                return r
            else
                error("unable to send the result to root")
            end
        end
    end
    return nothing
end

function _pmapreduce_local(f, op, iterators...; root::Integer = 0, comm::MPI.Comm = MPI.COMM_WORLD)
    if !(0 <= root < MPI.Comm_size(comm))
        throw(ArgumentError("root = $root does not satisfy  0 <= root < $(MPI.Comm_size(comm))"))
    end

    itzip = zip(iterators...)
    Niter = length(itzip)

    if Niter == 0
        throw(ArgumentError("reducing over an empty collection is not allowed"))
    end

    _split_iterators(itzip, Niter, comm)
end

# __mapreduce_singleprocess is almost identical to _mapreduce_local_zipsection, but we avoid using the zipped iterator
# This is because mapreduce(f, hcat, itrs...) and reduce(hcat, map(f, itrs...)) currently differ for single element arrays
# See issue https://github.com/JuliaLang/julia/issues/37917
__mapreduce_singleprocess(::Reduce, f, op, iterators...) = _mapreduce_zipsection(x -> f(x...), op, zip(iterators...))
__mapreduce_singleprocess(::Concat, f, op, iterators...) = mapreduce(f, op, iterators...)

function _mapreduce_singleprocess(m::Operation, root, comm, np_mapreduce, f, op, iterators...)
    # Either 1 process or 1 element in the iterator
    if isroot(comm, root)
        return __mapreduce_singleprocess(m, f, op, iterators...)
    else
        return nothing
    end
end

"""
    pmapreduce(f, op, iterators...; [root = 0], [comm = MPI.COMM_WORLD])

Apply function `f` to each element(s) in `iterators`, and then reduce the result using the elementwise
binary reduction operator `op`.
Both the `map` and the reduction are evaluated in parallel over the processes corresponding to the communicator `comm`.
The result of the operation is returned at `root`, while `nothing` is returned at the other processes.

The result of `pmapreduce` is equivalent to that of `mapreduce(f, (x, y) -> op.(x, y), iterators...)`.

!!! note
    Unlike the standard `mapreduce` operation in Julia, this operation is performed elementwise on the arrays returned
    from the various processes. The returned value must be compatible with the elementwise operation.

!!! note
    MPI reduction operators require each worker to return only one array with an eltype `T` that
    satisfies `isbitstype(T) == true`. Such a limitation does not exist for Julia operators.
"""
function pmapreduce(f, op, iterators...; root::Integer = 0, comm::MPI.Comm = MPI.COMM_WORLD)
    it, np_mapreduce, allactive, emptyiter  = _pmapreduce_local(f, op, iterators...; root = root, comm = comm)
    if np_mapreduce == 1
        return _mapreduce_singleprocess(Reduce(), root, comm, np_mapreduce, f, op, iterators...)
    end
    m = _mapreduce_local_zipsection(Reduce(), f, op, it)
    _collectroot(_reduce, m, op, comm, root, emptyiter, allactive)
end

"""
    pmapgatherv(f, op, iterators...; [root = 0], [comm = MPI.COMM_WORLD])

Apply function `f` to each element(s) in `iterators`, and then reduce the result using the concatenation operator `op`.
Both the `map` and the reduction are evaluated in parallel over the processes corresponding to the communicator `comm`.
The result of the operation is returned at `root`, while `nothing` is returned at the other processes.

The result of `pmapgatherv` is equivalent to that of `mapreduce(f, op, iterators...)`.

The concatenation operator `op` may be one of `vcat`, `hcat` and [`Cat`](@ref).

!!! warn
    Anonymous functions are not supported as the concatenation operator.
"""
function pmapgatherv(f, op, iterators...; root::Integer = 0, comm::MPI.Comm = MPI.COMM_WORLD)
    it, np_mapreduce, allactive, emptyiter  = _pmapreduce_local(f, op, iterators...; root = root, comm = comm)
    if np_mapreduce == 1
        return _mapreduce_singleprocess(Concat(), root, comm, np_mapreduce, f, op, iterators...)
    end
    m = _mapreduce_local_zipsection(Concat(), f, op, it)
    _collectroot(_gatherv, m, op, comm, root, emptyiter, allactive)
end

end


text:
using Pkg

Pkg.add(url="https://github.com/Qaintum/Qaintessent.jl", rev="master")


text:
module Day7

using ..Utils
using FFTW

pad_histogram(histogram) = begin
    len = length(histogram)
    # Make the length the next power of two which is at least 3*len
    padded_len = 2^(floor(log2(3*len)) + 1) |> Int
    padded = vcat(zeros(len), histogram, zeros(padded_len - 2len))
end

convolution_with(histogram, f) = begin
    len = length(histogram)
    padded = pad_histogram(histogram)
    padded_len = length(padded)

    abs_array = (Array(1:padded_len) .- padded_len / 2) .|> f

    conv = irfft(rfft(abs_array) .* rfft(padded), padded_len) |> fftshift
    conv = conv[len:2len-1]
    conv
end

function best_position(input, cost)
    histogram = comma_separated_to_histogram(input)
    conv = convolution_with(histogram, cost)
    pos = argmin(conv) - 1
    best = minimum(conv)
    return (pos, Int(round(best)))
end

part2_cost(n) = Int(round(abs(n) * (abs(n)+1) / 2))

function run()
    test = "16,1,2,0,4,2,7,1,2,14"
    @assert best_position(test, abs) == (2, 37)

    input = read("resources/day7_input.txt", String)
    @time @show best_position(input, abs)

    @assert best_position(test, part2_cost) == (5, 168)
    @time @show best_position(input, part2_cost)
end

end