text:
using Revise
using SnpArrays
using Knockoffs
using DelimitedFiles
using Random
using LinearAlgebra
using Distributions
using ProgressMeter
using MendelIHT
using VCFTools
using StatsBase
using CSV
using DataFrames
using Printf
using GLMNet
using GLM
using CodecZlib
BLAS.set_num_threads(1)

global qctools_exe = "/scratch/users/bbchu/qctool/build/release/qctool_v2.0.7"
global snpknock2_exe = "/scratch/users/bbchu/knockoffgwas/snpknock2/bin/snpknock2"
global rapid_exe = "/scratch/users/bbchu/RaPID/RaPID_v.1.7"
global partition_exe = "/scratch/users/bbchu/knockoffgwas/knockoffgwas/utils/partition.R" # modify this script to generate knockoffs of different group resolution
global propca_exe = "/scratch/users/bbchu/ProPCA/build/propca"
global bolt_exe = "/scratch/users/bbchu/BOLT-LMM_v2.2/bolt"
global hdf_to_vcf_exe = "/scratch/users/bbchu/scripts/hdf_to_vcf.py"
global prs_genotypes_exe = "/scratch/users/bbchu/PRS_Admixture_Simulation/simulate_genotypes.py"
global rapid_filter_exe = "/scratch/users/bbchu/scripts/filter_mapping_file.py"
global rapid_interpolate_exe = "/scratch/users/bbchu/scripts/interpolate_loci.py"

function create_plink(
    outfile::AbstractString,
    x1::AbstractMatrix,
    x2::AbstractMatrix;
    chr::Vector{String} = ["1" for _ in 1:size(x1, 2)],
    pos::Vector{Int} = collect(100:100:size(x1, 2))
    )
    n, p = size(x1)
    x = SnpArray(outfile * ".bed", n, p)
    for j in 1:p, i in 1:n
        c = x1[i, j] + x2[i, j]
        if c == 0
            x[i, j] = 0x00
        elseif c == 1
            x[i, j] = 0x02
        elseif c == 2
            x[i, j] = 0x03
        else
            error("matrix entries should be 0, 1, or 2 but was $c!")
        end
    end
    # create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(outfile * ".bim", "w") do f
        for i in 1:p
            println(f, "$(chr[i])\tsnp$i\t0\t$(pos[i])\t1\t2")
        end
    end
    # create .fam file structure: https://www.cog-genomics.org/plink2/formats#fam
    open(outfile * ".fam", "w") do f
        for i in 1:n
            println(f, "$i\t1\t0\t0\t1\t-9")
        end
    end
    return nothing
end

# todo: How to systematically reduce connected components in related graph?
# see: https://github.com/msesia/knockoffgwas/issues/3
function process_rapid_output(inputfile, outputfile)
    writer = open(outputfile, "w")
    println(writer, "CHR ID1 HID1 ID2 HID2 BP.start BP.end site.start site.end cM FAM1 FAM2")
    try
        df = readdlm(inputfile)
        total_ibd = size(df, 1)
        keep_prob = 1000 / total_ibd # want to keep on average 1000 IBD segments
        for r in eachrow(df)
            chr, id1, id2, hap1, hap2, start_pos, end_pos, genetic_len, start_site, end_site =
                Int(r[1]), Int(r[2]), Int(r[3]), Int(r[4]), Int(r[5]), Int(r[6]), Int(r[7]),
                r[8], Int(r[9]), Int(r[10])
            rand() < keep_prob && println(writer, chr, ' ', id1, ' ', hap1, ' ', id2, ' ', hap2, ' ',
                start_pos, ' ', end_pos, ' ', start_site, ' ', end_site, ' ',
                genetic_len, ' ', 1, ' ', 1)
        end
    catch
        println("0 IBD segments detected!")
    end
    close(writer)
end

function make_bgen_samplefile(filename, n)
    open(filename, "w") do io
        println(io, "ID_1 ID_2 missing sex")
        println(io, "0 0 0 D")
        for i in 1:n
            println(io, "$i 1 0 1")
        end
    end
end

# compute most ancestry informative markers
function find_500_AIMs(HCEU, HYRI)
    # we don't know which samples in all.vcf belongs to which population,
    # so create a VCF file for which we do know
    H1 = [HCEU[1:2:end, :]; HYRI[1:2:end, :]]
    H2 = [HCEU[2:2:end, :]; HYRI[2:2:end, :]]
    sample_ids = [string(i) for i in 1:size(H1, 1)]
    write_vcf("nonadmixed.vcf", H1, H2, sampleID=sample_ids)

    # create dictionary mapping sampleID to population
    nCEU = size(HCEU, 1) >> 1
    nYRI = size(HYRI, 1) >> 1
    sampleID_to_population = Dict{String, String}()
    [sampleID_to_population[string(i)] = "CEU" for i in 1:nCEU]
    [sampleID_to_population[string(i)] = "YRI" for i in (nCEU+1):(nCEU + nYRI)]

    # select aims
    pvals = VCFTools.aim_select("nonadmixed.vcf", sampleID_to_population)
    if count(iszero, pvals) ≥ 500
        idx = findall(iszero, pvals)
        shuffle!(idx)
        top_aims = idx[1:500]
    else
        top_aims = sortperm(pvals)[1:500]
    end

    # clean up
    rm("nonadmixed.vcf", force=true)

    return top_aims
end

