text:
# Atkinson.jl

###### atkinson #####
"""
    atkinson(v, ϵ)

Compute the Atkinson Index of a vector `v` at a specified inequality
aversion parameter `ϵ`.

# Examples
```julia
julia> using Inequality
julia> atkinson([8, 5, 1, 3, 5, 6, 7, 6, 3], 1.2)
0.1631765870035865
```
"""
function atkinson(v::AbstractVector{<:Real}, ϵ::Real)::Float64

    ϵ >= 0 ? nothing : throw(ArgumentError("`ϵ` must be larger or equal than 0"))

    if ϵ == 1
        v = v[v .!= 0]
        return 1 - (exp.(mean(log.(v)))/mean(v))
    elseif ϵ < 1
        return 1 - mean( (v./mean(v)).^(1-ϵ) ).^(1/(1-ϵ))
    else
        v = v[v .!= 0]
        return 1 - mean( (v./mean(v)).^(1-ϵ) ).^(1/(1-ϵ))
    end
end

###### weighted atkinson #####
"""
    atkinson(v, w, ϵ)

Compute the weighted Atkinson Index of a vector `v` at a specified inequality
aversion parameter `ϵ`, using weights given by a weight vector `w`.

Weights must not be negative, missing or NaN. The weights and data vectors must have the same length.

# Examples
```julia
julia> using Inequality
julia> atkinson([8, 5, 1, 3], [0.1,0.5,0.3,0.8], 1.2)
0.1681319821792493
```
"""
function atkinson(v::AbstractVector{<:Real}, w::AbstractVector{<:Real}, ϵ::Real)::Float64

    ϵ >= 0 ? nothing : throw(ArgumentError("`ϵ` must be larger or equal than 0"))
    checks_weights(v, w)

    v = v/Statistics.mean(v)
    w = w/sum(w)

    if ϵ == 1
        w = w[v .!= 0]
        v = v[v .!= 0]

        return 1 - (prod(exp.(w .* log.(v)))/sum(v .* w) )
    elseif ϵ < 1
        return 1-(sum(((v/sum(v.* w)).^(1-ϵ)).*w)).^(1/(1-ϵ))
    else
        w = w[v .!= 0]
        v = v[v .!= 0]

        return 1-(sum(((v/sum(v.* w)).^(1-ϵ)).*w)).^(1/(1-ϵ))
    end
end

function atkinson(v::AbstractVector{<:Real}, w::AbstractWeights, ϵ::Real)::Float64

    ϵ >= 0 ? nothing : throw(ArgumentError("`ϵ` must be larger or equal than 0"))
    checks_weights(v, w)

    v = v/Statistics.mean(v)
    w = w/w.sum

    if ϵ == 1
        w = w[v .!= 0]
        v = v[v .!= 0]

        return 1 - (prod(exp.(w .* log.(v)))/sum(v .* w) )
    elseif ϵ < 1
        return 1-(sum(((v/sum(v.* w)).^(1-ϵ)).*w)).^(1/(1-ϵ))
    else
        w = w[v .!= 0]
        v = v[v .!= 0]

        return 1-(sum(((v/sum(v.* w)).^(1-ϵ)).*w)).^(1/(1-ϵ))
    end
end

"""
    watkinson(v, w, p)

Compute the atkinson index of `v` with weights `w` and inequality aversion parameter 'ϵ'. See also [`atkinson`](@atkinson)
"""
watkinson(v::AbstractVector{<:Real}, w::AbstractVector{<:Real}, ϵ::Real) = atkinson(v, w, ϵ)

watkinson(v::AbstractVector{<:Real}, w::AbstractWeights, ϵ::Real) = atkinson(v, w, ϵ)


text:
export Bernoulli

import Distributions: Categorical, probs

vague(::Type{ <: Categorical }, dims::Int) = Categorical(ones(dims) ./ dims)

prod_analytical_rule(::Type{ <: Categorical }, ::Type{ <: Categorical }) = ProdAnalyticalRuleAvailable()

function prod(::ProdAnalytical, left::Categorical, right::Categorical)
    # Multiplication of 2 categorical PMFs: p(z) = p(x) * p(y)
    mvec = clamp.(probvec(left) .* probvec(right), tiny, huge)
    norm = sum(mvec)
    return Categorical(mvec ./ norm)
end

probvec(dist::Categorical) = probs(dist)

ext: jl


text:
__precompile__(true)
module Makie

using Colors, GeometryTypes, GLVisualize, GLAbstraction, ColorVectorSpace
using StaticArrays, GLWindow, ModernGL, Contour, Quaternions

using Base.Iterators: repeated, drop
using Base: RefValue

struct Backend{B} end

include("plotutils/utils.jl")

include("plotsbase/scene.jl")
include("plotsbase/conversions.jl")
include("plotutils/units.jl")

const makie = Scene{:makie}

include("plotutils/layout.jl")

include("plotsbase/atomics.jl")
    # The actual implementation
    include("atomics/shared.jl")
    include("atomics/scatter.jl")
    include("atomics/lines.jl")
    include("atomics/text.jl")
    include("atomics/surface.jl")
    include("atomics/wireframe.jl")
    include("atomics/mesh.jl")
    include("atomics/imagelike.jl")
    include("plotsbase/contour.jl")
    include("plotsbase/legend.jl")

include("plotsbase/axis.jl")
include("plotsbase/output.jl")
include("iodevices.jl")
include("camera2d.jl")
# include("camera3d.jl")

export Scene, Node

export scatter, lines, linesegment, mesh, surface, wireframe, axis
export @ref, @theme, @default, to_node, to_value, lift_node, to_world, save
export available_marker_symbols, available_gradients, render_frame

# conversion

export to_float, to_markersize2d, to_spritemarker, to_linestyle, to_pattern
export to_color, to_colormap, to_colornorm, to_array, to_mesh, to_surface

export to_scale
export to_offset
export to_rotation
export to_image
export to_bool
export to_index_buffer
export to_index_buffer
export to_positions
export to_position
export to_array
export to_scalefunc
export to_text
export to_font
export to_intensity
export to_surface
export to_spritemarker
export to_static_vec
export to_rotations
export to_markersize2d
export to_markersize3d
export to_linestyle
export to_normals
export to_faces
export to_attribut_id
export to_mesh
export to_float
export to_color
export to_colornorm
export to_colormap
export available_gradients
export to_spatial_order
export to_interval
export to_volume_algorithm
export to_3floats
export to_2floats
export to_textalign

end # module


text:
function test_complex_1()

    Random.seed!(0)
    f = x -> [exp(-30 / x), exp(-30 / x)]
    xtrue = 100
    y = f(xtrue)
    T = 1
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    kernel = GaussianKernel(λ * y)
    ρ = 2.0^-20
    xhat = perk(reshape(y, :, 1), T, xDists, noiseDist, signalModels, kernel, ρ)
    error_rel_real = abs(xhat[] - xtrue) / xtrue

    Random.seed!(0)
    f = x -> complex(exp(-30 / x), exp(-30 / x))
    xtrue = 100
    y = f(xtrue)
    T = 1
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    kernel = GaussianKernel([λ * mean(y)])
    ρ = 2.0^-20
    xhat = perk(y, T, xDists, noiseDist, signalModels, kernel, ρ)
    error_rel_complex = abs(xhat[] - xtrue) / xtrue

    return error_rel_complex ≈ error_rel_real

end

function test_complex_2()

    rng = StableRNG(0)
    f = x -> complex(exp(-30 / x), exp(-30 / x))
    xtrue = 100
    y = f(xtrue)
    T = 200
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    kernel = GaussianKernel([λ * mean(y)])
    ρ = 2.0^-20
    xhat = perk(rng, y, T, xDists, noiseDist, signalModels, kernel, ρ)

    error_rel = abs(xhat[] - xtrue) / xtrue
    return isapprox(error_rel, 0.04464689602051635, atol = 1e-7)

end

function test_complex_3()

    Random.seed!(0)
    f = (x, ν) -> [exp(-ν / x), exp(-ν / x)]
    xtrue = 100
    ν = 30
    N = 1
    y = reshape(f(xtrue, ν), :, N)
    T = 1
    xDists = [Uniform(10, 500)]
    νDists = Uniform(30, 30.000000001)
    noiseDist = Normal(0, 0.01)
    signalModels = [f]
    λ = 2.0^-1.5
    kernel = GaussianKernel(λ * vec(y), [λ * mean(ν)])
    ρ = 2.0^-20
    xhat = perk(y, ν, T, xDists, νDists, noiseDist, signalModels, kernel, ρ)
    error_rel_real = abs(xhat[] - xtrue) / xtrue

    Random.seed!(0)
    f = (x, ν) -> complex(exp(-ν / x), exp(-ν / x))
    xtrue = 100
    ν = 30
    N = 1
    y = fill(f(xtrue, ν), 1, N)
    T = 1
    xDists = [Uniform(10, 500)]
    νDists = Uniform(30, 30.000000001)
    noiseDist = Normal(0, 0.01)
    signalModels = [f]
    λ = 2.0^-1.5
    kernel = GaussianKernel([λ * mean(y)], [λ * mean(ν)])
    ρ = 2.0^-20
    xhat = perk(y, ν, T, xDists, νDists, noiseDist, signalModels, kernel, ρ)
    error_rel_complex = abs(xhat[] - xtrue) / xtrue

    return error_rel_complex ≈ error_rel_real

end

function test_complex_4()

    rng = StableRNG(0)
    f = (x, ν) -> complex(exp(-ν / x), exp(-ν / x))
    xtrue = 100
    ν = 30
    N = 1
    y = fill(f(xtrue, ν), 1, N)
    T = 200
    xDists = [Uniform(10, 500)]
    νDists = Uniform(30, 30.000000001)
    noiseDist = Normal(0, 0.01)
    signalModels = [f]
    λ = 2.0^-1.5
    kernel = GaussianKernel([λ * mean(y)], [λ * mean(ν)])
    ρ = 2.0^-20
    xhat = perk(rng, y, ν, T, xDists, νDists, noiseDist, signalModels, kernel, ρ)

    error_rel = abs(xhat[] - xtrue) / xtrue
    return isapprox(error_rel, 0.05111009427971681, atol = 1e-7)

end

function test_complex_5()

    Random.seed!(0)
    f = x -> [exp(-30 / x), exp(-30 / x)]
    xtrue = 100
    y = f(xtrue)
    T = 1
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    H = 2
    kernel = GaussianRFF(H, λ * y)
    ρ = 2.0^-20
    xhat = perk(reshape(y, :, 1), T, xDists, noiseDist, signalModels, kernel, ρ)
    error_rel_real = abs(xhat[] - xtrue) / xtrue

    Random.seed!(0)
    f = x -> complex(exp(-30 / x), exp(-30 / x))
    xtrue = 100
    y = f(xtrue)
    T = 1
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    H = 2
    kernel = GaussianRFF(H, [λ * mean(y)])
    ρ = 2.0^-20
    xhat = perk(y, T, xDists, noiseDist, signalModels, kernel, ρ)
    error_rel_complex = abs(xhat[] - xtrue) / xtrue

    return error_rel_complex ≈ error_rel_real

end

function test_complex_6()

    rng = StableRNG(0)
    f = x -> complex(exp(-30 / x), exp(-30 / x))
    xtrue = 100
    y = f(xtrue)
    T = 200
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    H = 40
    kernel = GaussianRFF(H, [λ * mean(y)])
    ρ = 2.0^-20
    xhat = perk(rng, y, T, xDists, noiseDist, signalModels, kernel, ρ)

    error_rel = abs(xhat[] - xtrue) / xtrue
    return isapprox(error_rel, 0.02231978386441426, atol = 1e-7)

end

function test_complex_7()

    Random.seed!(0)
    f = (x, ν) -> [exp(-ν / x), exp(-ν / x)]
    xtrue = 100
    ν = 30
    N = 1
    y = reshape(f(xtrue, ν), :, N)
    T = 1
    xDists = [Uniform(10, 500)]
    νDists = Uniform(30, 30.000000001)
    noiseDist = Normal(0, 0.01)
    signalModels = [f]
    λ = 2.0^-1.5
    H = 2
    kernel = GaussianRFF(H, λ * vec(y), [λ * mean(ν)])
    ρ = 2.0^-20
    xhat = perk(y, ν, T, xDists, νDists, noiseDist, signalModels, kernel, ρ)
    error_rel_real = abs(xhat[] - xtrue) / xtrue

    Random.seed!(0)
    f = (x, ν) -> complex(exp(-ν / x), exp(-ν / x))
    xtrue = 100
    ν = 30
    N = 1
    y = fill(f(xtrue, ν), 1, N)
    T = 1
    xDists = [Uniform(10, 500)]
    νDists = Uniform(30, 30.000000001)
    noiseDist = Normal(0, 0.01)
    signalModels = [f]
    λ = 2.0^-1.5
    H = 2
    kernel = GaussianRFF(H, [λ * mean(y)], [λ * mean(ν)])
    ρ = 2.0^-20
    xhat = perk(y, ν, T, xDists, νDists, noiseDist, signalModels, kernel, ρ)
    error_rel_complex = abs(xhat[] - xtrue) / xtrue

    return error_rel_complex ≈ error_rel_real

