#!/usr/bin/perl
use strict;


###################################################
#
# this script was created to allow checking of individuals of GWAS chips
#  without sharing of actual genotypes
#    the original script was written from Sarah Medland (ENIMGMA and Ben Neale (PGC)
#       after an idea from Stephan Ripke 
#         see also http://www.ncbi.nlm.nih.gov/pubmed/22302573?dopt=Abstract
#
# for internal checks, see /home/radon01/sripke/bakker_ripke/hapmap_ref/cksum_files/0114a/README          
#
################################################

####################################
## todos
##
##
##  check if all SNPs are present in dataset -> done
##  check if cksum is in path -> done
 
my $helper_filename = "CKS_batches_50_0114b.tar.gz";;
my $helper_version = "cks_batch_50_0114";;


###################################################
###  system call with test if successfull
###################################################

sub mysystem(){
    my ($systemstr)="@_";
    system($systemstr);
    my $status = ($? >> 8);
    die "$systemstr\n->system call failed: $status" if ($status != 0);
}


##########################################
# subroutine to split a plink-output-line with references
##########################################

sub split_line_ref {
    my ($line)=${$_[0]};
    chomp($line);
    $line =~ s/^[\s]+//g;
    my @cols=  split /\s+/, $line;
    \@cols;
}


my $sourcedir = "";
my $outname = "";
my $error = 0;

use Getopt::Long;
GetOptions( 
    "bfile=s"=> \my $bfile,
    "source=s"=> \$sourcedir,
    "name=s"=> \$outname,
    "keep=s"=> \my $keepf,
    "ploc=s"=> \my $ploc,
    "force"=> \my $force,
 
    "help"=> \my $help,
 );

my $plink_bin = "plink";
### try to start plink
#my $test_plink = "plink --noweb >/dev/null 2>&1";
my $test_plink = "plink --help >/dev/null 2>&1";
system($test_plink);
my $status = ($? >> 8);


if ($status != 0) {
    print "...... cannot find plink in path ($status), test options...\n";

    my $test_plink = "$ploc/plink --help >/dev/null 2>&1";
    system($test_plink);
    my $status = ($? >> 8);
    if ($status != 0) {
	print "------------------------------------------------------------------------\n";
	print "Error: no plink binary found please give location via --ploc, see --help\n";
    	print "    here is a plink binary at broad: --ploc /home/unix/sripke/plink_src/src/\n";
	print "------------------------------------------------------------------------\n";
	unless ($help) {
	    die;
	}
	$error = 1;
    }
    else {
	print "...... plink binary found in specified location\n";
	$plink_bin = "$ploc/plink";
    }
    
}



my $test_cksum = "echo l | cksum > /dev/null 2>&1";
system($test_cksum);
my $status = ($? >> 8);


if ($status != 0) {

    print "------------------------------------------------------------------------\n";
    print "Error: program <cksum> not found\n";
    print "    this is highly unusual on Unix platforms, plesase contact S.Ripke \n";
    print "    see --help \n";
    print "------------------------------------------------------------------------\n";
    unless ($help) {
	die;
    }
    $error = 1;
    
}

my $nbatch = 10;

