text:
include(joinpath(dirname(@__FILE__), "MacroDocs.jl"))
Docile.Cache.register_module(joinpath(dirname(@__FILE__), "MacroDocs.jl"))

import .MacroDocs: MacroDocs, MacroSpec

facts("MacroSpec.") do

    metadata = Docile.Interface.metadata(MacroSpec)
    entries  = Docile.Interface.entries(metadata)

    context("General.") do

        @fact length(entries) --> 28

        @fact Docile.Interface.metadata(metadata) --> @compat(
            Dict{Symbol, Any}(
                :format  => :md,
                :exports => Set([:MacroSpec]),
                :manual  => String[]
                )
            )

    end

    context("Methods.") do

        @fact rawdocs(entries, fmeth(MacroSpec.docs_and_meta))    --> "docs_and_meta/1"
        @fact rawdocs(entries, fmeth(MacroSpec.docs_no_meta))     --> "docs_no_meta/1"
        @fact rawdocs(entries, fmeth(MacroSpec.extern_docs_meta)) --> ""

        @fact docsmeta(entries, :key, fmeth(MacroSpec.docs_and_meta))    --> :docs_and_meta
        @fact docsmeta(entries, :key, fmeth(MacroSpec.extern_docs_meta)) --> :test_case

    end

    context("Single line methods.") do

        @fact rawdocs(entries, fmeth(MacroSpec.inline_docs_meta))        --> "inline_docs_meta/1"
        @fact rawdocs(entries, fmeth(MacroSpec.inline_docs_no_meta))     --> "inline_docs_no_meta/1"
        @fact rawdocs(entries, fmeth(MacroSpec.inline_extern_docs_meta)) --> ""

        @fact docsmeta(entries, :key, fmeth(MacroSpec.inline_docs_meta))        --> :inline_docs_meta
        @fact docsmeta(entries, :key, fmeth(MacroSpec.inline_extern_docs_meta)) --> :test_case

    end

    context("Generic functions.") do

        @fact rawdocs(entries, MacroSpec.inline_docs_meta)        --> "inline_docs_meta"
        @fact rawdocs(entries, MacroSpec.inline_docs_no_meta)     --> "inline_docs_no_meta"
        @fact rawdocs(entries, MacroSpec.inline_extern_docs_meta) --> ""

        @fact docsmeta(entries, :key, MacroSpec.inline_docs_meta)        --> :inline_docs_meta
        @fact docsmeta(entries, :key, MacroSpec.inline_extern_docs_meta) --> :test_case

    end

    context("Types.") do

        @fact rawdocs(entries, MacroSpec.DocsAndMeta)    --> "DocsAndMeta"
        @fact rawdocs(entries, MacroSpec.DocsNoMeta)     --> "DocsNoMeta"
        @fact rawdocs(entries, MacroSpec.ExternDocsMeta) --> ""

        @fact docsmeta(entries, :key, MacroSpec.DocsAndMeta)    --> :DocsAndMeta
        @fact docsmeta(entries, :key, MacroSpec.ExternDocsMeta) --> :test_case

    end

    context("Immutable types.") do

        @fact rawdocs(entries, MacroSpec.DocsAndMetaImm)    --> "DocsAndMetaImm"
        @fact rawdocs(entries, MacroSpec.DocsNoMetaImm)     --> "DocsNoMetaImm"
        @fact rawdocs(entries, MacroSpec.ExternDocsMetaImm) --> ""

        @fact docsmeta(entries, :key, MacroSpec.DocsAndMetaImm)    --> :DocsAndMetaImm
        @fact docsmeta(entries, :key, MacroSpec.ExternDocsMetaImm) --> :test_case

    end

    context("Abstract types.") do

        @fact rawdocs(entries, MacroSpec.DocsAndMetaAbs)    --> "DocsAndMetaAbs"
        @fact rawdocs(entries, MacroSpec.DocsNoMetaAbs)     --> "DocsNoMetaAbs"
        @fact rawdocs(entries, MacroSpec.ExternDocsMetaAbs) --> ""

        @fact docsmeta(entries, :key, MacroSpec.DocsAndMetaAbs)    --> :DocsAndMetaAbs
        @fact docsmeta(entries, :key, MacroSpec.ExternDocsMetaAbs) --> :test_case

    end

    context("Constants.") do

        @fact rawdocs(entries, :DocsAndMetaConst)    --> "DocsAndMetaConst"
        @fact rawdocs(entries, :DocsNoMetaConst)     --> "DocsNoMetaConst"
        @fact rawdocs(entries, :ExternDocsMetaConst) --> ""

        @fact docsmeta(entries, :key, :DocsAndMetaConst)    --> :DocsAndMetaConst
        @fact docsmeta(entries, :key, :ExternDocsMetaConst) --> :test_case

    end

    context("Globals.") do

        @fact rawdocs(entries, :DocsAndMetaGlobal)    --> "DocsAndMetaGlobal"
        @fact rawdocs(entries, :DocsNoMetaGlobal)     --> "DocsNoMetaGlobal"
        @fact rawdocs(entries, :ExternDocsMetaGlobal) --> ""

        @fact docsmeta(entries, :key, :DocsAndMetaGlobal)    --> :DocsAndMetaGlobal
        @fact docsmeta(entries, :key, :ExternDocsMetaGlobal) --> :test_case

    end

    context("Macros.") do

        @fact rawdocs(entries, macrofunc(MacroSpec, "docs_and_meta"))    --> "macro_docs_and_meta"
        @fact rawdocs(entries, macrofunc(MacroSpec, "docs_no_meta"))     --> "macro_docs_no_meta"
        @fact rawdocs(entries, macrofunc(MacroSpec, "extern_docs_meta")) --> ""

        @fact docsmeta(entries, :key, macrofunc(MacroSpec, "docs_and_meta"))    --> :macro_docs_and_meta
        @fact docsmeta(entries, :key, macrofunc(MacroSpec, "extern_docs_meta")) --> :test_case

    end

    context("Modules.") do

        @fact rawdocs(entries, MacroSpec) --> ""

        @fact docsmeta(entries, :key, MacroSpec) --> :test_case

    end

end

facts("MacroDocs.") do

    metadata = Docile.Interface.metadata(MacroDocs)
    entries  = Docile.Interface.entries(metadata)

    context("General.") do

        @fact length(entries) --> 58

        @fact Docile.Interface.metadata(metadata) --> @compat(
            Dict{Symbol, Any}(
                :format => :md,
                :exports => Set([:MacroDocs]),
                :manual  => ["../../docs/manual.md"]
                )
            )

    end

    context("Functions.") do

        @fact rawdocs(entries, fmeth(MacroDocs.f_1)) --> "f_1"
        @fact rawdocs(entries, fmeth(MacroDocs.f_2)) --> "f_2"
        @fact rawdocs(entries, fmeth(MacroDocs.f_3)) --> "f_3"
        @fact rawdocs(entries, fmeth(MacroDocs.f_4)) --> "f_4"

        @fact rawdocs(entries, fmeth(MacroDocs.f_5))  --> "f_5"
        @fact rawdocs(entries, fmeth(MacroDocs.f_6))  --> "f_6"
        @fact rawdocs(entries, fmeth(MacroDocs.f_7))  --> "f_7"
        @fact rawdocs(entries, fmeth(MacroDocs.f_8))  --> "f_8"
        @fact rawdocs(entries, fmeth(MacroDocs.f_9))  --> "f_9"
        @fact rawdocs(entries, fmeth(MacroDocs.f_10)) --> "f_10"
        @fact rawdocs(entries, fmeth(MacroDocs.f_11)) --> "f_11"
        @fact rawdocs(entries, fmeth(MacroDocs.f_12)) --> "f_12"
        @fact rawdocs(entries, fmeth(MacroDocs.f_13)) --> "f_13"
        @fact rawdocs(entries, fmeth(MacroDocs.f_14)) --> "f_14"
        @fact rawdocs(entries, fmeth(MacroDocs.f_15)) --> "f_15"
        @fact rawdocs(entries, fmeth(MacroDocs.f_16)) --> "f_16"
        @fact rawdocs(entries, fmeth(MacroDocs.f_17)) --> "f_17"
        @fact rawdocs(entries, fmeth(MacroDocs.f_18)) --> "f_18"
        @fact rawdocs(entries, fmeth(MacroDocs.f_19)) --> "f_19"
        @fact rawdocs(entries, fmeth(MacroDocs.f_20)) --> "f_20"
        @fact rawdocs(entries, fmeth(MacroDocs.f_21)) --> "f_21"
        @fact rawdocs(entries, fmeth(MacroDocs.f_22)) --> "f_22"
        @fact rawdocs(entries, fmeth(MacroDocs.f_23)) --> "f_23"
        @fact rawdocs(entries, fmeth(MacroDocs.f_24)) --> "f_24"
        @fact rawdocs(entries, fmeth(MacroDocs.f_25)) --> "f_25"
        @fact rawdocs(entries, fmeth(MacroDocs.f_26)) --> "f_26"
        @fact rawdocs(entries, fmeth(MacroDocs.f_27)) --> "f_27"
        @fact rawdocs(entries, fmeth(MacroDocs.f_28)) --> "f_28"
        @fact rawdocs(entries, fmeth(MacroDocs.f_29)) --> "f_29"

        @fact rawdocs(entries, MacroDocs.f_1) --> "g_f_1"
        @fact rawdocs(entries, MacroDocs.f_2) --> "g_f_2"

        @fact rawdocs(entries, MacroDocs.f_30)        --> "g_f_30"
        @fact rawdocs(entries, fmeth(MacroDocs.f_30)) --> "f_30"

        @fact rawdocs(entries, MacroDocs.f_31)        --> "g_f_31"
        @fact rawdocs(entries, fmeth(MacroDocs.f_31)) --> "f_31"

        @fact docsmeta(entries, :returns, MacroDocs.f_31)        --> Int
        @fact docsmeta(entries, :returns, fmeth(MacroDocs.f_31)) --> Bool

        @fact rawdocs(entries, fmeth(MacroDocs.f_36)) --> "f_36"
        @fact rawdocs(entries, fmeth(MacroDocs.f_37)) --> "f_37"

        @fact rawdocs(entries, first(methods(Base.Pkg.add, @compat(Tuple{MacroDocs.FooType})))) --> "f_38"

    end

    context("Globals.") do

        @fact rawdocs(entries, :G_M_1) --> "G_M_1"
        @fact rawdocs(entries, :G_M_2) --> "G_M_2"

        @fact rawdocs(entries, :G_C_1) --> "G_C_1"
        @fact rawdocs(entries, :G_C_2) --> "G_C_2"

    end

    context("Macros.") do

        @fact rawdocs(entries, macrofunc(MacroDocs, "m_1")) --> "m_1"
        @fact rawdocs(entries, macrofunc(MacroDocs, "m_2")) --> "m_2"

    end

    context("Types.") do

        @fact rawdocs(entries, MacroDocs.T_A_1) --> "T_A_1"
        @fact rawdocs(entries, MacroDocs.T_A_2) --> "T_A_2"

        @fact rawdocs(entries, MacroDocs.T_M_1) --> "T_M_1"
        @fact rawdocs(entries, MacroDocs.T_M_2) --> "T_M_2"

        @fact rawdocs(entries, MacroDocs.T_I_1) --> "T_I_1"
        @fact rawdocs(entries, MacroDocs.T_I_2) --> "T_I_2"

        @fact rawdocs(entries, MacroDocs.BT_1)              --> "BT_1"
        @fact docsmeta(entries, :category, MacroDocs.BT_1)  --> :bitstype

        @fact rawdocs(entries, MacroDocs.BT_2)              --> "BT_2"
        @fact docsmeta(entries, :category, MacroDocs.BT_2)  --> :bitstype

        @fact rawdocs(entries, :TA_1)              --> "TA_1"
        @fact docsmeta(entries, :category, :TA_1)  --> :typealias

        @fact rawdocs(entries, :TA_2)              --> "TA_2"
        @fact docsmeta(entries, :category, :TA_2)  --> :typealias

    end

end


text:
using HyperCubicRoots
# Creates HgF as an alias for HypergeometricFunctions._₂F₁
# import HyperCubicRoots: HyperCubicRoots, HgF, solve_real_cubic_roots, solve_all_cubic_roots

import Polynomials
import PolynomialRoots

function test_identity(f1,f2,test_array;tol=10*eps())
    A = f1.(test_array)
    B = f2.(test_array)
    D = abs.(A .- B)
    M = maximum(D)

    return (M < tol), M, argmax(D)
end

aa = 0.5:1e-2:100;

# # Testing on a known case:
L(z) = log(1+z)/z
G(z) = HgF(1,1,2,-Complex(z))

test_identity(L,G,aa)

## Cubic-related identity:
P(z) = ∛(1 + z) + ∛(1-z);
H(z) = 2HgF(-1/6, 1/3, 1/2, Complex(z)^2)

bb = 0:1e-3:0.999

test_identity(H,P,bb, tol=100*eps())

# Note: we need to ensure that |z| < 1 in the argument of HgF
P(1.01) ≈ H(1.01)  #  isapprox(P(1.01), H(1.01))   false
# (This is done via the Kummer transformation in the cubic roots derivation.)

##
cc = [-72, -36, 0, 1]

p1 = -12
q1 = -36
Del1 = q1^2 + p1^3
s1 = Del1/p1^3

r_num = PolynomialRoots.roots(cc)

r_hyp = [
     2q1/p1*HgF(1/3, 2/3, 1/2, s1),
   -(2q1/3p1)*HgF(2/3, 4/3, 3/2, 1/2*(1 + sqrt(s1))),
   -(2q1/3p1)*HgF(2/3, 4/3, 3/2, 1/2*(1 - sqrt(s1))),
]

r_trig = [
    2*sqrt(12)*cos(  π/18),
    2*sqrt(12)*cos(13π/18),
    2*sqrt(12)*cos(25π/18),
]

# # Leading coefficient is zero:
aa = [-417.47, 843.75, 521.67, 0.0]
try
    rr_hyp = HyperCubicRoots.solve_real_cubic_roots(aa)
catch e
    @warn e
end

real_to_tolerance(vec,ϵ) = real.(filter(z -> abs(imag(z)) < ϵ, vec))

# # Three real roots example:
# a = [-23.88, -8.04, 81.02, 13.17]

# # One real roots example:
# a = [81.80, -59.66, -76.47, -64.32]

# # General comparision:
# Compare the roots obtained from PolynomialRoots to this method.

# Notes: the below test compares the numerical roots from PolynomialRoots
# with those generated by the hypergeometric method.
# Experimental observations show that the number of real roots and their
# approximate value are in agreement,
# when using a tolerance (test_leading_tol) of 1e-5 or larger.
# However, the evaluation of the roots in the polynomial function indicate that the
# PolynomialRoots solutions are much more accurate:
# p = Polynomials.Polynomial([-590.92, 213.72, -419.35, 274.27])
# e.g. hypergeometric: p(1.7742160152725672) =  0.00263
#     PolynomialRoots: p(1.7742140121531602) = -3.96224e-13
# The trade-off is that the hypergeometric evaluation is about a factor of 2 faster.
# It may be worth "polishing" the root is greater accuracy is required.

