text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

"""
    absorption(glass::AbstractGlass, wavelength; temperature=20°C, pressure=1Atm)

Compute the intensity absorption per mm of `glass` at `wavelength`, optionally at specified `temperature` and `pressure`.
Transmission values are linearly interpolated from the adjacent values in the data table of `glass`, if `wavelength` is below the minimum or above the maximum in the table then the nearest value is taken.

Absorption is defined as ``\\frac{-\\log(t)}{\\tau}`` where ``t`` is the transmission value and ``\\tau`` is the thickness, both of which are provided in the data table.

If unitless, arguments are interpretted as μm, °C and Atm respectively.

# Examples
```julia-repl
julia> absorption(GlassCat.SUMITA.LAK7, 700u"nm")
0.0006018072325563021

julia> absorption(GlassCat.SCHOTT.N_BK7, 0.55, temperature = 22.0)
0.00016504471175660636

julia> absorption(GlassCat.SCHOTT.PSK3, 532u"nm", temperature = 25u"°C", pressure = 1.3)
0.00020855284788532435
```
"""
function absorption(glass::Glass, wavelength::Length; temperature::Temperature = TEMP_REF_UNITFUL, pressure::Float64 = PRESSURE_REF)::Float64
    λ = Float64(ustrip(u"μm", wavelength))
    return absorption(glass, λ, temperature = ustrip(Float64, u"°C", temperature), pressure = pressure)
end

function absorption(glass::Glass, λ::T; temperature::T = T(TEMP_REF), pressure::T = T(PRESSURE_REF))::T where {T<:Real}
    # if the glass has no transmission data then assume no absorption
    if glass.transmission === nothing
        return zero(T)
    end

    reference_temp = T(glass.temp)

    # to work out the wavelength at the reference temperature we need the RIs of air at system temp and at reference temp
    n_air_at_sys = absairindex(λ, temperature = temperature, pressure = pressure)
    n_air_at_ref = absairindex(λ, temperature = reference_temp)

    # scale the wavelength to air at the reference temperature/pressure
    λ = λ * (n_air_at_sys / n_air_at_ref)

    tdata = glass.transmission
    N = glass.transmissionN
    if λ < tdata[1][1]
        t = tdata[1][2]
        τ = tdata[1][3]
        return T(-log1p(t - 1.0) / τ)
    elseif λ > tdata[N][1]
        t = tdata[N][2]
        τ = tdata[N][3]
        return T(-log1p(t - 1.0) / τ)
    else
        let λlow = 0.0, tlow = 0.0, τlow = 0.0, λhigh = 0.0, thigh = 0.0, τhigh = 0.0
            for i in 2:N
                if λ <= tdata[i][1]
                    λlow, tlow, τlow = tdata[i - 1]
                    λhigh, thigh, τhigh = tdata[i]
                    break
                end
            end
            λhigh = T(λhigh)
            λlow = T(λlow)
            δλ = λhigh - λlow
            @assert τlow == τhigh
            t = (tlow * (λhigh - λ) / δλ) + (thigh * (λ - λlow) / δλ)
            return -log1p(t - 1.0) / τhigh
        end
    end
end

function absorption(::AirType, ::Length; temperature::Temperature = TEMP_REF_UNITFUL, pressure::Float64 = PRESSURE_REF)::Float64
    return 0.0
end

function absorption(::AirType, ::T; temperature::T = T(TEMP_REF), pressure::T = T(PRESSURE_REF))::T where {T<:Real}
    return zero(T)
end

"""
    index(glass::AbstractGlass, wavelength; temperature=20°C, pressure=1Atm)

Compute the refractive index of `glass` at `wavelength`, optionally at specified `temperature` and `pressure`.
Result is relative to the refractive index of air at given temperature and pressure.

If unitless, arguments are interpretted as μm, °C and Atm respectively.

**This is defined to always equal 1.0 for Air at any temperature and pressure**, use [`absairindex`](@ref) for the absolute refractive index of air at a given temperature and pressure.

# Examples
```julia-repl
julia> index(GlassCat.SUMITA.LAK7, 700u"nm")
1.646494204478318

julia> index(GlassCat.SCHOTT.N_BK7, 0.55, temperature = 22.0)
1.51852824383283

julia> index(GlassCat.HOYA.FF1, 532u"nm", temperature = 25u"°C", pressure = 1.3)
1.5144848290944655
```
"""
function index(glass::Glass, wavelength::Length; temperature::Temperature = TEMP_REF_UNITFUL, pressure::Float64 = PRESSURE_REF)::Float64
    λ = Float64(ustrip(uconvert(u"μm", wavelength)))
    return index(glass, λ, temperature = ustrip(Float64, u"°C", temperature), pressure = pressure)
end

function index(glass::Glass, λ::T; temperature::T = T(TEMP_REF), pressure::T = T(PRESSURE_REF))::T where {T<:Real}
    # all calculations for the material must be done at the refernce temperature
    reference_temp = T(glass.temp)

    # to work out the wavelength at the reference temperature we need the RIs of air at system temp and at reference temp
    n_air_at_sys = absairindex(λ, temperature = temperature, pressure = pressure)
    n_air_at_ref = absairindex(λ, temperature = reference_temp)

    # scale the wavelength to air at the reference temperature/pressure
    λabs = λ * n_air_at_sys
    λ = λabs / n_air_at_ref

    if (λ < glass.λmin) || (λ > glass.λmax)
        error("Cannot calculate an index for the specified wavelength: $λ, valid range: [$(glass.λmin), $(glass.λmax)].\n")
    end

    if glass.dispform == -2
        # Cauchy
        n_rel = T(glass.C1) + (glass.C2 * λ^(-2)) + (glass.C3 * λ^(-4)) + (glass.C4 * λ^(-6)) + (glass.C5 * λ^(-8)) + (glass.C6 * λ^(-10))
    elseif glass.dispform == -1
        # use fitted result from GOptical:
        n_rel = T(glass.Nd + (glass.Nd - one(T)) / glass.Vd * (glass.C1 + glass.C2 / λ + glass.C3 / λ^2 + glass.C4 / λ^3))
    elseif glass.dispform == 1
        # Schott
        formula_rhs = T(glass.C1) + (glass.C2 * λ^2) + (glass.C3 * λ^(-2)) + (glass.C4 * λ^(-4)) + (glass.C5 * λ^(-6)) + (glass.C6 * λ^(-8))
        n_rel = sqrt(formula_rhs)
    elseif glass.dispform == 2
        # Sellmeier1
        formula_rhs = (glass.C1 * λ^2 / (λ^2 - glass.C2)) + (glass.C3 * λ^2 / (λ^2 - glass.C4)) + (glass.C5 * λ^2 / (λ^2 - glass.C6))
        n_rel = sqrt(formula_rhs + one(T))
    elseif glass.dispform == 3
        # Herzberger
        L = one(T) / (λ^2 - T(0.028))
        n_rel = T(glass.C1) + (glass.C2 * L) + (glass.C3 * L^2) + (glass.C4 * λ^2) + (glass.C5 * λ^4) + (glass.C6 * λ^6)
    elseif glass.dispform == 4
        # Sellmeier2
        formula_rhs = T(glass.C1) + (glass.C2 * λ^2 / (λ^2 - (glass.C3)^2)) + (glass.C4 * λ^2 / (λ^2 - (glass.C5)^2))
        n_rel = sqrt(formula_rhs + one(T))
    elseif glass.dispform == 5
        # Conrady
        n_rel = T(glass.C1) + (glass.C2 / λ) + (glass.C3 / λ^3.5)
    elseif glass.dispform == 6
        # Sellmeier3
        formula_rhs = (glass.C1 * λ^2 / (λ^2 - glass.C2)) + (glass.C3 * λ^2 / (λ^2 - glass.C4)) + (glass.C5 * λ^2 / (λ^2 - glass.C6)) + (glass.C7 * λ^2 / (λ^2 - glass.C8))
        n_rel = sqrt(formula_rhs + one(T))
    elseif glass.dispform == 7
        # HandbookOfOptics1
        formula_rhs = T(glass.C1) + (glass.C2 / (λ^2 - glass.C3)) - (glass.C4 * λ^2)
        n_rel = sqrt(formula_rhs)
    elseif glass.dispform == 8
        # HandbookOfOptics2
        formula_rhs = T(glass.C1) + (glass.C2 * λ^2 / (λ^2 - glass.C3)) - (glass.C4 * λ^2)
        n_rel = sqrt(formula_rhs)
    elseif glass.dispform == 9
        # Sellmeier4
        formula_rhs = T(glass.C1) + (glass.C2 * λ^2 / (λ^2 - glass.C3)) + (glass.C4 * λ^2 / (λ^2 - glass.C5))
        n_rel = sqrt(formula_rhs)
    elseif glass.dispform == 10
        # Extended1
        formula_rhs = T(glass.C1) + (glass.C2 * λ^2) + (glass.C3 * λ^(-2)) + (glass.C4 * λ^(-4)) + (glass.C5 * λ^(-6)) + (glass.C6 * λ^(-8)) + (glass.C7 * λ^(-10)) + (glass.C8 * λ^(-12))
        n_rel = sqrt(formula_rhs)
    elseif glass.dispform == 11
        # Sellmeier5
        formula_rhs = (glass.C1 * λ^2 / (λ^2 - glass.C2)) + (glass.C3 * λ^2 / (λ^2 - glass.C4)) + (glass.C5 * λ^2 / (λ^2 - glass.C6)) + (glass.C7 * λ^2 / (λ^2 - glass.C8)) + (glass.C9 * λ^2 / (λ^2 - glass.C10))
        n_rel = sqrt(formula_rhs + one(T))
    elseif glass.dispform == 12
        # Extended2
        formula_rhs = T(glass.C1) + (glass.C2 * λ^2) + (glass.C3 * λ^(-2)) + (glass.C4 * λ^(-4)) + (glass.C5 * λ^(-6)) + (glass.C6 * λ^(-8)) + (glass.C7 * λ^4) + (glass.C8 * λ^6)
        n_rel = sqrt(formula_rhs)
    elseif glass.dispform == 13
        # Extended3
        formula_rhs = T(glass.C1) + (glass.C2 * λ^2) + (glass.C3 * λ^(4)) + (glass.C4 * λ^(-2)) + (glass.C5 * λ^(-4)) + (glass.C6 * λ^(-6)) + (glass.C7 * λ^(-8)) + (glass.C8 * λ^(-10)) + (glass.C9 * λ^(-12))
        n_rel = sqrt(formula_rhs)
    else
        @error "Invalid glass dispersion formula"
    end

    # get the absolute index of the material
    n_abs = n_rel * n_air_at_ref

    # If "TD" is included in the glass data, then include pressure and temperature dependence of the lens
    # environment. From Schott"s technical report "TIE-19: Temperature Coefficient of the Refractive Index".
    # The above "n_rel" data are assumed to be from the reference temperature T_ref. Now we add a small change
    # delta_n to it due to a change in temperature.
    ΔT = temperature - reference_temp
    if !isnan(glass.D₀) && abs(ΔT) > 0.0 && (glass.D₀ != 0 || glass.D₁ != 0 || glass.D₂ != 0 || glass.E₀ != 0 || glass.E₁ != 0)
        Sₜₖ = glass.λₜₖ < 0.0 ? -one(T) : one(T)
        Δn_abs = ((n_rel^2 - one(T)) / (2.0 * n_rel)) * (glass.D₀ * ΔT + glass.D₁ * ΔT^2 + glass.D₂ * ΔT^3 + ((glass.E₀ * ΔT + glass.E₁ * ΔT^2) / (λ^2 - Sₜₖ * glass.λₜₖ^2)))
        n_abs = n_abs + Δn_abs
    end

    # make the index relative to the RI of the air at the system temperature/pressure again
    n_rel = n_abs / n_air_at_sys
    return n_rel
end

function index(::AirType, ::Length; temperature::Temperature = TEMP_REF_UNITFUL, pressure::Float64 = PRESSURE_REF)::Float64
    return 1.0
end

function index(::AirType, ::T; temperature::T = T(TEMP_REF), pressure::T = T(PRESSURE_REF))::T where {T<:Real}
    return one(T)
end

"""
    absairindex(wavelength; temperature=20°C, pressure=1Atm)

Compute the absolute refractive index of air at `wavelength`, optionally at specified `temperature` and `pressure`. If unitless, arguments are interpretted as μm, °C and Atm respectively.

# Examples
```julia-repl
julia> absairindex(700u"nm")
1.000271074905147

julia> absairindex(0.7, temperature=27.0)
1.000264738846504

julia> absairindex(532u"nm", temperature = 25u"°C", pressure = 1.3)
1.0003494991178161
```
"""
function absairindex(wavelength::Length; temperature::Temperature = TEMP_REF_UNITFUL, pressure::Float64 = PRESSURE_REF)::Float64
    # convert to required units
    λ = Float64(ustrip(uconvert(u"μm", wavelength)))
    return absairindex(λ, temperature = ustrip(Float64, u"°C", temperature), pressure = pressure)
end

function absairindex(λ::T; temperature::T = T(TEMP_REF), pressure::T = T(PRESSURE_REF))::T where {T<:Real}
    # convert to required units
    n_ref = one(T) + ((6432.8 + ((2949810.0 * λ^2) / (146.0 * λ^2 - one(T))) + ((25540.0 * λ^2) / (41.0 * λ^2 - one(T)))) * 1e-8)
    n_rel = one(T) + ((n_ref - one(T)) / (one(T) + (temperature - 15.0) * 0.0034785)) * (pressure / PRESSURE_REF)
    return n_rel
end

"""
    polyfit_indices(wavelengths, n_rel; degree=5)

Fit a polynomial to `indices` at `wavelengths`, optionally specifying the `degree` of the polynomial.
Returns tuple of array of fitted indices at wavelengths and the polynomial.
"""
function polyfit_indices(wavelengths::Union{AbstractRange{<:Length},AbstractArray{<:Length,1}}, indices::AbstractArray{<:Number,1}; degree::Int = 5)
    w = ustrip.(uconvert.(u"μm", wavelengths))
    okay = (indices .> 0.0)
    if !any(okay)
        return (ones(Float64, size(w)) .* NaN, nothing)
    end
    xs = range(-1.0, stop = 1.0, length = length(w[okay]))
    poly = fit(xs, indices[okay], degree)
    interp_indices = poly.(xs)
    # ensure output has all entries
    out = ones(Float64, size(w)) .* NaN
    out[okay] = interp_indices
    return (out, poly)
