text:
using GNSSTools
using ProgressMeter
using JLD
using Statistics
using PyPlot
pygui(true)

function get_metric(trackresults; start_t=0.5)
    T = trackresults.T
    start_idx = floor(Int, start_t/T)
    σ²ᵩ = var(trackresults.dphi_meas[start_idx:end])
    return σ²ᵩ
end

function compare_tracking_cases(f_d_cases, fd_rate_cases, h_parms, CN0s, Ts,
                                sigtype, f_s, t_length,
                                qₐ, state_num, σω, acquisition_T, M,
                                fine_acq_T, dll_b;
                                doppler_bin_num=2, f_if=0, nADC=8,
                                iterations=1, p_fa=1e-7, p_d=0.9, prn=1,
                                attempts=3, save_to_file=missing, channel="I",
                                case_names=missing,
                                use_fd_rate_init=false)
    results = zeros(length(Ts), length(f_d_cases), length(h_parms), length(CN0s))
    signal = definesignal(sigtype, f_s, t_length; f_if=f_if, nADC=nADC)
    N = length(Ts)*length(f_d_cases)*length(h_parms)*length(CN0s)*iterations
    sample_num_per_1ms = floor(Int, f_s*1e-3)
    p = Progress(N, 1, "Processing...")
    for i in 1:length(Ts)
        T = Ts[i]
        for j in 1:length(f_d_cases)
            f_d = f_d_cases[j]
            fd_center = round(f_d*T)/T
            fd_range = doppler_bin_num/T
            fd_rate = fd_rate_cases[j]
            for k in 1:length(h_parms)
                h_parm = h_parms[1]
                h₀ = h_parm[3]
                h₋₁ = h_parm[2]
                h₋₂ = h_parm[1]
                for g in 1:length(CN0s)
                    CN0 = CN0s[g]
                    metric_val = 0.
                    for iteration in 1:iterations
                        P_d = 0
                        # attempt = 0
                        # while (P_d < p_d) || (attempt != attempts)
                            ϕ₀ = rand(0:0.0001:2π)
                            code_start_idx = rand(1:1:sample_num_per_1ms)
                            definesignal!(signal; f_d=f_d, fd_rate=fd_rate,
                                          phi=ϕ₀, code_start_idx=code_start_idx,
                                          receiver_h_parms=h_parm, CN0=CN0,
                                          new_thermal_noise=true,
                                          new_phase_noise=true)
                            generatesignal!(signal)
                            σᵩ² = phase_noise_variance(CN0, T)
                            if use_fd_rate_init
                                fd_rate_init = fd_rate + σω*randn()
                            else
                                fd_rate_init = 0
                            end
                            acqresults, trackresults, P_d, above_threshold = process(signal,
                                            sigtype,
                                            prn, channel;
                                            fine_acq_method=:fft,
                                            fd_center=fd_center,
                                            fd_range=fd_range,
                                            h₀=h₀,
                                            h₋₂=h₋₂,
                                            q_mult=1,
                                            cov_mult=1,
                                            R_mult=1,
                                            σᵩ²=σᵩ²,
                                            σω=σω,
                                            q_a=qₐ,
                                            dynamickf=true,
                                            state_num=state_num,
                                            dll_b=dll_b,
                                            acquisition_T=acquisition_T,
                                            fine_acq_T=fine_acq_T,
                                            tracking_T=T,
                                            M=M,
                                            fd_rate=fd_rate_init,
                                            show_plot=false,
                                            return_Pd=true)
                            # attempt += 1
                        # end
                        next!(p)
                        metric_val += get_metric(trackresults)
                    end
                    results[i,j,k,g] = sqrt(metric_val / iterations)
                end
            end
        end
    end
    if ~ismissing(save_to_file)
        save(string(save_to_file, ".jld"),
             "results", results,
             "f_d_cases", f_d_cases,
             "fd_rate_cases", fd_rate_cases,
             "h_parms", h_parms,
             "Ts", Ts,
             "CN0s", CN0s,
             "state_num", state_num,
             "q_a", qₐ,
             "acquisition_T", acquisition_T,
             "fine_acq_T", fine_acq_T,
             "M", M,
             "dll_b", dll_b,
             "σω", σω,
             "channel", channel,
             "sigtype_name", sigtype.name,
             "nADC", nADC,
             "t_length", t_length,
             "f_s", f_s,
             "doppler_bin_num", doppler_bin_num,
             "p_fa", p_fa,
             "p_d", p_d,
             "f_if", f_if,
             "attempts", attempts,
             "case_names", case_names,
             "use_fd_rate_init", use_fd_rate_init)
    end
    return results
end

function make_plot_for_tracking_cases(file_name; figsize=(8, 3),
                                      save_to_file=false,
                                      case_names=missing)
    data = load(string(file_name, ".jld"),
                "results",
                "f_d_cases",
                "fd_rate_cases",
                "h_parms",
                "Ts",
                "CN0s",
                "state_num",
                "q_a")
    results, f_d_cases, fd_rate_cases, h_parms, Ts, CN0s, state_num, qₐ = data
    cmap = get_cmap("viridis")
    fig = figure(figsize=figsize)
    total_subplots = length(Ts)
    cols = 2
    rows = ceil(Int, total_subplots/cols)
    for i in 1:length(Ts)
        T = Ts[i]
        ax = fig.add_subplot(rows, cols, i)
        for j in 1:length(f_d_cases)
            f_d = f_d_cases[j]
            fd_rate = fd_rate_cases[j]
            color = cmap(float(j)/length(f_d_cases))
            if ismissing(case_names)
                label = "Case $(j)"
            else
                label = case_names[j]
            end
            if i == 1
                label = string(label, "; TCXO")
                ax.plot(CN0s, results[i,j,1,:].*(180/π), "-", color=color,
                        label=label)
                ax.plot(CN0s, results[i,j,2,:].*(180/π), ":", color=color)
            elseif i == 2
                label = string(label, "; OCXO")
                ax.plot(CN0s, results[i,j,1,:].*(180/π), "-", color=color)
                ax.plot(CN0s, results[i,j,2,:].*(180/π), ":", color=color,
                        label=label)
            else
                ax.plot(CN0s, results[i,j,1,:].*(180/π), "-", color=color)
                ax.plot(CN0s, results[i,j,2,:].*(180/π), ":", color=color)
            end
        end
        xlabel("C/N₀ (dB⋅Hz)")
        ylabel("Unfiltered σᵩ (degrees)")
        letter = ('a' + (i-1))
        title("($(letter))")
        legend()
    end
    subplots_adjust(wspace=0.35, bottom=0.15, left=0.06, right=0.94, top=0.92)
    if save_to_file
        savefig(string(file_name, ".pdf"), dpi=300)
    end
end

low_fd = 4e3  # GPS max(f_d)
high_fd = 35.2e3  # Starlink max(f_d)

low_fd_rate = -0.730  # GPS min(fd_rate)
high_fd_rate = -475.0  # Starlink min(fd_rate)

f_d_cases     = [     low_fd,       low_fd,      high_fd,     high_fd]
fd_rate_cases = [low_fd_rate, high_fd_rate, high_fd_rate, low_fd_rate]
case_names = [string("Low ", L"f_d", "; Low ", L"\dot{f_d}"),
              string("Low ", L"f_d", "; High ", L"\dot{f_d}"),
              string("High ", L"f_d", "; High ", L"\dot{f_d}"),
              string("High ", L"f_d", "; Low ", L"\dot{f_d}")]
case_num = length(f_d_cases)

tcxo_h_parms = [1.9e-6, 2e-6, 9.7e-8, 9.1e-11, 9.9e-15] ./ 16.3676e6^2  # TCXO
ocxo_h_parms = [2.2e-9, 1.6e-10, 1.6e-10, 7.1e-13, 7.2e-16] ./ 10e6^2  # OCXO
h_parms = [tcxo_h_parms, ocxo_h_parms]

CN0s = Array(range(20, 50, step=1))

iterations = 10

channel = "I"
sigtype = define_l1ca_code_type()
# Ts = [1e-3 2e-3 5e-3 10e-3]
Ts = [1e-3 10e-3]
f_s = 2.5e6
name = "l1ca"

# channel = "Q"
# sigtype = define_l5_code_type(;channel=channel)
# Ts = [1e-3 20e-3]
# f_s = 25e6
# name = "l5"

t_length = 2
p_fa = 1e-7
qₐ = 0.1
state_num = 2
σω = 1
acquisition_T=20e-3
fine_acq_T=50e-3
M=10
dll_b=0.1

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"
file_name_2state = string(directory, "ch4_compare_tracking_with_diff_cases_$(state_num)-state_q_a=$(qₐ)")

results = compare_tracking_cases(f_d_cases, fd_rate_cases, h_parms, CN0s, Ts,
                                 sigtype, f_s, t_length,
                                 qₐ, state_num, σω, acquisition_T, M,
                                 fine_acq_T, dll_b;
                                 doppler_bin_num=2, f_if=0, nADC=8,
                                 iterations=iterations, p_fa=1e-7, p_d=0.9,
                                 attempts=3, save_to_file=file_name_2state,
                                 channel=channel, case_names=case_names)

t_length = 2
p_fa = 1e-7
qₐ = 0.1
state_num = 2
σω = 1
acquisition_T=20e-3
fine_acq_T=50e-3
M=10
dll_b=0.1

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"
file_name_2state_give_fd_rate = string(directory, "ch4_compare_tracking_with_diff_cases_$(state_num)-state_q_a=$(qₐ)_give_fd_rate")

results = compare_tracking_cases(f_d_cases, fd_rate_cases, h_parms, CN0s, Ts,
                                 sigtype, f_s, t_length,
                                 qₐ, state_num, σω, acquisition_T, M,
                                 fine_acq_T, dll_b;
                                 doppler_bin_num=2, f_if=0, nADC=8,
                                 iterations=iterations, p_fa=1e-7, p_d=0.9,
                                 attempts=3,
                                 save_to_file=file_name_2state_give_fd_rate,
                                 channel=channel, case_names=case_names,
                                 use_fd_rate_init=true)

t_length = 2
p_fa = 1e-7
qₐ = 0.1
state_num = 3
σω = 1
acquisition_T=20e-3
fine_acq_T=50e-3
M=10
dll_b=0.1

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"
file_name_3state_q_a_dont_give_fd_rate = string(directory, "ch4_compare_tracking_with_diff_cases_$(state_num)-state_q_a=$(qₐ)_dont_give_fd_rate")

results = compare_tracking_cases(f_d_cases, fd_rate_cases, h_parms, CN0s, Ts,
                                 sigtype, f_s, t_length,
                                 qₐ, state_num, σω, acquisition_T, M,
                                 fine_acq_T, dll_b;
                                 doppler_bin_num=2, f_if=0, nADC=8,
                                 iterations=iterations, p_fa=1e-7, p_d=0.9,
                                 attempts=3,
                                 save_to_file=file_name_3state_q_a_dont_give_fd_rate,
                                 channel=channel,
                                 case_names=case_names)

t_length = 2
p_fa = 1e-7
qₐ = 0.1
state_num = 3
σω = 1
acquisition_T=20e-3
fine_acq_T=50e-3
M=10
dll_b=0.1

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"
file_name_3state_q_a_give_fd_rate = string(directory, "ch4_compare_tracking_with_diff_cases_$(state_num)-state_q_a=$(qₐ)_give_fd_rate")

results = compare_tracking_cases(f_d_cases, fd_rate_cases, h_parms, CN0s, Ts,
                                 sigtype, f_s, t_length,
                                 qₐ, state_num, σω, acquisition_T, M,
                                 fine_acq_T, dll_b;
                                 doppler_bin_num=2, f_if=0, nADC=8,
                                 iterations=iterations, p_fa=1e-7, p_d=0.9,
                                 attempts=3,
                                 save_to_file=file_name_3state_q_a_give_fd_rate,
                                 channel=channel,
                                 case_names=case_names,
                                 use_fd_rate_init=true)

t_length = 2
p_fa = 1e-7
qₐ = 1000
state_num = 3
σω = 1
acquisition_T=20e-3
fine_acq_T=50e-3
M=10
dll_b=0.1

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"
file_name_3state = string(directory, "ch4_compare_tracking_with_diff_cases_$(state_num)-state_q_a=$(qₐ)")

results = compare_tracking_cases(f_d_cases, fd_rate_cases, h_parms, CN0s, Ts,
                                 sigtype, f_s, t_length,
                                 qₐ, state_num, σω, acquisition_T, M,
                                 fine_acq_T, dll_b;
                                 doppler_bin_num=2, f_if=0, nADC=8,
                                 iterations=iterations, p_fa=1e-7, p_d=0.9,
                                 attempts=3, save_to_file=file_name_3state,
                                 channel=channel,
                                 case_names=case_names)

make_plot_for_tracking_cases(file_name_2state; case_names=case_names,
                             save_to_file=true)
make_plot_for_tracking_cases(file_name_2state_give_fd_rate; case_names=case_names,
                             save_to_file=true)
make_plot_for_tracking_cases(file_name_3state_q_a_dont_give_fd_rate;
                             case_names=case_names, save_to_file=true)
make_plot_for_tracking_cases(file_name_3state_q_a_give_fd_rate;
                             case_names=case_names, save_to_file=true)
make_plot_for_tracking_cases(file_name_3state; case_names=case_names,
                             save_to_file=true)


text:
@testset "Generic Groups" begin
  include("Grp/GenGrp.jl")
  include("Grp/Morphisms.jl")
end

ext: jl


text:
# Dataset object
abstract type AbstractDataset{T <: Real} end

function nfeatures(data::AbstractDataset) end

function ndata(data::AbstractDataset) end

function predict!(data::AbstractDataset, x::AbstractVector, intercept) end

function loss(x::AbstractVector, data::AbstractDataset) end

function gradient!(grad::AbstractVector, x::AbstractVector,
                   data::AbstractDataset, fit_intercept::Bool) end

function hessian!(hess::AbstractVector, x::AbstractVector,
                   data::AbstractDataset, fit_intercept::Bool) end

function diaghess!(diagh::AbstractVector, x::AbstractVector,
                   data::AbstractDataset, fit_intercept::Bool) end