test_leading_tol = 1e-5
# Notes:
#  1e-5 is experimentally accurate for Float64 calculations
#  1e-15 works for BigFloat calculations, albeit much slower calculations.

test_lead_zero_tol = 1e-8
# test_lead_zero_tol = big"1e-8"
test_real_tol = 1e-8
N_iter = 100_000
global count = 0
for i in 1:N_iter
    global count
    a = rand(-1000:0.01:1000,4)  # randomly generate the coefficients of the cubic.
    # a = rand(big"-1000":0.01:big"1000",4)  # randomly generate the coefficients of the cubic: BigFloat
    if abs(a[4]) < test_leading_tol
        @info("The leading cubic coefficient is approximately zero. Skipping iteration $(i).")
        println("a =", a)
        continue
    end
    # fa = Polynomials.Polynomial(a)
    R_hyp = sort(HyperCubicRoots.solve_real_cubic_roots(a,leading_tol=test_lead_zero_tol))

    R_num = PolynomialRoots.roots(a)
    R_num_real = sort(real_to_tolerance(R_num,test_real_tol))

    try
        @assert all(isapprox.(sort(R_hyp), sort(R_num_real), atol=test_leading_tol))
    catch e
        @error e
        println(e)
        println("Coefficients: ")
        println(a)
        println("HyperCubicRoots : R_hyp = ", R_hyp)
        println("PolynomialRoots : R_num = ", R_num_real)
        break
    end

    ## Display examples of computing 3 real roots:
    if mod(i, N_iter ÷ 10)  == 0 || length(R_hyp) == 3
        length(R_hyp) == 3 && count <= 10 ? (count = count + 1) : continue
        println()
        println("Coefficients: ")
        println(a)
        println("HyperCubicRoots:")
        @show(R_hyp)
        println("PolynomialRoots")
        @show R_num
    end

end

function polyeval(aa, root)
    p = Polynomials.Polynomial(aa)
    return p.(root)
end

# # We see that the accuray of the root degenerates as the leading coefficient
# # becomes smaller. Hence we set the accuracy of the method to about 1e-7.
# TYPE = Float16 # errors due to overflow in calculation.
# TYPE = Float32
TYPE = Float64
# TYPE = BigFloat
for i in 1.0:1:8.0
    x = convert(Vector{TYPE},[-1.,2.1,-3.2,10^-i])
    p = Polynomials.Polynomial(x)
    r = HyperCubicRoots.solve_real_cubic_roots(x, leading_tol=TYPE(1e-16))
    println( (i, 10^-i, r[1]))
    println(p(r[1]))
end

# # The relative magnitudes of the polynomial coefficients plays a role here.
# # The various methods from
for i in 1.0:1:10.0
    x = [-1.,34.1,-34.2,10^-i]
    p = Polynomials.Polynomial(x)
    r = HyperCubicRoots.solve_real_cubic_roots(x, leading_tol=HyperCubicRoots.CUBIC_ATOL)
    # r = PolynomialRoots.roots(x)
    # r = PolynomialRoots.solve_cubic_eq(Vector{Complex{Float64}}(x))
    # println(r)
    if !isempty(r)
        println( (i, 10^-i, r[1], p(r[1])))
    end
end

## Timing comparison:
# # Cubics:
trials = 1_000_000
@time for _ in 1:trials
    a = rand(0.01:0.01:1000.01,4)
    R_hyp = HyperCubicRoots.solve_real_cubic_roots(a)
end

@time for _ in 1:trials
    a = rand(-1000:0.01:1000.01,4)
    R_num = PolynomialRoots.roots(a)
end

# # Quartics:
trials = 1_000_000
@time for _ in 1:trials
    a = rand(0.01:0.01:1000.01,5)
    R_hyp = HyperCubicRoots.solve_all_quartic_roots(a; warn_scaling=false)
end

@time for _ in 1:trials
    a = rand(-1000:0.01:1000.01,5)
    R_num = PolynomialRoots.roots(a)
end

## Comparing our solution functions to PolynomialRoots:
comp_tol = 1e-3  # comparison difference tolerance
t = Vector{Complex{Float64}}()  # vector to hold
r = Vector{Complex{Float64}}()
n = 4  # test n = 2 (quadratic), 3 (cubic), 4 (quartic)
rdvec = []

for i in 1:100_000
    rdvec = rand(-1000:0.01:1000.01,n+1)

    if n == 2
        t = solve_quadratic_roots(rdvec)
    elseif n == 3
        t = solve_all_cubic_roots(rdvec)
    elseif n == 4
        t = solve_all_quartic_roots(rdvec)
    end

    r = PolynomialRoots.roots(rdvec)

    # print("HyperCubicRoots: ")
    # display(t)
    # print("PolynomialRoots: ")
    # display(r)
    # println("")

    t_r = sort(real.(t))
    t_i = sort(imag.(t))
    P_r = sort(real.(r))
    P_i = sort(imag.(r))
    d_r = isapprox.(t_r,P_r; atol=comp_tol)
    d_i = isapprox.(t_i,P_i; atol=comp_tol)

    # if length(t) == 3 && sum(imag.(r)) > 1e-20
    if !all(d_r)
        @warn "Real component mismatch."
        display(t)
        display(r)
        display(t_r .- P_r)
        break
    elseif !all(d_i)
        @warn "Imaginary component mismatch."
        display(t)
        display(r)
        display(t_i .- P_i)
        break
    end

end

# # Evaluation comparison:
polyeval(rdvec,solve_all_quartic_roots(rdvec))
polyeval(rdvec,PolynomialRoots.roots(rdvec))

## Fail case:
# # Constant term is zero:
rdvec = [
    0.0
  388.54
 -442.51
  445.24
]

# # Largish error compared to PolynomialRoots:
rdvec = [
    -729.31
    346.23
   -544.86
    721.04
     -0.01
]

## Fail cases for symmetric roots bug:
# # They seem to just be the case of two real roots.
# # We needed to fix the symmetries:

polycoeff = [
    -93.12
    -738.32
    -477.59
    237.14
    487.27
]
polycoeff = [
  936.22
  165.74
  749.85
  229.61
 -553.5
 ]

t = solve_all_quartic_roots(polycoeff)
r = PolynomialRoots.roots(polycoeff)

t_r = sort(real.(t))
t_i = sort(imag.(t))
P_r = sort(real.(r))
P_i = sort(imag.(r))
d_r = isapprox.(t_r,P_r; atol=comp_tol)
d_i = isapprox.(t_i,P_i; atol=comp_tol)

t_r .- P_r
t_i .- P_i

a = polycoeff[1:4]./polycoeff[5]
-a[4]/4

warn_scaling = true; leading_tol=1e-8; coeff_tol=1.0/leading_tol;

rd = []
h =[]
println("Start\n")
for _ in 1:10
    rd = rand(0.01:0.01:100.01,5)
    h = solve_all_quartic_roots(rd; leading_tol=1e-8, coeff_tol=1e10)
    if !isapprox(imag(h), zero(imag(h))) && real(h[1]) > 0
        display(h)
    end
end

ext: jl

max_stars_repo_name: jd-foster/HyperCubicRoots

max_issues_repo_name: jd-foster/HyperCubicRoots

max_forks_repo_name: jd-foster/HyperCubicRoots


text:
function mantel(mat1, mat2; n = 1000)
    size(mat1) == size(mat2) || throw(ArgumentError("Matrices must be the same size"))
    all(issymmetric, (mat1, mat2)) || error("Matries are not symetric")

    ord = size(mat1, 1)

    v1 = view(mat1, [CartesianIndex(i,j) for i in 1:ord for j in 1:ord if i < j])
    v2 = view(mat2, [CartesianIndex(i,j) for i in 1:ord for j in 1:ord if i < j])

    c_real = cor(v1, v2)
    cs = Vector{Float64}(undef, n)

    @inbounds for i in range(1, n)
        cs[i] = cor(view(v1, randperm(ord)), view(v2, randperm(ord)))
    end

    return (c_real, sum(>(c_real), cs) / n)

end

max_stars_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_issues_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_forks_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis


text:
function eval_cooling(data_fill,cols)
    i_nonmissing = ismissing.(data_fill[!,cols["temp_col"]]).==false
    df = data_fill[i_nonmissing,[:datetime,Symbol(cols["temp_col"])]]

    # Eval cooling windows
    cooling_ranges = ECHOsched.cooling_onoff_segmentation(
        df.datetime,
        df.temperature_deg_c,
        θ_grad_onset=0.05,
    )

    # Map cooling ranges to data_fill dataframe (which includes missing timestamps)
    cooling_ranges_fill = ECHOsched.map_data_to_datafill.((df,),(data_fill,),cooling_ranges)
    data_fill[!,Symbol("$(cols["temp_col"])_neg_grad")] = fill!(Array{Bool}(undef,nrow(data_fill)),false)
    data_fill[!,Symbol("$(cols["temp_col"])_neg_grad")][vcat(collect.(cooling_ranges_fill)...)] .= true
    # Include "ac_on" bool in data_fill
    data_fill[!,:ac_on] = fill!(Array{Bool}(undef,nrow(data_fill)),false)
    data_fill[!,:ac_on][vcat(collect.(cooling_ranges_fill)...)] .= true

    # Find standby ranges
    standby_ranges_fill = ECHOsched.find_standby_segments(cooling_ranges_fill,nrow(data_fill))
    # Include "ac_standby" bool in data_fill
    data_fill[!,:ac_standby] = fill!(Array{Bool}(undef,nrow(data_fill)),false)
    data_fill.ac_standby[vcat(collect.(standby_ranges_fill)...)] .= true

    # Find gaps longer than 5 minutes
    gap_ranges_fill = ECHOsched.find_gaps(df,data_fill;gap_treshold_minutes=5)
    # Include gap bool in data_fill
    data_fill[!,:prolonged_gap] = fill!(Array{Bool}(undef,nrow(data_fill)),false)
    data_fill.prolonged_gap[vcat(collect.(gap_ranges_fill)...)] .= true

    # Impose gaps in cooling/standby segments
    data_fill.ac_on[data_fill.prolonged_gap] .= false
    data_fill.ac_standby[data_fill.prolonged_gap] .= false

    # Check for mischaracterised standby periods (high loads, even where temperature increases)
    standby_ranges_fill = ECHOsched.ranges_from_vec(findall(data_fill.ac_standby))
    i_cooling_corr = ECHOsched.check_standby_segment.((data_fill[!,cols["power_col"]],),standby_ranges_fill;power_threshold=50)
    standby_ranges_fill = standby_ranges_fill[Not(i_cooling_corr)]
    data_fill.ac_standby[Not(vcat(collect.(standby_ranges_fill)...))] .= false
    data_fill.ac_on[setdiff(findall(data_fill.prolonged_gap.==false),findall(data_fill.ac_standby))] .= true

    # Ground truth against power
    i_nonmissing = findall(ismissing.(data_fill[!,cols["power_col"]]).==false)
    i_stby = i_nonmissing[findall(data_fill[i_nonmissing,cols["power_col"]].<50)]
    i_ac_on = i_nonmissing[findall(data_fill[i_nonmissing,cols["power_col"]].>=50)]
    data_fill.ac_on[i_stby] .= false
    data_fill.ac_on[i_ac_on] .= true
    data_fill.ac_standby[i_stby] .= true
    data_fill.ac_standby[i_ac_on] .= false

    return data_fill, cooling_ranges_fill, standby_ranges_fill
end