end

"""
    plot_indices(glass::AbstractGlass; polyfit=false, fiterror=false, degree=5, temperature=20°C, pressure=1Atm, nsamples=300, sampling_domain="wavelength")

Plot the refractive index for `glass` for `nsamples` within its valid range of wavelengths, optionally at `temperature` and `pressure`.
`polyfit` will show a polynomial of optionally specified `degree` fitted to the data, `fiterror` will also show the fitting error of the result.
`sampling_domain` specifies whether the samples will be spaced uniformly in "wavelength" or "wavenumber".
"""
function plot_indices(glass::AbstractGlass; polyfit::Bool = false, fiterror::Bool = false, degree::Int = 5, temperature::Temperature = TEMP_REF_UNITFUL, pressure::Float64 = PRESSURE_REF, nsamples::Int = 300, sampling_domain::String = "wavelength")
    if isair(glass)
        wavemin = 380 * u"nm"
        wavemax = 740 * u"nm"
    else
        wavemin = glass.λmin * u"μm"
        wavemax = glass.λmax * u"μm"
    end

    if (sampling_domain == "wavelength")
        waves = range(wavemin, stop = wavemax, length = nsamples)      # wavelength in um
    elseif (sampling_domain == "wavenumber")
        sigma_min = 1.0 / wavemax
        sigma_max = 1.0 / wavemin
        wavenumbers = range(sigma_min, stop = sigma_max, length = nsamples) # wavenumber in um.^-1
        waves = 1.0 ./ wavenumbers
    else
        error("Invalid sampling domain, should be \"wavelength\" or \"wavenumber\"")
    end

    p = plot(xlabel = "wavelength (um)", ylabel = "refractive index")

    f = w -> begin
        try
            return index(glass, w, temperature = temperature, pressure = pressure)
        catch
            return NaN
        end
    end
    indices = [f(w) for w in waves]
    plot!(ustrip.(waves), indices, color = :blue, label = "From Data")

    if polyfit
        (p_indices, _) = polyfit_indices(waves, indices, degree = degree)
        plot!(ustrip.(waves), p_indices, color = :black, markersize = 4, label = "Polyfit")
    end

    if polyfit && fiterror
        err = p_indices - indices
        p2 = plot(xlabel = "wavelength (um)", ylabel = "fit error")
        plot!(ustrip.(waves), err, color = :red, label = "Fit Error")
        p = plot(p, p2, layout = 2)
    end

    plot!(title = "$(glassname(glass)) dispersion")

    gui(p)
end

"""
    drawglassmap(glasscatalog::Module; λ::Length = 550nm, glassfontsize::Integer = 3, showprefixglasses::Bool = false)

Draw a scatter plot of index vs dispersion (the derivative of index with respect to wavelength). Both index and
dispersion are computed at wavelength λ.

Choose glasses to graph using the glassfilterprediate argument. This is a function that receives a Glass object and returns true if the glass should be graphed.

If showprefixglasses is true then glasses with names like `F_BAK7` will be displayed. Otherwise glasses that have a
leading letter prefix followed by an underscore, such as `F_`, will not be displayed.

The index formulas for some glasses may give incorrect results if λ is outside the valid range for that glass. This can
give anomalous results, such as indices less than zero or greater than 6. To filter out these glasses set maximumindex
to a reasonable value such as 3.0.

example: plot only glasses that do not contain the strings "E_" and "J_"

drawglassmap(NIKON,showprefixglasses = true,glassfilterpredicate = (x) -> !occursin("J_",string(x)) && !occursin("E_",string(x)))
"""
function drawglassmap(glasscatalog::Module; λ::Length = 550nm, glassfontsize::Integer = 3, showprefixglasses::Bool = false, minindex = 1.0, maxindex = 3.0, mindispersion = -.3, maxdispersion = 0.0, glassfilterpredicate = (x)->true)
    wavelength = Float64(ustrip(uconvert(μm, λ)))
    indices = Vector{Float64}(undef,0)
    dispersions = Vector{Float64}(undef,0)
    glassnames = Vector{String}(undef,0)

    for name in names(glasscatalog)
        glass = eval(:($glasscatalog.$name))
        glassstring = String(name)
        hasprefix = occursin("_", glassstring)

        if typeof(glass) !== Module && (minindex <= index(glass, wavelength) <= maxindex)
            f(x) = index(glass,x)
            g = x -> ForwardDiff.derivative(f, x);
            dispersion = g(wavelength)

            # don't show glasses that have an _ in the name. This prevents cluttering the map with many glasses of
            # similar (index, dispersion).
            if glassfilterpredicate(glass) && (mindispersion <= dispersion <= maxdispersion) && (showprefixglasses || !hasprefix)
                push!(indices, index(glass, wavelength))
                push!(dispersions, dispersion)
                push!(glassnames, String(name))
            end
        end
    end

    font = Plots.font(family = "Sans", pointsize = glassfontsize, color = RGB(0.0,0.0,.4))
    series_annotations = Plots.series_annotations(glassnames, font)
    scatter(
        dispersions,
        indices;
        series_annotations,
        markeralpha = 0.0,
        legends = :none,
        xaxis = "dispersion @$λ",
        yaxis = "index",
        title = "Glass Catalog: $glasscatalog",
        xflip = true) #should use markershape = :none to prevent markers from being drawn but this option doesn't work. Used markeralpha = 0 so the markers are invisible. A hack which works.
end


text:
# Create example data
x = LinRange(-2, 2, 40)
y = x.^3 .+ x.^2 .+ x
# Plot x and y
stair(x, y)
# Plot y with indices for x values
stair(y)
# step directly after x each position
stair(y, where="pre")
# step between two x positions
stair(y, where="mid")
# step immediately before x each position
stair(y, where="post")

ext: jl


text:
abstract type AbstractBC end

# Periodic BC
struct PeriodicGridBC <: AbstractBC
  lower::Float64
  upper::Float64
end
PeriodicGridBC(x) = PeriodicGridBC(0.0, x)
lower(p::PeriodicGridBC) = p.lower
upper(p::PeriodicGridBC) = p.upper

Base.length(p::PeriodicGridBC) = p.upper - p.lower
Base.in(x::Number, p::PeriodicGridBC) = p.lower <= x < p.upper
Base.:<(x::Number, p::PeriodicGridBC) = x < p.lower
Base.:>(x::Number, p::PeriodicGridBC) = x > p.upper

function (p::PeriodicGridBC)(x::Number)
  @assert isfinite(x)
  x ∈ p && return x
  x < p && return p(x + length(p))
  x > p && return p(x - length(p))
  @show x, p
  throw(error("Should never reach here, $p, $x"))
end


text:
# test_problems.jl : collection of test problems

# example rhs from A Multigrid Tutorial, scnd. ed.
h(x,y) = 2((1-6x^2)*y^2*(1-y^2)+(1-6y^2)*x^2*(1-x^2))

# default test examples
g(x) = Float64[1 + sin(π*xᵢ) for xᵢ in x]
g(x,y) = Float64[1 + sin(π*xᵢ)*sin(π*yᵢ) for xᵢ in x, yᵢ in y]
g(x,y,z) = Float64[1 + sin(π*xᵢ)*sin(π*yᵢ)*sin(π*zᵢ) for xᵢ in x, yᵢ in y, zᵢ in z]

f(n) = g(range(1/n, stop=1-1/n, length=n-1))
f(n,m) = g(range(1/n, stop=1-1/n, length=n-1),range(1/m, stop=1-1/m, length=m-1))
f(n,m,l) = g(range(1/n, stop=1-1/n, length=n-1),range(1/m, stop=1-1/m, length=m-1),range(1/l, stop=1-1/l, length=l-1))

p₁₁(n) = laplace1d(n)
p₁₂(n) = laplace2d(n,n)
p₁₃(n) = laplace3d(n,n,n)

p₂₁(n) = elliptic1d(f(n))
p₂₂(n) = elliptic2d(f(n,n))
p₂₃(n) = elliptic3d(f(n,n,n))


text:
using Term
import Term.Tables: Table
import Term.Layout: PlaceHolder

"""
    This is an example showing how to create tables in Term.jl

There's MANY options to tables, so we can't cover all of them here, head to the docs
for more info.
"""

t = 1:5
data = hcat(t, ones(length(t)), rand(Int8, length(t)))

"""
You can create data from `Matrix` and `Vector` objects
"""

print(Table(data))

"""
You can justify the table headers
"""

print(
    Table(
        data;
        header = ["Num", "Const.", "Values"],
        header_style = "bold white",
        columns_style = ["dim", "bold", "red"],
    ),
)

"""
For most of the options, you can either pass a single value
(which gets applied to a whole column/row) or a vector of values
to e.g. specify the style of each column indpendently.

You can use vertical/horizontal padding to space out the table.
"""

print(
    Table(
        data;
        header = ["Num", "Const.", "Values"],
        header_style = "bold white",
        columns_style = ["dim", "bold", "red"],
        hpad = [1, 2, 5],
        columns_justify = [:center, :right, :left],
        vpad = 1,
    ),
)

"""
You can also specify footers
"""

print(Table(data; footer = ["get", "a", "footer"], footer_justify = :center))

"""
    Or even use a function that gets applied column wise
"""

print(Table(data; footer = sum, footer_justify = :center, footer_style = "dim bold"))

"""
You can use different `Box` types to change the table style.
And you can create layouts with tables as you would with any
other renderables.
"""

print(
    (Table(data; box = :ROUNDED, style = "red") * " " * Table(data; box = :HEAVY)) /
    Table(data; box = :SIMPLE_HEAVY, style = "dim blue"),
)

"""
Tables can also be created from `Dict` objects and can include other
renderables.
"""

ph1 = PlaceHolder(25, 5)
ph2 = PlaceHolder(23, 9)
ph3 = PlaceHolder(22, 11)

data = Dict(
    "first\ncol." => [ph1, ph2, ph3],
    "second\ncol." => [ph3, ph2, ph3],
    "third\ncol." => [ph2, ph2, ph1],
)

print(Table(data))


text:
function _randtensors(S)
    AL = TensorMap(randn, ComplexF64, S * S, S * S)
    M = TensorMap(randn, ComplexF64, S * S * S * S * S * S, S * S * S * S * S * S)
    return AL, M
    # return PEPOKit.leftenv(AL, M)
end

function leftgauge(
    A::Array{T,4},
    C::Matrix{T}=Matrix{eltype(A)}(I, size(A, 1), size(A, 1));
    tol=1e-12,
    maxiter=100,
    kwargs...,
) where {T}
    D = size(A)[1]
    d = size(A)[2]
    AL, C, λ = leftgauge(reshape(A, (D, d * d, D)), C; tol=tol, maxiter=maxiter, kwargs...)
    return reshape(AL, (D, d, d, D)), C, λ
end
function leftgauge(
    A::Array{T,3},
    C::Matrix{T}=Matrix{eltype(A)}(I, size(A, 1), size(A, 1));
    tol=1e-12,
    maxiter=100,
    kwargs...,
) where {T}
    λ2s, ρs, info =
        eigsolve(C' * C, 1, :LM; ishermitian=false, tol=tol, maxiter=1, eager=true) do ρ
            @tensor ρE[a, b] := ρ[a', b'] * A[b', s, b] * conj(A[a', s, a])
            return ρE
        end
    ρ = ρs[1] + ρs[1]'
    ρ ./= tr(ρ)
    # C = cholesky!(ρ).U
    # If ρ is not exactly positive definite, cholesky will fail
    F = svd!(ρ)
    C = lmul!(Diagonal(sqrt.(F.S)), F.Vt)
    _, C = qrpos!(C)

    D, d, = size(A)
    Q, R = qrpos!(reshape(C * reshape(A, D, d * D), D * d, D))
    AL = reshape(Q, D, d, D)
    λ = norm(R)
    rmul!(R, 1 / λ)
    numiter = 1
    while norm(C - R) > tol && numiter < maxiter
        # C = R
        λs, Cs, info =
            eigsolve(R, 1, :LM; ishermitian=false, tol=tol, eager=true, maxiter=1) do X
                @tensor Y[a, b] := X[a', b'] * A[b', s, b] * conj(AL[a', s, a])
                return Y
            end
        _, C = qrpos!(Cs[1])
        # The previous lines can speed up the process when C is still very far from the correct
        # gauge transform, it finds an improved value of C by finding the fixed point of a
        # 'mixed' transfer matrix composed of `A` and `AL`, even though `AL` is also still not
        # entirely correct. Therefore, we restrict the number of iterations to be 1 and don't
        # check for convergence
        Q, R = qrpos!(reshape(C * reshape(A, D, d * D), D * d, D))
        AL = reshape(Q, D, d, D)
        λ = norm(R)
        rmul!(R, 1 / λ)
        numiter += 1
    end
    C = R
    return AL, C, λ
end

function rightgauge(
    A::Array{T,4},
    C::Matrix{T}=Matrix{eltype(A)}(I, size(A, 1), size(A, 1));
    tol=1e-12,
    maxiter=100,
    kwargs...,
) where {T}
    D = size(A)[1]
    d = size(A)[2]
    AL, C, λ = rightgauge(reshape(A, (D, d * d, D)), C; tol=tol, maxiter=maxiter, kwargs...)
    return reshape(AL, (D, d, d, D)), C, λ
end
function rightgauge(
    A::Array{T,3},
    C::Matrix{T}=Matrix{eltype(A)}(I, size(A, 1), size(A, 1));
    tol=1e-12,
    maxiter=100,
    kwargs...,
) where {T}
    AL, C, λ = leftgauge(
        permutedims(A, (3, 2, 1)), permutedims(C, (2, 1)); tol=tol, maxiter=maxiter, kwargs
    )
    return permutedims(AL, (3, 2, 1)), permutedims(C, (2, 1)), λ