# the header row of rfmix output is #CHROM  POS     ID      REF     VAL
# it should be #CHROM  POS     ID      REF     ALT ...
function process_rfmix(samples, intputfile, outfile)
    reader = GzipDecompressorStream(open(intputfile))
    writer = GzipCompressorStream(open(outfile, "w"))
    # skip broken header lines
    for i in 1:5
        readline(reader)
    end
    # write minimum header into
    print(writer, "##fileformat=VCFv4.2\n")
    print(writer, "##source=rfmix\n")
    print(writer, "##FORMAT=<ID=GT,Number=1,Type=String,Description=\"Genotype\">\n")
    print(writer, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT")
    for id in ["sample$i" for i in 1:samples]
        print(writer, "\t", id)
    end
    print(writer, "\n")
    # write line by line
    while !eof(reader)
        println(writer, readline(reader))
    end
    close(reader)
    close(writer)
end

function simulate_genotypes(nYRI, nCEU, nmate, nadmixed, p, seed, outdir)
    # change directory to PRS_Admixture_Simulation
    # will copy data back to outdir at the end of this function
    cd("PRS_Admixture_Simulation")

    # simulate CEU and YRI populations and their mating, then convert to VCF format
    @time run(`python $prs_genotypes_exe --sim $seed --nYRI $nYRI --nCEU $nCEU --nmate $nmate --nadmixed $nadmixed --threads $(Threads.nthreads())`)
    cd("output/sim$seed/trees")
    run(pipeline(`python $hdf_to_vcf_exe tree_all.hdf`, stdout="all.full.vcf")) # 10000 samples
    run(pipeline(`python $hdf_to_vcf_exe tree_CEU_GWAS_nofilt.hdf`, stdout="CEU.full.vcf"))
    run(pipeline(`python $hdf_to_vcf_exe tree_YRI_GWAS_nofilt.hdf`, stdout="YRI.full.vcf"))
    # run(pipeline(`python $hdf_to_vcf_exe tree_mate.hdf`, stdout="mate.full.vcf")) # mate.full.vcf are CEU/YRI samples used to generate admixed samples, and then excluded from CEU/YRI.full.vcf files

    # admixed samples are made by rfmix, but their vcf format is old, so we need to process it
    process_rfmix(nadmixed, "../admixed_data/output/admix_afr_amer.query.vcf.gz", "admixed.full.vcf.gz")

    # keep p SNPs with maf at least 0.01 (todo: keep how many snps??)
    _, _, _, _, _, maf_by_record, _ = gtstats("all.full.vcf")
    idx = findall(x -> x > 0.01, maf_by_record)
    shuffle!(idx)
    record_mask = sort(idx[1:p])
    # maf_perm = sortperm(maf_by_record, rev=true)
    # record_mask = maf_perm[1:p]
    VCFTools.filter("CEU.full.vcf", record_mask, 1:VCFTools.nsamples("CEU.full.vcf"), des="CEU.vcf")
    VCFTools.filter("YRI.full.vcf", record_mask, 1:VCFTools.nsamples("YRI.full.vcf"), des="YRI.vcf")
    VCFTools.filter("admixed.full.vcf.gz", record_mask, 1:VCFTools.nsamples("admixed.full.vcf.gz"), des="admixed.vcf")

    # import haplotypes
    HCEU = convert_ht(Bool, "CEU.vcf", msg="importing CEU.vcf")
    HYRI = convert_ht(Bool, "YRI.vcf", msg="importing YRI.vcf")
    Hadmix,_,_,pos,_,_,_ = convert_ht(Bool, "admixed.vcf", save_snp_info=true, msg="importing admixed.vcf")

    # write 90% samples to training data
    chr = ["20" for i in 1:size(Hadmix, 2)]
    nCEUtrain = Int(0.9 * size(HCEU, 1))
    nYRItrain = Int(0.9 * size(HYRI, 1))
    nadmixedtrain = Int(0.9 * size(Hadmix, 1))
    H1_train = [HCEU[1:2:nCEUtrain, :]; HYRI[1:2:nYRItrain, :]; Hadmix[1:2:nadmixedtrain, :]]
    H2_train = [HCEU[2:2:nCEUtrain, :]; HYRI[2:2:nYRItrain, :]; Hadmix[2:2:nadmixedtrain, :]]
    create_plink("train", H1_train, H2_train, chr=chr, pos=pos)
    write_vcf("train.phased.vcf.gz", H1_train, H2_train, chr=chr, pos=pos)
    open("train.populations.txt", "w") do io
        [println(io, "CEU") for i in 1:nCEUtrain>>1]
        [println(io, "YRI") for i in 1:nYRItrain>>1]
        [println(io, "admixed") for i in 1:nadmixedtrain>>1]
    end

    # write 10% samples to testing data
    H1_test = [HCEU[nCEUtrain+1:2:end, :]; HYRI[nYRItrain+1:2:end, :]; Hadmix[nadmixedtrain+1:2:end, :]]
    H2_test = [HCEU[nCEUtrain+2:2:end, :]; HYRI[nYRItrain+2:2:end, :]; Hadmix[nadmixedtrain+2:2:end, :]]
    create_plink("test", H1_test, H2_test, chr=chr, pos=pos)
    open("test.populations.txt", "w") do io
        [println(io, "CEU") for i in 1:(size(HCEU, 1) - nCEUtrain)>>1]
        [println(io, "YRI") for i in 1:(size(HYRI, 1) - nYRItrain)>>1]
        [println(io, "admixed") for i in 1:(size(Hadmix, 1) - nadmixedtrain)>>1]
    end

    # find top 500 differentially expressed SNPs via AIM select
    diff_markers = find_500_AIMs(HCEU, HYRI)
    writedlm("diff_markers.txt", diff_markers)

    # also generate QC file that contains all SNPs and all samples
    snpdata = SnpData("train")
    snpIDs = snpdata.snp_info[!, :snpid]
    sampleIDs = Matrix(snpdata.person_info[!, 1:2])
    writedlm("variants_qc.txt", snpIDs)
    writedlm("samples_qc.txt", sampleIDs)

    # move relevant files to outdir
    mv("train.bed", outdir * "train.bed", force=true)
    mv("train.bim", outdir * "train.bim", force=true)
    mv("train.fam", outdir * "train.fam", force=true)
    mv("test.bed", outdir * "test.bed", force=true)
    mv("test.bim", outdir * "test.bim", force=true)
    mv("test.fam", outdir * "test.fam", force=true)
    mv("variants_qc.txt", outdir * "variants_qc.txt", force=true)
    mv("samples_qc.txt", outdir * "samples_qc.txt", force=true)
    mv("train.phased.vcf.gz", outdir * "train.phased.vcf.gz", force=true)
    mv("train.populations.txt", outdir * "train.populations.txt", force=true)
    mv("test.populations.txt", outdir * "test.populations.txt", force=true)
    mv("diff_markers.txt", outdir * "diff_markers.txt", force=true)

    # clean up
    rm("CEU.vcf", force=true)
    rm("YRI.vcf", force=true)
    rm("all.vcf", force=true)
    rm("admixed.vcf", force=true)
    rm("CEU.full.vcf", force=true)
    rm("YRI.full.vcf", force=true)
    rm("all.full.vcf", force=true)
    rm("admixed.full.vcf.gz", force=true)
end

function read_vcf_pos(vcffile)
    _,_,_,pos,_,_,_ = convert_ht(Bool, vcffile, save_snp_info=true, msg="importing..")
    return pos
end

function make_every_mapfile_needed()
    # make map file for rapid https://github.com/ZhiGroup/RaPID
    vcffile = "train.phased.vcf.gz"
    outfile = "train.rapid.map"
    run(`python $rapid_filter_exe /scratch/users/bbchu/RaPID/genetic_maps/genetic_map_GRCh38_chr20.txt raw_map.txt`)
    run(`python $rapid_interpolate_exe raw_map.txt $vcffile $outfile`)

    # make map file needed for knockoffgwas (todo: how to define map_cM and recombination rate for simulated data??)
    pos = read_vcf_pos(vcffile)
    map_cM = LinRange(0.0, round(pos[end] / 1e6), length(pos)) # even spacing
    recomb_rate = 0.01 .* rand(length(pos)) # random recomb rate
    open("sim.partition.map", "w") do io
        println(io, "Chromosome\tPosition(bp)\tRate(cM/Mb)\tMap(cM)")
        for i in 1:length(pos)
            println(io, "chr20\t", pos[i], '\t', recomb_rate[i], '\t', map_cM[i])
        end
    end

    # partition genotypes to groups (for making group knockoffs)
    partition(partition_exe, "train", "sim.partition.map", "variants_qc.txt", "sim.partition.txt")
end

function make_knockoffs(seed, outdir)
    cd(outdir)

    # make map files
    make_every_mapfile_needed()

    # run rapid to detect IBD segments
    Random.seed!(seed)
    vcffile = "train.phased.vcf.gz"
    mapfile = "train.rapid.map"
    outfolder = "rapid"
    d = 3    # minimum IBD length in cM
    w = 3    # number of SNPs per window
    r = 10   # number of runs
    s = 2    # Minimum number of successes to consider a hit
    @time rapid(rapid_exe, vcffile, mapfile, d, outfolder, w, r, s)

    # unzip and postprocess rapid output to suit snpknock2
    run(pipeline(`gunzip -c ./rapid/results.max.gz`, stdout="./rapid/results.max"))
    process_rapid_output("./rapid/results.max", "sim.snpknock.ibdmap")

    # convert VCF to BGEN format
    outfile = "train.bgen"
    run(`$qctools_exe -g $vcffile -og $outfile`)
    make_bgen_samplefile("train.sample", VCFTools.nsamples(vcffile))

    # snpknock2 arguments
    bgenfile = "train"
    sample_qc = "samples_qc.txt"
    variant_qc = "variants_qc.txt"
    mapfile = "sim.partition.map"
    partfile = "sim.partition.txt"
    ibdfile = "sim.snpknock.ibdmap"
    K = 10
    cluster_size_min = 1000
    cluster_size_max = 10000
    hmm_rho = 1
    hmm_lambda = 1e-3
    windows = 0
    n_threads = Threads.nthreads()
    compute_references = true
    generate_knockoffs = true
    outfile = "train.knockoffs"
    @time snpknock2(snpknock2_exe, bgenfile, sample_qc, variant_qc, mapfile, partfile, ibdfile,
        K, cluster_size_min, cluster_size_max, hmm_rho, hmm_lambda, windows, n_threads,
        seed, compute_references, generate_knockoffs, outfile)
end

"""
    simulate_beta_and_y!(xtrain, xtest, diff_markers, ::Normal)

Simulates β and y given indices of differentiated markers and random seem.
There are 100 causal snps, 50 on normally differentiated SNPs and 50 on
abnormally differentiated snps. If `add_interaction=true`, there will be 15
interaction terms as summarized above.
"""
function simulate_beta_and_y!(
    xtrain::AbstractMatrix,
    xtest::AbstractMatrix,
    diff_markers::Vector{Int},
    ::Normal
    )
    n, p = size(xtrain)
    ntest = size(xtest, 1)
    h2 = 0.5 # heritability
    K = 100 # number of causal snps
    d = Normal(0, sqrt(h2 / 2K)) # from paper: Efficient Implementation of Penalized Regression for Genetic Risk Prediction

    # simulate β with linear effect, half are on normally differentiated SNPs half on abnormal SNPs
    β = zeros(p)
    norm_markers = setdiff(1:p, diff_markers)
    norm_causal_snps = sample(norm_markers, K >> 1, replace=false)
    β[norm_causal_snps] .= rand(d, K >> 1)
    if length(diff_markers) > 0
        shuffle!(diff_markers)
        diff_causal_markers = diff_markers[1:K >> 1]
        β[diff_causal_markers] .= rand(d, K >> 1)
    end

    # simulate y (train and test)
    ϵ = Normal(0, 1 - h2)
    ytrain = xtrain * β + rand(ϵ, n)
    ytest = xtest * β + rand(ϵ, ntest)

    return ytrain, ytest, β, norm_markers, norm_causal_snps, diff_markers
end

function simulate_beta_and_y!(
    xtrain::AbstractMatrix,
    xtest::AbstractMatrix,
    diff_markers::Vector{Int},
    ::Bernoulli
    )
    n, p = size(xtrain)
    ntest = size(xtest, 1)
    h2 = 0.5 # heritability
    K = 100 # number of causal snps
    d = Normal(0, 0.3)

    # simulate β with linear effect, half are on normally differentiated SNPs half on abnormal SNPs
    β = zeros(p)
    norm_markers = setdiff(1:p, diff_markers)
    norm_causal_snps = sample(norm_markers, K >> 1, replace=false)
    β[norm_causal_snps] .= rand(d, K >> 1)
    if length(diff_markers) > 0
        shuffle!(diff_markers)
        diff_causal_markers = diff_markers[1:K >> 1]
        β[diff_causal_markers] .= rand(d, K >> 1)
    end

    # simulate y (train and test)
    ytrain_mean = GLM.linkinv.(LogitLink(), xtrain * β)
    ytest_mean = GLM.linkinv.(LogitLink(), xtest * β)
    ytrain = Float64.([rand(Bernoulli(i)) for i in ytrain_mean])
    ytest = Float64.([rand(Bernoulli(i)) for i in ytest_mean])

    return ytrain, ytest, β, norm_markers, norm_causal_snps, diff_markers
end

function get_diff_markers(data_dir)
    try
        return Vector{Int}(vec(readdlm(data_dir)))
    catch
        return Int[]
    end
end

function one_simulation(nYRI, nCEU, nmate, nadmixed, p, dist, seed;
    use_PCA=true,               # whether to compute, store, and use PCA in model selection
    simulate_interactions=true  # whether to add interacting SNPs in βtrue
    )
    cur_dir = pwd() * "/"

    # save all simulated data under data/sim$seed
    isdir("data") || mkdir("data")
    isdir("data/sim$seed") || mkdir("data/sim$seed")
    cd("data/sim$seed")
    data_dir = cur_dir * "data/sim$seed/"

    # simulate phased genotypes with pop structure and cryptic relatedness
    isdir("PRS_Admixture_Simulation") || cp("/scratch/users/bbchu/PRS_Admixture_Simulation", data_dir * "PRS_Admixture_Simulation", force=true)
    simulate_genotypes(nYRI, nCEU, nmate, nadmixed, p, seed, data_dir)

    # make knockoffs
    make_knockoffs(seed, data_dir)
    GC.gc()

    # compute PCA on original genotypes using ProPCA
    if use_PCA
        plinkfile = data_dir * "train"
        outfile = data_dir * "pca_"
        run(`$propca_exe -g $plinkfile -o $outfile -nt $(Threads.nthreads()) -seed $seed`)
        z = readdlm(data_dir * "pca_projections.txt")
        standardize!(z)
    end

    # loop over group resolution
    for group in 0:1
        cd(data_dir)

        # import genotypes, knockoff, and other info
        snpdata = SnpData(data_dir * "knockoffs/train.knockoffs_res$(group)")
        snpid = snpdata.snp_info.snpid
        knockoff_idx = endswith.(snpid, ".k")
        original = findall(knockoff_idx .== false)
        knockoff = findall(knockoff_idx)
        x = snpdata.snparray
        xla = convert(Matrix{Float64}, @view(x[:, original]), center=true, scale=true, impute=true)
        xla_test = convert(Matrix{Float64}, SnpArray(data_dir * "test.bed"), center=true, scale=true, impute=true)
        xko_la = convert(Matrix{Float64}, x, center=true, scale=true, impute=true)
        groups = repeat(Vector{Int}(readdlm(data_dir * "knockoffs/train.knockoffs_res$(group)_grp.txt",
            header=true)[1][:, 2]), inner=2) .+ 1
        diff_markers = get_diff_markers(data_dir * "diff_markers.txt")

        # If requested, concatenate PCs to genotype matrix (only needed for lasso)
        xla_full = use_PCA ? [xla z] : xla
        xko_la_full = use_PCA ? [xko_la z] : xko_la
        if use_PCA # must also asign groups to PCs
            max_g = maximum(groups)
            PC_groups = collect(max_g+1:max_g+5)
        end
        groups_full = use_PCA ? [groups; PC_groups] : groups
        GC.gc()

        #
        # simulate beta and y, form non-genetic covariates (used in IHT only)
        #
        Random.seed!(seed)
        if simulate_interactions
            ytrain, ytest, β, norm_markers, norm_causal_snps, diff_markers,
                γ, interact_pairs =
                simulate_beta_and_y_with_interaction(xla, xla_test, diff_markers)
        else
            ytrain, ytest, β, norm_markers, norm_causal_snps, diff_markers =
                simulate_beta_and_y!(xla, xla_test, diff_markers, dist)
        end
        covar = use_PCA ? [ones(size(xla, 1)) z] : ones(size(xla, 1))
        link = canonicallink(dist)
        init_beta = dist == Normal() ? true : false

        #
        # Run standard IHT
        #
        Random.seed!(seed)
        path = 10:10:250
        mses = cv_iht(ytrain, xla, covar, d=dist, l=link, path=path, init_beta=init_beta, max_iter=100)
        GC.gc()
        Random.seed!(seed)
        k_rough_guess = path[argmin(mses)]
        dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
        mses_new = cv_iht(ytrain, xla, covar, d=dist, l=link, path=dense_path, init_beta=init_beta, max_iter=100)
        GC.gc()
        Random.seed!(seed)
        result = fit_iht(ytrain, xla, covar, k=dense_path[argmin(mses_new)], d=dist, l=link,
            init_beta=init_beta, max_iter=500)
        β_iht = result.beta
        GC.gc()

        #
        # run knockoff IHT
        #
        Random.seed!(seed)
        path = 10:10:250
        mses = cv_iht(ytrain, xko_la, covar, d=dist, l=link, path=path, init_beta=init_beta, max_iter=100)
        GC.gc()
        Random.seed!(seed)
        k_rough_guess = path[argmin(mses)]
        dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
        mses_new = cv_iht(ytrain, xko_la, covar, d=dist, l=link, path=dense_path, init_beta=init_beta, max_iter=100)
        GC.gc()
        Random.seed!(seed)
        ko_iht_result = fit_iht(ytrain, xko_la, covar, d=dist, l=link,
            k=dense_path[argmin(mses_new)], init_beta=init_beta, max_iter=500)

        #
        # Run standard lasso
        #
        Random.seed!(seed)
        lasso_y = dist == Normal() ? ytrain : string.(ytrain)
        cv = glmnetcv(xla_full, lasso_y, nfolds=5, parallel=true)
        β_lasso = GLMNet.coef(cv)
        GC.gc()

        #
        # Run knockoff lasso
        #
        Random.seed!(seed)
        ko_lasso_cv = glmnetcv(xko_la_full, lasso_y, nfolds=5, parallel=true)
        GC.gc()

        #
        # Run bolt lmm (need to make covariate and phenotype file first)
        #
        bedfile = data_dir * "train"
        covfile = "cov_bolt.txt"
        phefile = "y_bolt.txt"
        outfile = "bolt_output.txt"
        open(phefile, "w") do io
            println(io, "FID IID trait1")
            for i in 1:size(xla, 1)
                println(io, "$i 1 $(ytrain[i])")
            end
        end
        if use_PCA
            open(covfile, "w") do io
                println(io, "FID IID PC1 PC2 PC3 PC4 PC5")
                for i in 1:size(xla, 1)
                    println(io, "$i 1 ", z[i, 1], ' ', z[i, 2], ' ',  z[i, 3], ' ', z[i, 4], ' ', z[i, 5])
                end
            end
            run(`$bolt_exe --bfile=$bedfile --covarFile=$covfile
                --phenoFile=$phefile --phenoCol=trait1
                --qCovarCol=PC\{1:5\} --lmmInfOnly --numLeaveOutChunks=2 --statsFile $outfile`)
        else
            run(`$bolt_exe --bfile=$bedfile --phenoFile=$phefile --phenoCol=trait1
                --lmmInfOnly --numLeaveOutChunks=2 --statsFile $outfile`)
        end

        # current directory
        group_dir = cur_dir * "res$group/"
        isdir(group_dir) || mkdir(group_dir)
        cd(group_dir)
        GC.gc()

        for fdr in [0.05, 0.1, 0.25, 0.5]
            top_dir = group_dir * "fdr$fdr/"
            new_dir = group_dir * "fdr$fdr/sim$seed/"
            isdir(top_dir) || mkdir(top_dir)
            isdir(new_dir) || mkdir(new_dir)
            cd(new_dir)

            #
            # save true (simulated) parameters
            #
            writedlm("y_train.txt", ytrain)
            writedlm("y_test.txt", ytest)
            writedlm("beta_true.txt", β)
            writedlm("normal_markers.txt", norm_markers)
            writedlm("normal_causal_markers.txt", norm_causal_snps)
            writedlm("diff_markers.txt", diff_markers)
            simulate_interactions && writedlm("gamma_true.txt", γ)
            simulate_interactions && writedlm("interacting_markers.txt", interact_pairs)

            #
            # save standard IHT/lasso results
            #
            writedlm("iht.beta", β_iht)
            writedlm("lasso.beta", β_lasso)

            #
            # knockoff IHT statistics
            #
            β_iht_knockoff = extract_beta(ko_iht_result.beta, fdr, groups, original, knockoff)
            writedlm("iht.knockoff.beta", ko_iht_result.beta)
            writedlm("iht.knockoff.postfilter.beta", β_iht_knockoff)
            GC.gc()

            #
            # run knockoff IHT with wrapped cross validation
            #
            # Random.seed!(seed)
            # path = 10:10:250
            # mses = cv_iht_knockoff(ytrain, xko_la, covar, original,
            #     knockoff, fdr, d=dist, l=link, path=path, init_beta=init_beta, group_ko=groups, max_iter=500)
            # GC.gc()
            # Random.seed!(seed)
            # k_rough_guess = path[argmin(mses)]
            # dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
            # mses_new = cv_iht_knockoff(ytrain, xko_la, covar, original, knockoff, fdr,
            #     d=dist, l=link, path=dense_path, init_beta=init_beta, group_ko=groups, max_iter=500)
            # GC.gc()
            # # adjust sparsity level so it best matches sparsity chosen by ko filter
            # Random.seed!(seed)
            # best_k = dense_path[argmin(mses_new)]
            # best_β = tune_k(ytrain, xko_la, covar, original, knockoff, groups, fdr, best_k, dist)
            # β_iht_knockoff_cv = extract_beta(best_β, fdr, groups, original, knockoff)
            # writedlm("iht.knockoff.cv.beta", best_β)
            # writedlm("iht.knockoff.cv.postfilter.beta", β_iht_knockoff_cv)
            # GC.gc()

            #
            # knockoff LASSO statistics
            #
            β_lasso_knockoff = extract_beta(GLMNet.coef(ko_lasso_cv), fdr, groups_full, original, knockoff)
            writedlm("lasso.knockoff.beta", GLMNet.coef(ko_lasso_cv))
            writedlm("lasso.knockoff.postfilter.beta", β_lasso_knockoff)
            GC.gc()

            #
            # copy bolt output
            #
            use_PCA && cp(data_dir * covfile, new_dir * covfile, force=true)
            cp(data_dir * phefile, new_dir * phefile, force=true)
            cp(data_dir * outfile, new_dir * outfile, force=true)
        end
        use_PCA && rm(data_dir * covfile, force=true)
        rm(data_dir * phefile, force=true)
        rm(data_dir * outfile, force=true)
        GC.gc()
    end
end

##
## Genotypes are simulated using msprime, following https://github.com/taylorcavazos/PRS_Admixture_Simulation
## Genotypes mimics chr20 linkage
## Founders used to make admixed samples are not included in the final sample used for training
## To train PRS, we use all samples (CEU + YRI + admixed)
## To test PRS, we use all testing samples (CEU + YRI + admixed) that are disjoint from training samples
##

nYRI = 5000
nCEU = 5000
nadmixed = 2000 # for admixed samples
nmate = 1000    # for generating founder population (1000 each from YRI and CEU) to make admixed samples; total samples = 5000+5000+2000-1000*2 = 10000
p = 50000
use_PCA = true
d = Normal()
seed = parse(Int, ARGS[1])
one_simulation(nYRI, nCEU, nmate, nadmixed, p, d, seed,
    use_PCA=use_PCA, simulate_interactions=false)


text:
using DataAPI

using CategoricalArrays: CategoricalVector, CategoricalArray, CategoricalPool

function compress_then_write(b::CategoricalVector{T,IntType}, io) where {T, IntType<:Integer}
    compress_refs = compress_then_write(b.refs, io)
    compress_poolindex = compress_then_write(DataAPI.levels(b), io)

    (
        type = CategoricalVector,
        refs = compress_refs,
        poolindex = compress_poolindex,
        ordered = b.pool.ordered,
    )
end

# function column_loader(::Type{CategoricalVector{Union{Missing, T}, I}}, io, metadata) where {T, I}
#     println("got here1")
#     refs_meta = metadata.refs
#     pi_meta = metadata.poolindex
#     ref = column_loader(refs_meta.type, io, refs_meta)
#     poolindex = column_loader(pi_meta.type, io, pi_meta)

#     return CategoricalArray{pi_meta.type,1}(
#         ref,
#         CategoricalPool{eltype(poolindex),eltype(ref)}(Array(poolindex), metadata.ordered),
#     )
# end

function column_loader(::Type{CategoricalVector}, io, metadata)
    refs_meta = metadata.refs
    pi_meta = metadata.poolindex
    ref = column_loader(refs_meta.type, io, refs_meta)
    poolindex = column_loader(pi_meta.type, io, pi_meta)

    # this checks for missing in the values which would be represented by ref = 0
    if any(==(0), ref)
        return CategoricalArray{Union{pi_meta.type, Missing},1}(
            ref,
            CategoricalPool{eltype(poolindex),eltype(ref)}(Array(poolindex), metadata.ordered),
        )
    else
        # no missing in the values, just return
        return CategoricalArray{pi_meta.type,1}(
            ref,
            CategoricalPool{eltype(poolindex),eltype(ref)}(Array(poolindex), metadata.ordered),
        )
    end
end

if false
    # # compress_then_write-cate.jl
    # using Revise
    # using JDF
    # using DataFrames, Debugger
    # b = categorical(["a", "b", "a", "c"])
    # io = open("c:/data/test.io", "w")
    # metadata = compress_then_write(a, io)
    # close(io)
    #
    # io = open("c:/data/test.io", "r")
    # aa = column_loader(CategoricalVector, io, metadata)
    # close(io)
    #
    # df = DataFrame(a = a)
    #
    # savejdf("c:/data/pls_del.jdf", df)
    #
    # loadjdf("c:/data/pls_del.jdf")
end

ext: jl

max_stars_repo_name: xiaodaigh/JDF

max_issues_repo_name: xiaodaigh/JDF

max_forks_repo_name: xiaodaigh/JDF


text:
#
# Datasets
#

# Use an Enum here when it doesn't make us allocate
const LC_COMPACT_STORAGE = 0x00
const LC_CONTIGUOUS_STORAGE = 0x01
const LC_CHUNKED_STORAGE = 0x02

function load_dataset(f::JLDFile, offset::RelOffset)
    if haskey(f.jloffset, offset)
        # Stored as WeakRefs and may no longer exist
        val = f.jloffset[offset].value
        val !== nothing && return val
    end

    io = f.io
    seek(io, fileoffset(f, offset))
     # Test for object header version
    header_version = jlread(io, UInt8)
    if header_version == 1
        cio = io
        jlread(cio, UInt8)
        num_messages = jlread(cio, UInt16)
        obj_ref_count = jlread(cio, UInt32)
        obj_header_size = jlread(cio, UInt32)
        # println("Reading $num_messages Messages")
        # println("obj_ref_count = $obj_ref_count")
        # println("obj_header_size = $obj_header_size")

        chunk_end = position(cio) + obj_header_size
        # Skip to nearest 8byte aligned position
        skip_to_aligned!(cio, fileoffset(f, offset))

    elseif header_version == 2
        seek(io, fileoffset(f, offset))
        cio = begin_checksum_read(io)
        sz, version = read_obj_start(cio)
        chunk_end = position(cio) + sz #- 4
    else
        throw(UnsupportedVersionException("Unknown object header version $header_version"))
    end

    # Messages
    dataspace = ReadDataspace()
    attrs = EMPTY_READ_ATTRIBUTES
    datatype_class::UInt8 = 0
    datatype_offset::Int64 = 0
    data_offset::Int64 = 0
    data_length::Int = -1
    chunked_storage::Bool = false
    filter_id::UInt16 = 0
    while (curpos = position(cio)) <= chunk_end-4
        if header_version == 1
            msg_type = jlread(cio, UInt16)
            msg_size = jlread(cio, UInt16)
            flags = jlread(cio, UInt8)
            jlread(cio, UInt8); jlread(cio, UInt16)
            endpos = curpos + 8 + msg_size
            endpos = endpos + 8 - mod1(endpos-fileoffset(f, offset), 8)
            msg = (; msg_type, size=msg_size, flags)
            # println("Message Type: $msg_type $(MESSAGE_TYPES[msg_type])")
            # println("Message Size: $msg_size")
            # println("Message flags: $flags")
            # println("Message ends at $endpos")
        else # header_version == 2
            msg = jlread(cio, HeaderMessage)
            endpos = curpos + jlsizeof(HeaderMessage) + msg.size
        end
        if msg.msg_type == HM_DATASPACE
            dataspace = read_dataspace_message(cio)
        elseif msg.msg_type == HM_DATATYPE
            datatype_class, datatype_offset = read_datatype_message(cio, f, (msg.flags & 2) == 2)
        elseif msg.msg_type == HM_FILL_VALUE_OLD
            #(jlread(cio, UInt8) == 3 && jlread(cio, UInt8) == 0x09) || throw(UnsupportedFeatureException())
        elseif msg.msg_type == HM_FILL_VALUE
            #(jlread(cio, UInt8) == 3 && jlread(cio, UInt8) == 0x09) || throw(UnsupportedFeatureException())
        elseif msg.msg_type == HM_DATA_LAYOUT
            version = jlread(cio, UInt8)
            if version == 4 || version == 3
                storage_type = jlread(cio, UInt8)
                if storage_type == LC_COMPACT_STORAGE
                    data_length = jlread(cio, UInt16)
                    data_offset = position(cio)
                elseif storage_type == LC_CONTIGUOUS_STORAGE
                    data_offset = fileoffset(f, jlread(cio, RelOffset))
                    data_length = jlread(cio, Length)
                elseif storage_type == LC_CHUNKED_STORAGE
                    # TODO: validate this
                    flags = jlread(cio, UInt8)
                    dimensionality = jlread(cio, UInt8)
                    dimensionality_size = jlread(cio, UInt8)
                    skip(cio, Int(dimensionality)*Int(dimensionality_size))

                    chunk_indexing_type = jlread(cio, UInt8)
                    chunk_indexing_type == 1 || throw(UnsupportedFeatureException("Unknown chunk indexing type"))
                    data_length = jlread(cio, Length)
                    jlread(cio, UInt32)
                    data_offset = fileoffset(f, jlread(cio, RelOffset))
                    chunked_storage = true
                else
                    throw(UnsupportedFeatureException("Unknown data layout"))
                end
            end
        elseif msg.msg_type == HM_FILTER_PIPELINE
            version = jlread(cio, UInt8)
            version == 2 || throw(UnsupportedVersionException("Filter Pipeline Message version $version is not implemented"))
            nfilters = jlread(cio, UInt8)
            nfilters == 1 || throw(UnsupportedFeatureException())
            filter_id = jlread(cio, UInt16)
            issupported_filter(filter_id) || throw(UnsupportedFeatureException("Unknown Compression Filter $filter_id"))
        elseif msg.msg_type == HM_ATTRIBUTE
            if attrs === EMPTY_READ_ATTRIBUTES
                attrs = ReadAttribute[read_attribute(cio, f)]
            else
                push!(attrs, read_attribute(cio, f))
            end
        elseif (msg.flags & 2^3) != 0
            throw(UnsupportedFeatureException())
        end
        seek(cio, endpos)
    end
    seek(cio, chunk_end)

    filter_id != 0 && !chunked_storage && throw(InvalidDataException("Compressed data must be chunked"))
    if header_version == 2
        # Checksum
        end_checksum(cio) == jlread(io, UInt32) || throw(InvalidDataException("Invalid Checksum"))
    end
    # TODO verify that data length matches
    val = read_data(f, dataspace, datatype_class, datatype_offset, data_offset, data_length,
                    filter_id, offset, attrs)
    val
end

"""
    read_attr_data(f::JLDFile, attr::ReadAttribute)

jlread data from an attribute.
"""
read_attr_data(f::JLDFile, attr::ReadAttribute) =
    read_data(f, attr.dataspace, attr.datatype_class, attr.datatype_offset,
              attr.data_offset)

"""
    read_attr_data(f::JLDFile, attr::ReadAttribute, expected_datatype::H5Datatype,
                   rr::ReadRepresentation)

jlread data from an attribute, assuming a specific HDF5 datatype and ReadRepresentation. If
the HDF5 datatype does not match, throws an `UnsupportedFeatureException`. This allows
better type stability while simultaneously validating the data.
"""
function read_attr_data(f::JLDFile, attr::ReadAttribute, expected_datatype::H5Datatype,
                        rr::ReadRepresentation)
    io = f.io
    if attr.datatype_class == class(expected_datatype)
        seek(io, attr.datatype_offset)
        dt = jlread(io, typeof(expected_datatype))
        if dt == expected_datatype
            seek(f.io, attr.data_offset)
            read_dataspace = (attr.dataspace, NULL_REFERENCE, -1, UInt16(0))
            return read_data(f, rr, read_dataspace)
        end
    end
    throw(UnsupportedFeatureException())
end

"""
    read_data(f::JLDFile, dataspace::ReadDataspace, datatype_class::UInt8,
              datatype_offset::Int64, data_offset::Int64[, filter_id::UInt16,
              header_offset::RelOffset, attributes::Vector{ReadAttribute}])

Read data from a file. If `datatype_class` is typemax(UInt8), the datatype is assumed to be
committed, and `datatype_offset` points to the offset of the committed datatype's header.
Otherwise, datatype_offset points to the offset of the datatype attribute.
"""
function read_data(f::JLDFile, dataspace::ReadDataspace,
                   datatype_class::UInt8, datatype_offset::Int64,
                   data_offset::Int64, data_length::Int=-1, filter_id::UInt16=UInt16(0),
                   header_offset::RelOffset=NULL_REFERENCE,
                   attributes::Union{Vector{ReadAttribute},Nothing}=nothing)
    # See if there is a julia type attribute
    io = f.io
    if datatype_class == typemax(UInt8) # Committed datatype
        rr = jltype(f, f.datatype_locations[h5offset(f, datatype_offset)])
        seek(io, data_offset)
        read_dataspace = (dataspace, header_offset, data_length, filter_id)
        read_data(f, rr, read_dataspace, attributes)
    else
        seek(io, datatype_offset)
        @read_datatype io datatype_class dt begin
            rr = jltype(f, dt)
            seek(io, data_offset)
            read_dataspace = (dataspace, header_offset, data_length, filter_id)
            read_data(f, rr, read_dataspace, attributes)
        end
    end
end

# Most types can only be scalars or arrays
function read_data(f::JLDFile,
     @nospecialize(rr),
     read_dataspace::Tuple{ReadDataspace,RelOffset,Int,UInt16},
     attributes::Union{Vector{ReadAttribute},Nothing}=nothing)

    dataspace, header_offset, data_length, filter_id = read_dataspace
    if dataspace.dataspace_type == DS_SCALAR
        filter_id != 0 && throw(UnsupportedFeatureException())
        read_scalar(f, rr, header_offset)
    elseif dataspace.dataspace_type == DS_SIMPLE
        read_array(f, dataspace, rr, data_length, filter_id, header_offset, attributes)
    elseif dataspace.dataspace_type == DS_V1 && dataspace.dimensionality == 0
        read_scalar(f, rr, header_offset)
    elseif dataspace.dataspace_type == DS_V1
        read_array(f, dataspace, rr, data_length, filter_id, header_offset, attributes)
    else
        throw(UnsupportedFeatureException())
    end
end

# Reference arrays can only be arrays or null dataspace (for Union{} case)
function read_data(f::JLDFile,
    rr::ReadRepresentation{Any,RelOffset},
    read_dataspace::Tuple{ReadDataspace,RelOffset,Int,UInt16},
    attributes::Vector{ReadAttribute})

    dataspace, header_offset, data_length, filter_id = read_dataspace
    filter_id != 0 && throw(UnsupportedFeatureException())
    if dataspace.dataspace_type == DS_SIMPLE
        # Since this is an array of references, there should be an attribute
        # informing us of the type
        io = f.io
        startpos = position(io)
        for x in attributes
            if x.name == :julia_type
                T = read_attr_data(f, x)
                if isa(T, UnknownType)
                    str = typestring(T)
                    @warn("type $(str) does not exist in workspace; interpreting Array{$str} as Array{Any}")
                    T = Any
                end
                seek(io, startpos)
                return read_array(f, dataspace, ReadRepresentation{T,RelOffset}(),
                                  -1, UInt16(0), header_offset, attributes)
            end
        end
    elseif dataspace.dataspace_type == DS_NULL
        return read_empty(ReadRepresentation{Union{},nothing}(), f,
                          attributes[find_dimensions_attr(attributes)],
                          header_offset)
    end
    throw(UnsupportedFeatureException())
end

# Types with no payload can only be null dataspace
function read_data(f::JLDFile,
                   rr::Union{ReadRepresentation{T,nothing} where T,
                             ReadRepresentation{T,CustomSerialization{S,nothing}} where {S,T}},
                   read_dataspace::Tuple{ReadDataspace,RelOffset,Int,UInt16},
                   attributes::Vector{ReadAttribute})
    dataspace, header_offset, data_length, filter_id = read_dataspace
    filter_id != 0 && throw(UnsupportedFeatureException())
    dataspace.dataspace_type == DS_NULL || throw(UnsupportedFeatureException())

    dimensions_attr_index = find_dimensions_attr(attributes)
    if dimensions_attr_index == 0
        jlconvert(rr, f, Ptr{Cvoid}(0), header_offset)
    else
        # dimensions attribute => array of empty type
        read_empty(rr, f, attributes[dimensions_attr_index], header_offset)
    end
end

function find_dimensions_attr(attributes::Vector{ReadAttribute})
    dimensions_attr_index = 0
    julia_type_attr_index = 0
    for i = 1:length(attributes)
        x = attributes[i]
        if x.name == :dimensions
            dimensions_attr_index = i
        end
    end
    dimensions_attr_index
end

function read_empty(rr::ReadRepresentation{T}, f::JLDFile,
                 dimensions_attr::ReadAttribute, header_offset::RelOffset) where T
    dimensions_attr.datatype_class == DT_FIXED_POINT || throw(UnsupportedFeatureException())

    io = f.io
    seek(io, dimensions_attr.dataspace.dimensions_offset)
    ndims = Int(jlread(io, Length))

    seek(io, dimensions_attr.datatype_offset)
    jlread(io, FixedPointDatatype) == h5fieldtype(f, Int64, Int64, Val{true}) || throw(UnsupportedFeatureException())

    seek(io, dimensions_attr.data_offset)
    v = construct_array(io, T, Val(ndims))
    if isconcretetype(T)
        for i = 1:length(v)
            @inbounds v[i] = jlconvert(rr, f, Ptr{Cvoid}(0), header_offset)
        end
    end
    header_offset !== NULL_REFERENCE && (f.jloffset[header_offset] = WeakRef(v))
    v
end

get_ndims_offset(f::JLDFile, dataspace::ReadDataspace, attributes::Nothing) =
    (dataspace.dimensionality, dataspace.dimensions_offset)

function get_ndims_offset(f::JLDFile, dataspace::ReadDataspace, attributes::Vector{ReadAttribute})
    ndims = dataspace.dimensionality
    offset = dataspace.dimensions_offset
    if !isempty(attributes)
        for x in attributes
            if x.name == :dimensions
                (x.dataspace.dataspace_type == DS_SIMPLE &&
                 x.dataspace.dimensionality == 1) || throw(InvalidDataException())
                seek(f.io, x.dataspace.dimensions_offset)
                ndims = UInt8(jlread(f.io, Length))
                offset = x.data_offset
            end
        end
    end
    (ndims, offset)
end

"""
    construct_array{T}(io::IO, ::Type{T}, ::Val{ndims})

Construct array by reading `ndims` dimensions from `io`. Assumes `io` has already been
seeked to the correct position.
"""
function construct_array(io::IO, ::Type{T}, ::Val{1}) where {T}
    n = jlread(io, Int64)
    Vector{T}(undef, n)
end

function construct_array(io::IO, ::Type{T}, ::Val{2}) where {T}
    d2 = jlread(io, Int64)
    d1 = jlread(io, Int64)
    Matrix{T}(undef, d1, d2)
end

function construct_array(io::IO, ::Type{T}, ::Val{3}) where {T}
    d3 = jlread(io, Int64)
    d2 = jlread(io, Int64)
    d1 = jlread(io, Int64)
    Array{T,3}(undef, d1, d2, d3)
end

function construct_array(io::IO, ::Type{T}, ::Val{N})::Array{T,N} where {T,N}
    ds = reverse(ntuple(i->jlread(io, Int64), Val(N)))
    Array{T,N}(undef, ds...)
end

function read_array(f::JLDFile, dataspace::ReadDataspace,
                    rr::ReadRepresentation{T,RR}, data_length::Int,
                    filter_id::UInt16, header_offset::RelOffset,
                    attributes::Union{Vector{ReadAttribute},Nothing}) where {T,RR}
    io = f.io
    data_offset = position(io)
    ndims, offset = get_ndims_offset(f, dataspace, attributes)

    seek(io, offset)
    v = construct_array(io, T, Val(Int(ndims)))
    n = length(v)
    seek(io, data_offset)
    if filter_id !=0
        read_compressed_array!(v, f, rr, data_length, filter_id)
    else
        read_array!(v, f, rr)
    end
    header_offset !== NULL_REFERENCE && (f.jloffset[header_offset] = WeakRef(v))
    v
end

function payload_size_without_storage_message(dataspace::WriteDataspace, datatype::H5Datatype)
    sz = 6 + 4 + jlsizeof(dataspace)::Int + 4 + jlsizeof(datatype) + length(dataspace.attributes)*jlsizeof(HeaderMessage)
    for attr in dataspace.attributes
        sz += jlsizeof(attr)::Int
    end
    sz
end

function write_dataset(
        f::JLDFile,
        dataspace::WriteDataspace,
        datatype::H5Datatype,
        odr::S,
        data::Array{T},
        wsession::JLDWriteSession,
        compress = f.compress,
        ) where {T,S}
    io = f.io
    datasz = odr_sizeof(odr)::Int * numel(dataspace)::Int
    #layout_class
    if datasz < 8192
        layout_class = LC_COMPACT_STORAGE
    elseif compress != false
        layout_class = LC_CHUNKED_STORAGE
    else
        layout_class = LC_CONTIGUOUS_STORAGE
    end
    psz = payload_size_without_storage_message(dataspace, datatype)::Int
    if datasz < 8192
        layout_class = LC_COMPACT_STORAGE
        psz += jlsizeof(CompactStorageMessage) + datasz
    elseif compress != false && isconcretetype(T) && isbitstype(T)
        # Only now figure out if the compression argument is valid
        invoke_again, filter_id, compressor = get_compressor(compress)
        if invoke_again
            return Base.invokelatest(write_dataset, f, dataspace, datatype, odr, data, wsession, compress)::RelOffset
        end
        layout_class = LC_CHUNKED_STORAGE
        psz += chunked_storage_message_size(ndims(data)) + pipeline_message_size(filter_id::UInt16)
    else
        layout_class = LC_CONTIGUOUS_STORAGE
        psz += jlsizeof(ContiguousStorageMessage)
    end
    fullsz = jlsizeof(ObjectStart) + size_size(psz) + psz + 4

    header_offset = f.end_of_data
    seek(io, header_offset)
    f.end_of_data = header_offset + fullsz

    if !isa(wsession, JLDWriteSession{Union{}})
        wsession.h5offset[objectid(data)] = h5offset(f, header_offset)
        push!(wsession.objects, data)
    end

    cio = begin_checksum_write(io, fullsz - 4)
    write_object_header_and_dataspace_message(cio, f, psz, dataspace)
    write_datatype_message(cio, datatype)

    # Data storage layout
    if layout_class == LC_COMPACT_STORAGE
        jlwrite(cio, CompactStorageMessage(datasz))
        if datasz != 0
            write_data(cio, f, data, odr, datamode(odr), wsession)
        end
        jlwrite(io, end_checksum(cio))
    elseif layout_class == LC_CHUNKED_STORAGE

        write_compressed_data(cio, f, data, odr, wsession, filter_id, compressor)

    else
        jlwrite(cio, ContiguousStorageMessage(datasz, h5offset(f, f.end_of_data)))
        jlwrite(io, end_checksum(cio))

        f.end_of_data += datasz
        write_data(io, f, data, odr, datamode(odr), wsession)
    end

    h5offset(f, header_offset)
end

function write_dataset(f::JLDFile, dataspace::WriteDataspace, datatype::H5Datatype, odr::S, data, wsession::JLDWriteSession) where S
    io = f.io
    datasz = (odr_sizeof(odr)::Int * numel(dataspace))
    psz = payload_size_without_storage_message(dataspace, datatype)

    # The simplest CompactStorageMessage only supports data sets < 2^16
    if datasz < typemax(UInt16)
        layout_class = LC_COMPACT_STORAGE
        storage_message = CompactStorageMessage
        psz += jlsizeof(CompactStorageMessage) + datasz
    else
        layout_class = LC_CONTIGUOUS_STORAGE
        storage_message = ContiguousStorageMessage
        psz += jlsizeof(ContiguousStorageMessage)
    end

    fullsz = jlsizeof(ObjectStart) + size_size(psz) + psz + 4

    header_offset = f.end_of_data
    seek(io, header_offset)
    f.end_of_data = header_offset + fullsz

    if ismutabletype(typeof(data)) && !isa(wsession, JLDWriteSession{Union{}})
        wsession.h5offset[objectid(data)] = h5offset(f, header_offset)
        push!(wsession.objects, data)
    end

    cio = begin_checksum_write(io, fullsz - 4)
    write_object_header_and_dataspace_message(cio, f, psz, dataspace)
    write_datatype_message(cio, datatype)

    # Data storage layout
    if layout_class == LC_COMPACT_STORAGE
        jlwrite(cio, CompactStorageMessage(datasz))
        if datasz != 0
            write_data(cio, f, data, odr, datamode(odr), wsession)
        end
        jlwrite(io, end_checksum(cio))
    else
        jlwrite(cio, ContiguousStorageMessage(datasz, h5offset(f, f.end_of_data)))
        jlwrite(io, end_checksum(cio))

        f.end_of_data += datasz
        write_data(io, f, data, odr, datamode(odr), wsession)
    end

    h5offset(f, header_offset)
end

function write_object_header_and_dataspace_message(cio::IO, f::JLDFile, psz::Int, dataspace::WriteDataspace)
    jlwrite(cio, ObjectStart(size_flag(psz)))
    write_size(cio, psz)

    # Fill value
    jlwrite(cio, HeaderMessage(HM_FILL_VALUE, 2, 0))
    jlwrite(cio, UInt8(3)) # Version
    jlwrite(cio, 0x09)     # Flags

    # Dataspace
    jlwrite(cio, HeaderMessage(HM_DATASPACE, jlsizeof(dataspace), 0))
    jlwrite(cio, dataspace)

    # Attributes
    for attr in dataspace.attributes
        jlwrite(cio, HeaderMessage(HM_ATTRIBUTE, jlsizeof(attr), 0))
        write_attribute(cio, f, attr, f.datatype_wsession)
    end
end

function write_datatype_message(cio::IO, datatype::H5Datatype)
    jlwrite(cio, HeaderMessage(HM_DATATYPE, jlsizeof(datatype), 1 | (2*isa(datatype, CommittedDatatype))))
    jlwrite(cio, datatype)
end

struct CompactStorageMessage
    hm::HeaderMessage
    version::UInt8
    layout_class::UInt8
    data_size::UInt16
end
define_packed(CompactStorageMessage)
@inline CompactStorageMessage(datasz::Int) =
    CompactStorageMessage(
            HeaderMessage(HM_DATA_LAYOUT, jlsizeof(CompactStorageMessage) - jlsizeof(HeaderMessage) + datasz, 0),
            4, LC_COMPACT_STORAGE, datasz
    )

struct ContiguousStorageMessage
    hm::HeaderMessage
    version::UInt8
    layout_class::UInt8
    address::RelOffset
    data_size::Length
end
define_packed(ContiguousStorageMessage)
@inline ContiguousStorageMessage(datasz::Int, offset::RelOffset) =
    ContiguousStorageMessage(
        HeaderMessage(HM_DATA_LAYOUT, jlsizeof(ContiguousStorageMessage) - jlsizeof(HeaderMessage), 0),
        4, LC_CONTIGUOUS_STORAGE, offset, datasz
    )

function write_dataset(f::JLDFile, x, wsession::JLDWriteSession)
    odr = objodr(x)
    write_dataset(f, WriteDataspace(f, x, odr), h5type(f, x), odr, x, wsession)::RelOffset
end

function write_ref_mutable(f::JLDFile, x, wsession::JLDWriteSession)
    offset = get(wsession.h5offset, objectid(x), RelOffset(0))
    offset != RelOffset(0) ? offset : write_dataset(f, x, wsession)::RelOffset
end

write_ref_mutable(f::JLDFile, x, wsession::JLDWriteSession{Union{}}) =
    write_dataset(f, x, wsession)

function write_ref(f::JLDFile, x, wsession::JLDWriteSession)
    if ismutabletype(typeof(x))
        write_ref_mutable(f, x, wsession)::RelOffset
    else
        write_dataset(f, x, wsession)::RelOffset
    end
end
write_ref(f::JLDFile, x::RelOffset, wsession::JLDWriteSession) = x

Base.delete!(f::JLDFile, x::AbstractString) = delete!(f.root_group, x)

function Base.delete!(g::Group, name::AbstractString)
    g.f.writable || throw(ArgumentError("Cannot delete in read-only mode"))
    if !haskey(g, name)
        @warn "No entry named $name was found"
        return
    end

    if '/' in name
        dir, dname = rsplit(name, '/'; limit=2)
        if isempty(dir)
            g = g.f.root_group
        else
            g = g[dir]
        end
        name = string(dname)
    end

    # Simple case first. If it hasn't been written yet,
    # the file doesn't need to be altered.
    if haskey(g.unwritten_links, name)
        delete!(g.unwritten_links, name)
        return
    elseif haskey(g.unwritten_child_groups, name)
        delete!(g.unwritten_child_groups, name)
        return
    end

    # Dataset must already exist in the file
    # Retrieve offset of group in file
    offset = group_offset(g)
    offset == NULL_REFERENCE && throw(InternalError("Group could not be found."))
    delete_written_link!(g.f, offset, name)
    delete!(g.written_links, name)
    return
end

function group_offset(g::Group)
    # Given a group, retrieve its offset inside the parent file
    # Return a null reference when it can't be found. However, that
    # should never be the case as the calling functions checks for this
    if g === g.f.root_group
        return g.f.root_group_offset
    end
    for (offset, lg) in pairs(g.f.loaded_groups)
        if lg === g
            return offset
        end
    end
    return NULL_REFERENCE
end

function delete_written_link!(f::JLDFile, roffset::RelOffset, name::AbstractString)
    # Location of written link in group structure is not known a priory
    # Instead, we walk through the structure similar to `load_group`
    # until the correct link message is found
    # The deletion is done by replacing that message by a placeholder nill message
    io = f.io
    chunk_start_offset::Int64 = fileoffset(f, roffset)
    seek(io, chunk_start_offset)

    sz = read_obj_start(io)
    chunk_checksum_offset::Int64 = position(io) + sz

    continuation_offset::Int64 = -1
    continuation_length::Length = 0
    link_deleted = false
    while !link_deleted
        if continuation_offset != -1
            seek(io, continuation_offset)
            chunk_start_offset = continuation_offset
            chunk_checksum_offset = continuation_offset + continuation_length - 4
            continuation_offset = -1
            jlread(io, UInt32) == OBJECT_HEADER_CONTINUATION_SIGNATURE || throw(InvalidDataException())
        end

        while (curpos = position(io)) <= chunk_checksum_offset - 4
            msg = jlread(io, HeaderMessage)
            endpos = curpos + jlsizeof(HeaderMessage) + msg.size

            if msg.msg_type == HM_LINK_MESSAGE
                dataset_name, loffset = read_link(io)
                if dataset_name == name
                    # delete link
                    seek(io, curpos)
                    jlwrite(io, HeaderMessage(HM_NIL, msg.size, 0))
                    link_deleted = true
                    break
                end
            elseif msg.msg_type == HM_OBJECT_HEADER_CONTINUATION
                continuation_offset = chunk_start_offset = fileoffset(f, jlread(io, RelOffset))
                continuation_length = jlread(io, Length)
            end
            seek(io, endpos)
        end

        continuation_offset == -1 && break
    end

    # Update the Checksum
    seek(io, chunk_start_offset)
    cio = begin_checksum_read(io)
    seek(cio, chunk_checksum_offset)
    seek(io, chunk_checksum_offset)
    jlwrite(io, end_checksum(cio))
    return
end


text:
struct LinearAnalysis
    options::Dict{Symbol, Any}
end

LinearAnalysis(; kwargs...) = LinearAnalysis(Dict{Symbol, Any}(kwargs))

add_intercept_column(x::AbstractVector{T}) where {T} = [ones(T, length(x)) x]

# TODO: add multidimensional version
function (l::LinearAnalysis)(le::Entry)
    entry = map(le) do p, n
        x, y = p
        weights = get(n, :weights, similar(x, 0))
        npoints = get(l.options, :npoints, 200)
        interval = get(l.options, :interval, length(weights) > 0 ? nothing : :confidence)
        dropcollinear = get(l.options, :dropcollinear, false)
        lin_model = GLM.lm(add_intercept_column(x), collect(y); wts=weights, dropcollinear)
        isnothing(lin_model) && return Entry[]
        x̂ = range(extrema(x)..., length=npoints)
        pred = GLM.predict(lin_model, add_intercept_column(x̂); interval)
        return if !isnothing(interval)
            ŷ, lower, upper = map(vec, pred) # GLM prediction returns matrices
            (x̂, ŷ), (; lower, upper)
        else
            ŷ = vec(pred) # GLM prediction returns matrix
            (x̂, ŷ), (;)
        end
    end
    default_plottype = isempty(entry.named) ? Lines : LinesFill
    plottype = Makie.plottype(entry.plottype, default_plottype)
    return Entry(entry; plottype)
end

"""
    linear(; interval)

Compute a linear fit of `y ~ 1 + x`. An optional named mapping `weights` determines the weights.
Use `interval` to specify what type of interval the shaded band should represent.
Valid values of interval are `:confidence` delimiting the uncertainty of the predicted
relationship, and `:prediction` delimiting estimated bounds for new data points.
"""
linear(; kwargs...) = transformation(LinearAnalysis(; kwargs...))


text:
relaxed_tols = (default_tol_relax = 1000,)
insts = OrderedDict()
insts["minimal"] = [
    ((:matpoly2, true, true),),
    ((:matpoly5, true, true),),
    ((:matpoly5, true, false),),
    ((:matpoly5, false, false), :SOCExpPSD),
    ]
insts["fast"] = [
    ((:matpoly1, true, true),),
    ((:matpoly1, true, false),),
    ((:matpoly1, false, false), :SOCExpPSD),
    ((:matpoly2, true, true),),
    ((:matpoly2, true, false),),
    ((:matpoly2, false, false), :SOCExpPSD, relaxed_tols),
    ((:matpoly3, true, true),),
    ((:matpoly3, true, false),),
    ((:matpoly3, false, false), :SOCExpPSD),
    ((:matpoly4, true, true),),
    ((:matpoly4, true, false),),
    ((:matpoly4, false, false), :SOCExpPSD),
    ((:matpoly6, true, true),),
    ((:matpoly6, true, false),),
    ((:matpoly6, false, false), :SOCExpPSD),
    ((:matpoly7, true, true), nothing, relaxed_tols),
    ((:matpoly7, true, false),),
    ((:matpoly7, false, false), :SOCExpPSD),
    ]
insts["various"] = insts["fast"]
return (SemidefinitePolyJuMP, insts)


text:
using JuMP, Plots, LinearAlgebra, SparseArrays, Ipopt

#Model initialization
mod = Model( optimizer_with_attributes(Ipopt.Optimizer , "max_iter" => 100,
                                        "mumps_mem_percent" => 500))
#The parameters of the 2-D body
λ = 45.0                              # Conductivity
c = 460.0                             # Specific heat capacitivity
ρ = 7800.0                            # Material density
α = λ / (c*ρ)                         # Diffusivity
L_y = 0.1                             #Width of rod, L_x = L_y / 3
N_y = 12                              #Discretization nodes in y-direction
N_x = Int(round(N_y / 3));            #Discretization nodes in x-direction
h = L_y / (N_y - 1)                   #step size Δx = Δy

#System , input and output matrices: θ' = Aθ(t) + Bu(t)
A = spdiagm(0 => -4*ones(N_x*N_y))
for i in 1:N_x * N_y
    if i % N_x == 1
        A[i+1,i] = 2
    elseif i % N_x == 0
        A[i-1,i] = 2
    else
        A[i-1,i] = 1
        A[i+1,i] = 1
    end
    if (i <= N_x)
        A[i+N_x,i] = 2
    elseif N_x * N_y - i < N_x
        A[i- N_x,i] = 2
    else
        A[i-N_x,i] = 1
        A[i+N_x,i] = 1
    end
end

A *= α / (h^2);

B = spzeros(3,N_x*N_y);
B[1,1] = 1;
B[2,2] = 1;
B[3,3] = 1;
B *= (2/ (c * ρ * h));

#y(t) = Cθ(t)
C = spzeros(N_x*N_y,1)
C[N_x*N_y] = 1.0;

#Integration from t0 = 0 to tf
tf = 30000.0;                   #Final time
dt = 1.5;                       #Sampling period Δt
steps = round(Int,tf / dt);     #number of time steps in simulation.

u_min = 0.0                     #lower input bound
u_max = 15000.0                 #upper input bound
@variable(mod,θ1[1:N_x*N_y,1:steps]);                                       # state vector plate1
@variable(mod,θ2[1:N_x*N_y,1:steps]);                                       # state vector plate2
@variable(mod,θ3[1:N_x*N_y,1:steps]);                                       # state vector plate3
@variable(mod,u_min <= u[i = 1:3,k = 1:steps - 1] <= u_max);                # Bounded Input u

θinit = 273.0                  #initial temperature of the body
θref = [400,600,500]           #Reference temperature for each plate

#Initial value added as constraints
@constraint(mod, θ1[:,1] .== θinit)
@constraint(mod, θ2[:,1] .== θinit)
@constraint(mod, θ3[:,1] .== θinit)
#System dynamics
for k in 1:steps - 1,i in 1:N_x*N_y
    @constraint(mod,θ1[i,k + 1] .==  θ1[i,k] + dt * ((A[:,i]' * θ1[:,k]) + B[1,i]' * u[1,k]))
    @constraint(mod,θ2[i,k + 1] ==  θ2[i,k] + dt * ((A[:,i]' * θ2[:,k]) + B[2,i]' * u[2,k]))
    @constraint(mod,θ3[i,k + 1] ==  θ3[i,k] + dt * ((A[:,i]' * θ3[:,k]) + B[3,i]' * u[3,k]))
end
#Q and R weighing matrices
Q = 1000 * C * C';
R = 0.001 * I;
#Cost function
@NLobjective(mod, Min, 0.5 * dt * sum(sum(Q[i,i] *(θref[1] - θ1[i,k])^2 for i in 1:N_x*N_y)
+sum(Q[i,i] *(θref[2] - θ2[i,k])^2 for i in 1:N_x*N_y) + sum(Q[i,i] *(θref[3] - θ3[i,k])^2 for i in 1:N_x*N_y)
+ sum(R[i,i] * u[i,k]^2 for i in 1:3) for k in 1:steps - 1))
#Optimization process
optimize!(mod)
#Plotting results
states_θstart1 = zeros(steps)
states_θstart2 = zeros(steps)
states_θstart3 = zeros(steps)
states_θend3 = zeros(steps)
states_θend2 = zeros(steps)
states_θend1 = zeros(steps)
input_u1 = zeros(steps - 1)
input_u2 = zeros(steps - 1)
input_u3 = zeros(steps - 1)
for i = 1 : steps-1
    states_θstart1[i] = JuMP.value(θ1[1,i])
    states_θstart2[i] = JuMP.value(θ2[2,i])
    states_θstart3[i] = JuMP.value(θ3[3,i])
    states_θend3[i] = JuMP.value(θ3[N_x*N_y,i])
    states_θend2[i] = JuMP.value(θ2[N_x*N_y,i])
    states_θend1[i] = JuMP.value(θ1[N_x*N_y,i])
    input_u1[i] = JuMP.value(u[1,i])
    input_u2[i] = JuMP.value(u[2,i])
    input_u3[i] = JuMP.value(u[3,i])
end

states_θend3[end] = JuMP.value(θ3[end,end])
states_θend2[end] = JuMP.value(θ2[end,end])
states_θend1[end] = JuMP.value(θ1[end,end])
states_θstart1[end] = JuMP.value(θ1[1,end])
states_θstart2[end] = JuMP.value(θ2[2,end])
states_θstart3[end] = JuMP.value(θ3[3,end])
#Plotting temperatures at the input nodes
p2 = plot(states_θstart1, label = "θ_in1", xlabel = "time[s]", ylabel = "Temperature[K]")
p2 = plot!(states_θstart2,label = "θ_in2")
p2 = plot!(states_θstart3,label = "θ_in3")
 savefig("lq_optimal_2D_distributed.png")
#Plotting temperatures at the output
p1 = plot(states_θend3,label = "output-3", xlabel = "time[s]", ylabel = "Temperature[K]")
p1 = plot!(states_θend2,label = "output-2")
p1 = plot!(states_θend1,label = "output-1")
savefig("lq_optimal_2D-output-bounded_distributed.png")
#Plotting the input signals
p3 = plot(input_u1,label = "in_1", xlabel = "time[s]", ylabel = "Input")
p3 = plot!(input_u2,label = "in_2")
p3 = plot!(input_u3,label = "in_3")
savefig("lq_optimal_2D_inputbounded_distributed.png")


text:
NIScope_Misc_Dict = Dict()

# Reset
#reset_instr(obj::T) where {T<:NIScope}
# IDN
#get_idn(obj::T) where {T<:NIScope}
# Default setup
#default_setup(obj::T) where {T<:NIScope}
# Clear sweeps - not available for NIScope
#clear_sweeps(obj::T) where {T<:NIScope}

"""
scope = connect!(obj::T) where {T<:NIScope}
"""
function connect!(obj::T) where {T<:NIScope}
	obj_ptr = ViPSession(0)
	obj_Int = [UInt8.(collect(obj.address)); UInt8(0)] # terminate with NULL char
	status = niScope_init(obj_Int, Int16(0), Int16(0), obj_ptr)
	obj.handle = obj_ptr[]
	obj.initialized = true
	return obj
end

# close
"""
status = disconnect!(obj::T) where {T<:NIScope}
"""
function disconnect!(obj::T) where {T<:NIScope}
	status = niScope_close(obj.handle)
	obj.initialized = false
	return status
end

# IDN
function get_idn(obj::T) where {T<:NIScope}
	chlist = [UInt8.(collect("")); UInt8(0)] # terminate with NULL char
	bufSize = Int32(256)
	res    = zeros(UInt8,bufSize)
	pres   = Ref(res)
	status = niScope_GetAttributeViString(obj.handle,
		chlist, NISCOPE_ATTR_INSTRUMENT_MODEL, bufSize, res)
	return (status, unsafe_string(pointer(res)))
end

# Auto set
"""
status = niScope_AutoSetup(obj::T) where {T<:NIScope}
"""
function niScope_AutoSetup(obj::T) where {T<:NIScope}
	status = niScope_AutoSetup(obj.handle)
	return (status, niScope_GetError(obj))
end

# Trigger immediate
"""
status = niScope_ConfigureTriggerImmediate(obj::T) where {T<:NIScope}
"""
function niScope_ConfigureTriggerImmediate(obj::T) where {T<:NIScope}
	status = niScope_ConfigureTriggerImmediate(obj.handle)
	return (status, niScope_GetError(obj))
end

# Init acquisition
"""
status = niScope_InitiateAcquisition(obj::T) where {T<:NIScope}
"""
function niScope_InitiateAcquisition(obj::T) where {T<:NIScope}
	status = niScope_InitiateAcquisition(obj.handle)
	return (status, niScope_GetError(obj))
end

# niScope_ActualRecordLength
"""
record_length = niScope_ActualRecordLength(obj::T) where {T<:NIScope}
"""
function niScope_ActualRecordLength(obj::T) where {T<:NIScope}
#ViStatus niScope_ActualRecordLength (ViSession vi, ViInt32* recordLength);
	res    = Int32(0)
	pres   = ViPInt32(res)
	status = niScope_ActualRecordLength(obj.handle,pres)
	return pres[]
end


text:
module StandardYoungTableaux

import Distributions

export IPartition
export StandardYoungTableau
export dualSYT
export SYT2ballot
export ballot2SYT
export allSYTx
export allBallotSequences
export firstSYT
export nextSYT
export hooks
export hooklengths
export countSYTx
export RS
export SYT2YoungPath
export YoungPath2SYT
export randomSYT
export randomYoungPath

"""
    IPartition(partition)

Defines an integer partition.

# Argument

- `partition`: vector of positive integers in decreasing order
"""
struct IPartition
  n::Int64
  partition::Vector{Int64}

  function IPartition(partition::Vector{Int64}, check::Bool = true)
    partition = filter(!=(0), partition)
    if check
      if any(diff(partition) .> 0)
        throw(ArgumentError("Partition entries must be decreasing."))
      end
      if !isempty(partition)
        partition[end] >= 1 ||
          throw(ArgumentError("Found non-positive entry in partition."))
      end
    end
    return new(sum(partition), partition)
  end
end

function _jaggedTranspose(aa)
  out = Vector{Vector{Int64}}(undef, 0)
  while !isempty(aa)
    push!(out, map(x -> x[1], aa))
    aa = filter(x -> !isempty(x), map(x -> x[2:end], aa))
  end
  return out
end

function _checkSYTrows(tableau)
  all(map(row -> all(diff(row) .> 0), tableau))
end

function _isSYT(tableau) # _isSYT([[1], [2,3]]) is true ! - does not check shape
  entries = vcat(tableau...)
  return issetequal(entries, 1:length(entries)) &&
         _checkSYTrows(tableau) &&
         _checkSYTrows(_jaggedTranspose(tableau))
end

"""
    StandardYoungTableau(tableau)

Defines a standard Young tableau.

# Argument
- `tableau`: a vector of vectors of positive integers
"""
struct StandardYoungTableau
  shape::IPartition
  tableau::Vector{Vector{Int64}}

  function StandardYoungTableau(tableau::Vector{Vector{Int64}}, check::Bool = true)
    shape = map(length, tableau)
    if check
      if any(shape .== 0) || any(diff(shape) .> 0)
        throw(ArgumentError("Invalid shape."))
      end
      if !_isSYT(tableau)
        throw(ArgumentError("Not a valid tableau."))
      end
    end
    return new(IPartition(shape, false), tableau)
  end
end

"""
    dualSYT(syt)

Dual (conjugate) standard Young tableau.

# Argument
- `syt`: a standard Young tableau
"""
function dualSYT(syt::StandardYoungTableau)
  return StandardYoungTableau(_jaggedTranspose(syt.tableau), false)
end

function _isBallotSequence(s::Vector{Int64}) # https://oeis.org/A000085
  if s[1] != 1 || any(s .<= 0)
    return false
  end
  for k = 1:length(s)
    init = s[1:k]
    for r = 1:maximum(init)
      if length(findall(==(r), init)) < length(findall(==(r + 1), init))
        return false
      end
    end
  end
  return true
end

function _syt2ballot(tableau)
  N = sum(map(length, tableau))
  a = Int64[]
  for i = 1:N
    k = 0
    j = Int64[]
    while isempty(j)
      k += 1
      j = findall(==(i), tableau[k])
      if !isempty(j)
        push!(a, k)
      end
    end
  end
  return a
end

"""
    SYT2ballot(syt)

The ballot sequence corresponding to a standard Young tableau.

# Argument
- `syt`: a standard Young tableau
"""
function SYT2ballot(syt::StandardYoungTableau)
  return _syt2ballot(syt.tableau)
end

function _ballot2syt(a::Vector{Int64})
  return map(i -> findall(==(i), a), 1:maximum(a))
end

"""
    ballot2SYT(b)

The standard Young tableau corresponding to a ballot sequence.

# Argument
- `b`: a ballot sequence (vector of integers)
"""
function ballot2SYT(b::Vector{Int64})
  if !_isBallotSequence(b)
    error("Not a ballot sequence.")
  end
  return StandardYoungTableau(_ballot2syt(b), false)
end

function _ballot(N, lambda, a, more)
  #N = sum(lambda)
  it = N
  @inbounds if more
    lambda = zeros(Int64, N)
    lambda[1] = 1
    isave = 0
    for i = 2:N
      lambda[a[i]] += 1
      if a[i] < a[i-1]
        isave = i
        break
      end
    end
    if isave == 0
      return (a = a, more = false)
    end
    it = lambda[a[isave]+1]
    for j = 0:(N-1)
      i = N - j
      if lambda[i] == it
        a[isave] = i
        lambda[i] -= 1
        it = isave - 1
        break
      end
    end
  end # end of 'if more'
  lambda0 = zeros(Int64, N)
  @inbounds lambda0[1:length(lambda)] = lambda
  lambda = lambda0
  k = 1
  ir = 1
  @inbounds while true
    if N < ir
      break
    end
    if lambda[ir] != 0
      a[k] = ir
      lambda[ir] -= 1
      k += 1
      ir += 1
      continue
    end
    if it < k
      break
    end
    ir = 1
  end
  if N == 1
    return (a = a, more = false)
  end
  @inbounds for j = 1:(N-1)
    if a[j+1] < a[j]
      return (a = a, more = true)
    end
  end
  return (a = a, more = false)
end

"""
    allBallotSequences(lambda)

All ballot sequences associated to a given integer partition.

# Argument
- `lambda`: an integer partition
"""
function allBallotSequences(lambda::IPartition)
  N = lambda.n
  partition = lambda.partition
  (a, more) = _ballot(N, partition, zeros(Int64, N), false)
  out = [copy(a)]
  while more
    (a, more) = _ballot(N, partition, a, true)
    push!(out, copy(a))
  end
  return out
end

"""
    allSYTx(lambda)

All standard Young tableaux of a given shape.

# Argument
- `lambda`: an integer partition

# Example

    lambda = IPartition([3, 1])
    allSYTx(lambda)
"""
function allSYTx(lambda::IPartition)
  N = lambda.n
  partition = lambda.partition
  (a, more) = _ballot(N, partition, zeros(Int64, N), false)
  out = [StandardYoungTableau(_ballot2syt(a), false)]
  while more
    (a, more) = _ballot(N, partition, a, true)
    push!(out, StandardYoungTableau(_ballot2syt(a), false))
  end
  return out
end

function _islastsyt(tableau)
  a = _syt2ballot(tableau)
  N = length(a)
  for j = 1:(N-1)
    if a[j+1] < a[j]
      return false
    end
  end
  return true
end

function _firstsyt(lambda, a)
  lambda = filter(!=(0), lambda)
  N = sum(lambda)
  if N == 1
    return [[1]]
  end
  it = N
  lambda0 = zeros(Int64, N)
  lambda0[1:length(lambda)] = lambda
  lambda = lambda0
  k = 1
  ir = 1
  while true
    if N < ir
      break
    end
    if lambda[ir] != 0
      a[k] = ir
      lambda[ir] -= 1
      k += 1
      ir += 1
      continue
    end
    if it < k
      break
    end
    ir = 1
  end
  return _ballot2syt(a)
end

function _nextsyt(tableau)
  if _islastsyt(tableau)
    return nothing
  end
  a = _syt2ballot(tableau)
  N = sum(map(length, tableau))
  lambda = zeros(Int64, N)
  lambda[1] = 1
  isave = 0
  for i = 2:N
    lambda[a[i]] += 1
    if a[i] < a[i-1]
      isave = i
      break
    end
  end
  if isave == 0
    return _ballot2syt(a)
  end
  it = lambda[a[isave]+1]
  for j = 0:(N-1)
    i = N - j
    if lambda[i] == it
      a[isave] = i
      lambda[i] -= 1
      it = isave - 1
      break
    end
  end
  return _firstsyt(lambda, a)
end

"""
    firstSYT(lambda)

The 'first' standard Young tableau of a given shape.

# Argument
- `lambda`: an integer partition
"""
function firstSYT(lambda::IPartition)
  return StandardYoungTableau(
    _firstsyt(lambda.partition, zeros(Int64, lambda.n)),
    false
  )
end

"""
    nextSYT(syt)

The 'next' standard Young tableau.

# Argument
- `syt`: a standard Young tableau

# Example

    y = firstSYT(IPartition([4, 2]))
    nextSYT(y)
"""
function nextSYT(syt::StandardYoungTableau)
  succ = _nextsyt(syt.tableau)
  if isnothing(succ)
    return nothing
  end
  return StandardYoungTableau(succ, false)
end

function _dualPartition(lambda)
  return map(i -> length(findall(>=(i), lambda)), 1:lambda[1])
end

"""
    hooks(lambda)

The hooks of a partition.

# Argument
- `lambda`: an integer partition
"""
function hooks(lambda::IPartition)
  partition = lambda.partition
  dlambda = _dualPartition(partition)
  l = length(partition)
  out = Vector{Vector{Vector{Int64}}}(undef, l)
  for i = 1:l
    out[i] = Vector{Vector{Int64}}(undef, partition[i])
    for j = 1:partition[i]
      out[i][j] = [0; partition[i] - j + 1]
    end
  end
  for j = 1:length(dlambda)
    for i = 1:dlambda[j]
      out[i][j][1] = dlambda[j] - i + 1
    end
  end
  return out
end

"""
    hooklengths(lambda)

The hook lengths of a partition.

# Argument
- `lambda`: an integer partition
"""
function hooklengths(lambda::IPartition)
  h = hooks(lambda)
  map(x -> vcat((map(y -> sum(y) - 1, x))...), h)
end

"""
    countSYTx(lambda)

Number of standard Young tableaux of a given shape.

# Argument
- `lambda`: an integer partition, the shape

# Example

    lambda = IPartition([4, 2])
    countSYTx(lambda)
    length(allSYTx(lambda))
"""
function countSYTx(lambda::IPartition)
  numterms = Int64[]
  denterms = filter(>=(2), vcat(hooklengths(lambda)...))
  for i = 2:lambda.n
    idx = findfirst(==(i), denterms)
    if isnothing(idx)
      push!(numterms, i)
    else
      deleteat!(denterms, idx)
    end
  end
  return numerator(prod(numterms) // prod(denterms))
end

function _bump(P, Q, e, i)
  if isempty(P)
    return (P = [[e]], Q = [[i]])
  end
  p = P[1]
  if e > last(p)
    P[1] = vcat(p, e)
    Q[1] = vcat(Q[1], i)
    return (P = P, Q = Q)
  else
    j = findfirst(!, p .< e)
    w = p[j]
    P[1][j] = e
    b = _bump(P[2:end], Q[2:end], w, i)
    return (P = vcat([P[1]], b.P), Q = vcat([Q[1]], b.Q))
  end
end

"""
    RS(sigma)

Pair of standard Young tableaux given by the Robinson-Schensted(-Knuth) correspondence from a permutation.

# Argument
- `sigma`: a permutation given as a vector of integers

# Example

    RS([3, 4, 1, 2])
"""
function RS(sigma)
  if !issetequal(sigma, 1:length(sigma))
    error("`sigma` is not a permutation")
  end
  out = _bump([], [], sigma[1], 1)
  for i = 2:length(sigma)
    out = _bump(out.P, out.Q, sigma[i], i)
  end
  return (
    P = StandardYoungTableau(out.P, false),
    Q = StandardYoungTableau(out.Q, false)
  )
end

function _whichRow(tableau, x)
  idx = findfirst(==(x), tableau[1])
  i = 1
  while isnothing(idx)
    i += 1
    idx = findfirst(==(x), tableau[i])
  end
  return i
end

"""
    SYT2YoungPath(syt)

Path of the Young graph corresponding to a standard Young tableau.

# Argument
- `syt`: a standard Young tableau

# Example

    y = StandardYoungTableau([[1,3,4], [2]])
    SYT2YoungPath(y)
"""
function SYT2YoungPath(syt::StandardYoungTableau)
  N = syt.shape.n
  path = Vector{Vector{Int64}}(undef, N)
  path[1] = vcat(1, zeros(Int64, length(syt.tableau) - 1))
  for k = 2:N
    part = copy(path[k-1])
    i = _whichRow(syt.tableau, k)
    part[i] += 1
    path[k] = part
  end
  return map(x -> IPartition(x, false), path)
end

"""
    YoungPath2SYT(path)

Standard Young tableau corresponding to a path of the Young graph.

# Argument
- `path`: a vector of integers partitions forming a path of the Young graph

# Example

    y = StandardYoungTableau([[1,3,4], [2]])
    path = SYT2YoungPath(y)
    YoungPath2SYT(path)
"""
function YoungPath2SYT(path::Vector{IPartition})
  Ns = map(x -> x.n, path)
  N = length(path)
  if any(Ns .!= 1:N)
    error(
      "Invalid path. The n-th element of the path must be a partition of n."
    )
  end
  if N == 1
    return StandardYoungTableau([[1]], false)
  end
  incrmts = diff(map(x -> length(x.partition), path))
  if any((incrmts .!= 0) .* (incrmts .!= 1))
    error("Invalid path.")
  end
  l = length(path[end].partition)
  path2 = zeros(Int64, l, N)
  for i = 1:N
    part_i = path[i].partition
    path2[1:length(part_i), i] = part_i
  end
  diffs = diff(path2, dims = 2)
  check = mapslices(
    x -> sum(x) == 1 && all((x .== 0) .| (x .== 1)), diffs, dims = 1
  )
  if !all(check)
    error("invalid path")
  end
  tableau = vcat([[1]], repeat([Int64[]], l - 1))
  for i = 1:(N-1)
    idx = findfirst(==(1), diffs[:, i])
    tableau[idx] = vcat(tableau[idx], i + 1)
  end
  return StandardYoungTableau(tableau, false)
end

"""
    randomSYT(lambda)

Uniformly random standard Young tableau of a given shape.

# Argument
- `lambda`: an integer partition, the shape
"""
function randomSYT(lambda::IPartition)
  partition = copy(lambda.partition)
  N = lambda.n
  a = zeros(Int64, N)
  i = 0
  k = 0
  while true
    i += 1
    for j in 1:partition[i]
      a[j] += 1
      k += 1
    end
    if N <= k
      break
    end
  end
  for m in 1:N
    local i, j
    while true
      i = rand(1:a[1])
      j = rand(1:partition[1])
      if i <= a[j] && j <= partition[i]
        break
      end
    end
    while true
      ih = a[j] + partition[i] - i - j
      if ih == 0
        break
      end
      k = rand(1:ih)
      if k <= (partition[i] - j)
        j += k
      else
        i += k - partition[i] + j
      end
    end
    partition[i] -= 1
    a[j] -= 1
    a[N - m + 1] = i
  end
  return StandardYoungTableau(_ballot2syt(a), false)
end

function _connectedPartitions(partition0)
  partition = copy(partition0)
  partition[1] += 1
  out = [partition]
  for i in 2:length(partition0)
    if partition0[i] < partition0[i - 1]
      partition = copy(partition0)
      partition[i] += 1
      push!(out, partition)
    end
  end
  return vcat(out, [vcat(partition0, 1)])
end

"""
    randomYoungPath(n)

Samples a path of the Young graph according to the Plancherel growth process.

# Argument
`n`: a positive integer, the size of the path to be sampled

# Example

    path = randomYoungPath(5)
    YoungPath2SYT(path)
"""
function randomYoungPath(n::Int64)
  path = [[1]]
  for i in 1:(n-1)
    set = _connectedPartitions(path[i])
    f_nu = countSYTx(IPartition(path[i], false))
    probs = map(p -> countSYTx(IPartition(p, false)), set) ./ ((i + 1) * f_nu)
    k = rand(Distributions.Categorical(probs))
    push!(path, set[k])
  end
  return map(x -> IPartition(x, false), path)
end

end