#!/usr/bin/env perl

# Author: NTM

#
# Copyright (C) Nicolas Thierry-Mieg, 2009.
#
#
# This is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This script is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this script; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA


# UPDATE 10/03/08: adapted this script to read the Empty/Small wells
# that we noticed when rebuilding the WASP2 and WASP2_384 arrays 
# for the Finishing dataset (WASP2_empty_wells_2008Feb.txt).
# All Empty and Small wells are marked as FAINT.

# takes a single arg: a dir containing a set of 
# colony3 data files corresponding to results 
# with the WASP6 (in 384 format), WASP2 (in 1536), or
# WASP2-384 (in 384) pools.
# NOTE: dir cannot contain 2 types (WASP2, WASP2-384, WASP6),
# you must have them in distinct subdirs and call this
# script once on each subdir.
# Read all files in this dir, create a subdir 
# of Sigs/ in the working dir for each different bait,
# and make an interpool sig file for each
# WASP batch.
# If some colony3 files are missing for a bait,
# any batch that is entirely in the missing files
# is ignored (no sig file is generated), and for
# batches that have some pools in the missing files
# we set these pools' values to FAINT.

use strict ;
use warnings ;




###############################################################################

# mappings between WASP plates+wells and batch+poolnums is
# in $waspCtoPfile* (generated by makeBatchMappings.*.pl).
# format is: "<plate>,<col>,<row>,<batch>,<poolnum>"
# we will only know whether this is WASP2 or WASP6 later.
# So, we make both arrays, and we will save the correct
# one to @plate2pool when we know the format.
# UPDATE 07/01/22: Xiaofeng rebuilt the WASP6 384 plates but switched
# B and C spots in plate 6. So, we use the original coords2pools.WASP6 
# mapping for the 4 baits done with the initial plates (ie T326, T380,
# T388 and T448), and default to the new mapping for other baits.
# UPDATE 07/07/09: now dealing with the screens for the paper; Xiaofeng
# rebuilt the WASP6 array for this, and this time he started from the
# stock in 384 plates, so the B and C spots are back to normal. However
# baits for the paper were previously screened in the "Next12" batch.
# Therefore I can't discriminate on bait name as I did before.
# I'm using the series: Next12 was always series <= 4, while forPaper is 5-8.
# UPDATE 07/08/10: to deal with WASP2-384 data.
my @plate2pool ;


# first, WASP6 original mapping
my $waspCtoPfile6 = "Mappings/coords2pools.WASP6.original" ;
open(C2P6, "$waspCtoPfile6") || 
    die "cannot open $waspCtoPfile6 for reading\n" ;
# parse this file and save the data in plate2pool6orig: "$batch,$poolnum" is 
# stored in $plate2pool6orig[($P-1)*384 + ($C-1)*16 + ($R-1)]
my @plate2pool6orig ;
# skip first 2 lines (header)
<C2P6> ;
<C2P6> ;
while(my $line =<C2P6>)
{
    chomp $line ;
    ($line =~ /^(\d+),(\d+),(\d+),(\d+),(\d+)$/) ||
	die "cannot parse line in $waspCtoPfile6:\n$line\n" ;
    my ($P,$C,$R,$B,$poolnum) = ($1,$2,$3,$4,$5) ;
    $plate2pool6orig[($P-1)*384 + ($C-1)*16 + ($R-1)] = "$B,$poolnum" ;
}
close(C2P6) ;

# next, WASP6 new (B and C switched) mapping
$waspCtoPfile6 = "Mappings/coords2pools.WASP6.BandCswitched" ;
open(C2P6, "$waspCtoPfile6") || 
    die "cannot open $waspCtoPfile6 for reading\n" ;
# parse this file and save the data in plate2pool6switched: "$batch,$poolnum" is 
# stored in $plate2pool6switched[($P-1)*384 + ($C-1)*16 + ($R-1)]
my @plate2pool6switched ;
# skip first 2 lines (header)
<C2P6> ;
<C2P6> ;
while(my $line =<C2P6>)
{
    chomp $line ;
    ($line =~ /^(\d+),(\d+),(\d+),(\d+),(\d+)$/) ||
	die "cannot parse line in $waspCtoPfile6:\n$line\n" ;
    my ($P,$C,$R,$B,$poolnum) = ($1,$2,$3,$4,$5) ;
    $plate2pool6switched[($P-1)*384 + ($C-1)*16 + ($R-1)] = "$B,$poolnum" ;
}
close(C2P6) ;