end

function applyH1(
    AC::Array{Complex{T},3},
    FL::Array{Complex{T},3},
    FR::Array{Complex{T},3},
    M::Union{Array{T,4},Array{Complex{T},4}},
) where {T}
    @tensor HAC[α, s, β] := FL[α, a, α'] * AC[α', s', β'] * M[a, s, b, s'] * FR[β', b, β]
end

function applyH0(C::Matrix{T}, FL::Array{T,3}, FR::Array{T,3}) where {T}
    @tensor HC[α, β] := FL[α, a, α'] * C[α', β'] * FR[β', a, β]
end

function leftenv(
    A::Array{Complex{T},3},
    M::Union{Array{T,4},Array{Complex{T},4}},
    FL::Array{Complex{T},3}=randn(eltype(A), size(A, 1), size(M, 1), size(A, 1));
    eager=true,
    kwargs...,
) where {T}
    λs, FLs, info = eigsolve(FL, 1, :LM; eager=true, ishermitian=false) do FL
        @tensor FL[α, a, β] :=
            FL[α', a', β'] * A[β', s', β] * M[a', s, a, s'] * conj(A)[α', s, α] # try a a'
    end
    return FLs[1], λs[1], info
end

function rightenv(
    A::Array{Complex{T},3},
    M::Union{Array{T,4},Array{Complex{T},4}},
    FR::Array{Complex{T},3}=randn(eltype(A), size(A, 1), size(M, 1), size(A, 1));
    eager=true,
    kwargs...,
) where {T}
    λs, FRs, info = eigsolve(FR, 1, :LM; eager=true, ishermitian=false) do FR
        @tensor FR[α, a, β] :=
            A[α, s', α'] * FR[α', a', β'] * M[a, s, a', s'] * conj(A[β, s, β'])
    end
    # ---> 1
    # FR-> 2
    # ---> 3
    return FRs[1], λs[1], info
end

safesign(x::Number) = iszero(x) ? one(x) : sign(x)

qrpos(A) = qrpos!(copy(A))
function qrpos!(A::Matrix{T}) where {T}
    F = qr!(A)
    Q = Matrix(F.Q)
    R = F.R
    phases = safesign.(diag(R))
    rmul!(Q, Diagonal(phases))
    lmul!(Diagonal(conj!(phases)), R)
    return Q, R
end

lqpos(A) = lqpos!(copy(A))
function lqpos!(A::Matrix{T}) where {T}
    F = qr!(Matrix(A'))
    Q = Matrix(Matrix(F.Q)')
    L = Matrix(F.R')
    phases = safesign.(diag(L))
    lmul!(Diagonal(phases), Q)
    rmul!(L, Diagonal(conj!(phases)))
    return L, Q
end

function vumps(
    A::Array{Complex{T},3},
    M::Union{Array{T,4},Array{Complex{T},4}};
    verbose=false,
    tol=1e-8,
    maxvumpsiters=100,
    kwargs...,
) where {T}
    AL, = leftgauge(A)
    AR, C = rightgauge(AL)

    FL, λL = leftenv(AL, M; kwargs...)
    FR, λR = rightenv(AR, M; kwargs...)

    verbose && println("Starting point has λ ≈ $λL ≈ $λR")

    λ, AL, C, AR, = vumpsstep(AL, C, AR, M, FL, FR; tol=tol / 10)
    AL, C, = leftgauge(AR, C; tol=tol / 10, kwargs...) # regauge MPS: not really necessary
    FL, λL = leftenv(AL, M, FL; tol=tol / 10, kwargs...)
    FR, λR = rightenv(AR, M, FR; tol=tol / 10, kwargs...)
    FR ./= @tensor scalar(FL[c, b, a] * C[a, a'] * conj(C[c, c']) * FR[a', b, c']) # normalize FL and FR: not really necessary

    # Convergence measure: norm of the projection of the residual onto the tangent space
    @tensor AC[a, s, b] := AL[a, s, b'] * C[b', b]
    MAC = applyH1(AC, FL, FR, M)
    @tensor MAC[a, s, b] -= AL[a, s, b'] * (conj(AL[a', s', b']) * MAC[a', s', b])
    err = norm(MAC)
    i = 1
    verbose && println("Step $i: λ ≈ $λ ≈ $λL ≈ $λR, err ≈ $err")

    while err > tol && i < maxvumpsiters
        λ, AL, C, AR, = vumpsstep(AL, C, AR, M, FL, FR; tol=tol / 10)
        AL, C, = leftgauge(AR, C; tol=tol / 10, kwargs...) # regauge MPS: not really necessary
        FL, λL = leftenv(AL, M, FL; tol=tol / 10, kwargs...)
        FR, λR = rightenv(AR, M, FR; tol=tol / 10, kwargs...)
        FR ./= @tensor scalar(FL[c, b, a] * C[a, a'] * conj(C[c, c']) * FR[a', b, c']) # normalize FL and FR: not really necessary

        # Convergence measure: norm of the projection of the residual onto the tangent space
        @tensor AC[a, s, b] := AL[a, s, b'] * C[b', b]
        MAC = applyH1(AC, FL, FR, M)
        @tensor MAC[a, s, b] -= AL[a, s, b'] * (conj(AL[a', s', b']) * MAC[a', s', b])
        err = norm(MAC)
        i += 1
        verbose && println("Step $i: λ ≈ $λ ≈ $λL ≈ $λR, err ≈ $err")
    end
    return λ, AL, C, AR, FL, FR
