#!/usr/bin/perl -w

use lib qw(./);
use strict;
use warnings;
use Data::Dumper;
use Switch;
use Benchmark;
use Getopt::Long;
#use global qw(:GetFH);

# User modules:
use hypergeom_allPerl;
#use Stat::multtest;

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

# Change these paths if needed
my $gff_all_genes_file = $ENV{"DATAFILE"} || "/disc8/data4projects/annot_dmel/dmel_mrna_names_gp_v4.gff" || die "Not defined the refseq.gff file\n";

print 

## Chromosome length
my $length_2L = 22407834;
my $length_2R = 20766785;
my $length_3L = 23771897;
my $length_3R = 27905053;
my $length_4  = 1281640;
my $length_X  = 22224390;


###########################################
## Options

my ($chr2L_ref, $chr2R_ref, $chr3L_ref, $chr3R_ref, $chr4_ref, $chrX_ref);
my $regulated_file;
my $array_file = "/home/miguel/projects/ash2_arrays/pos_chrom/cluster_genomic/data/all_genes_in_refseq_v4.refseqs";

#"~/projects/ash2_arrays/pos_chrom/cluster_genomic/data/all_genes_in_refseq_v4.refseqs";
my $window=30000;
my $gwindow;
my $offset=10000;
my $pval_lim = 0.005;
my $min_genes_in_cluster = 3;
my $help = 0;
my $options = GetOptions ("reg=s"         => \$regulated_file,
			  "ref:s"         => \$array_file,
			  "window:i"      => \$window,
			  "gwindow:i"     => \$gwindow,
			  "offset:i"      => \$offset,
			  "pval:f"        => \$pval_lim,
			  "min_genes:i"   => \$min_genes_in_cluster,
			  "help"          => \$help);

if ( (!defined $regulated_file) || ($help) ){
    print <<EOF;
cluster_genomic.pl [options] --reg <refseq_file>
Options:
    --ref [file]          : File containing a reference list of refseqs. If not present all annotated refseqs are used
    --window [integer]    : Width of the length to calculate the clusters (Default: 30000)
    --offset [integer]    : Offset between windows (Default: 10000)
    --gwindow [integer]   : Defines the window by number of genes instead of by number of nucleotides
    --pval [float]        : Max allowed pvalue per cluster (Default: 0.005)
    --min_genes [integer] : Min number of genes that can for a cluster
    --help                : Displays this document
EOF
exit(0);
}

if (!defined $gwindow){
    print "### Options:\n# Reg_file:$regulated_file\n# Ref_file:$array_file\n# window_size=$window\n# offset=$offset\n# pval=$pval_lim\n# min_genes=$min_genes_in_cluster\n\n";
} else {
print "# Reg_file:$regulated_file\n# Ref_file:$array_file\n# window_size=$gwindow genes\n# pval=$pval_lim\n# min_genes=$min_genes_in_cluster\n\n";
}

($chr2L_ref, $chr2R_ref, $chr3L_ref, $chr3R_ref, $chr4_ref, $chrX_ref) = 
    &get_all_genes_by_chrs ($gff_all_genes_file);

#my $arrayRegGenes_ref = &extract_genes ($regulated_file);
my $arrayAllGenes_ref = &extract_genes ($array_file);
my $arrayRegGenes_ref = &extract_genes ($regulated_file);

# Dumper code
#   my $msg = Dumper ($arrayAllGenes_ref);
#    print $msg."\n\n";
#   exit();

#my $n_reg_genes = &is_regulated ($arrayAllGenes_ref, $rarrayRegGenes_ref);

my ($n_reg_genes, $n_all_genes) = 
    (scalar @{$arrayRegGenes_ref},scalar @{$arrayAllGenes_ref});

print STDERR "\nLocating genes in chr 2L...";
my ($chr2L_all_ref,$n_genes_2L_ref, $n_genes_2L_reg) = &slice_hash ($arrayAllGenes_ref, $arrayRegGenes_ref, $chr2L_ref);
print STDERR "($n_genes_2L_ref/$n_all_genes (all), $n_genes_2L_reg/$n_reg_genes (reg)) Done\n";

