#!/usr/bin/env perl

# Copyright (c) Illumina 2008
# Author: Richard Shaw
# This software is covered by the "Illumina Genome Analyzer Software
# License Agreement" and the "Illumina Source Code License Agreement",
# and certain third party copyright/licenses, and any user of this
# source file is bound by the terms therein (see accompanying files
# Illumina_Genome_Analyzer_Software_License_Agreement.pdf and
# Illumina_Source_Code_License_Agreement.pdf and third party
# copyright/license notices).

use warnings;
use strict;

use lib '/home/psgendb/local/pkg/CASAVA_v1.8.2-build/lib/CASAVA-1.8.2/perl'; # substituted by CMake during install
use Casava::Common::IOLib qw(executeCmd);
use Casava::Common::Utils qw(gnuplotImage get_tiles_by_lane updateCycleOffsets);

#------------------------------------------------------------------------------


my $plot_data_ext_str = 'dat';
my $plot_script_ext_str = 'txt';

#------------------------------------------------------------------------------

sub check_num_cycles_found($$$)
{
    my ($num_cycles, $last_cycle_num, $table_desc) = @_;

    if (int($last_cycle_num) != int($num_cycles)) {
	warn("WARNING: Number of cycles found in $table_desc ($last_cycle_num) does "
	     . "not match number specified at top of file ($num_cycles).");
    }
}

#------------------------------------------------------------------------------

sub parse_rescore_file($$$$$$$$)
{
    my ($rescore_text_file_name,
	$num_reads_ref,
	$cycle_offsets_ref,
	$err_info_by_cycle_ref,
	$cum_errs_by_cycle_ref,
	$max_errs_to_count_ref,
	$num_unique_aligns_by_read_ref,
	$bad_num_cycles_found_ref) = @_;

    my $errs_by_cycle_hdr = "Breakdown of errors by cycle";
    my $in_errs_by_cycle = 0;
    my $in_cum_errs = 0;
    my $cycle_num = 0;
    my $num_cycles = 0;
    my $bad_num_cycles_found = 0;

    my $read_num = 0; # default

    my $cycle_offset = 0;

    if (open(RESCORE_TXT, $rescore_text_file_name))
    {

        while (<RESCORE_TXT>) {
    	if ($_ =~ m/\#/) { # Handle comments first - some contain relevant info
    	    if ($_ =~ /^\# Lane.*Read/) {
    		if ($_ =~ /^\# Lane.* : Read (\d+)/) {
    		    # Header for numbered individual Read
                        #$$num_reads_ref = $read_num = $1; # assumes Reads in order
                        $$num_reads_ref++;
                        $read_num++;
    
    		    if ($read_num > 1) {
    			if ($num_cycles == 0) {
    			    $bad_num_cycles_found = 1;
    			}
    
    			# We now know number of cycles in the previous read.
    			updateCycleOffsets(@{$cycle_offsets_ref},
    					   $read_num - 1, $num_cycles);
    			$cycle_offset = $cycle_offsets_ref->[$read_num];
    		    }
    		} else {
    		    # All Reads header - assumed after the individual Reads.
    		    last; # do not need any All Reads statistics
    		}
    	    }
    
    	    next;
    	}
    
    	chomp;
    
    	if ($in_errs_by_cycle == 1) {
                # Should happen num cycles times, so handle early.
    	    my @fields = split;
    
    	    if (@fields != 3) {
    		# First line after an errors-by-cycle data line (if any).
    		$in_errs_by_cycle = 0;
    
    		check_num_cycles_found($num_cycles, $cycle_num,
    				       'errors by cycle');
    		next;
    	    }
    
    	    $cycle_num = $fields[0];
    	    @{$err_info_by_cycle_ref->[$cycle_offset + $cycle_num - 1]} 
    	    = @fields;
    	} elsif ($_ =~ /^\!(\d+)\s(.+)/) {
                # Should happen num cycles times, so handle early.
    	    $in_cum_errs = 1;
    
    	    $cycle_num = $1;
    	    my @err_counts = split('\t', $2);
    	    @{$cum_errs_by_cycle_ref->[$cycle_offset + $cycle_num - 1]}
    	    = @err_counts;
    
    	    $$max_errs_to_count_ref = @err_counts;
    	} elsif ($in_cum_errs) {
    	    # First line after a cumulative errors line
    	    $in_cum_errs = 0;
    
    	    check_num_cycles_found($num_cycles, $cycle_num,
    				   'cumulative errors by cycle');
    	} elsif ($_ eq $errs_by_cycle_hdr) {
    	    # Skip the column headings line.
    	    die "Expected line after $errs_by_cycle_hdr"
    		unless defined(<RESCORE_TXT>);
    	    $in_errs_by_cycle = 1;
    	} elsif ($_ =~ /unique alignments : (\d+) \(total score \d+\)/) {
    	    $num_unique_aligns_by_read_ref->[$read_num] = $1;
    	} elsif ($_ =~ /cycles : (\d+)/) {
    	    # Assumed to be before the main tables so the value can be used
    	    # as a check on the highest cycle numbers found in them.
    	    $num_cycles = $1;
    	}
    
        } # loop over input (rescore) file
    
        close (RESCORE_TXT);
    }
    else
    {
        warn("WARNING: Problem opening $rescore_text_file_name: $!");
    }
    
    # We now know the number of cycles in the last individual Read.

    if ($num_cycles == 0) 
    {
        $$bad_num_cycles_found_ref = 1;
    }

    updateCycleOffsets(@{$cycle_offsets_ref}, $read_num, $num_cycles);
}