end

function test_complex_8()

    rng = StableRNG(0)
    f = (x, ν) -> complex(exp(-ν / x), exp(-ν / x))
    xtrue = 100
    ν = 30
    N = 1
    y = fill(f(xtrue, ν), 1, N)
    T = 200
    xDists = [Uniform(10, 500)]
    νDists = Uniform(30, 30.000000001)
    noiseDist = Normal(0, 0.01)
    signalModels = [f]
    λ = 2.0^-1.5
    H = 40
    kernel = GaussianRFF(H, [λ * mean(y)], [λ * mean(ν)])
    ρ = 2.0^-20
    xhat = perk(rng, y, ν, T, xDists, νDists, noiseDist, signalModels, kernel, ρ)

    error_rel = abs(xhat[] - xtrue) / xtrue
    return isapprox(error_rel, 0.06272802910855091, atol = 1e-7)

end

@testset "Complex PERK" begin

    @test test_complex_1()
    @test test_complex_2()
    @test test_complex_3()
    @test test_complex_4()
    @test test_complex_5()
    @test test_complex_6()
    @test test_complex_7()
    @test test_complex_8()

end


text:
using Oceananigans.Architectures
using Oceananigans.BoundaryConditions
using Oceananigans.TurbulenceClosures: calculate_diffusivities!
using Oceananigans.ImmersedBoundaries: mask_immersed_field!
using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!

import Oceananigans.TimeSteppers: update_state!

"""
    update_state!(model::HydrostaticFreeSurfaceModel)

Update peripheral aspects of the model (auxiliary fields, halo regions, diffusivities,
hydrostatic pressure) to the current model state.
"""
function update_state!(model::HydrostaticFreeSurfaceModel)

    # Mask immersed fields
    masking_events = Tuple(mask_immersed_field!(field) for field in merge(model.auxiliary_fields, prognostic_fields(model)))

    wait(device(model.architecture), MultiEvent(masking_events))

    # Fill halos for velocities and tracers. On the CubedSphere, the halo filling for velocity fields is wrong.
    fill_halo_regions!(prognostic_fields(model), model.architecture, model.clock, fields(model))

    # This _refills_ the halos for horizontal velocity fields when grid::ConformalCubedSphereGrid
    # For every other type of grid, fill_horizontal_velocity_halos! does nothing.
    fill_horizontal_velocity_halos!(model.velocities.u, model.velocities.v, model.architecture)

    compute_w_from_continuity!(model)

    fill_halo_regions!(model.velocities.w, model.architecture, model.clock, fields(model))

    # Compute auxiliary fields
    for aux_field in model.auxiliary_fields
        compute!(aux_field)
    end

    # Calculate diffusivities
    calculate_diffusivities!(model.diffusivities, model.architecture, model.grid, model.closure,
                             model.buoyancy, model.velocities, model.tracers)

    fill_halo_regions!(model.diffusivities, model.architecture, model.clock, fields(model))

    # Calculate hydrostatic pressure
    pressure_calculation = launch!(model.architecture, model.grid, Val(:xy), update_hydrostatic_pressure!,
                                   model.pressure.pHY′, model.grid, model.buoyancy, model.tracers,
                                   dependencies=Event(device(model.architecture)))

    # Fill halo regions for pressure
    wait(device(model.architecture), pressure_calculation)

    fill_halo_regions!(model.pressure.pHY′, model.architecture)

    return nothing
end


text:
using Test, OMEinsumContractionOrders, OMEinsum

@testset "save load" begin
    for code in [
        EinCode([[1,2], [2,3], [3,4]], [1,4]),
        EinCode([['a','b'], ['b','c'], ['c','d']], ['a','d'])
    ]
        for optcode in [optimize_code(code, uniformsize(code, 2), GreedyMethod()),
            optimize_code(code, uniformsize(code, 2), TreeSA(nslices=1))]
            filename = tempname()
            writejson(filename, optcode)
            code2 = readjson(filename)
            @test optcode == code2
        end
    end
end


text:
module TestExamples2D

using Test
using Trixi

include("test_trixi.jl")

# Start with a clean environment: remove Trixi output directory if it exists
outdir = "out"
isdir(outdir) && rm(outdir, recursive=true)

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
const EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "2d")

# Run basic tests
@testset "Examples 2D" begin
  @testset "parameters.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [9.144681765639205e-6],
            linf = [6.437440532547356e-5])
  end
  @testset "parameters.toml with polydeg=1" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [0.05264106093598111],
            linf = [0.08754218386076518],
            polydeg=1)
  end
  @testset "parameters.toml with carpenter_kennedy_erk43" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [8.908962577028364e-6],
            linf = [6.969419032576418e-5],
            time_integration_scheme = "timestep_carpenter_kennedy_erk43_2N!",
            cfl = 0.5)
  end
  @testset "parameters.toml with parsani_ketcheson_deconinck_erk94" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [7.932405161658336e-6],
            linf = [6.509399993848142e-5],
            time_integration_scheme = "timestep_parsani_ketcheson_deconinck_erk94_3Sstar!")
  end
  @testset "parameters.toml with parsani_ketcheson_deconinck_erk32" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [0.00440542760645958],
            linf = [0.012549162970726613],
            time_integration_scheme = "timestep_parsani_ketcheson_deconinck_erk32_3Sstar!")
  end
  @testset "parameters_alfven_wave.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_alfven_wave.toml"),
            l2   = [0.00011134513490658689, 5.880188909157728e-6, 5.880188909159547e-6, 8.432880997656317e-6, 1.2942387343501909e-6, 1.2238820298971968e-6, 1.2238820298896402e-6, 1.830621754702352e-6, 8.071881352562945e-7],
            linf = [0.00025632790161078667, 1.6379021163651086e-5, 1.637902116437273e-5, 2.58759953227633e-5, 5.327732286231068e-6, 8.118520269495555e-6, 8.118520269606577e-6, 1.2107354757678879e-5, 4.165806320060789e-6])
  end
  @testset "parameters_amr.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_amr.toml"),
            l2   = [0.12533080510721514],
            linf = [0.999980298294775])
  end
  @testset "parameters_amr_nonperiodic.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_amr_nonperiodic.toml"),
            l2   = [2.401816280408703e-5],
            linf = [0.0003102822951955575])
  end
  @testset "parameters_amr_vortex.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_amr_vortex.toml"),
            l2   = [2.0750351586876505e-6, 0.003281637561081054, 0.0032807189382436106, 0.0046470466205649425],
            linf = [4.625172721961501e-5, 0.0318570623352572, 0.031910329823320094, 0.04575283708569344])
  end
  @testset "parameters_blast_wave_shockcapturing.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_blast_wave_shockcapturing.toml"),
            l2   = [0.13910202327088322, 0.11538722576277083, 0.1153873048510009, 0.3387876385945495],
            linf = [1.454418325889352, 1.3236875559310013, 1.323687555933169, 1.8225476335086368],
            n_steps_max=30)
  end
  @testset "parameters_ec.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_ec.toml"),
            l2   = [0.06159341742582756, 0.05012484425381723, 0.05013298724507752, 0.22537740506116724],
            linf = [0.29912627861573327, 0.30886767304359375, 0.3088108573487326, 1.0657556075017878])
  end
  @testset "parameters_density_wave.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_density_wave.toml"),
            l2   = [0.001060077845747576, 0.00010600778457107525, 0.00021201556914875742, 2.6501946139091318e-5],
            linf = [0.0065356386867677085, 0.0006535638688170142, 0.0013071277374487877, 0.0001633909674296774],
            extra_analysis_quantities=["l2_error_primitive", "linf_error_primitive"], t_end=0.5)
  end
  @testset "parameters_ec_mhd.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_ec_mhd.toml"),
            l2   = [0.03607862694368351, 0.04281395008247395, 0.04280207686965749, 0.025746770192645763, 0.1611518499414067, 0.017455917249117023, 0.017456981264942977, 0.02688321120361229, 0.00015024027267648003],
            linf = [0.23502083666166018, 0.3156846367743936, 0.31227895161037256, 0.2118146956106238, 0.9743049414302711, 0.09050624115026618, 0.09131633488909774, 0.15693063355520998, 0.0038394720095667593])
  end
  @testset "parameters_hyp_diff_harmonic_nonperiodic.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_harmonic_nonperiodic.toml"),
            l2   = [8.618132353932638e-8, 5.619399844708813e-7, 5.619399845476024e-7],
            linf = [1.124861862326869e-6, 8.622436471483752e-6, 8.622436469707395e-6])
  end
  @testset "parameters_hyp_diff_llf.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_llf.toml"),
            l2   = [0.00015687751088073104, 0.0010259867353397119, 0.0010259867353398994],
            linf = [0.001198695640053704, 0.006423873515701395, 0.006423873515686296])
  end
  @testset "parameters_hyp_diff_nonperiodic.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_nonperiodic.toml"),
            l2   = [8.523077654037775e-6, 2.877932365308637e-5, 5.454942769137812e-5],
            linf = [5.484978959957587e-5, 0.00014544895979200218, 0.000324491268921534])
  end
  @testset "parameters_hyp_diff_upwind.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_upwind.toml"),
            l2   = [5.868147556488962e-6, 3.8051792732628014e-5, 3.8051792732620214e-5],
            linf = [3.70196549871471e-5, 0.0002072058411455302, 0.00020720584114464202])
  end
  @testset "parameters_mortar.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar.toml"),
            l2   = [0.022356422238096973],
            linf = [0.5043638249003257])
  end
  @testset "parameters_mortar_alfven_wave.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar_alfven_wave.toml"),
            l2   = [4.608223422391918e-6, 1.6891556053250136e-6, 1.6202140809698534e-6, 1.6994400213969954e-6, 1.4856807283318347e-6, 1.387768347373047e-6, 1.3411738859512443e-6, 1.7155298750074954e-6, 9.799762075600064e-7],
            linf = [3.52219535260101e-5, 1.534468550207224e-5, 1.426263439847919e-5, 1.4421456102198249e-5, 7.743399239257265e-6, 1.019242699840106e-5, 9.862935257842764e-6, 1.6018118328936515e-5, 5.563695788849475e-6],
            t_end=1.0)
  end
  @testset "parameters_mortar_vortex.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar_vortex.toml"),
            l2   = [2.1202421511973067e-6, 2.7929028341308907e-5, 3.7593065314592924e-5, 8.813423453465327e-5],
            linf = [5.93205509794581e-5, 0.0007486675478352023, 0.0008175405566226424, 0.002212267888996422])
  end
  @testset "parameters_mortar_vortex_split.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar_vortex_split.toml"),
            l2   = [2.1203040671963692e-6, 2.8053312800289536e-5, 3.761758762899687e-5, 8.840565162128428e-5],
            linf = [5.900575985384737e-5, 0.0007547236106317801, 0.000817616344069072, 0.0022090204216524967])
  end
  @testset "parameters_mortar_vortex_split.toml with flux_central" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar_vortex_split.toml"),
            l2   = [2.1202421512026147e-6, 2.7929028341288412e-5, 3.759306531457842e-5, 8.813423453452753e-5],
            linf = [5.932055097812583e-5, 0.0007486675478027838, 0.0008175405566221983, 0.0022122678889928693],
            volume_flux = "flux_central")
  end
  @testset "parameters_mortar_vortex_split.toml with flux_shima_etal" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar_vortex_split.toml"),
            l2   = [2.1200379425410095e-6, 2.805632600815787e-5, 3.759464715100376e-5, 8.84115216688531e-5],
            linf = [5.934112354222254e-5, 0.00075475390405777, 0.0008162778009123128, 0.002206991473730824],
            volume_flux = "flux_shima_etal")
  end
  @testset "parameters_mortar_vortex_split.toml with flux_ranocha" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar_vortex_split.toml"),
            l2   = [2.120037931908414e-6, 2.805632845562748e-5, 3.759465243706522e-5, 8.841157002762106e-5],
            linf = [5.934036929955422e-5, 0.0007547536380712039, 0.000816277844819191, 0.0022070017103743567],
            volume_flux = "flux_ranocha")
  end
  @testset "parameters_mortar_vortex_split_shockcapturing.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_mortar_vortex_split_shockcapturing.toml"),
            l2   = [2.1205855860697905e-6, 2.805356649496243e-5, 3.7617723084029226e-5, 8.841527980901164e-5],
            linf = [5.9005286894620035e-5, 0.0007547295163081724, 0.0008176139355887679, 0.0022089993378280326])
  end
  @testset "parameters_nonperiodic.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_nonperiodic.toml"),
            l2   = [2.3652137675654753e-6, 2.1386731303685556e-6, 2.138673130413185e-6, 6.009920290578574e-6],
            linf = [1.4080448659026246e-5, 1.7581818010814487e-5, 1.758181801525538e-5, 5.9568540361709665e-5])
  end
  @testset "parameters_source_terms.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_source_terms.toml"),
            l2   = [8.517783186497567e-7, 1.2350199409361865e-6, 1.2350199409828616e-6, 4.277884398786315e-6],
            linf = [8.357934254688004e-6, 1.0326389653148027e-5, 1.0326389654924384e-5, 4.4961900057316484e-5])
  end
  @testset "parameters_vortex.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_vortex.toml"),
            l2   = [3.6343138447409784e-6, 0.0032111379843728876, 0.0032111482778261658, 0.004545715889714643],
            linf = [7.901869034399045e-5, 0.030511158864742205, 0.030451936462313256, 0.04361908901631395])
  end
  @testset "parameters_vortex_split_shockcapturing.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_vortex_split_shockcapturing.toml"),
            l2   = [3.8034711509468997e-6, 5.561030973129845e-5, 5.563956603258559e-5, 0.00015706441614772137],
            linf = [8.493408680687597e-5, 0.0009610606296146518, 0.0009684675522437791, 0.003075812221315033])
  end
  @testset "parameters_weak_blast_wave_shockcapturing.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_weak_blast_wave_shockcapturing.toml"),
            l2   = [0.05365734539276933, 0.04683903386565478, 0.04684207891980008, 0.19632055541821553],
            linf = [0.18542234326379825, 0.24074440953554058, 0.23261143887822433, 0.687464986948263])
  end
  @testset "parameters_khi_amr.toml with t_end=0.2" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_khi_amr.toml"),
            l2   = [0.0016901662212296502, 0.005064145650514081, 0.004794017657493158, 0.0039877996168673525],
            linf = [0.027437774935491266, 0.02344577999610231, 0.016129408502293267, 0.018237901415986357],
            t_end = 0.2)
  end
  @testset "parameters_blob_amr.toml with t_end=0.12" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_blob_amr.toml"),
            l2   = [0.20167272820008805, 1.1836133229138053, 0.10165112533393011, 5.237361125542303],
            linf = [14.085801194734044, 71.07468448364403, 7.366158173410174, 297.2413787328775],
            t_end = 0.12)
  end
  @testset "parameters_orszag_tang.toml with t_end=0.09" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_orszag_tang.toml"),
            l2   = [0.2167754423280243, 0.26359006944514035, 0.31414664596581104, 0.0, 0.5121313736058165, 0.22945862419190388, 0.34334574741333446, 0.0, 0.006113771244085815],
            linf = [1.2713060246638115, 0.6733000091062668, 0.8764507997673942, 0.0, 2.819437219022322, 0.6756993525030384, 0.9970376054894377, 0.0, 0.05310784504107611],
            t_end = 0.09)
  end
  # second orszag-tang test added to exercise all components of flux_hll for GLM-MHD
  @testset "parameters_orszag_tang.toml with t_end=0.06" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_orszag_tang.toml"),
            l2   = [0.10797201158738569, 0.20183923675400306, 0.22965987220003012, 0.0, 0.2993862890890887, 0.156865945562645, 0.24280278839125313, 0.0, 0.0034424029849492947],
            linf = [0.5605977377102305, 0.509274304343078, 0.6612312950402195, 0.0, 0.9814914518422291, 0.39977774699748586, 0.6748004026413148, 0.0, 0.06802505466018915],
            t_end = 0.06, surface_flux = "flux_hll")
  end
  @testset "parameters_ec_mortar.toml with shock_capturing" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_ec_mortar.toml"),
            l2   = [0.04816136246215661, 0.03713041026830962, 0.03713130328181323, 0.1777051166244772],
            linf = [0.3118606868100966, 0.34614370128998007, 0.3460122144359348, 1.1085840270633454],
            volume_integral_type = "shock_capturing")
  end
  @testset "parameters.toml with restart and t_end=2" begin
    Trixi.run(joinpath(EXAMPLES_DIR, "parameters.toml"))
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [1.2148032444677485e-5],
            linf = [6.495644794757283e-5],
            t_end = 2,
            restart = true, restart_filename = "out/restart_000040.h5")
  end
  @test_nowarn Trixi.convtest(joinpath(EXAMPLES_DIR, "parameters.toml"), 3)
  @testset "parameters_blast_wave_shockcapturing_amr.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_blast_wave_shockcapturing_amr.toml"), t_end=1.0,
            l2   = [0.6776486969229697, 0.2813026529898539, 0.28130256451012314, 0.7174702524881598],
            linf = [2.8939055423031532, 1.7997630098946864, 1.799711865996927, 3.034122348258568])
  end
  @testset "parameters_sedov_blast_wave_shockcapturing_amr.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_sedov_blast_wave_shockcapturing_amr.toml"), t_end=1.0,
            l2   = [0.48181540798407435, 0.16553711811584917, 0.16553711811592348, 0.6436020727868234],
            linf = [2.4861229629790813, 1.2873838211418498, 1.2873838211478545, 6.473895863328632])
  end
