#!/usr/bin/env perl

use Getopt::Long;
use Pod::Usage;	
use FindBin;
use lib $FindBin::RealBin;
use rsem_perl_utils;

use Env qw(@PATH);
@PATH = ($FindBin::RealBin, @PATH);

use strict;

my $status;

my $gtfF = "";
my $mappingF = "";
my $polyAChoice = 0; # 0, default, pad polyA tails for all isoforms; 1, --no-polyA; 2, --no-polyA-subset
my $no_polyA = 0; # for option --no-polyA 
my $subsetFile = "";
my $polyALen = 125;
my $bowtie = 0;
my $bowtie_path = "";
my $bowtie2 = 0;
my $bowtie2_path = "";
my $quiet = 0;
my $help = 0;

my $alleleMappingF = "";

GetOptions("gtf=s" => \$gtfF,
	   "transcript-to-gene-map=s" => \$mappingF,
	   "allele-to-gene-map=s" => \$alleleMappingF,
	   "no-polyA" => \$no_polyA, 
	   "no-polyA-subset=s" => \$subsetFile,
	   "polyA-length=i" => \$polyALen,
	   "bowtie" => \$bowtie,
	   "bowtie-path=s" => \$bowtie_path,
	   "bowtie2" => \$bowtie2,
	   "bowtie2-path=s" => \$bowtie2_path,
	   "q|quiet" => \$quiet,
	   "h|help" => \$help) or pod2usage(-exitval => 2, -verbose => 2);

pod2usage(-verbose => 2) if ($help == 1);
pod2usage(-msg => "Set --no-polyA & --no-polyA-subset at the same time!", -exitval => 2, -verbose => 2) if ($no_polyA == 1 && $subsetFile ne '');
pod2usage(-msg => "--transcript-to-gene-map and --allele-to-gene-map are mutually exclusive!", -exitval => 2, -verbose => 2) if (($mappingF ne "") && ($alleleMappingF ne ""));
pod2usage(-msg => "--gtf and --allele-to-gene-map are mutually exclusive!", -exitval => 2, -verbose => 2) if (($gtfF ne "") && ($alleleMappingF ne ""));
pod2usage(-msg => "Invalid number of arguments!", -exitval => 2, -verbose => 2) if (scalar(@ARGV) != 2);

if (!$bowtie && ($bowtie_path ne "")) { print "Warning: If Bowtie is not used, no need to set --bowtie-path option!\n"; }
if (!$bowtie2 && ($bowtie2_path ne "")) { print "Warning: If Bowtie 2 is not used, no need to set --bowtie2-path option!\n"; }
 
my $type;

if ($gtfF ne "") { $type = 0; }
else { $type = 1; }

my @list = split(/,/, $ARGV[0]);
my $size = scalar(@list);