text:
module TestGrouping
    using Test, DataFrames, Random
    const ≅ = isequal

    @testset "colwise" begin
        Random.seed!(1)
        df = DataFrame(a = repeat(Union{Int, Missing}[1, 3, 2, 4], outer=[2]),
                       b = repeat(Union{Int, Missing}[2, 1], outer=[4]),
                       c = Vector{Union{Float64, Missing}}(randn(8)))

        missingfree = DataFrame([collect(1:10)], [:x1])

        @testset "::Function, ::AbstractDataFrame" begin
            cw = colwise(sum, df)
            answer = [20, 12, -0.4283098098931877]
            @test isa(cw, Vector{Real})
            @test size(cw) == (ncol(df),)
            @test cw == answer

            cw = colwise(sum, missingfree)
            answer = [55]
            @test isa(cw, Array{Int, 1})
            @test size(cw) == (ncol(missingfree),)
            @test cw == answer
        end

        @testset "::Function, ::GroupedDataFrame" begin
            gd = groupby(DataFrame(A = [:A, :A, :B, :B], B = 1:4), :A)
            @test colwise(length, gd) == [[2,2], [2,2]]
        end

        @testset "::Vector, ::AbstractDataFrame" begin
            cw = colwise([sum], df)
            answer = [20 12 -0.4283098098931877]
            @test isa(cw, Array{Real, 2})
            @test size(cw) == (length([sum]),ncol(df))
            @test cw == answer

            cw = colwise([sum, minimum], missingfree)
            answer = reshape([55, 1], (2,1))
            @test isa(cw, Array{Int, 2})
            @test size(cw) == (length([sum, minimum]), ncol(missingfree))
            @test cw == answer

            cw = colwise([Vector{Union{Int, Missing}}], missingfree)
            answer = reshape([Vector{Union{Int, Missing}}(1:10)], (1,1))
            @test isa(cw, Array{Vector{Union{Int, Missing}},2})
            @test size(cw) == (1, ncol(missingfree))
            @test cw == answer

            @test_throws MethodError colwise(["Bob", :Susie], DataFrame(A = 1:10, B = 11:20))
        end

        @testset "::Vector, ::GroupedDataFrame" begin
            gd = groupby(DataFrame(A = [:A, :A, :B, :B], B = 1:4), :A)
            @test colwise([length], gd) == [[2 2], [2 2]]
        end

        @testset "::Tuple, ::AbstractDataFrame" begin
            cw = colwise((sum, length), df)
            answer = Any[20 12 -0.4283098098931877; 8 8 8]
            @test isa(cw, Array{Real, 2})
            @test size(cw) == (length((sum, length)), ncol(df))
            @test cw == answer

            cw = colwise((sum, length), missingfree)
            answer = reshape([55, 10], (2,1))
            @test isa(cw, Array{Int, 2})
            @test size(cw) == (length((sum, length)), ncol(missingfree))
            @test cw == answer

            cw = colwise((CategoricalArray, Vector{Union{Int, Missing}}), missingfree)
            answer = reshape([CategoricalArray(1:10), Vector{Union{Int, Missing}}(1:10)],
                             (2, ncol(missingfree)))
            @test typeof(cw) == Array{AbstractVector,2}
            @test size(cw) == (2, ncol(missingfree))
            @test cw == answer

            @test_throws MethodError colwise(("Bob", :Susie), DataFrame(A = 1:10, B = 11:20))
        end

        @testset "::Tuple, ::GroupedDataFrame" begin
            gd = groupby(DataFrame(A = [:A, :A, :B, :B], B = 1:4), :A)
            @test colwise((length), gd) == [[2,2],[2,2]]
        end
    end

    @testset "by, groupby and map(::Function, ::GroupedDataFrame)" begin
        Random.seed!(1)
        df = DataFrame(a = repeat(Union{Int, Missing}[1, 3, 2, 4], outer=[2]),
                       b = repeat(Union{Int, Missing}[2, 1], outer=[4]),
                       c = Vector{Union{Float64, Missing}}(randn(8)))

        cols = [:a, :b]

        f1(df) = DataFrame(cmax = maximum(df[:c]))
        f2(df) = (cmax = maximum(df[:c]),)
        f3(df) = maximum(df[:c])
        f4(df) = [maximum(df[:c]), minimum(df[:c])]
        f5(df) = reshape([maximum(df[:c]), minimum(df[:c])], 2, 1)
        f6(df) = [maximum(df[:c]) minimum(df[:c])]
        f7(df) = (c2 = df[:c].^2,)
        f8(df) = DataFrame(c2 = df[:c].^2)

        res = unique(df[cols])
        res.cmax = [maximum(df[(df.a .== a) .& (df.b .== b), :c])
                    for (a, b) in zip(res.a, res.b)]
        res2 = unique(df[cols])[repeat(1:4, inner=2), :]
        res2.x1 = collect(Iterators.flatten([[maximum(df[(df.a .== a) .& (df.b .== b), :c]),
                                              minimum(df[(df.a .== a) .& (df.b .== b), :c])]
                                              for (a, b) in zip(res.a, res.b)]))
        res3 = unique(df[cols])
        res3.x1 = [maximum(df[(df.a .== a) .& (df.b .== b), :c])
                   for (a, b) in zip(res.a, res.b)]
        res3.x2 = [minimum(df[(df.a .== a) .& (df.b .== b), :c])
                   for (a, b) in zip(res.a, res.b)]
        res4 = df[cols]
        res4.c2 = df.c.^2
        sres = sort(res, cols)
        sres2 = sort(res2, cols)
        sres3 = sort(res3, cols)
        sres4 = sort(res4, cols)

        # by() without groups sorting
        @test sort(by(df, cols, identity)) ==
            sort(hcat(df, df[cols], makeunique=true))[[:a, :b, :a_1, :b_1, :c]]
        @test sort(by(df, cols, df -> DataFrameRow(df, 1))[[:a, :b, :c]]) ==
            sort(df[.!nonunique(df, cols), :])
        @test by(df, cols, f1) == res
        @test by(df, cols, f2) == res
        @test rename(by(df, cols, f3), :x1 => :cmax) == res
        @test by(df, cols, f4) == res2
        @test by(df, cols, f5) == res2
        @test by(df, cols, f6) == res3
        @test sort(by(df, cols, f7)) == sort(res4)
        @test sort(by(df, cols, f8)) == sort(res4)

        # by() with groups sorting
        @test by(df, cols, identity, sort=true) ==
            sort(hcat(df, df[cols], makeunique=true), cols)[[:a, :b, :a_1, :b_1, :c]]
        @test by(df, cols, df -> DataFrameRow(df, 1), sort=true)[[:a, :b, :c]] ==
            sort(df[.!nonunique(df, cols), :])
        @test by(df, cols, f1, sort=true) == sres
        @test by(df, cols, f2, sort=true) == sres
        @test rename(by(df, cols, f3, sort=true), :x1 => :cmax) == sres
        @test by(df, cols, f4, sort=true) == sres2
        @test by(df, cols, f5, sort=true) == sres2
        @test by(df, cols, f6, sort=true) == sres3
        @test by(df, cols, f7, sort=true) == sres4
        @test by(df, cols, f8, sort=true) == sres4

        @test by(df, [:a], f1) == by(df, :a, f1)
        @test by(df, [:a], f1, sort=true) == by(df, :a, f1, sort=true)

        # groupby() without groups sorting
        gd = groupby(df, cols)
        @test sort(combine(identity, gd)) ==
            sort(combine(gd)) ==
            sort(hcat(df, df[cols], makeunique=true))[[:a, :b, :a_1, :b_1, :c]]
        @test combine(f1, gd) == res
        @test combine(f2, gd) == res
        @test rename(combine(f3, gd), :x1 => :cmax) == res
        @test combine(f4, gd) == res2
        @test combine(f5, gd) == res2
        @test combine(f6, gd) == res3
        @test sort(combine(f7, gd)) == sort(res4)
        @test sort(combine(f8, gd)) == sort(res4)

        # groupby() with groups sorting
        gd = groupby(df, cols, sort=true)
        for i in 1:length(gd)
            @test all(gd[i].a .== sres.a[i])
            @test all(gd[i].b .== sres.b[i])
        end
        @test combine(identity, gd) ==
            combine(gd) ==
            sort(hcat(df, df[cols], makeunique=true), cols)[[:a, :b, :a_1, :b_1, :c]]
        @test combine(f1, gd) == sres
        @test combine(f2, gd) == sres
        @test rename(combine(f3, gd), :x1 => :cmax) == sres
        @test combine(f4, gd) == sres2
        @test combine(f5, gd) == sres2
        @test combine(f6, gd) == sres3
        @test combine(f7, gd) == sres4
        @test combine(f8, gd) == sres4

        # map() without and with groups sorting
        for sort in (false, true)
            gd = groupby(df, cols, sort=sort)
            v = map(d -> d[[:c]], gd)
            @test length(gd) == length(v)
            @test v[1] == gd[1] && v[2] == gd[2] && v[3] == gd[3] && v[4] == gd[4]
            v = map(f1, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f1, df, cols, sort=sort)
            v = map(f2, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f2, df, cols, sort=sort)
            v = map(f3, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f3, df, cols, sort=sort)
            v = map(f4, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f4, df, cols, sort=sort)
            v = map(f5, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f5, df, cols, sort=sort)
            v = map(f5, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f5, df, cols, sort=sort)
            v = map(f6, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f6, df, cols, sort=sort)
            v = map(f7, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f7, df, cols, sort=sort)
            v = map(f8, gd)
            @test vcat(v[1], v[2], v[3], v[4]) == by(f8, df, cols, sort=sort)
        end

        # testing pool overflow
        df2 = DataFrame(v1 = categorical(collect(1:1000)), v2 = categorical(fill(1, 1000)))
        @test groupby(df2, [:v1, :v2]).starts == collect(1:1000)
        @test groupby(df2, [:v2, :v1]).starts == collect(1:1000)

        # grouping empty table
        @test groupby(DataFrame(A=Int[]), :A).starts == Int[]
        # grouping single row
        @test groupby(DataFrame(A=Int[1]), :A).starts == Int[1]

        # issue #960
        x = CategoricalArray(collect(1:20))
        df = DataFrame(v1=x, v2=x)
        groupby(df, [:v1, :v2])

        df2 = by(e->1, DataFrame(x=Int64[]), :x)
        @test size(df2) == (0, 1)
        @test sum(df2[:x]) == 0

        # Check that reordering levels does not confuse groupby
        df = DataFrame(Key1 = CategoricalArray(["A", "A", "B", "B"]),
                       Key2 = CategoricalArray(["A", "B", "A", "B"]),
                       Value = 1:4)
        gd = groupby(df, :Key1)
        @test length(gd) == 2
        @test gd[1] == DataFrame(Key1=["A", "A"], Key2=["A", "B"], Value=1:2)
        @test gd[2] == DataFrame(Key1=["B", "B"], Key2=["A", "B"], Value=3:4)
        gd = groupby(df, [:Key1, :Key2])
        @test length(gd) == 4
        @test gd[1] == DataFrame(Key1="A", Key2="A", Value=1)
        @test gd[2] == DataFrame(Key1="A", Key2="B", Value=2)
        @test gd[3] == DataFrame(Key1="B", Key2="A", Value=3)
        @test gd[4] == DataFrame(Key1="B", Key2="B", Value=4)
        # Reorder levels, add unused level
        levels!(df[:Key1], ["Z", "B", "A"])
        levels!(df[:Key2], ["Z", "B", "A"])
        gd = groupby(df, :Key1)
        @test length(gd) == 2
        @test gd[1] == DataFrame(Key1=["B", "B"], Key2=["A", "B"], Value=3:4)
        @test gd[2] == DataFrame(Key1=["A", "A"], Key2=["A", "B"], Value=1:2)
        gd = groupby(df, [:Key1, :Key2])
        @test length(gd) == 4
        @test gd[1] == DataFrame(Key1="A", Key2="A", Value=1)
        @test gd[2] == DataFrame(Key1="A", Key2="B", Value=2)
        @test gd[3] == DataFrame(Key1="B", Key2="A", Value=3)
        @test gd[4] == DataFrame(Key1="B", Key2="B", Value=4)
        # Make first level unused too
        df.Key1[1:2] .= "B"
        gd = groupby(df, :Key1)
        @test length(gd) == 1
        @test gd[1] == DataFrame(Key1="B", Key2=["A", "B", "A", "B"], Value=1:4)
        gd = groupby(df, [:Key1, :Key2])
        @test length(gd) == 2
        @test gd[1] == DataFrame(Key1="B", Key2="A", Value=[1, 3])
        @test gd[2] == DataFrame(Key1="B", Key2="B", Value=[2, 4])

        # Check that CategoricalArray column is preserved when returning a value...
        res = combine(d -> DataFrame(x=d[1, :Key2]), groupby(df, :Key1))
        @test res.x isa CategoricalVector{String}
        res = combine(d -> (x=d[1, :Key2],), groupby(df, :Key1))
        @test res.x isa CategoricalVector{String}
        # ...and when returning an array
        res = combine(d -> DataFrame(x=d[:Key1]), groupby(df, :Key1))
        @test res.x isa CategoricalVector{String}

        # Check that CategoricalArray and String give a String...
        res = combine(d -> d.Key1 == ["A", "A"] ? DataFrame(x=d[1, :Key1]) : DataFrame(x="C"),
                     groupby(df, :Key1))
        @test res.x isa Vector{String}
        res = combine(d -> d.Key1 == ["A", "A"] ? (x=d[1, :Key1],) : (x="C",),
                      groupby(df, :Key1))
        @test res.x isa Vector{String}
        # ...even when CategoricalString comes second
        res = combine(d -> d.Key1 == ["B", "B"] ? DataFrame(x=d[1, :Key1]) : DataFrame(x="C"),
                      groupby(df, :Key1))
        @test res.x isa Vector{String}
        res = combine(d -> d.Key1 == ["B", "B"] ? (x=d[1, :Key1],) : (x="C",),
                      groupby(df, :Key1))
        @test res.x isa Vector{String}

        df = DataFrame(x = [1, 2, 3], y = [2, 3, 1])

        # Test function returning DataFrameRow
        res = by(d -> DataFrameRow(d, 1), df, :x)
        @test res == DataFrame(x=df.x, x_1=df.x, y=df.y)

        # Test function returning Tuple
        res = by(d -> (sum(d.y),), df, :x)
        @test res == DataFrame(x=df.x, x1=tuple.([2, 3, 1]))

        # Test with some groups returning empty data frames
        @test by(d -> d.x == [1] ? DataFrame(z=[]) : DataFrame(z=1), df, :x) ==
            DataFrame(x=[2, 3], z=[1, 1])
        v = map(d -> d.x == [1] ? DataFrame(z=[]) : DataFrame(z=1), groupby(df, :x))
        @test length(v) == 2
        @test vcat(v[1], v[2]) == DataFrame(x=[2, 3], z=[1, 1])

        # Test that returning values of different types works with NamedTuple
        res = by(d -> d.x == [1] ? 1 : 2.0, df, :x)
        @test res.x1 isa Vector{Float64}
        @test res.x1 == [1, 2, 2]
        # Two columns need to be widened at different times
        res = by(d -> (a=d.x == [1] ? 1 : 2.0, b=d.x == [3] ? missing : "a"), df, :x)
        @test res.a isa Vector{Float64}
        @test res.a == [1, 2, 2]
        @test res.b isa Vector{Union{String,Missing}}
        @test res.b ≅ ["a", "a", missing]
        # Corner case: two columns need to be widened at the same time
        res = by(d -> (a=d.x == [1] ? 1 : 2.0, b=d.x == [1] ? missing : "a"), df, :x)
        @test res.a isa Vector{Float64}
        @test res.a == [1, 2, 2]
        @test res.b isa Vector{Union{String,Missing}}
        @test res.b ≅ [missing, "a", "a"]

        # Test that returning values of different types works with DataFrame
        res = by(d -> DataFrame(x1 = d.x == [1] ? 1 : 2.0), df, :x)
        @test res.x1 isa Vector{Float64}
        @test res.x1 == [1, 2, 2]
        # Two columns need to be widened at different times
        res = by(d -> DataFrame(a=d.x == [1] ? 1 : 2.0, b=d.x == [3] ? missing : "a"), df, :x)
        @test res.a isa Vector{Float64}
        @test res.a == [1, 2, 2]
        @test res.b isa Vector{Union{String,Missing}}
        @test res.b ≅ ["a", "a", missing]
        # Corner case: two columns need to be widened at the same time
        res = by(d -> DataFrame(a=d.x == [1] ? 1 : 2.0, b=d.x == [1] ? missing : "a"), df, :x)
        @test res.a isa Vector{Float64}
        @test res.a == [1, 2, 2]
        @test res.b isa Vector{Union{String,Missing}}
        @test res.b ≅ [missing, "a", "a"]

        # Test return values with columns in different orders
        @test by(d -> d.x == [1] ? (x1=1, x2=3) : (x2=2, x1=4), df, :x) ==
            DataFrame(x=1:3, x1=[1, 4, 4], x2=[3, 2, 2])
        @test by(d -> d.x == [1] ? DataFrame(x1=1, x2=3) : DataFrame(x2=2, x1=4), df, :x) ==
            DataFrame(x=1:3, x1=[1, 4, 4], x2=[3, 2, 2])

        # Test with NamedTuple with columns of incompatible lengths
        @test_throws DimensionMismatch by(d -> (x1=[1], x2=[3, 4]), df, :x)
        @test_throws DimensionMismatch by(d -> d.x == [1] ? (x1=[1], x2=[3]) :
                                                            (x1=[1], x2=[3, 4]), df, :x)

        # Test with incompatible return values
        @test_throws ArgumentError by(d -> d.x == [1] ? (x1=1,) : DataFrame(x1=1), df, :x)
        @test_throws ArgumentError by(d -> d.x == [1] ? NamedTuple() : (x1=1), df, :x)
        @test_throws ArgumentError by(d -> d.x == [1] ? 1 : DataFrame(x1=1), df, :x)
        @test_throws ArgumentError by(d -> d.x == [1] ? DataFrame() : DataFrame(x1=1), df, :x)
        # Special case allowed due to how implementation works
        @test by(d -> d.x == [1] ? 1 : (x1=1), df, :x) == by(d -> 1, df, :x)

        # Test that columns names and types are respected for empty input
        df = DataFrame(x=Int[], y=String[])
        res = by(d -> 1, df, :x)
        @test size(res) == (0, 1)
        @test res.x isa Vector{Int}
    end

    @testset "grouping with missings" for df in
        (DataFrame(Key1 = ["A", missing, "B", "B", "A"],
                   Key2 = ["B", "A", "A", missing, "A"],
                   Value = 1:5),
         DataFrame(Key1 = categorical(["A", missing, "B", "B", "A"]),
                   Key2 = ["B", "A", "A", missing, "A"],
                   Value = 1:5),
         DataFrame(Key1 = ["A", missing, "B", "B", "A"],
                   Key2 = categorical(["B", "A", "A", missing, "A"]),
                   Value = 1:5),
         DataFrame(Key1 = categorical(["A", missing, "B", "B", "A"]),
                   Key2 = categorical(["B", "A", "A", missing, "A"]),
                   Value = 1:5))

        @testset "sort=false, skipmissing=false" begin
            gd = groupby(df, :Key1)
            @test length(gd) == 3
            if df.Key1 isa CategoricalVector # Order of missing changes
                @test gd[1] == DataFrame(Key1=["A", "A"], Key2=["B", "A"], Value=[1, 5])
                @test gd[2] ≅ DataFrame(Key1=["B", "B"], Key2=["A", missing], Value=3:4)
                @test gd[3] ≅ DataFrame(Key1=missing, Key2="A", Value=2)
            else
                @test gd[1] == DataFrame(Key1=["A", "A"], Key2=["B", "A"], Value=[1, 5])
                @test gd[2] ≅ DataFrame(Key1=missing, Key2="A", Value=2)
                @test gd[3] ≅ DataFrame(Key1=["B", "B"], Key2=["A", missing], Value=3:4)
            end

            gd = groupby(df, [:Key1, :Key2])
            @test length(gd) == 5
            @test gd[1] == DataFrame(Key1="A", Key2="B", Value=1)
            @test gd[2] ≅ DataFrame(Key1=missing, Key2="A", Value=2)
            @test gd[3] == DataFrame(Key1="B", Key2="A", Value=3)
            @test gd[4] ≅ DataFrame(Key1="B", Key2=missing, Value=4)
            @test gd[5] ≅ DataFrame(Key1="A", Key2="A", Value=5)
        end

        @testset "sort=false, skipmissing=true" begin
            gd = groupby(df, :Key1, skipmissing=true)
            @test length(gd) == 2
            @test gd[1] == DataFrame(Key1=["A", "A"], Key2=["B", "A"], Value=[1, 5])
            @test gd[2] ≅ DataFrame(Key1=["B", "B"], Key2=["A", missing], Value=3:4)

            gd = groupby(df, [:Key1, :Key2], skipmissing=true)
            @test length(gd) == 3
            @test gd[1] == DataFrame(Key1="A", Key2="B", Value=1)
            @test gd[2] == DataFrame(Key1="B", Key2="A", Value=3)
            @test gd[3] == DataFrame(Key1="A", Key2="A", Value=5)
        end

        @testset "sort=true, skipmissing=false" begin
            gd = groupby(df, :Key1, sort=true)
            @test length(gd) == 3
            @test gd[1] == DataFrame(Key1=["A", "A"], Key2=["B", "A"], Value=[1, 5])
            @test gd[2] ≅ DataFrame(Key1=["B", "B"], Key2=["A", missing], Value=3:4)
            @test gd[3] ≅ DataFrame(Key1=missing, Key2="A", Value=2)

            gd = groupby(df, [:Key1, :Key2], sort=true)
            @test length(gd) == 5
            @test gd[1] ≅ DataFrame(Key1="A", Key2="A", Value=5)
            @test gd[2] == DataFrame(Key1="A", Key2="B", Value=1)
            @test gd[3] == DataFrame(Key1="B", Key2="A", Value=3)
            @test gd[4] ≅ DataFrame(Key1="B", Key2=missing, Value=4)
            @test gd[5] ≅ DataFrame(Key1=missing, Key2="A", Value=2)
        end

        @testset "sort=true, skipmissing=true" begin
            gd = groupby(df, :Key1, sort=true, skipmissing=true)
            @test length(gd) == 2
            @test gd[1] == DataFrame(Key1=["A", "A"], Key2=["B", "A"], Value=[1, 5])
            @test gd[2] ≅ DataFrame(Key1=["B", "B"], Key2=["A", missing], Value=3:4)

            gd = groupby(df, [:Key1, :Key2], sort=true, skipmissing=true)
            @test length(gd) == 3
            @test gd[1] == DataFrame(Key1="A", Key2="A", Value=5)
            @test gd[2] == DataFrame(Key1="A", Key2="B", Value=1)
            @test gd[3] == DataFrame(Key1="B", Key2="A", Value=3)
        end
    end

    @testset "by, combine and map with pair interface" begin
        vexp = x -> exp.(x)
        Random.seed!(1)
        df = DataFrame(a = repeat([1, 3, 2, 4], outer=[2]),
                       b = repeat([2, 1], outer=[4]),
                       c = randn(8))

        # Only test that different by syntaxes work,
        # and rely on tests below for deeper checks
        @test by(:c => sum, df, :a) ==
            by(df, :a, :c => sum) ==
            by(df, :a, (:c => sum,)) ==
            by(df, :a, [:c => sum]) ==
            by(df, :a, c_sum = :c => sum) ==
            by(d -> (c_sum=sum(d.c),), df, :a)
            by(df, :a, d -> (c_sum=sum(d.c),))

        @test by(:c => vexp, df, :a) ==
            by(df, :a, :c => vexp) ==
            by(df, :a, (:c => vexp,)) ==
            by(df, :a, [:c => vexp]) ==
            by(df, :a, c_function = :c => vexp) ==
            by(d -> (c_function=vexp(d.c),), df, :a)
            by(df, :a, d -> (c_function=vexp(d.c),))

        @test by(df, :a, :b => sum, :c => sum) ==
            by(df, :a, (:b => sum, :c => sum,)) ==
            by(df, :a, [:b => sum, :c => sum]) ==
            by(df, :a, b_sum = :b => sum, c_sum = :c => sum) ==
            by(d -> (b_sum=sum(d.b), c_sum=sum(d.c)), df, :a)
            by(df, :a, d -> (b_sum=sum(d.b), c_sum=sum(d.c)))

        @test by(df, :a, :b => vexp, :c => identity) ==
            by(df, :a, (:b => vexp, :c => identity,)) ==
            by(df, :a, [:b => vexp, :c => identity]) ==
            by(df, :a, b_function = :b => vexp, c_identity = :c => identity) ==
            by(d -> (b_function=vexp(d.b), c_identity=identity(d.c)), df, :a)
            by(df, :a, d -> (b_function=vexp(d.b), c_identity=identity(d.c)))

        gd = groupby(df, :a)

        # Only test that different combine syntaxes work,
        # and rely on tests below for deeper checks
        @test combine(:c => sum, gd) ==
            combine(gd, :c => sum) ==
            combine(gd, (:c => sum,)) ==
            combine(gd, [:c => sum]) ==
            combine(gd, c_sum = :c => sum) ==
            combine(d -> (c_sum=sum(d.c),), gd)
            combine(gd, d -> (c_sum=sum(d.c),))

        @test combine(:c => vexp, gd) ==
            combine(gd, :c => vexp) ==
            combine(gd, (:c => vexp,)) ==
            combine(gd, [:c => vexp]) ==
            combine(gd, c_function = :c => vexp) ==
            combine(d -> (c_function=exp.(d.c),), gd)
            combine(gd, d -> (c_function=exp.(d.c),))

        @test combine(gd, :b => sum, :c => sum) ==
            combine(gd, (:b => sum, :c => sum,)) ==
            combine(gd, [:b => sum, :c => sum]) ==
            combine(gd, b_sum = :b => sum, c_sum = :c => sum) ==
            combine(d -> (b_sum=sum(d.b), c_sum=sum(d.c)), gd)
            combine(gd, d -> (b_sum=sum(d.b), c_sum=sum(d.c)))

        @test combine(gd, :b => vexp, :c => identity) ==
            combine(gd, (:b => vexp, :c => identity,)) ==
            combine(gd, [:b => vexp, :c => identity]) ==
            combine(gd, b_function = :b => vexp, c_identity = :c => identity) ==
            combine(d -> (b_function=vexp(d.b), c_identity=d.c), gd)
            combine(gd, d -> (b_function=vexp(d.b), c_identity=d.c))

        for f in (map, combine)
            for col in (:c, 3)
                @test f(col => sum, gd) == f(d -> (c_sum=sum(d.c),), gd)
                @test f(col => x -> sum(x), gd) == f(d -> (c_function=sum(d.c),), gd)
                @test f(col => x -> (z=sum(x),), gd) == f(d -> (z=sum(d.c),), gd)
                @test f(col => x -> DataFrame(z=sum(x),), gd) == f(d -> (z=sum(d.c),), gd)
                @test f(col => identity, gd) == f(d -> (c_identity=d.c,), gd)
                @test f(col => x -> (z=x,), gd) == f(d -> (z=d.c,), gd)

                @test f((xyz = col => sum,), gd) ==
                    f(d -> (xyz=sum(d.c),), gd)
                @test f((xyz = col => x -> sum(x),), gd) ==
                    f(d -> (xyz=sum(d.c),), gd)
                @test f((xyz = col => x -> (sum(x),),), gd) ==
                    f(d -> (xyz=(sum(d.c),),), gd)
                @test_throws ArgumentError f((xyz = col => x -> (z=sum(x),),), gd)
                @test_throws ArgumentError f((xyz = col => x -> DataFrame(z=sum(x),),), gd)
                @test_throws ArgumentError f((xyz = col => x -> (z=x,),), gd)
                @test_throws ArgumentError f(col => x -> (z=1, xzz=[1]), gd)

                for wrap in (vcat, tuple)
                    @test f(wrap(col => sum), gd) ==
                        f(d -> (c_sum=sum(d.c),), gd)
                    @test f(wrap(col => x -> sum(x)), gd) ==
                        f(d -> (c_function=sum(d.c),), gd)
                    @test f(wrap(col => x -> (sum(x),)), gd) ==
                        f(d -> (c_function=(sum(d.c),),), gd)
                    @test_throws ArgumentError f(wrap(col => x -> (z=sum(x),)), gd)
                    @test_throws ArgumentError f(wrap(col => x -> DataFrame(z=sum(x),)), gd)
                    @test_throws ArgumentError f(wrap(col => x -> (z=x,)), gd)
                    @test_throws ArgumentError f(wrap(col => x -> (z=1, xzz=[1])), gd)
                end
            end
            for cols in ((:b, :c), [:b, :c], (2, 3), 2:3, [2, 3], [false, true, true])
                @test f(cols => x -> (y=exp.(x.b), z=x.c), gd) ==
                    f(d -> (y=exp.(d.b), z=d.c), gd)
                @test f(cols => x -> [exp.(x.b) x.c], gd) ==
                    f(d -> [exp.(d.b) d.c], gd)

                @test f((xyz = cols => x -> sum(x.b) + sum(x.c),), gd) ==
                    f(d -> (xyz=sum(d.b) + sum(d.c),), gd)
                if eltype(cols) === Bool
                    cols2 = [[false, true, false], [false, false, true]]
                    @test_throws ArgumentError f((xyz = cols[1] => sum, xzz = cols2[2] => sum), gd)
                    @test_throws ArgumentError f((xyz = cols[1] => sum, xzz = cols2[1] => sum), gd)
                    @test_throws ArgumentError f((xyz = cols[1] => sum, xzz = cols2[2] => x -> first(x)), gd)
                else
                    cols2 = cols
                    @test f((xyz = cols2[1] => sum, xzz = cols2[2] => sum), gd) ==
                        f(d -> (xyz=sum(d.b), xzz=sum(d.c)), gd)
                    @test f((xyz = cols2[1] => sum, xzz = cols2[1] => sum), gd) ==
                        f(d -> (xyz=sum(d.b), xzz=sum(d.b)), gd)
                    @test f((xyz = cols2[1] => sum, xzz = cols2[2] => x -> first(x)), gd) ==
                        f(d -> (xyz=sum(d.b), xzz=first(d.c)), gd)
                    @test_throws ArgumentError f((xyz = cols2[1] => vexp, xzz = cols2[2] => sum), gd)
                end

                @test_throws ArgumentError f(cols => x -> (y=exp.(x.b), z=sum(x.c)), gd)
                @test_throws ArgumentError f((xyz = cols2 => x -> DataFrame(y=exp.(x.b), z=sum(x.c)),), gd)
                @test_throws ArgumentError f((xyz = cols2 => x -> [exp.(x.b) x.c],), gd)

                for wrap in (vcat, tuple)
                    @test f(wrap(cols => x -> sum(x.b) + sum(x.c)), gd) ==
                        f(d -> sum(d.b) + sum(d.c), gd)

                    if eltype(cols) === Bool
                        cols2 = [[false, true, false], [false, false, true]]
                        @test f(wrap(cols2[1] => x -> sum(x.b), cols2[2] => x -> sum(x.c)), gd) ==
                            f(d -> (x1=sum(d.b), x2=sum(d.c)), gd)
                        @test f(wrap(cols2[1] => x -> sum(x.b), cols2[2] => x -> first(x.c)), gd) ==
                            f(d -> (x1=sum(d.b), x2=first(d.c)), gd)
                    else
                        cols2 = cols
                        @test f(wrap(cols[1] => sum, cols[2] => sum), gd) ==
                            f(d -> (b_sum=sum(d.b), c_sum=sum(d.c)), gd)
                        @test f(wrap(cols[1] => sum, cols[2] => x -> first(x)), gd) ==
                            f(d -> (b_sum=sum(d.b), c_function=first(d.c)), gd)
                        @test_throws ArgumentError f(wrap(cols2[1] => vexp, cols2[2] => sum), gd)
                    end

                    @test_throws ArgumentError f(wrap(cols => x -> DataFrame(y=exp.(x.b), z=sum(x.c))), gd)
                    @test_throws ArgumentError f(wrap(cols => x -> [exp.(x.b) x.c]), gd)
                end
            end
        end
    end

    @testset "iteration protocol" begin
        gd = groupby(DataFrame(A = [:A, :A, :B, :B], B = 1:4), :A)
        for v in gd
            @test size(v) == (2,2)
        end
    end

    @testset "getindex" begin
        df = DataFrame(a = repeat([1, 2, 3, 4], outer=[2]),
                       b = 1:8)
        gd = groupby(df, :a)
        @test gd[1] isa SubDataFrame
        @test gd[1] == view(df, [1, 5], :)
        @test_throws BoundsError gd[5]
        if VERSION < v"1.0.0-"
            @test gd[true] == gd[1]
        else
            @test_throws ArgumentError gd[true]
        end
        @test_throws MethodError gd["a"]
        gd2 = gd[[true, false, false, false]]
        @test length(gd2) == 1
        @test gd2[1] == gd[1]
        @test_throws BoundsError gd[[true, false]]
        gd3 = gd[:]
        @test gd3 isa GroupedDataFrame
        @test length(gd3) == 4
        @test gd3 == gd
        for i in 1:4
            @test gd3[i] == gd[i]
        end
        gd4 = gd[[1,2]]
        @test gd4 isa GroupedDataFrame
        @test length(gd4) == 2
        for i in 1:2
            @test gd4[i] == gd[i]
        end
        @test_throws BoundsError gd[1:5]
    end

    @testset "== and isequal" begin
        df1 = DataFrame(a = repeat([1, 2, 3, 4], outer=[2]),
                        b = 1:8)
        df2 = DataFrame(a = repeat([1, 2, 3, 4], outer=[2]),
                        b = [1:7;missing])
        gd1 = groupby(df1, :a)
        gd2 = groupby(df2, :a)
        @test gd1 == gd1
        @test isequal(gd1, gd1)
        @test ismissing(gd1 == gd2)
        @test !isequal(gd1, gd2)
        @test ismissing(gd2 == gd2)
        @test isequal(gd2, gd2)
        df1.c = df1.a
        df2.c = df2.a
        @test gd1 != groupby(df2, :c)
        df2[7, :b] = 10
        @test gd1 != gd2
        df3 = DataFrame(a = repeat([1, 2, 3, missing], outer=[2]),
                        b = 1:8)
        df4 = DataFrame(a = repeat([1, 2, 3, missing], outer=[2]),
                        b = [1:7;missing])
        gd3 = groupby(df3, :a)
        gd4 = groupby(df4, :a)
        @test ismissing(gd3 == gd4)
        @test !isequal(gd3, gd4)
        gd3 = groupby(df3, :a, skipmissing = true)
        gd4 = groupby(df4, :a, skipmissing = true)
        @test gd3 == gd4
        @test isequal(gd3, gd4)
    end
