#**************************************************************************#
#      SMR test for CACNA2D4
#**************************************************************************#

smr_single_snp = function(bZX, seZX, bZY, seZY) {
  # SMR analysis using the top-associated SNP
  bXY = bZY/bZX
  bXY_se = sqrt( (seZY^2*bZX^2 + seZX^2*bZY^2) / bZX^4 )
  bXY_pvalue = pchisq( (bXY/bXY_se)^2, 1, lower.tail=F ) 
  return(list(bXY = bXY, bXY_se = bXY_se, bXY_pvalue = bXY_pvalue))
}


heidi_two_snps = function(bZX1, seZX1, bZY1, seZY1, bZX2, seZX2, bZY2, seZY2, ld_r) {
  # SMR to test bXY
  smr_result1 = smr_single_snp(bZX1, seZX1, bZY1, seZY1);
  bXY1 = smr_result1$bXY; bXY1_se = smr_result1$bXY_se;
  smr_result2 = smr_single_snp(bZX2, seZX2, bZY2, seZY2);
  bXY2 = smr_result2$bXY; bXY2_se = smr_result2$bXY_se;
  
  # Difference
  d = bXY2 - bXY1
  cov_bXY = ld_r * ( seZY1*seZY2 / (bZX1*bZX2)  + bXY1*bXY2*(seZX1/bZX1*seZX2/bZX2) )
  var_d = bXY1_se^2 + bXY2_se^2 - 2*cov_bXY
  d_pvalue = pchisq(d^2/var_d, 1, lower.tail=F)
  return(list(d = d, d_se = sqrt(var_d), d_pvalue = d_pvalue))
}

# rs1044825
bZX1 = 0.447; seZX1 = 0.0186;
bZY1 = -0.0377; seZY1 = 0.0087;
smr_result1 = smr_single_snp(bZX1, seZX1, bZY1, seZY1)

# rs6489330
bZX2 = 0.211; seZX2 = 0.02384;
bZY2 = -0.0378; seZY2 = 0.0108;
smr_result2 = smr_single_snp(bZX2, seZX2, bZY2, seZY2)

# Testing the difference
ld_r = 0.413
difference = heidi_two_snps(bZX1, seZX1, bZY1, seZY1, bZX2, seZX2, bZY2, seZY2, ld_r)