if ($help || $error){
    print "

######################################################################################
#
#  id_geno_checksum
#
# this script allows checking of individuals of GWAS chips without sharing of actual genotypes
#    the precursor shell script was written from Sarah Medland (ENIMGMA) and Ben Neale (PGC)
#    after an idea from Stephan Ripke 
#
#    this perl script has been written by Stephan Ripke after an idea of Stephan Ripke,
#            many thanks to Menachem Fromer for intellectual help
#
#  it uses 10 btaches with 50 SNPs each, which means 3^50 = 7.1e23 different outcomes with genotypes 0,1,2 at each SNP
#  the somehow outdated cksum command is used on purpose since it restricts to 32 bit differnt cksums: 2^32 = 4.3e10
#  this means for each checksum there are ~1.0e23 differnt genotype configurations, so the cksum does not provide a single key to the genotype configuration
#  still the number of different cksums is high enough that a random false positive overlap seems unlikely
#
#
#   as a published alternative consider also:
#     see also http://www.ncbi.nlm.nih.gov/pubmed/22302573?dopt=Abstract
######################################################################################

  usage: $0 --bfile BFILE

      options:
       
       mandatory:
        --bfile STRING   plink binary to check (root without bed/bim/fam)


       optional
        --source STRING directory for the location of the helper-files,
                            default is current working directory
                            if the files are not found they will be downloaded from the web
        --name STRING   name for this run of checksums, be default the bfile
        --keep STRING   specify keep file for individuals to keep
        --ploc STRING   directory that contains plink-binary
                           on broad: /home/unix/sripke/plink_src/src/

        --help          this help message


 for download of helper files use these two commands:

   - wget http://www.broadinstitute.org/~sripke/share_links/checksums_download/$helper_filename 
   - tar -xvf $helper_filename 

---------------------------------------------------
 created by Stephan Ripke 2014 at MGH, Boston, MA
 feel free to contact me:
    sripke (at) broadinstitute (dot) org
--------------------------------------------------
 \n";
    exit 2;
}

#   - wget http://www.broadinstitute.org/~sripke/share_links/checksums_download/CheckSum28Jan14.zip
#   - unzip CheckSum28Jan14.zip


#exit;

use Cwd;
use File::Path;


print "checking presence of all necessary files\n";
unless ($bfile) {
    print "------------------------------------------------------------------------\n";
    print "Error: no plink dataset specified, see --help\n";
    print "------------------------------------------------------------------------\n";
    die;
}

unless (-e "$bfile.bed" || -e "$bfile.bim" || -e "$bfile.fam") {
    print "------------------------------------------------------------------------\n";
    print "Error: plink dataset (bed/bim/fam) not found, see --help\n";
    print "------------------------------------------------------------------------\n";
    die;
}

my $rootdir = &Cwd::cwd();

if ($sourcedir eq "") {
    $sourcedir = "$rootdir";
}

#my $batch_templ =  "XXX_SNPs";
my $batch_templ =  "cks_batch_50_0114.XXX";
foreach my $i (0..$nbatch-1) {
    my $batch_file = "$sourcedir/$batch_templ";
    my $i_str = sprintf "%02d",$i;
    $batch_file =~ s/XXX/$i_str/;
    unless (-e "$batch_file") {
	print "------------------------------------------------------------------------\n";
	print "$batch_file not found, see --help for download instructions\n";


	print " - wget http://www.broadinstitute.org/~sripke/share_links/checksums_download/$helper_filename \n";
	print " - tar -xvf $helper_filename \n";


	print "------------------------------------------------------------------------\n";
	die;
    }
#    unless (-e "$sourcedir/$i"."_Alleles") {
#	print "------------------------------------------------------------------------\n";
#	print "$sourcedir/$i"."_Alleles not found, see --help for download instructions\n";
#	print "------------------------------------------------------------------------\n";
#	die;
#   }
}


#unless (-e "$sourcedir/SNP.list") {
#    print "------------------------------------------------------------------------\n";
#    print "$sourcedir/SNP.list not found, see --help for download instructions\n";
#    print "------------------------------------------------------------------------\n";
#    die;
#}

my $ra_file = "$sourcedir/$helper_version.recode_allele";

unless (-e "$ra_file") {
    print "------------------------------------------------------------------------\n";
    print "$ra_file not found, see --help for download instructions\n";
    print "------------------------------------------------------------------------\n";
    die;
}

my $ua_file = "$sourcedir/$helper_version.update_alleles";

unless (-e "$ua_file") {
    print "------------------------------------------------------------------------\n";
    print "$ua_file not found, see --help for download instructions\n";
    print "------------------------------------------------------------------------\n";
    die;
}