end


text:
# For linear algebra
using StaticArrays: SVector
using LinearAlgebra
using SparseArrays: SparseMatrixCSC, sparse

# Fast Cholesky
using KoLesky

# for logs
using Printf

# multiple dispatch
import IterativeSolvers: cg!, cg
import LinearAlgebra: mul!, ldiv!
import Base: size

using BenchmarkTools
using Profile

## PDEs type
abstract type AbstractPDEs end
struct NonlinElliptic2d{Tα,Tm,TΩ} <: AbstractPDEs
    # eqn: -Δu + α*u^m = f in [Ω[1,1],Ω[2,1]]*[Ω[1,2],Ω[2,2]]
    α::Tα
    m::Tm
    Ω::TΩ
    bdy::Function
    rhs::Function
end

## sample points
function sample_points_rdm(eqn::NonlinElliptic2d, N_domain, N_boundary)
    Ω = eqn.Ω
    x1l = Ω[1,1]
    x1r = Ω[2,1]
    x2l = Ω[1,2]
    x2r = Ω[2,2]

    X_domain = hcat(rand(Float64,(N_domain, 1))*(x1r-x1l).+x1l,rand(Float64,(N_domain, 1))*(x2r-x2l).+x2l)

    N_bd_each=convert(Int64, N_boundary/4)
    if N_boundary != 4 * N_bd_each
        println("[sample points] N_boundary not divided by 4, replaced by ", 4 * N_bd_each)
        N_boundary = 4 * N_bd_each
    end

    X_boundary = zeros((N_boundary, 2))
    # bottom face
    X_boundary[1:N_bd_each, :] = hcat((x1r-x1l)*rand(Float64,(N_bd_each,1)).+x1l, x2l*ones(N_bd_each))
    # right face
    X_boundary[N_bd_each+1:2*N_bd_each, :] = hcat(x1r*ones(N_bd_each),(x2r-x2l)*rand(Float64,(N_bd_each,1)).+x2l)
    # top face
    X_boundary[2*N_bd_each+1:3*N_bd_each, :] = hcat((x1r-x1l)*rand(Float64,(N_bd_each,1)).+x1l, x2r*ones(N_bd_each))

    # left face
    X_boundary[3*N_bd_each+1:N_boundary, :] = hcat(x1l*ones(N_bd_each), (x2r-x2l)*rand(Float64,(N_bd_each,1)).+x2l)
    return X_domain', X_boundary'