if ($size == 1 && (-d $list[0])) {
    my $dir = $list[0];
    @list = (<$dir/*.fa>, <$dir/*.fasta>);
    $size = scalar(@list);
}

pod2usage(-msg => "reference_fasta_file(s) is empty! Please check if you provide the correct folder name or file suffixes!", -exitval => 2, -verbose => 2) if ($size <= 0);

if ($no_polyA) { $polyAChoice = 1 }
elsif ($subsetFile ne "") { $polyAChoice = 2; }

if ($bowtie_path ne "") { $bowtie_path .= "/"; }
if ($bowtie2_path ne "") { $bowtie2_path .= "/"; }

my $command = "";

if ($type == 0) {
    $"=" ";
    $command = "rsem-extract-reference-transcripts $ARGV[1] $quiet $gtfF";
    if ($mappingF ne "") { $command .= " 1 $mappingF"; }
    else { $command .= " 0"; }
    $command .= " @list";
    &runCommand($command);
}
else {
    $"=" ";
    $command = "rsem-synthesis-reference-transcripts $ARGV[1] $quiet";
    if ($mappingF ne "") { $command .= " 1 $mappingF"; }
    elsif ($alleleMappingF ne "") { $command .= " 2 $alleleMappingF"; }
    else { $command .= " 0"; } 
    $command .= " @list";
    &runCommand($command);
}

$command = "rsem-preref $ARGV[1].transcripts.fa $polyAChoice $ARGV[1] -l $polyALen";
if ($polyAChoice == 2) { $command .= " -f $subsetFile"; }
if ($quiet) { $command .= " -q"; }

&runCommand($command);

if ($bowtie) {
    $command = $bowtie_path."bowtie-build -f";
    if ($quiet) { $command .= " -q"; }
    $command .= " $ARGV[1].n2g.idx.fa $ARGV[1]";

    &runCommand($command);
}

if ($bowtie2) { 
    $command = $bowtie2_path."bowtie2-build -f";
    if ($quiet) { $command .= " -q"; }
    $command .= " $ARGV[1].idx.fa $ARGV[1]";
    
    &runCommand($command);
}

__END__

=head1 NAME

rsem-prepare-reference

=head1 SYNOPSIS

rsem-prepare-reference [options] reference_fasta_file(s) reference_name

=head1 ARGUMENTS

=over

=item B<reference_fasta_file(s)>

Either a comma-separated list of Multi-FASTA formatted files OR a directory name. If a directory name is specified, RSEM will read all files with suffix ".fa" or ".fasta" in this directory.  The files should contain either the sequences of transcripts or an entire genome, depending on whether the --gtf option is used.

=item B<reference name>

The name of the reference used. RSEM will generate several reference-related files that are prefixed by this name. This name can contain path information (e.g. /ref/mm9).

=back

=head1 OPTIONS

=over 

=item B<--gtf> <file>

If this option is on, RSEM assumes that 'reference_fasta_file(s)' contains the sequence of a genome, and will extract transcript reference sequences using the gene annotations specified in <file>, which should be in GTF format.

If this option is off, RSEM will assume 'reference_fasta_file(s)' contains the reference transcripts. In this case, RSEM assumes that name of each sequence in the Multi-FASTA files is its transcript_id. 

(Default: off)

=item B<--transcript-to-gene-map> <file>

Use information from <file> to map from transcript (isoform) ids to gene ids.
Each line of <file> should be of the form:

gene_id transcript_id

with the two fields separated by a tab character.

If you are using a GTF file for the "UCSC Genes" gene set from the UCSC Genome Browser, then the "knownIsoforms.txt" file (obtained from the "Downloads" section of the UCSC Genome Browser site) is of this format.

If this option is off, then the mapping of isoforms to genes depends on whether the --gtf option is specified.  If --gtf is specified, then RSEM uses the "gene_id" and "transcript_id" attributes in the GTF file.  Otherwise, RSEM assumes that each sequence in the reference sequence files is a separate gene.

(Default: off)

=item B<--allele-to-gene-map> <file>

Use information from <file> to provide gene_id and transcript_id information for each allele-specific transcript.
Each line of <file> should be of the form:

gene_id transcript_id allele_id

with the fields separated by a tab character.

This option is designed for quantifying allele-specific expression. It is only valid if '--gtf' option is not specified. allele_id should be the sequence names presented in the Multi-FASTA-formatted files.  

(Default: off) 

=item B<--no-polyA>

Do not add poly(A) tails to the end of reference isoforms. (Default: adding poly(A) tails to all transcripts)

=item B<--no-polyA-subset> <file>

Add poly(A) tails to all transcripts except those listed in <file>. <file> is a file containing a list of transcript_ids. (Default: add poly(A) tails to all transcripts. This option cannot be used with '--no-polyA'.)

=item B<--polyA-length> <int>

The length of the poly(A) tails to be added. (Default: 125)

=item B<--bowtie>

Build Bowtie indices. (Default: off)

=item B<--bowtie-path> <path>

The path to the Bowtie executables. (Default: the path to Bowtie executables is assumed to be in the user's PATH environment variable)

=item B<--bowtie2>

Build Bowtie 2 indices. (Default: off)

=item B<--bowtie2-path>

The path to the Bowtie 2 executables. (Default: the path to Bowtie 2 executables is assumed to be in the user's PATH environment variable)

=item B<-q/--quiet>

Suppress the output of logging information. (Default: off)

=item B<-h/--help>

Show help information.

=back

=head1 DESCRIPTION

This program extracts/preprocesses the reference sequences for RSEM. It can optionally build Bowtie indices (with '--bowtie' option) and/or Bowtie 2 indices (with '--bowtie2' option) using their default parameters. If an alternative aligner is to be used, indices for that particular aligner can be built from either 'reference_name.idx.fa' or 'reference_name.n2g.idx.fa' (see OUTPUT for details). This program is used in conjunction with the 'rsem-calculate-expression' program.

=head1 OUTPUT

This program will generate 'reference_name.grp', 'reference_name.ti', 'reference_name.transcripts.fa', 'reference_name.seq', 'reference_name.chrlist' (if '--gtf' is on), 'reference_name.idx.fa', 'reference_name.n2g.idx.fa' and optional Bowtie/Bowtie 2 index files.

'reference_name.grp', 'reference_name.ti', 'reference_name.seq', and 'reference_name.chrlist' are used by RSEM internally.

B<'reference_name.transcripts.fa'> contains the extracted reference transcripts in Multi-FASTA format. Poly(A) tails are not added and it may contain lower case bases in its sequences if the corresponding genomic regions are soft-masked.

B<'reference_name.idx.fa' and 'reference_name.n2g.idx.fa'> are used by aligners to build their own indices. Poly(A) tails are added unless '--no-polyA' is set. In addition, all sequence bases are converted into upper case. The only difference between 'reference_name.idx.fa' and 'reference_name.n2g.idx.fa' is that 'reference_name.n2g.idx.fa' in addition converts all 'N' characters to 'G' characters. This conversion is in particular desired for aligners (e.g. Bowtie) that do not allow reads to overlap with 'N' characters in the reference sequences. Otherwise, 'reference_name.idx.fa' should be used to build the aligner's index files. RSEM uses 'reference_name.idx.fa' to build Bowtie 2 indices and 'reference_name.n2g.idx.fa' to build Bowtie indices. For visualizing the transcript-coordinate-based BAM files generated by RSEM in IGV, 'reference_name.idx.fa' should be imported as a "genome" (see Visualization section in README.md for details). 
 
=head1 EXAMPLES

1) Suppose we have mouse RNA-Seq data and want to use the UCSC mm9 version of the mouse genome. We have downloaded the UCSC Genes transcript annotations in GTF format (as mm9.gtf) using the Table Browser and the knownIsoforms.txt file for mm9 from the UCSC Downloads. We also have all chromosome files for mm9 in the directory '/data/mm9'.  We want to put the generated reference files under '/ref' with name 'mouse_125'. We'll add poly(A) tails with length 125. Please note that GTF files generated from UCSC's Table Browser do not contain isoform-gene relationship information.  For the UCSC Genes annotation, this information can be obtained from the knownIsoforms.txt file.  Suppose we want to build Bowtie indices and Bowtie executables are found in '/sw/bowtie'.

There are two ways to write the command:

 rsem-prepare-reference --gtf mm9.gtf \
                        --transcript-to-gene-map knownIsoforms.txt \
                        --bowtie \
                        --bowtie-path /sw/bowtie \                  
                        /data/mm9/chr1.fa,/data/mm9/chr2.fa,...,/data/mm9/chrM.fa \
                        /ref/mouse_125

OR

 rsem-prepare-reference --gtf mm9.gtf \
                        --transcript-to-gene-map knownIsoforms.txt \
                        --bowtie \
                        --bowtie-path /sw/bowtie \
                        /data/mm9 \
                        /ref/mouse_125

2) Suppose we also want to build Bowtie 2 indices in the above example and Bowtie 2 executables are found in '/sw/bowtie2', the command will be:

 rsem-prepare-reference --gtf mm9.gtf \
                        --transcript-to-gene-map knownIsoforms.txt \
                        --bowtie \
                        --bowtie-path /sw/bowtie \
                        --bowtie2 \
                        --bowtie2-path /sw/bowtie2 \
                        /data/mm9 \
                        /ref/mouse_125

3) Suppose we only have transcripts from EST tags in 'mm9.fasta'. In addition, we also have isoform-gene information in 'mapping.txt'. We do not want to add any poly(A) tails. The reference_name will be set to 'mouse_0'. In addition, we do not want to build Bowtie/Bowtie 2 indices, and will use an alternative aligner to align reads against either 'mouse_0.idx.fa' or 'mouse_0.idx.n2g.fa':

 rsem-prepare-reference --transcript-to-gene-map mapping.txt \
                        --no-polyA \
                        mm9.fasta \
                        mouse_0

=cut