end
function vumpsstep(
    AL::Array{Complex{T},3},
    C::Matrix{Complex{T}},
    AR::Array{Complex{T},3},
    M::Union{Array{T,4},Array{Complex{T},4}},
    FL::Array{Complex{T},3},
    FR::Array{Complex{T},3};
    kwargs...,
) where {T}
    D, d, = size(AL)
    @tensor AC[a, s, b] := AL[a, s, b'] * C[b', b]
    μ1s, ACs, info1 = eigsolve(
        x -> applyH1(x, FL, FR, M),
        AC,
        1,
        :LM;
        ishermitian=false,
        eager=true,
        maxiter=1,
        kwargs...,
    )
    μ0s, Cs, info0 = eigsolve(
        x -> applyH0(x, FL, FR), C, 1; ishermitian=false, eager=true, maxiter=1, kwargs...
    )
    λ = real(μ1s[1] / μ0s[1])
    AC = ACs[1]
    C = Cs[1]

    QAC, RAC = qrpos(reshape(AC, (D * d, D)))
    QC, RC = qrpos(C)
    AL = reshape(QAC * QC', (D, d, D))
    errL = norm(RAC - RC)

    LAC, QAC = lqpos(reshape(AC, (D, d * D)))
    LC, QC = lqpos(C)

    AR = reshape(QC' * QAC, (D, d, D))
    errR = norm(LAC - LC)

    return λ, AL, C, AR, errL, errR
end

function _getsize(a, order)
    dims = size(a)
    shape = Array{Integer,1}()
    shapeinv = Array{Integer,1}()
    perm = Array{Integer,1}()
    for group in order
        D = 1
        for index in group
            D = dims[index] * D
            push!(shapeinv, dims[index])
            push!(perm, index)
        end
        push!(shape, D)
    end
    return Tuple(shape), Tuple(shapeinv), Tuple(perm)
end

function _merge(a, order)
    shape, _, perm = _getsize(a, order)
    aa = permutedims(a, perm)
    return reshape(aa, shape)
end

# broken
function _unmerge(a, order)
    _, shapeinv, perm = _getsize(a, order)
    aa = reshape(a, shapeinv)
    return permutedims(aa, invperm(perm))
end

function statmechmpo(β, h, D)
    M = zeros(D, D, D, D)
    for i in 1:D
        M[i, i, i, i] = 1
    end
    X = zeros(D, D)
    for j in 1:D, i in 1:D
        X[i, j] = exp(-β * h(i, j))
    end
    Xsq = sqrt(X)
    @tensor M1[a, b, c, d] :=
        M[a', b', c', d'] * Xsq[c', c] * Xsq[d', d] * Xsq[a, a'] * Xsq[b, b']

    # For computing energy: M2 is a tensor across 2 nearest neighbor sites in the lattice, whose
    # expectation value in the converged fixed point of the transfer matrix represents the energy
    Y = zeros(D, D)
    for j in 1:D, i in 1:D
        Y[i, j] = h(i, j) * exp(-β * h(i, j))
    end
    @tensor M2[a, b1, b2, c, d2, d1] :=
        M[a', b1', c1, d1'] *
        Xsq[a, a'] *
        Xsq[b1, b1'] *
        Xsq[d1', d1] *
        Y[c1, c2] *
        M[c2, b2', c', d2'] *
        Xsq[b2, b2'] *
        Xsq[d2', d2] *
        Xsq[c', c]

    return M1, M2
end

function classicalisingmpo(β; J=1.0, h=0.0)
    return statmechmpo(β, (s1, s2) -> -J * (-1)^(s1 != s2) - h / 2 * (s1 == 1 + s2 == 1), 2)
end

#----------

function _leftenvapply(FL, AL, M)
    @tensor FLnew[-1 -2 -3; -4 -5 -6] :=
        AL[u 12; 6 -6] *
        FL[d 8 7; 2 1 u] *
        M[1 2 3 -2 -3 6; 7 8 9 -4 -5 12] *
        (AL')[9 -1; d 3]
    return FLnew
end

@testset "func: leftenv()" begin
    @testset "test ComplexSpace" begin
        S = ℂ^2
        AL, M = _randtensors(S)
        FL, λ, info = PEPOKit.leftenv(AL, M)
        @test λ * FL ≈ _leftenvapply(FL, AL, M)
        λs = []
        for i in 1:3
            _, λ, _ = PEPOKit.leftenv(AL, M)
            push!(λs, λ)
        end
        @test all(x -> x ≈ λs[1], λs)
    end
    @testset "test GradedSpace" begin
        S = Z2Space(0 => 1, 1 => 1)
        AL, M = _randtensors(S)
        FL, λ, info = PEPOKit.leftenv(AL, M)
        @test λ * FL ≈ _leftenvapply(FL, AL, M)
    end
end

@testset "trace vumps" begin
    tol=1e-8
    S = ℂ
    tA = TensorMap(randn, ComplexF64, S^10 * S^2, S^10)
    #tM = TensorMap(randn, ComplexF64, S^4 * S^4, S^4 * S^4)

    βc = log(1 + sqrt(2)) / 2
    β = 0.95 * βc

    aA = convert(Array, tA)
    #aM = convert(Array, tM)
    aM, _ = classicalisingmpo(β)
    tM = TensorMap(aM, ℂ^2 * ℂ^2, ℂ^2 * ℂ^2)

    tAL, tCL, tλL = PEPOKit.leftgauge(tA)
    aAL, aCL, aλL = leftgauge(aA)

    tAR, tCR, tλR = PEPOKit.rightgauge(tAL)
    aAR, aCR, aλR = rightgauge(aAL)

    @testset "gauging" begin
        @testset "left" begin
            @test aAL ≈ convert(Array, tAL) # Test AL correct
            @test aCL ≈ convert(Array, tCL) # Test C is correct
            @test aλL ≈ tλL # Test λ is correct
        end
        @testset "right" begin
            @test aAR ≈ convert(Array, tAR) # Test AL correct
            @test aCR ≈ convert(Array, tCR) # Test C is correct
            @test aλR ≈ tλR # Test λ is correct
        end
    end

    tC = tCR
    aC = aCR

    tFL, tλL, _ = PEPOKit.leftenv(tAL, tM)
    aFL, aλL, _ = leftenv(aAL, aM)

    tFR, tλR, _ = PEPOKit.rightenv(tAR, tM)
    aFR, aλR, _ = rightenv(aAR, aM)

    @testset "environments" begin
        @testset "left" begin
            @test PEPOKit._flsolve(tAL, tM, tFL) ≈ tλL * tFL
            @test tλL ≈ aλL
            @test aFL / aFL[1, 1, 1] ≈ convert(Array, tFL) / tFL[1, 1, 1]
            @info "tλL = $(tλL) \n aλL = $(aλL)"
        end
        @testset "right" begin
            @test PEPOKit._frsolve(tAR, tM, tFR) ≈ tλR * tFR
            @test tλR ≈ aλR
            @test aFR / aFR[1, 1, 1] ≈ permutedims(convert(Array, tFR) / tFR[1, 1, 1], (2, 1, 3))
            @info "tλR = $(tλR) \n aλR = $(aλR)"
        end
        @testset "invariance" begin
            @test aλR ≈ aλL
            @test tλR ≈ tλL
        end
    end

    tλs, tmps, tenv, tinfo = PEPOKit._vumpsupdate(tAL, tC, tAR, tM, tFL, tFR)

    tλ = tλs[1]
    tλL = real(tλs[2])
    tλR = real(tλs[3])

    aλ, aAL, aC, aAR, = vumpsstep(aAL, aC, aAR, aM, aFL, aFR; tol=tol / 10)

    aAL, aC, = leftgauge(aAR, aC; tol=tol / 10) # regauge MPS: not really necessary
    aFL, aλL = leftenv(aAL, aM, aFL; tol=tol / 10)
    aFR, aλR = rightenv(aAR, aM, aFR; tol=tol / 10)

    @testset "initial interation" begin
        @test tλ ≈ aλ
        @test tλL ≈ aλL
        @test tλR ≈ aλR
    end
end

@testset "vumps" begin
    S = ℂ
    tA = TensorMap(randn, ComplexF64, S^4 * S^2, S^2 * S^4)
    tM = TensorMap(
        randn,
        ComplexF64,
        S^2 * S^2 * S^2 * S^2 * S^2 * S^2,
        S^2 * S^2 * S^2 * S^2 * S^2 * S^2,
    )
    aA = convert(Array, tA)
    aM = convert(Array, tM)

    tAL, tC, tλ = PEPOKit.leftgauge(tA)
    aAL, aC, aλ = leftgauge(aA)

    tAR, tC, tλ = PEPOKit.rightgauge(tAL)
    aAR, aC, aλ = rightgauge(aAL)
    @testset "gauging" begin
        @testset "left" begin
            @test aAL ≈ convert(Array, tAL) # Test AL correct
            @test aC ≈ convert(Array, tC) # Test C is correct
            @test aλ ≈ tλ # Test λ is correct
        end
        @testset "right" begin
            @test aAR ≈ convert(Array, tAR) # Test AL correct
            @test aC ≈ convert(Array, tC) # Test C is correct
            @test aλ ≈ tλ # Test λ is correct
        end
    end
    order_A = ((1,), (2, 3), (4,))
    order_FL = ((1,), (2, 3, 4, 5), (6,))
    order_M = ((1, 2, 7, 8), (3, 9), (4, 5, 10, 11), (6, 12))
    tFL, tλL, _ = PEPOKit.leftenv(tAL, tM) # this is (2,2) (6,6)
    aFL, aλL, _ = leftenv(_merge(aAL, order_A), _merge(aM, order_M)) # this is rank 3, rank 4
    @testset "environments" begin
        tFL_arr = _merge(convert(Array, tFL), order_FL)
        @test tλL ≈ aλL
        @info "tλL = $(tλL) \n aλL = $(aλL)"
        # @test aFL ≈ tFL_arr
    end
end

@testset "thermal ising model" begin
    βc = log(1 + sqrt(2)) / 2
    β = 0.95 * βc

end

max_stars_repo_name: jack-dunham/PEPOKit

max_issues_repo_name: jack-dunham/PEPOKit

max_forks_repo_name: jack-dunham/PEPOKit


text:
using RefractiveIndexDatabase
using BenchmarkTools

suite = BenchmarkGroup()
suite["formula"] = BenchmarkGroup()
suite["tabulated"] = BenchmarkGroup()

let material, l, f
    material = get_material("other", "Mg-LiTaO3", "Moutzouris-o")
    l = range(material.n.λrange..., length=10_000)
    f = RefractiveIndex(material)
    suite["formula"]["MgLiTaO3"] = @benchmarkable $f.($l)
end

let material, l, f
    material = get_material("main", "Ag", "Johnson")
    l = range(extrema(material.λ)..., length=10_000)
    f = RefractiveIndex(material)
    suite["tabulated"]["Ag"] = @benchmarkable $f.($l)
end

tune!(suite)
results = run(suite; verbose=true, seconds = 1)

for (suite, group) in results
    @show suite
    for (benchmark, result) in group
        @show benchmark
        display(result)
    end
end

# # If a cache of tuned parameters already exists, use it, otherwise, tune and cache
# # the benchmark parameters. Reusing cached parameters is faster and more reliable
# # than re-tuning `suite` every time the file is included.
# paramspath = joinpath(dirname(@__FILE__), "params.json")
#
# if isfile(paramspath)
#     loadparams!(suite, BenchmarkTools.load(paramspath)[1], :evals)
# else
#     tune!(suite)
#     BenchmarkTools.save(paramspath, params(suite))
# end


text:
using AbstractGPs, KernelFunctions, Distributions, Soss, LinearAlgebra, Plots, Random

# Adapted from a post by Arthur Lui
# https://discourse.julialang.org/t/gaussian-process-model-with-turing/42453/32?u=cscherrer

# Define model.
m = @model x begin
    # Priors.
    α ~ LogNormal(0.0, 0.1)
    ρ ~ LogNormal(0.0, 1.0)
    σ ~ LogNormal(0.0, 1.0)

    # Covariance function.
    kernel = α^2 * transform(SEKernel(), 1/(ρ*√2))

    # GP (implicit zero-mean).
    gp = GP(kernel)

    # Sampling Distribution (MvNormal likelihood).
    y ~ gp(x, σ^2 + 1e-6)  # add 1e-6 for numerical stability.
end

Random.seed!(1)
# Get some fake data
x = randn(20)

y =  sinpi.(x) .+ 0.1 .* randn(20);

# Sample from the posterior
post = dynamicHMC(m(x=x), (;y=y))

pred = [posterior(p._y_dist, y) for p in predict(Soss.withdistributions(m)(x=x), post)]

# gps = [posterior(p.gp(x), y) for p in pred]

plt = scatter(x, y; legend=false, dpi=200)

function plotgps(pred, x, y)
    plt = plot(;dpi=200)
    xx = collect(range(minimum(x)-0.5, maximum(x)+0.5; length=200))
    for p in pred
        try
            gp = posterior(p._y_dist, y)
            plot!(plt, xx, rand(gp(xx,1e-9),2); alpha=0.01, color="black", legend=false)
        catch
        end
    end
    scatter!(plt, x, y; color=1)
    ylims!(plt, -2.5, 2.5)
    return plt
end

plotgps(pred, x, y)

particles(post)

[posterior(p._y_dist, y) for p in predict(Soss.withdistributions(m)(x=x), post)]


text:
# no need to run this to generate any jld2, but was used to get extreme indices

using TimeSeriesClustering

reference_results = Dict{String,Any}()

ts_input_data = load_timeseries_data(:CEP_GER18)

ev1 = SimpleExtremeValueDescr("wind-dena42","max","absolute")
ev2 = SimpleExtremeValueDescr("solar-dena42","min","integral")
ev3 = SimpleExtremeValueDescr("el_demand-dena21","max","integral")
ev4 = SimpleExtremeValueDescr("el_demand-dena21","min","absolute")
ev = [ev1, ev2, ev3]

ts_input_data_mod,extr_vals1,extr_idcs = simple_extr_val_sel(ts_input_data,ev1;rep_mod_method="feasibility")
println(extr_idcs)
ts_input_data_mod,extr_vals2,extr_idcs = simple_extr_val_sel(ts_input_data,ev2;rep_mod_method="feasibility")
println(extr_idcs)
ts_input_data_mod,extr_vals3,extr_idcs = simple_extr_val_sel(ts_input_data,ev3;rep_mod_method="feasibility")
println(extr_idcs)
ts_input_data_mod,extr_vals4,extr_idcs = simple_extr_val_sel(ts_input_data,ev4;rep_mod_method="feasibility")
println(extr_idcs)

#@save normpath(joinpath(dirname(@__FILE__),"extreme_vals.jld2")) reference_results


text:
export ThresholdPercentage

"""

    ThresholdPercentage(
        data::DataAccessor,
        th::Float64
    )

Recommend based on percentage of ratings which are greater than a certain threshold value `th`.
"""
struct ThresholdPercentage <: Recommender
    data::DataAccessor
    th::Float64
    scores::AbstractVector

    function ThresholdPercentage(data::DataAccessor, th::Float64)
        n_item = size(data.R, 2)
        new(data, th, vector(n_item))
    end
end

isbuilt(recommender::ThresholdPercentage) = isfilled(recommender.scores)

function build!(recommender::ThresholdPercentage)
    n_item = size(recommender.data.R, 2)

    for i in 1:n_item
        v = recommender.data.R[:, i]
        recommender.scores[i] = length(v[v .>= recommender.th]) / count(!isalmostzero, v) * 100.0
    end
end

function ranking(recommender::ThresholdPercentage, u::Int, i::Int)
    check_build_status(recommender)
    recommender.scores[i]
end


text:
# # Vibration of a cube of nearly incompressible material: alternative models

# Source code: [`unit_cube_modes_alt_tut.jl`](unit_cube_modes_alt_tut.jl)

# ## Description

# Compute the free-vibration spectrum of a unit cube of nearly
# incompressible isotropic material, E = 1, ν = 0.499, and ρ = 1 (refer to [1]).

# Here we show how alternative finite element models compare:
# The solution with the serendipity quadratic hexahedron is supplemented with
# solutions obtained with advanced finite elements: nodal-integration energy
# stabilized hexahedra and tetrahedra, and mean-strain hexahedra and
# tetrahedra.

# ## References

# [1] Puso MA, Solberg J (2006) A stabilized nodally integrated tetrahedral. International Journal for Numerical Methods in Engineering 67: 841-867.
# [2] P. Krysl, Mean-strain 8-node hexahedron with optimized energy-sampling
# stabilization, Finite Elements in Analysis and Design 108 (2016) 41–53.

# ![](unit_cube-mode7.png)

# ## Goals

# - Show how to set up a simulation loop that will run all the models and collect data.
# - Show how to present the computed spectrum curves.

##
# ## Definitions

# This is the finite element toolkit itself.
using FinEtools

# The linear stress analysis application is implemented in this package.
using FinEtoolsDeforLinear

# Convenience import.
using FinEtools.MeshExportModule

# The eigenvalue problem is solved with the Lanczos algorithm from this package.
using Arpack

using SymRCM

# The material properties and dimensions are defined with physical units.
E = 1*phun("PA");
nu = 0.499;
rho = 1*phun("KG/M^3");
a = 1*phun("M"); # length of the side of the cube
N = 8
neigvs = 20 # how many eigenvalues
OmegaShift = (0.01*2*pi)^2; # The frequency with which to shift

# The model is fully three-dimensional, and hence the material model  and the
# FEMM created below need to refer to an appropriate model-reduction scheme.
MR = DeforModelRed3D
material = MatDeforElastIso(MR, rho, E, nu, 0.0);

#
models = [
    ("H20", H20block, GaussRule(3,2), FEMMDeforLinear, 1),
    ("ESNICEH8", H8block, NodalTensorProductRule(3), FEMMDeforLinearESNICEH8, 2),
    ("ESNICET4", T4block, NodalSimplexRule(3), FEMMDeforLinearESNICET4, 2),
    ("MSH8", H8block, NodalTensorProductRule(3), FEMMDeforLinearMSH8, 2),
    ("MST10", T10block, TetRule(4), FEMMDeforLinearMST10, 1),
]

# Run  the simulation loop over all the models.
sigdig(n) = round(n * 10000) / 10000
results = let
    results = []
    for m in models
        fens, fes  = m[2](a, a, a, m[5]*N, m[5]*N, m[5]*N);
        @show count(fens)
        geom = NodalField(fens.xyz)
        u = NodalField(zeros(size(fens.xyz,1),3))
        numbering = let
            C = connectionmatrix(FEMMBase(IntegDomain(fes, m[3])), count(fens))
            numbering = symrcm(C)
        end
        numberdofs!(u, numbering);
        println("nfreedofs = $(u.nfreedofs)")

        femm = m[4](MR, IntegDomain(fes, m[3]), material);
        femm = associategeometry!(femm, geom)
        K = stiffness(femm, geom, u)
        M = mass(femm, geom, u);

        evals, evecs, nconv = eigs(K+OmegaShift*M, M; nev=neigvs, which=:SM)
        @show nconv == neigvs
        evals = evals .- OmegaShift;
        fs = real(sqrt.(complex(evals)))/(2*pi)
        println("$(m[1]) eigenvalues: $(sigdig.(fs)) [Hz]")
        push!(results, (m, fs))
    end
    results # return it
end

##
# ## Present the results graphically

using Gnuplot

@gp  "set terminal windows 0 "  :-
for r in results
    @gp  :- collect(1:length(r[2])) vec(r[2]) " lw 2 with lp title '$(r[1][1])' "  :-
end
@gp  :- "set xlabel 'Mode number [ND]'" :-
@gp  :- "set ylabel 'Frequency [Hz]'"


text:
# test creation of RawData, FFTData, CorrData

# create RawData from SeisData
fs = 100.
nx = Int(86400 * fs)
S = SeisData(2)
S.id[1] = "N1.SS1.00.HHZ"
S.name[1] = "New Station"
S.loc[1] = GeoLoc()
S.fs[1] = 100.0
S.gain[1] = 2.0
S.resp[1] = PZResp([0.0+0.0*im 1.0+0.767*im; 0.0+0.0*im 1.0-0.767*im])
S.units[1] = ""
S.src[1] = ""
S.notes[1] = ["asdfasdf","dfhdfgh"]
S.misc[1] = Dict{String,Any}("P" => 2.0)
S.t[1] = [1 round(Int, d2u(DateTime(Date(now())))*1e6); nx 0]
S.x[1] = rand(Float32,nx)

S.id[2] = "N2.SS2.00.HHZ"
S.name[2] = "Old Station"
S.loc[2] = GeoLoc()
S.fs[2] = fs
S.gain[2] = 22.0
S.resp[2] = PZResp64(z = [0.0+0.0*im, 0.0+0.0*im], p = [1.0+1.0*im, 1.0-1.0*im])
S.units[2] = "ms/2"
S.src[2] = "file"
S.notes[2] = ["0913840183","klnelgng"]
S.misc[2] = Dict{String,Any}("S" => 6.5)
S.t[2] = [1 round(Int, d2u(DateTime(Date(now())))*1e6); nx 0]
S.x[2] = rand(Float32,nx)

# test empty allocator
@testset "Empty RawData" begin
@test !isa(try RawData() catch ex ex end, Exception)    # Passes
R = RawData()
@test isempty(R)
end

# check SeisData to RawData with similar cc_len, cc_step
@testset "SeisData to RawData cc_len == cc_step" begin
cc_len = cc_step = 1800.
R = RawData(S,cc_len,cc_step)
@test !isa(try RawData(S,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin == 1 / R.cc_len         # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 48                 # check number of windows
@test R == R                            # test equality
@test u2d(R.t[1]) == u2d(S.t[1][1,2] * 1e-6) # test starttimes equal
@test eltype(R.x) == eltype(S.x[1])
end

# check uncommon cc_len & cc_step
@testset "SeisData odd cc_len/cc_step" begin
cc_len = π
cc_step = 1800.
@test_throws DomainError RawData(S,cc_len,cc_step)
cc_len = 1800.
cc_step = π
@test_throws DomainError RawData(S,cc_len,cc_step)

# check huge cc_len
cc_len = 86401.
cc_step = 1800.
@test_throws DomainError RawData(S,cc_len,cc_step)
end

# check SeisChannel to RawData with similar cc_len, cc_step
@testset "SeisChannel to RawData cc_len == cc_step" begin
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(Ch,cc_len,cc_step)
@test !isa(try RawData(Ch,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin == 1 / R.cc_len         # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 48                 # check number of windows
@test R == R                            # test equality
@test u2d(R.t[1]) == u2d(Ch.t[1,2] * 1e-6) # test starttimes equal
@test eltype(R.x) == eltype(Ch.x)
end

# check SeisChannel to RawData with different cc_len, cc_step
@testset "SeisChannel to RawData cc_len != cc_step" begin
cc_len = 1800.
cc_step = 450.
Ch = S[1]
R = RawData(Ch,cc_len,cc_step)
@test !isa(try RawData(Ch,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin ==  1 / R.cc_len        # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 189                # check number of windows
@test R == R                            # test equality
@test u2d(R.t[1]) == u2d(Ch.t[1,2] * 1e-6) # test starttimes equal
@test eltype(R.x) == eltype(Ch.x)
end

# check uncommon cc_len & cc_step
@testset "SeisChannel odd cc_len/cc_step" begin
Ch = S[1]
cc_len = π
cc_step = 1800.
@test_throws DomainError RawData(Ch,cc_len,cc_step)
cc_len = 1800.
cc_step = π
@test_throws DomainError RawData(Ch,cc_len,cc_step)

# check huge cc_len
cc_len = 86401.
cc_step = 1800.
@test_throws DomainError RawData(Ch,cc_len,cc_step)
end

## test when starttime is not aligned with cc_len
# this SeisData starts at 00:16:40
@testset "SeisChannel to RawData starttime offset" begin
S.t[1] = [1 round(Int, (d2u(DateTime(Date(now()))) + 1e3) *1e6 ); nx 0]
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(S,cc_len,cc_step)
@test !isa(try RawData(Ch,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin == 1 / R.cc_len         # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 47                 # check number of windows
@test eltype(R.x) == eltype(Ch.x)
end

# test CPU/ GPU
@testset "RawData CPU/GPU" begin
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(S,cc_len,cc_step)
@test R == cpu(R)
@test cpu(R) == cpu(R)
@test R == cpu(gpu(R))
end

# test in
@testset "RawData appends" begin
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(S,cc_len,cc_step)
@test in(R.id,R)

# test addition
R = RawData(S[1],cc_len,cc_step)
R1 = RawData()
@test append!(R1,R) == R
@test append!(R1,R) == R + R
end

## test FFTData

# test empty allocator
@testset "FFTData creation" begin
@test !isa(try FFTData() catch ex ex end, Exception)    # Passes
F = FFTData()
@test isempty(F)

# test FFTData creation
diveven(x) = Int(ceil(x / 2)) + iseven(Int(x)) # function to check rfft size
S.t[1] = [1 round(Int, d2u(DateTime(Date(now())))*1e6); nx 0]
cc_len = cc_step = 1800.
R = RawData(S,cc_len,cc_step)
FFT = rfft(R.x,1)

@test !isa(try FFTData(R.name, R.id, R.loc, R.fs, R.gain, R.freqmin,
                       R.freqmax, R.cc_len, R.cc_step, false, "", R.resp,
                       R.misc, R.notes, R.t, FFT) catch ex ex end, Exception)    # Passes
F = FFTData(R.name, R.id, R.loc, R.fs, R.gain, R.freqmin,
                       R.freqmax, R.cc_len, R.cc_step, false, "", R.resp,
                       R.misc, R.notes, R.t, FFT)
@test length(F.t) == size(F.fft,2)              # check number of windows
@test size(F.fft,1) == diveven(F.cc_len * F.fs) # check length of window
@test F.freqmin ==  1 / F.cc_len                # check freqmin
@test F.freqmax == F.fs / 2                     # check freqmax
@test size(F.fft,2) == 48                       # check number of windows
@test ndims(F.fft) == 2                         # check dimensions
@test F == F                                    # test equality
@test u2d(F.t[1]) == u2d(S.t[1][1,2] * 1e-6)    # test starttimes equal
@test eltype(F.fft) <: Complex

# test CPU/ GPU

@test F == cpu(F)
@test cpu(F) == cpu(F)
@test F == cpu(gpu(F))

# test in
@test in(F.id,F)

# test addition
F1 = FFTData()
@test append!(F1,F) == F
@test append!(F1,F) == F + F
@test size(F1.fft,2) == 2 * size(F.fft,2)
end
## test CorrData
# test empty allocator
@testset "CorrData creation" begin
@test !isa(try CorrData() catch ex ex end, Exception)    # Passes
C = CorrData()
@test isempty(C)

# test CorrData creation
maxlag = 100.
cc_len = cc_step = 1800.
R1 = RawData(S[1],cc_len,cc_step)
R2 = RawData(S[2],cc_len,cc_step)
F1 = rfft(R1)
F2 = rfft(R2)
comp = F1.name[end] * F2.name[end]
corr = correlate(F1.fft,F2.fft,Int(F1.cc_len * F1.fs),Int(maxlag*F1.fs))

@test !isa(try CorrData(F1.name, F1.id, F1.loc, comp, false, "CC", F1.fs, F1.gain,
  F1.freqmin, F1.freqmax, F1.cc_len, F1.cc_step, F1.whitened, F1.time_norm,
  F1.resp, F1.misc, F1.notes, 0.,0., 0., maxlag, F1.t, corr) catch ex ex end, Exception)    # Passes
C = CorrData(F1.name, F1.id, F1.loc, comp, false, "CC", F1.fs, F1.gain,
  F1.freqmin, F1.freqmax, F1.cc_len, F1.cc_step, F1.whitened, F1.time_norm,
  F1.resp, F1.misc, F1.notes, 0.,0., 0., maxlag, F1.t, corr)
@test length(C.t) == size(C.corr,2)             # check number of windows
@test size(C.corr,1) == Int(2 * C.maxlag * C.fs + 1) # check length of window
@test C.freqmin ==  1 / C.cc_len                # check freqmin
@test C.freqmax == C.fs / 2                     # check freqmax
@test size(C.corr,2) == 48                      # check number of windows
@test ndims(C.corr) == 2                        # check dimensions
@test C == C                                    # test equality
@test u2d(C.t[1]) == u2d(S.t[1][1,2] * 1e-6)    # test starttimes equal
@test eltype(C.corr) <: Real

# test CPU/ GPU
@test C == cpu(C)
@test cpu(C) == cpu(C)
@test C == cpu(gpu(C))

# test in
@test in(C.id,C)

# test addition
C1 = CorrData()
@test append!(C1,C) == C
@test append!(C1,C) == C + C
@test size(C1.corr,2) == 2 * size(C.corr,2)
end

## test indexing
@testset "Indexing" begin
    cc_len = cc_step = 1800.
    R1 = RawData(S[1],cc_len,cc_step)

    # test length
    @test length(R1) == length(R1.t)

    # test firstindex
    @test firstindex(R1) == 1

    # test lastindex
    @test lastindex(R1) == length(R1.t)

    # test indexing
    ind = [1,2,4,8,16,32]
    R2 = R1[ind]
    @test length(R2) == length(ind)
    @test size(R2.x,2) == length(ind)
    @test all(R2.x .== R1.x[:,ind])
    @test all(R2.t .== R1.t[ind])

    # test indexing with collection
    ind = 1:2:10
    R2 = R1[ind]
    @test length(R2) == length(ind)
    @test size(R2.x,2) == length(ind)
    @test all(R2.x .== R1.x[:,ind])
    @test all(R2.t .== R1.t[ind])

    # test indexing with DateTime
    ind = u2d(R1.t[1]):Hour(1):u2d(R1.t[end])
    R2 = R1[ind]
    @test length(R2) == length(ind)
    @test size(R2.x,2) == length(ind)

    # test setindex
    ind = 1:2:48
    R3 = deepcopy(R1)
    setindex!(R3,R2,ind)
    # test
    @test R3[ind] == R2

    # test sorting
    R4 = sort(R1[1:2:48] + R1[2:2:48])
    @test R4 == R1
end


text:
using Documenter, GeohashHilbert

makedocs(sitename = "GeohashHilbert documentation")

deploydocs(
    repo = "github.com/zobainc/GeohashHilbert.jl.git"
)


text:
using Random, Images, JLD2, LinearAlgebra
using JOLI, Statistics, FFTW
using JUDI4Cloud
using Printf

using ArgParse
include("../utils/parse_cmd.jl")
parsed_args = parse_commandline()
L = parsed_args["nv"]
nsrc = parsed_args["nsrc"]
vm = parsed_args["vm"]
nth = parsed_args["nth"]

Random.seed!(1234)

JLD2.@load "../models/Compass_tti_625m.jld2"
JLD2.@load "../models/timelapsevrho$(L)vint.jld2" vp_stack rho_stack
idx_wb = find_water_bottom(rho.-rho[1,1])
include("../utils/Jitter.jl")

m_stack = [(1f3./vp_stack[i]).^2f0 for i = 1:L]

model_stack = [Model(n,d,o,m_stack[i]; rho=rho_stack[i],nb=80) for i = 1:L]

dtS = dtR = 4f0
timeS = timeR = 2500f0

nrec = Int.(floor((n[1]-1)*d[1]))   # 1m

creds=joinpath(pwd(),"..","credentials.json")
init_culsterless(nsrc; credentials=creds, vm_size=vm, pool_name="Fq", verbose=1, nthreads=nth)

xsrc_stack = [convertToCell(ContJitter(n, d, nsrc)) for i = 1:L]
ysrc = convertToCell(range(0f0,stop=0f0,length=nsrc))
zsrc_stack = [convertToCell(range(10f0,stop=10f0,length=nsrc).+10f0*rand(Float32,nsrc).-5f0) for i = 1:L]

srcGeometry_stack = [Geometry(xsrc_stack[i], ysrc, zsrc_stack[i]; dt=dtS, t=timeS) for i = 1:L]
wavelet = ricker_wavelet(timeS,dtS,0.025f0)

q_stack = [judiVector(srcGeometry_stack[i], wavelet) for i = 1:L]

xrec = range(d[1],stop=(n[1]-1)*d[1],length=nrec)
yrec = 0f0
zrec = range((maximum(idx_wb)-1)*d[1]-2f0,stop=(maximum(idx_wb)-1)*d[1]-2f0,length=nrec)

recGeometry = Geometry(xrec,yrec,zrec; dt=dtR, t=timeR, nsrc=nsrc)

ntComp = get_computational_nt(srcGeometry_stack[1], recGeometry, model)
info = Info(prod(n), nsrc, ntComp)

opt = JUDI.Options(free_surface=true)

F_stack = [judiProjection(info, recGeometry)*judiModeling(info, model_stack[i]; options=opt)*adjoint(judiProjection(info, srcGeometry_stack[i])) for i = 1:L]

dobs_stack = Array{judiVector{Float32,Array{Float32,2}},1}(undef, L)
@sync for i = 1:L
    @async dobs_stack[i] = F_stack[i] * q_stack[i]
end
JLD2.@save "../data/dobs$(L)vint$(nsrc)nsrc.jld2" dobs_stack q_stack


text:
@doc raw"""
    PowerKernel([γ=1])

The Power Kernel is a negative definite kernel given by:
```
    κ(x,y) = ‖x-y‖²ᵞ   γ ∈ (0,1]
```
where `γ` is a shape parameter of the Euclidean distance.

# Examples

```jldoctest; setup = :(using MLKernels)
julia> PowerKernel()
PowerKernel{Float64}(1.0)

julia> PowerKernel(0.5f0)
PowerKernel{Float32}(0.5)
```
"""
struct PowerKernel{T<:AbstractFloat} <: NegativeDefiniteKernel{T}
    γ::T
    function PowerKernel{T}(γ::Real=T(1)) where {T<:AbstractFloat}
        @check_args(PowerKernel, γ, one(T) >= γ > zero(T), "γ ∈ (0,1]")
        new{T}(γ)
    end
end
PowerKernel(γ::T=1.0) where {T<:Real} = PowerKernel{promote_float(T)}(γ)

@inline basefunction(::PowerKernel) = SquaredEuclidean()

@inline kappa(κ::PowerKernel{T}, d²::T) where {T} = d²^κ.γ

function convert(::Type{K}, κ::PowerKernel) where {K>:PowerKernel{T}} where T
    return PowerKernel{T}(κ.γ)
end


text:
@testset "contains" begin
    html = "foo"
    block = PlutoStaticHTML.code_block(html)
    @test contains(block, "<b>foo</b>")

    notebook = Notebook([
        Cell("x = 1 + 1"),
        Cell("using Images: load"),
        Cell("PKGDIR = \"$PKGDIR\""),
        Cell("""im_file(ext) = joinpath(PKGDIR, "test", "im", "im.\$ext")"""),
        Cell("""load(im_file("png"))""")
    ])
    html, nb = notebook2html_helper(notebook)
    lines = split(html, '\n')

    @test contains(lines[1], "1 + 1")
    @test contains(lines[2], "2")
    @test contains(lines[2], "class='$(PlutoStaticHTML.OUTPUT_PRE_CLASS_DEFAULT)'")

    @test contains(lines[end-1], """src=\"data:image/png;base64,""")
    @test contains(lines[end-1], "<img")

    notebook = Notebook([
        Cell("struct A end"),
        Cell("""
            struct B
                x::Int
                a::A
            end
            """
            ),
        Cell("B(1, A())")
    ])
    use_distributed = false
    html, nb = notebook2html_helper(notebook; use_distributed)
    lines = split(html, '\n')
    @test contains(lines[end-1], "B(1, A())")

    notebook = Notebook([
        Cell("md\"my text\"")
    ])
    html, nb = notebook2html_helper(notebook, HTMLOptions(; hide_md_code=true); use_distributed)
    lines = split(html, '\n')
    @test lines[1] == ""

    html, nb = notebook2html_helper(notebook, HTMLOptions(; hide_md_code=false); use_distributed)
    lines = split(html, '\n')
    @test lines[1] != ""

    opts = HTMLOptions(; hide_md_code=false, hide_code=true)
    html, nb = notebook2html_helper(notebook, opts; use_distributed)
    lines = split(html, '\n')
    @test lines[1] == ""
end

@testset "use_distributed=false and pwd" begin
    dir = pwd()
    nb = Notebook([Cell("1 + 1")])
    _, _ = notebook2html_helper(nb, HTMLOptions(; hide_md_code=true); use_distributed=false)
    @test pwd() == dir
end

@testset "hide" begin
    nb = Notebook([
        Cell("x = 1 + 1 # hide"),
        Cell("""
            # hideall
            y = 3 + 3
            """)
        ])
    html, nb = notebook2html_helper(nb; use_distributed=false)
    @test contains(html, ">2<")
    @test !contains(html, "1 + 1")
    @test contains(html, ">6<")
    @test !contains(html, "3 + 3")
end

@testset "from_file" begin
    mktempdir() do dir
        file = joinpath(dir, "tmp.jl")
        content = pluto_notebook_content("x = 1 + 2")
        write(file, content)
        html = PlutoStaticHTML.notebook2html(file)
        @test contains(html, "3")
    end
end

@testset "run_notebook!_errors" begin
    # To avoid changing the current working directory; I don't know what causes it to change
    # exactly.
    cd(pwd()) do
        mktempdir() do dir
            text = pluto_notebook_content("sum(1, :b)")
            path = joinpath(dir, "notebook.jl")
            write(path, text)
            session = ServerSession()
            session.options.evaluation.workspace_use_distributed = false
            err = nothing
            try
                nb = PlutoStaticHTML.run_notebook!(path, session)
            catch err
            end

            @test err isa Exception
            msg = sprint(showerror, err)
            @test contains(msg, "notebook failed")
            @test contains(msg, "notebook.jl")
            @test contains(msg, "sum(1, :b)")
            @test contains(msg, "Closest candidates are")
            @test contains(msg, "_foldl_impl")
        end
    end
end

@testset "pluto-docs-binding" begin
    text = """
        "This is a docstring"
        foo(x) = x
        """
    nb = Notebook([
        Cell(text),
    ])
    html, nb = notebook2html_helper(nb; use_distributed=false)

    @test !contains(html, "pluto-docs-binding")
end

@testset "benchmark-hack" begin
    # Related to https://github.com/fonsp/Pluto.jl/issues/1664
    nb = Notebook([
        Cell("using BenchmarkTools"),
        Cell("@benchmark sum(x)"),
        Cell("x = [1, 2]")
    ])
    html, nb = notebook2html_helper(nb)
    @test contains(html, "BenchmarkTools.Trial")
end

@testset "show_output_above_code" begin
    nb = Notebook([
        Cell("x = 1 + 1020"),
    ])
    hopts = HTMLOptions(; show_output_above_code=true)
    html, _ = notebook2html_helper(nb, hopts; use_distributed=false)
    lines = split(html, '\n')

    @test contains(lines[1], "1021")
    @test contains(lines[2], "1 + 1020")
end

@testset "with_terminal" begin
    nb = Notebook([
        Cell("using PlutoUI"),
        Cell("f(x) = Base.inferencebarrier(x);"),
        Cell("""
            with_terminal() do
                @code_warntype f(1)
            end
            """)
    ])
    html, _ = notebook2html_helper(nb)
    # Basically, this only tests whether `_patch_with_terminal` is applied.
    @test contains(html, """
""")
end

@testset "replace_code_tabs" begin
    code = """
        		1
        	2
        """
    @test PlutoStaticHTML._replace_code_tabs(code) == """
                1
            2
        """

    nb = Notebook([
        Cell("	a = 1 + 1021;"),
        Cell("		b = 1 + 1021;"),
    ])
    hopts = HTMLOptions()
    html, _ = notebook2html_helper(nb, hopts; use_distributed=false)
    lines = split(html, '\n')

    filter!(!isempty, lines)
    @test contains(lines[1], "    a = 1 + 1021")
    @test contains(lines[2], "        b = 1 + 1021")
end

@testset "big-table" begin
    # Using DataFrames here instead of Tables to get the number of rows for long tables.
    nb = Notebook([
        Cell("using DataFrames: DataFrame"),
        Cell("DataFrame(rand(120, 20), :auto)")
    ])
    hopts = HTMLOptions()
    html, _ = notebook2html_helper(nb, hopts; use_distributed=false)
    # Use write(tmp.html, html) to look at this.

    # First column is empty in the header (because of indexes).
    @test contains(html, """

        """)
    # At least one of the rows contains three dots.
    @test contains(html, "...	\n
")

    # The number of rows is shown.
    @test contains(html, "120	")

    # If not being careful, the first element of the last column is taken, so "more" becomes "m".
    @test !contains(html, "m	")
end


text:
function fib(n::Int64)
	if n== 0
		return    0
	elseif n ==1
		return 1
	else
		return fib(n-1)+fib(n-2)
	end
end

println(fib(10))

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
using Tokens
using Test

function testtree()
end


text:
# reverse
Base.reverse(lookup::NoLookup) = lookup
Base.reverse(lookup::AutoLookup) = lookup
function Base.reverse(lookup::AbstractCategorical)
    i = reverse(index(lookup))
    o = reverse(order(lookup))
    rebuild(lookup; data=i, order=o)
end
function Base.reverse(lookup::AbstractSampled)
    i = reverse(parent(lookup))
    o = reverse(order(lookup))
    sp = reverse(span(lookup))
    rebuild(lookup; data=i, order=o, span=sp)
end

# Order
Base.reverse(::ReverseOrdered) = ForwardOrdered()
Base.reverse(::ForwardOrdered) = ReverseOrdered()
Base.reverse(o::Unordered) = Unordered()

# Span
Base.reverse(span::Irregular) = span
Base.reverse(span::Regular) = Regular(-step(span))
Base.reverse(span::Explicit) = Explicit(reverse(val(span), dims=2))


text:
using Pkg, MPI, Test
MPI.Init()

type=[Float64,Float32]

if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
    using CUDA
    A = CuArray
else
    A = Array
end

for T in type
    a=A{T}(reshape(collect(1:100),10,10))
    b = MPI.Allreduce(a, MPI.MIN, MPI.COMM_WORLD)
    @test isapprox(a,b)
end


text:
# Julia

function gcd(a, b)
    if (b > 0) return gcd(b, a % b)
    else return a
    end
end

println("test")

N = 1000
A = []
for i = 0:N-1
    push!(A, i)
end

ans = 0
for i in A
    for j in A
        global ans += gcd(i, j)
    end
end
println(ans)


text:
function value(obj::TimeScalarFunction, arg0::AbsoluteDate)
    return jcall(obj, "value", jdouble, (AbsoluteDate,), arg0)
end

function value(obj::TimeScalarFunction, arg0::FieldAbsoluteDate)
    return jcall(obj, "value", RealFieldElement, (FieldAbsoluteDate,), arg0)
end


text:
(f,fio) = mktemp()
# test :lg
@test save(f, p1) == 1
@test save(f, p1; compress=true) == 1
@test save(f, p2; compress=true) == 1
@test (ne(p2), nv(p2)) == (9, 10)
@test length(sprint(save, p1, "p1")) == 398
@test length(sprint(save, p2, "p2")) == 47
gs = load(joinpath(testdir, "testdata", "tutte-pathdigraph.jgz"), "pathdigraph")
@test gs == p2
@test_throws ErrorException load(joinpath(testdir, "testdata", "tutte-pathdigraph.jgz"), "badname")

d = Dict{AbstractString, SimpleGraph}("p1"=>p1, "p2"=>p2)
@test save(f,d) == 2

# test :graphml
@test save(f, p1, :graphml) == 1
gs = load(joinpath(testdir, "testdata", "grafo1853.13.graphml"), :graphml)
@test length(gs) == 1
@test haskey(gs, "G") #Name of graph
graphml_g = gs["G"]
@test nv(graphml_g) == 13
@test ne(graphml_g) == 15
gs = load(joinpath(testdir, "testdata", "twounnamedgraphs.graphml"), :graphml)
@test gs["graph"] == Graph(gs["digraph"])
@test save(f, g3, :graphml) == 1
@test_throws ErrorException load(joinpath(testdir, "testdata", "twounnamedgraphs.graphml"), "badname", :graphml)
# test a graphml load that results in a warning
# redirecting per https://thenewphalls.wordpress.com/2014/03/21/capturing-output-in-julia/
origSTDERR = STDERR
(outread, outwrite) = redirect_stderr()
gs = load(joinpath(testdir,"testdata","warngraph.graphml"), :graphml)
gsg = load(joinpath(testdir,"testdata","warngraph.graphml"), "G", :graphml)
@test_throws KeyError badgraph = load(joinpath(testdir, "testdata", "badgraph.graphml"), :graphml)
flush(outread)
flush(outwrite)
close(outread)
close(outwrite)
redirect_stderr(origSTDERR)
@test gs["G"] == graphml_g == gsg

# test :gml
gs = load(joinpath(testdir,"testdata", "twographs-10-28.gml"), :gml)
gml1 = gs["gml1"]
gml2 = gs["digraph"]
gml1a = load(joinpath(testdir,"testdata", "twographs-10-28.gml"), "gml1", :gml)
@test gml1a == gml1
@test nv(gml1) == nv(gml2) == 10
@test ne(gml1) == ne(gml2) == 28
gml1a = load(joinpath(testdir,"testdata", "twographs-10-28.gml"), "gml1", :gml)
@test gml1a == gml1
gs = load(joinpath(testdir,"testdata", "twounnamedgraphs.gml"), :gml)
gml1 = gs["graph"]
gml2 = gs["digraph"]
@test nv(gml1) == 4
@test ne(gml1) == 6
@test nv(gml2) == 4
@test ne(gml2) == 9
@test_throws ErrorException load(joinpath(testdir, "testdata", "twounnamedgraphs.gml"), "badname", :gml)

@test save(f, gml1, :gml) == 1
gml1 = load(f, :gml)["graph"]
@test nv(gml1) == 4
@test ne(gml1) == 6

gs = load(joinpath(testdir,"testdata", "twographs-10-28.gml"), :gml)
@test save(f, gs, :gml) == 2
gs = load(f, :gml)
gml1 = gs["gml1"]
gml2 = gs["digraph"]
@test nv(gml1) == nv(gml2) == 10
@test ne(gml1) == ne(gml2) == 28

# test :dot
gs = load(joinpath(testdir, "testdata", "twographs.dot"), :dot)
@test length(gs) == 2
@test gs["g1"] == CompleteGraph(6)
@test nv(gs["g2"]) == 4 && ne(gs["g2"]) == 6 && is_directed(gs["g2"])
@test_throws ErrorException load(joinpath(testdir, "testdata", "twographs.dot"), "badname", :dot)

# test :gexf
@test save(f, p1, :gexf) == 1
@test_throws ErrorException load(STDIN, :gexf)

#test :net
g10 = CompleteGraph(10)
fname,fio = mktemp()
close(fio)
@test save(fname, g10, :net) == 1
@test load(fname,:net)["g"] == g10
rm(fname)

g10 = PathDiGraph(10)
@test save(fname, g10, :net) == 1
@test load(fname,:net)["g"] == g10
rm(fname)

g10 = load(joinpath(testdir, "testdata", "kinship.net"), :net)["g"]
@test nv(g10) == 6
@test ne(g10) == 8

using JLD

function write_readback(path::AbstractString, g)
    jldfile = jldopen(path, "w")
    jldfile["g"] = g
    close(jldfile)

    jldfile = jldopen(path, "r")
    gs = read(jldfile, "g")
    return gs
end

function testjldio(path::AbstractString, g::Graph)
    gs = write_readback(path, g)
    gloaded = Graph(gs)
    @test gloaded == g
end

println("*** Running JLD IO tests")
graphs = [PathGraph(10), CompleteGraph(5), WheelGraph(7)]
for (i,g) in enumerate(graphs)
    path = joinpath(testdir,"testdata", "test.$i.jld")
    testjldio(path, g)
    #delete the file (it gets left on test failure so you could debug it)
    rm(path)
end

for (i,g) in enumerate(graphs)
    eprop = Dict{Edge,Char}([(e, Char(i)) for e in edges(g)])
    net = LightGraphs.Network{Graph, Int, Char}(g, 1:nv(g), eprop)
    path = joinpath(testdir,"testdata", "test.$i.jld")
    nsaved = write_readback(path, net)
    @test LightGraphs.Network(nsaved) == net
    #delete the file (it gets left on test failure so you could debug it)
    rm(path)
end
println("*** Finished JLD IO tests")


text:
module UMICORE
using ..AGFFileReader: Glass, GlassID, AGF
export GASIR1, GASIR2

const GASIR1 = Glass(GlassID(AGF, 1774), 3, 2.50730798, 0.0761612483, 0.0116256625, -0.000129196515, -6.01958096e-8, -1.78603241e-10, 0.0, 0.0, 0.0, 0.0, 2.0, 14.0, 4.9751e-5, 1.5398e-7, 0.0, 0.0, 0.0, 0.0, 20.0, 0.0, -1.0, -1.0, 17.0, -1.0, 0, -1.0, [(2.0, 0.9998, 2.0), (4.0, 0.9968, 2.0), (6.0, 0.9954, 2.0), (8.0, 0.9941, 2.0), (10.0, 0.9926, 2.0), (12.0, 0.9567, 2.0), (14.0, 0.9275, 2.0)], 1.0, -1.0, -1.0, 0, 0.0, 0, 4.4, 0)
const GASIR2 = Glass(GlassID(AGF, 1775), 3, 2.60073644, 0.0917946407, 0.0122684939, -0.000165642849, -7.8013525e-8, -1.51225803e-10, 0.0, 0.0, 0.0, 0.0, 2.0, 14.0, 4.6e-5, 0.0, 0.0, 0.0, 0.0, 0.0, 20.0, 0.0, -1.0, -1.0, 16.0, -1.0, 2, -1.0, [(2.0, 0.999, 2.0), (4.0, 0.999, 2.0), (6.0, 0.995, 2.0), (8.0, 0.9934, 2.0), (10.0, 0.9984, 2.0), (12.0, 0.9823, 2.0), (14.0, 0.9724, 2.0)], 1.0, -1.0, -1.0, 0, 0.0, 0, 4.7, 0)
end # module

max_stars_repo_name: rambunctiousapple/GlassCat

max_issues_repo_name: rambunctiousapple/GlassCat

max_forks_repo_name: rambunctiousapple/GlassCat


text:
import ..UncertainDatasets:
	AbstractUncertainIndexDataset,
	UncertainIndexDataset

import ..UncertainDatasets:ConstrainedUncertainIndexDataset

"""
    constrain(udata::UncertainIndexDataset,
        s::SamplingConstraint) -> ConstrainedUncertainIndexDataset

Return a uncertain dataset by applying the constraint `s` to each
uncertain value in `udata`.
"""
constrain(udata::UncertainIndexDataset, constraint::SamplingConstraint) =
	ConstrainedUncertainIndexDataset([constrain(uval, constraint) for uval in udata])

"""
    constrain(udata::UncertainIndexDataset,
        constraints::Vector{T}) where {T<:SamplingConstraint} -> ConstrainedUncertainIndexDataset

Return a uncertain dataset by applying a different sampling constraint to each uncertain
value in `udata`.
"""
function constrain(udata::UncertainIndexDataset, constraints::Vector{T}) where {T<:SamplingConstraint}
	if length(udata) != length(constraints)
		error("Number of sampling constraints must match length of dataset.")
	end

	n_vals = length(udata)

	ConstrainedUncertainIndexDataset([constrain(udata[i], constraints[i]) for i in 1:n_vals])
end

"""
    constrain(udata::UncertainIndexDataset,
        constraint::SamplingConstraint) ->  ConstrainedUncertainIndexDataset

Return a uncertain dataset by applying the `constraint` to each
uncertain value in `udata`.
"""
constrain(udata::ConstrainedUncertainIndexDataset, constraint::SamplingConstraint) =
	ConstrainedUncertainIndexDataset([constrain(uval, constraint) for uval in udata])

"""
    constrain(udata::UncertainIndexDataset,
        constraints::Vector{T}) where {T<:SamplingConstraint} -> ConstrainedUncertainIndexDataset

Return a uncertain dataset by applying a different sampling constraint to each uncertain
value in `udata`.
"""
function constrain(udata::ConstrainedUncertainIndexDataset,
        constraints::Vector{T}) where {T<:SamplingConstraint}

    if length(udata) != length(constraints)
        error("Number of sampling constraints must match length of dataset.")
    end

    n_vals = length(udata)

    ConstrainedUncertainIndexDataset([constrain(udata[i], constraints[i]) for i in 1:n_vals])
end

export constrain, verify_constraints


text:
#= TODO test append_axis!
@testset "append tests" begin
    @test append_axis!(CombineStack(), [1, 2], [3, 4]) == [1, 2, 3, 4]
    @test_throws ErrorException append_axis!(CombineStack(), 1:3, 3:4)
end
=#


text:
m = SimpleGridWorld()
p = FunctionPolicy(s->:up)
is = GWPos(3,4)
stage = first(stepthrough(m, p, is, max_steps=1))

c = POMDPModels.render(m, stage)
draw(SVG(tempname()*".svg", 10cm, 10cm), c)

c = POMDPModels.render(m, stage, color=s->reward(m,s))
draw(SVG(tempname()*".svg", 10cm, 10cm), c)

c = POMDPModels.render(m, stage, color=s->rand())
draw(SVG(tempname()*".svg", 10cm, 10cm), c)

c = POMDPModels.render(m, stage, color=s->"yellow")
draw(SVG(tempname()*".svg", 10cm, 10cm), c)

c = POMDPModels.render(m, stage, color=s->"yellow", policy=p)
draw(SVG(tempname()*".svg", 10cm, 10cm), c)


text:
"""
    LogNormalExpert(μ, σ)

PDF:

```math
f(x; \\mu, \\sigma) = \\frac{1}{x \\sqrt{2 \\pi \\sigma^2}}
\\exp \\left( - \\frac{(\\log(x) - \\mu)^2}{2 \\sigma^2} \\right),
\\quad x > 0
```

See also: [Lognormal Distribution](https://en.wikipedia.org/wiki/Log-normal_distribution) (Wikipedia)

"""
struct LogNormalExpert{T<:Real} <: NonZIContinuousExpert
    μ::T
    σ::T
    LogNormalExpert{T}(µ::T, σ::T) where {T<:Real} = new{T}(µ, σ)
end

function LogNormalExpert(μ::T, σ::T; check_args=true) where {T <: Real}
    check_args && @check_args(LogNormalExpert, σ >= zero(σ))
    return LogNormalExpert{T}(μ, σ)
end

## Outer constructors
LogNormalExpert(μ::Real, σ::Real) = LogNormalExpert(promote(μ, σ)...)
LogNormalExpert(μ::Integer, σ::Integer) = LogNormalExpert(float(μ), float(σ))
LogNormalExpert() = LogNormalExpert(0.0, 1.0)

## Conversion
function convert(::Type{LogNormalExpert{T}}, μ::S, σ::S) where {T <: Real, S <: Real}
    LogNormalExpert(T(μ), T(σ))
end
function convert(::Type{LogNormalExpert{T}}, d::LogNormalExpert{S}) where {T <: Real, S <: Real}
    LogNormalExpert(T(d.μ), T(d.σ), check_args=false)
end
copy(d::LogNormalExpert) = LogNormalExpert(d.μ, d.σ, check_args=false)

## Loglikelihood of Expert
logpdf(d::LogNormalExpert, x...) = Distributions.logpdf.(Distributions.LogNormal(d.μ, d.σ), x...)
pdf(d::LogNormalExpert, x...) = Distributions.pdf.(Distributions.LogNormal(d.μ, d.σ), x...)
logcdf(d::LogNormalExpert, x...) = Distributions.logcdf.(Distributions.LogNormal(d.μ, d.σ), x...)
cdf(d::LogNormalExpert, x...) = Distributions.cdf.(Distributions.LogNormal(d.μ, d.σ), x...)

## expert_ll, etc
expert_ll_exact(d::LogNormalExpert, x::Real) = LRMoE.logpdf(d, x)
function expert_ll(d::LogNormalExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_ll = (yl == yu) ? logpdf.(d, yl) : logcdf.(d, yu) + log1mexp.(logcdf.(d, yl) - logcdf.(d, yu))
    expert_ll = (tu == 0.) ? -Inf : expert_ll
    return expert_ll
end
function expert_tn(d::LogNormalExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu))
    expert_tn = (tu == 0.) ? -Inf : expert_tn
    return expert_tn
end
function expert_tn_bar(d::LogNormalExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn_bar = (tl == tu) ? 0.0 : log1mexp.(logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu)))
    return expert_tn_bar
end

exposurize_expert(d::LogNormalExpert; exposure = 1) = d

## Parameters
params(d::LogNormalExpert) = (d.μ, d.σ)
function params_init(y, d::LogNormalExpert)
    pos_idx = (y .> 0.0)
    μ_init, σ_init = mean(log.(y[pos_idx])), sqrt(var(log.(y[pos_idx])))
    μ_init = isnan(μ_init) ? 0.0 : μ_init
    σ_init = isnan(σ_init) ? 1.0 : σ_init
    return LogNormalExpert(μ_init, σ_init)
end

## KS stats for parameter initialization
function ks_distance(y, d::LogNormalExpert)
    p_zero = sum(y .== 0.0) / sum(y .>= 0.0)
    return max(abs(p_zero-0.0), (1-0.0)*HypothesisTests.ksstats(y[y .> 0.0], Distributions.LogNormal(d.μ, d.σ))[2])
end

## Simululation
sim_expert(d::LogNormalExpert) = Distributions.rand(Distributions.LogNormal(d.μ, d.σ), 1)[1]

## penalty
penalty_init(d::LogNormalExpert) = [Inf 1.0 Inf]
no_penalty_init(d::LogNormalExpert) = [Inf 1.0 Inf]
penalize(d::LogNormalExpert, p) = 0 # (d.μ/p[1])^2 + (p[2]-1)*log(d.σ) - d.σ/p[3]

## statistics
mean(d::LogNormalExpert) = mean(Distributions.LogNormal(d.μ, d.σ))
var(d::LogNormalExpert) = var(Distributions.LogNormal(d.μ, d.σ))
quantile(d::LogNormalExpert, p) = quantile(Distributions.LogNormal(d.μ, d.σ), p)
lev(d::LogNormalExpert, u) = exp(d.μ+0.5*d.σ^2) * cdf.(Normal(d.μ+d.σ^2, d.σ), log(u)) + u*(1-cdf.(Normal(d.μ, d.σ), log(u)))
excess(d::LogNormalExpert, u) = mean(d) - lev(d, u)

## Misc functions for E-Step
function _diff_dens_series(d::LogNormalExpert, yl, yu)
    return exp(-0.5*(log(yl)-d.μ)^2/(d.σ^2)) - exp(-0.5*(log(yu)-d.μ)^2/(d.σ^2))
end

function _diff_dist_series(d::LogNormalExpert, yl, yu)
    return ( 0.5 + 0.5*erf((log(yu)-d.μ)/(sqrt2*d.σ)) ) - ( 0.5 + 0.5*erf((log(yl)-d.μ)/(sqrt2*d.σ)) )
end

function _int_obs_logY(d::LogNormalExpert, yl, yu, expert_ll)
    if yl == yu
        return log(yl)
    else
        return exp(-expert_ll) * ( d.σ*invsqrt2π*0.5*_diff_dens_series(d, yl, yu) + d.μ*_diff_dist_series(d, yl, yu) )
    end
end

function _int_lat_logY(d::LogNormalExpert, tl, tu, expert_tn_bar)
    return exp(-expert_tn_bar) * (d.μ - ( d.σ*invsqrt2π*0.5*_diff_dens_series(d, tl, tu) + d.μ*_diff_dist_series(d, tl, tu) ) )
end

function _zdensz_series(z)
    return (isinf(z) ? 0.0 : z*exp(-0.5*z^2))
end

function _diff_zdensz_series(d::LogNormalExpert, yl, yu)
    return _zdensz_series((log(yl)-d.μ)/d.σ) - _zdensz_series((log(yu)-d.μ)/d.σ)
end

function _int_obs_logY_sq(d::LogNormalExpert, yl, yu, expert_ll)
    if yl == yu
        return (log(yl))^2
    else
        return exp(-expert_ll) * ( ((d.σ)^2 + (d.μ)^2)*_diff_dist_series(d, yl, yu) + d.σ*d.μ*invsqrt2π*_diff_dens_series(d, yl, yu) + (d.σ)^2*invsqrt2π*_diff_zdensz_series(d, yl, yu) )
    end
end

function _int_lat_logY_sq(d::LogNormalExpert, tl, tu, expert_tn_bar)
    return exp(-expert_tn_bar) * ( ((d.σ)^2 + (d.μ)^2) - ( ((d.σ)^2 + (d.μ)^2)*_diff_dist_series(d, tl, tu) + d.σ*d.μ*invsqrt2π*_diff_dens_series(d, tl, tu) + (d.σ)^2*invsqrt2π*_diff_zdensz_series(d, tl, tu) ))
end

## EM: M-Step
function EM_M_expert(d::LogNormalExpert,
                     tl, yl, yu, tu,
                     exposure,
                     z_e_obs, z_e_lat, k_e;
                     penalty = true, pen_pararms_jk = [Inf 1.0 Inf])

    if d.σ < 0.000001
        return d
    end

    expert_ll_pos = expert_ll.(d, tl, yl, yu, tu)
    expert_tn_bar_pos = expert_tn_bar.(d, tl, yl, yu, tu)

    # Further E-Step
    logY_e_obs = vec( _int_obs_logY.(d, yl, yu, expert_ll_pos) )
    logY_e_lat = vec( _int_lat_logY.(d, tl, tu, expert_tn_bar_pos) )
    nan2num(logY_e_lat, 0.0) # get rid of NaN

    logY_sq_e_obs = vec( _int_obs_logY_sq.(d, yl, yu, expert_ll_pos) )
    logY_sq_e_lat = vec( _int_lat_logY_sq.(d, tl, tu, expert_tn_bar_pos) )
    nan2num(logY_sq_e_lat, 0.0) # get rid of NaN

    # Update parameters
    pos_idx = (yu .!= 0.0)
    term_zkz = z_e_obs[pos_idx] .+ (z_e_lat[pos_idx] .* k_e[pos_idx])
    term_zkz_logY = (z_e_obs[pos_idx] .* logY_e_obs[pos_idx]) .+ (z_e_lat[pos_idx] .* k_e[pos_idx] .* logY_e_lat[pos_idx])
    term_zkz_logY_sq = (z_e_obs[pos_idx] .* logY_sq_e_obs[pos_idx]) .+ (z_e_lat[pos_idx] .* k_e[pos_idx] .* logY_sq_e_lat[pos_idx])

    μ_new = sum(term_zkz_logY)[1] / sum(term_zkz)[1]
    σ_new = sqrt( 1/sum(term_zkz)[1] * (sum(term_zkz_logY_sq)[1] - 2.0*μ_new*sum(term_zkz_logY)[1] + (μ_new)^2*sum(term_zkz)[1] ) )

    return LogNormalExpert(μ_new, σ_new)
end

## EM: M-Step, exact observations
function EM_M_expert_exact(d::LogNormalExpert,
                            ye, exposure,
                            z_e_obs;
                            penalty = true, pen_pararms_jk = [Inf 1.0 Inf])

    if d.σ < 0.000001
        return d
    end

    # Further E-Step
    logY_e_obs = log.(ye)

    logY_sq_e_obs = (log.(ye)).^2

    # Update parameters
    pos_idx = (ye .!= 0.0)
    term_zkz = z_e_obs[pos_idx]
    term_zkz_logY = (z_e_obs[pos_idx] .* logY_e_obs[pos_idx])
    term_zkz_logY_sq = (z_e_obs[pos_idx] .* logY_sq_e_obs[pos_idx])

    μ_new = sum(term_zkz_logY)[1] / sum(term_zkz)[1]
    σ_new = sqrt( 1/sum(term_zkz)[1] * (sum(term_zkz_logY_sq)[1] - 2.0*μ_new*sum(term_zkz_logY)[1] + (μ_new)^2*sum(term_zkz)[1] ) )

    return LogNormalExpert(μ_new, σ_new)
end

max_stars_repo_name: sparktseung/JuliaLRMoE

max_issues_repo_name: sparktseung/JuliaLRMoE

max_forks_repo_name: sparktseung/JuliaLRMoE


text:
# TODO: more tests!
using Test
using ChainRulesCore
using LinearAlgebra: Diagonal
using ChainRulesCore: extern, accumulate, accumulate!, store!, @scalar_rule,
    Wirtinger, wirtinger_primal, wirtinger_conjugate,
    Zero, One, DoesNotExist, Thunk
using Base.Broadcast: broadcastable

@testset "ChainRulesCore" begin
    include("differentials.jl")
    include("rules.jl")
end


text:
module TestConstraintIndicatorSOS1

using Test

using MathOptInterface
const MOI = MathOptInterface

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

include("../utilities.jl")

function test_indicator_by_SOS1()
    # linear problem with indicator constraints
    # max  2x1 + 3x2
    # s.t. x1 + x2 <= 10
    #      z1 == 1 ==> x2 <= 8
    #      z2 == 1 ==> x1 + x2 == 9
    #      z3 == 1 ==> x1 >= 5
    #      z1 + z2 >= 1
    model = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
    x1 = MOI.add_variable(model)
    x2 = MOI.add_variable(model)
    z1 = MOI.add_variable(model)
    z2 = MOI.add_variable(model)
    z3 = MOI.add_variable(model)

    vc2 = MOI.add_constraint(model, z2, MOI.ZeroOne())
    vc1 = MOI.add_constraint(model, z1, MOI.ZeroOne())
    vc3 = MOI.add_constraint(model, z3, MOI.ZeroOne())
    f1 = MOI.VectorAffineFunction(
        [
            MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z1)),
            MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0],
    )
    iset1 = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(8.0))

    f2 = MOI.VectorAffineFunction(
        [
            MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z2)),
            MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x1)),
            MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0],
    )
    iset2 = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.EqualTo(9.0))

    f3 = MOI.VectorAffineFunction(
        [
            MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z3)),
            MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x1)),
        ],
        [0.0, 0.0],
    )
    iset3 = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.GreaterThan(5.0))

    BT1 = MOI.Bridges.Constraint.concrete_bridge_type(
        MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64},
        typeof(f1),
        typeof(iset1),
    )
    bridge1 = MOI.Bridges.Constraint.bridge_constraint(BT1, model, f1, iset1)
    @test BT1 <:
          MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64,<:MOI.LessThan}
    @test bridge1 isa BT1

    BT2 = MOI.Bridges.Constraint.concrete_bridge_type(
        MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64},
        typeof(f2),
        typeof(iset2),
    )
    bridge2 = MOI.Bridges.Constraint.bridge_constraint(BT2, model, f2, iset2)
    @test BT2 <:
          MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64,<:MOI.EqualTo}
    @test bridge2 isa BT2

    BT3 = MOI.Bridges.Constraint.concrete_bridge_type(
        MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64},
        typeof(f3),
        typeof(iset3),
    )
    bridge3 = MOI.Bridges.Constraint.bridge_constraint(BT3, model, f3, iset3)
    @test BT3 <:
          MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64,<:MOI.GreaterThan}
    @test bridge3 isa BT3

    w1 = bridge1.slack
    @test MOI.get(model, MOI.ConstraintFunction(), bridge1.sos_index) ==
          MOI.VectorOfVariables([w1, z1])
    lin_cons1 = bridge1.affine_index
    lin_func1 = MOI.get(model, MOI.ConstraintFunction(), lin_cons1)
    @test lin_func1 ≈ MOI.ScalarAffineFunction(
        [MOI.ScalarAffineTerm(1.0, x2), MOI.ScalarAffineTerm(1.0, w1)],
        0.0,
    )
    @test MOI.get(model, MOI.ConstraintSet(), lin_cons1) == MOI.LessThan(8.0)

    w2 = bridge2.slack
    @test MOI.get(model, MOI.ConstraintFunction(), bridge2.sos_index) ==
          MOI.VectorOfVariables([w2, z2])
    lin_cons2 = bridge2.affine_index
    lin_func2 = MOI.get(model, MOI.ConstraintFunction(), lin_cons2)
    @test lin_func2 ≈ MOI.ScalarAffineFunction(
        [
            MOI.ScalarAffineTerm(1.0, x1),
            MOI.ScalarAffineTerm(1.0, x2),
            MOI.ScalarAffineTerm(1.0, w2),
        ],
        0.0,
    )
    @test MOI.get(model, MOI.ConstraintSet(), lin_cons2) == MOI.EqualTo(9.0)

    w3 = bridge3.slack
    MOI.GreaterThan(0.0)
    @test MOI.get(model, MOI.ConstraintFunction(), bridge3.sos_index) ==
          MOI.VectorOfVariables([w3, z3])
    lin_cons3 = bridge3.affine_index
    lin_func3 = MOI.get(model, MOI.ConstraintFunction(), lin_cons3)
    @test lin_func3 ≈ MOI.ScalarAffineFunction(
        [MOI.ScalarAffineTerm(1.0, x1), MOI.ScalarAffineTerm(1.0, w3)],
        0.0,
    )
    @test MOI.get(model, MOI.ConstraintSet(), lin_cons3) == MOI.GreaterThan(5.0)

    ## MOI.get on bridge
    @test MOI.get(model, MOI.ConstraintSet(), bridge3) == iset3
    @test MOI.get(model, MOI.ConstraintFunction(), bridge3) ≈ f3
    return