print STDERR "Locating genes in chr 2R...";
my ($chr2R_all_ref,$n_genes_2R_ref, $n_genes_2R_reg) = &slice_hash ($arrayAllGenes_ref, $arrayRegGenes_ref, $chr2R_ref);
print STDERR "($n_genes_2R_ref/$n_all_genes (all), $n_genes_2R_reg/$n_reg_genes (reg)) Done\n";

print STDERR "Locating genes in chr 3L...";
my ($chr3L_all_ref,$n_genes_3L_ref, $n_genes_3L_reg) = &slice_hash ($arrayAllGenes_ref, $arrayRegGenes_ref, $chr3L_ref);
print STDERR "($n_genes_3L_ref/$n_all_genes (all), $n_genes_3L_reg/$n_reg_genes (reg)) Done\n";

print STDERR "Locating genes in chr 3R...";
my ($chr3R_all_ref,$n_genes_3R_ref, $n_genes_3R_reg) = &slice_hash ($arrayAllGenes_ref, $arrayRegGenes_ref, $chr3R_ref);
print STDERR "($n_genes_3R_ref/$n_all_genes (all), $n_genes_3R_reg/$n_reg_genes (reg)) Done\n";

print STDERR "Locating genes in chr X...";
my ($chrX_all_ref,$n_genes_X_ref, $n_genes_X_reg) = &slice_hash ($arrayAllGenes_ref, $arrayRegGenes_ref, $chrX_ref);
print STDERR "($n_genes_X_ref/$n_all_genes (all), $n_genes_X_reg/$n_reg_genes (reg)) Done\n";

print STDERR "Locating genes in chr 4...";
my ($chr4_all_ref,$n_genes_4_ref, $n_genes_4_reg) = &slice_hash ($arrayAllGenes_ref, $arrayRegGenes_ref, $chr4_ref);
print STDERR "($n_genes_4_ref/$n_all_genes (all), $n_genes_4_reg/$n_reg_genes (reg)) Done\n";

my ($cluster2L, $cluster2R, $cluster3L, $cluster3R, $cluster4, $clusterX);
if (defined $gwindow){
    print STDERR "** Finding clusters in a window of $gwindow genes **\n";
    $cluster2L = &gene_window_clusters ($chr2L_all_ref, $gwindow, $n_genes_2L_ref, $n_genes_2L_reg, "chr2L");
    $cluster2R = &gene_window_clusters ($chr2R_all_ref, $gwindow, $n_genes_2R_ref, $n_genes_2R_reg, "chr2R");
    $cluster3L = &gene_window_clusters ($chr3L_all_ref, $gwindow, $n_genes_3L_ref, $n_genes_3L_reg, "chr3L");
    $cluster3R = &gene_window_clusters ($chr3R_all_ref, $gwindow, $n_genes_3R_ref, $n_genes_3R_reg, "chr3R");
    $cluster4  = &gene_window_clusters ($chr4_all_ref, $gwindow, $n_genes_4_ref, $n_genes_4_reg, "chr4");
    $clusterX  = &gene_window_clusters ($chrX_all_ref, $gwindow, $n_genes_X_ref, $n_genes_X_reg, "chrX");
}
else {
    print STDERR "** Finding clusters in a window of $window Kb with an offset of $offset Kb **\n";
    $cluster2L = &cluster_finder ($chr2L_all_ref, $length_2L, $n_genes_2L_ref, $n_genes_2L_reg, $window, $offset, "chr2L");
    $cluster2R = &cluster_finder ($chr2R_all_ref, $length_2R, $n_genes_2R_ref, $n_genes_2R_reg, $window, $offset, "chr2R");
    $cluster3L = &cluster_finder ($chr3L_all_ref, $length_3L, $n_genes_3L_ref, $n_genes_3L_reg, $window, $offset, "chr3L");
    $cluster3R = &cluster_finder ($chr3R_all_ref, $length_3R, $n_genes_3R_ref, $n_genes_3R_reg, $window, $offset, "chr3R");
    $cluster4 = &cluster_finder  ($chr4_all_ref, $length_4, $n_genes_4_ref, $n_genes_4_reg, $window, $offset, "chr4");
    $clusterX = &cluster_finder  ($chrX_all_ref, $length_X, $n_genes_X_ref, $n_genes_X_reg, $window, $offset, "chrX");
}
# Dumper code
#    my $msg = Dumper ($chr2L_all_ref);
#     print $msg."\n\n";
#    exit();

    print "### Processed genes:\n# total:$n_all_genes\n# reg:$n_reg_genes\n# chr2L:$n_genes_2L_reg/$n_genes_2L_ref\n# chr2R:$n_genes_2R_reg/$n_genes_2R_ref\n# chr3L:$n_genes_3L_reg/$n_genes_3L_ref\n# chr3R:$n_genes_3R_reg/$n_genes_3R_ref\n# chr4:$n_genes_4_reg/$n_genes_4_ref\n# chrX:$n_genes_X_reg/$n_genes_X_ref\n\n";

