text:
#################################################################################################
# AIM 2
#################################################################################################
## Deps
using Plots
using CodecZlib
using StatsPlots
using Distributions
using DelimitedFiles
using Plots.PlotMeasures

using Pkg
Pkg.activate("./")
using CpelNano

## Constants
const cov_levels = [5.0,10.0,15.0,20.0,25.0]
const cov_labels = ["$(Int(c))x" for c in cov_levels]
const noise_levels = [2.0,2.5,3.0,3.5]
const noise_labels = ["$(s)" for s in noise_levels]
const aim_dir = "/Users/jordiabante/OneDrive - Johns Hopkins/CpelNano/Data/Simulations/Aim-2/"
const gt_dir = "/Users/jordiabante/OneDrive - Johns Hopkins/CpelNano/Data/Simulations/ground_truth/"
const blind_friend_col = ["#999999","#E69F00","#56B4E9","#009E73","#F0E442","#0072B2","#D55E00","#CC79A7"]

## Default attributes
default(titlefont=(14, "arial"),guidefont=(16, "arial"),tickfont=(12, "arial"))

#################################################################################################
# Functions
#################################################################################################
function read_in_mod_file(path::String)::Tuple{Vector{Int64},Array{Float64,2}}

    # Get data
    stream = GzipDecompressorStream(open(path))
    in_data = readdlm(stream)
    close(stream)

    # Read in estimated regions and parameters
    cpel_reg = convert.(Int64, in_data[:,2])
    cpel_params = in_data[:,4]
    cpel_params = map(x -> parse.(Float64, split(x, ",")), cpel_params)

    # Get parameters
    a = map(x -> x[1], cpel_params)
    b = map(x -> x[2], cpel_params)
    c = map(x -> x[3], cpel_params)

    # Return vectors
    return cpel_reg, [a b c]

end
function read_in_mod_file_dic(path::String)::Dict{String,NTuple{2,Vector{Float64}}}

    # Get data
    stream = GzipDecompressorStream(open(path))
    in_data = readdlm(stream)
    close(stream)

    # Remove region with no α,β
    kp_ind = [in_data[i,5] != "" && in_data[i,6] != "" for i = 1:size(in_data)[1]]
    in_data = in_data[kp_ind,:]

    # Read in
    cpel_reg = ["$(in_data[i,1])-$(in_data[i,2])-$(in_data[i,3])" for i = 1:size(in_data)[1]]
    αs = string.(in_data[:,5])
    βs = string.(in_data[:,6])
    αs = map(x -> parse.(Float64, split(x, ",")), αs)
    βs = map(x -> parse.(Float64, split(x, ",")), βs)

    # Keep indices with more than two groups
    kp_ind = length.(βs) .> 2
    cpel_reg = cpel_reg[kp_ind]
    αs = αs[kp_ind]
    βs = βs[kp_ind]

    # Create dictionary
    out_dic = Dict{String,NTuple{2,Vector{Float64}}}()
    for (i, reg) in enumerate(cpel_reg)
        out_dic[reg] = (αs[i], βs[i])
    end

    # Return dictionary
    return out_dic

end
function read_in_ex_file(path::String)::Tuple{Vector{Int64},Vector{Float64}}

    # Get data
    stream = GzipDecompressorStream(open(path))
    in_data = readdlm(stream)
    close(stream)

    # Read in estimated regions and parameters
    cpel_reg = convert.(Int64, in_data[:,2])
    cpel_ex = convert.(Float64, in_data[:,4])

    # Return vectors
    return cpel_reg, cpel_ex

end
function read_in_exx_file(path::String)::Tuple{Vector{Int64},Vector{Float64}}

    # Get data
    stream = GzipDecompressorStream(open(path))
    in_data = readdlm(stream)
    close(stream)

    # Read in estimated regions and parameters
    cpel_cg1 = convert.(Int64, in_data[:,2])
    cpel_exx = convert.(Float64, in_data[:,4])

    # Return vectors
    return cpel_cg1, cpel_exx

end
function read_in_diff_file(path::String)::Tuple{Vector{Int64},Vector{Float64}}

    # Get data
    stream = GzipDecompressorStream(open(path))
    in_data = readdlm(stream)
    close(stream)

    # Read in estimated regions and parameters
    cpel_reg = convert.(Int64, in_data[:,2])
    cpel_stat = convert.(Float64, in_data[:,4])

    # Return vectors
    return cpel_reg, cpel_stat

end
function find_comm_regions(sts1::Vector{Int64}, sts2::Vector{Int64})::NTuple{2,Vector{Bool}}

    # Find common regions
    int_sts = intersect(Set(sts1), Set(sts2))
    sts1_ind = map(i -> sts1[i] in int_sts, 1:length(sts1))
    sts2_ind = map(i -> sts2[i] in int_sts, 1:length(sts2))

    # Return
    return sts1_ind, sts2_ind

end
function cos_sim(mat1::Array{Float64,2}, mat2::Array{Float64,2})::Vector{Float64}

    # Intermediate quantities
    num = sum(mat1 .* mat2, dims=2)
    den = sqrt.(sum(mat1.^2, dims=2)) .* sqrt.(sum(mat2.^2, dims=2))

    # Return cosine similarity
    return vec(num ./ den)

end
function mae(vec1::Vector{Float64}, vec2::Vector{Float64})::Vector{Float64}

    # Return mean absolute error
    return vec(abs.(vec1 - vec2))

end
function get_EX_EXX_logZ(αs::Vector{Float64}, βs::Vector{Float64})::Tuple{Vector{Float64},Vector{Float64},Float64}

    # Get matrices
    logu1 = CpelNano.get_log_u(αs[1])
    loguL = CpelNano.get_log_u(αs[end])
    logWs = [CpelNano.get_log_W(αs[l], αs[l + 1], βs[l]) for l = 1:(length(αs) - 1)]

    # Get log Z
    logZ = CpelNano.get_log_Z(logu1, loguL, logWs)

    # Get expectations
    ex = CpelNano.get_E_X_log(logu1, loguL, logWs, logZ)
    exx = CpelNano.get_E_XX_log(logu1, loguL, logWs, logZ)

    # Return
    return ex, exx, logZ

end
function comp_cmd(m1::NTuple{2,Vector{Float64}}, m2::NTuple{2,Vector{Float64}})::Float64

    # Create mixture
    αmix = 0.5 * (m1[1] + m2[1])
    βmix = 0.5 * (m1[2] + m2[2])

    # Compute intermediate quantities
    ex1, exx1, logZ1 = get_EX_EXX_logZ(m1[1], m1[2])
    ex2, exx2, logZ2 = get_EX_EXX_logZ(m2[1], m2[2])
    ex_mix, exx_mix, logZ_mix = get_EX_EXX_logZ(αmix, βmix)

    # Compute entropies
    h1 = logZ1 - m1[1]' * ex1 - m1[2]' * exx1
    h2 = logZ2 - m2[1]' * ex2 - m2[2]' * exx2
    c1 = logZ_mix - αmix' * ex1 - βmix' * exx1
    c2 = logZ_mix - αmix' * ex2 - βmix' * exx2

    # Compute CMD
    cmd = 1.0 - (h1 + h2) / (c1 + c2)

    # Return
    return min(1.0, max(0.0, cmd))

end
function comp_kl(p::NTuple{2,Vector{Float64}}, q::NTuple{2,Vector{Float64}})::Float64

    # Compute intermediate quantities
    ex_p, exx_p, logZp = get_EX_EXX_logZ(p[1], p[2])
    ex_q, exx_q, logZq = get_EX_EXX_logZ(q[1], q[2])

    # Compute entropy of true distribution p(x)
    hp = logZp - p[1]' * ex_p - p[2]' * exx_p

    # Compute cross-entropy of estimate q(x) wrt true p(x)
    c = logZq - q[1]' * ex_p - q[2]' * exx_p

    # Return H(P,Q)-H(P)
    return c - hp

end
function comp_cmd_set_keys(comm_keys::Vector{String}, dic1::Dict{String,NTuple{2,Vector{Float64}}}, dic2::Dict{String,NTuple{2,Vector{Float64}}})::Vector{Float64}

    # Loop over keys
    cmds = map(key -> comp_cmd(dic1[key], dic2[key]), comm_keys)

    # Return
    return cmds

end
function comp_kl_set_keys(comm_keys::Vector{String}, dic1::Dict{String,NTuple{2,Vector{Float64}}}, dic2::Dict{String,NTuple{2,Vector{Float64}}})::Vector{Float64}

    # Loop over keys
    kls = map(key -> comp_kl(dic1[key], dic2[key]), comm_keys)

    # Return
    return kls

