#!/usr/bin/env perl
# generate genotype matrix by querying a set of BAMs for read counts for
# a set of SNVs
# MNE 5/2013-
#
# TO DO:
# - optimization: only write 1 "sites" file, use symlinks for others

use strict;
use warnings;

use Getopt::Long;
use Carp qw(confess);

use PostprocessedVariantReport;

use Getopt::Long;

use SampleName;
use BAMMatrix;
use FileUtils qw(read_simple_file);
use DelimitedFile;
use BAMSet2;
use Cluster qw(run_now_or_later);
use FileUtils qw(newer_than);
use MiscUtils qw(dump_die);

my $NS_VARIANT_LABEL = "NTvariant";
my $NS_VARIANT_LABEL_AA = "AAvariant";
my $RAM = 4000;

my %FLAGS;
my @PASSTHROUGH_FIELDS;
my @BAM_LOOKUP_DIRS;
GetOptions(
	   \%FLAGS,
	   "-variant-file=s",
	   # single SNV report to analyze
	   "-variant-list=s",
	   # file containing list of SNV report names

	   "-bam-list=s",
	   "-sample-label",

	   "-flat",
	   "-sample-x",
	   "-sample-y",
	   "-all",
	   # matrix output options

	   "-nonsynonymous-only",
	   "-warn",

	   "-snvs-only",
	   # filter to SNVs only
	   "-direct",
	   "-nonsilent-only",
	   # filter to non-silent SNV event types only
	   "-non-bad",

	   "-bam-lookup=s" => \@BAM_LOOKUP_DIRS,
	   # use sample names in data to find BAMs.
	   # this is a directory to search.
	   "-require-bam",

	   "-passthrough=s" => \@PASSTHROUGH_FIELDS,
	   # annotation column(s) to pass through
	   "-passthrough-all",
	   # try to pass through ALL columns in source report
	   #
	   # The passthrough options are intended for static
	   # annotations (gene symbol, AA change, etc.);
	   # if variant appears in multiple source reports
	   # an arbitrary row will be chosen.
	   #

	   "-force-16",
	   # force 16-column interpretation of variant report

	   "-fraction-mode",

	   # SNV list formatting options:
	   "-bambino",
	   # Bambino report columns
	   "-snv4",
	   # SNVs in 4-field period delimited format,
	   # e.g. "chrMT.2696.A.G"

	   "-report-sample",
	   # report sample(s) the variant was originally observed in.
	   # requires Sample field in variant file input.

	   "-now",

	   "-clean",
	   # remove all intermediate and log files on completion

	   "-ram=i" => \$RAM,
	   "-name=s",

	   "-min-variant-reads-for=s",

	   "-idataplex",
	   # submit sub-jobs to idataplex

	   "-step1=s",
	   "-step2",
	   "-out-dir=s",

	   "-min-quality=i",

	  );

run_now_or_later("-memory" => $RAM) || exit(0);

my $bm = new BAMMatrix("-job_id" => $FLAGS{name} || "combined_matrix");
$bm->idataplex(1) if $FLAGS{idataplex};
$bm->working_dir($FLAGS{"out-dir"}) if $FLAGS{"out-dir"};
if (my $mq = $FLAGS{"min-quality"}) {
  $bm->min_quality($mq);
}

if (my $of = $FLAGS{step1}) {
  $bm->step_mode(1);
  $bm->step1_outfile($of);
} elsif ($FLAGS{step2}) {
  $bm->step_mode(2);
}

if (my $f = $FLAGS{"min-variant-reads-for"}) {
  my ($bam, $count) = split /,/, $f;
  die unless $count;
  die unless -s $bam;
  $bm->snv_filter_bam($bam);
  $bm->snv_filter_min_variant_reads($count);
}

#my $jram = $RAM / 3;
#$bm->java_ram($jram) if $jram > $bm->java_ram();

$bm->warn(1) if $FLAGS{warn};
$bm->bam_label_sample($FLAGS{"sample-label"} ? 1 : 0);
$bm->fraction_mode(1) if $FLAGS{"fraction-mode"};

if ($FLAGS{flat}) {
  $bm->matrix_format("simple");
} elsif ($FLAGS{"sample-x"}) {
  $bm->matrix_format("sample_x");
} elsif ($FLAGS{"sample-y"}) {
  $bm->matrix_format("sample_y");
}
# matrix output formatting