function hessvec!(hv::AbstractVector, x::AbstractVector,
                  data::AbstractDataset, fit_intercept::Bool) end

struct EmptyDataset{T <: Real}  <: AbstractDataset{T} end

ndata(::EmptyDataset) = 0
nfeatures(::EmptyDataset) = 0
loss(::EmptyDataset) = NaN


text:
export trial_result

@doc """ type trial_result
  Julia type that stores both the parameters and the results of a trial.
"""
mutable struct trial_result
  nn_simtype::Int64
  n::Int64    # sample size, must be <= N
  N::Int64    # popsize
  N_mu::Float64  # N*mu, population mutation rate.  theta = 2*N_mu
  ngens::Int64
  burn_in::Float64  # The number of generations of burn-in is N/mu + 50
  dfe::Function
  dfe_str::AbstractString
  use_poplist::Bool
  mean_fitness::Float64
  expected_richness::Float64  # sum_{i=0}^{n-1} theta/(theta+i) where theta = 2*N_mu.
  average_richness::Float64   # Average number of traits in populations
  stderr_richness::Float64
  expected_w_homoz::Float64
  w_homoz::Float64
  stderr_w_homoz::Float64
  IQV::Float64
  stderr_IQV::Float64
  entropy::Float64
  stderr_entropy::Float64
end

# Constructor that sets the parameters
function trial_result( nn_simtype::Int64, n::Int64, N::Int64, N_mu::Float64, ngens::Int64,
    burn_in::Float64=2.0, dfe::Function=dfe_neutral,
    dfe_str::AbstractString="neutral"; use_poplist=false )
  tr = trial_result( nn_simtype, n, N, N_mu, ngens, burn_in, dfe, dfe_str, use_poplist, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 )
  add_expected_richness( tr )
  tr
end

max_stars_repo_name: um-tech-evolution/nearly-neutral

max_issues_repo_name: um-tech-evolution/nearly-neutral

max_forks_repo_name: um-tech-evolution/nearly-neutral


text:
# -----------------------------------------------------------------------------
##### CYLINDER/SHELL #####-----------------------------------------------------
function shellregioncylinder(dataobject::ClumpDataType;
                            radius::Array{<:Real,1}=[0.,0.],
                            height::Real=0.,
                            center::Array{<:Any,1}=[0.,0.,0.],
                            range_unit::Symbol=:standard,
                            direction::Symbol=:z,
                            inverse::Bool=false,
                            verbose::Bool=verbose_mode)

    printtime("", verbose)

    radius_in  = radius[1]
    radius_out = radius[2]
    if radius_in == 0. || radius_out == 0. || height == 0. || in(0., center)
        error("[Mera]: given radius, height or center should be != 0.")
    end

    boxlen = dataobject.boxlen
    scale = dataobject.scale

    # convert given ranges and print overview on screen
    ranges, cx_shift, cy_shift, cz_shift, radius_in_shift, radius_out_shift, height_shift = prep_cylindrical_shellranges(dataobject.info, center, radius_in, radius_out, height, range_unit, verbose)

    if inverse == false
        sub_data = filter(p-> sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2)
                                    >= ( radius_in_shift*boxlen )  &&

                              sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2)
                                    <= ( radius_out_shift*boxlen ) &&

                            abs(p.peak_z - cz_shift*boxlen) <= ( height_shift*boxlen),
                                dataobject.data)
    elseif inverse == true
        sub_data = filter(p-> sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2)
                                    < ( radius_in_shift*boxlen )  ||

                              sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2)
                                    > ( radius_out_shift*boxlen ) ||

                            abs(p.peak_z - cz_shift*boxlen) > ( height_shift*boxlen),
                                dataobject.data)
        ranges = dataobject.ranges
    end

    printtablememory(sub_data, verbose)

    clumpdata = ClumpDataType()
    clumpdata.data = sub_data
    clumpdata.info = dataobject.info
    clumpdata.boxlen = dataobject.boxlen
    clumpdata.ranges = ranges
    clumpdata.selected_clumpvars = dataobject.selected_clumpvars
    clumpdata.used_descriptors = dataobject.used_descriptors
    clumpdata.scale = dataobject.scale
    return clumpdata
end

# -----------------------------------------------------------------------------
##### SPHERE/SHELL #####-------------------------------------------------------
function shellregionsphere(dataobject::ClumpDataType;
                            radius::Array{<:Real,1}=[0.,0.],
                            center::Array{<:Any,1}=[0.,0.,0.],
                            range_unit::Symbol=:standard,
                            inverse::Bool=false,
                            verbose::Bool=verbose_mode)

    printtime("", verbose)

    radius_in  = radius[1]
    radius_out = radius[2]
    if radius_in == 0. || radius_out == 0. || in(0., center)
        error("[Mera]: given inner and outer radius or center should be != 0.")
    end

    boxlen = dataobject.boxlen
    scale = dataobject.scale

    # convert given ranges and print overview on screen
    ranges, cx_shift, cy_shift, cz_shift, radius_in_shift, radius_out_shift = prep_spherical_shellranges(dataobject.info, center, radius_in, radius_out, range_unit, verbose)

    if inverse == false
        sub_data = filter(p-> sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2+
                                    (p.peak_z -  cz_shift*boxlen)^2 )
                                    >= ( radius_in_shift*boxlen ) &&

                                sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2+
                                    (p.peak_z -  cz_shift*boxlen)^2 )
                                    <= ( radius_out_shift*boxlen ),
                                dataobject.data)

    elseif inverse == true
        sub_data = filter(p-> sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2+
                                    (p.peak_z -  cz_shift*boxlen)^2 )
                                    < ( radius_in_shift*boxlen ) ||

                                sqrt( (p.peak_x -  cx_shift*boxlen)^2 +
                                    (p.peak_y -  cy_shift*boxlen )^2+
                                    (p.peak_z -  cz_shift*boxlen)^2 )
                                    > ( radius_out_shift*boxlen ),
                                dataobject.data)
        ranges = dataobject.ranges
    end

    printtablememory(sub_data, verbose)

    clumpdata = ClumpDataType()
    clumpdata.data = sub_data
    clumpdata.info = dataobject.info
    clumpdata.boxlen = dataobject.boxlen
    clumpdata.ranges = ranges
    clumpdata.selected_clumpvars = dataobject.selected_clumpvars
    clumpdata.used_descriptors = dataobject.used_descriptors
    clumpdata.scale = dataobject.scale
    return clumpdata
end


text:
using Example2018

@show Point(1,2) + Point(3,4)

q= MPoint(1,2)
@show q
q.x=9
@show q


text:
using Sockets: IPAddr, IPv4, IPv6

# Copy of `Sockets._sizeof_uv_interface_address`
const _sizeof_uv_interface_address = ccall(:jl_uv_sizeof_interface_address, Int32, ())

# https://github.com/JuliaLang/julia/pull/30349
if VERSION < v"1.2.0-DEV.56"
    using Base: uv_error

    function getipaddrs()
        addresses = IPv4[]
        addr_ref = Ref{Ptr{UInt8}}(C_NULL)
        count_ref = Ref{Int32}(1)
        lo_present = false
        err = ccall(
            :jl_uv_interface_addresses,
            Int32,
            (Ref{Ptr{UInt8}}, Ref{Int32}),
            addr_ref,
            count_ref,
        )
        uv_error("getlocalip", err)
        addr, count = addr_ref[], count_ref[]
        for i in 0:(count - 1)
            current_addr = addr + i * _sizeof_uv_interface_address
            if 1 == ccall(
                :jl_uv_interface_address_is_internal, Int32, (Ptr{UInt8},), current_addr
            )
                lo_present = true
                continue
            end
            sockaddr = ccall(
                :jl_uv_interface_address_sockaddr, Ptr{Cvoid}, (Ptr{UInt8},), current_addr
            )
            if ccall(:jl_sockaddr_in_is_ip4, Int32, (Ptr{Cvoid},), sockaddr) == 1
                push!(
                    addresses,
                    IPv4(ntoh(ccall(:jl_sockaddr_host4, UInt32, (Ptr{Cvoid},), sockaddr))),
                )
            end
        end
        ccall(:uv_free_interface_addresses, Cvoid, (Ptr{UInt8}, Int32), addr, count)
        return addresses
    end
else
    using Sockets: getipaddrs
end

"""
    is_link_local(ip::IPv4) -> Bool

Determine if the IP address is within the [link-local address]
(https://en.wikipedia.org/wiki/Link-local_address) block 169.254.0.0/16.
"""
is_link_local(ip::IPv4) = ip"169.254.0.0" <= ip <= ip"169.254.255.255"

# Julia structure mirroring `uv_interface_address_t`
# http://docs.libuv.org/en/v1.x/misc.html#c.uv_interface_address_t
struct InterfaceAddress{T<:IPAddr}
    name::String       # Name of the network interface
    is_internal::Bool  # Interface is a loopback device
    address::T
    # netmask::T  # No accessors available currently
end

# Julia commit 58bafe499b
const _jl_sockaddr_is_ip4, _jl_sockaddr_is_ip6 = let
    prefix = VERSION >= v"1.3.0-DEV.434" ? :jl_sockaddr_is_ : :jl_sockaddr_in_is_
    Symbol(prefix, :ip4), Symbol(prefix, :ip6)
end

# Based upon `getipaddrs` in "stdlib/Sockets/src/addrinfo.jl"
function get_interface_addrs()
    addresses = InterfaceAddress[]
    addr_ref = Ref{Ptr{UInt8}}(C_NULL)
    count_ref = Ref{Int32}(1)
    lo_present = false
    err = ccall(
        :jl_uv_interface_addresses,
        Int32,
        (Ref{Ptr{UInt8}}, Ref{Int32}),
        addr_ref,
        count_ref,
    )
    Base.uv_error("getlocalip", err)
    addr, count = addr_ref[], count_ref[]
    for i in 0:(count - 1)
        current_addr = addr + i * _sizeof_uv_interface_address
        # Note: Extracting interface name without a proper accessor
        name = unsafe_string(unsafe_load(Ptr{Cstring}(current_addr)))
        is_internal =
            ccall(
                :jl_uv_interface_address_is_internal, Int32, (Ptr{UInt8},), current_addr
            ) == 1
        sockaddr = ccall(
            :jl_uv_interface_address_sockaddr, Ptr{Cvoid}, (Ptr{UInt8},), current_addr
        )
        ip = if ccall(_jl_sockaddr_is_ip4, Int32, (Ptr{Cvoid},), sockaddr) == 1
            IPv4(ntoh(ccall(:jl_sockaddr_host4, UInt32, (Ptr{Cvoid},), sockaddr)))
        elseif ccall(_jl_sockaddr_is_ip6, Int32, (Ptr{Cvoid},), sockaddr) == 1
            addr6 = Ref{UInt128}()
            scope_id = ccall(
                :jl_sockaddr_host6, UInt32, (Ptr{Cvoid}, Ref{UInt128}), sockaddr, addr6
            )
            IPv6(ntoh(addr6[]))
        end
        push!(addresses, InterfaceAddress(name, is_internal, ip))
    end
    ccall(:uv_free_interface_addresses, Cvoid, (Ptr{UInt8}, Int32), addr, count)
    return addresses
end


text:
using Revise
using FRGRealTime
using FRGRealTime.MathFuns
using FRGRealTime.ThresholdFunctions
ThresholdFunctions.ImFb2(1.0,1.0,1.0,1.0,0.1)
Eb(1.0,1.0)
ThresholdFunction


text:
module DatasetManager

using Glob, DataFrames, CategoricalArrays, PrettyTables, StatsBase, Printf, ThreadPools,
    ProgressMeter, Crayons

export DataSubset, TrialConditions, Trial, DuplicateSourceError, Segment, SegmentResult,
    AbstractSource, Source, MissingSourceError, UnknownDeps

export findtrials, summarize, analyzedataset, export_trials, subject, conditions, sources, hassource,
    getsource, sourcepath, readsource, requiresource!, generatesource, dependencies,
    segment, source, readsegment, trial, results, resultsvariables

include("source.jl")
include("trial.jl")
include("segment.jl")
include("stack.jl")
include("summary.jl")

end


text:
#Merge methods together using metaprogramming?
function xarea!(ec::EChart, startval, endval; series::Int = 1, kwargs...)

	#set if not exists
    ec.series[series].markArea == nothing ? ec.series[series].markArea = MarkArea() : nothing

	#set endpoints
	push!(ec.series[series].markArea.data, [Dict("xAxis" => startval, "itemStyle" => ItemStyle()), Dict("xAxis" => endval)])

	for (k, v) in kwargs
		setfield!(ec.series[series].markArea.data[end][1]["itemStyle"], k, v)
	end

    return ec

end

function yarea!(ec::EChart, startval, endval; series::Int = 1, kwargs...)

	#set if not exists
    ec.series[series].markArea == nothing ? ec.series[series].markArea = MarkArea() : nothing

	#set endpoints
	push!(ec.series[series].markArea.data, [Dict("yAxis" => startval, "itemStyle" => ItemStyle()), Dict("yAxis" => endval)])

	for (k, v) in kwargs
		setfield!(ec.series[series].markArea.data[end][1]["itemStyle"], k, v)
	end

    return ec

end


text:
import ACE: OneParticleBasis, AbstractState
import ACE.OrthPolys: TransformedPolys

import NamedTupleTools
using NamedTupleTools: namedtuple

# ------------------------------------------------------------------------

@doc raw"""
`struct XScal1pBasis <: OneParticleBasis`

One-particle basis of the form $P_n(x_i)$ for a general scalar, invariant
input `x`. This type basically just translates the `TransformedPolys` into a valid
one-particle basis.
"""
mutable struct XScal1pBasis{ISYMS, FVAL, T, TT, TJ, NI, TRG} <: OneParticleBasis{T}
   P::TransformedPolys{T, TT, TJ}
   fval::FVAL
   rgs::TRG
   spec::Vector{NamedTuple{ISYMS, NTuple{NI, Int}}}
   coeffs::Matrix{T}
   label::String
end

xscal1pbasis(P::TransformedPolys, idxsyms, valsym::Symbol; kwargs...) =
      xscal1pbasis(P, idxsyms, GetVal{valsym}(); kwargs...)

