#!/usr/bin/env luajit

-----------------------------------
-- BEGIN: routines from klib.lua --
-----------------------------------

-- Description: getopt() translated from the BSD getopt(); compatible with the default Unix getopt()
--[[ Example:
	for o, a in os.getopt(arg, 'a:b') do
		print(o, a)
	end
]]--
function os.getopt(args, ostr)
	local arg, place = nil, 0;
	return function ()
		if place == 0 then -- update scanning pointer
			place = 1
			if #args == 0 or args[1]:sub(1, 1) ~= '-' then place = 0; return nil end
			if #args[1] >= 2 then
				place = place + 1
				if args[1]:sub(2, 2) == '-' then -- found "--"
					table.remove(args, 1);
					place = 0
					return nil;
				end
			end
		end
		local optopt = place <= #args[1] and args[1]:sub(place, place) or nil
		place = place + 1;
		local oli = optopt and ostr:find(optopt) or nil
		if optopt == ':' or oli == nil then -- unknown option
			if optopt == '-' then return nil end
			if place > #args[1] then
				table.remove(args, 1);
				place = 0;
			end
			return '?';
		end
		oli = oli + 1;
		if ostr:sub(oli, oli) ~= ':' then -- do not need argument
			arg = nil;
			if place > #args[1] then
				table.remove(args, 1);
				place = 0;
			end
		else -- need an argument
			if place <= #args[1] then  -- no white space
				arg = args[1]:sub(place);
			else
				table.remove(args, 1);
				if #args == 0 then -- an option requiring argument is the last one
					place = 0;
					if ostr:sub(1, 1) == ':' then return ':' end
					return '?';
				else arg = args[1] end
			end
			table.remove(args, 1);
			place = 0;
		end
		return optopt, arg;
	end
end