if (@PASSTHROUGH_FIELDS) {
  $bm->extra_fields(\@PASSTHROUGH_FIELDS);
  $bm->annotation_callback(\&passthrough_callback);
}

#
#  scan reports for variant sites:
#
my $vlist;
if (my $vf = $FLAGS{"variant-list"}) {
  $vlist = read_simple_file($vf);
} elsif (my $single = $FLAGS{"variant-file"}) {
  $vlist = [ $single ];
} else {
  die "no variant files: specify -variant-file [SNV file] | -variant-list [listfile]\n";
}

die "no variant files" unless @{$vlist};

if ($FLAGS{"passthrough-all"}) {
  my $pvr = new PostprocessedVariantReport();
  # parse various flavors of variant report
  $pvr->parse(
	      "-file" => $vlist->[0],
	      "-force-16" => $FLAGS{"force-16"}
	     );
  @PASSTHROUGH_FIELDS = @{$pvr->{headers}};
  $bm->extra_fields(\@PASSTHROUGH_FIELDS);
  $bm->annotation_callback(\&passthrough_callback);
}

if ($FLAGS{"report-sample"}) {
  $bm->add_extra_field("Sample");
  $bm->annotation_callback(\&sample_callback)
    unless $bm->annotation_callback();
}

my %annotations;
my %samples;
my %saw_bam;

my $cache_files = $bm->get_cache_files();
my %cache_files;
foreach my $cf (@{$cache_files}) {
  $cache_files{$cf} = 1;
}