end

# Coverage test for all initial conditions
@testset "Tests for initial conditions" begin
  # Linear scalar advection
  @testset "parameters.toml with initial_conditions_sin_sin" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [0.0001424424804667062],
            linf = [0.0007260692243250544],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_sin_sin")
  end
  @testset "parameters.toml with initial_conditions_constant" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [6.120436421866528e-16],
            linf = [1.3322676295501878e-15],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_constant")
  end
  @testset "parameters.toml with initial_conditions_linear_x_y" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [2.559042358408011e-16],
            linf = [6.8833827526759706e-15],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_linear_x_y",
            boundary_conditions = "boundary_conditions_linear_x_y",
            periodicity=false)
  end
  @testset "parameters.toml with initial_conditions_linear_x" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [1.5901063275642836e-16],
            linf = [1.5543122344752192e-15],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_linear_x",
            boundary_conditions = "boundary_conditions_linear_x",
            periodicity=false)
  end
  @testset "parameters.toml with initial_conditions_linear_y" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [1.597250146891042e-16],
            linf = [3.552713678800501e-15],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_linear_y",
            boundary_conditions = "boundary_conditions_linear_y",
            periodicity=false)
  end
  # Compressible Euler
  @testset "parameters_vortex.toml one step with initial_conditions_density_pulse" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_vortex.toml"),
            l2   = [0.003201074851451383, 0.0032010748514513724, 0.0032010748514513716, 0.0032010748514513794],
            linf = [0.043716393835876444, 0.043716393835876444, 0.043716393835876, 0.04371639383587578],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_density_pulse")
  end
  @testset "parameters_vortex.toml one step with initial_conditions_pressure_pulse" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_vortex.toml"),
            l2   = [0.00018950189533270512, 0.0020542290689775757, 0.002054229068977579, 0.01013381064979542],
            linf = [0.004763284475434837, 0.028439617580275578, 0.028439617580275467, 0.13640572175447918],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_pressure_pulse")
  end
  @testset "parameters_vortex.toml one step with initial_conditions_density_pressure_pulse" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_vortex.toml"),
            l2   = [0.0031880440066425803, 0.0050397619349217574, 0.005039761934921767, 0.014340770024960708],
            linf = [0.04279723800834989, 0.06783565847184869, 0.06783565847184914, 0.19291274039254347],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_density_pressure_pulse")
  end
  @testset "parameters_vortex.toml one step with initial_conditions_constant" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_vortex.toml"),
            l2   = [2.359732835648237e-16, 1.088770274131804e-16, 1.1814939065033234e-16, 1.980283448445849e-15],
            linf = [4.440892098500626e-16, 2.914335439641036e-16, 4.718447854656915e-16, 3.552713678800501e-15],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_constant")
  end
  @testset "parameters_sedov_blast_wave_shockcapturing_amr.toml one step" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_sedov_blast_wave_shockcapturing_amr.toml"),
            l2   = [0.002911075352335366, 0.01249799423742342, 0.01249799423742343, 0.11130739933709777],
            linf = [0.15341072072011042, 0.763322686048535, 0.7633226860485351, 5.184635785270958],
            n_steps_max = 1)
  end
  @testset "parameters_sedov_blast_wave_shockcapturing_amr.toml one step with initial_conditions_medium_sedov_blast_wave" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_sedov_blast_wave_shockcapturing_amr.toml"),
            l2   = [0.0029095199084281176, 0.012491250999308508, 0.012491250999308522, 0.11126623649275227],
            linf = [0.15334906997459008, 0.7629367729245761, 0.7629367729245761, 5.18264418672338],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_medium_sedov_blast_wave")
  end

  # GLM-MHD
  @testset "parameters_alfven_wave.toml one step with initial_conditions_constant" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_alfven_wave.toml"),
            l2   = [1.9377318494777845e-16, 2.0108417179968547e-16, 4.706803550379074e-16, 9.849916218369067e-17, 9.578096259273606e-15, 4.995499731290712e-16, 2.72017579525395e-16, 9.963303137205655e-17, 1.7656549191657418e-16],
            linf = [4.440892098500626e-16, 7.494005416219807e-16, 1.7763568394002505e-15, 2.220446049250313e-16, 2.1316282072803006e-14, 1.3322676295501878e-15, 8.881784197001252e-16, 2.220446049250313e-16, 7.414582366945819e-16],
            n_steps_max = 1,
            initial_conditions = "initial_conditions_constant")
  end
  @testset "parameters_rotor.toml one step" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_rotor.toml"),
            l2   = [0.019195558428459367, 0.0444345270981607, 0.04329215391420532, 0.0, 0.07639332626939939, 0.0044777850800556485, 0.007266820549785791, 0.0, 0.0008422441423441393],
            linf = [0.5171127312547106, 1.9513340343194656, 1.9404704272260034, 0.0, 2.993145341292035, 0.14424331829406456, 0.20606277176198587, 0.0, 0.06685635962044646],
            n_steps_max = 1)
  end
  @test_skip Trixi.run(joinpath(EXAMPLES_DIR, "parameters_mhd_blast_wave.toml"), n_steps_max=1)
end

# Only run extended tests if environment variable is set
if haskey(ENV, "TRIXI_TEST_EXTENDED") && lowercase(ENV["TRIXI_TEST_EXTENDED"]) in ("1", "on", "yes")
  @testset "Examples (long execution time)" begin
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_blob.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_blob_amr.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_ec_performance_test.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_khi.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_ec_mortar.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_khi_amr.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_mhd_blast_wave.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_orszag_tang.toml"))
    @test_nowarn Trixi.run(joinpath(EXAMPLES_DIR, "parameters_rotor.toml"))
  end
end

# Clean up afterwards: delete Trixi output directory
@test_nowarn rm(outdir, recursive=true)

end #module


text:
using QOI
using Test
using Downloads
using p7zip_jll
using ImageIO
using FileIO
using Scratch
using ColorTypes
using FixedPointNumbers

const DOWNLOAD_SHA = "869a6433a3af7ce84fc55fda6a5387d6c2113c3e8231153549a6407ed1e71696"

function get_testset()
    scratch = get_scratch!(Scratch, DOWNLOAD_SHA)
    if isempty(readdir(scratch))
        p7z = p7zip_jll.p7zip()
        file = Downloads.download("https://github.com/KristofferC/QOI.jl/releases/download/v0.0.0/qoi_test_images.zip")
        run(`$p7z e $file -o$scratch`)
    end
    return scratch
end

function check_roundtrip_qoi(file)
    header, image = QOI.qoi_decode_raw(file)
    data = QOI.qoi_encode_raw(image, header)
    return data == read(file)
end

testset = get_testset()

# Test roundtrip of raw decode/encode
for file in readdir(testset; join=true)
    if endswith(file, ".qoi")
        @info "testing roundtrip of $file"
        @test check_roundtrip_qoi(file)
    end
end

# Test correctness vs PNG
for file_name in Set(getindex.(splitext.(readdir(testset; join=true)), 1))
    qoi = file_name * ".qoi"
    isfile(qoi) || continue
    @info "comparing to png $qoi"
    png = file_name * ".png"
    img_qoi = QOI.qoi_decode(qoi)
    img_png = load(png)
    @test img_qoi == img_png

    t = tempname()
    QOI.qoi_encode(t, img_qoi)
    @test read(t) == read(qoi)
end

# Test saving imges not in N0f8 format
img_gray = rand(Gray{N0f8}, 5, 5)
t = tempname()
QOI.qoi_encode(t, img_gray)
img_qoi = QOI.qoi_decode(t)
@test convert(Matrix{RGB{N0f8}}, img_gray) == img_qoi

