text:
using Distributed: pmap, workers, WorkerPool

function prepare_stacks(
    igram_path;
    overwrite = false,
    ref_row = nothing,
    ref_col = nothing,
    ref_station = nothing,
    geo_path = nothing,
    window = 5,
    zero_masked = false,
)
    igram_path = abspath(igram_path)

    unw_stack_file = abspath(joinpath(igram_path, UNW_FILENAME))
    cc_stack_file = abspath(joinpath(igram_path, CC_FILENAME))
    mask_stack_file = abspath(joinpath(igram_path, MASK_FILENAME))

    # Step 1: create the masks.h5 file
    create_mask_stacks(
        igram_path,
        mask_filename = mask_stack_file,
        overwrite = overwrite,
        geo_path = geo_path,
    )

    # Step 2: use these masks to zero bad ares to zero in .int and .cc files
    # Note: this may have been run already after making .int files, before unwrapping
    if zero_masked
        zero_masked_areas(igram_path, input_exts = [".int", ".cc", ".unw"])
    end

    # TODO: Leave out cc stack for now
    # create_hdf5_stack(cc_stack_file, ".cc", overwrite=overwrite, do_repack=true)

    deramp_unws(
        igram_path,
        input_ext = ".unw",
        output_ext = ".unwflat",
        overwrite = overwrite,
    )

    create_hdf5_stack(
        unw_stack_file,
        ".unwflat",
        dset_name = STACK_FLAT_DSET,
        overwrite = overwrite,
    )

    # TODO: do we need this after I have "NaN"s in the deramp function?
    # zero_masked_areas(igram_path, input_exts=[".unwflat"])

    if isnothing(ref_station) && isnothing(ref_row) && isnothing(ref_col)
        # TODO: auto pick, redo that logic here
        # ref_row, ref_col, ref_station = find_reference_location(
        #     unw_stack_file=unw_stack_file,
        #     cc_stack_file=cc_stack_file,
        #     mask_stack_file=mask_stack_file,
        # )
        ref_row, ref_col, ref_station = (50, 50, "")
    end

    shift_unw_file(
        unw_stack_file,
        ref_row = ref_row,
        ref_col = ref_col,
        ref_station = ref_station,
        window = window,
        overwrite = overwrite,
    )
end

# Note: 6 workers for the .geo.mask file creation seems fastest
# More and they thrash while reading
function create_mask_stacks(
    igram_path;
    mask_filename = nothing,
    geo_path = nothing,
    overwrite = false,
)
    if isnothing(mask_filename)
        mask_filename = abspath(joinpath(igram_path, MASK_FILENAME))
    end
    if isnothing(geo_path)
        geo_path = dirname(dirname(abspath(igram_path)))
    end

    row_looks, col_looks = find_looks_taken(igram_path, geo_path = geo_path)
    loop_over_files(
        _get_geo_mask,
        geo_path,
        ".geo",
        ".geo.mask",
        looks = (row_looks, col_looks),
        out_dir = igram_path,
        max_procs = 4,
        overwrite = overwrite,
    )

    # Now with bigger geo files read in parallel,
    # write all to hdf5 file as stack
    save_masks(igram_path, geo_path, overwrite = overwrite)

    # Finall, add the aux. information
    dem_rsc = Sario.load(joinpath(igram_path, "dem.rsc"))
    Sario.save_dem_to_h5(mask_filename, dem_rsc, DEM_RSC_DSET, overwrite = overwrite)

    geo_date_list = Sario.find_geos(directory = geo_path, parse = true)
    Sario.save_geolist_to_h5(mask_filename, geo_date_list, overwrite = overwrite)

    igram_list = Sario.find_igrams(directory = igram_path, parse = true)
    Sario.save_intlist_to_h5(mask_filename, igram_list, overwrite = overwrite)
end

_get_geo_mask(arr) = (arr .== 0)

"""Creates igram masks by taking the logical-or of the two .geo files

Assumes save_geo_masks already run"""
function save_masks(
    igram_path,
    geo_path;
    overwrite = false,
    mask_filename = MASK_FILENAME,
    geo_dset_name = GEO_MASK_DSET,
    igram_dset_name = IGRAM_MASK_DSET,
)
    # TODO?: add checks for overwrite
    !Sario.check_dset(mask_filename, geo_dset_name, overwrite) && return
    !Sario.check_dset(mask_filename, igram_dset_name, overwrite) && return
    !Sario.check_dset(mask_filename, GEO_MASK_SUM_DSET, overwrite) && return

    geo_mask_stack = load_stack(directory = igram_path, file_ext = ".geo.mask")

    int_date_list = Sario.find_igrams(directory = igram_path)
    int_file_list = Sario.find_igrams(directory = igram_path, parse = false)
    geo_date_list = Sario.find_geos(directory = geo_path)

    nrows, ncols, _ = size(geo_mask_stack)
    int_mask_stack =
        Array{eltype(geo_mask_stack),3}(undef, (nrows, ncols, length(int_file_list)))
    for (idx, (early, late)) in enumerate(int_date_list)
        early_idx = findfirst(isequal(early), geo_date_list)
        late_idx = findfirst(isequal(late), geo_date_list)
        early_mask = geo_mask_stack[:, :, early_idx]
        late_mask = geo_mask_stack[:, :, late_idx]

        out_filename = int_file_list[idx] * ".mask"

        new_mask = early_mask .| late_mask
        int_mask_stack[:, :, idx] .= new_mask
        # save(out_filename, new_mask)  # If we want to save the igram mask itself
    end
    save_hdf5_stack(
        mask_filename,
        igram_dset_name,
        convert(Array{Bool}, int_mask_stack),
        overwrite = false,
    )
    save_hdf5_stack(
        mask_filename,
        geo_dset_name,
        convert(Array{Bool}, geo_mask_stack),
        overwrite = false,
    )
    # write(f, dset_name, permutedims(stack, (2, 1, 3)))
    # Also create one image of the total masks
    # TODO: any way to have sum reduce the dims?
    h5write(
        mask_filename,
        GEO_MASK_SUM_DSET,
        permutedims(dropdims(sum(geo_mask_stack, dims = 3), dims = 3)),
    )
end

function _remaining_files(
    in_files::Array{String,1},
    out_files::Array{String,1},
    overwrite::Bool,
)
    if overwrite
        return in_files, out_files, collect(1:length(in_files))
    end
    in_todo = Array{String,1}()
    out_todo = similar(in_todo)
    for (i, o) in zip(in_files, out_files)
        if !isfile(o)
            push!(in_todo, i)
            push!(out_todo, o)
        end
    end
    return in_todo, out_todo, indexin(in_todo, in_files)
end

# Step 2 functions: zeroing bad areas
function zero_masked_areas(directory; input_exts = [".int", ".cc"])
    for input_ext in input_exts
        in_files = find_files(input_ext, directory)
        out_files = in_files  # We are saving over the same file, just setting pixels to 0
        idxs = collect(1:length(in_files))

        # TODO: should we check for 0s in masked areas? maybe in the real function
        # in_todo, out_todo, idxs_todo = _remaining_files(in_files, out_files)

        wp = _get_workerpool(8)
        println("Starting zeroing areas for $input_ext")
        @time pmap(
            (name_in, name_out, mask_idx)->_load_and_run(
                zero_file,
                name_in,
                name_out,
                mask_idx,
                return_amp = true,
                do_permute = false,
            ),
            wp,
            in_files,
            out_files,
            idxs,
        )
        println("completed zeroing for $input_ext")
    end
end

# TODO: if i wanna make these true, default, i need to redo the _load_and_run arg passing
function zero_file(
    arr3d::AbstractArray{T,3},
    mask_idx;
    do_permute = false,
) where {T <: Number}
    mask = _read_mask(mask_idx, do_permute)
    data = @view arr3d[:, :, 2]
    data[mask] .= 0
    return arr3d
end

function zero_file(arr::AbstractArray{T,2}, mask_idx; do_permute = false) where {T <: Number}
    mask = _read_mask(mask_idx, do_permute)
    arr[mask] .= 0
    return arr
end

function _read_mask(mask_idx, do_permute = false)
    m = Bool.(h5read(MASK_FILENAME, IGRAM_MASK_DSET, (:, :, mask_idx))[:, :, 1])
    return do_permute ? permutedims(m) : m
end

"""Remove a linear ramp from .unw files, save as .unwflat"""
function deramp_unws(
    directory = ".";
    input_ext = ".unw",
    output_ext = ".unwflat",
    overwrite = true,
    do_permute = false,
)  # , highpass=false)
    println("Writing deramped unws to $output_ext files")
    in_files = find_files(input_ext, directory)
    out_files = [replace(fname, input_ext => output_ext) for fname in in_files]

    # If we already have these files made, skip the function
    in_todo, out_todo, idxs_todo = _remaining_files(in_files, out_files, overwrite)

    println("$(length(out_todo)) files remain, overwrite = $overwrite")
    out_todo = overwrite ? in_todo : out_todo  # If we have overwrite=true,
    if length(out_todo) == 0 && !overwrite
        println("skipping")
        return nothing
    end

    wp = _get_workerpool(8)
    # TODO: high pass failing due to edge effects for now
    # if highpass
    #     println("Starting stack deramp: High pass filter")
    #     pmap((name_in, name_out, mask_idx) -> _load_and_run(remove_lowpass, name_in, name_out, mask_idx, do_permute=do_permute),
    #     wp, in_todo, out_todo, idxs_todo)
    # else
    println("Starting stack deramp ")
    pmap(
        (name_in, name_out, mask_idx)->_load_and_run(
            remove_ramp,
            name_in,
            name_out,
            mask_idx,
            do_permute = do_permute,
        ),
        wp,
        in_todo,
        out_todo,
        idxs_todo,
    )
    # end
    println("Deramping stack complete")
end

sliceview(stack) = [view(stack, :, :, i) for i = 1:size(stack, 3)]

function remove_ramp(z, mask::AbstractArray{<:Number}; use_quadratic = false)
    # z_masked = copy(z)  # Do I really want a copy??
    z[mask] .= NaN

    if use_quadratic
        return z - quadratic_ramp(z)
    else
        return z - estimate_ramp(z)
    end
end

function remove_ramp(z, mask_idx::Int; do_permute = false, use_quadratic = false)
    mask = _read_mask(mask_idx, do_permute)
    return remove_ramp(z, mask, use_quadratic = use_quadratic)
end

function remove_lowpass(z, mask::AbstractArray{<:Number})
    # z_masked = copy(z)  # Do I really want a copy??
    z[mask] .= 0
    return z - lowpass(z)
end

function remove_lowpass(z, mask_idx::Int; do_permute = false)
    mask = _read_mask(mask_idx, do_permute)
    return remove_lowpass(z, mask)
end

import ImageFiltering: imfilter, Kernel
function lowpass(z::AbstractArray{<:AbstractFloat,2})
    # TODO: fix this for finding size
    g100 = Kernel.gaussian(100)
    return imfilter(z, g100)
end

"""Takes a 2D array an fits a linear plane to the data
    Ignores pixels that have nan or missing values

Since it is an order 1 surface, it will be 3 numbers, a, b, c from
     ax + by + c = z
"""
function estimate_ramp(z::AbstractArray{<:AbstractFloat,2})
    A = ones((length(z), 3))
    coeffs = Array{eltype(A),1}(undef, (3,))
    z_fit = similar(z)

    zidxs = CartesianIndices(z)
    for idx = 1:size(A, 1)
        # if good_idxs[idx]
        # row, col is equiv to y, x, but subtract 1 to start at 0
        y, x = zidxs[idx].I .- 1
        A[idx, 1] = x
        A[idx, 2] = y
        # end
    end

    good_idxs = .~isnan.(z)
    coeffs .= A[vec(good_idxs), :] \ z[vec(good_idxs)]
    a, b, c = coeffs

    for idx in CartesianIndices(z_fit)
        # again subtract 1 to keep consistent with the fitting
        y, x = idx.I .- 1
        z_fit[idx] = a * x + b * y + c
    end

    return z_fit
end

function quadratic_ramp(z::AbstractArray{<:AbstractFloat,2})
    A = ones((length(z), 6))
    coeffs = Array{eltype(A),1}(undef, (6,))
    z_fit = similar(z)

    zidxs = CartesianIndices(z)
    for idx = 1:size(A, 1)
        # if good_idxs[idx]
        # row, col is equiv to y, x, but subtract 1 to start at 0
        y, x = zidxs[idx].I .- 1
        A[idx, 1] = x
        A[idx, 2] = y
        A[idx, 3] = x * y
        A[idx, 4] = x^2
        A[idx, 5] = y^2
        # end
    end

    good_idxs = .~isnan.(z)
    coeffs .= A[vec(good_idxs), :] \ z[vec(good_idxs)]
    # coeffs will be 6 elements for the quadratic
    a, b, c, d, e, f = coeffs

    for idx in CartesianIndices(z_fit)
        # again subtract 1 to keep consistent with the fitting
        y, x = idx.I .- 1
        z_fit[idx] = a * x + b * y + c * x * y + d * x^2 + e * y^2 + f
    end

    return z_fit
end