end
function sample_points_grid(eqn::NonlinElliptic2d, h_in, h_bd)
    Ω = eqn.Ω
    x1l = Ω[1,1]
    x1r = Ω[2,1]
    x2l = Ω[1,2]
    x2r = Ω[2,2]
    x = x1l + h_in:h_in:x1r-h_in
    y = x2l + h_in:h_in:x2r-h_in
    X_domain = reduce(hcat,[[x[i], y[j]] for i in 1:length(x) for j in 1:length(x)])

    l = length(x1l:h_bd:x1r-h_bd)
    X_boundary = vcat([x1l:h_bd:x1r-h_bd x2l*ones(l)], [x1r*ones(l) x2l:h_bd:x2r-h_bd], [x1r:-h_bd:x1l+h_bd x2r*ones(l)], [x1l*ones(l) x2r:-h_bd:x1l+h_bd])
    return X_domain, X_boundary'
end

## exact algorithms
# assemby Gram matrices
function get_Gram_matrices(eqn::NonlinElliptic2d, cov::KoLesky.AbstractCovarianceFunction, X_domain, X_boundary, sol_now)
    d = 2
    N_domain = size(X_domain,2)
    N_boundary = size(X_boundary,2)
    Δδ_coefs = -1.0
    δ_coefs_int = eqn.α*eqn.m*(sol_now.^(eqn.m-1))

    # get linearized PDEs correponding measurements
    meas_δ = [KoLesky.PointMeasurement{d}(SVector{d,Float64}(X_boundary[:,i])) for i = 1:N_boundary]
    meas_Δδ = [KoLesky.ΔδPointMeasurement{Float64,d}(SVector{d,Float64}(X_domain[:,i]), Δδ_coefs, δ_coefs_int[i]) for i = 1:N_domain]
    meas_test_int = [KoLesky.PointMeasurement{d}(SVector{d,Float64}(X_domain[:,i])) for i = 1:N_domain]

    Theta_train = zeros(N_domain+N_boundary,N_domain+N_boundary)

    measurements = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,2)
    measurements[1] = meas_δ; measurements[2] = meas_Δδ
    cov(Theta_train, reduce(vcat,measurements))

    Theta_test = zeros(N_domain,N_domain+N_boundary)
    cov(view(Theta_test,1:N_domain,1:N_boundary), meas_test_int, meas_δ)
    cov(view(Theta_test,1:N_domain,N_boundary+1:N_domain+N_boundary), meas_test_int, meas_Δδ)
    return Theta_train, Theta_test

end
# iterative GPR
function iterGPR_exact(eqn, cov, X_domain, X_boundary, sol_init, nugget, GNsteps)
    # get the rhs and bdy data
    rhs = [eqn.rhs(X_domain[:,i]) for i in 1:N_domain]
    bdy = [eqn.bdy(X_boundary[:,i]) for i in 1:N_boundary]
    v = sol_init

    for _ in 1:GNsteps
        Theta_train, Theta_test = get_Gram_matrices(eqn, cov, X_domain, X_boundary, v)
        rhs_now = vcat(bdy, rhs+eqn.α*(eqn.m-1)*v.^eqn.m)

        v = Theta_test*(((Theta_train+nugget*diagm(diag(Theta_train))))\rhs_now)
    end
    return v
end

## algorithm using Kolesky and pcg
# struct that stores the factor of Theta_train
abstract type implicit_mtx end

# struct approx_Theta_train{Tv,Ti} <: implicit_mtx
#     P::Vector{Ti}
#     U::SparseMatrixCSC{Tv,Ti}
#     δ_coefs::Vector{Tv}
#     N_boundary::Ti
#     N_domain::Ti
# end

struct approx_Theta_train{Tv,Ti,Tmtx<:SparseMatrixCSC{Tv,Ti}} <: implicit_mtx
    P::Vector{Ti}
    U::Tmtx
    L::Tmtx
    δ_coefs::Vector{Tv}
    N_boundary::Ti
    N_domain::Ti
end

# struct precond_Theta_train{Tv,Ti} <: implicit_mtx
#     P::Vector{Ti}
#     U::SparseMatrixCSC{Tv,Ti}
# end
struct precond_Theta_train{Tv,Ti,Tmtx<:SparseMatrixCSC{Tv,Ti}} <: implicit_mtx
    P::Vector{Ti}
    U::Tmtx
    L::Tmtx
end

function size(A::approx_Theta_train, num)
    return size(A.U,num)
end

function mul!(x, Θtrain::approx_Theta_train, b)
    temp = vcat(b[1:Θtrain.N_boundary],Θtrain.δ_coefs.*b[Θtrain.N_boundary+1:end],b[Θtrain.N_boundary+1:end])
    temp = Θtrain.L\(Θtrain.U\temp[Θtrain.P])
    temp[Θtrain.P] = temp

    @views x[1:Θtrain.N_boundary] = temp[1:Θtrain.N_boundary]
    @views x[Θtrain.N_boundary+1:end] .= Θtrain.δ_coefs.*temp[Θtrain.N_boundary+1:Θtrain.N_boundary+Θtrain.N_domain] .+ temp[Θtrain.N_boundary+Θtrain.N_domain+1:end]
end

function ldiv!(x, precond::precond_Theta_train, b)
    x[precond.P] = precond.U*(precond.L*b[precond.P])
end