end

function test_basic()
    mock = MOI.Utilities.MockOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
    )
    MOI.Test.test_basic_VectorAffineFunction_Indicator_GreaterThan(
        MOI.Bridges.Constraint.IndicatortoSOS1{Float64}(mock),
        MOI.Test.Config(),
    )
    MOI.Test.test_basic_VectorAffineFunction_Indicator_LessThan(
        MOI.Bridges.Constraint.IndicatortoSOS1{Float64}(mock),
        MOI.Test.Config(),
    )
    return
end

function test_model_equality()
    mock = MOI.Utilities.MockOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
    )
    bridged_mock = MOI.Bridges.Constraint.IndicatortoSOS1{Float64}(mock)
    z, _ = MOI.add_constrained_variable(bridged_mock, MOI.ZeroOne())
    MOI.set(bridged_mock, MOI.VariableName(), z, "z")
    x = MOI.add_variable(bridged_mock)
    MOI.set(bridged_mock, MOI.VariableName(), x, "x")
    var_names = ["z", "x", "w"]
    f = MOI.VectorAffineFunction(
        [
            MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z)),
            MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x)),
        ],
        [0.0, 0.0],
    )
    iset = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(8.0))
    ci = MOI.add_constraint(bridged_mock, f, iset)
    @test length(MOI.get(mock, MOI.ListOfVariableIndices())) == 3
    MOI.set(
        mock,
        MOI.VariableName(),
        MOI.get(mock, MOI.ListOfVariableIndices()),
        var_names,
    )
    sos_cons_list = MOI.get(
        mock,
        MOI.ListOfConstraintIndices{MOI.VectorOfVariables,MOI.SOS1{Float64}}(),
    )
    @test length(sos_cons_list) == 1
    MOI.set(mock, MOI.ConstraintName(), sos_cons_list[1], "sos1")
    inequality_list = MOI.get(
        mock,
        MOI.ListOfConstraintIndices{
            MOI.ScalarAffineFunction{Float64},
            MOI.LessThan{Float64},
        }(),
    )
    @test length(inequality_list) == 1
    MOI.set(mock, MOI.ConstraintName(), inequality_list[1], "ineq")
    MOI.set(mock, MOI.ObjectiveFunction{MOI.VariableIndex}(), z)
    MOI.set(mock, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    s = """
    variables: z, x, w
    maxobjective: z
    sos1: [w, z] in MathOptInterface.SOS1([0.4, 0.6])
    ineq: x + w <= 8.0
    z in MathOptInterface.ZeroOne()
    """
    model = MOI.Utilities.Model{Float64}()
    MOI.Utilities.loadfromstring!(model, s)
    MOI.Test.util_test_models_equal(
        mock,
        model,
        var_names,
        ["sos1", "ineq"],
        [("z", MOI.ZeroOne())],
    )
    _test_delete_bridge(
        bridged_mock,
        ci,
        2,
        (
            (MOI.VariableIndex, MOI.ZeroOne, 1),
            (MOI.VectorOfVariables, MathOptInterface.SOS1{Float64}, 0),
            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}, 0),
        ),
        used_bridges = 1,
        num_bridged = 1,
    )

    model = MOI.Utilities.Model{Float64}()
    sbridged = """
    variables: x, z
    z in MathOptInterface.ZeroOne()
    maxobjective: z
    """
    MOI.Utilities.loadfromstring!(model, sbridged)
    MOI.Test.util_test_models_equal(
        bridged_mock,
        model,
        ["z", "x"],
        String[],
        [("z", MOI.ZeroOne())],
    )
    return