"""Make stack as hdf5 file from a group of existing files"""
function create_hdf5_stack(
    filename::String,
    file_ext::String;
    dset_name = STACK_DSET,
    directory = ".",
    geo_path = "../",
    overwrite = false,
    do_repack = false,
)

    filename = isabspath(filename) ? filename : abspath(joinpath(directory, filename))
    println("Creating stack file $filename")

    get_file_ext(filename) in (".h5", ".hdf5") ||
    ArgumentError("filename must end in .h5 or .hdf5")
    !Sario.check_dset(filename, dset_name, overwrite) && return

    file_list = find_files(file_ext, directory)

    testf = Sario.load(file_list[1], rsc_file = "dem.rsc")
    rows, cols = size(testf)
    # Note: since we're just loading/saving, don't permute on way in or way out
    shape = (cols, rows, length(file_list))
    _create_dset(filename, dset_name, shape, eltype(testf))
    arr_buf = zeros(eltype(testf), (cols, rows))
    mean_buf = zeros(eltype(testf), (cols, rows))

    h5open(filename, "r+") do hf
        dset = hf[dset_name]
        for (idx, f) in enumerate(file_list)
            arr_buf .= Sario.load(f, do_permute = false)
            dset[:, :, idx] = arr_buf
            mean_buf += arr_buf
        end
    end
    h5write(filename, dset_name * "_filenames", file_list)

    !Sario.check_dset(filename, STACK_MEAN_DSET, overwrite) && return
    mean_buf ./= length(file_list)
    h5write(filename, STACK_MEAN_DSET, mean_buf)

    # Now save dem rsc as well
    # dem_rsc = Sario.load(Sario.find_rsc_file(directory=directory))
    dem_rsc = Sario.load(joinpath(directory, "dem.rsc"))
    Sario.save_dem_to_h5(filename, dem_rsc; overwrite = overwrite)
    if isnothing(geo_path)
        geo_path = dirname(dirname(abspath(directory)))
    end
    geo_date_list = Sario.find_geos(directory = geo_path)
    Sario.save_geolist_to_h5(filename, geo_date_list, overwrite = overwrite)
    igram_list = Sario.find_igrams(directory = directory)
    Sario.save_intlist_to_h5(filename, igram_list, overwrite = overwrite)

    if do_repack
        repack(filename, dset_name)
    end

    return filename
end

function repack(filename, dset_name)
    # Now repack so that chunks are depth-wise (pixels for quick loading)
    nrows, ncols, nlayers = size(filename, dset_name)
    tmp_name = _make_tmp_name(filename)
    chunk_size = "$(nlayers)x1x1"
    println("Repacking $filename into size $chunk_size depth-wise chunks")
    cmd = `h5repack -v -f NONE -l $dset_name:CHUNK=$chunk_size $filename $tmp_name`
    println("Running:")
    println(cmd)
    @time run(cmd)
    run(`mv $tmp_name $filename`)
end
function _make_tmp_name(filename)
    fparts = splitpath(filename)
    return joinpath(fparts[1:end - 1]..., "tmp_" * fparts[end])
end

"""Subtracts reference pixel group from each layer

window is the size around the reference pixel to
average at each layer"""
function shift_unw_file(
    unw_stack_file::String;
    stack_flat_dset = nothing,
    ref_row = nothing,
    ref_col = nothing,
    window = 5,
    ref_station = nothing,
    overwrite::Bool = false,
    do_repack::Bool = true,
)
    """Runs a reference point shift on flattened stack of unw files stored in .h5

    do_repack will run h5repack to chunk"""
    if isnothing(stack_flat_dset)
        stack_flat_dset = STACK_FLAT_DSET
    end
    stack_flat_shifted_dset = stack_flat_dset * "_shifted"

    !Sario.check_dset(unw_stack_file, stack_flat_shifted_dset, overwrite) && return

    if (isnothing(ref_row) || isnothing(ref_col))
        isnothing(ref_station) &&
        throw(ArgumentError("Need ref_station if no ref_row/ref_col"))
        println("Using $ref_station as reference")
        demrsc = Sario.load_dem_from_h5(unw_stack_file)
        ref_row, ref_col = MapImages.station_rowcol(ref_station, demrsc)
    end
    println("Starting shift_stack: using ($ref_row, $ref_col) as reference")

    stack_size = size(unw_stack_file, stack_flat_dset)
    h5open(unw_stack_file, "r+") do f
        if !(stack_flat_dset in names(f))
            throw("Need $stack_flat_dset to be created in $unw_stack_file before shift stack can be run")
        end

        # stack_in = f[stack_flat_dset]
        d_create(
            f,
            stack_flat_shifted_dset,
            datatype(Float32),
            dataspace(stack_size),
            # "chunk", (10, 10, size(stack_in, 3)),  # Seems better to repack?
            # Chunking here and writing by images is like 100x slower
        )
    end
    h5open(unw_stack_file, "r+") do f
        stack_in = f[stack_flat_dset]
        stack_out = f[stack_flat_shifted_dset]

        # shift_stack(stack_in, stack_out, ref_row, ref_col, window=window)
        # Note: switching these so we don't have to permute dims upon loading HDF5Dset
        shift_stack(stack_in, stack_out, ref_col, ref_row, window = window)

    end

    if do_repack
        @time repack(unw_stack_file, stack_flat_shifted_dset)
    end

    h5writeattr(
        unw_stack_file,
        stack_flat_shifted_dset,
        Dict(REFERENCE_ATTR => [ref_row, ref_col]),
    )
    # Make sure we don't write Nothing
    station_str = isnothing(ref_station) ? "" : ref_station
    h5writeattr(
        unw_stack_file,
        stack_flat_shifted_dset,
        Dict(REFERENCE_STATION_ATTR => station_str),
    )

    println("Shifting stack complete")
end

"""
Note: window is size of the group around ref pixel to avg for reference.
    if window=1 or nothing, only the single pixel used to shift the group.
"""
# TODO: check later if worth doing one for HDF5Dataset, one for Array
function shift_stack(stack_in, stack_out, ref_row::Int, ref_col::Int; window::Int = 5)
    half_win = div(window, 2)
    winsize = (2half_win + 1, 2half_win + 1)  # Make sure it is odd sized
    patch = Array{Float32,2}(undef, winsize)

    nrows, ncols, nlayers = size(stack_in)
    layer = Array{Float32,2}(undef, (nrows, ncols))
    layer_out = similar(layer)

    @show typeof(stack_in)
    @show typeof(stack_out)

    @inbounds for k = 1:nlayers
        layer .= view(stack_in, :, :, k)
        layer_out .= _shift_layer(layer, patch, ref_row, ref_col, half_win)
        stack_out[:, :, k] = layer_out

        if k % 100 == 0
            println("Finished with $k layers")
        end
    end
    return stack_out
end
function shift_stack(stack_in, ref_row::Int, ref_col::Int; window::Int = 5)
    stack_out = Array{eltype(stack_in),ndims(stack_in)}(undef, size(stack_in))
    return shift_stack(stack_in, stack_out, ref_row, ref_col; window = window)
end

function _shift_layer(layer, patch, ref_row, ref_col, half_win)
    patch .= layer[ref_row - half_win:ref_row + half_win, ref_col - half_win:ref_col + half_win]

    # Adding the `view` to eliminate extra singleton dimension from HDF5
    return view(layer .- mean(patch), :, :)
end

function create_mean_hdf5(h5file::String; dset_name::String = STACK_DSET)
    h5open(h5file, "cw") do f
        dset = f[dset_name]
        nrows, ncols, nlayers = size(dset)
        mean_buf = zeros(eltype(dset), nrows, ncols)

        println("Creating mean from stack:")
        @show size(dset)
        for idx = 1:nlayers
            # Adding the `[:,:]` to eliminate extra singleton dimension from HDF5
            mean_buf .+= dset[:, :, idx][:, :]
        end
        mean_buf ./= nlayers
        write(f, STACK_MEAN_DSET, mean_buf)
    end
end

function _create_dset(h5file, dset_name, shape, dtype)
    # TODO: add a chunking option for ones we access depth-wise vs layer-wise
    # dset = d_create(g, "B", datatype(Float64), dataspace(1000,100,10), "chunk", (100,100,1))
    h5open(h5file, "r+") do f
        d_create(f, dset_name, datatype(dtype), dataspace(shape))
    end
end

"""Parallel way to run over many image files
Pass a function `f` to operate on each image

    With 8 workers, on 100 files of 1000x1000

    @time InsarTimeseries.loop_over_files(add1, ".", ".int", "_flat.int")

    runs in 1 second, vs about 8 seconds serially
"""
function loop_over_files(
    f,
    directory::String,
    input_ext::String,
    output_ext::String;
    looks = (1, 1),
    out_dir = nothing,
    max_procs::Union{Nothing,Int} = nothing,
    overwrite = false,
    do_permute = true,
)
    in_files = find_files(input_ext, directory)
    println("Looping over files: $in_files")
    out_files = [replace(fname, input_ext => output_ext) for fname in in_files]
    if !isnothing(out_dir)
        # Redo the path out if specified
        out_files = [joinpath(out_dir, splitpath(f)[end]) for f in out_files]
    end

    in_todo, out_todo, idxs_todo = _remaining_files(in_files, out_files, overwrite)
    println("$(length(out_todo)) files remain, overwrite = $overwrite")
    if length(out_todo) == 0 && !overwrite
        println("skipping")
        return nothing
    end

    wp = _get_workerpool(max_procs)
    pmap(
        (name_in, name_out)->_load_and_run(f, name_in, name_out, looks = looks, do_permute = do_permute),
        wp,
        in_todo,
        out_todo,
    )
    # @sync @distributed for fname in in_files
    # end
end

"""To manually pass which files you want to run on and save to"""
function loop_over_files(
    f,
    in_files::Array{String,1},
    out_files::Array{String,1};
    looks = (1, 1),
    max_procs::Union{Nothing,Int} = nothing,
    do_permute = true,
)
    wp = _get_workerpool(max_procs)
    pmap(
        (name_in, name_out)->_load_and_run(f, name_in, name_out, looks = looks, do_permute = do_permute),
        wp,
        in_files,
        out_files,
    )
end

function _load_and_run(f, name_in, name_out, args...; looks = (1, 1), kwargs...)
    println("Input: $name_in, out: $name_out")
    input_arr = Sario.load(name_in; looks = looks, kwargs...)
    out_arr = f(input_arr, args...)
    save(name_out, out_arr; kwargs...)
end

function _get_workerpool(max_procs)
    if isnothing(max_procs)
        return WorkerPool(workers())
    else
        m = min(max_procs, length(workers()))
        return WorkerPool(workers()[1:m])
    end
end

"""Calculates how many looks from .geo files to .int files"""
function find_looks_taken(
    igram_path;
    geo_path = "",
    igram_dem_file = "dem.rsc",
    geo_dem_file = "elevation.dem.rsc",
)
    if isempty(geo_path)
        geo_path = dirname(dirname(abspath(igram_path)))
    end

    geo_dem_rsc = Sario.load(joinpath(geo_path, geo_dem_file))
    igram_dem_rsc = Sario.load(joinpath(igram_path, igram_dem_file))

    row_looks = div(geo_dem_rsc.rows, igram_dem_rsc.rows)
    col_looks = div(geo_dem_rsc.cols, igram_dem_rsc.cols)
    return row_looks, col_looks
end


text:
"""
    log_text(logger::TBLogger, name::String, text::Any; step=step(logger))

Logs text with name `name` at step `step`
- text: If text is a 2-D or 3-D `Array`, it will be rendered as a table or a list. Any other data will be represented as string
"""
function log_text(logger::TBLogger, name::String, text::Any; step = nothing)
    summ = SummaryCollection(text_summary(name, text))
    write_event(logger.file, make_event(logger, summ, step=step))
end

function text_summary(name::String, text::Any)
    #Create a string tensor
    #shape of the tensor
    dims = Array{TensorShapeProto_Dim, 1}()
    if isa(text, AbstractArray)
        for dim in size(text)
            push!(dims, TensorShapeProto_Dim(size = dim))
        end
    else
        push!(dims, TensorShapeProto_Dim(size = 1))
    end
    texttensorshape = TensorShapeProto(dim = Vector(dims))

    #content of the tensor
    textstringval = Vector{Array{UInt8, 1}}()
    if isa(text, AbstractArray)
        if ndims(text) == 2
            text = permutedims(text, (2, 1))
        end
        for string in text
            string = markdown_repr(string)
            push!(textstringval, Array{UInt8, 1}(string))
        end
    else
        text = markdown_repr(text)
        push!(textstringval, Array{UInt8, 1}(text))
    end
    #metadata for the text
    textcontent = serialize_proto(TextPluginData(version = 0))
    plugindata = SummaryMetadata_PluginData(plugin_name = "text", content = textcontent)
    smd = SummaryMetadata(plugin_data = plugindata)
    #create tensor
    texttensor = TensorProto(dtype = _DataType.DT_STRING, string_val = textstringval, tensor_shape = texttensorshape)
    Summary_Value(tag = name, metadata = smd, tensor = texttensor)
end

"""
    markdown_repr(x)

Returns the most suitable method for displaying `x` in a markdown document.
"""
function markdown_repr(x)
    for mime_type in ("text/markdown", "text/html", "text/plain")
        try
            if showable(mime_type, x)
                return repr(mime_type, x)
            end
        catch
            @debug "It said it was showable, but it lied" mime_type x exception=err
            # we blindly catch everything, because who knows what evil
            # people are doing in their misimplemented `repr` functions
        end
    end
    return repr(x)
end

ext: jl


text:
#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

