#!/usr/local/bin/perl

# Author : Gavin Sherlock
# Date   : April 7th 2004

# This program will read in a gtr file, and a cdt file, and then cut
# the tree at a supplied correlation, outputting one file per node in
# the tree below the cut.

use strict;
use FindBin qw($Bin);
# $Bin is now the location of the script

use Getopt::Std;
use vars qw($opt_i $opt_c $opt_o $opt_p $opt_n);


if (!getopts('i:c:o:p:n:')) {
    &Usage;
}

#my $gtrFile     = $opt_i || &Usage;
my $correlation = $opt_c || &Usage;
my $outDir      = $opt_o;
my $partitionsFile = $opt_p;
my $minNodes = $opt_n;
my $gtrFile = shift;
my $cdtFile = shift;

my $partitionNo = 0;

# before changing our working directory, store the path of the gtr
if ($gtrFile !~ /^\//) {
    $gtrFile = $ENV{"PWD"}. "/" . $gtrFile;
}
#$gtrFile = "${Bin}/${gtrFile}";


unless ($opt_p || $opt_o) {
    print "At least one of -o or -p must be given.\n";
    &Usage;
}

if ($outDir) {
    mkdir($outDir) if (! -d);
    unless(-d $outDir && -w $outDir){
	print "Directory is either non-existent or non-writable.\n";
	&Usage;
#    }else{
#	chdir($outDir);
    }
}

# first parse the tree file

my $gtrTree;

if (-e $gtrFile){
	
    $gtrTree = &parseTreeFile($gtrFile);
	
}else{

    die "$gtrFile does not exist";

}

# now map the gene names to the gene identifiers in the
# cdt file

my %idToName;
my %idToData;
my $header;
my $weights;

&parseCDTFile($gtrFile, $cdtFile);

# now we can traverse the tree, and print out the subclusters

if ($partitionsFile) {
    open(PARTITIONS, ">".$opt_p) || die "cannot create partitions file: $!\n";
    print PARTITIONS "orf\tcluster\n";
}

&traverseTree($gtrTree, $correlation);

if ($partitionsFile) {
    close(PARTITIONS);
}

exit 0;


#####################################################################
sub Usage{
#####################################################################
# This subroutine simply prints out the usage information

    print "Usage:

pruneTree.pl -c <correlation> {-o <output dir>|-p <partitions file>} <gtr file>

pruneTree.pl will read in a gtr file (for which a cdt file must exist)
and, for each subcluster that exceeds the provided correlation, 
print out a PCL and gene list file (with -o) in the output dir and/or
print out a single file containing all genes and a partition number
(with -p). The partitions file is suitable for use with FIRE.\n";

    exit 1;

}

#####################################################################
sub parseTreeFile{
#####################################################################
# This subroutine reads in a tree file (such as a gtr file) and
# creates a structure that can be traversed.  The structure is
# probably not ideal, but it consists of two arrays, one for the
# leaves, and one for the nodes.  The left and right branches of a
# node are stored as references to entries in the relevant array.  For
# a gene, those fields are simply stored as zeros.  The index with
# each array is the name as the number of the node or gene in the gtr
# file, e.g NODE0X is stored in index 0 of the node array.  The
# subroutine then returns a reference to the very top node.

    my ($treeFile) = @_;

    my $gene = 0;

    my ($node, $leftType, $left, $rightType, $right, $corr);

    my (@leaves, @tree);
    
    open (TREEFILE, "$treeFile")||die "Can't open $treeFile:$!\n";
    
    while(<TREEFILE>){
	
	next unless (/NODE(\d+)X\t(ARRY|NODE|GENE)(\d+)X\t(ARRY|NODE|GENE)(\d+)X\t(-?\d\.\d+)/);
	
	$node      = $1;
	$leftType  = $2;
	$left      = $3;
	$rightType = $4;
	$right     = $5;
	$corr      = $6;
	
	if($rightType eq "NODE"){ 

	    $right = \$tree[$right];

	}else{

	    $leaves[$gene]={
		
		LEFT   => 0,
		RIGHT  => 0,
		CORR   => 0,
		NUMBER => $right,
		
	    };

	    $right = \$leaves[$gene];

	    $gene++;

	}
	
	if($leftType eq "NODE"){

	    $left = \$tree[$left];

	}else{

	    $leaves[$gene]={

		LEFT   => 0,
		RIGHT  => 0,
		CORR   => 0,
		NUMBER => $left,
		
	    };

	    $left=\$leaves[$gene];

	    $gene++;

	}
	
	$tree[$node] = {
	    
	    LEFT   => $left,
	    RIGHT  => $right,
	    CORR   => $corr,
	    NUMBER => $node,
	    
	};	
	
    }

    close TREEFILE;

    return (\$tree[$node]);    

}


