#!/usr/bin/perl
use warnings;
use Carp;
use Getopt::Long;

my $start_time = time();

our ($help);
our ($out, $fam, $mlinfo,  $mldose, $legend, $rsq , $chr, $legendfile, $batch);

GetOptions (	'help|h'=>\$help,  
# required 
		'faminfo|f=s'=>\$fam, 
		'mlinfo|i=s'=>\$mlinfo, 
		'mldose|d=s'=>\$mldose, 
		'legend|l=s'=>\$legendfile,

# optional 
		'chr=s'=>\$chr, 
		'out=s'=>\$out, 
		'batch|b=f'=>\$batch, 
		'rsq=f'=>\$rsq ) ;

print STDERR "\n********************************************************\n";
print STDERR "\nRestructuring Mach data to merlin-offline infered format\n\n";
print STDERR "       Sarah Medland - QIMR-GenEpi 2010\n";
print STDERR "       type -help for usage\n";
################ Last edit 17/12/10 - added batch options
print STDERR "\n********************************************************\n";


if ($help) {
  print " This perl script takes MACH format data and produces MERLIN-infer format files\n";
  print "\n Four input files are required\n";
  print "\n - fam file (\$fam or \$f)\n";
  print "        - assumed to have the following 6 columns\n";
  print "        - FID IID PID MID Sex Zygosity\n";
  print "        - Zygosity coding follows Merlin format, MZ=1, DZ=2, Singleton=0\n";
  print "\n - legend file (\$legendfile or \$l)\n";
  print "        - assumed to be a HapMapII (r22b36) legend file without any reformating\n";
  print "\n - mlinfo file (\$mlinfo or \$i)\n";
  print "        - assumed to have the come from Mach without any reformating\n";
  print "\n - mldose file (\$mldose or \$d)\n";
  print "        - assumed to have the come from Mach without any reformating\n";
  print "        - assumed to be gziped\n";
  print "\n Five gziped output files are produced\n";
  print "\n - map file (eg infer_format_22.map.gz)\n";
  print "        - Contains CHR SNP ~cM\n";
  print "        - ~cM position = BP/1000000\n";
  print "\n - dat file (eg infer_format_22.dat.gz)\n";
  print "        - This is an infer format dat file in which dose takes the following format\n";
  print "          T COUNT(C,rs11089130)\n";
  print "          T COUNT(A,rs738829)\n";
  print "\n - skip file (eg infer_format_22.skip.gz)\n";
  print "        - This is an infer format dat file in which snps with an rsq < .3 (or the defined rsq) are skiped\n";
  print "\n - freq file (eg infer_format_22.freq.gz)\n";
  print "        - This is an infer format freq file in which freq for each snp takes the following format\n";
  print "          M rs2738388\n";
  print "          A     C 0.7714\n";
  print "          A     A 0.2286\n";
  print "\n - ped file (eg infer_format_22.ped.gz)\n";
  print "        - This is an infer format ped file in containing zygosity and dosage data\n";
#  print "\n - idlist file (eg infer_format_22.idlist.gz)\n";
#  print "        - A list of FID and IID for those in the ped file\n";
  print "\n Optional arguements\n";
  print "\n - \$chr provide the chromosome number - default is taken from the name of the legend file\n";
  print "\n - \$rsq provide an alturnate rsq thresholds - default is .3 \n";
  print "\n - \$batch N cut the output files into batches of N markers\n";
  print "\n - \$out provide and alturnate prefix - default is infer_format_\$chr\n\n";

  exit;
  $help = 0;
}

#check if required files are present
if (!$fam)    {  die "FATAL ERROR: \nRequired file is missing: please specify -fam argument\n";}
if (!$mlinfo) {  die "FATAL ERROR: \nRequired file is missing: please specify -mlinfo argument\n";}
if (!$mldose) {  die "FATAL ERROR: \nRequired file is missing: please specify -mldose argument\n";}
#if (!$mlgeno) {  die "FATAL ERROR: \nRequired file is missing: please specify -mlgeno argument\n";}
#if (!$mlprob) {  die "FATAL ERROR: \nRequired file is missing: please specify -mlprob argument\n";}
if (!$legendfile) {  die "FATAL ERROR: \nRequired file is missing: please specify -legend argument\n";}