foreach my $vf (@{$vlist}) {
  printf STDERR "parsing %s...\n", $vf;

  foreach my $cf (keys %cache_files) {
    if (newer_than($vf, $cf)) {
      unlink $cf;
      delete $cache_files{$cf};
    }
  }

  die sprintf "ERROR: %s does not appear to be a file.  If you are using -variant-list, do you mean to be using -variant-file?\n", $vf unless -f $vf;

  my $rows;
  if ($FLAGS{"nonsynonymous-only"}) {
    # in this mode additional annotation columns have been added.
    # We can't use PostprocessedVariantReport() because this will
    # not be an expected format.  Assume headers are available.
    my $df = new DelimitedFile("-file" => $vf,
			       "-headers" => 1,
			      );
    # while (my $row = $df->next("-ref" => 1)) {  # headerless
    $rows = [];
    while (my $row = $df->get_hash()) {
      push @{$rows}, $row;
    }

    $bm->variant_label($NS_VARIANT_LABEL);
    $bm->annotation_callback(\&annotation_callback);
    $bm->extra_fields([ $NS_VARIANT_LABEL_AA ]);
  } elsif ($FLAGS{snv4}) {
    $rows = [];
    my $l = read_simple_file($vf);
    foreach my $v (@{$l}) {
      $v =~ s/\r$//;
      # strip DOS line end char, if present
      next unless $v =~ /\w/;
      # blank lines
      my @f = split /\./, $v;
      die sprintf "line \"%s\" doesn't contain exactly 4 fields", $v unless @f == 4;
      my %r;
      @r{qw(Chr WU_HG19_Pos ReferenceAllele MutantAllele)} = @f;
      push @{$rows}, \%r;
    }
  } elsif ($FLAGS{bambino}) {
    # file with Bambino-style headers
    my $df = new DelimitedFile("-file" => $vf,
			       "-headers" => 1,
			      );
    my $snvs_only = $FLAGS{"snvs-only"};
    my $direct_mode = $FLAGS{direct};
    my $count = 0;
    while (my $row = $df->get_hash()) {
      next if $snvs_only and $row->{Type} ne "SNP";
      if ($direct_mode) {
	$bm->add_snv($row->{Chr} || die,
		     $row->{Pos} || die,
		     $row->{Chr_Allele} || die,
		     $row->{Alternative_Allele} || die);
      } else {
	my %r;
	$r{Chr} = $row->{Chr} || die;
	$r{WU_HG19_Pos} = $row->{Pos} || die;
	$r{ReferenceAllele} = $row->{Chr_Allele} || die;
	$r{MutantAllele} = $row->{Alternative_Allele} || die;
	push @{$rows}, \%r;
      }
      $count++;
      printf STDERR "%d...\n", $count if $count % 100000 == 0;
    }


  } else {
    my $pvr = new PostprocessedVariantReport();
    # parse various flavors of variant report
    $pvr->parse(
		"-file" => $vf,
		"-force-16" => $FLAGS{"force-16"}
	       );
    $rows = $pvr->rows();

    if ($FLAGS{"snvs-only"}) {
      # filter to SNVs only
      # TO DO: should probably be the default
      my @filtered;
      my %want = map {$_, 1} qw(A C G T);
      foreach my $row (@{$rows}) {
	my $usable = 1;

	foreach my $f (qw(ReferenceAllele MutantAllele)) {
	  my $base = $row->{$f};
	  if ($base) {
	    if ($want{$base}) {
	      # OK
	    } elsif ($base eq "-") {
	      $usable = 0;
	    } elsif (length($base) > 1) {
	      # insertion
	      $usable = 0;
	    } else {
	      die "unhandled $base";
	    }
	  } else {
	    printf STDERR "ERROR: where is %s in %s, chr=%s pos=%s, keys=%s", $f, $vf, $row->{Chr}, $row->{WU_HG19_Pos}, join ",", sort keys %{$row};
	    $usable = 0;
	  }
	}

	if ($usable) {
	  push @filtered, $row;
	} else {
	  printf STDERR "skipping ref=%s var=%s\n", $row->{ReferenceAllele}, $row->{MutantAllele};
	}

      }
      $rows = \@filtered;
    }

    if ($FLAGS{"nonsilent-only"}) {
      # filter to non-silent variants only
      my @filtered;

      my %want = map {$_, 1} qw(
				 missense
				 nonsense
				 splice
			      );

      my %ignore = map {$_, 1} qw(
				   UTR_3
				   UTR_5
				   silent
				   exon
				   intron
				   splice_region
			      );

      foreach my $row (@{$rows}) {
	my $class = $row->{Class} || die;
	if ($want{$class}) {
	  # usable
	  push @filtered, $row;
	} elsif ($ignore{$class}) {
	  # don't care
	} else {
	  die sprintf "unhandled %s at %s.%s ref=%s mutant=%s\n",
	    $class,
	      $row->{Chr}, $row->{WU_HG19_Pos},
		$row->{ReferenceAllele}, $row->{MutantAllele};
	}
      }

      $rows = \@filtered;
    }

    if ($FLAGS{"non-bad"}) {
      my @filtered;
      my %wanted = map {$_, 1} qw(Good 2SNP);
      my %known = %wanted;
      $known{Bad} = 1;
      $known{GoodinSNPCluster} = 1;
      $known{SNPCluster} = 1;

      foreach my $row (@{$rows}) {
	my $state = $row->{State} || die "no State column";
	die "unknown state $state" unless $known{$state};
	push @filtered, $row if $wanted{$state};
      }

      $rows = \@filtered;
    }

    if (@BAM_LOOKUP_DIRS) {
      #
      # look up BAMs as needed from data
      #
      my @bs;
      foreach my $dir (@BAM_LOOKUP_DIRS) {
	my $bs = new BAMSet2();
	$bs->add_bams("-dir" => $dir);
	push @bs, $bs;
      }

      my @filtered;
      foreach my $row (@{$rows}) {
	my $sample = $row->{Sample} || die "Sample";

	$sample =~ /^(SJ[A-Z]+)/ || die;
	my $disease = $1;

	my $usable = 1;
	my ($subject, $type, $index);
	if ($sample =~ /_/) {
	  ($subject, $type) = split /_/, $sample;
	} elsif ($sample =~ /(SJ[A-Z]+\d+)([A-Z])$/) {
	  ($subject, $type) = ($1, $2);
	} elsif ($sample =~ /^SJ[A-Z]+\d+$/) {
	  $subject = $sample;
	  $type = "D";
	} elsif ($sample =~ /^(SJ[A-Z]+\d+)([A-Z])(\d)$/) {
	  # munged v2 sample name, e.g. SJHGG010324A1
	  ($subject, $type, $index) = ($1, $2, $3);
	  $subject =~ /^SJ[A-Z]+(\d+)$/ || die;
	  $subject = "SJ" . $1;
	  # reformat to v2 subject name
	} else {
	  printf STDERR "ERROR: can't parse sample %s\n", $sample;
	  $usable = 0;
	}

	if ($usable) {
	  die "no type for $subject" unless $type;
#	  printf STDERR "sample $sample $subject $type\n";
	  my $bam;
	  foreach my $bs (@bs) {
	    # use first BAM available in user-specified directory order
	    $bam = $bs->get_bam(
				"-disease" => $disease,
				"-subject" => $subject,
				"-strict-type" => $type,
#				"-index" => $index,
				"-warn" => 1,
				  );
	    if ($bam) {
#	      printf STDERR "found BAM %s\n", $bam;
	      last;
	    } else {
#	      printf STDERR "missed BAM for %s\n", $sample;
	    }
	  }

	  if ($bam) {
	    unless ($saw_bam{$bam}) {
	      printf STDERR "using BAM %s\n", $bam;
	      $bm->add_bam($bam);
	      $saw_bam{$bam} = 1;
	    }
	  } elsif ($FLAGS{"require-bam"}) {
	    # filter the input data only to rows where we actually
	    # have BAM data available.  This prevents variants from
	    # getting into the matrix unless we can reasonably expect
	    # to see them in at least one sample.  Without this it's
	    # possible to have variants where no samples have any
	    # observations, because the BAMs for the parent samples
	    # aren't available.
	    $usable = 0;
	  }

	}
	push @filtered, $row if $usable;
      }
      $rows = \@filtered;
    }
  }

  foreach my $row (@{$rows}) {
    my $chrom = $row->{Chr} || die "Chr";
    my $pos = $row->{WU_HG19_Pos} || $row->{Pos} || $row->{WU_HG18_Pos} || die "Pos";
    my $base_ref = $row->{ReferenceAllele} || die "ReferenceAllele";
    my $base_var = $row->{MutantAllele} || die "MutantAllele";

    foreach my $f (qw(ReferenceAllele MutantAllele)) {
      my $v = $row->{$f};
      dump_die($row, "invalid allele in field $f") unless $v =~ /^[ACGT]$/i;
    }

    my $key = join ".", $chrom, $pos, $base_ref, $base_var;
    $annotations{$key} = $row;

    if ($FLAGS{"report-sample"}) {
      my $sample = $row->{Sample} || die "no Sample column";
      $samples{$key}{$sample} = 1;
    }

    if ($FLAGS{"nonsynonymous-only"}) {
      die unless exists $row->{cds_nonsynonymous_flag};
      next unless $row->{cds_nonsynonymous_flag};
    }

    $bm->add_snv($chrom, $pos, $base_ref, $base_var);
  }
}

if ($FLAGS{"bam-list"}) {
  my $blist = read_simple_file($FLAGS{"bam-list"} || die "-bam-list");
  die unless @{$blist};

  my %saw;
  foreach my $bam (@{$blist}) {
    next unless $bam =~ /\w/;
    if (-s $bam) {
      if ($saw{$bam}) {
	printf STDERR "WARNING: skipping duplicate BAM %s\n", $bam;
      } else {
	$bm->add_bam($bam);
	$saw{$bam} = 1;
      }
    } else {
      die "ERROR: $bam is missing/empty";
    }
  }
} elsif (@BAM_LOOKUP_DIRS) {
  # OK, already dealt with during parsing
} else {
  die "how to get BAMs?";
}

if ($FLAGS{all}) {
  # generate all 3 matrix formats
  my @formats = qw(simple sample_x sample_y);
  foreach my $format (@formats) {
    $bm->matrix_format($format);
    $bm->generate_matrix();
  }
} else {
  $bm->generate_matrix();
}

$bm->cleanup() if $FLAGS{clean};

sub annotation_callback {
  my ($row) = @_;
  my $key = $row->{$NS_VARIANT_LABEL} || confess "no field $NS_VARIANT_LABEL";
  my $info = $annotations{$key} || die;
  $row->{$NS_VARIANT_LABEL_AA} = $info->{cds_change} || die "no cds_change field";
}

sub passthrough_callback {
  my ($row) = @_;
  my $key = $row->{variant} || die;
  my $ar = $annotations{$key} || die "can't find annotations";
  foreach my $f (@PASSTHROUGH_FIELDS) {
    $row->{$f} = $ar->{$f};
  }

  if ($FLAGS{"report-sample"}) {
    my $samples = $samples{$key} || die "where is sample list for $key?";
    $row->{Sample} = join ",", sort keys %{$samples};
  }
}

sub sample_callback {
  my ($row) = @_;
  my $key = $row->{variant} || die;
  my $samples = $samples{$key} || die "where is sample list for $key?";
  $row->{Sample} = join ",", sort keys %{$samples};
}