"""
    struct NoDuplicateDict{K, V} <: AbstractDict{K, V}
        dict::Dict{K, V}
    end

Same as `Dict{K, V}` but errors if constructed from an iterator with duplicate
keys.
"""
struct NoDuplicateDict{K,V} <: AbstractDict{K,V}
    dict::Dict{K,V}
    NoDuplicateDict{K,V}() where {K,V} = new{K,V}(Dict{K,V}())
end

# Implementation of the `AbstractDict` API.
function Base.empty(::NoDuplicateDict, ::Type{K}, ::Type{V}) where {K,V}
    return NoDuplicateDict{K,V}()
end
Base.iterate(d::NoDuplicateDict, args...) = iterate(d.dict, args...)
Base.length(d::NoDuplicateDict) = length(d.dict)
Base.haskey(dict::NoDuplicateDict, key) = haskey(dict.dict, key)
Base.getindex(dict::NoDuplicateDict, key) = getindex(dict.dict, key)
function Base.setindex!(dict::NoDuplicateDict, value, key)
    if haskey(dict, key)
        error("Repeated index ", key, ". Index sets must have unique elements.")
    end
    return setindex!(dict.dict, value, key)
end
function NoDuplicateDict{K,V}(it) where {K,V}
    dict = NoDuplicateDict{K,V}()
    for (k, v) in it
        dict[k] = v
    end
    return dict
end
function NoDuplicateDict(it)
    return Base.dict_with_eltype((K, V) -> NoDuplicateDict{K,V}, it, eltype(it))
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
####################################
# Abstract type
####################################

abstract type AbstractDelaunayTriangulation end

ADT = AbstractDelaunayTriangulation
Base.size(DT::ADT) = (length(DT.indices[1]) - 1, length(DT.indices))
Base.length(DT::ADT) = length(DT.indices)
dimension(DT::ADT) = length(DT.indices[1]) - 1
nsimplices(DT::ADT) = length(DT.indices)

# Indexing
Base.getindex(DT::ADT, i) = DT.indices[i]
Base.getindex(DT::ADT, i::Colon, j::Colon) = hcat(DT.indices...,)
Base.getindex(DT::ADT, i::Colon, j) = hcat(DT[j]...,)
Base.getindex(DT::ADT, i::Colon, j::Int) = hcat(DT[j])

Base.firstindex(DT::ADT) = 1
Base.lastindex(DT::ADT) = length(DT)

Base.eachindex(s::ADT) = Base.OneTo(length(s))
Base.iterate(s::ADT, state = 1) = iterate(s.indices, state)

function summarise(DT::AbstractDelaunayTriangulation)
    _type = typeof(DT)
    n_simplices = nsimplices(DT)
    D = dimension(DT)
    summary = "$D-dimensional $_type with $n_simplices simplices"
end

Base.show(io::IO, DT::AbstractDelaunayTriangulation) = println(io, summarise(DT))

export
AbstractDelaunayTriangulation,
dimension,
nsimplices


text:
export updateWd,multWd,setSources,setSourcesSame,setDobs,setWd
function updateWd(pMis::Array{RemoteChannel},Dc::Array{RemoteChannel})
@sync begin
		for k=1:length(pMis)
	@async begin
			pMis[k] = remotecall_fetch(updateWd,pMis[k].where,pMis[k],Dc[k]);
		end
	end
end
return pMis;
end

function updateWd(pMisRF::RemoteChannel,Dc::RemoteChannel)
pMis  = take!(pMisRF)
Dc = fetch(Dc);
pMis.Wd = 1.0./(real(Dc - pMis.dobs) + 1e-3*mean(abs(pMis.dobs[:]))) + 1im./(imag(Dc - pMis.dobs) + 1e-3*mean(abs(pMis.dobs[:])));
put!(pMisRF,pMis)
return pMisRF;
end

function multWd(pMis::Array{RemoteChannel},beta::Float64)
@sync begin
		for k=1:length(pMis)
	@async begin
			pMis[k] = remotecall_fetch(multWd,pMis[k].where,pMis[k],beta);
		end
	end
end
return pMis;
end

function multWd(pMisRF::RemoteChannel,beta::Float64)
pMis  = take!(pMisRF)
pMis.Wd *= beta;
put!(pMisRF,pMis)
return pMisRF;
end

export setSources, setDobs

function setWd(pMis::Array{RemoteChannel},newWd::Array)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setWd,pMis[k].where,pMis[k],newWd[k]);
		end
	end
end
return pMis;
end

function setWd(pMisRF::RemoteChannel,newWd)
pMis  = take!(pMisRF)
pMis.Wd = newWd;
put!(pMisRF,pMis)
return pMisRF;
end
function setDobs(pMis::Array{RemoteChannel},newDobs::Array)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setDobs,pMis[k].where,pMis[k],newDobs[k]);
		end
	end
end
return pMis;
end

function setDobs(pMisRF::RemoteChannel,newDobs)
pMis  = take!(pMisRF)
pMis.dobs = newDobs;
put!(pMisRF,pMis)
return pMisRF;
end

function setSourcesSame(pMis::Array{RemoteChannel},newSources)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setSources,pMis[k].where,pMis[k],newSources);
		end
	end
end
return pMis;
end

function setSources(pMis::Array{RemoteChannel},newSources)
@sync begin
		for k=1:length(pMis)
	@async begin
			## TODO: make sure sources are replicated for different frequencies
			pMis[k] = remotecall_fetch(setSources,pMis[k].where,pMis[k],newSources[k]);
		end
	end
end
return pMis;
end

function setSources(pMisRF::RemoteChannel,newSources)
pMis  = take!(pMisRF)
pMis.pFor.Sources = newSources;
put!(pMisRF,pMis)
return pMisRF;
end

export getWd
function getWd(pMis::Array{RemoteChannel})
Wd = Array{Array{ComplexF64}}(undef,length(pMis))
@sync begin
		for k=1:length(pMis)
	@async begin
			Wd[k] = remotecall_fetch(getWd,pMis[k].where,pMis[k]);
		end
	end
end
return Wd;
end

function getWd(pMisRF::RemoteChannel)
pMis  = fetch(pMisRF)
return pMis.Wd;
end

export getDobs
function getDobs(pMis::Array{RemoteChannel})
Dobs = Array{Array{ComplexF64}}(undef,length(pMis))
@sync begin
		for k=1:length(pMis)
	@async begin
			Dobs[k] = remotecall_fetch(getDobs,pMis[k].where,pMis[k]);
		end
	end
end
return Dobs;
end

function getDobs(pMisRF::RemoteChannel)
pMis  = fetch(pMisRF)
return pMis.dobs;
end


text:
using Test

include("../src/components/CO2cycle.jl")

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = page_model()
    addco2cycle(m, use_permafrost)

    set_param!(m, :CO2Cycle, :e_globalCO2emissions, readpagedata(m, "test/validationdata/$valdir/e_globalCO2emissions.csv"))
    set_param!(m, :CO2Cycle, :permte_permafrostemissions, readpagedata(m, "test/validationdata/$valdir/perm_tot_e_co2.csv"))
    set_param!(m, :CO2Cycle, :y_year,[2020.,2030.,2040.,2050.,2075.,2100.,2150.,2200.,2250.,2300.])#real values
    set_param!(m, :CO2Cycle, :y_year_0, 2015.)#real value
    set_param!(m, :CO2Cycle, :rt_g_globaltemperature, readpagedata(m, "test/validationdata/$valdir/rt_g_globaltemperature.csv"))
    p=load_parameters(m)
    set_leftover_params!(m,p) #important for setting left over component values
    ##running Model
    run(m)

    te = m[:CO2Cycle, :te_totalemissions]
    te_compare = readpagedata(m, "test/validationdata/$valdir/total_e_co2.csv")
    @test te ≈ te_compare rtol=1e-4

    asymp = m[:CO2Cycle, :asymptote_co2_proj]
    asymp_compare = readpagedata(m, "test/validationdata/$valdir/asymp_comp_co2_proj.csv")
    @test asymp ≈ asymp_compare rtol=1e-4

    oceanlong = m[:CO2Cycle, :ocean_long_uptake_component_proj]
    oceanlong_compare = readpagedata(m, "test/validationdata/$valdir/ocean_long_comp_co2_proj.csv")
    @test oceanlong ≈ oceanlong_compare rtol=1e-4

    renoccff = m[:CO2Cycle, :renoccf_remainCO2wocc]
    renoccff_compare = readpagedata(m, "test/validationdata/$valdir/re_co2_no_ccff.csv")
    @test renoccff ≈ renoccff_compare rtol=1e-4

    co2conc = m[:CO2Cycle,  :c_CO2concentration]
    co2conc_compare = readpagedata(m, "test/validationdata/$valdir/c_co2concentration.csv")

    @test co2conc ≈ co2conc_compare rtol=1e-4
end

max_stars_repo_name: lrennels/PAGE-ICE

max_issues_repo_name: lrennels/PAGE-ICE

max_forks_repo_name: lrennels/PAGE-ICE


text:
@with_kw immutable Vpara @deftype Float64
    # define model parameters

    # KinematicBicycle.jl only
    ax_min =  -2.
    ax_max = 2.
    # KinematicBicycle.jl only

    m          = 2.6887e+03
    Izz        = 4.1101e+03
    la         = 1.5775        # distance from CoG to front axle
    lb         = 1.7245        # distance from CoG to rear axle
    FzF0       = 1.3680e+04    # static front axle load
    FzR0       = 1.2696e+04    # static rear axle load

    # modified to match code: OCOA151019
    #KZX        = 805.    # KZX = dFZX_Coeff = dF: longitudinal load transfer coefficient
    #KZYR       = 1130.   # KZYR = dFZYR_Coeff: lateral load transfer coefficient - rear axle
    #KZYF       = 860.    # KZYF = dFZYF_Coeff: lateral load transfer coefficient - front axle

    # modified to match code: OCOA151019
    KZX        = 806.    # KZX = dFZX_Coeff = dF: longitudinal load transfer coefficient
    KZYR       = 1076.   # KZYR = dFZYR_Coeff: lateral load transfer coefficient - rear axle
    KZYF       = 675.    # KZYF = dFZYF_Coeff: lateral load transfer coefficient - front axle

    # defines polynominal for acceleration bounds
    AXC::Array{Float64,1} = [-0.000128015180401862,	0.00858618595422724,	-0.225657108071454,	3.08283259993589,	-0.000138537090018958,	0.00684702729623608,	-0.120391102052425,	-3.55886697370079]

    # vehicle Limits
    x_min    = 0.
    x_max    = 400.
    y_min    = 0.
    y_max    = 400.
    sa_min   = -30*pi/180
    sa_max   = 30*pi/180
    psi_min  = -2*pi
    psi_max  = 2*pi
    u_min    = 0.01   #5.
    u_max    = 29.
    sr_min   = -5*pi/180
    sr_max   = 5*pi/180
    jx_min   = -5.
    jx_max   = 5.

    # tire parameters
    FZ0     = 35000.0
    PCY1    = 1.5874               #Shape factor Cfy for lateral forces
    PDY1    = 0.73957              #Lateral friction Muy
    PDY2    = -0.075004            #Variation of friction Muy with load
    PEY1    = 0.37562              #Lateral curvature Efy at Fznom
    PEY2    = -0.069325            #Variation of curvature Efy with load
    PEY3    = 0.29168              #Zero order camber dependency of curvature Efy
    PKY1    = -10.289              #Maximum value of stiffness Kfy/Fznom
    PKY2    = 3.3343               #Load at which Kfy reaches maximum value
    PHY1    = 0.0056509            #Horizontal shift Shy at Fznom
    PHY2    = -0.0020257           #Variation of shift Shy with load
    PVY1    = 0.015216             #Vertical shift in Svy/Fz at Fznom
    PVY2    = -0.010365            #Variation of shift Svy/Fz with load
    PC1     = PCY1
    PD1     = PDY1 - PDY2
    PD2     = PDY2/FZ0
    PE1     = PEY1 - PEY2
    PE2     = PEY2/FZ0
    PE3     = PEY3
    PK1     = PKY1*FZ0
    PK2     = 1/(PKY2*FZ0)
    PH1     = PHY1 - PHY2
    PH2     = PHY2/FZ0
    PV1     = PVY1 - PVY2
    PV2     = PVY2/FZ0

    Caf =  -8.4138e+04	# cornering stiffness--front axle (N/rad)
    Car =  -7.8126e+04	# cornering stiffness-- rear axle (N/rad)
    #Fyf_min = -7500;Fyf_max = 7500;
    Fy_min = -7500
    Fy_max = 7500

    # constrained initial states
    x0_     = 200.
    y0_     = 0.
    psi0_   = pi/2
    v0_     = 0.
    u0_     = 15.
    sa0_    = 0.
    sr0_    = 0.
    ax0_    = 0.
    jx0_    = 0.
    r0_     = 0.

    # leave these parameters here
    Fz_min  = 1000.
    Fz_off  = 100.
    a_t     = Fz_min + 3*Fz_off  # soft tire force constraint constants
    b_t     = Fz_off
    EP      = 0.01
end


text:
# Automatically generated using Clang.jl