#print "fine\n";
#exit;

if ($outname eq "") {
    $outname = "$bfile";
}



my $workdir = "checkSum_$outname";

if (-e $workdir) {
    print "-----------------------------------------------------------------------------------------\n";
    print "Error: directory $workdir is already existing, please remove or choose another outname, see --help\n";
    print "-----------------------------------------------------------------------------------------\n";
    unless ($force) {
	die;
    }
}
if (-e "CheckSum.$outname.cs") {
    print "-----------------------------------------------------------------------------------------\n";
    print "Error: result file CheckSum.$outname.cs is already existing, please remove or choose another outname, see --help\n";
    print "-----------------------------------------------------------------------------------------\n";
    die;

}


my @created = mkpath(   ## $created ?
			"$workdir",
			{verbose => 0, mode => 0750},
    );


##################################################
### read snp list into hash
####################################################
print "testing dataset format...\n";
die "$ra_file: ".$! unless open RA, "< $ra_file";
my %recode_hash;
while (my $line = <RA>){
    my @cells = @{&split_line_ref(\$line)};
    $recode_hash{$cells[0]} = $cells[1];
}
close RA;

my $sin = 0;
die "$bfile.bim: ".$! unless open BIM, "< $bfile.bim";
while (my $line = <BIM>){
    my @cells = @{&split_line_ref(\$line)};
    if (exists $recode_hash{$cells[1]}) {
	$sin++;
    }
}
close BIM;

if ($sin < 10) {
    print "-----------------------------------------------------------------------------------------------\n";
    print "Error: less than 10 SNPs found being overlapping with dataset, please make sure:\n";
    print "1) snp names: rs-names\n";
    print "2) whole genome data required, this script will not work on a handful of SNPs\n";
    print "-----------------------------------------------------------------------------------------------\n";
    exit;
}






print "Extracting batch SNPs....\n...";

my $plinksys = "$plink_bin --noweb --bfile $bfile --extract $ra_file --make-bed --out $workdir/snps_for_checksum.$outname";

if ($keepf) {
    $plinksys .= " --keep $keepf";
}

unless (-e "$workdir/snps_for_checksum.$outname.bed") {
    &mysystem($plinksys);
}

print "...done\n";




print "Align batch SNPs....\n...";

my $plinksys = "$plink_bin --noweb --bfile $workdir/snps_for_checksum.$outname --update-alleles $ua_file --make-bed --out $workdir/snps_for_checksum.$outname.sa";

unless (-e "$workdir/snps_for_checksum.$outname.sa.bed") {
    &mysystem($plinksys);
}

print "...done\n";





###########################################################
### check if recode-allele is present in all SNPs (flips before alignment)
#########################################################


my $preflips = 0;
die "$workdir/snps_for_checksum.$outname.bim: ".$! unless open BIM, "< $workdir/snps_for_checksum.$outname.bim";

while (my $line = <BIM>){
    my @cells = @{&split_line_ref(\$line)};
    if ($cells[4] ne $recode_hash{$cells[1]} && $cells[5] ne $recode_hash{$cells[1]}) {
	$preflips++;
    }
}
close BIM;

#print "$preflips preflips\n";
#exit;

###########################################################
### check if recode-allele is present in all SNPs (no flips left) after alignment
#########################################################


my $strandflips = 0;
my $strand_ex = "SNP\tA1\tA2\tref_Allele";
die "$workdir/snps_for_checksum.$outname.sa.bim: ".$! unless open BIM, "< $workdir/snps_for_checksum.$outname.sa.bim";

while (my $line = <BIM>){
    my @cells = @{&split_line_ref(\$line)};
    if ($cells[4] ne $recode_hash{$cells[1]} && $cells[5] ne $recode_hash{$cells[1]}) {
	$strandflips++;
	if ($strandflips < 6) {
	    $strand_ex .= "$cells[1]\t$cells[4]\t$cells[5]\t$recode_hash{$cells[1]}\n";
	}
    }
}
close BIM;