foreach my $cluster_aref ($cluster2L, $cluster2R, $cluster3L, $cluster3R, $cluster4, $clusterX){
   &print_clusters ($cluster_aref);
}

### Preparar los datos para imprimir los genes y los clusters en formato GFF

# foreach my $cluster_aref ($cluster2L, $cluster2R, $cluster3L, $cluster3R, $cluster4, $clusterX){
#     &print_refGene ;
# }

########## SUBROUTINES ##########
sub get_all_genes_by_chrs
## Extracts all the drosophila refseqs and put them in a hash... 1 hash per chr
{
    print STDERR "Getting all the Drosophila RefSeqs...";
    my $f_all = shift @_;
    open ALL, $f_all or die "Can't find $f_all -> $!";
## One hash per chromosome
    my %chr2l;
    my %chr2r;
    my %chr3l;
    my %chr3r;
    my %chr4;
    my %chrx;

    while (my $gff_line = <ALL>){
	chomp $gff_line;
	my ($chr, $gpgff, $name, $begin, $end, $dot1, $strand, $dot2, $rest) = 
	    split("\t", $gff_line);
	my ($gene, @rest) = split (/\./, $rest);
	my %gene_hash =();
	
	if ($strand eq "+"){
	    $gene_hash{"begin"} = $begin;
	    $gene_hash{"end"} = $end;
	} else {
	    $gene_hash{"begin"} = $end;
	    $gene_hash{"end"} = $begin;
	}
	$gene_hash{"strand"} = $strand;
	$gene_hash{"name"} = $name;
	$gene_hash{"is_reg"} = 0;

	switch ($chr){
	    case "chr2L" { $chr2l{$gene} = {%gene_hash} }
	    case "chr2R" { $chr2r{$gene} = {%gene_hash} }
	    case "chr3L" { $chr3l{$gene} = {%gene_hash} }
	    case "chr3R" { $chr3r{$gene} = {%gene_hash} }
	    case "chr4"  { $chr4{$gene} = {%gene_hash} }
	    case "chrX"  { $chrx{$gene} = {%gene_hash} }
	}
    }
    print STDERR "Done\n";

    close (ALL);
    return (\%chr2l, \%chr2r, \%chr3l, \%chr3r, \%chr4, \%chrx);
}

sub print_hashref
{
    my $hashref = shift;

    foreach my $refseq (keys %{$hashref}){
	print "$refseq\t";
	foreach my $field (keys %{${$hashref}{$refseq}}) {
	    print ${${$hashref}{$refseq}}{$field}."\t";
	}
	print "\n";
    }
}