const IMF_MAGIC = 20000630
const IMF_VERSION_NUMBER = 2
const IMF_INCREASING_Y = 0
const IMF_DECREASING_Y = 1
const IMF_RANDOM_Y = 2
const IMF_NO_COMPRESSION = 0
const IMF_RLE_COMPRESSION = 1
const IMF_ZIPS_COMPRESSION = 2
const IMF_ZIP_COMPRESSION = 3
const IMF_PIZ_COMPRESSION = 4
const IMF_PXR24_COMPRESSION = 5
const IMF_B44_COMPRESSION = 6
const IMF_B44A_COMPRESSION = 7
const IMF_DWAA_COMPRESSION = 8
const IMF_DWAB_COMPRESSION = 9
const IMF_WRITE_R = 0x01
const IMF_WRITE_G = 0x02
const IMF_WRITE_B = 0x04
const IMF_WRITE_A = 0x08
const IMF_WRITE_Y = 0x10
const IMF_WRITE_C = 0x20
const IMF_WRITE_RGB = 0x07
const IMF_WRITE_YC = 0x30
const IMF_WRITE_YA = 0x18
const IMF_WRITE_YCA = 0x38
const IMF_ONE_LEVEL = 0
const IMF_MIPMAP_LEVELS = 1
const IMF_RIPMAP_LEVELS = 2
const IMF_ROUND_DOWN = 0
const IMF_ROUND_UP = 1
const ImfHeader = Cvoid
const ImfOutputFile = Cvoid
const ImfTiledOutputFile = Cvoid
const ImfInputFile = Cvoid
const ImfTiledInputFile = Cvoid
const ImfLut = Cvoid


text:
doc"""
  Verilog.verilator_adapter(id)

  creates the required C file for verilation.
"""
function verilator_adapter(ident)
  modname = __global_definition_cache[ident].module_name

  """
#include <verilated.h>
#include <iostream>
#include "V$(modname).h"

V$(modname) *$(modname);
vluint64_t main_time = 0;

extern "C" void init(){
  $(modname) = new V$(modname);
}

extern "C" void step(){
  $(modname)->eval();
}

extern "C" void finish(){
  $(modname)->final();
  delete $(modname);
}

$(verilator_setter(ident))

$(verilator_getter(ident))
"""
end

function verilator_setter(ident)
  inputs = __global_definition_cache[ident].inputs
  modname = __global_definition_cache[ident].module_name
  ilist   = join(["unsigned long long $inp" for inp in inputs], ",")
  setters = join(["  $(modname)->$(inp) = $inp;" for inp in inputs], "\n")
  """
extern "C" void set($ilist){
$setters
}
"""
end

function verilator_getter(ident)
  modname = __global_definition_cache[ident].module_name
  if length(__global_definition_cache[ident].outputlist) == 1
    #extract the output symbol for this module.
    oup = first(__global_definition_cache[ident].outputlist).first
    """
extern "C" unsigned long long get(){
  return $(modname)->$(oup);
}
    """
  else
    #in the case where the result is more complex, we'll have to output a
    #struct for each of output values, which will be encoded as a 64-bit integer.
    outputs = __global_definition_cache[ident].outputlist
    olist = join(["  unsigned long long $(oup.first);" for oup in outputs], "\n")
    getters = join(["  value->$(oup.first) = $(modname)->$(oup.first);" for oup in outputs], "\n")
    """
typedef struct{
$(olist)
} output_struct;

extern "C" void get(output_struct *value){
  $getters
}
    """
  end
end


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

module DensityRatioEstimation

using StatsBase
using Statistics
using LinearAlgebra
using Parameters
using Random

# implement fit for estimators
import StatsBase: fit

# API for density ratio estimation
include("api/optlibs.jl")
include("api/estimators.jl")
include("api/fitters.jl")

# utility functions
include("utils.jl")

# available estimators
include("kmm.jl")
include("kliep.jl")
include("lsif.jl")

# available estimator fitters
include("lcv.jl")

# pure Julia implementations
include("kmm/julia.jl")

# implementations that require extra dependencies
using Requires
function __init__()
  # KMM
  @require JuMP="4076af6c-e467-56ae-b986-b466b2749572" begin
    @require Ipopt="b6b21f68-93f8-5de0-b562-5493be1d77c9" include("kmm/jump.jl")
  end

  # KLIEP
  @require Optim="429524aa-4258-5aef-a3af-852621145aeb" include("kliep/optim.jl")
  @require Convex="f65535da-76fb-5f13-bab9-19810c17039a" begin
    @require ECOS="e2685f51-7e38-5353-a97d-a921fd2c8199" include("kliep/convex.jl")
  end

  # LSIF
  @require Optim="429524aa-4258-5aef-a3af-852621145aeb" include("lsif/optim.jl")
  @require JuMP="4076af6c-e467-56ae-b986-b466b2749572" begin
    @require Ipopt="b6b21f68-93f8-5de0-b562-5493be1d77c9" include("lsif/jump.jl")
  end

  # AD and GPU libs
  @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" include("utils/zygote.jl")
  @require CuArrays="3a865a2d-5b23-5a0f-bc46-62713ec82fae" include("utils/cuarrays.jl")
end

export
  # optim libs
  OptimizationLibrary,
  JuliaLib,
  OptimLib,
  ConvexLib,
  JuMPLib,

  # estimators
  DensityRatioEstimator,
  KMM, KLIEP, LSIF,
  available_optlib,
  default_optlib,
  densratiofunc,
  densratio,

  # estimator fitters
  EstimatorFitter,
  LCV,
  fit

end # module


text:
module SpectralTimeStepping
    using ApproxFun, DifferentialEquations

    export SpectralTimeSteppingProblem

    immutable SpectralTimeSteppingProblem{S,PROB}
        problem::PROB
        space::S
    end

    function SpectralTimeSteppingProblem(f,u0,tspan)
        const sp = space(u0)
        const n = ncoefficients(u0)
        prob = ODEProblem((t,u)->pad!(f(t,Fun(sp,u)).coefficients,n),u0.coefficients,tspan)
        SpectralTimeSteppingProblem(prob,sp)
    end

    immutable SpectralTimeSteppingSolution{S,SOL}
        solution::SOL
        space::S
    end

    (sol::SpectralTimeSteppingSolution)(t) = Fun(sol.space,sol.solution(t))
    (sol::SpectralTimeSteppingSolution)(t,x) = sol(t)(x)

    DifferentialEquations.solve(prob::SpectralTimeSteppingProblem,alg,args...;kwargs...) =
        SpectralTimeSteppingSolution(solve(prob.problem,alg,args...;kwargs...),prob.space)
end    # module


text:
abstract type AbstractReceiver end

struct Receiver{
        G <: Union{TG, SMatrix{N, N, Complex{TG}}, AbstractGainPhaseMismCrosstalk{TG}} where {N, TG <: AbstractFloat},
        A <: Union{RotXYZ{<:Real}, AbstractAttitude},
        C <: Unitful.Quantity{<:Real, Unitful.𝐓}
    } <: AbstractReceiver
    sampling_frequency::typeof(1.0Hz)
    intermediate_frequency::typeof(1.0Hz)
    gain_phase_mism_crosstalk::G
    attitude::A
    n0::C
end

function Receiver(
    sampling_frequency;
    intermediate_frequency = 0.0Hz,
    gain_phase_mism_crosstalk::G = 1.0,
    attitude::A = RotXYZ(0, 0, 0),
    n0::C = 1/Hz,
) where {
    N,
    G <: Union{TG, SMatrix{N, N, Complex{TG}}, AbstractGainPhaseMismCrosstalk{TG}} where TG <: AbstractFloat,
    A <: Union{RotXYZ{<:Real}, AbstractAttitude},
    C <: Unitful.Quantity{<:Real, Unitful.𝐓}
}
    Receiver{G, A, C}(
        sampling_frequency,
        intermediate_frequency,
        gain_phase_mism_crosstalk,
        attitude,
        n0
    )
end

function propagate(receiver::Receiver, num_samples, rng)
    Δt = num_samples / get_sampling_frequency(receiver)
    gain_phase_mism_crosstalk = propagate(receiver.gain_phase_mism_crosstalk, Δt, rng)
    attitude = propagate(receiver.attitude, Δt, rng)
    Receiver(
        receiver.sampling_frequency,
        receiver.intermediate_frequency,
        gain_phase_mism_crosstalk,
        attitude,
        receiver.n0
    )
end

function get_gain_phase_mism_crosstalk(receiver::Receiver)
    get_gain_phase_mism_crosstalk(receiver.gain_phase_mism_crosstalk)
end

function get_attitude(receiver::Receiver)
    get_attitude(receiver.attitude)
end

@inline function get_sampling_frequency(receiver::Receiver)
    receiver.sampling_frequency
end

@inline function get_intermediate_frequency(receiver::Receiver)
    receiver.intermediate_frequency
end

@inline get_noise_density(receiver::AbstractReceiver) = receiver.n0
@inline function get_noise_std(receiver::AbstractReceiver)
    sqrt(get_noise_density(receiver) * get_sampling_frequency(receiver))
end


text:
# ------------------------ Air ------------------------
"""
    γ_AIR = 1.4

Adiabatic index or ratio of specific heats (dry air at 20º C).
"""
const γ_AIR = 1.4

"""
    R_AIR = 287.05287  (J/(Kg·K))

Specific gas constant for dry air.
"""
const R_AIR = 287.05287

# ------------------------ Earth ------------------------
"""
    gD = 9.80665  (m/s²)

Down component of gravity acceleration at Earth surface at 45º geodetic latitude.

1. Stevens, B. L., Lewis, F. L., & Johnson, E. N. (2015). Aircraft control and simulation:
   dynamics, controls design, and autonomous systems. John Wiley & Sons. Equation (page 33)
"""
const gD = 9.80665  # m/s²

# ------------------------ UNIT CONVERSION ------------------------
# ANGLES
const RAD2DEG = 57.29578  # Radians (rad) to degrees (deg)
const DEG2RAD = 1 / RAD2DEG
# LENGTH
const FT2M = 0.3048  # Feet (ft) to meter (m)
const M2FT = 1 / FT2M
# TEMPERATURE
const KEL2RANK = 1.8  # Kelvin to Rankine
const RANK2KEL = 1 / KEL2RANK
# MASS
const KG2LB = 2.20462262185  # Kilogram (Kg) to pound (lb)
const LB2KG = 1 / KG2LB

const SLUG2KG = 14.5939029372  # Slug to Kilogram (kg)
const KG2SLUG = 1 / SLUG2KG
# PRESSURE
const PSF2PA = LB2KG * gD / FT2M^2  # Pounds per square foot (PSF) to Pascal (Pa)
const PA2PSF = 1 / PSF2PA

const SLUGFT32KGM3 = SLUG2KG / FT2M^3  # slug/ft³ to Kg/m³
const KGM32SLUGFT3 = 1 / SLUGFT32KGM3


text:
using PlanerTransformer
using Test

function testshow(x, verified::String)
  buf = IOBuffer()
  show(buf,x)
  @test String(take!(buf)) == verified
end

include("test_FerriteProperties.jl")
include("test_transformer_construction.jl")
include("test_TransformerAnalysis.jl")


text:
using BEAST
using CompScienceMeshes
using LinearAlgebra
using SparseArrays

import Plotly
import Plots
Plots.plotly()

function Base.:+(a::S, b::S) where {S<:BEAST.Space}
    @assert geometry(a) == geometry(b)
    S(geometry(a), [a.fns; b.fns], [a.pos; b.pos])
end

l = w = 1.0 #Length and width of capacitor plates
d = 0.1     #seperation of plates
h = 1/36    #size of meshes

Γ₀ = meshrectangle(l,w,h)
Γ₁ = CompScienceMeshes.translate(Γ₀, point(0.0,0.0,d))

γ₀ = meshsegment(l,h,3)
γ₁ = CompScienceMeshes.translate(γ₀, point(0.0,0.0,d))

Γ = weld(Γ₁, Γ₀)

κ = 2pi / 100
V₀ = 1.0
f = ScalarTrace(p -> V₀)

edges_all = skeleton(Γ, 1)
edges_int = submesh(!in(boundary(Γ)), edges_all)

on_γ₀ = overlap_gpredicate(γ₀)
edges_pt0 = submesh(edges_all) do edge
    ch = chart(edges_all, edge)
    return on_γ₀(ch)
end

on_γ₁ = overlap_gpredicate(γ₁)
edges_pt1 = submesh(edges_all) do edge
    ch = chart(edges_all, edge)
    return on_γ₁(ch)
end

RT_int = raviartthomas(Γ, edges_int)
RT_pt0 = raviartthomas(Γ, edges_pt0)
RT_pt1 = raviartthomas(Γ, edges_pt1)

verts_pt0_int = submesh(!in(boundary(edges_pt0)), skeleton(edges_pt0,0))
verts_pt1_int = submesh(!in(boundary(edges_pt1)), skeleton(edges_pt1,0))

C0 = connectivity(verts_pt0_int, edges_pt0)
C1 = connectivity(verts_pt1_int, edges_pt1)

X = RT_int
Y0 = RT_pt0 * C0
Y1 = RT_pt1 * C1

# This is ugly and will fail if the two ports are not equal...
D = sparse(reshape([fill(+1.0, length(edges_pt0)); fill(-1.0, length(edges_pt1))],length(edges_pt0)+length(edges_pt1),1))
RT_pt = RT_pt0 + RT_pt1
Z = RT_pt * D

@hilbertspace x y0 y1 z
@hilbertspace ξ η0 η1 ζ

T = Maxwell3D.singlelayer(wavenumber=κ)
trc = X->ntrace(X,γ₁)