img_24bit = rand(RGBA{Float64}, 5, 5)
t = tempname()
QOI.qoi_encode(t, img_24bit)
img_qoi = QOI.qoi_decode(t)
@test convert(Matrix{RGBA{N0f8}}, img_24bit) == img_qoi

# Invalid images

# Unexpected end
io = IOBuffer()
write(io, "qoif")
write(io, hton(UInt32(10)))
write(io, hton(UInt32(0)))
@test_throws QOI.QOIException("unexpected end of file") QOI.qoi_decode_raw(take!(io))

# Invalid width
io = IOBuffer()
write(io, "qoif")
write(io, hton.(UInt32.([0, 10])))
write(io, hton.(UInt8.([3, 1])))
@test_throws QOI.QOIException("invalid width in header, got 0") QOI.qoi_decode_raw(take!(io))

# Invalid height
io = IOBuffer()
write(io, "qoif")
write(io, hton.(UInt32.([10, 0])))
write(io, hton.(UInt8.([3, 1])))
@test_throws QOI.QOIException("invalid height in header, got 0") QOI.qoi_decode_raw(take!(io))

# Invalid channels
io = IOBuffer()
write(io, "qoif")
write(io, hton.(UInt32.([10, 5])))
write(io, hton.(UInt8.([5, 1])))
@test_throws QOI.QOIException("invalid channels in header, got 5") QOI.qoi_decode_raw(take!(io))

# Invalid colorspace
io = IOBuffer()
write(io, "qoif")
write(io, hton.(UInt32.([10, 5])))
write(io, hton.(UInt8.([3, 2])))
@test_throws QOI.QOIException("invalid colorspace in header, got 2") QOI.qoi_decode_raw(take!(io))

# Too little data after header
io = IOBuffer()
write(io, "qoif")
write(io, hton.(UInt32.([10, 5])))
write(io, hton.(UInt8.([3, 1])))
@test_throws QOI.QOIException("unexpected end of file") QOI.qoi_decode_raw(take!(io))

# Invalid magic bytes
io = IOBuffer()
write(io, "qoiz")
write(io, hton.(UInt32.([10, 5])))
write(io, hton.(UInt8.([3, 1])))
@test_throws QOI.QOIException("invalid magic bytes, got 0x716f697a, expected 0x716f6966") QOI.qoi_decode_raw(take!(io))

const f_logo = joinpath(@__DIR__, "qoi_logo.qoi")

# Test IO API
data_io = open(QOI.qoi_decode, f_logo)
data_file = QOI.qoi_decode(f_logo)
@test data_io == data_file

# Test FileIO integration (enable when https://github.com/JuliaIO/FileIO.jl/pull/354 is merged)
@test data_file == FileIO.load(f_logo)
t = tempname() * ".qoi"
FileIO.save(t, data_file)
@test read(f_logo) == read(t)


text:
include("test_trees.jl")
include("test_lvltree.jl")
include("test_blocktree.jl")


text:
using Documenter
using Diversity
using Phylo

makedocs(modules = [Diversity,
                    Diversity.Ecology, Diversity.Jost,
                    Diversity.Hill,
                    Diversity.ShortNames, Diversity.API],
         sitename = "Diversity.jl")

deploydocs(repo = "github.com/EcoJulia/Diversity.jl.git",
           devbranch = "dev",
           deps = Deps.pip("pygments",
                           "mkdocs",
                           "mkdocs-material",
                           "python-markdown-math"))


text:
function printBasisPool(diag, io = Base.stdout)
    printstyled(io, "Loop Basis ($(length(diag.loopBasis)) in total)\n", color = :blue)
    title = @sprintf("%5s%40s\n", "index", "loopBasis")
    printstyled(io, title, color = :green)
    for i = 1:length(diag.loopBasis)
        b = diag.loopBasis.basis[:, i]
        @printf(io, "%5i%40s\n", i, "$b")
    end
    println(io)
end

function printNodes(diag, io = Base.stdout)
    printstyled(io, "Node ($(length(diag.node)) in total)\n", color = :blue)
    title = @sprintf("%5s%5s%40s%40s\n", "index", "name", "para", "child")
    printstyled(io, title, color = :green)
    for (idx, n) in enumerate(diag.node.object)
        # site = isempty(p.siteBasis) ? "" : "$(p.siteBasis)"
        # loop = p.loopIdx <= 0 ? "" : "$(diag.basisPool[p.loopIdx])"
        # loop = p.loopIdx <= 0 ? "" : "$(p.loopIdx)"
        if n isa Propagator
            site = isempty(n.siteBasis) ? "" : "$(n.siteBasis)"
            loop = n.loopIdx <= 0 ? "" : "$(diag.loopBasis[n.loopIdx])"
            @printf(io, "%5i%5s%40s%40s%40s\n", idx, "$(n.name)", "$(n.para)", loop, site)
        else
            @printf(io, "%5i%5s%40s%40s\n", idx, "$(n.name)", "$(n.para)", "$(n.childNodes)")
        end
    end
    println(io)
end

"""
    showTree(diag::Diagrams, _root = diag.root[end]; verbose = 0, depth = 999)

    Visualize the diagram tree using ete3 python package

#Arguments
- `diag`: the Diagrams struct to visualize
- `_root`: the index of the root node to visualize
- `verbose=0`: the amount of information to show
- `depth=999`: deepest level of the diagram tree to show
"""
function showTree(tree::ExpressionTree, _root::Int; verbose = 0, depth = 999)

    # pushfirst!(PyVector(pyimport("sys")."path"), @__DIR__) #comment this line if no need to load local python module
    ete = PyCall.pyimport("ete3")

    function name_para(p)
        name = (p.name == :none) ? "" : " $(p.name)"
        para = isnothing(p.para) ? "" : " $(p.para)"
        return name * para
    end

    function factor(f)
        if f ≈ 1
            return ""
        end
        s = "$f"
        if length(s) <= 4
            return s
        else
            return @sprintf("%6.3e", f)
        end
    end

    function info(node, idx)
        s = "N$(idx)$(name_para(node)) = $(tree.node.current[idx])"
        # s *= sprint(show, node.para)
        # s *= ": "

        if node.operation == MUL
            s *= " = $(factor(node.factor))x"
        elseif node.operation == ADD
            s *= " = $(factor(node.factor))+"
        else
            error("not implemented!")
        end
        return s
    end

    function treeview(idx::Int, level, t = nothing)
        if isnothing(t)
            t = ete.Tree(name = " ")
        end

        if tree.node.object[idx] isa PropagatorId
            p = tree.node.object[idx]    #Propagator
            site = isempty(p.siteBasis) ? "" : " t$(p.siteBasis),"
            loop = p.loopIdx <= 0 ? "" : "k$(tree.loopBasis[p.loopIdx])"
            # loop = p.loopIdx <= 0 ? "" : "$(p.loopIdx)"
            nnt = t.add_child(name = "P$(idx)$(name_para(p)): $loop,$site $(factor(p.factor))")
        else # composite node
            nt = t.add_child(name = info(tree.node.object[idx], idx))
            name_face = ete.TextFace(nt.name, fgcolor = "black", fsize = 10)
            nt.add_face(name_face, column = 0, position = "branch-top")

            for child in tree.node.object[idx].children
                if child != -1
                    treeview(child, level + 1, nt)
                else
                    nnt = nt.add_child(name = "0")
                end
            end
        end

        return t
    end

    t = treeview(_root, 1)
    # style = ete.NodeStyle()
    # style["bgcolor"] = "Khaki"
    # t.set_style(style)

    ts = ete.TreeStyle()
    ts.show_leaf_name = true
    # ts.show_leaf_name = True
    # ts.layout_fn = my_layout
    ####### show tree vertically ############
    # ts.rotation = 90 #show tree vertically

    ####### show tree in an arc  #############
    # ts.mode = "c"
    # ts.arc_start = -180
    # ts.arc_span = 180
    # t.write(outfile="/home/kun/test.txt", format=8)
    t.show(tree_style = ts)
end
# function showTree(diag::ExpressionTree, _root::Component; kwargs...)
#     @assert _root.isNode "Can not visualize $_root, because it is not a Node!"
#     return showTree(diag, _root.index; kwargs...)
# end


text:
using LightGraphs, SimpleWeightedGraphs

function pearson_correlation(X, Y, n)
    correlation::Float64 = 0.0
    sumXY::Float64 = 0.0
    sumX2::Float64 = 0.0
    sumY2::Float64 = 0.0
    sumX::Float64 = 0.0
    sumY::Float64 = 0.0
    for j in 1:n
        sumX += X[j]
        sumX2 += X[j] * X[j]
        sumXY += X[j] * Y[j]
        sumY += Y[j]
        sumY2 += Y[j] * Y[j]
    end
    # Pearson Correlation
    correlation = ((n * sumXY) - (sumX * sumY) ) / (sqrt((n * sumX2 - sumX * sumX) * (n * sumY2 - sumY * sumY)))
    return correlation
end

function objective_insertion(bucket, fitness, item, orig_bucket, dest_bucket, instance)
    if orig_bucket < dest_bucket
        @inbounds for i in eachindex(bucket)
            if bucket[i] == orig_bucket
                fitness += ( -2 * instance.C[item,i] )
            elseif orig_bucket < bucket[i] < dest_bucket
                # fitness += (2  *(instance.C[i,item] - instance.C[item,i]))
                fitness += (2 * instance.D[i,item])
            elseif bucket[i] == dest_bucket
                fitness += (2 * instance.C[i,item])
            end
        end
    else
        @inbounds for i in eachindex(bucket)
            if bucket[i] == orig_bucket
                fitness += ( -2 * instance.C[i,item] )
            elseif dest_bucket < bucket[i] < orig_bucket
                # fitness += (2 *(instance.C[item,i] - instance.C[i,item]))
                fitness += (2 * instance.D[item,i])
            elseif bucket[i] == dest_bucket
                fitness += (2 * instance.C[item,i])
            end
        end
    end
    return fitness
end

function objective_swap(bucket, fitness, itemA, itemB, instance)
    if bucket[itemA] < bucket[itemB]
        @inbounds for i in eachindex(bucket)
            if bucket[i] == bucket[itemA]
                fitness += (-2 * instance.C[itemA,i])
                fitness += (2 * instance.C[itemB,i])
            elseif bucket[itemA] < bucket[i] < bucket[itemB]
                fitness += (2 * instance.D[i,itemA])
                fitness += (2 * instance.D[itemB,i])
            elseif bucket[i] == bucket[itemB]
                fitness += (2 * instance.C[i,itemA])
                fitness += (-2 * instance.C[i,itemB])
            end
        end
    else
        @inbounds for i in eachindex(bucket)
            if bucket[i] == bucket[itemB]
                fitness += (2 * instance.C[itemA,i])
                fitness += (-2 * instance.C[itemB,i])
            elseif bucket[itemB] < bucket[i] < bucket[itemA]
                fitness += (2 * instance.D[i,itemB])
                fitness += (2 * instance.D[itemA,i])
            elseif bucket[i] == bucket[itemA]
                fitness += (-2 * instance.C[i,itemA])
                fitness += (2 * instance.C[i,itemB])
            end
        end
    end
    return fitness
end

function neighbor_insertion(bucket, fitness, instance)

    best_fitness = fitness

    total_vizinhos = (2 * instance.total_itens) + 1

    # O(2n+1)
    buckets_testar = zeros(Int64, total_vizinhos)
    for i in 1:length(bucket)
        bucket[i] = 2 * bucket[i]
        buckets_testar[bucket[i]] += 1
    end

    ordem_visita = shuffle!(collect(1:instance.total_itens))

    for item in ordem_visita

        fitness = best_fitness
        best_bucket = orig_bucket = bucket[item]

        for id in 1:total_vizinhos

            if buckets_testar[id] > 0 && buckets_testar[id] != orig_bucket

                # Testa colocar na esquerda
                dest_bucket = id - 1
                fo = objective_insertion(bucket,
                                         fitness,
                                         item,
                                         orig_bucket,
                                         dest_bucket,
                                         instance)
                if fo > best_fitness
                    best_fitness = fo
                    best_bucket = dest_bucket
                end

                # Testa colocar no grupo
                dest_bucket = id
                fo = objective_insertion(bucket,
                                         fitness,
                                         item,
                                         orig_bucket,
                                         dest_bucket,
                                         instance)
                if fo > best_fitness
                    best_fitness = fo
                    best_bucket = dest_bucket
                end
            end
        end

        # Testa colocar no último lugar
        dest_bucket = total_vizinhos
        fo = objective_insertion(bucket,
                                 fitness,
                                 item,
                                 orig_bucket,
                                 dest_bucket,
                                 instance)
        if fo > best_fitness
            best_fitness = fo
            best_bucket = dest_bucket
        end

        bucket[item] = best_bucket
        buckets_testar[orig_bucket] -= 1
        buckets_testar[best_bucket] += 1

        if best_bucket % 2 == 1

            # O(2n+1)
            new_id = 1
            for i in 1:total_vizinhos
                if buckets_testar[i] > 0
                    buckets_testar[i] = new_id
                    new_id += 1
                end
            end

            # O(n)
            for i in 1:length(bucket)
                bucket[i] = 2 * (buckets_testar[bucket[i]])
            end

            # O(2n+1)
            fill!(buckets_testar, 0)

            # O(n)
            for i in 1:length(bucket)
                buckets_testar[bucket[i]] += 1
            end

        end
    end

    # O(2n+1)
    new_id = 1
    for i in 1:total_vizinhos
        if buckets_testar[i] > 0
            buckets_testar[i] = new_id
            new_id += 1
        end
    end

    # O(n)
    for i in 1:length(bucket)
        bucket[i] = buckets_testar[bucket[i]]
    end

    return best_fitness

