// procedure to compute either Wright's Fst or an unbiased estimator of if

#include "lib.h"
// Wright's Fst
static double Wright(double f1, double f2) {
	double
	  f, // frequency in the pooled population
	  H_ave, // average of HWE heterogosity in the two populations
	  H_all; // HWE heterozygosity in the pooled popuations

	H_ave = f1*(1.0 - f1) + f2*(1.0 - f2);
	f = (f1 + f2)/2.0;
	if (f == 0.0 || f == 1.0)
		return 0.0;
	H_all = 2.0*f*(1.0 - f);
	return (H_all - H_ave) / H_all;
}

/* unbiased estimator of Fst from:
  Weir, B.S. and Cockerham, C.C. 1984. Estimating F-statistics for the
  analysis of population structure. Evolution 38: 1358–1370.
as interpreted by:
  Akey, J.M., Zhang, G., Zhang, K., Jin, L., and Shriver, M.D. 2002.
  Interrogating a high-density SNP map for signatures of natural
  selection. Genome Res. 12: 1805–1814.
*/
static double Weir(int n1, double p1, int n2, double p2) {
	double F, p_bar, nc, MSP, MSG, N = n1 + n2; 

	if (p1 == p2)
		return 0.0;
	MSG = (n1*p1*(1.0-p1) + n2*p2*(1.0-p2))/(N-1.0);
	p_bar = (n1*p1 + n2*p2)/N;
	MSP = n1*(p1-p_bar)*(p1-p_bar) + n2*(p2-p_bar)*(p2-p_bar);
        nc = N - (double)(n1*n1 + n2*n2)/N;
	F = (MSP - MSG) / (MSP + (nc-1)*MSG);
	if (F < 0.0)
		F = 0.0;
	return F;
}

double Fst(int nA1, int na1, int nA2, int na2, int unbiased) {
	double p1, p2;

	p1 = (double)nA1 / (double)(nA1+na1);
	p2 = (double)nA2 / (double)(nA2+na2);

	return (unbiased ? Weir(nA1+na1, p1, nA2+na2, p2) : Wright(p1, p2));
}