if (!$chr) {
	if ($legendfile =~ m/chr(\d+)/) {
		print STDERR "  The -chr argument is automatically set as $1 based on the -legend argument\n";
		$chr = $1;
	} 
	else {  die ("FATAL ERROR: \nThe name of the legend file does not contain the chromosome number: please specify -chr argument\n"); 
	}
}

if (!$out) { 	$out= "infer_format_".$chr;
	print STDERR "  The -out argument is automatically set to $out \n";
}

if (!$rsq) { 	$rsq= 0.3;
	print STDERR "  The -rsq argument is automatically set to $rsq \n";
}


my ($LEGE, $INFO);
open($INFO, "$mlinfo")|| die "FATAL ERROR: \nCannot read the mlinfo file $mlinfo\n";


#####################################################
# Write files that are not affected by batch option #
#####################################################


$legend = readLegend ($legendfile);
writeMap  ($out, $mlinfo, $legend,  $chr);

sub readLegend {
	my ($legendfile) = @_;
	my (%legend);
	open (LEGEND, $legendfile) || die "FATAL ERROR: \nCannot read from LEGEND file $legendfile\n";
	my $countline=();
	while (<LEGEND>) {
		m/^(\S+)\s+(\S+)/ and $legend{$1} = $2;
		$countline++;
	}
	return (\%legend);
}


sub writeMap {
	my ($out, $mlinfo, $legend, $chr) = @_;
	open (MAP, "| gzip -c >$out.map.gz") || die "FATAL ERROR: \nCannot write to MAP file $out.map";
	open (INFO, $mlinfo) || die "FATAL ERROR: \nCannot read from mlinfo file $mlinfo \n";
	$_ = <INFO>;
	$_ =~ m/^SNP\s+Al1\s+Al2\s+Freq1\s+MAF\s+Quality\s+Rsq/ || die "FATAL ERROR: \nInvalid first line in info file: <$_>\n";
	my $countline=();
	while (<INFO>) {
		m/^(\S+)\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\S+)/ || die "FATAL ERROR: \nInvalid record found in info file: <$_>\n";
		if ($legend->{$1}) {
			print MAP "$chr\t$1\t", $legend->{$1}/1000000, "\n";
		} else {
			die "FATAL ERROR: \nThe SNP $1 is listed in the INFO file $mlinfo, but not annotated in the LEGEND file $legendfile \n";
		}
		$countline++;
	}
	print STDERR "\nFinished writting the map file $out.map.gz (with $countline snps)\n";
	print STDERR "Note:  An approxiate cM position has been calculated using BP/1000000 \n";
	close (INFO);
	close (MAP);
	$ndoses=$countline;
}
#	print STDERR " Debug nprobs is $ndoses\n" ;

writeFreq ($out, $mlinfo);
sub writeFreq {
	my ($out, $mlinfo ) = @_;
	open (DAT, "| gzip -c >$out.freq.gz") || die "FATAL ERROR: \nCannot write to FREQ file $out.freq\n";
	open (INFO, $mlinfo) || die "FATAL ERROR: \nCannot read from mlinfo file $mlinfo \n";
	$_ = <INFO>;
	while (<INFO>) {
		my ($snp, $A1, $A2, $f1,) = split ;
		my $f2 = sprintf("%.4f", 1-$f1);
		print DAT "M $snp\nA\t$A1 $f1\nA\t$A2 $f2\n";

	}
	close (INFO);
	close (DAT);
	print STDERR "\nFinished writting freq file $out.freq.gz\n";
}