sub link_genes
{
    my ($hash_ref, $ref_file) = @_;
    print STDERR "Giving names to the RefSeqs...\n";    
    open (REF, $ref_file) or die $!;
    
# Recorrer los hashes de hashes
    foreach my $refseq (keys %{$hash_ref}){
	print STDERR ."\t".$refseq."  ";
	while (my $ref_line = <REF>){
	    chomp $ref_line;
	    my @ref_rec = split ("\t", $ref_line);
	    if ($refseq eq $ref_rec[2]){
		${$hash_ref}{$refseq}{"name"} = $ref_rec[2];
		print STDERR "Ok\n";
	    }
	}
	seek (REF, 0, 0);
    }
    close (REF);
    print STDERR "Done\n";
}

sub extract_genes
{
#    print STDERR
    my $file = shift @_;
    open (F, $file) or die "Cannot open $file\n\n";

    my @gene_refseqs = (); 
    while (my $refseqs = <F>){
	chomp $refseqs;
	last if $refseqs=~/^$/;
	push (@gene_refseqs, $refseqs) if !&isin_array (\@gene_refseqs, $refseqs);
    }
    close (F);
    return \@gene_refseqs;
}


sub slice_hash
{
    my ($array_ref, $array_reg, $hash_ref) = @_;

    my %slice;
    my $n_genes_in_chr = 0;
    my $n_genes_reg_in_chr = 0;

    foreach my $gene (@{$array_ref}){
	if (defined ${$hash_ref}{$gene}){
	    $slice{$gene} = { %{${$hash_ref}{$gene}} };
	    $n_genes_in_chr++;
 	    if (&isin_array ($array_reg, $gene) ){
 		$slice{$gene}{"is_reg"} = 1;
		$n_genes_reg_in_chr++;
 	    }
	}
    }
# Dumper code
#    my $msg = Dumper (\%slice);
#     print $msg."\n\n";
#    exit();

    return (\%slice, $n_genes_in_chr, $n_genes_reg_in_chr);
}

sub cluster_finder
{
    ## Debe recibir las referencias a los hashes que tienen los genes (por cromosoma)
    ## y la longitud del chromosoma
    my ($ref_hash, $length, $N, $k, $window, $offset, $chr_name) = @_;
    my %pvals_h;
    my $n_clusters = 0;
    my @clusters; 

    print STDERR "Looking for clusters in $chr_name...";
    ## Bucle para desplazar la ventana:
    for (my $pos = 0; $pos<$length; $pos+=$offset){
	my ($M_genes,$x_genes) = &find_genes ($ref_hash, $pos, $window);
#	my $x_genes = &find_genes ($reg_hash, $pos, $window);
	my $M = scalar (@{$M_genes});
	my $x = scalar (@{$x_genes});
	
## Perform hypergeometric distribution test
	my $hpobj = Stat::hypergeom -> new ($N, $M, $k, $x);
	my $pvals = Stat::hypergeom -> over_represented_pval ($hpobj);
	if (($pvals < $pval_lim) && (scalar @{$x_genes} >= $min_genes_in_cluster)){ $n_clusters++;  }
#	if ($pvals < 0.01){
	if (scalar @{$x_genes} >= $min_genes_in_cluster ){
	    my %cluster;
	    $cluster{"chr"}    = $chr_name;
	    $cluster{"begin"}  = $pos;
	    $cluster{"end"}    = $pos+$window;
	    $cluster{"genes"}  = [@{$x_genes}];
	    $cluster{"pvalue"} = $pvals;
	    push @clusters, \%cluster;
#	    printf ("Cluster %d: from %d to %d => pval (%d %d %d %d) = %.3f (",
#		    $n_clusters++, $pos, $pos+$window, $N, $M, $k, $x, $pvals);
#	    foreach my $gene (@{$x_genes}){
#		printf ("%s ",$gene);
#	    }
#	    print "\n";
#	}
	    $pvals_h{$pos} = $pvals;  
	}

    }
## Performs multtest correction
#    my $qlimit = Stat::multtest->fdr_pval(\%pvals_h);
#    print "PVALUE limit: $qlimit\n";
#     foreach my $clus_ref (@clusters){
# 	my $key = ${$clus_ref}{"begin"};
# 	${$clus_ref}{"qval"} = ${${$qvobj}{$key}}[1];
#     }
# Dumper code
#    my $msg = Dumper ($qvobj);
#    print $msg."\n\n";
#    exit(0);
    print STDERR "Done ($n_clusters clusters)\n";
    return \@clusters;
}