#####################################################################
sub parseCDTFile{
#####################################################################
# This subroutine simply parses the cdt file, and populates a global
# hash of id to gene name

    my ($gtrFile, $cdtFile) = @_;

    if (!$cdtFile) {
	$gtrFile =~ /(.*\.)gtr$/;
	$cdtFile = $1."cdt";
    }

#    my $col2num = 0;
#    my $col3num = 0;
#    my $nomatch = 0;

    open (IN, $cdtFile) || die "Cannot open $cdtFile : $!";

    while (<IN>){

	chomp;

	die "This cdt file has no GID entry on the first line." if ($. == 1 && $_ !~ /^GID/);

	my @line = split("\t", $_, -1);

	shift (@line);

	if ($. == 1){

	    $header = join("\t", @line);
	    $weights = scalar(@line) - 3; # discarded one column, do 3 more

	}

	next unless /GENE(\d+)X\t(.+?)\t(.+?)\t/; #! shouldn't this be checked earlier?

	my $gene_no = $1;
	my $id = $2;
	my $name = $3;

	$idToName{$gene_no} = $id;

	$idToData{$gene_no} = join("\t", @line);

	$idToName{$gene_no} =~ s/\"//g;

	### added by mark 9/05
#	if ($id =~ /A_\d+_/) {
#	    # agilent id
#	    # for now look in the NAME column
#	    if ($name =~ /(Y.[LR]...[WC])/) {
#		$idToName{$gene_no} = $1;
#		$col3num++;
#	    } else {
#		$nomatch++;
#	    }
#	} else {
#	    $col2num++;
#	}

	if ($idToName{$gene_no} =~ /^\s*$/) {
	    die "FATAL ERROR: can't find gene name in line for $gene_no\n";
	}
    }

#    if ($col2num) {
#	print "Column 2 was used $col2num times for the gene name\n";
#    }
#    if ($col3num) {
#	print "Gene names were found $col3num times in column 3\n";
#    }
#    if ($nomatch) {
#	print "WARNING: In $nomatch instances, column 2 contained an Agilent identifier, but no gene name could be found in column 3. The Agilent identifier was used, which may not be what you want.\n";
#    }

    close IN;

}

#####################################################################
sub traverseTree{
#####################################################################
# This subroutine prints out the subclusters from a tree, such that
# only the largest subcluster that has a correlation that exceeds the
# cutoff is printed, and none of its subclusters are printed, ie there
# is no redundancy.  To do this it simply iterates over the tree in a
# depth first manner, and tests each node to see if its children are
# joined with a correlation greater than the cutoff.  If they are, the
# genes below that node are printed out, and it then returns, without
# traversing below that node.

    my ($tree, $correlation) = @_;

    if ($$tree->{LEFT} == 0 && $$tree->{RIGHT} == 0){

	# we have a gene that wasn't joined with anything,
	# so print it on its own (at least for debugging purposes)

	#print "GENE".$$tree->{NUMBER}."X on it's own\n\n";

    }elsif ($$tree->{CORR} > $correlation){

	# we have a node whose children were joined with a correlation
	# greater than the cutoff, so print all of it's genes

	if ($outDir) {
	    
	}

	my $nodeLines = [];
	my $pclLines = [];
	my $partitionsLines = [];

	&printTree($tree, $nodeLines, $pclLines, $partitionsLines);

	if ((!defined($minNodes)) || (($#$nodeLines + 1 >= $minNodes)))
 {

	    if ($outDir) {

		my $stem = "NODE".$$tree->{NUMBER}."X";

		my $nodefile = $stem.".list";
		my $pclfile  = $stem.".pcl";
		my $cdtfile  = $stem.".cdt";

		open (NODE, ">$outDir/$nodefile") || die "Cannot create $nodefile : $!";
		print NODE join("\n", @{$nodeLines}), "\n";
		close NODE;

		open (PCL,  ">$outDir/$pclfile")  || die "Cannot create $pclfile : $!";
		print PCL $header, "\n", "EWEIGHT\t\t", "\t1" x $weights, "\n";
		print PCL join("\n", @{$pclLines}), "\n";
		close PCL;
	    }

	    if ($partitionsFile) {

		print PARTITIONS join("\n", @{$partitionsLines}), "\n";

		$partitionNo++;
	    }

	}

	# john is commenting this out, as some files seg-fault it
	# anyway
#	my $working = 0;
#	if($working){
#	    # now cluster the pcl
#
#	    `$Bin/XCluster/cluster -f $pclfile -e 1`;
#
#	    # now make a png image
#
#	    `$Bin/makeCluster2.pl $cdtfile`;
#	}
#
#	`/Genomics/GO/bin/GOTermFinder.pl -a P -g gene_association.sgd -d /Genomics/GO/data/unparsed/ -o /Genomics/GO/data/parsed/ -t ${Bin}/$outDir $nodefile`

    }else{

	# we have a node whose children's joining correlation did not
	# exceed the threshold, so just keep processing it

	&traverseTree($$tree->{LEFT}, $correlation);
	&traverseTree($$tree->{RIGHT}, $correlation);

    }
		      
}

#####################################################################
sub printTree{
#####################################################################
# This subroutine prints out all the genes in a tree

    my ($tree, $nodeLines, $pclLines, $partitionsLines) = @_;

    if($$tree->{LEFT} == 0 && $$tree->{RIGHT} == 0){

	# we have a gene, so print it

	if ($outDir) {
#	    print NODE $idToName{$$tree->{NUMBER}}, "\n";
#	    print PCL  $idToData{$$tree->{NUMBER}}, "\n";
	    push(@{$nodeLines}, $idToName{$$tree->{NUMBER}});
	    push(@{$pclLines}, $idToData{$$tree->{NUMBER}});
	}
	if ($partitionsFile) {
#	    print PARTITIONS $idToName{$$tree->{NUMBER}}, "\t", $partitionNo, "\n";
	    push(@{$partitionsLines}, $idToName{$$tree->{NUMBER}}."\t".$partitionNo);
	}

    }else{
	
	# we have a node, so keep processing

	&printTree($$tree->{LEFT}, $nodeLines, $pclLines, $partitionsLines);
	&printTree($$tree->{RIGHT}, $nodeLines, $pclLines, $partitionsLines);

    }
		      
}