# now, WASP2x16
my $waspCtoPfile2 = "Mappings/coords2pools.WASP2x16" ;
open(C2P2, "$waspCtoPfile2") || 
    die "cannot open $waspCtoPfile2 for reading\n" ;
# parse this file and save the data in plate2pool2: "$batch,$poolnum" is 
# stored in $plate2pool2[($P-1)*1536 + ($C-1)*32 + ($R-1)]
my @plate2pool2 ;
# skip first 2 lines (header)
<C2P2> ;
<C2P2> ;
while(my $line =<C2P2>)
{
    chomp $line ;
    ($line =~ /^(\d+),(\d+),(\d+),(\d+),(\d+)$/) ||
	die "cannot parse line in $waspCtoPfile2:\n$line\n" ;
    my ($P,$C,$R,$B,$poolnum) = ($1,$2,$3,$4,$5) ;
    $plate2pool2[($P-1)*1536 + ($C-1)*32 + ($R-1)] = "$B,$poolnum" ;
}
close(C2P2) ;


# and finally, WASP2-384 which is WASP2x4
$waspCtoPfile2 = "Mappings/coords2pools.WASP2x4" ;
open(C2P2, "$waspCtoPfile2") || 
    die "cannot open $waspCtoPfile2 for reading\n" ;
# parse this file and save the data in plate2pool2_384: "$batch,$poolnum" is 
# stored in $plate2pool2_384[($P-1)*384 + ($C-1)*16 + ($R-1)]
my @plate2pool2_384 ;
# skip first 2 lines (header)
<C2P2> ;
<C2P2> ;
while(my $line =<C2P2>)
{
    chomp $line ;
    ($line =~ /^(\d+),(\d+),(\d+),(\d+),(\d+)$/) ||
	die "cannot parse line in $waspCtoPfile2:\n$line\n" ;
    my ($P,$C,$R,$B,$poolnum) = ($1,$2,$3,$4,$5) ;
    $plate2pool2_384[($P-1)*384 + ($C-1)*16 + ($R-1)] = "$B,$poolnum" ;
}
close(C2P2) ;


###############################################################################

# parse empties/small

my $emptyFile = "Mappings/WASP2_empty_wells_2008Feb.txt" ;
open(EMPTY, "$emptyFile") || die "cannot open empties $emptyFile\n" ;

# key is "$B,$poolnum", value is 1 for Empty wells and 2 for Small spots
my %empty = () ;

# skip first line (headers)
<EMPTY> ;

while(my $line = <EMPTY>)
{
    chomp $line ;

    # skip empty lines
    ($line =~ /^\s*$/) && next ;

    ($line =~ /^(\d+)\s+(\d+)\s+(\d+)\s+([ES])$/) ||
	die "cannot parse line in empties:\n$line\n" ;
    my ($p,$r,$c,$status) = ($1,$2,$3,$4) ;

    # WASP2_empty_wells_2008Feb.txt uses WASP2_384 coordinates although
    # it applies both to WASP2_384 and to WASP2.
    my ($emptySpot) = $plate2pool2_384[($p-1)*384 + ($c-1)*16 + ($r-1)] ;

    if ($status eq "E") { $empty{"$emptySpot"} = 1 ; }
    elsif ($status eq "S") { $empty{"$emptySpot"} = 2 ; }
    else { die "status illegal: $status\n" ; }
}


###############################################################################
# we need a numerical sort at some point
sub NUMERNIC { $a <=> $b ; }




###############################################################################

(@ARGV != 1) && 
    die "wrong number of args: requires a single arg, a dir holding colony3 files\n" ;
my $colonyDir = $ARGV[0] ;
opendir(COLDIR, "$colonyDir") ||
    die "cannot opendir $colonyDir, does this path exist?\n" ;
my @allFiles = grep(/\.dat$/, readdir(COLDIR)) ;
closedir(COLDIR) ;