sub find_genes
{
    my ($hash_ref, $p, $window) = @_;
    my @genes = ();
    my @genes_reg = ();
    foreach my $gene (keys (%{$hash_ref})){
	if ( (${$hash_ref}{$gene}{"begin"} > $p) && 
	     (${$hash_ref}{$gene}{"begin"} < ($p+$window)) ){
	    push @genes, ${${$hash_ref}{$gene}}{"name"};
	    if (${$hash_ref}{$gene}{"is_reg"} == 1){
		push @genes_reg, ${${$hash_ref}{$gene}}{"name"};
	    }
	}
    }
    return \@genes, \@genes_reg;
}

sub gene_window_clusters
{
    my ($hash_ref, $gwindow, $N, $k, $chr_name) = @_;
    my $n_clusters = 0;
    my @clusters; 
    my %pvals_h;
    my $M = $gwindow;
    my @g_ord = ();
    my $cont=0;
    $g_ord[$cont++] = ${$hash_ref}{$_} foreach (sort{ ${${$hash_ref}{$a}}{"begin"} <=> ${${$hash_ref}{$b}}{"begin"} } keys %{$hash_ref});

    my $tantos = scalar (@g_ord);
    print STDERR "Looking for clusters in $chr_name...";
    for (my $g = 0; $g < $tantos-$gwindow; $g++){
	my $x=0;
	my $gg = $g+$gwindow;
	my @reg_genes = ();
	my $beg_clust = $g_ord[$g]{begin};
	my $end_clust = $g_ord[$g+$gwindow]{begin};
	for (my $g_in_w = $g; $g_in_w < $g+$gwindow; $g_in_w++){
	    if ($g_ord[$g_in_w]{is_reg} == 1) {
		$x++;
		push @reg_genes, $g_ord[$g_in_w]{name};
	    }
	}
	my $hpobj = Stat::hypergeom -> new ($N, $M, $k, $x);
	my $pvals = Stat::hypergeom -> over_represented_pval ($hpobj);
	if (($pvals < $pval_lim) && (scalar @reg_genes >= $min_genes_in_cluster)){ $n_clusters++; }
	if (scalar @reg_genes >= $min_genes_in_cluster ){
	    my %cluster;
	    $cluster{"chr"}    = $chr_name;
	    $cluster{"begin"}  = $beg_clust;
	    $cluster{"end"}    = $end_clust;
	    $cluster{"genes"}  = [@reg_genes];
	    $cluster{"pvalue"} = $pvals;
	    push @clusters, \%cluster;
#	    printf ("Cluster %d: from %d to %d => pval (%d %d %d %d) = %.3f (",
#		    $n_clusters++, $pos, $pos+$window, $N, $M, $k, $x, $pvals);
#	    foreach my $gene (@{$x_genes}){
#		printf ("%s ",$gene);
#	    }
#	    print "\n";
#	}
	    $pvals_h{$beg_clust} = $pvals;  
	}
    }
    print STDERR "Done ($n_clusters clusters)\n";

# Dumper code
#      my $msg = Dumper (\@clusters);
#      print $msg."\n\n";
#      exit(0);


    return \@clusters;
}



#	for (my $nround = 0; $nround < scalar (keys %{$hash_ref})-$gwindow+1; $nround++){
#	    for (my $ng = 0; $ng < $gwindow; $ng++){