end

function neighbor_swap(bucket, fitness, instance)

    best_fitness = fitness

    ordem_visita_A = shuffle!(collect(1:instance.total_itens))
    ordem_visita_B = shuffle!(collect(1:instance.total_itens))

    for itemA in ordem_visita_A

        fitness = best_fitness

        best_item = itemA

        for itemB in ordem_visita_B

            # Não troca com o mesmo grupo
            if bucket[itemB] != bucket[itemA]
                fo = objective_swap(bucket, fitness, itemA, itemB, instance)
                if fo > best_fitness
                    best_fitness = fo
                    best_item = itemB
                end
            end

        end # For testa buckets

        if best_item != itemA
            # bucket[itemA],bucket[best_item] = bucket[best_item],bucket[itemA]
            temp = bucket[best_item]
            bucket[best_item] = bucket[itemA]
            bucket[itemA] = temp
        end

    end # For ordem visita

    return best_fitness

end

function local_search_vnd!(item::Int64,
    pop_keys::Array{Float64,2},
    pop_buckets::Array{Int64,2},
    pop_fitness::Array{Int64,1},
    instance::OBOPDataset,
    best_solution::OBOPSolution,
    generation::Int64,
    statistics::OBOPStatistics)

    # Calculate the difference between each key and corresponding bucket key start
    dif = [(pop_keys[i,item] - instance.interval_init[pop_buckets[i,item]]) for i in 1:instance.total_itens]

    # Add local search count
    statistics.total_local_search += 1

    f = [neighbor_swap, neighbor_insertion]

    fo_best = pop_fitness[item]

    vizinho = 1
    while vizinho <= 2

        println("##################")
        println("Neighbor: $vizinho")
        println("$fo_best $(pop_buckets[:,item])")

        #fo = f[vizinho](pop_buckets[:,item], fo_best, instance)
        fo = f[vizinho](view(pop_buckets,:,item), fo_best, instance)

        println("$fo $(pop_buckets[:,item])")

        if fo > fo_best
            fo_best = fo
            vizinho = 1
        else
            vizinho += 1
        end
    end

    # Checa se é melhor
    if fo_best > best_solution.objective
        statistics.total_local_search_effective += 1
        best_solution.total_time = time() - best_solution.start_time
        best_solution.objective = fo_best
        for i in 1:instance.total_itens
            best_solution.bucket[i] = pop_buckets[i,item]#z[pop_buckets[i,item]]
        end
        best_solution.generation = generation
        best_solution.local_search = false
    end

    # ============================================================= #
    #  Modificar as chaves do individuo respeitando a diferença     #
    #  inicial que ele tinha. Pegar o inicio do bucket que ele foi  #
    #  e soma a diferença que ele tinha                             #
    # ============================================================= #
    for i in 1:instance.total_itens
        pop_keys[i,item] = instance.interval_init[pop_buckets[i,item]] + dif[i]
    end
    pop_fitness[item] = fo_best

end

function clustering_search_new(elite_size::Int64,
    index_order::Array{Int64,1},
    pop_keys::Array{Float64,2},
    pop_buckets::Array{Int64,2},
    pop_fitness::Array{Int64,1},
    instance::OBOPDataset,
    best_solution::OBOPSolution,
    generation::Int64,
    statistics::OBOPStatistics)

    # Create weighted graph
    sigma = 0.7                 # Used to make graph sparse
    graph = SimpleWeightedGraph(elite_size)
    for i in 1:elite_size - 1
        for j in i + 1:elite_size
            weight = pearson_correlation(pop_keys[:,index_order[i]],
                       pop_keys[:,index_order[j]],
                       instance.total_itens)
            if weight >= sigma
                add_edge!(graph, i, j, weight)
            end
        end
    end

    # Return the index of the chromosomes that the LS will be applied
    index_local_search = weighted_label_propagation(graph)

    #ls = [local_search!,local_search_random!,local_search_parallel!]

    for item in index_local_search
        local_search_vnd!(index_order[item],
                    pop_keys,
                    pop_buckets,
                    pop_fitness,
                    instance,
                    best_solution,
                    generation,
                    statistics)
    end

end

function objective_partial(bucket, fitness, item, orig_bucket, dest_bucket, instance)
    if orig_bucket < dest_bucket
        @inbounds for i in eachindex(bucket)
            if bucket[i] == orig_bucket
                fitness += ( -2 * instance.C[item,i] )
            elseif orig_bucket < bucket[i] < dest_bucket
                # fitness += (2  *(instance.C[i,item] - instance.C[item,i]))
                fitness += (2 * instance.D[i,item])
            elseif bucket[i] == dest_bucket
                fitness += (2 * instance.C[i,item])
            end
        end
    else
        @inbounds for i in eachindex(bucket)
            if bucket[i] == orig_bucket
                fitness += ( -2 * instance.C[i,item] )
            elseif dest_bucket < bucket[i] < orig_bucket
                # fitness += (2 *(instance.C[item,i] - instance.C[i,item]))
                fitness += (2 * instance.D[item,i])
            elseif bucket[i] == dest_bucket
                fitness += (2 * instance.C[item,i])
            end
        end
    end
    return fitness
end

function weighted_label_propagation(graph)

    n = nv(graph)

    # Vertices communities
    labels = collect(1:n)

    # Index to start visiting
    visit_order = collect(1:n)

    # Total neighbor labels
    labels_total = Array{Float64}(undef, n)

    # Main loop WLP
    iteration = 1
    best_label = 0
    moviments = true
    while moviments
        moviments = false
        shuffle!(visit_order)
        for u in visit_order
            if size(neighbors(graph, u), 1) > 0
                fill!(labels_total, 0.0)
                for v in neighbors(graph, u)
                    labels_total[labels[v]] += graph.weights[u,v]
                end
                best_label = findmax(labels_total)[2]
                if best_label != labels[u]
                    moviments = true
                    labels[u] = best_label
                end
            end
        end
        iteration += 1
    end

    # Create groups and select one random representant from each group
    groups = [[] for _ in 1:n]
    for i in 1:n
        push!(groups[labels[i]], i)
    end

    # Select one random representant from each group
    index_local_search = []
    for i in 1:n
        if size(groups[i], 1) > 0
            push!(index_local_search, groups[i][rand(collect(1:size(groups[i], 1)))])
        end
    end

    return index_local_search
end

function local_search!(item::Int64,
                       pop_keys::Array{Float64,2},
                       pop_buckets::Array{Int64,2},
                       pop_fitness::Array{Int64,1},
                       instance::OBOPDataset,
                       best_solution::OBOPSolution,
                       generation::Int64,
                       statistics::OBOPStatistics)

    # Add local search count
    statistics.total_local_search += 1

    # Calculate the difference between each key and corresponding bucket key start
    dif = [(pop_keys[i,item] - instance.interval_init[pop_buckets[i,item]]) for i in 1:instance.total_itens]

    # Index sequence to start the local search
    indexes = sortperm(dif)

    # Create interval between buckets
    # bucket = 2 * bucket
    pop_buckets[:,item] = pop_buckets[:,item] * 2

    # =================== #
    # Inicia refinamento  #
    # =================== #
    originalObjective = currentObjective = bestObjective = pop_fitness[item]
    totalVizinhos = (2 * instance.total_itens) + 1
    bestBucket = nothing

    for j in 1:instance.total_itens
        j = indexes[j]
        bucketAtual = pop_buckets[j,item]
        bestBucket = bucketAtual
        for bucketVizinho in 1:totalVizinhos
            if bucketVizinho != bucketAtual
                currentObjective = objective_partial(pop_buckets[:,item],
                                   originalObjective,
                                   j,
                                   bucketAtual,
                                   bucketVizinho,
                                   instance)
                if currentObjective > bestObjective
                    bestObjective = currentObjective
                    bestBucket = bucketVizinho
                end
            end
        end

        if (bestBucket != bucketAtual)
            pop_buckets[j,item] = bestBucket
            originalObjective = bestObjective
        end
    end

    # ========================================================== #
    # Reindexa os buckets, pois podem estar em um intervalo > n  #
    # ========================================================== #
    z = zeros(Int64, totalVizinhos)
    id = 1
    for i in 1:instance.total_itens
        if z[ pop_buckets[i,item] ] == 0
            z[ pop_buckets[i,item] ] = id
            id += 1
        end
    end

    id = 1
    for i in 1:totalVizinhos
        if z[i] != 0
            z[i] = id
            id += 1
        end
    end

    # Checa se é melhor
    if bestObjective > best_solution.objective

        statistics.total_local_search_effective += 1

        best_solution.total_time = time() - best_solution.start_time
        best_solution.objective = bestObjective
        for i in 1:instance.total_itens
            best_solution.bucket[i] = z[pop_buckets[i,item]]
        end
        best_solution.generation = generation
        best_solution.local_search = false
    end

    # ============================================================= #
    #  Modificar as chaves do individuo respeitando a diferença     #
    #  inicial que ele tinha. Pegar o inicio do bucket que ele foi  #
    #  e soma a diferença que ele tinha                             #
    # ============================================================= #
    for i in 1:instance.total_itens
        pop_buckets[i,item] = z[pop_buckets[i,item]]
        pop_keys[i,item] = instance.interval_init[pop_buckets[i,item]] + dif[i]
    end
    pop_fitness[item] = bestObjective

    # println(pop_buckets[:,item])
    # println(pop_keys[:,item])
    # println(instance.interval_init)

    # println("ok partial LS")

end

function local_search_random!(item::Int64,
    pop_keys::Array{Float64,2},
    pop_buckets::Array{Int64,2},
    pop_fitness::Array{Int64,1},
    instance::OBOPDataset,
    best_solution::OBOPSolution,
    generation::Int64,
    statistics::OBOPStatistics)

    # Add local search count
    statistics.total_local_search += 1

    # Calculate the difference between each key and corresponding bucket key start
    dif = [(pop_keys[i,item] - instance.interval_init[pop_buckets[i,item]]) for i in 1:instance.total_itens]

    # Index sequence to start the local search
    indexes = shuffle!(collect(1:instance.total_itens))

    # Create interval between buckets
    pop_buckets[:,item] = pop_buckets[:,item] * 2

    # =================== #
    # Inicia refinamento  #
    # =================== #
    originalObjective = currentObjective = bestObjective = pop_fitness[item]
    totalVizinhos = (2 * instance.total_itens) + 1
    bestBucket = nothing

    for j in 1:instance.total_itens
        j = indexes[j]
        bucketAtual = pop_buckets[j,item]
        bestBucket = bucketAtual
        for bucketVizinho in 1:totalVizinhos
            if bucketVizinho != bucketAtual
                currentObjective = objective_partial(pop_buckets[:,item],
                originalObjective,
                j,
                bucketAtual,
                bucketVizinho,
                instance)
                if currentObjective > bestObjective
                    bestObjective = currentObjective
                    bestBucket = bucketVizinho
                end
            end
        end
        if (bestBucket != bucketAtual)
            pop_buckets[j,item] = bestBucket
            originalObjective = bestObjective
        end
    end

    # ========================================================== #
    # Reindexa os buckets, pois podem estar em um intervalo > n  #
    # ========================================================== #
    z = zeros(Int64, totalVizinhos)
    id = 1
    for i in 1:instance.total_itens
        if z[ pop_buckets[i,item] ] == 0
            z[ pop_buckets[i,item] ] = id
            id += 1
        end
    end

    id = 1
    for i in 1:totalVizinhos
        if z[i] != 0
            z[i] = id
            id += 1
        end
    end

    # Checa se é melhor
    if bestObjective > best_solution.objective
        statistics.total_local_search_effective += 1
        best_solution.total_time = time() - best_solution.start_time
        best_solution.objective = bestObjective
        for i in 1:instance.total_itens
            best_solution.bucket[i] = z[pop_buckets[i,item]]
        end
        best_solution.generation = generation
        best_solution.local_search = false
    end

    # ============================================================= #
    #  Modificar as chaves do individuo respeitando a diferença     #
    #  inicial que ele tinha. Pegar o inicio do bucket que ele foi  #
    #  e soma a diferença que ele tinha                             #
    # ============================================================= #

    for i in 1:instance.total_itens
        pop_buckets[i,item] = z[pop_buckets[i,item]]
        pop_keys[i,item] = instance.interval_init[pop_buckets[i,item]] + dif[i]
    end
    pop_fitness[item] = bestObjective