efie = @discretise( @varform(
    T[ξ,x]  + T[ξ,y0]  + T[ξ,y1]  + T[ξ,z]  +
    T[η0,x] + T[η0,y0] + T[η0,y1] + T[η0,z] +
    T[η1,x] + T[η1,y0] + T[η1,y1] + T[η1,z] +
    T[ζ,x]  + T[ζ,y0]  + T[ζ,y1]  + T[ζ,z] == f[trc(ζ)]),
    ξ∈X, η0∈Y0, η1∈Y1, ζ∈Z,
    x∈X, y0∈Y0, y1∈Y1, z∈Z)
u = solve(efie)

# You can access the current coefficients pertaining to subspaces
# using the Hilbert space 'placeholder'
@show length(u[x])
@show length(u[y0])
@show length(u[y1])
@show length(u[z])

S = X + Y0 + Y1 + Z
fcr, geo = facecurrents(u, S)
Plotly.plot(patch(geo, norm.(fcr))) |> display

# Compute the Scalar Potential across the ports. Note how a voltage
# gap of 1V comes out as a 'natural' condition on the solution.
zs = range(-0.5, stop=0.5, length=100)
pts = [point(0.5,0.5,z) for z ∈ zs]
Φ = potential(MWSingleLayerPotential3D(κ), pts, u, S, type=ComplexF64)
Plots.plot(zs, real(Φ), xlabel="height",ylabel="scalar potential",label=false) |> display

# Plot current along γ₀ and γ₁. Note: in this symmetric situation, we
# expect these to be opposite on corresponding points of the two ports.
# In general this is not true; in fact, the two ports could have different shape and size
traceS0 = ntrace(S, γ₀)
traceS1 = ntrace(S, γ₁)
fcr0, geo0 = facecurrents(u, traceS0)
fcr1, geo1 = facecurrents(u, traceS1)
plot()
plot!(imag.(fcr0))
plot!(imag.(fcr1)) |> display


text:
function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
    precompile(PDMats.call, (Type{PDMats.PDMat}, Array{Float64, 2}, Base.LinAlg.Cholesky{Float64, Array{Float64, 2}},))
    precompile(PDMats.call, (Type{PDMats.PDiagMat}, Array{Float64, 1}, Array{Float64, 1},))
    precompile(PDMats.full, (PDMats.PDiagMat{Float64, Array{Float64, 1}},))
end

max_stars_repo_name: waTeim/gpsFlow

max_issues_repo_name: waTeim/flying-monkey

max_forks_repo_name: waTeim/flying-monkey


text:
include("dependencies_for_runtests.jl")

using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary, GridFittedBottom
using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization

@inline surface_wind_stress(λ, φ, t, p) = p.τ₀ * cos(2π * (φ - p.φ₀) / p.Lφ)
@inline u_bottom_drag(i, j, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, 1]
@inline v_bottom_drag(i, j, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, 1]

@testset "Immersed boundaries with hydrostatic free surface models" begin
    @info "Testing immersed boundaries with hydrostatic free surface models..."

    for arch in archs

        arch_str = string(typeof(arch))

        @testset "GridFittedBoundary [$arch_str]" begin
            @info "Testing GridFittedBoundary with HydrostaticFreeSurfaceModel [$arch_str]..."

            underlying_grid = RectilinearGrid(arch, size=(8, 8, 8), x = (-5, 5), y = (-5, 5), z = (0, 2))

            bump(x, y, z) = z < exp(-x^2 - y^2)
            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBoundary(bump))

            for closure in (IsotropicDiffusivity(ν=1, κ=0.5),
                            IsotropicDiffusivity(ν=1, κ=0.5, time_discretization=VerticallyImplicitTimeDiscretization()))

                model = HydrostaticFreeSurfaceModel(grid = grid,
                                                    tracers = :b,
                                                    buoyancy = BuoyancyTracer(),
                                                    closure = closure)

                u = model.velocities.u
                b = model.tracers.b

                # Linear stratification
                set!(model, u = 1, b = (x, y, z) -> 4 * z)

                # Inside the bump
                @test b[4, 4, 2] == 0
                @test u[4, 4, 2] == 0

                simulation = Simulation(model, Δt = 1e-3, stop_iteration=2)

                run!(simulation)

                # Inside the bump
                @test b[4, 4, 2] == 0
                @test u[4, 4, 2] == 0
            end
        end

        @testset "Surface boundary conditions with immersed boundaries [$arch_str]" begin
            @info "  Testing surface boundary conditions with ImmersedBoundaries in HydrostaticFreeSurfaceModel [$arch_str]..."

            Nx = 60
            Ny = 60

            # A spherical domain
            underlying_grid = LatitudeLongitudeGrid(arch,
                                                    size = (Nx, Ny, 1),
                                                    longitude = (-30, 30),
                                                    latitude = (15, 75),
                                                    z = (-4000, 0))

            bathymetry = zeros(Nx, Ny) .- 4000
            view(bathymetry, 31:34, 43:47) .= 0
            bathymetry = arch_array(arch, bathymetry)

            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bathymetry))

            free_surface = ImplicitFreeSurface(gravitational_acceleration=0.1)
            coriolis = HydrostaticSphericalCoriolis(scheme = VectorInvariantEnstrophyConserving())

            surface_wind_stress_parameters = (τ₀ = 1e-4,
                                              Lφ = grid.Ly,
                                              φ₀ = 15)

            surface_wind_stress_bc = FluxBoundaryCondition(surface_wind_stress,
                                                           parameters = surface_wind_stress_parameters)

            μ = 1 / 60days

            u_bottom_drag_bc = FluxBoundaryCondition(u_bottom_drag,
                                                     discrete_form = true,
                                                     parameters = μ)

            v_bottom_drag_bc = FluxBoundaryCondition(v_bottom_drag,
                                                     discrete_form = true,
                                                     parameters = μ)

            u_bcs = FieldBoundaryConditions(top = surface_wind_stress_bc, bottom = u_bottom_drag_bc)
            v_bcs = FieldBoundaryConditions(bottom = v_bottom_drag_bc)

            νh₀ = 5e3 * (60 / grid.Nx)^2
            constant_horizontal_diffusivity = HorizontallyCurvilinearAnisotropicDiffusivity(νh=νh₀)

            model = HydrostaticFreeSurfaceModel(; grid,
                                                momentum_advection = VectorInvariant(),
                                                free_surface = free_surface,
                                                coriolis = coriolis,
                                                boundary_conditions = (u=u_bcs, v=v_bcs),
                                                closure = constant_horizontal_diffusivity,
                                                tracers = nothing,
                                                buoyancy = nothing)

            simulation = Simulation(model, Δt=3600, stop_iteration=1)

            run!(simulation)

            # If reached here it didn't error, so pass for now!
            @test true
        end

        @testset "Correct vertically-integrated lateral face areas with immersed boundaries [$arch_str]" begin
            @info "  Testing correct vertically-integrated lateral face areas with immersed boundaries [$arch_str]..."

            Nx = 5
            Ny = 5

            underlying_grid = RectilinearGrid(arch,
                                              size = (Nx, Ny, 3),
                                              extent = (Nx, Ny, 3),
                                              topology = (Periodic, Periodic, Bounded))

            # B for bathymetry
            B = [-3. for i=1:Nx, j=1:Ny ]
            B[2:Nx-1, 2:Ny-1] .= [-2 for i=2:Nx-1, j=2:Ny-1 ]
            B[3:Nx-2, 3:Ny-2] .= [-1 for i=3:Nx-2, j=3:Ny-2 ]

            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(B))

            model = HydrostaticFreeSurfaceModel(; grid,
                                                free_surface = ImplicitFreeSurface(),
                                                buoyancy = nothing,
                                                tracers = nothing,
                                                closure = nothing)

            x_ref = [0.0  0.0  0.0  0.0  0.0  0.0  0.0
                     0.0  3.0  3.0  3.0  3.0  3.0  0.0
                     0.0  3.0  2.0  2.0  2.0  2.0  0.0
                     0.0  3.0  2.0  1.0  1.0  2.0  0.0
                     0.0  3.0  2.0  2.0  2.0  2.0  0.0
                     0.0  3.0  3.0  3.0  3.0  3.0  0.0
                     0.0  0.0  0.0  0.0  0.0  0.0  0.0]'

            y_ref = [0.0  0.0  0.0  0.0  0.0  0.0  0.0
                     0.0  3.0  3.0  3.0  3.0  3.0  0.0
                     0.0  3.0  2.0  2.0  2.0  3.0  0.0
                     0.0  3.0  2.0  1.0  2.0  3.0  0.0
                     0.0  3.0  2.0  1.0  2.0  3.0  0.0
                     0.0  3.0  2.0  2.0  2.0  3.0  0.0
                     0.0  0.0  0.0  0.0  0.0  0.0  0.0]'

            fs = model.free_surface
            vertically_integrated_lateral_areas = fs.implicit_step_solver.vertically_integrated_lateral_areas

            ∫Axᶠᶜᶜ = vertically_integrated_lateral_areas.xᶠᶜᶜ
            ∫Ayᶜᶠᶜ = vertically_integrated_lateral_areas.yᶜᶠᶜ

            ∫Axᶠᶜᶜ = Array(parent(∫Axᶠᶜᶜ))
            ∫Ayᶜᶠᶜ = Array(parent(∫Ayᶜᶠᶜ))

            Ax_ok = ∫Axᶠᶜᶜ[:, :, 1] ≈ x_ref
            Ay_ok = ∫Ayᶜᶠᶜ[:, :, 1] ≈ y_ref

            @test (Ax_ok & Ay_ok)
        end
    end
end


text:
using DataFrames
using CSV
using StaticArrays
const ℝ = SVector{N,T} where {N,T}

"""
    readData(::Val{true}, filename)

Reads first passage time data from a file `filename` (assumed to be saved in a
correct format---no checks are performed).
"""
function readData(::Val{true}, filename)
    df = CSV.read(filename)
    x0 = ℝ{2}(df.upCross[1], df.x2[1])
    obs = ℝ{1}.(df.upCross)
    obsTime = Float64.(df.time)
    fpt = [FPTInfo((1,), (true,), (resetLvl,), (i==1,)) for
            (i, resetLvl) in enumerate(df.downCross[2:end])]
    fptOrPartObs = FPT()
    df, x0, obs, obsTime, fpt, fptOrPartObs
end

"""
    readData(::Val{false}, filename)

Reads partial observation type data from a file `filename` (assumed to be saved
in a correct format---no checks are performed)
"""
function readData(::Val{false}, filename)
    df = CSV.read(filename)
    obs = ℝ{1}.(df.x1)
    obsTime = Float64.(df.time)
    x0 = ℝ{2}(df.x1[1], df.x2[1])
    fpt = [NaN for _ in obsTime[2:end]]
    fptOrPartObs = PartObs()
    df, x0, obs, obsTime, fpt, fptOrPartObs
end

"""
    savePathsToFile(paths, time, filename)

Put diffusion paths and respective times to a dataframe and save to a file
"""
function savePathsToFile(paths, time, filename)
    idx = collect( Iterators.flatten([fill(i,length(X))
                            for (i,X) in enumerate(paths)]))
    time = collect(Iterators.flatten(fill(time,length(paths))))
    paths = collect(Iterators.flatten(paths))
    numDims = length(paths[1])
    df = DataFrame(idx=idx, time=time)
    for i in 1:numDims
        X = [pt[i] for pt in paths]
        df = hcat(df, X, makeunique=true)
    end
    CSV.write(filename, df)
    df
end

"""
    saveChainToFile(chains, filename)

Save MCMC chain to a file
"""
function saveChainToFile(chains, filename)
    df = DataFrame()
    for i in 1:length(chains[1])
        X = [pt[i] for pt in chains]
        df = hcat(df, X, makeunique=true)
    end
    CSV.write(filename, df)
    df
end


text:
@testset "crcfit with given h and q" begin

    h = [1/4, 1/3, 1/2, 2,3,4]
    q = [1/32, 2/27, 1/4, 8, 18, 32]

    crc = crcfit(h,q)

    rc₁ = crc.component[1]
    rc₂ = crc.component[2]

    @test crc.threshold[] ≈ 1.0

    @test rc₁.a ≈ 2.0
    @test rc₁.b ≈ 0.0 atol=1e-8
    @test rc₁.c ≈ 3.0

    @test rc₂.a ≈ 2.0
    @test rc₂.b ≈ 0.0 atol=1e-8
    @test rc₂.c ≈ 2.0
end

@testset "rcfit with given h and q" begin
    h = [2,3,4]
    q = [3,12,27]

    rc = rcfit(h,q)

    @test rc.gauging == Gauging.(h,q)
    @test rc.a ≈ 3.0
    @test rc.b ≈ 1.0
    @test rc.c ≈ 2.0

end


text:
# # Perfect CAKTE calibration with Ensemble Kalman Inversion

# ## Install dependencies

# ```julia
# using Pkg
# pkg"add ParameterEstimocean, Oceananigans, Distributions, CairoMakie"
# ```

using ParameterEstimocean, LinearAlgebra, CairoMakie

using ParameterEstimocean.Transformations: Transformation
using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: CATKEVerticalDiffusivity, MixingLength, SurfaceTKEFlux

# # Perfect observations of CATKE-driven mixing
#
# Our first task is to generate synthetic observations, using
# a one-dimensional model driven by surface fluxes and with
# turbulent mixing parameterized by CATKE. We use a simplified CATKE
# with no stability function (by setting `Cᴷuʳ = Cᴷcʳ = Cᴷeʳ = 0`)
# and "reasonable", but unrealistic parameters.
# We will only attempt to calibrate a subset of the parameters that we set
# to generate the observations.