sub print_clusters
{
    my ($cluster_aref, $header) = @_;

    my @good_clusters;

    foreach my $cluster_ref (@{$cluster_aref}){
	my $score = ${$cluster_ref}{"pvalue"};
	if ($score < $pval_lim){
	    push @good_clusters, {%{$cluster_ref}};
	}
    }
    
    &filter_overlaping_clusters (\@good_clusters);
    foreach my $cluster_ref (@good_clusters){
	if (defined $cluster_ref){
	    my $chr = ${$cluster_ref}{"chr"};
	    my $begin = ${$cluster_ref}{"begin"};
	    my $end = ${$cluster_ref}{"end"};
	    my $score = ${$cluster_ref}{"pvalue"};
	    $score = $score if ($score =~ /\d+/);
	    my @genes = @{${$cluster_ref}{"genes"}};
	    if ($score < $pval_lim){
		print "$chr\tcluster_genomic\tcluster\t$begin\t$end\t$score\t.\t.\t";
		foreach my $gene (@genes){
		    print $gene." ";
		}
		print "\n";
	    }
	}
    }

}

sub filter_overlaping_clusters  ### -f --filter_overlaping 
{
    my $good_clusters_aref = shift @_;
    my %non_overlaping;

    for (my $n=0; $n<(scalar (@{$good_clusters_aref})-1); $n++){
	if (defined (${$good_clusters_aref}[$n])){
#	    print $n." ".${${$good_clusters_aref}[$n]}{"end"}."\t".($n+1)." ".${${$good_clusters_aref}[($n+1)]}{"begin"}."\n";
	    if ( (${${$good_clusters_aref}[$n]}{"chr"}) eq ${${$good_clusters_aref}[($n+1)]}{"chr"}  && 
		 (${${$good_clusters_aref}[$n]}{"end"} > ${${$good_clusters_aref}[($n+1)]}{"begin"}) ){
		${${$good_clusters_aref}[($n+1)]}{"begin"} = ${${$good_clusters_aref}[$n]}{"begin"};
		${${$good_clusters_aref}[($n+1)]}{"pvalue"} = 
		    ${${$good_clusters_aref}[($n+1)]}{"pvalue"}>${${$good_clusters_aref}[($n)]}{"pvalue"}? ${${$good_clusters_aref}[$n]}{"pvalue"}:${${$good_clusters_aref}[($n+1)]}{"pvalue"};

		foreach my $name (@{${${$good_clusters_aref}[$n]}{"genes"}}){
		    unshift @{${${$good_clusters_aref}[($n+1)]}{"genes"}}, $name if !&isin_array (${${$good_clusters_aref}[($n+1)]}{"genes"}, $name);
		}
		undef (${$good_clusters_aref}[$n]);
	    }
	}
    }

#     for (my $n=0; $n<scalar @{$all_clusters_aref}; $n++){
# 	if (defined (${$all_clusters_aref}[$n])){
# 	    if (${${$all_clusters_aref}[$n]}{"end"} < ${${$all_clusters_aref}[$n+1]}{"begin"}){
# 		${${$all_clusters_aref}[$n]}{"end"} = ${${$all_clusters_aref}[$n+1]}{"end"};
# 		${${$all_clusters_aref}[$n]}{"pvalue"} = (${${$all_clusters_aref}[$n]}{"pvalue"} + ${${$all_clusters_aref}[$n+1]}{"pvalue"}) / 2;
# 		foreach my $name (@{${${$all_clusters_aref}[$n+1]}{"genes"}}){
# 		    push @{${${$all_clusters_aref}[$n]}{"genes"}}, $name if !&isin_array (${${$all_clusters_aref}[$n]}{"genes"}, $name);
# 		}
# 		undef (${$all_clusters_aref}[$n+1]);
# 	    }
# 	}
#     }

}

sub isin_array
{
    my ($ref_array, $goid) = @_;

    foreach my $id_inarray (@$ref_array){
	return 1 if $id_inarray eq $goid;
    }
    return 0;
}