xscal1pbasis(P::TransformedPolys, idxsyms, valsym::Symbol, validx::Int; kwargs...) =
      xscal1pbasis(P, idxsyms, GetVali{valsym, validx}(); kwargs...)

function xscal1pbasis(P::TransformedPolys, idxsyms, fval; label = "")
   ISYMS = tuple(keys(idxsyms)...)
   rgs = NamedTuple{ISYMS}( tuple([idxsyms[sym] for sym in ISYMS]...) )
   return XScal1pBasis(fval, ISYMS, rgs, P, label)
end

function XScal1pBasis(fval, ISYMS::NTuple{NI, Symbol},
                      rgs::TRG, P::TransformedPolys{T, TT, TJ},
                      label::String = "",
                      spec = NamedTuple{ISYMS, NTuple{NI, Int}}[],
                      coeffs = Matrix{T}(undef, (0,0))
                     ) where {NI, T, TT, TJ, TRG}
   return XScal1pBasis{ISYMS, typeof(fval), T, TT, TJ, NI, TRG}(P, fval, rgs, spec, coeffs, label)
end

_idxsyms(basis::XScal1pBasis{ISYMS}) where {ISYMS} = ISYMS

getval(X, basis::XScal1pBasis) = getval(X, basis.fval)

getval_d(X, basis::XScal1pBasis) = getval_d(X, basis.fval)

rand_radial(basis::XScal1pBasis) = rand_radial(basis.P)

# ---------------------- Implementation of Scal1pBasis

Base.length(basis::XScal1pBasis) = size(basis.coeffs, 1)

get_spec(basis::XScal1pBasis) = copy(basis.spec)

get_spec(basis::XScal1pBasis, i::Integer) = basis.spec[i]

function set_spec!(basis::XScal1pBasis, spec)
   basis.spec = identity.(spec)
   basis.coeffs = zeros(eltype(basis.coeffs), length(spec), length(basis.P))
   return basis
end

function isadmissible(b::NamedTuple{BSYMS}, basis::XScal1pBasis) where {BSYMS}
   ISYMS = _idxsyms(basis)
   @assert all(sym in BSYMS for sym in ISYMS)
   return all(b[sym] in basis.rgs[sym] for sym in ISYMS)
end

function sparsify!(basis::XScal1pBasis, spec)
   # spec is the part of the basis we keep
   inds = Vector{Int}(undef, length(spec))
   for (ib, b) in enumerate(spec)
      iold = findall(isequal(b), basis.spec)
      @assert length(iold) == 1
      inds[ib] = iold[1]
   end
   # keep the original order of the basis since we assume it is ordered
   # by some sensible notion of degree.
   p = sortperm(inds)
   basis.spec = spec[p]
   basis.coeffs = basis.coeffs[p, :]
   return basis
end

degree(b::NamedTuple, basis::XScal1pBasis) =
         getproperty(b, _idxsyms(basis)[1]) - 1

degree(b::NamedTuple, basis::XScal1pBasis, weight::Dict) =
      weight[_idxsyms(basis)[1]] * degree(b, basis)

"""
`fill_rand_coeffs!(basis::XScal1pBasis, f::Function)`

This fills the parameters of the XScal1pBasis with coefficients generated
by the function f. We say "random" because most typically, we expect
`f = rand` or `f = randn` or similar.
"""
function fill_rand_coeffs!(basis::XScal1pBasis, f::Function)
   for n = 1:length(basis.coeffs)
      basis.coeffs[n] = f()
   end
   return basis
end

"""
`fill_diag_coeffs!(basis::XScal1pBasis, sym = _idxsys(basis)[1])`

This sets XScal1p basis paramerers as follows:
```
P_{n,v} = J_n
```
where `sym = :n` and `v` is the rest of the symbols joined together.
I.e., this reduces the XScal basis to a standard Scal basis.
"""
function fill_diag_coeffs!(basis::XScal1pBasis, sym = _idxsyms(basis)[1])
   fill!(basis.coeffs, 0)
   for (ib, b) in enumerate(basis.spec)
      n = b[sym]
      basis.coeffs[ib, n] = 1
   end
   return basis
end

# =========================

==(P1::XScal1pBasis, P2::XScal1pBasis) = _allfieldsequal(P1, P2)

function write_dict(basis::XScal1pBasis{T}) where {T}
   ISYMS = _idxsyms(basis)
   return Dict(
      "__id__" => "ACE_XScal1pBasis",
          "P" => write_dict(basis.P),
          "fval" => write_dict(basis.fval),
          "syms" => [string.(ISYMS)...],
          "rgs" => Dict([ string(sym) => write_dict(collect(basis.rgs[sym]))
                          for sym in ISYMS]... ),
          "spec" => convert.(Dict, basis.spec),
          "coeffs" => write_dict(basis.coeffs),
          "label" => basis.label
      )
end

using NamedTupleTools: namedtuple

function read_dict(::Val{:ACE_XScal1pBasis}, D::Dict)
   P = read_dict(D["P"])
   fval = read_dict(D["fval"])
   ISYMS = tuple(Symbol.(D["syms"])...)
   rgs = namedtuple(Dict( [  sym => read_dict(D["rgs"][string(sym)])
                             for sym in ISYMS ]... )) |> NamedTuple{ISYMS}
   spec = NamedTuple{ISYMS}.(namedtuple.(D["spec"]))
   coeffs = read_dict(D["coeffs"])
   return XScal1pBasis(fval, ISYMS, rgs, P, D["label"], spec, coeffs)
end

valtype(basis::XScal1pBasis) =
      promote_type(valtype(basis.P), eltype(basis.coeffs))

valtype(basis::XScal1pBasis, cfg::AbstractConfiguration) =
      valtype(basis, zero(eltype(cfg)))

valtype(basis::XScal1pBasis, X::AbstractState) =
      promote_type(valtype(basis.P, getval(X, basis)), eltype(basis.coeffs))

gradtype(basis::XScal1pBasis, X::AbstractState) =
      dstate_type(valtype(basis, X), X)

argsyms(basis::XScal1pBasis) = ( _varsym(basis), )

symbols(basis::XScal1pBasis) = [ _idxsyms(basis)... ]

indexrange(basis::XScal1pBasis) = basis.rgs

# this is needed to generate the product 1p basis
function get_index(basis::XScal1pBasis, b::NamedTuple)
   ISYMS = _idxsyms(basis)
   b1 = select(b, ISYMS)
   idx = findall(isequal(b1), basis.spec)
   if length(idx) != 1
      @show b1
      @show basis.spec
   end
   if length(idx) == 0
      @error("XScal1pBasis : didn't find b in the spec")
      @show b
      @show basis.spec
      error("didn't find b in the spec")
   elseif length(idx) > 1
      error("b appears in spec more than once")
   end
   return idx[1]
end

# ---------------------------  Evaluation code
#

evaluate!(B, basis::XScal1pBasis, X::AbstractState) =
      evaluate!(B, basis, getval(X, basis))

function evaluate!(B, basis::XScal1pBasis, x::Number)
   P = evaluate(basis.P, x)
   mul!(B, basis.coeffs, P)
   return B
end

function evaluate_d!(dB, basis::XScal1pBasis, X::AbstractState)
   TDX = eltype(dB)
   x = getval(X, basis)
   dP = acquire_dB!(basis.P, x)
   evaluate_d!(dP, basis.P, x)
   dx = getval_d(X, basis)
   dB[:] .= TDX.( Ref(dx) .* (basis.coeffs * dP) )
   release_dB!(basis.P, dP)
   return dB
end

function evaluate_ed!(B, dB, basis::XScal1pBasis, X::AbstractState)
   TDX = eltype(dB)
   x = getval(X, basis)
   P = acquire_B!(basis.P, x)
   dP = acquire_dB!(basis.P, x)
   evaluate!(P, basis.P, x)
   evaluate_d!(dP, basis.P, x)
   mul!(B, basis.coeffs, P)
   dx = getval_d(X, basis)
   dB[:] .= TDX.( Ref(dx) .* (basis.coeffs * dP) )
   release_B!(basis.P, P)
   release_dB!(basis.P, dP)
   return B, dB
end

# *** TODO
# # this one we probably only need for training so can relax the efficiency a bit
# function evaluate_dd(basis::Scal1pBasis, X::AbstractState)
#    ddP = ForwardDiff.derivative(x -> evaluate_d(basis, _val(X, basis)))
#    TDX = gradtype(basis, X)
#    return _scal1pbasis_grad.(Ref(TDX), Ref(basis), ddP_n)
# end

#=   *** TODO
# -------------- AD codes

import ChainRules: rrule, ZeroTangent, NoTangent

function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState,
                         w::AbstractVector{<: Number})
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   a = _rrule_evaluate(basis.P, x, real.(w))
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

rrule(::typeof(evaluate), basis::Scal1pBasis, X::AbstractState) =
                  evaluate(basis, X),
                  w -> (NoTangent(), NoTangent(), _rrule_evaluate(basis, X, w))

function _rrule_evaluate_d(basis::Scal1pBasis, X::AbstractState,
                           w::AbstractVector)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   w1 = [ _val(w, basis) for w in w ]
   a = _rrule_evaluate_d(basis.P, x, w1)
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

function rrule(::typeof(evaluate_d), basis::Scal1pBasis, X::AbstractState)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   dB_ = evaluate_d(basis.P, x)
   TDX = dstate_type(valtype(basis, X), X)
   dB = [ TDX( NamedTuple{(_varsym(basis),)}( (dx,) ) )  for dx in dB_ ]
   return dB,
          w -> (NoTangent(), NoTangent(), _rrule_evaluate_d(basis, X, w))
end

=#


text:
#=
Provides a type called LQ for solving linear quadratic control
problems.

@author : Spencer Lyon <spencer.lyon@nyu.edu>
@author : Zac Cranko <zaccranko@gmail.com>

@date : 2014-07-05

References
----------

http://quant-econ.net/jl/lqcontrol.html

=#

"""
Linear quadratic optimal control of either infinite or finite horizon

The infinite horizon problem can be written

    min E sum_{t=0}^{infty} beta^t r(x_t, u_t)

with

    r(x_t, u_t) := x_t' R x_t + u_t' Q u_t + 2 u_t' N x_t

The finite horizon form is

    min E sum_{t=0}^{T-1} beta^t r(x_t, u_t) + beta^T x_T' R_f x_T

Both are minimized subject to the law of motion

    x_{t+1} = A x_t + B u_t + C w_{t+1}

Here x is n x 1, u is k x 1, w is j x 1 and the matrices are conformable for
these dimensions.  The sequence {w_t} is assumed to be white noise, with zero
mean and E w_t w_t' = I, the j x j identity.

For this model, the time t value (i.e., cost-to-go) function V_t takes the form

    x' P_T x + d_T

and the optimal policy is of the form u_T = -F_T x_T.  In the infinite horizon
case, V, P, d and F are all stationary.

##### Fields

- `Q::ScalarOrArray` : k x k payoff coefficient for control variable u. Must be
symmetric and nonnegative definite
- `R::ScalarOrArray` : n x n payoff coefficient matrix for state variable x.
Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : n x n coefficient on state in state transition
- `B::ScalarOrArray` : n x k coefficient on control in state transition
- `C::ScalarOrArray` : n x j coefficient on random shock in state transition
- `N::ScalarOrArray` : k x n cross product in payoff equation
- `bet::Real` : Discount factor in [0, 1]
- `capT::Union{Int, Void}` : Terminal period in finite horizon problem
- `rf::ScalarOrArray` : n x n terminal payoff in finite horizon problem. Must be
symmetric and nonnegative definite
- `P::ScalarOrArray` : n x n matrix in value function representation
V(x) = x'Px + d
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period

"""
type LQ
    Q::ScalarOrArray
    R::ScalarOrArray
    A::ScalarOrArray
    B::ScalarOrArray
    C::ScalarOrArray
    N::ScalarOrArray
    bet::Real
    capT::Union{Int, Void} # terminal period
    rf::ScalarOrArray
    P::ScalarOrArray
    d::Real
    F::ScalarOrArray # policy rule
end