# 07/08/10: we need a new way of distinguishing between WASP2-384 and WASP6.
# using $type, based on the number of colony3 files.
my $type = "" ;
if (@allFiles == 20)
{
    $type = "W2_384" ;
}
elsif (@allFiles == 6)
{
    $type = "W6" ;
}
elsif (@allFiles <= 5)
{
    # <= because some screens had only the first 2 plates
    $type = "W2" ;
}
else
{
    die "can't find type based on number of files:\n@allFiles\n\n" ;
}


my @baits = @allFiles ;
foreach my $i (0..$#baits)
{
    ($baits[$i] =~ s/_00\d\d\d\.dat$//) ||
	die "cannot remove tail from file $baits[$i]\n" ;
}
# sort and remove identical entries
@baits = sort(@baits) ;
{
    my @goodBaits = ("dummy") ;
    foreach my $bait (@baits)
    {
	($bait ne $goodBaits[$#goodBaits]) &&
	    push(@goodBaits,$bait) ;
    }
    ($#goodBaits == 0) && die "no good baits??\n" ;
    # remove dummy value
    shift(@goodBaits) ;
    @baits = @goodBaits ;
}

# make Sigs/ dir if needed
my $sigDir = "Sigs/" ;
(-d $sigDir) || 
    (mkdir($sigDir) && warn "created dir $sigDir\n") || 
    die "cannot make dir $sigDir\n" ;

# format will be either 384 or 1536, or 3840 for WASP2-384.
# Will be read from first file, then we check that
# every other file is in the same format.
my $format ;
my ($nbRows,$nbCols) ;

# OK, process each bait
foreach my $bait (@baits)
{
    # if this is the first file, format and (if WASP6 or WASP2_384) mapping will be
    # found later, and @plate2pool will be set correctly.
    # otherwise: if format==384, choose correct mapping based
    # on bait name and series
    if ((defined $format) && ($format == 384))
    {
	# sanity: 
	($type eq "W6") || (die "type is $type, should be W6 from format\n") ;
	# find series
	($bait =~ /^T(\d\d\d)[-_]s(\d+)[-_]/) || (die "\n\nERROR: cannot find baitnum and series in $bait\n\n") ;
	my ($baitnum,$ser) = ($1, $2) ;
	if ( ($ser <= 4) &&
	     ( ($baitnum == 196) || ($baitnum == 330) || ($baitnum == 370) || ($baitnum == 381) || 
	       ($baitnum == 382) || ($baitnum == 385) || ($baitnum == 386) || ($baitnum == 387) || 
	       ($baitnum == 392) || ($baitnum == 469) || ($baitnum == 474) || ($baitnum == 479) ) )
	{
	    warn "using WASP6 mapping with switched B and C in plate 6 for $bait\n" ;
	    @plate2pool = @plate2pool6switched ;
	}
	else
	{
	    warn "using original WASP6 mapping for $bait\n" ;
	    @plate2pool = @plate2pool6orig ;
	}
    }

    # results for this bait will be stored as:
    my (@strongs, @weaks, @faints, @nones) ;
    # eg $strongs[$B] is the list of poolnums of strong pos 
    # pools in batch $B, space-separated

    # find colony files for this bait
    my @goodFiles = grep(/$bait/, @allFiles) ;

    # remember plates that were found, so we can fill
    # missing plates with a default value of "faint"
    # index is plate num, value is 1 if plate was found.
    my @platesDone = () ;
    # also remember batches that have at least one
    # pool in a goodFile. Other batches don't have any data,
    # so we won't output any sig files for them.
    # index is batch num, value is 1 if batch was seen.
    my @batchesSeen = () ;

    foreach my $infile (@goodFiles)
    {
	open(COLFILE, "$colonyDir/$infile") ||
	    die "cannot read $colonyDir/$infile\n" ;

	($infile =~ /_(00\d\d\d)\.dat$/) ||
	    die "cannot find plate number in $infile\n" ;
	# + 0 to convert to number
	my $P = $1 + 0 ;

	$platesDone[$P] = 1 ;

	# parse file: begin by skipping 9 lines
	foreach my $i (1..9)
	{
	    <COLFILE> ;
	}
	# tenth line says if this is 384 or 1536
	{
	    my $line = <COLFILE> ;
	    #chop after chomp since these are in M$ format
	    chomp $line ; 
	    ((chop $line) ne "\r") &&
		die "chopping doesn't remove CR in $line\n" ;
	    if ($line =~ /^24 16 384$/)
	    {
		if (defined $format)
		{
		    ($format != 384) && ($format != 3840) &&
			die "$infile is 384 format, but previous files in $colonyDir/ were WASP2. Dying.\n" ;

		}
		else
		{
		    $format = 384 ;
		    $nbRows = 16 ;
		    $nbCols = 24 ;

		    if ($type eq "W6")
		    {
			# 07/01/22: use original WASP6 mapping for first 4 baits, switched mapping for others - NO:
			# 07/07/09: use switched only for series <= 4 and the Next12 baits.
			# find series
			($bait =~ /^T(\d\d\d)[-_]s(\d+)[-_]/) || (die "\n\nERROR: cannot find baitnum and series in $bait\n\n") ;
			my ($baitnum,$ser) = ($1, $2) ;
			if ( ($ser <= 4) &&
			     ( ($baitnum == 196) || ($baitnum == 330) || ($baitnum == 370) || ($baitnum == 381) || 
			       ($baitnum == 382) || ($baitnum == 385) || ($baitnum == 386) || ($baitnum == 387) || 
			       ($baitnum == 392) || ($baitnum == 469) || ($baitnum == 474) || ($baitnum == 479) ) )
			{
			    warn "using WASP6 mapping with switched B and C in plate 6 for $bait\n" ;
			    @plate2pool = @plate2pool6switched ;
			}
			else
			{
			    warn "using original WASP6 mapping for $bait\n" ;
			    @plate2pool = @plate2pool6orig ;
			}
		    }
		    elsif ($type eq "W2_384")
		    {
			$format *= 10 ;
			@plate2pool = @plate2pool2_384 ;
		    }
		    else
		    {
			die "type is $type, should be W6 or W2_384 from format\n" ;
		    }
		}
	    }
	    elsif ($line =~ /^48 32 1536$/)
	    {
		if (defined $format)
		{
		    ($format != 1536) && 
			die "$infile is WASP2, but previous files in $colonyDir/ were WASP6. Dying.\n" ;
		}
		else
		{
		    $format = 1536 ;
		    $nbRows = 32 ;
		    $nbCols = 48 ;
		    @plate2pool = @plate2pool2 ;
		}
	    }
	    else
	    {
		die "looking for format in $infile, can't parse line: :$line:\n" ;
	    }
	}

	# skip 2 more lines, then check third
	<COLFILE> ;
	<COLFILE> ;
	{
	    my $line = <COLFILE> ;
	    chomp $line ;
	    ((chop $line) ne "\r") &&
		die "chopping doesn't remove CR in $line\n" ;
	    ($line =~ /^rows  columns  size  circularity: $/) ||
		die "failed when checking last header line of $infile:\n$line\n" ;
	}

	# now comes the data, begin loop
	while(my $line = <COLFILE>)
	{
	    chomp $line ;
	    ((chop $line) ne "\r") &&
		die "chopping doesn't remove CR in $line\n" ;
	    ($line =~ /^\s+(\d+)\s+(\d+)\s+\d+\s+(-?[\d\.]+)\s*$/) ||
		die "cannot parse line from $infile:\n$line\n" ;
	    my ($R,$C,$circ) = ($1,$2,$3) ;

	    # skip line if it's an empty well
	    (defined $plate2pool[($P-1)*$nbRows*$nbCols + ($C-1)*$nbRows + ($R-1)]) ||
		next ;

	    my ($B,$poolnum) = 
		split(/,/, $plate2pool[($P-1)*$nbRows*$nbCols + ($C-1)*$nbRows + ($R-1)]) ;

	    ((defined $batchesSeen[$B]) && ($batchesSeen[$B]==1)) ||
		($batchesSeen[$B]=1) ;

	    # deal with empties straight away
	    # these occur only in Fin data, ie s21 and s22
	    
	    if (($bait =~ /[-_]s2[12][-_]/) && ($empty{"$B,$poolnum"}))
	    {
		# set the empty pools as FAINT
		if (defined $faints[$B])
		{
		    $faints[$B] .= "$poolnum " ;
		}
		else
		{
		    $faints[$B] = "$poolnum " ;
		}
	    }

	    elsif ($circ > 55)
	    {
		# strong pos
		if (defined $strongs[$B])
		{
		    $strongs[$B] .= "$poolnum " ;
		}
		else
		{
		    $strongs[$B] = "$poolnum " ;
		}
	    }
	    elsif ($circ > 35)
	    {
		# weak pos
		if (defined $weaks[$B])
		{
		    $weaks[$B] .= "$poolnum " ;
		}
		else
		{
		    $weaks[$B] = "$poolnum " ;
		}
	    }
	    elsif ($circ > 15)
	    {
		# faint neg
		if (defined $faints[$B])
		{
		    $faints[$B] .= "$poolnum " ;
		}
		else
		{
		    $faints[$B] = "$poolnum " ;
		}
	    }
	    elsif ($circ > -5)
	    {
		# totally neg
		if (defined $nones[$B])
		{
		    $nones[$B] .= "$poolnum " ;
		}
		else
		{
		    $nones[$B] = "$poolnum " ;
		}
	    }
	    else
	    {
		die "dying, circ is strange: $circ\n" ;
	    }
	}
	close(COLFILE) ;
    }

    # Deal with missing files: fill values with "faint"
    my $nbPlates ;
    if ($format==384) { $nbPlates = 6 ; }
    elsif ($format==1536) { $nbPlates = 5 ; }
    elsif ($format==3840) { $nbPlates = 20 ; }
    else { die "when dealing with missing files, format incorrect: $format\n" ; }

    foreach my $P (1..$nbPlates)
    {
	if ((defined $platesDone[$P]) && ($platesDone[$P]==1))
	{
	    next ;
	}
	else
	{
	    foreach my $R (1..$nbRows)
	    {
		foreach my $C (1..$nbCols)
		{
		    # skip line if it's an empty well
		    (defined $plate2pool[($P-1)*$nbRows*$nbCols + ($C-1)*$nbRows + ($R-1)]) ||
			next ;

		    my ($B,$poolnum) = 
			split(/,/, $plate2pool[($P-1)*$nbRows*$nbCols + ($C-1)*$nbRows + ($R-1)]) ;

		    if (defined $faints[$B])
		    {
			$faints[$B] .= "$poolnum " ;
		    }
		    else
		    {
			$faints[$B] = "$poolnum " ;
		    }
		}
	    }
	}
    }


    # OK, should be all done for this bait

    # sanity: there should be 169 pools per batch, except
    # for some WASP2 and WASP2_384 special batches:
    # 5-8 should have 263 pools, 9-10 should have 188,
    # and 37-38 should have 244.
    # WASP6 has 13 batches, WASP2 has 40.
    my $nbBatches = 13 ;
    (($format == 1536) || ($format == 3840)) && ($nbBatches = 40) ;
    foreach my $B (1..$nbBatches)
    {
	my $allPools = "";
	(defined $strongs[$B]) && ($allPools .= $strongs[$B]) ;
	(defined $weaks[$B]) && ($allPools .= $weaks[$B]) ;
	(defined $faints[$B]) && ($allPools .= $faints[$B]) ;
	(defined $nones[$B]) && ($allPools .= $nones[$B]) ;
	my @poolsTmp = split(/ /,$allPools) ;
	my $numPools = scalar(@poolsTmp) ;
	if (($format==1536) || ($format == 3840))
	{
	    if (($B >= 5) && ($B <= 8))
	    {
		($numPools != 263) &&
		    die "For bait $bait, wrong number of pools $numPools for batch $B\n" ;
	    }
	    elsif (($B >= 9) && ($B <= 10))
	    {
		($numPools != 188) &&
		    die "For bait $bait, wrong number of pools $numPools for batch $B\n" ;
	    }
	    elsif (($B >= 37) && ($B <= 38))
	    {
		($numPools != 244) &&
		    die "For bait $bait, wrong number of pools $numPools for batch $B\n" ;
	    }
	    else
	    {
		($numPools != 169) &&
		    die "For bait $bait, wrong number of pools $numPools for batch $B\n" ;
	    }
	}
	else
	{
	    # WASP6, all batches should have 169
	    ($numPools != 169) &&
		    die "For bait $bait, wrong number of pools $numPools for batch $B\n" ;
	}
    }
    

    # build $date, yy-mm-dd format
    my $date = "" ;
    {
	my @fulldate = localtime(time) ;
	my $day = $fulldate[3] ;
	($day < 10) && ($day = "0$day") ;
	# month starts at 0, adjust and add leading 0 if needed
	my $month = $fulldate[4]+1 ;
	($month < 10) && ($month = "0$month") ;
	# year starts at 1900
	my $year = $fulldate[5] - 100 ;
	($year < 10) && ($year = "0$year") ;
	$date = "$year-$month-$day" ;
    }

    # create dir for this bait
    my $waspType = "WASP6" ;
    ($format == 1536) && ($waspType = "WASP2") ;
    ($format == 3840) && ($waspType = "WASP2_384") ;
    my $baitDir = "$sigDir/$bait.$waspType/" ;
    (-e "$baitDir") && 
	(warn "dir (or file?) $baitDir already exists, skipping it\n") &&
	next ;
    (mkdir("$baitDir") && (warn "Created $baitDir\n")) || 
	die "cannot mkdir $baitDir\n" ;

    # build interpool files, named "$bait.WASPx.batch<$B>.sig"
    foreach my $B (1..$nbBatches)
    {
	# don't make any files for batches entirely in missing plates
	((defined $batchesSeen[$B]) && ($batchesSeen[$B]==1)) ||
	    next ;
	my $sigfile = "$baitDir/$bait.$waspType.batch$B.sig" ;
	(-e $sigfile) && die "$sigfile exists, impossible\n" ;
	open(SIGFILE, ">$sigfile") ||
	    die "cannot write to $sigfile\n" ;

	print SIGFILE "<signature version=\"2.0\">\n" ;
	print SIGFILE "   <header>\n" ;
	print SIGFILE "      <description>\n" ;
	print SIGFILE "         signature file for bait $bait, $waspType pools, batch $B.\n" ;
	print SIGFILE "         Generated by colonyToSig.pl\n" ;
	print SIGFILE "      </description>\n" ;
	print SIGFILE "      <date>\n" ;
	print SIGFILE "         $date\n" ;
	print SIGFILE "      </date>\n" ;
	print SIGFILE "   </header>\n" ;
	print SIGFILE "   <values>\n" ;
	print SIGFILE "      <STRONG>\n" ;
	if (defined $strongs[$B])
	{
	    # print strongs, sorted by poolnum
	    print SIGFILE "         " ;
	    my @sortedTmp = sort NUMERNIC split(/ /, $strongs[$B]) ;
	    print SIGFILE join(' ', @sortedTmp) ;
	    print SIGFILE "\n" ;
	}
	print SIGFILE "      </STRONG>\n" ;
	print SIGFILE "      <WEAK>\n" ;
	if (defined $weaks[$B])
	{
	    print SIGFILE "         " ;
	    my @sortedTmp = sort NUMERNIC split(/ /, $weaks[$B]) ;
	    print SIGFILE join(' ', @sortedTmp) ;
	    print SIGFILE "\n" ;
	}
	print SIGFILE "      </WEAK>\n" ;
	print SIGFILE "      <FAINT>\n" ;
	if (defined $faints[$B])
	{
	    print SIGFILE "         " ;
	    my @sortedTmp = sort NUMERNIC split(/ /, $faints[$B]) ;
	    print SIGFILE join(' ', @sortedTmp) ;
	    print SIGFILE "\n" ;
	}
	print SIGFILE "      </FAINT>\n" ;
	print SIGFILE "      <NEG>\n" ;
	if (defined $nones[$B])
	{
	    print SIGFILE "         " ;
	    my @sortedTmp = sort NUMERNIC split(/ /, $nones[$B]) ;
	    print SIGFILE join(' ', @sortedTmp) ;
	    print SIGFILE "\n" ;
	}
	print SIGFILE "      </NEG>\n" ;
	print SIGFILE "   </values>\n" ;
	print SIGFILE "</signature>\n" ;

	close(SIGFILE) ;
    }
}