end

function test_getting_primal_attributes()
    mock = MOI.Utilities.MockOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
    )
    z, _ = MOI.add_constrained_variable(mock, MOI.ZeroOne())
    x = MOI.add_variable(mock)
    f = MOI.VectorAffineFunction(
        [
            MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z)),
            MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x)),
        ],
        [0.0, 0.0],
    )
    iset = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(8.0))
    BT = MOI.Bridges.Constraint.concrete_bridge_type(
        MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64},
        typeof(f),
        typeof(iset),
    )
    bridge1 = MOI.Bridges.Constraint.bridge_constraint(BT, mock, f, iset)
    # w value should be defaulted to 0
    MOI.set(mock, MOI.VariablePrimalStart(), bridge1.slack, 0.0)
    affine_value = 6.0
    MOI.set(mock, MOI.ConstraintPrimalStart(), bridge1, [1.0, affine_value])
    @test MOI.get(mock, MOI.VariablePrimalStart(), z) ≈ 1.0
    @test MOI.get(mock, MOI.ConstraintPrimalStart(), bridge1.affine_index) ≈ 6.0

    # after setting the w value
    w_value = 3.0
    MOI.set(mock, MOI.VariablePrimalStart(), bridge1.slack, w_value)
    # linear function should not move
    @test all(
        MOI.get(mock, MOI.ConstraintPrimalStart(), bridge1) .≈
        (1.0, affine_value - w_value),
    )
    iseteq = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.EqualTo(8.0))
    BT = MOI.Bridges.Constraint.concrete_bridge_type(
        MOI.Bridges.Constraint.IndicatorSOS1Bridge{Float64},
        typeof(f),
        typeof(iseteq),
    )
    bridge_eq = MOI.Bridges.Constraint.bridge_constraint(BT, mock, f, iseteq)
    @test MOI.get(
        bridge_eq,
        MOI.NumberOfConstraints{
            MOI.ScalarAffineFunction{Float64},
            MOI.EqualTo{Float64},
        }(),
    ) == 1
    @test MOI.supports(
        mock,
        MOI.ConstraintPrimalStart(),
        MOI.Bridges.Constraint.IndicatorSOS1Bridge{
            Float64,
            MOI.LessThan{Float64},
        },
    )
    # VariablePrimal
    MOI.set(mock, MOI.VariablePrimal(), bridge1.slack, 33.0)
    z_value = 1.0
    MOI.set(mock, MOI.VariablePrimal(), bridge1.z, z_value)
    MOI.set(mock, MOI.VariablePrimal(), x, affine_value)

    # linear function should not move
    @test all(
        MOI.get(mock, MOI.ConstraintPrimal(), bridge1) .≈ (1.0, affine_value),
    )
    return