"""
Main constructor for LQ type

Specifies default argumets for all fields not part of the payoff function or
transition equation.

##### Arguments

- `Q::ScalarOrArray` : k x k payoff coefficient for control variable u. Must be
symmetric and nonnegative definite
- `R::ScalarOrArray` : n x n payoff coefficient matrix for state variable x.
Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : n x n coefficient on state in state transition
- `B::ScalarOrArray` : n x k coefficient on control in state transition
- `;C::ScalarOrArray{zeros(size(R}(1)))` : n x j coefficient on random shock in
state transition
- `;N::ScalarOrArray{zeros(size(B,1)}(size(A, 2)))` : k x n cross product in
payoff equation
- `;bet::Real(1.0)` : Discount factor in [0, 1]
- `capT::Union{Int, Void}(Void)` : Terminal period in finite horizon
problem
- `rf::ScalarOrArray{fill(NaN}(size(R)...))` : n x n terminal payoff in finite
horizon problem. Must be symmetric and nonnegative definite.

"""
function LQ(Q::ScalarOrArray,
            R::ScalarOrArray,
            A::ScalarOrArray,
            B::ScalarOrArray,
            C::ScalarOrArray=zeros(size(R, 1)),
            N::ScalarOrArray=zero(B'A);
            bet::ScalarOrArray=1.0,
            capT::Union{Int,Void}=nothing,
            rf::ScalarOrArray=fill(NaN, size(R)...))

    k = size(Q, 1)
    n = size(R, 1)
    F = k==n==1 ? zero(Float64) : zeros(Float64, k, n)
    P = copy(rf)
    d = 0.0

    LQ(Q, R, A, B, C, N, bet, capT, rf, P, d, F)
end

"""
Update `P` and `d` from the value function representation in finite horizon case

##### Arguments

- `lq::LQ` : instance of `LQ` type

##### Returns

- `P::ScalarOrArray` : n x n matrix in value function representation
V(x) = x'Px + d
- `d::Real` : Constant in value function representation

##### Notes

This function updates the `P` and `d` fields on the `lq` instance in addition to
returning them

"""
function update_values!(lq::LQ)
    # Simplify notation
    Q, R, A, B, N, C, P, d = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C, lq.P, lq.d

    # Some useful matrices
    s1 = Q + lq.bet * (B'P*B)
    s2 = lq.bet * (B'P*A) + N
    s3 = lq.bet * (A'P*A)

    # Compute F as (Q + B'PB)^{-1} (beta B'PA)
    lq.F = s1 \ s2

    # Shift P back in time one step
    new_P = R - s2'lq.F + s3

    # Recalling that trace(AB) = trace(BA)
    new_d = lq.bet * (d + trace(P * C * C'))

    # Set new state
    lq.P, lq.d = new_P, new_d
end

"""
Computes value and policy functions in infinite horizon model

##### Arguments

- `lq::LQ` : instance of `LQ` type

##### Returns

- `P::ScalarOrArray` : n x n matrix in value function representation
V(x) = x'Px + d
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period

##### Notes

This function updates the `P`, `d`, and `F` fields on the `lq` instance in
addition to returning them

"""
function stationary_values!(lq::LQ)
    # simplify notation
    Q, R, A, B, N, C = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C

    # solve Riccati equation, obtain P
    A0, B0 = sqrt(lq.bet) * A, sqrt(lq.bet) * B
    P = solve_discrete_riccati(A0, B0, R, Q, N)

    # Compute F
    s1 = Q + lq.bet * (B' * P * B)
    s2 = lq.bet * (B' * P * A) + N
    F = s1 \ s2

    # Compute d
    d = lq.bet * trace(P * C * C') / (1 - lq.bet)

    # Bind states
    lq.P, lq.F, lq.d = P, F, d
end

"""
Non-mutating routine for solving for `P`, `d`, and `F` in infinite horizon model

See docstring for stationary_values! for more explanation
"""
function stationary_values(lq::LQ)
    _lq = LQ(copy(lq.Q),
             copy(lq.R),
             copy(lq.A),
             copy(lq.B),
             copy(lq.C),
             copy(lq.N),
             bet=copy(lq.bet),
             capT=lq.capT,
             rf=copy(lq.rf))

    stationary_values!(_lq)
    return _lq.P, _lq.F, _lq.d
end

"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence{T}(lq::LQ, x0::T, policies)
    capT = length(policies)

    x_path = Array{T}(capT+1)
    u_path = Array{T}(capT)

    x_path[1] = x0
    u_path[1] = -(first(policies)*x0)
    w_path = lq.C * randn(capT+1)

    for t = 2:capT
        f = policies[t]
        x_path[t] = lq.A*x_path[t-1] + lq.B*u_path[t-1] + w_path[t]
        u_path[t] = -(f*x_path[t])
    end
    x_path[end] = lq.A*x_path[capT] + lq.B*u_path[capT] + w_path[end]

    x_path, u_path, w_path
end

"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence{T}(lq::LQ, x0::Vector{T}, policies)
    # Ensure correct dimensionality
    n, j, k = size(lq.C, 1), size(lq.C, 2), size(lq.B, 2)
    capT = length(policies)

    A, B, C = lq.A, reshape(lq.B, n, k), reshape(lq.C, n, j)

    x_path = Array{T}(n, capT+1)
    u_path = Array{T}(k, capT)
    w_path = C*randn(j, capT+1)

    x_path[:, 1] = x0
    u_path[:, 1] = -(first(policies)*x0)

    for t = 2:capT
        f = policies[t]
        x_path[:, t] = A*x_path[: ,t-1] + B*u_path[:, t-1] + w_path[:, t]
        u_path[:, t] = -(f*x_path[:, t])
    end
    x_path[:, end] = A*x_path[:, capT] + B*u_path[:, capT] + w_path[:, end]

    x_path, u_path, w_path
end

"""
Compute and return the optimal state and control sequence, assuming innovation N(0,1)

##### Arguments

- `lq::LQ` : instance of `LQ` type
- `x0::ScalarOrArray`: initial state
- `ts_length::Integer(100)` : maximum number of periods for which to return
process. If `lq` instance is finite horizon type, the sequenes are returned
only for `min(ts_length, lq.capT)`

##### Returns

- `x_path::Matrix{Float64}` : An n x T+1 matrix, where the t-th column
represents `x_t`
- `u_path::Matrix{Float64}` : A k x T matrix, where the t-th column represents
`u_t`
- `w_path::Matrix{Float64}` : A n x T+1 matrix, where the t-th column represents
`lq.C*N(0,1)`

"""
function compute_sequence(lq::LQ, x0::ScalarOrArray, ts_length::Integer=100)

    # Compute and record the sequence of policies
    if isa(lq.capT, Void)
        stationary_values!(lq)
        policies = fill(lq.F, ts_length)
    else
        capT = min(ts_length, lq.capT)
        policies = Array{typeof(lq.F)}(capT)
        for t = capT:-1:1
            update_values!(lq)
            policies[t] = lq.F
        end
    end

    _compute_sequence(lq, x0, policies)
end


text:
push!(LOAD_PATH, "../")

using Documenter

makedocs(sitename="Julia Notes")


text:
# ---
# title: 719. Find K-th Smallest Pair Distance
# id: problem719
# author: Qling
# date: 2021-04-07
# difficulty: Hard
# categories: Array, Binary Search, Heap
# link: <https://leetcode.com/problems/find-k-th-smallest-pair-distance/description/>
# hidden: true
# ---
#
# Given an integer array, return the k-th smallest **distance** among all the
# pairs. The distance of a pair (A, B) is defined as the absolute difference
# between A and B.
#
# **Example 1:**
#
#
#
#     Input:
#     nums = [1,3,1]
#     k = 1
#     **Output: 0**
#     Explanation:
#     Here are all the pairs:
#     (1,3) -> 2
#     (1,1) -> 0
#     (3,1) -> 2
#     Then the 1st smallest distance pair is (1,1), and its distance is 0.
#
#
# **Note:**
#
#   1. `2 <= len(nums) <= 10000`.
#   2. `0 <= nums[i] < 1000000`.
#   3. `1 <= k <= len(nums) * (len(nums) - 1) / 2`.
#
#
## @lc code=start
using LeetCode

function smallest_distance_pair(nums::Vector{Int}, k::Int)::Int
    sort!(nums)
    n = length(nums)
    function count_not_greater(diff)
        i, ans = 1, 0
        for j in 2: n
            while nums[j] - nums[i] > diff
                i += 1
            ans += j - i
            end
        end

        return ans
    end

    left, right = 0, nums[end] - nums[1]
    while left <= right
        mid = (left + right) >> 1
        if count_not_greater(mid) >= k
            right = mid - 1
        else
            left = mid + 1
        end
    end

    return left
end
## @lc code=end


text:
"""
A store that wraps any other AbstractStore but has access to the consolidated metadata
stored in the .zmetadata key. Whenever data attributes or metadata are accessed, the
data will be read from the dictionary instead.
"""
struct ConsolidatedStore{P} <: AbstractStore
  parent::P
  path::String
  cons::Dict{String,Any}
end
function ConsolidatedStore(s::AbstractStore, p)
  d = s[p, ".zmetadata"]
  if d === nothing
    throw(ArgumentError("Could not find consolidated metadata for store $s"))
  end
  ConsolidatedStore(s,p,JSON.parse(String(Zarr.maybecopy(d)))["metadata"])
end

function Base.show(io::IO,d::ConsolidatedStore)
    b = IOBuffer()
    show(b,d.parent)
    print(io, "Consolidated ", String(take!(b)))
end

storagesize(d::ConsolidatedStore,p) = storagesize(d.parent,p)
function Base.getindex(d::ConsolidatedStore,i::String)
    d.parent[i]
end
getmetadata(d::ConsolidatedStore,p) = Metadata(d.cons[_unconcpath(d,p,".zarray")])
getattrs(d::ConsolidatedStore, p) = get(d.cons,_unconcpath(d,p,".zattrs"), Dict{String,Any}())
function _unconcpath(d,p)
  startswith(p,d.path) || error("Requested key is not in consolidated path")
  lstrip(replace(p,d.path=>"", count=1),'/')
end
_unconcpath(d,p,s) = _concatpath(_unconcpath(d,p),s)
is_zarray(d::ConsolidatedStore,p) = haskey(d.cons,_unconcpath(d,p,".zarray"))
is_zgroup(d::ConsolidatedStore,p) = haskey(d.cons,_unconcpath(d,p,".zgroup"))
check_consolidated_write(i::String) = split(i,'/')[end] in (".zattrs",".zarray",".zgroup") &&
    throw(ArgumentError("Can not modify consolidated metadata, please re-open the dataset with `consolidated=false`"))

_pdict(d::ConsolidatedStore,p) = filter(((k,v),)->startswith(k,p),d.cons)
function subdirs(d::ConsolidatedStore,p)
  p2 = _unconcpath(d,p)
  d2 = _pdict(d,p2)
  _searchsubdict(d2,p,(sp,lp)->length(sp) > lp+1)
end

function subkeys(d::ConsolidatedStore,p)
  subkeys(d.parent,p)
end

function Base.setindex!(d::ConsolidatedStore,v,i::String)
  #Here we check that we don't overwrite consolidated information
  check_consolidated_write(i)
  d.parent[i] = v
end
function Base.delete!(d::ConsolidatedStore,i::String)
  check_consolidated_write(i)
  delete!(d.parent,i)
end

function consolidate_metadata(s::AbstractStore,d,prefix)
  for k in (".zattrs",".zarray",".zgroup")
    v = s[prefix,k]
    if v !== nothing
      d[_concatpath(prefix,k)] = JSON.parse(String(copy(v)))
    end
  end
  foreach(subdirs(s,prefix)) do subname
    consolidate_metadata(s,d,string(prefix,subname,"/"))
  end
  d
end
function consolidate_metadata(s::AbstractStore,p)
  d = consolidate_metadata(s,Dict{String,Any}(),p)
  buf = IOBuffer()
  JSON.print(buf,Dict("metadata"=>d),4)
  s[p,".zmetadata"] = take!(buf)
  ConsolidatedStore(s,p,d)
end

consolidate_metadata(s) = consolidate_metadata(zopen(s,"w"))


text:
"""

Struct to represent .h5 spectra produced by FIDASIM.
"""
struct FIDASIMSpectra{T<:Real}
    nchan::Int
    nlambda::Int
    lambda::Vector{T}
    brems::Union{Zeros{T,2},Matrix{T}}
    full::Union{Zeros{T,2},Matrix{T}}
    half::Union{Zeros{T,2},Matrix{T}}
    third::Union{Zeros{T,2},Matrix{T}}
    dcx::Union{Zeros{T,2},Matrix{T}}
    halo::Union{Zeros{T,2},Matrix{T}}
    cold::Union{Zeros{T,2},Matrix{T}}
    fida::Union{Zeros{T,2},Zeros{T,3},Matrix{T},Array{T,3}}
    pfida::Union{Zeros{T,2},Zeros{T,3},Matrix{T},Array{T,3}}
end

"""
    FIDASIMSpectra("path/to/FIDASIM/spectra/file.h5")

Create instance of FIDASIMSpectra struct from .h5 file. The data can then be
accessed via Julia struct syntax.

# Example
```julia-repl
julia> myFIDASIMSpectra = FIDASIMSpectra("path/to/FIDASIM/spectra/file.h5")
FIDASIMSpectra
    nchan = 81
    nlambda = 2000

julia> lambda_array = myFIDASIMSpectra.lambda
2000-element Array{Float64,1}:
647.005
647.015
647.025
647.035
⋮

julia> myfida_data = myFIDASIMSpectra.fida
2000×81 Array{Float64,2}:
0.0 0.0 ... 0.0 0.0
0.0 0.0 ... 0.0 0.0
⋮   ⋮    ⋮    ⋮   ⋮
0.0 0.0 ... 0.0 0.0
0.0 0.0 ... 0.0 0.0
```
"""
function FIDASIMSpectra(fname::String)
    isfile(fname) || error("File does not exist")

    f = h5open(fname)
    nchan = Int(read(f["nchan"]))
    lambda = read(f["lambda"])
    nlam = length(lambda)
    vars = keys(f)

    brems = "brems" in vars ? read(f["brems"]) : Zeros(nlam,nchan)
    full = "full" in vars ? read(f["full"]) : Zeros(nlam,nchan)
    half = "half" in vars ? read(f["half"]) : Zeros(nlam,nchan)
    third = "third" in vars ? read(f["third"]) : Zeros(nlam,nchan)
    dcx = "dcx" in vars ? read(f["dcx"]) : Zeros(nlam,nchan)
    halo = "halo" in vars ? read(f["halo"]) : Zeros(nlam,nchan)
    cold = "cold" in vars ? read(f["cold"]) : Zeros(nlam,nchan)
    if ("pfida" in vars) && !("fida" in vars)
        pfida = read(f["pfida"])
        fida = Zeros(size(pfida))
    end
    if !("pfida" in vars) && ("fida" in vars)
        fida = read(f["fida"])
        pfida = Zeros(size(fida))
    end
    if ("pfida" in vars) && ("fida" in vars)
        fida = read(f["fida"])
        pfida = read(f["pfida"])
    end
    if !("pfida" in vars) && !("fida" in vars)
        fida = Zeros(nlam,nchan)
        pfida = Zeros(size(fida))
    end

    return FIDASIMSpectra(Int(nchan),Int(nlam),lambda,brems,full,half,third,dcx,halo,cold,fida,pfida)
end

function Base.show(io::IO, s::FIDASIMSpectra)
    println(io, "FIDASIMSpectra")
    println(io, "  nchan   = $(s.nchan)")
    println(io, "  nlambda = $(s.nlambda)")
end

abstract type AbstractDistribution end

struct FIDASIMGuidingCenterFunction <: AbstractDistribution
    nr::Int
    nz::Int
    nenergy::Int
    npitch::Int
    r::Vector{Float64}
    z::Vector{Float64}
    energy::Vector{Float64}
    pitch::Vector{Float64}
    denf::Matrix{Float64}
    f::Array{Float64,4}
    nfast::Float64
end

function FIDASIMGuidingCenterFunction(f::HDF5.File)

    nr = read(f["nr"])
    nz = read(f["nz"])
    nenergy = read(f["nenergy"])
    npitch = read(f["npitch"])

    energy = read(f["energy"])
    pitch = read(f["pitch"])
    r = read(f["r"])
    z = read(f["z"])
    denf = read(f["denf"])
    dist = read(f["f"])

    fr = dist .* reshape(r,(1,1,length(r),1))
    dE = abs(energy[2] - energy[1])
    dp = abs(pitch[2] - pitch[1])
    dr = abs(r[2] - r[1])
    dz = abs(z[2] - z[1])
    nfast = sum(fr)*(2*pi*dE*dp*dr*dz)

    return FIDASIMGuidingCenterFunction(nr,nz,nenergy,npitch,r,z,energy,pitch,denf,dist,nfast)
end

function Base.show(io::IO, s::FIDASIMGuidingCenterFunction)
    println(io, "FIDASIMGuidingCenterFunction")
end

struct FIDASIMGuidingCenterParticles <: AbstractDistribution
    npart::Int
    nclass::Int
    class::Vector{Int}
    weight::Vector{Float64}
    r::Vector{Float64}
    z::Vector{Float64}
    energy::Vector{Float64}
    pitch::Vector{Float64}
    nfast::Float64
end

function FIDASIMGuidingCenterParticles(f::HDF5.File)

    npart = read(f["nparticle"])
    nclass = read(f["nclass"])
    class = read(f["class"])
    weight = read(f["weight"])

    energy = read(f["energy"])
    pitch = read(f["pitch"])
    r = read(f["r"])
    z = read(f["z"])

    return FIDASIMGuidingCenterParticles(npart,nclass,class,weight,r,z,energy,pitch,sum(weight))
end

function Base.show(io::IO, s::FIDASIMGuidingCenterParticles)
    println(io, "FIDASIMGuidingCenterParticles")
    println(io, "  nfast = $(s.nfast)")
    println(io, "  nparticles = $(s.npart)")
end

struct FIDASIMFullOrbitParticles <: AbstractDistribution
    npart::Int
    nclass::Int
    class::Vector{Int}
    weight::Vector{Float64}
    r::Vector{Float64}
    z::Vector{Float64}
    vr::Vector{Float64}
    vt::Vector{Float64}
    vz::Vector{Float64}
    nfast::Float64
end

function FIDASIMFullOrbitParticles(f::HDF5.File)

    npart = read(f["nparticle"])
    nclass = read(f["nclass"])
    class = read(f["class"])
    weight = read(f["weight"])

    r = read(f["r"])
    z = read(f["z"])
    vr = read(f["vr"])
    vt = read(f["vt"])
    vz = read(f["vz"])

    return FIDASIMFullOrbitParticles(npart,nclass,class,weight,r,z,vr,vt,vz,sum(weight))
end

function Base.show(io::IO, s::FIDASIMFullOrbitParticles)
    println(io, "FIDASIMFullOrbitParticles")
    println(io, "  nfast = $(s.nfast)")
    println(io, "  nparticles = $(s.npart)")
end

"""
    read_fidasim_distribution("path/to/FIDASIM/distribution/.h5/file")

Read FIDASIM distribution .h5 file. Return FIDASIMGuidingCenterFunction, FIDASIMGuidingCenterParticles
or FIDASIMFullOrbitParticles depending on dtype variable. The data can then be accessed via Julia struct
syntax.

# Example
```julia-repl
julia> myFIDASIM_distribution = read_fidasim_distribution("path/to/FIDASIM/distribution/.h5/file")
FIDASIMGuidingCenterFunction

julia> r_array = myFIDASIM_distribution.r
70-element Array{Float64,1}:
 100.0
 102.0
 104.0
 106.0
 108.0
 110.0
 112.0
 114.0
 ⋮

julia>
```
"""
function read_fidasim_distribution(filename::String)
    isfile(filename) || error("File does not exist")

    f = h5open(filename)
    dtype = read(f["type"])

    if dtype == 1
        d = FIDASIMGuidingCenterFunction(f)
    elseif dtype == 2
        d = FIDASIMGuidingCenterParticles(f)
    elseif dtype == 3
        d = FIDASIMFullOrbitParticles(f)
    else
        error("Unknown distribution type")
    end
    close(f)

    return d
end

"""
    write_fidasim_distribution(M,orbit_array)
    write_fidasim_distribution(M,orbit_array,filename="myOrbit_distribution.h35",chunksize=1000)

Write array of orbits (as found in Julia package GuidingCenterOrbits.jl/src/orbit.jl/Orbit) calculated in Julia to an
.h5 distribution file which can be used as input to FIDASIM. If the array of orbits is longer than 1000, it is highly
recommended to set the chunksize to approx 1000. This will then produce several .h5 distribution files. This is because
each distribution file should have a reasonable size to be read by FIDASIM. So if write_fidasim_distribution() creates
15 .h5 files for example, you will need to do 15 FIDASIM runs and then combine the 15 output spectras into one.
This can be done in numerous ways, taking care due to the likely large file amount of memory required.
"""
function write_fidasim_distribution(M::AbstractEquilibrium, orbits::Array; filename="orbits.h5",time=0.0,ntot=1e19,chunksize=0)

    if chunksize == 0
        orbs = (orbits,)
        nchunks = 1
    else
        orbs = partition(orbits,chunksize)
        nchunks = length(orbs)
    end

    #for io = 1:length(orbs)
    @sync @distributed for io = 1:length(orbs)
        oo = nth(orbs,io)
        norbs = length(oo)
        npart = sum(length(o.path.r) for o in oo)
        if chunksize != 0
            fname= splitext(filename)[1]*@sprintf("_%04d_%04d.h5",io,nchunks)
        else
            fname = filename
        end
        h5open(fname,"w") do file
            file["data_source"] = "Generated from GEQDSK"
            file["nclass"] = norbs
            file["nparticle"] = npart
            file["time"] = time
            file["type"] = 2
            file["tau_p","shuffle",(),"chunk", (norbs),"compress", 4] = [o.tau_p for o in oo]
            file["tau_t","shuffle",(),"chunk", (norbs),"compress", 4] = [o.tau_t for o in oo]

            if isa(oo[1].coordinate, EPRCoordinate)
                file["energy_c","shuffle",(),"chunk",(norbs),"compress",4] = [o.coordinate.energy for o in oo]
                file["pitch_c","shuffle",(),"chunk",(norbs),"compress",4] = [o.coordinate.pitch for o in oo]
                file["r_c","shuffle",(),"chunk",(norbs),"compress",4] = [o.coordinate.r for o in oo]
                file["z_c","shuffle",(),"chunk",(norbs),"compress",4] = [o.coordinate.z for o in oo]
            end

            file["energy","shuffle",(),"chunk",(npart),"compress",4] = vcat((o.path.energy for o in oo)...)
            file["pitch","shuffle",(),"chunk",(npart),"compress",4] = vcat((B0Ip_sign(M)*o.path.pitch for o in oo)...)
            file["r","shuffle",(),"chunk",(npart),"compress",4] = vcat((100*o.path.r for o in oo)...)
            file["z","shuffle",(),"chunk",(npart),"compress",4] = vcat((100*o.path.z for o in oo)...)
            file["class","shuffle",(),"chunk",(npart),"compress",4] = vcat((fill(i,length(o.path)) for (i,o) in enumerate(oo))...)
            file["weight","shuffle",(),"chunk",(npart),"compress",4] = vcat((o.path.dt.*(ntot/sum(o.path.dt)) for o in oo)...)

            ## create datasets
            #energy = d_create(file, "energy", datatype(Float64), (npart,), "shuffle", (), "chunk", (npart,), "compress", 4)
            #pitch = d_create(file, "pitch", datatype(Float64), (npart,), "shuffle", (), "chunk", (npart,), "compress", 4)
            #r = d_create(file, "r", datatype(Float64), (npart,), "shuffle", (), "chunk", (npart,), "compress", 4)
            #z = d_create(file, "z", datatype(Float64), (npart,), "shuffle", (), "chunk", (npart,), "compress", 4)
            #class = d_create(file, "class", datatype(Int32), (npart,), "shuffle", (), "chunk", (npart,), "compress", 4)
            #weight = d_create(file, "weight", datatype(Float64), (npart,), "shuffle", (), "chunk", (npart,), "compress", 4)

            ## incrementally write to dataset
            #istart = 1
            #for (i,o) in enumerate(oo)
            #    npath = length(o.path.energy)
            #    iend = istart + npath - 1
            #    energy[istart:iend] = o.path.energy
            #    pitch[istart:iend] = M.sigma*o.path.pitch
            #    r[istart:iend] = o.path.r*100
            #    z[istart:iend] = o.path.z*100
            #    class[istart:iend] = fill(i,npath)
            #    weight[istart:iend] = o.path.dt.*(ntot/sum(o.path.dt))
            #    istart = istart + npath
            #end
        end
    end
    nothing
end

struct FIDASIMBeamGeometry
    name::String
    src::Vector{Float64}
    axis::Vector{Float64}
    shape::Int
    widy::Float64
    widz::Float64
    divy::Vector{Float64}
    divz::Vector{Float64}
    focy::Float64
    focz::Float64
    naperture::Int
    ashape::Vector{Int}
    awidy::Vector{Float64}
    awidz::Vector{Float64}
    aoffy::Vector{Float64}
    aoffz::Vector{Float64}
    adist::Vector{Float64}
end

function FIDASIMBeamGeometry(filename::String)
    isfile(filename) || error("File does not exist")

    f = h5open(filename)
    if !("nbi" in keys(f))
        close(f)
        error("Beam geometry not in file")
    end

    name = read(f["nbi"]["name"])
    src = read(f["nbi"]["src"])
    axis = read(f["nbi"]["axis"])
    shape = Int(read(f["nbi"]["shape"]))
    widy = read(f["nbi"]["widy"])
    widz = read(f["nbi"]["widz"])
    divy = read(f["nbi"]["divy"])
    divz = read(f["nbi"]["divz"])
    focy = read(f["nbi"]["focy"])
    focz = read(f["nbi"]["focz"])
    naperture = Int(read(f["nbi"]["naperture"]))
    if naperture > 0
        ashape = Int.(read(f["nbi"]["ashape"]))
        awidy = read(f["nbi"]["awidy"])
        awidz = read(f["nbi"]["awidz"])
        aoffy = read(f["nbi"]["aoffy"])
        aoffz = read(f["nbi"]["aoffz"])
        adist = read(f["nbi"]["adist"])
    else
        ashape = Int[]
        awidy = Float64[]
        awidz = Float64[]
        aoffy = Float64[]
        aoffz = Float64[]
        adist = Float64[]
    end
    close(f)

    return FIDASIMBeamGeometry(name,src,axis,shape,widy,widz,divy,divz,focy,focz,
                               naperture, ashape, awidy, awidz, aoffy, aoffz, adist)
end

function Base.show(io::IO, s::FIDASIMBeamGeometry)
    println(io, "FIDASIMBeamGeometry")
    println(io, "  name: $(s.name)")
end

struct FIDASIMSpectraGeometry
    nchan::Int
    system::String
    id::Vector{String}
    radius::Vector{Float64}
    lens::Matrix{Float64}
    axis::Matrix{Float64}
    spot_size::Vector{Float64}
    sigma_pi::Vector{Float64}
end

function FIDASIMSpectraGeometry(filename::String)
    isfile(filename) || error("File does not exist")

    f = h5open(filename)
    if !("spec" in keys(f))
        close(f)
        error("Spectra geometry not in file")
    end

    nchan = Int(read(f["spec"]["nchan"]))
    system = read(f["spec"]["system"])
    id = read(f["spec"]["id"])
    radius = read(f["spec"]["radius"])
    lens = read(f["spec"]["lens"])
    axis = read(f["spec"]["axis"])
    spot_size = read(f["spec"]["spot_size"])
    sigma_pi = read(f["spec"]["sigma_pi"])

    close(f)

    return FIDASIMSpectraGeometry(nchan,system,id,radius,lens,axis,spot_size,sigma_pi)
end

function Base.show(io::IO, s::FIDASIMSpectraGeometry)
    println(io, "FIDASIMSpectraGeometry")
    println(io, "  nchan: $(s.nchan)")
    println(io, "  system: "*s.system)
end

struct FIDASIMNPAGeometry
    nchan::Int
    system::String
    id::Vector{String}
    radius::Vector{Float64}
    a_shape::Vector{Int}
    d_shape::Vector{Int}
    a_cent::Matrix{Float64}
    a_redge::Matrix{Float64}
    a_tedge::Matrix{Float64}
    d_cent::Matrix{Float64}
    d_redge::Matrix{Float64}
    d_tedge::Matrix{Float64}
end

function FIDASIMNPAGeometry(filename::String)
    isfile(filename) || error("File does not exist")

    f = h5open(filename)
    if !("npa" in keys(f))
        close(f)
        error("NPA geometry not in file")
    end

    nchan = Int(read(f["npa"]["nchan"]))
    system = read(f["npa"]["system"])
    id = read(f["npa"]["id"])
    radius = read(f["npa"]["radius"])
    a_shape = Int.(read(f["npa"]["a_shape"]))
    d_shape = Int.(read(f["npa"]["d_shape"]))
    a_cent = read(f["npa"]["a_cent"])
    a_redge = read(f["npa"]["a_redge"])
    a_tedge = read(f["npa"]["a_tedge"])
    d_cent = read(f["npa"]["d_cent"])
    d_redge = read(f["npa"]["d_redge"])
    d_tedge = read(f["npa"]["d_tedge"])

    close(f)

    return FIDASIMNPAGeometry(nchan, system, id, radius,
                              a_shape, d_shape,
                              a_cent, a_redge, a_tedge,
                              d_cent, d_redge, d_tedge)
end

function Base.show(io::IO, s::FIDASIMNPAGeometry)
    println(io, "FIDASIMNPAGeometry")
    println(io, "  nchan: $(s.nchan)")
    print(io, "  system: "*s.system)
end

function write_fidasim_geometry(nbi::FIDASIMBeamGeometry;
                                filename="geometry.h5",
                                data_source="Generated by OrbitTomography.jl",
                                spec::Union{FIDASIMSpectraGeometry,Nothing}=nothing,
                                npa::Union{FIDASIMNPAGeometry,Nothing}=nothing)

    h5open(filename,"w") do file
        file["/nbi/data_source"] = data_source
        file["/nbi/name"] = nbi.name
        file["/nbi/src"] = nbi.src
        file["/nbi/axis"] = nbi.axis
        file["/nbi/focy"] = nbi.focy
        file["/nbi/focz"] = nbi.focz
        file["/nbi/divy"] = nbi.divy
        file["/nbi/divz"] = nbi.divz
        file["/nbi/widy"] = nbi.widy
        file["/nbi/widz"] = nbi.widz
        file["/nbi/shape"] = nbi.shape
        file["/nbi/naperture"] = nbi.naperture
        if nbi.naperture > 0
            file["/nbi/ashape"] = nbi.ashape
            file["nbi/awidy"] = nbi.awidy
            file["nbi/awidz"] = nbi.awidz
            file["nbi/aoffy"] = nbi.aoffy
            file["nbi/aoffz"] = nbi.aoffz
            file["nbi/adist"] = nbi.adist
        end

        if spec != nothing
            file["spec/data_source"] = data_source
            file["spec/nchan"] = spec.nchan
            file["spec/system"] = spec.system
            file["spec/id"] = spec.id
            file["spec/radius"] = spec.radius
            file["spec/lens"] = spec.lens
            file["spec/axis"] = spec.axis
            file["spec/spot_size"] = spec.spot_size
            file["spec/sigma_pi"] = spec.sigma_pi
        end

        if npa != nothing
            file["npa/data_source"] = data_source
            file["npa/nchan"] = npa.nchan
            file["npa/system"] = npa.system
            file["npa/id"] = npa.id
            file["npa/radius"] = npa.radius
            file["npa/a_shape"] = npa.a_shape
            file["npa/d_shape"] = npa.d_shape
            file["npa/a_cent"] = npa.a_cent
            file["npa/a_redge"] = npa.a_redge
            file["npa/a_tedge"] = npa.a_tedge
            file["npa/d_cent"] = npa.d_cent
            file["npa/d_redge"] = npa.d_redge
            file["npa/d_tedge"] = npa.d_tedge
        end
    end
end

struct FIDASIMPlasmaParameters
    nr::Int
    nz::Int
    nphi::Int
    r::Vector{Float64}
    z::Vector{Float64}
    phi::Vector{Float64}
    mask::Union{Matrix{Int16},Array{Int16,3}}
    te::Union{Matrix{Float64},Array{Float64,3}}
    ti::Union{Matrix{Float64},Array{Float64,3}}
    dene::Union{Matrix{Float64},Array{Float64,3}}
    zeff::Union{Matrix{Float64},Array{Float64,3}}
    vr::Union{Matrix{Float64},Array{Float64,3}}
    vt::Union{Matrix{Float64},Array{Float64,3}}
    vz::Union{Matrix{Float64},Array{Float64,3}}
    denn::Union{Matrix{Float64},Array{Float64,3}}
end

function FIDASIMPlasmaParameters(filename::String)
    isfile(filename) || error("File does not exist")

    f = h5open(filename)

    fp = f["plasma"]
    nr = read(fp["nr"])
    nz = read(fp["nz"])
    if "nphi" in keys(fp)
        nphi = read(fp["nphi"])
    else
        nphi = 1
    end
    r = read(fp["r"])
    z = read(fp["z"])
    if nphi > 1
        phi = read(fp["phi"])
    else
        phi = [0.0]
    end

    mask = Int16.(read(fp["mask"]))
    te = read(fp["te"])
    ti = read(fp["ti"])
    dene = read(fp["dene"])
    zeff = read(fp["zeff"])

    vr = read(fp["vr"])
    vt = read(fp["vt"])
    vz = read(fp["vz"])

    if "denn" in keys(fp)
        denn = read(fp["denn"])
    else
        denn = zero(dene)
    end

    close(f)

    return FIDASIMPlasmaParameters(nr,nz,nphi,r,z,phi,mask,te,ti,dene,zeff,vr,vt,vz,denn)
end

function Base.show(io::IO, s::FIDASIMPlasmaParameters)
    print(io, "FIDASIMPlasmaParameters: $(s.nr)×$(s.nz)×$(s.nphi)")
end

mutable struct FIDAWeightFunction{TL<:AbstractVector,TE<:AbstractVector,
                                  TP<:AbstractVector,TW<:AbstractArray}
    lambda::TL
    energy::TE
    pitch::TP
    W::TW
end

function FIDAWeightFunction(filename::String,ichan)
    isfile(filename) || error("File does not exist")

    f = h5open(filename)

    W = dropdims(f["weight"][:,:,:,ichan],dims=4)
    lambda = read(f,"lambda")
    energy = read(f,"energy")
    pitch = read(f,"pitch")
    close(f)

    dE = energy[2]-energy[1]
    dP = abs(pitch[2]-pitch[1])
    return FIDAWeightFunction(lambda,energy,pitch,W*(1e4*dE*dP))
end


text:
using ShaleDrillingLikelihood: DrillingRevenueTimeTrend,
    DrillingRevenueNoTrend,
    DrillingRevenueTimeFE,
    DrillingRevenueTimeTrendDgt0,
    baseyear,
    tech,
    range_i_to_ip1,
    AbstractUnitProblem,
    DateQuarter,
    state_idx,
    InitialDrilling,
    j1chars,
    yearrange

export FormulaProduce, FormulaRoyalty, ThetaProduceStarting

using StatsBase: countmap
using FileIO: load
using Dates: year
using DataFrames, Query, StatsModels
using Base: OneTo

const cols = Dict(
    :j2ptr   => "j2ptr",
    :j1ptr   => "j1ptr",
    :iwptr   => "iw_ptr",
    :tptr    => "tptr",
    :qptr    => "q_ptr",

    :prices  => "prices",
    :ichars  => "ichars",
    :j1chars => "j1chars",
    :qchars  => "qchars",
    :tchars  => "tchars",
)

# ----------------------------
# Production data
# ----------------------------

# variables for regressions
function FormulaProduce(::DrillReward{<:DrillingRevenueTimeTrend}, startvalues=false)
    strt = @formula(logcumgas_r ~ 1 + log_ogip + α_t + (1|icat) + (1|iD))
    full = @formula(logcumgas_r ~ 1 + log_ogip + α_t)
    fm = startvalues ? strt : full
    return fm
end

function FormulaProduce(::DrillReward{<:DrillingRevenueTimeTrendDgt0}, startvalues=false)
    strt = @formula(logcumgas_r ~ 1 + log_ogip + α_t + Dgt0 + (1|icat) + (1|iD))
    full = @formula(logcumgas_r ~ 1 + log_ogip + α_t + Dgt0)
    fm = startvalues ? strt : full
    return fm
end

function FormulaProduce(::DrillReward{<:DrillingRevenueNoTrend}, startvalues=false)
    strt = @formula(logcumgas_r ~ 1 + log_ogip  + (1|icat) + (1|iD))
    full = @formula(logcumgas_r ~ 1 + log_ogip)
    fm = startvalues ? strt : full
    return fm
end
function FormulaProduce(::DrillReward{<:DrillingRevenueTimeFE}, startvalues=false)
    strt = @formula(logcumgas_r ~ 1 + log_ogip + well_start_fe  + (1|icat) + (1|iD))
    full = @formula(logcumgas_r ~ 1 + log_ogip + well_start_fe)
    fm = startvalues ? strt : full
    return fm
end

FormulaProduce(x::DynamicDrillModel) = FormulaProduction(reward(x))

function yearmonth_float(t)
    ym = yearmonth(t)
    return first(ym) + (last(ym)-1)/12
end

function ym_to_ur(ur::UnitRange, ym::Float64)
    baseyr = first(ur)
    lastyr = last(ur)
    y = floor(ym - baseyr) + baseyr
    return clamp(y, baseyr, lastyr)
end

ym_to_ur(ur::UnitRange, t::Date) = ym_to_ur(ur, yearmonth_float(t))
ym_to_ur(rwrd::DrillReward, x) = ym_to_ur(yearrange(tech(revenue(rwrd))), x)

function update_qchars!(_qchars::DataFrame, rwrd)
    _qchars[!,:α_t]  .= year.(_qchars[!,:well_start_date]) .- baseyear(rwrd)
    _qchars[!,:icat] .= CategoricalVector(_qchars[!,:i])
    wsy = map(t -> ym_to_ur(rwrd, t), _qchars[!,:well_start_date])
    _qchars[!,:well_start_fe] = CategoricalArray(wsy)
    return _qchars
end

function DataProduce(rwrd::DrillReward, path; model=ProductionModel(), kwargs...)

    # load in data
    alldata = load(path)
    _qchars = alldata[cols[:qchars]]
    update_qchars!(_qchars,rwrd)

    # formula
    startvalues = false
    fm = FormulaProduce(rwrd, startvalues)

    # y, X
    qmf = ModelFrame(fm, _qchars)
    y = response(qmf)
    X = Matrix(modelmatrix(qmf)')
    X[1:end-1,:] .= X[2:end,:]
    X[end,:] .= 1

    @assert all(@view(X[end,:]) .== 1)

    # data
    obs_ptr   = Int.(alldata[cols[:qptr]])
    group_ptr = Int.(alldata[cols[:iwptr]])
    data = DataProduce(y, X, obs_ptr, group_ptr)
    return data
end

function ThetaProduceStarting(rwrd::DrillReward, path)

    alldata = load(path)
    _qchars = alldata[cols[:qchars]]
    update_qchars!(_qchars,rwrd)

    startingvals = true
    fm = FormulaProduce(rwrd, startingvals)
    res = fit!(LinearMixedModel(fm, _qchars))
    σwell, σψ, σϵ = vcat(std(res)...)
    cf = coef(res)
    qcoef = vcat(σψ, cf[2:end], cf[1], σϵ, σwell)
    return qcoef
end

# ----------------------------
# Royalty data
# ----------------------------

FormulaRoyalty(x) = @formula(royalty_id ~ 0 + log_houseval_med + grantor_share_in_state_NO + pct_0_imperv_2001 + log_ogip)

function DataRoyalty(rwrd::DrillReward, path; model=RoyaltyModel(), kwargs...)
    alldata = load(path)
    _ichars = alldata[cols[:ichars]]
    fm = FormulaRoyalty(rwrd)
    rmf = ModelFrame(fm, _ichars)
    y = Int.(response(rmf))
    X = Matrix(modelmatrix(rmf)')

    lhsvar = Symbol(fm.lhs)

    roy_dict = sort(countmap(_ichars[!,:royalty_nearest]))
    royalty_rates = collect(keys(roy_dict))

    # do some checks that data are OK
    y_dict   = sort(countmap(_ichars[!,lhsvar]))
    collect(values(y_dict)) == collect(values(roy_dict)) || throw(error())

    data = DataRoyalty(model, y, X, royalty_rates)
    return data
end

# ----------------------------
# Drilling data primitives
# ----------------------------

itypes(r::DataFrameRow) = ( r[:log_ogip], r[:royalty_nearest], )

# ----------------------------
# Prices
# ----------------------------

validnum(x) = isfinite(x) && !ismissing(x)

logprice(  r::DataFrameRow) = log(r[:well_revenue] / r[:ppi2009])
logrigrate(r::DataFrameRow) = log(r[:dayrate_c]    / r[:ppi2009])
myyear(    r::DataFrameRow, baseyr=BaseYear()) = floor(r[:qtrdate] - baseyr) + baseyr
year_clamp(r::DataFrameRow, lb, ub) = clamp(myyear(r), lb, ub)

@deprecate logcost(r) logrigrate(r)

function year_clamp(r::DataFrameRow, d::DrillReward{R,<:AbstractDrillingCost_TimeFE}) where {R<:DrillingRevenue}
    c = cost(d)
    return year_clamp(r, ShaleDrillingLikelihood.start(c), ShaleDrillingLikelihood.stop(c))
end

function ztuple(d::DrillReward{R,C}) where {R, C<:Union{DrillingCost_constant, DrillingCost_dgt1}}
    f(r) = (logprice(r), )
    return f
end

function ztuple(d::DrillReward{R,C}) where {R, C<:AbstractDrillingCost_TimeFE}
    y(r) = myyear(r, ShaleDrillingLikelihood.start(cost(d)))
    f(r) = (logprice(r),  y(r))
    return f
end

function ztuple(d::DrillReward{R,C}) where {R, C<:Union{DrillingCost_TimeFE_rigrate, DrillingCost_TimeFE_rig_costdiffs}}
    y(r) = myyear(r, ShaleDrillingLikelihood.start(cost(d)))
    f(r) = (logprice(r), logrigrate(r), y(r), )
    return f
end

# ----------------------------
# Dates
# ----------------------------

RDate_to_Date(x::Real)::Date = Date(1970,1,1) + Dates.Day(x)

function Ryearmon_to_Date(x::AbstractFloat)
  y = floor(x)
  m = round((x - y) * 12.) + 1
  Date(y,m)
end

position_in_range(x::Real, rng::StepRangeLen) = Int( (x-first(rng)) / step(rng) ) + 1

# ----------------------------
# statespace
# ----------------------------

function (::Type{T})(tc::DataFrame) where {T<:AbstractUnitProblem}
    _dmax  = maximum(tc[!,:d])
    _Dmax  = maximum(tc[!,:D])
    _τ0max = maximum(tc[!,:tau0])
    _τ1max = maximum(tc[!,:tau1])
    _ext   = maximum(tc[tc[!,:tau1] .== 0, :tau0])
    wp_tuple = Int.((_dmax, max(_Dmax,1), _τ0max, _τ1max, _ext))
    println("Making well problem with tuple $(wp_tuple)")
    _wp = T(wp_tuple...)
    return _wp
end

tau1_tau0_D_d1(r::DataFrameRow) = (r[:tau1], r[:tau0], r[:D], r[:d1],)

function sum_to_one!(x::AbstractVector)
    invsumx = 1/sum(x)
    x .*= invsumx
end

# ----------------------------
# Form primitives from datadrill
# ----------------------------

function DataDrillPrimitive(rwrd::DrillReward, path; kwargs...)

    rdata = load(path)
    j1ptr   = Int.(rdata[ cols[:j1ptr] ])
    tptr    = Int.(rdata[ cols[:tptr] ])

    tchars  = rdata[ cols[:tchars] ]
    _j2ptr   = rdata[ cols[:j2ptr] ]
    _j1chars = rdata[ cols[:j1chars] ]
    _ichars  = rdata[ cols[:ichars] ]
    _qchars  = rdata[ cols[:qchars] ]

    last_well_drill_year = year(maximum(_qchars[!,:well_start_date]))
    first_tchars, last_tchars = extrema(tchars[!,:qtrdate])
    last_year = max(last_well_drill_year, floor(last_tchars))

    prices  = rdata[ cols[:prices] ] |>
        @filter(
            validnum(_.dayrate_c) &&
            validnum(_.well_revenue) &&
            floor(first_tchars) <= _.qtrdate &&
            _.qtrdate < last_year+1
        ) |> DataFrame

    z = map(ztuple(rwrd), eachrow(prices))
    minqtr, maxqtr = extrema(prices[!,:qtrdate])
    zrng = DateQuarter(minqtr) : Quarter(1) : DateQuarter(maxqtr)
    qrng = minqtr : 0.25 : maxqtr

    maxt = nrow(tchars)
    searchq(q) = position_in_range(q, qrng)
    jtstart = searchq.(tchars[min.(maxt, tptr[1:end-1]), :qtrdate])

    zchars = ExogTimeVars(z, zrng)
    y = Int.(tchars[!,:d])
    wp = LeasedProblem(tchars)
    x = map(r -> state_idx(wp, tau1_tau0_D_d1(r)...), eachrow(tchars))

    ichars = map(itypes, eachrow(_ichars))
    j2ptr = Int(first(_j2ptr)) : Int(last(_j2ptr))

    @assert length(j2ptr) == length(_j2ptr)

    jwtvec = _j1chars[!, :share_of_unit_lease_owns]

    if length(j1ptr) == 0
        newj1ptr = ones(length(ichars)+1)
    else
        newj1ptr = j1ptr
    end

    data = DataDrillPrimitive(
        rwrd, newj1ptr, j2ptr, tptr, jtstart,
        jwtvec, ichars, y, x, zchars, wp
    )

    for unit in data
        sum_to_one!(j1chars(unit))
    end

    return data
end


text:
#
# Factors Reduce
#

"""
    reducedim(op, ϕ::Factor, dims, [v0])

Reduces a dimension in `ϕ` using function `op`.
See `Base.reducedim` for more details
"""
Base.reducedim(op, ϕ::Factor, dim::Symbol, v0=nothing) = reducedim(op, ϕ, [dim], v0)

# a more efficient version than `reducedim!(deepcopy(ϕ))` ?
function Base.reducedim(op, ϕ::Factor, dims::Vector{Symbol}, v0=nothing)
    _check_dims_valid(dims, ϕ)
    # needs to be a tuple for squeeze
    inds = (sort!(indexin(dims, ϕ))...)

    isempty(inds) && return deepcopy(ϕ)

    dims_new = deepcopy(ϕ.dimensions)
    deleteat!(dims_new, inds)
    pot_new = _red_dim(op, ϕ, inds, v0)

    return Factor(dims_new, pot_new)
end

# reduce the dimension and then squeeze them out
_red_dim(op, ϕ::Factor, inds::Tuple, ::Void) = squeeze(reducedim(op, ϕ.potential, inds), inds)
_red_dim(op, ϕ::Factor, inds::Tuple, v0) = squeeze(reducedim(op, ϕ.potential, inds, v0), inds)

"""
    Z(ϕ)

Return the partition function (sum of all elements).
"""
Z(ϕ) = sum(ϕ)

Base.sum(ϕ::Factor, dims=names(ϕ)) = reducedim(+, ϕ, dims)
Base.prod(ϕ::Factor, dims=names(ϕ)) = reducedim(*, ϕ, dims)
Base.maximum(ϕ::Factor, dims=names(ϕ)) = reducedim(max, ϕ, dims)
Base.minimum(ϕ::Factor, dims=names(ϕ)) = reducedim(min, ϕ, dims)

"""
    marginalize(ϕ, dims)

Sum out all but the specified dimensions
"""
marginalize(ϕ::Factor, dim::Symbol) = marginalize(ϕ, [dim])
marginalize(ϕ::Factor, dims::AbstractVector{Symbol}) = sum(ϕ, setdiff(names(ϕ), dims))


text:
## This can be added to Gtk, but for now is here

## add these to Gtk lists.jl

## return tuple for a row
function Base.getindex(store::Gtk.GtkListStore, index::Int)
    iter = Gtk.iter_from_index(store, index)
    GtkTreeModel(store)[iter]
end

## return i,j cell from store -- may not be in view coordinates
function Base.getindex(lstore::Gtk.GtkListStore, index::Int, column::Int)
    iter = Gtk.iter_from_index(lstore, index)

    val = Gtk.mutable(Gtk.GValue())
    Gtk.G_.value(GtkTreeModel(lstore), iter, column - 1, val)
    val[Any]
end

function Base.setindex!(lstore::Gtk.GtkListStore, value, index::Int, column::Int)
    iter = Gtk.iter_from_index(lstore, index)
    Gtk.G_.value(lstore, iter, column - 1, Gtk.gvalue(value))
end

## Selection ...

function selected(selection::GtkTreeSelection)
    model = Gtk.mutable(Ptr{GtkTreeModel})
    iter = Gtk.mutable(GtkTreeIter)

    if !hasselection(selection)
        return (model[], {})

    elseif getproperty(selection, :mode, Int) == Gtk.GConstants.GtkSelectionMode.MULTIPLE
        rows = Gtk.GLib.GList(ccall((:gtk_tree_selection_get_selected_rows, Gtk.libgtk),
                                    Ptr{Gtk._GSList{Ptr{Gtk.GtkTreePath}}},
                                    (Ptr{GObject}, Ptr{GtkTreeModel}),
                                    selection, model))
        function path_to_iter(path)
            iter = Gtk.mutable(GtkTreeIter)
            ret = bool(ccall((:gtk_tree_model_get_iter, Gtk.libgtk),
                             Cint,
                             (Ptr{Gtk.GObject}, Ptr{Gtk.GtkTreeIter}, Ptr{Gtk.GtkTreePath}),
                             model[], iter, path))
            if !ret
                error("No selection of GtkTreeSelection")
            end
            iter[]
        end
        return (model[], [path_to_iter(row) for row in  rows])
    else
        ret = bool(ccall((:gtk_tree_selection_get_selected,Gtk.libgtk),Cint,
                         (Ptr{GObject},Ptr{Ptr{GtkTreeModel}},Ptr{GtkTreeIter}),
                         selection, model, iter))
        if !ret
            error("No selection of GtkTreeSelection")
        end
        return (model[],{iter[]})
    end
end

        ## This is a function of the view -- not the model
##
## return Int[] of selected indices
function gtk_jgui_selected(view::Gtk.GtkTreeView)
    selection = Gtk.G_.selection(view)
    ret = Int[]

    hasselection(selection) || return ret

    if getproperty(selection, :mode, Int) == Gtk.GConstants.GtkSelectionMode.MULTIPLE
        model = Gtk.mutable(GtkTreeModel(getproperty(view, :model, Gtk.GtkListStore)))
        rows = Gtk.GLib.GList(ccall((:gtk_tree_selection_get_selected_rows, Gtk.libgtk),
                                    Ptr{Gtk._GSList{Gtk.GtkTreePath}},
                                    (Ptr{GObject}, Ptr{GtkTreeModel}),
                                    selection, model))

        for path in rows
            i = ccall((:gtk_tree_path_to_string, Gtk.libgtk), Ptr{Uint8},
                      ( Ptr{GtkTreePath},),
                      path) |> bytestring |> int |> x -> x .+ 1
            push!(ret, i)
        end
    else
        ## Gtk.selected *should* work here, but last line gives
        ## coercion issue.
        ## m, iter = selected(selection) # issue with Gtk.selected
        m = Gtk.mutable(Ptr{GtkTreeModel})
        iter = Gtk.mutable(GtkTreeIter)
        res = bool(ccall((:gtk_tree_selection_get_selected,Gtk.libgtk),Cint,
                         (Ptr{GObject},Ptr{Ptr{GtkTreeModel}},Ptr{GtkTreeIter}),
                         selection,m,iter))

        !res && error("No selection of GtkTreeSelection")
        ## what is row?
        i = ccall((:gtk_tree_model_get_string_from_iter, Gtk.libgtk),
                  Ptr{Uint8},
                  (Ptr{GObject}, Ptr{GtkTreeIter}), m[], iter) |> bytestring |> int |> x -> x+1
        push!(ret, i)
    end
    ret
end

## return path of selected
function gtk_jgui_tree_selected(view::Gtk.GtkTreeView)
    selection = Gtk.G_.selection(view)
    ret = Any[]

    hasselection(selection) || return ret

    ## convert i:j:k -> [i+1, j+1, k+1]
    path_to_index(p) =  map(x -> x+1, map(int, split(p, ":") ))

    if getproperty(selection, :mode, Int) == Gtk.GConstants.GtkSelectionMode.MULTIPLE
        model = Gtk.mutable(GtkTreeModel(getproperty(view, :model, Gtk.GtkListStore)))
        rows = Gtk.GLib.GList(ccall((:gtk_tree_selection_get_selected_rows, Gtk.libgtk),
                                    Ptr{Gtk._GSList{Gtk.GtkTreePath}},
                                    (Ptr{GObject}, Ptr{GtkTreeModel}),
                                    selection, model))

        for path in rows
            i = ccall((:gtk_tree_path_to_string, Gtk.libgtk), Ptr{Uint8},
                      ( Ptr{GtkTreePath},),
                      path) |> bytestring |> path_to_index
            push!(ret, i)
        end
    else
        ## Gtk.selected *should* work here, but last line gives
        ## coercion issue.
        ## m, iter = selected(selection) # issue with Gtk.selected
        m = Gtk.mutable(Ptr{GtkTreeModel})
        iter = Gtk.mutable(GtkTreeIter)
        res = bool(ccall((:gtk_tree_selection_get_selected,Gtk.libgtk),Cint,
                         (Ptr{GObject},Ptr{Ptr{GtkTreeModel}},Ptr{GtkTreeIter}),
                         selection,m,iter))

        !res && error("No selection of GtkTreeSelection")
        ## what is row?
        i = ccall((:gtk_tree_model_get_string_from_iter, Gtk.libgtk),
                  Ptr{Uint8},
                  (Ptr{GObject}, Ptr{GtkTreeIter}), m[], iter) |> bytestring |> path_to_index
        push!(ret, i)
    end
    ## return just one for tree
    ret[1]
end

## Add to selection using index
function Base.select!(view::Gtk.GtkTreeViewLeaf, index::Int)
    selection = Gtk.G_.selection(view)
    store = getproperty(view, :model, Gtk.GtkListStoreLeaf)
    iter = Gtk.iter_from_index(store, index)
    Gtk.select!(selection, iter)
end

## Add to selection using index
function Base.select!(view::Gtk.GtkTreeViewLeaf, index::Int)
    selection = Gtk.G_.selection(view)
    store = getproperty(view, :model, Gtk.GtkListStoreLeaf)
    iter = Gtk.iter_from_index(store, index)
    Gtk.select!(selection, iter)
end

function Gtk.unselect!(view::Gtk.GtkTreeViewLeaf, index::Int)
    selection = Gtk.G_.selection(view)
    store = getproperty(view, :model, Gtk.GtkListStoreLeaf)
    iter = Gtk.iter_from_index(store, index)
    Gtk.unselect!(selection, iter)
end

function Gtk.unselectall!(view::Gtk.GtkTreeViewLeaf)
    selection = Gtk.G_.selection(view)
    Gtk.unselectall!(selection)
end

##

## This is Gtk.path_at_pos modified to get more
function tree_view_row_col_from_x_y(treeView::GtkTreeView, x::Integer, y::Integer)
    return(0,0)

    ## XXX THIS IS FAILING XXX
    pathPtr = Gtk.mutable(Ptr{GtkTreePath})
    vcol = Gtk.mutable(@GtkTreeViewColumn())
    path = GtkTreePath()

    ret = bool( ccall((:gtk_tree_view_get_path_at_pos,Gtk.libgtk),Cint,
                      (Ptr{GObject},Cint,Cint,Ptr{Ptr{Void}},Ptr{Ptr{Void}},Ptr{Cint},Ptr{Cint} ),
                       treeView,x,y,pathPtr,vcol,C_NULL,C_NULL) )
    if ret
        path = convert(GtkTreePath, pathPtr[])
        vcol = vcol[]
        row = int(string(path)) + 1

        columns = Gtk.GList(ccall((:gtk_tree_view_get_columns, Gtk.libgtk),
                                Ptr{Gtk._GSList{Gtk.GtkTreeViewColumn}}, (Ptr{GObject},), treeView))
        col = filter(i -> getproperty(columns[i], :title, String) == getproperty(vcol, :title, String), 1:length(columns))[1]

        return(row, col)
    end
    return(0,0)
end

## Names
function get_tree_view_names(treeView::GtkTreeView)
    columns = Gtk.GList(ccall((:gtk_tree_view_get_columns, Gtk.libgtk),
                              Ptr{Gtk._GSList{Gtk.GtkTreeViewColumn}}, (Ptr{GObject},), treeView))
    [getproperty(column, :title, String) for column in columns]
end

function set_tree_view_names(treeView::GtkTreeView, nms::Vector)
    columns = Gtk.GList(ccall((:gtk_tree_view_get_columns, Gtk.libgtk),
                              Ptr{Gtk._GSList{Gtk.GtkTreeViewColumn}}, (Ptr{GObject},), treeView))
    [setproperty!(column, :title, nm) for (column, nm) in zip(columns, nms)]
end

### TreeStores ...
##
function tree_store_set_values(treeStore::GtkTreeStoreLeaf, iter, values)
    for (i,value) in enumerate(values)
        ccall((:gtk_tree_store_set_value,Gtk.libgtk),Void,(Ptr{GObject},Ptr{GtkTreeIter},Cint,Ptr{Gtk.GValue}),
              treeStore,iter,i-1,Gtk.gvalue(value))
    end
    iter[]
end


text:
push!(LOAD_PATH,"../src")

using Test
using TypesParticles
using TypesBasis
using BuildBasis
using InputParser
using NuclearAttraction

"""
		run_buildnuclrattract_test()

""" function run_buildnuclrattract_test()
    @testset "Nuclear Attraction" begin
        system,basisset = getcalcsetup("../examples/HatreeFock.in")
        basisfunc,numelec = buildbasisfunc(system,basisset);
        nuclrattractcalc = buildnuclrattract(system,basisfunc,basisset);
        nuclrattracttest = [-2.7494e0 -1.3415e0 -2.1471e0 -1.2642e0;
                  -1.3415e0 -1.2431e0 -1.2642e0 -1.2027e0;
                  -2.1471e0 -1.2642e0 -2.7494e0 -1.3415e0;
                  -1.2642e0 -1.2027e0 -1.3415e0 -1.2431e0];

        tolerance = 1.0e-4;
        for i=1:4
            for j=1:4
                @test sign(nuclrattractcalc[i,j]) == sign(nuclrattracttest[i,j]);
                diff = abs(nuclrattractcalc[i,j]-nuclrattracttest[i,j]);
                @test diff <= tolerance;
	    end
        end #i

    end #testset
end #run_nuclrattract_test()


text:
### Types
abstract type AbstractEnv end

abstract type AbstractController end


text:
g = DiGraph{Node{String}, Edge{Node{String}}}()

add_node!(g, "a")
add_node!(g, "b")
add_edge!(g, "a", "b")

@test length(nodes(g)) == 2
@test length(edges(g)) == 1
@test out_edges(g, "a") == in_edges(g, "b")

g2 = EvolvingGraph{Node{String}, String}()
add_edge!(g2, "a", "b", "t1")
add_edge!(g2, "b", "c", "t1")
add_edge!(g2, "c", "d", "t2")
add_edge!(g2, "a", "b", "t2")

g3 = aggregate_graph(g2)
display(g3)

A = adjacency_matrix(g3)
@test A[1,2] == 1
@test A[2,3] == 1
@test num_nodes(g3) == 4
@test num_edges(g3) == 3

# add a graph to an evolving graph
g = DiGraph{Node{String}, Edge{Node{String}}}()
add_node!(g, "1")
add_node!(g, "2")
add_edge!(g, "1", "2")
eg = EvolvingGraph{Node{String}, Int}()
add_graph!(eg, g, 1)
@test length(nodes(eg)) == 2
@test length(edges(eg)) == 1


text:
"""
    fast_chunk(x::AbstractArray, ::Val{n}, ::Val{dim})

Type-stable and faster version of `MLUtils.chunk`
"""
@inline fast_chunk(h::Int, n::Int) = (1:h) .+ h * (n - 1)
@inline function fast_chunk(x::AbstractArray, h::Int, n::Int, ::Val{dim}) where {dim}
    selectdim(x, dim, fast_chunk(h, n))
end
@inline function fast_chunk(x::AbstractArray, ::Val{N}, d::Val{D}) where {N, D}
    fast_chunk.((x,), size(x, D) ÷ N, 1:N, d)
end

"""
    flatten_spatial(x::AbstractArray{T, 4})

Flattens the first 2 dimensions of `x`, and permutes the remaining dimensions to (2, 1, 3)
"""
@inline function flatten_spatial(x::AbstractArray{T, 4}) where {T}
    return permutedims(reshape(x, (:, size(x, 3), size(x, 4))), (2, 1, 3))
end

"""
    seconddimmean(x)

Computes the mean of `x` along dimension `2`
"""
@inline seconddimmean(x) = dropdims(mean(x, dims=2); dims=2)

"""
    normalise(x::AbstractArray, activation; dims=ndims(x), epsilon=ofeltype(x, 1e-5))

Normalises the array `x` to have a mean of 0 and standard deviation of 1, and applies the activation function `activation` to the result.
"""
@inline function normalise(x::AbstractArray, ::typeof(identity); dims=ndims(x),
                           epsilon=ofeltype(x, 1e-5))
    xmean = mean(x, dims=dims)
    xstd = std(x, dims=dims, mean=xmean, corrected=false)
    return @. (x - xmean) / (xstd + epsilon)
end

@inline function normalise(x::AbstractArray, activation; dims=ndims(x),
                           epsilon=ofeltype(x, 1e-5))
    xmean = mean(x, dims=dims)
    xstd = std(x, dims=dims, mean=xmean, corrected=false)
    return @. activation((x - xmean) / (xstd + epsilon))
end

# Model construction utilities
function assert_name_present_in(name, possibilities)
    @assert name in possibilities "`name` must be one of $(possibilities)"
end

get_pretrained_weights_path(name::Symbol) = get_pretrained_weights_path(string(name))
function get_pretrained_weights_path(name::String)
    try
        return @artifact_str(name)
    catch LoadError
        throw(ArgumentError("No pretrained weights available for `$name`"))
    end
end

function initialize_model(name::Symbol, model; pretrained::Bool=false, rng=nothing, seed=0,
                          kwargs...)
    if pretrained
        path = get_pretrained_weights_path(name)
        ps = load(joinpath(path, "$name.jld2"), "parameters")
        st = load(joinpath(path, "$name.jld2"), "states")
    else
        if rng === nothing
            rng = Random.default_rng()
            Random.seed!(rng, seed)
        end

        ps, st = Lux.setup(rng, model)
    end
    return model, ps, st
end


text:
using Test

@testset "AllTests" begin
    include("tests1.jl")
    @testset "Test2" begin
        @test true
        sleep(0.5)
    end
    @testset "Test3" begin
        @test true
        sleep(0.5)
    end
    @testset "Test4" begin
        @test_broken false
    end
    @testset "Test5" begin
        @test true
        sleep(0.5)
        @test true
        sleep(0.5)
    end
    @testset "Test6" begin
        @test true
        sleep(0.5)
        @test true
        sleep(0.5)
        @test true
    end
    include("tests2/tests.jl")
end


text:
using Revise
using AdFem
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using MAT
using ADCMEKit
np = pyimport("numpy")
using PyPlot
using SpecialFunctions

# simulation parameter setup
n = 20
NT = 400
ρ = 0.1 # design variable in α-schemes
m = 5n
h = 1/n
Δt = 10000. /NT

# coordinates
xo = zeros((m+1)*(n+1))
yo = zeros((m+1)*(n+1))
for i = 1:m+1
  for j = 1:n+1
    idx = (j-1)*(m+1)+i
    xo[idx] = (i-1)*h
    yo[idx] = (j-1)*h
  end
end

# Dirichlet boundary condition on three sides, and Neumann boundary condition (traction-free) on the top
bdnode = bcnode("left  | lower | right", m, n, h)

# viscoelasticity η and shear modulus μ

ηf = (x,y)->begin
  if y<=0.25
    return 10000.
  else
    return 1.
  end
end

η = constant(eval_f_on_gauss_pts(ηf, m, n, h))
μ = 0.001 * constant(ones(4m*n))

# linear elasticity matrix
coef = 2μ*η/(η + μ*Δt)
mapH = c->begin
  c * diagm(0=>ones(2))
end
H = map(mapH, coef)

Δt = Δt * ones(NT)

# generalized mass matrix and stiffness matrix
M = 100. * constant(compute_fem_mass_matrix1(m, n, h))
K = compute_fem_stiffness_matrix1(H, m, n, h)
C = spzero((m+1)*(n+1))

C = 0.1 * M + 0.1 * K

# fixed displacement
db = zeros((m+1)*(n+1))
for j = 1:n+1
  idx = (j-1)*(m+1)+1
  if j<=div(n, 4)
    db[idx] = 1.
  else
    db[idx] = (1-(j-div(n, 4))/(3div(n, 4)))
  end
  # y = (j-1)*h
  # db[idx] = y * (1-y) * 0.0
end

# cast the problem into homogeneous Dirichlet problem
idof = ones(Bool, (m+1)*(n+1))
idof[bdnode] .= false
idof = findall(idof)

d0 = db
v0 = zeros((m+1)*(n+1))
# v0[bcnode("left", m, n, h)] .= 1.0

a0 = vector(idof, M[idof, idof]\((-K*db)[idof]), (m+1)*(n+1))

function solver(A, rhs)
  A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)
  rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
  A\rhs
end

function antiplane_visco_αscheme(M::Union{SparseTensor, SparseMatrixCSC},
  K::Union{SparseTensor, SparseMatrixCSC},
  d0::Union{Array{Float64, 1}, PyObject},
  v0::Union{Array{Float64, 1}, PyObject},
  a0::Union{Array{Float64, 1}, PyObject},
  Δt::Array{Float64};
  ρ::Float64 = 1.0)
  nt = length(Δt)
  αm = (2ρ-1)/(ρ+1)
  αf = ρ/(1+ρ)
  γ = 1/2-αm+αf
  β = 0.25*(1-αm+αf)^2
  d = length(d0)

  M = isa(M, SparseMatrixCSC) ? constant(M) : M
  K = isa(K, SparseMatrixCSC) ? constant(K) : K
  d0, v0, a0, Δt = convert_to_tensor([d0, v0, a0, Δt], [Float64, Float64, Float64, Float64])

  A = (1-αm)*M + (1-αf)*K*β*Δt[1]^2
  A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)

  function equ(dc, vc, ac, dt, εc, σc, i)
    dn = dc + dt*vc + dt^2/2*(1-2β)*ac
    vn = vc + dt*((1-γ)*ac)

    df = (1-αf)*dn + αf*dc
    vf = (1-αf)*vn + αf*vc
    am = αm*ac

    σ_dt = (1-αf)*dt
    Σ = 2* repeat(μ.*η/(η+μ*σ_dt),1,2) .*εc - repeat(η/(η+μ*σ_dt), 1, 2) .* σc
    Force = compute_strain_energy_term1(Σ, m, n, h)

    rhs = - (M*am + C*vf + K*df) + Force

    rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
    A\rhs

  end

  function condition(i, tas...)
    return i<=nt
  end
  function body(i, tas...)
    dc_arr, vc_arr, ac_arr, σc_arr = tas
    dc = read(dc_arr, i)
    vc = read(vc_arr, i)
    ac = read(ac_arr, i)
    σc = read(σc_arr, i)
    εc = eval_strain_on_gauss_pts1(dc, m, n, h)

    y = equ(dc, vc, ac, Δt[i],εc, σc, i)
    dn = dc + Δt[i]*vc + Δt[i]^2/2*((1-2β)*ac+2β*y)
    vn = vc + Δt[i]*((1-γ)*ac+γ*y)

    εn = eval_strain_on_gauss_pts1(dn, m, n, h)
    σn = 2* repeat(μ.*η/(η+μ*Δt[i]),1,2) .* (εn - εc) + repeat(η/(η+μ*Δt[i]), 1, 2) .* σc

    i+1, write(dc_arr, i+1, dn), write(vc_arr, i+1, vn), write(ac_arr, i+1, y), write(σc_arr, i+1, σn)
  end

  dM = TensorArray(nt+1); vM = TensorArray(nt+1); aM = TensorArray(nt+1); σM = TensorArray(nt+1);
  dM = write(dM, 1, d0)
  vM = write(vM, 1, v0)
  aM = write(aM, 1, a0)

  ε0 = eval_strain_on_gauss_pts1(d0, m, n, h)
  σ0 = batch_matmul(H, ε0)
  σM = write(σM, 1, σ0)

  i = constant(1, dtype=Int32)
  _, d, v, a = while_loop(condition, body, [i,dM, vM, aM, σM])
  set_shape(stack(d), (nt+1, length(a0))), set_shape(stack(v), (nt+1, length(a0))), set_shape(stack(a), (nt+1, length(a0)))
end

# d, v, a = αscheme(M, C, K, zeros(NT, (m+1)*(n+1)), d0, v0, a0, Δt; solve = solver)

d, v, a = antiplane_visco_αscheme(M, K, d0, v0, a0, Δt, ρ=ρ)

sess = Session(); init(sess)
d_, v_, a_ = run(sess, [d, v, a])

# close("all")
# for (k,tid) in enumerate(LinRange{Int64}(1, NT+1, 5))
#   t = (tid-1)*Δt[1]
#   plot((d_[tid, :]+db)[(1:m+1)],"C$k-", label="$t", markersize=3)
# end
# legend()

# close("all")
# for tid in LinRange{Int64}(1, NT+1, 5)
#   plot((a_[tid, :])[1:m+1], label="$tid")
# end
# legend()

# plot(d_[end,1:m+1])
# plot(d_[div(NT+1,2),1:m+1])
# plot(d_[div(NT+1,3),1:m+1])

# pcolormesh(reshape(d_[end,:], m+1, n+1)')
figure()
pl, = plot([], [], "o-", markersize = 3)
t = title("0")
xi = (0:m)*h
xlim(-h, (m+1)*h)
xlabel("Distance")
ylim(-0.1, 1.1)
ylabel("Displacement")
tight_layout()
function update(i)
  pl.set_data(xi[:], d_[i,1:m+1])
  t.set_text("time = $(i*Δt[1])")
end
p = animate(update, [1:1:20;25:5:NT+1])
saveanim(p, "displacement.gif")

figure()
pl, = plot([], [], "o-", markersize = 3)
t = title("time = 0")
xi = (0:m)*h
xlim(-h, (m+1)*h)
xlabel("Distance")
ylim(-0.0001, 0.0005)
ylabel("Velocity")
tight_layout()
function update(i)
  pl.set_data(xi[:], v_[i,1:m+1])
  t.set_text("time = $(i*Δt[1])")
end
p = animate(update, [1:1:20;25:5:NT+1])
saveanim(p, "velocity.gif")

figure()
pl, = plot([], [], "o-", markersize = 3)
t = title("time = 0")
xi = (0:m)*h
xlim(-h, (m+1)*h)
xlabel("Distance")
ylim(-0.001, 0.001)
ylabel("Strain Rate")
tight_layout()
function update(i)
  pl.set_data(xi[1:end-1], (v_[i,2:m+1]-v_[i,1:m])/h)
  t.set_text("time = $(i*Δt[1])")
end
p = animate(update, [1:1:20;25:5:NT+1])
saveanim(p, "strain_rate.gif")

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages