#!/usr/bin/perl

#Copyright (C) 2006-2008 Keio University
#(Kris Popendorf) <comp@bio.keio.ac.jp> (2006)
#
#This file is part of Murasaki.
#
#Murasaki 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 3 of the License, or
#(at your option) any later version.
#
#Murasaki 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 Murasaki.  If not, see <http://www.gnu.org/licenses/>.

use strict;
use Getopt::Long;
use Pod::Usage;
use File::Basename;
use POSIX qw{floor};
#use Data::Dump qw{dump};

BEGIN {
  unshift(@INC,(fileparse($0))[1].'perlmodules');
}
use Murasaki;
use Murasaki::Ticker qw{resetTick tick};

our $root;
my $geneparse="$root/geneparse";
$geneparse="$root/geneparse.pl" unless -x $geneparse;
die "Need geneparse to be available" unless -x $geneparse;

my ($nodeTotal,$edgeTotal);

my ($help,$man);
my ($width,$height,$res)=(8,6,96);
my ($echo,$colors,$connect,$maxDraw,$colorby,$colorOpts,$title,$pointCex,$plotOrder,$density,$placeLegend,$legendPts,$squeeze,$pch);
$pointCex=.25;
$legendPts=8;
$pch=19;
$plotOrder='big';
my $err=
  GetOptions('help|?' => \$help, man => \$man,
	     'res=f'=>\$res,'width=f'=>\$width,'height=f'=>\$height,
	     echo=>\$echo,
	     'colors|palette=s'=>\$colors,
	     'colorby=s'=>\$colorby,
	     connect=>\$connect,'maxlines=i'=>\$maxDraw,
	     'colorOpts=s'=>\$colorOpts,'title=s'=>\$title,
	     'size=f'=>\$pointCex,
	     bigfirst=>sub {$plotOrder='big';},
	     smallfirst=>sub {$plotOrder='small';},
	     'density=i'=>\$density,
	     'legend=s'=>\$placeLegend,
	     'legendpts=i'=>\$legendPts,
	     'squeeze=s'=>\$squeeze
	  );
pod2usage(1) if !$err or $help or @ARGV<1 or !$colorby=~m/^size|id|density$/ or $legendPts<2;
pod2usage(-exitstatus => 0, -verbose => 2) if $man;

my $inf=shift @ARGV;
die "File not found: $inf" unless -f $inf;
die "Input must be a .repeats file" unless $inf=~m/^(.*)\.repeats$/;
my ($prefix)=$1;
die "Invalid file name for $inf" unless $prefix;
my $seqs="$prefix.seqs";
die "Seqs file ($seqs) not found" unless -f $seqs;

my @seqs=split(/\n/,slurp($seqs));
our $seqCount=scalar(@seqs);
our @seqLengths=map {my $l=`$geneparse -l $_`;chomp $l; $l} @seqs;

if($density){
  $colorby='density';
}elsif($colorby eq 'density'){
  $density=floor(max(@seqLengths)/($width*$res)*.8*($pointCex*40)); #approximate, but ok.
}

my $outf="$inf.png";
$outf="$inf.".join("-",$colorby ? ($colorby):(),$colors ? ($colors):()).".png" if $colors or $colorby;
if(@ARGV>0){
  $outf=shift @ARGV;
}
print "Writing output to $outf\n";

my $Rfh;
my ($type)= $outf=~m/\.(\w+)$/;
my $outputter=$type ne 'pdf' ?
  ($type=~/^e?ps$/ ?
   qq!postscript(file="$outf",width=$width,height=$height)!:
   qq!bitmap(file="$outf",type="png16m",width=$width,height=$height,res=$res)!):
  qq!pdf(file="$outf",width=$width,height=$height)!;
my $Ropts="--slave" unless $echo;
open($Rfh,"|R --vanilla $Ropts");
dor($outputter);
open(my $infh,"$inf");