end

end  # module

TestConstraintIndicatorSOS1.runtests()


text:
module Voronoi
    export Point, Segment, Vertex, Edge, get_edges

    @static if is_windows()
        const library_path =  joinpath(dirname(@__FILE__), "voronoi.dll")
    end

    @static if is_unix()
        const library_path = joinpath(dirname(@__FILE__), "voronoi.so")
    end

    type Point
        X::Int32
        Y::Int32
    end

    type Segment
      p0::Point
      p1::Point
    end

    type Vertex
        X::Float64
        Y::Float64
    end

    type Edge
        start_index::UInt64
        end_index::UInt64

        is_primary::Bool

        site1::UInt64
        site2::UInt64
    end

    function Base.show(io::IO, point::Point)
      print(io, "[$(point.X),$(point.Y)]")
    end

    function Base.show(io::IO, vertex::Vertex)
      print(io, "[$(vertex.X),$(vertex.Y)]")
    end

    function Base.show(io::IO, segment::Segment)
      print(io, "($(segment.p0),$(segment.p1))")
    end

    function appendVertex!(outputArray::Ptr{Void}, vertex::Vertex)
        ourArray = unsafe_pointer_to_objref(outputArray)::Vector{Vertex}

        push!(ourArray, vertex)
    end

    function appendEdge!(outputArray::Ptr{Void}, edge::Edge)
        ourArray = unsafe_pointer_to_objref(outputArray)::Vector{Edge}

        push!(ourArray, Edge(edge.start_index, edge.end_index, edge.is_primary,
                              edge.site1, edge.site2))
    end

    function get_edges(points::Vector{Point}, segments::Vector{Segment})
      vertices = Vector{Vertex}()
      edges = Vector{Edge}()

      ccall((:get_edges, library_path), Void, (Ptr{Ptr{Cint}}, Csize_t,
                                               Ptr{Ptr{Cint}}, Csize_t,
                                               Any, Ptr{Void},
                                               Any, Ptr{Void}),
            [[p.X, p.Y] for p in points],
            length(points),
            [[s.p0.X, s.p0.Y, s.p1.X, s.p1.Y] for s in segments],
            length(segments),
            vertices,
            cfunction(appendVertex!, Any, (Ptr{Void}, Vertex)),
            edges,
            cfunction(appendEdge!, Any, (Ptr{Void}, Edge)))

      vertices, edges
    end
end