## Load utilities
examples_path = joinpath(pathof(ParameterEstimocean), "..", "..", "examples")
include(joinpath(examples_path, "intro_to_inverse_problems.jl"))

mixing_length = MixingLength(Cᴬu  = 0.0,
                             Cᴬc  = 1.0,
                             Cᴬe  = 0.0,
                             Cᴷu⁻ = 0.1,
                             Cᴷc⁻ = 0.1,
                             Cᴷe⁻ = 0.1,
                             Cᴷuʳ = 0.0,
                             Cᴷcʳ = 0.0,
                             Cᴷeʳ = 0.0)

catke = CATKEVerticalDiffusivity(mixing_length=mixing_length)

## Specify both wind mixing and convection:
data_path = generate_synthetic_observations("catke",
                                            closure = catke,
                                            tracers = (:b, :e),
                                            Nz = 32,
                                            Lz = 64,
                                            Δt = 10.0,
                                            stop_time = 12hours,
                                            overwrite = true,
                                            Qᵘ = -1e-4,
                                            Qᵇ = 1e-8,
                                            N² = 1e-5)

# Next, we load and inspect the observations to make sure they're sensible:

transformation = (u = ZScore(),
                  v = ZScore(),
                  b = ZScore(),
                  e = RescaledZScore(1e-1))

observations = SyntheticObservations(data_path; field_names=(:u, :v, :b, :e), transformation)

fig = Figure()

ax_b = Axis(fig[1, 1], xlabel = "Buoyancy\n[10⁻⁴ m s⁻²]", ylabel = "z [m]")
ax_u = Axis(fig[1, 2], xlabel = "Velocities\n[cm s⁻¹]")
ax_e = Axis(fig[1, 3], xlabel = "Turbulent kinetic energy\n[cm² s⁻²]")

z = znodes(Center, observations.grid)

colorcycle = [:black, :red, :blue, :orange, :pink]

for i = 1:length(observations.times)
    b = observations.field_time_serieses.b[i]
    e = observations.field_time_serieses.e[i]
    u = observations.field_time_serieses.u[i]
    v = observations.field_time_serieses.v[i]
    t = observations.times[i]

    label = "t = " * prettytime(t)
    u_label = i == 1 ? "u, " * label : label
    v_label = i == 1 ? "v, " * label : label
    ## Note unit conversions below, eg m s⁻² -> 10⁻⁴ m s⁻²
    lines!(ax_b, 1e4 * interior(b)[1, 1, :], z; label, color=colorcycle[i])
    lines!(ax_u, 1e2 * interior(u)[1, 1, :], z; linestyle=:solid, color=colorcycle[i], label=u_label)
    lines!(ax_u, 1e2 * interior(v)[1, 1, :], z; linestyle=:dash, color=colorcycle[i], label=v_label)
    lines!(ax_e, 1e4 * interior(e)[1, 1, :], z; label, color=colorcycle[i])
end

axislegend(ax_b, position=:rb)
axislegend(ax_u, position=:lb, merge=true)
axislegend(ax_e, position=:rb)

##display(fig)

save("synthetic_catke_observations.svg", fig); nothing # hide

# ![](synthetic_catke_observations.svg)

# Well, that looks like a boundary layer, in some respects.
#
# # Calibration
#
# Next, we build a simulation of an ensemble of column models to calibrate
# CATKE using Ensemble Kalman Inversion.

architecture = CPU()
ensemble_simulation, closure★ = build_ensemble_simulation(observations, architecture; Nensemble=20)

# We choose to calibrate a subset of the CATKE parameters,

priors = (Cᴬu = lognormal(mean=0.05, std=0.05),
          Cᴬc = lognormal(mean=0.8,  std=0.1),
          Cᴬe = lognormal(mean=0.1,  std=0.05))

free_parameters = FreeParameters(priors)

# The handy utility function `build_ensemble_simulation` also tells us the optimal
# parameters that were used when generating the synthetic observations:

@show θ★ = (Cᴬu = closure★.mixing_length.Cᴬu,
            Cᴬc = closure★.mixing_length.Cᴬc,
            Cᴬe = closure★.mixing_length.Cᴬe)

# We construct the `InverseProblem` from `observations`, `ensemble_simulation`, and
# `free_parameters`,

calibration = InverseProblem(observations, ensemble_simulation, free_parameters)

# We can check that the first ensemble member of the mapped output, which was run with
# the "true" # parameters, is identical to the mapped observations:

G = forward_map(calibration, θ★)
y = observation_map(calibration)

@show G[:, 1] ≈ y

# # Ensemble Kalman Inversion
#
# Next, we construct an `EnsembleKalmanInversion` (EKI) object,
#
# The calibration is done here using Ensemble Kalman Inversion. For more information about the
# algorithm refer to
# [EnsembleKalmanProcesses.jl documentation](
# https://clima.github.io/EnsembleKalmanProcesses.jl/stable/ensemble_kalman_inversion/).

eki = EnsembleKalmanInversion(calibration;
                              noise_covariance = 1e-3,
                              resampler = Resampler(acceptable_failure_fraction=0.3))

# and perform few iterations to see if we can converge to the true parameter values.

iterate!(eki; iterations = 10)

# Last, we visualize the outputs of EKI calibration.

## Convert everything to a vector
optimal_θ = collect(values(θ★))
ensemble_mean_θ = map(summary -> collect(values(summary.ensemble_mean)), eki.iteration_summaries)
θ_variances = map(summary -> collect(values(summary.ensemble_var)), eki.iteration_summaries)

absolute_error = NamedTuple(name => map(θ -> θ[p] - θ★[p], ensemble_mean_θ) for (p, name) in enumerate(free_parameters.names))
relative_error = NamedTuple(name => abs.(absolute_error[name]) ./ θ★[name] for name in free_parameters.names)

output_distances = map(θ -> norm(forward_map(calibration, θ)[:, 1:1] - y), ensemble_mean_θ)

fig = Figure()

ax_error = Axis(fig[1, 1], title = "Parameter distance", xlabel = "Iteration", ylabel = "|⟨θₙ⟩ - θ★| / θ★")

for name in free_parameters.names
    lines!(ax_error, 0:eki.iteration, parent(relative_error[name]), linewidth=2, label=string(name))
end

axislegend(ax_error, position=:rt)

lines(fig[1, 2], 0:eki.iteration, parent(output_distances), color = :blue, linewidth = 2,
      axis = (title = "Output distance", xlabel = "Iteration", ylabel = "|G(⟨θₙ⟩) - y|"))

ax3 = Axis(fig[2, 1:2], title = "Parameter convergence", xlabel = "Iteration",
           ylabel = "Relative change ensemble variance", yscale = log10)

for (p, name) in enumerate(free_parameters.names)
    θp_variances = [θ_variances[iter][p] for iter = 0:eki.iteration]
    lines!(ax3, 0:eki.iteration, parent(θp_variances / θp_variances[1]), label = String(name), linewidth = 2)
end

axislegend(ax3, valign = :top, halign = :right)

##display(fig)

save("perfect_catke_calibration_summary.svg", fig); nothing #hide

# ![](perfect_catke_calibration_summary.svg)

final_mean_θ = eki.iteration_summaries[end].ensemble_mean
forward_run!(calibration, [θ★, final_mean_θ])

time_series_collector = calibration.time_series_collector
times = time_series_collector.times

## Extract last save point and plot each solution component
Nt = length(times)

b = time_series_collector.field_time_serieses.b[Nt]
e = time_series_collector.field_time_serieses.e[Nt]
u = time_series_collector.field_time_serieses.u[Nt]
v = time_series_collector.field_time_serieses.v[Nt]

t = times[Nt]
z = znodes(b)

## The ensemble varies along the first, or `x`-dimension:
b★ = 1e4 * interior(b)[1, 1, :]  # convert units m s⁻² -> 10⁻⁴ m s⁻²
b¹ = 1e4 * interior(b)[2, 1, :]  # convert units m s⁻² -> 10⁻⁴ m s⁻²

e★ = 1e4 * interior(e)[1, 1, :]  # convert units m² s⁻² -> cm² s⁻²
e¹ = 1e4 * interior(e)[2, 1, :]  # convert units m² s⁻² -> cm² s⁻²

u★ = 1e2 * interior(u)[1, 1, :]  # convert units m s⁻¹ -> cm s⁻¹
u¹ = 1e2 * interior(u)[2, 1, :]  # convert units m s⁻¹ -> cm s⁻¹

v★ = 1e2 * interior(v)[1, 1, :]  # convert units m s⁻¹ -> cm s⁻¹
v¹ = 1e2 * interior(v)[2, 1, :]  # convert units m s⁻¹ -> cm s⁻¹

fig = Figure()

ax = Axis(fig[1, 1], xlabel = "Buoyancy\n[10⁻⁴ m s⁻²]", ylabel = "z [m]")
b★_label = "true b at " * prettytime(t)
b¹_label = "b with ⟨θ⟩"
lines!(ax, b★, z; label=b★_label, linewidth=3)
lines!(ax, b¹, z; label=b¹_label, linewidth=2)
axislegend(ax, position=:lb)

ax = Axis(fig[1, 2], xlabel = "Turbulent kinetic energy\n[cm² s⁻²]")
e★_label = "true e at " * prettytime(t)
e¹_label = "e with ⟨θ⟩"
lines!(ax, e★, z; label=e★_label, linewidth=3)
lines!(ax, e¹, z; label=e¹_label, linewidth=2)
axislegend(ax, position=:lb)

ax = Axis(fig[1, 3], xlabel = "Velocities\n[cm s⁻¹]")
u★_label = "true u at " * prettytime(t)
u¹_label = "u with ⟨θ⟩"
v★_label = "true v"
v¹_label = "v with ⟨θ⟩"
lines!(ax, u★, z; label=u★_label, linewidth=3)
lines!(ax, u¹, z; label=u¹_label, linewidth=2)
lines!(ax, v★, z; label=v★_label, linestyle=:dash, linewidth=3)
lines!(ax, v¹, z; label=v¹_label, linestyle=:dash, linewidth=2)
axislegend(ax, position=:lb)

save("perfect_catke_calibration_particle_realizations.svg", fig); nothing # hide

# ![](perfect_catke_calibration_particle_realizations.svg)

##display(fig)

# And also we plot the the distributions of the various model ensembles for few EKI iterations to see
# if and how well they converge to the true diffusivity values.

fig = Figure()

ax1 = Axis(fig[1, 1])
ax2 = Axis(fig[2, 1], xlabel = "Cᴬu", ylabel = "Cᴬc")
ax3 = Axis(fig[2, 2])
scatters = []
labels = String[]

for iteration in [0, 1, 5, 10]
    ## Make parameter matrix
    parameters = eki.iteration_summaries[iteration].parameters
    Nensemble = length(parameters)
    parameter_ensemble_matrix = [parameters[i][j] for i=1:Nensemble, j=1:2]

    label = iteration == 0 ? "Initial ensemble" : "Iteration $iteration"
    push!(labels, label)
    push!(scatters, scatter!(ax2, parameter_ensemble_matrix))
    density!(ax1, parameter_ensemble_matrix[:, 1])
    density!(ax3, parameter_ensemble_matrix[:, 2], direction = :y)
end

vlines!(ax1, [θ★.Cᴬu], color = :red)
vlines!(ax2, [θ★.Cᴬu], color = :red)
hlines!(ax2, [θ★.Cᴬc], color = :red)
hlines!(ax3, [θ★.Cᴬc], color = :red)

colsize!(fig.layout, 1, Fixed(300))
colsize!(fig.layout, 2, Fixed(200))
rowsize!(fig.layout, 1, Fixed(200))
rowsize!(fig.layout, 2, Fixed(300))

Legend(fig[1, 2], scatters, labels, valign = :bottom, halign = :left)

hidedecorations!(ax1, grid = false)
hidedecorations!(ax3, grid = false)

##display(fig)

save("perfect_catke_calibration_parameter_distributions.svg", fig); nothing # hide

# ![](perfect_catke_calibration_parameter_distributions.svg)

# Hint: if using a REPL or notebook, try
# `using Pkg; Pkg.add("ElectronDisplay"); using ElectronDisplay; display(fig)`
# To see the figure in a window.


text:
using Libdl
const libs = (PETSc_jll.libpetsc, )

function initialize(libhdl::Ptr{Cvoid})
  PetscInitializeNoArguments_ptr = dlsym(libhdl, :PetscInitializeNoArguments)
  @chk ccall(PetscInitializeNoArguments_ptr, PetscErrorCode, ())
end

function DataTypeFromString(libhdl::Ptr{Cvoid}, name::AbstractString)
    PetscDataTypeFromString_ptr = dlsym(libhdl, :PetscDataTypeFromString)
    dtype_ref = Ref{PetscDataType}()
    found_ref = Ref{PetscBool}()
    @chk ccall(PetscDataTypeFromString_ptr, PetscErrorCode,
             (Cstring, Ptr{PetscDataType}, Ptr{PetscBool}),
             name, dtype_ref, found_ref)
    @assert found_ref[] == PETSC_TRUE
    return dtype_ref[]
end
function PetscDataTypeGetSize(libhdl::Ptr{Cvoid}, dtype::PetscDataType)
    PetscDataTypeGetSize_ptr = dlsym(libhdl, :PetscDataTypeGetSize)
    datasize_ref = Ref{Csize_t}()
    @chk ccall(PetscDataTypeGetSize_ptr, PetscErrorCode,
             (PetscDataType, Ptr{Csize_t}),
             dtype, datasize_ref)
    return datasize_ref[]
