#!/usr/bin/perl -w
use strict;
use Bio::DB::HTS::Tabix;
use Getopt::Long;
use File::Basename;
use List::Util qw[min max];
use File::Temp qw/ tempfile unlink0 /;

###########################
####### Documents #########
# The modified UKBB file should look like:
# 0. chr
# 1. pos
# 2. variant 
# 3. minor_allele    
# 4. minor_AF        
# 5. expected_case_minor_AC (to be removed from output)
# 6. low_confidence_variant  
# 7. n_complete_samples      
# 8. AC      
# 9. ytx     
# 10.beta    
# 11.se      
# 12.tstat   
# 13.pval
'chr','pos','variant','minor_allele','minor_AF','low_confidence_variant','n_complete_samples','AC','ytx','beta','se','tstat','pval'
# The original UKBB files were modified by:
# zcat Z87.gwas.imputed_v3.male.tsv.bgz | sed '1s/^/chr\tpos\t/' | awk 'NR==1 { print; next } { split($1, a, ":", seps); print a[1]"\t"a[2]"\t"$0; }' | sed '1d' | sort -k1,1V -k2,2n | bgzip > my-Z87.gwas.imputed_v3.male.tsv.bgz
# and indexed using Tabix:
# tabix -s 1 -b 2 -e 2 my-Z87.gwas.imputed_v3.male.tsv.bgz
#
# How it works:
# 1. read the exposure UKBB GWAS file and keep the common (MAF>=0.01) and significant SNPs (5E-8) in memory.
# 2. add the rsid using UKBB annotation file, variant.tsv.bgz.
# 3. clump the signficant SNPs (external rscript).
# 4. read the outcome GWAS list, and extract non-significant SNPs (p >5E-8, for MR) and significant SNPs (p<=5E-8, for pleotrophy analysis).
# 5. save non-significant SNPS and significant SNPS in separate files.
###########################
###########################
print "UKBB GWAS MR screening script v0.0.1 by Peng Chen\n";

## BEGIN processing params
#  See if requirements are met. Get parameters.
if (scalar(@ARGV) <2) {
	&usage();
	die();
}
my %opts = ();
Getopt::Long::GetOptions(\%opts,qw(
        exp=s
        outcome=s
        prefix=s
	p1=i
	p2=i
)) || die "Failed to parse options\n\n";
$opts{'prefix'} = "output" if (!defined($opts{'prefix'}));
$opts{'p1'} = 5E-8 if (!defined($opts{'p1'}));
$opts{'p2'} = 0.05 if (!defined($opts{'p2'}));
my $annot_file = '/home/peng/test/ukbb-variant-rsid.tsv.gz';

#  Test if exp files exist.
die "E:Unable to find $opts{'exp'}\n" unless ( -e $opts{'exp'} and -r $opts{'exp'} );
die "E:Unable to find the index file for $opts{'exp'}\n" unless ( -e $opts{'exp'}.".tbi" and -r $opts{'exp'}.".tbi" );
#  Test if outcome files exist.
my @outcomes = ();
open OUTCOME, "<$opts{'outcome'}" or die "E:Unable to read outcome list file.\n";
while (<OUTCOME>) {
	chomp;
	push @outcomes, $_;
	die "E:Unable to find outcome $_.\n" unless ( -e $_ and -r $_ );
	die "E:Unable to find outcome index file $_.tbi\n" unless ( -e $_.".tbi" and -r $_.".tbi" );
}
## END processing params

##print "--exp:".$opts{'exp'}."\n";
##print "--outcome:".$opts{'outcome'}."\n";
print STDERR "I:Reading exposure GWAS: ".$opts{'exp'}."...\n";
my $tabix = Bio::DB::HTS::Tabix->new( filename => $opts{'exp'} );

## BEGIN get seqnames
my $ref_seqnames = $tabix->seqnames;
## END get seqnames

## BEGIN loop through seqnames in exposure GWAS
#  put significant SNPs in @signif
my @signif = ();
my ($seqname, $iter);
foreach $seqname (@$ref_seqnames) {
	print STDERR "I:Processing seqname ".$seqname."...\n";
	$iter = $tabix->query($seqname);
	my $num_sig = 0;
	my $num_snp = 0;
	while ( my $n = $iter->next ) {
		# seems no need to chomp?
		$num_snp++;
		my @exp_line = split('\t', $n);
		if ($exp_line[scalar(@exp_line)-1]<=$opts{'p1'} and $exp_line[4] >=0.01) {
			if (scalar(@exp_line) ==14) {
				splice(@exp_line, 5, 1);
			}
			push @signif, \@exp_line;
			$num_sig++;
		}
	}
	print STDERR "I:Processed ".$num_snp." SNPs, Found ".$num_sig." significantly associated common SNPs.\n";
}
$tabix->close;
print STDERR "I:Got total SNPs: ".scalar(@signif)."\n";
die "W:Not sufficient SNPs for downstream processing. I will not proceed.\n\n" if ( scalar(@signif) <=0);
## END loop through seqnames in exposure GWAS