#------------------------------------------------------------------------------
# Write gnuplot data file for errors and blanks by cycle.

sub write_error_blank_data_file($$$$$)
{
    my ($plot_tmp_name, $num_reads, $cycle_offsets_ref,
	$err_info_by_cycle_ref, $num_unique_aligns_by_read_ref) = @_;
    my $data_filename = join('.', $plot_tmp_name, $plot_data_ext_str);

    open (OUT, ">$data_filename")
	|| die "Problem opening $data_filename for writing: $!";

    for (my $read_ind = 1; $read_ind <= $num_reads; ++$read_ind) {
	my $num_unique_aligns = $num_unique_aligns_by_read_ref->[$read_ind];
	my $next_read_start_cycle = $cycle_offsets_ref->[$read_ind + 1];

	for (my $cycle_ind = $cycle_offsets_ref->[$read_ind];
	     $cycle_ind < $next_read_start_cycle; ++$cycle_ind) {

	    printf OUT "%d\t", $cycle_ind + 1;

	    # If there are no unique alignments for a read,
	    # set err% 0, blank% 100.
	    if ($num_unique_aligns == 0) {
		print OUT "0.0\t100.0\n";
		next;
	    }
	    for (my $field_ind = 1; $field_ind <= 2; ++$field_ind) {
		print OUT "$err_info_by_cycle_ref->[$cycle_ind][$field_ind]\t";
	    }

	    print OUT "\n";
	}
    }

    close (OUT);
}

#------------------------------------------------------------------------------

sub draw_error_blank_graphs($$$$)
{
    my ($rescore_text_filename, $plot_tmp_name, $plot_image_filename, $plot_thumbnail_filename) = @_;

    my $plot_script_name = join('.', $plot_tmp_name, $plot_script_ext_str);
    my $data_filename = join('.', $plot_tmp_name, $plot_data_ext_str);

    open (OUT, ">$plot_script_name")
	|| die "Problem opening $plot_script_name for writing: $!";

# 1. "set output" doesn't work in multiplot mode, have to pipe
# the output to a file instead
# 2. Would like to autoscale the x axes, probably need a newer version
# of gnuplot
print OUT << "END";
#    set terminal png color;
set terminal postscript portrait color solid;
set nokey;
#    set yrange [0:];
set yrange [0:20];
set boxwidth 0.6;
set size 1.0, 1.0;
set origin 0.0, 0.0;
set multiplot;
set size 1.0, 0.5;
set origin 0.0, 0.5;
set title "Error rates by cycle for file $rescore_text_filename";
#    set xlabel "Cycle Number";
set ylabel "Percent";
plot '$data_filename' using 1:2 title 'error rate' with impulses lt 1 lw 24;
set origin 0.0, 0.0;
set title "Occurrence of blank bases by cycle for file $rescore_text_filename";
set xlabel "Cycle Number";
plot '$data_filename' using 1:3 title 'blank rate' with impulses lt 3 lw 24;
set nomultiplot;

END

    close (OUT);
    # gif driver does not support linewidth on gnuplot 4.0 (works on 4.2 though)
    gnuplotImage($plot_script_name, $plot_image_filename);

    open (OUT, ">$plot_script_name")
    || die "Problem opening $plot_script_name for writing: $!";

# 1. "set output" doesn't work in multiplot mode, have to pipe
# the output to a file instead
# 2. Would like to autoscale the x axes, probably need a newer version
# of gnuplot
print OUT << "END";
#    set terminal png color;
set terminal gif medium size 84,84 \\
                       xffffff x000000 x404040 \\
                       xff0000 x00ff00 \\
                       x0000ff x000000 \\
                       x000000 x0000ff \\
                       x000000 x000000 \\
                       x777777  
set nokey;
set lmargin 0;
set rmargin 0;
set tmargin 0;
set bmargin 0;
set noytics;
set noxtics;
#    set yrange [0:];
set yrange [0:20];
set boxwidth 0.6;
set size 1.0, 1.0;
set origin 0.0, 0.0;
set multiplot;
set size 1.0, 0.5;
set origin 0.0, 0.5;
plot '$data_filename' using 1:2 with impulses lt 1 lw 4;
set origin 0.0, 0.0;
plot '$data_filename' using 1:3 with impulses lt 3 lw 4;
set nomultiplot;

END

    close (OUT);
    executeCmd("gnuplot $plot_script_name >$plot_thumbnail_filename");
}

#------------------------------------------------------------------------------
# Write gnuplot data file for cumulative errors.

sub write_cum_error_data_file($$$$$$)
{
    my ($plot_tmp_name, $num_reads, $cycle_offsets_ref,
	$cum_errs_by_cycle_ref, $max_errs_to_count,
	$num_unique_aligns_by_read_ref) = @_;

    my $data_filename = join('.', $plot_tmp_name, $plot_data_ext_str);

    open (OUT, ">$data_filename")
	|| die "Problem opening $data_filename for writing: $!";

    for (my $read_ind = 1; $read_ind <= $num_reads; ++$read_ind) {
	my $num_unique_aligns = $num_unique_aligns_by_read_ref->[$read_ind];
	my $next_read_start_cycle = $cycle_offsets_ref->[$read_ind + 1];

	for (my $cycle_ind = $cycle_offsets_ref->[$read_ind];
	     $cycle_ind < $next_read_start_cycle; ++$cycle_ind) {
   
	    printf OUT "%d\t", $cycle_ind + 1;

	    # If there are no unique alignments for a read, set count of reads
	    # with < N errors to be 0 for all N considered.
	    if ($num_unique_aligns == 0) {
		for (my $num_errs_ind = 0;
		     $num_errs_ind < $max_errs_to_count; ++$num_errs_ind) {
		    print OUT "0\t";
		}

		print OUT "\n";
		next;
	    }
  
	    for (my $num_errs_ind = 0; $num_errs_ind < $max_errs_to_count;
		 ++$num_errs_ind) {
		if (!defined($cum_errs_by_cycle_ref->
			     [$cycle_ind][$num_errs_ind])) {
		    print("result not defined for cycle-1 $cycle_ind, ",
			  "j $num_errs_ind\n");
		}

		$cum_errs_by_cycle_ref->[$cycle_ind][$num_errs_ind]
		    = (($cum_errs_by_cycle_ref->[$cycle_ind][$num_errs_ind]
			/ $num_unique_aligns) * 100);
		print(OUT $cum_errs_by_cycle_ref->[$cycle_ind][$num_errs_ind],
		      "\t");
	    }

	    print OUT "\n";
	}
    }

    close(OUT);
}

#------------------------------------------------------------------------------

sub draw_cum_err_graph($$$$$)
{
    my ($rescore_text_filename, $plot_tmp_name, $plot_image_filename,
	$plot_thumbnail_filename, $num_unique_aligns_by_read_ref) = @_;

    my $plot_script_name = join('.', $plot_tmp_name, $plot_script_ext_str);
    my $data_filename = join('.', $plot_tmp_name, $plot_data_ext_str);

    my $lane = 1;

    # Make label
    my @num_unique_aligns_by_read = @{$num_unique_aligns_by_read_ref};
    shift(@num_unique_aligns_by_read); # Removes undefined 'read 0' value.
    my $num_unique_aligns_str = join(', ', @num_unique_aligns_by_read);
    my $label = "CLUSTERS = $num_unique_aligns_str";

    my $title = $rescore_text_filename;
    $title =~ s/rescore/realign/; # to match old plots

    my $width = 24;

    open (OUT, ">$plot_script_name")
	|| die "Problem opening $plot_script_name for writing: $!";

    print OUT << "END";
# set terminal png color;
set terminal postscript portrait color solid;
# set nokey;
set yrange [0:100];
set ytics 10;
    
# set boxwidth 0.6;
    
set size 1.0, 1.0;
set origin 0.0, 0.0;
set title "$title";
set xlabel "Cycle Number";
set ylabel "% reads with 'x' errors or less";
set key left bottom below
set label "$label" at -5,-6,0

plot '$data_filename'  using 1:6 title '4 or less' with impulses lt 7 lw $width, '$data_filename' using 1:5 title '3 or less' with impulses lt 2 lw $width,'$data_filename' using 1:4 title '2 or less' with impulses lt 3 lw $width,'$data_filename' using 1:3 title '1 or less' with impulses lt 1 lw $width, '$data_filename' using 1:2 title '0 errors' with impulses lt 9 lw $width; 

END
    close (OUT);
    gnuplotImage($plot_script_name, $plot_image_filename);

    $width = 4;
    open (OUT, ">$plot_script_name")
    || die "Problem opening $plot_script_name for writing: $!";

    print OUT << "END";
# set terminal png color;
set terminal gif medium size 84,84 \\
                       xffffff x000000 x404040 \\
                       xff0000 x00ff00 \\
                       x0000ff x000000 \\
                       x000000 x0000ff \\
                       x000000 x000000 \\
                       x777777  

set yrange [0:100];
set noytics;
set noxtics;

set size 1.0, 1.0;
set origin 0.0, 0.0;
set nokey;
set lmargin 0;
set rmargin 0;
set tmargin 0;
set bmargin 0;

plot '$data_filename'  using 1:6 with impulses lt 7 lw $width, '$data_filename' using 1:5 with impulses lt 2 lw $width,'$data_filename' using 1:4 with impulses lt 3 lw $width,'$data_filename' using 1:3 with impulses lt 1 lw $width, '$data_filename' using 1:2 with impulses lt 9 lw $width; 

END
    close (OUT);
    executeCmd("gnuplot $plot_script_name 1> $plot_thumbnail_filename");
}

#------------------------------------------------------------------------------

my $app_name_len = length($0);

