using GZip

"""
    make_ref(src, ref)
---
With plink file `src.{bed,bim,fam}`, this function:
1. imputes the few missing genotypes in this dataset
2. phasing is also done by `beagle.jar`
3. write results to `ref.{bed,bim,fam}`
"""
function make_ref(src, ref; window=300)
    title("Make imputation reference")
    item("Check parameters")
    if !isfile(src * ".bed")
        warning("$src.bed doesn't exist")
        return
    end
    pathto, target = splitdir(ref)
    if length(pathto) == 0
        warning("Creating reference dataset in the current directory")
    else
        isdir(pathto) || mkdir(pathto)
    end
    done()

    item("Create reference")
    tmp = mktempdir(".")
    run(`$plink --sheep --bfile $src --recode vcf-iid bgz --out $tmp/mid`)
    run(pipeline(`java -jar $beagle
                    window=$window
                    gt=$tmp/mid.vcf.gz
                    ne=$nsNe
                    out=$ref`,
                 `grep Window`))
    rm(tmp, recursive=true, force=true)
    done()

    item("Convert to plink format")
    run(`$plink --sheep --make-bed --vcf $ref.vcf.gz --out $ref`)
    done()
end


"""
    make_ref(ref, fmp, plk...)
---
Make a reference `ref`, using linkage map file `fmp`.
This function will
- exclude the extra SNP in `plk...`, but not in `lmp`,
- merge the result `plk...`
- impute the *few* (usually are) missing SNP and create a reference file.
"""
function make_ref(ref, fmp, plk...)
    title("Merging several platforms into one reference")
    lmp = read_map(fmp)
    sst = Set(lmp.snp)          # SNP set
    dir, _ = splitdir(ref)      # write all the temp files and results in the ref dir.
    isdir(dir) || mkpath(dir)
    open("$dir/list", "w") do io
        for p in plk
            println(io, "$p-s")
        end
    end

    # check if a SNP needs flip.  I just remove them (308 SNP)
    alsnp = Dict{String, Set{Char}}()
    for p in plk
        bim = read_bim("$p.bim")
        for (_, snp, _, _, a1, a2) in eachrow(bim)
            if haskey(alsnp, snp)
                push!(alsnp[snp], a1)
            else
                alsnp[snp] = Set(a1)
            end
            push!(alsnp[snp], a2)
        end
    end
    excl = Set{String}()
    for (snp, val) in alsnp
        if length(val) > 2
            push!(excl, snp)
        end
    end
    
    for p in plk
        open("$dir/keep.snp", "w") do io
            for line in eachline("$p.bim")
                snp = split(line)[2]
                if snp ∈ sst && snp ∉ excl
                    println(io, snp)
                end
            end
        end
        run(`$plink --sheep --bfile $p 
		--extract $dir/keep.snp
		--make-bed
		--out $p-s`)
    end
    run(`$plink --sheep --merge-list $dir/list --out $dir/1`)
    run(`$plink --sheep --bfile $dir/1 --list-duplicate-vars suppress-first --out $dir/1`)
    run(`$plink --sheep --bfile $dir/1 --exclude $dir/1.dupvar --make-bed --out $ref`)
    #make_ref("$dir/2", ref)
end

"""
    function impute(cur, ref, out; nthread=7)
---
Impute `cur`.{bed,bim,fam} with reference/training set `ref`.{bed,bim,fam}, and put
the result to `out`.{bed,bim,fam}.

**Note**:

- SNP only in `cur` will be removed.
- beagle was assigned 7 threads by default to avoid some weird errors.
"""
function impute(cur, ref, out; window=300)
    title("Impute $cur.{bed,bim,fam} with training set $ref")
    item("Check parameters")
    if !isfile(cur * ".bed")
        warning("$cur.bed doesn't exist")
        return
    end
    if !isfile("$ref.vcf.gz")
        warning("Reference $ref.vcf.gz is not ready")
        return
    end
    pathto, target = splitdir(out)
    if(length(pathto) == 0)
        warning("Creating target in the current directory")
    else
        isdir(pathto) || mkdir(pathto)
    end
    done()

    item("Prepare $cur and remove extra SNP in $cur")
    tmp = mktempdir(".")
    refsnp = begin
        df = read_bim("$ref.bim")
        Set(df.snp)
    end
    cursnp = begin
        df = read_bim("$cur.bim")
        Set(df.snp)
    end
    exclude = setdiff(cursnp, refsnp)
    write("$tmp/exclude.snp", join(exclude, '\n'), '\n')
    run(`$plink
            --sheep
            --bfile $cur
            --exclude $tmp/exclude.snp
            --recode vcf-iid bgz
            --out $tmp/mid`)
    done()
    
    item("Impute $cur with $ref")
    run(pipeline(`java -jar $beagle
                    window=$window
                    ne=$nsNe
                    ref=$ref.vcf.gz
                    gt=$tmp/mid.vcf.gz
                    out=$out`,
                 `grep Window`))
    done()
    
    item("Convert result to plink format")
    run(`$plink --sheep --make-bed --vcf $out.vcf.gz --out $out`)
    done()
    
    rm(tmp, recursive=true, force=true)
end


"""
    function merge_vcfs(tgt, refs...)
---
Merge several VCF files of the same platform,
or imputed to the same platform,
to create one reference.

This is a very dirty function.  The refs must be of the same platform.
"""
function merge_vcfs(tgt, refs...)
    out = gzopen("$tgt.vcf.gz", "w")
    gz = Any[]
    header = String[]
    
    for r in refs
        z = gzopen(r)
        push!(gz, z)
    end

    for line in eachline(gz[1])
        if line[2] == '#'
            println(out, line)
        else
            print(out, line)
            break
        end
    end

    for i in 2:length(gz)
        for line in eachline(gz[i])
            if line[2] != '#'
                print(out, '\t', join(split(line)[10:end], '\t'))
                break
            end
        end
    end
    print(out, '\n')

    for line in eachline(gz[1])
        print(out, line)
        for i in 2:length(gz)
            l = readline(gz[i])
            print(out, '\t', join(split(l)[10:end], '\t'))
        end
        print(out, '\n')
    end
    
    for z in gz
        close(z)
    end
    close(out)

    run(`$plink --sheep --make-bed --vcf $tgt.vcf.gz --out $tgt`)
end