## BEGIN add rsid and clump snps
print STDERR "I:Matching rsid...\n";
open ANNOT, "zcat $annot_file|" or die "E:Unable to open annotation file.\n";
my %annot;
while (<ANNOT>) {
	chomp;
	my @line = split("\t",$_);
	$annot{$line[0]} = $line[1];
}
my $num_uniq_annot = scalar(keys(%annot));
print STDERR "I:Got $num_uniq_annot annotations.\n";
close ANNOT;
my $annotated = 0;
my ($tmp_fh, $tmp_file) = tempfile('cleanukbb-XXXXX', DIR=>'/tmp/');
my $tmp_fh_start = $tmp_fh->getpos;
for (my $i=0;$i<scalar(@signif);$i++) {
	my $ref_arr = $signif[$i];
	my @arr = @$ref_arr;
	if (defined($annot{$arr[2]})) {
		$annotated++;
		push @arr, $annot{$arr[2]};
		$signif[$i] = \@arr;
		print $tmp_fh join("\t", @arr)."\n";
	}
}
print STDERR "I:Added $annotated rsids in $tmp_file.\n\n";
# clump snps
print STDERR "I:Clump SNPs...\n";
my @clump_cmd = ("Rscript", "--vanilla", "~/bin/clump.r", $tmp_file);
@signif = ();
if (system(join(" ", @clump_cmd)) ==0) {
	$tmp_fh->setpos($tmp_fh_start);
	#open TMP, "<$tmp_file" or die "E:Unable to find the temporary file.\n";
	while (<$tmp_fh>) {
		next if ( $. == 1);
		chomp;
		my @line = split("\t",$_);
		push @signif, \@line;
	}
	print STDERR "I:Read ".scalar(@signif)." lines from $tmp_file.\n";
	#close TMP;
} else {
	die "E:Clump SNPS failed, check if the rscript works.\n"
}
print STDERR "\n";
unlink0($tmp_fh, $tmp_file);
## END add rsid and clump snps

for (my $i=0;$i<scalar(@outcomes);$i++) {
	## BEGIN get pvalues in outcome GWASs
	print STDERR "I:Scanning outcome GWAS files: ".$outcomes[$i]."...\n";
	$tabix = Bio::DB::HTS::Tabix->new( filename => $outcomes[$i] );
	my $exp_line;
	my @exp_mr = ();
	my @out_mr = ();
	my @exp_pleo = ();
	my @out_pleo = ();
	foreach $exp_line (@signif) {
		my $outcome_iter = $tabix->query(join('', $$exp_line[0],":",$$exp_line[1]."-".$$exp_line[1]));
		while ( my $outcome_line = $outcome_iter->next ) {
			my @outcome_line = split('\t', $outcome_line);
			next if ( $outcome_line[scalar(@outcome_line)-1] eq 'NaN');
			if ( $outcome_line[2] eq $$exp_line[2] ) {
				if ( scalar(@outcome_line) ==14 ) {
					splice(@outcome_line, 5, 1);
				}
				if (  $outcome_line[scalar(@outcome_line)-1] > 5e-8 ) {
					push @exp_mr, $exp_line;
					push @out_mr, \@outcome_line;
				} elsif ($outcome_line[scalar(@outcome_line)-1] <= 5e-8) {
					push @exp_pleo, $exp_line;
					push @out_pleo, \@outcome_line;
				}
			}
		}
	}
	$tabix->close;
	## END get pvalues in outcome GWAS

	## BEGIN summary and output
	print STDERR "I:Found ".scalar(@exp_mr)." instrument SNPs at 5E-8.\n";
	print STDERR "I:Found ".scalar(@exp_pleo)." pleotrophic SNPs at 5E-8.\n";
	my @exp_file_parse = fileparse($opts{'exp'}, qr/\.[^.]*/);
	my @out_file_parse = fileparse($outcomes[$i], qr/\.[^.]*/);
	my @exp_file = split(/\./,$exp_file_parse[0]);
	my @out_file = split(/\./,$out_file_parse[0]);
	my $prefix = $exp_file[0]."_".$out_file[0]."_".$exp_file[3];
	open EXP_IV, ">$prefix.exp.tsv" or die "E:Unable to write to $prefix.exp.tsv.\n";
	open OUT_IV, ">$prefix.out.tsv" or die "E:Unable to write to $prefix.out.tsv.\n";
	open EXP_PLEO, ">$prefix.exp.pleo.tsv" or die "E:Unable to write to $prefix.exp.pleo.tsv.\n";
	open OUT_PLEO, ">$prefix.out.pleo.tsv" or die "E:Unabel to write to $prefix.out.pleo.tsv.\n";
	for (my $j=0; $j<scalar(@exp_mr); $j++) {
		my $ref_arr = $exp_mr[$j];
		my $line = join("\t", @$ref_arr)."\n";
		print EXP_IV $line;
		
		$ref_arr = $out_mr[$j];
		$line = join("\t", @$ref_arr)."\n";
		print OUT_IV $line;
	}
	
	for (my $j=0; $j<scalar(@exp_pleo); $j++) {
		my $ref_arr = $exp_pleo[$j];
		my $line = join("\t", @$ref_arr)."\n";
		print EXP_PLEO $line;
		
		$ref_arr = $out_pleo[$j];
		$line = join("\t", @$ref_arr)."\n";
		print OUT_PLEO $line;
	}

	close EXP_IV; close OUT_IV; close EXP_PLEO; close OUT_PLEO;
	## END summary and output
	
	print STDERR "\n";
}
print STDERR "Analysis finished! Have a nice day!\n\n";


sub usage {
	print "This program's minimal requirements are:\n";
	print "\t--exp <exposure GWAS summary file>;\n";
	print "\t--outcome <outcome GWAS summary file>;\n";
	print "Both files are modified UKBB GWAS round2 files, compressed by bgzip and indexed by tabix.\n";
}