if ($strandflips > 0) {
    print "-----------------------------------------------------------------------------------------------------------\n";
    print "Error: $strandflips SNPs have mismatching allele names (e.g. CheckSum.$outname.cs.error), please make sure:\n";
    print "1) alleles being A,C,G,T\n";
    print "-----------------------------------------------------------------------------------------------------------\n";


    die "CheckSum.$outname.cs.error: ".$! unless open ERR, "> CheckSum.$outname.cs.error";
    
    print ERR "$strand_ex\n";
    
    close ERR;
    die;
}
else {
    if (-e "CheckSum.$outname.cs.error") {
	&mysystem ("rm CheckSum.$outname.cs.error");
    }
}


#exit;

###########################################################
### creating checksums
#########################################################

print "...working on the checksums\n";


my %id_cs;
my $failed_batches = 0;
my %fb ;

foreach my $i (0..$nbatch-1) {

    my $batch_file = "$sourcedir/$batch_templ";
    my $i_str = sprintf "%02d",$i;
    $batch_file =~ s/XXX/$i_str/;

   #### extract subset
    my $plinksys = "$plink_bin --noweb --bfile $workdir/snps_for_checksum.$outname.sa --recodeA --extract $batch_file --recode-allele $ra_file --out $workdir/CheckSum.$i.$outname.sa";
    
    unless (-e "$workdir/CheckSum.$i.$outname.sa.raw") {
	&mysystem($plinksys);
    }

    my $i1 = $i + 1;
    print "read $workdir/CheckSum.$i.$outname.sa.raw: batch $i1 out of $nbatch\n";
    die "$workdir/CheckSum.$i.$outname.raw: ".$! unless open RAW, "< $workdir/CheckSum.$i.$outname.sa.raw";

    my $ll=0;
    my $head = <RAW>;
    while (my $line = <RAW>){
	my $geno_str = "";
	my @cells = @{&split_line_ref(\$line)};
	my $cc = 0;
	foreach my $g (6..$#cells) {
	    if ($cells[$g] eq "NA") {
		$cells[$g] = 9;
	    }
	    $geno_str .= "$cells[$g]";
	    $cc++;
	}
	my $cks = `echo $geno_str | cksum`;
	if ($cc != 50) {
	    $cks = "---";
	    unless (exists $fb{$i}){
		$failed_batches++;
	    }
	    $fb{$i} = 1;
	}
	my @cells_cs = @{&split_line_ref(\$cks)};
	my $id_tmp = $cells[0]." ". $cells[1];
	$id_cs{$id_tmp} .= " ".$cells_cs[0];

	$ll++;


    }
    close RAW;

}

die "$workdir/CheckSum.$outname.cs: ".$! unless open CS, "> $workdir/CheckSum.$outname.cs";

foreach my $k (keys %id_cs) {
    print CS $k.$id_cs{$k}."\n";
}

close CS;


die "CheckSum.$outname.cs.log: ".$! unless open LOG, "> CheckSum.$outname.cs.log";

print LOG "successfully created checksums for $nbatch batches with 50 SNPs each\n";
print LOG "\nnumber of failed batches (not all SNPs found): $failed_batches\n";
print LOG "\nnumber of strand flipped aligned SNPs: $preflips\n";
print LOG "\nthese were the helper_files: $helper_version\n";

close LOG;

&mysystem ("cp $workdir/CheckSum.$outname.cs .");
    
    print "######################################################################\n";
    print "####################     SUCCESS    ##################################\n";
    print "######################################################################\n";
    print "## please share file CheckSum.$outname.cs with your collaborator    ##\n";
    print "## have a look at CheckSum.$outname.cs.log for more details         ##\n";
    print "## use id_geno_checktest for convenient overlap test / summary      ##\n";
    print "######################################################################\n";

exit;