my $usage
    = "Usage: $0 <tiles.txt> <s_1> <_rescore.txt>\n"
    . (' ' x $app_name_len)
    . "        <_rescore.gif> <_rescore_thumb.gif> <_errors.gif> <_errors_thumb.gif> <score_dir> <temp_dir> <plots_dir>\n"
    . "For each tile in the specified tile list file that matches the lane\n"
    . "prefix, reads the corresponding rescore text file to produce image\n"
    . "files :-\n"
    . "<TILE>_rescore.gif : graphs of errors and uncalled bases by cycle\n"
    . "<TILE>_rescore_thumb.gif : graphs of errors and uncalled bases by cycle\n"
    . "<TILE>_errors_thumb.png  : graph of cumulative errors by cycle\n"
    . "<TILE>_errors.png  : graph of cumulative errors by cycle\n";

die $usage unless (@ARGV == 10);

my ($tile_file_path, $lane_prefix, $rescore_text_suffix,
    $rescore_png_suffix, $rescore_thumbnail_suffix, $errors_png_suffix, $errors_thumbnail_suffix, 
    $score_dir, $tmpdir, $plotdir) = @ARGV;

my @tile_list;

if (!get_tiles_by_lane($tile_file_path, $lane_prefix, @tile_list))
{
    die("Failed to get tile list from file '$tile_file_path'.");
}


foreach my $tile (@tile_list) {

    my $plot_tmp_name = File::Spec->catfile($tmpdir, "$tile\_peg");
    my $rescore_text_filename = File::Spec->catfile($score_dir, $tile . $rescore_text_suffix);
    my $rescore_png_filename = File::Spec->catfile($plotdir,
						   $tile . $rescore_png_suffix);
    my $rescore_thumbnail_filename = File::Spec->catfile($plotdir,
                           $tile . $rescore_thumbnail_suffix);
    my $errors_png_filename = File::Spec->catfile($plotdir,
						  $tile . $errors_png_suffix);
	my $errors_thumbnail_filename = File::Spec->catfile($plotdir,
                          $tile . $errors_thumbnail_suffix);

    my $num_reads = 0; # default

    # @cycle_offsets is indexed by 1-offset $read_num.
    # Initialise to zero for Read 1.
    my @cycle_offsets = (0, 0);

    my @err_info_by_cycle;
    my @cum_errs_by_cycle;
    my $max_errs_to_count = 5; # default value - overridden by num cols found
    my @num_unique_aligns_by_read;
    my $bad_num_cycles_found = 0;

    parse_rescore_file($rescore_text_filename,
		       \$num_reads,
		       \@cycle_offsets,
		       \@err_info_by_cycle,
		       \@cum_errs_by_cycle,
		       \$max_errs_to_count,
		       \@num_unique_aligns_by_read,
		       \$bad_num_cycles_found);

    # If no data found, create a blank file and continue 
    # (rather than stopping the pipeline).
    if ($bad_num_cycles_found) 
    {
        warn("WARNING: No clusters found in one or more reads for tile $tile, "
         . "creating empty $rescore_png_filename and "
         . "$errors_png_filename.");
        for my $out_file ($rescore_png_filename, $errors_png_filename) 
        {
            open (OUT, ">$out_file") 
            || die "Problem creating blank file $out_file: $!";
            close (OUT);
        }
    }
    else
    {


        write_error_blank_data_file($plot_tmp_name, $num_reads, 
                                    \@cycle_offsets,
                                    \@err_info_by_cycle,
                                    \@num_unique_aligns_by_read);

        draw_error_blank_graphs($rescore_text_filename, $plot_tmp_name,
                                    $rescore_png_filename, $rescore_thumbnail_filename);

        write_cum_error_data_file($plot_tmp_name, $num_reads, 
                                    \@cycle_offsets,
                                    \@cum_errs_by_cycle, $max_errs_to_count,
                                    \@num_unique_aligns_by_read);

        draw_cum_err_graph($rescore_text_filename, $plot_tmp_name,
                          $errors_png_filename, $errors_thumbnail_filename,
                          \@num_unique_aligns_by_read);
    }
}

#------------------------------------------------------------------------------