function iterGPR_fast_pcg(eqn, cov, X_domain, X_boundary, sol_init, nugget, GNsteps; ρ_big =4.0, ρ_small=6.0, k_neighbors = 4, lambda = 1.5, alpha = 1.0)

    @printf "[Algorithm]: iterative GPR + fast cholesky factorization + pcg\n"
    N_domain = size(X_domain,2)
    N_boundary = size(X_boundary,2)

    # get the rhs and bdy data
    rhs = [eqn.rhs(X_domain[:,i]) for i in 1:N_domain]
    bdy = [eqn.bdy(X_boundary[:,i]) for i in 1:N_boundary]

    sol_now = sol_init

    # form the fast Cholesky part that can be used to compute mtx-vct mul for Theta_test
    d = 2
    Δδ_coefs = -1.0
    δ_coefs = 0.0
    meas_δ = [KoLesky.PointMeasurement{d}(SVector{d,Float64}(X_boundary[:,i])) for i = 1:N_boundary]
    meas_δ_int = [KoLesky.PointMeasurement{d}(SVector{d,Float64}(X_domain[:,i])) for i = 1:N_domain]
    meas_Δδ = [KoLesky.ΔδPointMeasurement{Float64,d}(SVector{d,Float64}(X_domain[:,i]), Δδ_coefs, δ_coefs) for i = 1:N_domain]
    measurements = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,3)
    measurements[1] = meas_δ; measurements[2] = meas_δ_int
    measurements[3] = meas_Δδ

    @printf("[Big Theta: implicit factorization] time")
    @time implicit_bigΘ = KoLesky.ImplicitKLFactorization_FollowDiracs(cov, measurements, ρ_big, k_neighbors; lambda = lambda, alpha = alpha)

    # @printf("[Big Theta: implicit factorization] length of row indices")
    # @show length.(KoLesky.row_indices.(implicit_bigΘ.supernodes.supernodes))
    # @printf("[Big Theta: implicit factorization] length of column indices")
    # @show length.(KoLesky.column_indices.(implicit_bigΘ.supernodes.supernodes))

    @printf("[Big Theta: explicit factorization] time")
    @time explicit_bigΘ = KoLesky.ExplicitKLFactorization(implicit_bigΘ; nugget = nugget)
    U_bigΘ = explicit_bigΘ.U
    L_bigΘ = sparse(U_bigΘ')
    P_bigΘ = explicit_bigΘ.P

    Θtrain = approx_Theta_train(P_bigΘ, U_bigΘ, L_bigΘ,zeros(N_domain),N_boundary,N_domain)

    for step in 1:GNsteps

        @printf "[Current GN step] %d\n" step
        # get Cholesky of Theta_train
        Δδ_coefs = -1.0
        δ_coefs_int = eqn.α*eqn.m*sol_now.^(eqn.m-1)

        meas_Δδ = [KoLesky.ΔδPointMeasurement{Float64,d}(SVector{d,Float64}(X_domain[:,i]), Δδ_coefs, δ_coefs_int[i]) for i = 1:N_domain]
        measurements = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,2)
        measurements[1] = meas_δ; measurements[2] = meas_Δδ

        @printf("[Theta Train: implicit factorization] time")
        @time implicit_factor = KoLesky.ImplicitKLFactorization(cov, measurements, ρ_small, k_neighbors; lambda = lambda, alpha = alpha)

        # @printf("[Theta Train:implicit factorization] length of row indices")
        # @show length.(KoLesky.row_indices.(implicit_factor.supernodes.supernodes))
        # @printf("[Theta Train:implicit factorization] length of column indices")
        # @show length.(KoLesky.column_indices.(implicit_factor.supernodes.supernodes))

        @printf("[Theta Train: explicit factorization] time")
        @time explicit_factor = KoLesky.ExplicitKLFactorization(implicit_factor; nugget = nugget)

        U = explicit_factor.U
        L = sparse(U')
        P = explicit_factor.P
        rhs_now = vcat(bdy, rhs.+eqn.α*(eqn.m-1)*sol_now.^eqn.m)

        # use the approximate solution as the initial point for the pCG iteration
        sol = U*(L*rhs_now[P])
        sol[P] = sol
        # pcg step for Theta_train\rhs

        Θtrain.δ_coefs .= δ_coefs_int
        # Θtrain = approx_Theta_train(P_bigΘ, U_bigΘ, δ_coefs_int,N_boundary,N_domain)
        # x = similar(rhs_now)
        # @time mul!(x, Θtrain::approx_Theta_train, rhs_now)

        precond = precond_Theta_train(P,U,L)
        @printf "[pcg started]\n"
        @time sol, ch = cg!(sol, Θtrain, rhs_now; Pl = precond, log=true)

        println("[pcg finished]",ch)

        # get approximated sol_now = Theta_test * sol
        temp = vcat(sol[1:N_boundary], δ_coefs_int .* sol[N_boundary+1:end], sol[N_boundary+1:end])
        temp = L_bigΘ\(U_bigΘ\temp[P_bigΘ])
        temp[P_bigΘ] = temp
        @views sol_now = temp[N_boundary+1:N_domain+N_boundary]
    end
    return sol_now
end

############# main ################
const α = 1.0;
const m = 3;
const Ω = [[0,1] [0,1]]
# ground truth solution
const freq = 20
const s = 3
function fun_u(x)
    ans = 0
    for k = 1:freq
        ans += sin(pi*k*x[1])*sin(pi*k*x[2])/k^s
        # H^t norm squared is sum 1/k^{2s-2t}
    end
    return ans
end
# right hand side
function fun_rhs(x)
    ans = 0
    for k = 1:freq
        ans += (2*k^2*pi^2)*sin(pi*k*x[1])*sin(pi*k*x[2])/k^s
        # H^t norm squared is sum 1/k^{2s-2t}
    end
    return ans + α*fun_u(x)^m
end

# boundary value
function fun_bdy(x)
    return fun_u(x)
end

@printf "[solver started] NonlinElliptic2d\n"
eqn = NonlinElliptic2d(α,m,Ω,fun_bdy,fun_rhs)

h_in = 0.005; h_bd = h_in
X_domain, X_boundary = sample_points_grid(eqn, h_in, h_bd)
# X_domain, X_boundary = sample_points_rdm(eqn, 900, 124)
N_domain = size(X_domain,2)
N_boundary = size(X_boundary,2)
@printf "[sample points] N_domain is %d, N_boundary is %d\n" N_domain N_boundary

lengthscale = 0.2
cov = KoLesky.MaternCovariance5_2(lengthscale)
# cov = KoLesky.GaussianCovariance(lengthscale)
@printf "[kernel] kernel is MaternCovariance5_2, lengthscale %.2f\n" lengthscale

nugget = 1e-14
@printf "[nugget] nugget term %e\n" nugget

GNsteps_approximate = 3
@printf "[total GN steps] %d\n" GNsteps_approximate

ρ_big = 3.0
ρ_small = 3.0

# ρ_big = 5.0
# ρ_small = 5.0
k_neighbors = 2
@printf "[Fast Cholesky] ρ_big = %.2f, ρ_small = %.2f, k_neighbors = %d\n" ρ_big ρ_small k_neighbors

sol_init = randn(N_domain) # initial solution
truth = [fun_u(X_domain[:,i]) for i in 1:N_domain]
# @profview
@time sol = iterGPR_fast_pcg(eqn, cov, X_domain, X_boundary, sol_init, nugget, GNsteps_approximate; ρ_big = ρ_big, ρ_small = ρ_small, k_neighbors=k_neighbors);
pts_accuracy = sqrt(sum((truth-sol).^2)/N_domain)
println("[L2 accuracy: pCG method]",pts_accuracy)
pts_max_accuracy = maximum(abs.(truth-sol))
println("[Linf accuracy: pCG method]",pts_max_accuracy)

# GNsteps_exact = 3
# @time sol_exact = iterGPR_exact(eqn, cov, X_domain, X_boundary, sol_init, nugget, GNsteps_exact)
# pts_accuracy_exact = sqrt(sum((truth-sol_exact).^2)/N_domain)
# println("[L2 accuracy: exact method]", pts_accuracy_exact)
# pts_max_accuracy_exact = maximum(abs.(truth-sol_exact))
# println("[Linf accuracy: exact method]", pts_max_accuracy_exact)

# using PProf
# pprof(;webport=58699)


text:
"""
    WSVarScoreTestInvariant

A base object for the score test of within-subject variance linear mixed model
data instance, independent of the input data.
Numbers of test variables should be predetermined.

H0: β1 = 0 and τ1 = 0, H1: β1 ≠ 0 or τ1 ≠ 0, for the full model of
WiSER (with parameters β = [β1, β2], τ = [τ1, τ2], and Lγ).
We make use of the fitted null model.
"""
struct WSVarScoreTestInvariant{T <: BlasReal}
    nullmodel  :: WSVarLmmModel{T}

    # dimensions
    p          :: Int       # #mean parameters in linear regression
    q          :: Int       # #random effects
    q◺         :: Int       # triangle number of q
    l          :: Int       # #parameters for modeling WS variability
    m          :: Int       # #individuals/clusters
    r_X1                :: Int      # #test variables in X1
    r_W1                :: Int      # #test variables in W1
    r                   :: Int      # #total test variables, = r_X1 + r_W1.

    # working arrays
    A_21_β2β1_rowsums :: Matrix{T} # for time-invariant X1, p x m.
    A_21_τ2τ1_rowsums :: Matrix{T} # for time-invariant W1, l x m.
    A_21_Lγτ1_rowsums :: Matrix{T} # for time-invariant W1, q◺ x m.

    #ψ_1pre :: Matrix{T} # r x m,
    ψ_1  :: Vector{T} # length-r, sum_i
    ψ_1obs :: Matrix{T} # r x m
    ψ_β1 :: AbstractVector{T} # = @view(ψ_1[1:r_X1])
    ψ_τ1 :: AbstractVector{T} # = @view(ψ_1[r_X1+1:end])

    # ψ_β1 = transpose(X1) * ψ_β1_pre.
    ψ_β1_pre :: Vector{T} # length-m.
    # ψ_τ1 = transpose(W1) * ψ_τ1_pre.
    ψ_τ1_pre :: Vector{T} # length-m.
    ψ_βτ_pre :: Vector{T}

    var_β1_pre :: T
    var_τ1_pre :: T
    var_βτ_pre :: T

    ψ_2obs  :: Matrix{T} # (p + l + q◺) x m.

    B_11 :: Matrix{T} # size r x r
    B_21 :: Matrix{T} # p + l + q◺ x r
    A_21 :: Matrix{T} # p + l + q◺ x r
    A_21_β2β1   :: AbstractMatrix{T} # p x r_X1, testbaseobs.A_21_β2β1_pre * X1
    A_21_τ2τ1   :: AbstractMatrix{T} # l x r_W1, testbaseobs.A_21_τ2τ1_pre * W1
    A_21_Lγτ1   :: AbstractMatrix{T} # q◺ x r_W1, testbaseobs.A_21_Lγτ1_pre * W1
    AinvBAinv   :: AbstractMatrix{T} # r x r.
    tmp_sr      :: AbstractMatrix{T} # p + l + q◺ x r.
    tmp_srx1    :: AbstractMatrix{T} # p + l + q◺ x r_X1.
    tmp_srw1    :: AbstractMatrix{T} # p + l + q◺ x r_W1.
    tmp_rr      :: AbstractMatrix{T} # r x r.
    tmp_rx1rx1  :: AbstractMatrix{T} # r_X1 x r_X1.
    tmp_rw1rw1  :: AbstractMatrix{T} # r_W1 x r_W1.
    tmp_r       :: AbstractVector{T}
    tmp_rx1     :: AbstractVector{T}
    tmp_rw1     :: AbstractVector{T}
end

function WSVarScoreTestInvariant(
    nullmodel::WSVarLmmModel{T},
    r_X1::Int, r_W1::Int;
    testbaseobsvec::Union{Vector{WSVarScoreTestBaseObs{T}}, Nothing} = nothing
) where {T <: BlasReal}
    @assert r_X1 >= 0 && r_W1 >= 0
    @assert r_X1 > 0 || r_W1 > 0
    if testbaseobsvec === nothing
        testbaseobsvec = [WSVarScoreTestBaseObs(obs) for obs in nullmodel.data]
    end
    @assert nullmodel.isfitted[1] "Please fit the model first."

    p, q, l, m = nullmodel.p, nullmodel.q, nullmodel.l, nullmodel.m
    q◺ = ◺(q)
    r = r_X1 + r_W1

    A_21_β2β1_rowsums = Matrix{T}(undef, p,  m)
    A_21_τ2τ1_rowsums = Matrix{T}(undef, l,  m)
    A_21_Lγτ1_rowsums = Matrix{T}(undef, q◺, m)
    for (i, baseobs) in enumerate(testbaseobsvec)
        A_21_β2β1_rowsums[:, i] .= baseobs.A_21_β2β1_rowsum
        A_21_τ2τ1_rowsums[:, i] .= baseobs.A_21_τ2τ1_rowsum
        A_21_Lγτ1_rowsums[:, i] .= baseobs.A_21_Lγτ1_rowsum
    end

    ψ_1  = Vector{T}(undef, r)
    ψ_1obs = Matrix{T}(undef, r, m)
    ψ_β1 = @view(ψ_1[1:r_X1])
    ψ_τ1 = @view(ψ_1[(r_X1 + 1):end])

    # ψ_β1 = transpose(X1) * ψ_β1_pre.
    ψ_β1_pre = Vector{T}(undef, m)
    # ψ_τ1 = transpose(W1) * ψ_τ1_pre.
    ψ_τ1_pre = Vector{T}(undef, m)
    # for joint testing
    ψ_βτ_pre = Vector{T}(undef, m)
    for (i, obs) in enumerate(nullmodel.data)
        ψ_β1_pre[i] = sum(obs.Dinv_r - transpose(obs.rt_UUt))
        ψ_τ1_pre[i] = -sum(obs.diagDVRV)
        ψ_βτ_pre[i] = ψ_β1_pre[i] + ψ_τ1_pre[i]
    end

    var_β1_pre = var(ψ_β1_pre)
    var_τ1_pre = var(ψ_τ1_pre)
    var_βτ_pre = var(ψ_βτ_pre)

    # println("ψ_β1_pre: ", ψ_β1_pre)
    # println("ψ_τ1_pre: ", ψ_τ1_pre)

    # build ψ_2obs
    ψ_2obs = Matrix{T}(undef, p + l + q◺, m)
    fill!(ψ_2obs, zero(T))
    for (i, obs) in enumerate(nullmodel.data)
        ψ_β2 = @view ψ_2obs[1:p, i]
        ψ_τ2 = @view ψ_2obs[(p + 1):(p + l), i]
        ψ_Lγ = @view ψ_2obs[(p + l + 1):end, i]
        ψ_β2 .= obs.∇β
        ψ_τ2 .= obs.∇τ
        offset = 1
        @inbounds for j in 1:q, i in j:q
            ψ_Lγ[offset] = obs.∇Lγ[i, j]
            offset += 1
        end
    end

    B_11 = Matrix{T}(undef, r, r)
    B_21 = Matrix{T}(undef, p + l + q◺, r)
    A_21 = Matrix{T}(undef, p + l + q◺, r)
    A_21_β2β1 = @view A_21[1:p, 1:r_X1]
    A_21_τ2τ1 = @view A_21[(p + 1):(p + l), (r_X1 + 1):r]
    A_21_Lγτ1 = @view A_21[(p + l + 1):end, (r_X1 + 1):r]

    AinvBAinv = nullmodel.Ainv * nullmodel.B * nullmodel.Ainv
    tmp_sr      = Matrix{T}(undef, p + l + q◺, r)
    tmp_srx1    = Matrix{T}(undef, p + l + q◺, r_X1)
    tmp_srw1    = Matrix{T}(undef, p + l + q◺, r_W1)
    tmp_rr      = Matrix{T}(undef, r, r)
    tmp_rx1rx1  = Matrix{T}(undef, r_X1, r_X1)
    tmp_rw1rw1  = Matrix{T}(undef, r_W1, r_W1)
    tmp_r       = Vector{T}(undef, r)
    tmp_rx1     = Vector{T}(undef, r_X1)
    tmp_rw1     = Vector{T}(undef, r_W1)

    WSVarScoreTestInvariant{T}(nullmodel, p, q, q◺, l, m, r_X1, r_W1, r,
        A_21_β2β1_rowsums, A_21_τ2τ1_rowsums, A_21_Lγτ1_rowsums,
        ψ_1, ψ_1obs, ψ_β1, ψ_τ1, ψ_β1_pre, ψ_τ1_pre, ψ_βτ_pre,
        var_β1_pre, var_τ1_pre, var_βτ_pre,
        ψ_2obs, B_11, B_21, A_21, A_21_β2β1, A_21_τ2τ1, A_21_Lγτ1,
        AinvBAinv, tmp_sr, tmp_srx1, tmp_srw1, tmp_rr, tmp_rx1rx1, tmp_rw1rw1,
        tmp_r, tmp_rx1, tmp_rw1
    )
end

"""
    test!(st::WSVarScoreTestInvariant, X1, W1)

Performs the score test, and returns the three p-values on the
time-invariant test data.

# Arguments

- `st::WSVarScoreTestInvariant`: data structure for the test
- `X1::Union{Nothing, AbstractMatrix{T}}`: m x r_X1 data, nothing if r_X1 == 0
- `W1::Union{Nothing, AbstractMatrix{T}}`: m x r_W1 data, nothing if r_W1 == 0
"""
function test!(st::WSVarScoreTestInvariant,
    X1::Union{Nothing, AbstractVecOrMat{T}},
    W1::Union{Nothing, AbstractVecOrMat{T}}) where {T <: BlasReal}

    if st.r_X1 > 0
        @assert size(X1, 1) == st.m
        @assert size(X1, 2) == st.r_X1
    end
    if st.r_W1 > 0
        @assert size(W1, 1) == st.m
        @assert size(W1, 2) == st.r_W1
    end

    nm = st.nullmodel
    p, q, l, m = nm.p, nm.q, nm.l, nm.m
    r_X1, r_W1, r = st.r_X1, st.r_W1, st.r

    # build ψ_1: sum_i testobs.ψ_1
    if r_X1 > 0
        st.ψ_1obs[1:r_X1, :] .= transpose(X1) .* transpose(st.ψ_β1_pre)
        mul!(st.ψ_β1, transpose(X1), st.ψ_β1_pre)
    end
    if r_W1 > 0
        st.ψ_1obs[(r_X1 + 1):end, :] .= transpose(W1) .* transpose(st.ψ_τ1_pre)
        mul!(st.ψ_τ1, transpose(W1), st.ψ_τ1_pre)
    end

    # println("ψ_β1: ", st.ψ_β1)
    # println("ψ_τ1: ", st.ψ_τ1)

    # build B_11: using BLAS.syrk!()
    BLAS.syrk!('U', 'N', one(T), st.ψ_1obs, zero(T), st.B_11)
    copytri!(st.B_11, 'U')
    lmul!(one(T) / m, st.B_11)

    # build B_21
    mul!(st.B_21, st.ψ_2obs, transpose(st.ψ_1obs), one(T) / m, zero(T))

    # build A_21: 1/m sum_i Ai_21.
    fill!(st.A_21, zero(T))
    if r_X1 > 0
        mul!(st.A_21_β2β1, st.A_21_β2β1_rowsums, X1, one(T) / m, zero(T))
    end
    if r_W1 > 0
        mul!(st.A_21_τ2τ1, st.A_21_τ2τ1_rowsums, W1, one(T) / m, zero(T))
        mul!(st.A_21_Lγτ1, st.A_21_Lγτ1_rowsums, W1, one(T) / m, zero(T))
    end

    pvals = pvalues!(st)
    return pvals
end


text:
using CompScienceMeshes, BEAST

Γ = readmesh(joinpath(dirname(@__FILE__),"sphere2.in"))
X, Y = raviartthomas(Γ), buffachristiansen(Γ)

ϵ, μ, ω = 1.0, 1.0, 1.0; κ = ω * √(ϵ*μ)
K, N = Maxwell3D.doublelayer(wavenumber=κ), NCross()
E = Maxwell3D.planewave(direction=ẑ, polarization=x̂, wavenumber=κ)
H = -1/(im*μ*ω)*curl(E)
h = (n × H) × n

@hilbertspace j
@hilbertspace m
mfie = @discretise (K+0.5N)[m,j] == h[m]  j∈X m∈Y
u = gmres(mfie)

include("utils/postproc.jl")
include("utils/plotresults.jl")


text:
using HDF5
using Base.Test
using DSGE

include("../util.jl")
path = dirname(@__FILE__)

h5 = h5open("$path/../reference/solve.h5")
TTT_expected = read(h5, "TTT")
CCC_expected = reshape(read(h5, "CCC"), 72, 1)
RRR_expected = read(h5, "RRR")
close(h5)

m = Model990()
TTT, RRR, CCC = solve(m)
@test_matrix_approx_eq TTT_expected TTT
@test_matrix_approx_eq RRR_expected RRR
@test_matrix_approx_eq CCC_expected CCC


text:
module DeepCFR

using CounterfactualRegret
const CFR = CounterfactualRegret
using Random
using ProgressMeter
using Flux
using FileIO
using RecipesBase

include("memory.jl")
include("solver.jl")
include("ESCFR.jl")
include("OSCFR.jl")
include("train.jl")
include("check.jl")
include("callback.jl")
include("fitting.jl")

export DeepCFRSolver

end


text:
tests = [
    "types",
    "fpt",
    "sim"]

println("Running tests:")

for t in tests
    test_fn = "$t.jl"
    println(" * $test_fn")
    include(test_fn)
end


text:
"""
	NVModel  <: AbstractNewsvendorProblem
Description of a newsvendor problem in its typical form.
"""

"""
	NVModel(demand, cost, price, salvage=0)

Captures a [Newsvendor model](https://en.wikipedia.org/wiki/Newsvendor_model)
with unit `cost`, unit selling `price`, and
`demand` distribution. The demand can be any univariate distribution from the
package [Distributions.jl](https://juliastats.org/Distributions.jl/latest/univariate/).
# Examples
```julia-repl
julia> using Distributions
```
```jldoctest nvm; setup = :(using Distributions, NewsvendorModel)
julia> nvm = NVModel(demand = Normal(50, 20), cost = 5, price = 7)
Data of the Newsvendor Model
 * Demand distribution: Normal{Float64}(μ=50.0, σ=20.0)
 * Unit cost: 5.00
 * Unit selling price: 7.00
```
This defines a model with unit cost 5 and unit price 7, where uncertain demand
draws from a normal distribution with mean 50 and standard deviation 20.

Optional keyword arguments and their defaults:

	NVModel(demand [; kwargs])

- `cost` for a unit; defaults to `0.0`
- `price` for selling a unit; defaults to `0.0`
- `salvage` value obtained from scraping a leftover unit; defaults to `0.0`
- `holding` cost induced by a leftover unit, e.g., extra captial cost or warehousing cost; essentially a negative salvage value; defaults to `0.0`
- `backorder` penalty for being short a unit, e.g., contractual penalty for missing delivery targets or missed future profit of an unserved customer; defaults to `0.0`
- `substitute` benefit received from selling an alternative to an unserved customer, e.g., when selling another product or serving in the future; essentially a negative backorder penalty; defaults to `0.0`
- `fixcost` fixed cost of operations; defaults to `0.0`
- `q_min` minimal feasible quantity, e.g., due to production limits; must be nonnegative; defaults to `0`
- `q_max` maximal feasible quantity, e.g., due to production limits; must be greater than or equal to `q_min`; defaults to `Inf`

# Examples
Define a newsvendor problem with unit cost 5, unit price 7, uniform demand
between 50 and 80, where a unit salvages for 0.5, and backorder comes at a penalty
of 2 per unit, and the operations incur a fixed cost of 100, as follows:
```jldoctest nvm
julia> nvm2 = NVModel(demand = Uniform(50, 80), cost = 5, price = 7, salvage = 0.5, backorder = 2, fixcost = 100)
Data of the Newsvendor Model
 * Demand distribution: Uniform{Float64}(a=50.0, b=80.0)
 * Unit cost: 5.00
 * Unit selling price: 7.00
 * Unit salvage value: 0.50
 * Unit backorder penalty: 2.00
 * Fixed cost: 100.00
```

Note that demand is a necessary argument that can be passed
without keyword in first place. Moreover, only values that differ from the default will be shown.

Define a newsvendor problem with unit cost 5, unit price 7, uniform demand
between 50 and 80, where a unit salvages for 0.5, and backorder comes at a penalty
of 2 per unit, and the operations incur a fixed cost of 100, as follows:
```jldoctest nvm
julia> nvm3 = NVModel(Uniform(50, 80), 5, 7, 0.5, backorder = 2, fixcost = 100, q_min=0)
Data of the Newsvendor Model
 * Demand distribution: Uniform{Float64}(a=50.0, b=80.0)
 * Unit cost: 5.00
 * Unit selling price: 7.00
 * Unit salvage value: 0.50
 * Unit backorder penalty: 2.00
 * Fixed cost: 100.00
julia> nvm3 == nvm2
true
```

Holding cost is essentially overage cost.
Backorder penalty is essentially underage cost.
The beer game has holding cost of 0.5 USD and backlog costs 1 USD per unit.
Demand is assumed to be uniform betwen 0 and 300.
```jldoctest nvm
julia> beer = NVModel(Uniform(0, 300), backorder = 1, holding = 0.5)
Data of the Newsvendor Model
 * Demand distribution: Uniform{Float64}(a=0.0, b=300.0)
 * Unit cost: 0.00
 * Unit selling price: 0.00
 * Unit holding cost: 0.50
 * Unit backorder penalty: 1.00
```

"""
struct NVModel <: AbstractNewsvendorProblem
    demand::UnivariateDistribution
    cost::Real
    price::Real
    salvage::Real
    holding::Real
    backorder::Real
    substitute::Real
    fixcost::Real
    q_min::Real
    q_max::Real
    function NVModel(d, c, p, s, h, b, r, f, q_min, q_max)
        if q_min > q_max return  error("q_min > q_max not feasible") end
        if q_min < 0.0 return error("q_min must be nonnegative") end
        return new(d, c, p, s, h, b, r, f, q_min, q_max)
    end
end

# Outer constructor
function NVModel(d; cost = 0.0, price = zero(cost), salvage = zero(cost),
    holding = zero(cost), backorder = zero(cost), substitute = zero(cost),
    fixcost = zero(cost), q_min = 0, q_max = Inf)
    c, p, s, h, b, r, f = promote(cost, price, salvage, holding, backorder, substitute, fixcost)
    NVModel(d, c, p, s, h, b, r, f, q_min, q_max)
end
function NVModel(d, cost, price, salvage = zero(cost);
    holding = zero(cost), backorder = zero(cost), substitute = zero(cost),
    fixcost = zero(cost), q_min = 0, q_max = Inf)
    c, p, s, h, b, r, f = promote(cost, price, salvage, holding, backorder, substitute, fixcost)
    NVModel(d, c, p, s, h, b, r, f, q_min, q_max)
end
function NVModel(;demand, cost = 0.0, price = zero(cost), salvage = zero(cost),
    holding = zero(cost), backorder = zero(cost), substitute = zero(cost),
    fixcost = zero(cost), q_min = 0, q_max = Inf)
    c, p, s, h, b, r, f = promote(cost, price, salvage, holding, backorder, substitute, fixcost)
    NVModel(demand, c, p, s, h, b, r, f, q_min, q_max)
end

# Characteristic functions
underage_cost(nvm::NVModel) = nvm.price - nvm.cost + nvm.backorder - nvm.substitute
overage_cost(nvm::NVModel) = nvm.cost - nvm.salvage + nvm.holding
distr(nvm::NVModel) = nvm.demand
"At q=0, expected profit = μ × (substitute - backorder) - fixed cost"
profit_shift(nvm::NVModel) = ( mean(nvm.demand) * (nvm.substitute - nvm.backorder)
                              - nvm.fixcost )
q_min(nvm::NVModel) = nvm.q_min
q_max(nvm::NVModel) = nvm.q_max

function Base.show(io::IO, nvm::NVModel)
    @printf io "Data of the Newsvendor Model\n"
    @printf io " * Demand distribution: "
    println(io, nvm.demand)
    @printf io " * Unit cost: %.2f\n" nvm.cost
    @printf io " * Unit selling price: %.2f" nvm.price
    if nvm.salvage != zero(nvm.cost)
        @printf io "\n * Unit salvage value: %.2f" nvm.salvage
    end
    if nvm.holding != zero(nvm.cost)
        @printf io "\n * Unit holding cost: %.2f" nvm.holding
    end
    if nvm.backorder != zero(nvm.cost)
        @printf io "\n * Unit backorder penalty: %.2f" nvm.backorder
    end
    if nvm.substitute != zero(nvm.cost)
        @printf io "\n * Unit substitute margin: %.2f" nvm.substitute
    end
    if nvm.fixcost != zero(nvm.cost)
        @printf io "\n * Fixed cost: %.2f" nvm.fixcost
    end
    if nvm.q_min != zero(nvm.cost)
        print(io, "\n * Minimal feasible quanitity: $(nvm.q_min) units")
    end
    if nvm.q_max != Inf
        print(io, "\n * Maximal feasible quanitity: $(nvm.q_max) units")
    end
    @printf io "\n"
    return
end

## Storage of results

mutable struct NVResult
    nvm::NVModel
    rounded::Bool
    q_opt::Real
    profit::Real
    underage_cost::Real
    overage_cost::Real
    critical_fractile::Real
    sales::Real
    leftover::Real
    lost_sales::Real
    salvage_revenue::Real
    holding_cost::Real
    backorder_penalty::Real
    substitute_profit::Real
end
"""
	nvm(res::NVResult)
Get underlying model from a stored result.
"""
nvmodel(res::NVResult) = res.nvm

"""
	rounded(res::NVResult)
Get information whether round-up rule was applied to get the stored result.
"""
rounded(res::NVResult) = res.rounded

"""
	q_opt(res::NVResult)
Get optimal quantity from a stored result.
"""
q_opt(res::NVResult) = res.q_opt

"""
	profit(res::NVResult)
Get expected profit from a stored result.
"""
profit(res::NVResult) = res.profit

"""
	underage_cost(res::NVResult)
Get cost of underage from a stored result.
"""
underage_cost(res::NVResult) = res.underage_cost

"""
	overage_cost(res::NVResult)
Get cost of overage from a stored result.
"""
overage_cost(res::NVResult) = res.overage_cost

"""
	critical_fractile(res::NVResult)
Get critical fractile from a stored result.
"""
critical_fractile(res::NVResult) = res.critical_fractile

"""
	sales(res::NVResult)
Get expected sales from a stored result.
"""
sales(res::NVResult) = res.sales

"""
	lost_sales(res::NVResult)
Get expected lost sales model from a stored result.
"""
lost_sales(res::NVResult) = res.lost_sales

"""
	leftover(res::NVResult)
Get expected leftover from a stored result.
"""
leftover(res::NVResult) = res.leftover

"""
    salvage_revenue(res::NVResult)
Get expected salvage revenue from a stored result.
"""
salvage_revenue(res::NVResult) = res.salvage_revenue

"""
    holding_cost(res::NVResult)
Get expected holding cost from a stored result.
"""
holding_cost(res::NVResult) = res.holding_cost

"""
	backorder_penalty(res::NVResult)
Get expected backorder penalty from a stored result.
"""
backorder_penalty(res::NVResult) = res.backorder_penalty

"""
    substitute_profit(res::NVResult)
Get expected substitute profit from a stored result.
"""
substitute_profit(res::NVResult) = res.substitute_profit

function Base.show(io::IO, r::NVResult)
    @printf io "=====================================\n"
    @printf io "Results of maximizing expected profit\n"
    if rounded(r)
        @printf io " * Optimal quantity: %d units\n" q_opt(r)
    else
        @printf io " * Optimal quantity: %.2f\n" q_opt(r)
    end
    @printf io " * Expected profit: %.2f\n" profit(r)
    @printf io "=====================================\n"
    @printf io "This is a consequence of\n"
    @printf io " * Cost of underage:  %.2f\n" underage_cost(r)
    if nvmodel(r).price != zero(nvmodel(r).cost)
        @printf io "   ╚ + Price:               %.2f\n" nvmodel(r).price
    end
    if nvmodel(r).cost != zero(nvmodel(r).cost)
        @printf io "   ╚ - Cost:                %.2f\n" nvmodel(r).cost
    end
    if nvmodel(r).backorder != zero(nvmodel(r).cost)
        @printf io "   ╚ + Backorder penalty:   %.2f\n" nvmodel(r).backorder
    end
    if nvmodel(r).substitute != zero(nvmodel(r).cost)
        @printf io "   ╚ - Substitute profit:   %.2f\n" nvmodel(r).substitute
    end
    @printf io " * Cost of overage:   %.2f\n" overage_cost(r)
    if nvmodel(r).cost != zero(nvmodel(r).cost)
        @printf io "   ╚ + Cost:                %.2f\n" nvmodel(r).cost
    end
    if nvmodel(r).salvage != zero(nvmodel(r).cost)
        @printf io "   ╚ - Salvage value:       %.2f\n" nvmodel(r).salvage
    end
    if nvmodel(r).holding != zero(nvmodel(r).cost)
        @printf io "   ╚ + Holding cost:        %.2f\n" nvmodel(r).holding
    end
    @printf io " * Critical fractile: %.2f\n" critical_fractile(r)
    if 0 >= critical_fractile(r)
        @printf io " ! No benefit from trade !\n"
    end
    if 1 <= critical_fractile(r)
        @printf io " ! No damage from excess quanitiy !\n"
        if isinf(q_opt(r))
            @printf io " ! No upper bound on quanitity !\n"
        end
    else
        @printf io " * Rounded to nearest integer: %s\n" rounded(r)
    end
    @printf io "-------------------------------------\n"
    @printf io "Ordering the optimal quantity yields\n"
    @printf io " * Expected sales: %.2f units\n" sales(r)
    @printf io " * Expected lost sales: %.2f units\n" lost_sales(r)
    @printf io " * Expected leftover: %.2f units\n" leftover(r)
    if nvmodel(r).salvage != zero(nvmodel(r).cost)
        @printf io " * Expected salvage revenue: %.2f\n" salvage_revenue(r)
    end
    if nvmodel(r).backorder != zero(nvmodel(r).cost)
        @printf io " * Expected backorder penalty: %.2f\n" backorder_penalty(r)
    end
    if nvmodel(r).holding != zero(nvmodel(r).cost)
        @printf io " * Expected holding cost: %.2f\n" holding_cost(r)
    end
    if nvmodel(r).substitute != zero(nvmodel(r).cost)
        @printf io " * Expected substitute profit: %.2f\n" substitute_profit(r)
    end
    @printf io "-------------------------------------\n"
    return
end

## Optimization

"""
	solve(nvm::NVModel [; rounded = true])

Compute the stocking quantity `q_opt(nvm)` that maximizes the expected profit
as well as further the metrics when using it.
The optimal quantity is rounded up to the next integer by default. To get the
exact real number, set the keyword arguement `rounded = false`.
"""
function solve(nvm::NVModel; rounded = true)
    q = q_opt(nvm, rounded = rounded)
    salvage_revenue = nvm.salvage * leftover(nvm, q)
    holding_cost = nvm.holding * leftover(nvm, q)
    backorder_penalty = nvm.backorder * lost_sales(nvm, q)
    substitute_profit = nvm.substitute * lost_sales(nvm, q)

    NVResult(nvm,
        rounded,
        q,
        profit(nvm, q),
        underage_cost(nvm),
        overage_cost(nvm),
        critical_fractile(nvm),
        sales(nvm, q),
        leftover(nvm, q),
        lost_sales(nvm, q),
        salvage_revenue,
        holding_cost,
        backorder_penalty,
        substitute_profit,
    )
end


text:
using Distributions
using Klara
using PGUManifoldMC

include("rv_model.jl")

using RvModelKeplerian

DATADIR = "../../data"
SUBDATADIR = "mala"

dataset = readdlm(joinpath(DATADIR, "example1.csv"), ',', header=false); # read observational data
obs_times = dataset[:,1]
obs_rv = dataset[:,2]
sigma_obs = dataset[:,3]
set_times(obs_times);     # set data to use for model evaluation
set_obs( obs_rv);
set_sigma_obs(sigma_obs);

include("utils_ex.jl")
param_true = make_param_true_ex1()
param_perturb_scale = make_param_perturb_scale(param_true)
param_init = 0
param_init = param_true.+0.01*param_perturb_scale.*randn(length(param_true))
println("param_init= ",param_init)

nmcmc = 50000
nburnin = 10000
mcrange = BasicMCRange(nsteps=nmcmc, burnin=nburnin)

outopts = Dict{Symbol, Any}(:monitor=>[:value, :logtarget], :diagnostics=>[:accept])

v0 = Dict(:p=>param_init)

p = BasicContMuvParameter(
  :p,
  logtarget=plogtarget,
  diffopts=DiffOptions(mode=:forward)
)

model = likelihood_model(p, false)
#model = likelihood_model([p], isindexed=false)

# target_accept_rates = [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.5, 0.75, 0.9]
target_accept_rates = [0.574]
#target_accept_rates = [0.25, 0.5, 0.75, 0.9]  # 0.574 for MALA
#target_accept_rates = [0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]
#target_accept_rates = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8]
#target_accept_rates = [0.574]
jobs = Array(Any,length(target_accept_rates))
mala_chains = Array(Any,length(target_accept_rates))
mala_times = Array(Float64, length(target_accept_rates))
mala_stepsizes = Array(Float64, length(target_accept_rates))
mala_acceptance = Array(Float64, length(target_accept_rates))
mala_esses = Array(Float64, length(target_accept_rates))
mala_iacts = Array(Float64, length(target_accept_rates))

for i in 1:length(target_accept_rates)
  target_accept_rate = target_accept_rates[i]
  println("# i= ", i, ": ",target_accept_rate)
  println("# p= ",param_init)
  sampler = MALA(0.02)
  mctuner = AcceptanceRateMCTuner(target_accept_rate, score=x -> logistic_rate_score(x, 3.), verbose=true)
  # mctuner = AcceptanceRateMCTuner(target_accept_rate, score=x -> logistic_rate_score(x, 0.3), verbose=true)
  #mctuner = VanillaMCTuner(verbose=true)
  v0 = Dict(:p=>param_init)
  job = BasicMCJob(
    model,
    sampler,
    mcrange,
    v0,
    tuner=mctuner,
    outopts=outopts
  )

  tic()
  run(job)
  runtime = toc()

  chain = output(job)
  ratio = acceptance(chain)

  println("# mean[",i,"] = ",mean(chain))
  println("# std[",i,"] =",vec(std(chain.value,2)))
  println("# ess[",i,"] = ",ess(chain))
  println("# iact[",i,"] = ",iact(chain))

  jobs[i] = job
  mala_times[i] = runtime
  mala_stepsizes[i] = job.sstate.tune.step
  mala_acceptance[i] = ratio
  mala_esses[i] = minimum(ess(chain))
  mala_iacts[i] = maximum(iact(chain))
  mala_chains[i] = chain
end


text:
const singular = let res_init=zeros(4), jac_init=zeros(4,4), x_init=[3.0, -1.0, 0.0, 1.0]

    function res(x, r)
        r[1] = x[1] + 10.0 * x[2]
        r[2] = sqrt(5.0) * (x[3] - x[4])
        r[3] = (x[2] - 2x[3])^2
        r[4] = sqrt(10.0) * (x[1] - x[4])^2
        return r
    end

    function jac(x, jac)
        jac[:] .= 0.0
        jac[1, 1] = 1.0
        jac[1, 2] = 10.0
        jac[2, 3] = sqrt(5.0)
        jac[2, 4] = -jac[2, 3]
        jac[3, 2] = 2.0 * (x[2] - 2.0 * x[3])
        jac[3, 3] = -2.0 * jac[3, 2]
        jac[4, 1] = sqrt(40.0) * (x[1] - x[4])
        jac[4, 4] = -jac[4, 1]
        return jac
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("singular", res, jac,
                                                   res_init, jac_init, x_init;
                                                   scale=scale, verbose=verbose,
                                                   print_steps=print_steps)
end


text:
using Measurements

function add_uncertainty_stack_first(res, port1, port2,
                                     D_rel, D_abs,
                                     Q_rel, Q_abs,
                                     p_abs, p_std_fac, p_rel,
                                     d_deltaL,
                                     stacksize, inds=res.resl.irun)
    @unpack Q, press, press_std, trun  = res.resl

    ii = find_low_std(press_std, [port1, port2], inds);
    trun = trun[ii];
    inds = inds[ii];

    D = get_time_series(res.resi, trun)[1]
    Q = Q[inds]
    p1, p2 = press[port1][inds], press[port2][inds]
    p1_std, p2_std = press_std[port1][inds], press[port2][inds]
    deltaL = port_dist(port1, port2, res)

    D, Q, p1, p2, grad_p = add_uncertainty(D, Q, p1, p2, deltaL,
                                           D_rel, D_abs,
                                           Q_rel, Q_abs,
                                           p_abs,
                                           p1_std*p_std_fac, p2_std*p_std_fac,
                                           p_rel)
    D, Q, p1, p2, p1_std, p2_std = [stack_it(vec,stacksize) for vec in (D,Q,p1,p2,p1_std,p2_std)];
    tstac = stack_it_time(trun, stacksize)
    grad_p = calculate_press_grad(p1, p2, 0, 0, deltaL)[1]
    return D, Q, p1, p2, grad_p, tstac
end

function add_uncertainty(res::RcRes, port1, port2,
                         D_rel, D_abs,
                         Q_rel, Q_abs,
                         p_abs, p_std_fac, p_rel,
                         d_deltaL)
    @unpack Q, press, press_std, trun, irun  = res.resl
    p1, p2 = press[port1][irun], press[port2][irun]
    p1_std, p2_std = press_std[port1][irun], press[port2][irun]
    Q = Q[irun]
    D = get_time_series(res.resi, trun)[1]

    deltaL = port_dist(port1, port2, res) ± d_deltaL
    add_uncertainty(D, Q, p1, p2, deltaL,
                    D_rel, D_abs,
                    Q_rel, Q_abs,
                    p_abs, p1_std*p_std_fac, p2_std*p_std_fac,
                    p_rel)
end
function add_uncertainty(D, Q, p1, p2, deltaL,
                         D_rel, D_abs,
                         Q_rel, Q_abs,
                         p_abs, p1_std, p2_std, p_rel)
    D = relabs.(D, D_rel, D_abs)
    Q = relabs.(Q, Q_rel, Q_abs)
    p1 = relabs.(p1, p_rel, p_abs+p1_std)
    p2 = relabs.(p2, p_rel, p_abs+p2_std)
    grad_p = calculate_press_grad(p1, p2, 0, 0, deltaL)[1]
    return D, Q, p1, p2, grad_p
end

relabs(v, rel, abs) = v±(abs + v*rel)

function P.plot(t, v::Vector{<:Measurements.Measurement}; kwargs...)
    @show "here"
    va = Measurements.value.(v)
    err = Measurements.uncertainty.(v)
    P.fill_between(t, va-err, va+err, alpha=0.3)
    P.plot(t, va; kwargs...)
end
P.plot(v::Vector{<:Measurements.Measurement}; kwargs...) = P.plot(0:length(v)-1, v; kwargs...)

# https://matplotlib.org/examples/pylab_examples/ellipse_demo.html
function P.scatter(x::Vector{<:Measurements.Measurement}, y::Vector{<:Measurements.Measurement}, col::Vector;
                   scale_size=10, kwargs...)
    xx = Measurements.value.(x)
    xe = Measurements.uncertainty.(x)
    yy = Measurements.value.(y)
    ye = Measurements.uncertainty.(y)
    sizes = max.(xe,ye)*scale_size
    P.scatter(xx,yy,sizes,col; kwargs...)
end
function P.scatter(x::Vector{<:Measurements.Measurement}, y::Vector{<:Measurements.Measurement}; scale_size=10, kwargs...)
    xx = Measurements.value.(x)
    xe = Measurements.uncertainty.(x)
    yy = Measurements.value.(y)
    ye = Measurements.uncertainty.(y)
    sizes = max.(xe,ye)*scale_size
    P.scatter(xx,yy,sizes; kwargs...)
end

function P.errorbar(x::Vector{<:Measurements.Measurement}, y::Vector{<:Measurements.Measurement}; kwargs...)
    xx = Measurements.value.(x)
    xe = Measurements.uncertainty.(x)
    yy = Measurements.value.(y)
    ye = Measurements.uncertainty.(y)
    P.errorbar(xx,yy,ye,xe; kwargs...)
end


text:
const millisecond = 1.0 / 31556952000.0
const second = 1.0 / 31556952.0
const minute = 1.0 / 525949.2
const hour = 1.0 / 8765.82
const day = 1.0 / 365.2425
const week = 7.0 / 365.2425
const fortnight = 14.0 / 365.2425
const month = 1.0 / 12.0


text:
using Documenter
using ODEHybrid

makedocs(
    sitename = "ODEHybrid.jl",
    format = Documenter.HTML(prettyurls = false),
    pages = [
        "Introduction" => "index.md",
        # "Getting Started"
        # "Main Functions"
        # "Time Series Logger"
        # "Helper Functions"    # and conversions?
        "API" => "api.md"
    ]
)

deploydocs(
    repo = "github.com/benjjensen/ODEHybrid.jl.git",
    devbranch = "main"
)


text:
# Input data
data = hcat(map(v -> map(x->parse(Int64, x), v), map(line->split(line, ""), readlines("Day09.input")))...)';

# TODO

# Part 1
# ans = mindistance((x, data) -> sum(map(v -> abs(v - x), data)), data);
# println("Part 1: $ans")

cols = size(data, 1)
rows = size(data, 2)

countrows = (x, y) -> begin
  #  @show ((x, y), data[x, y])
    if x < 1 || y < 1 || x > cols || y > rows || data[x, y] == 9
        return 0
    end
    #println("OK", x, y, " ", data[x, y])

    data[x, y] = 9

    return 1 + countrows(x-1, y) + countrows(x+1, y) + countrows(x, y-1) + countrows(x, y+1)
end

v = []
for i in 1:cols
    for j in 1:rows
        push!(v, countrows(i, j))
    end
end

h = prod(reverse(sort(filter(>(0), v)))[1:3])

# Part 2
println("Part 2: $h")


text:
"wavenumber from Challis, R. E., et al. Ultrasound techniques for characterizing colloidal dispersions. Reports on progress in physics 68.7 (2005): 1541."
wavenumber_challis(ωs::AbstractArray,medium::Acoustic{T,2}, species::Species{T,2}; kws...) where T<:AbstractFloat = [wavenumber_challis(ω, medium, species; kws...) for ω in ωs]

# will sum Hankel orders only up to n =2 as in the paper "Ultrasound techniques for characterizing colloidal dispersions"
function wavenumber_challis(ω::Number, medium::Acoustic{T,2}, species::Species{T,2};
    radius_multiplier = 1.005, verbose = false, basis_order=2) where T<:AbstractFloat
  # tol=0.004; radius_multiplier = 1.005
  # a12 = radius_multiplier(a1+a2)

  volume_fraction = sum(sp.volume_fraction for sp in species)
  if volume_fraction >= 0.4
    @warn("the volume fraction $(round(100*volume_fraction))% is a bit too high, expect a relative error of approximately $(round(100*volume_fraction^3.0))%")
  end

  kT2 = (ω/medium.c)^2.0
  kT2 += 4.0im*sum(number_density(sp) * Zn(ω,sp,medium,m) for sp in species
  , m in -basis_order:basis_order)

  # second order number fraction, sum up too same hankel order
  kT2 += -(8.0/(ω/medium.c)^2)*sum(
    number_density(sp)^2 * abs(m2-m1)*Zn(ω,sp1,medium,m1)*Zn(ω,sp1,medium,m2)
  for sp1 in species, m1 = -basis_order:basis_order, m2 = -basis_order:basis_order)

  return sqrt(kT2)
end

"low volfrac and low wavenumber. Fails badly for strong scatterers"
function one_species_low_wavenumber(ωs, medium::Acoustic{T,2}, sp::Specie{T,2}) where T<:AbstractFloat
  φ = sp.volume_fraction
  kS = 1.0/sp.particle.medium.c
  k = 1.0/medium.c
  P = 1 - kS^2*medium.ρ/(k^2*medium.ρ)
  Q = (medium.ρ - sp.particle.medium.ρ)/(medium.ρ + sp.particle.medium.ρ)
  (ωs./medium.c).*(1 - (φ/2)*(P+2*Q)-(φ^2/8)*(2*P^2-(P+2*Q)^2))
end

function wavenumber_far_field_low_volumefraction(ω::Complex{T}, medium::Acoustic{T,2}, species::Species{T,2};
        tol=0.0002, #radius_multiplier = 1.005,
        verbose = false) where T<:AbstractFloat
  # tol=0.004; radius_multiplier = 1.005
  # a12 = radius_multiplier(a1+a2)

  volume_fraction = sum(sp.volume_fraction for sp in species)
  if volume_fraction >= 0.4
    @warn("the volume fraction $(round(100*volume_fraction))% is a bit too high, expect a relative error of approximately $(round(100*volume_fraction^3.0))%")
  end
  kT2 = (ω/medium.c)^2.0
  next_order = 4.0im*sum(number_density(sp) * Zn(ω,sp,medium,0) for sp in species)
  basis_order=1

  # sum more hankel orders until the relative error < tol
  while abs(next_order/kT2) > tol
    kT2 += next_order
    basis_order +=1
    next_order = 4.0im*sum(number_density(sp) * Zn(ω,sp,medium,m) for sp in species, m in (-basis_order,basis_order))
  end
  kT2 += next_order
  if verbose println("max Hankel order = $basis_order") end

  Zns_vec = [- t_matrix(s, medium, ω, basis_order) for s in species]

  # second order number fraction, sum up too same hankel order
  kT2 += -8.0*sum(
    number_density(species[l1]) * number_density(species[l2]) * abs(m2-m1) / (ω/medium.c)^2 * Zns_vec[l1][m1+basis_order+1,m1+basis_order+1] * Zns_vec[l2][m2+basis_order+1,m2+basis_order+1]
  for l1 in eachindex(species), l2 in eachindex(species), m1 = -basis_order:basis_order, m2 = -basis_order:basis_order)

  return sqrt(kT2)
end

max_stars_repo_name: arturgower/EffectiveWaves