#!/bin/bash

module load plink2
GCTA="./gcta_nr_robust"
PLINK="plink --allow-no-sex"
GEMMA="./gemma"
FUSION="./FUSION/"

PRE="Adipose_Subcutaneous"
PRE="Prostate"
PRE="Whole_Blood"
NR=10

PRE=$1
NR=${SLURM_ARRAY_TASK_ID}

### Input files needed:
# "$PRE.v7.normalized_expression.bed.gz.HEADER" is generated by `cat [bed.gz] | head -n1 | tr '\t' '\n'`
# "$PRE.v7.covariates.txt.covar" is a PLINK format covariates file
# "$PRE.v7.normalized_expression.bed.gz" is the gzipped bed format expression matrix
# ---

mkdir tmp/$PRE.$NR
rm -f HSQ/$PRE.$NR.hsq HSQ/TSPEC.$PRE.$NR.hsq

COVAR="$PRE.v7.covariates.txt.covar"
zcat $PRE.v7.normalized_expression.bed.gz | tail -n+2 | awk -v i=$NR 'NR > (i-1)*100 && NR <= i*100' |  while read PARAM; do

# PARAM contains the expression values for this gene
GNAME=`echo $PARAM | awk '{ print $4 }'`
CHR=`echo $PARAM | awk '{ print $1 }'`
P0=`echo $PARAM | awk '{ p=$2 - 500e3; if(p<0) p=0; print p; }'`
P1=`echo $PARAM | awk '{ print $3 + 500e3 }'`

OUT="tmp/$PRE.$NR/$PRE.$GNAME"

# Conver the expression matrix to a PLINK format phenotype
echo $GNAME $CHR $P0 $P1
echo $PARAM | tr ' ' '\n' | paste $PRE.v7.normalized_expression.bed.gz.HEADER $PRE.v7.normalized_expression.bed.gz.HEADER - | tail -n+5 > $OUT.pheno

# Extract the locus around this gene
rm -f $OUT.bed
$PLINK --silent --bfile GTEx.v7.HM3 --chr $CHR --from-bp $P0 --to-bp $P1 --make-bed --out $OUT --pheno $OUT.pheno --keep $OUT.pheno --maf 0.0001
if [ ! -f $OUT.bed ]; then
continue
fi

# Run FUSION
FINAL_OUT="WEIGHTS/$PRE/$PRE.$GNAME"
Rscript $FUSION/FUSION.compute_weights.R \
--bfile $OUT --tmp $OUT.tmp --out $FINAL_OUT --verbose 0 --save_hsq --PATH_gcta $GCTA --PATH_gemma $GEMMA --models blup,lasso,top1,enet --covar $PRE.v7.covariates.txt.covar --hsq_p 1.0
cat $FINAL_OUT.hsq | awk -vw="$PRE $w" '{ print w,$0 }' >> HSQ/$PRE.$NR.hsq
rm -f $FINAL_OUT.hsq

rm $OUT.*
done

rm -fr tmp/$PRE.$NR
touch HSQ/$PRE.$NR.done