end

const libtypes = map(libs) do lib
    libhdl = dlopen(lib)
    initialize(libhdl)
    PETSC_REAL = DataTypeFromString(libhdl, "Real")
    PETSC_SCALAR = DataTypeFromString(libhdl, "Scalar")
    PETSC_INT_SIZE = PetscDataTypeGetSize(libhdl, PETSC_INT)

    PetscReal =
        PETSC_REAL == PETSC_DOUBLE ? Cdouble :
        PETSC_REAL == PETSC_FLOAT ? Cfloat :
        error("PETSC_REAL = $PETSC_REAL not supported.")

    PetscScalar =
        PETSC_SCALAR == PETSC_REAL ? PetscReal :
        PETSC_SCALAR == PETSC_COMPLEX ? Complex{PetscReal} :
        error("PETSC_SCALAR = $PETSC_SCALAR not supported.")

    PetscInt =
        PETSC_INT_SIZE == 4 ? Int32 :
        PETSC_INT_SIZE == 8 ? Int64 :
        error("PETSC_INT_SIZE = $PETSC_INT_SIZE not supported.")

    # TODO: PetscBLASInt, PetscMPIInt ?
    return (lib, PetscScalar, PetscReal, PetscInt)
end

const scalar_types = map(x -> x[2], libtypes)

macro for_libpetsc(expr)
  quote
    for (libpetsc, PetscScalar, PetscReal, PetscInt) in libtypes
      @eval esc($expr)
    end
  end
end

@for_libpetsc inttype(::Type{$PetscScalar}) = $PetscInt


text:
module TwoBit

import BioGenerics
import BioSequences

include("record.jl")
include("reader.jl")
include("writer.jl")

end # module


text:
using Distributions: UnivariateDistribution,
                     MultivariateDistribution,
                     MatrixDistribution,
                     Distribution

const AMBIGUITY_MSG = "Ambiguous `LHS .~ RHS` or `@. LHS ~ RHS` syntax. The broadcasting " *
    "can either be column-wise following the convention of Distributions.jl or " *
    "element-wise following Julia's general broadcasting semantics. Please make sure " *
    "that the element type of `LHS` is not a supertype of the support type of " *
    "`AbstractVector` to eliminate ambiguity."

alg_str(spl::Sampler) = string(nameof(typeof(spl.alg)))

# utility funcs for querying sampler information
require_gradient(spl::Sampler) = false
require_particles(spl::Sampler) = false

_getindex(x, inds::Tuple) = _getindex(x[first(inds)...], Base.tail(inds))
_getindex(x, inds::Tuple{}) = x

# assume
function tilde(rng, ctx::DefaultContext, sampler, right, vn::VarName, _, vi)
    return _tilde(rng, sampler, right, vn, vi)
end
function tilde(rng, ctx::PriorContext, sampler, right, vn::VarName, inds, vi)
    if ctx.vars !== nothing
        vi[vn] = vectorize(right, _getindex(getfield(ctx.vars, getsym(vn)), inds))
        settrans!(vi, false, vn)
    end
    return _tilde(rng, sampler, right, vn, vi)
end
function tilde(rng, ctx::LikelihoodContext, sampler, right, vn::VarName, inds, vi)
    if ctx.vars isa NamedTuple && haskey(ctx.vars, getsym(vn))
        vi[vn] = vectorize(right, _getindex(getfield(ctx.vars, getsym(vn)), inds))
        settrans!(vi, false, vn)
    end
    return _tilde(rng, sampler, NoDist(right), vn, vi)
end
function tilde(rng, ctx::MiniBatchContext, sampler, right, left::VarName, inds, vi)
    return tilde(rng, ctx.ctx, sampler, right, left, inds, vi)
end

"""
    tilde_assume(rng, ctx, sampler, right, vn, inds, vi)

Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
accumulate the log probability, and return the sampled value.

Falls back to `tilde(rng, ctx, sampler, right, vn, inds, vi)`.
"""
function tilde_assume(rng, ctx, sampler, right, vn, inds, vi)
    value, logp = tilde(rng, ctx, sampler, right, vn, inds, vi)
    acclogp!(vi, logp)
    return value
end

function _tilde(rng, sampler, right, vn::VarName, vi)
    return assume(rng, sampler, right, vn, vi)
end
function _tilde(rng, sampler, right::NamedDist, vn::VarName, vi)
    return _tilde(rng, sampler, right.dist, right.name, vi)
end

# observe
function tilde(ctx::DefaultContext, sampler, right, left, vi)
    return _tilde(sampler, right, left, vi)
end
function tilde(ctx::PriorContext, sampler, right, left, vi)
    return 0
end
function tilde(ctx::LikelihoodContext, sampler, right, left, vi)
    return _tilde(sampler, right, left, vi)
end
function tilde(ctx::MiniBatchContext, sampler, right, left, vi)
    return ctx.loglike_scalar * tilde(ctx.ctx, sampler, right, left, vi)
end

"""
    tilde_observe(ctx, sampler, right, left, vname, vinds, vi)

Handle observed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
accumulate the log probability, and return the observed value.

Falls back to `tilde(ctx, sampler, right, left, vi)` ignoring the information about variable name
and indices; if needed, these can be accessed through this function, though.
"""
function tilde_observe(ctx, sampler, right, left, vname, vinds, vi)
    logp = tilde(ctx, sampler, right, left, vi)
    acclogp!(vi, logp)
    return left
end

"""
    tilde_observe(ctx, sampler, right, left, vi)

Handle observed constants, e.g., `1.0 ~ Normal()`, accumulate the log probability, and
return the observed value.

Falls back to `tilde(ctx, sampler, right, left, vi)`.
"""
function tilde_observe(ctx, sampler, right, left, vi)
    logp = tilde(ctx, sampler, right, left, vi)
    acclogp!(vi, logp)
    return left
end

_tilde(sampler, right, left, vi) = observe(sampler, right, left, vi)

function assume(rng, spl::Sampler, dist)
    error("DynamicPPL.assume: unmanaged inference algorithm: $(typeof(spl))")
end

function observe(spl::Sampler, weight)
    error("DynamicPPL.observe: unmanaged inference algorithm: $(typeof(spl))")
end

function assume(
    rng,
    spl::Union{SampleFromPrior,SampleFromUniform},
    dist::Distribution,
    vn::VarName,
    vi,
)
    if haskey(vi, vn)
        # Always overwrite the parameters with new ones for `SampleFromUniform`.
        if spl isa SampleFromUniform || is_flagged(vi, vn, "del")
            unset_flag!(vi, vn, "del")
            r = init(rng, dist, spl)
            vi[vn] = vectorize(dist, r)
            settrans!(vi, false, vn)
            setorder!(vi, vn, get_num_produce(vi))
        else
            r = vi[vn]
        end
    else
        r = init(rng, dist, spl)
        push!(vi, vn, r, dist, spl)
        settrans!(vi, false, vn)
    end
    return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn))
end

function observe(
    spl::Union{SampleFromPrior, SampleFromUniform},
    dist::Distribution,
    value,
    vi,
)
    increment_num_produce!(vi)
    return Distributions.loglikelihood(dist, value)
end

# .~ functions

# assume
function dot_tilde(rng, ctx::DefaultContext, sampler, right, left, vn::VarName, _, vi)
    vns, dist = get_vns_and_dist(right, left, vn)
    return _dot_tilde(rng, sampler, dist, left, vns, vi)
end
function dot_tilde(
    rng,
    ctx::LikelihoodContext,
    sampler,
    right,
    left,
    vn::VarName,
    inds,
    vi,
)
    if ctx.vars isa NamedTuple && haskey(ctx.vars, getsym(vn))
        var = _getindex(getfield(ctx.vars, getsym(vn)), inds)
        vns, dist = get_vns_and_dist(right, var, vn)
        set_val!(vi, vns, dist, var)
        settrans!.(Ref(vi), false, vns)
    else
        vns, dist = get_vns_and_dist(right, left, vn)
    end
    return _dot_tilde(rng, sampler, NoDist.(dist), left, vns, vi)
end
function dot_tilde(rng, ctx::MiniBatchContext, sampler, right, left, vn::VarName, inds, vi)
    return dot_tilde(rng, ctx.ctx, sampler, right, left, vn, inds, vi)
end
function dot_tilde(
    rng,
    ctx::PriorContext,
    sampler,
    right,
    left,
    vn::VarName,
    inds,
    vi,
)
    if ctx.vars !== nothing
        var = _getindex(getfield(ctx.vars, getsym(vn)), inds)
        vns, dist = get_vns_and_dist(right, var, vn)
        set_val!(vi, vns, dist, var)
        settrans!.(Ref(vi), false, vns)
    else
        vns, dist = get_vns_and_dist(right, left, vn)
    end
    return _dot_tilde(rng, sampler, dist, left, vns, vi)
end

"""
    dot_tilde_assume(rng, ctx, sampler, right, left, vn, inds, vi)

Handle broadcasted assumed variables, e.g., `x .~ MvNormal()` (where `x` does not occur in the
model inputs), accumulate the log probability, and return the sampled value.

Falls back to `dot_tilde(rng, ctx, sampler, right, left, vn, inds, vi)`.
"""
function dot_tilde_assume(rng, ctx, sampler, right, left, vn, inds, vi)
    value, logp = dot_tilde(rng, ctx, sampler, right, left, vn, inds, vi)
    acclogp!(vi, logp)
    return value
end

function get_vns_and_dist(dist::NamedDist, var, vn::VarName)
    return get_vns_and_dist(dist.dist, var, dist.name)
end
function get_vns_and_dist(dist::MultivariateDistribution, var::AbstractMatrix, vn::VarName)
    getvn = i -> VarName(vn, (vn.indexing..., (Colon(), i)))
    return getvn.(1:size(var, 2)), dist

end
function get_vns_and_dist(
    dist::Union{Distribution, AbstractArray{<:Distribution}},
    var::AbstractArray,
    vn::VarName
)
    getvn = ind -> VarName(vn, (vn.indexing..., Tuple(ind)))
    return getvn.(CartesianIndices(var)), dist
end

function _dot_tilde(rng, sampler, right, left, vns::AbstractArray{<:VarName}, vi)
    return dot_assume(rng, sampler, right, vns, left, vi)
end

# Ambiguity error when not sure to use Distributions convention or Julia broadcasting semantics
function _dot_tilde(
    rng,
    sampler::AbstractSampler,
    right::Union{MultivariateDistribution, AbstractVector{<:MultivariateDistribution}},
    left::AbstractMatrix{>:AbstractVector},
    vn::AbstractVector{<:VarName},
    vi,
)
    throw(DimensionMismatch(AMBIGUITY_MSG))
end

function dot_assume(
    rng,
    spl::Union{SampleFromPrior, SampleFromUniform},
    dist::MultivariateDistribution,
    vns::AbstractVector{<:VarName},
    var::AbstractMatrix,
    vi,
)
    @assert length(dist) == size(var, 1)
    r = get_and_set_val!(rng, vi, vns, dist, spl)
    lp = sum(Bijectors.logpdf_with_trans(dist, r, istrans(vi, vns[1])))
    var .= r
    return var, lp
end
function dot_assume(
    rng,
    spl::Union{SampleFromPrior, SampleFromUniform},
    dists::Union{Distribution, AbstractArray{<:Distribution}},
    vns::AbstractArray{<:VarName},
    var::AbstractArray,
    vi,
)
    r = get_and_set_val!(rng, vi, vns, dists, spl)
    # Make sure `r` is not a matrix for multivariate distributions
    lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans(vi, vns[1])))
    var .= r
    return var, lp
end
function dot_assume(
    rng,
    spl::Sampler,
    ::Any,
    ::AbstractArray{<:VarName},
    ::Any,
    ::Any,
)
    error("[DynamicPPL] $(alg_str(spl)) doesn't support vectorizing assume statement")
end

function get_and_set_val!(
    rng,
    vi,
    vns::AbstractVector{<:VarName},
    dist::MultivariateDistribution,
    spl::Union{SampleFromPrior,SampleFromUniform},
)
    n = length(vns)
    if haskey(vi, vns[1])
        # Always overwrite the parameters with new ones for `SampleFromUniform`.
        if spl isa SampleFromUniform || is_flagged(vi, vns[1], "del")
            unset_flag!(vi, vns[1], "del")
            r = init(rng, dist, spl, n)
            for i in 1:n
                vn = vns[i]
                vi[vn] = vectorize(dist, r[:, i])
                settrans!(vi, false, vn)
                setorder!(vi, vn, get_num_produce(vi))
            end
        else
            r = vi[vns]
        end
    else
        r = init(rng, dist, spl, n)
        for i in 1:n
            vn = vns[i]
            push!(vi, vn, r[:,i], dist, spl)
            settrans!(vi, false, vn)
        end
    end
    return r
end

function get_and_set_val!(
    rng,
    vi,
    vns::AbstractArray{<:VarName},
    dists::Union{Distribution, AbstractArray{<:Distribution}},
    spl::Union{SampleFromPrior,SampleFromUniform},
)
    if haskey(vi, vns[1])
        # Always overwrite the parameters with new ones for `SampleFromUniform`.
        if spl isa SampleFromUniform || is_flagged(vi, vns[1], "del")
            unset_flag!(vi, vns[1], "del")
            f = (vn, dist) -> init(rng, dist, spl)
            r = f.(vns, dists)
            for i in eachindex(vns)
                vn = vns[i]
                dist = dists isa AbstractArray ? dists[i] : dists
                vi[vn] = vectorize(dist, r[i])
                settrans!(vi, false, vn)
                setorder!(vi, vn, get_num_produce(vi))
            end
        else
            r = reshape(vi[vec(vns)], size(vns))
        end
    else
        f = (vn, dist) -> init(rng, dist, spl)
        r = f.(vns, dists)
        push!.(Ref(vi), vns, r, dists, Ref(spl))
        settrans!.(Ref(vi), false, vns)
    end
    return r