#####################################################
#      Write files if there is no batch option      #
if (!$batch) { 
#####################################################

writeDat  ($out, $mlinfo);
writeSkip ($out, $mlinfo);

sub writeDat {
	my ($out, $mlinfo ) = @_;
	open (DAT, "| gzip -c >$out.dat.gz") || die "FATAL ERROR: \nCannot write to DAT file $out.dat\n";
	open (INFO, $mlinfo) || die "FATAL ERROR: \nCannot read from mlinfo file $mlinfo \n";
	print DAT "Z Zygosity\n";
	$_ = <INFO>;
	while (<INFO>) {
		my ($snp, $A1, $A2,) = split ;
		print DAT "T COUNT($A1\,$snp)\n";

	}
	close (INFO);
	close (DAT);
#	print STDERR "\nDebug batch option is in not use\n";
	print STDERR "\nFinished writting a dat file $out.dat.gz containing all snps\n";
}

sub writeSkip {
	my ($out, $mlinfo ) = @_;
	open (DAT, "| gzip -c >$out.skip.gz") || die "FATAL ERROR: \nCannot write to DAT file $out.dat\n";
	open (INFO, $mlinfo) || die "FATAL ERROR: \nCannot read from mlinfo file $mlinfo \n";
	print DAT "Z Zygosity\n";
	$_ = <INFO>;
	while (<INFO>) {
		my ($snp, $A1, $A2, $freq,$maf,$qual,$RSQ) = split ;
		if ($RSQ>$rsq) {
		print DAT "T COUNT($A1\,$snp)\n";
		} else {
		print DAT "S COUNT($A1\,$snp)\n";
	}
}
	close (INFO);
	close (DAT);
	print STDERR "\nFinished writting a dat file $out.skip.gz skipping all snps with an rsq<$rsq\n";
}


$pedinfo = readPedInfo ($fam);
sub readPedInfo {
	my ($fam) = @_;
	my $pedinfo = {};
	open (PED, $fam) || die "FATAL ERROR: \nCannot read from FAM file $fam \n";
	while (<PED>) {
		m/^((\S+)\s+(\S+)\s+\S+\s+\S+\s+\S+\s+\S+)/ || die "FATAL ERROR: Error: invalid record found: <$_>\n";
		$pedinfo->{$2}{$3} = $1;

	}
	close (PED);
	return ($pedinfo);
}

writePed ($out, $mldose,  $fam) ;

sub writePed {
	my ($out,  $mldose,  $fam) = @_;
	open (IDS, "| gzip -c >$out.idlist.gz") || die "FATAL ERROR: \nCannot write to idlist";
	open (PED, "| gzip -c >$out.ped.gz") || die "FATAL ERROR: \nCannot write to PED file $out.ped.gz";
	open (DOSE, " gzip -cd  $mldose | ") || die "FATAL ERROR: \nConnot read the mldose file $mldose\n";
	my $countple=();

	while (1) {
		$_ = readline (DOSE);
		$_ or last;
		
		my ($famid, $indid, $geno, @geno  );
		s/[\r\n]+$//;
		m/^(\S+)\->(\S+)\s+ML_DOSE\s+(.+)/ || die "FATAL ERROR: \nInvalid record (famid, indid, ML_DOSE expected): <$_>\n";
		($famid, $indid, $geno) = ($1, $2, $3);
		@geno = split (/\s+/, $geno);
		$size = @geno;

		if (@geno !=$ndoses) { die "FATAL ERROR: \nFewer dose calls for family $famid individual $indid in $mldose (" . scalar (@geno) . ") versus $ndoses snps in the mlinfo file\n";}
		
		$pedinfo->{$famid}{$indid} || die "FATAL ERROR: \nIndividual $famid $indid is listed in DOSE file $mldose, but not annotated in FAM file $fam\n";
		print PED $pedinfo->{$famid}{$indid};
		print PED " $geno\n";
		#Print out a list of ids for records
		print IDS $pedinfo->{$famid}{$indid};
		print IDS "\n";
		$countple++;

}
	close (IDS);
	close (DOSE);
	close (PED);	
	print STDERR "\nFinished writting the ped file $out.ped.gz (with $countple individuals and $size snps) \n ";
	print STDERR "\nFinished writting the id list to $out.idlist.gz  \n ";

}
#####################################################
#           End of no batch option code             #
}
#####################################################


