#!/bin/bash

export job=$1
export TMPDIR=/tmp
export dir=${INF}/sentinels
export list=${INF}/work/INF1.merge
export p=$(awk -v job=${job} 'NR==job+1{print $5}' ${list})
export flanking=1e6
export chr=$(sed '1d' ${list} | awk -v job=${job} 'NR==job{print $8}')
export pos=$(sed '1d' ${list} | awk -v job=${job} 'NR==job{print $9}')
export start=$(sed '1d' ${list} | awk -v job=${job} -v d=${flanking} 'NR==job{start=$2-d;if(start<0) start=0; print start}')
export end=$(sed '1d' ${list} | awk -v job=${job} -v d=${flanking} 'NR==job{print $3+d}')
export r=$(sed '1d' ${list} | awk -v job=${job} 'NR==job{print $6}')
export pr=${p}-${r}
export sumstats=${INF}/sumstats/INTERVAL/INTERVAL.${p}.gz
export bfile=${INF}/EUR
export sample=${INF}/INTERVAL/o5000-inf1-outlier_in-r2.sample
export study=INTERVAL
export N=4994
export n=1000

# module load gcc/5.4.0 lapack/3.8.0 qctool/2.0.6
# module load bgen/20180807

cd work
# z0 - NLRP2 region if overlap
(
  echo MarkerName Chromosome Position Allele2 Allele1 MAF Effect StdErr
  zcat ${sumstats} | \
  awk 'NR > 1' | \
  awk -vchr=${chr} -vstart=${start} -vend=${end} '
  {
    if (($2==chr && $3 >= start && $3 < end && $11 < 1e-3) && !(chr==19 && $3 >= 5329685 && $3 <= 54500000)) {
       if ($8 < 0.5) maf = $8; else maf = 1-$8
       if (maf > 0 && maf <= 0.5 && $9 != "NA" && $10 != "NA") print $1, $2, $3, toupper($6), toupper($7), maf, $9, $10
    }
  } '
) > ${pr}.tmp
awk 'count[$1]++>0 {print $1}' ${pr}.tmp | grep -v -f - ${pr}.tmp > ${pr}.z0
awk 'NR > 1{print $1}' ${pr}.z0 > ${pr}.incl

# bgen
qctool -g ${INF}/INTERVAL/INTERVAL-${chr}.bgen -og ${pr}.bgen -ofiletype bgen -incl-rsids ${pr}.incl

# bgi
bgenix -g ${pr}.bgen -index -clobber
ln -sf ${pr}.bgen.bgi ${pr}.bgi

# z
(
  join <(awk 'NR > 1' ${pr}.z0 | cut -d' ' -f1,4,5  | sort -k1,1) \
       <(bgenix -g ${pr}.bgen -list 2>&1 | awk 'NR>9 && NF==7'| cut -f2,6,7 | sort -k1,1) | \
  awk '{print $1, ($2!=$4)}'
) > ${pr}.flip
(
  awk 'BEGIN {print "rsid", "chromosome", "position", "allele1", "allele2", "maf", "beta", "se", "flip"}'
  join <(awk 'NR > 1' ${pr}.z0) ${pr}.flip | awk '{if($9==1) {t=$4;$4=$5;$5=t};$7=-$7; print}' | awk 'a[$1]++==0'
) > ${pr}.z

function ld_finemap()
# ldstore 1.1 and finemap 1.3.1.
{
  (
    echo "z;ld;snp;config;cred;log;n_samples"
    echo "${pr}.z;${pr}.ld;${pr}.snp;${pr}.config;${pr}.cred;${pr}.log;$N"
  ) > ${pr}.master
  ldstore_v1.1 --bcor ${pr}-1 --bgen ${pr}.bgen --n-threads 1
  ldstore_v1.1 --bcor ${pr}-1 --merge 1
  ldstore_v1.1 --bcor ${pr}-1 --matrix ${pr}.ld
  rm ${pr}-1_*
  mv ${pr}-1 ${pr}.bcor
  rm -rf ${pr}.cred* ${pr}.dose ${pr}.snp ${pr}.config
  if [ ! -f ${dir}/${pr}.jma.cojo ]; then
     export k=1
  else
     export k=$(awk "END{print NR-1}" ${dir}/${pr}.jma.cojo)
  fi
  finemap_v1.3.1 --sss --in-files ${pr}.master --log --n-causal-snps $k --corr-group 0.7 --group-snps
}

function ld_finemap2()
# ldstore 2 and finemap 1.4
{
  (
    echo "z;ld;bgen;bgi;bcor;bdose;snp;config;cred;log;n_samples"
    echo "${pr}.z;${pr}.ld;${pr}.bgen;${pr}.bgi;${pr}.bcor;${pr}.bdose;${pr}.snp;${pr}.config;${pr}.cred;${pr}.log;$N"
  ) > ${pr}.master2
  ldstore_v2.0_x86_64 --in-files ${pr}.master2 --write-bcor
  ldstore_v2.0_x86_64 --in-files ${pr}.master2 --bcor-to-text
  finemap_v1.4_x86_64 --sss --in-files ${pr}.master2 --log --n-causal-snps $k --corr-group 0.7 --group-snps
}

# xlsx
ld_finemap
gzip -f ${pr}.ld
R -q --no-save < ${INF}/csd3/finemap.R > /dev/null

cd -