end

function set_val!(
    vi,
    vns::AbstractVector{<:VarName},
    dist::MultivariateDistribution,
    val::AbstractMatrix,
)
    @assert size(val, 2) == length(vns)
    foreach(enumerate(vns)) do (i, vn)
        vi[vn] = val[:,i]
    end
    return val
end
function set_val!(
    vi,
    vns::AbstractArray{<:VarName},
    dists::Union{Distribution, AbstractArray{<:Distribution}},
    val::AbstractArray,
)
    @assert size(val) == size(vns)
    foreach(CartesianIndices(val)) do ind
        dist = dists isa AbstractArray ? dists[ind] : dists
        vi[vns[ind]] = vectorize(dist, val[ind])
    end
    return val
end

# observe
function dot_tilde(ctx::DefaultContext, sampler, right, left, vi)
    return _dot_tilde(sampler, right, left, vi)
end
function dot_tilde(ctx::PriorContext, sampler, right, left, vi)
    return 0
end
function dot_tilde(ctx::LikelihoodContext, sampler, right, left, vi)
    return _dot_tilde(sampler, right, left, vi)
end
function dot_tilde(ctx::MiniBatchContext, sampler, right, left, vi)
    return ctx.loglike_scalar * dot_tilde(ctx.ctx, sampler, right, left, vi)
end

"""
    dot_tilde_observe(ctx, sampler, right, left, vname, vinds, vi)

Handle broadcasted observed values, e.g., `x .~ MvNormal()` (where `x` does occur the model inputs),
accumulate the log probability, and return the observed value.

Falls back to `dot_tilde(ctx, sampler, right, left, vi)` ignoring the information about variable
name and indices; if needed, these can be accessed through this function, though.
"""
function dot_tilde_observe(ctx, sampler, right, left, vn, inds, vi)
    logp = dot_tilde(ctx, sampler, right, left, vi)
    acclogp!(vi, logp)
    return left
end

"""
    dot_tilde_observe(ctx, sampler, right, left, vi)

Handle broadcasted observed constants, e.g., `[1.0] .~ MvNormal()`, accumulate the log
probability, and return the observed value.

Falls back to `dot_tilde(ctx, sampler, right, left, vi)`.
"""
function dot_tilde_observe(ctx, sampler, right, left, vi)
    logp = dot_tilde(ctx, sampler, right, left, vi)
    acclogp!(vi, logp)
    return left
end

function _dot_tilde(sampler, right, left::AbstractArray, vi)
    return dot_observe(sampler, right, left, vi)
end
# Ambiguity error when not sure to use Distributions convention or Julia broadcasting semantics
function _dot_tilde(
    sampler::AbstractSampler,
    right::Union{MultivariateDistribution, AbstractVector{<:MultivariateDistribution}},
    left::AbstractMatrix{>:AbstractVector},
    vi,
)
    throw(DimensionMismatch(AMBIGUITY_MSG))
end

function dot_observe(
    spl::Union{SampleFromPrior, SampleFromUniform},
    dist::MultivariateDistribution,
    value::AbstractMatrix,
    vi,
)
    increment_num_produce!(vi)
    @debug "dist = $dist"
    @debug "value = $value"
    return Distributions.loglikelihood(dist, value)
end
function dot_observe(
    spl::Union{SampleFromPrior, SampleFromUniform},
    dists::Distribution,
    value::AbstractArray,
    vi,
)
    increment_num_produce!(vi)
    @debug "dists = $dists"
    @debug "value = $value"
    return Distributions.loglikelihood(dists, value)
end
function dot_observe(
    spl::Union{SampleFromPrior, SampleFromUniform},
    dists::AbstractArray{<:Distribution},
    value::AbstractArray,
    vi,
)
    increment_num_produce!(vi)
    @debug "dists = $dists"
    @debug "value = $value"
    return sum(zip(dists, value)) do (d, v)
        Distributions.loglikelihood(d, v)
    end
end
function dot_observe(
    spl::Sampler,
    ::Any,
    ::Any,
    ::Any,
)
    error("[DynamicPPL] $(alg_str(spl)) doesn't support vectorizing observe statement")
end


text:
module ConjugateGradients
    include("genericblas.jl")
    include("reader.jl")
    include("cg.jl")
    include("bicgstab.jl")
end


text:
using TransportBasedInference: ncoeff

@testset "Test optimization max_terms = nothing without/with QR" begin
    Nx = 2

    m = 20
    Ne = 500
    X = randn(Nx, Ne).^2 + 0.1*randn(Nx, Ne)

    L = LinearTransform(X; diag = true)
    transform!(L, X);

    idx = [0 0; 0 1; 0 2; 0 3; 1 1]
    coeff =  [   -0.7708538710735897;
              0.1875006767230035;
              1.419396079869706;
              0.2691388018878241;
             -2.9879004455954723];

    C = HermiteMapComponent(m, Nx, idx, coeff);
    C0 = deepcopy(C)

    C_noqr, error_noqr = TransportBasedInference.optimize(C, X, nothing; withqr = false, verbose = false);
    C_qr, error_qr = TransportBasedInference.optimize(C, X, nothing; withqr = true, verbose = false);

    S = Storage(C0.I.f, X)
    coeff0 = getcoeff(C0)
    precond = zeros(ncoeff(C0), ncoeff(C0))
    precond!(precond, coeff0, S, C0, X)

    res = Optim.optimize(Optim.only_fg!(negative_log_likelihood(S, C0, X)), coeff0,
          Optim.LBFGS(; m = 20, P = Preconditioner(precond)))

    @test norm(getcoeff(C_noqr) - res.minimizer) <1e-6
    @test norm(getcoeff(C_qr) - res.minimizer) <1e-6

    error = res.minimum

    @test abs(error_noqr - error) < 1e-8
    @test abs(error_qr - error) < 1e-8

end

@testset "Test optimization max_terms is an integer without/with QR" begin
    Nx = 2

    m = 20
    Ne = 500
    X = randn(Nx, Ne).^2 + 0.1*randn(Nx, Ne)

    L = LinearTransform(X; diag = true)
    transform!(L, X);

    C = HermiteMapComponent(m, Nx);

    C_noqr, error_noqr = TransportBasedInference.optimize(C, X, 4; withqr = false, verbose = false);
    C_qr, error_qr = TransportBasedInference.optimize(C, X, 4; withqr = true, verbose = false);

    idx0 = getidx(C_noqr)

    C0 = HermiteMapComponent(m, Nx, idx0, zeros(size(idx0,1)));

    S = Storage(C0.I.f, X)
    coeff0 = getcoeff(C0)
    precond = zeros(ncoeff(C0), ncoeff(C0))
    precond!(precond, coeff0, S, C0, X)

    res = Optim.optimize(Optim.only_fg!(negative_log_likelihood(S, C0, X)), coeff0,
          Optim.LBFGS(; m = 20, P = Preconditioner(precond)))

    @test norm(getcoeff(C_noqr) - res.minimizer) <5e-5
    @test norm(getcoeff(C_qr) - res.minimizer) <5e-5

    error = res.minimum

    @test abs(error_noqr[end] - error) < 5e-5
    @test abs(error_qr[end] - error) < 5e-5

end

@testset "Test optimization max_terms is kfold without/with qr" begin

    Nx = 2

    m = 20
    Ne = 200
    X = randn(Nx, Ne).^2 + 0.1*randn(Nx, Ne)

    L = LinearTransform(X; diag = true)
    transform!(L, X)

    C = HermiteMapComponent(m, Nx)

    n_folds = 5
    folds = kfolds(1:size(X,2), k = n_folds)

    # Run greedy approximation
    max_iter = min(m-1,ceil(Int64, sqrt(size(X,2))))

    valid_error_noqr = zeros(max_iter+1, n_folds)
    valid_error_qr = zeros(max_iter+1, n_folds)

    @inbounds for i=1:n_folds
        idx_train, idx_valid = folds[i]
        @test size(idx_valid,1) == 40
        @test size(idx_train,1) == 160
        @test size(idx_train,1) < size(X,2)
        C_noqr, error_noqr = greedyfit(m, Nx, X[:,idx_train], X[:,idx_valid], max_iter; withqr = false, verbose  = false)
        C_qr, error_qr = greedyfit(m, Nx, X[:,idx_train], X[:,idx_valid], max_iter; withqr = true, verbose  = false)

        # error[2] contains the history of the validation error
        valid_error_noqr[:,i] .= deepcopy(error_noqr[2])
        valid_error_qr[:,i] .= deepcopy(error_qr[2])

    end

    mean_valid_error_noqr = mean(valid_error_noqr, dims  = 2)[:,1]
    mean_valid_error_qr = mean(valid_error_qr, dims  = 2)[:,1]

    @test size(mean_valid_error_noqr,1) == max_iter+1
    @test size(mean_valid_error_noqr,1) == max_iter+1

    value_noqr, opt_nterms_noqr = findmin(mean_valid_error_noqr)
    value_qr, opt_nterms_qr = findmin(mean_valid_error_qr)

    @test value_noqr == mean_valid_error_noqr[opt_nterms_noqr]
    @test value_qr == mean_valid_error_qr[opt_nterms_qr]

    # Run greedy fit up to opt_nterms on all the data
    C_opt_noqr, error_opt_noqr = greedyfit(m, Nx, X, opt_nterms_noqr; withqr = false, verbose  = false)
    C_opt_qr, error_opt_qr = greedyfit(m, Nx, X, opt_nterms_qr; withqr = false, verbose  = false)

    @test size(getcoeff(C_opt_noqr),1) == opt_nterms_noqr
    @test size(getcoeff(C_opt_qr),1) == opt_nterms_qr

    C_new_noqr, error_new_noqr = TransportBasedInference.optimize(C, X, "kfold"; withqr = false)
    C_new_qr, error_new_qr = TransportBasedInference.optimize(C, X, "kfold"; withqr = true)

    @test norm(getcoeff(C_opt_noqr) - getcoeff(C_new_noqr))<1e-5
    @test norm(error_opt_noqr - error_new_noqr)<1e-5

    @test norm(getcoeff(C_opt_qr) - getcoeff(C_new_qr))<1e-5
    @test norm(error_opt_qr - error_new_qr)<1e-5

end

@testset "Test optimization max_terms is split with/without QR" begin

    Nx = 2

    m = 20
    Ne = 200
    X = randn(Nx, Ne).^2 + 0.1*randn(Nx, Ne)

    L = LinearTransform(X; diag = true)
    transform!(L, X)

    C = HermiteMapComponent(m, Nx)

    nvalid = ceil(Int64, floor(0.2*size(X,2)))
    X_train = X[:,nvalid+1:end]
    X_valid = X[:,1:nvalid]

    @test size(X_train, 2) == 160
    @test size(X_valid, 2) == 40

    # Set maximum patience for optimization
    maxpatience = 20

    # Run greedy approximation
    max_iter = ceil(Int64, sqrt(size(X,2)))

    C_opt_noqr, error_opt_noqr = greedyfit(m, Nx, X_train, X_valid, max_iter;
                                           withqr = false, maxpatience = maxpatience,
                                           verbose  = false)
   C_opt_qr, error_opt_qr = greedyfit(m, Nx, X_train, X_valid, max_iter;
                                          withqr = true, maxpatience = maxpatience,
                                          verbose  = false)

    C_new_noqr, error_new_noqr = TransportBasedInference.optimize(C, X, "split"; withqr = false)

    C_new_qr, error_new_qr = TransportBasedInference.optimize(C, X, "split"; withqr = true)

    @test norm(getcoeff(C_opt_noqr) - getcoeff(C_new_noqr))<1e-8
    @test norm(error_opt_noqr[1] - error_new_noqr[1])<1e-8
    @test norm(error_opt_noqr[2] - error_new_noqr[2])<1e-8

    @test size(getcoeff(C_new_noqr),1) == max_iter

    @test norm(getcoeff(C_opt_qr) - getcoeff(C_new_qr))<1e-8
    @test norm(error_opt_qr[1] - error_new_qr[1])<1e-8
    @test norm(error_opt_qr[2] - error_new_qr[2])<1e-8

    @test size(getcoeff(C_new_qr),1) == max_iter
end


text:
using CSV
using DataFrames
using Plots

function visualize()
  data = DataFrame(
    CSV.File("pulitzer-circulation-data.csv", normalizenames = true))

  sort!(data, :Pulitzer_Prize_Winners_and_Finalists_1990_2014, rev = true)
  println(data)
  sort!(data, :Pulitzer_Prize_Winners_and_Finalists_1990_2014)

  x = data.Newspaper
  y = data.Pulitzer_Prize_Winners_and_Finalists_1990_2014
  # y = getproperty(data,
  #   Symbol("Pulitzer Prize Winners and Finalists, 1990-2014"))

  p = plot(x, y, seriestype = :bar, direction = :h, size = (1000, 2000),
    legend = false)
  savefig(p, "plot.png")
end

visualize()