# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.3.1
#     language: julia
#     name: julia-1.3
# ---

# # Test and Benchmark

using Revise
using VCFTools
using LinearAlgebra
using MendelImpute
using GeneticVariation
using Random
using StatsBase
using CodecZlib
using ProgressMeter
using JLD2, FileIO, JLSO
using BenchmarkTools
using GroupSlices

# ## Haplotyping timing varying number of haplotypes

Random.seed!(123)
ns = [100, 1000, 10000] # number of individuals
ps = [100, 1000, 10000]    # number of typed SNPs
ds = [100, 1000, 10000]  # number of reference haplotypes
for n in ns, p in ps, d in ds
    H = convert(Matrix{Float32}, rand(0:1, p, d))
    X = convert(Matrix{Float32}, rand(0:2, p, n))
    M = Transpose(H) * H
    for j in 1:d, i in 1:(j - 1) # off-diagonal
        M[i, j] = 2M[i, j] + M[i, i] + M[j, j]
    end
    for j in 1:d # diagonal
        M[j, j] *= 4
    end
    N = Transpose(X) * H
    for I in eachindex(N)
        N[I] *= 2
    end
    happair  = zeros(Int, n), zeros(Int, n)
    hapscore = zeros(eltype(N), n)
    t = @elapsed haplopair!(happair, hapscore, M, N)
    println("n = $n, p = $p, d = $d time = $t")
end

# # Try using Distance package

# +
Random.seed!(123)

n = 1000
p = 100
d = 1000

H = bitrand(p, d)
X = Matrix{Union{UInt8, Missing}}(rand(UInt8, p, n));
# -

@time happairs, hapscore = haplopair(X, H)

@time happairs, hapscore = haplopair_thin(X, H, keep = 1000) #blas 2

@time happairs, hapscore = haplopair_thin2(X, H, keep = 1000) #blas 3

# +
linsp = 100:1000:8000
t1 = zeros(length(linsp))
t2 = zeros(length(linsp))
t3 = zeros(length(linsp))
t4 = zeros(length(linsp))
n = 1000
p = 512
keep = 100
r = 100

Random.seed!(2020)
for (i, d) in enumerate(linsp)
    H = bitrand(p, d)
    X = Matrix{Union{UInt8, Missing}}(rand(UInt8, p, n));
    t1[i] = @elapsed haplopair(X, H)
    t2[i] = @elapsed haplopair_thin(X, H, keep = keep)
    t3[i] = @elapsed haplopair_thin2(X, H, keep = keep)
    t4[i] = @elapsed haplopair_lasso(X, H, r = r)
end
# -

using Plots
x = plot(linsp, t1, title="1000 samples, 512 typed SNPs per window",
    label="Original search time", ylabel="Seconds", xlabel="Number of haplotypes")
plot!(x, linsp, t2, label="Haplotype thinning r = $keep (BLAS 2)")
plot!(x, linsp, t3, label="Haplotype thinning r = $keep (BLAS 3)")
plot!(x, linsp, t4, label="Lasso r = $r")

# +
# ZOOM in

linsp = 100:100:2500
t1 = zeros(length(linsp))
t2 = zeros(length(linsp))
t3 = zeros(length(linsp))
t4 = zeros(length(linsp))
n = 1000
p = 512
keep = 100
r = 100

Random.seed!(2020)
for (i, d) in enumerate(linsp)
    H = bitrand(p, d)
    X = Matrix{Union{UInt8, Missing}}(rand(UInt8, p, n));
    t1[i] = @elapsed haplopair(X, H)
    t2[i] = @elapsed haplopair_thin(X, H, keep = keep)
    t3[i] = @elapsed haplopair_thin2(X, H, keep = keep)
    t4[i] = @elapsed haplopair_lasso(X, H, r = r)
end
y = plot(linsp, t1, title = "Zoomed in, n = $n, p = $p", label="Original search time", 
    ylabel="Seconds", xlabel="Number of haplotypes")
plot!(y, linsp, t2, label="Haplotype thinning r = $keep (BLAS 2)")
plot!(y, linsp, t3, label="Haplotype thinning r = $keep (BLAS 3)")
plot!(y, linsp, t4, label="Lasso r = $r")
# -

savefig(y, "zoom_in.png")
savefig(x, "zoom_out.png")

# # Benchmark original code

# +
Random.seed!(123)

n = 100
p = 100
d = 1000

M = rand(Float32, d, d)
N = rand(Float32, n, d)
H = convert(Matrix{Float32}, rand(0:1, p, d))
X = convert(Matrix{Float32}, rand(0:2, p, n))
happair1 = zeros(Int, n)
happair2 = zeros(Int, n)
hapscore = zeros(Float32, n)

@btime haplopair!(happair1, happair2, hapscore, M, N) setup=(M=rand(Float32, d, d), N=rand(Float32, n, d))

# +
Random.seed!(123)

n = 100
p = 100
d = 1000

M = rand(Float32, d, d)
N = rand(Float32, n, d)
H = convert(Matrix{Float32}, rand(0:1, p, d))
X = convert(Matrix{Float32}, rand(0:2, p, n))
happairs = zeros(Int, n), zeros(Int, n)
hapscore = zeros(Float32, n)

@btime haplopair!(happairs, hapscore, M, N) setup=(M=rand(Float32, d, d), N=rand(Float32, n, d))
# -







StarcoderdataJupyterScriptsDedupFiltered