$title=(fileparse($inf))[0] unless $title;
my $subtitle="Colored according to repeat ".($colorby ? $colorby:"id");
my ($fg,$bg)=("'black'","'white'");
if($colors eq 'heat'){ #because heat ends in white, which obviously doesn't show up well on white
  ($fg,$bg)=($bg,$fg);
  dor("par(bg=$bg,fg=$fg,col.axis=$fg,col.lab=$fg,col.main=$fg,col.sub=$fg)");
}
dor("par(yaxt='n',las=1)");
do {
  my ($squeezeLeft,$squeezeRight)=(($squeeze=~m/(\S+),(\S+)/) ? ($1,1-$2):($squeeze,1));
  dor("par('fig'=c($squeezeLeft,$squeezeRight,0,1))")
} if $squeeze;
my $seqLines="c(".join(",",map {(0,$seqLengths[$_])} (0..$#seqs))."),c(".
  join(",",map {($_,$_)} (0..$#seqs)).")";
dor("plot($seqLines,'n',ylim=c(-.5,".($#seqs+.5)."),col='black',xlab='Sequence Position',ylab='',main='$title',sub='$subtitle')"); #sets the range automatically
foreach my $i (0..$#seqs){
  my $y=($#seqs-$i);
  dor("lines(c(0,$seqLengths[$i]),c($y,$y),lwd=2)");
}
dor("par(yaxt='s')");
do {
  my $num=$seqCount+1;
  dor("axis(side=2,at=c(".join(",",0..$#seqs)."),labels=c(".
      join(",",map {$num--;"\"#$num: ".((fileparse($_,qr/\.[^.]+$/))[0]).'"'} reverse(@seqs))."),tick=TRUE)");
};

#start the parsing!

my (@clusters,@sizes);

my $maxPos;
my @seqMax;
my $clusterId;
my ($length,@locs,$size);
while (<$infh>) {
  chomp;
  unless($_){
    push(@clusters,{size=>$size,locs=>[@locs],id=>$clusterId});
    push(@sizes,$size);
    #start fresh
    undef $length;
    undef $size;
    @locs=();
    $clusterId++;
    next;
  }
  my ($seq,$dat)=m/^([^:]+): (.*)$/ or die "Weird line: $_";
  if ($seq eq 'R') {
    $length=length($dat);
  } else {
    my @poses=split(/\s/,$dat);
    $size+=scalar @poses;
    $locs[$seq]=[@poses];
  }
}

@sizes=sort @sizes;
resampleDensity() if($colorby eq 'density');

my $paletteMaker='rainbow';
$paletteMaker=$colors if $colors;
if($paletteMaker eq 'rainbow'){
  $colorOpts="end=5/6" unless $colorOpts or $colorby eq 'id';
}else{
  if($colors=~m/^rainbow|heat|terrain|topo|cm$/){ #built in palettes
    $paletteMaker.=".colors"
  }else{
    $colors="blue,yellow,red" if $colors eq 'cool';
    dor("require(graphics);custom.colors=colorRampPalette(c(".
	join(',',map {'"'.$_.'"'} split(/,/,$colors)).
	"),space='Lab')");
    $paletteMaker="custom.colors";
  }
}
$colorOpts=','.$colorOpts if $colorOpts;
print "Using $paletteMaker to make colors for ".(scalar @clusters)." clusters\n";
dor("pal<-${paletteMaker}(".scalar(@clusters).($colorOpts ? ",$colorOpts":"").")");
print "Initial set up completed. Plotting clusters.\n";
resetTick(scalar @clusters);

my ($totalLines,$linesDrawn,$totalPoints,$clustersLined);
foreach my $cluster (sort {$plotOrder eq 'big' ? $b->{size} <=> $a->{size}:$a->{size} <=> $b->{size}} @clusters){
  local $"=',';
  my $lines;
  my @locs=@{$cluster->{locs}};
  my (@xlist,@ylist);
  my (@prevx,@prevy);
  my $color=$cluster->{id};
  $color=getSizeRank($cluster->{size}) if $colorby eq 'size';
  foreach my $seq (0..$#locs) {
    my (@myx,@myy);
    foreach my $loci (0..$#{$locs[$seq]}) {
      my $x=abs($locs[$seq][$loci]);
      my $y=$#seqs-($seq+($locs[$seq][$loci]>0 ? -.25:.25));
      push(@myx,$x);
      push(@myy,$y);
      $totalPoints++;
    }

    if($connect and $seq>0){ #draw lines between all pairs in each adjacent seq
      my (@cx,@cy);
      foreach my $i (0..$#myx){
	foreach my $j (0..$#prevx){
	  push(@cx,$myx[$i],$prevx[$j],'NA');
	  push(@cy,$myy[$i],$prevy[$j],'NA');
	  $totalLines++;
	  $lines++;
	}
      }
      if(!$maxDraw or $lines<=$maxDraw){
	dor("lines(c(@cx),c(@cy),col=pal[$color])");
	$linesDrawn+=$lines;
	$clustersLined++;
      }
      @prevx=@myx;
      @prevy=@myy;
    }
    push(@xlist,@myx);
    push(@ylist,@myy);
  }
  dor("points(c(@xlist),c(@ylist),col=pal[$color],cex=$pointCex,pch=$pch)");
  tick();
}

if($placeLegend){
  dor("legend('$placeLegend',c(".join(",",map {'"'.$_.'"'} ('high',(('') x ($legendPts-2)),'low'))."),pch=$pch,col=rev(${paletteMaker}($legendPts)))");
}

print "\nDone\n";
print "Drew $totalPoints points.\n";
if($connect){
  print "Drew $linesDrawn lines (out of $totalLines possible (".percent($linesDrawn,$totalLines).") for $clustersLined (".percent($clustersLined,scalar(@clusters)*($seqCount-1)).") clusters-sequences\n";
}
close $Rfh;

exit;

do {
  my %rankMemo;
sub getSizeRank {
  my ($size)=@_;
  return $rankMemo{$size} if exists $rankMemo{$size};
  my $rank;
  foreach my $i (0..$#sizes){
#    $rank=$i+1 and last if $sizes[$i]==$size;
    $rank=($paletteMaker eq 'rainbow' ? $#sizes-$i:$i)+1 and last if $sizes[$i]==$size;
  }
  $rankMemo{$size}=$rank;
  return $rank;
}

};

sub percent {
  return 'N/A%' unless $_[1];
  return sprintf("%.2f%%",$_[0]/$_[1]*100);
}

sub dor {
  my @cmds=@_;
  foreach my $cmd (@cmds){
    $cmd.=";" unless $cmd=~m/;\s+$/;
    $cmd.="\n" unless $cmd=~m/\n$/;
    print $Rfh $cmd;
    print $cmd if $echo;
  }
}

sub max {
  my $max=$_[0];
  foreach (@_){
    $max=$_ if $_>$max;
  }
  return $max;
}

sub slurp {
  local $/;
  open(my $fh,"@_") or return;
  return <$fh>;
}

sub resampleDensity {
  my @pixels;
  my $pixels=floor(max(@seqLengths)/$density);
  print "Calculating density every $density bp (into $pixels points) from ".scalar(@clusters)." clusters.\n";
  foreach my $cluster (@clusters){
    my @locs=@{$cluster->{locs}};
    foreach my $seq (0..$#locs){
      foreach my $loci (0..$#{$locs[$seq]}) {
	my $pi=round($locs[$seq][$loci]/$density);
	$pixels[$seq]->{$pi}++;
      }
    }
  }

  my %sizes;
  my $tick;
  foreach my $seq (0..$#seqs){
    my $size;
    foreach my $pixel (ref $pixels[$seq] ? keys(%{$pixels[$seq]}):()){ #it's possible we don't have any repeats
      $size+=$pixels[$seq]->{$pixel};
      $sizes{$size}[$seq]=[] unless $sizes{$size};
      push(@{$sizes{$size}[$seq]},$pixel*$density);
      $tick++;
    }
  }

  @clusters=(); @sizes=(); #reset existing clusters and sizes
  my $id=0;
  foreach my $size (sort {$plotOrder eq 'big' ? $b<=>$a : $a<=>$b} keys %sizes){
    push(@clusters,{size=>$size,locs=>$sizes{$size},id=>$id});
  }continue{$id++}
}

sub round {
  my ($number)=@_;
  return int($number + .5 * ($number <=> 0));
}

__END__

=head1 NAME

repeatviz.pl -- visualize a bunch of repeats using graphviz

=head1 SYNOPSIS

repeatviz.pl <input> [output]

=head1 OPTIONS

output is defined based on the input name unless otherwise specified

 Options:
--echo shows all input to R (very spammy).
--colors={rainbow,heat,terrain,topo,cm,cool} uses an alternate color palette (default is rainbow). Custom palettes can also be created by providing a list of colors (like red,yellow,blue).
--colorby={size,id,density} specify what to select colors based on (default is id)
--density={N} sample repeat density at every N bp (default calculates a value based on width*res and genome size)
--coloropts={extra options to pass to R's color generator}
--connect draw connecting edges for points in  adjacent sequences
--maxlines={N} don't draw lines to connect cluster-sequence pairs with more than N edges
--size={N} size of dots
--bigfirst plot big clusters first
--smallfirst plot small clusters first(default)
--res={DPI} defaults to 96
--width|height={N} graph size in inches
--legend={position} position to place legend (R keywords like "right" or "topright" are good)
--legendpts={N} number of dots to put in the legend (default is 8). must be at least 2
--squeeze={[0,1)} Y axis labels not fitting in the frame? squeeze the graph over with this.
--squeeze={[0,1),F} an amount to change the right margin can also be specified by F like this.