#####################################################
#      Write files if there is a batch option       #
if ($batch) { 
#####################################################
use POSIX;
 my ($cuts) ;
	$cuts = ceil($ndoses/$batch);
#	print STDERR " Debug:  ndoses $ndoses : nbatches $batch : ncuts $cuts \n" ;
print STDERR "\n********************************************************\n";
	print STDERR " Parallelisation mode is enabled\n" ;
	print STDERR "* Chromosome $chr has $ndoses snps $cuts ped, dat and skip files will be produced \n" ;
	print STDERR "* Each set of files contains $batch snps\n" ;
 $stop=1;
#	open (DOSE, " gzip -cd  $mldose | ") || die "FATAL ERROR: \nConnot read the mldose file $mldose\n";


$pedinfo = readPedInfo1 ($fam);
sub readPedInfo1 {
	my ($fam) = @_;
	my $pedinfo = {};
	open (PED, $fam) || die "FATAL ERROR: \nCannot read from FAM file $fam \n";
	while (<PED>) {
		m/^((\S+)\s+(\S+)\s+\S+\s+\S+\s+\S+\s+\S+)/ || die "FATAL ERROR: Error: invalid record found: <$_>\n";
		$pedinfo->{$2}{$3} = $1;
	}
	close (PED);
	return ($pedinfo);
}

#####################################################
#             Start of  batch loop                  #
for ($curent = 1; $curent <= $cuts ; $curent++) {
#####################################################
 our $start=();
 $start = $stop + 1 ;
 $startsnp =$start -1;
 $startcol =$startsnp -1;
 if ((($curent*$batch)+1) <= ($ndoses+1)) { $stop = ($curent*$batch)+1 ;}
 if ((($curent*$batch)+1) > ($ndoses+1)) { $stop = ((($curent-1)*$batch)+$ndoses%$batch)+1 ;}

# $stop  = ($curent*$batch)+1 ;
 $stopsnp =$stop -1;
 $stopcol =$stopsnp -1;


#	print STDERR "\n Debug:  start $start : stop $stop " ;
#	print STDERR "\nDebug batch option is in use\n";

writeDat1  ($out, $mlinfo);
writeSkip1 ($out, $mlinfo);

sub writeDat1 {
	my ($out, $mlinfo ) = @_;
	open (DAT, "| gzip -c >$out.$curent.dat.gz") || die "FATAL ERROR: \nCannot write to DAT file $out.dat\n";
	open (INFO, $mlinfo) || die "FATAL ERROR: \nCannot read from mlinfo file $mlinfo \n";
	print DAT "Z Zygosity\n";
	$_ = <INFO>;
	while (<INFO>) {
		my ($snp, $A1, $A2,) = split ;
	if (($. >= $start )  && ($. <= $stop)) {
		print DAT "T COUNT($A1\,$snp)\n";
#	print STDERR "Debug: T COUNT($A1\,$snp)\n";
} 
	}
	close (INFO);
	close (DAT);
	print STDERR "\nFinished writting a dat file $out.$curent.dat.gz containing snps $startsnp to $stopsnp \n";
}

sub writeSkip1 {
	my ($out, $mlinfo ) = @_;
	open (DAT, "| gzip -c >$out.$curent.skip.gz") || die "FATAL ERROR: \nCannot write to DAT file $out.dat\n";
	open (INFO, $mlinfo) || die "FATAL ERROR: \nCannot read from mlinfo file $mlinfo \n";
	print DAT "Z Zygosity\n";
	$_ = <INFO>;
	while (<INFO>) {
		my ($snp, $A1, $A2, $freq,$maf,$qual,$RSQ) = split ;
	if (($. >= $start )  && ($. <= $stop)) {
		if ($RSQ>$rsq) {
		print DAT "T COUNT($A1\,$snp)\n";
#	print STDERR "Debug: $snp\n";
		} else {
		print DAT "S COUNT($A1\,$snp)\n";
#	print STDERR "Debug: $snp\n";
	}
} }
	close (INFO);
	close (DAT);
	print STDERR "Finished writting a dat file $out.$curent.skip.gz skipping snps with an rsq<$rsq\n";}
}
# end of batch loop

	print STDERR "\n";
writePed1 ($out, $mldose, $fam) ;
sub writePed1 {
	my ($out,  $mldose,  $fam) = @_;

	my $countple=();
	open (DOSE, " gzip -cd  $mldose | ") && print STDERR "Reading dosage file $mldose\n" ;
#	open (DOSE, " gzip -cd  $mldose | ") || die "FATAL ERROR: \nConnot read the mldose file $mldose\n";

#debugging	open (IDS, "| gzip -c >$out.$curent.idlist.gz") || die "FATAL ERROR: \nCannot write to idlist";

for ($curent = 1; $curent <= $cuts ; $curent++) {
	open (PED.$curent, "| gzip -c >$out.$curent.ped.gz") || die "FATAL ERROR: \nCannot write to PED file $out.$curent.ped.gz";
}
	while (1) {
		$_ = readline (DOSE);
		$_ or last;
		
		my ($famid, $indid, $geno, @geno  );
		my @slicegeno=();
		s/[\r\n]+$//;
		m/^(\S+)\->(\S+)\s+ML_DOSE\s+(.+)/ || die "FATAL ERROR: \nInvalid record (famid, indid, ML_DOSE expected): <$_>\n";
		($famid, $indid, $geno) = ($1, $2, $3);
		@geno = split (/\s+/, $geno);
		$size = @geno;

		if (@geno !=$ndoses) { die "FATAL ERROR: \nFewer dose calls for family $famid individual $indid in $mldose (" . scalar (@geno) . ") versus $ndoses snps in the mlinfo file\n";}
		
		$pedinfo->{$famid}{$indid} || die "FATAL ERROR: \nIndividual $famid $indid is listed in DOSE file $mldose, but not annotated in FAM file $fam\n";
 $stop=1;
for ($curent = 1; $curent <= $cuts ; $curent++) {
 our $start=();
 $start = $stop + 1 ;
 $startsnp =$start -1;
 $startcol =$startsnp -1;
 if ((($curent*$batch)+1) <= ($ndoses+1)) { $stop = ($curent*$batch)+1 ;}
 if ((($curent*$batch)+1) > ($ndoses+1)) { $stop = ((($curent-1)*$batch)+$ndoses%$batch)+1 ;}
 $stopsnp =$stop -1;
 $stopcol =$stopsnp -1;
 $file =PED.$curent;
		@slicegeno = @geno[$startcol..$stopcol] ;
		print $file $pedinfo->{$famid}{$indid};

		print $file " @slicegeno\n" ;

}
#		Print out a list of ids for records - Debugging
#		print IDS $pedinfo->{$famid}{$indid};
#		print IDS "\n";

		$countple++;

}
#	close (IDS);
 $stop=1;
for ($curent = 1; $curent <= $cuts ; $curent++) {
 our $start=();
 $start = $stop + 1 ;
 $startsnp =$start -1;
 $startcol =$startsnp -1;
 if ((($curent*$batch)+1) <= ($ndoses+1)) { $stop = ($curent*$batch)+1 ;}
 if ((($curent*$batch)+1) > ($ndoses+1)) { $stop = ((($curent-1)*$batch)+$ndoses%$batch)+1 ;}
 $stopsnp =$stop -1;
 $stopcol =$stopsnp -1;

	close (PED.$curent);	
	print STDERR "Finished writting the ped file $out.$curent.ped.gz (with snps $startsnp to $stopsnp for $. individuals) \n";
#	print STDERR "Finished writting the id list of $. individuals to $out.$curent.idlist.gz  \n ";
}

#####################################################
#             End of  batch loop                    #
}
#####################################################

	close (DOSE);

#####################################################
#             End of  batch option code             #
}
#####################################################

print "\n";

my  $end_time = time();
my  $duration = $end_time - $start_time ;
my  $minutes =0;
my  $remainder=0;
my  $hours=0;
my  $time ;

if($duration > 59) {
    $minutes=int($duration/60);
    $remainder=$duration % 60;
    $time = "$minutes minutes $remainder seconds";
}
else {
     $time="$duration seconds";
}
if($minutes > 59) {
     $hours=int($minutes/60);
     $remainder=$minutes % 60;
     $time = "$hours hours $remainder minutes";
}
print "\n Run time: $time\n";