end

function local_search_parallel!(item::Int64,
    pop_keys::Array{Float64,2},
    pop_buckets::Array{Int64,2},
    pop_fitness::Array{Int64,1},
    instance::OBOPDataset,
    best_solution::OBOPSolution,
    generation::Int64,
    statistics::OBOPStatistics)

    # Add local search count
    statistics.total_local_search += 1

    # Calculate the difference between each key and corresponding bucket key start
    dif = [(pop_keys[i,item] - instance.interval_init[pop_buckets[i,item]]) for i in 1:instance.total_itens]

    # Index sequence to start the local search
    indexes = sortperm(dif)

    # Create interval between buckets
    # bucket = 2 * bucket
    pop_buckets[:,item] = pop_buckets[:,item] * 2

    # =================== #
    # Inicia refinamento  #
    # =================== #
    originalObjective = currentObjective = bestObjective = pop_fitness[item]
    totalVizinhos = (2 * instance.total_itens) + 1
    bestBucket = nothing

    neighbor_objective = zeros(totalVizinhos)

    for j in 1:instance.total_itens

        j = indexes[j]
        bucketAtual = pop_buckets[j,item]
        bestBucket = bucketAtual

        neighbor_objective[bucketAtual] = bestObjective

        Threads.@threads for bucketVizinho in 1:totalVizinhos
            if bucketVizinho != bucketAtual
                neighbor_objective[bucketVizinho] = objective_partial(pop_buckets[:,item],
                                                                      originalObjective,
                                                                      j,
                                                                      bucketAtual,
                                                                      bucketVizinho,
                                                                      instance)
            end
        end

        bestBucket = argmax(neighbor_objective)

        if (bestBucket != bucketAtual)
            pop_buckets[j,item] = bestBucket
            originalObjective = neighbor_objective[bestBucket]
        end

    end

    # ========================================================== #
    # Reindexa os buckets, pois podem estar em um intervalo > n  #
    # ========================================================== #
    z = zeros(Int64, totalVizinhos)
    id = 1
    for i in 1:instance.total_itens
        if z[ pop_buckets[i,item] ] == 0
            z[ pop_buckets[i,item] ] = id
            id += 1
        end
    end

    id = 1
    for i in 1:totalVizinhos
        if z[i] != 0
            z[i] = id
            id += 1
        end
    end

    # Checa se é melhor
    if bestObjective > best_solution.objective

        statistics.total_local_search_effective += 1

        best_solution.total_time = time() - best_solution.start_time
        best_solution.objective = bestObjective
        for i in 1:instance.total_itens
            best_solution.bucket[i] = z[pop_buckets[i,item]]
        end
        best_solution.generation = generation
        best_solution.local_search = false
    end

    # ============================================================= #
    #  Modificar as chaves do individuo respeitando a diferença     #
    #  inicial que ele tinha. Pegar o inicio do bucket que ele foi  #
    #  e soma a diferença que ele tinha                             #
    # ============================================================= #
    for i in 1:instance.total_itens
        pop_buckets[i,item] = z[pop_buckets[i,item]]
        pop_keys[i,item] = instance.interval_init[pop_buckets[i,item]] + dif[i]
    end
    pop_fitness[item] = bestObjective

end

function clustering_search(elite_size::Int64,
                           index_order::Array{Int64,1},
                           pop_keys::Array{Float64,2},
                           pop_buckets::Array{Int64,2},
                           pop_fitness::Array{Int64,1},
                           instance::OBOPDataset,
                           best_solution::OBOPSolution,
                           generation::Int64,
                           statistics::OBOPStatistics,
                           ls_type)

    # Create weighted graph
    sigma = 0.7                 # Used to make graph sparse
    graph = SimpleWeightedGraph(elite_size)
    for i in 1:elite_size - 1
        for j in i + 1:elite_size
            weight = pearson_correlation(pop_keys[:,index_order[i]],
                                              pop_keys[:,index_order[j]],
                                              instance.total_itens)
            if weight >= sigma
                add_edge!(graph, i, j, weight)
            end
        end
    end

    # Return the index of the chromosomes that the LS will be applied
    index_local_search = weighted_label_propagation(graph)

    ls = [local_search!,local_search_random!,local_search_parallel!]

    for item in index_local_search
        ls[ls_type](index_order[item],
        pop_keys,
        pop_buckets,
        pop_fitness,
        instance,
        best_solution,
        generation,
        statistics)
        #=
        local_search!(index_order[item],
                      pop_keys,
                      pop_buckets,
                      pop_fitness,
                      instance,
                      best_solution,
                      generation,
                      statistics) =#
    end

end

max_stars_repo_name: LuizHNLorena/AbrkgaObop

max_issues_repo_name: LuizHNLorena/AbrkgaObop

max_forks_repo_name: LuizHNLorena/AbrkgaObop


text:
using CvxCompress, Test

@testset "CvxCompress, 2D" begin
    nz,nx,bz,bx,σ = 2048,4096,32,32,1e-3
    x = rand(Float32, nz, nx)
    y = zeros(Float32, nz, nx)
    z = zeros(UInt32, nz*nx)
    n = CvxCompress.compress!(z, CvxCompressor((bz,bx),σ), x)
    CvxCompress.decompress!(y, CvxCompressor((bz,bx),σ), z, n)

    @test x ≈ y rtol=1e-3
end

@testset "CvxCompress, 3D" begin
    nz,ny,nx,bz,by,bx,σ = 128,64,256,32,32,32,1e-3
    x = rand(Float32, nz, ny, nx)
    y = zeros(Float32, nz, ny, nx)
    z = zeros(UInt32, nz*ny*nx)
    n = CvxCompress.compress!(z, CvxCompressor((bz,by,bx),σ), x)
    CvxCompress.decompress!(y, CvxCompressor((bz,by,bx),σ), z, n)

    @test x ≈ y rtol=1e-3
end


text:
using DiscreteElasticRods
DER = DiscreteElasticRods

segments = [10, 10]
bcs = Vector{DER.rod_bc}()
push!(bcs, DER.free_bc(1,0))
push!(bcs, DER.static_clamp_bc(1,1))
push!(bcs, DER.free_bc(2,0))
push!(bcs, DER.static_clamp_bc(2,1))

nq, param_2_rod = DER.get_specification_layer(segments, bcs)
q = rand(80)

param_2_rod(q)

using Zygote
out, back = Zygote.pullback(param_2_rod, rand(nq))

back(out)


text:
using Revise
using LinearAlgebra, Setfield, Parameters, KrylovKit

using BifurcationKit
const BK = BifurcationKit

TY = Float64
AF = Array{TY}
#################################################################
# if 1==1
# 	ENV["GKSwstype"] = "nul"    # needed for the GR backend on headless servers
# 	using Plots
# 	using SixelTerm
# else
# 	using Plots
# 	GR.inline("iterm")
# end
####################################################################################################

using CUDA
CUDA.allowscalar(false)
import LinearAlgebra: mul!, axpby!
mul!(x::CuArray, y::CuArray, α::T) where {T <: Number} = (x .= α .* y)
mul!(x::CuArray, y::T, α::CuArray) where {T <: Number} = (x .= α .* y)
axpby!(a::T, X::CuArray, b::T, Y::CuArray) where {T <: Number} = (Y .= a .* X .+ b .* Y)