-- Description: string split
function string:split(sep, n)
	local a, start = {}, 1;
	sep = sep or "%s+";
	repeat
		local b, e = self:find(sep, start);
		if b == nil then
			table.insert(a, self:sub(start));
			break
		end
		a[#a+1] = self:sub(start, b - 1);
		start = e + 1;
		if n and #a == n then
			table.insert(a, self:sub(start));
			break
		end
	until start > #self;
	return a;
end

-- Description: smart file open
function io.xopen(fn, mode)
	mode = mode or 'r';
	if fn == nil then return io.stdin;
	elseif fn == '-' then return (mode == 'r' and io.stdin) or io.stdout;
	elseif fn:sub(-3) == '.gz' then return (mode == 'r' and io.popen('gzip -dc ' .. fn, 'r')) or io.popen('gzip > ' .. fn, 'w');
	elseif fn:sub(-4) == '.bz2' then return (mode == 'r' and io.popen('bzip2 -dc ' .. fn, 'r')) or io.popen('bgzip2 > ' .. fn, 'w');
	else return io.open(fn, mode) end
end

-- Description: log gamma function
-- Required by: math.lbinom()
-- Reference: AS245, 2nd algorithm, http://lib.stat.cmu.edu/apstat/245
function math.lgamma(z)
	local x;
	x = 0.1659470187408462e-06     / (z+7);
	x = x + 0.9934937113930748e-05 / (z+6);
	x = x - 0.1385710331296526     / (z+5);
	x = x + 12.50734324009056      / (z+4);
	x = x - 176.6150291498386      / (z+3);
	x = x + 771.3234287757674      / (z+2);
	x = x - 1259.139216722289      / (z+1);
	x = x + 676.5203681218835      / z;
	x = x + 0.9999999999995183;
	return math.log(x) - 5.58106146679532777 - z + (z-0.5) * math.log(z+6.5);
end

-- Description: regularized incomplete gamma function
-- Dependent on: math.lgamma()
--[[
  Formulas are taken from Wiki, with additional input from Numerical
  Recipes in C (for modified Lentz's algorithm) and AS245
  (http://lib.stat.cmu.edu/apstat/245).
 
  A good online calculator is available at:
 
    http://www.danielsoper.com/statcalc/calc23.aspx
 
  It calculates upper incomplete gamma function, which equals
  math.igamma(s,z,true)*math.exp(math.lgamma(s))
]]--
function math.igamma(s, z, complement)

	local function _kf_gammap(s, z)
		local sum, x = 1, 1;
		for k = 1, 100 do
			x = x * z / (s + k);
			sum = sum + x;
			if x / sum < 1e-14 then break end
		end
		return math.exp(s * math.log(z) - z - math.lgamma(s + 1.) + math.log(sum));
	end

	local function _kf_gammaq(s, z)
		local C, D, f, TINY;
		f = 1. + z - s; C = f; D = 0.; TINY = 1e-290;
		-- Modified Lentz's algorithm for computing continued fraction. See Numerical Recipes in C, 2nd edition, section 5.2
		for j = 1, 100 do
			local d;
			local a, b = j * (s - j), j*2 + 1 + z - s;
			D = b + a * D;
			if D < TINY then D = TINY end
			C = b + a / C;
			if C < TINY then C = TINY end
			D = 1. / D;
			d = C * D;
			f = f * d;
			if math.abs(d - 1) < 1e-14 then break end
		end
		return math.exp(s * math.log(z) - z - math.lgamma(s) - math.log(f));
	end

	if complement then
		return ((z <= 1 or z < s) and 1 - _kf_gammap(s, z)) or _kf_gammaq(s, z);
	else 
		return ((z <= 1 or z < s) and _kf_gammap(s, z)) or (1 - _kf_gammaq(s, z));
	end
end

function math.brent(func, a, b, tol)
	local gold1, gold2, tiny, max_iter = 1.6180339887, 0.3819660113, 1e-20, 100

	local fa, fb = func(a, data), func(b, data)
	if fb > fa then -- swap, such that f(a) > f(b)
		a, b, fa, fb = b, a, fb, fa
	end
	local c = b + gold1 * (b - a)
	local fc = func(c) -- golden section extrapolation
	while fb > fc do
		local bound = b + 100.0 * (c - b) -- the farthest point where we want to go
		local r = (b - a) * (fb - fc)
		local q = (b - c) * (fb - fa)
		if math.abs(q - r) < tiny then -- avoid 0 denominator
			tmp = q > r and tiny or 0.0 - tiny
		else tmp = q - r end
		u = b - ((b - c) * q - (b - a) * r) / (2.0 * tmp) -- u is the parabolic extrapolation point
		if (b > u and u > c) or (b < u and u < c) then -- u lies between b and c
			fu = func(u)
			if fu < fc then -- (b,u,c) bracket the minimum
				a, b, fa, fb = b, u, fb, fu
				break
			elseif fu > fb then -- (a,b,u) bracket the minimum
				c, fc = u, fu
				break
			end
			u = c + gold1 * (c - b)
			fu = func(u) -- golden section extrapolation
		elseif (c > u and u > bound) or (c < u and u < bound) then -- u lies between c and bound
			fu = func(u)
			if fu < fc then -- fb > fc > fu
				b, c, u = c, u, c + gold1 * (c - b)
				fb, fc, fu = fc, fu, func(u)
			else -- (b,c,u) bracket the minimum
				a, b, c = b, c, u
				fa, fb, fc = fb, fc, fu
				break
			end
		elseif (u > bound and bound > c) or (u < bound and bound < c) then -- u goes beyond the bound
			u = bound
			fu = func(u)
		else -- u goes the other way around, use golden section extrapolation
			u = c + gold1 * (c - b)
			fu = func(u)
		end
		a, b, c = b, c, u
		fa, fb, fc = fb, fc, fu
	end
	if a > c then a, c = c, a end -- swap

	-- now, a<b<c, fa>fb and fb<fc, move on to Brent's algorithm
	local e, d = 0, 0
	local w, v, fw, fv
	w, v = b, b
	fw, fv = fb, fb
	for iter = 1, max_iter do
		local mid = 0.5 * (a + c)
		local tol1 = tol * math.abs(b) + tiny
		local tol2 = 2.0 * tol1
		if math.abs(b - mid) <= tol2 - 0.5 * (c - a) then return fb, b end -- found
		if math.abs(e) > tol1 then
			-- related to parabolic interpolation
			local r = (b - w) * (fb - fv)
			local q = (b - v) * (fb - fw)
			local p = (b - v) * q - (b - w) * r
			q = 2.0 * (q - r)
			if q > 0.0 then p = 0.0 - p
			else q = 0.0 - q end
			eold, e = e, d
			if math.abs(p) >= math.abs(0.5 * q * eold) or p <= q * (a - b) or p >= q * (c - b) then
				e = b >= mid and a - b or c - b
				d = gold2 * e
			else
				d, u = p / q, b + d -- actual parabolic interpolation happens here
				if u - a < tol2 or c - u < tol2 then
					d = mid > b and tol1 or 0.0 - tol1
				end
			end
		else -- golden section interpolation
			e = b >= min and a - b or c - b
			d = gold2 * e
		end
		u = fabs(d) >= tol1 and b + d or b + (d > 0.0 and tol1 or -tol1);
		fu = func(u)
		if fu <= fb then -- u is the minimum point so far
			if u >= b then a = b
			else c = b end
			v, w, b = w, b, u
			fv, fw, fb = fw, fb, fu
		else -- adjust (a,c) and (u,v,w)
			if u < b then a = u
			else c = u end
			if fu <= fw or w == b then
				v, w = w, u
				fv, fw = fw, fu
			elseif fu <= fv or v == b or v == w then
				v, fv = u, fu;
			end
		end
	end
	return fb, b
end

matrix = {}

-- Description: chi^2 test for contingency tables
-- Dependent on: math.igamma()
function matrix.chi2(a)
	if #a == 2 and #a[1] == 2 then -- 2x2 table
		local x, z
		x = (a[1][1] + a[1][2]) * (a[2][1] + a[2][2]) * (a[1][1] + a[2][1]) * (a[1][2] + a[2][2])
		if x == 0 then return 0, 1, false end
		z = a[1][1] * a[2][2] - a[1][2] * a[2][1]
		z = (a[1][1] + a[1][2] + a[2][1] + a[2][2]) * z * z / x
		return z, math.igamma(.5, .5 * z, true), true
	else -- generic table
		local rs, cs, n, m, N, z = {}, {}, #a, #a[1], 0, 0
		for i = 1, n do rs[i] = 0 end
		for j = 1, m do cs[j] = 0 end
		for i = 1, n do -- compute column sum and row sum
			for j = 1, m do cs[j], rs[i] = cs[j] + a[i][j], rs[i] + a[i][j] end
		end
		for i = 1, n do N = N + rs[i] end
		for i = 1, n do -- compute the chi^2 statistics
			for j = 1, m do
				local E = rs[i] * cs[j] / N;
				z = z + (a[i][j] - E) * (a[i][j] - E) / E
			end
		end
		return z, math.igamma(.5 * (n-1) * (m-1), .5 * z, true), true;
	end
end

---------------------------------
-- END: routines from klib.lua --
---------------------------------


--------------------------
-- BEGIN: misc routines --
--------------------------

-- precompute an array for PL->probability conversion
-- @param m maximum PL
function algo_init_q2p(m)
	local q2p = {}
	for i = 0, m do
		q2p[i] = math.pow(10, -i / 10)
	end
	return q2p
end

-- given the haplotype frequency, compute r^2
-- @param f 4 haplotype frequencies; f[] is 0-indexed.
-- @return r^2
function algo_r2(f)
	local p = { f[0] + f[1], f[0] + f[2] }
	local D = f[0] * f[3] - f[1] * f[2]
	return (p[1] == 0 or p[2] == 0 or 1-p[1] == 0 or 1-p[2] == 0) and 0 or D * D  / (p[1] * p[2] * (1 - p[1]) * (1 - p[2]))
end

-- parse a VCF line to get PL
-- @param q2p is computed by algo_init_q2p()
function text_parse_pl(t, q2p, parse_GT)
	parse_GT = parse_GT == nil and true or false
	local ht, gt, pl = {}, {}, {}
	local s, j0 = t[9]:split(':'), 0
	for j = 1, #s do
		if s[j] == 'PL' then j0 = j break end
	end
	local has_GT = (s[1] == 'GT' and parse_GT) and true or false
	for i = 10, #t do
		if j0 > 0 then
			local s = t[i]:split(':')
			local a, b = 1, s[j0]:find(',')
			pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, b - 1))]
			a, b = b + 1, s[j0]:find(',', b + 1)
			pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, b - 1))]
			a, b = b + 1, s[j0]:find(',', b + 1)
			pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, (b and b - 1) or nil))]
		end
		if has_GT then
			if t[i]:sub(1, 1) ~= '.' then
				local g = tonumber(t[i]:sub(1, 1)) + tonumber(t[i]:sub(3, 3));
				gt[#gt+1] = 1e-6; gt[#gt+1] = 1e-6; gt[#gt+1] = 1e-6
				gt[#gt - 2 + g] = 1
				ht[#ht+1] = tonumber(t[i]:sub(1, 1)); ht[#ht+1] = tonumber(t[i]:sub(3, 3));
			else
				gt[#gt+1] = 1; gt[#gt+1] = 1; gt[#gt+1] = 1
				ht[#ht+1] = -1; ht[#ht+1] = -1;
			end
		end
--		print(t[i], pl[#pl-2], pl[#pl-1], pl[#pl], gt[#gt-2], gt[#gt-1], gt[#gt])
	end
	if #pl == 0 then pl = nil end
	local x = has_GT and { t[1], t[2], ht, gt, pl } or { t[1], t[2], nil, nil, pl }
	return x
end

-- Infer haplotype frequency
-- @param pdg  genotype likelihoods P(D|g) generated by text_parse_pl(). pdg[] is 1-indexed.
-- @param eps  precision [1e-5]
-- @return 2-locus haplotype frequencies, 0-indexed array
function algo_hapfreq2(pdg, eps)
	eps = eps or 1e-5
	local n, f = #pdg[1] / 3, {[0]=0.25, 0.25, 0.25, 0.25}
	for iter = 1, 100 do
		local F = {[0]=0, 0, 0, 0}
		for i = 0, n - 1 do
			local p1, p2 = {[0]=pdg[1][i*3+1], pdg[1][i*3+2], pdg[1][i*3+3]}, {[0]=pdg[2][i*3+1], pdg[2][i*3+2], pdg[2][i*3+3]}
			local u = { [0]=
				f[0] * (f[0] * p1[0] * p2[0] + f[1] * p1[0] * p2[1] + f[2] * p1[1] * p2[0] + f[3] * p1[1] * p2[1]),
				f[1] * (f[0] * p1[0] * p2[1] + f[1] * p1[0] * p2[2] + f[2] * p1[1] * p2[1] + f[3] * p1[1] * p2[2]),
				f[2] * (f[0] * p1[1] * p2[0] + f[1] * p1[1] * p2[1] + f[2] * p1[2] * p2[0] + f[3] * p1[2] * p2[1]),
				f[3] * (f[0] * p1[1] * p2[1] + f[1] * p1[1] * p2[2] + f[2] * p1[2] * p2[1] + f[3] * p1[2] * p2[2])
			}
			local s = u[0] + u[1] + u[2] + u[3]
			s = 1 / (s * n)
			F[0] = F[0] + u[0] * s
			F[1] = F[1] + u[1] * s
			F[2] = F[2] + u[2] * s
			F[3] = F[3] + u[3] * s
		end
		local e = 0
		for k = 0, 3 do
			e = math.abs(f[k] - F[k]) > e and math.abs(f[k] - F[k]) or e
		end
		for k = 0, 3 do f[k] = F[k] end
		if e < eps then break end
--		print(f[0], f[1], f[2], f[3])
	end
	return f
end

------------------------
-- END: misc routines --
------------------------


---------------------
-- BEGIN: commands --
---------------------

-- CMD vcf2bgl: convert PL tagged VCF to Beagle input --
function cmd_vcf2bgl()
	if #arg == 0 then
		print("\nUsage: vcf2bgl.lua <in.vcf>")
		print("\nNB: This command finds PL by matching /(\\d+),(\\d+),(\\d+)/.\n");
		os.exit(1)
	end
	
	local lookup = {}
	for i = 0, 10000 do lookup[i] = string.format("%.4f", math.pow(10, -i/10)) end
	
	local fp = io.xopen(arg[1])
	for l in fp:lines() do
		if l:sub(1, 2) == '##' then -- meta lines; do nothing
		elseif l:sub(1, 1) == '#' then -- sample lines
			local t, s = l:split('\t'), {}
			for i = 10, #t do s[#s+1] = t[i]; s[#s+1] = t[i]; s[#s+1] = t[i] end
			print('marker', 'alleleA', 'alleleB', table.concat(s, '\t'))
		else -- data line
			local t = l:split('\t');
			if t[5] ~= '.' and t[5]:find(",") == nil and #t[5] == 1 and #t[4] == 1 then -- biallic SNP
				local x, z = -1, {};
				if t[9]:find('PL') then
					for i = 10, #t do
						local AA, Aa, aa = t[i]:match('(%d+),(%d+),(%d+)')
						AA = tonumber(AA); Aa = tonumber(Aa); aa = tonumber(aa);
						if AA ~= nil then
							z[#z+1] = lookup[AA]; z[#z+1] = lookup[Aa]; z[#z+1] = lookup[aa];
						else z[#z+1] = 1; z[#z+1] = 1; z[#z+1] = 1; end
					end
					print(t[1]..':'..t[2], t[4], t[5], table.concat(z, '\t'))
				elseif t[9]:find('GL') then
					print('Error: not implemented')
					os.exit(1)
				end
			end
		end
	end
	fp:close()
end

-- CMD bgl2vcf: convert Beagle output to VCF
function cmd_bgl2vcf()
	if #arg < 2 then
		print('Usage: bgl2vcf.lua <in.phased> <in.gprobs>')
		os.exit(1)
	end
	
	local fpp = io.xopen(arg[1]);
	local fpg = io.xopen(arg[2]);
	for lg in fpg:lines() do
		local tp, tg, a = fpp:read():split('%s'), lg:split('%s', 4), {}
		if tp[1] == 'I' then
			for i = 3, #tp, 2 do a[#a+1] = tp[i] end
			print('#CHROM', 'POS', 'ID', 'REF', 'ALT', 'QUAL', 'FILTER', 'INFO', 'FORMAT', table.concat(a, '\t'))
		else
			local chr, pos = tg[1]:match('(%S+):(%d+)$')
			a = {chr, pos, '.', tg[2], tg[3], 30, '.', '.', 'GT'}
			for i = 3, #tp, 2 do
				a[#a+1] = ((tp[i] == tg[2] and 0) or 1) .. '|' .. ((tp[i+1] == tg[2] and 0) or 1)
			end
			print(table.concat(a, '\t'))
		end
	end
	fpg:close(); fpp:close();
end

-- CMD freq: count alleles in each population
function cmd_freq()
	-- parse the command line
	local site_only = true; -- print site allele frequency or not
	for c in os.getopt(arg, 's') do
		if c == 's' then site_only = false end
	end
	if #arg == 0 then
		print("\nUsage: vcfutils.lua freq [-s] <in.vcf> [samples.txt]\n")
		print("NB: 1) This command only considers biallelic variants.")
		print("    2) Apply '-s' to get the allele frequency spectrum.")
		print("    3) 'samples.txt' is TAB-delimited with each line consisting of sample and population.")
		print("")
		os.exit(1)
	end
	
	-- read the sample-population pairs
	local pop, sample = {}, {}
	if #arg > 1 then
		local fp = io.xopen(arg[2]);
		for l in fp:lines() do
			local s, p = l:match("^(%S+)%s+(%S+)"); -- sample, population pair
			sample[s] = p; -- FIXME: check duplications
			if pop[p] then table.insert(pop[p], s)
			else pop[p] = {s} end
		end
		fp:close();
	end
	pop['NA'] = {}
	
	-- parse VCF
	fp = (#arg >= 2 and io.xopen(arg[1])) or io.stdin;
	local col, cnt = {}, {};
	for k in pairs(pop) do
		col[k], cnt[k] = {}, {[0]=0};
	end
	for l in fp:lines() do
		if l:sub(1, 2) == '##' then -- meta lines; do nothing
		elseif l:sub(1, 1) == '#' then -- the sample line
			local t, del_NA = l:split('\t'), true;
			for i = 10, #t do
				local k = sample[t[i]]
				if k == nil then
					k, del_NA = 'NA', false
					table.insert(pop[k], t[i])
				end
				table.insert(col[k], i);
				table.insert(cnt[k], 0);
				table.insert(cnt[k], 0);
			end
			if del_NA then pop['NA'], col['NA'], cnt['NA'] = nil, nil, nil end
		else -- data lines
			local t = l:split('\t');
			if t[5] ~= '.' and t[5]:find(",") == nil then -- biallic
				if site_only == true then io.write(t[1], '\t', t[2], '\t', t[4], '\t', t[5]) end
				for k, v in pairs(col) do
					local ac, an = 0, 0;
					for i = 1, #v do
						local a1, a2 = t[v[i]]:match("^(%d).(%d)");
						if a1 ~= nil then ac, an = ac + a1 + a2, an + 2 end
					end
					if site_only == true then io.write('\t', k, ':', an, ':', ac) end
					if an == #cnt[k] then cnt[k][ac] = cnt[k][ac] + 1 end
				end
				if site_only == true then io.write('\n') end
			end
		end
	end
	fp:close();
	
	-- print
	if site_only == false then
		for k, v in pairs(cnt) do
			io.write(k .. "\t" .. #v);
			for i = 0, #v do io.write("\t" .. v[i]) end
			io.write('\n');
		end
	end
end

function cmd_vcf2chi2()
	if #arg < 3 then
		print("Usage: vcfutils.lua vcf2chi2 <in.vcf> <group1.list> <group2.list>");
		os.exit(1)
	end
	
	local g = {};
	
	-- read the list of groups
	local fp = io.xopen(arg[2]);
	for l in fp:lines() do local x = l:match("^(%S+)"); g[x] = 1 end -- FIXME: check duplicate
	fp:close()
	fp = io.xopen(arg[3]);
	for l in fp:lines() do local x = l:match("^(%S+)"); g[x] = 2 end
	fp:close()
	
	-- process VCF
	fp = io.xopen(arg[1])
	local h = {{}, {}}
	for l in fp:lines() do
		if l:sub(1, 2) == '##' then print(l) -- meta lines; do nothing
		elseif l:sub(1, 1) == '#' then -- sample lines
			local t = l:split('\t');
			for i = 10, #t do
				if g[t[i]] == 1 then table.insert(h[1], i)
				elseif g[t[i]] == 2 then table.insert(h[2], i) end
			end
			while #t > 8 do table.remove(t) end
			print(table.concat(t, "\t"))
		else -- data line
			local t = l:split('\t');
			if t[5] ~= '.' and t[5]:find(",") == nil then -- biallic
				local a = {{0, 0}, {0, 0}}
				for i = 1, 2 do
					for _, k in pairs(h[i]) do
						if t[k]:find("^0.0") then a[i][1] = a[i][1] + 2
						elseif t[k]:find("^1.1") then a[i][2] = a[i][2] + 2
						elseif t[k]:find("^0.1") or t[k]:find("^1.0") then
							a[i][1], a[i][2] = a[i][1] + 1, a[i][2] + 1
						end
					end
				end
				local chi2, p, succ = matrix.chi2(a);
				while #t > 8 do table.remove(t) end
				--print(a[1][1], a[1][2], a[2][1], a[2][2], chi2, p);
				if succ then print(table.concat(t, "\t") .. ";PCHI2=" .. string.format("%.3g", p)
						.. string.format(';AF1=%.4g;AF2=%.4g,%.4g', (a[1][2]+a[2][2]) / (a[1][1]+a[1][2]+a[2][1]+a[2][2]),
						a[1][2]/(a[1][1]+a[1][2]), a[2][2]/(a[2][1]+a[2][2])))
				else print(table.concat(t, "\t")) end
			end
		end
	end
	fp:close()
end

-- CMD: compute r^2
function cmd_r2()
	local w, is_ht, is_gt = 1, false, false
	for o, a in os.getopt(arg, 'w:hg') do
		if o == 'w' then w = tonumber(a)
		elseif o == 'h' then is_ht, is_gt = true, true
		elseif o == 'g' then is_gt = true
		end
	end
	if #arg == 0 then
		print("Usage: vcfutils.lua r2 [-hg] [-w 1] <in.vcf>")
		os.exit(1)
	end
	local stack, fp, q2p = {}, io.xopen(arg[1]), algo_init_q2p(1023)
	for l in fp:lines() do
		if l:sub(1, 1) ~= '#' then
			local t = l:split('\t')
			local x = text_parse_pl(t, q2p)
			if #t[5] == 1 and t[5] ~= '.' then -- biallelic
				local r2 = {}
				for k = 1, w do
					if is_gt == false then -- use PL
						if stack[k] then
							local pdg = { stack[k][5], x[5] }
							r2[#r2+1] = algo_r2(algo_hapfreq2(pdg))
						else r2[#r2+1] = 0 end
					elseif is_ht == false then -- use unphased GT
						if stack[k] then
							local pdg = { stack[k][4], x[4] }
							r2[#r2+1] = algo_r2(algo_hapfreq2(pdg))
						else r2[#r2+1] = 0 end
					else -- use phased GT
						if stack[k] then
							local f, ht = { [0]=0, 0, 0, 0 }, { stack[k][3], x[3] }
							for i = 1, #ht[1] do
								local j = ht[1][i] * 2 + ht[2][i]
								f[j] = f[j] + 1
							end
							local sum = f[0] + f[1] + f[2] + f[3]
							for k = 0, 3 do f[k] = f[k] / sum end
							r2[#r2+1] = algo_r2(f)
						else r2[#r2+1] = 0 end
					end
				end
				for k = 1, #r2 do
					r2[k] = string.format('%.3f', r2[k])
				end
				print(x[1], x[2], table.concat(r2, '\t'))
				if #stack == w then table.remove(stack, 1) end
				stack[#stack+1] = x
			end
		end
	end
	fp:close()
end

-------------------
-- END: commands --
-------------------


-------------------
-- MAIN FUNCTION --
-------------------

if #arg == 0 then
	print("\nUsage:   vcfutils.lua <command> <arguments>\n")
	print("Command: freq        count biallelic alleles in each population")
	print("         r2          compute r^2")
	print("         vcf2chi2    compute 1-degree chi-square between two groups of samples")
	print("         vcf2bgl     convert PL annotated VCF to Beagle input")
	print("         bgl2vcf     convert Beagle input to VCF")
	print("")
	os.exit(1)
end

local cmd = arg[1]
table.remove(arg, 1)
if cmd == 'vcf2bgl' then cmd_vcf2bgl()
elseif cmd == 'bgl2vcf' then cmd_bgl2vcf()
elseif cmd == 'freq' then cmd_freq()
elseif cmd == 'r2' then cmd_r2()
elseif cmd == 'vcf2chi2' then cmd_vcf2chi2()
else
	print('ERROR: unknown command "' .. cmd .. '"')
	os.exit(1)
end