end
function gen_kl_plt_array()

    println("Generating: KL plots")

    # Read in ground-truth
    gt_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug.txt.gz"
    gt_αβ_dic = read_in_mod_file_dic(gt_file)

    # Create plot array
    plt_arr = Vector{Any}()
    for (i, sig) in enumerate(noise_levels)
        println("Noise level: $(sig)")
        # marg_box_mat = fill(NaN, (length(keys(gt_αβ_dic)), length(cov_labels)))
        nonoise_box_mat = fill(NaN, (length(keys(gt_αβ_dic)), length(cov_labels)))
        cpelnano_box_mat = fill(NaN, (length(keys(gt_αβ_dic)), length(cov_labels)))
        for (j, cove) in enumerate(cov_labels)
            println("Coverage level: $(cove)")
            # Get filenames
            # marg_file = "$(aim_dir)/sample_averages/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_theta.txt.gz"
            nonoise_file = "$(aim_dir)/cpelnano_nonoise/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_theta.txt.gz"
            cpelnano_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_theta.txt.gz"
            # Read in files
            # marg_αβ_dic = read_in_mod_file_dic(marg_file)
            nonoise_αβ_dic = read_in_mod_file_dic(nonoise_file)
            cpelnano_αβ_dic = read_in_mod_file_dic(cpelnano_file)
            # Record performance for CpelNano
            comm_keys = collect(intersect(Set(keys(gt_αβ_dic)), Set(keys(cpelnano_αβ_dic))))
            key_ind = [x in comm_keys for x in keys(gt_αβ_dic)]
            cpelnano_box_mat[key_ind,j] .= comp_kl_set_keys(comm_keys, gt_αβ_dic, cpelnano_αβ_dic)
            # Record performance for CpelNano no-noise
            comm_keys = collect(intersect(Set(keys(gt_αβ_dic)), Set(keys(nonoise_αβ_dic))))
            key_ind = [x in comm_keys for x in keys(gt_αβ_dic)]
            nonoise_box_mat[key_ind,j] .= comp_kl_set_keys(comm_keys, gt_αβ_dic, nonoise_αβ_dic)
            # Record performance for marginal
            # comm_keys = collect(intersect(Set(keys(gt_αβ_dic)), Set(keys(marg_αβ_dic))))
            # key_ind = [x in comm_keys for x in keys(gt_αβ_dic)]
            # marg_box_mat[key_ind,j] .= comp_kl_set_keys(comm_keys, gt_αβ_dic, marg_αβ_dic)
        end

        # Arrange CpelNano data
        cpelnano_data = vcat(cpelnano_box_mat...)
        cpelnano_labels = fill("M1", length(cpelnano_data))
        cpelnano_covs = vcat([fill(cove, size(cpelnano_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(cpelnano_data)
        cpelnano_data = cpelnano_data[keepind]
        cpelnano_labels = cpelnano_labels[keepind]
        cpelnano_covs = cpelnano_covs[keepind]

        # Arrange CpelNano no-noise data
        nonoise_data = vcat(nonoise_box_mat...)
        nonoise_labels = fill("M2", length(nonoise_data))
        nonoise_covs = vcat([fill(cove, size(nonoise_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(nonoise_data)
        nonoise_data = nonoise_data[keepind]
        nonoise_labels = nonoise_labels[keepind]
        nonoise_covs = nonoise_covs[keepind]

        # Arrange marginal data
        # marg_data = vcat(marg_box_mat...)
        # marg_labels = fill("CMD (Emp.)", length(marg_data))
        # marg_covs = vcat([fill(cove, size(marg_box_mat)[1]) for cove in cov_labels]...)
        # keepind = .!isnan.(marg_data)
        # marg_data = marg_data[keepind]
        # marg_labels = marg_labels[keepind]
        # marg_covs = marg_covs[keepind]

        # Push boxplot to array
        plt_data = vcat(cpelnano_data, nonoise_data) # , marg_data)
        plt_labels = vcat(cpelnano_labels, nonoise_labels) # , marg_labels)
        plt_covs = vcat(cpelnano_covs, nonoise_covs) # , marg_covs)
        ylim = (0.0, 25)
        if i == 1
            plt = groupedboxplot(plt_covs,plt_data,group=plt_labels,bar_width=0.5,ylim=ylim,
                title="\\sigma=$(noise_labels[i])",xlab="Coverage",ylab="KL",outliers=false)
        else
            plt = groupedboxplot(plt_covs,plt_data,group=plt_labels,bar_width=0.5,ylim=ylim,
                title="\\sigma=$(noise_labels[i])",xlab="Coverage",ylab="KL",label="",outliers=false)
        end
        push!(plt_arr, plt)

    end

    # Return array
    return plt_arr

end
function gen_param_plt_array()

    println("Generating: params")

    # Plot base
    xlab = "Coverage"
    ylab = "Cosine similarity"
    palet = fill("#999999", length(cov_labels))
    plt_base = plot(seriestype=:boxplot, color_palette=palet, xticks=(1:5, cov_labels),
        ylim=(-1, 1), xlab=xlab, ylab=ylab, labels="", alpha=0.5)

    ## Read in ground-truth
    gt_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug.txt.gz"
    gt_sts, gt_params = read_in_mod_file(gt_file)

    # Create plot array
    plt_arr = Vector{Any}()
    for (i, sig) in enumerate(noise_levels)
        println("Noise level: $(sig)")
        plt_aux = deepcopy(plt_base)
        box_mat = fill(NaN, (length(gt_sts), length(cov_labels)))
        for (j, cove) in enumerate(cov_labels)
            # Get filename of sample
            model_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_theta.txt.gz"
            # Read in model file
            nano_sts, nano_params = read_in_mod_file(model_file)
            # Find intersect
            gt_sts_ind, nano_sts_ind = find_comm_regions(gt_sts, nano_sts)
            # Record metric in matrix
            box_mat[gt_sts_ind,j] .= cos_sim(gt_params[gt_sts_ind,:], nano_params[nano_sts_ind,:])
            # Add to plot
            boxplot!(plt_aux, box_mat[gt_sts_ind,j], labels="", title=noise_labels[i], outliers=false)
        end
        # Push boxplot to array
        push!(plt_arr, plt_aux)
    end

    # Return array
    return plt_arr

end
function gen_stat_plt_array(stat, sig)

    cpelnano_data = []
    nonoise_data = []
    cpelnano_cove = []
    nonoise_cove = []
    for (j, cove) in enumerate(cov_labels)
        # Get filenames
        cpelnano_file = "$(aim_dir)/differential/two_sample_comp_gt_vs_cpelnano_$(sig)_$(cov_labels[j])_$(stat).txt.gz"
        nonoise_file = "$(aim_dir)/differential/two_sample_comp_gt_vs_nonoise_$(sig)_$(cov_labels[j])_$(stat).txt.gz"
        # Read in files
        cpelnano_regs, cpelnano_stat = read_in_diff_file(cpelnano_file)
        nonoise_regs, nonoise_stat = read_in_diff_file(nonoise_file)
        # Record performance metric for E[X] in matrix
        append!(cpelnano_data, abs.(cpelnano_stat))
        append!(cpelnano_cove, fill(cove, length(cpelnano_stat)))
        # Record performance metric for E[X] (no noise modeling) in matrix
        append!(nonoise_data, abs.(nonoise_stat))
        append!(nonoise_cove, fill(cove, length(nonoise_stat)))
    end

    # Labels
    cpelnano_labels = fill("M1", length(cpelnano_data))
    nonoise_labels = fill("M2", length(nonoise_data))

    # Plot
    plt_data = vcat(cpelnano_data, nonoise_data)
    plt_labels = vcat(cpelnano_labels, nonoise_labels)
    plt_covs = vcat(cpelnano_cove, nonoise_cove)
    ylab = stat == "Tcmd" ? "CMD" : "Absolute Error"
    if stat == "Tcmd"
        yticks = 0.00:0.05:0.15
        ylim = (0.0, 0.15)
    elseif stat == "Tnme"
        yticks = 0.0:0.1:0.50
        ylim = (0.0, 0.50)
    else
        yticks = 0.0:0.05:0.25
        ylim = (0.0, 0.25)
    end
    plt = groupedboxplot(plt_covs,plt_data,group=plt_labels,bar_width=0.5,ylim=ylim,yticks=yticks,
        title="$(stat)",xlab="Coverage",ylab=ylab,outliers=false,top_margin=10px, bottom_margin=20px,
        left_margin=40px, right_margin=20px)

    # Return array
    return plt

end
function gen_ey_plt_array()

    println("Generating: E[Y]=E[0.5*(X+1)]")

    ## Read in ground-truth
    gt_ex_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug_ex.txt.gz"
    gt_ex_sts, gt_ex = read_in_ex_file(gt_ex_file)

    # Create plot array
    plt_arr = Vector{Any}()
    for (i, sig) in enumerate(noise_levels)
        println("Noise level: $(sig)")
        ex_box_mat = fill(NaN, (length(gt_ex_sts), length(cov_labels)))
        smp_ex_box_mat = fill(NaN, (length(gt_ex_sts), length(cov_labels)))
        nonoise_ex_box_mat = fill(NaN, (length(gt_ex_sts), length(cov_labels)))
        for (j, cove) in enumerate(cov_labels)
            # Get filenames
            ex_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_ex.txt.gz"
            smp_ex_file = "$(aim_dir)/sample_averages/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_ex.txt.gz"
            nonoise_ex_file = "$(aim_dir)/cpelnano_nonoise/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_ex.txt.gz"
            # Read in files
            ex_sts, ex = read_in_ex_file(ex_file)
            smp_ex_sts, smp_exs = read_in_ex_file(smp_ex_file)
            nonoise_ex_sts, nonoise_ex = read_in_ex_file(nonoise_ex_file)
            # Record performance metric for E[X] in matrix
            gt_sts_ind, ex_sts_ind = find_comm_regions(gt_ex_sts, ex_sts)
            ex_box_mat[gt_sts_ind,j] .= mae(gt_ex[gt_sts_ind], ex[ex_sts_ind]) ./ 2.0
            # Record performance metric for E[X] (no noise modeling) in matrix
            gt_sts_ind, nonoise_ex_sts_ind = find_comm_regions(gt_ex_sts, nonoise_ex_sts)
            nonoise_ex_box_mat[gt_sts_ind,j] .= mae(gt_ex[gt_sts_ind], nonoise_ex[nonoise_ex_sts_ind]) ./ 2.0
            # Record performance metric for X̄ in matrix
            gt_sts_ind, smp_ex_sts_ind = find_comm_regions(gt_ex_sts, smp_ex_sts)
            smp_ex_box_mat[gt_sts_ind,j] .= mae(gt_ex[gt_sts_ind], smp_exs[smp_ex_sts_ind]) ./ 2.0
        end

        # Arrange E[X] data
        ex_data = vcat(ex_box_mat...)
        ex_labels = fill("E[X] (M1)", length(ex_data))
        ex_covs = vcat([fill(cove, size(ex_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(ex_data)
        ex_data = ex_data[keepind]
        ex_labels = ex_labels[keepind]
        ex_covs = ex_covs[keepind]

        # Arrange no noise modeling E[X] data
        nonoise_ex_data = vcat(nonoise_ex_box_mat...)
        nonoise_ex_labels = fill("E[X] (M2)", length(nonoise_ex_data))
        nonoise_ex_covs = vcat([fill(cove, size(nonoise_ex_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(nonoise_ex_data)
        nonoise_ex_data = nonoise_ex_data[keepind]
        nonoise_ex_labels = nonoise_ex_labels[keepind]
        nonoise_ex_covs = nonoise_ex_covs[keepind]

        # Arrange X̄ data
        smp_ex_data = vcat(smp_ex_box_mat...)
        smp_ex_labels = fill("Xbar", length(smp_ex_data))
        smp_ex_covs = vcat([fill(cove, size(smp_ex_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(smp_ex_data)
        smp_ex_data = smp_ex_data[keepind]
        smp_ex_labels = smp_ex_labels[keepind]
        smp_ex_covs = smp_ex_covs[keepind]

        # Push boxplot to array
        plt_data = vcat(ex_data, nonoise_ex_data, smp_ex_data)
        plt_labels = vcat(ex_labels, nonoise_ex_labels, smp_ex_labels)
        plt_covs = vcat(ex_covs, nonoise_ex_covs, smp_ex_covs)
        if i == 1
            plt = groupedboxplot(plt_covs,plt_data,group=plt_labels,bar_width=0.5,ylim=(0.0, 0.5),
                title="\\sigma=$(noise_labels[i])",xlab="Coverage",ylab="Absolute error",outliers=false)
        else
            plt = groupedboxplot(plt_covs,plt_data,group=plt_labels,bar_width=0.5,ylim=(0.0, 0.5),
                title="\\sigma=$(noise_labels[i])",xlab="Coverage",ylab="Absolute error",label="",outliers=false)
        end
    push!(plt_arr, plt)

    end

    # Return array
    return plt_arr

end
function comp_eyy(ex_cg::Vector{Int64}, ex::Vector{Float64}, exx_cg1::Vector{Int64}, exx::Vector{Float64})::Vector{Float64}

    # Compute E[YnYn+1]
    eyy = fill(NaN, length(exx))
    @inbounds for n = 1:length(eyy)
        # Find corresponding CG sites
        cg1_ind = findfirst(isequal(exx_cg1[n]), ex_cg)
        isnothing(cg1_ind) && continue
        # E[YY] = 0.25 *(E[XnXn+1]+E[Xn]+E[Xn+1]+1)
        eyy[n] = 0.25 * (exx[n] + ex[cg1_ind] + ex[cg1_ind + 1] + 1.0)
    end

    # Return
    return min.(1.0, max.(0.0, eyy))

end
function gen_eyy_plt_array()

    println("Generating: E[YY']=E[0.5(X+1)0.5(X'+1)]")

    ## Read in ground-truth
    gt_ex_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug_ex.txt.gz"
    gt_exx_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug_exx.txt.gz"
    gt_ex_cg, gt_ex = read_in_ex_file(gt_ex_file)
    gt_exx_cg1, gt_exx = read_in_exx_file(gt_exx_file)
    gt_eyy = comp_eyy(gt_ex_cg, gt_ex, gt_exx_cg1, gt_exx)

    # Create plot array
    plt_arr = Vector{Any}()
    for (i, sig) in enumerate(noise_levels)
        println("Noise level: $(sig)")
        eyy_box_mat = fill(NaN, (length(gt_exx_cg1), length(cov_labels)))
        smp_eyy_box_mat = fill(NaN, (length(gt_exx_cg1), length(cov_labels)))
        nonoise_eyy_box_mat = fill(NaN, (length(gt_exx_cg1), length(cov_labels)))
        for (j, cove) in enumerate(cov_labels)
            println("Coverage: $(cov_labels[j])")
            # Get filenames
            ex_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_ex.txt.gz"
            exx_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_exx.txt.gz"
            smp_ex_file = "$(aim_dir)/sample_averages/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_ex.txt.gz"
            smp_exx_file = "$(aim_dir)/sample_averages/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_exx.txt.gz"
            nonoise_ex_file = "$(aim_dir)/cpelnano_nonoise/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_ex.txt.gz"
            nonoise_exx_file = "$(aim_dir)/cpelnano_nonoise/sigma_$(sig)/gm12878_chr22_$(cov_labels[j])_exx.txt.gz"
            # Read in files
            ex_cg1, ex = read_in_ex_file(ex_file)
            smp_ex_cg1, smp_ex = read_in_ex_file(smp_ex_file)
            nonoise_ex_cg1, nonoise_ex = read_in_ex_file(nonoise_ex_file)
            exx_cg1, exx = read_in_exx_file(exx_file)
            smp_exx_cg1, smp_exx = read_in_exx_file(smp_exx_file)
            nonoise_exx_cg1, nonoise_exx = read_in_exx_file(nonoise_exx_file)
            # Compute E[YY] where Y=0.5*(X+1)
            eyy = comp_eyy(ex_cg1, ex, exx_cg1, exx)
            smp_eyy = comp_eyy(smp_ex_cg1, smp_ex, smp_exx_cg1, smp_exx)
            nonoise_eyy = comp_eyy(nonoise_ex_cg1, nonoise_ex, nonoise_exx_cg1, nonoise_exx)
            # Record metric for E[YY] in matrix
            gt_sts_ind, exx_sts_ind = find_comm_regions(gt_exx_cg1, exx_cg1)
            eyy_box_mat[gt_sts_ind,j] .= mae(gt_eyy[gt_sts_ind], eyy[exx_sts_ind])
            # Record metric for E[YY] (no noise) in matrix
            gt_sts_ind, nonoise_exx_sts_ind = find_comm_regions(gt_exx_cg1, nonoise_exx_cg1)
            nonoise_eyy_box_mat[gt_sts_ind,j] .= mae(gt_eyy[gt_sts_ind], nonoise_eyy[nonoise_exx_sts_ind])
            # Record metric for X̄ in matrix
            gt_sts_ind, smp_exx_sts_ind = find_comm_regions(gt_exx_cg1, smp_exx_cg1)
            smp_eyy_box_mat[gt_sts_ind,j] .= mae(gt_eyy[gt_sts_ind], smp_eyy[smp_exx_sts_ind])
        end

        # Arrange E[YY] data
        eyy_data = vcat(eyy_box_mat...)
        eyy_labels = fill("E[XX] (M1)", length(eyy_data))
        eyy_covs = vcat([fill(cove, size(eyy_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(eyy_data)
        eyy_data = eyy_data[keepind]
        eyy_labels = eyy_labels[keepind]
        eyy_covs = eyy_covs[keepind]

        # Arrange E[YY] (no noise) data
        nonoise_eyy_data = vcat(nonoise_eyy_box_mat...)
        nonoise_eyy_labels = fill("E[XX] (M2)", length(nonoise_eyy_data))
        nonoise_eyy_covs = vcat([fill(cove, size(nonoise_eyy_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(nonoise_eyy_data)
        nonoise_eyy_data = nonoise_eyy_data[keepind]
        nonoise_eyy_labels = nonoise_eyy_labels[keepind]
        nonoise_eyy_covs = nonoise_eyy_covs[keepind]

        # Arrange X̄X̄ data
        smp_eyy_data = vcat(smp_eyy_box_mat...)
        smp_eyy_labels = fill("XXbar", length(smp_eyy_data))
        smp_eyy_covs = vcat([fill(cove, size(smp_eyy_box_mat)[1]) for cove in cov_labels]...)
        keepind = .!isnan.(smp_eyy_data)
        smp_eyy_data = smp_eyy_data[keepind]
        smp_eyy_labels = smp_eyy_labels[keepind]
        smp_eyy_covs = smp_eyy_covs[keepind]

        # Push boxplot to array
        plt_data = vcat(eyy_data, nonoise_eyy_data, smp_eyy_data)
        plt_covs = vcat(eyy_covs, nonoise_eyy_covs, smp_eyy_covs)
        plt_labels = vcat(eyy_labels, nonoise_eyy_labels, smp_eyy_labels)
        if i == 1
            plt = groupedboxplot(plt_covs,plt_data,group=plt_labels,bar_width=0.5,ylim=(0.0, 0.5),
                title="\\sigma=$(noise_labels[i])",xlab="Coverage",ylab="Absolute error",outliers=false)
        else
            plt = groupedboxplot(plt_covs,plt_data,group=plt_labels,bar_width=0.5,ylim=(0.0, 0.5),
                title="\\sigma=$(noise_labels[i])",xlab="Coverage",ylab="Absolute error",label="",outliers=false)
        end
        push!(plt_arr, plt)

    end

    # Return array
    return plt_arr

end
function plt_params_scatter(cove, sig)

    # Find index of coverage label
    cov_ind = findfirst(isequal(cove), cov_levels)

    # Files
    gt_mod_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug.txt.gz"
    nano_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[cov_ind])_theta.txt.gz"

    # Params
    wgbs_sts, wgbs_params = read_in_mod_file(gt_mod_file)
    nano_sts, nano_params = read_in_mod_file(nano_file)

    # Find common regions
    wgbs_sts_ind, nano_sts_ind = find_comm_regions(wgbs_sts, nano_sts)
    wgbs_params = wgbs_params[wgbs_sts_ind,:]
    nano_params = nano_params[nano_sts_ind,:]

    # Plot scatter
    xlab = "Fine-grain"
    ylab = "Coarse-grain"
    p_a = plot(wgbs_params[:,1],nano_params[:,1],seriestype=:scatter,ylabel=ylab,markersize=0.2,
        markeralpha=0.1,xlim=(-10.0, 10.0),ylim=(-10.0, 10.0),label="",title="a");
    p_b = plot(wgbs_params[:,2],nano_params[:,2],seriestype=:scatter,ylabel=ylab,markersize=0.2,
        markeralpha=0.1,xlim=(-100.0, 100.0),ylim=(-100.0, 100.0),label="",title="b");
    p_c = plot(wgbs_params[:,3],nano_params[:,3],seriestype=:scatter,xlabel=xlab,ylabel=ylab,
        markersize=0.2,markeralpha=0.1,xlim=(-20.0, 20.0),ylim=(-20.0, 20.0),label="",title="c");
    plot(p_a, p_b, p_c, layout=(3, 1), size=(500, 800))
    savefig("$(aim_dir)/Scatter-Params-cov-$(cove)-s-$(sig)-Aim-2.pdf")
    savefig("$(aim_dir)/Scatter-Params-cov-$(cove)-s-$(sig)-Aim-2.png")

    # Return nothing
    return nothing

    end
function get_htmp_mat(vec1, vec2, axis_rng)

    outmat = fill(0.0, (length(axis_rng) - 1, length(axis_rng) - 1))
    for i in (length(axis_rng) - 1):-1:1
            for j in 1:(length(axis_rng) - 1)
            ind_1 = [axis_rng[i] <= x < axis_rng[i + 1] for x in vec1]
            ind_2 = [axis_rng[j] <= x < axis_rng[j + 1] for x in vec2]
            outmat[i,j] = sum(ind_1 .& ind_2)
        end
    end

    outmat .+= 1
    outmat ./= sum(outmat)
    outmat = -log10.(outmat)

    return outmat

end
function plt_params_heatmap(cove, sig)

    # Find index of coverage label
    cov_ind = findfirst(isequal(cove), cov_levels)

    # Files
    gt_mod_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug.txt.gz"
    nano_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[cov_ind])_theta.txt.gz"

    # Params
    wgbs_sts, wgbs_params = read_in_mod_file(gt_mod_file)
    nano_sts, nano_params = read_in_mod_file(nano_file)

    # Find common regions
    wgbs_sts_ind, nano_sts_ind = find_comm_regions(wgbs_sts, nano_sts)
    wgbs_params = wgbs_params[wgbs_sts_ind,:]
    nano_params = nano_params[nano_sts_ind,:]

    # Ranges
    a_axis_rng = collect(-10.0:1:10.0)
    b_axis_rng = collect(-100.0:10:100.0)
    c_axis_rng = collect(-20.0:2:20.0)

    # Construct matrix
    a_mat = get_htmp_mat(nano_params[:,1], wgbs_params[:,1], a_axis_rng)
    b_mat = get_htmp_mat(nano_params[:,2], wgbs_params[:,2], b_axis_rng)
    c_mat = get_htmp_mat(nano_params[:,3], wgbs_params[:,3], c_axis_rng)

    # Plot scatter
    xlab = "Fine-grain"
    ylab = "Coarse-grain"
    a_axis_labs = [isodd(i) ? "$(a_axis_rng[i])" : "" for i = 1:(length(a_axis_rng))]
    a_axis_ticks = (1:length(a_axis_labs), a_axis_labs)
    b_axis_labs = [isodd(i) ? "$(b_axis_rng[i])" : "" for i = 1:(length(b_axis_rng))]
    b_axis_ticks = (1:length(b_axis_labs), b_axis_labs)
    c_axis_labs = [isodd(i) ? "$(c_axis_rng[i])" : "" for i = 1:(length(c_axis_rng))]
    c_axis_ticks = (1:length(c_axis_labs), c_axis_labs)
    p_a = heatmap(a_mat,xlabel=xlab, ylabel=ylab, xticks=a_axis_ticks, yticks=a_axis_ticks,
        label="", c=cgrad([:white, :black]), clims=(1, 4), xrotation=45,
        title="\\alpha (\\rho=$(round(cor(nano_params[:,1], wgbs_params[:,1]);digits=2)))");
    p_b = heatmap(b_mat,xlabel=xlab, ylabel=ylab, xticks=b_axis_ticks, yticks=b_axis_ticks,
        label="", c=cgrad([:white, :black]), clims=(1, 4), xrotation=45,
        title="\\beta (\\rho=$(round(cor(nano_params[:,2], wgbs_params[:,2]);digits=2)))");
    p_c = heatmap(c_mat,xlabel=xlab, ylabel=ylab, xticks=c_axis_ticks, yticks=c_axis_ticks,
        label="", c=cgrad([:white, :black]), clims=(1, 4), xrotation=45,
        title="\\gamma (\\rho=$(round(cor(nano_params[:,3], wgbs_params[:,3]);digits=2)))");
    plot(p_a, p_b, p_c, layout=(3, 1), size=(800, 1500),top_margin=10px, bottom_margin=20px,
        left_margin=40px, right_margin=20px)
    savefig("$(aim_dir)/Heatmap-Params-cov-$(cove)-s-$(sig)-Aim-2.pdf")

    # Return nothing
    return nothing

end
function plt_exp_heatmap(cove, sig)

    # Ground-truth quantities
    gt_exx_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug_exx.txt.gz"
    gt_ex_file = "$(gt_dir)/GM12878_wgbs_cpelnano_theta_aug_ex.txt.gz"
    gt_ex_sts, gt_ex = read_in_ex_file(gt_ex_file)
    gt_exx_sts, gt_exx = read_in_exx_file(gt_exx_file)
    gt_ey = 0.5 .* (gt_ex .+ 1.0)
    gt_eyy = comp_eyy(gt_ex_sts, gt_ex, gt_exx_sts, gt_exx)

    # Find index of coverage label
    cov_ind = findfirst(isequal(cove), cov_levels)

    # Estimated quantities
    ex_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[cov_ind])_ex.txt.gz"
    exx_file = "$(aim_dir)/cpelnano/sigma_$(sig)/gm12878_chr22_$(cov_labels[cov_ind])_exx.txt.gz"
    ex_sts, ex = read_in_ex_file(ex_file)
    exx_sts, exx = read_in_exx_file(exx_file)
    ey = 0.5 .* (ex .+ 1.0)
    eyy = comp_eyy(ex_sts, ex, exx_sts, exx)

    # Find common regions E[X]
    wgbs_sts_ex_ind, ex_sts_ind = find_comm_regions(gt_ex_sts, ex_sts)
    gt_ey = gt_ey[wgbs_sts_ex_ind]
    ey = ey[ex_sts_ind]

    # Find common regions E[XX]
    wgbs_sts_exx_ind, exx_sts_ind = find_comm_regions(gt_exx_sts, exx_sts)
    gt_eyy = gt_eyy[wgbs_sts_exx_ind]
    eyy = eyy[exx_sts_ind]

    # Construct matrix
    res_grid = 0.05
    axis_rng = collect(0:res_grid:1.0)
    mat_y = fill(0.0, (length(axis_rng) - 1, length(axis_rng) - 1))
    mat_yy = fill(0.0, (length(axis_rng) - 1, length(axis_rng) - 1))
    for i in (length(axis_rng) - 1):-1:1
            for j in 1:(length(axis_rng) - 1)
            ind_1 = [axis_rng[i] <= x < axis_rng[i + 1] for x in ey]
            ind_2 = [axis_rng[j] <= x < axis_rng[j + 1] for x in gt_ey]
            mat_y[i,j] = sum(ind_1 .& ind_2)
            ind_1 = [axis_rng[i] <= x < axis_rng[i + 1] for x in eyy]
            ind_2 = [axis_rng[j] <= x < axis_rng[j + 1] for x in gt_eyy]
    mat_yy[i,j] = sum(ind_1 .& ind_2)
    end
    end

    # Plot
    mat_y .+= 1
    mat_yy .+= 1
    mat_y ./= sum(mat_y)
    mat_yy ./= sum(mat_yy)
    mat_y_log10 = -log10.(mat_y)
    mat_yy_log10 = -log10.(mat_yy)
    xlab = "Fine-grain (WGBS)"
    ylab = "Coarse-grain (Nanopore)"
    axis_labs = [isodd(i) ? "$(axis_rng[i])" : "" for i = 1:(length(axis_rng))]
    axis_ticks = (1:length(axis_labs), axis_labs)
    p_ex = heatmap(mat_y_log10, xlabel=xlab, ylabel=ylab, xticks=axis_ticks, yticks=axis_ticks,
        label="", c=cgrad([:white, :black]), clims=(1, 5), xrotation=45,
        title="E[X] (\\rho=$(round(cor(gt_ey, ey);digits=2)))");
    p_exx = heatmap(mat_yy_log10, xlabel=xlab, ylabel=ylab, xticks=axis_ticks, yticks=axis_ticks,
        label="", c=cgrad([:white, :black]), clims=(1, 5), xrotation=45,
        title="E[XX] (\\rho=$(round(cor(gt_eyy, eyy);digits=2)))");
    plot(p_ex, p_exx, layout=(1, 2), size=(1100, 500), top_margin=10px, bottom_margin=20px,
        left_margin=20px, right_margin=20px)
    savefig("$(nano_dir)/Heatmap-Exp-cov-$(cove)-s-$(sig)-Aim-2-Log-Res-$(res_grid).pdf")

    # Return nothing
    return nothing

end
#################################################################################################
# Calls
#################################################################################################

## Parameter estimates

# Generate plot array
plt_arr = gen_param_plt_array()

# Make plot
plot(plt_arr...,layout=(2, 2),size=(1000, 1000))
savefig("$(aim_dir)/Boxplots-Params-Aim-2.pdf")

# # Generate plot array
# plt_arr = gen_kl_plt_array()

# # Make plot
# plot(plt_arr...,layout=(2, 2),size=(1000, 1000))
# savefig("$(aim_dir)/Boxplots-KL-Params-Aim-2.pdf")

## E[X] vs X̄

# # Generate plot array
# plt_arr = gen_ey_plt_array()

# # Make plot
# plot(plt_arr...,layout=(2, 2),size=(1000, 1000))
# savefig("$(aim_dir)/Boxplots-EX-Aim-2.pdf")

## E[XX] vs X̄

# # Generate plot array
# plt_arr = gen_eyy_plt_array()

# # Make plot
# plot(plt_arr...,layout=(2, 2),size=(1000, 1000))
# savefig("$(aim_dir)/Boxplots-EXX-Aim-2.pdf")

## Heatmap plots for specific noise and cov

# Heatmap parameters
println("Generating: Scatter parameters")
plt_params_heatmap(10.0,3.0)
plt_params_heatmap(20.0,3.0)

# Heatmap Expectations
println("Generating: Scatter expectations")
plt_exp_heatmap(10.0,2.0)
plt_exp_heatmap(20.0,2.0)
plt_exp_heatmap(10.0,3.0)
plt_exp_heatmap(20.0,3.0)
plt_exp_heatmap(10.0,3.5)
plt_exp_heatmap(20.0,3.5)

## Analysis regions

Tmml_plt = gen_stat_plt_array("Tmml", 3.0)
Tnme_plt = gen_stat_plt_array("Tnme", 3.0)
Tcmd_plt = gen_stat_plt_array("Tcmd", 3.0)
plot(Tmml_plt,Tnme_plt,Tcmd_plt,layout=(3, 1),size=(1000, 2000))
savefig("$(aim_dir)/Boxplots-Analysis-Regions-Aim-2.pdf")


text:
include("../src/PowerFlowNetworks.jl")
using .PowerFlowNetworks
using SQLite
using Dates
using DataFrames
using ArgParse

ArgParse.parse_item(::Type{Vector{String}}, x::AbstractString) = split(x, ",")

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--dbpath", "-p"
            help = "Path of the sqlite database file."
            default = "data/PowerFlowNetworks.sqlite"
        "--source_type"
            help = "The source types to check."
            default = ["RAWGO", "MATPOWERM"]
            arg_type = Vector{String}
        "--rethrow"
            help = "Wether to rethrow the errors. (will stop the execution)"
            default = false
            action = :store_true
        "--one"
            help = "Check only one instance."
            default = false
            action = :store_true
        "--name"
            help = "Name of the instance to test. (Only used if `--one` is set)"
            default = ""
        "--scenario"
            help = "Scenario of the instance to test. (Only used if `--one` is set)"
            default = ""
    end
    return parse_args(s)
end

function check_network(name, scenario, source_path, source_type; rethrow_error)
    try
        network = PowerFlowNetwork(source_path, source_type)
    catch e
        println("$name scenario $scenario throwed: $e \n  source_path: $source_path\n  source_type: $source_type")
        rethrow_error && rethrow()
    end
end

check_network_dfrow(row; rethrow_error) = check_network(row[:name], row[:scenario], row[:source_path], row[:source_type]; rethrow_error=rethrow_error)

function check_networks(db, source_type; rethrow_error)
    source_type = map(s -> "'$s'", source_type)
    query = "SELECT name, scenario, source_path, source_type FROM instances WHERE source_type IN ($(join(source_type, ',')))"
    results = DBInterface.execute(db, query) |> DataFrame
    check_function(row) = check_network_dfrow(row; rethrow_error=rethrow_error)
    check_function.(eachrow(results[!, [:name, :scenario, :source_path, :source_type]]))
end

function main()
    parsed_args = parse_commandline()
    db = SQLite.DB(parsed_args["dbpath"])
    if parsed_args["one"]
        name, scenario = parsed_args["name"], parsed_args["scenario"]
        query = "SELECT source_path, source_type FROM instances WHERE name = '$name' AND scenario = $scenario"
        result = DBInterface.execute(db, query) |> DataFrame
        source_path, source_type = result[1, :source_path], result[1, :source_type]
        check_network(name, scenario, source_path, source_type; rethrow_error=true)
    else
        check_networks(db, parsed_args["source_type"]; rethrow_error=parsed_args["rethrow"])
    end
end

main()

ext: jl

max_stars_repo_name: charlyalizadeh/PowerFlowNetworks

max_issues_repo_name: charlyalizadeh/PowerFlowNetworks

max_forks_repo_name: charlyalizadeh/PowerFlowNetworks


text:
"""
    WeightedBinaryLoss{L,W} <: SupervisedLoss

Can an be used to represent a re-weighted version of some type of
binary loss `L`. The weight-factor `W`, which must be in `[0, 1]`,
denotes the relative weight of the positive class, while the
relative weight of the negative class will be `1 - W`.
For example: To create a typealias for a 1:4 weighted version of
`L2HingeLoss`, type:

```julia
const WeightedL2HingeLoss = LossFunctions.WeightedBinaryLoss{L2HingeLoss,0.2}
```

This new loss-type can then be instantiated in the same manner and
with the same parameters as the original unscaled loss-type.

In contrast, in order to only create a re-weighted instance of some
specific loss you can use `weightedloss(L2HingeLoss(), Val(0.2))`.
See `?weightedloss` for more information.
"""
struct WeightedBinaryLoss{L<:MarginLoss,W} <: SupervisedLoss
    loss::L
    WeightedBinaryLoss{L,W}(loss::L) where {L<:MarginLoss, W} = new{L,W}(loss)
end

@generated function (::Type{WeightedBinaryLoss{L,W}})(args...) where {L<:MarginLoss, W}
    typeof(W) <: Number && 0 <= W <= 1 || _werror()
    :(WeightedBinaryLoss{L,W}(L(args...)))
end

_werror() = throw(ArgumentError("The given \"weight\" has to be a number in the interval [0, 1]"))

@generated function WeightedBinaryLoss(loss::L, ::Val{W}) where {L<:MarginLoss,W}
    typeof(W) <: Number && 0 <= W <= 1 || _werror()
    :(WeightedBinaryLoss{L,W}(loss))
end

function WeightedBinaryLoss(loss::SupervisedLoss, w::Number)
    WeightedBinaryLoss(loss, Val(w))
end

@generated function WeightedBinaryLoss(s::WeightedBinaryLoss{T,W1}, ::Val{W2}) where {T,W1,W2}
    :(WeightedBinaryLoss(s.loss, Val($(W1*W2))))
end

for fun in (:value, :deriv, :deriv2)
    @eval function ($fun)(l::WeightedBinaryLoss{L,W}, target::Number, output::Number) where {L,W}
        # We interpret the W to be the weight of the positive class
        if target > 0
            W * ($fun)(l.loss, target, output)
        else
            (1-W) * ($fun)(l.loss, target, output)
        end
    end
end

"""
    weightedloss(loss, weight)

Returns a weighted version of `loss` for which the value of the
positive class is changed to be `weight` times its original, and the
negative class `1 - weight` times its original respectively.

Note: If `typeof(weight) <: Number`, then this method will poison the
type-inference of the calling scope. This is because `weight` will be
promoted to a type parameter. For a typestable version use the
following signature: `weightedloss(loss, Val(weight))`
"""
weightedloss(loss::Loss, weight::Number) = WeightedBinaryLoss(loss, weight)
weightedloss(loss::Loss, ::Val{W}) where {W} = WeightedBinaryLoss(loss, Val(W))

# An α-weighted version of a classification callibrated margin loss is
# itself classification callibrated if and only if α == 1/2
isclasscalibrated(l::WeightedBinaryLoss{T,W}) where {T,W} = W == 0.5 && isclasscalibrated(l.loss)

# TODO: Think about this semantic
issymmetric(::WeightedBinaryLoss) = false

for prop in [:isminimizable, :isdifferentiable,
             :istwicedifferentiable,
             :isconvex, :isstrictlyconvex,
             :isstronglyconvex, :isnemitski,
             :isunivfishercons, :isfishercons,
             :islipschitzcont, :islocallylipschitzcont,
             :isclipable, :ismarginbased,
             :isdistancebased]
    @eval ($prop)(l::WeightedBinaryLoss) = ($prop)(l.loss)
end

for prop_param in (:isdifferentiable, :istwicedifferentiable)
    @eval ($prop_param)(l::WeightedBinaryLoss, at) = ($prop_param)(l.loss, at)
end


text:
"""
    escape(val)

Escape backslash, dollar and quote from string `val`.
"""
function escape(val)
     valescaped = val
     # backslash must come first
     for c in ['\\','$','"']
        valescaped = replace(valescaped,c => "\\$c")
    end
	return valescaped
end

"""
    escape(val)

Escape string for variable names.
"""
function escapevar(val)
    return replace(val," " => "_")
end

function ncgen(io::IO,fname; newfname = "filename.nc")
    ds = NCDataset(fname)
    unlimited_dims = unlimited(ds.dim)
    print(io,"using NCDatasets, DataStructures\n")
    print(io,"ds = NCDataset(\"$(escape(newfname))\",\"c\"")
    ncgen_setattrib(io,ds.attrib)
    print(io,")\n\n")

    print(io,"# Dimensions\n\n")
    for (d,v) in ds.dim
        if d in unlimited_dims
            print(io,"ds.dim[\"$d\"] = Inf # unlimited dimension\n")
        else
            print(io,"ds.dim[\"$d\"] = $v\n")
        end
    end

    print(io,"\n# Declare variables\n\n")

    for (d,v) in ds
        print(io,"nc$(escapevar(d)) = defVar(ds,\"$d\", $(eltype(v.var)), $(dimnames(v))")
        ncgen_setattrib(io,v.attrib)
        print(io,")\n\n")
    end

    print(io,"\n# Define variables\n\n")

    for d in keys(ds)
        print(io,"# nc$(escapevar(d))[:] = ...\n")
    end

    print(io,"\nclose(ds)\n")
    close(ds)
end

"""
    ncgen(fname; ...)
    ncgen(fname,jlname; ...)

Generate the Julia code that would produce a NetCDF file with the same metadata
as the NetCDF file `fname`. The code is placed in the file `jlname` or printed
to the standard output. By default the new NetCDF file is called `filename.nc`.
This can be changed with the optional parameter `newfname`.
"""
ncgen(fname; kwargs...)  = ncgen(stdout,fname; kwargs...)

function ncgen(fname,jlname; kwargs...)
    open(jlname,"w") do io
        ncgen(io, fname; kwargs...)
    end
end
export ncgen

litteral(val::String) = "\"$(escape(val))\""
litteral(val::Float64) = val
litteral(val::Number) = "$(eltype(val))($(val))"
litteral(val) = "$(val)" # for arrays

function ncgen_setattrib(io,attrib)
    if length(attrib) == 0
        return
    end

    print(io,", attrib = OrderedDict(\n")

    for (d,val) in attrib
        litval = litteral(val)
        print(io,"    \"$d\"" * (" "^max(0,(25-length(d)))) * " => $litval,\n");
    end
    print(io,")")

end


text:
baremodule TaskDAGAnalyzers

function dag end
function span end
function work end
function parallelism end
function summary end

module Internal

using ..TaskDAGAnalyzers: TaskDAGAnalyzers

using TaskDAGRecorders
using TaskDAGRecorders.Internal:
    AnyContext, SyncContext, TaskContext, TaskStat, _get!, get_sync_context

import ShowGraphviz

include("utils.jl")
include("dag.jl")
include("graphviz.jl")
include("summary.jl")

end  # module Internal

module Examples
include("examples.jl")
end  # module Examples

Internal.define_docstring()

end  # baremodule TaskDAGAnalyzers


text:
# Compute RSS, which is the squared Euclidean distance between every point
# and its assigned center
function rss{S <: Real, T <: Real}(x::Matrix{S},
                                   centers::Matrix{T},
                                   assignments::Vector{Int},
                                   n::Integer,
                                   p::Integer)
    residuals = 0.0

    for i = 1:n
        for j in 1:p
            residuals += (x[i, j] - centers[assignments[i], j])^2
        end
    end

    return residuals
end

# Find shortest distance from x to any of n_centers existing clusters
function dsquared{S <: Real, T <: Real}(x::Matrix{S},
                                        centers::Matrix{T},
                                        n_centers::Integer,
                                        p::Integer)
    m = Inf

    for i in 1:n_centers
        dst = 0.0
        for j in 1:p
            dst += (centers[i, j] - x[1, j])^2
        end
        if dst < m
            m = dst
        end
    end

    return m
end

# The k-means rule
function k_means_initialization!{S <: Real, T <: Real}(x::Matrix{S},
                                                       centers::Matrix{T},
                                                       n::Integer,
                                                       p::Integer,
                                                       k::Integer)
    for cluster in 1:k
        centers[cluster, :] = x[rand(1:n), :]
    end
end

# The k-means++ rule
function k_meanspp_initialization!{S <: Real, T <: Real}(x::Matrix{S},
                                                         centers::Matrix{T},
                                                         n::Integer,
                                                         p::Integer,
                                                         k::Integer)
    # Pre-allocate variables
    dsquareds = Array(Float64, n)
    dstr = Categorical(n)

    # Choose an initial cluster center uniformly at random from data points
    c1 = rand(1:n)
    centers[1, :] = x[c1, :]

    for cluster in 2:k
        s = 0.0
        for i in 1:n
            Dsq = dsquared(x[i, :], centers, cluster - 1, p)
            dsquareds[i] = Dsq
            s += Dsq
        end
        dstr.prob = dsquareds / s
        selected = rand(dstr)
        centers[cluster, :] = x[selected, :]
    end
end

# Update cluster assignments
function update_assignments!{S <: Real, T <: Real}(x::Matrix{S},
                                                   centers::Matrix{T},
                                                   distances::Vector{Float64},
                                                   assignments::Vector{Int},
                                                   n::Integer,
                                                   p::Integer,
                                                   k::Integer)
    for i = 1:n
        for cluster in 1:k
            d = 0.0
            for j in 1:p
                d += (x[i, j] - centers[cluster, j])^2
            end
            distances[cluster] = d
        end

        assignments[i] = indmin(distances)
    end
end

# Update cluster centers
function update_centers!(x::Matrix,
                         centers::Matrix,
                         assignments::Vector{Int},
                         n::Integer,
                         p::Integer,
                         k::Integer)
    for cluster = 1:k
        for j in 1:p
            centers[cluster, j] = 0.0
        end

        n_obs = 0
        for i in 1:n
            if assignments[i] == cluster
                n_obs += 1
                for j in 1:p
                    centers[cluster, j] += x[i, j]
                end
            end
        end

        if n_obs > 0
            for j in 1:p
                centers[cluster, j] /= n_obs
            end
        end
    end
end

function k_means{T <: Real}(x::Matrix{T}, k::Integer, opts::Options)
    # Set default options
    @defaults opts tolerance => 1.0e-8
    @defaults opts max_iter => 1000
    @defaults opts initialize_centers! => k_meanspp_initialization!

    # Keep track of the number of data points and their dimensionality
    n, p = size(x)

    # Run until convergence or until the number of iterations is too high
    converged, iter = false, 0

    # Use current_rss and previous_rss to measure convergence
    previous_rss, current_rss = Inf, Inf

    # Pre-allocate arrays that are used repeatedly
    centers = Array(Float64, k, p)
    assignments = Array(Int, n)
    distances = Array(Float64, k)

    # Initialize the cluster centers
    initialize_centers!(x, centers, n, p, k)

    # Initialize the assignment of points to clusters
    update_assignments!(x, centers, distances, assignments, n, p, k)

    # Iterate until convergence
    while !converged && iter < max_iter
        # Increment the iteration counter
        iter += 1

        # Recompute cluster assignments given current centers
        update_centers!(x, centers, assignments, n, p, k)

        # Reassign points to the closest cluster
        update_assignments!(x, centers, distances, assignments, n, p, k)

        # Update the RSS values for previous and current cluster assignments
        previous_rss = current_rss
        current_rss = rss(x, centers, assignments, n, p)

        # Assess convergence
        if abs(previous_rss - current_rss) <= tolerance
            converged = true
        end
    end

    return KMeansOutput(assignments,
                        centers,
                        iter,
                        current_rss,
                        converged)
end

function k_means{T <: Real}(x::Matrix{T}, k::Integer)
    k_means(x, k, Options())
end


text:
using .Networks

using Dates

################################################################################
####       TSO
################################################################################

"""
utilisé pour le mode 3:
Prend des décisions fermes vu que c'est la dernière execution du TSO
Décide de la reserve
"""
struct TSOAtFOBiLevel <: AbstractTSO
end
function run(runnable::TSOAtFOBiLevel, ech::Dates.DateTime, firmness, TS::Vector{Dates.DateTime}, context::AbstractContext)
    error("unimplemented")
    return #result
end

#### TSO COMMON :
function update_tso_schedule!(context::AbstractContext, ech, result, firmness,
                            runnable::AbstractTSO)
    tso_schedule = get_tso_schedule(context)
    tso_schedule.decider_type = DeciderType(runnable)
    tso_schedule.decision_time = ech
    error("unimplemented")
end
function update_tso_actions!(context::AbstractContext, ech, result, firmness,
                            runnable::AbstractTSO)
    error("unimplemented")
end

################################################################################
####       Firmness
################################################################################

function compute_firmness(runnable::AbstractRunnable,
                    ech::Dates.DateTime, next_ech::Union{Nothing,Dates.DateTime},
                    TS::Vector{Dates.DateTime}, context::AbstractContext)
    @debug "compute decisions firmness"
    generators = collect(Networks.get_generators(get_network(context)))
    return compute_firmness(ech, next_ech, TS, generators)
end

"""
    All decisions are Firm (DECIDED or TO_DECIDE)
"""
function compute_firmness(runnable::Union{EnergyMarketAtFO,TSOAtFOBiLevel},
                    ech::Dates.DateTime, next_ech::Union{Nothing,Dates.DateTime},
                    TS::Vector{Dates.DateTime}, context::AbstractContext)
    @debug "compute decisions firmness : only firm decisions (DECIDED or TO_DECIDE)"
    next_ech = nothing
    generators = collect(Networks.get_generators(get_network(context)))
    return compute_firmness(ech, next_ech, TS, generators)
end

################################################################################
####       Utils
################################################################################

struct Assessment <: AbstractRunnable
end
function run(runnable::Assessment, ech::Dates.DateTime, firmness, TS::Vector{Dates.DateTime}, context::AbstractContext)
    return #result
end

struct EnterFO <: AbstractRunnable
end
function run(runnable::EnterFO, ech::Dates.DateTime, firmness, TS::Vector{Dates.DateTime}, context::AbstractContext)
    println("-----Entrée dans la fenêtre opérationnelle-----")
    return nothing #result
end
function compute_firmness(runnable::EnterFO,
                        ech::Dates.DateTime, next_ech::Union{Nothing,Dates.DateTime},
                        TS::Vector{Dates.DateTime}, context::AbstractContext)
    return Firmness()
end

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf


text:
# Control function for 2D Fisher-Stefan level-set solutions
# Alex Tam, 02/11/2021

# Load packages
using Parameters
using DifferentialEquations
using DelimitedFiles

# Include external files
include("domain.jl")
include("fkpp.jl")
include("velocity.jl")
include("level-set.jl")
include("reinitialisation.jl")

"Data structure for model parameters"
@with_kw struct Params
    D::Float64 = 1.0 # [-] Diffusion coefficient
    λ::Float64 = 1.0 # [-] Reaction rate
    κ::Float64 = 0.1 # [-] Stefan parameter
    u_f::Float64 = 0.0 # [-] Density at interface
    u_b::Float64 = 0.0 # [-] Density at computational boundary
    α::Float64 = 0.5 # [-] Maximum initial density
    β::Float64 = 10.0 # [-] Initial interface radius
    θb::Float64 = 0.01 # [-] Threshold for whether a grid point is close to interface
    θ::Float64 = 1.99 # [-] Parameter for minmod flux-limiter
    Lx::Float64 = 20.0 # [-] Spatial domain limit (x)
    Ly::Float64 = 20.0 # [-] Spatial domain limit (y)
    T::Float64 = 100.0 # [-] End time
    Nx::Int = 201 # [-] Number of grid points (x)
    Ny::Int = 201 # [-] Number of grid points (y)
    Nt::Int = 50001 # [-] Number of time steps
    V_Iterations::Int = 15 # [-] Number of iterations for velocity extrapolation PDE
    ϕ_Iterations::Int = 20 # [-] Number of iterations for reinitialisation PDE
end

"Generate initial density and level-set function"
function ic(par::Params, x, y)
    U = Array{Float64}(undef, par.Nx, par.Ny) # Pre-allocate 2D array of U
    ϕ = Array{Float64}(undef, par.Nx, par.Ny) # Pre-allocate 2D array of ϕ
    # Loop over grid points
    for i = 1:length(x)
        for j = 1:length(y)
            # Circle
            if sqrt((x[i]-par.Lx/2)^2 + (y[j]-par.Ly/2)^2) <= par.β
                U[i,j] = par.α # Constant density (inside Ω)
            else
                U[i,j] = par.u_f # Initial density (outside Ω)
            end
            ϕ[i,j] = (sqrt((x[i]-par.Lx/2)^2 + (y[j]-par.Ly/2)^2) - par.β) # Initial signed-distance
            # # Rectangle
            # if (x[i] >= par.Lx/2 - 4.0/2) && (x[i] <= par.Lx/2 + 4.0/2) && (y[j] >= par.Ly/2 - 4.0/2) && (y[j] <= par.Ly/2 + 4.0/2)
            #     U[i,j] = par.α
            #     ϕ[i,j] = -0.5
            # else
            #     U[i,j] = par.u_f
            #     ϕ[i,j] = 0.5
            # end
        end
    end
    return U, ϕ
end

"Build a vector of 2D data, ordered by entries in D"
function build_vector(U::Array{Float64}, D)
    u = Vector{Float64}() # Pre-allocate empty vector
    for gp in D
        push!(u, U[gp.xInd, gp.yInd])
    end
    return u
end

"Build 2D array of data from vector"
function build_u_matrix(u::Vector, par, D)
    U = par.u_b*ones(par.Nx, par.Ny) # Pre-allocate (incorporate a Dirichlet condition on computational boundary)
    for i = 1:length(D)
        U[D[i].xInd, D[i].yInd] = u[i]
    end
    return U
end

"Build 2D array of data from vector"
function build_v_matrix(v::Vector, par, D)
    V = zeros(par.Nx, par.Ny) # Pre-allocate (incorporate a Dirichlet condition on computational boundary)
    for i = 1:length(D)
        V[D[i].xInd, D[i].yInd] = v[i]
    end
    return V
end

function front_position(x, y, ϕ, nx, ny, dx, dy)
    L = Vector{Float64}() # Initialise
    # Find front position using x-direction slice
    ϕv = ϕ[:,ny]
    for i = 1:length(ϕv)
        if (ϕv[i] < 0) && (ϕv[i+1] >= 0)
            θ = ϕv[i]/(ϕv[i] - ϕv[i+1])
            push!(L, x[i] + θ*dx)
        end
    end
    # Find front position using y-direction slice
    ϕv = ϕ[nx,:]
    for i = 1:length(ϕv)
        if (ϕv[i] < 0) && (ϕv[i+1] >= 0)
            θ = ϕv[i]/(ϕv[i] - ϕv[i+1])
            push!(L, y[i] + θ*dy)
        end
    end
    return L[1], L[2]
end

"Compute a solution"
function fisher_stefan_2d()
    # Parameters and domain
    par = Params() # Initialise data structure of model parameters
    nx::Int = (par.Nx-1)/2; ny::Int = (par.Ny-1)/2; # Indices for slice plots
    x = range(0, par.Lx, length = par.Nx); dx = x[2] - x[1] # Computational domain (x)
    y = range(0, par.Ly, length = par.Ny); dy = y[2] - y[1] # Computational domain (y)
    t = range(0, par.T, length = par.Nt); dt = t[2] - t[1] # Time domain
    writedlm("x.csv", x); writedlm("y.csv", y); writedlm("t.csv", t) # Write data to files
    # Initial condition
    U, ϕ = ic(par, x, y) # Obtain initial density and ϕ
    ϕ = reinitialisation(ϕ, par, dx, dy, 100)
    writedlm("U-0.csv", U); writedlm("Phi-0.csv", ϕ) # Write data to files
    Lx = Vector{Float64}() # Preallocate empty vector of interface position
    Ly = Vector{Float64}() # Preallocate empty vector of interface position
    plot_times = Vector{Int}() # Vector of time-steps at which data is obtained
    push!(Lx, par.Lx/2 + par.β); push!(Ly, par.Ly/2 + par.β)
    # Time stepping
    for i = 1:par.Nt-1
        # 1. Find Ω, dΩ, and irregular grid points
        D = find_domain(par, ϕ)
        dΩ = find_interface(D, ϕ)
        # 2. Solve FKPP equation on Ω
        U = fkpp(D, U, ϕ, par, dx, dy, dt, i)
        # 3. Compute extension velocity field
        V = extend_velocity(D, dΩ, U, ϕ, par, dx, dy)
        # 4. Solve level-set equation
        ϕ = level_set(V, ϕ, par, dx, dy, dt)
        # 5. Re-initialise level-set function as a signed-distance
        if mod(i, 10) == 0
            ϕ = reinitialisation(ϕ, par, dx, dy, par.ϕ_Iterations)
        end
        # Optional: Post-processing
        if mod(i, 5000) == 0
            writedlm("ux-$i.csv", U[:,ny])
            writedlm("uy-$i.csv", U[nx,:])
            writedlm("U-$i.csv", U)
            writedlm("V-$i.csv", V)
            writedlm("Phi-$i.csv", ϕ)
            push!(plot_times, i)
            writedlm("plot_times.csv", plot_times)
        end
        px, py = front_position(x, y, ϕ, nx, ny, dx, dy)
        push!(Lx, px); push!(Ly, py)
        writedlm("Lx.csv", Lx)
        writedlm("Ly.csv", Ly)
    end
end

@time fisher_stefan_2d()


text:
include("../../src/Random-Pertubation-on-Constrained-Manifolds/system.jl")
include("pd_basin_stability.jl")

#bsode, varode = BasinStability(powergrid,28,:ω,100, [-π,π], [-100.0,100.0], ProjectionMethod = "ODE")
#bsoptim, varoptim = BasinStability(powergrid,28,:ω,100, [-π,π], [-100.0,100.0], ProjectionMethod = "Optim")

bs, var = BasinStability(powergrid,28.0,:ω,1, [-π,π], [-100.0,100.0], ProjectionMethod = "Optim", dae = true)
odebs, varode = BasinStability(powergrid,28,:ω,1, [-π,π], [-100.0,100.0], ProjectionMethod = "Optim")

array = RemoveNodes(powergrid,:θ)
PlotBasinStability(array, bs, var, labtext = "Optim Projection, ODE")
PlotBasinStability!(array, odebs, varode, labtext = "Optim Projection, DAE")
png("BasinStability_Plot_ode_dae")

using DelimitedFiles
writedlm("var_ode.txt", varode)
writedlm("var_optim.txt", varoptim)

using GraphPlot

using LightGraphs
using Cairo, Compose
draw(PNG("powergrid.png", 16cm, 16cm), gplot(powergrid.graph))

vec = powergrid.nodes |> collect

types = typeof.(vec)

max_stars_repo_name: luap-pik/PowerDynIC

max_issues_repo_name: luap-pik/PowerDynIC

max_forks_repo_name: luap-pik/PowerDynIC


text:
# TODO
# - find need allocations in attractive case
# - decimate allocations here
# - I guess maybe it's additional Diagonal's etc?

################################################################################
### Type and general utility
################################################################################

"""
    BlockDiagonal(blocks...)
    BlockDiagonal{eltype, N_blocks, MatrixType}(init, n, m)

Creates a block diagonal matrix from a collection of blocks or an initializer,
i.e. `undef` or `I`. In the latter case `n, m` specifies the size of the full
block diagonal matrix.

Notes:
- `vmul!` and friends are implemented for real and complex `BlockDiagonal`
- some base methods are implemented - `Matrix`, `getidnex`, `setindex`, `copyto!`,
`size`, `copy`, `log`, `exp`, `det`, `*`
- BlockDiagonal of 2 equally sized square matrices is heavily tested and used
- BlockDiagonal of more equally sized square matrices should work
- irregularly sized and non-square matrices sometimes work (TODO)
- methods generally assume correct size and use @inbounds for performance.
"""
struct BlockDiagonal{T, N, AT <: AbstractMatrix{T}} <: AbstractMatrix{T}
    blocks::NTuple{N, AT}

    function BlockDiagonal(blocks::NTuple{N, AT}) where {T, N, AT <: AbstractMatrix{T}}
        n, m = size(first(blocks))
        n == m || throw(ErrorException("Each block must be square."))
        all(b -> size(b) == (n, n), blocks) || throw(ErrorException(
            "All Blocks must have the same size."
        ))
        # One matrix with zeros is faster than multiple 3x3 matrices on my
        # machine (avx2). This will most likely depend on AVX version...
        # if n < 4
        #     full = zeros(T, n*N, n*N)
        #     for i in 1:N
        #         @views copyto!(full[(i-1)*n+1 : i*n, (i-1)*n+1 : i*n], blocks[i])
        #     end
        #     # new{T, 1, AT}((full,))
        #     # Is this bad? - Yes, because the rest of the code doesn't know about Matrix returns
        #     return full
        # else
            return new{T, N, AT}(blocks)
        # end
    end
end

BlockDiagonal(blocks::AT...) where {T, AT <: AbstractMatrix{T}} = BlockDiagonal(blocks)
function BlockDiagonal{T, N, AT}(::UndefInitializer, n, m) where {T, N, AT}
    n == m || throw(ErrorException("Each block must be square."))
    (n, m) = divrem(n, N)
    m == 0 || throw(ErrorException(
        "The size of the Matrix is incompatible with the blocking ($n % $N != 0)"
    ))
    BlockDiagonal([AT(undef, n, n) for _ in 1:N]...)
end
function BlockDiagonal{T, N, AT}(::UniformScaling, n, m) where {T, N, AT}
    n == m || throw(ErrorException("Each block must be square."))
    (n, m) = divrem(n, N)
    m == 0 || throw(ErrorException(
        "The size of the Matrix is incompatible with the blocking ($n % $N != 0)"
    ))
    BlockDiagonal([AT(I, n, n) for _ in 1:N]...)
end

function Base.Matrix(B::BlockDiagonal{T, N}) where {T, N}
    n = size(B.blocks[1], 1)
    M = N * n
    output = zeros(T, M, M)
    for i in 1:N
        @views copyto!(output[(i-1)*n+1 : i*n, (i-1)*n+1 : i*n], B.blocks[i])
    end
    output
end

function Base.checkbounds(B::BlockDiagonal, i::Integer, j::Integer)
    N = M = 0
    @inbounds for b in B.blocks
        N += size(b, 1)
        M += size(b, 2)
    end
    (1 ≤ i ≤ N) && (1 ≤ j ≤ M)
end

function Base.getindex(B::BlockDiagonal{T, N}, i, j) where {T, N}
    @boundscheck checkbounds(B, i, j)
    @inbounds n = size(B.blocks[1], 1)

    b, k = divrem(i-1, n)
    _b, l = divrem(j-1, n)

    if b != _b
        return zero(T)
    else
        @inbounds return B.blocks[b+1][k+1, l+1]
    end
end

function Base.setindex!(B::BlockDiagonal, val, i, j)
    i0 = j0 = 0
    for block in B.blocks
        i1 = size(block, 1) + i0
        j1 = size(block, 2) + j0
        if (i0 < i <= i1) && (j0 < j <= j1)
            @inbounds block[i - i0, j - j0] = val
            return B
        end
        i0 = i1; j0 = j1
    end
    throw(BoundsError(B, (i, j)))
end

function Base.copyto!(B1::BlockDiagonal{T, N}, B2::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        copyto!(B1.blocks[i], B2.blocks[i])
    end
    return B1
end
function Base.copyto!(B::BlockDiagonal, ::UniformScaling)
    @inbounds for block in B.blocks
        copyto!(block, I)
    end
    return B
end
function Base.copyto!(B::BlockDiagonal, D::Diagonal)
    n = 0
    @inbounds for block in B.blocks
        s = 1 + n
        n += size(block, 1)
        @views copyto!(block, Diagonal(D.diag[s:n]))
    end
    return B
end

function Base.copyto!(output::Matrix, B::BlockDiagonal{T, N}) where {T, N}
    n = size(B.blocks[1], 1)
    for i in 1:N
        @views copyto!(output[(i-1)*n+1 : i*n, (i-1)*n+1 : i*n], B.blocks[i])
    end
    output
end

# for Base.show
Base.size(B::BlockDiagonal) = mapreduce(b -> size(b), (a, b) -> a .+ b, B.blocks)

# Needed for tests
Base.copy(B::BlockDiagonal) = deepcopy(B)

# Needed for stack building (shouldn't be performance critical)
function Base.:(*)(x::Number, B::BlockDiagonal)
    BlockDiagonal(map(block -> x * block, B.blocks)...)
end
Base.:(*)(B::BlockDiagonal, x::Number) = x * B
function Base.:(*)(B1::BlockDiagonal{T1, N}, B2::BlockDiagonal{T2, N}) where {T1, T2, N}
    BlockDiagonal(map(*, B1.blocks, B2.blocks)...)
end

function Base.exp(B::BlockDiagonal)
    BlockDiagonal(map(block -> exp(block), B.blocks)...)
end

# This takes super long to compile...?
# Dot syntax is another 6x slower
# function Base.log(B::BlockDiagonal)
#     BlockDiagonal(map(block -> log(block), B.blocks))
# end

# function LinearAlgebra.det(B::BlockDiagonal{T}) where {T}
#     output = T(1)
#     for b in B.blocks
#         output *= det(b)
#     end
#     output
# end

# I thought this would be needed for greens(k, l), but it's not?
function LinearAlgebra.transpose!(A::BlockDiagonal{T, N}, B::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        transpose!(A.blocks[i], B.blocks[i])
    end
    A
end
# This however is used for greens(k, l)
function LinearAlgebra.rmul!(B::BlockDiagonal, f::Number)
    @inbounds for block in B.blocks
        rmul!(block, f)
    end
    B
end

################################################################################
### AVX-powered matrix multiplications
################################################################################

function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, B::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        vmul!(C.blocks[i], A.blocks[i], B.blocks[i])
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, B::Diagonal) where {T, N}
    # Assuming correct size
    n = 0
    @inbounds for i in 1:N
        c = C.blocks[i]
        a = A.blocks[i]
        s = 1+n
        n += size(a, 1)
        @views vmul!(c, a, Diagonal(B.diag[s:n]))
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::Diagonal, B::BlockDiagonal{T, N}) where {T, N}
    # Assuming correct size
    n = 0
    @inbounds for i in 1:N
        c = C.blocks[i]
        b = B.blocks[i]
        s = 1+n
        n += size(b, 1)
        @views vmul!(c, Diagonal(A.diag[s:n]), b)
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, X::Adjoint{T}) where {T, N}
    B = X.parent
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        c = C.blocks[i]
        vmul!(c, a, Adjoint(b))
    end
end
function vmul!(C::BlockDiagonal{T, N}, X::Adjoint{T}, B::BlockDiagonal{T, N}) where {T, N}
    A = X.parent
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        c = C.blocks[i]
        vmul!(c, Adjoint(a), b)
    end
end
function vmul!(C::BD, X1::Adjoint{T, BD}, X2::Adjoint{T, BD}) where {T, N, BD <: BlockDiagonal{T, N}}
    A = X1.parent
    B = X2.parent
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        c = C.blocks[i]
        vmul!(c, Adjoint(a), Adjoint(b))
    end
end
function rvmul!(A::BlockDiagonal, B::Diagonal)
    # Assuming correct size
    n = 0
    @inbounds for block in A.blocks
        s = 1+n
        n += size(block, 1)
        @views rvmul!(block, Diagonal(B.diag[s:n]))
    end
end
function lvmul!(A::Diagonal, B::BlockDiagonal)
    # Assuming correct size
    n = 0
    @inbounds for block in B.blocks
        s = 1+n
        n += size(block, 1)
        @views lvmul!(Diagonal(A.diag[s:n]), block)
    end
end

# used in greens(k, l)
function rvadd!(A::BlockDiagonal{T, N}, B::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        rvadd!(a, b)
    end
end
# used in equal time greens
function rvadd!(B::BlockDiagonal{T, N}, D::Diagonal) where {T, N}
    # Assuming correct size
    n = 0
    @inbounds for block in B.blocks
        s = 1+n
        n += size(block, 1)
        @views rvadd!(block, Diagonal(D.diag[s:n]))
    end
end
# used in CombinedGreensIterator
function vsub!(O::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, ::UniformScaling) where {T, N}
    @inbounds for i in 1:N
        vsub!(O.blocks[i], A.blocks[i], I)
    end
end

function rdivp!(A::BD, T::BD, O::BD, pivot) where {ET<:Real, N, BD <: BlockDiagonal{ET, N}}
    # assume Diagonal is ±1!
    @inbounds begin
        n = size(A.blocks[1], 1)

        # Apply pivot
        for b in 1:N
            o = O.blocks[b]
            a = A.blocks[b]
            t = T.blocks[b]
            offset = (b-1)*n

            for j in 1:n
                p = pivot[j + offset]
                @turbo for i in 1:n
                    o[i, j] = a[i, p]
                end
            end

            # do the rdiv
            # @turbo will segfault on `k in 1:0`, so pull out first loop
            invt11 = 1.0 / t[1, 1]
            @turbo for i in 1:n
                a[i, 1] = o[i, 1] * invt11
            end
            for j in 2:n
                invtjj = 1.0 / t[j, j]
                @turbo for i in 1:n
                    x = o[i, j]
                    for k in 1:j-1
                        x -= a[i, k] * t[k, j]
                    end
                    a[i, j] = x * invtjj
                end
            end
        end
    end
    A
end

################################################################################
### Fallbacks
################################################################################

function rdivp!(A::BD, T::BD, O::BD, pivot) where {ET, N, BD <: BlockDiagonal{ET, N}}
    # Assuming correct size
    n = 0
    @inbounds for i in 1:N
        s = 1+n
        n += size(A.blocks[i], 1)
        @views rdivp!(A.blocks[i], T.blocks[i], O.blocks[i], pivot[s:n])
    end
    A
end

################################################################################
### UDT
################################################################################

function udt_AVX_pivot!(
        U::BlockDiagonal{Float64, N, AT},
        D::AbstractArray{Float64, 1},
        input::BlockDiagonal{Float64, N, AT},
        pivot::AbstractArray{Int64, 1} = Vector(UnitRange(1:size(input, 1))),
        temp::AbstractArray{Float64, 1} = Vector{Float64}(undef, length(D)),
        apply_pivot::Val = Val(true)
    ) where {N, AT <: AbstractMatrix{Float64}}
    # Assuming correct size
    n = 0
    @inbounds for i in 1:N
        s = 1+n
        n += size(U.blocks[i], 1)
        @views udt_AVX_pivot!(
            U.blocks[i], D[s:n], input.blocks[i],
            pivot[s:n], temp[s:n], apply_pivot
        )
    end
end
function udt_AVX_pivot!(
        U::BlockDiagonal{ComplexF64, N, AT},
        D::AbstractArray{Float64, 1},
        input::BlockDiagonal{ComplexF64, N, AT},
        pivot::AbstractArray{Int64, 1} = Vector(UnitRange(1:size(input, 1))),
        temp::AbstractArray{ComplexF64, 1} = Vector{ComplexF64}(undef, length(D)),
        apply_pivot::Val = Val(true)
    ) where {N, AT <: AbstractMatrix{ComplexF64}}
    # Assuming correct size
    n = 0
    @inbounds for i in 1:N
        s = 1+n
        n += size(U.blocks[i], 1)
        @views udt_AVX_pivot!(
            U.blocks[i], D[s:n], input.blocks[i],
            pivot[s:n], temp[s:n], apply_pivot
        )
    end
end


text:
using Core.Compiler: AbstractInterpreter, NativeInterpreter, InferenceState,
    OptimizationState, CodeInfo, CodeInstance, InferenceResult, WorldRange,
    IRCode, SSAValue, inlining_policy

struct InferredSource
    src::CodeInfo
    stmt_infos::Vector{Any}
    rt::Any
end

struct OptimizedSource
    ir::IRCode
    isinlineable::Bool
end

mutable struct CthulhuInterpreter <: AbstractInterpreter
    native::NativeInterpreter

    unopt::Dict{Union{MethodInstance, InferenceResult}, InferredSource}
    opt::Dict{MethodInstance, CodeInstance}

    msgs::Dict{MethodInstance, Vector{Pair{Int, String}}}
end

CthulhuInterpreter() = CthulhuInterpreter(
    NativeInterpreter(),
    Dict{MethodInstance, InferredSource}(),
    Dict{MethodInstance, CodeInstance}(),
    Dict{MethodInstance, Vector{Tuple{Int, String}}}()
)

import Core.Compiler: InferenceParams, OptimizationParams, get_world_counter,
    get_inference_cache, code_cache,
    WorldView, lock_mi_inference, unlock_mi_inference, InferenceState

InferenceParams(interp::CthulhuInterpreter) = InferenceParams(interp.native)
OptimizationParams(interp::CthulhuInterpreter) = OptimizationParams(interp.native)
get_world_counter(interp::CthulhuInterpreter) = get_world_counter(interp.native)
get_inference_cache(interp::CthulhuInterpreter) = get_inference_cache(interp.native)

# No need to do any locking since we're not putting our results into the runtime cache
lock_mi_inference(interp::CthulhuInterpreter, mi::MethodInstance) = nothing
unlock_mi_inference(interp::CthulhuInterpreter, mi::MethodInstance) = nothing

code_cache(interp::CthulhuInterpreter) = interp.opt
Core.Compiler.get(a::Dict, b, c) = Base.get(a,b,c)
Core.Compiler.get(a::WorldView{<:Dict}, b, c) = Base.get(a.cache,b,c)
Core.Compiler.haskey(a::Dict, b) = Base.haskey(a, b)
Core.Compiler.haskey(a::WorldView{<:Dict}, b) =
    Core.Compiler.haskey(a.cache, b)
Core.Compiler.setindex!(a::Dict, b, c) = setindex!(a, b, c)
Core.Compiler.may_optimize(interp::CthulhuInterpreter) = true
Core.Compiler.may_compress(interp::CthulhuInterpreter) = false
Core.Compiler.may_discard_trees(interp::CthulhuInterpreter) = false
Core.Compiler.verbose_stmt_info(interp::CthulhuInterpreter) = true

function Core.Compiler.add_remark!(interp::CthulhuInterpreter, sv::InferenceState, msg)
    push!(get!(interp.msgs, sv.linfo, Tuple{Int, String}[]),
        sv.currpc => msg)
end

function Core.Compiler.finish(state::InferenceState, interp::CthulhuInterpreter)
    r = invoke(Core.Compiler.finish, Tuple{InferenceState, AbstractInterpreter}, state, interp)
    interp.unopt[Core.Compiler.any(state.result.overridden_by_const) ? state.result : state.linfo] = InferredSource(
        copy(isa(state.src, OptimizationState) ?
            state.src.src : state.src),
        copy(state.stmt_info),
        state.result.result)
    return r
end

function Core.Compiler.transform_result_for_cache(interp::CthulhuInterpreter, linfo::MethodInstance,
        valid_worlds::Core.Compiler.WorldRange, @nospecialize(inferred_result))
    if isa(inferred_result, OptimizationState)
        opt = inferred_result
        if isdefined(opt, :ir)
            return OptimizedSource(opt.ir, opt.src.inlineable)
        end
    end
    return inferred_result
end

function Core.Compiler.inlining_policy(interp::CthulhuInterpreter)
    function (src)
        @assert isa(src, OptimizedSource)
        return src.isinlineable ? src.ir : nothing
    end
end

function Core.Compiler.finish!(interp::CthulhuInterpreter, caller::InferenceResult)
    if isa(caller.src, OptimizationState)
        opt = caller.src
        if isdefined(opt, :ir)
            caller.src = OptimizedSource(opt.ir, opt.src.inlineable)
        end
    end
end


text:
type sdb
  susceptible::Array{Bool,2}
  infectious::Array{Bool,2}
  recovered::Array{Bool,2}
  new_infection::Array{Bool,2}
  new_recovery::Array{Bool,2}
  event_times::Array{Float64,1}
  event_type::Array{ASCIIString,1}
end

"""
Save the likelihood function from repetively determining state,
by producing an array which contains this information for all
relevant time steps
"""
function state_array(event_db::edb)
  sa=sdb(fill(false, (size(event_db.events)[1], length(unique(event_db.event_times[event_db.event_times.<Inf])))), fill(false, (size(event_db.events)[1], length(unique(event_db.event_times[event_db.event_times.<Inf])))), fill(false, (size(event_db.events)[1], length(unique(event_db.event_times[event_db.event_times.<Inf])))), fill(false, (size(event_db.events)[1], length(unique(event_db.event_times[event_db.event_times.<Inf])))), fill(false, (size(event_db.events)[1], length(unique(event_db.event_times[event_db.event_times.<Inf])))), unique(event_db.event_times[event_db.event_times.<Inf]), fill("I", length(unique(event_db.event_times[event_db.event_times.<Inf]))))
  for i = 1:length(sa.event_times)
    sa.susceptible[:,i] = find_state(event_db, sa.event_times[i], "S")
    sa.infectious[:,i] = find_state(event_db, sa.event_times[i], "I")
    if size(event_db.events)[2]==4
      sa.recovered[:,i] = find_state(event_db, sa.event_times[i], "R")
    end
  end
  for i = 2:length(sa.event_times)
    if sa.susceptible[:,i] == sa.susceptible[:,i-1]
      sa.event_type[i] = "R"
    end
    sa.new_infection[:,i] = sa.infectious[:,i] & sa.susceptible[:,i-1]
    if size(event_db.events)[2]==4
      sa.new_recovery[:,i] = sa.recovered[:,i] & sa.infectious[:,i-1]
    end
  end
  return sa
end

"""
Determine recovery times for individuals
"""
function find_recovery_times(event_db::edb, narm=true)
  if narm == false
    return event_db.events[:,4]  - event_db.events[:,3]
  end
  if narm == true
    times = event_db.events[:,4]  - event_db.events[:,3]
    return times[isnan(times).==false]
  end
end

function create_loglikelihood(pop_db::DataFrame, event_db::edb)
  """
  Create an efficient log likelihood function for continuous and discrete
  SI and SIR models
  """
  sa = state_array(event_db)
  distance_mat = create_dist_mtx(pop_db)
  if event_db.cd == "continuous" && size(event_db.events)[2]==4
    """
    loglikleihood for continuous SIR model
    """
    function cSIR_ll(parameters)
      if parameters[1] <= 0 || parameters[2] <= 0 || parameters[3] <= 0
        return -Inf
      else
        dist_ab = dist_ab_mtx(distance_mat, parameters[1], parameters[2])
        ilm_ll = 0.0
        for i = 2:length(sa.event_times)
          if sum(sa.new_infection[:,i]) > 0
            exp_rates=sum(dist_ab[sa.new_infection[:,i], sa.infectious[:,i-1]], 2).^-1.0
            for j = 1:length(exp_rates)
              ilm_ll += loglikelihood(Exponential(exp_rates[j]), [sa.event_times[i]-sa.event_times[i-1]])
            end
          end
          if sum(sa.new_recovery[:,i]) > 0
            ilm_ll += loglikelihood(Exponential(parameters[3]), fill(sa.event_times[i]-sa.event_times[i-1], sum(sa.new_recovery[:,i])))
          end
          ilm_ll += sum(dist_ab[sa.susceptible[:,i], sa.infectious[:,i-1]])*(sa.event_times[i-1]-sa.event_times[i])
          ilm_ll += sum(sa.infectious[:,i]) * (1/parameters[3]) * (sa.event_times[i-1]-sa.event_times[i])
        end
        return ilm_ll
      end
    end
    return cSIR_ll
  end

  if event_db.cd == "continuous" && size(event_db.events)[2]==3
    """
    loglikleihood for continuous SI model
    """
    function cSI_ll(parameters)
      if parameters[1] <= 0 || parameters[2] <= 0
        return -Inf
      else
        dist_ab = dist_ab_mtx(distance_mat, parameters[1], parameters[2])
        ilm_ll = 0.0
        for i = 2:length(sa.event_times)
            exp_rates=sum(dist_ab[sa.new_infection[:,i], sa.infectious[:,i-1]], 2).^-1.0
            for j = 1:length(exp_rates)
              ilm_ll += loglikelihood(Exponential(exp_rates[j]), [sa.event_times[i]-sa.event_times[i-1]])
            end
          ilm_ll += sum(dist_ab[sa.susceptible[:,i], sa.infectious[:,i-1]])*(sa.event_times[i-1]-sa.event_times[i])
        end
        return ilm_ll
      end
    end
    return cSI_ll
  end

  if event_db.cd == "discrete" && size(event_db.events)[2]==4
    rt = find_recovery_times(event_db)
    """
    loglikleihood for discrete SIR model
    """
    function dSIR_ll(parameters)
      if parameters[1] <= 0 || parameters[2] <= 0 || parameters[3] <= 0
        return -Inf
      else
        dist_ab = dist_ab_mtx(distance_mat, parameters[1], parameters[2])
        ilm_ll = 0.0
        for i = 2:maximum(event_db.events[:,3])
          ilm_ll += -sum(dist_ab[sa.susceptible[:,i], sa.infectious[:,i-1]])
          ilm_ll += sum(log(1-exp(-sum(dist_ab[sa.new_infection[:,i], sa.infectious[:,i-1]], 2))))
        end
        ilm_ll += loglikelihood(Geometric(1/parameters[3]), rt-1)
        return ilm_ll
      end
    end
    return dSIR_ll
  end

  if event_db.cd == "discrete" && size(event_db.events)[2]==3
    """
    loglikleihood for discrete SI model
    """
    function dSI_ll(parameters)
      if parameters[1] <= 0 || parameters[2] <= 0
        return -Inf
      else
        dist_ab = dist_ab_mtx(distance_mat, parameters[1], parameters[2])
        ilm_ll = 0.0
        for i = 2:maximum(event_db.events[:,3])
          ilm_ll += -sum(dist_ab[sa.susceptible[:,i], sa.infectious[:,i-1]])
          ilm_ll += sum(log(1-exp(-sum(dist_ab[sa.new_infection[:,i], sa.infectious[:,i-1]], 2))))
        end
        return ilm_ll
      end
    end
    return dSI_ll
  end
end


text:
## pyplot.jl  PyPlot plotting routines for ROC and DET
## (c) 2016 David A. van Leeuwen
##
## Licensed under the MIT software license, see LICENSE.md

@require PyPlot begin
    import PyPlot

    function detplot(r::Roc, nr=1; xlabel="Pfa (%)", ylabel="Pmiss (%)", title="DET plot")
        if nr==1
            fig = PyPlot.figure("DET plot")
            PyPlot.title("DET plot (must fix axes)")
            grid = [0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 40]
            ticklabels = AbstractString[]
            for x in grid
                if x<1
                    push!(ticklabels, @sprintf("%3.1f", x))
                else
                    push!(ticklabels, @sprintf("%1.0f", x))
                end
            end
            grid ./= 100
            PyPlot.axes(aspect="equal")
            ## This should be done with matplotlib scale, but I don't know the Julia interface.
            PyPlot.axis([qnorm(0.8e-3), qnorm(0.55), qnorm(0.8e-3), qnorm(0.55)])
            PyPlot.plot(qnorm(r.pfa), qnorm(r.pmiss))
            PyPlot.xlabel(xlabel)
            PyPlot.ylabel(ylabel)

        end
    end

end


text:
using Documenter, SauterSchwabQuadrature

makedocs(
    modules=[SauterSchwabQuadrature],
    doctest=false,
    sitename="SauterSchwabQuadrature",
)

# deploydocs(
#     deps = Deps.pip("pygments","mkdocs", "python-markdown-math"),
#     repo = "github.com/ga96tik/SauterSchwabQuadrature.jl.git",
#     julia  = "0.6",
#     osname = "linux")

deploydocs(
    # deps = Deps.pip("pygments","mkdocs", "python-markdown-math"),
    repo = "github.com/ga96tik/SauterSchwabQuadrature.jl.git",
    # julia  = "0.6",
    # osname = "linux",
)


text:
function render_all()
    return dbc_container([
        dbc_row([
            render_info()
        ]),
        dbc_row([
            render_func_init(),
            dbc_col([
                dbc_button("Sample", color="success", id="dtu-btn-turing")
            ], md =2),
            render_plots()
        ])
    ])
end
function render_info()
    return dbc_col([
        dbc_jumbotron([
            html_h1("Dash Turing App", className="dtu-display-3"),
            html_p(
                "Web app for Bayesian inference using Dash.jl and Turing.jl ",
                className="lead",
            ),
            html_hr(className="my-2"),
            html_p(
                "User can define a Turing model. The dataset's column names are available as variable names.",
            )
        ])
    ])
end
function render_func_init()
    return dbc_col([
        html_div([
            render_uploader(),
            render_function(),
            render_init()
        ]),
    ], md=5)
end
function render_function()
    return dbc_card([
        dbc_cardbody([
            html_h4("Model specification:", className="card-title"),
            html_div([
                dbc_textarea(
                    id = "dtu-user-funcs",
                    bs_size="lg",
                    className="mb-3",
                    rows  = 6,
                    value="""
                        @model regression_model(x, y) = begin
                            a ~ Normal()
                            b ~ Normal()
                            for i in 1:10
                                y[i] ~ Normal(a + b * x[i], 1.0)
                            end
                        end
                    """,
                )
            ]),
        ])
    ])
end
function render_init()
    return dbc_card([
            dbc_cardbody([
                html_h4("Conditional model:", className="card-title"),
                html_div([
                    dbc_input(
                        id = "dtu-user-func-init",
                        bs_size="lg",
                        className="mb-3",
                        value="regression_model(x, y)",
                    )
                ]),
            ])
        ])
end
function render_uploader()
    return dbc_card([
            dbc_cardbody([
                html_h4("Select Data", className="card-title"),
                html_div([
                    html_div(
                        id="dtu-def-upload-data",
                        [
                        html_label("Select a dataset:")
                        dcc_dropdown(
                            id="dtu-def-data-drop",
                            options =[
                                Dict("label" => "Select Data", "value" => "0"),
                                Dict("label" => "Data 1", "value" => "1"),
                                Dict("label" => "Data 2", "value" => "2"),
                                Dict("label" => "Data 3", "value" => "3")
                            ],
                            value="0"
                            )
                        ],
                    ),
                    html_h3("or"),
                    dcc_upload(
                        id="dtu-upload-datan",
                        children=html_div([
                            "Drag and Drop or ",
                            html_a("Select Files")
                        ]),
                        style=Dict(
                            "width" => "100%",
                            "height" => "60px",
                            "lineHeight" => "60px",
                            "borderWidth" => "1px",
                            "borderStyle" => "dashed",
                            "borderRadius" => "5px",
                            "textAlign" => "center",
                            "margin" => "10px"
                        ),
                        # Allow multiple files to be uploaded
                        multiple=false
                    )
                ]),
            ])
        ])
end

function render_liveupdates()
    global x_graph
    global y_graph
    return html_div([
        html_h3("Plots"),
        dcc_interval(
            id="interval-component",
            interval=500, # in milliseconds
            n_intervals=0,
            disabled=false
        ),
        dcc_graph(id="dtu-params-state1",
            figure=Dict(
                "data" => [
                    Dict("x" => x_graph, "y" => y_graph),
                ]
            )
        ),
        dcc_graph(id="dtu-params-state2",
            figure=Dict(
                "data" => [
                    Dict("x" => x_graph, "y" => y_graph1),
                ]
            )
        )
    ])
end
function render_plots()
    return dbc_col(id="dtu-params-plot-loc", md =5)
end


text:
using LinearAlgebra

function e_med(u, u_est)
    return median(abs.(LAP_julia.vec_len.(u .- u_est))[:])
end

function e_mean(u, u_est)
    return mean(abs.(LAP_julia.vec_len.(u .- u_est))[:])
end

## init

# lena
using TestImages
img = testimage("lena_gray")
img = Float32.(img)
wanted_max_displacement = 15
constant_flow = false

imgshow(img)
save("../PolyFilter_LAP/target.png", img)

##
# create 3 flows
if constant_flow == false
    # one
    base1 = [1 + im 2 - im;  2 - im 1 + im]
    tile_size1 = round(Int64, size(img, 1)/2)
    base_flow1 = repeat(base1, inner=(tile_size1, tile_size1))
    flow1 = LAP_julia.smooth_with_gaussian(base_flow1, 40)
    mag_ratio1 = wanted_max_displacement/LAP_julia.max_displacement(flow1)
    flow1 = flow1 .* mag_ratio1
    showflow(flow1)

    # two
    base2 = [1 + im 2 - 2im;  2 - im 1]
    tile_size2 = round(Int64, size(img, 1)/2)
    base_flow2 = repeat(base2, inner=(tile_size2, tile_size2))
    flow2 = LAP_julia.smooth_with_gaussian(base_flow2, 50)
    mag_ratio2 = wanted_max_displacement/LAP_julia.max_displacement(flow2)
    flow2 = flow2 .* mag_ratio2
    showflow(flow2)

    # three
    base3 = [-1 + im 2 + 2im;  2 - im im]
    tile_size3 = round(Int64, size(img, 1)/2)
    base_flow3 = repeat(base3, inner=(tile_size3, tile_size3))
    flow3 = LAP_julia.smooth_with_gaussian(base_flow3, 70)
    mag_ratio3 = wanted_max_displacement/LAP_julia.max_displacement(flow3)
    flow3 = flow3 .* mag_ratio3
    showflow(flow3)
else
     # one
    flow1 = repeat([-1 + 1im], inner=(size(img)))
    mag_ratio1 = wanted_max_displacement/LAP_julia.max_displacement(flow1);
    flow1 = flow1 .* mag_ratio1;

    # two
    flow2 = repeat([-12 + 1im], inner=(size(img)))
    mag_ratio2 = wanted_max_displacement/LAP_julia.max_displacement(flow2);
    flow2 = flow2 .* mag_ratio2;

    # one
    flow3 = repeat([-1 - 4im], inner=(size(img)))
    mag_ratio3 = wanted_max_displacement/LAP_julia.max_displacement(flow3);
    flow3 = flow3 .* mag_ratio3;
end

## test lap methods

#results
pflap_times = zeros(3)
pflap_e_meds = zeros(3)
pflap_e_means = zeros(3)

# one
flow = flow1; num = 1;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est, source_reg = pflap(img, imgw, display=false);
bench = @benchmark pflap(img, imgw, display=false)
pflap_times[num] = median(bench.times) / 10^9
pflap_e_meds[num] = e_med(flow, u_est)
pflap_e_means[num] = e_mean(flow, u_est)

# two
flow = flow2; num = 2;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est, source_reg = pflap(img, imgw, display=false);
bench = @benchmark pflap(img, imgw, display=false)
pflap_times[num] = median(bench.times) / 10^9
pflap_e_meds[num] = e_med(flow, u_est)
pflap_e_means[num] = e_mean(flow, u_est)

# three
flow = flow3; num = 3;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est, source_reg = pflap(img, imgw, display=false);
bench = @benchmark pflap(img, imgw, display=false)
pflap_times[num] = median(bench.times) / 10^9
pflap_e_meds[num] = e_med(flow, u_est)
pflap_e_means[num] = e_mean(flow, u_est)
showflow(u_est)

pflap_times
pflap_e_meds
pflap_e_means

## single_lap

# specifics
fhs = 23
window = [47, 47]

#results
lap_times = zeros(3)
lap_e_meds = zeros(3)
lap_e_means = zeros(3)

# one
flow = flow1; num = 1;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est = lap(img, imgw, fhs, window)
bench = @benchmark lap(img, imgw, fhs, window)
lap_times[num] = median(bench.times) / 10^9
lap_e_meds[num] = e_med(flow, u_est)
lap_e_means[num] = e_mean(flow, u_est)

# two
flow = flow2; num = 2;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est = lap(img, imgw, fhs, window)
bench = @benchmark lap(img, imgw, fhs, window)
lap_times[num] = median(bench.times) / 10^9
lap_e_meds[num] = e_med(flow, u_est)
lap_e_means[num] = e_mean(flow, u_est)

# three
flow = flow3; num = 3;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est = lap(img, imgw, fhs, window)
bench = @benchmark lap(img, imgw, fhs, window)
lap_times[num] = median(bench.times) / 10^9
lap_e_meds[num] = e_med(flow, u_est)
lap_e_means[num] = e_mean(flow, u_est)

lap_times
lap_e_meds
lap_e_means

## sparse single lap

# specifics
fhs = 23
window = [47, 47]

#results
splap_times = zeros(3)
splap_e_meds = zeros(3)
splap_e_means = zeros(3)

# one
flow = flow1; num = 1;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est = new_alg(img, imgw, fhs, window)
bench = @benchmark new_alg(img, imgw, fhs, window)
splap_times[num] = median(bench.times) / 10^9
splap_e_meds[num] = e_med(flow, u_est)
splap_e_means[num] = e_mean(flow, u_est)

# two
flow = flow2; num = 2;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est = new_alg(img, imgw, fhs, window)
bench = @benchmark new_alg(img, imgw, fhs, window)
splap_times[num] = median(bench.times) / 10^9
splap_e_meds[num] = e_med(flow, u_est)
splap_e_means[num] = e_mean(flow, u_est)

# three
flow = flow3; num = 3;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est = new_alg(img, imgw, fhs, window)
bench = @benchmark new_alg(img, imgw, fhs, window)
splap_times[num] = median(bench.times) / 10^9
splap_e_meds[num] = e_med(flow, u_est)
splap_e_means[num] = e_mean(flow, u_est)

splap_times
splap_e_meds
splap_e_means

# classic
function lap(img, imgw, fhs, window_size)
    classic_estim = single_lap(img, imgw, fhs, window_size)
    LAP_julia.inpaint_nans!(classic_estim)
    LAP_julia.smooth_with_gaussian(classic_estim, window_size)
    return classic_estim
end

# new
function new_alg(img, imgw, fhs, window_size)
    mask = parent(padarray(trues(size(img).-(2*fhs, 2*fhs)), Fill(false, (fhs, fhs), (fhs, fhs))))
    inds = find_edge_points(img, mask=mask)
    points = LAP_julia.inds_to_points(inds)
    new_estim = single_lap_at_points(img, imgw, fhs, window_size, points, 3)
    full_new_estim = interpolate_flow(new_estim, inds)
    return full_new_estim
end

## sparse pf lap

#results
sppflap_times = zeros(3)
sppflap_e_meds = zeros(3)
sppflap_e_means = zeros(3)

# one
flow = flow1; num = 1;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est, source_reg = sparse_pflap(img, imgw, display=false)
bench = @benchmark new_alg(img, imgw, fhs, window)
sppflap_times[num] = median(bench.times) / 10^9
sppflap_e_meds[num] = e_med(flow, u_est)
sppflap_e_means[num] = e_mean(flow, u_est)

# two
flow = flow2; num = 2;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est, source_reg = sparse_pflap(img, imgw, display=false)
bench = @benchmark new_alg(img, imgw, fhs, window)
sppflap_times[num] = median(bench.times) / 10^9
sppflap_e_meds[num] = e_med(flow, u_est)
sppflap_e_means[num] = e_mean(flow, u_est)

# three
flow = flow3; num = 3;
imgw = warp_img(img, -real(flow), -imag(flow));
u_est, source_reg = sparse_pflap(img, imgw, display=false)
bench = @benchmark new_alg(img, imgw, fhs, window)
sppflap_times[num] = median(bench.times) / 10^9
sppflap_e_meds[num] = e_med(flow, u_est)
sppflap_e_means[num] = e_mean(flow, u_est)
showflow(u_est)
showflow(flow .- u_est)

imgshow(img .- source_reg)
imgoverlay(img, source_reg)

sppflap_times
sppflap_e_meds
sppflap_e_means

## make comparation table:

lap = [lap_times, lap_e_meds, lap_e_means]
pflap = [pflap_times, pflap_e_meds, pflap_e_means]
splap = [splap_times, splap_e_meds, splap_e_means]
sppflap = [sppflap_times, sppflap_e_meds, sppflap_e_means]

columns = ["Time", "e_med", "e_mean"]
lap_avg = round.(map(mean, lap), digits=3)
pflap_avg = round.(map(mean, pflap), digits=3)
splap_avg = round.(map(mean, splap), digits=3)
sppflap_avg = round.(map(mean, sppflap), digits=3)

println(columns)
map(println, [pflap_avg, lap_avg, sppflap_avg, splap_avg])

## results
map(println, [lap_avg, pflap_avg, splap_avg, sppflap_avg])
"LENA"
["Time", "e_med", "e_mean"]
[1.198, 1.384, 1.806]
[4.668, 0.852, 1.012]
[0.232, 1.331, 1.835]
[0.205, 0.546, 0.619]

"const"
map(println, [pflap_avg, lap_avg, sppflap_avg, splap_avg])
["Time", "e_med", "e_mean"]
6.664 & 0.116 & 0.721
1.019 & 2.553 & 3.149
0.194 & 0.428 & 0.988
0.196 & 2.285 & 2.524


text:
export MACLearner

"""
    MACLearner(;kwargs...)

[Mean Actor-Critic](https://arxiv.org/pdf/1709.00503.pdf)

# Keyword arguments

- `approximator`::[`ActorCritic`](@ref)
- `γ::Float32`, reward discount rate
-  `bootstrap::bool`, if false then Q function is approximated using monte carlo returns.
"""

Base.@kwdef mutable struct MACLearner{A<:ActorCritic} <: AbstractLearner
    approximator::A
    γ::Float32
    max_grad_norm::Union{Nothing,Float32} = nothing
    norm::Float32 = 0.0f0
    actor_loss::Float32 = 0.0f0
    critic_loss::Float32 = 0.0f0
    loss::Float32 = 0.0f0
    bootstrap::Bool = true
    update_freq::Int
    update_step::Int = 0
end

Flux.functor(x::MACLearner) = (app = x.approximator,), y -> @set x.approximator = y.app

function (learner::MACLearner)(env::MultiThreadEnv)
    learner.approximator.actor(send_to_device(device(learner.approximator), state(env))) |>
    send_to_host
end

function (learner::MACLearner)(env)
    s = state(env)
    s = Flux.unsqueeze(s, ndims(s) + 1)
    s = send_to_device(device(learner.approximator), s)
    learner.approximator.actor(s) |> vec |> send_to_host
end

function RLBase.update!(learner::MACLearner, t::CircularArraySARTTrajectory, ::AbstractEnv, ::PreActStage)
    length(t) == 0 && return  # in the first update, only state & action is inserted into trajectory
    learner.update_step += 1
    if learner.update_step % learner.update_freq == 0
        _update!(learner, t)
    end
end

function _update!(learner::MACLearner, t::CircularArraySARTTrajectory)
    n = length(t)
    states = select_last_dim(t[:state], 1:n)
    actions = select_last_dim(t[:action], 1:n)
    rewards = t[:reward]
    terminals = t[:terminal]

    AC = learner.approximator
    γ = learner.γ
    D = device(AC)

    states = send_to_device(D, states)
    states_flattened = flatten_batch(states) # (state_size..., n_thread * update_freq)

    actions = flatten_batch(actions)
    actions = CartesianIndex.(actions, 1:length(actions))

    if learner.bootstrap
        next_state = select_last_frame(t[:state])
        next_state = send_to_device(D, next_state)
        next_state_values = AC.critic(next_state)

        gains = discount_rewards(
            rewards,
            γ;
            dims = 2,
            init = send_to_host(next_state_values),
            terminal = terminals,
        )
        gains = send_to_device(D, gains)
    else
        next_state_flattened = flatten_batch(select_last_dim(t[:state], 2:n+1))
        next_state_flattened = send_to_device(D, next_state_flattened)
        rewards_flattened = flatten_batch(rewards)
        rewards_flattened = send_to_device(D, rewards_flattened)
    end

    action_values = AC.critic(states_flattened)

    ps1 = Flux.params(AC.actor)
    gs1 = gradient(ps1) do
        logits = AC.actor(states_flattened)
        probs = softmax(logits)
        actor_loss = -mean(sum((probs .* Zygote.dropgrad(action_values)), dims = 1))
        loss = actor_loss
        ignore() do
            learner.actor_loss = actor_loss
        end
        loss
    end
    if !isnothing(learner.max_grad_norm)
        learner.norm = clip_by_global_norm!(gs1, ps1, learner.max_grad_norm)
    end
    update!(AC.actor, gs1)

    ps2 = Flux.params(AC.critic)
    gs2 = gradient(ps2) do
        if learner.bootstrap
            critic_loss = mean((vec(gains) .- vec(action_values[actions])) .^ 2)
        else
            next_state_values = AC.critic(next_state_flattened)
            target_action_values =
                vec(rewards_flattened) .+
                γ * vec(Zygote.dropgrad(sum(
                    next_state_values .* softmax(AC.actor(next_state_flattened)),
                    dims = 1,
                )))
            critic_loss =
                mean((vec(target_action_values) .- vec(action_values[actions])) .^ 2)
        end

        loss = critic_loss
        ignore() do
            learner.critic_loss = critic_loss
        end
        loss
    end
    if !isnothing(learner.max_grad_norm)
        learner.norm = clip_by_global_norm!(gs2, ps2, learner.max_grad_norm)
    end
    update!(AC.critic, gs2)
end

RLCore.check(::QBasedPolicy{<:MACLearner}, ::MultiThreadEnv) = nothing