TY = Float64
AF = CuArray{TY}
####################################################################################################
using Plots
# to simplify plotting of the solution
plotsol(x; k...) = heatmap(reshape(Array(x), Nx, Ny)'; color=:viridis, k...)
plotsol!(x; k...) = heatmap!(reshape(Array(x), Nx, Ny)'; color=:viridis, k...)
norminf(x) = maximum(abs.(x))
####################################################################################################
using AbstractFFTs, FFTW, KrylovKit
import Base: *, \

# Making the linear operator a subtype of BK.LinearSolver is handy as we will use it
# in the Newton iterations.
struct SHLinearOp{Treal, Tcomp, Tl1, Tplan, Tiplan} <: BK.AbstractIterativeLinearSolver
	tmp_real::Treal         # temporary
	tmp_complex::Tcomp      # temporary
	l1::Tl1
	fftplan::Tplan
	ifftplan::Tiplan
end

struct SHEigOp{Tsh <: SHLinearOp, Tσ} <: BK.AbstractEigenSolver
	sh::Tsh
	σ::Tσ
end

function SHLinearOp(Nx, lx, Ny, ly; AF = Array{TY})
	# AF is a type, it could be CuArray{TY} to run the following on GPU
	k1 = vcat(collect(0:Nx/2), collect(Nx/2+1:Nx-1) .- Nx)
	k2 = vcat(collect(0:Ny/2), collect(Ny/2+1:Ny-1) .- Ny)
	d2 = [(1-(pi/lx * kx)^2 - (pi/ly * ky)^2)^2 + 1. for kx in k1, ky in k2]
	tmpc = Complex.(AF(zeros(Nx, Ny)))
	return SHLinearOp(AF(zeros(Nx, Ny)), tmpc, AF(d2), plan_fft!(tmpc), plan_ifft!(tmpc))
end

function apply(c::SHLinearOp, u, multiplier, op = *)
	c.tmp_complex .= Complex.(u)
	c.fftplan * c.tmp_complex
	c.tmp_complex .= op.(c.tmp_complex, multiplier)
	c.ifftplan * c.tmp_complex
	c.tmp_real .= real.(c.tmp_complex)
	return copy(c.tmp_real)
end

*(c::SHLinearOp, u) = apply(c, u, c.l1)
\(c::SHLinearOp, u) = apply(c, u, c.l1, /)
####################################################################################################
Nx = 2^8
Ny = 2^8
lx = 8pi * 2
ly = 2*2pi/sqrt(3) * 2

X = -lx .+ 2lx/Nx * collect(0:Nx-1)
Y = -ly .+ 2ly/Ny * collect(0:Ny-1)

sol0 = [(cos(x) .+ cos(x/2) * cos(sqrt(3) * y/2) ) for x in X, y in Y]
		sol0 .= sol0 .- minimum(vec(sol0))
		sol0 ./= maximum(vec(sol0))
		sol0 = sol0 .- 0.25
		sol0 .*= 1.7
		# heatmap(sol0, color=:viridis)

function (sh::SHLinearOp)(J, rhs; shift = 0., rtol =  1e-9)
	u, l, ν = J
	udiag = l .+ 1 .+ 2ν .* u .- 3 .* u.^2 .- shift
	res, info = KrylovKit.linsolve( du -> -du .+ sh \ (udiag .* du), sh \ rhs, rtol = rtol, maxiter = 6, ishermitian = true)
	return res, true, info.numops
end

function (sheig::SHEigOp)(J, nev::Int; kwargs...)
	u, l, ν = J
	sh = sheig.sh
	σ = sheig.σ
	A = du -> sh(J, du; shift = σ)[1]

	# we adapt the krylov dimension as function of the requested eigenvalue number
	vals, vec, info = KrylovKit.eigsolve(A, sh \ AF(rand(eltype(u), size(u))), nev, :LM, tol = 1e-10, maxiter = 20, verbosity = 2, ishermitian = true, krylovdim = max(30, nev + 10))
	@show 1 ./vals .+ σ
	return 1 ./vals .+ σ, vec, true, info.numops
end

function F_shfft(u, p)
	@unpack l, ν, L = p
	return -(L * u) .+ ((l+1) .* u .+ ν .* u.^2 .- u.^3)
end

J_shfft(u, p) = (u, p.l, p.ν)

L = SHLinearOp(Nx, lx, Ny, ly, AF = AF)
Leig = SHEigOp(L, 0.1) # for eigenvalues computation
# Leig((sol_hexa, -0.1, 1.3), 20; σ = 0.5)

par = (l = -0.1, ν = 1.3, L = L)

@time F_shfft(AF(sol0), par); # 0.008022 seconds (12 allocations: 1.500 MiB)

opt_new = NewtonPar(verbose = true, tol = 1e-6, linsolver = L, eigsolver = Leig)
	sol_hexa, hist, flag = @time newton(
				F_shfft, J_shfft,
				AF(sol0), par,
				opt_new, normN = norminf)
	println("--> norm(sol) = ", norminf(sol_hexa))

plotsol(sol_hexa)
####################################################################################################
# trial using IterativeSolvers

# function (sh::SHLinearOp)(J, rhs)
# 	u, l, ν = J
# 	udiag = l .+ 1 .+ 2ν .* u .- 3 .* u.^2
# 	res, info = res, info = KrylovKit.linsolve( u -> -u .+ sh \ (udiag .* u), sh \ rhs, tol = 1e-9, maxiter = 6)
# 	return res, true, info.numops
# end
####################################################################################################
# computation of normal form
# we collect the matrix-free derivatives
jet = BK.getJet(F_shfft)
nf = computeNormalForm(jet..., br, 2)
####################################################################################################
deflationOp = DeflationOperator(2, 1.0, [sol_hexa])

opt_new = @set opt_new.maxIter = 250
outdef, _, flag, _ = @time newton(F_shfft, J_shfft,
			0.4 .* sol_hexa .* AF([exp(-1(x+0lx)^2/25) for x in X, y in Y]),
			par,
			opt_new, deflationOp, normN = x-> maximum(abs.(x)))
		println("--> norm(sol) = ", norm(outdef))
		plotsol(outdef) |> display
		flag && push!(deflationOp, outdef)

####################################################################################################
opts_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.007, ds= -0.005, pMax = 0.005, pMin = -1.0, plotEveryStep = 10, newtonOptions = setproperties(opt_new; tol = 1e-6, maxIter = 15), maxSteps = 88,
	detectBifurcation = 0,
	precisionStability = 1e-5,
	saveEigenvectors = false,
	nev = 11 )

	br, u1 = @time continuation(
		F_shfft, J_shfft,
		deflationOp[1], par, (@lens _.l),
		opts_cont;
		linearAlgo = BorderingBLS(solver = L, checkPrecision = false),
		# linearAlgo = MatrixFreeBLS(L),
		plot = true, verbosity = 3,
		plotSolution = (x, p;kwargs...) -> plotsol!(x; color=:viridis, kwargs...),
		recordFromSolution = (x, p) -> norm(x), normC = norminf,
		)


text:
samples = Int(1e3)
networks = Dict([
  ("G", [0 0 1; 1 0 0; 0 -1 0]),
  ("R", [0 0 -1; -1 0 0; 0 -1 0]),
  ("S", [1 -1; 1 0]),
  ("G+IPFC", [0 0 1; 1 0 0; 0 -1 1]),
  ("G+INFB", [0 0 1; 1 -1 0; 0 -1 0]),
  ("F", [-1 0 0 1; 1 -1 0 0; 0 1 -1 0; 0 0 -1 1]),
  ("G+R", [0 -1 1; 1 0 -1; -1 -1 0]),
  ("Cdc25", [0 0 -1 1; -1 0 0 0; 0 -1 0 0; 1 0 0 0]),
  ("Wee1", [0 0 1 1; 1 0 0 0; 0 -1 0 0; -1 0 0 0]),
])
hparams = Dict([("peak_num_thresh", 30), ("freq_tolerance", 0.01),
                ("equil_tscales", 50), ("sim_tscales", 40),
                ("power_threshold", 1e-7), ("amp_cv_thresh", 0.05),
                ("abstol", 1e-12), ("reltol", 1e-6)])
param_limits = Dict([
  ("α", [0.0, 1.0]),
  ("β", [10.0, 100.0]),
  ("γ", [100.0, 10000.0]),
  ("κ", [0.0, 1.0]),
  ("η", [3.0, 10.0]),
])
for (name, connectivity) in networks
  println("Analyzing network $name, Connectivity: $connectivity")
  local model = create_model(connectivity)
  @time sim, p_sample = find_oscillations(model, samples, param_limits, hparams)
  @test size(p_sample, 1) > 1
end

networks_not_oscil = Dict([
  ("P3", [0 0 1; 1 0 0; 0 1 0]),
  ("RP", [0 0 1; -1 0 0; 0 -1 0]),
  ("P2", [0 1; 1 0]),
  ("P4", [0 0 0 1; 1 0 0 0; 0 1 0 0; 0 0 1 0]),
  ("P5", [0 0 0 -1; 1 0 0 0; 0 -1 0 0; 0 0 1 0]),
])

for (name, connectivity) in networks_not_oscil
  println("Analyzing network $name, Connectivity: $connectivity")
  local model = create_model(connectivity)
  @time sim, p_sample = find_oscillations(model, samples, param_limits, hparams)
  @test size(p_sample, 1) == 0
end


text:
#==
Functions related to calculating electromagnetic field components at any height within the
ionosphere.

This includes calculating wavefields vector from the Booker quartic and Pitteway
integration of wavefields.
==#

"""
    Wavefields{H<:AbstractRange}

Struct containing wavefield components for each mode of `eas` at `heights`.

The six electromagnetic field components are stored in the `v` field of the struct and
accessed as a `Matrix` of `SVector{6,ComplexF64}` corresponding to `[height, mode]`.
"""
struct Wavefields{H<:AbstractRange}
    v::Matrix{SVector{6,ComplexF64}}  # v[height,mode]  ; type is derived from typeof T, which is from M
    heights::H
    eas::Vector{EigenAngle}
end

"""
    Wavefields(heights, eas::Vector{EigenAngle})

A constructor for initializing `Wavefields` with an appropriately sized `undef` Matrix
given eigenangles `eas` and `heights`.
"""
function Wavefields(heights, eas::Vector{EigenAngle})
    v = Matrix{SVector{6,ComplexF64}}(undef, length(heights), length(eas))
    return Wavefields(v, heights, eas)
end

Base.size(A::Wavefields) = size(A.v)
Base.size(A::Wavefields, I...) = size(A.v, I...)
Base.getindex(A::Wavefields, i) = A.v[i]
Base.getindex(A::Wavefields, I...) = A.v[I...]
Base.setindex!(A::Wavefields, x, i) = (A.v[i] = x)
Base.similar(A::Wavefields) = Wavefields(A.heights, A.eas)
Base.copy(A::Wavefields) = Wavefields(copy(A.v), copy(A.heights), copy(A.eas))
Base.view(A::Wavefields, I...) = view(A.v, I...)

"""
    ==(A::Wavefields, B::Wavefields)

Check for equality `==` between `A.v` and `B.v`, `heights`, and `eas`.
"""
function (==)(A::Wavefields, B::Wavefields)
    A.heights == B.heights || return false
    A.eas == B.eas || return false
    A.v == B.v || return false
    return true
end

heights(A::Wavefields) = A.heights
numheights(A::Wavefields) = length(A.heights)
eigenangles(A::Wavefields) = A.eas
nummodes(A::Wavefields) = length(A.eas)

function Base.isvalid(A::Wavefields)
    vzs, veas = size(A.v)
    zlen = numheights(A)
    vzs == zlen || return false
    ealen = nummodes(A)
    veas == ealen || return false
    return true
end

"""
    WavefieldIntegrationParams{S,T,T2,H}

Parameters passed to Pitteway integration of wavefields [Pitteway1965].

# Fields

- `z::Float64`: height in meters.
- `bottomz::Float64`: bottom height of integration in meters.
- `ortho_scalar::Complex{Float64}`: scalar for Gram-Schmidt orthogonalization.
- `e1_scalar::Float64`: scalar for wavefield vector 1.
- `e2_scalar::Float64`: scalar for wavefield vector 2.
- `ea::EigenAngle`: wavefield eigenangle.
- `frequency::Frequency`: electromagentic wave frequency in Hz.
- `bfield::BField`: Earth's magnetic field in Tesla.
- `species::S`: species in the ionosphere.
- `params::LMPParams{T,T2,H}`: module-wide parameters.
"""
struct WavefieldIntegrationParams{S,T,T2,H}
    z::Float64
    bottomz::Float64
    ortho_scalar::ComplexF64
    e1_scalar::Float64
    e2_scalar::Float64
    ea::EigenAngle
    frequency::Frequency
    bfield::BField
    species::S
    params::LMPParams{T,T2,H}
end

"""
    WavefieldIntegrationParams(topheight, ea, frequency, bfield, species, params)

Initialize a `WavefieldIntegrationParams` for downward Pitteway scaled integration
[Pitteway1965].

Automatically set values are:

- `bottomz = BOTTOMHEIGHT`
- `ortho_scalar = zero(ComplexF64)`
- `e1_scalar = one(Float64)`
- `e2_scalar = one(Float64)`
"""
function WavefieldIntegrationParams(topheight, ea, frequency, bfield, species::S,
    params::LMPParams{T,T2,H}) where {S,T,T2,H}
    return WavefieldIntegrationParams{S,T,T2,H}(topheight, BOTTOMHEIGHT, zero(ComplexF64),
        1.0, 1.0, ea, frequency, bfield, species, params)
end

"""
    ScaleRecord

Struct for saving wavefield scaling information used in a callback during Pitteway
integration of wavefields [Pitteway1965].

# Fields

- `z::Float64`: current height in meters.
- `e::SMatrix{4,2,Complex{Float64},8}`: wavefield matrix.
- `ortho_scalar::Complex{Float64}`: scalar for Gram-Schmidt orthogonalization.
- `e1_scalar::Float64`: scalar for wavefield vector 1.
- `e2_scalar::Float64`: scalar for wavefield vector 2.
"""
struct ScaleRecord
    z::Float64
    e::SMatrix{4,2,Complex{Float64},8}
    ortho_scalar::Complex{Float64}
    e1_scalar::Float64
    e2_scalar::Float64
end

"""
    dedz(e, k, T::Tmatrix)

Compute ``de/dz = -ikTe`` where ``e = (Ex, -Ey, Z₀Hx, Z₀Hy)ᵀ``.
"""
dedz(e, k, T::TMatrix) = -1im*k*(T*e)  # `(T*e)` uses specialized TMatrix math

"""
    dedz(e, p, z)

Compute derivative of field components vector `e` at height `z`.

The parameter `p` should be a `WavefieldIntegrationParams`.
"""
function dedz(e, p, z)
    @unpack ea, frequency, bfield, species, params = p

    M = susceptibility(z, frequency, bfield, species, params=params)
    T = tmatrix(ea, M)

    return dedz(e, frequency.k, T)
end

"""
    scalingcondition(e, z, int)

Return `true` if wavefields should be scaled, otherwise `false` for wavefields `e` at height
`z` and integrator `int`.

Wavefields should be scaled if any component of `real(e)` or `imag(e)` are `>= 1`. In
addition, force scaling at `z <= bottomz` to ensure initial upgoing wave is unit amplitude.
"""
scalingcondition(e, z, int) = any(x->(real(x) >= 1 || imag(x) >= 1), e) || z <= int.p.bottomz

"""
    scale!(integrator)

Apply wavefield scaling with [`scalewavefields`](@ref) to the integrator.
"""
function scale!(integrator)
    new_e, new_orthos, new_e1s, new_e2s = scalewavefields(integrator.u)

    # Last set of scaling values
    @unpack bottomz, ea, frequency, bfield, species, params = integrator.p

    #==
    NOTE: `integrator.t` is the "time" of the _proposed_ step. Therefore, integrator.t`
    might equal `0.0`, for example, before it's actually gotten to the bottom.
    `integrator.prevt` is the last `t` on the "left" side of the `integrator`, which covers
    the local interval [`tprev`, `t`]. The "condition" is met at `integrator.t` and
    `integrator.t` is the time at which the affect occurs.
    However, it is not guaranteed that each (`tprev`, `t`) directly abuts the next (`tprev`,
    `t`).
    ==#

    integrator.p = WavefieldIntegrationParams(integrator.t,
                                              bottomz,
                                              new_orthos,
                                              new_e1s, new_e2s,
                                              ea, frequency, bfield, species, params)

    integrator.u = new_e

    return nothing
end

"""
    scalewavefields(e1, e2)
    scalewavefields(e::SMatrix{4,2})

Orthonormalize the vectors `e1` and `e2` or the columns of `e`, and also return the scaling
terms `a`, `e1_scale_val`, and `e2_scale_val` applied to the original vectors.

This first applies Gram-Schmidt orthogonalization and then scales the vectors so they each
have length 1, i.e. `norm(e1) == norm(e2) == 1`. This is the technique suggested by
[Pitteway1965] to counter numerical swamping during integration of wavefields.

# References

[Pitteway1965]: M. L. V. Pitteway, “The numerical calculation of wave-fields, reflexion
    coefficients and polarizations for long radio waves in the lower ionosphere. I.,” Phil.
    Trans. R. Soc. Lond. A, vol. 257, no. 1079, pp. 219–241, Mar. 1965,
    doi: 10.1098/rsta.1965.0004.
"""
scalewavefields

function scalewavefields(e1, e2)
    # Orthogonalize vectors `e1` and `e2` (Gram-Schmidt process)
    # `dot` for complex vectors automatically conjugates first vector
    e1_dot_e1 = real(dot(e1, e1))  # == sum(abs2.(e1)), imag(dot(e1,e1)) == 0
    a = dot(e1, e2)/e1_dot_e1  # purposefully unsigned
    e2 -= a*e1

    # Normalize `e1` and `e2`
    e1_scale_val = 1/(sqrt(e1_dot_e1))
    e2_scale_val = 1/(norm(e2))  # == 1/sqrt(dot(e2,e2))
    e1 *= e1_scale_val
    e2 *= e2_scale_val  # == normalize(e2)

    return e1, e2, a, e1_scale_val, e2_scale_val
end

function scalewavefields(e::SMatrix{4,2})
    e1, e2, a, e1_scale_val, e2_scale_val = scalewavefields(e[:,1], e[:,2])

    e = hcat(e1, e2)

    return e, a, e1_scale_val, e2_scale_val
end

"""
    unscalewavefields!(e, saved_values::SavedValues)

Unscale the integrated wavefields, a vector of fields at each integration step `e`, in place.

See also: [`unscalewavefields`](@ref)
"""
function unscalewavefields!(e, saved_values::SavedValues)
    zs = saved_values.t
    records = saved_values.saveval

    # Initialize the "reference" scaling altitude
    ref_z = last(records).z

    # Usually `osum = 0`, `prod_e1 = 1`, and `prod_e2 = 1` at initialization,
    # but we set up the fields at the ground (`last(records)`) outside the loop
    osum = last(records).ortho_scalar
    prod_e1 = last(records).e1_scalar
    prod_e2 = last(records).e2_scalar

    # Unscaling we go from the bottom up
    @inbounds for i in reverse(eachindex(e))

        # Unpack variables
        record_z = records[i].z
        scaled_e = records[i].e
        ortho_scalar = records[i].ortho_scalar
        e1_scalar = records[i].e1_scalar
        e2_scalar = records[i].e2_scalar

        z = zs[i]

        # Only update ref_z when there is both:
        # 1) we have reached the height where a new ref_z should go into effect
        # 2) there is a new ref_z
        if (z > record_z) & (record_z > ref_z)
            ref_z = record_z

            osum *= e1_scalar/e2_scalar
            osum += ortho_scalar
            prod_e1 *= e1_scalar
            prod_e2 *= e2_scalar
        end

        if i == lastindex(e)
            # Bottom doesn't require correction
            e[i] = scaled_e
        elseif z > ref_z
            # From the bottom, the first correction may not need to be applied
            # until some higher altitude
            e2 = (scaled_e[:,2] - osum*scaled_e[:,1])*prod_e2
            e1 = scaled_e[:,1]*prod_e1
            e[i] = hcat(e1,e2)
        else
            e[i] = scaled_e
        end
    end

    return nothing
end

"""
    unscalewavefields(saved_values::SavedValues)

Return the unscaled integrated wavefields originally scaled by [`scalewavefields`](@ref).

The bottom level does not get unscaled. We reference the higher levels to the bottom. The
level above the bottom level is additionally scaled by the amount that was applied to
originally get from this level down to the bottom level. The next level up (2 above the
bottom level) is scaled by the amount applied to the next level and then the bottom
level, i.e. we keep track of a cumulative correction on the way back up.

Assumes fields have been scaled by [`scalewavefields`](@ref) during integration.

See also: [`unscalewavefields!`](@ref)
"""
function unscalewavefields(saved_values::SavedValues)
    e = Vector{SArray{Tuple{4,2},Complex{Float64},2,8}}(undef, length(saved_values.saveval))

    unscalewavefields!(e, saved_values)

    return e
end

"""
    savevalues(u, t, integrator)

Return a `ScaleRecord` from `u`, `t`, and `integrator`.
"""
savevalues(u, t, integrator) = ScaleRecord(integrator.p.z,
                                           u,
                                           integrator.p.ortho_scalar,
                                           integrator.p.e1_scalar,
                                           integrator.p.e2_scalar)

"""
    integratewavefields(zs, ea::EigenAngle, frequency::Frequency, bfield::BField,
        species; params=LMPParams())

Compute wavefields vectors `e` at `zs` by downward integration over heights `zs`.

`params.wavefieldintegrationparams` is used by this function rather than
`params.integrationparams`.
"""
function integratewavefields(zs, ea::EigenAngle, frequency::Frequency, bfield::BField,
    species; params=LMPParams(), unscale=true)
    # TODO: version that updates output `e` in place

    issorted(zs; rev=true) ||
        throw(ArgumentError("`zs` should go from top to bottom of the ionosphere."))

    @unpack wavefieldintegrationparams = params
    @unpack tolerance, solver, dt, force_dtmin, maxiters = wavefieldintegrationparams

    # Initial conditions
    Mtop = susceptibility(first(zs), frequency, bfield, species; params=params)
    Ttop = tmatrix(ea, Mtop)
    e0 = bookerwavefields(Ttop)

    # Works best if `e0` fields are normalized at top height
    # Don't want to orthogonalize `e0[:,2]` here because it won't be undone!
    e0 = hcat(normalize(e0[:,1]), normalize(e0[:,2]))

    # saved_positions=(true, true) because we discontinuously modify `u`.
    # This is independent of `saveat` and `save_everystep`
    cb = DiscreteCallback(scalingcondition, scale!; save_positions=(true, true))

    saved_values = SavedValues(Float64, ScaleRecord)
    scb = SavingCallback(savevalues, saved_values; save_everystep=false, saveat=zs, tdir=-1)

    p = WavefieldIntegrationParams(params.topheight, ea, frequency, bfield, species, params)

    # WARNING: Without `lazy=false` (necessary since we're using DiscreteCallback) don't
    # use continuous solution output! Also, we're only saving at zs.
    prob = ODEProblem{false}(dedz, e0, (first(zs), last(zs)), p)
    sol = solve(prob, solver; callback=CallbackSet(cb, scb),
                save_everystep=false, save_start=false, save_end=false,
                dt=dt, force_dtmin=force_dtmin, maxiters=maxiters,
                atol=tolerance, rtol=tolerance)

    if unscale
        e = unscalewavefields(saved_values)
    else
        records = saved_values.saveval
        e = [records[i].e for i in eachindex(records)]
    end

    return e
end

"""
    boundaryscalars(R, Rg, e1, e2, isotropic::Bool=false)
    boundaryscalars(R, Rg, e, isotropic::Bool=false)

Compute coefficients `(b1, b2)` required to sum the two wavefields vectors `e1` and `e2` or
both columns of `e` for the total wavefield at the ground as ``e = b1*e1 + b2*e2``.

The `b1` and `b2` that satisfy the waveguide boundary conditions are only valid for true
eigenangles of the waveguide.

!!! note

    This function assumes that the reflection coefficients `R` and `Rg` and the wavefield
    vectors `e1`, `e2` are at the ground.

    These boundary conditions
"""
function boundaryscalars(R, Rg, e1, e2, isotropic::Bool=false)
    # This is similar to excitation factor calculation, using the waveguide mode condition

    ex1, ey1, hx1, hy1 = e1[1], -e1[2], e1[3], e1[4]  # the ey component was originally -Ey
    ex2, ey2, hx2, hy2 = e2[1], -e2[2], e2[3], e2[4]

    # NOTE: Because `R`s and `e`s are at the ground, `hy0` and `ey0` are the same regardless
    # of `earthcurvature` being `true` or `false`.
    hy0 = 1  # == (1 + Rg[1,1])/(1 + Rg[1,1])
    ey0 = 1  # == (1 + Rg[2,2])/(1 + Rg[2,2])

    paral = 1 - R[1,1]*Rg[1,1]
    perp = 1 - R[2,2]*Rg[2,2]

    if isotropic
        if abs2(paral) < abs2(perp)
            b1 = hy0/hy1
            b2 = zero(hy1)
            # eyg = 0
        else
            b1 = zero(ey2)
            b2 = ey0/ey2
            # hyg = 0
        end
    else
        # polarization ratio Ey/Hy (often `f` or `fofr` in papers)
        if abs2(paral) < abs2(perp)
            EyHy = ((1 + Rg[2,2])*R[2,1]*Rg[1,1])/((1 + Rg[1,1])*perp)
        else
            EyHy = ((1 + Rg[2,2])*paral)/((1 + Rg[1,1])*R[1,2]*Rg[2,2])
        end
        a = (-ey1 + EyHy*hy1)/(ey2 - EyHy*hy2)
        hysum = hy1 + a*hy2
        b1 = hy0/hysum
        b2 = b1*a
    end

    return b1, b2
end

boundaryscalars(R, Rg, e, isotropic::Bool=false) =
    boundaryscalars(R, Rg, e[:,1], e[:,2], isotropic)

"""
    fieldstrengths!(EH, zs, me::ModeEquation; params=LMPParams())
    fieldstrengths!(EH, zs, ea::EigenAngle, frequency::Frequency, bfield::BField,
        species, ground::Ground; params=LMPParams())

Compute ``(Ex, Ey, Ez, Hx, Hy, Hz)ᵀ`` wavefields vectors as elements of `EH` by fullwave
integration at each height in `zs`.

The wavefields are scaled to satisfy the waveguide boundary conditions, which is only valid
at solutions of the mode equation.
"""
function fieldstrengths!(EH, zs, ea::EigenAngle, frequency::Frequency, bfield::BField,
    species, ground::Ground; params=LMPParams())

    @assert length(EH) == length(zs)
    zs[end] == 0 || @warn "Waveguide math assumes fields and reflection coefficients are
        calculated at the ground (`z = 0`)."

    e = integratewavefields(zs, ea, frequency, bfield, species; params=params)
    R = bookerreflection(ea, e[end])
    Rg = fresnelreflection(ea, ground, frequency)
    b1, b2 = boundaryscalars(R, Rg, e[end], isisotropic(bfield))

    S = ea.sinθ
    @inbounds for i in eachindex(EH)
        M = susceptibility(zs[i], frequency, bfield, species; params=params)

        # Scale to waveguide boundary conditions
        w = e[i][:,1]*b1 + e[i][:,2]*b2

        ex = w[1]
        ey = -w[2]
        ez = -(w[4]*S + M[3,1]*w[1] - M[3,2]*w[2])/(1 + M[3,3])
        hx = w[3]
        hy = w[4]
        hz = -w[2]*S

        EH[i] = SVector(ex, ey, ez, hx, hy, hz)
    end

    return nothing
end

function fieldstrengths!(EH, zs, me::ModeEquation; params=LMPParams())
    @unpack ea, frequency, waveguide = me
    @unpack bfield, species, ground = waveguide

    fieldstrengths!(EH, zs, ea, frequency, bfield, species, ground; params=params)
end

"""
    fieldstrengths(zs, me::ModeEquation; params=LMPParams())

Preallocate vector of wavefields `EH`, then call [`fieldstrengths!`](@ref) and return `EH`.

Each element of `EH` is an `SVector` of ``ex, ey, ez, hx, hy, hz``.
"""
function fieldstrengths(zs, me::ModeEquation; params=LMPParams())
    EH = Vector{SVector{6, ComplexF64}}(undef, length(zs))
    fieldstrengths!(EH, zs, me; params=params)
    return EH
end

"""
    calculate_wavefields!(wavefields, adjoint_wavefields, frequency, waveguide,
        adjoint_waveguide; params=LMPParams())

Compute fields of `wavefields` in-place scaled to satisfy the `waveguide` boundary
conditions.

This function implements the method of integrating wavefields suggested by
[Pitteway1965].

# References

[Pitteway1965]: M. L. V. Pitteway, “The numerical calculation of wave-fields, reflexion
    coefficients and polarizations for long radio waves in the lower ionosphere. I.,” Phil.
    Trans. R. Soc. Lond. A, vol. 257, no. 1079, pp. 219–241, Mar. 1965,
    doi: 10.1098/rsta.1965.0004.
"""
function calculate_wavefields!(wavefields, adjoint_wavefields, frequency, waveguide,
    adjoint_waveguide; params=LMPParams())

    @unpack bfield, species, ground = waveguide

    ground == adjoint_waveguide.ground ||
        throw(ArgumentError("waveguide and adjoint_waveguide should have same Ground"))
    species == adjoint_waveguide.species ||
        throw(ArgumentError("waveguide and adjoint_waveguide should have same Species"))

    adjoint_bfield = adjoint_waveguide.bfield

    zs = heights(wavefields)
    modes = eigenangles(wavefields)
    adjoint_modes = eigenangles(adjoint_wavefields)
    modes == adjoint_modes || @warn "Full mode conversion physics assumes adjoint
        wavefields are calculated for eigenangles of the original waveguide."

    for m in eachindex(modes)
        fieldstrengths!(view(wavefields,:,m), zs, modes[m], frequency, bfield, species,
                        ground; params=params)
        fieldstrengths!(view(adjoint_wavefields,:,m), zs, adjoint_modes[m], frequency,
                        adjoint_bfield, species, ground; params=params)
    end

    return nothing
end